const CuDense{ElT,VecT}                 = Dense{ElT,VecT} where {VecT<:CuVector}
const CuDenseTensor{ElT,N,StoreT,IndsT} = Tensor{ElT,N,StoreT,IndsT} where {StoreT<:CuDense}

Dense{T, SA}(x::Dense{T, SB}) where {T<:Number, SA<:CuArray, SB<:Array} = Dense{T, SA}(CuArray(x))
Dense{T, SA}(x::Dense{T, SB}) where {T<:Number, SA<:Array, SB<:CuArray} = Dense{T, SA}(collect(x.data))
Dense{T, S}(size::Integer) where {T, S<:CuArray{<:T}} = Dense{T, S}(CUDA.zeros(T, size))
function Dense{T, S}(x::T, size::Integer) where {T, S<:CuArray{<:T}}
    arr = CuArray{T}(undef, size)
    fill!(arr, x)
    Dense{T, S}(arr)
end
cpu(x::CuDense{T}) where {T<:Number} = Dense(collect(x.data))
Base.complex(::Type{Dense{ElT, VT}}) where {ElT, VT<:CuArray} = Dense{complex(ElT),CuVector{complex(ElT)}}

similartype(::Type{<:CuArray{<:Any,N}}, eltype::Type) where {N} = CuArray{eltype,N}
NDTensors.similar(::Type{<:CuArray{T}}, dims) where {T} = CuArray{T,length(dims)}(undef, dims)

CuArray(x::CuDense{ElT}) where {ElT} = CuVector{ElT}(data(x))
CuArray{ElT, N}(x::CuDenseTensor{ElT, N}) where {ElT, N} = CuArray{ElT, N}(reshape(data(store(x)), dims(inds(x))...))
CuArray(x::CuDenseTensor{ElT, N}) where {ElT, N} = CuArray{ElT, N}(x)

*(D::Dense{T, AT},x::S) where {T,AT<:CuArray,S<:Number} = Dense(x .* data(D))

Base.:(==)(::Type{<:CuDense{ElT1,CVec1}}, ::Type{<:CuDense{ElT2,CVec2}}) where {ElT1,ElT2,CVec1,CVec2} = (ElT1 == ElT2)
Base.getindex(D::CuDense{<:Number})          = collect(data(D))[]
Base.getindex(D::CuDenseTensor{<:Number, 0}) = store(D)[]
LinearAlgebra.norm(T::CuDenseTensor) = norm(data(store(T)))

# This is for type promotion for Scalar*Dense
function Base.promote_rule(::Type{<:Dense{ElT1,CuVector{ElT1}}},
                           ::Type{ElT2}) where {ElT1,
                                                ElT2<:Number}
  ElR  = promote_type(ElT1,ElT2)
  VecR = CuVector{ElR}
  return Dense{ElR,VecR}
end

function Base.permutedims(T::CuDenseTensor{<:Number,N},
                          perm::NTuple{N,Int}) where {N}
  Tp = NDTensors.similar(T,ITensors.NDTensors.permute(inds(T),perm))
  #Tp = permute(T,perm; always_copy=true)
  permute!(Tp, T)
  return Tp
end

function Base.permutedims!(R::CuDenseTensor{<:Number,N},
                           T::CuDenseTensor{<:Number,N},
                           perm::NTuple{N,Int}) where {N}
    return permutedims!!(R, T, perm)
end

function permutedims!!(B::Tensor{ElT,N,StoreT,IndsB},
                       A::Tensor{ElT,N,StoreT,IndsA},
                       perm::NTuple{N,Int},
                       f::Function=(r,t)->permute!(r,t)) where {N,ElT,IndsB,IndsA,StoreT<:CuDense{ElT}}
  Ais = inds(A)
  Bis = ITensors.NDTensors.permute(inds(A), perm)
  B = f(B, A)
  return B
end

function Base.similar(::Type{<:CuDenseTensor{ElT}},
                      inds::Dims) where {ElT}
    storage_arr = CuVector{ElT}(undef,dim(inds))
    return Tensor(Dense(storage_arr),inds)
end

function Base.similar(::Type{<:CuDenseTensor{ElT}},
                      inds) where {ElT}
    storage_arr = CuVector{ElT}(undef,dim(inds))
    return Tensor(Dense(storage_arr),inds)
end

function outer!(R::CuDenseTensor,
                T1::CuDenseTensor,
                T2::CuDenseTensor)
  R_dat = vec(array(T1))*transpose(vec(array(T2)))
  copyto!(data(store(R)), vec(R_dat))
  inds_outer = unioninds(inds(T1),inds(T2))
  return R
end

function contract!!(R::CuDenseTensor{<:Number,NR},
                    labelsR::NTuple{NR},
                    T1::CuDenseTensor{<:Number,N1},
                    labelsT1::NTuple{N1},
                    T2::CuDenseTensor{<:Number,N2},
                    labelsT2::NTuple{N2}) where {NR,N1,N2}
  if N1==0
    # TODO: replace with an add! function?
    # What about doing `R .= T1[] .* PermutedDimsArray(T2,perm)`?
    perm = getperm(labelsR,labelsT2)
    newT2 = Tensor(Dense(data(store(T1)).*data(store(T2))), inds(T2))
    permute!(R,newT2)
  elseif N2==0
    perm = getperm(labelsR,labelsT1)
    newT1 = Tensor(Dense(data(store(T2)).*data(store(T1))), inds(T1))
    permute!(R,newT1)
  elseif N1+N2==NR
    # TODO: permute T1 and T2 appropriately first (can be more efficient
    # then permuting the result of T1⊗T2)
    # TODO: implement the in-place version directly
    R = outer!!(R,T1,T2)
    inds_outer = unioninds(inds(T1),inds(T2))
    R = Tensor(store(R), inds_outer)
  else
    R = _contract!!(R,labelsR,T1,labelsT1,T2,labelsT2)
  end
  return R
end

function permutedims!!(B::CuDenseTensor{ElT,0},
                       A::CuDenseTensor{ElT,0},
                       perm::NTuple{0,Int},
                       f=(r,t)->permute!(r,t)) where {ElT<:Number}
    Cs = f(B, A)
    return Tensor(Dense(vec(Cs)), IndexSet{0}())
end

function permutedims!!(B::CuDenseTensor{ElT,N},
                       A::CuDenseTensor{ElT,0},
                       perm::NTuple{N,Int},
                       f=(r,t)->permute!(r,t)) where {N, ElT<:Number}
    Cis = ITensors.NDTensors.permute(inds(B), perm)
    Cs = f(B, A)
    return Tensor(Dense(vec(Cs)), Cis)
end

function _contract_scalar!(R::CuDenseTensor{ElR}, labelsR,
                           T1::Number, labelsT1,
                           T2::Number, labelsT2,
                           α = one(ElR), β = zero(ElR)) where {ElR}
  if iszero(β)
    copyto!(data(R), [α * T1 * T2])
  elseif iszero(α)
    copyto!(data(R), β.*data(R))
  else
    copyto!(data(R), [α * T1 * T2] .+ β.*data(R))
  end
  return R
end

function _contract_scalar!(R::CuDenseTensor{ElR,NR}, labelsR,
                           T₁::CuDenseTensor, labelsT₁,
                           T₂::CuDenseTensor, labelsT₂,
                           α = one(ElR), β=zero(ElR)) where {ElR,NR}
    if nnz(T₁) == nnz(T₂) == 1
        new_R = Tensor(Dense(data(store(T₁)).*data(store(T₂))), inds(R))
        copyto!(store(R), store(new_R))
    elseif nnz(T₁) == 1
        props = ContractionProperties(labelsT₁, labelsT₂, labelsR)
        compute_contraction_properties!(props, T₁, T₂, R)
        R = _contract!(R, T₁, T₂, props, α, β)
        #perm = getperm(labelsR,labelsT₂)
        #newT2 = Tensor(Dense(data(store(T₁)).*data(store(T₂))), inds(T₂))
        #permute!(R,newT2)
    elseif nnz(T₂) == 1
        props = ContractionProperties(labelsT₁, labelsT₂, labelsR)
        compute_contraction_properties!(props, T₁, T₂, R)
        R = _contract!(R, T₁, T₂, props, α, β)
        #perm = getperm(labelsR,labelsT₁)
        #newT1 = Tensor(Dense(data(store(T₁)).*data(store(T₂))), inds(T₁))
        #permute!(R,newT1)
    else
        error("In _contract_scalar!, one tensor must be a scalar")
    end
    return R
end

function _gemm_contract!(CT::DenseTensor{El,NC},
                         AT::DenseTensor{El,NA},
                         BT::DenseTensor{El,NB},
                         props::ContractionProperties,
                         α::Number=one(El),
                         β::Number=zero(El)) where {El,NC,NA,NB}
    # TODO: directly use Tensor instead of Array
    C = array(CT)
    A = array(AT)
    B = array(BT)

    tA = 'N'
    if props.permuteA
        pA = NTuple{NA,Int}(props.PA)
        Ap = permutedims(A, pA)
        AM = reshape(Ap, props.dmid, props.dleft)
        tA = 'T'
    else
        #A doesn't have to be permuted
        if Atrans(props)
            AM = reshape(A,props.dmid,props.dleft)
            tA = 'T'
        else
            AM = reshape(A,props.dleft,props.dmid)
        end
    end

    tB = 'N'
    if props.permuteB
        pB = NTuple{NB,Int}(props.PB)
        Bp = permutedims(B, pB)
        BM = reshape(Bp, props.dmid, props.dright)
    else
        if Btrans(props)
            BM = reshape(B,props.dright,props.dmid)
            tB = 'T'
        else
            BM = reshape(B,props.dmid,props.dright)
        end
    end

    #TODO: this logic may be wrong
    if props.permuteC
        #Need to copy here since we will be permuting
        #into C later
        CM = reshape(copy(C), props.dleft, props.dright)
    else
        if Ctrans(props)
            CM = reshape(C,props.dleft,props.dright)
            (AM,BM) = (BM,AM)
            if tA==tB
                tA = tB = (tA == 'T' ? 'N' : 'T')
            end
        else
            CM = reshape(C, props.dleft, props.dright)
        end
    end

    CM = CUBLAS.gemm!(tA,tB,El(α),AM,BM,El(β),CM)

    if props.permuteC
        pC = NTuple{NC,Int}(props.PC)
        Cr = reshape(CM,props.newCrange...)
        @strided C .= permutedims(Cr, pC)
    end
    return C
end

function _contract!(CT::CuDenseTensor{El1,NC},
                    AT::CuDenseTensor{El2,NA},
                    BT::CuDenseTensor{El3,NB},
                    props::ContractionProperties,
                    α::Number=one(Ela),β::Number=zero(El3)) where {El1,El2,El3,NC,NA,NB}
  if ndims(CT) > 12 || ndims(BT) > 12 || ndims(AT) > 12
    return _gemm_contract!(CT, AT, BT, props, α, β)
  end
  Ainds = inds(AT)
  Adims = dims(Ainds)
  Binds = inds(BT)
  Bdims = dims(Binds)
  Cinds = inds(CT)
  Cdims = dims(Cinds)
  Adata = reshape(data(store(AT)),Adims...)
  Bdata = reshape(data(store(BT)),Bdims...)
  Cdata = reshape(data(store(CT)),Cdims...)
  contracted = commoninds(Ainds, Binds)
  A_only = uniqueinds(Ainds, Binds)
  B_only = uniqueinds(Binds, Ainds)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(contracted)
      push!(ind_dict, i)
  end
  if length(A_only) > 0
      for (idx, i) in enumerate(A_only)
          push!(ind_dict, i)
      end
  end
  if length(B_only) > 0
      for (idx, i) in enumerate(B_only)
          push!(ind_dict, i)
      end
  end
  ctainds = zeros(Int, length(Ainds))
  ctbinds = zeros(Int, length(Binds))
  ctcinds = zeros(Int, length(Cinds))
  for (ii, ia) in enumerate(Ainds)
      ctainds[ii] = findfirst(x->x==ia, ind_dict)
  end
  for (ii, ib) in enumerate(Binds)
      ctbinds[ii] = findfirst(x->x==ib, ind_dict)
  end
  for (ii, ic) in enumerate(Cinds)
      ctcinds[ii] = findfirst(x->x==ic, ind_dict)
  end
  id_op    = CUDA.CUTENSOR.CUTENSOR.CUTENSOR_OP_IDENTITY
  dict_key = ""
  for cc in zip(ctcinds, Cdims)
      dict_key *= string(cc[1]) * "," * string(cc[2]) * ","
  end
  for aa in zip(ctainds, Adims)
      dict_key *= string(aa[1]) * "," * string(aa[2]) * ","
  end
  for bb in zip(ctbinds, Bdims)
      dict_key *= string(bb[1]) * "," * string(bb[2]) * ","
  end
  if haskey(ENV, "CUTENSOR_AUTOTUNE") && tryparse(Int, ENV["CUTENSOR_AUTOTUNE"]) == 1
      if haskey(ContractionPlans, dict_key)
          dict_val = ContractionPlans[dict_key]
          algo  = dict_val
          #plan  = dict_val[2]
          Cdata = CUDA.CUTENSOR.contraction!(α, Adata, Vector{Char}(ctainds), id_op, Bdata, Vector{Char}(ctbinds), id_op, β, Cdata, Vector{Char}(ctcinds), id_op, id_op; algo=algo)
      else
          # loop through all algos
          # pick the fastest one
          # store that plan!
          best_time = 1e6
          best_plan = nothing
          best_algo = nothing
          max_algos = Ref{Int32}(C_NULL)
          CUDA.CUTENSOR.cutensorContractionMaxAlgos(max_algos)
          # fix once the other options are documented
          #algos = collect(Cint(CUDA.CUTENSOR.CUTENSOR_ALGO_GETT):Cint(max_algos[] - 1))
          algos = collect(Cint(CUDA.CUTENSOR.CUTENSOR_ALGO_GETT):Cint(-1))
          for algo in reverse(algos)
              try
                  Cdata, this_time, bytes, gctime, memallocs = @timed CUDA.CUTENSOR.contraction!(α, Adata, Vector{Char}(ctainds), id_op, Bdata, Vector{Char}(ctbinds), id_op, β, Cdata, Vector{Char}(ctcinds), id_op, id_op; algo=CUDA.CUTENSOR.cutensorAlgo_t(algo))
                  if this_time < best_time
                      best_time = this_time
                      #best_plan = this_plan
                      best_algo = CUDA.CUTENSOR.cutensorAlgo_t(algo)
                  end
              catch err
                  @warn "Algorithm $algo not supported"
              end
          end
          ContractionPlans[dict_key] = best_algo
      end
  else
      Cdata = CUDA.CUTENSOR.contraction!(α, Adata, Vector{Char}(ctainds), id_op, Bdata, Vector{Char}(ctbinds), id_op, β, Cdata, Vector{Char}(ctcinds), id_op, id_op)
  end
  return parent(Cdata)
end

function Base.:+(B::CuDenseTensor, A::CuDenseTensor)
  opC  = CUTENSOR.CUTENSOR_OP_IDENTITY
  opA  = CUTENSOR.CUTENSOR_OP_IDENTITY
  opAC = CUTENSOR.CUTENSOR_OP_ADD
  Ais = inds(A)
  Bis = inds(B)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(inds(A))
      push!(ind_dict, i)
  end
  Adata = data(store(A))
  Bdata = data(store(B))
  reshapeBdata = reshape(Bdata,dims(Bis)...)
  reshapeAdata = reshape(Adata,dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
      ctainds[ii] = findfirst(x->x==ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
      ctbinds[ii] = findfirst(x->x==ib, ind_dict)
  end
  ctcinds = copy(ctbinds)
  C = CUDA.zeros(eltype(Bdata), dims(Bis)...)
  CUTENSOR.elementwiseBinary!(one(eltype(Adata)), reshapeAdata, ctainds, opA, one(eltype(Bdata)), reshapeBdata, ctbinds, opC, C, ctcinds, opAC)
  copyto!(data(store(B)), vec(C))
  return B
end

function Base.:+(B::CuDense, Bis::IndexSet, A::CuDense, Ais::IndexSet)
  opA  = CUTENSOR.CUTENSOR_OP_IDENTITY
  opC  = CUTENSOR.CUTENSOR_OP_IDENTITY
  opAC = CUTENSOR.CUTENSOR_OP_ADD
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(Ais)
      push!(ind_dict, i)
  end
  Adata = data(A)
  Bdata = data(B)
  reshapeBdata = reshape(Bdata,dims(Bis)...)
  reshapeAdata = reshape(Adata,dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
      ctainds[ii] = findfirst(x->x==ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
      ctbinds[ii] = findfirst(x->x==ib, ind_dict)
  end
  ctcinds = copy(ctbinds)
  C = CUDA.zeros(eltype(Bdata), dims(Bis)...)
  Cis = Bis
  C = CUTENSOR.elementwiseBinary!(1, reshapeAdata, ctainds, opA, 1, reshapeBdata, ctbinds, opC, C, ctcinds, opAC)
  copyto!(data(B), vec(C))
  return C
end

function Base.:-(B::CuDenseTensor, A::CuDenseTensor)
  opC  = CUTENSOR.CUTENSOR_OP_IDENTITY
  opA  = CUTENSOR.CUTENSOR_OP_IDENTITY
  opAC = CUTENSOR.CUTENSOR_OP_ADD
  Ais = inds(A)
  Bis = inds(B)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(inds(A))
      push!(ind_dict, i)
  end
  Adata = data(store(A))
  Bdata = data(store(B))
  reshapeBdata = reshape(Bdata,dims(Bis)...)
  reshapeAdata = reshape(Adata,dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
      ctainds[ii] = findfirst(x->x==ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
      ctbinds[ii] = findfirst(x->x==ib, ind_dict)
  end
  ctcinds = copy(ctbinds)
  C = CUDA.zeros(eltype(Bdata), dims(Bis))
  CUTENSOR.elementwiseBinary!(one(eltype(Adata)), reshapeAdata, ctainds, opA, -one(eltype(Bdata)), reshapeBdata, ctbinds, opC, C, ctcinds, opAC)
  copyto!(data(store(B)), vec(C))
  return B
end

function Base.:-(A::CuDense, Ais::IndexSet, B::CuDense, Bis::IndexSet)
  opA  = CUTENSOR.CUTENSOR_OP_IDENTITY
  opC  = CUTENSOR.CUTENSOR_OP_IDENTITY
  opAC = CUTENSOR.CUTENSOR_OP_ADD
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(Ais)
      push!(ind_dict, i)
  end
  Adata = data(A)
  Bdata = data(B)
  reshapeBdata = reshape(Bdata,dims(Bis)...)
  reshapeAdata = reshape(Adata,dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
      ctainds[ii] = findfirst(x->x==ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
      ctbinds[ii] = findfirst(x->x==ib, ind_dict)
  end
  ctcinds = copy(ctbinds)
  C = CUDA.zeros(eltype(Bdata), dims(Bis)...)
  Cis = Bis
  C = CUTENSOR.elementwiseBinary!(one(eltype(Adata)), reshapeAdata, ctainds, opA, -one(eltype(Bdata)), reshapeBdata, ctbinds, opC, C, ctcinds, opAC)
  copyto!(data(B), vec(C))
  return C
end

function Base.permute!(B::CuDenseTensor, A::CuDenseTensor)
  Ais = inds(A)
  Bis = inds(B)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(Ais)
      push!(ind_dict, i)
  end
  Adata = data(store(A))
  Bdata = data(store(B))
  reshapeBdata = reshape(Bdata,dims(Bis)...)
  reshapeAdata = reshape(Adata,dims(Ais)...)
  if ndims(A) < 12 # use CUTENSOR
      ctainds = zeros(Int, length(Ais))
      ctbinds = zeros(Int, length(Bis))
      for (ii, ia) in enumerate(Ais)
          ctainds[ii] = findfirst(x->x==ia, ind_dict)
      end
      for (ii, ib) in enumerate(Bis)
          ctbinds[ii] = findfirst(x->x==ib, ind_dict)
      end
      CUDA.CUTENSOR.permutation!(one(eltype(Adata)), reshapeAdata, Vector{Char}(ctainds), reshapeBdata, Vector{Char}(ctbinds))
  else # use GPUArrays
      perm = Int[]
      for aix in Ais
        b_pos = findfirst(bix->bix==aix, Bis)
        push!(perm, b_pos)
      end
      @assert isperm(perm)
      permutedims!(reshapeBdata, reshapeAdata, invperm(perm))
  end
  return vec(reshapeBdata)
end

function Base.permute!(B::CuDense, Bis::IndexSet, A::CuDense, Ais::IndexSet)
  ind_dict = Vector{Index}()
  for (idx, i) in enumerate(Ais)
      push!(ind_dict, i)
  end
  Adata = data(A)
  Bdata = data(B)
  reshapeBdata = reshape(Bdata,dims(Bis)...)
  reshapeAdata = reshape(Adata,dims(Ais)...)
  ctainds = zeros(Int, length(Ais))
  ctbinds = zeros(Int, length(Bis))
  for (ii, ia) in enumerate(Ais)
      ctainds[ii] = findfirst(x->x==ia, ind_dict)
  end
  for (ii, ib) in enumerate(Bis)
      ctbinds[ii] = findfirst(x->x==ib, ind_dict)
  end

  CUDA.CUTENSOR.permutation!(one(eltype(Adata)), reshapeAdata, Vector{Char}(ctainds), reshapeBdata, Vector{Char}(ctbinds))
  return vec(reshapeBdata)
end

Base.:/(A::CuDenseTensor, x::Number) = A * inv(x)


#NOTE: May also want to check multiple years
const YEAR_PAIRS = [(2009,2010),(2013,2014)]

function singleYearPlots(data::NCCSData; yearPairs::Vector{Tuple{Int,Int}} = YEAR_PAIRS,
  outputPath::String = GRAPH_PATH)
  local plotNames::Vector{String} = Vector{String}()
  local plots::Vector{PlotContainer} = Vector{PlotContainer}()
  local subdf::SubDataFrame

  for p ∈ YEAR_PAIRS
    @assert p[2] == p[1] + 1 #only support one year difference for now
    subdf = view(data.df, data.df[:fisyr] .== p[2], :)
    subdf = view(subdf, (!ismissing).(subdf[:lpreturnbyyr]), :)
    subdf = view(subdf, (!ismissing).(subdf[:preturnbyyr]), :)

    push!(plotNames, "Performance Percentile in $(p[2]) v. $(p[1])")
    push!(plots, plot(subdf, x=:lpreturnbyyr, y=:preturnbyyr,
      #size=:lnetassets,
      Guide.xlabel("Return in $(p[1])"),Guide.ylabel("Return in $(p[2])"),
      Geom.point,
      Coord.cartesian(xmin = 0.0, xmax=1.0, ymin=0.0, ymax = 1.0)))
  end

  for i ∈ 1:length(plots) #write the graphs
    draw(PNG("$outputPath\\$(plotNames[i]).png", 9inch, 7inch), plots[i])
    #savefig(plots[i], "$outputPath\\$(plotNames[i])$(graphsuffix).$ext")
    #println("Graph $(plotNames[i]) written.")
  end

  return nothing
end

function persistancePlots(data::NCCSData)

  singleYearPlots(data)

  println("completed persistance plots")
end


cfg = get_config("test.yaml")
cfg["innovation_max"] = 5

@testset "Individual" begin

    ind = NEATInd(cfg)

    n_nodes = cfg["n_in"] + cfg["n_out"]
    @test ind.n_hidden == 0
    @test length(ind.connections) == cfg["n_in"] * cfg["n_out"]
    weights = []
    for c in ind.connections
        @test c.in_node < c.out_node
        @test c.weight < Inf && c.weight > -Inf
        push!(weights, c.weight)
        @test c.enabled || ~c.enabled
        @test c.innovation <= cfg["innovation_max"]
        @test c.innovation <= length(ind.connections) + cfg["innovation_max"]
    end

    @test length(unique(weights)) > 1
end

@testset "Reconstruct individual" begin

    ind = NEATInd(cfg)
    for c in ind.connections
        c.weight = rand()
        if rand() < 0.5
            c.enabled = false
        end
    end

    string_ind = string(ind)
    @test typeof(string_ind) == String

    ind2 = NEATInd(cfg, string_ind)
    test_identical(ind, ind2)
end


using Test, IonSim
using IonSim.PhysicalConstants: m_ca40
using Suppressor

@suppress_err begin

@testset "ion_configurations -- LinearChain" begin
    C = Ca40()
    lc = LinearChain(
            ions=[C, C, C, C], com_frequencies=(x=5, y=5, z=1),
            vibrational_modes=(y=[1], z=[4])
        )
    @test ions(lc) == lc.ions
    # test get_vibrational_modes, which should return an array of the selected
    # VibrationalModes in the linear chain
    vms = lc.vibrational_modes
    @test get_vibrational_modes(lc) == [vms.x..., vms.y..., vms.z...]

    # make sure ion numbers are updated
    @test [I.number for I in lc.ions] == [1, 2, 3, 4]

    # cmode -> pre-evaluated equilibrium positions for four ion chain
    cmode = [-0.06392393573, -0.0202155287427, 0.0202155287427, 0.0639239357]
    if lc.ions[1].position > 0
        cmode .*= -1
    end
    @test [I.position for I in lc.ions] ≈ cmode rtol=1e-6

    # should get warning if same ion is input multiple times to ion kwarg
    warning = "Some ions point to the same thing. Making copies."
    @test_logs (:warn, warning) LinearChain(
        ions=[C, C], com_frequencies=(x=4,y=4,z=1), vibrational_modes=(x=[],y=[],z=[1,2])
    )
    # and copies should be made of the repeated ions, so that they are no longer the same
    chain1 = LinearChain(
        ions=[C, C], com_frequencies=(x=4,y=4,z=1), vibrational_modes=(x=[],y=[],z=[1,2])
    )
    @test !(chain1.ions[1] ≡ chain1.ions[2])

    # Make sure there are no errors with print/show
    @suppress print(lc); show(lc)
end

@testset "ion_configurations -- general" begin
    # test calculation of equilibrium positions for a linear chain of equal mass ions against
    # known values
    posL = [-2.8708, -2.10003, -1.4504, -0.85378, -0.2821]
    pos = [posL; -1 .* reverse(posL)]
    @test any(isapprox.(linear_equilibrium_positions(10), pos, rtol=1e-4))

    # and test calculation of characterstic length scale for linear chain, equal mass
    @test characteristic_length_scale(m_ca40, 1e6) ≈ 4.449042804354e-6

    # and do the same for Anm, which computes the normal modes
    @test_throws AssertionError Anm(2, (x=0.5, y=0.5, z=1), (x=1, y=0, z=0))
    cst = [-0.2132, 0.6742, -0.6742, 0.2132]
    freq, mode = Anm(4, (x=2, y=2, z=1), (x=0, y=0, z=1))[end]
    @test freq ≈ √9.308 rtol=1e-4
    if mode[1] > 0
        cst = -1 .* cst
    end
    @test any(isapprox.(mode, cst, rtol=1e-4))

    # test '_sparsify', which should drop small values from an array
    x = [1e-6, 1e-5]
    IonSim._sparsify!(x, 2e-6)
    @test any(x .== [0, 1e-5])
end
end  # end suppress


# CSCI5573Final/daggen.jl
# Licensed under the MIT License. See LICENSE.md file in the project root for
# full license information.

# This is all a bit ugly. Lets just not talk about it.
struct Edge
    ID1::Int64
    ID2::Int64
end

function ListToDictDAG(tasklist, dot_file="")
    dag = Dict{Int64,_Task}()
    comm_cost = Dict{Edge, Int64}()

    for task in tasklist[1:end-1]
        new_task = ParseTask(task)
        dag[new_task.ID] = new_task
    end

    new_task = ParseTask(tasklist[end], true)
    dag[new_task.ID] = new_task

    for (ID, _task) in dag
        for child in _task.Children
            add_dependency!(dag[child], ID)
        end
    end

    # writeDOT("pre", dag)

    # Extract edges
    for (ID, _task) in copy(dag)
        if _task.Type === :TRANSFER

            parent = _task.Dependencies[begin]
            child = _task.Children[begin]

            # Add edges. Just for giggles we do both directions
            comm_cost[Edge(parent, child)] = _task.Cost
            comm_cost[Edge(child, parent)] = _task.Cost

            # Remove edge from the parent child relationship
            filter!(e->e!==ID, dag[parent].Children)
            filter!(e->e!==ID, dag[child].Dependencies)

            # Reconnect the parent and child without edge
            push!(dag[parent].Children, child)
            push!(dag[child].Dependencies, parent)

            # Delete Edge from dag
            delete!(dag, ID)
        end
    end

    # Finally add in all of the non-comm edges
    for (ID, _task) in copy(dag)
        for child in _task.Children
            if !haskey(comm_cost, Edge(ID, child))
                comm_cost[Edge(ID, child)] = 0
                comm_cost[Edge(child, ID)] = 0
            end
        end
    end

    # writeDOT("post", dag)

    if length(dot_file) > 0
        writeDOT(dot_file, dag)
    end

    return dag, comm_cost
end

function daggen_init()
    name = "daggen"
    dir = "daggen"
    touch("daggen_init")

    init_script = """
                #!/bin/bash
                pushd $dir
                make
                mv $(name) ../dag_generator
                rm *.o
                popd"""

    open("daggen_init", "w") do file
        write(file, init_script)
    end

    chmod("daggen_init", 2000)
    run(`./daggen_init`)
    rm("daggen_init")
end

"""
- `outfile`: output file
- `n`: number of tasks
- `mindata`: minimum data size
- `maxdata`: maximum data size
- `minalpha`: minimum Amdahl's law parameter value
- `maxalpha`: maximum Amdahl's law parameter value
- `fat`: dag shape - 1.0=max parallel, 0.0=min parallel
- `density`: number of dependencies
- `regular`: regularity for num tasks per level
- `ccr`: communication(MBytes) to computation(sec) ratio
- `jump`: number of levels spanned by communications
- `dot`: output generated DAG in the DOT format
"""
function daggen(;
    outfile="", num_tasks=-1, mindata=-1, maxdata=-1, minalpha=-1, maxalpha=-1,
    fat=-1, density=-1, regular=-1, ccr=-1, jump=-1, dot=false, verbose=false)

    err = verbose ? "daggen_err.txt" : nothing

    base = `./dag_generator`
    out = (outfile == "") ? `` : ` -o $outfile`
    num = (num_tasks == -1) ? `` : ` -n $(Int64(round(num_tasks/2-sqrt(num_tasks+1))))` # Weird Stuff here
    mind = (mindata == -1) ? `` : ` --mindata $mindata`
    maxd = (maxdata == -1) ? `` : ` --maxdata $maxdata`
    mina = (minalpha == -1) ? `` : ` --minalpha $minalpha`
    maxa = (maxalpha == -1) ? `` : ` --maxalpha $maxalpha`
    fat = (fat == -1) ? `` : ` --fat $fat`
    cr = (ccr == -1) ? `` : ` --ccr $ccr`
    jmp = (jump == -1) ? `` : ` --jump $jump`
    dt = !dot ? `` : ` --dot`

    command = `$base $out $num $mind $maxd $mina $maxa $fat $cr $jmp $dt`
    daggen_init()

    buffer = IOBuffer()
    ps = run(pipeline(`$(command)`, stdout=buffer, stderr=err), wait=true)

    daggen_clean()

    # This processes the IOBuffer as a String array
    tasklist = buffer |> take! |> String |> l->split(l,"\n") |> l->deleteat!(l,[1,2,length(l)])

    num_task = tasklist[1] |> c->split(c, " ") |> n->n[2] |> i->parse(Int, i)

    deleteat!(tasklist, 1) # delete count from tasklist

    return String.(tasklist), num_task
end

function read_daggen(file)
    config = readlines(file)

    # This processes the IOBuffer as a String array
    tasklist = config |> l->deleteat!(l,[1,2])

    num_task = tasklist[1] |> c->split(c, " ") |> n->n[2] |> i->parse(Int, i)

    deleteat!(tasklist, 1) # delete count from tasklist

    return String.(tasklist), num_task
end

function daggen_clean()
    return run(`rm dag_generator`)
end

function writeDOT(Filename, dag)
    open(Filename, "w") do file
        write(file, "digraph DAG_Schedule {\n")

        for (_, task) in sort(dag) # Write self
            color = task.Type == :TRANSFER ? "grey" : "black"

            node = "  T$(task.ID) [size=\"$(task.Cost)\", "
            node *= "overhead=\"$(task.Complexity)\", "
            node *= "color=\"$(color)\"]\n"
            write(file, node)

            for ID in task.Children
                child = dag[ID]
                edge = "  T$(task.ID) -> T$(child.ID) "
                edge *= "[size=\"$(task.Cost)\", "
                edge *= "color=\"$(color)\"]\n"
                write(file, edge)
            end
        end

        write(file, "}\n")
    end
end


export rhotrace,
    rhotrace!,
    steptrace,
    steptrace!,
    radiustrace,
    radiustrace!,
    objtrace,
    objtrace!,
    gradtrace,
    gradtrace!

function log_display(v, name)
    if (m = minimum(v)) > 0.0
        return v
    elseif m ≈ 0.0
        @warn "$name: shifted by $(eps(Float64)) to display on Log Scale"
        return v .+ eps(Float64)
    end
    @warn "$name: shifted by $m to display on Log Scale"
    return v .+ m
end

"""
    objtrace

Provide any number of terminal simulations as arguments and `objtrace` displays
a plot of each arguments' objective function values at each succussful iterate.

To add a simulation to the current `objtrace` plot use `objtrace!`.
"""
function objtrace end

"""
    gradtrace

Provide any number of terminal simulations as arguments and `gradtrace` displays
a plot of each arguments' normed gradient values at each succussful iterate.

To add a simulation to the current `gradtrace` plot use `gradtrace!`.
"""
function gradtrace end

"""
    radiustrace

Provide any number of terminal simulations as arguments and `radiustrace` displays
a plot of each arguments' trust region subproblem radius values at each succussful iteration.

To add a simulation to the current `radiustrace` plot use `radiustrace!`.
"""
function radiustrace end

"""
    steptrace

Provide any number of terminal simulations as arguments and `steptrace` displays
a plot of each arguments' step measure at each succussful iteration.

To add a simulation to the current `steptrace` plot use `steptrace!`.
"""
function steptrace end

"""
    rhotrace

Provide any number of terminal simulations as arguments and `rhotrace` displays
a plot of each arguments' ρ ratio at each successful iteration.

To add a simulation to the current `rhotrace` plot use `rhotrace!`.
"""
function rhotrace end

@userplot ObjTrace

@recipe function f(result::ObjTrace)
    num_args = length(result.args)

    m = model(result.args[1])

    y_lab = (isa(formula(m), Missing) || num_args > 1 ? "\$ f \$" : formula(m))

    xguide := "Iterates (gHS evaluations)"

    yguide := y_lab

    yguidefontsize := 7

    xguidefontsize := 7

    title --> "Objective Trace"

    titlefontsize --> 11

    yaxis --> :log

    legend --> :bottomright

    @series begin
        x = []

        labels = Matrix{String}(undef, 1, num_args)

        # Current idea is to accept trace, backend, simulation objects (may need to include model, driver in simulation.jl)
        for i ∈ 1:num_args

            s = result.args[i]

            qn_update = string(QN_update(driver(s)))

            s_update = "6.1." * string(S_update(driver(s)))[end]

            two_w = samples(driver(s))

            dim = dimension(model(s))

            id = name(model(s))

            v = log_display(f_vals(s), id)

            labels[
                1,
                i,
            ] = "\$ f \\quad \\mathrm{$id}: ~ n=$dim \\quad w=$(Int(two_w/2.0)) \\quad (\\mathrm{$s_update}) \\quad \\mathrm{$qn_update} \$"

            push!(x, v)
        end

        label --> labels

        x
    end
end

@userplot GradTrace

@recipe function f(result::GradTrace)
    num_args = length(result.args)

    m = model(result.args[1])

    xguide := "Iterates (gHS evaluations)"

    yguide := "\$ ||\\nabla f ||\$"

    yguidefontsize := 7

    xguidefontsize := 7

    title --> "Normed Gradient Trace"

    titlefontsize --> 11

    yaxis --> :log

    legend --> :bottomright
    @series begin
        x = []

        labels = Matrix{String}(undef, 1, num_args)

        for i ∈ 1:num_args

            s = result.args[i]

            qn_update = string(QN_update(driver(s)))

            s_update = "6.1." * string(S_update(driver(s)))[end]

            two_w = samples(driver(s))

            dim = dimension(model(s))

            id = name(model(s))

            v = log_display(∇f_norms(s), id)

            labels[
                1,
                i,
            ] = "\$ ||\\nabla f|| ~ \\mathrm{$id}: ~ n=$dim \\quad w=$(Int(two_w/2.0)) \\quad (\\mathrm{$s_update}) \\quad \\mathrm{$qn_update} \$"

            push!(x, v)
        end

        label --> labels

        x
    end
end

@userplot RadiusTrace

@recipe function f(result::RadiusTrace)

    xguide := "Successful Steps"

    yguide := "\$ \\Delta_k \$"

    yguidefontsize := 7

    xguidefontsize := 7

    title --> "Trust-Region Subproblem Radius Trace"

    titlefontsize --> 11

    yaxis --> :log

    legend --> :bottomright

    @series begin
        x = []

        labels = Matrix{String}(undef, 1, length(result.args))

        for i ∈ 1:length(result.args)

            s = result.args[i]

            qn_update = string(QN_update(driver(s)))

            s_update = "6.1." * string(S_update(driver(s)))[end]

            two_w = samples(driver(s))

            dim = dimension(model(s))

            id = name(model(s))

            v = log_display(Δ_vals(s), id)

            labels[
                1,
                i,
            ] = "\$ \\Delta_k ~ \\mathrm{$id}: ~ n=$dim \\quad w=$(Int(two_w/2.0)) \\quad (\\mathrm{$s_update}) \\quad \\mathrm{$qn_update} \$"

            push!(x, v)
        end

        label --> labels

        x
    end
end

@userplot StepTrace

@recipe function f(result::StepTrace)

    xguide := "Successful Steps"

    yguide := "\$ ||p_k|| \$"

    yguidefontsize := 7

    xguidefontsize := 7

    title --> "Step Distance Trace"

    titlefontsize --> 11

    yaxis --> :log

    legend --> :bottomright

    @series begin
        x = []

        labels = Matrix{String}(undef, 1, length(result.args))

        for i ∈ 1:length(result.args)

            s = result.args[i]

            qn_update = string(QN_update(driver(s)))

            s_update = "6.1." * string(S_update(driver(s)))[end]

            two_w = samples(driver(s))

            dim = dimension(model(s))

            id = name(model(s))

            v = log_display(p_norms(s), id)

            labels[
                1,
                i,
            ] = "\$ ||p_k|| ~ \\mathrm{$id}: ~ n=$dim \\quad w=$(Int(two_w/2.0)) \\quad (\\mathrm{$s_update}) \\quad \\mathrm{$qn_update} \$"

            push!(x, v)
        end

        label --> labels

        x
    end

end

@userplot RhoTrace

@recipe function f(result::RhoTrace)

    xguide := "Successful Steps"

    yguide := "\$ \\rho \$"

    yguidefontsize := 7

    xguidefontsize := 7

    title --> "Actual Reduction to Model Reduction Ratio Trace"

    titlefontsize --> 11

    yaxis --> :log

    legend --> :bottomright

    @series begin
        x = []

        labels = Matrix{String}(undef, 1, length(result.args))

        for i ∈ 1:length(result.args)

            s = result.args[i]

            qn_update = string(QN_update(driver(s)))

            s_update = "6.1." * string(S_update(driver(s)))[end]

            two_w = samples(driver(s))

            dim = dimension(model(s))

            id = name(model(s))

            v = log_display(ρ_vals(s), id)

            labels[
                1,
                i,
            ] = "\$ \\rho ~ \\mathrm{$id}: ~ n=$dim \\quad w=$(Int(two_w/2.0)) \\quad (\\mathrm{$s_update}) \\quad \\mathrm{$qn_update} \$"

            push!(x, v)
        end

        label --> labels

        x
    end
end


module ExtendableGrids

using DocStringExtensions
using ElasticArrays
using AbstractTrees

using InteractiveUtils
using SparseArrays
using Printf
using Random
using Dates
using LinearAlgebra

include("adjacency.jl")
export Adjacency,VariableTargetAdjacency,FixedTargetAdjacency
export atranspose,num_targets,num_sources,num_links,append!, max_num_targets_per_source
export asparse,tryfix,makevar

include("vectorofconstants.jl")
export VectorOfConstants

include("typehierarchy.jl")
export AbstractExtendableGridApexType
export typehierarchy

include("elementgeometry.jl")
export AbstractElementGeometry, ElementInfo
export elementgeometries

export AbstractElementGeometry0D
export Vertex0D

export AbstractElementGeometry1D
export Edge1D

export AbstractElementGeometry2D
export Polygon2D,Triangle2D,Quadrilateral2D,Pentagon2D,Hexagon2D,Parallelogram2D,Circle2D

export AbstractElementGeometry3D
export Polyhedron3D,Tetrahedron3D, Hexahedron3D,Parallelepiped3D,Prism3D,TrianglePrism3D,Sphere3D

export AbstractElementGeometry4D
export Polychoron4D,HyperCube4D

export dim_element

include("coordinatesystem.jl")
export coordinatesystems
export AbstractCoordinateSystem
export Cartesian1D,Cartesian2D,Cartesian3D
export Cylindrical2D,Cylindrical3D
export Polar2D,Polar1D ,Spherical3D,Spherical1D

include("extendablegrid.jl")
export ExtendableGrid
export instantiate, veryform
export AbstractGridComponent
export AbstractGridAdjacency,AbstractElementGeometries,AbstractElementRegions
export Coordinates,CellNodes,BFaceNodes,CellGeometries,BFaceGeometries,CellRegions,BFaceRegions
export NumCellRegions,NumBFaceRegions,CoordinateSystem
export AbstractGridFloatArray1D,AbstractGridFloatArray2D
export AbstractGridIntegerArray1D,AbstractGridIntegerArray2D
export index_type, coord_type
export dim_space, dim_grid
export num_nodes, num_cells, num_bfaces, num_cellregions, num_bfaceregions
export gridcomponents
export seemingly_equal

include("subgrid.jl")
export subgrid

include("more.jl")
export EdgeNodes, CellEdges,EdgeCells,BFaceCells,BFaceNormals
export local_celledgenodes,num_edges

include("regionedit.jl")
export cellmask!,bfacemask!

include("simplexgrid.jl")
export simplexgrid, geomspace,glue
export XCoordinates, YCoordinates, ZCoordinates
export writefile

include("tokenstream.jl")
export TokenStream, gettoken, expecttoken,trytoken

end # module


const ContactImplicitMPC = Main
include(joinpath(@__DIR__, "..", "src/dynamics", "pushbot", "visuals.jl"))
vis = Visualizer()
# render(vis)
open(vis)

s = get_simulation("pushbot", "flat_2D_lc", "flat")
model = s.model
env = s.env

# time
h = 0.04
H = 100

# reference trajectory
ref_traj = contact_trajectory(model, env, H, h)
ref_traj.h
qref = [0.0; 0.0]
ur = zeros(model.dim.u)
γr = zeros(model.dim.c)
br = zeros(model.dim.c * friction_dim(env))
ψr = zeros(model.dim.c)
ηr = zeros(model.dim.c * friction_dim(env))
wr = zeros(model.dim.w)

# set reference
for t = 1:H
	ref_traj.z[t] = pack_z(model, env, qref, γr, br, ψr, ηr)
	ref_traj.θ[t] = pack_θ(model, qref, qref, ur, wr, model.μ_world, ref_traj.h)
end

# test reference
for t = 1:H
	r = ContactImplicitMPC.residual(model, env, ref_traj.z[t], ref_traj.θ[t], 0.0)#[model.dim.q .+ (1:model.dim.c)]
	@test norm(r) < 1.0e-4
end

# initial conditions
q0 = @SVector [0.0 * π, 0.0]
q1 = @SVector [0.0 * π, 0.0]

# simulator
sim = ContactImplicitMPC.simulator(s, q0, q1, h, H,
	ip_opts = ContactImplicitMPC.InteriorPointOptions(
		r_tol = 1.0e-6, κ_tol = 1.0e-5),
	sim_opts = ContactImplicitMPC.SimulatorOptions(warmstart = false))

# simulate
status = ContactImplicitMPC.simulate!(sim)
@test status

# MPC
N_sample = 2
H_mpc = 40
h_sim = h / N_sample
H_sim = 1000

# barrier parameter
κ_mpc = 1.0e-4

# SLOW RECOVERY
obj = TrackingVelocityObjective(model, env, H_mpc,
	q = [Diagonal([12*(t/H_mpc)^2; 2.0*(t/H_mpc)^4]) for t = 1:H_mpc-0],
	v = [Diagonal([1; 0.01] ./ (h^2.0)) for t = 1:H_mpc-0],
	u = [Diagonal([100; 1]) for t = 1:H_mpc-0],
	γ = [Diagonal(1.0e-100 * ones(model.dim.c)) for t = 1:H_mpc-0],
	b = [Diagonal(1.0e-100 * ones(model.dim.c * friction_dim(env))) for t = 1:H_mpc])
# FAST RECOVERY
obj = TrackingVelocityObjective(model, env, H_mpc,
    q = [Diagonal([12*(t/H_mpc)^2; 12*(t/H_mpc)^2]) for t = 1:H_mpc-0],
	v = [Diagonal([1; 0.01] ./ (h^2.0)) for t = 1:H_mpc-0],
    u = [Diagonal([100; 1]) for t = 1:H_mpc-0],
    γ = [Diagonal(1.0e-100 * ones(model.dim.c)) for t = 1:H_mpc-0],
    b = [Diagonal(1.0e-100 * ones(model.dim.c * friction_dim(env))) for t = 1:H_mpc])

p = linearized_mpc_policy(ref_traj, s, obj,
    H_mpc = H_mpc,
    N_sample = N_sample,
    κ_mpc = κ_mpc,
    n_opts = NewtonOptions(
		r_tol = 3e-4,
		max_iter = 10,
		max_time = ref_traj.h/2, # HARD REAL TIME
		# verbose = true,
		),
    mpc_opts = LinearizedMPCOptions(),
    )

idx_d1 = 20
idx_d2 = idx_d1 + 200
idx_d3 = idx_d2 + 80
idx_d4 = idx_d3 + 200
idx_d5 = idx_d4 + 30
idx = [idx_d1, idx_d2, idx_d3, idx_d4, idx_d5]
impulses = [[-5.5; 0.0], [+5.5; 0.0], [+5.5; 0.0], [-1.5; 0.0], [-6.5; 0.0]]
d = impulse_disturbances(impulses, idx)

q1_sim = SVector{model.dim.q}([0.0, 0.0])
q0_sim = SVector{model.dim.q}([0.0, 0.0])

sim = ContactImplicitMPC.simulator(s, q0_sim, q1_sim, h_sim, H_sim,
    p = p,
	d = d,
    sim_opts = ContactImplicitMPC.SimulatorOptions(warmstart = true))

status = ContactImplicitMPC.simulate!(sim, verbose = true)
# @profiler status = ContactImplicitMPC.simulate!(sim, verbose = true)

################################################################################
# Timing result
################################################################################
process!(sim)
# Time budget
ref_traj.h # 0.04
# Time used on average
sim.stats.μ_dt # 0.0138
sim.stats.σ_dt # 0.0138
# Speed ratio
H_sim * h_sim / sum(sim.stats.dt) # 2.90

l = 1
lu = 1
nc = model.dim.c
plt = plot(layout=(3,1), legend=false)
plot!(plt[1,1], hcat(Vector.(vcat([fill(ref_traj.q[i], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[1,1], hcat(Vector.([q[l:l] for q in sim.traj.q])...)', color=:blue, linewidth=1.0)
plot!(plt[2,1], hcat(Vector.(vcat([fill(ref_traj.u[i][lu:lu], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[3,1], hcat(Vector.(vcat([fill(ref_traj.γ[i][1:nc], N_sample) for i=1:H]...))...)',
    color=:red, linewidth=3.0)
plot!(plt[2,1], hcat(Vector.([u[lu:lu] for u in sim.traj.u]*N_sample)...)', color=:blue, linewidth=1.0)

sample = 1

anim = visualize_robot!(vis, model, sim.traj, sample = 1)#, sample = sample, α=1.0)
pθ_right = generate_pusher_traj(d, sim.traj, side=:right)
pθ_left  = generate_pusher_traj(d, sim.traj, side=:left)
visualize_disturbance!(vis, model, pθ_right, anim=anim, sample=sample, offset=0.05, name=:PusherRight)
visualize_disturbance!(vis, model, pθ_left,  anim=anim, sample=sample, offset=0.05, name=:PusherLeft)

γ_max = maximum(hcat(sim.traj.γ...))
u_max = maximum(hcat(sim.traj.u...))

plot((hcat(sim.traj.γ...) ./ γ_max)', linetype = :steppost)
plot!((hcat(sim.traj.u...) ./ u_max)[2:2, :]', linetype = :steppost)
plot((hcat(sim.traj.u...) ./ u_max)[1:1, :]', linetype = :steppost)
plot(hcat(sim.traj.q...)')

plot([pθ[1] for pθ in pθ_right])
plot!([pθ[1] for pθ in pθ_left])
plot!([q[1] for q in sim.traj.q[3:end]])

# Display highlights
t_highlights = [20,23,35,216]
α_highlights = ones(4)
α_pusher = [1.0, 0.0, 0.0, 0.0]

# # Initial push
# i = 1
# t = t_highlights[i]
# α = i/(length(t_highlights)+1)
# name = Symbol("High$i")
# pusher_name = Symbol("HighPusherRight$i")
# build_robot!(vis, model, name=name, α=α_highlights[i])
# build_disturbance!(vis, model, name=pusher_name, α=α_highlights[i])
# set_robot!(vis, model, sim.traj.q[t+2], name=name)
# set_disturbance!(vis, model, pθ_left[t], name=pusher_name, offset=0.05)

# # Contact with wall
# i = 2
# t = t_highlights[i]
# α = i/(length(t_highlights)+1)
# name = Symbol("High$i")
# pusher_name = Symbol("HighPusherRight$i")
# build_robot!(vis, model, name=name, α=α_highlights[i])
# build_disturbance!(vis, model, name=pusher_name, α=α_pusher[i])
# set_robot!(vis, model, sim.traj.q[t+2], name=name)
# set_disturbance!(vis, model, pθ_left[t], name=pusher_name, offset=0.05)

# # Recovery main body
# i = 3
# t = t_highlights[i]
# α = i/(length(t_highlights)+1)
# name = Symbol("High$i")
# pusher_name = Symbol("HighPusherRight$i")
# build_robot!(vis, model, name=name, α=α_highlights[i])
# build_disturbance!(vis, model, name=pusher_name, α=α_pusher[i])
# set_robot!(vis, model, sim.traj.q[t+2], name=name)
# set_disturbance!(vis, model, pθ_left[t], name=pusher_name, offset=0.05)

# Recovery hand
i = 4
t = t_highlights[i]
α = i/(length(t_highlights)+1)
name = Symbol("High$i")
pusher_name = Symbol("HighPusherRight$i")
build_robot!(vis, model, name=name, α=α_highlights[i])
build_disturbance!(vis, model, name=pusher_name, α=α_pusher[i])
set_robot!(vis, model, sim.traj.q[t+2], name=name)
set_disturbance!(vis, model, pθ_left[t], name=pusher_name, offset=0.05)

# filename = "pushbot_multiple_pushes2"
# MeshCat.convert_frames_to_video(
#     "/home/simon/Downloads/$filename.tar",
#     "/home/simon/Documents/video/$filename.mp4", overwrite=true)
#
# convert_video_to_gif(
#     "/home/simon/Documents/video/$filename.mp4",
#     "/home/simon/Documents/video/$filename.gif", overwrite=true)


const JOBS = Channel{String}(32)
const OUT = Channel{String}(32)

"""
    runjobs()
"""
function runjobs()
    while true
        job = take!(JOBS)
        if job == "job_rand_sleep"
            put!(OUT, """{ "payload": "Job started. This may take some time..." }""")
            result = 3 * rand()
            sleep(result)
            put!(OUT, """{ "payload": $result }""")
        end
    end
end

"""
    startws(port::Integer)
"""
function startws(port::Integer=3020, nthreads=Threads.nthreads())
    @async HTTP.listen(HTTP.Sockets.localhost, port) do http
        if HTTP.WebSockets.is_upgrade(http.message)
            HTTP.WebSockets.upgrade(http) do ws
                @sync begin
                    @async while !eof(ws)
                        msg = String(readavailable(ws))
                        put!(JOBS, msg)
                    end
                    @async while isopen(ws)
                        write(ws, take!(OUT))
                    end
                end
            end
        end
    end
    println("WebSocket server is running at http://$(HTTP.Sockets.localhost):$port")

    for _ = 1:nthreads
        @async runjobs()
    end
end


#=
We want to simulate a number of machines that each produce a specific product and
put it in a warehouse (a `Store`). There is a combination process that needs one
of each generated products and combines it into another. The simulation ends when
the fictive container of combined goods is full.

We create our own type `Product` with two fields that allow to identify the kind of
product and to identify its "creator" by means of a serial number.
We also create a type `Machine` that has an associated process. This machine has
an ID an make one type of products. All generated product are put into the same
store.

The combiner process works as follows:
* generate the events matching the availability of one of each product
* wait until all of these events are realised. *Note*: if you only have two events,
you can simply use `ev1` & `ev2`.
* generate the event of putting a fictive combined product in a container
* @yield the event (i.e. time-out until done)
* verify if the container is full and if so stop the simulation on this event.

This application illustrates how you can wait on multiple other events before
continuing the simulation. Keep in mind that this requires ALL events to be
processed.
=#

using Logging
using SimJulia
using ResumableFunctions
import Base.show

struct Product
    kind::Symbol
    serial_number::String
end

show(io::Core.IO, p::Product) =  print(io, "$(p.kind) with serial N° $(p.serial_number)")

mutable struct Machine
    id::String
    prod_kind::Symbol
    prod_time::Int
    n_prod::Int
    prod_store::Store{Product}
    proc::Process
    function Machine(id::String,
                     prod_kind::Symbol, prod_time::Int,
                     prod_store::Store{Product}, n_prod::Int=0)
        m = new()
        m.id = id
        m.prod_kind = prod_kind
        m.prod_time = prod_time
        m.n_prod = n_prod
        m.prod_store = prod_store
        m.proc = @process produce(m.prod_store.env, m)
        return m
    end
end

@resumable function produce(env::Environment, m::Machine)
    while true
        @yield timeout(env, m.prod_time)
        m.n_prod += 1
        newprod = Product(m.prod_kind, "$(m.id)-"*lpad("$(m.n_prod)", 4, "0"))
        @debug "$(env.time) - Machine $(m.id) produced a $(newprod)"
        @yield put(m.prod_store, newprod)
    end
end

@resumable function combiner(env::Environment, prod_store::Store{Product},
                             products::Array{Symbol, 1}, result_container::Container)
    while true
        # obtain a request for each product kind
        requests = map(x::Symbol -> get(prod_store, p::Product -> p.kind == x), products)
        @debug "$(env.time) - Combiner requests: $(requests)"
        # all requests must be matched
        @yield Operator(SimJulia.eval_and, requests...)
        @debug "$(env.time) - Combiner got all required products!"
        newcomb = put(result_container, 1)
        @yield newcomb
        # stop the simulation after the final combination was made
        if result_container.level == result_container.capacity
            @info "Combiner made all products in $(env.time) time units"
            SimJulia.stop_simulation(newcomb)
        end
        @debug "$(env.time) - Current container level $(result_container.level)"
    end

end

function mysim()
    # setup simulation
    @info "\n$("-"^70)\nMachine application\n$("-"^70)\n"
    sim = Simulation()
    # products
    prod_store = Store{Product}(sim)
    products = [:nut, :bolt, :rivet, :beam]
    prod_times = [1;2;4;6]
    # combined results
    result_container = Container(sim, 10)
    # machines
    machines = [Machine("Machine $(i)",products[i], prod_times[i],prod_store) for i in 1:length(products)]
    @process combiner(sim, prod_store, products, result_container)
    run(sim)
end

# logging settings
Logging.disable_logging(LogLevel(-1001)) # (de-)activate debug messages
logger = Logging.ConsoleLogger(stdout, Logging.Debug)

# run simulation
with_logger(logger) do
    mysim()
end


"""
    figsize(width_in_inch, height_width_ratio=0.68)

    Convert figure size in inches (which is often required by publishers) to that in points (which is the unit for CairoMakie). Note that 1 point in CairoMakie is equal to 1/72 inch.
"""
function figsize(width_in_inch, height_width_ratio=HWRATIO)
    width_in_point = floor(Int, 72width_in_inch)
    height_in_point = floor(Int, width_in_point*height_width_ratio)
    return width_in_point, height_in_point
end

savefig(f, fig) = save(f, fig, pt_per_unit=1)


using DataFrames
using CSV
using FileIO
using Statistics

include("experiments/experimentalutils.jl")

resultsFolder = mainfolder * "experiments/twostage_vae_scoretest/"
files = readdir(resultsFolder)

results = []
for f in files
    if isfile(resultsFolder * f) && occursin(".csv", f)
        push!(results, CSV.read(resultsFolder * f))
    end
end
results = vcat(results...)

using Plots
using StatsPlots
plotlyjs()

# p1 = plot()
# for (i, n) in enumerate(names(results)[7:13])
#     scatter!(p1, ones(4) .* i, results[n], ylim = (0,1), label = String(n))
# end

# p2 = plot()
# for i in 1:4
#     scatter!(p2, 1:7, [results[i, c] for c in 7:13], ylim = (0,1))
# end

# for n in names(results)[7:end]
#     @df results boxplot(:dataset, n, size = [700, 700])
# end
pp = []
for n in names(results)[8:14]
    i = 1
    p = plot(title = String(n))
    for d in unique(results[:dataset])
        data = vec(results[results[:dataset] .== d, n])
        scatter!(p, ones(size(data)) .* i, data, ylim = (0,1), label = d)
        # scatter!(p, ones(size(data)) .* i, data, zcolor = vec(results[results[:dataset] .== d, 14]), ylim = (0,1), label = d)
        i += 1
    end
    push!(pp, p)
end
push!(pp, plot())
plot(pp..., layout = (2,4), size = (1200, 700))

aucs = names(results)[8:14]
# fit = names(results)[15:20]
fit = names(results)[[15, 16, 19, 20]]

pp = []
for (i, f) in enumerate(fit)
    for (j, a) in enumerate(aucs)
        p = plot()
        for d in unique(results[:dataset])
            p = scatter!(results[results[:dataset] .== d, a], results[results[:dataset] .== d, f])
        end
        if j == 1
            ylabel!(p, String(f))
        end
        if i == length(fit)
            xlabel!(p, String(a))
        end
        push!(pp, p)
    end
end
plot(pp..., layout = (length(fit), length(aucs)), size = (1400, 800), legend = nothing)

ps = plot()
for d in unique(results[:dataset])
    scatter3d!(ps, results[results[:dataset] .== d, :rec_err_test], results[results[:dataset] .== d, :u_test_dst], results[results[:dataset] .== d, :auc_pxv_pz_jacobian_dec])
end
xlabel!(ps, "rec_err_test")
ylabel!(ps, "u_test_dst")
plot(ps, title = "auc_pxv_pz_jacobian_dec")

ps = plot()
d = "blood-transfusion"
res = results[results[:dataset] .== d, :]
for i in unique(results[:i])
    scatter3d!(ps, res[res[:i] .== i, :rec_err_test], res[res[:i] .== i, :u_test_dst], res[res[:i] .== i, :auc_pxv_pz_jacobian_dec])
end
xlabel!(ps, "rec_err_test")
ylabel!(ps, "u_test_dst")
plot(ps, title = "auc_pxv_pz_jacobian_dec")

gdfs = [by(results, [:dataset, :i], df -> DataFrame([[maximum(df[n])]],[n])) for n in names(results)[8:14]]
all = reduce((a...) -> join(a..., on = [:dataset, :i]), gdfs)

pp = []
for n in names(all)[3:9]
    i = 1
    p = plot(title = String(n))
    for d in unique(all[:dataset])
        data = vec(all[all[:dataset] .== d, n])
        scatter!(p, ones(size(data)) .* i, data, ylim = (0,1), label = d)
        # scatter!(p, ones(size(data)) .* i, data, zcolor = vec(all[all[:dataset] .== d, 14]), ylim = (0,1), label = d)
        i += 1
    end
    push!(pp, p)
end
push!(pp, plot())
plot(pp..., layout = (2,4), size = (1200, 700))

# aggres = []
# for d in unique(results[:dataset])
#     ddf = results[results[:dataset] .== d, :]
#     mean_auc_pxv = mean(ddf[:auc_pxv])
#     mean_auc_pz = mean(ddf[:auc_pz])
#     push!(aggres, DataFrame(dataset = d, auc_pxv = mean_auc_pxv, auc_pz = mean_auc_pz))
# end
# aggres = vcat(aggres...)


"""
    struct GlobalMatrices
Stores global matrices and some other things.
Everything that is very expensive to compute which stays
constant throughout the simulation should be stored here.
    GlobalMatrices(basis::Basis, filter::Filter, dimensions)
Initialize GlobalMatrices for `basis`, `filter` and `dimensions`.
"""
 struct GlobalMatrices
     normalsigns::Dict{Face, Int64}
     normalidxs::Dict{Face, Int64}
     oppositefaces::Dict{Face, Face}

     function GlobalMatrices(dimensions)
        normalsigns = Dict(left => -1, right => 1, top => 1, bottom => -1)
        normalidxs = Dict(left => 1, right => 1, top => 2, bottom => 2)
        oppositefaces = Dict(left => right, right => left, top => bottom, bottom => top)

         new(normalsigns,
            normalidxs,
            oppositefaces)
     end
 end


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

Node-specific code

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

#===========================
Node constructors
===========================#

function Node{T,M}(
    vals,
    height;
    flags = 0x0,
    capacity = DEFAULT_NODE_CAPACITY,
    max_height = DEFAULT_MAX_HEIGHT
) where {T,M}

    height = min(height, max_height)
    next = Vector{Node{T,M}}(undef, height)
    Node{T,M}(vals, next, capacity, flags)
end

LeftSentinel{T,M}(; max_height = DEFAULT_MAX_HEIGHT, kws...) where {T,M} =
    Node{T,M}(Vector{T}(undef,0), max_height; flags = FLAG_IS_LEFT_SENTINEL, kws...)

RightSentinel{T,M}(; max_height = DEFAULT_MAX_HEIGHT, kws...) where {T,M} =
    Node{T,M}(Vector{T}(undef,0), max_height; flags = FLAG_IS_RIGHT_SENTINEL, kws...)

#===========================
Node external API
===========================#

function Base.string(node::Node)
    result = "$(node.vals), height = $(height(node))"
    "Node($result)"
end

Base.length(node::Node) = length(node.vals)

key(node::Node) = node.vals[1]

capacity(node::Node) = node.capacity
Base.isempty(node::Node) = (length(node) == 0)
isfull(node::Node) = is_sentinel(node) || (length(node) == capacity(node))

Base.in(val, node::Node) =
    searchsorted(node.vals, val) |>
    idx -> first(idx) ≤ last(idx)

@generated function Base.insert!(node::Node{T,M}, val) where {T,M}
    quote
        if length(node) ≥ capacity(node)
            "Node size exceeds capacity ($(capacity(node)))" |>
            ErrorException |>
            throw
        end

        idx = searchsorted(node.vals, val)

        $(
            if M == :Set
                quote
                    if first(idx) > last(idx)     # Value not found in node
                        insert!(node.vals, first(idx), val)
                    end
                end
            else
                quote
                    insert!(node.vals, first(idx), val)
                end
            end
        )
    end
end

Base.delete!(node::Node, val) =
    searchsorted(node.vals, val) |>
    idx -> if first(idx) ≤ last(idx)
        deleteat!(node.vals, first(idx))
    end

function split!(node::Node{T,M}; kws...) where {T,M}
    median_idx = div(length(node), 2)
    right_vals = splice!(node.vals, median_idx+1:length(node))
    node, Node{T,M}(right_vals; kws...)
end

"""
Insert a new node between a list of predecessor and successor nodes
"""
function interpolate_node!(predecessors, successors, node)
    for level = 1:height(node)
        link_nodes!(predecessors[level], node, level)
        link_nodes!(node, successors[level], level)
    end
end


#NOTE FILE TO BE DELETED

SRC_DIR = joinpath(Base.source_dir(), "..", "src")
AUX_DIR = joinpath(SRC_DIR, "auxiliary")
OUT_DIR = joinpath(Base.source_dir(), "..", "output")
mkpath(OUT_DIR)

#include(joinpath(SRC_DIR, "BridgeSDEInference.jl"))
#using Main.BridgeSDEInference
include(joinpath(SRC_DIR, "fitzHughNagumo.jl"))
include(joinpath(SRC_DIR, "conjugateUpdt.jl"))

include(joinpath(SRC_DIR, "types.jl"))
include(joinpath(SRC_DIR, "vern7.jl"))
include(joinpath(SRC_DIR, "tsit5.jl"))
include(joinpath(SRC_DIR, "rk4.jl"))
include(joinpath(SRC_DIR, "ralston3.jl"))
include(joinpath(SRC_DIR, "priors.jl"))
include(joinpath(SRC_DIR, "guid_prop_bridge.jl"))

include(joinpath(SRC_DIR, "bounded_diffusion_domain.jl"))
include(joinpath(SRC_DIR, "euler_maruyama_dom_restr.jl"))
include(joinpath(SRC_DIR, "lorenz_system.jl"))
include(joinpath(SRC_DIR, "lorenz_system_const_vola.jl"))

include(joinpath(SRC_DIR, "random_walk.jl"))
include(joinpath(SRC_DIR, "blocking_schedule.jl"))
include(joinpath(SRC_DIR, "starting_pt.jl"))
include(joinpath(SRC_DIR, "mcmc.jl"))
include(joinpath(SRC_DIR, "path_to_wiener.jl"))
include(joinpath(SRC_DIR, "metropolis_adjusted_langevin_kernel.jl"))

using Distributions # to define priors
using Random        # to seed the random number generator

param = :complexConjug
# Initial parameter guess.
θ₀ = [10.0, -8.0, 15.0, 0.0, 3.0]
# Target law
P˟ = FitzhughDiffusion(param, θ₀...)
# Auxiliary law
P̃ = FitzhughDiffusionAux(param, θ₀..., 0.0, ℝ(0.0), 2.0, ℝ(0.5))

L = @SMatrix [1. 0.]
Σdiagel = 10^(-4)
Σ = @SMatrix [Σdiagel]

using ForwardDiff
import ForwardDiff: Dual, Tag
tt = 0.0:0.01:2.0
Wnr = Wiener{Float64}()
W = Bridge.samplepath(tt, zero(Float64))
sample!(W, Wnr)
y = ℝ(0.0, 0.0)
CT = Tag{Val{:custom_tag}, Float64}
y_dual = ℝ(Dual{CT}(0.0,0.0), Dual{CT}(0.0,0.0))
X = SamplePath(tt, zeros(typeof(y_dual), length(tt)))
# compute log likelihood for a given wiener path and parameter vector
function ll(θ)
    for t in θ
        print(t, "\n\n")
    end
    P˟ = FitzhughDiffusion(param, θ...)
    P̃ = FitzhughDiffusionAux(param, θ..., 0.0, ℝ(Dual{CT}(0.0,0.0)),
                                          2.0, ℝ(Dual{CT}(0.5,0.0)))
    L = @SMatrix [Dual{CT}(1.0,0.0) Dual{CT}(0.0,0.0)]
    Σ = @SMatrix [Dual{CT}(10^(-4),0.0)]
    P = GuidPropBridge(Dual{CT,Float64,1}, tt, P˟, P̃, L, ℝ(Dual{CT}(0.5,0.0)), Σ;
                       changePt=NoChangePt(), solver=Vern7())
    solve!(Euler(), X, y_dual, W, P)
    loglik = llikelihood(LeftRule(), X, P)
    loglik
end

chunkSize = 1
result = DiffResults.GradientResult(θ₀)
cfg = ForwardDiff.GradientConfig(ll, θ₀, ForwardDiff.Chunk{chunkSize}(), CT())
@time ForwardDiff.gradient!(result, ll, θ₀, cfg, Val{false}()) # turn off tag checking..., this needs to be solved somehow in the future

result.value
result.derivs

Dual{CT}.(L, 0.0)

using Interpolations
using StaticArrays
itp = LinearInterpolation([0.0, 1.0, 3.0], [ℝ(1.0, 2.0), ℝ(3.0,-2.0), ℝ(2.0,0.0)])

itp = LinearInterpolation([0.0, 1.0], [ℝ(0.0), ℝ(0.0)], extrapolation_bc = Line())
itp(2.0)

using Plots
tt = collect(0.0:0.01:3.0)
xx₁ = [x[1] for x in itp(tt)]
xx₂ = [x[2] for x in itp(tt)]
plot(tt, xx₁)
plot!(tt, xx₂)

LinearInterpolation([0.0, 1.0, 3.0], [ℝ(1.0, 2.0), ℝ(3.0,-2.0), ℝ(2.0,0.0)])


module SpikeSynchrony

include("SPIKEdistance.jl")
include("vanRossum.jl")

end # module


__precompile__()

module IntervalTrees

export
    IntervalTree,
    IntervalMap,
    AbstractInterval,
    Interval,
    IntervalValue,
    depth,
    hasintersection,
    from,
    value

include("slice.jl")

"""
An `AbstractInterval{T}` must have a `first` and `last` function each returning
a value of type `T`, and `first(i) <= last(i)` must always be true.
"""
abstract AbstractInterval{T}

function Base.isless(u::AbstractInterval, v::AbstractInterval)
    return first(u) < first(v) || (first(u) == first(v) && last(u) < last(v))
end

"""
A basic interval.
"""
immutable Interval{T} <: AbstractInterval{T}
    first::T
    last::T
end

Base.first{T}(i::Interval{T}) = i.first
Base.last{T}(i::Interval{T}) = i.last

"""
An interval with some associated data.
"""
immutable IntervalValue{K, V} <: AbstractInterval{K}
    first::K
    last::K
    value::V
end

Base.first{K, V}(i::IntervalValue{K, V}) = i.first
Base.last{K, V}(i::IntervalValue{K, V}) = i.last
value{K, V}(i::IntervalValue{K, V}) = i.value

Base.print(io::IO, x::Interval) = print(io, "\n($(first(x)),$(last(x)))")
function Base.show(io::IO, x::Interval)
    show(io, typeof(x))
    print(io, x)
end

Base.print(io::IO, x::IntervalValue) = print(io, "\n($(first(x)),$(last(x))) => $(value(x))")
function Base.show(io::IO, x::IntervalValue)
    show(io, typeof(x))
    print(io, x)
end

# Each of these types is indexes by K, V, B, where
#   K : Interval type. Intervals are represented as (K, K) tuples.
#   V : Interval value type. This an interval type that may have associated
#        data. It must have `first`, `last`, and `isless` methods.
#   B : Integer giving the B-tree order.

abstract Node{K, V, B}

type InternalNode{K, V, B} <: Node{K, V, B}
    # Internal nodes are keyed by the minimum interval in the right subtree.  We
    # need internal node keys to be intervals themselves, since ordering by
    # start value alone only works if start values are unique. We don't
    # force that restriction, so we must break ties in order to split nodes that
    # are full of intervals with the same start value.
    keys::Slice{Interval{K}, B}

    # Maximum interval end-point in this subtree.
    maxend::K

    children::Slice{Node{K, V, B}, B}

    parent::Nullable{InternalNode{K, V, B}}

    # Sibling/cousin pointers.
    left::Nullable{InternalNode{K, V, B}}
    right::Nullable{InternalNode{K, V, B}}

    function InternalNode()
        t = new(Slice{Interval{K}, B}(), zero(K), Slice{Node{K, V, B}, B}(),
                Nullable{InternalNode{K, V, B}}(),
                Nullable{InternalNode{K, V, B}}(),
                Nullable{InternalNode{K, V, B}}())
        return t
    end
end

function minstart(t::InternalNode)
    return t.keys[1].first
end

type LeafNode{K, V, B} <: Node{K, V, B}
    entries::Slice{V, B}

    # maximum interval end-point in this leaf node.
    maxend::K

    parent::Nullable{InternalNode{K, V, B}}

    # Sibling/cousin pointers.
    left::Nullable{LeafNode{K, V, B}}
    right::Nullable{LeafNode{K, V, B}}

    function LeafNode()
        return new(
            Slice{V, B}(),
            zero(K),
            Nullable{InternalNode{K, V, B}}(),
            Nullable{LeafNode{K, V, B}}(),
            Nullable{LeafNode{K, V, B}}())
    end
end

function minstart(t::LeafNode)
    return first(t.entries[1])
end

type IntervalBTree{K, V, B}
    root::Node{K, V, B}
    n::Int # Number of entries

    function IntervalBTree()
        return new(LeafNode{K, V, B}(), 0)
    end

    # Construct an interval tree from a sorted array of intervals.
    #
    # This is generally much more efficient than constructing the same tree by
    # inserting intervals one by one.
    #
    # Args:
    #   entries: Interval entry values in sorted order.
    #
    function IntervalBTree(entries::AbstractVector{V})
        if !issorted(entries)
            error("Intervals must be sorted to construct an IntervalTree")
        end

        # Here's the plan: figure out how many leaf nodes we need, allocate
        # them all up front. Copy in the keys and values, then work up towards
        # the root.

        n = length(entries)

        if n == 0
            return new(LeafNode{K, V, B}(), 0)
        end

        numleaves = cld(n, B - 2)
        leaves = [LeafNode{K, V, B}() for _ in 1:numleaves]

        maxends = Array(K, numleaves)
        minkeys = Array(Interval{K}, numleaves)

        # divy up the keys and values among the leaves
        keys_per_leaf = cld(n, numleaves)
        for i in 1:numleaves
            u = (i - 1) * keys_per_leaf + 1
            v = min(n, i * keys_per_leaf)
            minkeys[i] = Interval{K}(first(entries[u]), last(entries[u]))
            maxends[i] = last(entries[u])
            for j in u:v
                push!(leaves[i].entries, entries[j])
                maxends[i] = max(maxends[i], last(entries[j]))
            end
            leaves[i].maxend = maxends[i]
        end

        children = leaves
        while length(children) > 1
            # sibling pointers
            for i in 1:length(children)-1
                children[i].right = children[i+1]
                children[i+1].left = children[i]
            end

            # make parents
            numparents = cld(length(children), B - 2)
            parents = [InternalNode{K, V, B}() for _ in 1:numparents]

            # divy up children among parents
            children_per_parent = cld(length(children), numparents)
            for i in 1:numparents
                u = (i - 1) * keys_per_leaf + 1
                v = min(length(children), i * keys_per_leaf)
                maxend = maxends[u]
                for j in u:v
                    push!(parents[i].children, children[j])
                    children[j].parent = parents[i]
                    maxend = max(maxend, maxends[j])
                    if j > u
                        push!(parents[i].keys, minkeys[j])
                    end
                end
                minkeys[i] = minkeys[u]
                parents[i].maxend = maxend
                maxends[i] = maxend
            end

            children = parents
        end
        @assert length(children) == 1

        return new(children[1], n)
    end
end

# Default B-tree order
typealias IntervalTree{K, V} IntervalBTree{K, V, 64}

# Show

function Base.show(io::IO, it::IntervalTree)
    show(io, typeof(it))
    n = length(it)
    if length(it) > 6
        # Hacky random access ...
        for (i, x) in enumerate(it)
            i < 4 && print(x)
        end
        print("\n\u22EE") # Vertical ellipsis
        for (i, x) in enumerate(it)
            i > (n-3) && print(x)
        end
    else
        for x in it
            print(x)
        end
    end
end

# Length
# ------

function Base.length(t::IntervalBTree)
    return t.n
end

function depth(t::IntervalBTree)
    k = 1
    node = t.root
    while !isa(node, LeafNode)
        k += 1
        node = node.children[1]
    end
    return k
end

function Base.length(t::InternalNode)
    return length(t.children)
end

function Base.length(t::LeafNode)
    return length(t.entries)
end

function Base.isempty(t::IntervalBTree)
    return t.n == 0
end

function Base.isempty(t::InternalNode)
    return isempty(t.children)
end

function Base.isempty(t::LeafNode)
    return isempty(t.entries)
end

# Iterating
# ---------

immutable IntervalBTreeIteratorState{K, V, B}
    leaf::Nullable{LeafNode{K, V, B}}
    i::Int
end

function Base.start{K, V, B}(t::IntervalBTree{K, V, B})
    # traverse to the first leaf node
    node = t.root
    while !isa(node, LeafNode{K, V, B})
        node = node.children[1]
    end
    return IntervalBTreeIteratorState(Nullable(node), 1)
end

function Base.next{K, V, B}(t::IntervalBTree{K, V, B},
                            state::IntervalBTreeIteratorState{K, V, B})
    leaf = get(state.leaf)
    entry = leaf.entries[state.i]
    if state.i < length(leaf)
        state = IntervalBTreeIteratorState{K, V, B}(leaf, state.i + 1)
    else
        state = IntervalBTreeIteratorState{K, V, B}(leaf.right, 1)
    end
    return entry, state
end

function Base.done{K, V, B}(t::IntervalBTree{K, V, B},
                            state::IntervalBTreeIteratorState{K, V, B})
    return isnull(state.leaf) || isempty(get(state.leaf))
end

# Iterate from a given starting from the first interval that intersects a given
# point.

immutable IntervalFromIterator{K, V, B}
    t::IntervalBTree{K, V, B}
    p::K
end

function from{K, V, B}(t::IntervalBTree{K, V, B}, p)
    return IntervalFromIterator{K, V, B}(t, convert(K, p))
end

function Base.start{K, V, B}(it::IntervalFromIterator{K, V, B})
    node, i = firstfrom(it.t, it.p)
    if i == 0
        return IntervalBTreeIteratorState{K, V, B}(Nullable{LeafNode{K, V, B}}(), i)
    else
        return IntervalBTreeIteratorState{K, V, B}(node, i)
    end
end

function Base.next{K, V, B}(it::IntervalFromIterator{K, V, B},
                            state::IntervalBTreeIteratorState{K, V, B})
    return next(it.t, state)
end

function Base.done{K, V, B}(it::IntervalFromIterator{K, V, B},
                            state::IntervalBTreeIteratorState{K, V, B})
    return done(it.t, state)
end

# Inserting
# ---------

# Split a leaf into two, returning (leftnode, rightnode)
function split!{K, V, B}(left::LeafNode{K, V, B})
    right = LeafNode{K, V, B}()
    right.right = left.right
    right.left = left
    left.right = right
    if !isnull(right.right)
        get(right.right).left = right
    end
    right.parent = left.parent

    m = length(left)
    resize!(right.entries, m - div(m, 2))
    copy!(right.entries, 1, left.entries, div(m, 2) + 1, length(right.entries))
    resize!(left.entries, div(m, 2))

    left.maxend = last(left.entries[1])
    for entry in left.entries
        left.maxend = max(left.maxend, last(entry))
    end

    right.maxend = last(right.entries[1])
    for entry in right.entries
        right.maxend = max(right.maxend, last(entry))
    end

    return left, right
end

# Split an internal node in two, returning (leftnode, rightnode)
function split!{K, V, B}(left::InternalNode{K, V, B})
    right = InternalNode{K, V, B}()
    right.right = left.right
    right.left = left
    left.right = right
    if !isnull(right.right)
        get(right.right).left = right
    end
    right.parent = left.parent

    m = length(left)

    resize!(right.children, m - div(m, 2))
    resize!(right.keys, m - div(m, 2) - 1)

    copy!(right.children, 1, left.children, div(m, 2)+1, length(right.children))
    copy!(right.keys, 1, left.keys, div(m, 2)+1, length(right.keys))

    resize!(left.children, div(m, 2))
    resize!(left.keys, div(m, 2) - 1)

    left.maxend = left.children[1].maxend
    for child in left.children
        left.maxend = max(left.maxend, child.maxend)
    end

    right.maxend = right.children[1].maxend
    for child in right.children
        right.maxend = max(right.maxend, child.maxend)
    end

    for child in right.children
        child.parent = right
    end

    return left, right
end

# Find the maximum interval end point is a subtree
function nodemaxend{K, V, B}(t::InternalNode{K, V, B})
    maxend = zero(K)
    for i in 1:length(t.children)
        maxend = max(maxend, t.children[i].maxend)
    end
    return maxend
end

function nodemaxend{K, V, B}(t::LeafNode{K, V, B})
    maxend = zero(K)
    for i in 1:length(t.entries)
        maxend = max(maxend, last(t.entries[i]))
    end
    return maxend
end

function sameparent{K, V, B}(u::Node{K, V, B}, v::Node{K, V, B})
    return (isnull(u.parent) && isnull(v.parent)) ||
           (!isnull(u.parent) && !isnull(v.parent) && get(u.parent) == get(v.parent))
end

# Find the first leaf node in the tree
function firstleaf(t::IntervalBTree)
    return firstleaf(t.root)
end

function firstleaf(t::InternalNode)
    return firstleaf(t.children[1])
end

function firstleaf(t::LeafNode)
    return t
end

# TODO: Having to construct Nullables here is killng us.
# How can we avoid that?
#
function nextleafkey{K, V, B}(t::LeafNode{K, V, B}, i::Int)
    if i < length(t)
        return (Nullable{LeafNode{K, V, B}}(t), i + 1)
    else
        return (t.right, 1)
    end
end

# This version of nextleafkey is used in nextintersection at aggressively avoid
# allocation, and dereferecing nullables more than once.
macro nextleafkey(leaf, nleaf, i)
    quote
        if $(esc(i)) < length($(esc(leaf)))
            $(esc(i)) += 1
        else
            $(esc(i)) = 1
            $(esc(nleaf)) = $(esc(leaf)).right
        end
    end
end

# Find the minimum interval in a subtree
function minkey(t::InternalNode)
    return minkey(t.children[1])
end

function minkey{K, V, B}(t::LeafNode{K, V, B})
    return Interval{K}(first(t.entries[1]), last(t.entries[1]))
end

function Base.push!{K, V, B}(t::IntervalBTree{K, V, B}, entry::V,
                             unique_key::Bool=false, update::Bool=true)
    return _push!(t, t.root, entry, unique_key, update)
end

function _push!{K, V, B}(t::IntervalBTree{K, V, B},
                         node::InternalNode{K, V, B},
                         entry::V, unique_key::Bool, update::Bool)
    i = findidx(node, entry) # key index
    j = i <= length(node) - 1 && entry >= node.keys[i] ? i + 1 : i # child index
    return _push!(t, node.children[j], entry, unique_key, update)
end

function _push!{K, V, B}(t::IntervalBTree{K, V, B},
                         node::LeafNode{K, V, B}, entry::V,
                         unique_key::Bool, update::Bool)
    i = max(1, findidx(node, entry))

    # key exists and we are replacing it
    if i <= length(node) && unique_key &&
        first(node.entries[i]) == first(entry) && last(node.entries[i]) == last(entry)
        if update
            node.entries[i] = entry
            return entry
        else
            return node.entries[i]
        end
    end

    #unsafe_insert!(node.entries, i, entry)
    insert!(node.entries, i, entry)
    if length(node) == 1 || last(entry) > node.maxend
        node.maxend = last(entry)
    end
    t.n += 1

    # split when full
    if length(node) == B
        maxend = node.maxend
        leftleaf, rightleaf = split!(node)
        median = Interval{K}(first(rightleaf.entries[1]),
                             last(rightleaf.entries[1]))

        # travel back up the tree setting maxend values and splitting as needed
        parent = node.parent
        leftnode = leftleaf
        rightnode = rightleaf
        hassplit = true
        while !isnull(parent)
            p = get(parent)
            if hassplit
                i = findidx(p, entry) # key index
                j = i <= length(p) - 1 && entry >= p.keys[i] ? i + 1 : i # child index
                insert!(p.children, j + 1, rightnode)
                p.maxend = max(p.maxend, maxend)
                insert!(p.keys, j, median)

                # split when full
                if length(p) == B
                    leftnode, rightnode = split!(p)
                    median = minkey(rightnode)
                    maxend = max(leftnode.maxend, rightnode.maxend)
                else
                    hassplit = false
                end
            else
                p.maxend = max(p.maxend, last(entry))
            end
            parent = p.parent
        end

        if hassplit
            t.root = InternalNode{K, V, B}()
            push!(t.root.keys, median)
            push!(t.root.children, leftnode)
            push!(t.root.children, rightnode)
            t.root.maxend = maxend
            leftnode.parent = t.root
            rightnode.parent = t.root
        end
    else
        # travel back up the tree setting maxend values
        parent = node.parent
        while !isnull(parent)
            p = get(parent)
            p.maxend = max(p.maxend, last(entry))
            parent = p.parent
        end
    end

    return entry
end

# Deleting
# --------

function deletefirst!{K, V, B}(t::IntervalBTree{K, V, B}, first::K, last::K)
    return deletefirst!(t, Interval{K}(first, last))
end

function deletefirst!{K, V, B}(t::IntervalBTree{K, V, B}, key::Tuple{K, K})
    return deletefirst!(t, Interval{K}(key[1], key[2]))
end

function deletefirst!{K, V, B}(t::IntervalBTree{K, V, B}, key::Interval{K})
    result = _deletefirst!(t.root, key)
    if result.keyfound
        t.n -= 1
    end

    # if the root has only one child, promote the child
    if isa(t.root, InternalNode) && length(t.root) == 1
        t.root = t.root.children[1]
        t.root.parent = Nullable{InternalNode{K, V, B}}()
    end

    return t
end

# Indicate what steps are need to account for an updated child.
immutable KeyFate
    value::UInt8
end

const KEYFATE_NONE         = KeyFate(0) # no changes
const KEYFATE_DELETE       = KeyFate(1) # delete the child node
const KEYFATE_DELETE_RIGHT = KeyFate(2) # delete the child's right sibling
const KEYFATE_UPDATE_LEFT  = KeyFate(3) # update the key separating the node
                                        # from its left sibling
const KEYFATE_UPDATE_RIGHT = KeyFate(4) # update the key separating the node
                                        # from its right sibling

# Information returned from a call to _delete
immutable DeleteResult
    keyfound::Bool # true if the key was found

    # Indicate what steps are need to account for an updated child. One of:
    fate::KeyFate
end

# Delete key from the subtree is present. Return a DeleteResult object.
function _deletefirst!{K, V, B}(t::InternalNode{K, V, B}, key::Interval{K})
    i = findidx(t, key)
    j = i <= length(t) - 1 && key >= t.keys[i] ? i + 1 : i # child index

    ans = _deletefirst!(t.children[j], key)

    if !ans.keyfound
        return DeleteResult(false, KEYFATE_NONE)
    end

    if ans.fate == KEYFATE_UPDATE_LEFT
        t.maxend = nodemaxend(t)
        if j > 1
            t.keys[j - 1] = minkey(t.children[j]);
            return DeleteResult(true, KEYFATE_NONE)
        else
            return DeleteResult(true, KEYFATE_UPDATE_LEFT)
        end
    elseif ans.fate == KEYFATE_UPDATE_RIGHT
        t.maxend = nodemaxend(t)
        if j < length(t)
            t.keys[j] = minkey(t.children[j+1])
        end

        if j > 1
            t.keys[j - 1] = minkey(t.children[j]);
            return DeleteResult(true, KEYFATE_NONE)
        else
            return DeleteResult(true, KEYFATE_UPDATE_LEFT)
        end
    elseif ans.fate == KEYFATE_DELETE || ans.fate == KEYFATE_DELETE_RIGHT
        deleteidx = ans.fate == KEYFATE_DELETE ? j : j + 1

        # deleteidx == 1 would only happen if we had only one child, but if
        # that were the case it would have been promoted.
        @assert deleteidx > 1

        splice!(t.keys, deleteidx - 1)
        splice!(t.children, deleteidx)
        t.maxend = nodemaxend(t)

        # not underfull
        minsize = div(B, 2)
        if length(t) >= minsize
            return DeleteResult(true, deleteidx == 2 ? KEYFATE_UPDATE_LEFT : KEYFATE_NONE)
        end

        # borrow right
        if !isnull(t.right)
            right = get(t.right)
            if sameparent(right, t) && length(right) > minsize
                splice!(right.keys, 1)
                push!(t.children, splice!(right.children, 1))
                push!(t.keys, minkey(t.children[end]))
                t.maxend = max(t.maxend, nodemaxend(t.children[end]))
                right.maxend = nodemaxend(right)
                t.children[end].parent = t

                return DeleteResult(true, KEYFATE_UPDATE_RIGHT)
            end
        end

        # borrow left
        if !isnull(t.left)
            left = get(t.left)
            if sameparent(left, t) && length(left) > minsize
                insert!(t.children, 1, pop!(left.children))
                pop!(left.keys)
                insert!(t.keys, 1, minkey(t.children[2]))
                t.maxend = max(t.maxend, nodemaxend(t.children[1]))
                left.maxend = nodemaxend(left)
                t.children[1].parent = t

                return DeleteResult(true, KEYFATE_UPDATE_LEFT)
            end
        end

        # merge with left
        if !isnull(t.left)
            left = get(t.left)
            if sameparent(left, t)
                merge!(left, t)
                left.right = t.right
                if !isnull(t.right)
                    get(t.right).left = left
                end

                return DeleteResult(true, KEYFATE_DELETE)
            end
        end

        # merge with right
        if !isnull(t.right)
            right = get(t.right)
            if sameparent(right, t)
                merge!(t, right)
                if !isnull(right.right)
                    get(right.right).left = t
                end
                t.right = right.right

                return DeleteResult(true, KEYFATE_DELETE_RIGHT)
            end
        end

        # Allow the root node to be underfull
        return DeleteResult(true, KEYFATE_NONE)
    else
        t.maxend = nodemaxend(t)
        return DeleteResult(true, KEYFATE_NONE)
    end
end

function _deletefirst!{K, V, B}(t::LeafNode{K, V, B}, key::Interval{K})
    i = findidx(t, key)

    # do nothing if the key isn't present
    if i < 1 || i > length(t) ||
        first(t.entries[i]) != key.first || last(t.entries[i]) != key.last
        return DeleteResult(false, KEYFATE_NONE)
    end

    splice!(t.entries, i)

    # This is the root node. Allow it to be empty.
    if isempty(t)
        return DeleteResult(true, KEYFATE_NONE)
    end

    minsize = div(B, 2)

    t.maxend = nodemaxend(t)

    # not underfull
    if length(t) >= minsize
        return DeleteResult(true, i == 1 ? KEYFATE_UPDATE_LEFT : KEYFATE_NONE)
    end

    # borrow right
    if !isnull(t.right)
        right = get(t.right)
        if sameparent(right, t) && length(right) > minsize
            push!(t.entries, splice!(right.entries, 1))
            t.maxend = max(t.maxend, last(t.entries[end]))
            right.maxend = nodemaxend(right)

            return DeleteResult(true, KEYFATE_UPDATE_RIGHT)
        end
    end

    # borrow left
    if !isnull(t.left)
        left = get(t.left)
        if sameparent(left, t) && length(left) > minsize
            insert!(t.entries, 1, pop!(left.entries))
            t.maxend = max(t.maxend, last(t.entries[1]))
            left.maxend = nodemaxend(left)

            return DeleteResult(true, KEYFATE_UPDATE_LEFT)
        end
    end

    # merge with left
    if !isnull(t.left)
        left = get(t.left)
        if sameparent(left, t)
            merge!(left, t)
            left.right = t.right
            if !isnull(t.right)
                get(t.right).left = left
            end
            return DeleteResult(true, KEYFATE_DELETE)
        end
    end

    # merge with right
    if !isnull(t.right)
        right = get(t.right)
        if sameparent(right, t)
            merge!(t, right)
            if !isnull(right.right)
                get(right.right).left = t
            end
            t.right = right.right
            return DeleteResult(true, KEYFATE_DELETE_RIGHT)
        end
    end

    # This must be the root node. Allow it to be underfull.
    return DeleteResult(true, KEYFATE_NONE)
end

# Join two leaf nodes into one.
function merge!{K, V, B}(left::LeafNode{K, V, B}, right::LeafNode{K, V, B})
    leftlen, rightlen = length(left), length(right)
    @assert leftlen + rightlen <= B
    resize!(left.entries, leftlen + rightlen)
    copy!(left.entries, leftlen+1, right.entries, 1, length(right.entries))
    left.maxend = nodemaxend(left)
    return
end

# Join two internal nodes into one
function merge!{K, V, B}(left::InternalNode{K, V, B}, right::InternalNode{K, V, B})
    leftlen, rightlen = length(left), length(right)
    @assert length(left) + length(right) <= B
    resize!(left.keys, leftlen + rightlen - 1)
    resize!(left.children, leftlen + rightlen)
    copy!(left.children, leftlen+1, right.children, 1, length(right.children))
    for i in leftlen+1:length(left.children)
        left.children[i].parent = left
    end
    left.keys[leftlen] = minkey(left.children[leftlen+1])
    copy!(left.keys, leftlen+1, right.keys, 1, length(right.keys))
    left.maxend = nodemaxend(left)
    return
end

# Searching
# ---------

function findidx{K, V, B}(t::LeafNode{K, V, B}, key::AbstractInterval{K})
    if isempty(t.entries)
        return 0
    end
    return searchsortedfirst(t.entries, key)
end

function findidx{K, V, B}(t::InternalNode{K, V, B}, key::AbstractInterval{K})
    if isempty(t.keys)
        return 0
    end
    i = searchsortedfirst(t.keys, key)
    return min(i, length(t.keys))
end

function Base.haskey{K, V, B}(t::LeafNode{K, V, B}, key::AbstractInterval{K})
    i = findidx(t, key)
    return 1 <= i <= length(t) && first(t.entries[i]) == key.first &&
           last(t.entries[i]) == key.last
end

function Base.haskey{K, V, B}(t::InternalNode{K, V, B}, key::AbstractInterval{K})
    i = findidx(t, key)
    if i <= length(t) - 1 && key >= t.keys[i]
        return haskey(t.children[i+1], key)
    else
        return haskey(t.children[i], key)
    end
end

function Base.haskey{K, V, B}(t::IntervalBTree{K, V, B}, key0::Tuple{Any, Any})
    key = Interval{K}(key0[1], key0[2])
    return haskey(t.root, key)
end

# Intersection
# ------------

# There are two ways we can tackle intersection: traversal (search the tree for
# intervals that intersect) and iteration (iterate through the tree's intervals
# in sorted order).
#
# Traversal is roughly O(log(n)) while iteration is roughly O(n). Obviously
# if we are doing a single intersection, we would choose traversal, but when
# intersecting two trees it becomes more interesting.
#
# If the second tree is of size m, then intersection by traversal is
# O(min(m * log(n), n * log(m))) while interesction by iteration is O(m + n).
# When m and n are both large, it becomes more efficient to intersect by
# iteration, but deciding when requires some heuristics.
#

# Return true iff two key1 and key2 intersect.
@inline function intersects{K}(key1::AbstractInterval{K}, key2::AbstractInterval{K})
    return first(key1) <= last(key2) && first(key2) <= last(key1)
end

# Return true if the tree has an intersection with the given position
function hasintersection{K, V, B}(t::IntervalBTree{K, V, B}, query)
    return hasintersection(t.root, convert(K, query))
end

function hasintersection{K, V, B}(t::InternalNode{K, V, B}, query::K)
    if isempty(t) || t.maxend < query
        return false
    end

    for (i, child) in enumerate(t.children)
        if child.maxend >= query && (i == 1 || t.keys[i-1].first <= query)
            if hasintersection(child, query)
                return true
            end
        elseif minstart(child) > query
            break
        end
    end

    return false
end

function hasintersection{K, V, B}(t::LeafNode{K, V, B}, query::K)
    if isempty(t) || t.maxend < query
        return false
    end

    for i in 1:length(t)
        if first(t.entries[i]) <= query <= last(t.entries[i])
            return true
        elseif query < first(t.entries[i])
            break
        end
    end

    return false
end

# Represent an intersection in an IntervalTree by pointing to a LeafNode
# and an index within that leaf node. No intersection is represented with
# index == 0.
type Intersection{K, V, B}
    index::Int
    node::LeafNode{K, V, B}

    Intersection(index, node) = new(index, node)
    Intersection() = new(0)
end

# Find the first interval in the tree that intersects the query and return
# as a (leafnode, index) pair, indicating that leafnode.keys[index] intersects.
# If no intersection is found, index is 0 and leafnode is the last node
# searched.
function firstintersection!{K, V, B}(t::IntervalBTree{K, V, B},
                                     query::AbstractInterval{K},
                                     out::Intersection{K, V, B})
    return firstintersection!(t.root, query, out)
end

function firstintersection!{K, V, B}(t::InternalNode{K, V, B},
                                     query::AbstractInterval{K},
                                     out::Intersection{K, V, B})
    if isempty(t) || t.maxend < first(query)
        out.index = 0
        return
    end

    for (i, child) in enumerate(t.children)
        if child.maxend >= first(query) && (i == 1 || t.keys[i-1].first <= last(query))
            firstintersection!(child, query, out)
            if out.index > 0
                return
            end
        elseif minstart(child) > last(query)
            break
        end
    end

    out.index = 0
    return
end

function firstintersection!{K, V, B}(t::LeafNode{K, V, B},
                                     query::AbstractInterval{K},
                                     out::Intersection{K, V, B})
    if isempty(t) || t.maxend < first(query)
        out.index = 0
        return
    end

    for i in 1:length(t)
        if intersects(t.entries[i], query)
            out.index = i
            out.node = t
            return
        elseif query.last < first(t.entries[i])
            break
        end
    end

    out.index = 0
    return
end

# Find the first key with an end point >= query
function firstfrom{K, V, B}(t::IntervalBTree{K, V, B}, query::K)
    return firstfrom(t.root, query)
end

function firstfrom{K, V, B}(t::InternalNode{K, V, B}, query::K)
    if isempty(t)
        return (t, 0)
    end

    for (i, child) in enumerate(t.children)
        if child.maxend >= query
            return firstfrom(child, query)
        end
    end

    return (t, 0)
end

function firstfrom{K, V, B}(t::LeafNode{K, V, B}, query::K)
    if isempty(t)
        return (t, 0)
    end

    for i in 1:length(t)
        if last(t.entries[i]) >= query
            return (t, i)
        end
    end

    return (t, 0)
end

# If query intersects node.keys[i], return the next intersecting key as
# a (leafnode, index) pair.
function nextintersection!{K, V, B}(t::LeafNode{K, V, B}, i::Integer,
                                    query::AbstractInterval{K},
                                    out::Intersection{K, V, B})
    j = i + 1
    while true
        while j <= length(t)
            if intersects(unsafe_getindex(t.entries, j), query)
                out.index = j
                out.node = t
                return
            end
            j += 1
        end
        j = 1
        if isnull(t.right)
            break
        end
        t = get(t.right)
        if minstart(t) > last(query)
            break
        end
    end

    out.index = 0
    return
end

type IntervalIntersectionIterator{K, V, B}
    intersection::Intersection{K, V, B}
    t::IntervalBTree{K, V, B}
    query::AbstractInterval{K}

    function IntervalIntersectionIterator(intersection, t, query)
        return new(intersection, t, query)
    end

    function IntervalIntersectionIterator()
        return new(Intersection{K, V, B}())
    end
end

# Intersect an interval tree t with a single interval, returning an iterator
# over the intersecting (key, value) pairs in t.
function Base.intersect{K, V, B}(t::IntervalBTree{K, V, B}, query0::Tuple{Any, Any})
    query = Interval{K}(query0[1], query0[2])
    return intersect(t, query)
end

function Base.intersect{K, V, B}(t::IntervalBTree{K, V, B}, first::K, last::K)
    query = Interval{K}(first, last)
    return intersect(t, query)
end

function Base.intersect{K, V, B}(t::IntervalBTree{K, V, B}, query::AbstractInterval{K})
    return IntervalIntersectionIterator{K, V, B}(Intersection{K, V, B}(), t, query)
end

function Base.start{K, V, B}(it::IntervalIntersectionIterator{K, V, B})
    return firstintersection!(it.t, it.query, it.intersection)
end

function Base.next{K, V, B}(it::IntervalIntersectionIterator{K, V, B}, ::Void)
    intersection = it.intersection
    entry = intersection.node.entries[intersection.index]
    nextintersection!(intersection.node, intersection.index,
                      it.query, intersection)
    return entry, nothing
end

function Base.done{K, V, B}(it::IntervalIntersectionIterator{K, V, B}, ::Void)
    return it.intersection.index == 0
end

type IntersectionIterator{K, V1, B1, V2, B2}
    t1::IntervalBTree{K, V1, B1}
    t2::IntervalBTree{K, V2, B2}

    # if true, use successive intersection, if false iterative.
    successive::Bool

    # true if done
    isdone::Bool

    # successive intersection state
    intersection::Intersection{K, V2, B2}
    t1_state::IntervalBTreeIteratorState{K, V1, B1}
    t1_value::V1

    # iterative intersection state
    u::Nullable{LeafNode{K, V1, B1}}
    v::Nullable{LeafNode{K, V2, B2}}
    w::Nullable{LeafNode{K, V2, B2}}
    i::Int
    j::Int
    k::Int

    function IntersectionIterator(t1, t2, successive::Bool)
        return new(t1, t2, successive)
    end
end

function Base.start{K, V1, B1, V2, B2}(it::IntersectionIterator{K, V1, B1, V2, B2})
    it.isdone = true
    if it.successive
        # Successive Query Intersection: intersect by repeatedly performing
        # queries against on tree ot the other.
        it.intersection = Intersection{K, V2, B2}()
        t1_state = start(it.t1)
        while !done(it.t1, t1_state)
            t1_value, t1_state = next(it.t1, t1_state)
            firstintersection!(it.t2, t1_value, it.intersection)
            if it.intersection.index != 0
                it.t1_state = t1_state
                it.t1_value = t1_value
                it.isdone = false
                return nothing
            end
        end
    else
        # Iterative Intersection: Intersect by iterating through the two
        # collections in unison.
        #
        # Notation notes. We proceed by finding all the the intervals
        # in t2 that intersect u.keys[i], before proceeding to the next
        # key in t1.
        #
        # We use v to keep track of the first leaf node in t2 that might contain an
        # intersecting interval, # and with (w, j) the current intersecting entry in
        # t2. We need to hang onto v since we will need to jump back when the
        # entry in t1 gets incremented.
        #
        # The thing to keep in mind is that this is just like the merge operation in
        # mergesort, except that some backtracking is needed since intersection
        # isn't as simple as ordering.

        it.u = isempty(it.t1) ? Nullable{LeafNode{K, V1, B1}}() : firstleaf(it.t1)
        it.v = isempty(it.t2) ? Nullable{LeafNode{K, V2, B2}}() : firstleaf(it.t2)
        it.w = it.v
        it.i, it.j, it.k = 1, 1, 1
        nextintersection!(it)
    end
    return nothing
end

function nextintersection!{K, V1, B1, V2, B2}(it::IntersectionIterator{K, V1, B1, V2, B2})
    u, v, w, i, j, k = it.u, it.v, it.w, it.i, it.j, it.k
    it.isdone = true
    while !isnull(u) && !isnull(v)
        unode = get(u)
        vnode = get(v)
        if isnull(w)
            @nextleafkey(unode, u, i)
            w, k = v, j
            continue
        end
        wnode = get(w)

        uentry = unsafe_getindex(unode.entries, i)
        ventry = unsafe_getindex(vnode.entries, j)
        wentry = unsafe_getindex(wnode.entries, k)

        if last(uentry) < first(ventry)
            @nextleafkey(unode, u, i)
        elseif last(ventry) < first(uentry)
            @nextleafkey(vnode, v, j)
        elseif first(wentry) <= last(uentry)
            if intersects(uentry, wentry)
                it.isdone = false
                break
            end
            @nextleafkey(wnode, w, k)
        else
            @nextleafkey(unode, u, i)
            w, k = v, j
        end
    end

    it.u = u
    it.v = v
    it.w = w
    it.i = i
    it.j = j
    it.k = k
    return
end

function Base.next{K, V1, B1, V2, B2}(it::IntersectionIterator{K, V1, B1, V2, B2}, state)
    if it.successive
        intersection = it.intersection
        entry = unsafe_getindex(intersection.node.entries, intersection.index)
        value = (it.t1_value, entry)
        nextintersection!(intersection.node, intersection.index, it.t1_value,
                          intersection)

        t1 = it.t1
        t2 = it.t2
        t1_state = it.t1_state
        t1_value = it.t1_value
        while intersection.index == 0 && !done(t1, t1_state)
            t1_value, t1_state = next(t1, t1_state)
            firstintersection!(t2, t1_value, intersection)
        end
        it.isdone = intersection.index == 0
        it.t1_state = t1_state
        it.t1_value = t1_value
    else
        u = get(it.u)
        w = get(it.w)
        value = (unsafe_getindex(u.entries, it.i),
                 unsafe_getindex(w.entries, it.k))
        @nextleafkey(w, it.w, it.k)
        nextintersection!(it)
    end
    return (value, nothing)
end

function Base.done{K, V1, B1, V2, B2}(it::IntersectionIterator{K, V1, B1, V2, B2}, state)
    return it.isdone
end

# Intersect two interval trees, returning an iterator yielding values of the
# form:
#   ((key1, value1), (key2, value2))
#
# Where key1 is from the first tree and key2 from the second, and they
# intersect.
function Base.intersect{K, V1, B1, V2, B2}(t1::IntervalBTree{K, V1, B1},
                                           t2::IntervalBTree{K, V2, B2};
                                           method=:auto)
    # We decide heuristically which intersection algorithm to use.
    n = length(t1)
    m = length(t2)

    iterative_cost  = n + m
    successive_cost = n * log(1 + m)

    if method == :auto
        iterative_cost  = n + m
        successive_cost = n * log(1 + m)
        if iterative_cost < successive_cost
            return IntersectionIterator{K, V1, B1, V2, B2}(t1, t2, false)
        else
            return IntersectionIterator{K, V1, B1, V2, B2}(t1, t2, true)
        end
    elseif method == :successive
        return IntersectionIterator{K, V1, B1, V2, B2}(t1, t2, true)
    elseif method == :iterative
        return IntersectionIterator{K, V1, B1, V2, B2}(t1, t2, false)
    else
        error("No such intersection method: $method")
    end
end

# Diagnostics
# -----------

# Dumb tree printing, useful only for debuging.
#function showtree(io::IO, t::IntervalBTree)
    #showtree(io, t.root, 0)
#end

#function showtree(t::IntervalBTree)
    #showtree(STDOUT, t)
#end

#function showtree(io::IO, t::InternalNode, indent::Int)
    #for (i, child) in enumerate(t.children)
        #showtree(io, child, indent+1)
        #if i <= length(t.keys)
            #print(io, repeat("  ", indent), t.keys[i], ":\n")
        #end
    #end
#end

#function showtree(io::IO, t::LeafNode, indent::Int)
    #for (k, v) in zip(t.keys, t.values)
        #print(io, repeat("  ", indent), k, ": ", v, "\n")
    #end
#end

include("map.jl")

if isdefined(Base, :iteratorsize)
    Base.iteratorsize(::IntervalFromIterator)         = Base.SizeUnknown()
    Base.iteratorsize(::IntervalIntersectionIterator) = Base.SizeUnknown()
    Base.iteratorsize(::IntersectionIterator)         = Base.SizeUnknown()
    Base.iteratorsize(::IntervalKeyIterator)          = Base.SizeUnknown()
    Base.iteratorsize(::IntervalValueIterator)        = Base.SizeUnknown()
end

end  # module IntervalTrees