using ImagineInterface

import Unitful: μm, s

############LOAD A RIG-SPECIFIC COMMAND TEMPLATE#################
sample_rate = 50000s^-1 #analog output samples per second
				#Currently all output channels must use the same sample rate
				#Currently this also sets the analog input rate in Imagine
				#High-framerate recordings should use a high sample rate.
@show ImagineInterface.RIGS #show supported rigs
rig = "ocpi-2"
ocpi2 = rigtemplate(rig; sample_rate = sample_rate)
@show ocpi2 #This is a vector where each element is an empty ImagineSignal
#We can edit commands individually by extracting them from this array.
#Note however that after you finish editing and are ready to write a command file, all commands must have an equal number of samples (or remain empty)

#Let's look at a positioner trace
positioners = getpositioners(ocpi2)
pos = positioners[1]
#"pos" is an ImagineSignal.  We can append sequences of samples to the (currently empty) signal to use as a command for Imagine.
#We can append vectors of raw samples, voltage units (since this is an analog output device), or world units.
#If you're not sure what type os samples you can append, or what ranges of values are allowed, you can check with this command
@show intervals(pos) #this signal has 3 (equivalent) representations, one uses raw sample units, one uses Voltage units, and the other uses Length units
#While you can append vectors using any of the 3 representations, in practice it's most intuitive to use "world" units (μm for the positioner)  when interacting with ImagineSignals

#You can also query the intervals individually
@show interval_raw(pos)
@show interval_volts(pos)
@show interval_world(pos)

#We can create a sample vector describing motion in microns like this
#If the appended vector does not lie within its corresponding interval (as queried by the `intervals` function)
#then you will get an error
sweep_up = [0.0:0.1:800.0...] * Unitful.μm #sweep once from the min to max
#If you want to add these samples to the positioner command, you can do it like this:
append!(pos, "sweep_up", sweep_up)
#When appending a sequence, the sequence is automatically compressed to save storage space.
#Repeated sequences cost almost zero additional storage space.
#In order to retrieve the sequence again, you must decompress it.
#You can decompress using its key, in this case "sweep_up":
sweep_up2 = get_samples(pos, "sweep_up") #now sweep_up == sweep_up2
#Or you can use sample indices
some_samps = get_samples(pos, 50, 100)
#OR you can use time
more_samps = get_samples(pos, 0.1s, 0.13s)

#You can also remove the last set of samples appended to the list like this:
pop!(pos)
#This package provides easier-to-use functions to get commonly used sequences of samples.
#Currently you can generate an entire stack of positioner, camera, and laser pulse signals in the fashion shown below

############STACK PARAMETERS#################
pmin = 0.0*μm #Piezo start position in microns
pmax = 200.0*μm #Piezo stop position in microns
stack_img_time = 1.0s #Time to complete the imaging sweep with the piezo (remember we also need to reset it to its starting position)
reset_time = 0.5s #Time to reset piezo to starting position.  This time plus "stack_img_time" determines how long it takes to complete an entire stack and be ready to start a new stack
z_spacing = 3.1μm #The space between slices in the z-stack.
z_pad = 5.0μm #Set this greater than 0 if you want to ignore the edges of the positioner sweep (only take slices in a central region)
			#This is helpful for high-speed acquisitions where the piezo positioner may oscillate at the edges of sweeps
#Now we need to set the exposure time.  With constraints imposed by the above pmin, pmax, stack_img_time, and z_spacing parameters there is actually a maximum exposure time.
#If you exceed this time then you will get en error.  The helper function below will tell you what is the maximum time possible with these parameters
#(We add/subtract z_pad because it shrinks the effective range of imaging sweep)
#@show mx_exp = max_exp_time((pmax-z_pad) - (pmin+z_pad), z_spacing, stack_img_time) #TODO: implement this function

#There is also a minimum exposure time that the camera can achieve.  This depends on the camera model.  For now all OCPI microscopes have similar cameras, but the helper
#function still requires you to specify the target rig in case cameras change in the future
#Note also that the camera's maximum frame-per-second rate depends on the ROI size.  For more details see the PCO.Edge camera manuals.
#If you're not sure of the maximal ROI size on an OCPI rig, retrieve it with this function:
@show hmax, vmax = chip_size(rig)
#You can check the maximal framerate when using a smaller rectangular section of the chip
#Note that current PCO.Edge camera framerates only depend on the vertical size of the ROI
h = 1000 #doesn't matter with current cameras
v = 1000
@show mx_f = max_framerate(rig, h,v)
mn_exp = 1/mx_f #This is the minimum possible exposure time
exp_time = 0.011s  #Exposure time of the camera. Make sure this is greater than mn_exp and less than mx_exp above
flash_frac = 0.1 #fraction of time to keep laser on during exposure.  If you set this greater than 1 then the laser will stay on constantly during the imaging sweep

############STACK GENERATION EXAMPLE#################
#The below function generates 3 sequences of samples, stored in a dictionary
d = gen_unidirectional_stack(pmin, pmax, z_spacing, stack_img_time, reset_time, exp_time, sample_rate, flash_frac; z_pad = z_pad)
#d["positioner"] is a vector of piezo samples (in length units)
#d["camera"]  is vector of true-or-false values encoding a camera exposure trigger sequence
#d["laser"]  is vector of true-or-false values encoding laser flash sequence (these will be center-aligned with the exposure pulses)
@show nframes = d["nframes"]  # is an integer keeping track of the number of frames in this sequence, useful when creating the command file

#Similarly we can also generate a bidirectional imaging cycle (which actually includes two stacks)
#Note the command is the same except "reset_time" is missing because imaging occurs during what was the reset period in the previous example
#...and the forward and backward imaging sweeps are of equal duration
d2 = gen_bidirectional_stack(pmin, pmax, z_spacing, stack_img_time, exp_time, sample_rate, flash_frac; z_pad = z_pad)
@show d2["nframes"] #note that there are twice as many frames as in the unidirectional example because one cycle of the positioner includes two stacks

#Let's append our newly-created sample vectors to their respective commands in the template
append!(pos, "uni_stack_pos", d["positioner"])
lasers = getlasers(ocpi2)
las1 = lasers[1] #You could also append to any of the other lasers
append!(las1, "uni_stack_las1", d["laser"])
cams = getcameras(ocpi2)
cam1 = cams[1] #You could also append to the other camera
append!(cam1, "uni_stack_cam1", d["camera"])

using Plots, ImaginePlots #You will have to install ImaginePlots with Pkg.clone before it will work

#There is a special plot method that lets you visualize all commands.  Currently it groups commands in multiple plot windows base on hardware
plot(ocpi2) #note that this currently won't display when you run the script via include("workflow.jl") Type or paste the command in the terminal to see the plot

#If you are satisfied with the waveform created and would like to duplicate it many times for a multi-stack acquisition with Imagine...
replicate!(pos, 9) #appends the current contents of pos 9 times to the end of pos
#In this case the replicate! command is equivalent to calling append!(pos, "uni_stack_pos") 99 times
replicate!(las1, 9) #appends the current contents of pos 9 times to the end of pos
replicate!(cam1, 9) #appends the current contents of pos 9 times to the end of pos

#If you plot again you'll see a total of 10 stacks
#However plotting is pretty slow at the moment, so it's recommended to visualize one stack at a time
#plot(ocpi2)

#When you are ready, export the set of commands to a .json file readable by Imagine

#All commands passed to this function must have an equal number of samples OR no samples
#In addition to the commands themselves, currently you must provide additional metadata so that the .imagine header gets written correctly during a recording
nstacks = 10
write_commands("test.json", ocpi2, nstacks, nframes, exp_time; isbidi = false)


mutable struct TimeManagement
    move_overhead::Int # UCI can set a move overhead value.
    isinfinite::Bool # Has UCI told us we have infinite thinking time?
    movetime::Int # Has UCI specified a move time?
    depth::Int # Has UCI specified a depth?

    start_time::Float64 # Grab the search start time.
    clock_time::Int # Grab the clock time of the turn to move.
    inc_time::Int # Get the time increment
    movestogo::Int # Get the number of moves to go.

    ideal_time::Int # Ideal time we should aim for
    max_time::Int # Maximum time available
end

TimeManagement() = TimeManagement(100, false, 0, 0, 0, 0, 0, 0, 0, 0)

elapsedtime(timeman::TimeManagement) = (time() - timeman.start_time) * 1000

# Not all time management features are supported yet.
function initTimeManagement!(timeman::TimeManagement)
    if timeman.movestogo == -1
        timeman.ideal_time = fld(timeman.clock_time + 20*timeman.inc_time, 45)
        timeman.max_time = 10 * fld(timeman.clock_time + 20*timeman.inc_time, 45)
    else
        timeman.ideal_time = fld(timeman.clock_time, timeman.movestogo + 5) + timeman.inc_time
        timeman.max_time = 10 * fld(timeman.clock_time, timeman.movestogo + 10) + timeman.inc_time
    end
    timeman.ideal_time = min(timeman.ideal_time, timeman.clock_time - timeman.move_overhead)
    timeman.max_time = min(timeman.max_time, timeman.clock_time - timeman.move_overhead)
    if timeman.movetime > 0
        timeman.ideal_time = timeman.movetime
        timeman.max_time = timeman.movetime
    end
end

function istermination(timeman::TimeManagement)
    timeman.depth == 0 && elapsedtime(timeman) > timeman.ideal_time
end


#=
# FV solver
=#
module FVSolver

export solve, nonlinear_solve

import ..Femshop: JULIA, CPP, MATLAB, DENDRO, HOMG, CUSTOM_GEN_TARGET,
                SQUARE, IRREGULAR, UNIFORM_GRID, TREE, UNSTRUCTURED,
                CG, DG, HDG, FV,
                NODAL, MODAL, CELL, LEGENDRE, UNIFORM, GAUSS, LOBATTO,
                NONLINEAR_NEWTON, NONLINEAR_SOMETHING,
                EULER_EXPLICIT, EULER_IMPLICIT, CRANK_NICHOLSON, RK4, LSRK4, ABM4,
                DEFAULT_SOLVER, PETSC,
                VTK, RAW_OUTPUT, CUSTOM_OUTPUT,
                DIRICHLET, NEUMANN, ROBIN, NO_BC, FLUX,
                MSH_V2, MSH_V4,
                SCALAR, VECTOR, TENSOR, SYM_TENSOR, VAR_ARRAY,
                LHS, RHS,
                LINEMESH, QUADMESH, HEXMESH
import ..Femshop: log_entry, printerr
import ..Femshop: config, prob, variables, mesh_data, grid_data, refel, time_stepper, elemental_order, genfunctions, indexers
import ..Femshop: Variable, Coefficient, GenFunction
import ..Femshop: GeometricFactors, geo_factors, geometric_factors, geometric_factors_face, build_deriv_matrix
# specific to FV
import ..Femshop: FVInfo, fv_info, FV_cell_to_node, FV_node_to_cell
import ..Femshop: ParentMaps, parent_maps, fv_grid, fv_refel, fv_geo_factors, build_local_patch

using LinearAlgebra, SparseArrays

include("fv_boundary.jl");

function linear_solve(var, source_lhs, source_rhs, flux_lhs, flux_rhs, stepper=nothing, assemble_func=nothing)
    # If more than one variable
    if typeof(var) <: Array
        # multiple variables being solved for simultaneously
        dofs_per_node = 0;
        dofs_per_loop = 0;
        for vi=1:length(var)
            dofs_per_loop += length(var[vi].symvar);
            dofs_per_node += var[vi].total_components;
        end
    else
        # one variable
        dofs_per_loop = length(var.symvar);
        dofs_per_node = var.total_components;
    end
    if fv_grid === nothing
        nel = size(grid_data.loc2glb, 2);
        nfaces = size(grid_data.face2element, 2);
    else
        nel = size(fv_grid.loc2glb, 2);
        nfaces = size(fv_grid.face2element, 2);
    end

    Nn = dofs_per_node * nel;
    Nf = dofs_per_node * nfaces

    # Allocate arrays that will be used by assemble
    # These vectors will hold the integrated values(one per cell).
    # They will later be combined.
    sourcevec = zeros(Nn);
    fluxvec = zeros(Nn);
    facefluxvec = zeros(Nf);
    face_done = zeros(Int, nfaces); # Increment when the corresponding flux value is computed.
    allocated_vecs = [sourcevec, fluxvec, facefluxvec, face_done];

    if prob.time_dependent && !(stepper === nothing)
        log_entry("Beginning "*string(stepper.Nsteps)*" time steps.");
        t = 0;
        sol = get_var_vals(var);

        # allocate storage used by steppers
        if stepper.type == LSRK4
            tmppi = zeros(size(sol));
            tmpki = zeros(size(sol));
        elseif stepper.type == RK4
            tmpki = zeros(length(sol), stepper.stages);
        end

        start_t = Base.Libc.time();
        last2update = 0;
        last10update = 0;
        print("Time stepping progress(%): 0");
        for i=1:stepper.Nsteps
            if stepper.stages > 1
                # LSRK4 is a special case, low storage
                if stepper.type == LSRK4
                    # Low storage RK4:
                    # p0 = u
                    #   ki = ai*k(i-1) + dt*f(p(i-1), t+ci*dt)
                    #   pi = p(i-1) + bi*ki
                    # u = p5

                    tmppi = get_var_vals(var, tmppi);
                    tmpki = zeros(size(sol));
                    for rki=1:stepper.stages
                        rktime = t + stepper.c[rki]*stepper.dt;
                        # p(i-1) is currently in u

                        sol = assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, stime, stepper.dt, assemble_loops=assemble_func);

                        if rki == 1 # because a1 == 0
                            tmpki = stepper.dt .* sol;
                        else
                            tmpki = stepper.a[rki].*tmpki + stepper.dt.*sol;
                        end
                        tmppi = tmppi + stepper.b[rki].*tmpki

                        place_sol_in_vars(var, tmppi, stepper);
                    end

                else
                    # Explicit multi-stage methods:
                    # x = x + dt*sum(bi*ki)
                    # ki = rhs(t+ci*dt, x+dt*sum(aij*kj)))   j < i

                    # will hold the final result
                    sol = get_var_vals(var, sol);
                    # will be placed in var.values for each stage
                    tmpvals = sol;
                    for stage=1:stepper.stages
                        stime = t + stepper.c[stage]*stepper.dt;

                        tmpki[:,stage] = assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, stime, stepper.dt, assemble_loops=assemble_func);

                        tmpvals = sol;
                        for j=1:(stage-1)
                            if stepper.a[stage, j] > 0
                                tmpvals += stepper.dt * stepper.a[stage, j] .* tmpki[:,j];
                            end
                        end

                        place_sol_in_vars(var, tmpvals, stepper);
                    end
                    for stage=1:stepper.stages
                        sol += stepper.dt * stepper.b[stage] .* tmpki[:, stage];
                    end
                    place_sol_in_vars(var, sol, stepper);
                end

            elseif stepper.type == EULER_EXPLICIT
                sol = sol .+ stepper.dt .* assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, t, stepper.dt, assemble_loops=assemble_func);

                place_sol_in_vars(var, sol, stepper);

            else
                printerr("Only explicit time steppers for FV are ready. TODO")
                return sol;
            end

            # ########## uncomment to return after one time step
            # return sol
            # ############

            t += stepper.dt;

            progressPercent = Int(floor(i*100.0/stepper.Nsteps));
            if progressPercent - last2update >= 2
                last2update = progressPercent;
                if progressPercent - last10update >= 10
                    print(progressPercent);
                    last10update = progressPercent;
                else
                    print(".");
                end
            end
        end
        println("");
        end_t = Base.Libc.time();

        log_entry("Stepping took "*string(end_t-start_t)*" seconds.");
        #log_entry("Stepping took "*string(end_t-start_t)*" seconds. ("*string(assemble_t)*" for assembly, "*string(linsolve_t)*" for linear solve)");
        #display(sol);
		# outfile = "linear_sol.txt"
		# open(outfile, "w") do f
  		# 	for ii in sol
    	# 		println(f, ii)
  		# 	end
		# end # the file f is automatically closed after this block finishes
        return sol;

    else
        # Does it make any sense to do this for time-independent problems?
        return sol;
    end
end

# macro loops(indices, ranges, content)
#     n = length(indices.args);
#     if n == 1
#         this_ind = indices.args[1];
#         this_range = ranges.args[1];
#         return esc(quote
#             for $this_ind in $this_range
#                 $content
#             end
#         end)
#     else
#         this_ind = indices.args[1];
#         this_range = ranges.args[1];
#         sub_ind = copy(indices);
#         sub_ind.args = sub_ind.args[2:end];
#         sub_range = copy(ranges);
#         sub_range.args = sub_range.args[2:end];
#         return esc(:(@loops([$this_ind], [$this_range], @loops($sub_ind, $sub_range, $content))));
#     end
# end

#

function assemble(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node=1, dofs_per_loop=1, t=0, dt=0; assemble_loops=nothing)
    # If an assembly loop function was provided, use it
    if !(assemble_loops === nothing)
        return assemble_loops.func(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, t, dt);
    end
    # If parent maps were created for high order flux, use a slightly different function
    if !(parent_maps === nothing)
        return assemble_using_parent_child(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node, dofs_per_loop, t, dt);
    end

    nel = size(grid_data.loc2glb, 2);

    # Label things that were allocated externally
    sourcevec = allocated_vecs[1];
    fluxvec = allocated_vecs[2];
    facefluxvec = allocated_vecs[3];
    face_done = allocated_vecs[4];

    face_done .= 0;

    # Elemental loop
    for ei=1:nel
        eid = elemental_order[ei]; # The index of this element
        # Zero the result vectors for this element
        sourcevec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] .= 0;
        fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] .= 0;

        ##### Source integrated over the cell #####
        # Compute RHS volume integral
        if !(source_rhs === nothing)
            #sourceargs = prepare_args(var, eid, 0, RHS, "volume", t, dt); # (var, e, nodex, loc2glb, refel, detj, J, t, dt)
            sourceargs = (var, eid, 0, grid_data, geo_factors, fv_info, refel, t, dt);
            source = source_rhs.func(sourceargs) ./ geo_factors.volume[eid];
            # Add to global source vector
            sourcevec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] = source;
        end

        ##### Flux integrated over the faces #####
        # Loop over this element's faces.
        for i=1:refel.Nfaces
            fid = grid_data.element2face[i, eid];

            if !(flux_rhs === nothing)
                if face_done[fid] == 0
                    face_done[fid] = 1; # Possible race condition, but in the worst case it will be computed twice.

                    #fluxargs = prepare_args(var, eid, fid, RHS, "surface", t, dt); #(var, (e, neighbor), refel, vol_loc2glb, nodex, cellx, frefelind, facex, face2glb, normal, fdetj, face_area, (J, vol_J_neighbor), t, dt);
                    fluxargs = (var, eid, fid, grid_data, geo_factors, fv_info, refel, t, dt);
                    flux = flux_rhs.func(fluxargs) .* geo_factors.area[fid];
                    # Add to global flux vector for faces
                    facefluxvec[((fid-1)*dofs_per_node + 1):(fid*dofs_per_node)] = flux;
                    # Combine all flux for this element
                    fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] += flux ./ geo_factors.volume[eid];

                else
                    # This flux has either been computed or is being computed by another thread.
                    # The state will need to be known before paralellizing, but for now assume it's complete.
                    fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] -= facefluxvec[((fid-1)*dofs_per_node + 1):(fid*dofs_per_node)] ./ geo_factors.volume[eid];
                end
            end

            # Boundary conditions are applied to flux
            fbid = grid_data.facebid[fid]; # BID of this face
            if fbid > 0
                facex = grid_data.allnodes[:, grid_data.face2glb[:,1,fid]];  # face node coordinates

                if typeof(var) <: Array
                    dofind = 0;
                    for vi=1:length(var)
                        for compo=1:length(var[vi].symvar)
                            dofind = dofind + 1;
                            if prob.bc_type[var[vi].index, fbid] == NO_BC
                                # do nothing
                            elseif prob.bc_type[var[vi].index, fbid] == FLUX
                                # compute the value and add it to the flux directly
                                # Qvec = (refel.surf_wg[grid_data.faceRefelInd[1,fid]] .* geo_factors.face_detJ[fid])' * (refel.surf_Q[grid_data.faceRefelInd[1,fid]])[:, refel.face2local[grid_data.faceRefelInd[1,fid]]]
                                # Qvec = Qvec ./ geo_factors.area[fid];
                                # bflux = FV_flux_bc_rhs_only(prob.bc_func[var[vi].index, fbid][compo], facex, Qvec, t, dofind, dofs_per_node) .* geo_factors.area[fid];
                                bflux = FV_flux_bc_rhs_only_simple(prob.bc_func[var[vi].index, fbid][compo], facex, t) .* geo_factors.area[fid];

                                fluxvec[(eid-1)*dofs_per_node + dofind] += (bflux - facefluxvec[(fid-1)*dofs_per_node + dofind]) ./ geo_factors.volume[eid];
                                facefluxvec[(fid-1)*dofs_per_node + dofind] = bflux;
                            else
                                printerr("Unsupported boundary condition type: "*prob.bc_type[var[vi].index, fbid]);
                            end
                        end
                    end
                else
                    for d=1:dofs_per_node
                        dofind = d;
                        if prob.bc_type[var.index, fbid] == NO_BC
                            # do nothing
                        elseif prob.bc_type[var.index, fbid] == FLUX
                            # compute the value and add it to the flux directly
                            # Qvec = (refel.surf_wg[grid_data.faceRefelInd[1,fid]] .* geo_factors.face_detJ[fid])' * (refel.surf_Q[grid_data.faceRefelInd[1,fid]])[:, refel.face2local[grid_data.faceRefelInd[1,fid]]]
                            # Qvec = Qvec ./ geo_factors.area[fid];
                            # bflux = FV_flux_bc_rhs_only(prob.bc_func[var.index, fbid][d], facex, Qvec, t, dofind, dofs_per_node) .* geo_factors.area[fid];
                            bflux = FV_flux_bc_rhs_only_simple(prob.bc_func[var.index, fbid][d], facex, t) .* geo_factors.area[fid];

                            fluxvec[(eid-1)*dofs_per_node + dofind] += (bflux - facefluxvec[(fid-1)*dofs_per_node + dofind]) ./ geo_factors.volume[eid];
                            facefluxvec[(fid-1)*dofs_per_node + dofind] = bflux;
                        else
                            printerr("Unsupported boundary condition type: "*prob.bc_type[var.index, fbid]);
                        end
                    end
                end
            end# BCs

        end# face loop
    end# element loop

    return sourcevec + fluxvec;
end

function assemble_using_parent_child(var, source_lhs, source_rhs, flux_lhs, flux_rhs, allocated_vecs, dofs_per_node=1, dofs_per_loop=1, t=0, dt=0)
    debug = false;

    # Label things that were allocated externally
    sourcevec = allocated_vecs[1];
    fluxvec = allocated_vecs[2];
    facefluxvec = allocated_vecs[3];
    face_done = allocated_vecs[4];

    face_done .= 0;
    fluxvec .= 0; # need to do this here

    # Loop over parent elements
    # Since flux depends only on this and nearest neighbor parent cells, it is safe
    # to partition between parents using only the nearest parents as ghosts.
    Nparents = size(parent_maps.parent2child,2);
    nchildren = size(parent_maps.parent2child,1); # (assumes one element type)
    npfaces = size(grid_data.element2face,1); # Outer faces of parent. Parents are elements in grid_data (assumes one element type)
    ncfaces = size(parent_maps.parent2face,1); # All faces of children in a parent (assumes one element type)
    dim = size(grid_data.allnodes,1);

    for parentid=1:Nparents
        # Get the cells in the local patch
        patch_cells = parent_maps.patches[:, parentid];

        # Split the work into two parts:
        # - source terms that don't depend on neighbors and can be done independently (loop over elements)
        # - flux terms that depend on neighbors (loop over faces)

        # The loop over elements for source terms.
        for childid = 1:nchildren
            eid = parent_maps.parent2child[childid, parentid];
            # Zero the result vectors for this element
            sourcevec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] .= 0;

            # Use source_rhs for RHS volume integral
            if !(source_rhs === nothing)
                sourceargs = (var, eid, 0, fv_grid, fv_geo_factors, fv_info, fv_refel, t, dt);
                source = source_rhs.func(sourceargs; patch_cells=patch_cells) ./ fv_geo_factors.volume[eid];
                # Add to global source vector
                sourcevec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] = source;
            end
        end

        # Loop over this parent's faces for flux terms.
        for pfaceid = 1:ncfaces
            fid = parent_maps.parent2face[pfaceid, parentid];
            els = fv_grid.face2element[:,fid];
            # # The eid passed to the flux func. will be one inside the parent.
            # if els[1] in parent_maps.parent2child[:,parentid]
            #     eid = els[1];
            #     neighbor = els[2];
            # else
            #     eid = els[2];
            #     neighbor = els[1];
            # end
            eid = els[1];
            neighbor = els[2];

            if debug && parentid < 3
                println("parent: "*string(parentid)*", fid: "*string(fid)*"("*string(pfaceid)*"), eid/neigh: "*string(eid)*"/"*string(neighbor));
                println("patch: "*string(patch_cells));
            end

            if !(flux_rhs === nothing)
                if face_done[fid] == 0
                    face_done[fid] = 1; # Possible race condition, but in the worst case it will be computed twice.\

                    (left_cells, right_cells) = get_left_right_cells(patch_cells, pfaceid, parent_maps, dim);

                    # leftness and rightness depends on the normal vector direction
                    # normal points left to right.
                    leftcell = fv_grid.face2element[1,fid];
                    if leftcell == right_cells[1]
                        #oops, swap them
                        tmp = left_cells;
                        left_cells = right_cells;
                        right_cells = tmp;
                    end

                    # remove zero values
                    left_cells = remove_zero_cells(left_cells);
                    right_cells = remove_zero_cells(right_cells);

                    #fluxargs = prepare_args(var, eid, fid, RHS, "surface", t, dt); #(var, (e, neighbor), refel, vol_loc2glb, nodex, cellx, frefelind, facex, face2glb, normal, fdetj, face_area, (J, vol_J_neighbor), t, dt);
                    fluxargs = (var, eid, fid, fv_grid, fv_geo_factors, fv_info, fv_refel, t, dt);
                    flux = flux_rhs.func(fluxargs; left_cells=left_cells, right_cells=right_cells) .* fv_geo_factors.area[fid];

                    # Add to global flux vector for faces
                    facefluxvec[((fid-1)*dofs_per_node + 1):(fid*dofs_per_node)] = flux;
                    # Contribute to flux for elements on both sides
                    fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] += flux ./ fv_geo_factors.volume[eid];
                    if neighbor > 0 # not a boundary, so the other side exists
                        fluxvec[((neighbor-1)*dofs_per_node + 1):(neighbor*dofs_per_node)] -= flux ./ fv_geo_factors.volume[neighbor];
                    end

                    if debug && parentid < 3
                        println("cells left: "*string(left_cells)*", right: "*string(right_cells));
                        println("flux: "*string(flux));
                        println("contribute "*string(flux ./ fv_geo_factors.volume[eid])*" to "*string(((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)));
                        if neighbor > 0
                            println("contribute "*string(-flux ./ fv_geo_factors.volume[neighbor])*" to "*string(((neighbor-1)*dofs_per_node + 1):(neighbor*dofs_per_node)));
                        end
                        println("fluxvec "*string(fluxvec[1:8]));
                        println("");
                    end

                else
                    # This flux has either been computed or is being computed by another thread.
                    # The state will need to be known before paralellizing, but for now assume it's complete.
                    # fluxvec[((eid-1)*dofs_per_node + 1):(eid*dofs_per_node)] .-= facefluxvec[((fid-1)*dofs_per_node + 1):(fid*dofs_per_node)] ./ fv_geo_factors.volume[eid];
                end
            end

            # Boundary conditions are applied to flux
            fbid = fv_grid.facebid[fid]; # BID of this face
            if fbid > 0
                if debug
                    println("applying bdry cond"*string(fbid)*" to "*string(fid));
                end
                facex = fv_grid.allnodes[:, fv_grid.face2glb[:,1,fid]];  # face node coordinates

                if typeof(var) <: Array
                    dofind = 0;
                    for vi=1:length(var)
                        for compo=1:length(var[vi].symvar)
                            dofind = dofind + 1;
                            if prob.bc_type[var[vi].index, fbid] == NO_BC
                                # do nothing
                            elseif prob.bc_type[var[vi].index, fbid] == FLUX
                                # compute the value and add it to the flux directly
                                Qvec = (fv_refel.surf_wg[fv_grid.faceRefelInd[1,fid]] .* fv_geo_factors.face_detJ[fid])' * (fv_refel.surf_Q[fv_grid.faceRefelInd[1,fid]])[:, fv_refel.face2local[fv_grid.faceRefelInd[1,fid]]]
                                Qvec = Qvec ./ fv_geo_factors.area[fid];
                                bflux = FV_flux_bc_rhs_only(prob.bc_func[var[vi].index, fbid][compo], facex, Qvec, t, dofind, dofs_per_node) .* fv_geo_factors.area[fid];

                                fluxvec[(eid-1)*dofs_per_node + dofind] += (bflux - facefluxvec[(fid-1)*dofs_per_node + dofind]) ./ fv_geo_factors.volume[eid];
                                facefluxvec[(fid-1)*dofs_per_node + dofind] = bflux;
                            else
                                printerr("Unsupported boundary condition type: "*prob.bc_type[var[vi].index, fbid]);
                            end
                        end
                    end
                else
                    for d=1:dofs_per_node
                        dofind = d;
                        if prob.bc_type[var.index, fbid] == NO_BC
                            # do nothing
                        elseif prob.bc_type[var.index, fbid] == FLUX
                            # compute the value and add it to the flux directly
                            Qvec = (fv_refel.surf_wg[fv_grid.faceRefelInd[1,fid]] .* fv_geo_factors.face_detJ[fid])' * (fv_refel.surf_Q[fv_grid.faceRefelInd[1,fid]])[:, fv_refel.face2local[fv_grid.faceRefelInd[1,fid]]]
                            Qvec = Qvec ./ fv_geo_factors.area[fid];
                            bflux = FV_flux_bc_rhs_only(prob.bc_func[var.index, fbid][d], facex, Qvec, t, dofind, dofs_per_node) .* fv_geo_factors.area[fid];

                            fluxvec[(eid-1)*dofs_per_node + dofind] += (bflux - facefluxvec[(fid-1)*dofs_per_node + dofind]) ./ fv_geo_factors.volume[eid];
                            facefluxvec[(fid-1)*dofs_per_node + dofind] = bflux;
                        else
                            printerr("Unsupported boundary condition type: "*prob.bc_type[var.index, fbid]);
                        end
                    end
                end
            end# BCs
        end# face loop
    end# element loop
    if debug
        println("fluxvec: "*string(fluxvec[1:10]));
    end
    return sourcevec + fluxvec;
end

# Returns a copy of a with zeros removed.
function remove_zero_cells(a)
    b = similar(a);
    nnz = 0;
    for ai in a
        if !(ai == 0)
            nnz += 1;
            b[nnz] = ai;
        end
    end
    b = b[1:nnz];

    return b;
end

# Returns arrays of left and right cells
function get_left_right_cells(patch, face, maps, dim)
    # Provide all of the cells that can be used for this face in a particular order.
    # Make two arrays, one for left one for right, starting from the nearest neighbor.
    left_cells = [];
    right_cells = [];
    if dim == 1
        left_cell_table_1d = [
            [ # children=1
                [2],
                [1, 2]
            ],
            [ # children=2
                [4, 3],
                [1, 4, 3],
                [2, 1, 4, 3]
            ],
            [ # children=3
                [6, 5, 4],
                [1, 6, 5, 4],
                [2, 1, 6, 5, 4],
                [3, 2, 1, 6, 5, 4]
            ],
            [ # children=4
                [8, 7, 6, 5],
                [1, 8, 7, 6, 5],
                [2, 1, 8, 7, 6, 5],
                [3, 2, 1, 8, 7, 6, 5],
                [4, 3, 2, 1, 8, 7, 6, 5],
            ]
        ]
        right_cell_table_1d = [
            [ # children=1
                [1, 3],
                [3]
            ],
            [ # children=2
                [1, 2, 5, 6],
                [2, 5, 6],
                [5, 6]
            ],
            [ # children=3
                [1, 2, 3, 7, 8, 9],
                [2, 3, 7, 8, 9],
                [3, 7, 8, 9],
                [7, 8, 9]
            ],
            [ # children=4
                [1, 2, 3, 4, 9, 10, 11, 12],
                [2, 3, 4, 9, 10, 11, 12],
                [3, 4, 9, 10, 11, 12],
                [4, 9, 10, 11, 12],
                [9, 10, 11, 12]
            ]
        ]
        nchildren = size(maps.parent2child,1);
        left_cells = left_cell_table_1d[nchildren][face];
        right_cells = right_cell_table_1d[nchildren][face];

    elseif dim == 2
        if size(maps.parent2neighbor,1) == 3 # triangles
            # Triangle parents have 9 faces, patches have 16 cells
            # Here Left means toward the center of the central parent
            left_cell_table_triangle = [
                [1, 4, 3, 15, 16, 13, 2, ],
                [2, 4, 3, 9, 12, 11, 1],
                [2, 4, 1, 7, 8, 5, 3],
                [3, 4, 1, 13, 16, 15, 2],
                [3, 4, 2, 11, 12, 9, 1],
                [1, 4, 2, 5, 8, 7, 3],
                [4, 2, 3, 9, 11, 12],
                [4, 1, 3, 15, 13, 16],
                [4, 1, 2, 5, 7, 8]
            ]
            right_cell_table_triangle = [
                [5, 8, 6, 7],
                [7, 8, 6, 5],
                [9, 12, 10, 11],
                [11, 12, 10, 9],
                [13, 16, 14, 15],
                [15, 16, 14, 13],
                [1, 5, 15, 8, 16],
                [2, 7, 9, 8, 12],
                [3, 11, 13, 12, 16]
            ]
            left_cells = left_cell_table_triangle[face];
            right_cells = right_cell_table_triangle[face];
        else # quads
            #TODO
        end

    elseif dim == 3
        #TODO
    end

    return (patch[left_cells], patch[right_cells]);
end

# Insert the single dof into the greater vector
function insert_linear!(b, bel, glb, dof, Ndofs)
    # group nodal dofs
    for d=1:length(dof)
        ind = glb.*Ndofs .- (Ndofs-dof[d]);
        ind2 = ((d-1)*length(glb)+1):(d*length(glb));

        b[ind] = b[ind] + bel[ind2];
    end
end

function place_sol_in_vars(var, sol, stepper)
    # place the values in the variable value arrays
    if typeof(var) <: Array
        tmp = 0;
        totalcomponents = 0;
        for vi=1:length(var)
            totalcomponents = totalcomponents + var[vi].total_components;
        end
        for vi=1:length(var)
            components = var[vi].total_components;
            for compi=1:components
                if stepper.type == EULER_EXPLICIT
                    var[vi].values[compi,:] = sol[(compi+tmp):totalcomponents:end];
                else
                    var[vi].values[compi,:] = sol[(compi+tmp):totalcomponents:end];
                end
                tmp = tmp + 1;
            end
        end
    else
        components = var.total_components;
        for compi=1:components
            if stepper.type == EULER_EXPLICIT
                var.values[compi,:] = sol[compi:components:end];
            else
                var.values[compi,:] = sol[compi:components:end];
            end
        end
    end
end

function get_var_vals(var, vect=nothing)
    # place the variable values in a vector
    if typeof(var) <: Array
        tmp = 0;
        totalcomponents = 0;
        for vi=1:length(var)
            totalcomponents = totalcomponents + var[vi].total_components;
        end
        if vect === nothing
            vect = zeros(totalcomponents * length(var[1].values[1,:]));
        end

        for vi=1:length(var)
            components = var[vi].total_components;
            for compi=1:components
                vect[(compi+tmp):totalcomponents:end] = var[vi].values[compi,:];
                tmp = tmp + 1;
            end
        end
    else
        components = var.total_components;
        if vect === nothing
            vect = zeros(components * length(var.values[1,:]));
        end
        for compi=1:components
            vect[compi:components:end] = var.values[compi,:];
        end
    end

    return vect;
end

end #module


const Accessor = Function;

abstract type Averager end;

# TODO: consider creating an averager which uses arbitrary precision arithmetic
mutable struct StandardAverager{T,N} <: Averager
  value::T;
  normalizer::N;
  accessor::Accessor;
end

function StandardAverager(accessor::Accessor, chain::EAPChain)
  return StandardAverager(
                          0.0 * accessor(chain),
                          0.0,
                          accessor
                         );
end

function StandardAverager(dt::DataType, accessor::Accessor, chain::EAPChain)
  return StandardAverager(
                          dt(0.0 * accessor(chain)),
                          dt(0.0),
                          accessor
                         );
end

function StandardAverager(dt1::DataType, dt2::DataType, accessor::Accessor, chain::EAPChain)
  return StandardAverager(
                          dt2(0.0 * accessor(chain)),
                          dt1(0.0),
                          accessor
                         );
end

@inline get_avg(stdavg::StandardAverager) = stdavg.value / stdavg.normalizer;

function record!(stdavg::StandardAverager, chain::EAPChain)
  stdavg.value += stdavg.accessor(chain);
  stdavg.normalizer += 1;
end

function record!(stdavg::StandardAverager{<:Vector}, chain::EAPChain)
  stdavg.value[:] += stdavg.accessor(chain);
  stdavg.normalizer += 1;
end

abstract type WeightFunction end;

struct UmbrellaAverager{T,N} <: Averager
  stdavg::StandardAverager{T,N};
  weightFunction::WeightFunction;
end

function UmbrellaAverager(ac::Accessor, wf::WeightFunction, chain::EAPChain)
  return UmbrellaAverager(StandardAverager(ac, chain), wf);
end

@inline get_avg(ua::UmbrellaAverager) = get_avg(ua.stdavg);

function record!(ua::UmbrellaAverager, chain::EAPChain)
  expw = exp(ua.weightFunction(chain));
  ua.stdavg.value += ua.stdavg.accessor(chain) / expw;
  ua.stdavg.normalizer += 1.0 / expw;
end

function record!(ua::UmbrellaAverager{Vector{<:Real}}, chain::EAPChain)
  expw = exp(ua.weightFunction(chain));
  ua.stdavg.value[:] += ua.stdavg.accessor(chain) / expw;
  ua.stdavg.normalizer += 1.0 / expw;
end

function record!(ua::UmbrellaAverager{Dec128}, chain::EAPChain)
  expw = exp(Dec128(ua.weightFunction(chain)));
  ua.stdavg.value += ua.stdavg.accessor(chain) / expw;
  ua.stdavg.normalizer += 1.0 / expw;
end

function record!(ua::UmbrellaAverager{Vector{Dec128}}, chain::EAPChain)
  expw = exp(Dec128(ua.weightFunction(chain)));
  ua.stdavg.value[:] += ua.stdavg.accessor(chain) / expw;
  ua.stdavg.normalizer += 1.0 / expw;
end

function record!(ua::UmbrellaAverager{BigFloat}, chain::EAPChain)
  expw = exp(BigFloat(ua.weightFunction(chain)));
  ua.stdavg.value += ua.stdavg.accessor(chain) / expw;
  ua.stdavg.normalizer += 1.0 / expw;
end

function record!(ua::UmbrellaAverager{Vector{BigFloat}}, chain::EAPChain)
  expw = exp(BigFloat(ua.weightFunction(chain)));
  ua.stdavg.value[:] += ua.stdavg.accessor(chain) / expw;
  ua.stdavg.normalizer += 1.0 / expw;
end

struct WeightlessFunction <: WeightFunction
end

(wf::WeightlessFunction)(::EAPChain) = 1.0;

struct AntiDipoleWeightFunction <: WeightFunction
  log_gauge::Float64;
end

# forward call
@inline AntiDipoleWeightFunction(chain::EAPChain) = AntiDipoleWeightFunction(chain, chain.μ);
function AntiDipoleWeightFunction(chain::EAPChain, dr::DielectricResponse)
  return AntiDipoleWeightFunction(-(dr.K1 + 2 * dr.K2) * chain.E0*chain.E0 *
                                   n(chain) / (3 * chain.kT) + chain.Ω);
end

function AntiDipoleWeightFunction(chain::EAPChain, dr::PolarResponse)
  return AntiDipoleWeightFunction(-eigvals(dr.M)[end] * chain.E0 *
                                   n(chain) / (3 * chain.kT) + chain.Ω);
end

function (wf::AntiDipoleWeightFunction)(chain::EAPChain)
  return (sum(chain.us) / chain.kT *
          (0.2 + 0.8*(exp(-(chain.Fx^2+chain.Fz^2)/chain.kT)))
          - wf.log_gauge);
end


module XML

using LabelMe
using LightXML

export load, annotation, source, image_size, object, polygon, element

function load(filename::AbstractString)
    e = root(parse_file(filename))
    return annotation(e)
end

function annotation(e::XMLElement)
    a = Annotation()
    a.filename = element_content(e, "filename")
    a.folder = element_content(e, "folder")
    a.source = source(find_element(e, "source"))
    a.image_size = image_size(find_element(e, "imagesize"))
    for object_element in get_elements_by_tagname(e, "object")
        push!(a.objects, object(object_element))
    end
    return a
end

function source(e::Void)
    return Source()
end

function source(e::XMLElement)
    s = Source()
    s.sourceImage = element_content(e, "sourceImage")
    s.sourceAnnotation = element_content(e, "sourceAnnotation")
    return s
end

function image_size(e::XMLElement)
    nrows = tryparse(Int64, element_content(e, "nrows"))
    ncols = tryparse(Int64, element_content(e, "ncols"))
    if isnull(nrows) || isnull(ncols)
        return nothing
    end
    return (get(nrows), get(ncols))
end

function object(e::XMLElement)
    o = Object()
    o.id = tryparse(Int64, element_content(e, "id"))
    o.name = element_content(e, "name")
    o.deleted = tryparse(UInt8, element_content(e, "deleted"))
    o.verified = tryparse(UInt8, element_content(e, "verified"))
    o.occluded = tryparse(UInt8, element_content(e, "occluded"))
    o.attributes = element_content(e, "attributes")
    o.parts_parent = tryparse(Int64, element_content(find_element(e, "parts"), "ispartof"))
    hasparts = element_content(find_element(e, "parts"), "hasparts")
    o.parts_children = [parse(Int64, p) for p in split(hasparts, ','; keep=false)]
    o.date = DateTime(element_content(e, "date"), "d-u-yyyy H:M:S")
    o.polygon = polygon(find_element(e, "polygon"))
    return o
end

function polygon(e::XMLElement)
    p = Polygon()
    p.username = element_content(e, "username")
    for pt_element in get_elements_by_tagname(e, "pt")
        push!(p.points, (
            parse(Float32, element_content(pt_element, "x")),
            parse(Float32, element_content(pt_element, "y"))
            ))
    end
    return p
end

function element_content(e::XMLElement, tag::AbstractString)
    first = find_element(e, tag)
    if first == nothing
        return ""
    end
    return content(first)
end

function element(a::Annotation)
    e = new_element("annotation")
    add_text(new_child(e, "filename"), a.filename)
    add_text(new_child(e, "folder"), a.folder)
    add_child(e, element(a.source))
    i = new_child(e, "imagesize")
    add_text(new_child(i, "nrows"), isnull(a.image_size) ? "" : string(get(a.image_size)[1]))
    add_text(new_child(i, "ncols"), isnull(a.image_size) ? "" : string(get(a.image_size)[2]))
    for o in a.objects
        add_child(e, element(o))
    end
    return e
end

function element(s::Source)
    e = new_element("source")
    add_text(new_child(e, "sourceImage"), s.sourceImage)
    add_text(new_child(e, "sourceAnnotation"), s.sourceAnnotation)
    return e
end

function element(o::Object)
    e = new_element("object")
    if !isnull(o.id)
        add_text(new_child(e, "id"), string(get(o.id)))
    end
    add_text(new_child(e, "name"), o.name)
    if !isnull(o.deleted)
        add_text(new_child(e, "deleted"), string(UInt8(get(o.deleted))))
    end
    if !isnull(o.verified)
        add_text(new_child(e, "verified"), string(UInt8(get(o.verified))))
    end
    if !isnull(o.occluded)
        add_text(new_child(e, "occluded"), string(UInt8(get(o.occluded))))
    end
    add_text(new_child(e, "attributes"), o.attributes)
    p = new_child(e, "parts")
    if !isnull(o.parts_parent)
        add_text(new_child(p, "ispartof"), string(get(o.parts_parent)))
    end
    add_text(new_child(p, "hasparts"), join(o.parts_children, ","))
    add_text(new_child(e, "date"), Dates.format(o.date, "d-u-yyyy H:M:S"))
    add_child(e, element(o.polygon))
    return e
end

function element(p::Polygon)
    e = new_element("polygon")
    add_text(new_child(e, "username"), p.username)
    for pt in p.points
        pt_e = new_child(e, "pt")
        add_text(new_child(pt_e, "x"), string(pt[1]))
        add_text(new_child(pt_e, "y"), string(pt[2]))
    end
    return e
end

end # module


using XYZ
using StaticArrays
using Base.Test

# use test file from LasIO
filename = "libLAS_1.2.las" # point format 0
filepath = joinpath(Pkg.dir("LasIO"), "test", filename)
workdir = dirname(@__FILE__)

@testset "Cloud contruction and indexing" begin
    cloud, header = XYZ.read_pointcloud(filepath)

    attr = cloud.attributes
    @test isa(attr, Dict{Symbol,Vector})
    @test Set(keys(attr)) == Set(Symbol[:number_of_returns, :edge_of_flight_line, :intensity,
        :key_point, :scan_direction, :withheld, :pt_src_id, :classification, :scan_angle, :synthetic,
        :return_number, :user_data])

    @test XYZ.boundingbox(cloud) == XYZ.calc_bbox(cloud)

    cloudp1 = cloud[1:1]
    @test length(cloudp1) == 1
    @test XYZ.positions(cloudp1) == [SVector{3, Float64}(1.44013394e6, 375000.23, 846.66)]
    @test cloudp1[:intensity] == [0x00fa]
    @test cloudp1[:scan_angle] == [Int8(0)]
    @test cloudp1[:user_data] == [0x00]
    @test cloudp1[:pt_src_id] == [0x001d]
    @test cloudp1[:return_number] == [0x00]
    @test cloudp1[:number_of_returns] == [0x00]
    @test cloudp1[:scan_direction] == [false]
    @test cloudp1[:edge_of_flight_line] == [false]
    @test cloudp1[:classification] == [0x02]
    @test cloudp1[:synthetic] == [false]
    @test cloudp1[:key_point] == [false]
    @test cloudp1[:withheld] == [false]
end

@testset "Outlier classification" begin
    cloud, header = XYZ.read_pointcloud(filepath)

    XYZ.reset_class!(cloud)
    @test all(c -> c == 0x00, cloud[:classification])

    XYZ.classify_outliers!(cloud;
        cellsize = 100.0,   # cellsize used for evaluation of points
        dz = 1.0,           # threshold in vertical distance between low points
        max_outliers = 15,  # max number of outliers in one cell = low  points to be evaluated
        max_height = 100.0) # threshold in vertical distance for high points

    low_noise = UInt8(7) # from ASPRS Standard LIDAR Point Classes
    high_noise = UInt8(18) # from ASPRS Standard LIDAR Point Classes (LAS 1.4)
    @test count(c -> c == low_noise, cloud[:classification]) == 13
    @test count(c -> c == high_noise, cloud[:classification]) == 1948
end

@testset "Rasterization" begin
    cloud, header = XYZ.read_pointcloud(filepath)
    cellsize = 10.0
    nground = 402812
    nodata = -9999.0
    rasterindex = XYZ.define_raster(cloud, cellsize)
    @test rasterindex.npoints == 497536
    @test typeof(rasterindex) == XYZ.Raster{Int32}
    @test rasterindex.nrow == rasterindex.ncol == 500

    raster = XYZ.rasterize(cloud, rasterindex; reducer=XYZ.reducer_minz)
    @test size(raster) == (500, 500, 1)
    @test eltype(raster) == Float32
    @test count(x -> x == nodata, raster) == 3336

    rasterindex_ground = XYZ.filter_raster(rasterindex, cloud, XYZ.ground)
    @test rasterindex_ground.npoints == nground

    # pointfilter
    @test count(i -> XYZ.ground(cloud, i), 1:length(cloud)) == nground
end


using SpecialFunctions

function test(f, α, u0, T, h)
    N=Int64(T/h)
end


using CUDA
using Random
using Test
using VUMPS
using VUMPS: vumps, vumps_env

@testset "$(Ni)x$(Nj) VUMPSRuntime with $atype{$dtype}" for Ni = [1,2,3], Nj = [1,2,3], atype = [Array, CuArray], dtype = [Float64, ComplexF64]
    @test SquareLattice <: AbstractLattice

    M = Array{atype{dtype,4},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        M[i,j] = atype(rand(dtype,2,2,2,2))
    end
    rt = SquareVUMPSRuntime(M, Val(:random), 2)
end

@testset "$(Ni)x$(Nj) vumps with $atype{$dtype}" for Ni = [2], Nj = [2], atype = [Array, CuArray], dtype = [Float64, ComplexF64]
    M = Array{atype{dtype,4},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        M[i,j] = atype(rand(dtype,2,2,2,2))
    end
    rt = SquareVUMPSRuntime(M, Val(:random), 2)
    env = vumps(rt)
    @test env !== nothing
end

@testset "$(Ni)x$(Nj) vumps_env with $atype{$dtype}" for Ni = [2], Nj = [2], atype = [Array, CuArray], dtype = [Float64, ComplexF64]
    M = Array{atype{dtype,4},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        M[i,j] = atype(rand(dtype,2,2,2,2))
    end
    env = vumps_env(M; χ = 10)
    @test env !== nothing
end

@testset "$(Ni)x$(Nj) obs_env" for Ni = [2], Nj = [2], atype = [Array, CuArray], dtype = [Float64, ComplexF64]
    M = Array{atype{dtype,4},2}(undef, Ni, Nj)
    for j = 1:Nj, i = 1:Ni
        M[i,j] = atype(rand(dtype,2,2,2,2))
    end
    env = obs_env(M; χ = 10, verbose = true, savefile = true, infolder = "./data1", outfolder = "./data2", show_every = 3)
    @test env !== nothing
end


module PipsNlpInterface

#TODO: update pips-nlp interface to work on optigraphs in addition to optinodes for subproblems

using SparseArrays
using LinearAlgebra
using DataStructures
using PipsNLP
import MPI

import JuMP
import MathOptInterface
const MOI = MathOptInterface
using Plasmo

include("pipsnlp_c_interface.jl")
using .PipsNlpSolver

include("pips_utils.jl")
include("nlp_evaluator.jl")

export pipsnlp_solve

function pipsnlp_solve(graph::OptiGraph) #Assume graph variables and constraints are first stage
    #Check structure: 3 possibilities
    #1.) optigraph without subgraphs (only link constraints, no first stage -> second stage constraints)
    #2.) optigraph with single subgraph (first stage and linkconstraints between optinodes)
    #3.) Not yet supported: optigraph with multiple subgraphs (first stage and linkconstraints between subgraphs)

    comm = MPI.COMM_WORLD
    if MPI.Comm_rank(comm) == 0
        println("Building Model for PIPS-NLP")
    end

    #NOTE: either the interface figures out the linking structure, or the user provided it (e.g. using distribute_optigraph)
    if haskey(graph.ext,:user_pips_data)
        worker_data = graph.ext[:user_pips_data]
    else
        worker_data = PipsNLP._setup_pips_nlp_data!(graph)
    end

    first_stage = worker_data.first_stage
    submodels = setdiff(all_nodes(graph),[first_stage])
    n_sub_models = length(submodels)
    model_list = [first_stage; submodels]

    #unpack data
    first_stage_data = PipsNLP._get_pips_data(first_stage)
    nlinkeq = worker_data.n_linkeq_cons
    nlinkineq = worker_data.n_linkineq_cons
    ineqlink_lb = worker_data.link_ineq_lower
    ineqlink_ub = worker_data.link_ineq_upper
    eqlink_lb = worker_data.link_eq_lower
    eqlink_ub = worker_data.link_eq_upper

    #initialize values
    function str_init_x0(nodeid, x0)
        node = model_list[nodeid+1]
        jump_initval = copy(JuMP.start_value.(JuMP.all_variables(node)))
        local_initval = Vector{Float64}(undef,JuMP.num_variables(node))

        nothing_values = isa.(jump_initval,Nothing)
        float_values = .!(nothing_values)
        local_initval[nothing_values] .= 0  #set to 0 by default
        local_initval[float_values] .= jump_initval[float_values]
        original_copy(local_initval,x0)
    end

    #get problem info
    function str_prob_info(nodeid,flag,mode,col_lb,col_ub,row_lb,row_ub)
        if flag != 1
            node = model_list[nodeid+1]
            local_data = PipsNLP._get_pips_data(node)
            #do a warm start on the node
            if !(local_data.loaded)
                local_data.loaded = true
                if (nodeid > 0)
                    if haskey(node.ext, :warmStart)
    	                if node.ext[:warmStart] == true
              	      	   optimize!(node)
                   	   	end
                    end
            	end

                nlp_lb, nlp_ub = constraintbounds(node.ext[:constraint_data])       #This is every constraint in the model
         		local_data.local_m  = length(nlp_lb)          #number of local constraints (rows)
    			newRowId = Array{Int}(undef,local_data.local_m)
    			eqId = 1
    			ineqId = 1

                #go through all constraints, check if they are equality of inequality
                for c in 1:local_data.local_m
                    if nlp_lb[c] == nlp_ub[c]
                        push!(local_data.eq_idx, c)
                        newRowId[c] = eqId
                        eqId +=  1
                    else
        				push!(local_data.ineq_idx, c)
        				newRowId[c] = ineqId
        				ineqId += 1
                    end
         		end

                #Local node data
                local_data.m  = local_data.local_m + local_data.num_eqconnect + local_data.num_ineqconnect
                local_data.n = JuMP.num_variables(node)
                local_data.x_sol = zeros(Float64,local_data.n)
                local_data.firstJeqmat = sparse(local_data.firstIeq, local_data.firstJeq, local_data.firstVeq, local_data.num_eqconnect, first_stage_data.n)
                local_data.secondJeqmat = sparse(local_data.secondIeq, local_data.secondJeq, local_data.secondVeq, local_data.num_eqconnect, local_data.n)
                local_data.firstJineqmat = sparse(local_data.firstIineq, local_data.firstJineq, local_data.firstVineq, local_data.num_ineqconnect, first_stage_data.n)
                local_data.secondJineqmat = sparse(local_data.secondIineq, local_data.secondJineq, local_data.secondVineq, local_data.num_ineqconnect, local_data.n)

                if node.nlp_data == nothing
                    JuMP._init_NLP(node)
                end
                local_data.d = JuMP.NLPEvaluator(node)
                MOI.initialize(local_data.d, [:Grad,:Jac, :Hess])
                jac_structure = pips_jacobian_structure(local_data.d)
                Ijac = [jac_structure[i][1] for i = 1:length(jac_structure)]
                Jjac = [jac_structure[j][2] for j = 1:length(jac_structure)]
                Ijaceq = Int[]
                Jjaceq = Int[]
                Ijacineq = Int[]
                Jjacineq = Int[]
                jac_eq_index = Int[]
                jac_ineq_index = Int[]
                for i in 1:length(Ijac)
                    c = Ijac[i]
                    if nlp_lb[c] == nlp_ub[c]
                        modifiedrow = newRowId[c]
                        push!(Ijaceq, modifiedrow)
                        push!(Jjaceq, Jjac[i])
                        push!(jac_eq_index, i)
                    else
                        modifiedrow = newRowId[c]
                        push!(Ijacineq, modifiedrow)
                        push!(Jjacineq, Jjac[i])
                        push!(jac_ineq_index,i)
                    end
                end

        		node.ext[:Ijaceq] = Ijaceq
         		node.ext[:Jjaceq] = Jjaceq
                node.ext[:Ijacineq] = Ijacineq
         		node.ext[:Jjacineq] = Jjacineq
         		node.ext[:jac_eq_index] = jac_eq_index
         		node.ext[:jac_ineq_index] = jac_ineq_index
                hess_structure = pips_hessian_lagrangian_structure(local_data.d)
                Ihess = [hess_structure[i][1] for i = 1:length(hess_structure)]
                Jhess = [hess_structure[j][2] for j = 1:length(hess_structure)]

         		Hmap = Bool[]
         		node_Hrows = Int[]
         		node_Hcols = Int[]
         		for i in 1:length(Ihess)
            	    if Jhess[i] <= Ihess[i]
               	        push!(node_Hrows, Ihess[i])
                   		push!(node_Hcols, Jhess[i])
                   		push!(Hmap, true)
            	    else
                        push!(Hmap, false)
                    end
                end
         		val = ones(Float64,length(node_Hrows))
         		mat = sparse(node_Hrows,node_Hcols,val, local_data.n, local_data.n)
         		node.ext[:Hrows] = node_Hrows
         		node.ext[:Hcols] = node_Hcols
         		node.ext[:Hmap] = Hmap
         		local_hessnnz = length(mat.rowval)
         		local_data.local_unsym_hessnnz = length(Ihess)
         		local_data.hessnnz = local_hessnnz
		    end

	 	    if mode == :Values
                nlp_lb, nlp_ub = constraintbounds(node.ext[:constraint_data])
    			eq_lb=Float64[]
    			eq_ub=Float64[]
    			ineq_lb=Float64[]
    			ineq_ub=Float64[]
    			for i in 1: length(nlp_lb)
    			    if nlp_lb[i] == nlp_ub[i]
    			       push!(eq_lb, nlp_lb[i])
    			       push!(eq_ub, nlp_ub[i])
    			    else
    			       push!(ineq_lb, nlp_lb[i])
    			       push!(ineq_ub, nlp_ub[i])
    			    end
    			end

    			if nodeid !=  0 #0 is the master
    			   eq_lb = [eq_lb; local_data.eqconnect_lb]
    			   eq_ub = [eq_ub; local_data.eqconnect_ub]
    			   ineq_lb = [ineq_lb; local_data.ineqconnect_lb]
    			   ineq_ub = [ineq_ub; local_data.ineqconnect_ub]
    			end

    			original_copy([eq_lb;ineq_lb], row_lb)
    			original_copy([eq_ub;ineq_ub], row_ub)
                colLower = variablelowerbounds(node)
                colUpper = variableupperbounds(node)
                original_copy(colLower, col_lb)
    			original_copy(colUpper, col_ub)
            end
		    return (local_data.n,local_data.m)
		else
		    if mode == :Values
                original_copy([eqlink_lb; ineqlink_lb], row_lb)
                original_copy([eqlink_ub; ineqlink_ub], row_ub)
            end
            return (0, nlinkeq + nlinkineq)
		end
    end

    #evaluate objective
    #x0 is first stage variable values, x1 is local values
    function str_eval_f(nodeid,x0,x1)
    	node = model_list[nodeid+1] #Julia doesn't start index at 0
        local_data = PipsNLP._get_pips_data(node)
        local_d = PipsNLP._get_pips_data(node).d
        if nodeid ==  0
            local_x = x0
        else
            local_x = x1
        end
        #check objective sign
        local_scl = (JuMP.objective_sense(node) == MOI.MAX_SENSE) ? -1.0 : 1.0
        f = local_scl*pips_eval_objective(local_d,local_x)
        return f
    end

    #evaluate constraints
    function str_eval_g(nodeid,x0,x1,new_eq_g, new_inq_g)
        node = model_list[nodeid+1]
        local_data = PipsNLP._get_pips_data(node)
        local_d = PipsNLP._get_pips_data(node).d
        if nodeid ==  0
            local_x = x0
        else
            local_x = x1
        end
        local_g = Array{Float64}(undef,local_data.local_m)
        pips_eval_constraint(local_d, local_g, local_x)

        #evaluates local constraints and connecting constraints with first stage
        new_eq_g[1:end] = [local_g[local_data.eq_idx]; local_data.firstJeqmat*x0+local_data.secondJeqmat*x1]
        new_inq_g[1:end] = [local_g[local_data.ineq_idx]; local_data.firstJineqmat*x0+local_data.secondJineqmat*x1]
	    return Int32(1)
    end

    #get solution
    function str_write_solution(id::Integer, x::Vector{Float64}, y_eq::Vector{Float64}, y_ieq::Vector{Float64})
        node = model_list[id+1]
        local_data = PipsNLP._get_pips_data(node)
        local_data.x_sol = copy(x)

        #TODO: Grab dual values
        rank = MPI.Comm_rank(comm)
        local_data.coreid = rank
    end

    #evaluate gradient
    function str_eval_grad_f(rowid,colid,x0,x1,new_grad_f)
        node = model_list[rowid+1]
        if rowid == colid
            local_data = PipsNLP._get_pips_data(node)
            local_d = PipsNLP._get_pips_data(node).d
            if colid ==  0
                local_x = x0
            else
                local_x = x1
            end
            local_grad_f = Array{Float64}(undef,local_data.n)
            pips_eval_objective_gradient(local_d, local_grad_f, local_x)
            #eval_grad_f(local_d, local_grad_f, local_x)
            #local_scl = (node.objSense == :Min) ? 1.0 : -1.0
            local_scl = (JuMP.objective_sense(node) == MOI.MAX_SENSE) ? -1.0 : 1.0
            #scale!(local_grad_f,local_scl)
            rmul!(local_grad_f,local_scl)
            original_copy(local_grad_f, new_grad_f)
        elseif colid == 0
            new_grad_f[1:end] .= 0
        else
            @assert(false)
        end
        return Int32(1)
    end

    #copy row/column indices to C pointer
    function array_copy(src,dest)
        @assert(length(src)==length(dest))
        for i in 1:length(src)
            dest[i] = src[i]-1
        end
    end

    #copy array elements
    function original_copy(src,dest)
        @assert(length(src)==length(dest))
        for i in 1:length(src)
            dest[i] = src[i]
        end
    end

    #evaluate jacobian
    function str_eval_jac_g(rowid,colid,flag, x0,x1,mode,e_rowidx,e_colptr,e_values,i_rowidx,i_colptr,i_values)
        if flag != 1 #populate parent child structure
            node = model_list[rowid+1]
            local_data = PipsNLP._get_pips_data(node)
            local_m_eq = length(local_data.eq_idx)
            local_m_ineq = length(local_data.ineq_idx)
            if mode == :Structure
                if (rowid == colid)
                    Ieq=[node.ext[:Ijaceq];local_data.secondIeq .+ local_m_eq]
                    Jeq=[node.ext[:Jjaceq];local_data.secondJeq]
                    Veq= ones(Float64, length(Ieq))
                    Iineq=[node.ext[:Ijacineq];local_data.secondIineq .+ local_m_ineq]
                    Jineq=[node.ext[:Jjacineq];local_data.secondJineq]
                    Vineq=ones(Float64, length(Iineq))
                    eqmat = sparse(Ieq, Jeq, Veq, local_m_eq + local_data.num_eqconnect, local_data.n)
                    ineqmat = sparse(Iineq, Jineq, Vineq, local_m_ineq + local_data.num_ineqconnect, local_data.n)
                else
                    eqmat = sparse(local_m_eq .+ local_data.firstIeq, local_data.firstJeq, local_data.firstVeq, local_m_eq .+ local_data.num_eqconnect, first_stage_data.n)
                    ineqmat = sparse(local_m_ineq .+ local_data.firstIineq, local_data.firstJineq, local_data.firstVineq, local_m_ineq .+ local_data.num_ineqconnect, first_stage_data.n)
                end
                return(length(eqmat.rowval), length(ineqmat.rowval))
            else #mode = :Values
                if rowid == colid    #evaluate Jacobian for block
                    if colid ==  0   #if it's the root node
                        local_x = x0
                    else
                        local_x = x1
                    end
                    local_values = Array{Float64}(undef,length(node.ext[:Ijaceq])+length(node.ext[:Ijacineq]))
                    #eval_jac_g(local_data.d, local_values, local_x)
                    pips_eval_constraint_jacobian(local_data.d, local_values, local_x)
                    jac_eq_index = node.ext[:jac_eq_index]
                    jac_ineq_index = node.ext[:jac_ineq_index]
                    Ieq=[node.ext[:Ijaceq];local_data.secondIeq .+ local_m_eq]
                    Jeq=[node.ext[:Jjaceq];local_data.secondJeq]
                    Veq=[local_values[jac_eq_index];local_data.secondVeq]
                    Iineq=[node.ext[:Ijacineq];local_data.secondIineq .+ local_m_ineq]
                    Jineq=[node.ext[:Jjacineq];local_data.secondJineq]
                    Vineq=[local_values[jac_ineq_index];local_data.secondVineq]
                    eqmat = _sparse_keep_zero(Ieq, Jeq, Veq, local_m_eq + local_data.num_eqconnect, local_data.n)
                    ineqmat = _sparse_keep_zero(Iineq, Jineq, Vineq, local_m_ineq + local_data.num_ineqconnect, local_data.n)
                else #evaluate
                    eqmat = _sparse_keep_zero(local_m_eq .+ local_data.firstIeq, local_data.firstJeq, local_data.firstVeq, local_m_eq .+ local_data.num_eqconnect, first_stage_data.n)
                    ineqmat = _sparse_keep_zero(local_m_ineq .+ local_data.firstIineq, local_data.firstJineq, local_data.firstVineq, local_m_ineq .+ local_data.num_ineqconnect, first_stage_data.n)
                end
                if length(eqmat.nzval) > 0
                    array_copy(eqmat.rowval,e_rowidx)
                    array_copy(eqmat.colptr,e_colptr)
                    original_copy(eqmat.nzval,e_values)
                end
                if length(ineqmat.nzval) > 0
                    array_copy(ineqmat.rowval,i_rowidx)
                    array_copy(ineqmat.colptr,i_colptr)
                    original_copy(ineqmat.nzval, i_values)
                end
            end
        else #populate linkconstraint structure
            node = model_list[rowid+1]
            local_data = PipsNLP._get_pips_data(node)
            linkIeq = local_data.linkIeq
            linkJeq = local_data.linkJeq
            linkVeq = local_data.linkVeq
            linkIineq = local_data.linkIineq
            linkJineq = local_data.linkJineq
            linkVineq = local_data.linkVineq

            if mode == :Structure
                return(length(linkVeq), length(linkVineq))
            else
                eqmat = sparse(linkIeq, linkJeq, linkVeq, nlinkeq, local_data.n)
                ineqmat = sparse(linkIineq, linkJineq, linkVineq, nlinkineq, local_data.n)
                if length(eqmat.nzval) > 0
                    array_copy(eqmat.rowval,e_rowidx)
                    array_copy(eqmat.colptr,e_colptr)
                    original_copy(eqmat.nzval,e_values)
                end
                if length(ineqmat.nzval) > 0
                    array_copy(ineqmat.rowval,i_rowidx)
                    array_copy(ineqmat.colptr,i_colptr)
                    original_copy(ineqmat.nzval, i_values)
                end
            end
        end
        return Int32(1)
    end

    #evaluate hessian
    function str_eval_h(rowid,colid,x0,x1,obj_factor,lambda,mode,rowidx,colptr,values)
        node = model_list[colid+1]
        local_data = PipsNLP._get_pips_data(node)
        if mode == :Structure
            if rowid == colid
                node_Hrows  = node.ext[:Hrows]
                node_Hcols = node.ext[:Hcols]
                node_val = ones(Float64,length(node_Hrows))
                mat = _sparse_keep_zero(node_Hrows,node_Hcols,node_val, local_data.n, local_data.n)
                return length(mat.rowval)
            elseif colid == 0
                node_Hrows  = node.ext[:Hrows]
                node_Hcols = node.ext[:Hcols]
                node_Hrows,node_Hcols = _exchange(node_Hrows,node_Hcols)
                node_val = ones(Float64,length(node_Hrows))
                mat = _sparse_keep_zero(node_Hrows,node_Hcols,node_val, local_data.n, local_data.n)
                return length(mat.rowval)
            else
                return 0
            end
        else #fill in values
            if rowid == colid
                node_Hmap = node.ext[:Hmap]
                node_Hrows  = node.ext[:Hrows]
                node_Hcols = node.ext[:Hcols]
                if colid ==  0
                    local_x = x0
                else
                    local_x = x1
                end
                local_unsym_values = Array{Float64}(undef,local_data.local_unsym_hessnnz)
                node_val = ones(Float64,length(node_Hrows))
                local_scl = (JuMP.objective_sense(node) == MOI.MAX_SENSE) ? -1.0 : 1.0
                local_m_eq = length(local_data.eq_idx)
                local_m_ineq = length(local_data.ineq_idx)
                local_lambda = zeros(Float64, local_data.local_m)
                for i in 1:local_m_eq
                    local_lambda[local_data.eq_idx[i]] = lambda[i]
                end
                for i in 1:local_m_ineq
                    local_lambda[local_data.ineq_idx[i]] = lambda[i+local_m_eq]
                end

                #eval_hesslag(local_data.d, local_unsym_values, local_x, obj_factor*local_scl, local_lambda)
                pips_eval_hessian_lagrangian(local_data.d, local_unsym_values, local_x, obj_factor*local_scl, local_lambda)
                local_sym_index=1
                for i in 1:local_data.local_unsym_hessnnz
                    if node_Hmap[i]
                        node_val[local_sym_index] = local_unsym_values[i]
                        local_sym_index +=1
                    end
                end
                node_Hrows,node_Hcols = _exchange(node_Hrows,node_Hcols)
                mat = _sparse_keep_zero(node_Hrows,node_Hcols,node_val, local_data.n, local_data.n)
                array_copy(mat.rowval,rowidx)
                array_copy(mat.colptr,colptr)
                original_copy(mat.nzval,values)
            elseif colid ==	0
                node_Hrows  = node.ext[:Hrows]
                node_Hcols = node.ext[:Hcols]
                node_val = zeros(Float64,length(node_Hrows))
                node_Hrows,node_Hcols = _exchange(node_Hrows,node_Hcols)
                mat = _sparse_keep_zero(node_Hrows,node_Hcols,node_val, local_data.n, local_data.n)
                array_copy(mat.rowval,rowidx)
                array_copy(mat.colptr,colptr)
                original_copy(mat.nzval,values)
            else
            end
        end
        return Int32(1)
    end

    #initalize MPI if not already
    if !MPI.Initialized()
        MPI.Init()
    end

    # comm = MPI.COMM_WORLD
    if (MPI.Comm_rank(comm) == 0)
        t1 = time()
    end

    root = 0
    rank = MPI.Comm_rank(comm)

    ###################################
    #PIPS-NLP FUNCTIONS
    ###################################
    #Create PipsModel (The PIPS interface model) and pass all the functions it requires
    model = PipsModel(:Min, 0,
    n_sub_models,
    str_init_x0,
    str_prob_info,
    str_eval_f,
    str_eval_g,
    str_eval_grad_f,
    str_eval_jac_g,
    str_eval_h,
    str_write_solution)

    #create the C++ problem struct
    prob = createProblemStruct(comm, model, true)
    if rank == 0
        println("Created PIPS-NLP Problem Struct")
    end

    #probably don't need this barrier
    MPI.Barrier(comm)

    #solve the problem struct.  MPI will do communication in the solver.
    ret = solveProblemStruct(prob)
    ###################################

    if rank == 0
        println("Solved PIPS-NLP Problem Struct")
    end

    if rank == 0
        println("Timing Results:")
        println("init_x0  ",prob.t_jl_init_x0)
        println("str_init_x0  ", prob.t_jl_str_prob_info)
        println("eval_f  ", prob.t_jl_eval_f)
        println("eval_g0  ",prob.t_jl_eval_g)
        println("eval_grad_f  ", prob.t_jl_eval_grad_f)
        println("eval_jac  ", prob.t_jl_eval_jac_g)
        println("str_eval_jac  ", prob.t_jl_str_eval_jac_g)
        println("eval_h  ",  prob.t_jl_eval_h)
        println("str_eval_h ",  prob.t_jl_str_eval_h)
        println("eval_write_solution  ",  prob.t_jl_write_solution)
        println("PIPS-NLP time:   ",  time() - t1, " (s)")
    end

    status_dict = Dict(
    0 => MOI.LOCALLY_SOLVED,        #:Optimal
    1 => MOI.OTHER_LIMIT,           #:Not_Finished
    2 => MOI.ITERATION_LIMIT,       #:Iteration_Limit
    3 => MOI.LOCALLY_INFEASIBLE,    #:Infeasible_Problem_Detected
    4 => MOI.NUMERICAL_ERROR)       #:Restoration_needed

    if ret in keys(status_dict)
        status = status_dict[ret]
    else
        status = MOI.OTHER_ERROR
    end

    #set solution
    for (idx,node) in enumerate(model_list)  #set solution values for each model
        local_data = PipsNLP._get_pips_data(node)
        if idx != 1 #all cores have the first stage
            coreid = zeros(Int, 1)
            sc = MPI.Reduce(local_data.coreid, MPI.SUM, root, comm)

            if rank == root
                coreid[1] = sc
            end

            MPI.Bcast!(coreid, length(coreid), root, comm)

            n = zeros(Int, 1)
            n[1] = local_data.n
            MPI.Bcast!(n,length(n),coreid[1], comm)

            if rank != coreid[1]
                local_data.n = n[1]
                local_data.x_sol = zeros(Float64,local_data.n)
            end
            MPI.Bcast!(local_data.x_sol, local_data.n, coreid[1], comm)
        end

        vars = Plasmo.all_variables(node)
        vals = local_data.x_sol
        Plasmo.set_node_primals(node,vars,vals)
        Plasmo.set_node_status(node,status)
        #TODO: set duals
        #TODO: load results into a graph backend optimizer
    end

    if rank == 0
        println(status)
    end

    return status
end  #end pips_nlp_solve

end #end module


#############################################################################
# ECOS.jl
# Wrapper around the ECOS solver https://github.com/ifa-ethz/ecos
# See http://github.com/JuliaOpt/ECOS.jl
#############################################################################
# types.jl
# Julia implementations of the types defined in ecos.h
# Types are mapped as follows
# pfloat -> Cdouble
# idxint -> SuiteSparse_long -> Clong
#############################################################################

# Solve status flags
# From: https://github.com/ifa-ethz/ecos/blob/master/include/ecos.h
const ECOS_OPTIMAL      =  0  # Problem solved to optimality
const ECOS_PINF         =  1  # Found certificate of primal infeasibility
const ECOS_DINF         =  2  # Found certificate of dual infeasibility
const ECOS_INACC_OFFSET = 10  # Offset exitflag at inaccurate results
const ECOS_MAXIT        = -1  # Maximum number of iterations reached
const ECOS_NUMERICS     = -2  # Search direction unreliable
const ECOS_OUTCONE      = -3  # s or z got outside the cone, numerics?
const ECOS_FATAL        = -7  # Unknown problem in solver

immutable Clpcone
    p::Clong
    w::Ptr{Cdouble}
    v::Ptr{Cdouble}
    kkt_idx::Ptr{Clong}
end

immutable Csocone
    p::Clong
    skbar::Ptr{Cdouble}
    zkbar::Ptr{Cdouble}
    a::Clong
    d1::Clong
    w::Cdouble
    eta::Cdouble
    eta_square::Cdouble
    q::Ptr{Cdouble}
    Didx::Ptr{Clong}
    u0::Cdouble
    u1::Cdouble
    v1::Cdouble
end

immutable Ccone
    lpcone::Ptr{Clpcone}
    socone::Ptr{Csocone}
    nsoc::Clong
end

immutable Cspmat
    jc::Ptr{Clong}
    ir::Ptr{Clong}
    pr::Ptr{Cdouble}
    n::Clong
    m::Clong
    nnz::Clong
end

immutable Ckkt
    PKPt::Ptr{Cspmat}
    L::Ptr{Cspmat}
    D::Ptr{Cdouble}
    work1::Ptr{Cdouble}
    work2::Ptr{Cdouble}
    work3::Ptr{Cdouble}
    work4::Ptr{Cdouble}
    work5::Ptr{Cdouble}
    work6::Ptr{Cdouble}
    RHS1::Ptr{Cdouble}
    RHS2::Ptr{Cdouble}
    dx1::Ptr{Cdouble}
    dx2::Ptr{Cdouble}
    dy1::Ptr{Cdouble}
    dy2::Ptr{Cdouble}
    dz1::Ptr{Cdouble}
    dz2::Ptr{Cdouble}
    P::Ptr{Clong}
    Pinv::Ptr{Clong}
    PK::Ptr{Clong}
    Parent::Ptr{Clong}
    Sign::Ptr{Clong}
    Pattern::Ptr{Clong}
    Flag::Ptr{Clong}
    Lnz::Ptr{Clong}
    delta::Cdouble
end

immutable Cstats
    pcost::Cdouble
    dcost::Cdouble
    pres::Cdouble
    dres::Cdouble
    pinf::Cdouble
    dinf::Cdouble
    pinfres::Cdouble
    dinfres::Cdouble
    gap::Cdouble
    relgap::Cdouble
    sigma::Cdouble
    mu::Cdouble
    step::Cdouble
    step_aff::Cdouble
    kapovert::Cdouble
    iter::Clong
    nitref1::Clong
    nitref2::Clong
    nitref3::Clong
    tsetup::Cdouble
    tsolve::Cdouble
    tfactor::Cdouble
    tkktsolve::Cdouble
    torder::Cdouble
    tkktcreate::Cdouble
    ttranspose::Cdouble
    tperm::Cdouble
    tfactor_t1::Cdouble
    tfactor_t2::Cdouble
end

immutable Csettings
    gamma::Cdouble
    delta::Cdouble
    eps::Cdouble
    feastol::Cdouble
    abstol::Cdouble
    reltol::Cdouble
    feastol_inacc::Cdouble
    abstol_inacc::Cdouble
    reltol_inacc::Cdouble
    nitref::Clong
    maxit::Clong
    verbose::Clong
end

immutable Cpwork
    # Dimensions
    n::Clong
    m::Clong
    p::Clong
    D::Clong

    # Variables
    x::Ptr{Cdouble}
    y::Ptr{Cdouble}
    z::Ptr{Cdouble}
    s::Ptr{Cdouble}
    lambda::Ptr{Cdouble}
    kap::Cdouble
    tau::Cdouble

    # Best iterates seen so far
    best_x::Ptr{Cdouble}
    best_y::Ptr{Cdouble}
    best_z::Ptr{Cdouble}
    best_s::Ptr{Cdouble}
    best_kap::Cdouble
    best_tau::Cdouble
    best_cx::Cdouble
    best_by::Cdouble
    best_hz::Cdouble
    best_info::Ptr{Cstats}

    # Temporary variables
    dsaff::Ptr{Cdouble}
    dzaff::Ptr{Cdouble}
    W_times_dzaff::Ptr{Cdouble}
    dsaff_by_W::Ptr{Cdouble}
    saff::Ptr{Cdouble}
    zaff::Ptr{Cdouble}

    # Cone
    C::Ptr{Ccone}
    A::Ptr{Cspmat}
    G::Ptr{Cspmat}
    c::Ptr{Cdouble}
    b::Ptr{Cdouble}
    h::Ptr{Cdouble}

    # equilibration vector
    xequil::Ptr{Cdouble}
    Aequil::Ptr{Cdouble}
    Gequil::Ptr{Cdouble}

    # scalings of problem data
    resx0::Cdouble
    resy0::Cdouble
    resz0::Cdouble

    # residuals
    rx::Ptr{Cdouble}
    ry::Ptr{Cdouble}
    rz::Ptr{Cdouble}
    rt::Cdouble
    hresx::Cdouble
    hresy::Cdouble
    hresz::Cdouble

    # temporary storage
    cx::Cdouble
    by::Cdouble
    hz::Cdouble
    sz::Cdouble

    # KKT System
    KKT::Ptr{Ckkt}

    # info struct
    info::Ptr{Cstats}

    # settings struct
    stgs::Ptr{Csettings}
end


module UnitfulNLopt

using Unitful
import Unitful: Quantity, Units
import NLopt
import NLopt: Opt

export UnitfulOpt

struct UnitfulOpt
    opt::Opt
    xunit::Units
    funit::Units
end

function UnitfulOpt(s::Symbol, N::Int, xu::Units, fu::Units)
    UnitfulOpt(Opt(s, N), xu, fu)
end

function NLopt.min_objective!(uopt::UnitfulOpt, f)
    g(x,y) = ustrip(f(x*uopt.xunit, y*uopt.funit*uopt.xunit^-1))
    NLopt.min_objective!(uopt.opt, g)
end

for f in (:equality_constraint!, :inequality_constraint!)
    @eval function NLopt.$f(uopt::UnitfulOpt, f, tol::Real)
        g(x,y) = ustrip(f(x*uopt.xunit, y*uopt.funit*uopt.xunit^-1))
        NLopt.$f(uopt.opt, g, tol)
    end
end

for f in (:optimize, :optimize!)
    @eval function NLopt.$f(uopt::UnitfulOpt, x::AbstractVector)
        fmin, xmin, res = NLopt.$f(uopt.opt, ustrip(x))
        fmin*uopt.funit, xmin*uopt.xunit, res
    end
end

for f in (:algorithm,)
    @eval function NLopt.$f(uopt::UnitfulOpt)
        NLopt.$f(uopt.opt)
    end
end

fset(f::Symbol) = Symbol(f, '!')

for f in (:ftol_rel, :xtol_rel, :maxeval, :maxtime)
    @eval function NLopt.$f(uopt::UnitfulOpt)
        NLopt.$f(uopt.opt)
    end
    @eval function NLopt.$(fset(f))(uopt::UnitfulOpt, x)
        NLopt.$(fset(f))(uopt.opt, x)
    end
end

for f in (:stopval, :ftol_abs)
    @eval function NLopt.$f(uopt::UnitfulOpt)
        NLopt.$f(uopt.opt) * uopt.funit
    end
    @eval function NLopt.$(fset(f))(uopt::UnitfulOpt, x)
        NLopt.$(fset(f))(uopt.opt, ustrip(uconvert(uopt.funit, x)))
    end
end

for f in (:xtol_abs, :lower_bounds, :upper_bounds, :initial_step)
    @eval function NLopt.$f(uopt::UnitfulOpt)
        NLopt.$f(uopt.opt) * uopt.xunit
    end
    @eval function NLopt.$(fset(f))(uopt::UnitfulOpt, x)
        NLopt.$(fset(f))(uopt.opt, ustrip(uconvert(uopt.xunit, x)))
    end
end

end # module


immutable PolygonGeometry <: Gadfly.GeometryElement
    order::Int
    fill::Bool
    preserve_order::Bool

    function PolygonGeometry(; order::Int=0, fill::Bool=false,
                               preserve_order::Bool=false)
        return new(order, fill, preserve_order)
    end
end

const polygon = PolygonGeometry

function element_aesthetics(::PolygonGeometry)
    return [:x, :y, :color, :group]
end

function polygon_points(xs, ys, preserve_order)
    T = (@compat Tuple{eltype(xs), eltype(ys)})
    if preserve_order
        return T[(x, y)for (x, y) in zip(xs, ys)]
    else
        centroid_x, centroid_y = mean(xs), mean(ys)
        θ = atan2(xs - centroid_x, ys - centroid_y)
        perm = sortperm(θ)
        return T[(x, y)for (x, y) in zip(xs[perm], ys[perm])]
    end
end

# Render polygon geometry.
function render(geom::PolygonGeometry, theme::Gadfly.Theme,
                aes::Gadfly.Aesthetics)
    Gadfly.assert_aesthetics_defined("Geom.polygon", aes, :x, :y)

    default_aes = Gadfly.Aesthetics()
    default_aes.color = PooledDataArray(RGBA{Float32}[theme.default_color])
    aes = inherit(aes, default_aes)

    ctx = context(order=geom.order)
    T = (eltype(aes.x), eltype(aes.y))

    if aes.group != nothing
        XT, YT = eltype(aes.x), eltype(aes.y)
        xs = DefaultDict(Any, Vector{XT}, () -> XT[])
        ys = DefaultDict(Any, Vector{YT}, () -> YT[])
        for (x, y, c, g) in zip(aes.x, aes.y, cycle(aes.color), cycle(aes.group))
            push!(xs[(c,g)], x)
            push!(ys[(c,g)], y)
        end

        compose!(ctx,
            Compose.polygon([polygon_points(xs[(c,g)], ys[(c,g)], geom.preserve_order)
                             for (c,g) in keys(xs)]))
        cs = [c for (c,g) in keys(xs)]
        if geom.fill
            compose!(ctx, fill(cs),
                     stroke(map(theme.discrete_highlight_color, cs)))
        else
            compose!(ctx, fill(nothing), stroke(cs))
        end
    elseif length(aes.color) == 1 &&
            !(isa(aes.color, PooledDataArray) && length(levels(aes.color)) > 1)
        compose!(ctx, Compose.polygon(polygon_points(aes.x, aes.y, geom.preserve_order)))
        if geom.fill
            compose!(ctx, fill(aes.color[1]),
                     stroke(theme.discrete_highlight_color(aes.color[1])))
        else
            compose!(ctx, fill(nothing),
                     stroke(aes.color[1]))
        end
    else
        XT, YT = eltype(aes.x), eltype(aes.y)
        xs = DefaultDict(ColorValue, Vector{XT}, () -> XT[])
        ys = DefaultDict(ColorValue, Vector{YT}, () -> YT[])
        for (x, y, c) in zip(aes.x, aes.y, cycle(aes.color))
            push!(xs[c], x)
            push!(ys[c], y)
        end

        compose!(ctx,
            Compose.polygon([polygon_points(xs[c], ys[c], geom.preserve_order)
                             for c in keys(xs)]))
        cs = collect(keys(xs))
        if geom.fill
            compose!(ctx, fill(cs),
                     stroke(map(theme.discrete_highlight_color, cs)))
        else
            compose!(ctx, fill(nothing), stroke(cs))
        end
    end

    return compose!(ctx, linewidth(theme.line_width), svgclass("geometry"))
end