"""

  DepthwiseConvAllowNinChangeStrategy(newoutputsmax::Integer, multipliersmax::Integer, base, [fallback])
  DepthwiseConvAllowNinChangeStrategy(allowed_new_outgroups::AbstractVector{<:Integer}, allowed_multipliers::AbstractVector{<:Integer}, base, [fallback])

`DecoratingJuMPΔSizeStrategy` which allows both nin and nout of `DepthwiseConv` layers to change independently.

Might cause optimization to take very long time so use with care! Use [`DepthwiseConvSimpleΔSizeStrategy`](@ref)
if `DepthwiseConvAllowNinChangeStrategy` takes too long.

The elements of `allowed_new_outgroups` determine how many extra elements in the output dimension of the weight
shall be tried for each existing output element. For example, for a `DepthwiseConv((k1,k2), nin=>nout))` there
are `nout / nin` elements in the output dimension. With `allowed_new_outgroups = 0:3` it is allowed to insert
0, 1, 2 or 3 new elements in the output dimension between each already existing element (so with `nout / nin`
elements the maximum increase is `3 * nout / nin`).

The elements of `allowed_multipliers` determine the total number of allowed output elements, i.e the allowed
ratios of `nout / nin`.

If `fallback` is not provided, it will be derived from `base`.
"""
struct DepthwiseConvAllowNinChangeStrategy{S,F} <: DecoratingJuMPΔSizeStrategy
  allowed_new_outgroups::Vector{Int}
  allowed_multipliers::Vector{Int}
  base::S
  fallback::F
end
DepthwiseConvAllowNinChangeStrategy(newoutputsmax::Integer, multipliersmax::Integer,base,fb...) = DepthwiseConvAllowNinChangeStrategy(0:newoutputsmax, 1:multipliersmax, base, fb...)

function DepthwiseConvAllowNinChangeStrategy(
  allowed_new_outgroups::AbstractVector{<:Integer},
  allowed_multipliers::AbstractVector{<:Integer},
  base, fb= recurse_fallback(s -> DepthwiseConvAllowNinChangeStrategy(allowed_new_outgroups, allowed_multipliers, s), base))
  return DepthwiseConvAllowNinChangeStrategy(collect(Int, allowed_new_outgroups), collect(Int, allowed_multipliers), base, fb)
end

NaiveNASlib.base(s::DepthwiseConvAllowNinChangeStrategy) = s.base
NaiveNASlib.fallback(s::DepthwiseConvAllowNinChangeStrategy) = s.fallback

NaiveNASlib.add_participants!(s::DepthwiseConvAllowNinChangeStrategy, vs=AbstractVertex[]) = NaiveNASlib.add_participants!(base(s), vs)

"""
  DepthwiseConvSimpleΔSizeStrategy(base, [fallback])

`DecoratingJuMPΔSizeStrategy` which only allows nout of `DepthwiseConv` layers to change.

Use if [`DepthwiseConvAllowNinChangeStrategy`](@ref) takes too long to solve.

The elements of `allowed_multipliers` determine the total number of allowed output elements, i.e the allowed
ratios of `nout / nin`.

If `fallback` is not provided, it will be derived from `base`.
"""
struct DepthwiseConvSimpleΔSizeStrategy{S, F} <: DecoratingJuMPΔSizeStrategy
  allowed_multipliers::Vector{Int}
  base::S
  fallback::F
end

DepthwiseConvSimpleΔSizeStrategy(maxms::Integer, base, fb...) = DepthwiseConvSimpleΔSizeStrategy(1:maxms, base, fb...)
function DepthwiseConvSimpleΔSizeStrategy(ms::AbstractVector{<:Integer}, base, fb=recurse_fallback(s -> DepthwiseConvSimpleΔSizeStrategy(ms, s), base))
  return DepthwiseConvSimpleΔSizeStrategy(collect(Int, ms), base, fb)
end
NaiveNASlib.base(s::DepthwiseConvSimpleΔSizeStrategy) = s.base
NaiveNASlib.fallback(s::DepthwiseConvSimpleΔSizeStrategy) = s.fallback

NaiveNASlib.add_participants!(s::DepthwiseConvSimpleΔSizeStrategy, vs=AbstractVertex[]) = NaiveNASlib.add_participants!(base(s), vs)

recurse_fallback(f, s::AbstractJuMPΔSizeStrategy) = wrap_fallback(f, NaiveNASlib.fallback(s))
recurse_fallback(f, s::NaiveNASlib.DefaultJuMPΔSizeStrategy) = s
recurse_fallback(f, s::NaiveNASlib.ThrowΔSizeFailError) = s
recurse_fallback(f, s::NaiveNASlib.ΔSizeFailNoOp) = s
recurse_fallback(f, s::NaiveNASlib.LogΔSizeExec) = NaiveNASlib.LogΔSizeExec(s.msgfun, s.level, f(s.andthen))

wrap_fallback(f, s) = f(s)
wrap_fallback(f, s::NaiveNASlib.LogΔSizeExec) = NaiveNASlib.LogΔSizeExec(s.msgfun, s.level, f(s.andthen))

function NaiveNASlib.compconstraint!(case, s, ::FluxLayer, data) end
function NaiveNASlib.compconstraint!(case, s::DecoratingJuMPΔSizeStrategy, lt::FluxLayer, data)
   NaiveNASlib.compconstraint!(case, NaiveNASlib.base(s), lt, data)
end
# To avoid ambiguity
function NaiveNASlib.compconstraint!(case::NaiveNASlib.ScalarSize, s::DecoratingJuMPΔSizeStrategy, lt::FluxDepthwiseConv, data)
  NaiveNASlib.compconstraint!(case, NaiveNASlib.base(s), lt, data)
end
function NaiveNASlib.compconstraint!(::NaiveNASlib.ScalarSize, s::AbstractJuMPΔSizeStrategy, ::FluxDepthwiseConv, data, ms=allowed_multipliers(s))

  # Add constraint that nout(l) == n * nin(l) where n is integer
  ins = filter(vin -> vin in keys(data.noutdict), inputs(data.vertex))

   # "data.noutdict[data.vertex] == data.noutdict[ins[i]] * x" where x is an integer variable is not linear
   # Instead we use the good old big-M strategy to set up a series of "or" constraints (exactly one of them must be true).
   # This is combined with the abs value formulation to force equality.
   # multiplier[j] is false if and only if data.noutdict[data.vertex] == ms[j]*data.noutdict[ins[i]]
   # all but one of multiplier must be true
   # Each constraint below is trivially true if multiplier[j] is true since 1e6 is way bigger than the difference between the two variables
  multipliers = @variable(data.model, [1:length(ms)], Bin)
  @constraint(data.model, sum(multipliers) == length(multipliers)-1)
  @constraint(data.model, [i=1:length(ins),j=1:length(ms)], data.noutdict[data.vertex] - ms[j]*data.noutdict[ins[i]] + multipliers[j] * 1e6 >= 0)
  @constraint(data.model, [i=1:length(ins),j=1:length(ms)], data.noutdict[data.vertex] - ms[j]*data.noutdict[ins[i]] - multipliers[j] * 1e6 <= 0)

  # Inputs which does not have a variable, possibly because all_in_Δsize_graph did not consider it to be part of the set of vertices which may change
  # We will constrain data.vertex to have integer multiple of its current size
  fixedins = filter(vin -> vin ∉ ins, inputs(data.vertex))
  if !isempty(fixedins)
    fv_out = @variable(data.model, integer=true)
    @constraint(data.model, [i=1:length(fixedins)], data.noutdict[data.vertex] == nout(fixedins[i]) * fv_out)
  end
end

allowed_multipliers(s::DepthwiseConvAllowNinChangeStrategy) = s.allowed_multipliers
allowed_multipliers(s::DepthwiseConvSimpleΔSizeStrategy) = s.allowed_multipliers
allowed_multipliers(::AbstractJuMPΔSizeStrategy) = 1:10

function NaiveNASlib.compconstraint!(case::NaiveNASlib.NeuronIndices, s::DecoratingJuMPΔSizeStrategy, t::FluxDepthwiseConv, data)
  NaiveNASlib.compconstraint!(case, base(s), t, data)
end
function NaiveNASlib.compconstraint!(case::NaiveNASlib.NeuronIndices, s::AbstractJuMPΔSizeStrategy, t::FluxDepthwiseConv, data)
  # Fallbacks don't matter here since we won't call it from below here, just add default so we don't accidentally crash due to some
  # strategy which hasn't defined a fallback
  if 15 < sum(keys(data.outselectvars)) do v
      layertype(v) isa FluxDepthwiseConv || return 0
      return log2(nout(v)) # Very roughly determined...
  end
    return NaiveNASlib.compconstraint!(case, DepthwiseConvSimpleΔSizeStrategy(10, s, NaiveNASlib.DefaultJuMPΔSizeStrategy()), t, data)
  end
  # The number of allowed multipliers can probably be better tuned, perhaps based on current size.
  return NaiveNASlib.compconstraint!(case, DepthwiseConvAllowNinChangeStrategy(10, 10, s, NaiveNASlib.DefaultJuMPΔSizeStrategy()), t, data)
  #=
  For benchmarking:
    using NaiveNASflux, Flux, NaiveNASlib.Advanced
    function timedwc(ds, ws)
      for w in ws
        for d in ds
          iv = conv2dinputvertex("in", w)
          dv = reduce((v, i) -> fluxvertex("dv$i", DepthwiseConv((1,1), nout(v) => fld1(i, 3) * nout(v)), v), 1:d; init=iv)
          metric = sum(ancestors(dv)) do v
            layer(v) isa DepthwiseConv || return 0
            return log2(nout(v))
          end
          res = @timed Δsize!(ΔNoutRelaxed(outputs(iv)[1] => w; fallback = ΔSizeFailNoOp()))
          println((width=w, depth = d, metric=metric, time=res.time))
          res.time < 5 || break
        end
      end
    end
  =#
end

function NaiveNASlib.compconstraint!(::NaiveNASlib.NeuronIndices, s::DepthwiseConvSimpleΔSizeStrategy, t::FluxDepthwiseConv, data)
  model = data.model
  v = data.vertex
  select = data.outselectvars[v]
  insert = data.outinsertvars[v]

  nin(v)[] == 1 && return # Special case, no restrictions as we only need to be an integer multple of 1

  ngroups = div(nout(v), nin(v)[])
  # Neurons mapped to the same weight are interleaved, i.e layer.weight[:,:,1,:] maps to y[1:ngroups:end] where y = layer(x)
  for group in 1:ngroups
    neurons_in_group = select[group : ngroups : end]
    @constraint(model, neurons_in_group[1] == neurons_in_group[end])
    @constraint(model, [i=2:length(neurons_in_group)], neurons_in_group[i] == neurons_in_group[i-1])

    insert_in_group = insert[group : ngroups : end]
    @constraint(model, insert_in_group[1] == insert_in_group[end])
    @constraint(model, [i=2:length(insert_in_group)], insert_in_group[i] == insert_in_group[i-1])
  end

  NaiveNASlib.compconstraint!(NaiveNASlib.ScalarSize(), s, t, data, allowed_multipliers(s))
end

function NaiveNASlib.compconstraint!(case::NaiveNASlib.NeuronIndices, s::DepthwiseConvAllowNinChangeStrategy, t::FluxDepthwiseConv, data)
  model = data.model
  v = data.vertex
  select = data.outselectvars[v]
  insert = data.outinsertvars[v]

  nin(v)[] == 1 && return # Special case, no restrictions as we only need to be an integer multple of 1?

  # Step 1:
  # Neurons mapped to the same weight are interleaved, i.e layer.weight[:,:,1,:] maps to y[1:ngroups:end] where y = layer(x)
  # where ngroups = nout / nin. For example, nout = 12 and nin = 4 mean size(layer.weight) == (..,3, 4)
  # Thus we have ngroups = 3 and with groupsize of 4.
  # Now, if we were to change nin to 3, we would still have ngroups = 3 but with groupsize 3 and thus nout = 9
  #
  ins = filter(vin -> vin in keys(data.noutdict), inputs(v))
  # If inputs to v are not part of problem we have to keep nin(v) fixed!
  isempty(ins) && return NaiveNASlib.compconstraint!(case, DepthwiseConvSimpleΔSizeStrategy(allowed_multipliers(s), base(s)), t, data)
  # TODO: Check if input is immutable and do simple strat then too?
  inselect = data.outselectvars[ins[]]
  ininsert = data.outinsertvars[ins[]]

  #ngroups = div(nout(v), nin(v)[])
  ningroups = nin(v)[]
  add_depthwise_constraints(model, inselect, ininsert, select, insert, ningroups, s.allowed_new_outgroups, s.allowed_multipliers)
end

function add_depthwise_constraints(model, inselect, ininsert, select, insert, ningroups, allowed_new_outgroups, allowed_multipliers)

  # ningroups is the number of "input elements" in the weight array, last dimension at the time of writing
  # noutgroups is the number of "output elements" in the weight array, second to last dimension at time of writing
  # nin is == ningroups while nout == ningroups * noutgroups

  # When changing nin by Δ, we will get ningroups += Δ which just means that we insert/remove Δ input elements.
  # Inserting one new input element at position i will get us noutgroups new consecutive outputputs at position i
  # Thus nout change by Δ * noutgroups.

  # Example:

  # dc = DepthwiseConv((1,1), 3 => 9; bias=false);

  # size(dc.weight)
  # (1, 1, 3, 3)

  # indata = randn(Float32, 1,1,3, 1);

  # reshape(dc(indata),:)'
  # 1×9 adjoint(::Vector{Float32}) with eltype Float32:
  #  0.315234  0.267166  -0.227757  0.03523  0.0516438  -0.124007  -0.63409  0.0273361  -0.457146

  # Inserting zeros instead of adding/removing just as it is less verbose to do so.
  # dc.weight[:,:,:,2] .= 0;

  # reshape(dc(indata),:)'
  # 1×9 adjoint(::Vector{Float32}) with eltype Float32:
 #  0.315234  0.267166  -0.227757  0.0  0.0  0.0  -0.63409  0.0273361  -0.457146

  # When changing nout by Δ where Δ ==  k * ningroups where k is an integer we can simply just add/remove k output elements
  # Output elements are interleaved in the output of the layer, so if we add/remove output element i, the elements i:ningroups:end
  # will be added/dropped from the layer output.
  # When Δ is not an integer multiple things get pretty hairy and I'm not sure the constraints below handle all causes
  # We then basically need to change nin so that the new nout is an integer multiple of the new nin.

  # Example

  # dc = DepthwiseConv((1,1), 3 => 9; bias=false);

  # reshape(dc(indata),:)'
  # 1×9 adjoint(::Vector{Float32}) with eltype Float32:
  #  0.109549  0.273408  -0.299883  0.0131096  0.00240822  -0.0461159  -0.138605  -0.557134  0.0246111

  # julia> dc.weight[:,:,2,:] .= 0;

  # reshape(dc(indata),:)'
  # 1×9 adjoint(::Vector{Float32}) with eltype Float32:
  # 0.109549  0.0  -0.299883  0.0131096  0.0  -0.0461159  -0.138605  0.0  0.0246111

  # One thing which makes this a bit tricky is that while noutgruops and ningroups can change
  # we are here stuck with the select and insert arrays that correspond to the original sizes
  # As such, keep in mind that noutgroups and ningroups in this function always relate to the
  # lengths of the arrays of JuMP variables that we have right now.

  M = 1e6
  noutgroups = div(length(select), ningroups)

  # select_none_in_group[g] == 1 if we drop group g
  # TODO: Replace by just sum(select_in_group[g,:]) below?
  select_none_in_group = @variable(model, [1:noutgroups], Bin)
  select_in_group = reshape(select, noutgroups, ningroups)
  @constraint(model, sum(select_in_group, dims=2) .<= M .* (1 .- select_none_in_group))

  # Tie selected input indices to selected output indices. If we are to drop any layer output indices, we need to drop the whole group
  # as one whole group is tied to a single element in the output element dimension of the weight array.
  @constraint(model, [g=1:noutgroups, i=1:ningroups], inselect[i] - select_in_group[g, i] + select_none_in_group[g] * M >= 0)
  @constraint(model, [g=1:noutgroups, i=1:ningroups], inselect[i] - select_in_group[g, i] - select_none_in_group[g] * M <= 0)

  # This variable handles the constraint that if we add a new input, we need to add noutgroups new outputs
  # Note that for now, it is only connected to the insert variable through the sum constraint below.
  insert_new_inoutgroups = @variable(model, [[1], 1:ningroups], Int, lower_bound=0, upper_bound=ningroups * maximum(allowed_multipliers))

  insize = @expression(model, sum(inselect) + sum(ininsert))
  outsize = @expression(model, sum(select) + sum(insert))

  # inmultipliers[j] == 1 if nout(v) == allowed_multipliers[j] * nin(v)[]
  inmultipliers = @variable(model, [1:length(allowed_multipliers)], Bin)
  #SOS1 == Only one can be non-zero. Not strictly needed, but it seems like it speeds up the solver
  @constraint(model, inmultipliers in SOS1(1:length(inmultipliers)))
  @constraint(model, sum(inmultipliers) == 1)

  @constraint(model,[i=1:length(ininsert), j=1:length(inmultipliers)], allowed_multipliers[j] * ininsert[i] - insert_new_inoutgroups[1, i] + (1-inmultipliers[j]) * M >= 0)
  @constraint(model,[i=1:length(ininsert), j=1:length(inmultipliers)], allowed_multipliers[j] * ininsert[i] - insert_new_inoutgroups[1, i] - (1-inmultipliers[j]) * M <= 0)

  @constraint(model, [j=1:length(inmultipliers)], allowed_multipliers[j] * insize - outsize + (1-inmultipliers[j]) * M >= 0)
  @constraint(model, [j=1:length(inmultipliers)], allowed_multipliers[j] * insize - outsize - (1-inmultipliers[j]) * M <= 0)

  insert_new_inoutgroups_all_inds = vcat(zeros(noutgroups-1,ningroups), insert_new_inoutgroups)

  # This variable handles the constraint that if we want to increase the output size without changing
  # the input size, we need to insert whole outgroups since all noutgroup outputs are tied to a single
  # output element in the weight array.
  insert_new_outgroups = @variable(model, [1:noutgroups, 1:ningroups], Int, lower_bound = 0, upper_bound=noutgroups * maximum(allowed_new_outgroups))

  # insert_no_outgroup[g] == 1 if we don't insert a new output group after group g
  # TODO: Replace by just sum(insert_new_outgroups[g,:]) below?
  insert_no_outgroup = @variable(model, [1:noutgroups], Bin)
  @constraint(model, sum(insert_new_outgroups, dims=2) .<= M .* (1 .- insert_no_outgroup))

  # When adding a new output group, all inserts must be identical
  # If we don't add any, all inserts are just 0
  @constraint(model, [g=1:noutgroups, i=2:ningroups], insert_new_outgroups[g,i] == insert_new_outgroups[g,i-1])
  @constraint(model, [g=1:noutgroups], insert_new_outgroups[g,1] == insert_new_outgroups[g,end])

  # new_outgroup[g,j] == 1 if we are inserting allowed_new_outgroups[j] new output groups after output group g
  noutmults = 1:length(allowed_new_outgroups)
  new_outgroup = @variable(model, [1:noutgroups, noutmults], Bin)
  #SOS1 == Only one can be non-zero. Not strictly needed, but it seems like it speeds up the solver
  @constraint(model,[g=1:noutgroups], new_outgroup[g,:] in SOS1(1:length(allowed_new_outgroups)))
  @constraint(model,[g=1:noutgroups], sum(new_outgroup[g,:]) == 1)

  groupsum = @expression(model, [g=1:noutgroups], sum(insert_new_outgroups[g,:]) - sum(insert_new_inoutgroups_all_inds[g,:]))
  @constraint(model, [g=1:noutgroups, j=noutmults], groupsum[g] - allowed_new_outgroups[j]*insize + (1-new_outgroup[g,j] + insert_no_outgroup[g]) * M >= 0)
  @constraint(model, [g=1:noutgroups, j=noutmults], groupsum[g] - allowed_new_outgroups[j]*insize - (1-new_outgroup[g,j] + insert_no_outgroup[g]) * M <= 0)

  # Finally, we say what the inserts shall be: the sum of inserts from new inputs and new outputs
  # I think this allows for them to be somewhat independent, but there are still cases when they
  # can't change simultaneously. TODO: Check those cases
  @constraint(model, insert .== reshape(insert_new_outgroups, :) .+ reshape(insert_new_inoutgroups_all_inds,:))

end


let

    roadway = gen_straight_roadway(3, 1000.0, lane_width=1.0)
    rec = SceneRecord(1, 0.1, 5)
    update!(rec, Scene([
            Vehicle(VehicleState(VecSE2( 0.0,0.0,0.0), roadway, 10.0), VehicleDef(AgentClass.CAR, 5.0, 2.0), 1),
            Vehicle(VehicleState(VecSE2(10.0,0.0,0.0), roadway, 10.0), VehicleDef(AgentClass.CAR, 5.0, 2.0), 2),
        ]))

    @test convert(Float64, get(IS_COLLIDING, rec, roadway, 1)) == 0.0
    @test convert(Float64, get(IS_COLLIDING, rec, roadway, 2)) == 0.0

    @test isapprox(convert(Float64, get(MARKERDIST_LEFT, rec, roadway, 1)), 0.5)
    @test isapprox(convert(Float64, get(MARKERDIST_LEFT, rec, roadway, 2)), 0.5)
    @test isapprox(convert(Float64, get(MARKERDIST_RIGHT, rec, roadway, 1)), 0.5)
    @test isapprox(convert(Float64, get(MARKERDIST_RIGHT, rec, roadway, 2)), 0.5)
    @test isapprox(convert(Float64, get(MARKERDIST_LEFT_LEFT, rec, roadway, 1)), 1.5)
    @test isapprox(convert(Float64, get(MARKERDIST_LEFT_LEFT, rec, roadway, 2)), 1.5)
    @test isnan(convert(Float64, get(MARKERDIST_RIGHT_RIGHT, rec, roadway, 1)))
    @test isnan(convert(Float64, get(MARKERDIST_RIGHT_RIGHT, rec, roadway, 2)))

    @test isapprox(convert(Float64, get(DIST_FRONT, rec, roadway, 1)), 10.0 - 5.0)
    @test convert(Float64, get(DIST_FRONT, rec, roadway, 2, censor_hi=100.0)) == 100.0

    update!(rec, Scene([
            Vehicle(VehicleState(VecSE2( 1.0,0.0,0.0), roadway, 10.0), VehicleDef(AgentClass.CAR, 5.0, 2.0), 1),
            Vehicle(VehicleState(VecSE2(10.0,0.0,0.0), roadway, 10.0), VehicleDef(AgentClass.CAR, 5.0, 2.0), 2),
            Vehicle(VehicleState(VecSE2(12.0,1.0,0.0), roadway, 10.0), VehicleDef(AgentClass.CAR, 5.0, 2.0), 3),
            Vehicle(VehicleState(VecSE2( 0.0,1.0,0.0), roadway, 10.0), VehicleDef(AgentClass.CAR, 5.0, 2.0), 4),
        ]))

    @test isapprox(convert(Float64, get(MARKERDIST_LEFT, rec, roadway, 3)), 0.5)
    @test isapprox(convert(Float64, get(MARKERDIST_RIGHT, rec, roadway, 3)), 0.5)
    @test isapprox(convert(Float64, get(MARKERDIST_LEFT_LEFT, rec, roadway, 3)), 1.5)
    @test isapprox(convert(Float64, get(MARKERDIST_RIGHT_RIGHT, rec, roadway, 3)), 1.5)

    @test isapprox(convert(Float64, get(DIST_FRONT, rec, roadway, 1)), 9.0 - 5.0)
    @test isapprox(convert(Float64, get(DIST_FRONT_LEFT, rec, roadway, 1)), 11.0 - 5.0)
    @test convert(Float64, get(DIST_FRONT_RIGHT, rec, roadway, 1)) == 100.0
    @test isapprox(convert(Float64, get(DIST_FRONT, rec, roadway, 4)), 12.0 - 5.0)
    @test convert(Float64, get(DIST_FRONT_LEFT, rec, roadway, 4)) == 100.0
    @test isapprox(convert(Float64, get(DIST_FRONT_RIGHT, rec, roadway, 4)), 1.0 - 5.0)

    @test convert(Float64, get(IS_COLLIDING, rec, roadway, 1)) == 1.0
    @test convert(Float64, get(IS_COLLIDING, rec, roadway, 2)) == 1.0
    @test convert(Float64, get(IS_COLLIDING, rec, roadway, 3)) == 1.0
    @test convert(Float64, get(IS_COLLIDING, rec, roadway, 4)) == 1.0

    update!(rec, Scene([
            Vehicle(VehicleState(VecSE2(  0.0,0.0,0.0), roadway, 10.0), VehicleDef(AgentClass.CAR, 5.0, 2.0), 1),
            Vehicle(VehicleState(VecSE2(-10.0,1.0,0.0), roadway, 10.0), VehicleDef(AgentClass.CAR, 5.0, 2.0), 2),
            Vehicle(VehicleState(VecSE2(  0.0,1.0,0.5), roadway, 10.0), VehicleDef(AgentClass.CAR, 5.0, 2.0), 3),
        ]))

    @test convert(Float64, get(IS_COLLIDING, rec, roadway, 1)) == 1.0
    @test convert(Float64, get(IS_COLLIDING, rec, roadway, 2)) == 0.0
    @test convert(Float64, get(IS_COLLIDING, rec, roadway, 3)) == 1.0

    # super-simple test
    for f in allfeatures()
        get(f, rec, roadway, 3)
    end
end


using Distributed
using Dates
using DelimitedFiles
using PyPlot

if (nprocs() < 2 && isdefined(Main, :usePrecaching) && usePrecaching)
    println("Adding process for file precaching!")
    addprocs(1)
end

scriptspath =  pwd()
if ! isdefined(Main, :FunctionsLoaded)
    #load includes
    include("includes/BaselineFunctions.jl")
    include("includes/InterpolationFunctions.jl")
    include("includes/DatasetPreloader.jl")
    include("includes/MasslistFunctions.jl")
    include("includes/PeakshapeFunctions.jl")
    include("includes/TofFunctions.jl")
    include("includes/MultipeakFunctions.jl")
    include("includes/ResultFileFunctions.jl")
    #misc module leftovers
    include("includes/TofTracer.jl")

    #Load Personal MassLibrary Definitions
    include("$(scriptspath)/manualMassLibrary.jl")

    #Load workflows
    include("$(scriptspath)/workflows/combinedMassScaleAndExtractSumSpec.jl")
    include("$(scriptspath)/workflows/peakShape.jl")
    include("$(scriptspath)/workflows/deconvolutionMatrix.jl")
    FunctionsLoaded = true
end


#Delete clusters
function b1_cb_template(widgetptr::Ptr,user_data::Tuple{Single_Channel})

    sc, = user_data
    clus=sc.buf.selected_clus

    if (clus<1) #do nothing if zeroth cluster selected
    else
        delete_cluster(sc.temp,clus)
        deleteat!(sc.sort_list,sc.total_clus+1)
        sc.total_clus -= 1
        sc.buf.selected_clus = 0
        selmodel=Gtk.GAccessor.selection(sc.sort_tv)
        Gtk.select!(selmodel, Gtk.iter_from_index(sc.sort_list,1))
        sc.buf.c_changed=true
    end
    nothing
end

function delete_cluster(c::ClusterTemplate,n)

    for i=1:size(c.templates,1)
        c.templates[i,n] = 0.0
        c.sig_max[i,n] = 0.0
        c.sig_min[i,n] = 0.0
    end

    c.tol[n]=0.0

    if n == c.num
        c.num -= 1
    else
        for i=n:(c.num-1)
            for j=1:size(c.templates,1)
                c.templates[j,i]=c.templates[j,i+1]
                c.sig_max[j,i]=c.sig_max[j,i+1]
                c.sig_min[j,i]=c.sig_min[j,i+1]
            end
            c.tol[i]=c.tol[i+1]
        end
        c.num -= 1
    end

    nothing
end

#Add Unit
function b2_cb_template(widget::Ptr,user_data::Tuple{Single_Channel})

    sc, = user_data

    #Add total number of units and go to that unit
    sc.total_clus += 1
    sc.temp.num += 1

    sc.buf.selected_clus = sc.total_clus
    push!(sc.sort_list,(sc.total_clus,))

    selmodel=Gtk.GAccessor.selection(sc.sort_tv)
    Gtk.select!(selmodel, Gtk.iter_from_index(sc.sort_list, sc.total_clus+1))

    nothing
end

function b3_cb_template(widget::Ptr,user_data::Tuple{Single_Channel,SortView})

    sc, sortview_widgets = user_data

    mybutton = convert(Button, widget)

    if sc.pause
        Gtk.GAccessor.active(sc.pause_button,false)
    end

    if !sc.hold
        clear_c2(sc.c2,sc.spike)
        sc.ctx2=Gtk.getgc(sc.c2)
        sc.ctx2s=copy(sc.ctx2)
        sc.buf.ind=1
        sc.buf.count=1
        sc.hold=true
        sc.pause=false
        change_button_label(mybutton,"Stop Collection")
    else

        Gtk.GAccessor.active(sc.pause_button,true)
        change_button_label(mybutton,"Collect Templates")

        if sc.buf.count==size(sc.buf.spikes,2)
            sc.buf.ind=sc.buf.count
        end

        if visible(sortview_widgets.win)
            recalc_features(sortview_widgets)
            replot_sort(sortview_widgets)
        end
    end

    nothing
end

function b4_cb_template(widgetptr::Ptr,user_data::Tuple{Single_Channel})

    sc, = user_data
    widget = convert(ToggleButton, widgetptr)

    #Untoggle
    if !getproperty(widget,:active,Bool)
        sc.buf.replot=true
    else
        #Toggling, so draw template
        if sc.pause
            draw_template_bounds(sc)
        end
    end

    nothing
end

function draw_template_bounds(sc::Single_Channel)

    ctx = sc.ctx2
    clus=sc.buf.selected_clus

     if clus>0
            s=sc.s
            o=sc.o

            Cairo.translate(ctx,0.0,sc.h2/2)
            Gtk.scale(ctx,sc.w2/sc.wave_points,s)

            move_to(ctx,1.0,sc.temp.templates[1,clus]+(sc.temp.sig_max[1,clus]*sc.temp.tol[clus])-o)

            for i=2:size(sc.temp.templates,1)
                y=sc.temp.templates[i,clus]+(sc.temp.sig_max[i,clus]*sc.temp.tol[clus])-o
                line_to(ctx,i,y)
            end

            y=sc.temp.templates[end,clus]-(sc.temp.sig_min[end,clus]*sc.temp.tol[clus])-o

            line_to(ctx,size(sc.temp.templates,1),y)

            for i=(size(sc.temp.templates,1)-1):-1:1
                y=sc.temp.templates[i,clus]-(sc.temp.sig_min[i,clus]*sc.temp.tol[clus])-o
                line_to(ctx,i,y)
            end

            close_path(ctx)

            select_color(ctx,clus+1)
            set_line_width(ctx,3.0)
            stroke_preserve(ctx)

            select_color(ctx,clus+1,.5)
            fill(ctx)
        end

    identity_matrix(ctx)
end

function check_cb_template(widget::Ptr,user_data::Tuple{Single_Channel})

    sc, = user_data

    mycheck=convert(CheckButton,widget)

    if getproperty(mycheck,:active,Bool)
        sc.sort_cb=true
    else
        sc.sort_cb=false
    end

    if sc.sort_cb
        draw_templates(sc)
    end

    nothing
end

function draw_templates(sc::Single_Channel)

    ctx = sc.ctx2s
    mywidth=width(ctx)
    myheight=height(ctx)

    s=sc.s
    o=sc.o

    Cairo.translate(ctx,0.0,myheight/2)
    Gtk.scale(ctx,mywidth/sc.wave_points,s)

    for clus=1:sc.total_clus

        move_to(ctx,1.0,(sc.temp.templates[1,clus])-o)

        for i=2:size(sc.temp.sig_max,1)
            y=sc.temp.templates[i,clus]-o
            line_to(ctx,i,y)
        end

        select_color(ctx,clus+1)
        set_line_width(ctx,3.0)
        stroke(ctx)
    end
    identity_matrix(ctx)
    nothing
end

function get_template_dims(sc::Single_Channel,clus)

    ctx=Gtk.getgc(sc.c3)
    mywidth=width(ctx)

    total_clus = max(sc.total_clus+1,5)

    xbounds=linspace(0.0,mywidth,total_clus+1)

    (xbounds[clus],xbounds[clus+1],0.0,130.0)
end

#=
Calculates the mean and bounds for a collection of spikes that
1) are not hidden by the mask vector
2) meet some condition

examples of this condition could be whether the spikes are contained within a rubberband, or if they are part of a certain cluster already
=#
function make_cluster(input::Array{T,2},mask,count,condition) where T

    hits=0
    mymean=zeros(Float64,size(input,1)-1)
    mysum=zeros(Int64,size(input,1)-1)
    mybounds=zeros(Float64,size(input,1)-1,2)

    for i=1:count
        if (condition[i])&&(mask[i])
            hits+=1
            for ii=1:length(mymean)
                mysum[ii] += input[ii,i]
                if hits==1
                    mybounds[ii,1]=input[ii,i]
                    mybounds[ii,2]=input[ii,i]
                else
                    if input[ii,i]>mybounds[ii,1]
                        mybounds[ii,1]=input[ii,i]
                    end
                    if input[ii,i]<mybounds[ii,2]
                        mybounds[ii,2]=input[ii,i]
                    end
                end
            end
        end
    end

    if hits==0
        hits=1
    end

    for i=1:length(mymean)
        mymean[i] = mysum[i]/hits
        mybounds[i,2] = abs(mymean[i]-mybounds[i,2])
        mybounds[i,1] = abs(mybounds[i,1] - mymean[i])
    end

    (mymean,mybounds)
end

function change_cluster(c::ClusterTemplate,mymean::Array{Float64,1},mystd::Array{Float64,2},n)

    for i=1:length(mymean)
        c.templates[i,n] = mymean[i]
        c.sig_max[i,n] = mystd[i,1]
        c.sig_min[i,n] = mystd[i,2]
    end

    nothing
end

function template_cluster(sc::Single_Channel,clus,mymean::Array{Float64,1},mymin::Array{Float64,1},mymax::Array{Float64,1},tol::Float64)

    @inbounds for i=1:sc.buf.ind

        mymisses=0
        for j=1:length(mymean)
            if (sc.buf.spikes[j,i]<(mymean[j]-(mymin[j]*tol)))|(sc.buf.spikes[j,i]>(mymean[j]+(mymax[j]*tol)))
                mymisses+=1
                if mymisses>1
                    break
                end
            end
        end
        if mymisses<1 #If passes template matching, set as unit
            sc.buf.clus[i]=clus
        elseif sc.buf.clus[i]==clus #If did not pass, but was previously, set to noise cluster
            sc.buf.clus[i]=0
        end
    end

    nothing
end

function send_clus(s::Array{T,1},sc::Single_Channel) where T<:Sorting
    s[sc.spike].c=deepcopy(sc.temp)
    nothing
end

function send_clus(s::DArray{T,1,Array{T,1}},sc::Single_Channel) where T<:Sorting
    (nn,mycore)=get_thres_id(s,sc.spike)
    remotecall_wait(((x,tt,num)->localpart(x)[num].c=tt),mycore,s,sc.temp,nn)
    nothing
end


for OP in (:(&), :(|), :(⊻))
    @eval begin

       @inline function $OP(x::T, y::T) where T<:SafeInteger
           ix = baseint(x)
           iy = baseint(y)
           result = $OP(ix, iy)
           return safeint(result)
       end

       @inline function $OP(x::T1, y::T2) where {T1<:SafeInteger, T2<:SafeInteger}
           xx, yy = promote(x, y)
           return $OP(xx, yy)
       end

       @inline function $OP(x::T1, y::T2) where {T1<:SafeInteger, T2<:Integer}
           xx, yy = promote(x, y)
           return $OP(xx, yy)
       end

       @inline function $OP(x::T2, y::T1) where {T1<:SafeInteger, T2<:Integer}
           xx, yy = promote(x, y)
           return $OP(xx, yy)
       end

   end
end

for OP in (:(<), :(<=), :(>=), :(>), :(!=), :(==), :isless, :isequal)
    @eval begin

       @inline function $OP(x::T, y::T) where T<:SafeInteger
           ix = baseint(x)
           iy = baseint(y)
           result = $OP(ix, iy)
           return safeint(result)
       end

       @inline function $OP(x::T1, y::T2) where {T1<:SafeInteger, T2<:SafeInteger}
           xx, yy = promote(x, y)
           return $OP(xx, yy)
       end

       @inline function $OP(x::T1, y::T2) where {T1<:SafeInteger, T2<:Integer}
           xx, yy = promote(x, y)
           return $OP(xx, yy)
       end

       @inline function $OP(x::T1, y::T2) where {T1<:Integer, T2<:SafeInteger}
           xx, yy = promote(x, y)
           return $OP(xx, yy)
       end
   end
end

for OP in (:(>>>), :(>>), :(<<))
    @eval begin

       @inline function $OP(x::T, y::T) where T<:SafeInteger
            r1 = baseint(x)
            r2 = baseint(y)
            bitsof(T) < abs(r2) && throw(OverflowError("cannot shift $T by $y"))
            result = $OP(r1, r2)
            return reinterpret(T, result)
        end

        @inline function $OP(x::T1, y::T2) where {T1<:SafeInteger, T2<:SafeInteger}
            xx = baseint(x)
            yy = baseint(y)
            bitsof(T1) < abs(yy) && throw(OverflowError("cannot shift $T1 by $yy"))
            return reinterpret(T1, $OP(xx, yy))
        end

        @inline function $OP(x::T1, y::T2) where {T1<:SafeInteger, T2<:Integer}
            xx = baseint(x)
            bitsof(T1) < abs(y) && throw(OverflowError("cannot shift $T1 by $y"))
            return reinterpret(T1, $OP(xx, y))
        end

        @inline function $OP(x::T1, y::Int) where {T1<:SafeInteger}
            xx = baseint(x)
            bitsof(T1) < abs(y) && throw(OverflowError("cannot shift $T1 by $y"))
            return reinterpret(T1, $OP(xx, y))
        end

        @inline function $OP(x::T1, y::T2) where {T1<:Integer, T2<:SafeInteger}
            yy = baseint(y)
            bitsof(T1) < abs(yy) && throw(OverflowError("cannot shift $T1 by $yy"))
            return reinterpret(T1, $OP(x, yy))
        end

   end
end


module RoMEPlotting

using Cairo, Fontconfig

using Reexport
@reexport using Gadfly
@reexport using Colors

using Statistics, LinearAlgebra
using StatsBase
using Compose
using Dates
using DistributedFactorGraphs
using KernelDensityEstimate, KernelDensityEstimatePlotting
using IncrementalInference, RoME
using DocStringExtensions
using ApproxManifoldProducts
using TensorCast
using Requires
# import ApproxManifoldProducts: mmd! # future dependency

import Gadfly: plot
import KernelDensityEstimatePlotting: plot, drawHorDens, plotKDE
import KernelDensityEstimatePlotting: getColorsByLength

# assuming this is a good size for everybody
@info "For larger plots when using a browswer, run Gadfly.set_default_plot_size(30cm,20cm)"
# Gadfly.set_default_plot_size(30cm,20cm)

export
  # Associated with IncrementalInference
  investigateMultidimKDE,
  drawHorDens, # from KDEPlotting
  plotKDE,
  plotVariable2D,
  plotKDEofnc,
  plotKDEresiduals,
  plotMCMC,
  plotUpMsgsAtCliq,
  plotPriorsAtCliq,
  investigateMultidimKDE,
  plot,
  whosWith,
  plotUpMsgAtCliq,
  dwnMsgsAtCliq,
  plotPose2DMC!,
  mcmcPose2D!,
  plotLbl,
  predCurrFactorBeliefs,
  plotFactorBeliefs,
  localProduct,
  plotLocalProduct,
  plotLocalProductCylinder,
  plotTreeProductUp,
  plotTreeProductDown,
  plotCliqUpMsgs,
  saveplot,
  animateVertexBelief,
  getColorsByLength,

  # Associated with RoME
  plotTrajectoryArrayPose2,
  togglePrtStbLines,
  plotLsrScanFeats,
  saveImgSeq,
  stbPrtLineLayers!,
  plotSLAM2D,
  plotSLAM2DPoses,
  plotSLAM2DLandmarks,
  plotPose,
  plotMarginalContour,
  accumulateMarginalContours,
  plotPose3Pairs,
  progressExamplePlot,
  plotTrckStep,
  plotPose2Vels,
  plotProductVsKDE,
  plotPairVariables,
  plotPairPose2,
  plotVariableGivenFactor,
  plotCliqDownMsgs,
  plotFactor,
  plotFactorValues,
  plotFactorMeasurements,
  reportFactors

# EXPERIMENTAL
const AbstractMatrix__{T} = Union{AbstractArray{T,2}, Adjoint{T,<:AbstractArray{T,2}}}

using Requires

# will be overwritten if flux is present (dont make const)
PlotTypesPose2 = Union{Type{Pose2Pose2}, Type{Pose2Point2BearingRange}, Type{Pose2Point2Range}, Type{Pose2Point2Bearing}}
ExtendedPose2Pose2Types = Pose2Pose2

include("SolverVisualization.jl")
include("RobotViz.jl")
include("services/PlotPose.jl")
include("CurrentlyUnmaintained.jl")
include("PlotHexUtils.jl")
include("PlotFactors.jl")
include("PlotVariables.jl")
include("PlotFactorsReload.jl")
include("Deprecated.jl")
include("NeedsFixing.jl")

function __init__()
  @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" include("FluxSpecificFeatures.jl")
  @require ImageMagick="6218d12a-5da1-5696-b52f-db25d2ecc6d1" include("imagemagick.jl")
end

end


###
# Methods to transform of heta-compiler generated files into Platform content
###

const SUPPORTED_VERSIONS = ["0.6.5", "0.6.6"]

# transformation of tuple to Platform
function Platform(
    models::NamedTuple,
    scenarios::Tuple,
    version::String
)
    # TODO: semver approach might be better
    @assert length(indexin(version, SUPPORTED_VERSIONS)) > 0 "Heta compiler of the version \"$version\" is not supported."

    print("Loading platform... ")
    model_pairs = [pair[1] => Model(pair[2]...) for pair in pairs(models)]

    platform = Platform(
        OrderedDict{Symbol,Model}(model_pairs),
        OrderedDict{Symbol,Scenario}()
    )
    println("OK!")

    return platform
end

# transformation of tuple to Model
function Model(
    init_func::Function,
    ode_func::Function,
    time_events::NamedTuple,
    c_events::NamedTuple,
    stop_events::NamedTuple,
    saving_generator::Function,
    constants_num::NamedTuple,
    events_active::NamedTuple,
    records_output::NamedTuple
)
    events = Pair[]
    ## FIXME : remove event name from heta-compiler
    for (name,event_tuple) in pairs(time_events) # time events
        evt = name => TimeEvent(event_tuple[1], event_tuple[2], event_tuple[4])
        push!(events, evt)
    end
    for (name,event_tuple) in pairs(c_events) # c events
        evt = name => CEvent(event_tuple[1], event_tuple[2], event_tuple[4])
        push!(events, evt)
    end
    for (name,event_tuple) in pairs(stop_events) # stop events
        evt = name => StopEvent(event_tuple[1], event_tuple[3])
        push!(events, evt)
    end

    #observable_pairs = filter((p) -> p[2], pairs(records_output)) # from records_output
    #observables = Symbol[p[1] for p in observable_pairs]
    records_output_ = collect(Pair{Symbol,Bool}, pairs(records_output))
    events_active_ = collect(Pair{Symbol,Bool}, pairs(events_active))

    ### fake run, disabled because it slows model loading down
    #=
    _u0, _p0 = init_func(constants_num)
    constants = LVector(constants_num)
    _params = Params(constants, _p0)
    prob = ODEProblem(ode_func, _u0, (0.,1.), _params)

    # check if default alg can solve the prob
    integrator = init(prob, DEFAULT_ALG)
    step!(integrator)
    ret = check_error(integrator)
    ret != :Success && @warn "Default algorithm returned $ret status. Consider using a different algorithm."
    =#
    model = Model(
        init_func,            # init_func
        ode_func,             # ode_func
        NamedTuple(events),   # events :: Changed to NamedTuple
        saving_generator,     # saving_generator
        records_output_,
        constants_num,        # constants :: Changed to NamedTuple
        NamedTuple(events_active_) # events_active :: Changed to NamedTuple
    )

    return model
end


module Experiment

export trialList2DataFrame
export data_of_run
export mat2jdl

using DataFrames
using MAT
using JLD
using Dates

using ..MyUtils
using ..Neurons

##################################################
#
#
#
##################################################
struct data_of_run
    paradigm::String
    date::Date
    time::DateTime
    sf::Float64
    trial_list::DataFrame
end

##################################################
#
#
#
##################################################
function trialList2DataFrame( data_dict )

    ### orderd according to intended column order in Julia
    ### (f1, f2): f1 -> column index in .mat.TrialList,
    ordered_column_name_pairs = [
        ([ Neurons.recording() for ii in 1:size(data_dict["TrialList"],1) ], "Neurons"),
        ([ [ data_dict["EyeX"][.!isnan.(data_dict["EyeX"][:,trl]),trl], data_dict["EyeY"][.!isnan.(data_dict["EyeY"][:,trl]),trl] ] for trl in 1:size(data_dict["EyeY"],2) ], "EyeXY"),
        (fill([[]],size(data_dict["TrialList"],1)), "Saccades"),
        (fill([[]],size(data_dict["TrialList"],1)), "Saccades_clean"),
        (fill([[]],size(data_dict["TrialList"],1)), "Fixations"),
        (fill([],size(data_dict["TrialList"],1)), "LookedAtTarget"),
        (data_dict["TrialList"][:,9], "TrialStartTime"),
        (data_dict["TrialList"][:,10], "TrialEndTime"),
        (fill(false, size(data_dict["TrialList"],1)), "Valid"),
        (data_dict["TrialList"][:,24].==1, "Hit"),
        (fill(false, size(data_dict["TrialList"],1)), "FA"),
        (fill(false, size(data_dict["TrialList"],1)), "Invalid"),
        (Array{Float64,1}(undef,size(data_dict["TrialList"],1)), "RTfirstFixation"),
        (Array{Float64,1}(undef,size(data_dict["TrialList"],1)), "RTcorrectFixation"),
        (data_dict["TrialList"][:,23].==1, "WasMissed"),                    ### !!! if this row changes, changes the part in mat2jdl() where missed trials are filtered out !!!
        (data_dict["TrialList"][:,22].==1, "WasTimeout"),
        (data_dict["TrialList"][:,16].+1, "ValidTargetIndex"),
        (data_dict["TrialList"][:,12].+1, "CueTargetIndex"),
        (data_dict["TrialList"][:,17].==1, "WasCueValid"),
        # (data_dict["TrialList"][:,18].==1, "WasFixationFailureFixPt"),    ### redundant with WasMissed
        # (data_dict["TrialList"][:,19].==1, "WasFixationFailureTarget"),   ### redundant with WasTimeout
        (data_dict["TrialList"][:,14], "LuminanceChange"),
        (data_dict["TrialList"][:,11], "LuminanceChangeTime"),
        (data_dict["TrialList"][:,21], "PauseDuration"),
        (data_dict["TrialList"][:,13], "ITI"),
        (data_dict["TrialList"][:,15], "NumberOfNonFixFailTrials"),
        (data_dict["TrialList"][:,20].==1, "WasLEDChangeActive")]

    ### -- Now in Behavior.jl --- ###
    # ### extract eye trace to store events in TrialList
    # xy = [[data_dict["EyeX"][:,trl], data_dict["EyeY"][:,trl]] for trl in 1:size(data_dict["EyeY"], 2)];
    # [deleteat!(xy[trl][dim], findall(isnan.(xy[trl][dim]))) for dim in [1,2], trl in 1:size(data_dict["EyeY"], 2)] # remove NaNs
    # sf = data_dict["samplingRate"]

    ### generate dataframe from ordered_column_name_pairs
    trial_list = DataFrame([el[1] for el in ordered_column_name_pairs], [Symbol(el[2]) for el in ordered_column_name_pairs])

    return trial_list
end

##################################################
#
# reads .mat file and generates raw_data stuct
#
##################################################
function mat2jdl(ifi, ofi)
    for (idx, mf_) in enumerate(ifi)

        ### read mat files of runs specified in paths_to_mat_files and store the in data_dicts array
        data_dict = matread(mf_)
        # @show data_dict
        #data_dict_ = deepcopy(data_dict)

        ### remove missed trials in data_dict
        missed_trials_idx = findall(data_dict["TrialList"][:, 23].==1)
        ks = ["TargetY", "VoltX", "TargetX", "EyeT", "TargetTime", "EyeX", "EyeY", "VoltY"]
        for k in ks
            data_dict[k] = data_dict[k][:, setdiff(1:end, missed_trials_idx)]
        end
        data_dict["TrialList"] = data_dict["TrialList"][setdiff(1:end, missed_trials_idx), :]

        ### remove trials with all NaN
        xy = [[data_dict["EyeX"][:,trl], data_dict["EyeY"][:,trl]] for trl in 1:size(data_dict["EyeY"], 2)];
        only_nans = [trl for trl in 1:size(xy, 1) if all(isnan.(xy[trl][1]))] # find all NaN trials
        for k in ks
            data_dict[k] = data_dict[k][:, setdiff(1:end, only_nans)]
        end
        data_dict["TrialList"] = data_dict["TrialList"][setdiff(1:end, only_nans), :];

        print(".mat file loaded: ", idx,"/",size(ifi)[1])

        ### convert trialList from mat file to julia table
        trial_list = trialList2DataFrame(data_dict)
        print(" | trial_list extracted")

        ## get indices of valid and invalid trials
        # valid_trials_idx = findall( trial_list[!, :Valid] )

        raw_data = data_of_run(
            "...",
            Date(mf_[end-26:end-17],"y-m-d"),
            DateTime(mf_[end-15:end-8],"H-M-S"),
            data_dict["samplingRate"],
            trial_list
        )

        ### store processed data to file
        save(joinpath(ofi, mf_[end-41:end-4]*".jld"), "raw_data", raw_data)
        println(" | .jdl saved")

        #return eye_dat
    end
end

end


# Clique types

##==============================================================================
## BayesTreeNodeData
##==============================================================================

"""
    $TYPEDEF

Cache messages being passed on the tree, one container per clique.

Notes
- See model 2 (?) on IIF #674
"""
mutable struct MessageBuffer
  # up receive message buffer (multiple children, multiple messages)
  upRx::Dict{Int, LikelihoodMessage}
  # down receive message buffer (one parent)
  downRx::Union{Nothing, LikelihoodMessage}
  # RESERVED up outgoing message buffer (one parent)
  upTx::Union{Nothing, LikelihoodMessage}
  # RESERVED down outgoing message buffer (multiple children but one message)
  downTx::Union{Nothing, LikelihoodMessage}
end
MessageBuffer() = MessageBuffer(Dict{Int, LikelihoodMessage}(), nothing, nothing, nothing)

"""
$(TYPEDEF)

Data structure for each clique in the Bayes (Junction) tree.
"""
mutable struct BayesTreeNodeData
  status::CliqStatus
  frontalIDs::Vector{Symbol}
  separatorIDs::Vector{Symbol}
  inmsgIDs::Vector{Symbol} # Int
  potIDs::Vector{Symbol} # Int # this is likely redundant TODO -- remove
  potentials::Vector{Symbol}
  partialpotential::Vector{Bool}

  dwnPotentials::Vector{Symbol}
  dwnPartialPotential::Vector{Bool}

  cliqAssocMat::Array{Bool,2}
  cliqMsgMat::Array{Bool,2}
  directvarIDs::Vector{Symbol}
  directFrtlMsgIDs::Vector{Symbol}
  msgskipIDs::Vector{Symbol}
  itervarIDs::Vector{Symbol}
  directPriorMsgIDs::Vector{Symbol}
  debug
  debugDwn

  allmarginalized::Bool
  initialized::Symbol
  upsolved::Bool
  downsolved::Bool
  isCliqReused::Bool             # holdover

  # JT Local messages saved for cache and debugging, see IIF #675
  messages::MessageBuffer
end

function BayesTreeNodeData(;status::CliqStatus=NULL,
                            frontalIDs=Symbol[],
                            separatorIDs=Symbol[],
                            inmsgIDs=Symbol[],
                            potIDs=Symbol[],
                            potentials=Symbol[],
                            partialpotential=Bool[],
                            dwnPotentials=Symbol[],
                            dwnPartialPotential=Bool[],
                            cliqAssocMat=Array{Bool}(undef, 0,0),
                            cliqMsgMat=Array{Bool}(undef, 0,0),
                            directvarIDs=Int[],
                            directFrtlMsgIDs=Int[],
                            msgskipIDs=Int[],
                            itervarIDs=Int[],
                            directPriorMsgIDs=Int[],
                            debug=nothing,
                            debugDwn=nothing,
                            allmarginalized=false,
                            initialized=:NULL,
                            upsolved=false,
                            downsolved=false,
                            isCliqReused=false,
                            messages = MessageBuffer()
                          )
  btnd = BayesTreeNodeData(status,
                        frontalIDs,
                        separatorIDs,
                        inmsgIDs,
                        potIDs,
                        potentials,
                        partialpotential,
                        dwnPotentials,
                        dwnPartialPotential,
                        cliqAssocMat,
                        cliqMsgMat,
                        directvarIDs,
                        directFrtlMsgIDs,
                        msgskipIDs,
                        itervarIDs,
                        directPriorMsgIDs,
                        debug,
                        debugDwn,
                        allmarginalized,
                        initialized,
                        upsolved,
                        downsolved,
                        isCliqReused,
                        messages  )
  #
  return btnd
end
#

function compare( c1::BayesTreeNodeData,
                  c2::BayesTreeNodeData;
                  skip::Vector{Symbol}=[] )
  #
  TP = true

  TP = TP && c1.frontalIDs == c2.frontalIDs
  TP = TP && c1.separatorIDs == c2.separatorIDs
  TP = TP && c1.inmsgIDs == c2.inmsgIDs
  TP = TP && c1.potIDs == c2.potIDs
  TP = TP && c1.potentials == c2.potentials
  TP = TP && c1.partialpotential == c2.partialpotential
  TP = TP && c1.dwnPotentials == c2.dwnPotentials
  TP = TP && c1.dwnPartialPotential == c2.dwnPartialPotential
  TP = TP && c1.cliqAssocMat == c2.cliqAssocMat
  TP = TP && c1.cliqMsgMat == c2.cliqMsgMat
  TP = TP && c1.directvarIDs == c2.directvarIDs
  TP = TP && c1.directFrtlMsgIDs == c2.directFrtlMsgIDs
  TP = TP && c1.msgskipIDs == c2.msgskipIDs
  TP = TP && c1.itervarIDs == c2.itervarIDs
  TP = TP && c1.directPriorMsgIDs == c2.directPriorMsgIDs
  TP = TP && c1.debug == c2.debug
  TP = TP && c1.debugDwn == c2.debugDwn
  TP = TP && c1.allmarginalized == c2.allmarginalized
  TP = TP && c1.initialized == c2.initialized
  TP = TP && c1.upsolved == c2.upsolved
  TP = TP && c1.downsolved == c2.downsolved
  TP = TP && c1.isCliqReused == c2.isCliqReused

  return TP
end

## Packed types for serialization

mutable struct PackedBayesTreeNodeData
  frontalIDs::Vector{Symbol}
  separatorIDs::Vector{Symbol}
  inmsgIDs::Vector{Symbol} # Int
  potIDs::Vector{Symbol} # Int # this is likely redundant TODO -- remove
  potentials::Vector{Symbol}
  partialpotential::Vector{Bool}
  dwnPotentials::Vector{Symbol}
  dwnPartialPotential::Vector{Bool}
  cliqAssocMat::Array{Bool,2}
  cliqMsgMat::Array{Bool,2}
  directvarIDs::Vector{Symbol} # Int
  directFrtlMsgIDs::Vector{Symbol} # Int
  msgskipIDs::Vector{Symbol} # Int
  itervarIDs::Vector{Symbol} # Int
  directPriorMsgIDs::Vector{Symbol} # Int
end

function convert(::Type{PackedBayesTreeNodeData}, btnd::BayesTreeNodeData)
  return PackedBayesTreeNodeData(
    btnd.frontalIDs,
    btnd.separatorIDs,
    btnd.inmsgIDs,
    btnd.potIDs,
    btnd.potentials,
    btnd.partialpotential,
    btnd.dwnPotentials,
    btnd.dwnPartialPotential,
    btnd.cliqAssocMat,
    btnd.cliqMsgMat,
    btnd.directvarIDs,
    btnd.directFrtlMsgIDs,
    btnd.msgskipIDs,
    btnd.itervarIDs,
    btnd.directPriorMsgIDs  )
end

function convert(::Type{BayesTreeNodeData}, pbtnd::PackedBayesTreeNodeData)
  btnd = BayesTreeNodeData()
    btnd.frontalIDs = pbtnd.frontalIDs
    btnd.separatorIDs = pbtnd.separatorIDs
    btnd.inmsgIDs = pbtnd.inmsgIDs
    btnd.potIDs = pbtnd.potIDs
    btnd.potentials = pbtnd.potentials
    btnd.partialpotential = pbtnd.partialpotential
    btnd.dwnPotentials = pbtnd.dwnPotentials
    btnd.dwnPartialPotential = pbtnd.dwnPartialPotential
    btnd.cliqAssocMat = pbtnd.cliqAssocMat
    btnd.cliqMsgMat = pbtnd.cliqMsgMat
    btnd.directvarIDs = pbtnd.directvarIDs
    btnd.directFrtlMsgIDs = pbtnd.directFrtlMsgIDs
    btnd.msgskipIDs = pbtnd.msgskipIDs
    btnd.itervarIDs = pbtnd.itervarIDs
    btnd.directPriorMsgIDs = pbtnd.directPriorMsgIDs
  return btnd
end

##==============================================================================
## Cliques
## TreeClique
##==============================================================================
struct CliqueId{T}
  value::T
end

Base.getindex(cId::CliqueId) = cId.value
Base.show(io::IO, ::MIME"text/plain", x::CliqueId) = print(io, x.value)
Base.show(io::IO, x::CliqueId) = print(io, x.value)

"""
    $(TYPEDEF)
Structure to store clique data
DEV NOTES: To replace TreeClique completely
    $(FIELDS)
"""
mutable struct TreeClique
  "Interger id unique within a tree with userId, robotId, sessionId"
  id::CliqueId{Int64} # not to be confused with the underlying index used by LightGraphs.jl, see issue #540
  "Data as `BayesTreeNodeData`"
  data::BayesTreeNodeData
  "Drawing attributes"
  attributes::Dict{String, Any}
  #solveInProgress #on a clique level a "solve in progress" might be very handy
end

getId(c::TreeClique) = c.id

function Base.getproperty(x::TreeClique,f::Symbol)
  if f == :index
    Base.depwarn("`TreeCliqe` field `index` is deprecated, use `id`", :getproperty)
    f = :id
  end
  getfield(x,f)
end

function Base.setproperty!(x::TreeClique, f::Symbol, val)
  if f == :index
    Base.depwarn("`TreeCliqe` field `index` is deprecated, use `id`", :setproperty!)
    f = :id
  end
  return setfield!(x, f, convert(fieldtype(typeof(x), f), val))
end

TreeClique(i::Int) = TreeClique(CliqueId(i), BayesTreeNodeData(), Dict{String,Any}())
TreeClique(id::CliqueId) = TreeClique(id, BayesTreeNodeData(), Dict{String,Any}())

DFG.getLabel(cliq::TreeClique) = cliq.attributes["label"]
function setLabel!(cliq::TreeClique, lbl::String)
  cliq.attributes["label"] = lbl
  lbl
end

## end Cliques


function SparseFieldVector(arg0::Field)
    return SparseFieldVector((Field,), arg0)
end

function SparseFieldVector(arg0::Field, arg1::Vector{FieldElement})
    return SparseFieldVector((Field, Vector{FieldElement}), arg0, arg1)
end

function SparseFieldVector(arg0::Field, arg1::jint)
    return SparseFieldVector((Field, jint), arg0, arg1)
end

function SparseFieldVector(arg0::Field, arg1::jint, arg2::jint)
    return SparseFieldVector((Field, jint, jint), arg0, arg1, arg2)
end

function SparseFieldVector(arg0::SparseFieldVector)
    return SparseFieldVector((SparseFieldVector,), arg0)
end

function add(obj::SparseFieldVector, arg0::FieldVector)
    return jcall(obj, "add", FieldVector, (FieldVector,), arg0)
end

function add(obj::SparseFieldVector, arg0::SparseFieldVector)
    return jcall(obj, "add", FieldVector, (SparseFieldVector,), arg0)
end

function append(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "append", FieldVector, (FieldElement,), arg0)
end

function append(obj::SparseFieldVector, arg0::FieldVector)
    return jcall(obj, "append", FieldVector, (FieldVector,), arg0)
end

function append(obj::SparseFieldVector, arg0::SparseFieldVector)
    return jcall(obj, "append", FieldVector, (SparseFieldVector,), arg0)
end

function copy(obj::SparseFieldVector)
    return jcall(obj, "copy", FieldVector, ())
end

function dot_product(obj::SparseFieldVector, arg0::FieldVector)
    return jcall(obj, "dotProduct", FieldElement, (FieldVector,), arg0)
end

function ebe_divide(obj::SparseFieldVector, arg0::FieldVector)
    return jcall(obj, "ebeDivide", FieldVector, (FieldVector,), arg0)
end

function ebe_multiply(obj::SparseFieldVector, arg0::FieldVector)
    return jcall(obj, "ebeMultiply", FieldVector, (FieldVector,), arg0)
end

function equals(obj::SparseFieldVector, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function get_dimension(obj::SparseFieldVector)
    return jcall(obj, "getDimension", jint, ())
end

function get_entry(obj::SparseFieldVector, arg0::jint)
    return jcall(obj, "getEntry", FieldElement, (jint,), arg0)
end

function get_field(obj::SparseFieldVector)
    return jcall(obj, "getField", Field, ())
end

function get_sub_vector(obj::SparseFieldVector, arg0::jint, arg1::jint)
    return jcall(obj, "getSubVector", FieldVector, (jint, jint), arg0, arg1)
end

function hash_code(obj::SparseFieldVector)
    return jcall(obj, "hashCode", jint, ())
end

function map_add(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "mapAdd", FieldVector, (FieldElement,), arg0)
end

function map_add_to_self(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "mapAddToSelf", FieldVector, (FieldElement,), arg0)
end

function map_divide(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "mapDivide", FieldVector, (FieldElement,), arg0)
end

function map_divide_to_self(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "mapDivideToSelf", FieldVector, (FieldElement,), arg0)
end

function map_inv(obj::SparseFieldVector)
    return jcall(obj, "mapInv", FieldVector, ())
end

function map_inv_to_self(obj::SparseFieldVector)
    return jcall(obj, "mapInvToSelf", FieldVector, ())
end

function map_multiply(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "mapMultiply", FieldVector, (FieldElement,), arg0)
end

function map_multiply_to_self(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "mapMultiplyToSelf", FieldVector, (FieldElement,), arg0)
end

function map_subtract(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "mapSubtract", FieldVector, (FieldElement,), arg0)
end

function map_subtract_to_self(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "mapSubtractToSelf", FieldVector, (FieldElement,), arg0)
end

function outer_product(obj::SparseFieldVector, arg0::FieldVector)
    return jcall(obj, "outerProduct", FieldMatrix, (FieldVector,), arg0)
end

function outer_product(obj::SparseFieldVector, arg0::SparseFieldVector)
    return jcall(obj, "outerProduct", FieldMatrix, (SparseFieldVector,), arg0)
end

function projection(obj::SparseFieldVector, arg0::FieldVector)
    return jcall(obj, "projection", FieldVector, (FieldVector,), arg0)
end

function set(obj::SparseFieldVector, arg0::FieldElement)
    return jcall(obj, "set", void, (FieldElement,), arg0)
end

function set_entry(obj::SparseFieldVector, arg0::jint, arg1::FieldElement)
    return jcall(obj, "setEntry", void, (jint, FieldElement), arg0, arg1)
end

function set_sub_vector(obj::SparseFieldVector, arg0::jint, arg1::FieldVector)
    return jcall(obj, "setSubVector", void, (jint, FieldVector), arg0, arg1)
end

function subtract(obj::SparseFieldVector, arg0::FieldVector)
    return jcall(obj, "subtract", FieldVector, (FieldVector,), arg0)
end

function subtract(obj::SparseFieldVector, arg0::SparseFieldVector)
    return jcall(obj, "subtract", SparseFieldVector, (SparseFieldVector,), arg0)
end

function to_array(obj::SparseFieldVector)
    return jcall(obj, "toArray", Vector{FieldElement}, ())
end

function walk_in_default_order(obj::SparseFieldVector, arg0::FieldVectorChangingVisitor)
    return jcall(obj, "walkInDefaultOrder", FieldElement, (FieldVectorChangingVisitor,), arg0)
end

function walk_in_default_order(obj::SparseFieldVector, arg0::FieldVectorChangingVisitor, arg1::jint, arg2::jint)
    return jcall(obj, "walkInDefaultOrder", FieldElement, (FieldVectorChangingVisitor, jint, jint), arg0, arg1, arg2)
end

function walk_in_default_order(obj::SparseFieldVector, arg0::FieldVectorPreservingVisitor)
    return jcall(obj, "walkInDefaultOrder", FieldElement, (FieldVectorPreservingVisitor,), arg0)
end

function walk_in_default_order(obj::SparseFieldVector, arg0::FieldVectorPreservingVisitor, arg1::jint, arg2::jint)
    return jcall(obj, "walkInDefaultOrder", FieldElement, (FieldVectorPreservingVisitor, jint, jint), arg0, arg1, arg2)
end

function walk_in_optimized_order(obj::SparseFieldVector, arg0::FieldVectorChangingVisitor)
    return jcall(obj, "walkInOptimizedOrder", FieldElement, (FieldVectorChangingVisitor,), arg0)
end

function walk_in_optimized_order(obj::SparseFieldVector, arg0::FieldVectorChangingVisitor, arg1::jint, arg2::jint)
    return jcall(obj, "walkInOptimizedOrder", FieldElement, (FieldVectorChangingVisitor, jint, jint), arg0, arg1, arg2)
end

function walk_in_optimized_order(obj::SparseFieldVector, arg0::FieldVectorPreservingVisitor)
    return jcall(obj, "walkInOptimizedOrder", FieldElement, (FieldVectorPreservingVisitor,), arg0)
end

function walk_in_optimized_order(obj::SparseFieldVector, arg0::FieldVectorPreservingVisitor, arg1::jint, arg2::jint)
    return jcall(obj, "walkInOptimizedOrder", FieldElement, (FieldVectorPreservingVisitor, jint, jint), arg0, arg1, arg2)
end


################################################################################
# MIT License
#
# Copyright (c) 2021 Hajime Nakagami<nakagami@gmail.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
################################################################################

mutable struct Arc4
    state::Vector{UInt8}
    x::UInt64
    y::UInt64

    function Arc4(key::Vector{UInt8})
        state::Vector{UInt8} = []
        for i in 0:255
            push!(state, i)
        end
        @assert length(state) == 256

        index1::UInt64 = 0
        index2::UInt64 = 0

        for i in 0:255
            index2 = (key[index1+1] + state[i+1] + index2) % 256
            (state[i+1], state[index2+1]) = (state[index2+1], state[i+1])
            index1 = (index1 + 1) % length(key)
        end

        new(state, 0, 0)
    end
end

function translate(arc4::Arc4, plain::Vector{UInt8})
    state = arc4.state
    enc::Vector{UInt8} = []
    for i in 1:length(plain)
        arc4.x = (arc4.x + 1) %256
        arc4.y = (arc4.y + state[arc4.x+1]) % 256
        state[arc4.x+1], state[arc4.y+1] = state[arc4.y+1], state[arc4.x+1]
        push!(enc, xor(plain[i], state[((state[arc4.x+1]+state[arc4.y+1]) % 256) + 1]))
    end
    enc
end


# Container data structure (structure-of-arrays style) for DG elements
# TODO: Taal refactor, remove NVARS, POLYDEG?
mutable struct ElementContainer2D{RealT<:Real, uEltype<:Real, NVARS, POLYDEG} <: AbstractContainer
  inverse_jacobian::Vector{RealT}        # [elements]
  node_coordinates::Array{RealT, 4}      # [orientation, i, j, elements]
  surface_flux_values::Array{uEltype, 4} # [variables, i, direction, elements]
  cell_ids::Vector{Int}                  # [elements]
  # internal `resize!`able storage
  _node_coordinates::Vector{RealT}
  _surface_flux_values::Vector{uEltype}
end

nvariables(::ElementContainer2D{RealT, uEltype, NVARS, POLYDEG}) where {RealT, uEltype, NVARS, POLYDEG} = NVARS
polydeg(::ElementContainer2D{RealT, uEltype, NVARS, POLYDEG}) where {RealT, uEltype, NVARS, POLYDEG} = POLYDEG
Base.eltype(::ElementContainer2D{RealT, uEltype}) where {RealT, uEltype} = uEltype

# Only one-dimensional `Array`s are `resize!`able in Julia.
# Hence, we use `Vector`s as internal storage and `resize!`
# them whenever needed. Then, we reuse the same memory by
# `unsafe_wrap`ping multi-dimensional `Array`s around the
# internal storage.
function Base.resize!(elements::ElementContainer2D, capacity)
  n_nodes = polydeg(elements) + 1
  nvars = nvariables(elements)
  @unpack _node_coordinates, _surface_flux_values,
          inverse_jacobian, cell_ids = elements

  resize!(inverse_jacobian, capacity)

  resize!(_node_coordinates, 2 * n_nodes * n_nodes * capacity)
  elements.node_coordinates = unsafe_wrap(Array, pointer(_node_coordinates),
                                          (2, n_nodes, n_nodes, capacity))

  resize!(_surface_flux_values, nvars * n_nodes * 2 * 2 * capacity)
  elements.surface_flux_values = unsafe_wrap(Array, pointer(_surface_flux_values),
                                             (nvars, n_nodes, 2 * 2, capacity))

  resize!(cell_ids, capacity)

  return nothing
end

function ElementContainer2D{RealT, uEltype, NVARS, POLYDEG}(capacity::Integer) where {RealT<:Real, uEltype<:Real, NVARS, POLYDEG}
  n_nodes = POLYDEG + 1
  nan_RealT = convert(RealT, NaN)
  nan_uEltype = convert(uEltype, NaN)

  # Initialize fields with defaults
  inverse_jacobian = fill(nan_RealT, capacity)

  _node_coordinates = fill(nan_RealT, 2 * n_nodes * n_nodes * capacity)
  node_coordinates = unsafe_wrap(Array, pointer(_node_coordinates),
                                 (2, n_nodes, n_nodes, capacity))

  _surface_flux_values = fill(nan_uEltype, NVARS * n_nodes * 2 * 2 * capacity)
  surface_flux_values = unsafe_wrap(Array, pointer(_surface_flux_values),
                                    (NVARS, n_nodes, 2 * 2, capacity))

  cell_ids = fill(typemin(Int), capacity)

  return ElementContainer2D{RealT, uEltype, NVARS, POLYDEG}(
    inverse_jacobian, node_coordinates, surface_flux_values, cell_ids,
    _node_coordinates, _surface_flux_values)
end

# Return number of elements
@inline nelements(elements::ElementContainer2D) = length(elements.cell_ids)
# TODO: Taal performance, 1:nelements(elements) vs. Base.OneTo(nelements(elements))
@inline eachelement(elements::ElementContainer2D) = Base.OneTo(nelements(elements))

# Create element container and initialize element data
function init_elements(cell_ids, mesh::TreeMesh2D,
                       equations::AbstractEquations{2, NVARS},
                       basis::LobattoLegendreBasis{T, NNODES}, ::Type{RealT}, ::Type{uEltype}) where {RealT<:Real, uEltype<:Real, NVARS, T, NNODES}
  # Initialize container
  n_elements = length(cell_ids)
  elements = ElementContainer2D{RealT, uEltype, NVARS, NNODES-1}(n_elements)

  init_elements!(elements, cell_ids, mesh, basis.nodes)
  return elements
end

function init_elements!(elements, cell_ids, mesh::TreeMesh2D, nodes)
  n_nodes = length(nodes)

  # Store cell ids
  elements.cell_ids .= cell_ids

  # Calculate inverse Jacobian and node coordinates
  for element in eachelement(elements)
    # Get cell id
    cell_id = cell_ids[element]

    # Get cell length
    dx = length_at_cell(mesh.tree, cell_id)

    # Calculate inverse Jacobian as 1/(h/2)
    elements.inverse_jacobian[element] = 2/dx

    # Calculate node coordinates
    for j in 1:n_nodes, i in 1:n_nodes
      elements.node_coordinates[1, i, j, element] = (
          mesh.tree.coordinates[1, cell_id] + dx/2 * nodes[i])
      elements.node_coordinates[2, i, j, element] = (
          mesh.tree.coordinates[2, cell_id] + dx/2 * nodes[j])
    end
  end

  return elements
end

# Container data structure (structure-of-arrays style) for DG interfaces
# TODO: Taal refactor, remove NVARS, POLYDEG?
mutable struct InterfaceContainer2D{uEltype<:Real, NVARS, POLYDEG} <: AbstractContainer
  u::Array{uEltype, 4}      # [leftright, variables, i, interfaces]
  neighbor_ids::Matrix{Int} # [leftright, interfaces]
  orientations::Vector{Int} # [interfaces]
  # internal `resize!`able storage
  _u::Vector{uEltype}
  _neighbor_ids::Vector{Int}
end

nvariables(::InterfaceContainer2D{uEltype, NVARS, POLYDEG}) where {uEltype, NVARS, POLYDEG} = NVARS
polydeg(::InterfaceContainer2D{uEltype, NVARS, POLYDEG}) where {uEltype, NVARS, POLYDEG} = POLYDEG
Base.eltype(::InterfaceContainer2D{uEltype}) where {uEltype} = uEltype

# See explanation of Base.resize! for the element container
function Base.resize!(interfaces::InterfaceContainer2D, capacity)
  n_nodes = polydeg(interfaces) + 1
  nvars = nvariables(interfaces)
  @unpack _u, _neighbor_ids, orientations = interfaces

  resize!(_u, 2 * nvars * n_nodes * capacity)
  interfaces.u = unsafe_wrap(Array, pointer(_u),
                             (2, nvars, n_nodes, capacity))

  resize!(_neighbor_ids, 2 * capacity)
  interfaces.neighbor_ids = unsafe_wrap(Array, pointer(_neighbor_ids ),
                                        (2, capacity))

  resize!(orientations, capacity)

  return nothing
end

function InterfaceContainer2D{uEltype, NVARS, POLYDEG}(capacity::Integer) where {uEltype<:Real, NVARS, POLYDEG}
  n_nodes = POLYDEG + 1
  nan = convert(uEltype, NaN)

  # Initialize fields with defaults
  _u = fill(nan, 2 * NVARS * n_nodes * capacity)
  u = unsafe_wrap(Array, pointer(_u),
                  (2, NVARS, n_nodes, capacity))

  _neighbor_ids = fill(typemin(Int), 2 * capacity)
  neighbor_ids = unsafe_wrap(Array, pointer(_neighbor_ids ),
                             (2, capacity))

  orientations = fill(typemin(Int), capacity)

  return InterfaceContainer2D{uEltype, NVARS, POLYDEG}(
    u, neighbor_ids, orientations,
    _u, _neighbor_ids)
end

# Return number of interfaces
@inline ninterfaces(interfaces::InterfaceContainer2D) = length(interfaces.orientations)

# Create interface container and initialize interface data in `elements`.
function init_interfaces(cell_ids, mesh::TreeMesh2D,
                         elements::ElementContainer2D{RealT, uEltype, NVARS, POLYDEG}) where {RealT<:Real, uEltype<:Real, NVARS, POLYDEG}
  # Initialize container
  n_interfaces = count_required_interfaces(mesh, cell_ids)
  interfaces = InterfaceContainer2D{uEltype, NVARS, POLYDEG}(n_interfaces)

  # Connect elements with interfaces
  init_interfaces!(interfaces, elements, mesh)
  return interfaces
end

# Count the number of interfaces that need to be created
function count_required_interfaces(mesh::TreeMesh2D, cell_ids)
  count = 0

  # Iterate over all cells
  for cell_id in cell_ids
    for direction in eachdirection(mesh.tree)
      # Only count interfaces in positive direction to avoid double counting
      if direction % 2 == 1
        continue
      end

      # If no neighbor exists, current cell is small or at boundary and thus we need a mortar
      if !has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # Skip if neighbor has children
      neighbor_cell_id = mesh.tree.neighbor_ids[direction, cell_id]
      if has_children(mesh.tree, neighbor_cell_id)
        continue
      end

      # Skip if neighbor is on different rank -> create MPI interface instead
      if mpi_isparallel() && !is_own_cell(mesh.tree, neighbor_cell_id)
        continue
      end

      count += 1
    end
  end

  return count
end

# Initialize connectivity between elements and interfaces
function init_interfaces!(interfaces, elements, mesh::TreeMesh2D)
  # Construct cell -> element mapping for easier algorithm implementation
  tree = mesh.tree
  c2e = zeros(Int, length(tree))
  for element in eachelement(elements)
    c2e[elements.cell_ids[element]] = element
  end

  # Reset interface count
  count = 0

  # Iterate over all elements to find neighbors and to connect via interfaces
  for element in eachelement(elements)
    # Get cell id
    cell_id = elements.cell_ids[element]

    # Loop over directions
    for direction in eachdirection(mesh.tree)
      # Only create interfaces in positive direction
      if direction % 2 == 1
        continue
      end

      # If no neighbor exists, current cell is small and thus we need a mortar
      if !has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # Skip if neighbor has children
      neighbor_cell_id = mesh.tree.neighbor_ids[direction, cell_id]
      if has_children(mesh.tree, neighbor_cell_id)
        continue
      end

      # Skip if neighbor is on different rank -> create MPI interface instead
      if mpi_isparallel() && !is_own_cell(mesh.tree, neighbor_cell_id)
        continue
      end

      # Create interface between elements (1 -> "left" of interface, 2 -> "right" of interface)
      count += 1
      interfaces.neighbor_ids[2, count] = c2e[neighbor_cell_id]
      interfaces.neighbor_ids[1, count] = element

      # Set orientation (x -> 1, y -> 2)
      interfaces.orientations[count] = div(direction, 2)
    end
  end

  @assert count == ninterfaces(interfaces) ("Actual interface count ($count) does not match " *
                                            "expectations $(ninterfaces(interfaces))")
end

# Container data structure (structure-of-arrays style) for DG boundaries
# TODO: Taal refactor, remove NVARS, POLYDEG?
mutable struct BoundaryContainer2D{RealT<:Real, uEltype<:Real, NVARS, POLYDEG} <: AbstractContainer
  u::Array{uEltype, 4}              # [leftright, variables, i, boundaries]
  neighbor_ids::Vector{Int}         # [boundaries]
  orientations::Vector{Int}         # [boundaries]
  neighbor_sides::Vector{Int}       # [boundaries]
  node_coordinates::Array{RealT, 3} # [orientation, i, elements]
  n_boundaries_per_direction::SVector{4, Int} # [direction]
  # internal `resize!`able storage
  _u::Vector{uEltype}
  _node_coordinates::Vector{RealT}
end

nvariables(::BoundaryContainer2D{RealT, uEltype, NVARS, POLYDEG}) where {RealT, uEltype, NVARS, POLYDEG} = NVARS
polydeg(::BoundaryContainer2D{RealT, uEltype, NVARS, POLYDEG}) where {RealT, uEltype, NVARS, POLYDEG} = POLYDEG
Base.eltype(::BoundaryContainer2D{RealT, uEltype}) where {RealT, uEltype} = uEltype

# See explanation of Base.resize! for the element container
function Base.resize!(boundaries::BoundaryContainer2D, capacity)
  n_nodes = polydeg(boundaries) + 1
  nvars = nvariables(boundaries)
  @unpack _u, _node_coordinates,
          neighbor_ids, orientations, neighbor_sides = boundaries

  resize!(_u, 2 * nvars * n_nodes * capacity)
  boundaries.u = unsafe_wrap(Array, pointer(_u),
                             (2, nvars, n_nodes, capacity))

  resize!(_node_coordinates, 2 * n_nodes * capacity)
  boundaries.node_coordinates = unsafe_wrap(Array, pointer(_node_coordinates),
                                            (2, n_nodes, capacity))

  resize!(neighbor_ids, capacity)

  resize!(orientations, capacity)

  resize!(neighbor_sides, capacity)

  return nothing
end

function BoundaryContainer2D{RealT, uEltype, NVARS, POLYDEG}(capacity::Integer) where {RealT<:Real, uEltype<:Real, NVARS, POLYDEG}
  n_nodes = POLYDEG + 1
  nan_RealT = convert(RealT, NaN)
  nan_uEltype = convert(uEltype, NaN)

  # Initialize fields with defaults
  _u = fill(nan_uEltype, 2 * NVARS * n_nodes * capacity)
  u = unsafe_wrap(Array, pointer(_u),
                  (2, NVARS, n_nodes, capacity))

  neighbor_ids = fill(typemin(Int), capacity)

  orientations = fill(typemin(Int), capacity)

  neighbor_sides = fill(typemin(Int), capacity)

  _node_coordinates = fill(nan_RealT, 2 * n_nodes * capacity)
  node_coordinates = unsafe_wrap(Array, pointer(_node_coordinates),
                                 (2, n_nodes, capacity))

  n_boundaries_per_direction = SVector(0, 0, 0, 0)

  return BoundaryContainer2D{RealT, uEltype, NVARS, POLYDEG}(
    u, neighbor_ids, orientations, neighbor_sides,
    node_coordinates, n_boundaries_per_direction,
    _u, _node_coordinates)
end

# Return number of boundaries
@inline nboundaries(boundaries::BoundaryContainer2D) = length(boundaries.orientations)

# Create boundaries container and initialize boundary data in `elements`.
function init_boundaries(cell_ids, mesh::TreeMesh2D,
                         elements::ElementContainer2D{RealT, uEltype, NVARS, POLYDEG}) where {RealT<:Real, uEltype<:Real, NVARS, POLYDEG}
  # Initialize container
  n_boundaries = count_required_boundaries(mesh, cell_ids)
  boundaries = BoundaryContainer2D{RealT, uEltype, NVARS, POLYDEG}(n_boundaries)

  # Connect elements with boundaries
  init_boundaries!(boundaries, elements, mesh)
  return boundaries
end

# Count the number of boundaries that need to be created
function count_required_boundaries(mesh::TreeMesh2D, cell_ids)
  count = 0

  # Iterate over all cells
  for cell_id in cell_ids
    for direction in eachdirection(mesh.tree)
      # If neighbor exists, current cell is not at a boundary
      if has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # If coarse neighbor exists, current cell is not at a boundary
      if has_coarse_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # No neighbor exists in this direction -> must be a boundary
      count += 1
    end
  end

  return count
end

# Initialize connectivity between elements and boundaries
function init_boundaries!(boundaries, elements, mesh::TreeMesh2D)
  # Reset boundaries count
  count = 0

  # Initialize boundary counts
  counts_per_direction = MVector(0, 0, 0, 0)

  # OBS! Iterate over directions first, then over elements, and count boundaries in each direction
  # Rationale: This way the boundaries are internally sorted by the directions -x, +x, -y etc.,
  #            obviating the need to store the boundary condition to be applied explicitly.
  # Loop over directions
  for direction in eachdirection(mesh.tree)
    # Iterate over all elements to find missing neighbors and to connect to boundaries
    for element in eachelement(elements)
      # Get cell id
      cell_id = elements.cell_ids[element]

      # If neighbor exists, current cell is not at a boundary
      if has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # If coarse neighbor exists, current cell is not at a boundary
      if has_coarse_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # Create boundary
      count += 1
      counts_per_direction[direction] += 1

      # Set neighbor element id
      boundaries.neighbor_ids[count] = element

      # Set neighbor side, which denotes the direction (1 -> negative, 2 -> positive) of the element
      if direction in (2, 4)
        boundaries.neighbor_sides[count] = 1
      else
        boundaries.neighbor_sides[count] = 2
      end

      # Set orientation (x -> 1, y -> 2)
      if direction in (1, 2)
        boundaries.orientations[count] = 1
      else
        boundaries.orientations[count] = 2
      end

      # Store node coordinates
      enc = elements.node_coordinates
      if direction == 1 # -x direction
        boundaries.node_coordinates[:, :, count] .= enc[:, 1,   :,   element]
      elseif direction == 2 # +x direction
        boundaries.node_coordinates[:, :, count] .= enc[:, end, :,   element]
      elseif direction == 3 # -y direction
        boundaries.node_coordinates[:, :, count] .= enc[:, :,   1,   element]
      elseif direction == 4 # +y direction
        boundaries.node_coordinates[:, :, count] .= enc[:, :,   end, element]
      else
        error("should not happen")
      end
    end
  end

  @assert count == nboundaries(boundaries) ("Actual boundaries count ($count) does not match " *
                                            "expectations $(nboundaries(boundaries))")
  @assert sum(counts_per_direction) == count

  boundaries.n_boundaries_per_direction = SVector(counts_per_direction)

  return boundaries.n_boundaries_per_direction
end

# Container data structure (structure-of-arrays style) for DG L2 mortars
# Positions/directions for orientations = 1, large_sides = 2:
# mortar is orthogonal to x-axis, large side is in positive coordinate direction wrt mortar
#           |    |
# upper = 2 |    |
#           |    |
#                | 3 = large side
#           |    |
# lower = 1 |    |
#           |    |
# TODO: Taal refactor, remove NVARS, POLYDEG?
# TODO: Taal refactor, mutable struct or resize! for AMR?
mutable struct L2MortarContainer2D{uEltype<:Real, NVARS, POLYDEG} <: AbstractContainer
  u_upper::Array{uEltype, 4}  # [leftright, variables, i, mortars]
  u_lower::Array{uEltype, 4}  # [leftright, variables, i, mortars]
  neighbor_ids::Matrix{Int}   # [position, mortars]
  # Large sides: left -> 1, right -> 2
  large_sides::Vector{Int}  # [mortars]
  orientations::Vector{Int} # [mortars]
  # internal `resize!`able storage
  _u_upper::Vector{uEltype}
  _u_lower::Vector{uEltype}
  _neighbor_ids::Vector{Int}
end

nvariables(::L2MortarContainer2D{uEltype, NVARS, POLYDEG}) where {uEltype, NVARS, POLYDEG} = NVARS
polydeg(::L2MortarContainer2D{uEltype, NVARS, POLYDEG}) where {uEltype, NVARS, POLYDEG} = POLYDEG
Base.eltype(::L2MortarContainer2D{uEltype}) where {uEltype} = uEltype

# See explanation of Base.resize! for the element container
function Base.resize!(mortars::L2MortarContainer2D, capacity)
  n_nodes = polydeg(mortars) + 1
  nvars = nvariables(mortars)
  @unpack _u_upper, _u_lower, _neighbor_ids,
          large_sides, orientations = mortars

  resize!(_u_upper, 2 * nvars * n_nodes * capacity)
  mortars.u_upper = unsafe_wrap(Array, pointer(_u_upper),
                                (2, nvars, n_nodes, capacity))

  resize!(_u_lower, 2 * nvars * n_nodes * capacity)
  mortars.u_lower = unsafe_wrap(Array, pointer(_u_lower),
                                (2, nvars, n_nodes, capacity))

  resize!(_neighbor_ids, 3 * capacity)
  mortars.neighbor_ids = unsafe_wrap(Array, pointer(_neighbor_ids ),
                                        (3, capacity))

  resize!(large_sides, capacity)

  resize!(orientations, capacity)

  return nothing
end

function L2MortarContainer2D{uEltype, NVARS, POLYDEG}(capacity::Integer) where {uEltype<:Real, NVARS, POLYDEG}
  n_nodes = POLYDEG + 1
  nan = convert(uEltype, NaN)

  # Initialize fields with defaults
  _u_upper = fill(nan, 2 * NVARS * n_nodes * capacity)
  u_upper = unsafe_wrap(Array, pointer(_u_upper),
                        (2, NVARS, n_nodes, capacity))

  _u_lower = fill(nan, 2 * NVARS * n_nodes * capacity)
  u_lower = unsafe_wrap(Array, pointer(_u_lower),
                        (2, NVARS, n_nodes, capacity))

  _neighbor_ids = fill(typemin(Int), 3 * capacity)
  neighbor_ids = unsafe_wrap(Array, pointer(_neighbor_ids),
                             (3, capacity))

  large_sides  = fill(typemin(Int), capacity)

  orientations = fill(typemin(Int), capacity)

  return L2MortarContainer2D{uEltype, NVARS, POLYDEG}(
    u_upper, u_lower, neighbor_ids, large_sides, orientations,
    _u_upper, _u_lower, _neighbor_ids)
end

# Return number of L2 mortars
@inline nmortars(l2mortars::L2MortarContainer2D) = length(l2mortars.orientations)

# Allow printing container contents
function Base.show(io::IO, ::MIME"text/plain", c::L2MortarContainer2D)
  @nospecialize c # reduce precompilation time

  println(io, '*'^20)
  for idx in CartesianIndices(c.u_upper)
    println(io, "c.u_upper[$idx] = $(c.u_upper[idx])")
  end
  for idx in CartesianIndices(c.u_lower)
    println(io, "c.u_lower[$idx] = $(c.u_lower[idx])")
  end
  println(io, "transpose(c.neighbor_ids) = $(transpose(c.neighbor_ids))")
  println(io, "c.large_sides = $(c.large_sides)")
  println(io, "c.orientations = $(c.orientations)")
  print(io,   '*'^20)
end

# Create mortar container and initialize mortar data in `elements`.
function init_mortars(cell_ids, mesh::TreeMesh2D,
                      elements::ElementContainer2D{RealT, uEltype, NVARS, POLYDEG},
                      ::LobattoLegendreMortarL2) where {RealT<:Real, uEltype<:Real, NVARS, POLYDEG}
  # Initialize containers
  n_mortars = count_required_mortars(mesh, cell_ids)
  mortars = L2MortarContainer2D{uEltype, NVARS, POLYDEG}(n_mortars)

  # Connect elements with mortars
  init_mortars!(mortars, elements, mesh)
  return mortars
end

# Count the number of mortars that need to be created
function count_required_mortars(mesh::TreeMesh2D, cell_ids)
  count = 0

  # Iterate over all cells and count mortars from perspective of coarse cells
  for cell_id in cell_ids
    for direction in eachdirection(mesh.tree)
      # If no neighbor exists, cell is small with large neighbor or at boundary -> do nothing
      if !has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # If neighbor has no children, this is a conforming interface -> do nothing
      neighbor_id = mesh.tree.neighbor_ids[direction, cell_id]
      if !has_children(mesh.tree, neighbor_id)
        continue
      end

      count +=1
    end
  end

  return count
end

# Initialize connectivity between elements and mortars
function init_mortars!(mortars, elements, mesh::TreeMesh2D)
  # Construct cell -> element mapping for easier algorithm implementation
  tree = mesh.tree
  c2e = zeros(Int, length(tree))
  for element in eachelement(elements)
    c2e[elements.cell_ids[element]] = element
  end

  # Reset interface count
  count = 0

  # Iterate over all elements to find neighbors and to connect via interfaces
  for element in eachelement(elements)
    # Get cell id
    cell_id = elements.cell_ids[element]

    for direction in eachdirection(mesh.tree)
      # If no neighbor exists, cell is small with large neighbor -> do nothing
      if !has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # If neighbor has no children, this is a conforming interface -> do nothing
      neighbor_cell_id = mesh.tree.neighbor_ids[direction, cell_id]
      if !has_children(mesh.tree, neighbor_cell_id)
        continue
      end

      # Create mortar between elements:
      # 1 -> small element in negative coordinate direction
      # 2 -> small element in positive coordinate direction
      # 3 -> large element
      count += 1
      mortars.neighbor_ids[3, count] = element
      if direction == 1
        mortars.neighbor_ids[1, count] = c2e[mesh.tree.child_ids[2, neighbor_cell_id]]
        mortars.neighbor_ids[2, count] = c2e[mesh.tree.child_ids[4, neighbor_cell_id]]
      elseif direction == 2
        mortars.neighbor_ids[1, count] = c2e[mesh.tree.child_ids[1, neighbor_cell_id]]
        mortars.neighbor_ids[2, count] = c2e[mesh.tree.child_ids[3, neighbor_cell_id]]
      elseif direction == 3
        mortars.neighbor_ids[1, count] = c2e[mesh.tree.child_ids[3, neighbor_cell_id]]
        mortars.neighbor_ids[2, count] = c2e[mesh.tree.child_ids[4, neighbor_cell_id]]
      elseif direction == 4
        mortars.neighbor_ids[1, count] = c2e[mesh.tree.child_ids[1, neighbor_cell_id]]
        mortars.neighbor_ids[2, count] = c2e[mesh.tree.child_ids[2, neighbor_cell_id]]
      else
        error("should not happen")
      end

      # Set large side, which denotes the direction (1 -> negative, 2 -> positive) of the large side
      if direction in [2, 4]
        mortars.large_sides[count] = 1
      else
        mortars.large_sides[count] = 2
      end

      # Set orientation (x -> 1, y -> 2)
      if direction in [1, 2]
        mortars.orientations[count] = 1
      else
        mortars.orientations[count] = 2
      end
    end
  end

  @assert count == nmortars(mortars) ("Actual mortar count ($count) does not match " *
                                      "expectations $(nmortars(mortars))")
end

# Container data structure (structure-of-arrays style) for DG MPI interfaces
mutable struct MPIInterfaceContainer2D{uEltype<:Real, NVARS, POLYDEG} <: AbstractContainer
  u::Array{uEltype, 4}           # [leftright, variables, i, interfaces]
  local_element_ids::Vector{Int} # [interfaces]
  orientations::Vector{Int}      # [interfaces]
  remote_sides::Vector{Int}      # [interfaces]
  # internal `resize!`able storage
  _u::Vector{uEltype}
end

nvariables(::MPIInterfaceContainer2D{uEltype, NVARS, POLYDEG}) where {uEltype, NVARS, POLYDEG} = NVARS
polydeg(::MPIInterfaceContainer2D{uEltype, NVARS, POLYDEG}) where {uEltype, NVARS, POLYDEG} = POLYDEG
Base.eltype(::MPIInterfaceContainer2D{uEltype}) where {uEltype} = uEltype

# See explanation of Base.resize! for the element container
function Base.resize!(mpi_interfaces::MPIInterfaceContainer2D, capacity)
  n_nodes = polydeg(mpi_interfaces) + 1
  nvars = nvariables(mpi_interfaces)
  @unpack _u, local_element_ids, orientations, remote_sides = mpi_interfaces

  resize!(_u, 2 * nvars * n_nodes * capacity)
  mpi_interfaces.u = unsafe_wrap(Array, pointer(_u),
                                 (2, nvars, n_nodes, capacity))

  resize!(local_element_ids, capacity)

  resize!(orientations, capacity)

  resize!(remote_sides, capacity)

  return nothing
end

function MPIInterfaceContainer2D{uEltype, NVARS, POLYDEG}(capacity::Integer) where {uEltype<:Real, NVARS, POLYDEG}
  n_nodes = POLYDEG + 1
  nan = convert(uEltype, NaN)

  # Initialize fields with defaults
  _u = fill(nan, 2 * NVARS * n_nodes * capacity)
  u = unsafe_wrap(Array, pointer(_u),
                  (2, NVARS, n_nodes, capacity))

  local_element_ids = fill(typemin(Int), capacity)

  orientations = fill(typemin(Int), capacity)

  remote_sides = fill(typemin(Int), capacity)

  return MPIInterfaceContainer2D{uEltype, NVARS, POLYDEG}(
    u, local_element_ids, orientations, remote_sides,
    _u)
end

# TODO: Taal, rename to ninterfaces?
# Return number of interfaces
nmpiinterfaces(mpi_interfaces::MPIInterfaceContainer2D) = length(mpi_interfaces.orientations)

# Create MPI interface container and initialize MPI interface data in `elements`.
function init_mpi_interfaces(cell_ids, mesh::TreeMesh2D,
                             elements::ElementContainer2D{RealT, uEltype, NVARS, POLYDEG}) where {RealT<:Real, uEltype<:Real, NVARS, POLYDEG}
  # Initialize container
  n_mpi_interfaces = count_required_mpi_interfaces(mesh, cell_ids)
  mpi_interfaces = MPIInterfaceContainer2D{uEltype, NVARS, POLYDEG}(n_mpi_interfaces)

  # Connect elements with interfaces
  init_mpi_interfaces!(mpi_interfaces, elements, mesh)
  return mpi_interfaces
end

# Count the number of MPI interfaces that need to be created
function count_required_mpi_interfaces(mesh::TreeMesh2D, cell_ids)
  count = 0

  # Iterate over all cells
  for cell_id in cell_ids
    for direction in eachdirection(mesh.tree)
      # If no neighbor exists, current cell is small or at boundary and thus we need a mortar
      if !has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # Skip if neighbor has children
      neighbor_cell_id = mesh.tree.neighbor_ids[direction, cell_id]
      if has_children(mesh.tree, neighbor_cell_id)
        continue
      end

      # Skip if neighbor is on this rank -> create regular interface instead
      if mpi_isparallel() && is_own_cell(mesh.tree, neighbor_cell_id)
        continue
      end

      count += 1
    end
  end

  return count
end

# Initialize connectivity between elements and interfaces
function init_mpi_interfaces!(mpi_interfaces, elements, mesh::TreeMesh2D)
  # Reset interface count
  count = 0

  # Iterate over all elements to find neighbors and to connect via mpi_interfaces
  for element in eachelement(elements)
    # Get cell id
    cell_id = elements.cell_ids[element]

    # Loop over directions
    for direction in eachdirection(mesh.tree)
      # If no neighbor exists, current cell is small and thus we need a mortar
      if !has_neighbor(mesh.tree, cell_id, direction)
        continue
      end

      # Skip if neighbor has children
      neighbor_cell_id = mesh.tree.neighbor_ids[direction, cell_id]
      if has_children(mesh.tree, neighbor_cell_id)
        continue
      end

      # Skip if neighbor is on this MPI rank -> create regular interface instead
      if mpi_isparallel() && is_own_cell(mesh.tree, neighbor_cell_id)
        continue
      end

      # Create interface between elements
      count += 1
      mpi_interfaces.local_element_ids[count] = element

      if direction in (2, 4) # element is "left" of interface, remote cell is "right" of interface
        mpi_interfaces.remote_sides[count] = 2
      else
        mpi_interfaces.remote_sides[count] = 1
      end

      # Set orientation (x -> 1, y -> 2)
      if direction in (1, 2) # x-direction
        mpi_interfaces.orientations[count] = 1
      else # y-direction
        mpi_interfaces.orientations[count] = 2
      end
    end
  end

  @assert count == nmpiinterfaces(mpi_interfaces) ("Actual interface count ($count) does not match "
                                                   * "expectations $(nmpiinterfaces(mpi_interfaces))")
end


# Unit tests for JUDI Geometry structure
# Philipp Witte (pwitte.slim@gmail.com)
# May 2018
#

using JUDI.TimeModeling, SeisIO, Test, LinearAlgebra

@testset "Geometry Unit Test" begin

    # Number of sources
    nsrc = 2

    # Constructor if nt is not passed
    xsrc = convertToCell(range(100f0, stop=1100f0, length=nsrc))
    ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
    zsrc = convertToCell(range(20f0, stop=20f0, length=nsrc))

    geometry =  Geometry(xsrc, ysrc, zsrc; dt=2f0, t=1000f0)

    @test isequal(typeof(geometry.xloc), Array{Any, 1})
    @test isequal(typeof(geometry.yloc), Array{Any, 1})
    @test isequal(typeof(geometry.zloc), Array{Any, 1})
    @test isequal(typeof(geometry.nt), Array{Any, 1})
    @test isequal(typeof(geometry.dt), Array{Any, 1})
    @test isequal(typeof(geometry.t), Array{Any, 1})

    # Constructor if coordinates are not passed as a cell arrays
    xsrc = range(100f0, stop=1100f0, length=nsrc)
    ysrc = range(0f0, stop=0f0, length=nsrc)
    zsrc = range(20f0, stop=20f0, length=nsrc)

    geometry = Geometry(xsrc, ysrc, zsrc; dt=4f0, t=1000f0, nsrc=nsrc)

    @test isequal(typeof(geometry.xloc), Array{Any, 1})
    @test isequal(typeof(geometry.yloc), Array{Any, 1})
    @test isequal(typeof(geometry.zloc), Array{Any, 1})
    @test isequal(typeof(geometry.nt), Array{Any, 1})
    @test isequal(typeof(geometry.dt), Array{Any, 1})
    @test isequal(typeof(geometry.t), Array{Any, 1})

    # Set up source geometry object from in-core data container
    block = segy_read("../data/unit_test_shot_records.segy")
    src_geometry = Geometry(block; key="source", segy_depth_key="SourceSurfaceElevation")
    rec_geometry = Geometry(block; key="receiver", segy_depth_key="RecGroupElevation")

    @test isequal(typeof(src_geometry), GeometryIC)
    @test isequal(typeof(rec_geometry), GeometryIC)
    @test isequal(get_header(block, "SourceSurfaceElevation")[1], src_geometry.zloc[1])
    @test isequal(get_header(block, "RecGroupElevation")[1], rec_geometry.zloc[1][1])
    @test isequal(get_header(block, "SourceX")[1], src_geometry.xloc[1])
    @test isequal(get_header(block, "GroupX")[1], rec_geometry.xloc[1][1])

    # Set up geometry summary from out-of-core data container
    container = segy_scan("../data/", "unit_test_shot_records", ["GroupX", "GroupY", "RecGroupElevation", "SourceSurfaceElevation", "dt"])
    src_geometry = Geometry(container; key="source", segy_depth_key="SourceSurfaceElevation")
    rec_geometry = Geometry(container; key="receiver", segy_depth_key="RecGroupElevation")

    @test isequal(typeof(src_geometry), GeometryOOC)
    @test isequal(typeof(rec_geometry), GeometryOOC)
    @test isequal(src_geometry.key, "source")
    @test isequal(rec_geometry.key, "receiver")
    @test isequal(src_geometry.segy_depth_key, "SourceSurfaceElevation")
    @test isequal(rec_geometry.segy_depth_key, "RecGroupElevation")
    @test isequal(prod(size(block.data)), sum(rec_geometry.nsamples))

    # Set up geometry summary from out-of-core data container passed as cell array
    container_cell = Array{SeisIO.SeisCon}(undef, nsrc)
    for j=1:nsrc
        container_cell[j] = split(container, j)
    end

    src_geometry = Geometry(container_cell; key="source", segy_depth_key="SourceSurfaceElevation")
    rec_geometry = Geometry(container_cell; key="receiver", segy_depth_key="RecGroupElevation")

    @test isequal(typeof(src_geometry), GeometryOOC)
    @test isequal(typeof(rec_geometry), GeometryOOC)
    @test isequal(src_geometry.key, "source")
    @test isequal(rec_geometry.key, "receiver")
    @test isequal(src_geometry.segy_depth_key, "SourceSurfaceElevation")
    @test isequal(rec_geometry.segy_depth_key, "RecGroupElevation")
    @test isequal(prod(size(block.data)), sum(rec_geometry.nsamples))

    # Load geometry from out-of-core Geometry container
    src_geometry_ic = Geometry(src_geometry)
    rec_geometry_ic = Geometry(rec_geometry)

    @test isequal(typeof(src_geometry_ic), GeometryIC)
    @test isequal(typeof(rec_geometry_ic), GeometryIC)
    @test isequal(get_header(block, "SourceSurfaceElevation")[1], src_geometry_ic.zloc[1])
    @test isequal(get_header(block, "RecGroupElevation")[1], rec_geometry_ic.zloc[1][1])
    @test isequal(get_header(block, "SourceX")[1], src_geometry_ic.xloc[1])
    @test isequal(get_header(block, "GroupX")[1], rec_geometry_ic.xloc[1][1])

    # Subsample in-core geometry structure
    src_geometry_sub = subsample(src_geometry_ic, 1)
    @test isequal(typeof(src_geometry_sub), GeometryIC)
    @test isequal(length(src_geometry_sub.xloc), 1)

    src_geometry_sub = subsample(src_geometry_ic, 1:2)
    @test isequal(typeof(src_geometry_sub), GeometryIC)
    @test isequal(length(src_geometry_sub.xloc), 2)

    # Subsample out-of-core geometry structure
    src_geometry_sub = subsample(src_geometry, 1)
    @test isequal(typeof(src_geometry_sub), GeometryOOC)
    @test isequal(length(src_geometry_sub.dt), 1)
    @test isequal(src_geometry_sub.segy_depth_key, "SourceSurfaceElevation")

    src_geometry_sub = subsample(src_geometry, 1:2)
    @test isequal(typeof(src_geometry_sub), GeometryOOC)
    @test isequal(length(src_geometry_sub.dt), 2)
    @test isequal(src_geometry_sub.segy_depth_key, "SourceSurfaceElevation")

    # Compare if geometries match
    @test compareGeometry(src_geometry_ic, src_geometry_ic)
    @test compareGeometry(rec_geometry_ic, rec_geometry_ic)

    @test compareGeometry(src_geometry, src_geometry)
    @test compareGeometry(rec_geometry, rec_geometry)

end


# Object-based wrapper for level-1 BLIS matrix-diagonal routines.
#

macro blis_object_api_level1d_form1(funcname)
    return quote
        @blis_ccall_object($funcname,
                           Cvoid,
                           a, xObj,
                           b, xObj)
    end
end

@blis_object_api_level1d_form1 addd
@blis_object_api_level1d_form1 copyd
@blis_object_api_level1d_form1 subd

macro blis_object_api_level1d_form2(funcname)
    return quote
        @blis_ccall_object($funcname,
                           Cvoid,
                           α, xObj,
                           a, xObj,
                           b, xObj)
    end
end

@blis_object_api_level1d_form2 axpyd
@blis_object_api_level1d_form2 scal2d

macro blis_object_api_level1d_form3(funcname)
    return quote
        @blis_ccall_object($funcname,
                           Cvoid,
                           α, xObj,
                           a, xObj)
    end
end

@blis_object_api_level1d_form3 scald
@blis_object_api_level1d_form3 setd
@blis_object_api_level1d_form3 setid
@blis_object_api_level1d_form3 shiftd

@blis_ccall_object(invertd,
                   Cvoid,
                   a, xObj)
@blis_ccall_object(xpbyd,
                   Cvoid,
                   a, xObj,
                   β, xObj,
                   b, xObj)


module JuMPLPUtils
using SparseArrays
using MathProgBase
using JuMP
include("lp_course_utils.jl")
export getslack
export getoptranges
export getoptrange

end # module