# Stuff in here should probably move to NaiveNASflux once mature enough...

struct SizePseudoTransparent{T} <: NaiveNASlib.DecoratingTrait
    base::T
end
NaiveNASlib.base(t::SizePseudoTransparent) = t.base

NaiveNASlib.all_in_Δsize_graph(mode, ::SizePseudoTransparent, args...) = NaiveNASlib.all_in_Δsize_graph(mode, SizeInvariant(), args...)

NaiveNASlib.nout(::SizePseudoTransparent, v::AbstractVertex) = calc_outsize(v)
NaiveNASlib.nin(::SizePseudoTransparent, v::AbstractVertex) = nin(SizeInvariant(), v)

calc_outsize(v::AbstractVertex, vo::AbstractVertex=v) = calc_outsize(base(v), vo)
calc_outsize(v::CompVertex, vo::AbstractVertex) = calc_outsize(v.computation, vo)

"""
    reshape_keepshape(x, shape)

Same as `reshape` except that `shape[i] = 0` results in `size(y, i+offs) == size(x, i)` where `y` is the output of this function and `offs = ndims(x) - length(shape)`.

This is basically trying to be compliant to the ONNX Reshape operator although the description there of how to interpret a shape of `0` is a bit vague.
"""
function reshape_keepshape(x, shape)
    offs = ndims(x) - length(shape)
    newshape = Flux.Zygote.@ignore map(enumerate(shape)) do (ind, new)
        new == 0 && return size(x, ind+offs)
        return new
    end
    return reshape(x, newshape...)
end

"""
    Reshape{T}
    Reshape(dims...; activation_dim=actdim(guess_layertype(length(dims))))
    Reshape(dims; activation_dim=actdim(guess_layertype(length(dims))))

Reshape operation wrapped in a struct for the sole purpose of handling size changes in neighbouring vertices.

Ensures that changes in output size of the input vertex are possible to reshape into the shape given by `dims`.

Will treat `Integer`s in dims as fixed, meaning that if `dims[activation_dim] isa Integer` the input size of the output vertices will be fixed as well.
"""
mutable struct Reshape{T}
    adim::Int
    dims::T
end
Reshape(dims...; activation_dim=actdim(guess_layertype(length(dims)))) = Reshape(activation_dim, dims)
Reshape(dims; activation_dim=actdim(guess_layertype(length(dims)))) = Reshape(activation_dim, dims)

(r::Reshape)(x) = any(==(0), r.dims) ? reshape_keepshape(x, r.dims) : reshape(x, r.dims)

function calc_outsize(r::Reshape, v)
    invertex = first(inputs(v))
    outshape = r.dims[r.adim]
    outshape == 0 && return nout(invertex)
    outshape isa Integer && return outshape
    outshape isa Colon && return 0 # Must be set later...
    error("Unknown outshape: " + outshape)
end

function NaiveNASlib.Δsize!(r::Reshape{<:Tuple}, ins::AbstractVector, outs::AbstractVector; kwargs...)
    r.dims = Tuple(map(enumerate(r.dims)) do (i, s)
        s isa Colon && return s
        i == r.adim && return length(outs)
        return s
    end)
end

NaiveNASflux.layer(r::Reshape) = r
NaiveNASflux.actdim(r::Reshape) = r.adim
NaiveNASflux.actrank(r::Reshape) = length(r.dims)

# If Reshape is wrapped in an AbstractMutableComp we will hit this method instead due to how NaiveNASflux unwraps things
NaiveNASlib.compconstraint!(case, s::NaiveNASlib.AbstractJuMPΔSizeStrategy, ::Type{<:Reshape}, data) = NaiveNASlib.compconstraint!(case, s, layer(data.vertex), data)

# Special case: Reshape to 2D with variable batch size. Maybe a more general case when this is ok is hiding here somewhere...
NaiveNASlib.compconstraint!(case, s::NaiveNASlib.AbstractJuMPΔSizeStrategy, r::Reshape{Tuple{Int, Colon}}, data) = reshape_nout_constraint(s, Colon(), r, filter(vin -> vin in keys(data.noutdict), inputs(data.vertex)), data)

function NaiveNASlib.compconstraint!(case, s::NaiveNASlib.AbstractJuMPΔSizeStrategy, r::Reshape, data)
    ins = filter(vin -> vin in keys(data.noutdict), inputs(data.vertex))

    reshape_nout_constraint(s, r.dims[r.adim], r, ins, data)
    isempty(ins) && return

    fixeddims = filter(dim -> dim isa Integer && dim != 0, collect(r.dims))
    isempty(fixeddims) && return

    if length(fixeddims) == length(r.dims)
        # No size change possible!
        @constraint(data.model, [i=1:length(ins), j=1:length(fixeddims)], data.noutdict[ins[i]] == nout(ins[i]))
        return
    end

    # Make sure that nout of inputs is an integer multiple of at least one of the output sizes
    # It is unfortunately too restrictive to just enforce fv_dims[j] * fixeddims[j] == nout[ins[i]] ∀ i, j
    # as it is enough that this is fulfilled for one j

    # Unfortunately, something like  fv_dims[j] * fixeddims[j] == nout[ins[i]] * b[i] where b is a binary variable is not possible as it is nonlinear.
    # Instead we use i slack(ish) variables with lower bound 0 and force the smallest of them to be 0 (i.e at least one dimension has zero slack).
    fv_dims = @variable(data.model, [1:length(fixeddims)], integer=true)
    slack = @variable(data.model, [1:length(fixeddims)], integer=true, lower_bound=0)
    @constraint(data.model, [i=1:length(ins), j=1:length(fixeddims)], fv_dims[j] * fixeddims[j] ==  data.noutdict[ins[i]] + slack[j])

    # Force the smallest slack to be 0 through a big-M strategy
    atleast_one = @variable(data.model, [1:length(fixeddims)], binary=true)
    min_slack = @variable(data.model, integer=true)
    M = 10000
    @constraint(data.model,[i=1:length(fixeddims)], -M * (1 - atleast_one[i]) <= min_slack - slack[i])
    @constraint(data.model,[i=1:length(fixeddims)], min_slack - slack[i] <= M * (1-atleast_one[i]))
    @constraint(data.model, sum(atleast_one) == 1)
    @constraint(data.model, min_slack == 0)
end

# Case 1: Output size is fixed so we will only put constraints on the inputs to ensure reshaping is possible
function reshape_nout_constraint(s, outsize::Integer, r, ins, data)
    if outsize == 0

        fix_size = true
        for iv in ins
            inrank = unique(actrank(iv))[] # Should be one element or else we are fked
            inadim = unique(actdim(iv))[] # Should be one element or else we are fked

            if inrank - length(r.dims) + r.adim + 1== inadim
                # Dimension to keep size of happens to be input layers activation dimension.
                # This is basically a SizeInvariant vertex
                @constraint(data.model, [i=1:length(ins)], data.noutdict[data.vertex] == data.noutdict[iv])
                fix_size = false
            end
        end

        if fix_size
            @constraint(data.model, [i=1:length(ins)], data.noutdict[data.vertex] == nout(data.vertex))
        end

    else
        @constraint(data.model, data.noutdict[data.vertex] == outsize)
    end
end

# Case 2: Outsize is not fixed so we need to change it so that the output vertices change their input size
function reshape_nout_constraint(s, outsize::Colon, r, ins, data)
    @constraint(data.model, [i=1:length(ins)], data.noutdict[data.vertex] / nout(data.vertex) == data.noutdict[ins[i]] / nout(ins[i]))
end

"""
    Flatten
    Flatten(dim)

Flattens input array along `dim` to a 2D array.

If input has shape `(d_1, d_2, ... d_n)` then the output will have shape `(d_1 X d_2 ... d_(dim-1), d_dim X d_(dim+1) ... X dn)`.

If `dim = 0`, the shape of the output is `(d_1 X d_2 ... d_n, 1)`, where the shape of the input is `(d_0, d_1, ... d_n)`.
"""
struct Flatten
    dim::Int
end

(f::Flatten)(x) = flatten(x, f.dim)
function flatten(x, dim)
    dim == 0 && return reshape(x, :, 1)
    xs = size(x)
    absdim = dim < 0 ? length(xs) + dim : dim
    return reshape(x, prod(xs[1:absdim]), prod(xs[absdim+1:end]))
end

NaiveNASflux.layer(f::Flatten) = f
NaiveNASflux.actdim(f::Flatten) = 1
NaiveNASflux.actrank(f::Flatten) = 1

calc_outsize(::Flatten, v) = 0

# If Flatten is wrapped in an AbstractMutableComp we will hit this method instead due to how NaiveNASflux unwraps things
NaiveNASlib.compconstraint!(case, s::NaiveNASlib.AbstractJuMPΔSizeStrategy, ::Type{<:Flatten}, data) = NaiveNASlib.compconstraint!(case, s, layer(data.vertex), data)

function NaiveNASlib.compconstraint!(case, ::AbstractJuMPΔSizeStrategy, f::Flatten, data)
    ins = filter(vin -> vin in keys(data.noutdict), inputs(data.vertex))
    for iv in ins
        inadim = unique(actdim(iv))[] # Should be one element or else we are fked

        absdim = f.dim > 0 ? f.dim : f.dim + unique(actrank(iv))[] + 1 # Should be one element or else we are fked

        if f.dim == 0 || inadim <= absdim
            # Size of input affect output size
            @constraint(data.model, [i=1:length(ins)], data.noutdict[data.vertex] / nout(data.vertex) ==  data.noutdict[iv] / nout(iv))
        end
        # else: Size of input affects batch size => no constraint needed
    end
end


#include("../ReadWrite/Header.jl")

"""
**SeisWindow**
*Window a seis file using header words.*
**IN**
* in
* out
* key=[]
* minval=[]
* maxval=[]
* note that windowing along the time axis is achieved by using the key "t".
**OUT**
*Credits: AS, 2015*
"""

function SeisWindow(in::String,out::String;key=[],minval=[],maxval=[])
    DATAPATH = get(ENV,"DATAPATH",join([pwd(),"/"]))
	extent = ReadTextHeader(in)
	tmin = extent.o1
	tmax = extent.o1 + extent.d1*(extent.n1-1)
	for ikey=1:length(key)
		if key[ikey] == "t"
			tmin = minval[ikey]
			tmax = maxval[ikey]
		end
	end
	itmin = convert(Int32,round((tmin - extent.o1)/extent.d1) + 1)
	if (itmin < 1)
		itmin = 1
	end
	itmax = convert(Int32,round((tmax - extent.o1)/extent.d1) + 1)
	if (itmax) > extent.n1
		itmax = extent.n1
	end
	SeisWindowHeaders(in,out;key=key,minval=minval,maxval=maxval,tmin=tmin,tmax=tmax)
	FetchTraces(in,out;itmin=itmin,itmax=itmax)
	tmp = join(["tmp_SeisWindow_",string(round(Int,rand()*100000))])
	@compat SeisProcessHeaders(out,tmp,[UpdateHeader],[Dict(:itmin=>itmin,:itmax=>itmax)])
	filename_h_tmp = join([DATAPATH tmp "@headers@"])
	filename_h_out = join([DATAPATH out "@headers@"])
	cp(filename_h_tmp,filename_h_out,remove_destination=true);
	rm(filename_h_tmp);
	#rm(tmp);
end

function FetchTraces(in::String,out::String;ntrace=500,itmin=round(Int,1),itmax=round(Int,9999999999))

	NX = GetNumTraces(out)
	itrace = round(Int,1)
	filename_data_in = ParseDataName(in)
    DATAPATH = get(ENV,"DATAPATH",join([pwd(),"/"]))
	filename_data_out = join([DATAPATH out "@data@"])
	stream_in  = open(filename_data_in)
	stream_out  = open(filename_data_out,"w")
	extent = ReadTextHeader(in)
	nt = extent.n1
	if itmin < 1
		itmin = round(Int,1)
	end
	if itmax > nt
		itmax = round(Int,nt)
	end
	while itrace <= NX
		if (itrace > 1)
			stream_out = open(filename_data_out,"a")
		end
		nx = NX - itrace + 1
		ntrace = nx > ntrace ? ntrace : nx
		h = SeisReadHeaders(out,group="some",itrace=itrace,ntrace=ntrace)
		d = zeros(Float32,itmax-itmin+1,ntrace)
		SeekTraces!(d,stream_in,h,itmin,itmax,nt,round(Int,ntrace))
		write(stream_out,d)
		close(stream_out)
		itrace += ntrace
	end
	close(stream_in)
end

function SeekTraces!{T}(d::AbstractArray{T,2},stream_in::IOStream,h::Array{Header,1},itmin,itmax,nt,ntrace)

	d1 = zeros(Float32,nt)
	for ix = 1 : ntrace
			position = 4*nt*(h[ix].tracenum-1)
			seek(stream_in,position)
			d1 = read(stream_in,Float32,nt)
			for it = itmin : itmax
				d[it - itmin + 1,ix] = d1[it]
			end
	end
	nothing
end

function UpdateHeader(h;itrace=1,itmin=1,itmax=9e9)
	ot = (itmin-1)*h[1].d1 + h[1].o1
	nt = itmax - itmin + 1
	for j = 1 : length(h)
		h[j].o1 = ot
		h[j].n1 = nt
		h[j].tracenum = j + itrace - 1
	end

	return h
end


using Iterators, FileIO, Colors
using Visualize, GeometryTypes, GtkReactive, ModernGL

using Visualize.GLRasterization
using Visualize.GLRasterization: swapbuffers!, destroy!, GtkWindow, GtkParent, OpenGLVersion, GLCanvas, make_context_current
using Visualize: VertexN, Light, Shading, vert_mesh, frag_mesh, Visible, NativeWindow
using Visualize: Resolution, View, ProjectionView, EyePosition, Projection, Name
using Visualize: VertexN, Light, MeshUniforms, Shading, vert_mesh, frag_mesh
using Visualize: Model, Color, ShadingFunction, Debugging, Canvas, add!
using FieldTraits: Partial, default, on

parent = Gtk.Window("Test", 800, 600, true, true)
box = Gtk.Box(:v)
tb = textbox(String)
sl = GtkReactive.slider(linspace(0.0, 20.0, 100))
mesh_color = map(sl) do val
    RGBA{Float32}(val, 0,0,1)
end
push!(box, tb)
push!(box, sl)

push!(parent, box)
window = GLRasterization.GtkWindow(OpenGLVersion => v"3.3.0", GtkParent => box)
Gtk.setproperty!(box, :expand, window[NativeWindow], true)

Gtk.showall(parent)
# for event in Visualize.NativeWindowEvents
#     add!(window, event)
# end
# add!(window, Mouse.Drag)
add!(window, Mouse.Buttons)
add!(window, Mouse.Position)
add!(window, Mouse.Drag)

cam = PerspectiveCamera(
    TranslationSpeed => 1f0,
    LookAt => Vec3f0(0),
    EyePosition => Vec3f0(6, 6, 8),
    Rotation => Vec3f0(0),
    Area => window[Area],
    RotationSpeed => 0.1f0
)

add!(cam, Translation, window,
    Mouse.right, Set(Keyboard.Button[])
)
add!(cam, Rotation, window,
    Mouse.left, Set(Keyboard.Button[])
)
add!(cam, Projection, View)
add!(cam, EyePosition, LookAt)

make_context_current(window[NativeWindow])
on(cam, ProjectionView) do pv
    Gtk.queue_render(window[NativeWindow])
end

catmesh = normalmesh(load(Pkg.dir("GLVisualize", "assets", "cat.obj")))
catmesh = Base.view(
    reinterpret(Visualize.VertexN, catmesh.parent), catmesh.indexes[1]
)
vbo = VertexArray(catmesh);

args = MeshUniforms(
    ShadingFunction => (V, N, L, color, shading, light) -> return L
)

canvas = window[Canvas];
for field in FieldTraits.Fields(eltype(canvas))
    if haskey(cam, field)
        FieldTraits.link!(field, cam => canvas)
    end
end
light = Light(
    Vec3f0(10),
    Vec3f0(0.0),
    Vec3f0(0.9),
    1.0f0,
    Vec3f0(0.2),
    1.0f0
)
shading = Shading(
    Vec3f0(0.0),
    Vec3f0(0.3),
    8.0f0
)
uniforms = (canvas, UniformBuffer(light), UniformBuffer(shading), UniformBuffer(args));

draw_cat = GLRasterizer(
    vbo,
    uniforms,
    vert_mesh, Visualize.frag_mesh
)

function gtk_render(gl_area, gdk_context)
    global draw_cat, uniforms, args, vbo, window
    glEnable(GL_DEPTH_TEST)
    glDepthMask(GL_TRUE)
    glDepthFunc(GL_LEQUAL)
    glDisable(GL_CULL_FACE)
    GLAbstraction.enabletransparency()
    glClearColor(1,1,1,1)
    glViewport(0, 0, widths(window[Area])...)
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    draw_cat(vbo, uniforms)
    glFlush()
    return false
end

Gtk.signal_connect(gtk_render, window[NativeWindow], "render")
cam[View] = eye(Mat4f0)
# signal_connect(parent, :destroy) do widget
#     Gtk.gtk_quit()
# end
# Gtk.gtk_main()


using Base.Test
using GLPlot
using GLAbstraction

using ModernGL

using GLFW
using Color
using Images
using Reactive
using ImmutableArrays
using FixedPointNumbers

N = 54

window = createdisplay(w=iround(N*11.1), h=N)
camera = OrthographicPixelCamera(window.inputs)

immutable GLGlyph{T <: Real} <: AbstractVector{T}
  glyph::T
end
Base.getindex(collection::GLGlyph, I::Integer) = I == 1 ? collection.glyph : error("Out of bounds")
Base.length{T}(::GLGlyph{T}) 		= 1
Base.length{T}(::Type{GLGlyph{T}})  = 1
Base.eltype{T}(::GLGlyph{T}) 		= T
Base.eltype{T}(::Type{GLGlyph{T}}) 	= T
Base.convert{T}(::Type{GLGlyph{T}}, x::Char) = (int(x) >= 0 && int(x) <= 256) ? GLGlyph(convert(T, x)) : error("This char: ", x, " can't be converted to GLGlyph")

test_images = {
	[float32(33) for i=1:N, j=1:N],
	[GLGlyph{UInt8}(33) for i=1:N, j=1:N],
	[Vector3{Ufixed8}(0, 1, 0) for i=1:N, j=1:N],
	[Vec4(0,1,0,1) for i=1:N, j=1:N],
	[Vec3(0,1,0)   for i=1:N, j=1:N],
	[AlphaColorValue(RGB(0f0,1f0,0f0), 1f0) for i=1:N, j=1:N],
	Images.ColorTypes.AlphaColor{RGB{Float32}, Float32}[Images.ColorTypes.AlphaColor(RGB(0f0,1f0,0f0), 1f0) for i=1:N, j=1:N],
	[AlphaColorValue(Images.ColorTypes.BGR(0f0,1f0, 0f0), 1f0) for i=1:N, j=1:N],
	"test.png",
	"test.tif",
	"test.jpg",
	"test.bmp",
}
scale_matrix 	= scalematrix(Vec3(1, 1, 1))
trans_matrix	= translationmatrix(Vec3(0))

test_renderobjects = map(test_images) do image
	global scale_matrix, trans_matrix
	obj = toopengl(Texture(image), model=scale_matrix*trans_matrix, camera=camera)
	trans_matrix *= translationmatrix(Vec3(N+1, 0, 0))
	obj
end

map(glplot, test_renderobjects)

renderloop(window)


"""
    all_libs()
"""
function all_libs()
    libs = Symbol[
        :CUBLAS,
        :GemmKernels,
        :Tullio,
    ]
    return libs
end

function getbenchfunc(lib::Symbol)
    if lib === :CUBLAS
        return benchmark_gemm_cublas!
    elseif lib === :GemmKernels
        return benchmark_gemm_gemmkernels!
    elseif lib === :Tullio
        return benchmark_gemm_tullio!
    end
    throw(ArgumentError("Unknown library \"$(lib)\""))
end

function gemm_cublas!(C, A, B)
    CUDA.CUBLAS.cublasSetMathMode(CUDA.CUBLAS.handle(), CUDA.CUBLAS.CUBLAS_TENSOR_OP_MATH)
    CUDA.CUBLAS.gemmEx!('N', 'N', 1, A, B, 0, C)
    return C
end

function benchmark_gemm_cublas!(C, A, B)
    CUDA.CUBLAS.cublasSetMathMode(CUDA.CUBLAS.handle(), CUDA.CUBLAS.CUBLAS_TENSOR_OP_MATH)
    return BenchmarkTools.@benchmark CUDA.@sync begin
        CUDA.CUBLAS.gemmEx!('N', 'N', 1, $A, $B, 0, $C)
    end
end

function gemm_gemmkernels!(C, A, B)
    GemmKernels.BLAS.gemmEx!('N', 'N', 1, A, B, 0, C)
    return C
end

function benchmark_gemm_gemmkernels!(C, A, B)
    return BenchmarkTools.@benchmark CUDA.@sync begin
        GemmKernels.BLAS.gemmEx!('N', 'N', 1, $A, $B, 0, $C)
    end
end

function gemm_tullio!(C, A, B)
    Tullio.@tullio C[i, j] = A[i, k] * B[k, j] cuda=256 grad=false
    return C
end

function benchmark_gemm_tullio!(C, A, B)
    return BenchmarkTools.@benchmark CUDA.@sync begin
        gemm_tullio!($C, $A, $B)
    end
end


using InteractiveUtils
using PlutoUI
using TickTock
using ConfParser
using Dates
using Random
using StaticArrays
using Distributed
using OrdinaryDiffEq
using LinearAlgebra
using JLD2
using Plots
using LoopVectorization
using BenchmarkTools
using StatsPlots
using DataFrames
using CSV
using LaTeXStrings
using Plots.PlotMeasures
plot(); # dummy plot just to get the dumb pkg to compile

#######################
## Constants n stuff ##
#######################

const Re   = 6370e3;        # Earth radius, f64
const c    = 3e8;           # speedo lite, f64
const Beq  = 3.e-5;         # B field at equator (T), f64

##################
## Loading data ##
##################

struct Resultant_Matrix
    label::String
    numParticles::Int64
    endTime::Float64
    allZ::Vector{Vector{Float64}}
    allPZ::Vector{Vector{Float64}}
    allT::Vector{Vector{Float64}}
    allPA::Vector{Vector{Float64}}
    allE::Vector{Vector{Float64}}
    lostParticles::Matrix{Float64}
    tVec::Vector{Float64}
    Zmatrix::Matrix{Float64}
    PZmatrix::Matrix{Float64}
    Ematrix::Matrix{Float64}
    PAmatrix::Matrix{Float64}
end

struct Precipitating_Particles
    label::String
    precipitating_fluxes_mean::Vector{Float64}
    precipitating_fluxes_plus::Vector{Float64}
    precipitating_fluxes_minus::Vector{Float64}
end

function parse_conf_file(directoryname::String, conffile::String)
    # Parsing the conf file
    conf = ConfParse(directoryname*"/"*conffile)
    parse_conf!(conf)
    numParticles = parse(Int64, retrieve(conf, "numberOfParticles"));
    startTime = parse(Float64, retrieve(conf, "startTime"));
    endTime = parse(Float64, retrieve(conf, "endTime"));
    tspan = (startTime, endTime); # integration time
    lossConeAngle = parse(Float64, retrieve(conf, "lossConeAngle"));
    saveDecimation = parse(Float64, retrieve(conf, "saveDecimation"));
    L = parse(Float64, retrieve(conf, "L"));
    omegam = parse(Float64, retrieve(conf, "omegam"));
    Omegape = parse(Float64, retrieve(conf, "Omegape"));
    z0 = parse(Float64, retrieve(conf, "z0"));
    lambda0 = parse(Float64, retrieve(conf, "lambda0"));
    waveAmplitudeModifier = parse(Float64, retrieve(conf, "waveAmplitudeModifier"));
    ELo = parse(Float64, retrieve(conf, "ELo"));
    EHi = parse(Float64, retrieve(conf, "EHi"));
    Esteps = parse(Float64, retrieve(conf, "Esteps"));
    PALo = parse(Float64, retrieve(conf, "PALo"));
    PAHi = parse(Float64, retrieve(conf, "PAHi"));
    PAsteps = parse(Float64, retrieve(conf, "PAsteps"));
    ICrange = [ELo, EHi, Esteps, PALo, PAHi, PAsteps];
    batches = parse(Int64, retrieve(conf, "batches"));
    numThreads = parse(Int64, retrieve(conf, "numberOfThreads"))
    @info "Parsed Config file: $conffile"
    return numParticles, endTime
end

function load_resultant_matrix(label::String, directoryname::String, basename::String, conffile::String, num_batches::Int64)
    # Parse conf file
    @time numParticles, endTime = parse_conf_file(directoryname, conffile)
    # Load in data
    @time allZ, allPZ, allT, allPA, allE = loadData(directoryname, basename, num_batches);
    # Count lost particles
    @time lostParticles = countLostParticles(allT, endTime);
    # Convert into a workable matrix
    @time tVec, Zmatrix, PZmatrix, Ematrix, PAmatrix = postProcessor(allT, allZ, allPZ, allE, allPA);
    return Resultant_Matrix(label, numParticles, endTime, allZ, allPZ, allT, allPA, allE, lostParticles,tVec, Zmatrix, PZmatrix, Ematrix, PAmatrix)
end

function export_results(label::String, precipitating_flux_timeseries)
    ###
    # extract data and save it into
    ###
    prec_flux_matrix = hcat(precipitating_flux_timeseries...)
    prec_flux_mean = Vector{Float64}()
    prec_flux_plus = Vector{Float64}()
    prec_flux_minus = Vector{Float64}()
    for energy_bin in eachrow(prec_flux_matrix)
        # remove zeros from each row
        row = filter(!iszero,energy_bin)
        if isempty(row)
            push!(prec_flux_mean, 1)
            push!(prec_flux_plus, 1)
            push!(prec_flux_minus, 1)
        elseif length(row) <= 2
            avg_val = mean(row)
            push!(prec_flux_mean, avg_val)
            push!(prec_flux_plus, avg_val*.25)
            push!(prec_flux_minus, avg_val*.25)
        else
            avg_val = mean(row)
            push!(prec_flux_mean, avg_val)
            push!(prec_flux_plus, maximum(row)-avg_val)
            push!(prec_flux_minus, avg_val-minimum(row))
        end
    end
    # test plot
    # plot(Egrid, prec_flux_mean, yerror=(prec_flux_minus, prec_flux_plus), ylim =(1e2,1e9), xlim=(50,800), yscale=:log10)

    return Precipitating_Particles(label, prec_flux_mean, prec_flux_plus, prec_flux_minus)
end
# themis_lolat = export_results("210429_themis_lolat", equatorial_fluxes_042921, prec_flux_timeseries_042921)
# @save "210429_data_storage.jld2" themis_lolat themis_hilat

###################
## PSD Functions ##
###################
# define dist function here

# generic
f0 = ((C::Float64, E::Float64, PA::Float64) -> ((C*(E/70)^-3)*(sin(deg2rad(PA)))))
# mms fit for 9/22/20
f0_092220 = function (E::Float64, PA)
    A = 4
    B0 = .25
    B1 = .3
    C = 1e5
    E0 = 300
    E1 = 70
    if E < E1
        B = B0
    elseif E <= E0
        B = B0 + (B1-B0)*(E-E0)/(E1-E0)
    elseif E > E0
        B = B1
    end
    return ((C*(E/70)^-A)*(sin(deg2rad(PA))-sin(deg2rad(8)))^B)
end
# mms fit for 10/27/20
f0_102720 = function (E::Float64, PA)
    A = 3
    B0 = .2
    B1 = .3
    C = 1.5e5
    E0 = 300
    E1 = 70
    if E < 70
        B = B0
    elseif E <= 300
        B = B0 + (B1-B0)*(E-E0)/(E1-E0)
    elseif E > 300
        B = B1
    end
    return ((C*(E/70)^-A)*(sin(deg2rad(PA))-sin(deg2rad(8)))^B)
end
# themis fit for 4/29/21
f0_042921 = function (E::Float64, PA)
    a1 = 5e-13
    a2 = 8e-16
    a3 = 2e-19
    b1 = .2
    b2 = 9
    b3 = 200
    if E < 5
        b = 0.5
    elseif E <= 20
        b = 0.5
    elseif E > 40
        b = 0
    end
    Epsd = (a1 * (1 + E/b1)^-5) + (a2 * (1 + E/b2)^-5) + (a3 * (1 + E/b3)^-3)
    PApsd = (sin(deg2rad(PA))-sin(deg2rad(8)))^b
    # km --> cm and 1/v^4 --> 1/E^2 conversion
    # since v = 438 * sqrt(1836 KeV)
    return (1e5*438^4*1836^2)*(Epsd*PApsd)
end

#################
## Plot Colors ##
#################

using Plots.PlotMeasures
function hexcolor(r::UInt8, g::UInt8, b::UInt8)
    return RGB(Int(r)/255,Int(g)/255,Int(b)/255)
end
bipride_pink = RGB(234/255, 2/255, 112/255);
bipride_orange = RGB(255/255, 79/255, 0/255);
bipride_lavender = RGB(155/255, 79/255, 150/255);
bipride_blue = RGB(0/255, 56/255, 168/255);
c1 = hexcolor(0xff,0x4f,0x00)
c2 = hexcolor(0xf8,0x00,0x4d)
c3 = hexcolor(0xd3,0x00,0x7d)
c4 = hexcolor(0x8f,0x19,0x9f)
c5 = hexcolor(0x00,0x38,0xa8)

###############################
## Post Processing functions ##
###############################

calcb!(b::Vector{Float64}, lambda::SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}) = @. b = sqrt(1+3*sin(lambda)^2)/(cos(lambda)^6)
calcGamma!(gamma::Vector{Float64}, pz::SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, mu::SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, b::Vector{Float64}) = @. gamma = sqrt(1 + pz^2 + 2*mu*b)
calcAlpha!(alpha::Vector{Float64}, mu::SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}, gamma::Vector{Float64}) = @. alpha = rad2deg(asin(sqrt((2*mu)/(gamma^2 - 1))))

function loadData(directory::String, basename::String, num_batches::Int64)
    #=
    Loads in the JLD2 file and creates timeseries of all the data.
    Takes in filename in string, returns 5 TxN vectors for time,
    z, pz, pitch angle, and energy timeseries for all N particles.
    =#

    allZ = Vector{Vector{Float64}}();
    allPZ = Vector{Vector{Float64}}();
    allE = Vector{Vector{Float64}}();
    allPA = Vector{Vector{Float64}}();
    allT = Vector{Vector{Float64}}();

    for i in 1:num_batches
        JLD2.@load directory*"/"*basename*"_$i.jld2" sol
        for traj in sol
            vars = Array(traj')
            timesteps = length(traj.t)
            b = zeros(timesteps)
            gamma = zeros(timesteps)
            Alpha = zeros(timesteps)

            @views calcb!(b,vars[:,5])
            @views calcGamma!(gamma,vars[:,2],vars[:,4],b)
            @views calcAlpha!(Alpha,vars[:,4],gamma)
            @views push!(allT, traj.t);
            @views push!(allZ, vars[:,1]);
            @views push!(allPZ, vars[:,2]);
            @views push!(allPA, Alpha);
            @views push!(allE, @. (511*(gamma - 1)));
        end
        @info "$(length(sol)) particles loaded in from $(basename*"_$i.jld2")"
        @info "Total $(length(allT)) particles loaded so far..."
    end
    @info "Loaded total $(length(allT)) particles!"
    return allZ, allPZ, allT, allPA, allE
end

function countLostParticles(allT::Vector{Vector{Float64}}, endTime::Float64)
    #=
    Based on time vectors, counts which ones were lost
    and at what time. Returns a Nx2 array where the first
    column is the time at which the particle was lost, and
    the 2nd column denotes the number lost at that point.
    =#
    lossCounter = []; # initialize vector

    for ntime in allT # loop thru each particle
        if maximum(ntime) != endTime # particle lost if ended early
            push!(lossCounter, maximum(ntime)); # the final entry in the time vector is when the particle got lost
        end
    end
    if isempty(lossCounter) # if particle wasn't lost, then throw in a NaN
        push!(lossCounter, NaN)
        @warn "All particles trapped!" # since all particles trapped
    end
    lostParticles = [sort(lossCounter) collect(1:length(lossCounter))] # sort it by time, so you can see when particles got lost

    if maximum(lostParticles[:,1]) != endTime # this adds in a final hline from last particle lost to end of simulation
        lostParticles = vcat(lostParticles, [endTime (maximum(lostParticles[:,2]))]);
    end

    @info "Total of $(lostParticles[end,end]) particles lost during sim"
    return lostParticles
end

function postProcessor(allT::Vector{Vector{Float64}}, allZ::Vector{Vector{Float64}}, allPZ::Vector{Vector{Float64}}, allE::Vector{Vector{Float64}}, allPA::Vector{Vector{Float64}})
    #=
    This function will take the output of the model and convert them into usable m x n matrices
    where m is max number of timesteps for the longest trajectory and N is number of particles
    Arrays that are not m long are filled with NaNs
    =#
    N = length(allT); # num particles from data
    tVec = allT[findall(i->i==maximum(length.(allT)),length.(allT))[1]]; # turns all time vectors into a single time vector spanning over the longest trajectory
    timeseriesLength = length(tVec); # all vectors must be this tall to ride
    Zmatrix = fill(NaN,timeseriesLength,N);
    PZmatrix = fill(NaN,timeseriesLength,N);
    Ematrix = fill(NaN,timeseriesLength,N);
    PAmatrix = fill(NaN,timeseriesLength,N);
    # iterate over each matrix column and fill with each vector
    for i = 1:N
        @views Zmatrix[1:length(allT[i]),i] = allZ[i]
        @views PZmatrix[1:length(allT[i]),i] = allPZ[i]
        @views Ematrix[1:length(allT[i]),i] = allE[i]
        @views PAmatrix[1:length(allT[i]),i] = allPA[i]
    end

    return tVec, Zmatrix, PZmatrix, Ematrix, PAmatrix
end

function animatePAD(gifFileName="PADanimation.gif", maxParticles=1000, binwidth=5)
    pyplot()
    animDec = 10; # make a png for animation every 10 points
    animScale = 10; # i.e. animscale = 10 means every 10 seconds in animation is 1 second of simulation time (increase for longer animation)
    anim = @animate for i=1:length(tVec)
        histogram(PAmatrix[i,:],bins=(0:binwidth:90), ylim = [0, 1.1*maxParticles], #xminorticks = 4, yminorticks = 2,
            legend = false, dpi = 100,
            xlabel="Pitch Angle (deg)", ylabel="Number of Particles", title="Time evolution of Pitch Angle Distribution");
        vline!([lossConeAngle]);
        annotate!(lossConeAngle+.5, 1.05*maxParticles, text("Lost Particle Count:", :left));
        annotate!(lossConeAngle+.5, .95*maxParticles, text((count(i->isnan(i), Ematrix[i,:])), :left));
        annotate!(80, 1.05*maxParticles, text("t = $(round(tVec[i]*R*L/(c),digits=4)) s"), :right)
    end every animDec
    savename = string("results/",gifFileName)
    gif(anim, savename, fps = (length(tVec)/animDec)/(animScale*endTime*Re*L/(c)))
end

function animateESD(gifFileName="ESDanimation.gif", maxParticles=1000, binwidth=50)
    pyplot()
    animDec = 10; # make a png for animation every 100 points
    animScale = 5; # i.e. animscale = 10 means every 10 seconds in animation is 1 second of simulation time (increase for longer animation)
    anim = @animate for i=1:length(tVec)
        histogram(Ematrix[i,:],bins=(0:binwidth:2000), ylim = [1, 1.1*maxParticles], #xminorticks = 4, yminorticks = 2,
            legend = false, dpi = 100, xscale=:log10, yscale =:log10, xticks = ([10,100,1000],["10", "100", "1000"]), xlim = [10,2000],
            xlabel="Energy (KeV)", ylabel="Number of Particles", title="Time evolution of Energy Spectra");
        annotate!(20, 1.05*maxParticles, text("t = $(round(tVec[i]*Re*L/(c),digits=4)) s"), :left)
    end every animDec
    savename = string("results/",gifFileName)
    gif(anim, savename, fps = (length(tVec)/animDec)/(animScale*endTime*Re*L/(c)))
end

function animatePSD(gifFileName="PSDanimation.gif", PAbinwidth=3, Ebinwidth=100, maxParticles=50)
    pyplot()
    animDec = 10; # make a png for animation every 10 points
    animScale = 1; # i.e. animscale = 10 means every 10 seconds in animation is 1 second of simulation time (increase for longer animation)
    anim = @animate for i=1:length(tVec)
        histogram2d(PAmatrix[i,:], Ematrix[i,:], bins=(0:PAbinwidth:90,0:Ebinwidth:2000),
            xlabel="Pitch Angle (deg)",ylabel="Energy (keV)",title="Time evolution of Phase Space Density",
            yscale=:log10, clims = (1,maxParticles), colorbar_scale=:log10, minorgrid=true,
            yticks = ([1, 10,100,1000,10000],["1", "10", "100", "1000", "10000"]));
        annotate!(lossConeAngle, 5, text("Lost Particle Count:", :left));
        annotate!(lossConeAngle, 3, text((count(i->isnan(i), Ematrix[i,:])), :left))
        annotate!(70, 5, text("t = $(round(tVec[i]*Re*L/(c),digits=4)) s"), :right)
        # count(i->isnan(i), Ematrix[1,:]) # counts number of NaNs (i.e particles lost) at certain time
    end every animDec
    savename = string("results/",gifFileName)
    gif(anim, savename, fps = (length(tVec)/animDec)/(animScale*endTime*Re*L/(c)))
end

function recalcDistFunc(Ematrix::Array{Float64,2},PAmatrix::Array{Float64,2},initial::Int64,final::Int64,distFunc,
    Egrid::Vector{Float64}, PAgrid::StepRange{Int64,Int64}, whistler_occurence_rate::Float64)
    #=
    Takes in matrix of Energy and PA, initial and final indices, and grid values as stepranges.
    Recalculates out a new PSD given a a distribution function at the initial and final indices.
    =#
    N = @views length(Ematrix[1,:]) # num particles
    EPA0i = @views round.([Ematrix[1,:] PAmatrix[1,:]])
    Nkl = @views N/(length(unique(EPA0i[:,1]))*length(unique(EPA0i[:,2]))) # obtain number of particles/bin

    EPAinitial = @views @. round([Ematrix[initial,:] PAmatrix[initial,:]])
    EPAfinal = @views @. round([Ematrix[final,:] PAmatrix[final,:]])
    EPAfinal_next = @views @. round([Ematrix[final+1,:] PAmatrix[final+1,:]])

    # initialize matrices to be used based on grid size
    f = zeros(length(Egrid), length(PAgrid));
    psd_init = zeros(length(Egrid), length(PAgrid));
    psd_final = zeros(length(Egrid), length(PAgrid));
    # initialize vectors to be filled in
    f0Vec = Vector{Float64}();
    indices = Vector{Tuple{Int64,Int64}}();
    psd_prec = zeros(length(Egrid));
    excludedParticles = 0;

    ### HANDLE NANS HERE ###
    # diff = @. abs(EPAinitial - EPAfinal)
    valid_rows = @. ~isnan(EPAfinal)[:,1] & ~isnan(EPAfinal[:,2]); # non-lost particles from final state
    valid_rows_next = @. ~isnan(EPAfinal_next)[:,1] & ~isnan(EPAfinal_next[:,2]); # non-lost particles from state at next timestep
    # if valid_rows!=valid_rows_next # this means that at least one of the particles will precipitate
        prec_rows = valid_rows .!= valid_rows_next # these are the indices of particles who are about to precipitate
        prec_indices = (ones(N).*vec(prec_rows))[vec(valid_rows)]
        # @info "$(length(ones(N)[prec_rows])) particles are about to precipitate"
    # end
    EPAinitial = EPAinitial[vec(valid_rows), :]; # new vec with no NaNs
    EPAfinal = EPAfinal[vec(valid_rows), :]; # new vec with no NaNs

    lostParticles = length(valid_rows) - length(EPAfinal[:,1]);
    # @info "$lostParticles particles lost at index $final."

    @inbounds for EPAi in eachrow(EPAinitial)
        k,l = 1, 1;
        while Egrid[k] < EPAi[1]; k+=1; end # energy
        while PAgrid[l] < EPAi[2]; l+=1; end # PA
        push!(indices, (k-1,l-1))
        push!(f0Vec, distFunc(EPAi[1], EPAi[2]))
        f[k-1,l-1] += 1;
    end

    psdVec = [(f0Vec[i]/f[indices[i][1],indices[i][2]]) for i in eachindex(f0Vec[:,1])]

    @inbounds for i in eachindex(EPAinitial[:,1]) # i iterates over each particle
        if ~(EPAfinal[i,1]>maximum(Egrid) || EPAfinal[i,2]>maximum(PAgrid)) # skip loop if data is outside of range
            k,l = 1, 1;
            while Egrid[k] < EPAinitial[i,1]; k+=1; end
            while PAgrid[l] < EPAinitial[i,2]; l+=1; end
            if k!=1 && l!=1; psd_init[k-1,l-1] += psdVec[i]; end
            k,l = 1, 1;
            while Egrid[k] < EPAfinal[i,1]; k+=1; end
            while PAgrid[l] < EPAfinal[i,2]; l+=1; end
            if k!=1 && l!=1; psd_final[k-1,l-1] += psdVec[i]; end
            if k!=1 && prec_indices[i] == 1.0;  # handle precipitaitng particles here
                psd_prec[k-1] += psdVec[i]*whistler_occurence_rate;
            end
        else
            excludedParticles += 1;
        end
    end
    # @info "Excluded $excludedParticles particles due to out of range."

    return f, psd_init, psd_final, psd_prec
end

function make_psd_timeseries(Ematrix,PAmatrix,tVec, dist_func, Egrid, PAgrid, whistler_occurence_rate)
    # recalc psd at every
    psd_timeseries = Vector{Matrix{Float64}}()
    f_timeseries = copy(psd_timeseries)
    psd_prec_timeseries = Vector{Vector{Float64}}()
    @inbounds for time_index in eachindex(tVec[1:end-1])
        f, _, psd_final, psd_prec = recalcDistFunc(Ematrix, PAmatrix, 1, time_index, dist_func, Egrid, PAgrid, whistler_occurence_rate);
        push!(f_timeseries, f)
        push!(psd_timeseries, psd_final)
        push!(psd_prec_timeseries, psd_prec)
    end
    return f_timeseries, psd_timeseries, psd_prec_timeseries
end

function bin_psd_prec_timeseries(psd_prec_timeseries, indexArray)
    # given indexArray (an array with the start indices of each time bin)
    # iterate through and sum PSD to produce PSD of given time bins
    binned_psd_prec_timeseries = Vector{Vector{Float64}}()
    index = 0
    # loop through each time bin
    @inbounds for time_bin_index in indexArray
        binned_psd = zeros(length(psd_prec_timeseries[1]))
        index += 1
        while index < time_bin_index # at every index, add in the
            @. binned_psd += psd_prec_timeseries[index]
            index += 1
        end
        push!(binned_psd_prec_timeseries, binned_psd)
    end
    return binned_psd_prec_timeseries
end

function calc_equatorial_fluxes(result_matrix::Resultant_Matrix, dist_func)
    # initial electron PSD is PSD_0i = flux(E0)/E0 of the recalced distribution
    _, psd_init, _, _ = recalcDistFunc(result_matrix.Ematrix,result_matrix.PAmatrix, 1, 1,dist_func, Egrid, PAgrid, 1.);
    @info "Equatorial fluxes calculated from $(result_matrix.label)."
    return [sum(Erow) for Erow in eachrow(psd_init)]  .* Egrid .* 1000 ./ (length(PAgrid))
end

function calc_prec_flux(result_matrix::Resultant_Matrix,timeBin::Int64,dist_func::Function,whistler_occurence_rate::Float64)
    # calc only precipitaitng fluxes
    allPrecip, indexArray, allPrecipInitial = precipitatingParticles(result_matrix, timeBin);
    f_timeseries, psd_timeseries, psd_prec_timeseries = make_psd_timeseries(result_matrix.Ematrix,result_matrix.PAmatrix,result_matrix.tVec, dist_func, Egrid, PAgrid, whistler_occurence_rate);
    binned_psd_prec_timeseries = bin_psd_prec_timeseries(psd_prec_timeseries, indexArray);
    @info "Precipitating fluxes calculated from $(result_matrix.label)."

    return calc_precipitating_flux_timeseries(binned_psd_prec_timeseries);
end

function calc_precipitating_flux_timeseries(binned_psd_prec_timeseries)
    # convert prec flux into units of 1/cm^2/s/MeV where [s] is actually the width of timebin
    return [prec_psd.*Egrid.*1000 for prec_psd in binned_psd_prec_timeseries]
end

function animate_flux_comparison(gifFileName::String, equatorial_flux, elfin_measurements, thing::Vector{Vector{Float64}}, minY, maxY)
    pyplot()
    animDec = 1; # make a png for animation every 10 points
    animScale = 50; # i.e. animscale = 10 means every 10 seconds in animation is 1 second of simulation time (increase for longer animation)
    maxEnergy=1000
    anim = @animate for i in eachindex(thing[1:end-1])
        plot(Egrid, equatorial_flux)
        plot!(elfin_measurements)
        plot!(Egrid, thing[1:end-1], color = :gray, alpha = .5);
        plot!(Egrid,thing[i], color = :orange);
        plot!(ylim =(minY,maxY), xlim=(10,maxEnergy), yscale=:log10, legend=false);
        plot!(xlabel="Energy (keV)", ylabel="Flux (1/cm^2/s/sr/MeV)", title="Energy Fluxes of Precipitating Particles");
        annotate!(300, 0.1*maxY, text("t = $(round(tVec[indexArray[i]]*Re*L/(c),digits=3)) s"), :left)
    end every animDec
    savename = string("results/",gifFileName)
    gif(anim, savename, fps = (length(tVec)/animDec)/(animScale*endTime*Re*L/(c)))
    return anim
end

function animate_a_thing(gifFileName::String, thing::Vector{Matrix{Float64}})
    pyplot()
    animDec = 10; # make a png for animation every 10 points
    animScale = 1; # i.e. animscale = 10 means every 10 seconds in animation is 1 second of simulation time (increase for longer animation)
    initial, final = 1, 1
    anim = @animate for i in eachindex(thing)
        heatmap(PAgrid, Egrid, log10.(thing[i]), fc = :plasma, colorbar = false,
            xlabel="Pitch Angle (deg)",ylabel="Energy (keV)", yscale=:log10, title = "Recalculated PSD at t = $(round(tVec[i]*Re*L/(c),digits=2)) s")
    end every animDec
    savename = string("results/",gifFileName)
    gif(anim, savename, fps = (length(tVec)/animDec)/(animScale*endTime*Re*L/(c)))
    return anim
end

function animate_a_thing(gifFileName::String, thing::Vector{Vector{Float64}}, minY, maxY)
    pyplot()
    animDec = 1; # make a png for animation every 10 points
    animScale = 50; # i.e. animscale = 10 means every 10 seconds in animation is 1 second of simulation time (increase for longer animation)
    maxEnergy=1000
    anim = @animate for i in eachindex(thing[1:end-1])
        plot(Egrid, thing[1:end-1], color = :gray, alpha = .5);
        plot!(Egrid,thing[i], color = :orange);
        plot!(ylim =(minY,maxY), xlim=(10,maxEnergy), yscale=:log10, xscale=:log10, legend=false);
        plot!(xlabel="Energy (keV)", ylabel="PSD", title="Energy Spectra of Precipitating Particles");
        annotate!(300, 0.1*maxY, text("t = $(round(tVec[indexArray[i]]*Re*L/(c),digits=3)) s"), :left)
    end every animDec
    savename = string("results/",gifFileName)
    gif(anim, savename, fps = (length(tVec)/animDec)/(animScale*endTime*Re*L/(c)))
    return anim
end

@userplot MarginalHist
@recipe function f(h::MarginalHist)
    if length(h.args) != 2 || !(typeof(h.args[1]) <: AbstractVector) ||
        !(typeof(h.args[2]) <: AbstractVector)
        error("Marginal Histograms should be given two vectors.  Got: $(typeof(h.args))")
    end
    x, y = h.args

    # set up the subplots
    legend := false
    link := :both
    framestyle := [:none :axes :none]
    grid := false
    layout := @layout [tophist           _
                       hist2d{0.9w,0.9h} righthist]

    # main histogram2d
    @series begin

        seriestype := :histogram2d
        subplot := 2
        x, y
    end

    # these are common to both marginal histograms
    fillcolor := :black
    fillalpha := 0.3
    linealpha := 0.3
    seriestype := :histogram

    # upper histogram
    @series begin
        subplot := 1
        x
    end

    # right histogram
    @series begin
        orientation := :h
        subplot := 3
        y
    end
end

function checkDistFunc(f::Array{Float64,2}, psd_init::Array{Float64,2}, psd_final::Array{Float64,2}, initial::Int64, final::Int64, Egrid::Vector{Float64}, PAgrid::StepRange{Int64,Int64})
    origPlot =  heatmap(PAgrid, Egrid, f, fc = :plasma,
    ylim = (35,1000), yscale=:log10, xlim = (8,90),
    xlabel="Pitch Angle (deg)",ylabel="Energy (keV)", title = "Original Flat PSD at t = 0");

    initPlot =  heatmap(PAgrid, Egrid, psd_init, fc = :plasma, colorbar_scale=:log10, clim = (1e3,1e4),
    ylim = (35,1000), yscale=:log10, xlim = (8,90),
    xlabel="Pitch Angle (deg)",ylabel="Energy (keV)", title = "Recalculated PSD at t = $(round(tVec[initial]*Re*L/(c),digits=3)) s");

    # finalPlot = heatmap(PAgrid, Egrid, log10.(psd_final), fc = :plasma, colorbar = false,
    #     xlabel="Pitch Angle (deg)",ylabel="Energy (keV)", title = "Recalculated PSD at t = $(round(tVec[final]*Re*L/(c),digits=2)) s");

    # origPlotAnime = marginalhist(round.(PAmatrix[final,:]), log10.(round.(Ematrix[final,:])),
    # xlim = (0,90), yscale = :identity ,yticks = ([1, 2,3,3.70],["10", "100", "1000", "5000"]),
    # fc =:plasma, bins= 30);

    initial_dist = plot(origPlot, initPlot, layout = (1,2), size = (1500,600)) # want this for just first two plots
    # savefig(initial_dist, string("results/plots/initial_recalced_dist3.png"))

    # plot(origPlot, initPlot, origPlotAnime, finalPlot, layout = (2,2), size = (1200,1200)) # want this for all plots
    # plot(finalPlot) # want this for just one plot

end

function animatePSD(gifFileName, psd_timeseries)
    pyplot()
    animDec = 5; # make a png for animation every 10 points
    animScale = 10; # i.e. animscale = 10 means every 10 seconds in animation is 1 second of simulation time (increase for longer animation)
    maxval=minimum(maximum.(psd_timeseries))
    anim = @inbounds @animate for i in eachindex(psd_timeseries)
        heatmap(PAgrid, Egrid,(psd_timeseries[i]), fc = :plasma, colorbar = true,
        colorbar_scale=:log10, clims=(1e3, maxval), ylim = (35,1000), yscale=:log10,
        xlim = (8,90),
        xlabel="Pitch Angle (deg)",ylabel="Energy (keV)", title = "Recalculated PSD at t = $(round(tVec[i]*Re*L/(c),digits=2)) s")
    end every animDec
    savename = string("results/",gifFileName)
    gif(anim, savename, fps = (length(tVec)/animDec)/(animScale*endTime*Re*L/(c)))
end

function get_Nloss_Ntotal_per_Energy(tVec, Ematrix)
    # get Nloss and Ntotal for each energy at each timestep
    # should return a length(tVec) long vector of vectors, each w/ dimension length(Egrid)
    for index in eachindex(tVec)
        Ematrix()
        #wip
    end
end

function precipitatingParticles(rm::Resultant_Matrix, timeBin=10)
    unitSimTime = mean(diff(rm.tVec[begin:end-1]));
    simTimeBin = convert(Int64, round(timeBin/unitSimTime, digits=0))

    # produce a vector with indices demarcating time bin boundaries
    indexArray = [minimum([i*simTimeBin length(rm.tVec)]) for i in 1:convert(Int64,floor(rm.endTime/timeBin))]
    if length(rm.tVec)!=indexArray[end] push!(indexArray, length(rm.tVec)-1) end #guarantee last index corresponds with proper end
    # remove all particles that didn't get lost
    Ematrixprime = @views Matrix(rm.Ematrix[:,findall(isnan, rm.Ematrix[indexArray[end-1],:])]) # this is a new matrix that only has particles that have precipitated
    # replace last value with second to last value to prevent off by one problems
    if indexArray[end] == length(rm.tVec)
        indexArray[end] = indexArray[end]-1
    end

    allPrecipInitial = Vector{Vector{Float64}}()
    allPrecip = Vector{Vector{Float64}}()
    for i in indexArray[1:end-1]
        @debug "$(length(Ematrixprime[1,:])) remaining particles" size(Ematrixprime)
        precipitatingIndices = @views findall(isnan, Ematrixprime[simTimeBin,:])
        notPrecipitatingIndices = @views findall(!isnan, Ematrixprime[simTimeBin,:])
        @debug "$(length(precipitatingIndices)) particles precipitated by index $i" precipitatingIndices
        # push!(allPrecip,[@views mean(Ematrixprime[:,j][findall(!isnan, Ematrixprime[:,j])]) for j in precipitatingIndices]) # provides mean of time bin
        push!(allPrecip,[Ematrixprime[:,j][last_index_before_nan(Ematrixprime[:,j])] for j in precipitatingIndices]) # provides last energy before precipitating
        push!(allPrecipInitial,[Ematrixprime[:,j][1] for j in precipitatingIndices]) # provides the initial energy of particle about to precipitate

        @debug "added stuff" size(Ematrixprime)
        if !isempty(notPrecipitatingIndices)
            Ematrixprime = @views Matrix(Ematrixprime[simTimeBin:end,notPrecipitatingIndices]) # purposely including one extra index in case particle is lost right after
        else
            break
        end
    end

    return allPrecip, indexArray, allPrecipInitial

end

function animatePrecipitatingParticles(gifFileName, rm::Resultant_Matrix, timeBin=10, maxFraction=0.004, maxEnergy=1000)
    allPrecip, indexArray = precipitatingParticles(rm, timeBin)
    animDec = 1; # make a png for animation every 10 points
    animScale = 50; # i.e. animscale = 10 means every 10 seconds in animation is 1 second of simulation time (increase for longer animation)
    anim = @animate for i in eachindex(allPrecip)
        density(allPrecip, color = :gray, alpha = .5);
        density!(allPrecip[i], color = :orange);
        plot!(ylim =(0.,maxFraction), yscale=:identity, xlim=(20,maxEnergy), xscale=:log10, legend=false);
        plot!(xlabel="Energy (keV)", ylabel="Fraction of $(rm.numParticles) total particles", title="Energy Spectra of Precipitating Particles");
        annotate!(40, 0.1*maxFraction, text("t = $(round(rm.tVec[indexArray[i]]*Re*L/(c),digits=3)) s"), :left)
    end every animDec
    savename = string("results/",gifFileName)
    gif(anim, savename, fps = (length(rm.tVec)/animDec)/(animScale*rm.endTime*Re*L/(c)))
end

function trajectoryChecking(index)
    converT = Re*L/(c)
    @info "starting PA: $(PAmatrix[1,index]), startin E: $(Ematrix[1,index])"
    plot(converT*tVec, PAmatrix[:,index]);
    plot!(converT*tVec, Ematrix[:,index])
    # plot(Zmatrix[:,index], PZmatrix[:,index])
end

function trajectoryTracing(trajectories, animDec)
    converT = Re*L/(c)

    PAanim = @animate for i in eachindex(tVec)
        plot(xlim=(0,12), ylim = (0,90));
        plot!(converT*tVec, PAmatrix[:,trajectories], color=:gray, alpha = .5, label = "");
        for j in trajectories
            plot!(converT*tVec[1:i], PAmatrix[:,j][1:i], label = "")
            scatter!((converT*tVec[i], PAmatrix[:,j][i]), label = "")
        end
    end every animDec

    savename1 = string("results/plots/","PAgif.gif")
    gif(PAanim, savename1, fps = 20)

    Eanim = @animate for i in eachindex(tVec)
        Eplot = plot(xlim=(0,12), ylim = (10,5000), yscale = :log10)
        plot!(converT*tVec, Ematrix[:,trajectories], color=:gray, alpha = .5, label = "");
        for j in trajectories
            Eplot = plot!(converT*tVec[1:i], Ematrix[:,j][1:i], label = "")
            Eplot = scatter!((converT*tVec[i], Ematrix[:,j][i]), label = "")
        end
    end every animDec

    savename2 = string("results/plots/","Egif.gif")
    gif(Eanim, savename2, fps = 20)

    ZPZanim = @animate for i in eachindex(tVec)
        ZPZplot = plot(xlim=(-1,1), ylim = (-7,7))
        for j in trajectories
            ZPZplot = plot!(Zmatrix[:,j][1:i], PZmatrix[:,j][1:i], alpha = max.((1:i) .+ 100 .- i, 0) / 100, label = "")
            ZPZplot = scatter!((Zmatrix[:,j][i], PZmatrix[:,j][i]), label = "")
        end
    end every animDec

    savename3 = string("results/plots/","ZPZgif.gif")
    gif(ZPZanim, savename3, fps = 20)

end

function last_index_before_nan(x::Vector{Float64})
    k = 1
    @inbounds for i in eachindex(x)
        if isnan(x[i])
            k = i-1
            return k
        end
    end
    return
 end

function find_lost_particles(row1::Vector{Float64}, row2::Vector{Float64})
    # compares two rows and returns the index of all particles from row 1 that will be lost in row 2

    # valid_rows = @. ~isnan(EPAfinal)[:,1] & ~isnan(EPAfinal[:,2]); # non-lost particles from final state
    # valid_rows_next = @. ~isnan(EPAfinal_next)[:,1] & ~isnan(EPAfinal_next[:,2]); # non-lost particles from state at next timestep
    # # if valid_rows!=valid_rows_next # this means that at least one of the particles will precipitate
    # prec_rows = valid_rows .!= valid_rows_next # these are the indices of particles who are about to precipitate
    # prec_indices = (ones(N).*vec(prec_rows))[vec(valid_rows)]
    # @info "$(length(ones(N)[prec_rows])) particles are about to precipitate"

end

function extract_idl_csv(
    time_name::String,
    data_name::String,
    ebin_name::String,
    start::DateTime, stop::DateTime)

    time_csv_name = "idl_csvs/"*time_name
    data_csv_name = "idl_csvs/"*data_name
    ebins_csv_name = "idl_csvs/"*ebin_name

    times_df =  CSV.File(time_csv_name; header=false, delim=',', types=Float64) |> DataFrame
    time = unix2datetime.(times_df.Column1)
    indices = findall((time.>start).&(time.<stop)) # these are the indices corresponding to the time range to sum over
    time_of_interest = time[indices]
    @info "Summing over $(time_of_interest[end]-time_of_interest[1])"

    # import particle flux data
    data_df  =  CSV.File(data_csv_name; header=false, delim=',', types=Float64) |> DataFrame
    data = [[data_df[row,col] for col in 1:length(data_df[1,:])] for row in 1:16]
    # get rid of NaNs and Infs
    data_of_interest = data
    for i = 1:16
        data_of_interest[i][findall(.!isfinite.(data[i]))] .= 0.0
    end
    flux = [sum(data_of_interest[energy][indices]) for energy in 1:16]

    ebins_df = CSV.File(ebins_csv_name; header=false, delim=',', types=Float64) |> DataFrame
    ebins = ebins_df.Column1

    return (ebins, flux)
end

# time_name = "092220_time.csv"
# data_name = "092220_prec.csv"
# error_name = "092220_precerror.csv"
# error_time_name = "092220_precerror_time.csv"
# ebin_name = "ebins.csv"
# start = DateTime(2020,9,22,9,16,38)
# stop = DateTime(2020,9,22,9,16,45)

function extract_idl_csv(
    time_name::String,
    data_name::String,
    error_name::String,
    error_time_name::String,
    ebin_name::String,
    start::DateTime, stop::DateTime)

    time_csv_name = "idl_csvs/"*time_name
    data_csv_name = "idl_csvs/"*data_name
    ebins_csv_name = "idl_csvs/"*ebin_name
    error_csv_name = "idl_csvs/"*error_name
    errortime_name = "idl_csvs/"*error_time_name

    times_df =  CSV.File(time_csv_name; header=false, delim=',', types=Float64) |> DataFrame
    time = unix2datetime.(times_df.Column1)
    indices = findall((time.>start).&(time.<stop)) # these are the indices corresponding to the time range to sum over
    time_of_interest = time[indices]
    @info "Summing over $(time_of_interest[end]-time_of_interest[1])"

    errtimes_df = CSV.File(errortime_name; header=false, delim=',', types=Float64) |> DataFrame
    errtime = unix2datetime.(errtimes_df.Column1)
    errindices = findall((errtime.>start).&(errtime.<stop)) # these are the indices corresponding to the time range to sum over
    if length(errindices) != length(indices)
        @error "Length mismatch between time and error time."
        return
    end

    # import particle flux data
    data_df  =  CSV.File(data_csv_name; header=false, delim=',', types=Float64) |> DataFrame
    data = [[data_df[row,col] for col in 1:length(data_df[1,:])] for row in 1:16]
    error_df  =  CSV.File(error_csv_name; header=false, delim=',', types=Float64) |> DataFrame
    error = [[error_df[row,col] for col in 1:length(error_df[1,:])] for row in 1:16]
    # get rid of NaNs and Infs
    data_of_interest = data
    error_of_interest = error
    for i = 1:16
        data_of_interest[i][findall(.!isfinite.(data[i]))] .= 0.0
        error_of_interest[i][findall(.!isfinite.(error[i]))] .= 0.0
    end
    flux = [sum(data_of_interest[energy][indices]) for energy in 1:16]

    error = [sqrt(sum((data_of_interest[energy][indices].*error_of_interest[energy][errindices]).^2)) for energy in 1:16]

    ebins_df = CSV.File(ebins_csv_name; header=false, delim=',', types=Float64) |> DataFrame
    ebins = ebins_df.Column1

    return (ebins, flux), error
end

function generate_flux_comparison(result_matrix::Resultant_Matrix, timeBin, dist_func, whistler_occurence_rate, Egrid, PAgrid,time_name, data_name, ebin_name, start, stop)
    allPrecip, indexArray, allPrecipInitial = precipitatingParticles(result_matrix.tVec, result_matrix.Ematrix, result_matrix.endTime, timeBin);
    f_timeseries, psd_timeseries, psd_prec_timeseries = make_psd_timeseries(result_matrix.Ematrix,result_matrix.PAmatrix,result_matrix.tVec, dist_func, Egrid, PAgrid, whistler_occurence_rate);
    binned_psd_prec_timeseries = bin_psd_prec_timeseries(psd_prec_timeseries, indexArray);
    equatorial_fluxes = calc_equatorial_fluxes(result_matrix.Ematrix,result_matrix.PAmatrix, dist_func, Egrid, PAgrid);
    prec_flux_timeseries = calc_precipitating_flux_timeseries(binned_psd_prec_timeseries);
    elfin_measurements = extract_idl_csv(time_name, data_name, ebin_name, start, stop);
    return equatorial_fluxes, elfin_measurements, prec_flux_timeseries
end

function generate_trapped_psd(dist_func, whistler_occurence_rate::Float64)
    f_timeseries, psd_timeseries, psd_prec_timeseries = make_psd_timeseries(Ematrix, PAmatrix, tVec, dist_func, Egrid, 8:2:90, whistler_occurence_rate)
    return calc_precipitating_flux_timeseries(bin_psd_prec_timeseries([psd_timeseries[particle][:,1] for particle in 1:length(tVec)-1], indexArray))
end

# Useful helpers
logrange(x1, x2, n::Int64) = [10^y for y in range(log10(x1), log10(x2), length=n)]

# df = DataFrame([allPrecip], :auto)
# CSV.write("for_james.csv",df)


include("graphs.jl")
using DataFrames
using Plots
using Statistics
using HypothesisTests
using CSV
using Weave
using IJulia
using Random

function benchmark(exp::Experiment)
    workspace = exp.workspace

    if !ispath(workspace * "/notebooks")
        mkpath(workspace * "/notebooks")
    end

    convert_doc(
        "./src/benchmark/notebooks/overview.jmd",
        workspace * "/notebooks/overview.ipynb",
    )

    IJulia.jupyterlab(dir = workspace * "/notebooks")
end

function benchmark(path::String = "_results")
    df = DataFrames.DataFrame()
    for item in readdir(path)
        println(item)
        maindf = DataFrame(CSV.File(path * '/' * item))
        rename!(maindf, :iteration => :runtime, :size => :problem_size)
        df = vcat(df, maindf[:, [:algorithm, :runtime, :problem_size]])
    end
    resultpvalue(df)
    return benchmark(df)
end

function benchmark(traces::Array{Trace,1})
    df = tracestodataframe(traces)
    resultpvalue(df)
    return benchmark(df)
end

function tracestodataframe(traces::Array{Trace,1})
    result = DataFrames.DataFrame()
    for trace in traces
        push!(
            result,
            (
                algorithm = string(typeof(trace.algorithm)),
                runtime = trace.optimum[2],
                problem_size = length(trace.individual.problem),
            ),
        )
    end
    result
end

function jobstodataframe(jobs::Array{Job,1}) end

function benchmark(data::DataFrames.DataFrame)
    algorithms = keys(groupby(data, "algorithm"))
    plts = []

    append!(plts, [plot()])
    for algo in algorithms
        classical_0 = filter(row -> row["algorithm"] == algo[1], data)
        classical_0 = combine(groupby(classical_0, :problem_size), :runtime => mean)
        y = classical_0[!, :runtime_mean]
        x = classical_0[!, :problem_size]
        # v = classical_0[!,4]
        plot!(
            plts[end],
            x,
            y,
            fillalpha = 0.3,
            formatter = :plain,
            dpi = 300,
            fontfamily = "Times New Roman Bold",
            legend = :right,
            linewidth = 2,
            markersize = 3,
            markershape = :circle,
            linealpha = 0.5,
            size = (600, 300),
            label = algo,
        )
    end
    return plts[1]
end

function resultpvalue(df::DataFrames.DataFrame)
    algo = keys(groupby(df, :algorithm))
    problem_size = keys(groupby(df, :ProblemParam_n))
    allowmissing!(df)

    for item in problem_size
        println("Problem Size: ", item["ProblemParam_n"])
        data2 = groupby(
            filter(row -> row["ProblemParam_n"] == item["ProblemParam_n"] && !ismissing(row["runtime"]), df),
            :algorithm,
        )
        for item in data2
            for item2 in data2
                p_value = pvalue(MannWhitneyUTest(item[!, "runtime"], item2[!, "runtime"]))
                println(item[1][1], " vs ", item2[1][1], ", Pvalue = ", p_value)
            end
        end
    end
end

function runtimes(jobs::Array{Job,1})
    result = DataFrames.DataFrame()
    allowmissing!(result)
    for job in jobs
        trace = job.trace
        if !isnothing(trace) && job.status == finished
            thing = (
                seed = trace.seed,
                algorithm = name(trace.algorithm),
                problem = string(trace.individual.problem),
                runtime = trace.optimum[2],
            )
            push!(result, thing)
        end
    end

    for job in shuffle(jobs)
        trace = job.trace
        if !isnothing(trace) && job.status == finished && size(trace.rows)[1]>1
            thing = (seed = trace.seed, algorithm = name(trace.algorithm))
            paramgroups = (
                ProblemParam = deepparameters(trace.individual.problem),
                AlgorithmParams = deepparameters(trace.algorithm),
            )
            for paramgroupkey in keys(paramgroups)
                for p in paramgroups[paramgroupkey]
                    label = Symbol(string(paramgroupkey) * "_" * string(p[1]))
                    if string(label) in names(result)
                        result[result[!,:seed].==trace.seed, label] .= p[2]
                    else
                        df = DataFrame([(seed = trace.seed, l = p[2])])
                        rename!(df, :l => label)
                        result = outerjoin(result, df, on = [:seed], makeunique = true)
                    end
                end
            end
        end
    end

    result
end

runtimes(exp::Experiment) = runtimes(exp.jobs)


#!/usr/bin/env julia

# Adapted from bin/cellfishing in the CellFishing.jl repo

using CellFishing
using HDF5
using DocOpt

using Dates
using Random
using SparseArrays
using StatsBase
using Statistics

args = docopt("""
Usage:
    cellfishing [--n-bits=<n-bits> --n-lshashes=<n-lshashes> --infer-stats=<infer-stats> --genes=<genes> --n-neighbors=<n-neighbors> --cutoff=<cutoff> --min-hits=<min-hits> --majority=<majority> --annotation=<annotation> --seed=<seed> --clean=<clean> --subsample-ref=<subsample-ref>] <reference> <query> <output>

CellFishing ― Cell search engine for scRNA-seq.

Options:
    -h --help                        Show this help.
    --version                        Show version.
    --n-bits=<n-bits>                The number of bits for hashing [default: 128].
    --n-lshashes=<n-lshashes>        The number of locality-sensitive hashes [default: 4].
    --infer-stats=<infer-stats>      Infer statistics from the query, the database, or both [default: both].
    --genes=<genes>                  Specify genes to use instead of selecting by CellFishing [default: cf_genes].
    --n-neighbors=<n-neighbors>      The number of nearest neighbors to search for [default: 10].
    --cutoff=<cutoff>...             Hamming distance cutoff [default: 120].
    --min-hits=<min-hits>            Minimal number of hits below which query will be rejected [default: 2].
    --majority=<majority>            Minimal fraction required to pass majority vote [default: 0.5].
    --annotation=<annotation>        Choose annotation column in meta table [default: cell_ontology_class].
    --clean=<clean>                  Clean data based on certain meta column [default: None].
    --subsample-ref=<subsample-ref>  Subsample reference data [default: None].
    --seed=<seed>                    Set random seed [default: None].
    --self                           Self-projection.
""")

function loadh5(datafile, clean=nothing, subsample=nothing, annotation=nothing, genes=nothing)
    HDF5.h5open(datafile) do file
        counts = read(file, "/exprs")
        if isa(counts, Dict)
            counts = SparseMatrixCSC(
                counts["shape"][2], counts["shape"][1],
                counts["indptr"] .+ 1, counts["indices"] .+ 1,
                counts["data"]
            )
        else
            counts = Matrix(counts)
        end
        featurenames = read(file, "/var_names")
        cellnames = read(file, "/obs_names")
        if annotation != nothing
            annotations = read(file, "/obs/$(annotation)")
        else
            annotations = nothing
        end
        if clean != nothing
            clean = read(file, "/obs/$(clean)")
            mask = clean .∈ [["" "na" "NA" "nan" "NaN"]]
            if count(mask) > 0
                print("Cleaning removed $(count(mask)) cells. ")
                counts = counts[:, .!mask]
                cellnames = cellnames[.!mask]
                if annotations != nothing
                    annotations = annotations[.!mask]
                end
            end
        end
        if subsample != nothing
            subsample = sample(1:size(counts, 2), subsample, replace=false)
            counts = counts[:, subsample]
            cellnames = cellnames[subsample]
            if annotations != nothing
                annotations = annotations[subsample]
            end
        end
        if genes != nothing
            genes = read(file, "/uns/$(genes)")
        else
            genes = nothing
        end
        return (counts=counts, featurenames=featurenames, cellnames=cellnames, annotations=annotations, genes=genes)
    end
end

macro message(msg, block)
    quote
        print(stderr, "  ", rpad(string($(esc(msg)), " "), 25, '\u2015'), " ")
        starttime = now()
        $(esc(block))
        println(stderr, Dates.canonicalize(Dates.CompoundPeriod(now() - starttime)))
    end
end

function writeresult(file, predictions, neighbors, time_per_cell)
    HDF5.h5open(file, "w") do file
        g = g_create(file, "prediction")
        for item in predictions
            k, v = item
            g[String(k)] = v
        end
        file["indexes"] = neighbors.indexes
        file["hammingdistances"] = neighbors.hammingdistances
        file["time"] = time_per_cell
    end
end

using Random: seed!
if args["--seed"] != "None"
    seed!(parse(Int, args["--seed"]))
end

if args["--clean"] == "None"
    args["--clean"] = nothing
end
if args["--subsample-ref"] == "None"
    args["--subsample-ref"] = nothing
else
    args["--subsample-ref"] = parse(Int, args["--subsample-ref"])
end
if args["--genes"] == "cf_genes"
    args["--genes"] = nothing
end

@message "Loading data" begin
    ref = loadh5(args["<reference>"], args["--clean"], args["--subsample-ref"], args["--annotation"], args["--genes"])
end
@message "Selecting features" begin
    if ref.genes == nothing
        features = CellFishing.selectfeatures(ref.counts, ref.featurenames)
    else
        nzv_features = ref.featurenames[dropdims(std(ref.counts, dims=2) .> 0, dims=2)]
        println("Removing $(size(setdiff(ref.genes, nzv_features), 1)) genes with zero variance.")
        ref_genes = intersect(nzv_features, ref.genes)
        features = CellFishing.Features(ref.featurenames, ref.featurenames .∈ [ref_genes])
    end
end
@message "Creating a database" begin
    n_bits = parse(Int, args["--n-bits"])
    n_lshashes = parse(Int, args["--n-lshashes"])
    database = CellFishing.CellIndex(ref.counts, features, metadata=ref.annotations, n_bits=n_bits, n_lshashes=n_lshashes)
end

k = parse(Float64, args["--n-neighbors"])
if k < 1
    k = round(k * size(ref.counts, 2))
    if args["--self"]
        k += 1
    end
end
k = Int64(k)

@message "Loading query data" begin
    query = loadh5(args["<query>"], args["--clean"])
end

@message "Searching the database" begin
    inferstats = args["--infer-stats"]
    inferstats = inferstats == "query"    ? :query :
                    inferstats == "database" ? :database :
                    inferstats == "both"     ? :both :
                    error("invalid --infer-stats option")
    start_time = now()
    neighbors = CellFishing.findneighbors(k, query.counts, query.featurenames, database, inferstats=inferstats)
end

time_per_cell = nothing
@message "Doing annotation" begin
    predictions = Dict()
    majority = parse(Float64, args["--majority"])
    min_hits = parse(Int, args["--min-hits"])
    cutoffs = split(args["--cutoff"], ",")
    for cutoff in cutoffs
        global time_per_cell
        predictions[cutoff] = Vector{String}()
        _cutoff = parse(Int, cutoff)
        for qid in 1:size(neighbors.indexes, 2)
            mask = neighbors.hammingdistances[:, qid] .≤ _cutoff
            annotations = database.metadata[neighbors.indexes[mask, qid]]
            if size(annotations, 1) < min_hits
                push!(predictions[cutoff], "rejected")
                continue
            end
            max_k, max_v = nothing, 0
            for item in countmap(annotations)
                k, v = item
                if v > max_v
                    max_k, max_v = k, v
                end
            end
            if max_v ≤ majority * size(annotations, 1)
                push!(predictions[cutoff], "ambiguous")
            else
                push!(predictions[cutoff], string(max_k))
            end
        end
        if time_per_cell == nothing
            time_per_cell = (now() - start_time).value / size(query.counts, 2)
        end
    end
end

@message "Saving results" begin
    writeresult(args["<output>"], predictions, neighbors, time_per_cell)
end


module GLSLBuildins

# not in base, so inferable stump must be defined
gl_FragCoord() = zero(Vec2f0)
# replace rule for transpiler
gl_FragCoord(backend::GL) = TGlobal("gl_FragCoord", Vec2f0)

gl_InstanceID() = Int32(0)
gl_InstanceID(backend::GL) = TGlobal("gl_InstanceID", Vec2f0)

#Base.sin(x) already defined
sin{T}(backend::GL, x::Type{T}) = TExpr("sin", T)

EmitVertex() = nothing
sin{T}(backend::GL) = TExpr("EmitVertex", T)

end


module TestCore

using Test
using MLJIteration
using IterationControl
using MLJBase
using MLJModelInterface
using ..DummyModel

X, y = make_dummy(N=20)
model = DummyIterativeModel(n=0)

@testset "integration: resampling=nothing" begin

    model = DummyIterativeModel(n=0)
    controls=[Step(2), Threshold(0.01), TimeLimit(0.005)]

    # using IterationControl.jl directly:
    mach = machine(deepcopy(model), X, y);
    function IterationControl.train!(mach::Machine{<:DummyIterativeModel},
                                     n::Int)
        mlj_model = mach.model
        mlj_model.n = mlj_model.n + n
        fit!(mach, verbosity=0)
    end
    IterationControl.loss(mach::Machine{<:DummyIterativeModel}) =
        last(training_losses(mach))
    IterationControl.train!(mach, controls..., verbosity=0)
    losses1 = report(mach).training_losses
    yhat1 = predict(mach, X)

    # using IteratedModel wrapper:
    imodel = IteratedModel(model=model,
                           controls=controls,
                           resampling=nothing)
    mach = machine(imodel, X, y)
    @test_logs((:info, r"Training"),
               (:info, MLJIteration.info_unspecified_iteration_parameter(:n)),
               (:info, r"final loss"),
               (:info, r"final training loss"),
               (:info, r"Stop"),
               (:info, r"Total"),
               fit!(mach, verbosity=1))
    imodel.resampling = nothing
    @test_logs fit!(mach, verbosity=0)

    losses2 = report(mach).model_report.training_losses
    yhat2 = predict(mach, X);
    @test imodel.model == DummyIterativeModel(n=0) # hygeine check

    # compare:
    @test losses1 ≈ losses2
    @test yhat1 ≈ yhat2

    # report:
    r = report(mach)
    @test r.model_report.training_losses ≈ losses1
    @test :controls in keys(r)
    i = r.n_iterations
    @test i == 38

    # warm restart when changing controls:
    noise(n) = fill((:info, r"43"), n)
    imodel.controls = [Step(1),
                       WithIterationsDo(i-> i>39,
                                        stop_if_true=true),
                       Info(x->"43")]
    @test_logs((:info, r"Updating"),
               noise(2)...,
               (:info, r""),
               (:info, r""),
               (:info, r""),
               (:info, r""),
               fit!(mach))
    @test report(mach).n_iterations == i + 2

    # warm restart when changing iteration parameter (trains one more
    # iteration because stopping control comes after `Step(...)`):
    imodel.model.n = 1
    @test_logs((:info, r"Updating"),
               noise(1)...,
               (:info, r""),
               (:info, r""),
               (:info, r""),
               (:info, r""),
               fit!(mach))
    @test report(mach).n_iterations == i + 3

    # cold restart when changing anything else:
    imodel.check_measure=false
    imodel.controls = [Step(1),
                       WithIterationsDo(i-> i>4,
                                        stop_if_true=true),
                       Info(x->"43")]
    @test_logs((:info, r"Updating"),
               (:info, MLJIteration.info_unspecified_iteration_parameter(:n)),
               noise(5)...,
               (:info, r""),
               (:info, r""),
               (:info, r""),
               (:info, r""),
               fit!(mach))
    @test report(mach).n_iterations == 5
end

@testset "integration: resampling=Holdout()" begin

    controls=[Step(2), Patience(4), TimeLimit(0.001)]

    # using IterationControl.jl directly:
    mach = machine(deepcopy(model), X, y);
    train, test = partition(eachindex(y), 0.7)
    function IterationControl.train!(mach::Machine{<:DummyIterativeModel},
                                     n::Int)
        mlj_model = mach.model
        mlj_model.n = mlj_model.n + n
        fit!(mach, rows=train, verbosity=0)
    end
    function IterationControl.loss(mach::Machine{<:DummyIterativeModel})
        yhat = predict(mach, rows=test)
        return mae(yhat, y[test]) |> mean
    end
    IterationControl.train!(mach, controls...; verbosity=0)
    losses1 = report(mach).training_losses;
    yhat1 = predict(mach, X[test]);
    niters = mach.model.n
    @test niters == length(losses1)

    # using IteratedModel wrapper:
    imodel = IteratedModel(model=model,
                           resampling=Holdout(fraction_train=0.7),
                           controls=controls)
    mach = machine(imodel, X, y)
    @test_logs((:info, r"Training"),
               (:info, MLJIteration.info_unspecified_iteration_parameter(:n)),
               (:info, MLJIteration.info_unspecified_measure(rms)),
               (:info, r"final loss"),
               (:info, r"final train"),
               (:info, r"Stop"),
               (:info, r"Total"),
               fit!(mach, verbosity=1))
    imodel.measure = mae
    @test_logs fit!(mach, verbosity=0)
    losses2 = report(mach).model_report.training_losses;
    yhat2 = predict(mach, X[test]);

    # compare:
    @test losses1 ≈ losses2
    @test yhat1 ≈ yhat2

    # now repeat wrapper train with retrain=true:
    imodel = IteratedModel(model=model,
                           resampling=Holdout(fraction_train=0.7),
                           controls=controls,
                           retrain=true,
                           measure=mae)
    mach = machine(imodel, X, y)
    fit!(mach, verbosity=0)
    losses3 = report(mach).model_report.training_losses
    yhat3 = predict(mach, X[test])
    @test fitted_params(mach).machine.model.n == niters

    # to be compared with atomic model trained on all data for same
    # number of iterations:
    model2 = DummyIterativeModel(n=niters)
    mach2 = machine(model2, X, y)
    fit!(mach2, verbosity=0)
    losses4 = report(mach2).training_losses
    yhat4 = predict(mach2, X[test])

    # check report:
    r = report(mach2)

    # compare:
    @test losses3 ≈ losses4
    @test yhat3 ≈ yhat4
end

@testset "integration: cyclic learning rates" begin

    X, y = make_dummy(N=100)
    ss = 3
    model = DummyIterativeModel()

    # constant learning rate:
    controls = [CycleLearningRate(stepsize=ss,
                                  lower=1,
                                  upper=1,
                                  learning_rate_parameter=:learning_rate),
                Step(1),
                NumberLimit(5*ss)]

    imodel = IteratedModel(model=model,
                           resampling=Holdout(fraction_train=0.7),
                           controls=controls,
                           measure=mae)

    mach = machine(imodel, X, y)
    fit!(mach, verbosity=0)
    yhat1 = predict(mach, X)

    # no cycling:
    imodel.controls = [Step(1), NumberLimit(5*ss)]
    fit!(mach, force=true, verbosity=0)
    yhat2 = predict(mach, X)

    # with real cycling:
    controls = [CycleLearningRate(stepsize=ss,
                                  lower=0.5,
                                  upper=1.5,
                                  learning_rate_parameter=:learning_rate),
                Step(1),
                NumberLimit(5*ss)]
    imodel.controls = controls
    fit!(mach, force=true, verbosity=0)
    yhat3 = predict(mach, X)

    # compare:
    @test yhat1 ≈ yhat2
    @test !(yhat3 ≈ yhat2)

end

@testset "issue #40" begin
    # change trait to an incorrect value:
    MLJModelInterface.iteration_parameter(::Type{<:DummyIterativeModel}) = :junk
    @test iteration_parameter(model) == :junk

    # user specifies correct value:
    iterated_model = IteratedModel(model, measure=mae, iteration_parameter=:n)
    mach = machine(iterated_model, X, y)

    # and model still runs:
    fit!(mach, verbosity=0)

    # change trait back:
    MLJModelInterface.iteration_parameter(::Type{<:DummyIterativeModel}) = :n
    @test iteration_parameter(model) == :n
end

end
true


# Load the grammar type.
include("grammar_derivation/grammar.jl")

"""
Provides a definition for a grammar derivation representation.
"""
type GrammarDerivationRepresentationDefinition <: RepresentationDefinition
  """
  The maximum number of times that the genome is allowed to wrap around itself
  if it has run out of codons to use. By default, wrapping is disabled.
  """
  max_wraps::Int

  """
  The grammar that should be used to produce the derivations.
  """
  grammar::Grammar

  GrammarDerivationRepresentationDefinition() = new(0)
end

type GrammarDerivationRepresentation <: Representation
  max_wraps::Int
  grammar::Grammar

  GrammarDerivationRepresentation(max_wraps::Int, grammar::Grammar) =
    new(max_wraps, grammar)
end

chromosome(::GrammarDerivationRepresentation) = AbstractString

"""
TODO: Document grammar derivation representations.
"""
function grammar(def::Function)
  g = GrammarDerivationRepresentationDefinition()
  def(g)
  g
end

"""
Composes a grammar derivation representation from a provided definition.
"""
function compose!(def::GrammarDerivationRepresentationDefinition)
  GrammarDerivationRepresentation(def.max_wraps, def.grammar)
end

"""
Converts from a vector of codons to a grammar derivation.
"""
convert!(
  r_from::IntVectorRepresentation,
  r_to::GrammarDerivationRepresentation,
  c_from::IndividualStage{Vector{Int}},
  c_to::IndividualStage{AbstractString}
) =
  c_to = derive(r_to, c_from.value)

"""
Produces a grammar derivation from a given sequence of codons, according to
a provided grammar derivation representation.
"""
derive(r::GrammarDerivationRepresentation, s::Vector{Int}) =
  derive(r.grammar, s, r.max_wraps)


# This file was generated, do not modify it. # hide
using Random # hide
rng = Random.GLOBAL_RNG
rand(rng, Float32, Normal())


using Base: depwarn
using Flux: Params

check_decay(opt, decay) = decay == 0 ? opt : Optimiser(opt, InvDecay(decay))

# legacy update rule
updaterule(opt, ps) = () -> update!(opt, ps)

function SGD(params::Union{AbstractArray, Params}, η = 0.1; decay = 0.)
  depwarn("SGD(params) is deprecated; use Descent(η::Float64) instead", :SGD)

  ps = params
  opt = Descent(η)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function Momentum(params::Union{AbstractArray, Params}, η = 0.01; ρ = 0.9, decay = 0.)
  depwarn("Momentum(params) is deprecated; use Momentum(η::Float64) instead", :Momentum)

  ps = params
  opt = Momentum(η, ρ)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function Nesterov(params::Union{AbstractArray, Params}, η = 0.001; ρ = 0.9, decay = 0.)
  depwarn("Nesterov(params) is deprecated; use Nesterov(η::Float64) instead", :Nesterov)

  ps = params
  opt = Nesterov(η, ρ)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function RMSProp(params::Union{AbstractArray, Params}, η = 0.001; ρ = 0.9, decay = 0.)
  depwarn("RMSProp(params) is deprecated; use RMSProp(η::Float64) instead", :RMSProp)

  ps = params
  opt = RMSProp(η, ρ)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function ADAM(params::Union{AbstractArray, Params}, η = 0.001; β1 = 0.9, β2 = 0.999, decay = 0.)
  depwarn("ADAM(params) is deprecated; use ADAM(η::Float64) instead", :ADAM)

  ps = params
  β = (β1, β2)
  opt = ADAM(η, β)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function ADAGrad(params::Union{AbstractArray, Params}, η::Float64 = 0.1; decay = 0.)
  depwarn("ADAGrad(params) is deprecated; use ADAGrad(η::Float64) instead", :ADAGrad)

  ps = params
  opt = ADAGrad(η)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function ADADelta(params::Union{AbstractArray, Params}, ρ::Float64 = 0.9; decay = 0.)
  depwarn("ADADelta(params) is deprecated; use ADADelta(η::Float64) instead", :ADADelta)

  ps = params
  opt = ADADelta(ρ)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function AdaMax(params::Union{AbstractArray, Params}, η = 0.001; β1 = 0.9, β2 = 0.999, decay = 0.)
  depwarn("AdaMax(params) is deprecated; use AdaMax(η::Float64) instead", :AdaMax)

  ps = params
  β = (β1, β2)
  opt = AdaMax(η, β)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function AMSGrad(params::Union{AbstractArray, Params}, η = 0.001; β1 = 0.9, β2 = 0.999, decay = 0.)
  depwarn("AMSGrad(params) is deprecated; use AMSGrad(η::Float64) instead", :AMSGrad)

  ps = params
  β = (β1, β2)
  opt = AMSGrad(η, β)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function NADAM(params::Union{AbstractArray, Params}, η = 0.001; β1 = 0.9, β2 = 0.999, decay = 0.)
  depwarn("NADAM(params) is deprecated; use NADAM(η::Float64) instead", :NADAM)

  ps = params
  β = (β1, β2)
  opt = NADAM(η, β)
  opt = check_decay(opt, decay)
  updaterule(opt, ps)
end

function ADAMW(params::Union{AbstractArray, Params}, η = 0.001; β1 = 0.9, β2 = 0.999, decay = 0.)
  depwarn("ADAMW(params) is deprecated; use ADAMW(η::Float64) instead", :ADAMW)

  ps = params
  β = (β1, β2)
  opt = ADAMW(η, β)
  opt = check_decay(opt, decay)
  decay != 0 && (opt = Optimiser(opt, WeightDecay(decay)))
  updaterule(opt, ps)
end

# Old training loop

struct OldOptimiser
  func
end

update!(opt::OldOptimiser, ps) = opt.func()

# Train function
function train!(loss, data, opt; cb = () -> ())
  depwarn("train!(loss, data, opt) is deprecated; use train!(loss, params, data, opt) instead", :train!)
  train!(loss, (), data, OldOptimiser(opt); cb = cb)
end


include( "transitionflux.jl" )
include( "nodeflux.jl" )
include( "nodepop.jl" )
include( "nodecomposition.jl" )
include( "subpop.jl" )
include( "subpopage.jl" )
include( "initpop.jl" )

const extensions = ["svg", "html", "png"]
const plotTypes = Dict{Symbol,NTuple{2,Function}}(
    :normal => (normalFluxPlot, normalCompPlot),
    :stacked => (stackedFluxPlot, stackedCompPlot) )
const labelInfix = Dict{Symbol,NTuple{2,String}}( :in => (" into ", " from "),
    :out => (" out of ", " to "), :within => (" within ", " to ") )


@testset "Custom Fourier Types" begin
    N = 5
    x = randn((N, N))
    fs = FourierTools.FourierSplit(x, 2, 2, 4)
    @test FourierTools.parenttype(fs) == typeof(x)

    fj = FourierTools.FourierJoin(x, 2, 2, 4)

    @test FourierTools.parenttype(fj) == typeof(x)
    @test FourierTools.parenttype(typeof(fj)) == typeof(x)

    @test FourierTools.IndexStyle(typeof(fj)) == IndexStyle(typeof(fj))
end


function define_indices_d(::Type{NoInterp}, d, pad)
    symix, symx = Symbol("ix_",d), Symbol("x_",d)
    :($symix = convert(Int, $symx))
end

function coefficients(::Type{NoInterp}, N, d)
    :()
end

function index_gen{IT<:DimSpec}(::Type{NoInterp}, ::Type{IT}, N::Integer, offsets...)
    if (length(offsets) < N)
        return :($(index_gen(IT, N, offsets..., 0)))
    else
        indices = [offsetsym(offsets[d], d) for d = 1:N]
        return :(itp.coefs[$(indices...)])
    end
end

padding(::Type{NoInterp}) = Val{0}()

# How many non-NoInterp dimensions are there?
count_interp_dims(::Type{NoInterp}, N) = 0
count_interp_dims{IT<:InterpolationType}(::Type{IT}, N) = N
function count_interp_dims{IT<:Tuple{Vararg{InterpolationType}}}(it::Type{IT}, N)
    n = 0
    for p in it.parameters
        n += count_interp_dims(p, 1)
    end
    n
end


function Random.Sampler(T::Type{<:Random.AbstractRNG}, x::Consistent, ::Val{Inf})
    # sampler for sampling indices of predictions
    predictions = x.predictions
    idxsampler = Random.Sampler(T, eachindex(predictions))
    return Random.SamplerSimple(x, (idxsampler, make_samplers(T, x)))
end

function make_samplers(T::Type{<:Random.AbstractRNG}, x::Consistent{<:Real})
    return Random.Sampler(T, float(eltype(x.predictions)))
end

# set up alias samplers
function make_samplers(
    T::Type{<:Random.AbstractRNG}, x::Consistent{<:AbstractVector{<:Real}}
)
    predictions = x.predictions
    prediction = first(predictions)
    floatsampler = Random.Sampler(T, float(eltype(prediction)))
    integersampler = Random.Sampler(T, Base.OneTo(length(prediction)))
    samplers = let floatsampler = floatsampler, integersampler = integersampler
        map(predictions) do prediction
            return AliasSampler(prediction, floatsampler, integersampler)
        end
    end
    return samplers
end

function make_samplers(T::Type{<:Random.AbstractRNG}, x::Consistent)
    samplers = let T = T
        map(x.predictions) do prediction
            return Random.Sampler(T, prediction)
        end
    end
    return samplers
end

# out-of-place sampling
function Base.rand(
    rng::Random.AbstractRNG, sampler::Random.SamplerSimple{<:Consistent{<:Real}}
)
    idxsampler, floatsampler = sampler.data
    idx = rand(rng, idxsampler)
    @inbounds prediction = sampler[].predictions[idx]
    @inbounds target = rand(rng, floatsampler) < prediction
    return (prediction, target)
end
function Base.rand(rng::Random.AbstractRNG, sampler::Random.SamplerSimple{<:Consistent})
    idxsampler, samplers = sampler.data
    idx = rand(rng, idxsampler)
    @inbounds prediction = sampler[].predictions[idx]
    @inbounds target = rand(rng, samplers[idx])
    return (prediction, target)
end

# in-place sampling
function Random.rand!(
    rng::Random.AbstractRNG,
    out::AbstractArray{<:Tuple{<:Any,<:Any}},
    sampler::Random.SamplerSimple{<:Consistent{<:Real}},
)
    predictions = sampler[].predictions
    idxsampler, floatsampler = sampler.data
    for i in eachindex(out)
        idx = rand(rng, idxsampler)
        @inbounds prediction = predictions[idx]
        @inbounds target = rand(rng, floatsampler) < prediction
        @inbounds out[i] = (prediction, target)
    end
    return out
end
function Random.rand!(
    rng::Random.AbstractRNG,
    out::AbstractArray{<:Tuple{<:Any,<:Any}},
    sampler::Random.SamplerSimple{<:Consistent},
)
    predictions = sampler[].predictions
    idxsampler, samplers = sampler.data
    for i in eachindex(out)
        idx = rand(rng, idxsampler)
        @inbounds prediction = predictions[idx]
        @inbounds target = rand(rng, samplers[idx])
        @inbounds out[i] = (prediction, target)
    end
    return out
end


@doc raw"""
    function replay(f, t :: T) where T <: Trace

Given a stochastic function `f` and a trace `t`, makes `sample` calls behave as though
they had sampled the values in `t` at the corresponding addresses.

Returns a tuple `(t_new, g)`, where `t_new` is a trace and `g` is a function.
The function signature of `g` is the same as that of `f` with the first argument removed; that
is, if `f(t :: Trace, params...)`, then `g(params...)`. Computation is delayed, so each of the
latent nodes in `t_new` has `interpretation = REPLAYED`.
Calling `g(params...)` executes the computation and each latent node in `t_new` reverts to
its original interpretation.
"""
function replay(f, t :: T) where T <: Trace
    t_new = deepcopy(t)
    interpret_latent!(t_new, REPLAYED)
    g(params...) = f(t_new, params...)
    (t_new, g)
end

@doc raw"""
    function replace(f, r :: Dict)

Given a mapping `r` from addresses to distribution-like (currently `Distributions`
or `Array{Any, 1}`s), replaces the current distributions at that set of addresses
with this set of distributions. Returns a function `g` that has return signature
`(t :: Trace, rval)` where `rval` is a return value of `f`.
"""
function replace(f, r :: Dict)
    function g(t :: Trace, params...)
        rval = f(t, params...)
        t = replace(t, r)
        (t, rval)
    end
    g
end

@doc raw"""
    function replace(t :: Trace, r :: Dict)

Given a mapping `r` from addresses to distribution-like (currently `Distributions`
or `Array{Any, 1}`s), replaces the current distributions at that set of addresses
with this set of distributions. Returns the modified trace.
"""
function replace(t :: T, r :: Dict) where T <: Trace
    for (a, d) in r
        old_n = deepcopy(t[a])
        t[a] = node(old_n.value, old_n.address, d, old_n.observed, old_n.interpretation)
    end
    t
end

@doc raw"""
    function rewrite(f, t :: T, r :: Dict) where T <: Trace

Rewrites the history of the trace to make it appear as if the values in the trace
were sampled at the addresses in the keys of `r` from the corresponding distributions in the values of `r`.
Returns a function with call signature `g(params...)` that returns `(t :: Trace, rval)`, where `rval`
is the return type of `f`.
"""
function rewrite(f, t :: T, r :: Dict) where T <: Trace
    function g(params...)
        new_t, h = replay(f, t)
        rval = h(params...)
        new_t = replace(new_t, r)
        (new_t, rval)
    end
    g
end

@doc raw"""
    function block(f, t :: T, addresses) where T <: Trace

Given a stochastic function `f`, a trace `t`, and an iterable of addresses, converts traced
randomness into untraced randomness.

Returns a tuple `(t_new, g)`, where `t_new` is a trace and `g` is a function.
The function signature of `g` is the same as that of `f` with the first argument removed; that
is, if `f(t :: Trace, params...)`, then `g(params...)`. Computation is delayed, so each of
the latent nodes in `t_new` has `interpretation = BLOCKED`. Calling `g(params...)` executes the
computation and each latent node in `t_new` with an address in `addresses` is removed.
"""
function block(f, t :: T, addresses) where T <: Trace
    t_new = deepcopy(t)
    for a in addresses
        t_new[a].interpretation = BLOCKED
    end
    g(params...) = f(t_new, params...)
    (t_new, g)
end

function block(f, addresses)
    function g(t::Trace, params...)
        r = f(t, params...)
        for a in addresses
            delete!(t.trace, a)
        end
        logprob!(t)
        r
    end
    g
end

@doc raw"""
    block(f :: F, t :: Trace) where F <: Function

Converts all traced randomness into untraced randomness.
"""
block(f, t :: Trace) = block(f, t, keys(t))

@doc raw"""
    function update(f, r :: SamplingResults{I}) where I <: InferenceType

Given a stochastic function `f` and a `SamplingResults` `r`, update the prior
predictive to the posterior predictive by jointly replacing all latent sample sites with the joint empirical
posterior. Returns a stochastic function `g` with the same call signature as `f`. This function will modify in
place the trace passed into it as the first argument.
"""
function update(f, r :: SamplingResults{I}) where I <: InferenceType
    function g(t :: Trace, params...)
        t_sampled = StatsBase.sample(r.traces)  # maintains correct dims throughout
        kts = keys(t_sampled)
        marginals = Dict()
        for a in kts
            sampled_not_obs = t_sampled[a].observed == false
            not_in_t = !(a in keys(t))
            in_t_and_not_obs = !not_in_t && t[a].observed == false
            if (sampled_not_obs && not_in_t) || in_t_and_not_obs
                sample(t, a, r)
                marginals[a] = r[a]
            end
        end
        h = rewrite(f, t, marginals)
        r = h(params...)
    end
    g
end

@doc raw"""
    function condition(f, evidence :: Dict)

Condition a trace modified by `f` on `evidence`, which maps addresses to observed
evidence associated with that address. Returns a function with call signature
identical to that of `f` and return signature `(t :: Trace, rtype)` where `rtype`
is the return type of `f`.
"""
function condition(f, evidence :: Dict; bare=false)
    function g(t :: Trace, params...)
        for (a, e) in evidence
            t[a] = node(e, a, [CONDITIONED], true, CONDITIONED)
        end
        r = f(t, params...)
        !bare ? (t, r) : r
    end
    g
end

export block, replay, update, replace, rewrite, condition


function write_stl(dataset::AbstractStaticVTKData, filepath_no_ext::String)
    polydata = PyVTK(trianulate(dataset))
    return write_stl(polydata, filepath_no_ext)
end

function write_stl(polydata::PyCall.PyObject, filepath_no_ext::String)
    writer = vtk.vtkSTLWriter()
    writer[:SetFileName](string(filepath_no_ext, ".stl"))
    writer[:SetInputData](polydata)
    writer[:Update]()
    return string(filepath_no_ext, ".stl")
end

function read_stl(filepath::String)
    if splitext(filepath)[2] == ".stl"
        reader = vtk.vtkSTLReader()
        reader[:SetFileName](filepath)
        reader[:Update]()

        polydata = reader[:GetOutput]()
        return extract_simple_block(polydata)
    else
        throw("The file input is not an stl file.")
    end
end