"""
    TrajectoryOptimization
Primary module for setting up and solving trajectory optimization problems.
"""
module TrajectoryOptimization

using RigidBodyDynamics
using ForwardDiff
using DocStringExtensions
using Interpolations
using RecipesBase
using LinearAlgebra
using Statistics
using Random
using SparseArrays
using SuiteSparse
using StaticArrays
using Logging
using Formatting
using Plots
using BenchmarkTools
using PartedArrays
using Parameters
using Rotations
using BlockArrays
using MathOptInterface
using TimerOutputs

export
    Dynamics,
    Problems

# Primary types
export
    Model,
    QuadraticCost,
    LQRCost,
    LQRObjective,
    GenericCost,
    Trajectory

export
    Problem,
    iLQRSolver,
    iLQRSolverOptions,
    AugmentedLagrangianSolver,
    AugmentedLagrangianSolverOptions,
    AugmentedLagrangianProblem,
    ALTROSolverOptions,
    DIRCOLSolver,
    DIRCOLSolverOptions,
    ProjectedNewtonSolver,
    ProjectedNewtonSolverOptions,
    SequentialNewtonSolver,
    DIRTRELSolver,
    Discrete,
    Continuous,
    Constraint,
    BoundConstraint,
    Equality,
    Inequality,
    ConstraintSet,
    StageConstraintSet,
    TerminalConstraintSet,
    ConstraintSet,
    Objective,
    Constraints

export
    rk3,
    rk4,
    midpoint,
    add_constraints!,
    goal_constraint,
    initial_controls!,
    initial_state!,
    circle_constraint,
    sphere_constraint

# Primary methods
export
    solve,
    solve!,
    rollout!,
    rollout,
    forwardpass!,
    backwardpass!,
    cost,
    max_violation,
    max_violation_direct,
    infeasible_control,
    line_trajectory,
    evaluate!,
    jacobian!

export
    get_sizes,
    num_constraints,
    get_num_constraints,
    get_num_controls,
    init_results,
    to_array,
    get_N,
    to_dvecs,
    quat2rot,
    sphere_constraint,
    circle_constraint,
    plot_trajectory!,
    plot_vertical_lines!,
    convergence_rate,
    plot_obstacles,
    evals,
    reset,
    reset_evals,
    final_time,
    total_time,
    count_constraints,
    inequalities,
    equalities,
    bounds,
    labels,
    terminal,
    stage,
    interp_rows

# Trajectory Types
Trajectory{T} = Vector{T} where T <: AbstractArray
VectorTrajectory{T} = Vector{Vector{T}} where T <: Real
MatrixTrajectory{T} = Vector{Matrix{T}} where T <: Real
AbstractVectorTrajectory{T} = Vector{V} where {V <: AbstractVector{T}, T <: Real}
DiagonalTrajectory{T} = Vector{Diagonal{T,Vector{T}}} where T <: Real
PartedVecTrajectory{T} = Vector{PartedVector{T,Vector{T}}}
PartedMatTrajectory{T} = Vector{PartedMatrix{T,Matrix{T}}}

include("constraints.jl")
include("constraint_sets.jl")
include("cost.jl")
include("model.jl")
include("integration.jl")
include("utils.jl")
include("objective.jl")
include("problem.jl")
include("solvers.jl")
include("rollout.jl")
include("dynamics.jl")
include("problems.jl")
include("logger.jl")

write_ipopt_options()
end


# beefblup
# Julia package for performing single-variate BLUP to find beef cattle
# breeding values
# (C) 2021 Thomas A. Christensen II
# Licensed under BSD-3-Clause License
# cSpell:includeRegExp #.*
# cSpell:includeRegExp ("""|''')[^\1]*\1

module BeefBLUP

# Import the required packages
using CSV
using DataFrames
using LinearAlgebra
using Dates
using Gtk

# Main entry-level function - acts just like the script
function beefblup()

    # Ask for an input spreadsheet
    path = open_dialog_native(
    "Select a beefblup worksheet",
    GtkNullContainer(),
    ("*.csv", GtkFileFilter("*.csv", name="beefblup worksheet"))
    )

    # Ask for an output text filename
    savepath = save_dialog_native(
    "Save your beefblup results",
    GtkNullContainer(),
    (GtkFileFilter("*.txt", name="Results file"),
    "*.txt")
    )

    # Ask for heritability
    print("What is the heritability for this trait?> ")
    h2 = parse(Float64, readline(stdin))

    beefblup(path, savepath, h2)

end

function beefblup(datafile::String, h2::Float64)
    # Assume the data is named the same as the file without the trailing extension
    dataname = join(split(datafile, ".")[1:end - 1])

    # Create a new results name
    resultsfile = string(dataname, "_results.txt")

    # Pass this info on to the worker
    beefblup(datafile, resultsfile, h2)
end

# Main worker function, can perform all the work if given all the user input
function beefblup(path::String, savepath::String, h2::Float64)

    # Import data from a suitable spreadsheet
    data = DataFrame(CSV.File(path))

    # Make sure the data is in the proper format
    renamecolstospec!(data)

    # Sort the array by date
    sort!(data, :birthdate)

    # Define fields to hold id values for animals and their parents
    numanimals = length(data.id)

    # Calculate the relationship matrix
    A = additiverelationshipmatrix(data.id, data.dam, data.sire)

    # Extract all of the fixed effects
    fixedeffectdata = data[:,5:end-1]

    (X, fixedeffects) = fixedeffectmatrix(fixedeffectdata)

    # Extract the observed data
    Y = convert(Array{Float64}, data[:,end])

    # The random effects matrix
    Z = Matrix{Int}(I, numanimals, numanimals)

    # Remove items where there is no data
    nullobs = findall(isnothing, Y)
    Z[nullobs, nullobs] .= 0

    # Calculate heritability
    λ = (1 - h2) / h2

    # Use the mixed-model equations
    MME = [X' * X X' * Z; Z' * X (Z' * Z) + (inv(A) .* λ)]
    MMY = [X' * Y; Z' * Y]
    solutions = MME \ MMY

    # Find the accuracies
    diaginv = diag(inv(MME))
    reliability = ones(Float64, length(diaginv)) .- diaginv .* λ

    # Find how many traits we found BLUE for
    numgroups = length(reliability) - numanimals

    # Split the BLUP and BLUE results
    β = solutions[1:numgroups]
    μ = solutions[numgroups+1:end]
    μ_reliability = reliability[numgroups+1:end]

    # Extract the names of the traits
    traitname = names(data)[end]

    # Start printing results to output
    fileID = open(savepath, "w")
    write(fileID, "beefblup Results Report\n")
    write(fileID, "Produced using beefblup (")
    write(fileID, "https://github.com/millironx/beefblup")
    write(fileID, ")\n\n")
    write(fileID, "Input:\t")
    write(fileID, path)
    write(fileID, "\nAnalysis performed:\t")
    write(fileID, string(Dates.today()))
    write(fileID, "\nTrait examined:\t")
    write(fileID, traitname)
    write(fileID, "\n\n")

    # Print base population stats
    write(fileID, "Base Population:\n")
    for i in 1:length(fixedeffects)
        effect = fixedeffects[i]
        write(fileID, "\t")
        write(fileID, string(effect.name))
        write(fileID, ":\t")
        write(fileID, string(effect.basetrait))
        write(fileID, "\n")
    end

    write(fileID, "\tMean ")
    write(fileID, traitname)
    write(fileID, ":\t")
    write(fileID, string(solutions[1]))
    write(fileID, "\n\n")

    # Contemporary group adjustments
    counter = 2
    write(fileID, "Contemporary Group Effects:\n")
    for i in 1:length(fixedeffects)
        effect = fixedeffects[i]
        write(fileID, "\t")
        write(fileID, effect.name)
        write(fileID, "\tEffect\n")
        for j in 1:length(effect.alltraits)
            types = effect.alltraits[j]
            write(fileID, "\t")
            write(fileID, string(types))
            write(fileID, "\t")
            write(fileID, string(β[counter]))
            write(fileID, "\n")

            counter = counter + 1
        end
        write(fileID, "\n")
    end
    write(fileID, "\n")

    # Expected breeding values
    write(fileID, "Expected Breeding Values:\n")
    write(fileID, "\tID\tEBV\tReliability\n")
    for i in 1:numanimals
        write(fileID, "\t")
        write(fileID, string(data.id[i]))
        write(fileID, "\t")
        write(fileID, string(μ[i]))
        write(fileID, "\t")
        write(fileID, string(μ_reliability[i]))
        write(fileID, "\n")
    end

    write(fileID, "\n - END REPORT -")
    close(fileID)

end

"""
    fixedeffectmatrix(fixedeffectdata::DataFrame)

Creates contemporary groupings and the fixed-effect incidence matrix based on the fixed
effects listed in `fixedeffectdata`.

Returns a tuple `(X::Matrix{Int}, fixedeffects::Array{FixedEffect})` in which `X` is the
actual matrix, and `fixedeffects` is the contemporary groupings.
"""
function fixedeffectmatrix(fixedeffectdata::DataFrame)
    # Declare an empty return matrix
    fixedeffects = FixedEffect[]

    # Add each trait to the array
    for i in 1:size(fixedeffectdata)[2]
        name = names(fixedeffectdata)[i]
        traits = eachcol(fixedeffectdata)[i]

        if length(unique(traits)) > 1
            push!(fixedeffects, FixedEffect(name, traits))
        else
            @warn string("column '", name, "' does not have any unique animals and will be dropped from analysis")
            pname = propertynames(fixedeffectdata)[i]
            DataFrames.select!(fixedeffectdata, Not(pname))
        end
    end

    X = ones(Int64, (size(fixedeffectdata)[1], 1))

    for i in 1:length(fixedeffects)
        trait = fixedeffects[i]
        for phenotype in trait.alltraits
            X = cat(X, Int64.(fixedeffectdata[:,i] .== phenotype), dims=2)
        end
    end

    return X, fixedeffects
end

"""
    additiverelationshipmatrix(id, dam, sire)

Returns the additive numerator relationship matrix based on the pedigree provided in `dam`
and `sire` for animals in `id`.

"""
function additiverelationshipmatrix(id::AbstractVector, damid::AbstractVector, sireid::AbstractVector)
    # Sanity-check for valid pedigree
    if !(length(id) == length(damid) && length(damid) == length(sireid))
        throw(ArgumentError("id, dam, and sire must be of the same length"))
    end

    # Convert to positions
    dam = indexin(damid, id)
    sire = indexin(sireid, id)

    # Calculate loop iterations
    numanimals = length(dam)

    # Create an empty matrix for the additive relationship matrix
    A = zeros(numanimals, numanimals)

    # Create the additive relationship matrix by the FORTRAN method presented by
    # Henderson
    for i in 1:numanimals
        if !isnothing(dam[i]) && !isnothing(sire[i])
            for j in 1:(i - 1)
                A[j,i] = 0.5 * (A[j,sire[i]] + A[j,dam[i]])
                A[i,j] = A[j,i]
            end
            A[i,i] = 1 + 0.5 * A[sire[i], dam[i]]
        elseif !isnothing(dam[i]) && isnothing(sire[i])
        for j in 1:(i - 1)
            A[j,i] = 0.5 * A[j,dam[i]]
            A[i,j] = A[j,i]
        end
        A[i,i] = 1
    elseif isnothing(dam[i]) && !isnothing(sire[i])
        for j in 1:(i - 1)
            A[j,i] = 0.5 * A[j,sire[i]]
            A[i,j] = A[j,i]
        end
        A[i,i] = 1
    else
        for j in 1:(i - 1)
            A[j,i] = 0
            A[i,j] = 0
        end
        A[i,i] = 1
        end
    end

    return A
end

"""
    renamecolstospec(::DataFrame)

Renames the first four columns of the beefblup data sheet so that they can be referred to by
name instead of by column index, regardless of user input.
"""
function renamecolstospec!(df::DataFrame)
    # Pull out the fixed-effect and observation name
    othernames = propertynames(df)[5:end]

    # Put specification column names and user-defined names together
    allnames = cat([:id, :birthdate, :dam, :sire], othernames, dims=1)

    # Rename in the DataFrame
    rename!(df, allnames, makeunique=true)
    return df
end

struct FixedEffect
    name::String
    basetrait::Any
    alltraits::AbstractArray{Any}
end

function FixedEffect(name::String, incidences)
    basetrait = last(unique(incidences))
    types = unique(incidences)[1:end-1]
    return FixedEffect(name, basetrait, types)
end

end


using Pkg
Pkg.activate(".")
using DataFrames,Distances,  BenchmarkTools
x = rand(1:4,100)
y = rand(100,5) .+ (x .* (x .==3))

preds = [["a","b","c","d"][i] for i in x]
preds2 = [rand(["a","b","c","d"]) for i in x]
preds3 = [rand(["a","b","c","d"]) for i in x]
df = DataFrame([preds,preds2,preds3],[:X,:Z,:W])

using SnoopCompile,StatsModels,ProfileView
tinf1 =@snoopi_deep begin
    using PERMANOVA
end
ProfileView.view(flamegraph(tinf1))
tinf2 =@snoopi_deep begin
    Hyde = hydra(df,y,BrayCurtis,@formula(1~X+Z+W),10000)
end
ProfileView.view(flamegraph(tinf2))
ttot, pcs = SnoopCompile.parcel(tinf2)
SnoopCompile.write("./precomp/precompiles", pcs)


function DopamineDQNBaseline(env; minimal_action_set = false)
    # Replicate Dopamine's Results!

    γ=0.99
    batch_size=32
    buffer_size = 1000000
    tn_update_freq= 8000
    hist_length = 4
    update_freq = 4
    min_mem_size = 20000

    learning_rate = 0.00025
    momentum_term = 0.00
    squared_grad_term = 0.95
    min_grad_term = 1e-5

    example_state = MinimalRLCore.get_state(env)
    action_set = if minimal_action_set
        get_minimal_actions(env)
    else
        get_actions(env)
    end

    init_f = Flux.glorot_uniform

    model = Chain(
        Conv((8,8), hist_length=>32, relu, stride=4, init=init_f),
        Conv((4,4), 32=>64, relu, stride=2, init=init_f),
        Conv((3,3), 64=>64, relu, stride=1, init=init_f),
        (x)->reshape(x, :, size(x, 4)),
        Dense(3136, 512, relu, initW=init_f),
        Dense(512, length(action_set), identity, initW=init_f)) |> gpu

    target_network  = deepcopy(model)

    return DQNAgent(
        model,
        target_network,
        QLearningHuberLoss(γ),
        DeepRL.RMSPropTFCentered(learning_rate,
                                 squared_grad_term,
                                 momentum_term,
                                 min_grad_term),
        DeepRL.ϵGreedyDecay((1.0, 0.01), 250000, min_mem_size, action_set),
        buffer_size,
        hist_length,
        example_state,
        batch_size,
        tn_update_freq,
        update_freq,
        min_mem_size;
        hist_squeeze = Val{false}(),
        state_preproc = DeepRL.image_manip_atari,
        state_postproc = DeepRL.image_norm,
        rew_transform = (r)->clamp(Float32(r), -1.0f0, 1.0f0),
        device = Flux.use_cuda[] ? Val{:gpu}() : Val{:cpu}())
end

function RevisitingALEDQNBaseline(env)
    # Replicate Dopamine's Results!

    γ=0.99
    batch_size=32
    buffer_size = 1000000
    tn_update_freq= 10000
    hist_length = 4
    update_freq = 4
    min_mem_size = 50000

    learning_rate = 0.00025
    momentum_term = 0.00
    squared_grad_term = 0.95
    min_grad_term = 1e-5

    example_state = MinimalRLCore.get_state(env)

    init_f = Flux.glorot_uniform

    model = Chain(
        Conv((8,8), hist_length=>32, relu, stride=4, init=init_f),
        Conv((4,4), 32=>64, relu, stride=2, init=init_f),
        Conv((3,3), 64=>64, relu, stride=1, init=init_f),
        (x)->reshape(x, :, size(x, 4)),
        Dense(3136, 512, relu, initW=init_f),
        Dense(512, length(get_actions(env)), identity, initW=init_f)) |> gpu

    target_network  = deepcopy(model)

    return DQNAgent(
        model,
        target_network,
        QLearningHuberLoss(γ),
        DeepRL.RMSPropTFCentered(learning_rate,
                                 squared_grad_term,
                                 momentum_term,
                                 min_grad_term),
        DeepRL.ϵGreedyDecay((1.0, 0.01), 1000000, min_mem_size, get_actions(env)),
        buffer_size,
        hist_length,
        example_state,
        batch_size,
        tn_update_freq,
        update_freq,
        min_mem_size;
        hist_squeeze = Val{false}(),
        state_preproc = DeepRL.image_manip_atari,
        state_postproc = DeepRL.image_norm,
        rew_transform = (r)->clamp(Float32(r), -1.0f0, 1.0f0),
        device = Flux.use_cuda[] ? Val{:gpu}() : Val{:cpu}())
end


struct DiffRuleOperation
    diffrule::DiffRule
    operations::OffsetVector{Operation,Vector{Operation}}
end
function DiffRuleOperation(dr::DiffRule, nparents::Int = num_parents(dr))
    ops = OffsetArray{Operation}(undef, -nparents:last(last(dr.sections)))
    DiffRuleOperation(
        dr, ops
    )
end

function DiffRuleOperation(op::Operation, newops::Vector{Operation})
    vparents = parents(op)
    nargs = length(vparents)
    dr = DERIVATIVERULES[InstructionArgs(instruction(op), nargs)]
    ops = OffsetArray{Operation}(undef, -nargs:last(last(dr.sections)))
    for i ∈ -nargs:-1
        ops[i] = newops[identifier(vparents[i + nargs + 1])]
    end
    DiffRuleOperation( dr, ops )
end
function LoopVectorization.parents(dro::DiffRuleOperation)
    nparents = num_parents(dro.diffrule)
    view(dro.operations[-nparents:-1])
end
LoopVectorization.instruction(dro::DiffRuleOperation, i::Int) = instruction(dro.diffrule, i)
LoopVectorization.operations(dro::DiffRuleOperation) = dro.operations
function returned_ind(dro::DiffRuleOperation, i::Int)
    i == 2 && return typemin(Int)
    dro.diffrule.returns[i - (i > 2)]
end
returned_inds(dro::DiffRuleOperation) = dro.diffrule.returns
section(dro::DiffRuleOperation, i::Int) = dro.diffrule.sections[i]
sections(dro::DiffRuleOperation) = dro.diffrule.sections
dependencies(dro::DiffRuleOperation) = dro.diffrule.dependencies
dependencies(dro::DiffRuleOperation, i::Int) = dro.diffrule.dependencies[i]

function Base.getindex(dro::DiffRuleOperation, i::Int)
    # i -= first(dro.operations.offsets)
    instruction(dro, i), dependencies(dro, i)
end


"""
   Wrapper for Sundials `N_Vector` that
   uses Julia `Vector{realtype}` as the data container.

   Implements `DenseVector` interface and
   manages automatic destruction of the referenced `N_Vector` when it is
   no longer in use.
"""
struct NVector <: DenseVector{realtype}
    ref_nv::Ref{N_Vector}   # reference to N_Vector
    v::Vector{realtype}     # array that is referenced by N_Vector

    function NVector(v::Vector{realtype})
        # note that N_VMake_Serial() creates N_Vector doesn't own the data,
        # so calling N_VDestroy_Serial() would not deallocate v
        nv = new(Ref{N_Vector}(N_VMake_Serial(length(v), v)), v)
        finalizer(release_handle, nv.ref_nv)
        return nv
    end

    function NVector(nv::N_Vector)
        # wrap N_Vector into NVector and get non-owning access to `nv` data
        # via `v`, but don't register finalizer for `nv`
        return new(Ref{N_Vector}(nv), asarray(nv))
    end
end

N_Vector(x::NVector) = convert(N_Vector,x)

release_handle(ref_nv::Ref{N_Vector}) = N_VDestroy_Serial(ref_nv[])

Base.size(nv::NVector, d...) = size(nv.v, d...)
Base.stride(nv::NVector, d::Integer) = stride(nv.v, d)

Base.getindex(nv::NVector, i::Real) = getindex(nv.v, i)
Base.getindex(nv::NVector, i::AbstractArray) = getindex(nv.v, i)
Base.getindex(nv::NVector, inds...) = getindex(nv.v, inds...)

Base.setindex!(nv::NVector, X, i::Real) = setindex!(nv.v, X, i)
Base.setindex!(nv::NVector, X, i::AbstractArray) = setindex!(nv.v, X, i)
Base.setindex!(nv::NVector, X, inds...) = setindex!(nv.v, X, inds...)

Base.pointer(nv::NVector) = Sundials.__N_VGetArrayPointer_Serial(nv.ref_nv[])

##################################################################
#
# Methods to convert between Julia Vectors and Sundials N_Vectors.
#
##################################################################

Base.convert(::Type{NVector}, v::Vector{realtype}) = NVector(v)
Base.convert(::Type{NVector}, v::Vector{T}) where {T<:Real} = NVector(copy!(similar(v, realtype), v))
Base.convert(::Type{NVector}, nv::NVector) = nv
Base.convert(::Type{NVector}, nv::N_Vector) = NVector(nv)
Base.convert(::Type{N_Vector}, nv::NVector) = nv.ref_nv[]
Base.convert(::Type{Vector{realtype}}, nv::NVector)= nv.v
Base.convert(::Type{Vector}, nv::NVector)= nv.v

""" `N_Vector(v::Vector{T})`

    Converts Julia `Vector` to `N_Vector`.

    Implicitly creates `NVector` object that manages automatic
    destruction of `N_Vector` object when no longer in use.
"""
Base.convert(::Type{N_Vector}, v::Vector{realtype}) = N_Vector(NVector(v))
Base.convert(::Type{N_Vector}, v::Vector{T}) where {T<:Real} = N_Vector(NVector(v))

Base.similar(nv::NVector) = NVector(similar(nv.v))

nvlength(x::N_Vector) = unsafe_load(unsafe_load(convert(Ptr{Ptr{Clong}}, x)))
# asarray() creates an array pointing to N_Vector data, but does not take the ownership
@inline asarray(x::N_Vector) = unsafe_wrap(Array, __N_VGetArrayPointer_Serial(x), (nvlength(x),), own=false)
@inline asarray(x::N_Vector, dims::Tuple) = unsafe_wrap(Array, __N_VGetArrayPointer_Serial(x), dims, own=false)
asarray(x::Vector{realtype}) = x
asarray(x::Ptr{realtype}, dims::Tuple) = unsafe_wrap(Array, x, dims, own=false)
@inline Base.convert(::Type{Vector{realtype}}, x::N_Vector) = asarray(x)
@inline Base.convert(::Type{Vector}, x::N_Vector) = asarray(x)

nvector(x::Vector{realtype}) = NVector(x)
#nvector(x::N_Vector) = x


#!/usr/bin/env julia

# Draw lines with various dash styles.

using Compose
using Color

function draw_lines(dash_patterns)
    if length(dash_patterns) == 0
        canvas()
    else
        compose(compose(canvas(), lines((0, 0), (1, 0)), strokedash(dash_patterns[1])),
                compose(canvas(0, 0.1), draw_lines(dash_patterns[2:end])))
    end
end

patterns = Array[
    [5mm, 5mm],
    [5mm, 10mm],
    [10mm, 5mm],
    [5mm, 1mm],
    [1mm, 5mm],
    [0.9mm],
    [15mm, 10mm, 5mm],
    [15mm, 10mm, 5mm, 10mm],
    [15mm, 10mm, 5mm, 10mm, 15mm],
    [5mm, 5mm, 1mm, 5mm]]

c = draw_lines(patterns)

img = PDF("dash.pdf", 4inch, 4(sqrt(3)/2)inch)
draw(img, compose(c, stroke(color("black")), linewidth(1mm)))


export
    # vehicles
    Frenet,
    VehicleState,
    Vehicle,

    NULL_FRENET,

    get_posG,
    get_vel_s,
    get_vel_t,
    get_center,
    get_footpoint,
    get_front,
    get_rear,
    get_lane_width,
    get_markerdist_left,
    get_markerdist_right,

    # trajdata
    TrajdataFrame,
    TrajdataState,
    Trajdata,
    TrajdataVehicleIterator,

    # scene
    Scene,

    NeighborLongitudinalResult,
    get_neighbor_fore_along_lane,
    get_neighbor_fore_along_left_lane,
    get_neighbor_fore_along_right_lane,
    get_neighbor_rear_along_lane,
    get_neighbor_rear_along_left_lane,
    get_neighbor_rear_along_right_lane,

    FrenetRelativePosition,
    get_frenet_relative_position,

    # scene record
    SceneRecord,

    pastframe_inbounds

include("frenet.jl")
include("vehicles.jl")
include("trajdatas.jl")
include("scenes.jl")
include("scene_records.jl")


################################################################################
# Starting Area
################################################################################

mutable struct StartingArea{T}
    # This is a rectangular starting area on state space
    x_nom::VehicleState{T} # Nominal state
    x_min::VehicleState{T} # Bound on the starting state
    x_max::VehicleState{T} # Bound on the starting state
end

function StartingArea()
    return StartingArea(VehicleState(), VehicleState(), VehicleState())
end

function StartingArea(x_nom::VehicleState{T}) where {T}
    return StartingArea{T}(copy(x_nom), copy(x_nom), copy(x_nom))
end

function StartingArea(x_min::VehicleState{T}, x_max::VehicleState{T}) where {T}
    return StartingArea{T}(0.5*(x_min+x_max), x_min, x_max)
end

function randstate(start::StartingArea{T}) where {T}
    x_min = start.x_min
    x_max = start.x_max
    return VehicleState{T}(
        x_min.x + rand(T)*(x_max.x - x_min.x),
        x_min.y + rand(T)*(x_max.y - x_min.y),
        x_min.θ + rand(T)*(x_max.θ - x_min.θ),
        x_min.v + rand(T)*(x_max.v - x_min.v),
        )
end


import DataStructures
import XMLDict

function getXML(path)
    s = open(path) do file
        read(file,String)
    end
    return xml = XMLDict.xml_dict(s)

end

# get potentially missing value out of dict
function getVal(dict,key)
    try
        return val = parse(Float64,dict[key])
    catch y
        if isa(y,KeyError)
            return val = missing
        else
            throw(y)
        end
    end
end

struct XTbMLTable
    select::DataStructures.DefaultOrderedDict
    ultimate::DataStructures.DefaultOrderedDict
end

function loadXTbMLTable(path)

    x = getXML(path)
    name = x["XTbML"]["ContentClassification"]["TableName"]
    tbl = XTbMLTable(DataStructures.DefaultOrderedDict(missing),DataStructures.DefaultOrderedDict(missing))

    if isa(x["XTbML"]["Table"],Vector)
        # for a select and ultimate table, will have multiple tables
        # parsed into a vector of tables
        for ai in x["XTbML"]["Table"][1]["Values"]["Axis"]
            issue_age = parse(Int,ai[:t])
            tbl.select[issue_age] = DataStructures.DefaultOrderedDict(missing)
            for aj in ai["Axis"]["Y"]
                duration = parse(Int,aj[:t])
                rate = getVal(aj,"")
                if !ismissing(rate)
                    tbl.select[issue_age][duration]= rate
                end
            end
        end

        for ai in x["XTbML"]["Table"][2]["Values"]["Axis"]["Y"]
            age = parse(Int,ai[:t])
            rate = getVal(ai,"")
            if !ismissing(rate)
                tbl.ultimate[age]= rate
            end
        end
    elseif isa(x["XTbML"]["Table"],DataStructures.OrderedDict)
        # a table without select period will just have one set of values, which
        # are loaded into the OrderedDict

        for ai in x["XTbML"]["Table"]["Values"]["Axis"]["Y"]
            age = parse(Int,ai[:t])
            rate = getVal(ai,"")
            if !ismissing(rate)
                tbl.ultimate[age]= rate
            end
        end
    else
        error("don't know how to handle table: "  *  name)
    end

    return (tbl, name)
end


module Graph2OpCode

import Basic.push!

function push!{C}(ins::Vector{Instruction{C}}, s::Sequence{C})
    for p in s.patterns
        push!(ins, p)
    end
end

function push!{C}(ins::Vector{Instruction{C}}, r::Range{C})
    if r.lo == r.hi
        push!(ins, Instruction{C}(EQ, 0, r.lo))
    else
        push!(ins, Instruction{C}(GE, 0, r.lo))
        push!(ins, Instruction{C}(LE, 0, r.hi))
    end
end

function push!{C}(ins::Vector{Instruction{C}}, c::Choice{C})
    jlast = length(c.patterns)
    for (j, p) in enumerate(c.patterns)
        if j != jlast
            i = Instruction{C}(ALT, 0, zero(C))
            push!(ins, p)
            i.n = length(ins) + 1
        else
            push!(ins, p)
        end
    end
end

function push!{C}(ins::Vector{Instruction{C}}, r::Repeat{C})

end

end


include("../src/iPredict.jl")
#include("src/iPredict.jl")
using .iPredict
using Test

#An extremely general systems test - make sure the output is a dictionary!
result = globalPredict("test/data/small_counts.txt",
"test/data/small_gene_isoform_relationships.txt", 1, true)

@test isa(result, Dict) == true

#Check an error is thrown when the gene names don't match
@test_throws ArgumentError globalPredict("test/data/test_counts_matrix.txt",
"test/data/all_isoforms.txt", 1, false)

#Do some general systems checks for oneGenePrediction too
result = oneGenePrediction("test/data/small_counts.txt",
"test/data/small_gene_isoform_relationships.txt", "G1", 10, true)

@test isa(result, Dict) == true
@test length(result) == 1

result = oneGenePrediction("test/data/alt_small_counts.txt",
"test/data/small_gene_isoform_relationships.txt", "G1", 10, true)

@test isa(result, Dict) == true
@test length(result) == 2


module Experimental

include("AlgebraicNets.jl")
include("MathFormulas.jl")
include("Markov.jl")

end


module PersistenceDiagramsBase

using Compat
using Tables

export PersistenceDiagram,
    PersistenceInterval,
    birth,
    death,
    persistence,
    midlife,
    dim,
    threshold,
    birth_simplex,
    death_simplex,
    representative

include("intervals.jl")
include("diagrams.jl")
include("tables.jl")

end


function iscompatible(system::SystemNamelist, cell_parameters::CellParametersCard)
    ibrav, celldm = system.ibrav, system.celldm
    if iszero(ibrav)
        if optionof(cell_parameters) in ("bohr", "angstrom")
            return all(iszero, celldm)
        else  # "alat"
            return !iszero(first(celldm))  # first(celldm) != 0
        end
    else
        return false
    end
end # function iscompatible
iscompatible(x::CellParametersCard, y::SystemNamelist) = iscompatible(y, x)


include("./ADMM.jl")


# Not in FiniteDiff because `u` -> scalar isn't used anywhere else,
# but could be upstreamed.
mutable struct UGradientWrapper{fType,tType,P} <: Function
  f::fType
  t::tType
  p::P
end

(ff::UGradientWrapper)(uprev) = ff.f(uprev,ff.p,ff.t)

mutable struct ParamGradientWrapper{fType,tType,uType} <: Function
  f::fType
  t::tType
  u::uType
end

(ff::ParamGradientWrapper)(p) = ff.f(ff.u,p,ff.t)

# the next four definitions are only needed in case of non-diagonal SDEs

mutable struct ParamNonDiagNoiseGradientWrapper{fType,tType,uType} <: Function
  f::fType
  t::tType
  u::uType
end

(ff::ParamNonDiagNoiseGradientWrapper)(p) = vec(ff.f(ff.u,p,ff.t))

mutable struct ParamNonDiagNoiseJacobianWrapper{fType,tType,uType,duType} <: Function
  f::fType
  t::tType
  u::uType
  du::duType
end

function (ff::ParamNonDiagNoiseJacobianWrapper)(p)
  du1 = similar(p, size(ff.du))
  ff.f(du1,ff.u,p,ff.t)
  return vec(du1)
end

function (ff::ParamNonDiagNoiseJacobianWrapper)(du1,p)
  ff.f(du1,ff.u,p,ff.t)
  return vec(du1)
end

mutable struct UNonDiagNoiseGradientWrapper{fType,tType,P} <: Function
  f::fType
  t::tType
  p::P
end

(ff::UNonDiagNoiseGradientWrapper)(uprev) = vec(ff.f(uprev,ff.p,ff.t))

mutable struct UNonDiagNoiseJacobianWrapper{fType,tType,P,duType} <: Function
  f::fType
  t::tType
  p::P
  du::duType
end

(ff::UNonDiagNoiseJacobianWrapper)(uprev) = (du1 = similar(ff.du); ff.f(du1,uprev,ff.p,ff.t); vec(du1))

function (ff::UNonDiagNoiseJacobianWrapper)(du1,uprev)
  ff.f(du1,uprev,ff.p,ff.t)
  return vec(du1)
end

Base.@pure function determine_chunksize(u,alg::DiffEqBase.AbstractSensitivityAlgorithm)
  determine_chunksize(u,get_chunksize(alg))
end

Base.@pure function determine_chunksize(u,CS)
  if CS != 0
    return CS
  else
    return ForwardDiff.pickchunksize(length(u))
  end
end

function jacobian(f, x::AbstractArray{<:Number}, alg::DiffEqBase.AbstractSensitivityAlgorithm)
  if alg_autodiff(alg)
    J = ForwardDiff.jacobian(f, x)
  else
    J = FiniteDiff.finite_difference_jacobian(f, x)
  end
  return J
end

function jacobian!(J::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number},
                   fx::Union{Nothing,AbstractArray{<:Number}}, alg::DiffEqBase.AbstractSensitivityAlgorithm, jac_config)
  if alg_autodiff(alg)
    if fx === nothing
      ForwardDiff.jacobian!(J, f, x)
    else
      ForwardDiff.jacobian!(J, f, fx, x, jac_config)
    end
  else
    FiniteDiff.finite_difference_jacobian!(J, f, x, jac_config)
  end
  nothing
end

function gradient!(df::AbstractArray{<:Number}, f,
                   x::Union{Number,AbstractArray{<:Number}},
                   alg::DiffEqBase.AbstractSensitivityAlgorithm, grad_config)
    if alg_autodiff(alg)
        ForwardDiff.gradient!(df, f, x, grad_config)
    else
        FiniteDiff.finite_difference_gradient!(df, f, x, grad_config)
    end
    nothing
end

"""
  jacobianvec!(Jv, f, x, v, alg, (buffer, seed)) -> nothing

``Jv <- J(f(x))v``
"""
function jacobianvec!(Jv::AbstractArray{<:Number}, f, x::AbstractArray{<:Number},
                      v, alg::DiffEqBase.AbstractSensitivityAlgorithm, config)
  if alg_autodiff(alg)
    buffer, seed = config
    TD = typeof(first(seed))
    T  = typeof(first(seed).partials)
    @. seed = TD(x, T(tuple(v)))
    f(buffer, seed)
    Jv .= ForwardDiff.partials.(buffer, 1)
  else
    buffer1, buffer2 = config
    f(buffer1,x)
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    @. x += ϵ*v
    f(buffer2,x)
    @. x -= ϵ*v
    @. du = (buffer2 - buffer1)/ϵ
  end
  nothing
end

function vecjacobian!(dλ, y, λ, p, t, S::SensitivityFunction;
                      dgrad=nothing, dy=nothing)
  _vecjacobian!(dλ, y, λ, p, t, S, S.sensealg.autojacvec, dgrad, dy)
  return
end

function _vecjacobian!(dλ, y, λ, p, t, S::SensitivityFunction, isautojacvec::Bool, dgrad, dy)
  @unpack sensealg, f = S
  prob = getprob(S)
  if isautojacvec isa Bool && !isautojacvec
    @unpack J, uf, f_cache, jac_config = S.diffcache
    if !(prob isa DiffEqBase.SteadyStateProblem)
      if DiffEqBase.has_jac(f)
        f.jac(J,y,p,t) # Calculate the Jacobian into J
      else
        uf.t = t
        uf.p = p
        jacobian!(J, uf, y, f_cache, sensealg, jac_config)
      end
      mul!(dλ',λ',J)
    end
    if dgrad !== nothing
      @unpack pJ, pf, paramjac_config = S.diffcache
      if DiffEqBase.has_paramjac(f)
        # Calculate the parameter Jacobian into pJ
        f.paramjac(pJ,y,prob.p,t)
      else
        pf.t = t
        pf.u = y
        if DiffEqBase.isinplace(prob)
          jacobian!(pJ, pf, prob.p, f_cache, sensealg, paramjac_config)
        else
          temp = jacobian(pf, prob.p, sensealg)
          pJ .= temp
        end
      end
      mul!(dgrad',λ',pJ)
    end
    if dy !== nothing
      if DiffEqBase.isinplace(prob)
         f(dy, y, p, t)
      else
        dy[:] .= vec(f(y, p, t))
      end
    end
  elseif DiffEqBase.isinplace(prob)
    _vecjacobian!(dλ, y, λ, p, t, S, ReverseDiffVJP(), dgrad, dy)
  else
    _vecjacobian!(dλ, y, λ, p, t, S, ZygoteVJP(), dgrad, dy)
  end
  return
end

function _vecjacobian!(dλ, y, λ, p, t, S::SensitivityFunction, isautojacvec::TrackerVJP, dgrad, dy)
  @unpack sensealg, f = S
  prob = getprob(S)
  isautojacvec = get_jacvec(sensealg)
  if DiffEqBase.isinplace(prob)
    _dy, back = Tracker.forward(y, prob.p) do u, p
      out_ = map(zero, u)
      f(out_, u, p, t)
      Tracker.collect(out_)
    end
    tmp1, tmp2 = Tracker.data.(back(λ))
    dλ[:] .= vec(tmp1)
    dgrad !== nothing && (dgrad[:] .= vec(tmp2))
    dy !== nothing && (dy[:] .= vec(Tracker.data(_dy)))
  else
    _dy, back = Tracker.forward(y, prob.p) do u, p
      Tracker.collect(f(u, p, t))
    end
    tmp1, tmp2 = Tracker.data.(back(λ))
    dλ[:] .= vec(tmp1)
    dgrad !== nothing && (dgrad[:] .= vec(tmp2))
    dy !== nothing && (dy[:] .= vec(Tracker.data(_dy)))
  end
  return
end

function _vecjacobian!(dλ, y, λ, p, t, S::SensitivityFunction, isautojacvec::ReverseDiffVJP, dgrad, dy)
  @unpack sensealg, f = S
  prob = getprob(S)
  isautojacvec = get_jacvec(sensealg)

  if typeof(prob) <: SteadyStateProblem || (eltype(λ) <: eltype(prob.u0) && typeof(t) <: eltype(prob.u0) && compile_tape(sensealg.autojacvec))
    tape = S.diffcache.paramjac_config

  ## These other cases happen due to autodiff in stiff ODE solvers
  elseif DiffEqBase.isinplace(prob)
    _y = eltype(y) === eltype(λ) ? y : convert.(promote_type(eltype(y),eltype(λ)),y)
    tape = ReverseDiff.GradientTape((_y, prob.p, [t])) do u,p,t
      du1 = similar(u, size(u))
      f(du1,u,p,first(t))
      return vec(du1)
     end
  else
    _y = eltype(y) === eltype(λ) ? y : convert.(promote_type(eltype(y),eltype(λ)),y)
    tape = ReverseDiff.GradientTape((_y, prob.p, [t])) do u,p,t
      vec(f(u,p,first(t)))
    end
  end

  if prob isa DiffEqBase.SteadyStateProblem
    tu, tp = ReverseDiff.input_hook(tape)
  else
    tu, tp, tt = ReverseDiff.input_hook(tape)
  end
  output = ReverseDiff.output_hook(tape)
  ReverseDiff.unseed!(tu) # clear any "leftover" derivatives from previous calls
  ReverseDiff.unseed!(tp)
  if !(prob isa DiffEqBase.SteadyStateProblem)
    ReverseDiff.unseed!(tt)
  end
  ReverseDiff.value!(tu, y)
  ReverseDiff.value!(tp, prob.p)
  if !(prob isa DiffEqBase.SteadyStateProblem)
    ReverseDiff.value!(tt, [t])
  end
  ReverseDiff.forward_pass!(tape)
  ReverseDiff.increment_deriv!(output, λ)
  ReverseDiff.reverse_pass!(tape)
  copyto!(vec(dλ), ReverseDiff.deriv(tu))
  dgrad !== nothing && copyto!(vec(dgrad), ReverseDiff.deriv(tp))
  ReverseDiff.pull_value!(output)
  dy !== nothing && copyto!(vec(dy), ReverseDiff.value(output))
  return
end

function _vecjacobian!(dλ, y, λ, p, t, S::SensitivityFunction, isautojacvec::ZygoteVJP, dgrad, dy)
  @unpack sensealg, f = S
  prob = getprob(S)

  isautojacvec = get_jacvec(sensealg)
  if DiffEqBase.isinplace(prob)
    _dy, back = Zygote.pullback(y, prob.p) do u, p
      out_ = Zygote.Buffer(similar(u))
      f(out_, u, p, t)
      vec(copy(out_))
    end
    tmp1,tmp2 = back(λ)
    dλ[:] .= vec(tmp1)
    dgrad !== nothing && tmp2 != nothing && (dgrad[:] .= vec(tmp2))
    dy !== nothing && (dy[:] .= vec(_dy))
  else
    _dy, back = Zygote.pullback(y, prob.p) do u, p
      vec(f(u, p, t))
    end
    tmp1,tmp2 = back(λ)
    dλ[:] .= vec(tmp1)

    dy !== nothing && (dy[:] .= vec(_dy))
    dgrad !== nothing && tmp2 != nothing && (dgrad[:] .= vec(tmp2))
  end
  return
end

function jacNoise!(λ, y, p, t, S::SensitivityFunction;
                      dgrad=nothing, dλ=nothing, dy=nothing)
  _jacNoise!(λ, y, p, t, S, S.sensealg.noise, dgrad, dλ, dy)
  return
end

function _jacNoise!(λ, y, p, t, S::SensitivityFunction, isnoise::Bool, dgrad, dλ, dy)
  @unpack sensealg, f = S
  prob = getprob(S)
  if isnoise isa Bool && !isnoise
    if dgrad !== nothing
      @unpack pJ, pf, f_cache, paramjac_noise_config = S.diffcache
      if DiffEqBase.has_paramjac(f)
        # Calculate the parameter Jacobian into pJ
        f.paramjac(pJ,y,prob.p,t)
      else
        pf.t = t
        pf.u = y
        if DiffEqBase.isinplace(prob)
          jacobian!(pJ, pf, prob.p, nothing, sensealg, nothing)
          #jacobian!(pJ, pf, prob.p, f_cache, sensealg, paramjac_noise_config)
        else
          temp = jacobian(pf, prob.p, sensealg)
          pJ .= temp
        end
      end

      if StochasticDiffEq.is_diagonal_noise(prob)
        pJt = transpose(λ).*transpose(pJ)
        dgrad[:] .= vec(pJt)
      else
        m = size(prob.noise_rate_prototype)[2]
        for i in 1:m
          tmp = λ'*pJ[(i-1)*m+1:i*m,:]
          dgrad[:,i] .= vec(tmp)
        end
      end
    end

    if dλ !== nothing && (isnoisemixing(sensealg) || !StochasticDiffEq.is_diagonal_noise(prob))
      @unpack J, uf, f_cache, jac_noise_config = S.diffcache
      if dy!=nothing
        if DiffEqBase.isinplace(prob)
          f(dy, y, p, t)
        else
          dy .= f(y, p, t)
        end
      end

      if DiffEqBase.has_jac(f)
        f.jac(J,y,p,t) # Calculate the Jacobian into J
      else
        if DiffEqBase.isinplace(prob)
          if dy != nothing
            ForwardDiff.jacobian!(J,uf,dy,y)
          else
            if StochasticDiffEq.is_diagonal_noise(prob)
              dy = similar(y)
            else
              dy = similar(prob.noise_rate_prototype)
              f(dy, y, p, t)
              ForwardDiff.jacobian!(J,uf,dy,y)
            end
            f(dy, y, p, t)
            ForwardDiff.jacobian!(J,uf,dy,y)
          end
        else
          tmp = ForwardDiff.jacobian(uf,y)
          J .= tmp
        end
      #  uf.t = t
      #  uf.p = p
      #  jacobian!(J, uf, y, nothing, sensealg, nothing)
      end

      if StochasticDiffEq.is_diagonal_noise(prob)
        Jt = transpose(λ).*transpose(J)
        dλ[:] .= vec(Jt)
      else
        for i in 1:m
          tmp = λ'*J[(i-1)*m+1:i*m,:]
          dλ[:,i] .= vec(tmp)
        end
      end
    end
  elseif DiffEqBase.isinplace(prob)
    _jacNoise!(λ, y, p, t, S, ReverseDiffNoise(), dgrad, dλ, dy)
  else
    _jacNoise!(λ, y, p, t, S, ZygoteNoise(), dgrad, dλ, dy)
  end
  return
end

function _jacNoise!(λ, y, p, t, S::SensitivityFunction, isnoise::ReverseDiffNoise, dgrad, dλ, dy)
  @unpack sensealg, f = S
  prob = getprob(S)

  for (i, λi) in enumerate(λ)
    tapei = S.diffcache.paramjac_noise_config[i]
    tu, tp, tt = ReverseDiff.input_hook(tapei)
    output = ReverseDiff.output_hook(tapei)
    ReverseDiff.unseed!(tu) # clear any "leftover" derivatives from previous calls
    ReverseDiff.unseed!(tp)
    ReverseDiff.unseed!(tt)
    ReverseDiff.value!(tu, y)
    ReverseDiff.value!(tp, p)
    ReverseDiff.value!(tt, [t])
    ReverseDiff.forward_pass!(tapei)
    if StochasticDiffEq.is_diagonal_noise(prob)
      ReverseDiff.increment_deriv!(output, λi)
    else
      ReverseDiff.increment_deriv!(output, λ)
    end
    ReverseDiff.reverse_pass!(tapei)

    deriv = ReverseDiff.deriv(tp)
    dgrad[:,i] .= vec(deriv)
    ReverseDiff.pull_value!(output)

    if StochasticDiffEq.is_diagonal_noise(prob)
      dλ !== nothing && (dλ[:,i] .= vec(ReverseDiff.deriv(tu)))
      dy !== nothing && (dy[i] = ReverseDiff.value(output))
    else
      dλ !== nothing && (dλ[:,i] .= vec(ReverseDiff.deriv(tu)))
      dy !== nothing && (dy[:,i] .= vec(ReverseDiff.value(output)))
    end
  end
  return
end

function _jacNoise!(λ, y, p, t, S::SensitivityFunction, isnoise::ZygoteNoise, dgrad, dλ, dy)
  @unpack sensealg, f = S
  prob = getprob(S)

  if StochasticDiffEq.is_diagonal_noise(prob)
    if DiffEqBase.isinplace(prob)
      for (i, λi) in enumerate(λ)
        _dy, back = Zygote.pullback(y, prob.p) do u, p
          out_ = Zygote.Buffer(similar(u))
          f(out_, u, p, t)
          copy(out_[i])
        end
        tmp1,tmp2 = back(λi) #issue: tmp2 = zeros(p)
        dgrad[:,i] .= vec(tmp2)
        dλ !== nothing && (dλ[:,i] .= vec(tmp1))
        dy !== nothing && (dy[i] = _dy)
      end
    else
      for (i, λi) in enumerate(λ)
        _dy, back = Zygote.pullback(y, prob.p) do u, p
          f(u, p, t)[i]
        end
        tmp1,tmp2 = back(λi)
        dgrad[:,i] .= vec(tmp2)
        dλ !== nothing && (dλ[:,i] .= vec(tmp1))
        dy !== nothing && (dy[i] = _dy)
      end
    end
  else
    if DiffEqBase.isinplace(prob)
      for (i, λi) in enumerate(λ)
        _dy, back = Zygote.pullback(y, prob.p) do u, p
          out_ = Zygote.Buffer(similar(prob.noise_rate_prototype))
          f(out_, u, p, t)
          copy(out_[:,i])
        end
        tmp1,tmp2 = back(λ)#issue with Zygote.Buffer
        dgrad[:,i] .= vec(tmp2)
        dλ !== nothing && (dλ[:,i] .= vec(tmp1))
        dy !== nothing && (dy[:,i] .= vec(_dy))
      end
    else
      for (i, λi) in enumerate(λ)
         _dy, back = Zygote.pullback(y, prob.p) do u, p
          f(u, p, t)[:,i]
        end
        tmp1,tmp2 = back(λ)
        dgrad[:,i] .= vec(tmp2)
        dλ !== nothing && (dλ[:,i] .= vec(tmp1))
        dy !== nothing && (dy[:,i] .= vec(_dy))
      end
    end
  end
  return
end

function accumulate_cost!(dλ, y, p, t, S::SensitivityFunction, dgrad=nothing)
  @unpack dg, dg_val, g, g_grad_config = S.diffcache
  if dg != nothing
    if !(dg isa Tuple)
      dg(dg_val,y,p,t)
      dλ .+= vec(dg_val)
    else
      dg[1](dg_val[1],y,p,t)
      dλ .+= vec(dg_val[1])
      if dgrad !== nothing
        dg[2](dg_val[2],y,p,t)
        dgrad .-= vec(dg_val[2])
      end
    end
  else
    g.t = t
    gradient!(dg_val, g, y, S.sensealg, g_grad_config)
    dλ .+= vec(dg_val)
  end
  return nothing
end

function build_jac_config(alg,uf,u)
  if alg_autodiff(alg)
    jac_config = ForwardDiff.JacobianConfig(uf,u,u,
                 ForwardDiff.Chunk{determine_chunksize(u,alg)}())
  else
    if diff_type(alg) != Val{:complex}
      jac_config = FiniteDiff.JacobianCache(similar(u),similar(u),
                                                 similar(u),diff_type(alg))
    else
      tmp = Complex{eltype(u)}.(u)
      du1 = Complex{eltype(u)}.(du1)
      jac_config = FiniteDiff.JacobianCache(tmp,du1,nothing,diff_type(alg))
    end
  end
  jac_config
end

function build_param_jac_config(alg,uf,u,p)
  if alg_autodiff(alg)
    jac_config = ForwardDiff.JacobianConfig(uf,u,p,
                 ForwardDiff.Chunk{determine_chunksize(p,alg)}())
  else
    if diff_type(alg) != Val{:complex}
      jac_config = FiniteDiff.JacobianCache(similar(p),similar(u),
                                                 similar(u),diff_type(alg))
    else
      tmp = Complex{eltype(p)}.(p)
      du1 = Complex{eltype(u)}.(u)
      jac_config = FiniteDiff.JacobianCache(tmp,du1,nothing,diff_type(alg))
    end
  end
  jac_config
end

function build_grad_config(alg,tf,du1,t)
  if alg_autodiff(alg)
    grad_config = ForwardDiff.GradientConfig(tf,du1,
                                    ForwardDiff.Chunk{determine_chunksize(du1,alg)}())
  else
    grad_config = FiniteDiff.GradientCache(du1,t,diff_type(alg))
  end
  grad_config
end


using HierarchicalUtils: encode, stringify

"""
    AbstractMillModel

Supertype for any model defined in `Mill.jl`.
"""
abstract type AbstractMillModel end

const MillFunction = Union{Dense, Chain, Function}

_make_array_model(m::Union{MillFunction, Function}) = ArrayModel(m)
_make_array_model(m) = m

include("arraymodel.jl")
include("bagmodel.jl")
include("productmodel.jl")
include("lazymodel.jl")
include("reflectinmodel.jl")


module GW
    include("../../examples/gridworld.jl")
end

@testset "gridworld" begin
    env = GW.GridWorld()
    reset!(env)
    while !terminated(env)
        act!(env, rand(actions(env)))
        render(env)
    end
    @test haskey(env.rewards, state(env))

    env2 = GW.GridWorld()
    @test state(clone(env2)) == state(env2)
    setstate!(env, state(env2))
    @test state(env) == state(env2)
    @test observe(env) == observe(env2)
    @test observations(env) == observations(env2)
end


"""
    AbstractDistributionFitOptimizer
    OptimOptimizer <: AbstractDistributionFitOptimizer
    set_optimizer(::AbstractDistributionFitOptimizer)

`DistributionFits.jl` uses the following interface to opimize an univariate
function `f` on bounded interval `[lower,upper]`:

    optimize(f, ::AbstractDistributionFitOptimizer, lower, upper)

Returning an object with fields `minimizer` and `converged`.

`DistributionFits.set_optimizer(::AbstractDistributionFitOptimizer)`
sets the optimizer used in calling the optimize function.
Specializing this function with a concrete type, allows using different
optimization packages.

By default, when the [Optim.jl](https://julianlsolvers.github.io/Optim.jl/stable/)
package is in scope, the
`OptimOptimizer` is set, which implements the interface by using Optim's
[optimize](https://julianlsolvers.github.io/Optim.jl/stable/#user/minimization/#minimizing-a-univariate-function-on-a-bounded-interval) function.

Hence, the module using DistribuitonFits.jl has to either

- explicitly invoke `using Optim`, or
- call `set_optimizer` with the concrete subtype of `AbstractDistributionFitOptimizer`
  for which the corresponding optimize method is implemented.
"""
abstract type AbstractDistributionFitOptimizer end
struct OptimOptimizer <:  AbstractDistributionFitOptimizer; end
struct NotSetOptimizer <: AbstractDistributionFitOptimizer; end

optimize(f, ::NotSetOptimizer, lower, upper) = error(
    "Optimizer not set yet. Either invoke 'using Optim' or 'DistributionFits.set_optimizer(...)'.")

optimizer = NotSetOptimizer();
optimize(f, lower, upper) = optimize(f, optimizer, lower, upper)
set_optimizer(opt::AbstractDistributionFitOptimizer) = (global optimizer = opt)


function _find_ranges_for_closeback!(ranges, x, y, _fl::F1, _fr::F2, ::Val{T1}, ::Val{T2}; threads = true) where {T1, T2, F1, F2}
    @_threadsfor threads for i in 1:length(x)
        curr_start = ranges[i].start
        if length(ranges[i]) == 0
            ranges[i] = 0:0
        else
            ranges[i] = 1:searchsortedlast_join(_fr, y, _fl(x[i])::T1, ranges[i].start, ranges[i].stop, Base.Order.Forward, Val(T2))
            if ranges[i].stop < curr_start
                ranges[i] = 0:curr_start
            end
        end
    end
end

function _find_ranges_for_closefor!(ranges, x, y, _fl::F1, _fr::F2, ::Val{T1}, ::Val{T2}; threads = true) where {T1, T2, F1, F2}
    @_threadsfor threads for i in 1:length(x)
        cur_stop = ranges[i].stop
        if length(ranges[i]) == 0
            ranges[i] = 0:0
        else
            ranges[i] = 1:searchsortedfirst_join(_fr, y, _fl(x[i])::T1, ranges[i].start, ranges[i].stop, Base.Order.Forward, Val(T2))
            if ranges[i].stop > cur_stop
                ranges[i] = 0:cur_stop
            end
        end
    end
end

function _find_ranges_for_closenearest!(ranges, x, y, _fl::F1, _fr::F2, ::Val{T1}, ::Val{T2}; threads = true) where {T1, T2, F1, F2}
    @_threadsfor threads for i in 1:length(x)
        cur_stop = ranges[i].stop
        curr_start = ranges[i].start
        if length(ranges[i]) == 0
            ranges[i] = 0:0
        else
            #TODO we don't need to search whole range for the second search
            fval = searchsortedfirst_join(_fr, y, _fl(x[i])::T1, ranges[i].start, ranges[i].stop, Base.Order.Forward, Val(T2))
            lval = searchsortedlast_join(_fr, y, _fl(x[i])::T1, ranges[i].start, ranges[i].stop, Base.Order.Forward, Val(T2))

            if fval > cur_stop
                ranges[i] = 0:lval
            elseif lval < curr_start
                ranges[i] = 0:fval
            else
                ranges[i] = min(fval, lval):max(fval, lval)
            end
        end
    end
end

_IST_(x,y)=true

function _close_notol_nn!(_res, ranges, x, rnn, lnn, bordervalue; threads = true)
    @_threadsfor threads for i in 1:length(ranges)
        if ranges[i] == 0:0
            _res[i] = missing
        else
            r1 = rnn[ranges[i].stop]
            if ranges[i].start == 0
                if ismissing(bordervalue)
                    _res[i] = missing
                else
                    _res[i] = x[ranges[i].stop]
                end
            else
                r2 = rnn[ranges[i].start]
                if isless(abs(r2-lnn[i]), abs(r1-lnn[i]))
                    _res[i] = x[ranges[i].start]
                else
                    _res[i] = x[ranges[i].stop]
                end
            end
        end
    end
end
# instead of finding the nearest, set res = op(x[left], x[right]), in border it sets res = op(missing, x[border])
function _close_notol_nn!(_res, ranges, x, rnn, lnn, bordervalue, op::F; threads = true) where F
    @_threadsfor threads for i in 1:length(ranges)
        if ranges[i] == 0:0
            _res[i] = missing
        else
            if ranges[i].start == 0
                _res[i] = op(missing, x[ranges[i].stop])
            else
                _res[i] = op(x[ranges[i].start],x[ranges[i].stop])
            end
        end
    end
end

function _close_notol_noborder!(_res, ranges, x, direction; threads = true)
    @_threadsfor threads for i in 1:length(ranges)
        if ranges[i] == 0:0
            _res[i] = missing
        else
            if ranges[i].start == 0
                _res[i] = missing
            else

                if direction == :backward
                    _res[i] = x[ranges[i].start]
                else
                    _res[i] = x[ranges[i].stop]
                end
            end
        end
    end
end
function _close_notol_border!(_res, ranges, x, bordervalue; threads = true)
    @_threadsfor threads for i in 1:length(ranges)
        if ranges[i] == 0:0
            _res[i] = missing
        else
            _res[i] = x[ranges[i].stop]
            if !bordervalue && ranges[i].start == 0
                _res[i] = missing
            end
        end
    end
end

function _close_tol_nn!(_res, ranges, x, rnn, lnn, tol, aem::F, bordervalue; threads = true) where F
    @_threadsfor threads for i in 1:length(ranges)
        if ranges[i] == 0:0
            _res[i] = missing
        else
            r1 = rnn[ranges[i].stop]
            if ranges[i].start == 0
                if ismissing(bordervalue)
                    _res[i] = missing
                else
                    if isless(abs(r1-lnn[i]), tol) && aem(r1, lnn[i])
                        _res[i] = x[ranges[i].stop]
                    else
                        _res[i] = missing
                    end
                end
            else
                r2 = rnn[ranges[i].start]
                if isless(abs(r2-lnn[i]), abs(r1-lnn[i]))
                    if isless(abs(r2-lnn[i]), tol) && aem(r2, lnn[i])
                        _res[i] = x[ranges[i].start]
                    else
                        _res[i] = missing
                    end
                else
                    if isless(abs(r1-lnn[i]), tol) && aem(r1, lnn[i])
                        _res[i] = x[ranges[i].stop]
                    else
                        _res[i] = missing
                    end
                end
            end
        end
    end
end

# instead of finding the nearest, set res = op(x[left], x[right]), in border it sets res = op(missing, x[border]), op(missing, missing) must be defined
function _close_tol_nn!(_res, ranges, x, rnn, lnn, tol, aem::F1, bordervalue, op::F2; threads = true) where {F1, F2}
    @_threadsfor threads for i in 1:length(ranges)
        if ranges[i] == 0:0
            _res[i] = missing
        else
            r1 = rnn[ranges[i].stop]
            if ranges[i].start == 0
                if isless(abs(r1-lnn[i]), tol) && aem(r1, lnn[i])
                    _res[i] = op(missing, x[ranges[i].stop])
                else
                    _res[i] = op(missing, missing)
                end
            else
                r2 = rnn[ranges[i].start]
                if isless(abs(r2-lnn[i]), tol) && aem(r2, lnn[i])
                    if isless(abs(r1-lnn[i]), tol) && aem(r1, lnn[i])
                        _res[i] = op(x[ranges[i].start], x[ranges[i].stop])
                    else
                        _res[i] = op(x[ranges[i].start], missing)
                    end
                else
                    if isless(abs(r1-lnn[i]), tol) && aem(r1, lnn[i])
                        _res[i] = op(missing, x[ranges[i].stop])
                    else
                        _res[i] = op(missing, missing)
                    end
                end
            end
        end
    end
end

function _close_tol_noborder!(_res, ranges, x, rnn, lnn, tol, aem::F, direction; threads = true) where F
    @_threadsfor threads for i in 1:length(ranges)
        if ranges[i] == 0:0
            _res[i] = missing
        else
            if ranges[i].start == 0
                _res[i] = missing
            else
                if direction == :backward
                    r1 = rnn[ranges[i].start]
                    if isless(abs(r1-lnn[i]), tol) && aem(r1, lnn[i])
                        _res[i] = x[ranges[i].start]
                    else
                        _res[i] = missing
                    end
                else
                    r1 = rnn[ranges[i].stop]
                    if isless(abs(r1-lnn[i]), tol) && aem(r1, lnn[i])
                        _res[i] = x[ranges[i].stop]
                    else
                        _res[i] = missing
                    end
                end
            end
        end
    end
end
function _close_tol_border!(_res, ranges, x, rnn, lnn, tol, aem::F, bordervalue; threads = true) where F
    @_threadsfor threads for i in 1:length(ranges)
        if ranges[i] == 0:0
            _res[i] = missing
        else
            r1 = rnn[ranges[i].stop]
            if isless(abs(r1-lnn[i]), tol) && aem(r1, lnn[i])
                _res[i] = x[ranges[i].stop]
            else
                _res[i] = missing
            end
            if !bordervalue && ranges[i].start == 0
                _res[i] = missing
            end
        end
    end
end

function  _fill_right_cols_table_close!(_res, x, ranges, total, borderval, fill_val, direction; op = nothing, nn = false, rnn = nothing, lnn = nothing, tol = nothing, aem = _IST_, threads = true)
    if borderval == :nearest
        bordervalue = true
    elseif borderval == :none
        bordervalue = missing
    elseif borderval == :missing
        bordervalue = false
    end
    if tol === nothing
        if nn
            if op === nothing
                _close_notol_nn!(_res, ranges, x, rnn, lnn, bordervalue; threads = threads)
            else
                _close_notol_nn!(_res, ranges, x, rnn, lnn, bordervalue, op; threads = threads)
            end
        else
            if ismissing(bordervalue)
                _close_notol_noborder!(_res, ranges, x, direction; threads = threads)
            else
                _close_notol_border!(_res, ranges, x, bordervalue; threads = threads)
            end
        end
    else
        if nn
            if op === nothing
                _close_tol_nn!(_res, ranges, x, rnn, lnn, tol, aem, bordervalue; threads = threads)
            else
                _close_tol_nn!(_res, ranges, x, rnn, lnn, tol, aem, bordervalue, op; threads = threads)
            end
        else
            if ismissing(bordervalue)
                _close_tol_noborder!(_res, ranges, x, rnn, lnn, tol, aem, direction; threads = threads)
            else
                _close_tol_border!(_res, ranges, x, rnn, lnn, tol, aem, bordervalue; threads = threads)
            end
        end
    end

end

function _change_refpool_find_range_for_close!(ranges, dsl, dsr, r_perms, oncols_left, oncols_right, direction, lmf, rmf, j; nsfpaj = true, threads = true)
    var_l = _columns(dsl)[oncols_left[j]]
    var_r = _columns(dsr)[oncols_right[j]]
    l_idx = oncols_left[j]
    r_idx = oncols_right[j]
    if lmf
        _fl = getformat(dsl, l_idx)
    else
        _fl = identity
    end
    if rmf
        _fr = getformat(dsr, r_idx)
    else
        _fr = identity
    end

    T1 = Core.Compiler.return_type(_fl, Tuple{eltype(var_l)})

    if DataAPI.refpool(var_r) !== nothing && nsfpaj
        true && throw(ErrorException("we shouldn't end up here"))
    else
        T2 = Core.Compiler.return_type(_fr, Tuple{eltype(var_r)})
        if direction == :backward
            _find_ranges_for_closeback!(ranges, var_l, view(var_r, r_perms), _fl, _fr, Val(T1), Val(T2); threads = threads)
        elseif direction == :forward
            _find_ranges_for_closefor!(ranges, var_l, view(var_r, r_perms), _fl, _fr, Val(T1), Val(T2); threads = threads)
        elseif direction == :nearest
            _find_ranges_for_closenearest!(ranges, var_l, view(var_r, r_perms), _fl, _fr, Val(T1), Val(T2); threads = threads)
        end
    end
end

# border = :nearest | :missing | :none
function _join_closejoin(dsl, dsr::AbstractDataset, ::Val{T}; onleft, onright, makeunique = false, border = :nearest, mapformats = [true, true], stable = false, alg = HeapSort, accelerate = false, direction = :backward, inplace = false, tol = nothing,  allow_exact_match = true, op = nothing, method = :sort, threads = true) where T
    isempty(dsl) && return copy(dsl)
    if !allow_exact_match
        #aem is the function to check allow_exact_match
        aem = !isequal
        # missing is greater than anything
        if tol == nothing
            tol = missing
        end
    else
        aem = _IST_
    end
    oncols_left = onleft
    oncols_right = onright
    right_cols = setdiff(1:length(index(dsr)), oncols_right)
    if !makeunique && !isempty(intersect(_names(dsl), _names(dsr)[right_cols]))
        throw(ArgumentError("duplicate column names, pass `makeunique = true` to make them unique using a suffix automatically." ))
    end

    nsfpaj = true
    # if the column for close join is a PA we cannot use the fast path
    if DataAPI.refpool(_columns(dsr)[oncols_right[end]]) !== nothing
        nsfpaj = false
    end
    if length(oncols_left) > 1 && method == :hash
        ranges, a, idx, minval, reps, sz, right_cols_2= _find_ranges_for_join_using_hash(dsl, dsr, onleft[1:end-1], onright[1:end-1], mapformats, true, Val(T), threads = threads)
        filter!(!=(0), reps)
        pushfirst!(reps, 1)
        cumsum!(reps, reps)
        pop!(reps)
        grng = GIVENRANGE(idx, reps, Int[], length(reps))
        starts, idx, last_valid_range = _sort_for_join_after_hash(dsr, oncols_right[end], stable, alg, mapformats, nsfpaj, grng, threads = threads)
    else
        ranges = Vector{UnitRange{T}}(undef, nrow(dsl))
        idx, uniquemode = _find_permute_and_fill_range_for_join!(ranges, dsr, dsl, oncols_right, oncols_left, stable, alg, mapformats, accelerate && length(oncols_right) > 1; nsfpaj = nsfpaj, threads = threads)

        for j in 1:(length(oncols_left) - 1)
            _change_refpool_find_range_for_join!(ranges, dsl, dsr, idx, oncols_left, oncols_right, mapformats[1], mapformats[2], j; nsfpaj = nsfpaj, threads = threads)
        end
    end

    # if border = :none , we should use :nearest direction
    _change_refpool_find_range_for_close!(ranges, dsl, dsr, idx, oncols_left, oncols_right, border == :none ? :nearest : direction, mapformats[1], mapformats[2], length(oncols_left); nsfpaj = nsfpaj, threads = threads)
    total_length = nrow(dsl)

    if inplace
        newds = dsl
    else
        res = []
        for j in 1:length(index(dsl))
            push!(res,  _columns(dsl)[j])
        end
        if dsl isa SubDataset
            newds = Dataset(res, copy(index(dsl)))
        else
            newds = Dataset(res, Index(copy(index(dsl).lookup), copy(index(dsl).names), copy(index(dsl).format)))
        end
    end

    for j in 1:length(right_cols)
        _res = allocatecol(_columns(dsr)[right_cols[j]], total_length)
        if DataAPI.refpool(_res) !== nothing
            fill_val = DataAPI.invrefpool(_res)[missing]
            _fill_right_cols_table_close!(_res.refs, view(DataAPI.refarray(_columns(dsr)[right_cols[j]]), idx), ranges, total_length, border, fill_val, direction; nn = direction == :nearest, rnn = view(_columns(dsr)[oncols_right[end]], idx), lnn = _columns(dsl)[oncols_left[end]], tol = tol, aem = aem, op = op, threads = threads)
        else
            _fill_right_cols_table_close!(_res, view(_columns(dsr)[right_cols[j]], idx), ranges, total_length, border, missing, direction; nn = direction == :nearest, rnn = view(_columns(dsr)[oncols_right[end]], idx), lnn = _columns(dsl)[oncols_left[end]], tol = tol, aem = aem, op = op, threads = threads)
        end
        push!(_columns(newds), _res)
        new_var_name = make_unique([_names(dsl); _names(dsr)[right_cols[j]]], makeunique = makeunique)[end]
        push!(index(newds), new_var_name)
        setformat!(newds, index(newds)[new_var_name], getformat(dsr, _names(dsr)[right_cols[j]]))
    end

    if inplace
        _modified(_attributes(newds))
    end
    newds

end


module MoireHamiltonian

using Plots

include("lattice.jl")
include("folding.jl")
include("bilayer.jl")

end


using Test
using DelimitedFiles

@testset "Find neighbors" begin

    include("../src/init.jl")

    @testset "Neighbors inside domain" begin
        actual = neighbors(4, 12, 3, 4)
        expected = Dict{String,Int64}("S" => 3, "W" => 1, "N" => 5, "E" => 7)
        @test (actual == expected)
    end;

    @testset "Neighbors near boundary" begin
        actual = neighbors(11, 12, 3, 4)
        expected = Dict{String,Int64}("S" => 10, "W" => 8, "N" => -1, "E" => -1)
        @test (actual == expected)
    end;

end;

@testset "End-to-end tests" begin

    reference_8x4_output = readdlm("resources/reference_8x4_output.dat")

    @testset "1 process, 8x4 grid" begin
        cmd = `mpirun -np 1 julia ../src/Diffusion_run.jl resources/test1_params.yml`
        run(cmd)
        expected = reference_8x4_output
        actual = readdlm("resources/test1.dat")
        @test all(actual .≈ expected)
    end;

    @testset "1 process, 4x8 grid" begin
        cmd = `mpirun -np 1 julia ../src/Diffusion_run.jl resources/test2_params.yml`
        run(cmd)
        nx = 4
        ny = 8
        expected = transpose(reshape(reference_8x4_output, nx + 2, ny + 2))
        actual = reshape(readdlm("resources/test2.dat"), ny + 2, nx + 2)
        @test all(actual .≈ expected)
    end;

    @testset "2x1 processes, 8x4 grid" begin
        cmd = `mpirun -np 2 julia ../src/Diffusion_run.jl resources/test3_params.yml`
        run(cmd)
        expected = reference_8x4_output
        actual = readdlm("resources/test3.dat")
        @test all(actual .≈ expected)
    end;

    @testset "2x1 processes, 4x8 grid" begin
        cmd = `mpirun -np 2 julia ../src/Diffusion_run.jl resources/test4_params.yml`
        run(cmd)
        nx = 4
        ny = 8
        expected = transpose(reshape(reference_8x4_output, nx + 2, ny + 2))
        actual = reshape(readdlm("resources/test4.dat"), ny + 2, nx + 2)
        @test all(actual .≈ expected)
    end;

    @testset "1x2 processes, 8x4 grid" begin
        cmd = `mpirun -np 2 julia ../src/Diffusion_run.jl resources/test5_params.yml`
        run(cmd)
        expected = reference_8x4_output
        actual = readdlm("resources/test5.dat")
        @test all(actual .≈ expected)
    end;

    @testset "1x2 processes, 4x8 grid" begin
        cmd = `mpirun -np 2 julia ../src/Diffusion_run.jl resources/test6_params.yml`
        run(cmd)
        nx = 4
        ny = 8
        expected = transpose(reshape(reference_8x4_output, nx + 2, ny + 2))
        actual = reshape(readdlm("resources/test6.dat"), ny + 2, nx + 2)
        @test all(actual .≈ expected)
    end;

end;


function read_single(filename::String)
    ret = Int64[]
    open(filename) do f
        for line in eachline(f)
            push!(ret, parse(Int64, line))
        end
    end
    return ret
end

function read_node_map(filename::String)
    nm = Dict{Int64,Int64}()
    open(filename) do f
        for (i, line) in enumerate(eachline(f))
            if i == 1; continue; end
            new, orig = [parse(Int64, v) for v in split(line)]
            nm[orig] = new
        end
    end
    return nm
end

function read_labels(filename::String, nm, n::Int64, incl_gender::Bool)
    labels = zeros(Int64, n)
    label_map = Dict{AbstractString,Int64}()

    open(filename) do f
        for line in eachline(f)
            data = split(line)
            orig_id = parse(Int64, data[1])
            # Some nodes do not participate in any edge
            if !haskey(nm, orig_id); continue; end
            node = nm[orig_id]

            # Label is either class or class and gender (note: some genders listed as "UNKNOWN")
            label = data[2]
            if incl_gender
                label = string(label, " ", data[3])
            end
            if !haskey(label_map, label)
                label_map[label] = length(label_map) + 1
            end
            labels[node] = label_map[label]
        end
    end

    return labels, label_map
end

function main()
    for dataset in ["high-school", "primary-school"]
        for incl_gender_in_label in [true, false]
            # Read data
            nverts = read_single("raw-$dataset/contact-$dataset-nverts.txt")
            simplices = read_single("raw-$dataset/contact-$dataset-simplices.txt")
            nm = read_node_map("raw-$dataset/contact-$dataset-nodemap.txt")
            metadata =
                if     dataset == "high-school";    "metadata_2013"
                elseif dataset == "primary-school"; "metadata_primaryschool"
                else                                ""
                end
            labels, label_map = read_labels("raw-$dataset/$metadata.txt", nm, maximum(simplices),
                                            incl_gender_in_label)

            # Get all unique hyperedges
            hedges = Set{Set{Int64}}()
            let curr_ind = 0
	        for nvert in nverts
    	            simplex = simplices[(curr_ind + 1):(curr_ind + nvert)]
                    push!(hedges, Set{Int64}(simplex))
	            curr_ind += nvert
                end
            end

            # setup directory
            dir = string("contact-", dataset, "-classes")
            if incl_gender_in_label
                dir = string(dir, "-gender")
            end
            if !isdir(dir)
                mkdir(dir)
            end

            # hyperedge data
            open("$dir/hyperedges-$dir.txt", "w") do f
                for hedge in hedges
                    write(f, join(sort(collect(hedge)), ','))
                    write(f, "\n")
                end
            end

            # names of the integer labels
            open("$dir/label-names-$dir.txt", "w") do f
                output = sort([(v, k) for (k, v) in label_map])
                for (v, k) in output
                    write(f, "$k\n")
                end
            end

            # labels for each node
            open("$dir/node-labels-$dir.txt", "w") do f
                for l in labels
                    write(f, "$l\n")
                end
            end
        end
    end
end


@testset "GymEnv" begin

function test_random_rollout(name)
    env = GymEnv(name)
    agent = env -> sample(action_space(env))
    ob = reset!(env)
    for _ in 1:10
        ac = agent(env)
        @test occursin(ac, action_space(env))
        step!(env, ac)
        @test occursin(get_state(env), observation_space(env))
        if is_end(env)
            break
        end
    end
end

test_random_rollout("CartPole-v0")
test_random_rollout("FrozenLake-v0")
end


using Test
using CausalDiscovery

function construct_data()
  Dict([("external" => []),
        ("initnext" => []),
        ("lifted" => []),
        ("types" => Dict())])
end

mod = nothing

function isinferred(f, args...; allow = Union{})
  ret = f(args...)
  inftypes = Base.return_types(f, Base.typesof(args...))
  rettype = ret isa Type ? Type{ret} : typeof(ret)
  rettype <: allow || rettype == Test.typesubtract(inftypes[1], allow)
end

function test_compile_if()
  data = construct_data()
  aexpr = au"""(if (== x 3) then (= y 4) else (= y 5))"""
  @test string(compile(aexpr, data)) == "if x == 3\n    y = 4\nelse\n    y = 5\nend"
end

function test_compile_assign()
  data = construct_data()
  aexpr = au"""(= x 3)"""
  @test compile(aexpr, data) == :(x = 3)
end

function test_compile_typedecl()
  data = construct_data()
  aexpr = au"""(: x Int)"""
  @test compile(aexpr, data) == :(local x::Int)
end

function test_compile_external()
  data = construct_data()
  aexpr = au"""(external (: click Click))"""
  @test compile(aexpr, data) == :()
end

function test_compile_let()
  data = construct_data()
  aexpr = au"""(let ((= y 3) (= x y)) x)"""

  @test compile(aexpr, data).args[end - 2: end] == [:(y = 3), :(x = y), :x]
end

function test_compile_list()
  data = construct_data()
  aexpr = au"""(list 1 2 3 4 5)"""
  @test compile(aexpr, data) == :([1, 2, 3, 4, 5])
end

function test_compile_call()
  data = construct_data()
  aexpr = au"""(f 1 2 3) """
  @test compile(aexpr, data) == :(f(1, 2, 3))
end

function test_compile_field()
  data = construct_data()
  aexpr = au"""(.. position x)"""
  @test compile(aexpr, data) == :(position.x)
end

function test_compile_particles()
  a = au"""(program
    (: GRID_SIZE Int)
    (= GRID_SIZE 16)

    (type alias Position ((: x Int) (: y Int)))
    (type alias Particle ((: position Position)))

    (external (: click Click))

    (: particles (List Particle))
    (= particles (initnext (list) (if (occurred click)
                                then (push! (prev particles) (particleGen (Position 1 1)))
                                else (map nextParticle (prev particles)))))
    (: nparticles Int)
    (= nparticles (length particles))

    (: isFree (-> Position Bool))
    (= isFree (fn (position)
                  (== (length (filter (--> particle (== (.. particle position) position)) (prev particles))) 0)))

    (: isWithinBounds (-> Position Bool))
    (= isWithinBounds (fn (position)
                          (& (>= (.. position x) 0) (& (< (.. position x) GRID_SIZE) (& (>= (.. position y) 0) (< (.. position y) GRID_SIZE))))))

    (: adjacentPositions (-> Position (List Position)))
    (= adjacentPositions (fn (position)
                            (let ((= x (.. position x))
                                  (= y (.. position y))
                                  (= positions (filter isWithinBounds (list (Position (+ x 1) y) (Position (- x 1) y) (Position x (+ y 1)) (Position x (- y 1))))))
                                  positions)))

    (: nextParticle (-> Particle Particle))
    (= nextParticle (fn (particle)
                        (let ((= freePositions (filter isFree (adjacentPositions (.. particle position)))))
                            (case freePositions
                                  (=> (list) particle)
                                  (=> _ (Particle (uniformChoice freePositions)))))))

    (: particleGen (-> Position Particle))
    (= particleGen (fn (initPosition) (Particle initPosition)))
  )"""
  @test string(compiletojulia(a)) == "\$(Expr(:toplevel, :(module CompiledProgram\n  export init, next\n  import Base.min\n  using Distributions\n  using MLStyle: @match\n  begin\n      function occurred(click)\n          click !== nothing\n      end\n  end\n  begin\n      function uniformChoice(freePositions)\n          freePositions[rand(Categorical(ones(length(freePositions)) / length(freePositions)))]\n      end\n  end\n  begin\n      function uniformChoice(freePositions, n)\n          map((idx->begin\n                      freePositions[idx]\n                  end), rand(Categorical(ones(length(freePositions)) / length(freePositions)), n))\n      end\n  end\n  begin\n      function min(arr)\n          min(arr...)\n      end\n  end\n  begin\n      struct Click\n          x::BigInt\n          y::BigInt\n      end\n  end\n  begin\n      function range(start::BigInt, stop::BigInt)\n          [start:stop;]\n      end\n  end\n  begin\n      struct Position\n          x::Int\n          y::Int\n      end\n  end\n  begin\n      struct Particle\n          position::Position\n      end\n  end\n  begin\n      function isFree(position::Position)::Bool\n          length(filter((particle->begin\n                              particle.position == position\n                          end), particlesPrev())) == 0\n      end\n  end\n  begin\n      function isWithinBounds(position::Position)::Bool\n          (position.x >= 0) & ((position.x < GRID_SIZE) & ((position.y >= 0) & (position.y < GRID_SIZE)))\n      end\n  end\n  begin\n      function adjacentPositions(position::Position)::Array{Position}\n          begin\n              x = position.x\n              y = position.y\n              positions = filter(isWithinBounds, [Position(x + 1, y), Position(x - 1, y), Position(x, y + 1), Position(x, y - 1)])\n              positions\n          end\n      end\n  end\n  begin\n      function nextParticle(particle::Particle)::Particle\n          begin\n              freePositions = filter(isFree, adjacentPositions(particle.position))\n              begin\n                  @match freePositions begin\n                          [] => particle\n                          _ => Particle(uniformChoice(freePositions))\n                      end\n              end\n          end\n      end\n  end\n  begin\n      function particleGen(initPosition::Position)::Particle\n          Particle(initPosition)\n      end\n  end\n  begin\n      mutable struct STATE\n          time::Int\n          particlesHistory::Dict{Int64, Array{Particle}}\n          GRID_SIZEHistory::Dict{Int64, Int}\n          nparticlesHistory::Dict{Int64, Int}\n          clickHistory::Dict{Int64, Union{Click, Nothing}}\n      end\n  end\n  state = STATE(0, Dict{Int64, Array{Particle}}(), Dict{Int64, Int}(), Dict{Int64, Int}(), Dict{Int64, Union{Click, Nothing}}())\n  begin\n      function particlesPrev(n::Int=1)::Array{Particle}\n          state.particlesHistory[if state.time - n >= 0\n                  state.time - n\n              else\n                  0\n              end]\n      end\n  end\n  begin\n      function GRID_SIZEPrev(n::Int=1)::Int\n          state.GRID_SIZEHistory[if state.time - n >= 0\n                  state.time - n\n              else\n                  0\n              end]\n      end\n  end\n  begin\n      function nparticlesPrev(n::Int=1)::Int\n          state.nparticlesHistory[if state.time - n >= 0\n                  state.time - n\n              else\n                  0\n              end]\n      end\n  end\n  begin\n      function clickPrev(n::Int=1)::Union{Click, Nothing}\n          state.clickHistory[if state.time - n >= 0\n                  state.time - n\n              else\n                  0\n              end]\n      end\n  end\n  begin\n      function init(click::Union{Click, Nothing})::STATE\n          state = STATE(0, Dict{Int64, Array{Particle}}(), Dict{Int64, Int}(), Dict{Int64, Int}(), Dict{Int64, Union{Click, Nothing}}())\n          particles = []\n          GRID_SIZE = 16\n          nparticles = length(particles)\n          state.clickHistory[state.time] = click\n          state.particlesHistory[state.time] = particles\n          state.GRID_SIZEHistory[state.time] = GRID_SIZE\n          state.nparticlesHistory[state.time] = nparticles\n          deepcopy(state)\n      end\n  end\n  begin\n      function next(old_state::STATE, click::Union{Click, Nothing})::STATE\n          global state = deepcopy(old_state)\n          state.time = state.time + 1\n          particles = if occurred(click)\n                  push!(particlesPrev(), particleGen(Position(1, 1)))\n              else\n                  map(nextParticle, particlesPrev())\n              end\n          GRID_SIZE = 16\n          nparticles = length(particles)\n          state.clickHistory[state.time] = click\n          state.particlesHistory[state.time] = particles\n          state.GRID_SIZEHistory[state.time] = GRID_SIZE\n          state.nparticlesHistory[state.time] = nparticles\n          deepcopy(state)\n      end\n  end\n  end)))"
  global mod = eval(compiletojulia(a))
end

function test_compile_types_inferred()
  @test isinferred(mod.init, nothing)
  @test isinferred(mod.init, mod.Click(5,5))
  state = mod.init(nothing)
  @test isinferred(mod.next, state, mod.Click(5,5))
  @test isinferred(mod.next, state, nothing)
end

@testset "compile" begin
  test_compile_if()
  test_compile_assign()
  test_compile_typedecl()
  test_compile_external()
  test_compile_let()
  test_compile_list()
  test_compile_call()
  test_compile_field()
  test_compile_particles()
  test_compile_types_inferred()
end


abstract type TracerBC end

"""
    ImpermeableTracer() :: TracerBC

No tracer diffusion across boundary
"""
struct ImpermeableTracer <: TracerBC end
function atmos_tracer_boundary_state!(
    nf,
    bc_tracer::ImpermeableTracer,
    atmos,
    args...,
)
    nothing
end
function atmos_tracer_normal_boundary_flux_second_order!(
    nf,
    bc_tracer::ImpermeableTracer,
    atmos,
    args...,
)
    nothing
end


import Base: chop

# BLAS/linear algebra overrides

@inline dot(x...) = LinearAlgebra.dot(x...)
@inline dot(M::Int,a::Ptr{T},incx::Int,b::Ptr{T},incy::Int) where {T<:Union{Float64,Float32}} =
    BLAS.dot(M,a,incx,b,incy)
@inline dot(M::Int,a::Ptr{T},incx::Int,b::Ptr{T},incy::Int) where {T<:Union{ComplexF64,ComplexF32}} =
    BLAS.dotc(M,a,incx,b,incy)

dotu(f::StridedVector{T},g::StridedVector{T}) where {T<:Union{ComplexF32,ComplexF64}} =
    BLAS.dotu(f,g)
dotu(f::AbstractVector{Complex{Float64}},g::AbstractVector{N}) where {N<:Real} = dot(conj(f),g)
dotu(f::AbstractVector{N},g::AbstractVector{T}) where {N<:Real,T<:Number} = dot(f,g)

normalize!(w::AbstractVector) = rmul!(w,inv(norm(w)))
normalize!(w::Vector{T}) where {T<:BlasFloat} = normalize!(length(w),w)
normalize!(n,w::Union{Vector{T},Ptr{T}}) where {T<:Union{Float64,Float32}} =
    BLAS.scal!(n,inv(BLAS.nrm2(n,w,1)),w,1)
normalize!(n,w::Union{Vector{T},Ptr{T}}) where {T<:Union{ComplexF64,ComplexF32}} =
    BLAS.scal!(n,T(inv(BLAS.nrm2(n,w,1))),w,1)

flipsign(x,y) = Base.flipsign(x,y)
flipsign(x,y::Complex) = y==0 ? x : x*sign(y)

# Used for spaces not defined yet
struct UnsetNumber <: Number  end
promote_rule(::Type{UnsetNumber},::Type{N}) where {N<:Number} = N
promote_rule(::Type{Bool}, ::Type{UnsetNumber}) = Bool

# Test the number of arguments a function takes
hasnumargs(f,k) = applicable(f,zeros(k)...)

# fast implementation of isapprox with atol a non-keyword argument in most cases
isapprox_atol(a,b,atol;kwds...) = isapprox(a,b;atol=atol,kwds...)
isapprox_atol(a::Vec,b::Vec,atol::Real=0;kwds...) = isapprox_atol(collect(a),collect(b),atol;kwds...)
function isapprox_atol(x::Number, y::Number, atol::Real=0; rtol::Real=Base.rtoldefault(x,y))
    x == y || (isfinite(x) && isfinite(y) && abs(x-y) <= atol + rtol*max(abs(x), abs(y)))
end
function isapprox_atol(x::AbstractArray{T}, y::AbstractArray{S},atol::Real=0; rtol::Real=Base.rtoldefault(T,S), norm::Function=vecnorm) where {T<:Number,S<:Number}
    d = norm(x - y)
    if isfinite(d)
        return d <= atol + rtol*max(norm(x), norm(y))
    else
        # Fall back to a component-wise approximate comparison
        return all(ab -> isapprox(ab[1], ab[2]; rtol=rtol, atol=atol), zip(x, y))
    end
end

# The second case handles zero
isapproxinteger(x) = isapprox(x,round(Int,x))  || isapprox(x+1,round(Int,x+1))

# This creates ApproxFunBase.real, ApproxFunBase.eps and ApproxFunBase.dou
# which we override for default julia types
real(x...) = Base.real(x...)
real(::Type{UnsetNumber}) = UnsetNumber
real(::Type{Array{T,n}}) where {T<:Real,n} = Array{T,n}
real(::Type{Array{T,n}}) where {T<:Complex,n} = Array{real(T),n}
real(::Type{Vec{N,T}}) where {N,T<:Real} = Vec{N,T}
real(::Type{Vec{N,T}}) where {N,T<:Complex} = Vec{N,real(T)}

float(x) = Base.float(x)
Base.float(::UnsetNumber) = UnsetNumber()
Base.float(::Type{UnsetNumber}) = UnsetNumber
float(::Type{Array{T,N}}) where {T,N} = Array{float(T),N}
float(::Type{SVector{N,T}}) where {T,N} = SVector{N,float(T)}

eps(x...) = Base.eps(x...)
eps(x) = Base.eps(x)

eps(::Type{T}) where T<:Integer = zero(T)
eps(::Type{T}) where T<:Rational = zero(T)
eps(::T) where T<:Integer = eps(T)

eps(::Type{Complex{T}}) where {T<:Real} = eps(real(T))
eps(z::Complex{T}) where {T<:Real} = eps(abs(z))
eps(::Type{Dual{Complex{T}}}) where {T<:Real} = eps(real(T))
eps(z::Dual{Complex{T}}) where {T<:Real} = eps(abs(z))

eps(::Type{Vector{T}}) where {T<:Number} = eps(T)
eps(::Type{Vec{k,T}}) where {k,T<:Number} = eps(T)

isnan(x) = Base.isnan(x)
isnan(x::Vec) = map(isnan,x)

# BLAS

# implement muladd default
muladd(a,b,c) = a*b+c
muladd(a::Number,b::Number,c::Number) = Base.muladd(a,b,c)

for TYP in (:Float64,:Float32,:ComplexF64,:ComplexF32)
    @eval scal!(n::Integer,cst::$TYP,ret::DenseArray{T},k::Integer) where {T<:$TYP} =
            BLAS.scal!(n,cst,ret,k)
end

scal!(n::Integer,cst::BlasFloat,ret::DenseArray{T},k::Integer) where {T<:BlasFloat} =
    BLAS.scal!(n,convert(T,cst),ret,k)

function scal!(n::Integer,cst::Number,ret::AbstractArray,k::Integer)
    @assert k*n ≤ length(ret)
    @simd for j=1:k:k*(n-1)+1
        @inbounds ret[j] *= cst
    end
    ret
end

scal!(cst::Number,v::AbstractArray) = scal!(length(v),cst,v,1)

# Helper routines

function reverseeven!(x::AbstractVector)
    n = length(x)
    if iseven(n)
        @inbounds @simd for k=2:2:n÷2
            x[k],x[n+2-k] = x[n+2-k],x[k]
        end
    else
        @inbounds @simd for k=2:2:n÷2
            x[k],x[n+1-k] = x[n+1-k],x[k]
        end
    end
    x
end

function negateeven!(x::AbstractVector)
    @inbounds @simd for k = 2:2:length(x)
        x[k] *= -1
    end
    x
end

#checkerboard, same as applying negativeeven! to all rows then all columns
function negateeven!(X::AbstractMatrix)
    for j = 1:2:size(X,2)
        @inbounds @simd for k = 2:2:size(X,1)
            X[k,j] *= -1
        end
    end
    for j = 2:2:size(X,2)
        @inbounds @simd for k = 1:2:size(X,1)
            X[k,j] *= -1
        end
    end
    X
end

const alternatesign! = negateeven!

alternatesign(v::AbstractVector) = alternatesign!(copy(v))

alternatingvector(n::Integer) = 2*mod([1:n],2) .- 1

function alternatingsum(v::AbstractVector)
    ret = zero(eltype(v))
    s = 1
    @inbounds for k=1:length(v)
        ret+=s*v[k]
        s*=-1
    end

    ret
end

# Sum Hadamard product of vectors up to minimum over lengths
function mindotu(a::AbstractVector,b::AbstractVector)
    ret,m = zero(promote_type(eltype(a),eltype(b))),min(length(a),length(b))
    @inbounds @simd for i=m:-1:1 ret += a[i]*b[i] end
    ret
end

# efficiently resize a Matrix.  Note it doesn't change the input ptr
function unsafe_resize!(W::AbstractMatrix,::Colon,m::Integer)
    if m == size(W,2)
        W
    else
        n=size(W,1)
        reshape(resize!(vec(W),n*m),n,m)
    end
end

function unsafe_resize!(W::AbstractMatrix,n::Integer,::Colon)
    N=size(W,1)
    if n == N
        W
    elseif n < N
        W[1:n,:]
    else
        m=size(W,2)
        ret=Matrix{eltype(W)}(n,m)
        ret[1:N,:] = W
        ret
    end
end

function unsafe_resize!(W::AbstractMatrix,n::Integer,m::Integer)
    N=size(W,1)
    if n == N
        unsafe_resize!(W,:,m)
    else
        unsafe_resize!(unsafe_resize!(W,n,:),:,m)
    end
end

function pad!(f::AbstractVector{T},n::Integer) where T
	if n > length(f)
		append!(f,zeros(T,n - length(f)))
	else
		resize!(f,n)
	end
end

function pad(f::AbstractVector{T},n::Integer) where T
	if n > length(f)
	   ret=Vector{T}(undef, n)
	   ret[1:length(f)]=f
	   for j=length(f)+1:n
	       ret[j]=zero(T)
	   end
       ret
	else
        f[1:n]
	end
end

function pad(f::AbstractVector{Any},n::Integer)
	if n > length(f)
        Any[f...,zeros(n - length(f))...]
	else
        f[1:n]
	end
end

pad(x::Number, n::Int) = [x; zeros(typeof(x), n-1)]

function pad(v::AbstractVector,n::Integer,m::Integer)
    @assert m==1
    pad(v,n)
end

function pad(A::AbstractMatrix,n::Integer,m::Integer)
    T=eltype(A)
	if n <= size(A,1) && m <= size(A,2)
        A[1:n,1:m]
	elseif n==0 || m==0
	   Matrix{T}(undef,n,m)  #fixes weird julia bug when T==None
    else
        ret = Matrix{T}(undef,n,m)
        minn=min(n,size(A,1))
        minm=min(m,size(A,2))
        for k=1:minn,j=1:minm
            @inbounds ret[k,j]=A[k,j]
        end
        for k=minn+1:n,j=1:minm
            @inbounds ret[k,j]=zero(T)
        end
        for k=1:n,j=minm+1:m
            @inbounds ret[k,j]=zero(T)
        end
        for k=minn+1:n,j=minm+1:m
            @inbounds ret[k,j]=zero(T)
        end

        ret
	end
end

pad(A::AbstractMatrix,::Colon,m::Integer) = pad(A,size(A,1),m)
pad(A::AbstractMatrix,n::Integer,::Colon) = pad(A,n,size(A,2))

function pad(v, ::PosInfinity)
    if isinf(length(v))
        v
    else
        Vcat(v, Zeros{Int}(∞))
    end
end

function pad(v::AbstractVector{T}, ::PosInfinity) where T
    if isinf(length(v))
        v
    else
        Vcat(v, Zeros{T}(∞))
    end
end

#TODO:padleft!

function padleft(f::AbstractVector,n::Integer)
	if (n > length(f))
        [zeros(n - length(f)); f]
	else
        f[end-n+1:end]
	end
end

##chop!
function chop!(c::AbstractVector,tol::Real)
    @assert tol >= 0

    for k=length(c):-1:1
        if abs(c[k]) > tol
            resize!(c,k)
            return c
        end
    end

    resize!(c,0)
    c
end

chop(f::AbstractVector,tol) = chop!(copy(f),tol)
chop!(f::AbstractVector) = chop!(f,eps())

function chop!(A::AbstractArray,tol)
    for k=size(A,1):-1:1
        if norm(A[k,:])>tol
            A=A[1:k,:]
            break
        end
    end
    for k=size(A,2):-1:1
        if norm(A[:,k])>tol
            A=A[:,1:k]
            break
        end
    end
    return A
end
chop(A::AbstractArray,tol)=chop!(A,tol)#replace by chop!(copy(A),tol) when chop! is actually in-place.

## interlace

function interlace(v::Union{Vector{Any},Tuple})
    #determine type
    T=Float64
    for vk in v
        if isa(vk,Vector{Complex{Float64}})
            T=Complex{Float64}
        end
    end
    b=Vector{Vector{T}}(undef, length(v))
    for k=1:length(v)
        b[k]=v[k]
    end
    interlace(b)
end

function interlace(a::AbstractVector{S},b::AbstractVector{V}) where {S<:Number,V<:Number}
    na=length(a);nb=length(b)
    T=promote_type(S,V)
    if nb≥na
        ret=zeros(T,2nb)
        ret[1:2:1+2*(na-1)]=a
        ret[2:2:end]=b
        ret
    else
        ret=zeros(T,2na-1)
        ret[1:2:end]=a
        if !isempty(b)
            ret[2:2:2+2*(nb-1)]=b
        end
        ret
    end
end

function interlace(a::AbstractVector,b::AbstractVector)
    na=length(a);nb=length(b)
    T=promote_type(eltype(a),eltype(b))
    if nb≥na
        ret=Vector{T}(undef, 2nb)
        ret[1:2:1+2*(na-1)]=a
        ret[2:2:end]=b
        ret
    else
        ret=Vector{T}(undef, 2na-1)
        ret[1:2:end]=a
        if !isempty(b)
            ret[2:2:2+2*(nb-1)]=b
        end
        ret
    end
end

### In-place O(n) interlacing

function highestleader(n::Int)
    i = 1
    while 3i < n i *= 3 end
    i
end

function nextindex(i::Int,n::Int)
    i <<= 1
    while i > n
        i -= n + 1
    end
    i
end

function cycle_rotate!(v::AbstractVector, leader::Int, it::Int, twom::Int)
    i = nextindex(leader, twom)
    while i != leader
        idx1, idx2 = it + i - 1, it + leader - 1
        @inbounds v[idx1], v[idx2] = v[idx2], v[idx1]
        i = nextindex(i, twom)
    end
    v
end

function right_cyclic_shift!(v::AbstractVector, it::Int, m::Int, n::Int)
    itpm = it + m
    itpmm1 = itpm - 1
    itpmpnm1 = itpmm1 + n
    reverse!(v, itpm, itpmpnm1)
    reverse!(v, itpm, itpmm1 + m)
    reverse!(v, itpm + m, itpmpnm1)
    v
end

"""
This function implements the algorithm described in:

    P. Jain, "A simple in-place algorithm for in-shuffle," arXiv:0805.1598, 2008.
"""
function interlace!(v::AbstractVector,offset::Int)
    N = length(v)
    if N < 2 + offset
        return v
    end

    it = 1 + offset
    m = 0
    n = 1

    while m < n
        twom = N + 1 - it
        h = highestleader(twom)
        m = h > 1 ? h÷2 : 1
        n = twom÷2

        right_cyclic_shift!(v,it,m,n)

        leader = 1
        while leader < 2m
            cycle_rotate!(v, leader, it, 2m)
            leader *= 3
        end

        it += 2m
    end
    v
end

## slnorm gives the norm of a slice of a matrix

function slnorm(u::AbstractMatrix,r::AbstractRange,::Colon)
    ret = zero(real(eltype(u)))
    for k=r
        @simd for j=1:size(u,2)
            #@inbounds
            ret=max(norm(u[k,j]),ret)
        end
    end
    ret
end

function slnorm(m::AbstractMatrix,kr::AbstractRange,jr::AbstractRange)
    ret=zero(real(eltype(m)))
    for j=jr
        nrm=zero(real(eltype(m)))
        for k=kr
            @inbounds nrm+=abs2(m[k,j])
        end
        ret=max(sqrt(nrm),ret)
    end
    ret
end

slnorm(m::AbstractMatrix,kr::AbstractRange,jr::Integer) = slnorm(m,kr,jr:jr)
slnorm(m::AbstractMatrix,kr::Integer,jr::AbstractRange) = slnorm(m,kr:kr,jr)

function slnorm(B::BandedMatrix{T},r::AbstractRange,::Colon) where T
    ret = zero(real(T))
    m=size(B,2)
    for k=r
        @simd for j=max(1,k-B.l):min(k+B.u,m)
            #@inbounds
            ret=max(norm(B[k,j]),ret)
        end
    end
    ret
end

slnorm(m::AbstractMatrix,k::Integer,::Colon) = slnorm(m,k,1:size(m,2))
slnorm(m::AbstractMatrix,::Colon,j::Integer) = slnorm(m,1:size(m,1),j)

## Infinity

Base.isless(x::Block{1}, y::PosInfinity) = isless(Int(x), y)
Base.isless(x::PosInfinity, y::Block{1}) = isless(x, Int(y))

## BandedMatrix

pad!(A::BandedMatrix,n,::Colon) = pad!(A,n,n+A.u)  # Default is to get all columns
columnrange(A,row::Integer) = max(1,row-bandwidth(A,1)):row+bandwidth(A,2)

## Store iterator
mutable struct CachedIterator{T,IT}
    iterator::IT
    storage::Vector{T}
    state
    length::Int
end

CachedIterator{T,IT}(it::IT, state) where {T,IT} = CachedIterator{T,IT}(it,T[],state,0)
CachedIterator(it::IT) where IT = CachedIterator{eltype(it),IT}(it, ())

function resize!(it::CachedIterator,n::Integer)
    m = it.length
    if n > m
        if n > length(it.storage)
            resize!(it.storage,2n)
        end

        @inbounds for k = m+1:n
            xst = iterate(it.iterator,it.state...)
            if xst == nothing
                it.length = k-1
                return it
            end
            it.storage[k] = xst[1]
            it.state = (xst[2],)
        end

        it.length = n
    end
    it
end

eltype(it::CachedIterator{T}) where {T} = T

iterate(it::CachedIterator) = iterate(it,1)
function iterate(it::CachedIterator,st::Int)
    if  st == it.length + 1 && iterate(it.iterator,it.state...) == nothing
        nothing
    else
        (it[st],st+1)
    end
end

function getindex(it::CachedIterator, k)
    mx = maximum(k)
    if mx > length(it) || mx < 1
        throw(BoundsError(it,k))
    end
    resize!(it,isempty(k) ? 0 : mx).storage[k]
end

function findfirst(f::Function,A::CachedIterator)
    k=1
    for c in A
        if f(c)
            return k
        end
        k+=1
    end
    return 0
end

function findfirst(A::CachedIterator,x)
    k=1
    for c in A
        if c == x
            return k
        end
        k+=1
    end
    return 0
end

length(A::CachedIterator) = length(A.iterator)

## nocat
vnocat(A...) = Base.vect(A...)
hnocat(A...) = Base.typed_hcat(mapreduce(typeof,promote_type,A),A...)
hvnocat(rows,A...) = Base.typed_hvcat(mapreduce(typeof,promote_type,A),rows,A...)
macro nocat(x)
    ex = expand(x)
    if ex.args[1] == :vcat
        ex.args[1] = :(ApproxFunBase.vnocat)
    elseif ex.args[1] == :hcat
        ex.args[1] = :(ApproxFunBase.hnocat)
    else
        @assert ex.args[1] == :hvcat
        ex.args[1] = :(ApproxFunBase.hvnocat)
    end
    esc(ex)
end

## Dynamic functions

struct DFunction <: Function
    f
end
(f::DFunction)(args...) = f.f(args...)

hasnumargs(f::DFunction, k) = hasnumargs(f.f, k)

dynamic(f) = f
dynamic(f::Function) = DFunction(f) # Assume f has to compile every time

# Matrix inputs

## conv

conv(x::AbstractVector, y::AbstractVector) = DSP.conv(x, y)
@generated function conv(x::SVector{N}, y::SVector{M}) where {N,M}
    NM = N+M-1
    quote
        convert(SVector{$NM}, DSP.conv(Vector(x), Vector(y)))
    end
end

conv(x::SVector{1}, y::SVector{1}) = x.*y
conv(x::AbstractVector, y::SVector{1}) = x*y[1]
conv(y::SVector{1}, x::AbstractVector) = y[1]*x
conv(x::AbstractFill, y::SVector{1}) = x*y[1]
conv(y::SVector{1}, x::AbstractFill) = y[1]*x
conv(x::AbstractFill, y::AbstractFill) = DSP.conv(x, y)

## BlockInterlacer
# interlaces coefficients by blocks
# this has the property that all the coefficients of a block of a subspace
# are grouped together, starting with the first bloc
#
# TODO: cache sums

struct BlockInterlacer{DMS<:Tuple}
    blocks::DMS
end

const TrivialInterlacer{d} = BlockInterlacer{NTuple{d,<:Ones}}

BlockInterlacer(v::AbstractVector) = BlockInterlacer(tuple(v...))

Base.eltype(it::BlockInterlacer) = Tuple{Int,Int}

dimensions(b::BlockInterlacer) = map(sum,b.blocks)
dimension(b::BlockInterlacer,k) = sum(b.blocks[k])
Base.length(b::BlockInterlacer) = mapreduce(sum,+,b.blocks)

# the state is always (whichblock,curblock,cursubblock,curcoefficients)
# start(it::BlockInterlacer) = (1,1,map(start,it.blocks),ntuple(zero,length(it.blocks)))

# are all Ints, so finite dimensional
function done(it::BlockInterlacer,st)
    for k=1:length(it.blocks)
        if st[end][k] < sum(it.blocks[k])
            return false
        end
    end
    return true
end

iterate(it::BlockInterlacer) =
    iterate(it, (1,1,ntuple(_ -> tuple(), length(it.blocks)),
            ntuple(zero,length(it.blocks))))

function iterate(it::BlockInterlacer, (N,k,blkst,lngs))
    done(it, (N,k,blkst,lngs)) && return nothing

    if N > length(it.blocks)
        # increment to next block
        blkst = map(function(blit,blst)
                xblst = iterate(blit, blst...)
                xblst == nothing ? blst : (xblst[2],)
            end,it.blocks,blkst)
        return iterate(it,(1,1,blkst,lngs))
    end

    Bnxtb = iterate(it.blocks[N],blkst[N]...)  # B is block size

    if Bnxtb == nothing
        # increment to next N
        return iterate(it,(N+1,1,blkst,lngs))
    end

    B,nxtb = Bnxtb

    if k > B
        #increment to next N
        return iterate(it,(N+1,1,blkst,lngs))
    end

    lngs = tuple(lngs[1:N-1]...,lngs[N]+1,lngs[N+1:end]...)
    return (N,lngs[N]),(N,k+1,blkst,lngs)
end

cache(Q::BlockInterlacer) = CachedIterator(Q)


import Base: +,-,/,*,==,!==

mutable struct Measurement
    Name::AbstractString
    Data::Array
    Labbook::Dict{Any,Any}
end

Base.size(m::Measurement) = size(m.Data)
Base.copy(m::Measurement)  = Measurement(copy(m.Name),copy(m.Data),copy(m.Labbook))

==(m1::Measurement,m2::Measurement) = (m1.Name == m2.Name) && (m1.Data == m2.Data)
!==(m1::Measurement,m2::Measurement) = (m1.Name !== m2.Name) || (m1.Data !== m2.Data)
-(m1::Measurement,m2::Measurement) = m1.Data - m2.Data
+(m1::Measurement,m2::Measurement) = m1.Data + m2.Data
*(m::Measurement,n::Number) = m.Data * n
*(m::Measurement,n::Number) = n * m.Data
/(m::Measurement,n::Number) = m.Data /n
/(m::Measurement,n::Number) = n / m.Data


@testset "Constructs_test" begin

    mt_a = Mod("a")
    mt_b = Mod("b")
    mt_c = Mod("c")
    mt_d = Mod("d")
    mt_e = Mod("e")
    # some test constructs
    oc1 = OrderedConstruct([mt_a,mt_d,mt_b])
    oc2 = OrderedConstruct([mt_a,mt_b,mt_e,mt_d])
    unoc1 = UnorderedConstruct([mt_e,mt_d,mt_c])
    unoc2 = UnorderedConstruct([mt_e,mt_a,mt_d,mt_e])
     @testset "Base" begin
        @test length(oc1) == 3
        @test length(oc2) == 4
        @test length(unoc1) == 3
        @test length(unoc2) == 4

        @test oc1[2] == mt_d
        @test oc2[3] == mt_e

        @test isequal(oc1[1:2],OrderedConstruct([mt_a,mt_d])) == true
        @test isequal(oc1[1:2],OrderedConstruct([mt_a,mt_d])) == true
        @test length(collect(oc1)) == length(oc1)
        @test length(collect(unoc1)) == length(unoc1)
        @test eltype(collect(oc1)) == typeof(mt_a)
    end

    @testset "*_ordered" begin
        @test isequal(mt_a*mt_b,OrderedConstruct([mt_a,mt_b])) == true
        @test isequal(OrderedConstruct([mt_a,mt_b])*mt_c,OrderedConstruct([mt_a,mt_b,mt_c])) == true
        @test isequal(mt_c*OrderedConstruct([mt_a,mt_b]),OrderedConstruct([mt_c,mt_a,mt_b])) == true
        @test isequal((OrderedConstruct([mt_a,mt_b])*OrderedConstruct([mt_d,mt_e])),OrderedConstruct([mt_a,mt_b,mt_d,mt_e])) == true
    end

    @testset "+_unordered" begin
        @test isequal((mt_a+mt_b),UnorderedConstruct([mt_a,mt_b])) == true
        @test isequal(UnorderedConstruct([mt_a,mt_b])+mt_c,UnorderedConstruct([mt_a,mt_b,mt_c])) == true
        @test isequal(mt_c+UnorderedConstruct([mt_a,mt_b]),UnorderedConstruct([mt_c,mt_a,mt_b])) == true
        @test isequal((UnorderedConstruct([mt_a,mt_b])+UnorderedConstruct([mt_d,mt_e])),UnorderedConstruct([mt_a,mt_b,mt_d,mt_e])) == true
    end

end