# # Routines for extracting segment/face data for tracks (rays) overlaid on a mesh
#
# # Return the HRPM/face indices in which each segment resides
# # Type-stable, other than warning
# function find_segment_faces(segment_points::Vector{Vector{Vector{Point_2D{F}}}},
#                             HRPM::HierarchicalRectangularlyPartitionedMesh{F, U},
#                             template_vec::MVector{N, U}
#                            ) where {F <: AbstractFloat, U <: Unsigned, N}
#
#     @info "    - Finding faces for each segment"
#     if !has_materialized_faces(HRPM)
#         @warn "Faces are not materialized for this mesh. This will be VERY slow"
#     end
#     nγ = length(segment_points)
#     bools = fill(false, nγ)
#     # Preallocate indices in the most frustrating way
#     indices =   [
#                     [
#                         [
#                             MVector{N, U}(zeros(U, N))
#                                 for i = 1:length(segment_points[iγ][it])-1 # Segments
#                         ] for it = 1:length(segment_points[iγ]) # Tracks
#                     ] for iγ = 1:nγ # Angles
#                 ]
#     Threads.@threads for iγ = 1:nγ
#         bools[iγ] = find_segment_faces_in_angle!(segment_points[iγ], indices[iγ], HRPM)
#     end
#     if !all(bools)
#         iγ_bad = findall(x->!x, bools)
#         @error "Failed to find indices for some points in segment_points$iγ_bad"
#     end
#     return indices
# end
#
# Return the face id in which each segment resides
function find_segment_faces(segment_points::Vector{Vector{Vector{Point_2D}}}, mesh::UnstructuredMesh_2D)
    @info "    - Finding faces for each segment"
    if 0 == length(mesh.materialized_faces)
        @warn "Faces are not materialized for this mesh. This will be VERY slow"
    end
    nγ = length(segment_points)
    # Preallocate indices in the most frustrating way
    segment_faces =   [
                          [
                            fill(UInt32(0), length(segment_points[iγ][it])-1) # Segments
                            for it = 1:length(segment_points[iγ]) # Tracks
                          ] for iγ = 1:nγ # Angles
                      ]
    Threads.@threads for iγ = 1:nγ
        find_segment_faces_in_angle!(segment_points[iγ], segment_faces[iγ], mesh)
        for it = 1:length(segment_points[iγ])
            if any(x->x == 0x00000000, segment_faces[iγ][it])
                @warn "Segment face not found for face in [$iγ][$it]"
            end
        end
    end
    return segment_faces
end

# # Ray trace an HRPM given the ray spacing and angular quadrature
# # Not type-stable
# function ray_trace(tₛ::F,
#                    ang_quad::ProductAngularQuadrature{nᵧ, nₚ, F},
#                    HRPM::HierarchicalRectangularlyPartitionedMesh{F, U}
#                    ) where {nᵧ, nₚ, F <: AbstractFloat, U <: Unsigned}
#     @info "Ray tracing"
#     tracks = generate_tracks(tₛ, ang_quad, HRPM)
#     segment_points = segmentize(tracks, HRPM)
#     ind_size = node_height(HRPM) + 1
#     @info "  - Using the naive segmentize + find face algorithm"
#     template_vec = MVector{ind_size, U}(zeros(U, ind_size))
#     segment_faces = find_segment_faces(segment_points, HRPM, template_vec)
#     return segment_points, segment_faces
# end
#
# Ray trace a mesh given the ray spacing and angular quadrature
function ray_trace(tₛ::Float64,
                   ang_quad::ProductAngularQuadrature{nᵧ, nₚ},
                   mesh::UnstructuredMesh_2D) where {nᵧ, nₚ}
    @info "Ray tracing"
    tracks = generate_tracks(tₛ, ang_quad, mesh, boundary_shape = "Rectangle")
    if use_E2E_raytracing(mesh)
        @info "  - Using the edge-to-edge algorithm"
        segment_points, segment_faces = ray_trace_edge_to_edge(tracks, mesh)
        validate_ray_tracing_data(segment_points, segment_faces, mesh, plot = enable_visualization)
        return segment_points, segment_faces
    else
        @info "  - Using the naive segmentize + find face algorithm"
        segment_points = segmentize(tracks, mesh)
        segment_faces = find_segment_faces(segment_points, mesh)
        validate_ray_tracing_data(segment_points, segment_faces, mesh, plot = enable_visualization)
        return segment_points, segment_faces
    end
end

# Get the segment points and the face which the segment lies in for all segments,
# in all tracks, in all angles, using the edge-to-edge ray tracing method.
# Assumes a rectangular boundary
function ray_trace_edge_to_edge(tracks::Vector{Vector{LineSegment_2D}}, mesh::UnstructuredMesh_2D)
    if visualize_ray_tracing
        @error "visualize_ray_tracing = true. Please reset to false in constants.jl"
    end
    if length(mesh.boundary_edges) != 4
        @error "Mesh does not have 4 boundary edges needed for edge-to-edge ray tracing!"
    end
    # index 1 = γ
    # index 2 = track
    # index 3 = point/segment
    nγ = length(tracks)
    segment_points =[
                        [
                            Point_2D[] for it = 1:length(tracks[iγ]) # Tracks
                        ] for iγ = 1:nγ # Angles
                    ]
    segment_faces = [
                        [
                            UInt32[] for it = 1:length(tracks[iγ]) # Tracks
                        ] for iγ = 1:nγ # Angles
                    ]
    # For each angle, get the segments and segment faces for each track
    Threads.@threads for iγ = 1:nγ
        ray_trace_angle_edge_to_edge!(tracks[iγ],
                                      segment_points[iγ],
                                      segment_faces[iγ],
                                      mesh)
    end
    return segment_points, segment_faces
end

# # Return the points of intersection between the tracks and the mesh edges.
# # Returns a Vector{Vector{Vector{Point_2D{F}}}}.
# #   index 1 = γ
# #   index 2 = track
# #   index 3 = point/segment
# # Type-stable other than the error message
# function segmentize(tracks::Vector{Vector{LineSegment_2D{F}}},
#                     HRPM::HierarchicalRectangularlyPartitionedMesh{F, U}
#                     ) where {F <: AbstractFloat, U <: Unsigned}
#
#     # Give info about intersection algorithm being used
#     int_alg = get_intersection_algorithm(HRPM)
#     @info "    - Segmentizing using the '$int_alg' algorithm"
#     # index 1 = γ
#     # index 2 = track
#     # index 3 = point/segment
#     nγ = length(tracks)
#     segment_points = Vector{Vector{Vector{Point_2D{F}}}}(undef, nγ)
#     Threads.@threads for iγ = 1:nγ
#         # for each track, intersect the track with the mesh
#         segment_points[iγ] = tracks[iγ] .∩ HRPM
#     end
#     return segment_points
# end
#
# Return the points of intersection between the tracks and the mesh edges.
# Returns a Vector{Vector{Vector{Point_2D{F}}}}.
#   index 1 = γ
#   index 2 = track
#   index 3 = point/segment
function segmentize(tracks::Vector{Vector{LineSegment_2D}}, mesh::UnstructuredMesh_2D)
    # Give info about intersection algorithm being used
    # index 1 = γ
    # index 2 = track
    # index 3 = point/segment
    nγ = length(tracks)
    segment_points = Vector{Vector{Vector{Point_2D}}}(undef, nγ)
    if length(mesh.materialized_edges) !== 0
        @info "    - Segmentizing using the 'Edges - Explicit' algorithm"
        Threads.@threads for iγ = 1:nγ
            segment_points[iγ] = intersect_edges_explicit.(tracks[iγ], Ref(mesh.materialized_edges))
        end
    elseif length(mesh.edges) !== 0
        @info "    - Segmentizing using the 'Edges - Implicit' algorithm"
        Threads.@threads for iγ = 1:nγ
            segment_points[iγ] = intersect_edges_implicit.(tracks[iγ], Ref(mesh.edges), Ref(mesh.points))
        end
    elseif length(mesh.materialized_faces) !== 0
        @info "    - Segmentizing using the 'Faces - Explicit' algorithm"
        Threads.@threads for iγ = 1:nγ
            segment_points[iγ] = intersect_faces_explicit.(tracks[iγ], Ref(mesh.materialized_faces))
        end
    else
        @info "    - Segmentizing using the 'Faces - Implicit' algorithm"
        Threads.@threads for iγ = 1:nγ
            segment_points[iγ] = intersect_faces_implicit.(tracks[iγ], Ref(mesh.faces), Ref(mesh.points))
        end
    end
    return segment_points
end

# Assumes materialized faces
function validate_ray_tracing_data(segment_points::Vector{Vector{Vector{Point_2D}}},
                                   segment_faces::Vector{Vector{Vector{UInt32}}},
                                   mesh::UnstructuredMesh_2D;
                                   plot::Bool = false)
    @info "  - Validating ray tracing data"
    nthreads = Threads.nthreads()
    nsegs = zeros(Int64, nthreads)
    nsegs_problem = zeros(Int64, nthreads) # Problem segment if the face doesn't match, and it's over 10 μm
    problem_indices = [ SVector{3, Int64}[] for i = 1:nthreads ]
    plot_segs_face = [ LineSegment_2D[] for i = 1:nthreads ]
    plot_points_face = [ Point_2D[] for i = 1:nthreads ]

    # Validate faces
    Threads.@threads for iγ = 1:length(segment_faces)
        for it = 1:length(segment_faces[iγ])
            nsegs[Threads.threadid()] += length(segment_faces[iγ][it])
            for iseg = 1:length(segment_faces[iγ][it])
                p_midpoint = midpoint(segment_points[iγ][it][iseg], segment_points[iγ][it][iseg+1])
                face = segment_faces[iγ][it][iseg]
                if p_midpoint ∉  mesh.materialized_faces[face]
                    p1 = segment_points[iγ][it][iseg]
                    p2 = segment_points[iγ][it][iseg + 1]
                    l = LineSegment_2D(p1, p2)
                    problem_length = 1e-3 < arclength(l)
                    if problem_length
                        nsegs_problem[Threads.threadid()] += 1
                        push!(problem_indices[Threads.threadid()], SVector(iγ, it, iseg))
                    end
                    # Append the points and line if we want to plot them.
                    # We only want to plot if actually a problem.
                    if enable_visualization && plot && problem_length
                        append!(plot_points_face[Threads.threadid()], [p1, p2])
                        push!(plot_segs_face[Threads.threadid()], l)
                    end
                end
            end
        end
    end

    # Attempt to fix problem segments if used E2E ray tracing
    if use_E2E_raytracing(mesh)
        @info "    - Attempting to fix $(sum(nsegs_problem)) segments with incorrect faces"
        fixed = [ Int64[] for i = 1:nthreads ]
        Threads.@threads for i = 1:nthreads
            for (iprob, problem_index) in enumerate(problem_indices[i])
                # Ray trace the track in reverse, then test to see if that fixed things
                new_info_ok = true
                iγ, it, iseg = problem_index
                p_midpoint = midpoint(segment_points[iγ][it][iseg], segment_points[iγ][it][iseg+1])
                face = segment_faces[iγ][it][iseg]
                # Check to see if still a problem. May have been fixed by a previous segment updating the
                # points/faces
                if p_midpoint ∈  mesh.materialized_faces[face]
                    nsegs_problem[i] -= 1
                    continue
                end
                npoints = length(segment_points[iγ][it])
                problem_line = LineSegment_2D(segment_points[iγ][it][1], segment_points[iγ][it][npoints])
                line_reversed = LineSegment_2D(problem_line.points[2], problem_line.points[1])
                reversed_points, reversed_faces = ray_trace_track_edge_to_edge(line_reversed,  mesh)
                # Check
                for iseg = 1:length(reversed_faces)
                    p_midpoint = midpoint(reversed_points[iseg], reversed_points[iseg+1])
                    face = reversed_faces[iseg]
                    if p_midpoint ∉  mesh.materialized_faces[face]
                        l = LineSegment_2D(reversed_points[iseg], reversed_points[iseg + 1])
                        problem_length = 1e-3 < arclength(l)
                        if problem_length
                            new_info_ok = false
                        end
                    end
                end
                if new_info_ok
                    push!(fixed[i], iprob)
                    nsegs_problem[i] -= 1
                    segment_points[iγ][it] = reverse(reversed_points)
                    segment_faces[iγ][it] = reverse(reversed_faces)
                else
                    new_face = findface(p_midpoint, mesh)
                    if new_face == 0
                        @warn "Could not find new segment face for segment $(problem_index[i])"
                    else
                        segment_faces[iγ][it][iseg] = new_face
                    end
                end
            end
        end
    end

    # Visualize
    problem_segs = sum(nsegs_problem)
    if enable_visualization && plot && 0 < problem_segs
        f = Figure()
        ax = Axis(f[1, 1], aspect = 1)
        display(f)
        mesh!(mesh.materialized_faces, color = (:black, 0.15))
        if 0 < length(mesh.materialized_edges)
            linesegments!(mesh.materialized_edges, color = :blue)
        else
            linesegments!(materialize_edges(mesh), color = :blue)
        end
        for i = 1:nthreads
            if 0 < length(fixed[i])
                deleteat!(plot_segs_face[i], fixed[i])
            end
            if 0 < length(plot_segs_face[i])
                linesegments!(plot_segs_face[i], color = :red)
            end
            if 0 < length(fixed[i])
                deleteat!(plot_points_face[i], sort!(vcat(2 .* fixed[i], 2 .* fixed[i] .- 1)))
            end
            if 0 < length(plot_points_face[i])
                scatter!(plot_points_face[i], color = :red)
            end
        end
    end
    nsegs_total = sum(nsegs)
    prob_percent = 100*problem_segs/nsegs_total
    @info "    - Problem segments: $problem_segs, Total segments: $nsegs_total"
    if 0 < problem_segs
        @info "      - Check for mesh overlap (darker gray) around areas of many problem segments."
    end
    return problem_segs == 0
end

# Plot
# -------------------------------------------------------------------------------------------------
# Plot ray tracing data one angle at a time.
if enable_visualization
#    function linesegments!(segment_points::Vector{Vector{Vector{Point_2D}}},
#                           segment_faces::Vector{Vector{Vector{UInt32}}})
#        println("Press enter to plot the segments in the next angle")
#        colormap = ColorSchemes.tab20.colors
#        lines_by_color = Vector{Vector{LineSegment_2D{T}}}(undef, 20)
#        nγ = length(segment_points)
#        for iγ = 1:nγ
#            f = Figure()
#            ax = Axis(f[1, 1], aspect = 1)
#            display(f)
#            for icolor = 1:20
#                lines_by_color[icolor] = LineSegment_2D{T}[]
#            end
#            for it = 1:length(segment_points[iγ])
#                for iseg = 1:length(segment_points[iγ][it])-1
#                    l = LineSegment_2D(segment_points[iγ][it][iseg], segment_points[iγ][it][iseg+1])
#                    face = segment_faces[iγ][it][iseg]
#                    if face == 0
#                        @error "Segment [$iγ][$it][$iseg] has a face id of 0"
#                    end
#                    push!(lines_by_color[face % 20 + 1], l)
#                end
#            end
#            for icolor = 1:20
#                linesegments!(lines_by_color[icolor], color = colormap[icolor])
#            end
#            s = readline()
#            println(iγ)
#        end
#    end

    # Set visualize_ray_tracing = true in constants.jl to get this to work.
    function plot_track_edge_to_edge(track::LineSegment_2D, mesh::UnstructuredMesh_2D)
        @info "Plotting ray tracing of track. Press enter to advance the ray"
        f = Figure()
        ax = Axis(f[1, 1], aspect = 1)
        display(f)
        linesegments!(mesh.materialized_edges, color = :blue)
        linesegments!(track, color = :orange)
        segment_points, segment_faces = ray_trace_track_edge_to_edge(track, mesh)
        return segment_points, segment_faces
    end
end


function NaivePDE.periodic_makes_boundary!{T1 <: Number,n}(
  x:: Array{T1,n}, p :: MPI_areaConfig)
  return NaivePDE.periodic_makes_boundary!(x, p.OL)
end

function periodic_expand_boundary{T1 <: Number,n}(
  xcore:: Array{T1,n}, p :: MPI_areaConfig)
  return NaivePDE.periodic_expand_boundary(xcore, p.OL)
end

function NaivePDE.get_core{T1 <: Number,n}(
  x:: Array{T1,n}, p :: MPI_areaConfig)
  return NaivePDE.get_core(x, p.OL)
end

function NaivePDE.NaN_expand{T1 <: Number,n}(
  xcore:: Array{T1,n}, p :: MPI_areaConfig)
  return NaivePDE.NaN_expand(xcore, p.OL)
end


# Make jump variables and a dict mapping hold vars to jump vars
function _initialize_hold_variables(trans_model::JuMP.Model,
                                      inf_model::InfiniteOpt.InfiniteModel)
    # search inf_model for hold vars and make a jump var for one that is used
    for index in sort(collect(keys(inf_model.vars)))
        var = inf_model.vars[index]
        if isa(var, InfiniteOpt.HoldVariable)
            hvref = InfiniteOpt.HoldVariableRef(inf_model, index)
            if is_used(hvref)
                vref = JuMP.add_variable(trans_model,
                                         JuMP.ScalarVariable(var.info),
                                         JuMP.name(hvref))
                transcription_data(trans_model).hold_to_var[hvref] = vref
            end
        end
    end
    return
end

# Return a vector of arrays containing the supports
function _list_supports(prefs::Tuple)::Vector
    return [InfiniteOpt.supports(pref) for pref in prefs]
end

# Make an index mapping for parameter support combinations for an infinite variable
function _make_supports(prefs::Tuple)::Vector
    support_list = _list_supports(prefs)
    combos = Iterators.product(support_list...)
    return [combo for combo in Iterators.take(combos, length(combos))]
end

# Make jump variables and a dict mapping infinite/point vars to jump vars
function _initialize_infinite_variables(trans_model::JuMP.Model,
                                        inf_model::InfiniteOpt.InfiniteModel)
    # search inf_model for infinite vars and make a jump var for all of its supports
    for index in sort(collect(keys(inf_model.vars)))
        var = inf_model.vars[index]
        if isa(var, InfiniteOpt.InfiniteVariable)
            ivref = InfiniteOpt.InfiniteVariableRef(inf_model, index)
            if InfiniteOpt.is_used(ivref)
                prefs = InfiniteOpt.parameter_refs(ivref)
                supports = _make_supports(prefs)
                transcription_data(trans_model).infvar_to_supports[ivref] = supports
                vrefs = Vector{JuMP.VariableRef}(undef, length(supports))
                name = InfiniteOpt._root_name(ivref)
                for i in eachindex(vrefs)
                    # TODO Perhaps add different naming options...
                    vrefs[i] = JuMP.add_variable(trans_model,
                                                 JuMP.ScalarVariable(var.info),
                                                 string(name, "(support: ", i, ")"))
                end
                transcription_data(trans_model).infinite_to_vars[ivref] = vrefs
            end
        end
    end
    return
end

# Map the point variable reference to a transcribed variable based off of the support
function _update_point_mapping(trans_model::JuMP.Model,
                               pvref::InfiniteOpt.PointVariableRef,
                               ivref::InfiniteOpt.InfiniteVariableRef,
                               support::Tuple)
    supps = InfiniteOpt.variable_supports(trans_model, ivref)
    # search for the support and update mapping
    for i in eachindex(supps)
       if all(isapprox.(support, supps[i]))
           vrefs = transcription_variable(trans_model, ivref)
           transcription_data(trans_model).point_to_var[pvref] = vrefs[i]
           return
       end
    end
    # fallback that should not be needed
    error("Couldn't find variable to map $pvref to.")
end

# Override the info of the jump variable with the point variable's if any is provided
function _update_point_info(trans_model::JuMP.Model,
                            pvref::InfiniteOpt.PointVariableRef)
    vref = transcription_variable(trans_model, pvref)
    if JuMP.has_lower_bound(pvref) && !JuMP.has_lower_bound(vref)
        if JuMP.is_fixed(vref)
            JuMP.unfix(vref)
        end
        JuMP.set_lower_bound(vref, JuMP.lower_bound(pvref))
    end
    if JuMP.has_upper_bound(pvref) && !JuMP.has_upper_bound(vref)
        if JuMP.is_fixed(vref)
            JuMP.unfix(vref)
        end
        JuMP.set_upper_bound(vref, JuMP.upper_bound(pvref))
    end
    if JuMP.is_fixed(pvref) && !JuMP.is_fixed(vref)
        JuMP.fix(vref, JuMP.fix_value(pvref), force = true)
    end
    if JuMP.is_binary(pvref) && !JuMP.is_binary(vref)
        if JuMP.is_integer(vref)
            JuMP.unset_integer(vref)
        end
        JuMP.set_binary(vref)
    elseif JuMP.is_integer(pvref) && !JuMP.is_integer(vref)
        if JuMP.is_binary(vref)
            JuMP.unset_binary(vref)
        end
        JuMP.set_integer(vref)
    end
    if JuMP.start_value(pvref) != JuMP.start_value(vref)
        JuMP.set_start_value(vref, JuMP.start_value(pvref))
    end
    return
end

# Map point variables to the correct transcribed infinite variable
function _map_point_variables(trans_model::JuMP.Model,
                              inf_model::InfiniteOpt.InfiniteModel)
    # search inf_model for point vars and map them to jump vars if they are used
    for (index, var) in inf_model.vars
        if isa(var, InfiniteOpt.PointVariable)
            pvref = InfiniteOpt.PointVariableRef(inf_model, index)
            if InfiniteOpt.is_used(pvref)
                ivref = InfiniteOpt.infinite_variable_ref(pvref)
                support = InfiniteOpt.parameter_values(pvref)
                _update_point_mapping(trans_model, pvref, ivref, support)
                _update_point_info(trans_model, pvref)
            end
        end
    end
    return
end

# Return the tuple index of a parameter based off of group_id
function _parameter_tuple_index(pref::InfiniteOpt.ParameterRef,
                                prefs::Tuple)::Int
    group = InfiniteOpt.group_id(pref)
    groups = InfiniteOpt._group.(prefs)
    if !(group in groups)
        return -1
    else
        return findfirst(isequal(group), groups)
    end
end

# Return the support value corresponding to a parameter reference
# return NaN is the parameter is not contained in prefs
function _parameter_value(pref::InfiniteOpt.ParameterRef, support::Tuple,
                          prefs::Tuple)::Float64
    pref_index = _parameter_tuple_index(pref, prefs)
    if pref_index == -1
        return NaN
    elseif isa(prefs[pref_index], InfiniteOpt.ParameterRef)
        return support[pref_index]
    else
        for (k, v) in prefs[pref_index].data
            if v == pref
                return support[pref_index].data[k]
            end
        end
    end
    return NaN
end

## Helper function for mapping InfiniteOpt variables to jump variables
# FiniteVariableRef
function _map_to_variable(fvref::InfiniteOpt.FiniteVariableRef, support::Tuple,
                          prefs::Tuple, trans_model::JuMP.Model)::JuMP.VariableRef
    return transcription_variable(trans_model, fvref)
end

# InfiniteVariableRef
function _map_to_variable(ivref::InfiniteOpt.InfiniteVariableRef, support::Tuple,
                          prefs::Tuple, trans_model::JuMP.Model)::JuMP.VariableRef
    # reduce support to only include the relavent parameter id groups
    ivref_groups = InfiniteOpt._group.(InfiniteOpt.parameter_refs(ivref))
    support_groups = InfiniteOpt._group.(prefs)
    reduced_support = Tuple(support[findfirst(isequal(group), support_groups)]
                            for group in ivref_groups)
    # find the jump variable associated with the support
    supps = InfiniteOpt.variable_supports(trans_model, ivref)
    for i in eachindex(supps)
        if all(isapprox.(reduced_support, supps[i]))
            return transcription_variable(trans_model, ivref)[i]
        end
    end
    # this shouldn't be needed, but provided as a backup
    error("Couldn't find JuMP variable corresponding to $ivref.")
end

# ReducedInfiniteVariableRef
function _map_to_variable(rvref::InfiniteOpt.ReducedInfiniteVariableRef,
                          support::Tuple, prefs::Tuple,
                          trans_model::JuMP.Model)::JuMP.VariableRef
    # parse the reduced parameters and modify the support to include them
    ivref = InfiniteOpt.infinite_variable_ref(rvref)
    orig_groups = InfiniteOpt._group.(InfiniteOpt.parameter_refs(ivref))
    support_groups = InfiniteOpt._group.(prefs)
    support = Any[i for i in support]
    prefs = Any[pref for pref in prefs]
    for (index, value) in InfiniteOpt.eval_supports(rvref)
        if orig_groups[index] in support_groups
            support[findfirst(isequal(orig_groups[index]), support_groups)] = value
        else
            push!(support, value)
            push!(prefs, InfiniteOpt.parameter_refs(ivref)[index])
        end
    end
    prefs = Tuple(pref for pref in prefs)
    support = Tuple(i for i in support)
    # call the infinite variable function with the updated support and prefs
    return _map_to_variable(ivref, support, prefs, trans_model)
end

# ParameterRef
function _map_to_variable(pref::InfiniteOpt.ParameterRef, support::Tuple,
                          prefs::Tuple, trans_model::JuMP.Model)::Float64
    # find pref in prefs and return associated support value
    value = _parameter_value(pref, support, prefs)
    # this error is a backup that should not be needed
    value === NaN && error("Couldn't find support corresponding to $pref.")
    return value
end

# Return a Bool whether the support satisfies the bounds or not
function _support_in_bounds(support::Tuple, prefs::Tuple, bounds::Dict)::Bool
    for (pref, set) in bounds
        value = _parameter_value(pref, support, prefs)
        if value === NaN
            continue
        end
        if value < set.lower_bound || value > set.upper_bound
            return false
        end
    end
    return true
end

# Return a bool array indicating which supports are in bounds
function _supports_in_bounds(supports::Vector, prefs::Tuple, bounds::Dict)::Vector
    return [_support_in_bounds(supports[i], prefs, bounds)
            for i in eachindex(supports)]
end

## Convert jump scalar expressions with InfiniteOpt variables into transcribed relations
# PointVariableRef and HoldVariableRef --> return scalar jump object
function _make_transcription_function(vref::InfiniteOpt.FiniteVariableRef,
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple
    return transcription_variable(trans_model, vref), ()
end

# InfiniteVariableRef --> return tuple of expressions, prefs, and support mapping
function _make_transcription_function(vref::InfiniteOpt.InfiniteVariableRef,
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple
    # Update the supports to include only those in the bounds and return vars
    if length(bounds) != 0
        supports = InfiniteOpt.variable_supports(trans_model, vref)
        prefs = InfiniteOpt.parameter_refs(vref)
        old_support_indices = _supports_in_bounds(supports, prefs, bounds)
        new_supports = supports[old_support_indices]
        return transcription_variable(trans_model, vref)[old_support_indices],
               InfiniteOpt.parameter_refs(vref), new_supports
    # easy case just return the info
    else
        return transcription_variable(trans_model, vref),
               InfiniteOpt.parameter_refs(vref),
               InfiniteOpt.variable_supports(trans_model, vref)
    end
end

# ParameterRef --> return tuple of numbers, pref, and support mapping
function _make_transcription_function(pref::InfiniteOpt.ParameterRef,
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple
    supports = unique(InfiniteOpt.supports(pref))
    # truncate the supports if out of bounds
    if pref in keys(bounds)
        old_support_indices = bounds[pref].lower_bound .<= supports .<= bounds[pref].upper_bound
        new_supports = supports[old_support_indices]
        return new_supports, (pref, ), new_supports
    # easy case just return the supports
    else
        return supports, (pref, ), supports
    end
end

# MeasureRef --> return depends if finite or infinite
function _make_transcription_function(mref::InfiniteOpt.MeasureRef,
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple
    # expand the measure
    func = InfiniteOpt.measure_function(mref)
    data = InfiniteOpt.measure_data(mref)
    new_func = InfiniteOpt._possible_convert(InfiniteOpt.FiniteVariableRef,
                                         InfiniteOpt.expand_measure(func, data,
                                            trans_model, _update_point_mapping))
    # will either call finite variable function or general variable function
    return _make_transcription_function(new_func, trans_model, bounds)
end

# GenericAffExpr of FiniteVariableRefs --> return scalar jump object
function _make_transcription_function(expr::JuMP.GenericAffExpr{C,
                                               <:InfiniteOpt.FiniteVariableRef},
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple where {C}
    # replace finite vars with jump vars
    pairs = [transcription_variable(trans_model, var) => coef
             for (var, coef) in expr.terms]
    return JuMP.GenericAffExpr(expr.constant,
                         JuMP._new_ordered_dict(JuMP.VariableRef, C, pairs)), ()
end

# GenericQuadExpr of FiniteVariableRefs --> return scalar jump object
function _make_transcription_function(expr::JuMP.GenericQuadExpr{C,
                                               <:InfiniteOpt.FiniteVariableRef},
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple where {C}
    # replace finite vars with jump vars
    pairs = Vector{Pair{JuMP.UnorderedPair{JuMP.VariableRef}, C}}(undef,
                                                             length(expr.terms))
    counter = 1
    for k in keys(expr.terms)
        a = transcription_variable(trans_model, k.a)
        b = transcription_variable(trans_model, k.b)
        pairs[counter] = JuMP.UnorderedPair(a, b) => expr.terms[k]
        counter += 1
    end
    aff = _make_transcription_function(expr.aff, trans_model)[1]
    return JuMP.GenericQuadExpr(aff,
     JuMP._new_ordered_dict(JuMP.UnorderedPair{JuMP.VariableRef}, C, pairs)), ()
end

# GenericAffExpr of GeneralVariableRefs --> return tuple of numbers, pref,
# and support mapping
function _make_transcription_function(expr::JuMP.GenericAffExpr{C,
                                              <:InfiniteOpt.GeneralVariableRef},
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple where {C}
    # check if there is only 1 var to dispatch to that transcription method
    if length(expr.terms) == 1
        var = first(keys(expr.terms))
        results = _make_transcription_function(var, trans_model, bounds)
        # results[2] contains prefs if is infinite or it is empty otherwise
        if length(results[2]) != 0
            exprs = expr.terms[var] * results[1] .+ expr.constant
            return exprs, results[2], results[3]
        else
            return expr.terms[var] * results[1] + expr.constant, ()
        end
    end
    # check to see if there are measures and expand them
    # TODO rework paradigm so expressions don't need to be checked for measures
    has_measure = false
    for var in keys(expr.terms)
        if var isa InfiniteOpt.MeasureRef
            has_measure = true
            break
        end
    end
    if has_measure
        expr = InfiniteOpt._expand_measures(expr, trans_model,
                                            _update_point_mapping)
    end
    # dispatch to quadratic method if the measures contained quadratic terms
    if isa(expr, JuMP.GenericQuadExpr)
        return _make_transcription_function(expr, trans_model, bounds)
    end
    # determine the common set of prefs and make all of the support combos
    prefs = InfiniteOpt._all_parameter_refs(expr)
    supports = _make_supports(prefs)
    if length(bounds) != 0
        old_support_indices = _supports_in_bounds(supports, prefs, bounds)
        supports = supports[old_support_indices]
    end
    exprs = [zero(JuMP.GenericAffExpr{C, JuMP.VariableRef})
             for i in eachindex(supports)]
    # make an expression for each support
    for i in eachindex(exprs)
        exprs[i].constant = expr.constant
        for (var, coef) in expr.terms
            # replace each variable with appropriate jump var
            new_var = _map_to_variable(var, supports[i], prefs,
                                       trans_model)
            if isa(new_var, JuMP.VariableRef)
                JuMP.add_to_expression!(exprs[i], coef, new_var)
            else
                JuMP.add_to_expression!(exprs[i], coef * new_var)
            end
        end
    end
    return exprs, prefs, supports
end

# GenericQuadExpr of GeneralVariableRefs
function _make_transcription_function(expr::JuMP.GenericQuadExpr{C,
                                              <:InfiniteOpt.GeneralVariableRef},
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple where {C}
    # check if there is only 1 var to dispatch to that transcription method
    if length(expr.terms) == 0 && length(expr.aff.terms) == 1
      var = first(keys(expr.aff.terms))
      results = _make_transcription_function(var, trans_model, bounds)
      # results[2] contains prefs if is infinite or it is empty otherwise
      if length(results[2]) != 0
          exprs = expr.aff.terms[var] * results[1] .+ expr.aff.constant
          return exprs, results[2], results[3]
      else
          return expr.aff.terms[var] * results[1] + expr.aff.constant, ()
      end
    end
    # check to see if there are measures and expand them
    # TODO rework paradigm so expressions don't need to be checked for measures
    all_variables = InfiniteOpt._all_function_variables(expr)
    has_measure = false
    for var in all_variables
        if var isa InfiniteOpt.MeasureRef
            has_measure = true
            break
        end
    end
    if has_measure
        # TODO expand measures in place by deleting entry and appending to the end
        # TODO thus the measure search can be done by expand measures
        expr = InfiniteOpt._expand_measures(expr, trans_model,
                                            _update_point_mapping)
    end
    # determine the common set of prefs and make all of the support combos
    prefs = InfiniteOpt._all_parameter_refs(expr)
    supports = _make_supports(prefs)
    if length(bounds) != 0
        old_support_indices = _supports_in_bounds(supports, prefs, bounds)
        supports = supports[old_support_indices]
    end
    exprs = [zero(JuMP.GenericQuadExpr{C, JuMP.VariableRef})
             for i in eachindex(supports)]
    # make an expression for each support
    for i in eachindex(exprs)
        exprs[i].aff.constant = expr.aff.constant
        for (var, coef) in expr.aff.terms
            # replace each variable with appropriate jump var
            new_var = _map_to_variable(var, supports[i], prefs,
                                       trans_model)
            if isa(new_var, JuMP.VariableRef)
                JuMP.add_to_expression!(exprs[i], coef, new_var)
            else
                JuMP.add_to_expression!(exprs[i], coef * new_var)
            end
        end
        for (pair, coef) in expr.terms
            # replace each variable with appropriate jump var
            var_a = _map_to_variable(pair.a, supports[i], prefs,
                                     trans_model)
            var_b = _map_to_variable(pair.b, supports[i], prefs,
                                     trans_model)
            if isa(var_a, JuMP.VariableRef) && isa(var_b, JuMP.VariableRef)
                JuMP.add_to_expression!(exprs[i], coef, var_a, var_b)
            elseif isa(var_b, JuMP.VariableRef)
                JuMP.add_to_expression!(exprs[i], coef * var_a, var_b)
            elseif isa(var_a, JuMP.VariableRef)
                JuMP.add_to_expression!(exprs[i], coef * var_b, var_a)
            else
                JuMP.add_to_expression!(exprs[i], coef * var_a * var_b)
            end
        end
    end
    # TODO Check if there are still quadratic terms and convert to AffExpr if necessary
    return exprs, prefs, supports
end

# GenericAffExpr and GenericQuadExpr of MeasureFiniteVariableRefs --> return
# depends on whether finite
function _make_transcription_function(expr::Union{JuMP.GenericAffExpr{C,
                                        <:InfiniteOpt.MeasureFiniteVariableRef},
                                        JuMP.GenericQuadExpr{C,
                                        <:InfiniteOpt.MeasureFiniteVariableRef}},
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple where {C}
    expr = InfiniteOpt._possible_convert(FiniteVariableRef,
                                         InfiniteOpt._expand_measures(expr,
                                         trans_model, _update_point_mapping))
    return _make_transcription_function(expr, trans_model, bounds)
end

# Empty jump variable expr (for constraints of form number <= number)
function _make_transcription_function(expr::JuMP.GenericAffExpr{C,
                                                              JuMP.VariableRef},
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple where {C}
    return expr, ()
end

# Fall back function for other jump objects
function _make_transcription_function(expr::JuMP.AbstractJuMPScalar,
                                      trans_model::JuMP.Model,
                                      bounds::Dict = Dict())::Tuple
    type = typeof(expr)
    error("Unsupported transcription of expression of type $type.")
end

## Construct the objective and error is contains non finite variables
function _set_objective(trans_model::JuMP.Model,
                        inf_model::InfiniteOpt.InfiniteModel)
    trans_obj, = _make_transcription_function(JuMP.objective_function(inf_model),
                                              trans_model)
    isa(trans_obj, Vector) && error("Objective is not finite, ensure all " *
                                    "infinite variables/parameters in measures " *
                                    "are evaluated completely.")
    JuMP.set_objective(trans_model, JuMP.objective_sense(inf_model), trans_obj)
    return
end

## Define helper functions for setting constraint mappings
# InfiniteConstraintRef
function _set_mapping(icref::InfiniteOpt.InfiniteConstraintRef,
                      crefs::Vector{<:JuMP.ConstraintRef})
    transcription_data(JuMP.owner_model(crefs[1])).infinite_to_constrs[icref] = crefs
    return
end

# MeasureConstraintRef (infinite)
function _set_mapping(mcref::InfiniteOpt.MeasureConstraintRef,
                      crefs::Vector{<:JuMP.ConstraintRef})
    transcription_data(JuMP.owner_model(crefs[1])).measure_to_constrs[mcref] = crefs
    return
end

# MeasureConstraintRef (finite)
function _set_mapping(mcref::InfiniteOpt.MeasureConstraintRef,
                      cref::JuMP.ConstraintRef)
    transcription_data(JuMP.owner_model(cref)).measure_to_constrs[mcref] = [cref]
    return
end

# FiniteConstraintRef
function _set_mapping(fcref::InfiniteOpt.FiniteConstraintRef,
                      cref::JuMP.ConstraintRef)
    transcription_data(JuMP.owner_model(cref)).finite_to_constr[fcref] = cref
    return
end

## Define helper functions for setting constraint supports
# InfiniteConstraintRef
function _set_supports(trans_model::JuMP.Model,
                       icref::InfiniteOpt.InfiniteConstraintRef,
                       supports::Vector)
    transcription_data(trans_model).infconstr_to_supports[icref] = supports
    return
end

# MeasureConstraintRef
function _set_supports(trans_model::JuMP.Model,
                       mcref::InfiniteOpt.MeasureConstraintRef,
                       supports::Vector)
    transcription_data(trans_model).measconstr_to_supports[mcref] = supports
    return
end

## Define helper functions for setting constraint parameter reference tuples
# InfiniteConstraintRef
function _set_parameter_refs(trans_model::JuMP.Model,
                             icref::InfiniteOpt.InfiniteConstraintRef,
                             prefs::Tuple)
    transcription_data(trans_model).infconstr_to_params[icref] = prefs
    return
end

# MeasureConstraintRef
function _set_parameter_refs(trans_model::JuMP.Model,
                             mcref::InfiniteOpt.MeasureConstraintRef,
                             prefs::Tuple)
    transcription_data(trans_model).measconstr_to_params[mcref] = prefs
    return
end

# Extract the root name of a constraint
function InfiniteOpt._root_name(cref::InfiniteOpt.GeneralConstraintRef)
    name = JuMP.name(cref)
    if length(name) == 0
        return "noname"
    else
        first_bracket = findfirst(isequal('['), name)
        if first_bracket == nothing
            return name
        else
            return name[1:first_bracket-1]
        end
    end
end

## leverage _make_transcription_function to transcribe the constraints
function _set_constraints(trans_model::JuMP.Model, inf_model::InfiniteOpt.InfiniteModel)
    # transform and add constraints that haven't already been added through add_variable
    ordered_indices = sort(collect(keys(inf_model.constrs)))
    for index in ordered_indices
        constr = inf_model.constrs[index]
        if !inf_model.constr_in_var_info[index]
            # extract the reference and transcribe the jump object function
            icref = InfiniteOpt._make_constraint_ref(inf_model, index)
            if isa(constr, BoundedScalarConstraint)
                results = _make_transcription_function(constr.func, trans_model,
                                                       constr.bounds.intervals)
            else
                results = _make_transcription_function(constr.func, trans_model)
            end
            # extract the name and create a transcribed constraint
            root_name = InfiniteOpt._root_name(icref)
            # results[2] contains prefs if is infinite or it is empty otherwise
            if length(results[2]) != 0
                crefs = Vector{JuMP.ConstraintRef}(undef, length(results[1]))
                for i in eachindex(crefs)
                    con = JuMP.build_constraint(error, results[1][i], constr.set)
                    # TODO Perhaps improve naming
                    name = string(root_name, "(Support: ", i, ")")
                    crefs[i] = JuMP.add_constraint(trans_model, con, name)
                end
                # update the mappings
                _set_mapping(icref, crefs)
                _set_parameter_refs(trans_model, icref, results[2])
                _set_supports(trans_model, icref, results[3])
            else
                # We have a finite constraint, just add normally and update mapping.
                con = JuMP.build_constraint(error, results[1], constr.set)
                cref = JuMP.add_constraint(trans_model, con, root_name)
                _set_mapping(icref, cref)
            end
        end
    end
    return
end

## Helper functions for mapping InfiniteOpt var info constrs to the transcribed ones
# FiniteVariableRef
function _map_info_constraints(fvref::InfiniteOpt.FiniteVariableRef,
                               trans_model::JuMP.Model)
    vref = transcription_variable(trans_model, fvref)
    # Check if variables have info constraints and map if they do
    if JuMP.has_lower_bound(fvref)
        _set_mapping(JuMP.LowerBoundRef(fvref), JuMP.LowerBoundRef(vref))
    end
    if JuMP.has_upper_bound(fvref)
        _set_mapping(JuMP.UpperBoundRef(fvref), JuMP.UpperBoundRef(vref))
    end
    if JuMP.is_fixed(fvref)
        _set_mapping(JuMP.FixRef(fvref), JuMP.FixRef(vref))
    end
    if JuMP.is_integer(fvref)
        _set_mapping(JuMP.IntegerRef(fvref), JuMP.IntegerRef(vref))
    end
    if JuMP.is_binary(fvref)
        _set_mapping(JuMP.BinaryRef(fvref), JuMP.BinaryRef(vref))
    end
    return
end

# InfiniteVariableRef
function _map_info_constraints(ivref::InfiniteOpt.InfiniteVariableRef,
                               trans_model::JuMP.Model)
    vrefs = transcription_variable(trans_model, ivref)
    # TODO Prealocate the size
    # Check if both variables have a constraint and map if they do
    if JuMP.has_lower_bound(ivref)
        crefs = JuMP.ConstraintRef[]
        supports = Tuple[]
        for i in eachindex(vrefs)
            if JuMP.has_lower_bound(vrefs[i])
                push!(crefs, JuMP.LowerBoundRef(vrefs[i]))
                push!(supports, InfiniteOpt.variable_supports(trans_model, ivref)[i])
            end
        end
        if length(crefs) != 0
            _set_mapping(JuMP.LowerBoundRef(ivref), crefs)
            _set_parameter_refs(trans_model, JuMP.LowerBoundRef(ivref),
                                InfiniteOpt.parameter_refs(ivref))
            _set_supports(trans_model, JuMP.LowerBoundRef(ivref), supports)
        end
    end
    if JuMP.has_upper_bound(ivref)
        crefs = JuMP.ConstraintRef[]
        supports = Tuple[]
        for i in eachindex(vrefs)
            if JuMP.has_upper_bound(vrefs[i])
                push!(crefs, JuMP.UpperBoundRef(vrefs[i]))
                push!(supports, InfiniteOpt.variable_supports(trans_model, ivref)[i])
            end
        end
        if length(crefs) != 0
            _set_mapping(JuMP.UpperBoundRef(ivref), crefs)
            _set_parameter_refs(trans_model, JuMP.UpperBoundRef(ivref),
                                InfiniteOpt.parameter_refs(ivref))
            _set_supports(trans_model, JuMP.UpperBoundRef(ivref), supports)
        end
    end
    if JuMP.is_fixed(ivref)
        crefs = JuMP.ConstraintRef[]
        supports = Tuple[]
        for i in eachindex(vrefs)
            if JuMP.is_fixed(vrefs[i])
                push!(crefs, JuMP.FixRef(vrefs[i]))
                push!(supports, InfiniteOpt.variable_supports(trans_model, ivref)[i])
            end
        end
        if length(crefs) != 0
            _set_mapping(JuMP.FixRef(ivref), crefs)
            _set_parameter_refs(trans_model, JuMP.FixRef(ivref),
                                InfiniteOpt.parameter_refs(ivref))
            _set_supports(trans_model, JuMP.FixRef(ivref), supports)
        end
    end
    if JuMP.is_integer(ivref)
        crefs = JuMP.ConstraintRef[]
        supports = Tuple[]
        for i in eachindex(vrefs)
            if JuMP.is_integer(vrefs[i])
                push!(crefs, JuMP.IntegerRef(vrefs[i]))
                push!(supports, InfiniteOpt.variable_supports(trans_model, ivref)[i])
            end
        end
        if length(crefs) != 0
            _set_mapping(JuMP.IntegerRef(ivref), crefs)
            _set_parameter_refs(trans_model, JuMP.IntegerRef(ivref),
                                InfiniteOpt.parameter_refs(ivref))
            _set_supports(trans_model, JuMP.IntegerRef(ivref), supports)
        end
    end
    if JuMP.is_binary(ivref)
        crefs = JuMP.ConstraintRef[]
        supports = Tuple[]
        for i in eachindex(vrefs)
            if JuMP.is_binary(vrefs[i])
                push!(crefs, JuMP.BinaryRef(vrefs[i]))
                push!(supports, InfiniteOpt.variable_supports(trans_model, ivref)[i])
            end
        end
        if length(crefs) != 0
            _set_mapping(JuMP.BinaryRef(ivref), crefs)
            _set_parameter_refs(trans_model, JuMP.BinaryRef(ivref),
                                InfiniteOpt.parameter_refs(ivref))
            _set_supports(trans_model, JuMP.BinaryRef(ivref), supports)
        end
    end
    return
end

## Map the variable info constraints between the two models
function _map_variable_info_constraints(trans_model::JuMP.Model,
                                        inf_model::InfiniteOpt.InfiniteModel)
    for index in keys(inf_model.vars)
        ivref = InfiniteOpt._make_variable_ref(inf_model, index)
        if InfiniteOpt.is_used(ivref)
            _map_info_constraints(ivref, trans_model)
        end
    end
    return
end

# Fill in a given trans_model with the reformulation of the inf_model
# This assumes that trans_model is empty
function _build_transcription_model!(trans_model::JuMP.Model,
                                     inf_model::InfiniteOpt.InfiniteModel)
    InfiniteOpt.fill_in_supports!(inf_model)
    _initialize_hold_variables(trans_model, inf_model)
    _initialize_infinite_variables(trans_model, inf_model)
    _map_point_variables(trans_model, inf_model)
    if JuMP.objective_sense(inf_model) != MOI.FEASIBILITY_SENSE
        _set_objective(trans_model, inf_model)
    end
    _map_variable_info_constraints(trans_model, inf_model)
    # TODO optimize performance --> bottlenecked with _make_transcription_function
    _set_constraints(trans_model, inf_model)
end

"""
    TranscriptionModel(model::InfiniteModel, [optimizer_constructor;
                       caching_mode::MOIU.CachingOptimizerMode = MOIU.AUTOMATIC,
                       bridge_constraints::Bool = true])

Return a `TranscriptionModel` of `model`. This transcribes all of the variables,
constraints, and objective.

**Example**
```julia-repl
julia> TranscriptionModel(model)
A JuMP Model
Feasibility problem with:
Variables: 130
`GenericAffExpr{Float64,VariableRef}`-in-`MathOptInterface.EqualTo{Float64}`: 25 constraints
`GenericAffExpr{Float64,VariableRef}`-in-`MathOptInterface.GreaterThan{Float64}`: 100 constraint
`GenericAffExpr{Float64,VariableRef}`-in-`MathOptInterface.LessThan{Float64}`: 84 constraints
`VariableRef`-in-`MathOptInterface.EqualTo{Float64}`: 40 constraints
`VariableRef`-in-`MathOptInterface.GreaterThan{Float64}`: 40 constraints
`VariableRef`-in-`MathOptInterface.LessThan{Float64}`: 25 constraints
`VariableRef`-in-`MathOptInterface.Integer`: 40 constraints
`VariableRef`-in-`MathOptInterface.ZeroOne`: 40 constraints
Model mode: AUTOMATIC
CachingOptimizer state: NO_OPTIMIZER
Solver name: No optimizer attached.
```
"""
function TranscriptionModel(inf_model::InfiniteOpt.InfiniteModel,
                            optimizer_constructor = nothing;
                            kwargs...)::JuMP.Model
    if isa(optimizer_constructor, Nothing)
        trans_model = TranscriptionModel(; kwargs...)
    else
        trans_model = TranscriptionModel(optimizer_constructor; kwargs...)
    end
    _build_transcription_model!(trans_model, inf_model)
    return trans_model
end


## Here we show how to create the moments data bunnytime.mnts.jld and bunnytimetest.mnts, from the full simulated data (which is kept in a Zenodo repo).

Pkg.clone("https://github.com/jondea/MultipleScattering.jl.git")
using MultipleScattering

Pkg.add("JLD")
using JLD

model_to_label(f) = (f.volfrac,f.particles[1].r)
impulse = get_gaussian_freq_impulse(1.0,2.48)

# Download full simulated data. Note, can take a long time.
  #training set
  download("https://zenodo.org/record/1126642/files/bunny.mst.jld", "bunny.mst.jld")
  #test set
  download("https://zenodo.org/record/1126642/files/bunnytest.mst.jld", "bunnytest.mst.jld")

# There is also a copy of moments data on Zenodo, which can de download using:
  # # Download data from permenant repo
  # download("https://zenodo.org/record/1126642/files/train.mnts.jld", "train.mnts.jld")
  # download("https://zenodo.org/record/1126642/files/test.mnts.jld", "test.mnts.jld")

# test set: create the moments of the backscattered waves in time
  test = collect(values(load("bunnytest.mst.jld")))[1]
  test = sort(test, by = model_to_label)
  test_batched = collect(groupby(model_to_label, test));

  tmp_labels = [model_to_label(bs[1]) for bs in test_batched];
  if length(union(tmp_labels)) != length(tmp_labels)
    error("batched the test data inccorectly")
  end

  test_moments = map(test_batched) do bs
    time_arr = cut_interval(ω_to_t(bs[1].k_arr),[10.,96.])
    StatisticalMoments(
        [TimeSimulation(b;time_arr = time_arr, impulse = impulse) for b in bs],
        4;
        response_apply=real
    )
  end;
  save("train.mnts.jld", "Array{StatisticalMoments}", test_moments)

  test_batched = 0
  test_moments = 0
  gc() # garbage collection

# training set: create the moments of the backscattered waves in time
  train = collect(values(load("bunny.mst.jld")))[1]
  train = sort(train, by = model_to_label)
  train_batched = collect(groupby(model_to_label, train));
  train = 0
  gc()

  tmp_labels = [model_to_label(bs[1]) for bs in train_batched];
  if length(union(tmp_labels)) != length(tmp_labels)
    error("batched the test data inccorectly")
  end

  train_moments = map(train_batched) do bs
    time_arr = cut_interval(ω_to_t(bs[1].k_arr),[10.,96.])
    StatisticalMoments(
        [TimeSimulation(b;time_arr = time_arr, impulse = impulse) for b in bs],
        4;
        response_apply=real
    )
  end;
  save("test.mnts.jld", "Array{StatisticalMoments}", train_moments)


using Oceananigans.Utils: prettysummary

struct FunctionField{LX, LY, LZ, C, P, F, G, T} <: AbstractField{LX, LY, LZ, G, T, 3}
          func :: F
          grid :: G
         clock :: C
    parameters :: P

    @doc """
        FunctionField{LX, LY, LZ}(func, grid; clock=nothing, parameters=nothing) where {LX, LY, LZ}

    Returns a `FunctionField` on `grid` and at location `LX, LY, LZ`.

    If `clock` is not specified, then `func` must be a function with signature
    `func(x, y, z)`. If clock is specified, `func` must be a function with signature
    `func(x, y, z, t)`, where `t` is internally determined from `clock.time`.

    A `FunctionField` will return the result of `func(x, y, z [, t])` at `LX, LY, LZ` on
    `grid` when indexed at `i, j, k`.
    """
    function FunctionField{LX, LY, LZ}(func::F, grid::G; clock::C=nothing, parameters::P=nothing) where {LX, LY, LZ, F, G, C, P}
        T = eltype(grid)
        return new{LX, LY, LZ, C, P, F, G, T}(func, grid, clock, parameters)
    end

    @doc """
        FunctionField{LX, LY, LZ}(func::FunctionField, grid; clock) where {LX, LY, LZ}

    Adds `clock` to an existing `FunctionField` and relocates it to `(LX, LY, LZ)` on `grid`.
    """
    function FunctionField{LX, LY, LZ}(f::FunctionField, grid::G; clock::C=nothing) where {LX, LY, LZ, G, C}
        P = typeof(f.parameters)
        T = eltype(grid)
        F = typeof(f.func)
        return new{LX, LY, LZ, C, P, F, G, T}(f.func, grid, clock, f.parameters)
    end
end

"""Return `a`, or convert `a` to `FunctionField` if `a::Function`"""
fieldify_function(L, a, grid) = a
fieldify_function(L, a::Function, grid) = FunctionField(L, a, grid)

"""
    FunctionField(L::Tuple, func, grid)

Returns a stationary `FunctionField` on `grid` and at location `L = (LX, LY, LZ)`,
where `func` is callable with signature `func(x, y, z)`.
"""
FunctionField(L::Tuple, func, grid) = FunctionField{L[1], L[2], L[3]}(func, grid)

# Various possibilities for calling FunctionField.func:
@inline call_func(clock, parameters, func, x, y, z)     = func(x, y, z, clock.time, parameters)
@inline call_func(::Nothing, parameters, func, x, y, z) = func(x, y, z, parameters)
@inline call_func(clock, ::Nothing, func, x, y, z)      = func(x, y, z, clock.time)
@inline call_func(::Nothing, ::Nothing, func, x, y, z)  = func(x, y, z)

# For setting ReducedField
@inline call_func(::Nothing, ::Nothing, func, x, y)     = func(x, y)
@inline call_func(::Nothing, ::Nothing, func, x)        = func(x)

@inline Base.getindex(f::FunctionField{LX, LY, LZ}, i, j, k) where {LX, LY, LZ} =
    call_func(f.clock, f.parameters, f.func, node(LX(), LY(), LZ(), i, j, k, f.grid)...)

@inline (f::FunctionField)(x...) = call_func(f.clock, f.parameters, f.func, x...)

Adapt.adapt_structure(to, f::FunctionField{LX, LY, LZ}) where {LX, LY, LZ} =
    FunctionField{LX, LY, LZ}(Adapt.adapt(to, f.func),
                           Adapt.adapt(to, f.grid),
                           clock = Adapt.adapt(to, f.clock),
                           parameters = Adapt.adapt(to, f.parameters))

Base.show(io::IO, field::FunctionField) =
    print(io, "FunctionField located at ", show_location(field), '\n',
          "├── func: $(prettysummary(field.func))", '\n',
          "├── grid: $(summary(field.grid))\n",
          "├── clock: $(summary(field.clock))\n",
          "└── parameters: $(field.parameters)")


using ITensors
using Test

@testset "Sweeps constructor" begin

  #Sweeps
  #1 cutoff=1.0E-12, maxdim=50, mindim=10, noise=1.0E-07
  #2 cutoff=1.0E-12, maxdim=100, mindim=20, noise=1.0E-08
  #3 cutoff=1.0E-12, maxdim=200, mindim=20, noise=1.0E-10
  #4 cutoff=1.0E-12, maxdim=400, mindim=20, noise=0.0E+00
  #5 cutoff=1.0E-12, maxdim=800, mindim=20, noise=1.0E-11
  #6 cutoff=1.0E-12, maxdim=800, mindim=20, noise=0.0E+00

  sweep_args =
    [
     "maxdim" "mindim" "cutoff" "noise"
      50       10       1e-12    1e-7
      100      20       1e-12    1e-8
      200      20       1e-12    1e-10
      400      20       1e-12    0
      800      20       1e-12    1e-11
      800      20       1e-12    0
     ]

  @testset "Don't specify nsweep" begin
    nsw = size(sweep_args, 1) - 1
    sw = Sweeps(sweep_args)

    @test nsweep(sw) == nsw

    @test maxdim(sw, 1) == 50
    @test maxdim(sw, 2) == 100
    @test maxdim(sw, 3) == 200
    @test maxdim(sw, 4) == 400
    for n in 5:nsw
      @test maxdim(sw, n) == 800
    end

    @test mindim(sw, 1) == 10
    for n in 2:nsw
      @test mindim(sw, n) == 20
    end

    for n in 1:nsw
      @test cutoff(sw, n) == 1e-12
    end

    @test noise(sw, 1) == 1e-7
    @test noise(sw, 2) == 1e-8
    @test noise(sw, 3) == 1e-10
    @test noise(sw, 4) == 0
    @test noise(sw, 5) == 1e-11
    @test noise(sw, 6) == 0
  end

  @testset "Specify nsweep, more than data" begin
    nsw = 7
    sw = Sweeps(nsw, sweep_args)

    @test nsweep(sw) == nsw

    @test maxdim(sw, 1) == 50
    @test maxdim(sw, 2) == 100
    @test maxdim(sw, 3) == 200
    @test maxdim(sw, 4) == 400
    for n in 5:nsw
      @test maxdim(sw, n) == 800
    end

    @test mindim(sw, 1) == 10
    for n in 2:nsw
      @test mindim(sw, n) == 20
    end

    for n in 1:nsw
      @test cutoff(sw, n) == 1e-12
    end

    @test noise(sw, 1) == 1e-7
    @test noise(sw, 2) == 1e-8
    @test noise(sw, 3) == 1e-10
    @test noise(sw, 4) == 0
    @test noise(sw, 5) == 1e-11
    @test noise(sw, 6) == 0
    @test noise(sw, 7) == 0
  end

  @testset "Specify nsweep, less than data" begin
    nsw = 5
    sw = Sweeps(nsw, sweep_args)

    @test nsweep(sw) == nsw

    @test maxdim(sw, 1) == 50
    @test maxdim(sw, 2) == 100
    @test maxdim(sw, 3) == 200
    @test maxdim(sw, 4) == 400
    for n in 5:nsw
      @test maxdim(sw, n) == 800
    end

    @test mindim(sw, 1) == 10
    for n in 2:nsw
      @test mindim(sw, n) == 20
    end

    for n in 1:nsw
      @test cutoff(sw, n) == 1e-12
    end

    @test noise(sw, 1) == 1e-7
    @test noise(sw, 2) == 1e-8
    @test noise(sw, 3) == 1e-10
    @test noise(sw, 4) == 0
    @test noise(sw, 5) == 1e-11
  end

  @testset "Variable types of input" begin
    sw = Sweeps(5)
    setnoise!(sw,1E-8,0)
    @test noise(sw,1) ≈ 1E-8
    @test noise(sw,2) ≈ 0.0
    @test noise(sw,3) ≈ 0.0
    setcutoff!(sw,0,1E-8,0,1E-12)
    @test cutoff(sw,1) ≈ 0.0
    @test cutoff(sw,2) ≈ 1E-8
    @test cutoff(sw,3) ≈ 0.0
    @test cutoff(sw,4) ≈ 1E-12
  end
end

nothing


module Verifier

using Random
using ..Utils
using ..Optimizer

using ProgressMeter: @showprogress
using nGraph

export verify

# We need to be able to verify that the output of modified graphs matches that of a standard
# graph.
#
# To do that, we use an auxiliary function that accepts the constructor function for a model
# as well as any passes to be performed on the model and checks that the base computation
# matches the modifed computation.

# Get all of the input and output parameters of an executable.
io_arrays(fex) = vflatten(nGraph.splat_inputs(fex), nGraph.splat_outputs(fex))

function verify(backend::nGraph.Backend{nGraph.CPU}, f, opt, cache;
        seed = 8086,
        inner_iterations = 5,
        outer_iterations = 2,
    )

    success = true
    for i in 1:outer_iterations
        println("Working on iteration $i")

        # Create a new seed for each outer iteration.
        outer_seed = rand(UInt64)

        # Wrapped the callable in a seed generator
        seeded_f = @closure begin
            Random.seed!(outer_seed)
            return f()
        end

        # Create a baseline copy and an optimized copy of the function.
        fex_baseline = actualize(backend, seeded_f)
        fex_optimized = first(Optimizer.factory(
            backend,
            seeded_f,
            opt;
            cache = cache,
            search_ratio = false
        ))

        # Sanity check - make sure all input tensors are equivalent.
        # Because some output tensors may be undefined - don't check those until the
        # executables are run for the first time.
        for (a, b) in zip(nGraph.splat_inputs(fex_baseline), nGraph.splat_inputs(fex_optimized))
            if parent(a) != parent(b)
                equal_error(i, 0, seed, outer_seed)
                return false
            end
        end

        for j in 1:inner_iterations
            println("    Inner Iteration $j")

            # Run both functions - make sure ALL input and output tensors match each itertion.
            @time fex_baseline()
            @time fex_optimized()
            for (a,b) in zip(io_arrays(fex_baseline), io_arrays(fex_optimized))
                # Check for NaN's
                if any(isnan, parent(a)) || any(isnan, parent(b))
                    nan_error(i, j, seed, outer_seed)
                    return false
                end

                # Check for Subnormal
                if any(issubnormal, parent(a)) || any(issubnormal, parent(b))
                    subnormal_error(i, j, seed, outer_seed)
                    return false
                end

                # Check equality
                #
                # Use `isapprox` because of non-determinism in GPU computations.
                if !isapprox(parent(a), parent(b))
                    equal_error(i, j, seed, outer_seed)
                    return (a, b)
                end
            end
        end
    end
    return true
end

# For the GPU, nondeterminism means we have to trace over a period of time.
function verify(backend::nGraph.Backend{nGraph.GPU}, f, opt, cache;
        seed = 8086,
        inner_iterations = 100,
        outer_iterations = 2,
    )

    success = true
    for i in 1:outer_iterations
        println("Working on iteration $i")

        # Create a new seed for each outer iteration.
        outer_seed = rand(UInt64)

        # Wrapped the callable in a seed generator
        seeded_f = @closure begin
            Random.seed!(outer_seed)
            return f()
        end

        # GC fence everything to ensure we don't run out of memory.
        GC.gc()
        baseline = run_baseline(backend, seeded_f, opt, cache, inner_iterations)
        GC.gc()
        optimized = run_optimized(backend, seeded_f, opt, cache, inner_iterations)
        GC.gc()
        return baseline, optimized
    end
end

# Inner functions for GC purposes.
function run_baseline(backend, f, opt, cache, inner_iterations)
    fex = actualize(backend, f)
    vals = []
    @showprogress 1 for j in 1:inner_iterations
        push!(vals, collect(parent(fex())))
    end
    return vals
end

function run_optimized(backend, f, opt, cache, inner_iterations)
    #fex = actualize(backend, f)
    fex = first(Optimizer.factory(
        backend,
        f,
        opt;
        cache = cache,
    ))

    vals = []
    @showprogress 1 for j in 1:inner_iterations
        push!(vals, collect(parent(fex())))
    end
    return vals
end

#####
##### Error Generation.
#####

nan_error(x...) = generic_error("NaN Values", x...)
subnormal_error(x...) = generic_error("Subnormal Values", x...)
equal_error(x...) = generic_error("Match Error", x...)

function generic_error(prefix, outer_iteration, inner_iteration, seed, outer_seed)
    message = """
    $prefix
    Outer Iteration: $outer_iteration
    Inner Iteration: $inner_iteration
    Master seed: $seed
    Seed for this round: $outer_seed
    """
    printstyled(message; color = :red)
    println()
    return nothing
end

end


using Robotics

using Test

include("checkpuma.jl")


using IncrementalInference
using Test

## a new factor that is broken
struct BrokenFactor{T<: SamplableBelief} <: AbstractRelativeRoots
    Z::T
end

function IIF.getSample(cf::CalcFactor{<:BrokenFactor})
  return rand(cf.factor.Z, 1)
end

function (s::CalcFactor{<:BrokenFactor})(z,
                                        wxi,
                                        wxj)
    #
    error("User factor has a bug.")
end

# # FIXME consolidate with CalcFactor according to #467
# function (s::BrokenFactor{<:IIF.ParametricTypes})(X1::AbstractArray{<:Real},
#                                                   X2::AbstractArray{<:Real};
#                                                   userdata::Union{Nothing,FactorMetadata}=nothing )
#     error("User factor has a bug -- USE NEW CalcFactor API INSTEAD, v0.21.")
# end

##

@testset "Test CSM monitor/watchdog on errors" begin

##

# create a factor graph
fg = generateGraph_LineStep(10;
                            poseEvery=1,
                            landmarkEvery=5,
                            posePriorsAt=[0],
                            sightDistance=4,
                            solverParams=SolverParams(algorithms=[:default, :parametric]))
#

initAll!(fg)

##

#TODO test FSM watchdog
# add a broken factor - mid
addFactor!(fg, [:x9, :lm10], BrokenFactor(Normal()); graphinit=false)
smtasks = Task[]
@test_throws CompositeException tree = IIF.solveTree!(fg; smtasks=smtasks);
sleep(0.1)

## Test parametric solve also

addFactor!(fg, [:x9, :lm10], BrokenFactor(Normal()); graphinit=false)
##
# IIF.solveTree!(fg; smtasks=smtasks, algorithm = :parametric)
##
@test_throws CompositeException tree2 = IIF.solveTree!(fg; smtasks=smtasks, algorithm = :parametric)
sleep(0.1)

deleteFactor!(fg, :x9lm10f2)

## add a broken factor - leave

addFactor!(fg, [:x10, :lm10], BrokenFactor(Normal()); graphinit=false)
@test_throws CompositeException tree2 = IIF.solveTree!(fg; smtasks=smtasks, algorithm = :parametric)
sleep(0.1)

deleteFactor!(fg, :x10lm10f2)

## add a broken factor - root
addFactor!(fg, [:x7, :lm10], BrokenFactor(Normal()); graphinit=false)
@test_throws CompositeException tree2 = IIF.solveTree!(fg; smtasks=smtasks, algorithm = :parametric)

##

end

@testset "test CSM debug options work" begin

## create a factor graph

fg = generateGraph_LineStep(10;
                            poseEvery=1,
                            landmarkEvery=5,
                            posePriorsAt=[0],
                            sightDistance=4 )
#

smtasks = Task[];
solveTree!(fg, smtasks=smtasks, recordcliqs=ls(fg));

## make sure we fetch the results

hists = fetchCliqHistoryAll!(smtasks);

printCSMHistoryLogical(hists)

printCSMHistorySequential(hists)

## test CSM resolve steps

@info "test repeatCSMStep"
csmc_ = repeatCSMStep!(hists, 1, 1);

##

end


import FLOWNoise
using Test
include("testingfunctions.jl")

@testset "Post-processing: Time to frequency domain" begin
    @testset "Case 1" begin # currently only converts time-domain output to frequency-domain output and compares to WOPWOP frequency domain

        filepath = "./data/case1/"
        windowtype = "None"

        thickness_dB_calc, loading_dB_calc, total_dB_calc,
            thickness_dBA_calc, loading_dBA_calc, total_dBA_calc,
            thickness_dB_actual, loading_dB_actual, total_dB_actual,
            thickness_dBA_actual, loading_dBA_actual, total_dBA_actual,
            thickness_OASPL_calc, loading_OASPL_calc, total_OASPL_calc,
            thickness_OASPLA_calc, loading_OASPLA_calc, total_OASPLA_calc,
            thickness_OASPL_actual, loading_OASPL_actual, total_OASPL_actual,
            thickness_OASPLA_actual, loading_OASPLA_actual, total_OASPLA_actual = prepare_test(filepath, windowtype)

        tol_case1 = 1e-2 # 1% error

        @testset "OASPL" begin
            @test isapprox(percent_error(thickness_OASPL_actual[1], thickness_OASPL_calc), 0.0, atol=tol_case1)
            @test isapprox(percent_error(loading_OASPL_actual[1], loading_OASPL_calc), 0.0, atol=tol_case1)
            @test isapprox(percent_error(total_OASPL_actual[1], total_OASPL_calc), 0.0, atol=tol_case1)
        end
        @testset "OASPLA" begin
            @test isapprox(percent_error(thickness_OASPLA_actual[1], thickness_OASPLA_calc), 0.0, atol=tol_case1)
            @test isapprox(percent_error(loading_OASPLA_actual[1], loading_OASPLA_calc), 0.0, atol=tol_case1)
            @test isapprox(percent_error(total_OASPLA_actual[1], total_OASPLA_calc), 0.0, atol=tol_case1)
        end
        @testset "Thickness SPL" begin #for loops are split up to make the tests more distinct, easier to see what's wrong instead of lumping together
            for i in 1:length(thickness_dB_actual)
                @test isapprox(percent_error(thickness_dB_actual[i], thickness_dB_calc[i]), 0.0, atol=tol_case1)
            end
        end
        @testset "Loading SPL" begin
            for i in 1:length(thickness_dB_actual)
                @test isapprox(percent_error(loading_dB_actual[i], loading_dB_calc[i]), 0.0, atol=tol_case1)
            end
        end
        @testset "Total SPL" begin
            for i in 1:length(thickness_dB_actual)
                @test isapprox(percent_error(total_dB_actual[i], total_dB_calc[i]), 0.0, atol=tol_case1)
            end
        end
        @testset "Thickness SPLA" begin
            for i in 1:length(thickness_dBA_actual)
                @test isapprox(percent_error(thickness_dBA_actual[i], thickness_dBA_calc[i]), 0.0, atol=tol_case1)
            end
        end
        @testset "Loading SPLA" begin
            for i in 1:length(thickness_dBA_actual)
                @test isapprox(percent_error(loading_dBA_actual[i], loading_dBA_calc[i]), 0.0, atol=tol_case1)
            end
        end
        @testset "Total SPLA" begin
            for i in 1:length(thickness_dBA_actual)
                @test isapprox(percent_error(total_dBA_actual[i], total_dBA_calc[i]), 0.0, atol=tol_case1)
            end
        end
    end

    # @testset "Case 2" begin #- Look at low/high pass filtering

    #     filepath = "./data/case2/"
    #     windowtype = "Hann"

    #     thickness_OASPL_calc, loading_OASPL_calc, total_OASPL_calc,
    #         thickness_OASPLA_calc, loading_OASPLA_calc, total_OASPLA_calc,
    #         thickness_OASPL_actual, loading_OASPL_actual, total_OASPL_actual,
    #         thickness_OASPLA_actual, loading_OASPLA_actual, total_OASPLA_actual = prepare_test(filepath, windowtype)

    #     tol_case1 = 1e-2

    #     @test isapprox(OASPL_thickness_dB, thickness_OASPL[1], atol=tol_case1)
    #     @test isapprox(OASPL_loading_dB, loading_OASPL[1], atol=tol_case1)
    #     @test isapprox(OASPL_total_dB, total_OASPL[1], atol=tol_case1)
    #     @test isapprox(OASPL_thickness_dBA, thickness_OASPLA[1], atol=tol_case1)
    #     @test isapprox(OASPL_loading_dBA, loading_OASPLA[1], atol=tol_case1)
    #     @test isapprox(OASPL_total_dBA, total_OASPLA[1], atol=tol_case1)

    # end

    # @testset "Case 4" begin - need to look at filtering

    # filepath = "./data/case5/"
    #     windowtype = "None"

    #     thickness_OASPL_calc, loading_OASPL_calc, total_OASPL_calc,
    #         thickness_OASPLA_calc, loading_OASPLA_calc, total_OASPLA_calc,
    #         thickness_OASPL_actual, loading_OASPL_actual, total_OASPL_actual,
    #         thickness_OASPLA_actual, loading_OASPLA_actual, total_OASPLA_actual = prepare_test(filepath, windowtype)

    #     tol_case1 = 1e-2

    #     @test isapprox(OASPL_thickness_dB, thickness_OASPL[1], atol=tol_case1)
    #     @test isapprox(OASPL_loading_dB, loading_OASPL[1], atol=tol_case1)
    #     @test isapprox(OASPL_total_dB, total_OASPL[1], atol=tol_case1)
    #     @test isapprox(OASPL_thickness_dBA, thickness_OASPLA[1], atol=tol_case1)
    #     @test isapprox(OASPL_loading_dBA, loading_OASPLA[1], atol=tol_case1)
    #     @test isapprox(OASPL_total_dBA, total_OASPLA[1], atol=tol_case1)

    # end

    # @testset "Case 5" begin - SPL found only for two specific ranges, not each frequency in the spectrum

    # filepath = "./data/case5/"
    #     windowtype = "None"

    #     thickness_OASPL_calc, loading_OASPL_calc, total_OASPL_calc,
    #         thickness_OASPLA_calc, loading_OASPLA_calc, total_OASPLA_calc,
    #         thickness_OASPL_actual, loading_OASPL_actual, total_OASPL_actual,
    #         thickness_OASPLA_actual, loading_OASPLA_actual, total_OASPLA_actual = prepare_test(filepath, windowtype)

    #     tol_case1 = 1e-2

    #     @test isapprox(OASPL_thickness_dB, thickness_OASPL[1], atol=tol_case1)
    #     @test isapprox(OASPL_loading_dB, loading_OASPL[1], atol=tol_case1)
    #     @test isapprox(OASPL_total_dB, total_OASPL[1], atol=tol_case1)
    #     @test isapprox(OASPL_thickness_dBA, thickness_OASPLA[1], atol=tol_case1)
    #     @test isapprox(OASPL_loading_dBA, loading_OASPLA[1], atol=tol_case1)
    #     @test isapprox(OASPL_total_dBA, total_OASPLA[1], atol=tol_case1)

    # end

end


# to_inplace.jl - transform ExGraph to buffered/in-place operations

const Num = Number
const Vec = AbstractVector
const Mat = AbstractMatrix

const INPLACE_PHS = Set([:u, :v, :w, :x, :y, :z, :X, :Y, :Z])

const INPLACE_RULES = [
    Type[Mat, Mat] => :(Z = X' * Y) => :(mul!(Z, transpose(X), Y)),
    Type[Mat, Mat] => :(Z = X * Y') => :(mul!(Z, X, transpose(Y))),
    Type[Mat, Mat] => :(Z = X * Y) => :(mul!(Z, X, Y)),
]

function to_inplace(g::ExGraph)
    evaluate!(g)
    g = expand_fixed_sequences(g)
    g = fuse_broadcasting(g)
    res = :(begin end)
    for nd in g.tape
        if getcategory(nd) != :input
            vex = to_inplace(g, nd)
            # push!(res.args, simplify(subs_size(vex, sizes)))
            push!(res.args, vex)
        end
    end
    res = subs_bcast_with_dot(sanitize(res))
    return res
end

function to_inplace(g::ExGraph, nd::Union{ExNode{:call}, ExNode{:bcast}, ExNode{:opaque}})
    ex = expand_const(g, to_expr_kw(nd)) |> simplify
    if isa(nd, ExNode{:call}) && !iscall(ex.args[2])
        return to_inplace(g, convert_call(g, nd))
    end
    # try patterns
    dep_vals = [getvalue(g[dep]) for dep in dependencies(nd)]
    for (types, (pat, rpat)) in INPLACE_RULES
        if all(isa(val, T) for (val, T) in zip(dep_vals, types))
            pat = sanitize(pat)
            rex = tryrewrite(ex, pat, rpat; phs=INPLACE_PHS)
            if rex != nothing
                return rex
            end
        end
    end
    # try broadcasting
    if is_bcast_vec(nd)
        lhs_is_scalar = isa(getvalue(nd), Number)
        return make_elementwise(to_expr_kw(nd); lhs_is_scalar=lhs_is_scalar)
    end
    # if LHS is an array, use .= instead of =
    if isa(getvalue(nd), AbstractArray)
        rex = to_expr_kw(nd)
        rex.head = :.=
        return rex
    end
    return to_expr_kw(nd)
end

function to_inplace(g::ExGraph, nd::ExNode{:ctor})
    ex = to_expr_kw(nd)
    insert!(ex.args[2].args, 3, :mem)
    insert!(ex.args[2].args, 4, QuoteNode(varname(nd)))
    return ex
end

function to_inplace(g::ExGraph, nd::ExNode)
    return to_expr_kw(nd)
end

to_buffered = to_inplace

## @inplacerule

# TODO: copied from XGrad, fix
isparameters(a) = isa(a, Expr) && a.head == :parameters

function without_types(pat)
    rpat = copy(pat)
    for i=2:length(rpat.args)
        a = rpat.args[i]
        if !isparameters(a)  # parameters aren't modified
            rpat.args[i] = isa(a, Expr) ? a.args[1] : a
        end
    end
    return rpat
end

function get_arg_names(pat)
    return [isa(a, Expr) ? a.args[1] : a for a in pat.args[2:end] if !isparameters(a)]
end

function get_arg_types(pat)
    return [isa(a, Expr) ? eval(Base, a.args[2]) : Any for a in pat.args[2:end] if !isparameters(a)]
end

function add_inplace_rule(mod::Module, tpat, rpat)
    @assert tpat.head == :(=)
    @assert rpat.head == :call
    op = canonical(mod, tpat.args[2].args[1])
    rop = canonical(mod, rpat.args[1])
    types = get_arg_types(tpat.args[2])
    pat = :($(tpat.args[1]) = $(without_types(tpat.args[2])))
    # replace function names with canonical names
    pat = subs(pat, Dict(tpat.args[2].args[1] => op))
    rpat = subs(rpat, Dict(rpat.args[1] => rop))
    rule = types => pat => rpat
    push!(INPLACE_RULES, rule)
end

macro inplacerule(tpat, rpat)
    mod = @__MODULE__
    add_inplace_rule(mod, tpat, rpat)
    nothing
end


abstract type RinexBase end
abstract type RinexHeader <: RinexBase end
abstract type RinexData <: RinexBase end

#= =========================== =#

@with_kw struct RinexNavHeader <: RinexHeader
    version::Tuple{Float64,String,String,SatelliteSystemType}
    runBy::Tuple{String,String,String}
    comment::String
    ionoCorr::String
    timeSysCorr::Dict{TimeCorrectionType,TimeSystemCorrection}
    leapSeconds::Tuple{Int64,Int64,Int64,Int64}
    corrSysTime::String
    deltaUTC::String
    dUTC::String
    ionAlpha::String
    ionBeta::String
    eoH::String
end

@with_kw struct EphEpoch
    si::SatID
    dt::DateTime
    d::Array{Float64,1}
end

@with_kw mutable struct EphEpochStore
    si::SatID
    keys::Array{DateTime,1}
    eph::Dict{DateTime,EphEpoch}
end

@with_kw struct RinexNavData <: RinexData
    data::Dict{SatID,EphEpochStore}
end

@with_kw struct RinexNav <: RinexBase
    filename::String
    header::RinexHeader
    data::RinexNavData
end

# TODO multiple data files

#= =========================== =#

D_SystemNumObs = Dict{SatelliteSystemType,Tuple{Int64,Array{ObsID,1}}}

@with_kw struct RinexObsHeader <: RinexHeader
    version::Tuple{Float64,Char,SatelliteSystemType}
    runBy::Tuple{String,String,String}
    comment::String
    markerName::String
    markerNumber::String
    markerType::String
    observer::Tuple{String,String}
    receiver::String
    antennaType::String
    antennaPosition::ECEF{Float64}
    antennaDeltaHEN::String
    antennaDeltaXYZ::String
    antennaPhaseCtr::String
    antennaBsightXYZ::String
    antennaZeroDirAzi::String
    antennaZeroDirXYZ::String
    centerOfMass::String
    numObs::String
    systemNumObs::D_SystemNumObs
    waveFact::String
    sigStrengthUnit::String
    interval::Float64
    firstTime::Tuple{DateTime,Millisecond}
    lastTime::String
    receiverOffset::String
    systemDCBSapplied::String
    systemPCVSapplied::String
    systemScaleFac::String
    systemPhaseShift::String
    glonassSlotFreqNo::String
    glonassCodPhsBias::String
    leapSeconds::String
    numSats::String
    prnObs::String
    eoH::String
    iscrx::Bool = false
    # recommend_length::Int64
    # depended on first obs last obs and interval
    # 2400 => 1 days 15 s interval
end

struct RinexDatum
    obs::Float64 # empty -> NaN
    lli::Char
    ssi::Char
    function RinexDatum(str::Union{String,SubString})
        v = strip(SubString(str, 1, 13))
        if isempty(v)
            obs = NaN
        else
            obs = parse(Float64, v)
        end
        llic = str[14]
        ssic = str[15]
        new(obs, llic, ssic)
    end
end

struct EpochHeader
    dt::DateTime
    flag::Char
    numObs::Int64
    RCO::Float64
end

@with_kw struct SatelliteObsData
    satid::SatID
    obsids::Array{ObsID,1}
    dts::Array{DateTime,1}
    obs::Array{Array{RinexDatum,1},1}
end
@with_kw struct RinexObsData <: RinexData
    epochs::Dict{DateTime,EpochHeader}
    comments::Dict{DateTime,String}
    obs::Dict{SatID,SatelliteObsData}
end

struct RinexObs <: RinexBase
    filename::String
    header::RinexObsHeader
    data::RinexObsData
end

# TODO multiple data files => cat?


module pureSubstanceThermo
    import Printf.@sprintf
    include("data.jl")
    include("SolidsEnthalpy.jl")
    include("SolidsCp.jl")
    include("SolidsEntropy.jl")
    include("GasThermo.jl")
    include("MolarMass.jl")
    export GasEntropy, GasEnthalpy, GasCp
    export SolidsCp, SolidsEntropy, SolidsEnthalpy
    export Mm
end  # module SolidsThermo


module TestMeasures

using MLJBase, Test
import Distributions
using CategoricalArrays
using Statistics
import LossFunctions
using StableRNGs
using OrderedCollections: LittleDict

rng  = StableRNGs.StableRNG(123)

@testset "aggregation" begin
    v = rand(5)
    @test aggregate(v, mae) ≈ mean(v)
    @test aggregate(v, TruePositive()) ≈ sum(v)
    @test aggregate(v, rms) ≈ sqrt(mean(v.^2))
    λ = rand()
    @test aggregate(λ, rms) === λ
    @test aggregate(aggregate(v, l2), l2) == aggregate(v, l2)
    m = LittleDict([0, 1, 2, 3, 4], v)
    @test aggregate(m, MTPR()) == mean(v)
end

@testset "metadata" begin
    measures()
    measures(m -> m.target_scitype <: AbstractVector{<:Finite} &&
             m.supports_weights)
    info(rms)
    @test true
end

@testset "coverage" begin
    # just checking that the traits work not that they're correct
    @test orientation(BrierScore()) == :score
    @test orientation(auc) == :score
    @test orientation(rms) == :loss

    @test reports_each_observation(auc) == false
    @test is_feature_dependent(auc) == false

    @test MLJBase.distribution_type(auc) == MLJBase.UnivariateFinite
end

@testset "MLJBase.value" begin
    yhat = randn(rng,5)
    X = (weight=randn(rng,5), x1 = randn(rng,5))
    y = randn(rng,5)
    w = randn(rng,5)

    @test MLJBase.value(mae, yhat, nothing, y, nothing) ≈ mae(yhat, y)
    @test MLJBase.value(mae, yhat, nothing, y, w) ≈ mae(yhat, y, w)

    spooky(yhat, y) = abs.(yhat - y) |> mean
    @test MLJBase.value(spooky, yhat, nothing, y, nothing) ≈ mae(yhat, y)

    cool(yhat, y, w) = abs.(yhat - y) .* w |> mean
    MLJBase.supports_weights(::Type{typeof(cool)}) = true
    @test MLJBase.value(cool, yhat, nothing, y, w) ≈ mae(yhat, y, w)

    funky(yhat, X, y) = X.weight .* abs.(yhat - y) |> mean
    MLJBase.is_feature_dependent(::Type{typeof(funky)}) = true
    @test MLJBase.value(funky, yhat, X, y, nothing) ≈ mae(yhat, y, X.weight)

    weird(yhat, X, y, w) = w .* X.weight .* abs.(yhat - y) |> mean
    MLJBase.is_feature_dependent(::Type{typeof(weird)}) = true
    MLJBase.supports_weights(::Type{typeof(weird)}) = true
    @test MLJBase.value(weird, yhat, X, y, w) ≈ mae(yhat, y, X.weight .* w)
end

mutable struct DRegressor <: Deterministic end
MLJBase.target_scitype(::Type{<:DRegressor}) =
    AbstractVector{<:Continuous}

mutable struct D2Regressor <: Deterministic end
MLJBase.target_scitype(::Type{<:D2Regressor}) =
    AbstractVector{Continuous}

mutable struct DClassifier <: Deterministic end
MLJBase.target_scitype(::Type{<:DClassifier}) =
    AbstractVector{<:Finite}

mutable struct PClassifier <: Probabilistic end
MLJBase.target_scitype(::Type{<:PClassifier}) =
    AbstractVector{<:Finite}

mutable struct PRegressor <: Probabilistic end
MLJBase.target_scitype(::Type{<:PRegressor}) =
    AbstractVector{<:Continuous}

mutable struct PCountRegressor <: Probabilistic end
MLJBase.target_scitype(::Type{<:PCountRegressor}) =
    AbstractVector{<:Count}

@testset "default_measure" begin
    @test MLJBase.default_measure(DRegressor()) == rms
    @test MLJBase.default_measure(D2Regressor()) == rms
    @test MLJBase.default_measure(DClassifier()) == misclassification_rate
    @test MLJBase.default_measure(PClassifier()) == log_loss

    @test MLJBase.default_measure(DRegressor) == rms
    @test MLJBase.default_measure(D2Regressor) == rms
    @test MLJBase.default_measure(DClassifier) == misclassification_rate
    @test MLJBase.default_measure(PClassifier) == log_loss

    @test MLJBase.default_measure(PRegressor) == log_loss
    @test MLJBase.default_measure(PCountRegressor) == log_loss
end

include("confusion_matrix.jl")
include("roc.jl")
include("continuous.jl")
include("finite.jl")
include("probabilistic.jl")
include("loss_functions_interface.jl")

@testset "show method for measures" begin
    io = IOBuffer()
    for meta in measures()
        m = eval(Meta.parse("$(meta.name)()"))
        show(io, MIME("text/plain"), m)
    end
end

@testset "missing and NaN values in aggregation" begin
    v =[1, 2, missing, 5, NaN]
    @test MLJBase.Sum()(v) == 8
    @test MLJBase.RootMeanSquare()(v) ≈ sqrt((1 + 4 + 25)/3)
    @test_throws MLJBase.ERR_NOTHING_LEFT_TO_AGGREGATE MLJBase.Mean()(Float32[])
end

end
true


using Revise
using Flipping
##
Directory_path = "/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/Datasets/Stimulations/Rbp4_gen_control"
Directory_path = "/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/Datasets/Stimulations/WT_new_protocols"
pokes, streaks, DataIndex  = Flipping.create_exp_dataframes(Directory_path)
d = Flipping.find_behavior(Directory_path)
names(pokes)
###
filepath = "/home/beatriz/mainen.flipping.5ht@gmail.com/Flipping/Datasets/Stimulations/WT_new_protocols/WJ16_200105a.csv"
Flipping.session_info(filepath)
findmax(streaks[:,:Day])
curr_data = FileIO.load(filepath) |> DataFrame
names(curr_data)
DataFrames.rename!(curr_data, Symbol("") => :Poke)
##
pokes, streaks, DataIndex = Flipping.create_exp_dataframes(d[1:2,:])
###
###


__precompile__()

module AutomotiveDrivingModels

using Compat
using Reexport
using Discretizers
using Parameters

@reexport using DataFrames
@reexport using Distributions
@reexport using Vec

include("core/AutoCore.jl")
@reexport using .AutoCore

include(Pkg.dir("AutomotiveDrivingModels", "src", "utils", "minkowski.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "features", "features.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "features", "feature_extractors.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "features", "aggregate.jl"))

include(Pkg.dir("AutomotiveDrivingModels", "src", "simulation", "actions.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "behaviors", "driver_models.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "simulation", "simulation.jl"))

include(Pkg.dir("AutomotiveDrivingModels", "src", "behaviors", "behaviors.jl"))

include(Pkg.dir("AutomotiveDrivingModels", "src", "utils", "roadway_generation.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "utils", "trajdata_cleaning.jl"))

include(Pkg.dir("AutomotiveDrivingModels", "src", "evaluation", "metrics.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "evaluation", "foldsets.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "evaluation", "fold_assigners.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "evaluation", "trajdata_segments.jl"))
include(Pkg.dir("AutomotiveDrivingModels", "src", "evaluation", "evaluation_data.jl"))

end # module


using MolecularGraph: bondsymbol!, bond!

@testset "smarts.bond" begin

@testset "bondsymbol" begin
    state = SmartsParser("", false)
    implicit1 = bondsymbol!(state)
    @test implicit1 === nothing

    state = SmartsParser("-", false)
    explicit1 = bondsymbol!(state)
    @test explicit1 == QueryFormula(:bondorder, 1)

    state = SmartsParser("\\?", false)
    stereo4 = bondsymbol!(state)
    @test stereo4 == QueryFormula(:not, QueryFormula(:stereo, :up))
    @test state.pos == 3
end

@testset "bond" begin
    state = SmilesParser("#", true)
    triple = bond!(state)
    @test triple.order == 3

    state = SmilesParser(":", true)
    arom = bond!(state)
    @test arom.isaromatic == true
end

@testset "smartsbond" begin
    state = SmartsParser("~", false)
    anyb = bond!(state)
    @test anyb.query == QueryFormula(:any, true)

    state = SmartsParser("-!@", false)
    notring = bond!(state)
    @test notring.query == QueryFormula(:and, Set([
        QueryFormula(:bondorder, 1),
        QueryFormula(:isringbond, false)
    ]))
end

end # smiles.bond


export gather!

import MPI

@doc """
    gather!(A, A_global)
    gather!(A, A_global; root=0)

Gather a CPU-array `A` from each member of the Cartesian grid of MPI processes into a one large CPU-array `A_global` on the root process (default: `0`).

!!! note "Memory usage note"
    `gather!` allocates at first call an internal buffer of the size of `A_global` and keeps it alive until `finalize_global_grid` is called. A (re-)allocation occurs only if `gather!` is called with a larger `A_global` than in any previous call since the call to `init_global_grid`. This is an optimisation to minimize (re-)allocation, which is very important as `gather!` is typically called in the main loop of a simulation and its performance is critical for the overall application performance.
"""
gather!

let
    global gather!, free_gather_buffer
    A_all_buf = zeros(0);

    "Free the buffer used by gather!."
    function free_gather_buffer()
        A_all_buf = nothing;
        GC.gc();
        A_all_buf = zeros(0);
    end

    function gather!(A::Array{T}, A_global::Union{Array{T}, Nothing}; root::Integer=0) where T <: GGNumber
        check_initialized();
        cart_gather!(A, A_global, me(), global_grid().dims, comm(); root=root);
        return nothing
    end

    function cart_gather!(A::Array{T}, A_global::Union{Array{T}, Nothing}, me::Integer, dims::Array{T2}, comm_cart::MPI.Comm; root::Integer=0, tag::Integer=0, ndims::Integer=NDIMS_MPI) where T <: GGNumber where T2 <: Integer
        nprocs = prod(dims);
        if me != root
            req = MPI.REQUEST_NULL;
            req = MPI.Isend(A, root, tag, comm_cart);
            MPI.Wait!(req);
        else # (me == root)
            A_global === nothing && error("The input argument A_global can't be `nothing` on the root")
            if length(A_global) != nprocs*length(A) error("The input argument A_global must be of length nprocs*length(A)") end
            if (eltype(A_all_buf) != T)
                A_all_buf = reinterpret(T, A_all_buf);
            end
            if length(A_all_buf) < nprocs*length(A)                       # Allocate only if the buffer is not large enough
                free_gather_buffer();                                     # Free the memory of the old buffer immediately as it can typically go up to the order of the total available memory.
                A_all_buf = zeros(T, Int(ceil(nprocs*length(A)/GG_ALLOC_GRANULARITY))*GG_ALLOC_GRANULARITY);  # Ensure that the amount of allocated memory is a multiple of GG_ALLOC_GRANULARITY*sizeof(T). So, we can always correctly reinterpret A_all_buf even if next time sizeof(T) is greater.
            end
            A_all_flat = view(A_all_buf,1:nprocs*length(A));              # Create a 1D-view on the amount of memory needed from A_all_buf.
            reqs = fill(MPI.REQUEST_NULL, nprocs);
            for p in [0:root-1; root+1:nprocs-1]
                cs = Cint[-1,-1,-1];
                MPI.Cart_coords!(comm_cart, p, cs);
                offset = cs[1]*length(A) + cs[2]*dims[1]*length(A) + cs[3]*dims[1]*dims[2]*length(A)
                A_c = view(A_all_flat, 1+offset:length(A)+offset);
                reqs[p+1] = MPI.Irecv!(A_c, p, tag, comm_cart);           # Irev! requires a contigous (SubArray) buffer (that is not both reshaped and reinterpreted)...
            end
            cs = MPI.Cart_coords(comm_cart);
            A_all = reshape(A_all_flat, (length(A), dims[1], dims[2], dims[3])); # Create a 4D-view on the amount of memory needed from A_all_buf.
            A_all[:,cs[1]+1,cs[2]+1,cs[3]+1] .= A[:];
            if (nprocs>1) MPI.Waitall!(reqs); end
            nx, ny, nz = size(view(A,:,:,:));
            for cz = 0:size(A_all,4)-1, cy = 0:size(A_all,3)-1, cx = 0:size(A_all,2)-1
                A_global[(1:nx).+cx*nx, (1:ny).+cy*ny, (1:nz).+cz*nz] .= reshape(A_all[:,cx+1,cy+1,cz+1],nx,ny,nz); # Store the data at the right place in A_global (works for 1D-3D, e.g. if 2D: nz=1, cz=0...)
            end
        end
    end
end