function DiffEqBase.__solve(prob::DiffEqBase.AbstractDDEProblem,
                            alg::AbstractMethodOfStepsAlgorithm, args...;
                            kwargs...)
  integrator = DiffEqBase.__init(prob, alg, args...; kwargs...)
  DiffEqBase.solve!(integrator)
  integrator.sol
end

function DiffEqBase.__init(prob::DiffEqBase.AbstractDDEProblem,
                           alg::AbstractMethodOfStepsAlgorithm,
                           timeseries_init = (),
                           ts_init = (),
                           ks_init = ();
                           saveat = (),
                           tstops = (),
                           d_discontinuities = (),
                           save_idxs = nothing,
                           save_everystep = isempty(saveat),
                           save_on = true,
                           save_start = save_everystep || isempty(saveat) || saveat isa Number || prob.tspan[1] in saveat,
                           save_end = nothing,
                           callback = nothing,
                           dense = save_everystep && isempty(saveat),
                           calck = (callback !== nothing && callback != CallbackSet()) || # Empty callback
                           dense, # and no dense output
                           dt = zero(eltype(prob.tspan)),
                           dtmin = DiffEqBase.prob2dtmin(prob; use_end_time=false),
                           dtmax = eltype(prob.tspan)(prob.tspan[end]-prob.tspan[1]),
                           force_dtmin = false,
                           adaptive = DiffEqBase.isadaptive(alg),
                           gamma = OrdinaryDiffEq.gamma_default(alg.alg),
                           abstol = nothing,
                           reltol = nothing,
                           qmin = OrdinaryDiffEq.qmin_default(alg.alg),
                           qmax = OrdinaryDiffEq.qmax_default(alg.alg),
                           qsteady_min = OrdinaryDiffEq.qsteady_min_default(alg.alg),
                           qsteady_max = OrdinaryDiffEq.qsteady_max_default(alg.alg),
                           qoldinit = DiffEqBase.isadaptive(alg) ? 1//10^4 : 0,
                           controller = nothing,
                           fullnormalize = true,
                           failfactor = 2,
                           beta1 = nothing,
                           beta2 = nothing,
                           maxiters = adaptive ? 1000000 : typemax(Int),
                           internalnorm = DiffEqBase.ODE_DEFAULT_NORM,
                           internalopnorm = LinearAlgebra.opnorm,
                           isoutofdomain = DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN,
                           unstable_check = DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK,
                           verbose = true,
                           timeseries_errors = true,
                           dense_errors = false,
                           advance_to_tstop = false,
                           stop_at_next_tstop = false,
                           initialize_save = true,
                           progress = false,
                           progress_steps = 1000,
                           progress_name = "DDE",
                           progress_message = DiffEqBase.ODE_DEFAULT_PROG_MESSAGE,
                           userdata = nothing,
                           allow_extrapolation = OrdinaryDiffEq.alg_extrapolates(alg),
                           initialize_integrator = true,
                           alias_u0 = false,
                           # keyword arguments for DDEs
                           discontinuity_interp_points::Int = 10,
                           discontinuity_abstol = eltype(prob.tspan)(1//Int64(10)^12),
                           discontinuity_reltol = 0,
                           kwargs... )
  if haskey(kwargs, :initial_order)
    @warn "initial_order has been deprecated. Please specify order_discontinuity_t0 in the DDEProblem instead."
    order_discontinuity_t0::Int = kwargs[:initial_order]
  else
    order_discontinuity_t0 = prob.order_discontinuity_t0
  end

  if alg.alg isa CompositeAlgorithm && alg.alg.choice_function isa AutoSwitch
    auto = alg.alg.choice_function
    alg = MethodOfSteps(CompositeAlgorithm(alg.alg.algs,
                             OrdinaryDiffEq.AutoSwitchCache(
                                             0,0,
                                             auto.nonstiffalg,
                                             auto.stiffalg,
                                             auto.stiffalgfirst,
                                             auto.maxstiffstep,
                                             auto.maxnonstiffstep,
                                             auto.nonstifftol,
                                             auto.stifftol,
                                             auto.dtfac,
                                             auto.stiffalgfirst,
                                             auto.switch_max
                                            )))
  end

  if haskey(kwargs, :minimal_solution)
    @warn "minimal_solution is ignored"
  end

  if !isempty(saveat) && dense
    @warn("Dense output is incompatible with saveat. Please use the SavingCallback from the Callback Library to mix the two behaviors.")
  end

  progress && @logmsg(-1, progress_name, _id = _id = :DelayDiffEq, progress = 0)

  isdae = prob.f.mass_matrix !== I && !(prob.f.mass_matrix isa Tuple) &&
      ArrayInterface.issingular(prob.f.mass_matrix)

  # unpack problem
  @unpack f, u0, h, tspan, p, neutral, constant_lags, dependent_lags = prob

  # determine type and direction of time
  tType = eltype(tspan)
  t0 = first(tspan)
  tdir = sign(last(tspan) - t0)
  tTypeNoUnits = typeof(one(tType))

  # Allow positive dtmax, but auto-convert
  dtmax > zero(dtmax) && tdir < zero(tdir) && (dtmax *= tdir)

  # no fixed-point iterations for constrained algorithms,
  # and thus `dtmax` should match minimal lag
  if isconstrained(alg) && has_constant_lags(prob)
    dtmax = tdir * min(abs(dtmax), minimum(abs, constant_lags))
  end

  # get absolute and relative tolerances
  abstol_internal = get_abstol(u0, tspan, alg.alg; abstol = abstol)
  reltol_internal = get_reltol(u0, tspan, alg.alg; reltol = reltol)

  # get rate prototype
  rate_prototype = rate_prototype_of(u0, tspan)

  # create a history function
  history = build_history_function(prob, alg, rate_prototype, reltol_internal;
                                   dt = dt, dtmin = dtmin, calck = false,
                                   adaptive = adaptive, internalnorm = internalnorm)
  f_with_history = ODEFunctionWrapper(f, history)

  # get states (possibly different from the ODE integrator!)
  u, uprev, uprev2 = u_uprev_uprev2(u0, alg;
                                    alias_u0 = alias_u0,
                                    adaptive = adaptive,
                                    allow_extrapolation = allow_extrapolation,
                                    calck = calck)
  uEltypeNoUnits = recursive_unitless_eltype(u)
  uBottomEltypeNoUnits = recursive_unitless_bottom_eltype(u)

  # initialize output arrays of the solution
  k = typeof(rate_prototype)[]
  ts, timeseries, ks = solution_arrays(u, tspan, rate_prototype;
                                       timeseries_init = timeseries_init,
                                       ts_init = ts_init,
                                       ks_init = ks_init,
                                       save_idxs = save_idxs,
                                       save_start = save_start)

  # build cache
  ode_integrator = history.integrator
  cache = OrdinaryDiffEq.alg_cache(alg.alg, u, rate_prototype, uEltypeNoUnits,
                                   uBottomEltypeNoUnits, tTypeNoUnits, uprev, uprev2,
                                   f_with_history, t0, zero(tType), reltol_internal, p, calck,
                                   Val(isinplace(prob)))

  # separate statistics of the integrator and the history
  destats = DDEStats(0)

  # create solution
  if iscomposite(alg)
    id = OrdinaryDiffEq.CompositeInterpolationData(f_with_history, timeseries, ts, ks,
                                                   Int[], dense, cache)
    sol = DiffEqBase.build_solution(prob, alg.alg, ts, timeseries;
                                    dense = dense, k = ks, interp = id,
                                    alg_choice = id.alg_choice, calculate_error = false,
                                    destats = destats)
  else
    id = OrdinaryDiffEq.InterpolationData(f_with_history, timeseries, ts, ks, dense, cache)
    sol = DiffEqBase.build_solution(prob, alg.alg, ts, timeseries;
                                    dense = dense, k = ks, interp = id,
                                    calculate_error = false, destats = destats)
  end

  # retrieve time stops, time points at which solutions is saved, and discontinuities
  tstops_internal =
    OrdinaryDiffEq.initialize_tstops(tType, tstops, d_discontinuities, tspan)
  saveat_internal = OrdinaryDiffEq.initialize_saveat(tType, saveat, tspan)
  d_discontinuities_internal = OrdinaryDiffEq.initialize_d_discontinuities(
    Discontinuity{tType,Int}, d_discontinuities, tspan
  )

  maximum_order = OrdinaryDiffEq.alg_maximum_order(alg)
  tstops_propagated, d_discontinuities_propagated =
    initialize_tstops_d_discontinuities_propagated(tType, tstops, d_discontinuities,
                                                   tspan, order_discontinuity_t0,
                                                   maximum_order, constant_lags, neutral)

  # reserve capacity for the solution
  sizehint!(sol, alg, tspan, tstops_internal, saveat_internal;
            save_everystep = save_everystep, adaptive = adaptive, dt = tType(dt),
            dtmin = dtmin, internalnorm = internalnorm)

  # create array of tracked discontinuities
  # used to find propagated discontinuities with callbacks and to keep track of all
  # passed discontinuities
  tracked_discontinuities = Discontinuity{tType,Int}[]
  if order_discontinuity_t0 ≤ maximum_order
    push!(tracked_discontinuities, Discontinuity(tdir * t0, order_discontinuity_t0))
  end

  # Create set of callbacks and its cache
  callback_set, callback_cache = callback_set_and_cache(prob, callback)

  # separate options of integrator and of dummy ODE integrator since ODE integrator always saves
  # every step and every index (necessary for history function)
  QT = tTypeNoUnits <: Integer ? typeof(qmin) : typeof(internalnorm(u, t0))

  # Setting up the step size controller
  if (beta1 !== nothing || beta2 !== nothing) && controller !== nothing
    throw(ArgumentError(
      "Setting both the legacy PID parameters `beta1, beta2 = $((beta1, beta2))` and the `controller = $controller` is not allowed."))
  end

  if (beta1 !== nothing || beta2 !== nothing)
    message = "Providing the legacy PID parameters `beta1, beta2` is deprecated. Use the keyword argument `controller` instead."
    Base.depwarn(message, :init)
    Base.depwarn(message, :solve)
  end

  if controller === nothing
    controller = OrdinaryDiffEq.default_controller(alg.alg, cache, convert(QT, qoldinit)::QT, beta1, beta2)
  end

  save_end_user = save_end
  save_end = save_end === nothing ? save_everystep || isempty(saveat) || saveat isa Number || prob.tspan[2] in saveat : save_end

  opts = OrdinaryDiffEq.DEOptions{typeof(abstol_internal),typeof(reltol_internal),
                                  QT,tType,typeof(controller),
                                  typeof(internalnorm),typeof(internalopnorm),
                                  typeof(save_end_user),
                                  typeof(callback_set),
                                  typeof(isoutofdomain),
                                  typeof(progress_message),typeof(unstable_check),
                                  typeof(tstops_internal),
                                  typeof(d_discontinuities_internal),typeof(userdata),
                                  typeof(save_idxs),
                                  typeof(maxiters),typeof(tstops),
                                  typeof(saveat),typeof(d_discontinuities)}(
                                    maxiters,save_everystep,
                                    adaptive,abstol_internal,reltol_internal,
                                    QT(gamma),QT(qmax),QT(qmin),
                                    QT(qsteady_max),QT(qsteady_min),
                                    QT(qoldinit),QT(failfactor),
                                    tType(dtmax),tType(dtmin),
                                    controller,
                                    internalnorm,internalopnorm,
                                    save_idxs,tstops_internal,saveat_internal,
                                    d_discontinuities_internal,
                                    tstops,saveat,d_discontinuities,
                                    userdata,progress,progress_steps,
                                    progress_name,progress_message,
                                    timeseries_errors,dense_errors,dense,
                                    save_on,save_start,save_end,save_end_user,
                                    callback_set,
                                    isoutofdomain,unstable_check,
                                    verbose,calck,force_dtmin,
                                    advance_to_tstop,stop_at_next_tstop)

  # create fixed point solver
  fpsolver = build_fpsolver(alg, alg.fpsolve, u, uEltypeNoUnits, uBottomEltypeNoUnits,
                            Val(isinplace(prob)))

  # initialize indices of u(t) and u(tprev) in the dense history
  prev_idx = 1
  prev2_idx = 1

  # create integrator combining the new defined problem function with history
  # information, the new solution, the parameters of the ODE integrator, and
  # parameters of fixed-point iteration
  # do not initialize fsalfirst and fsallast
  # rate/state = (state/time)/state = 1/t units, internalnorm drops units
  eigen_est = inv(one(tType))
  tprev = t0
  dtcache = tType(dt)
  dtpropose = tType(dt)
  iter = 0
  kshortsize = 0
  reeval_fsal = false
  u_modified = false
  EEst = QT(1)
  just_hit_tstop = false
  do_error_check = true
  isout = false
  accept_step = false
  force_stepfail = false
  last_stepfail = false
  event_last_time = 0
  vector_event_last_time = 1
  last_event_error = zero(uBottomEltypeNoUnits)
  dtchangeable = OrdinaryDiffEq.isdtchangeable(alg.alg)
  q11 = QT(1)
  success_iter = 0
  erracc = QT(1)
  dtacc = tType(1)

  integrator = DDEIntegrator{typeof(alg.alg),isinplace(prob),typeof(u0),tType,typeof(p),
                             typeof(eigen_est),QT,typeof(tdir),typeof(k),typeof(sol),
                             typeof(f_with_history),typeof(cache),
                             typeof(ode_integrator),typeof(fpsolver),
                             typeof(opts),typeof(discontinuity_abstol),
                             typeof(discontinuity_reltol),typeof(history),
                             typeof(tstops_propagated),typeof(d_discontinuities_propagated),
                             OrdinaryDiffEq.fsal_typeof(alg.alg, rate_prototype),
                             typeof(last_event_error),typeof(callback_cache)}(
                               sol, u, k, t0, tType(dt), f_with_history, p,
                               uprev, uprev2, tprev, prev_idx, prev2_idx,
                               fpsolver,
                               order_discontinuity_t0, tracked_discontinuities,
                               discontinuity_interp_points, discontinuity_abstol,
                               discontinuity_reltol, tstops_propagated,
                               d_discontinuities_propagated, alg.alg,
                               dtcache, dtchangeable, dtpropose, tdir,
                               eigen_est, EEst, QT(qoldinit), q11,
                               erracc, dtacc, success_iter,
                               iter, length(ts), length(ts), cache, callback_cache,
                               kshortsize, force_stepfail, last_stepfail,
                               just_hit_tstop, do_error_check, event_last_time, vector_event_last_time,
                               last_event_error, accept_step,
                               isout, reeval_fsal,
                               u_modified, isdae, opts, destats, history,
                               ode_integrator)

  # initialize DDE integrator
  if initialize_integrator
    initialize_solution!(integrator)
    OrdinaryDiffEq.initialize_callbacks!(integrator, initialize_save)
    OrdinaryDiffEq.initialize!(integrator)
  end

  # take care of time step dt = 0 and dt with incorrect sign
  OrdinaryDiffEq.handle_dt!(integrator)

  integrator
end

function DiffEqBase.solve!(integrator::DDEIntegrator)
  @unpack tdir, opts, sol = integrator
  @unpack tstops = opts

  # step over all stopping time points, similar to solving with ODE integrators
  @inbounds while !isempty(tstops)
    while tdir * integrator.t < first(tstops)
      # apply step or adapt step size
      OrdinaryDiffEq.loopheader!(integrator)

      # abort integration following same criteria as for ODEs:
      # maxiters exceeded, dt <= dtmin, integration unstable
      DiffEqBase.check_error!(integrator) === :Success || return sol

      # calculate next step
      OrdinaryDiffEq.perform_step!(integrator)

      # calculate proposed next step size, handle callbacks, and update solution
      OrdinaryDiffEq.loopfooter!(integrator)

      isempty(tstops) && break
    end

    # remove hit or passed stopping time points
    OrdinaryDiffEq.handle_tstop!(integrator)
  end

  # clean up solution
  DiffEqBase.postamble!(integrator)

  f = sol.prob.f

  if DiffEqBase.has_analytic(f)
    DiffEqBase.calculate_solution_errors!(sol;
                                          timeseries_errors = opts.timeseries_errors,
                                          dense_errors = opts.dense_errors)
  end
  sol.retcode === :Default || return sol

  integrator.sol = DiffEqBase.solution_new_retcode(sol, :Success)
end

function OrdinaryDiffEq.initialize_callbacks!(integrator::DDEIntegrator,
                                              initialize_save = true)
  callbacks = integrator.opts.callback
  prob = integrator.sol.prob

  # set up additional initial values of newly created DDE integrator
  # (such as fsalfirst) and its callbacks

  integrator.u_modified = true

  u_modified = initialize!(callbacks, integrator.u, integrator.t, integrator)

  # if the user modifies u, we need to fix previous values before initializing
  # FSAL in order for the starting derivatives to be correct
  if u_modified

    if isinplace(prob)
      recursivecopy!(integrator.uprev, integrator.u)
    else
      integrator.uprev = integrator.u
    end

    if OrdinaryDiffEq.alg_extrapolates(integrator.alg)
      if isinplace(prob)
        recursivecopy!(integrator.uprev2, integrator.uprev)
      else
        integrator.uprev2 = integrator.uprev
      end
    end

    # update heap of discontinuities
    # discontinuity is assumed to be of order 0, i.e. solution x is discontinuous
    add_next_discontinuities!(integrator, 0, integrator.t)

    # reset this as it is now handled so the integrators should proceed as normal
    reeval_internals_due_to_modification!(integrator, Val{false})

    if initialize_save &&
      (any((c)->c.save_positions[2],callbacks.discrete_callbacks) ||
       any((c)->c.save_positions[2],callbacks.continuous_callbacks))
      savevalues!(integrator, true)
    end
  end

  # reset this as it is now handled so the integrators should proceed as normal
  integrator.u_modified = false
end

function initialize_tstops_d_discontinuities_propagated(
  ::Type{T}, tstops, d_discontinuities, tspan, order_discontinuity_t0, alg_maximum_order,
  constant_lags, neutral,
) where T
  # create heaps for propagated discontinuities and corresponding time stops
  tstops_propagated = BinaryMinHeap{T}()
  d_discontinuities_propagated = BinaryMinHeap{Discontinuity{T,Int}}()

  # add discontinuities and time stops propagated from initial discontinuity
  if constant_lags !== nothing && !isempty(constant_lags) &&
    order_discontinuity_t0 ≤ alg_maximum_order

    sizehint!(tstops_propagated, length(constant_lags))
    sizehint!(d_discontinuities_propagated, length(constant_lags))

    t0, tf = tspan
    tdir = sign(tf - t0)
    maxlag = tdir * (tf - t0)
    next_order = neutral ? order_discontinuity_t0 : order_discontinuity_t0 + 1
    for lag in constant_lags
      if tdir * lag < maxlag
        t = tdir * (t0 + lag)
        push!(tstops_propagated, t)

        d = Discontinuity{T,Int}(t, next_order)
        push!(d_discontinuities_propagated, d)
      end
    end
  end

  return tstops_propagated, d_discontinuities_propagated
end


result = [
    2.32551,
    -0.155293,
    0.469214,
    -0.0870457,
    0.271364,
    -0.222398,
    1.01423,
    -0.179009,
    1.01724,
    -0.152972,
    0.216039,
    -0.251637,
    1.05034,
    -0.231289,
    0.245388,
    -0.231289,
    1.03371,
    -0.258872,
    1.03256,
    -0.262519,
    1.04748,
    -0.259143,
    1.0535,
    -0.266484,
    1.04711,
    -0.267177,
    1.02131,
    -0.280381,
]

include(joinpath(BASE_DIR, "test", "data_5bus_pu.jl"))
include(joinpath(BASE_DIR, "test", "data_14bus_pu.jl"))

nodes_5 = nodes5()
nodes_14 = nodes14()
c_sys14() = System(
    nodes_14,
    thermal_generators14(nodes_14),
    loads14(nodes_14),
    branches14(nodes_14),
    nothing,
    100.0,
    nothing,
    nothing,
);

c_sys5_re() = System(
    nodes_5,
    vcat(thermal_generators5(nodes_5), renewable_generators5(nodes_5)),
    loads5(nodes_5),
    nothing,
    nothing,
    100.0,
    nothing,
    nothing,
)

@testset begin
    using NLsolve
    # This is a negative test. The data passed for sys5_re is known to be infeasible.
    @test_logs(
        (:error, "The powerflow solver returned convergence = false"),
        match_mode = :any,
        @test !solve_powerflow!(c_sys5_re(), nlsolve)
    )
    @test solve_powerflow!(c_sys14(), nlsolve, method = :newton)
end


mutable struct RegisteredFunction
    s         :: Symbol
    dimension :: Integer
    f         :: Function
    gradf     :: Union{Nothing, Function}
    grad2f    :: Union{Nothing, Function}
    autodiff  :: Bool
end

# Options for the solver (more details like defaults in solver.jl)
mutable struct SolverOptions
    nl_solver                           :: Any# needs to be set
    log_levels                          :: Vector{Symbol}
    silent                              :: Bool
    atol                                :: Float64
    num_resolve_root_relaxation         :: Int64
    branch_strategy                     :: Symbol
    gain_mu                             :: Float64
    strong_branching_perc               :: Float64
    strong_branching_nsteps             :: Int64
    strong_branching_time_limit  :: Float64
    strong_restart                      :: Bool
    reliability_branching_threshold     :: Int64
    reliability_branching_perc          :: Float64
    incumbent_constr                    :: Bool
    obj_epsilon                         :: Float64
    time_limit                          :: Float64
    mip_gap                             :: Float64
    best_obj_stop                       :: Float64
    solution_limit                      :: Int64
    all_solutions                       :: Bool
    list_of_solutions                   :: Bool
    processors                          :: Int64
    traverse_strategy                   :: Symbol
    feasibility_pump                    :: Bool
    feasibility_pump_time_limit         :: Float64
    feasibility_pump_tolerance_counter  :: Int64
    tabu_list_length                    :: Int64
    num_resolve_nlp_feasibility_pump    :: Int64
    mip_solver                          :: Any
    allow_almost_solved                 :: Bool
    allow_almost_solved_integral        :: Bool
    registered_functions                :: Union{Nothing, Vector{RegisteredFunction}}
    seed                                :: Int64

    # only for testing
    force_parallel                      :: Bool
    debug                               :: Bool
    debug_write                         :: Bool
    debug_file_path                     :: String

    fixed_gain_mu                       :: Bool
end

mutable struct SolutionObj
    solution    :: Vector{Float64}
    objval      :: Float64
end

# Juniper MOI struct

mutable struct JuniperProblem
    nl_solver           :: Any
    nl_solver_options   :: Vector{Pair}

    model               :: JuMP.Model

    relaxation_status   :: MOI.TerminationStatusCode
    relaxation_objval   :: Float64
    relaxation_solution :: Vector{Float64}

    status              :: MOI.TerminationStatusCode
    objval              :: Float64
    best_bound          :: Float64

    x                   :: Vector{JuMP.VariableRef}
    primal_start        :: Vector{Real}
    num_constr          :: Int64
    num_nl_constr       :: Int64
    num_q_constr        :: Int64
    num_l_constr        :: Int64
    num_var             :: Int64
    l_var               :: Vector{Float64}
    u_var               :: Vector{Float64}

    has_nl_objective    :: Bool
    nlp_evaluator       :: MOI.AbstractNLPEvaluator

    objective           :: Union{SVF, SAF, SQF, Nothing}

    disc2var_idx        :: Vector{Int64}
    var2disc_idx        :: Vector{Int64}

    var_type            :: Vector{Symbol}
    obj_sense           :: Symbol
    num_disc_var        :: Int64

    solution            :: Vector{Float64}

    soltime             :: Float64
    options             :: SolverOptions
    solutions           :: Vector{SolutionObj}
    nsolutions          :: Int64

    mip_solver          :: Any
    mip_solver_options  :: Vector{Pair}

    relaxation_time     :: Float64
    start_time          :: Float64

    # Info
    nintvars            :: Int64
    nbinvars            :: Int64
    nnodes              :: Int64
    ncuts               :: Int64
    nbranches           :: Int64
    nlevels             :: Int64

    fpump_info          :: Dict{Symbol,Float64}

    # debug
    debugDict           :: Dict{Symbol,Any}

    JuniperProblem() = new()
end

###########################################################################
########################## FPump ##########################################
###########################################################################
mutable struct TabuList
    sols      :: Vector{Vector{Float64}}
    length    :: Int64
    pointer   :: Int64

    TabuList() = new()
end

############################################################################
######################## Tree structure ####################################
############################################################################
mutable struct BnBNode
    idx                 :: Int64
    level               :: Int64
    l_var               :: Vector{Float64}
    u_var               :: Vector{Float64}
    solution            :: Vector{Float64}
    var_idx             :: Int64
    state               :: Symbol
    relaxation_state    :: MOI.TerminationStatusCode
    best_bound          :: Float64
    path                :: Vector{String}   # list of parent hashes
    hash                :: String
end

mutable struct Incumbent
    objval      :: Float64
    solution    :: Vector{Float64}
    only_almost :: Bool
end

mutable struct GainObj
    minus           :: Vector{Float64} # gain of objective per variable on left node
    plus            :: Vector{Float64} # gain of objective per variable on right node
    minus_counter   :: Vector{Int64} # obj_gain_m / obj_gain_mc => average gain on left node
    plus_counter    :: Vector{Int64} # obj_gain_p / obj_gain_pc => average gain on right node
    # counter of how often one child is infeasible (>= 0) subtract 1 if both Optimal
    # can be negative temporately in step_obj but not in tree.obj_gain_
    inf_counter     :: Vector{Int64}
end

mutable struct BnBTreeObj
    m               :: Juniper.JuniperProblem
    incumbent       :: Incumbent
    limit           :: Symbol
    obj_gain        :: GainObj
    disc2var_idx    :: Vector{Int64}
    var2disc_idx    :: Vector{Int64}
    options         :: Juniper.SolverOptions
    obj_fac         :: Int64 # factor for objective 1 if max -1 if min
    start_time      :: Float64
    nsolutions      :: Int64
    branch_nodes    :: Vector{BnBNode}
    best_bound      :: Float64
    global_solver   :: Bool

    BnBTreeObj() = new()
end

mutable struct StrongBranchStep
    var_idx             :: Int64
    l_relaxation_state  :: MOI.TerminationStatusCode
    r_relaxation_state  :: MOI.TerminationStatusCode
    init_restart        :: Bool

    StrongBranchStep() = new()
end

# the object holds information for the current step
mutable struct StepObj
    node                :: BnBNode # current branch node
    var_idx             :: Int64   # variable to branch on
    state               :: Symbol  # if infeasible => break (might be set by strong branching)
    nrestarts           :: Int64
    gain_gap            :: Float64
    obj_gain            :: GainObj
    idx_time            :: Float64
    node_idx_time       :: Float64
    upd_gains_time      :: Float64
    node_branch_time    :: Float64
    branch_time         :: Float64
    integral            :: Vector{BnBNode}
    branch              :: Vector{BnBNode}
    l_nd                :: BnBNode
    r_nd                :: BnBNode
    counter             :: Int64
    upd_gains           :: Symbol
    strong_disc_vars    :: Vector{Int64}
    branch_strategy     :: Symbol # the strategy that is used in this step
    strong_branching    :: Union{Nothing,Vector{StrongBranchStep}} # used only if debug=true and branch_strategy of this step is :Strong

    StepObj() = new()
end

mutable struct TimeObj
    solve_leaves_get_idx :: Float64
    solve_leaves_branch :: Float64
    branch :: Float64
    get_idx :: Float64
    upd_gains :: Float64
end


# script to fit model to data
#
# The script is called as
#
# fitmodel.jl [model] [objective] [[rat]]
#
# and writes data to ../fitdata/fits/fit_[model]_[objective]_[rat]. Here, [rat]
# is optional and indicates the rat identifier (1-4). If not provided, the data
# for all rats is collapsed into a single rat.

using Random, CSV

# include required scripts
include("common.jl")
include("models.jl")
include("simmodel.jl")
include("behaviorstats.jl")
include("modelfitting.jl")

# settings
const blockreps = 10
const simtrials = 10^5
const burnin = 1000
const nrats = 1 # (do not divide dataset by 4 rats, nrats = 4)
const datapath = "../fitdata/fits"
const paraminiobjfn = "psychchronseq1"

# returns model for given name
function createmodel(modelname)
    modeldict = Dict("opt" => OptimalModel,
                     "optlapse" => OptimalLapseModel,
                     "optcollapse" => OptimalColLapseModel,
                     "optbiaslapse" => OptimalBiasLapseModel,
                     "optcolbiaslapse" => OptimalColBiasLapseModel,
                     "adflapse" => ADFLapseModel,
                     "adfcollapse" => ADFColLapseModel,
                     "adfbiaslapse" => ADFBiasLapseModel,
                     "adfcolbiaslapse" => ADFColBiasLapseModel,
                     "randomcollapse" => RandomColLapseModel,
                     "randomcolbiaslapse" => RandomColBiasLapseModel,
                     "gammalapse" => GammaLapseModel,
                     "gammacollapse" => GammaColLapseModel,
                     "gammabiaslapse" => GammaBiasLapseModel,
                     "gammacolbiaslapse" => GammaColBiasLapseModel,
                     "deltalapse" => DeltaLapseModel,
                     "deltacollapse" => DeltaColLapseModel,
                     "deltabiaslapse" => DeltaBiasLapseModel,
                     "deltacolbiaslapse" => DeltaColBiasLapseModel,
                     "deltaexpcolbiaslapse" => DeltaExpColBiasLapseModel,
                     "lcalapse" => LCALapseModel,
                     "lcacollapse" => LCAColLapseModel)
    if modelname in keys(modeldict)
        return modeldict[modelname]()
    else
        error("Uknown model identifier $modelname")
    end
end

# returns fitting objective for given name and model
function objectivefn(objfnname, m, ratid::Int=0)
    # load data for individual / all rats
    d_ident = load_data(:ident, ratid)
    d_categ = load_data(:categ, ratid)
    d_itlvd = load_data(:itlvd, ratid)
    # simulated trials
    Random.seed!(0)
    evidence, ids = catevidence(ident_cs, categ_cs, 10, 10^5)
    itlvdevidence, itlvdids = catitlvdevidence(itlvd_cs, 10, 10^5)
    # create objective function
    if objfnname == "psychchron1"
        o = PsychChron1ObjFn(m, evidence, ids, burnin, d_ident, d_categ, nrats)
    elseif objfnname == "psychchron1ident"
        o = PsychChron1IdentObjFn(m, evidence, ids, burnin, d_ident, d_categ, nrats)
    elseif objfnname == "psychchron1categ"
        o = PsychChron1CategObjFn(m, evidence, ids, burnin, d_ident, d_categ, nrats)
    elseif objfnname == "psychchron2"
        o = PsychChron2ObjFn(m, evidence, ids, burnin, d_ident, d_categ, nrats)
    elseif objfnname == "psychchron3"
        o = PsychChron3ObjFn(m, evidence, ids, burnin, d_ident, d_categ, nrats)
    elseif objfnname == "psychchronseq1"
        o = PsychChronSeq1ObjFn(m, evidence, ids, burnin, d_ident, d_categ, nrats)
    elseif objfnname == "psychchronseq2"
        o = PsychChronSeq2ObjFn(m, evidence, ids, burnin, d_ident, d_categ, nrats)
    elseif objfnname == "psychchronseq3"
        o = PsychChronSeq3ObjFn(m, evidence, ids, burnin, d_ident, d_categ, nrats)
    elseif objfnname == "psychchron1itlvd"
        o = PsychChron1ItlvdObjFn(m, itlvdevidence, itlvdids, burnin, d_itlvd, nrats)
    elseif objfnname == "psychchron1itlvdedge"
        # only fit stimuli that also appear in ident/categ condition
        o = PsychChron1ItlvdObjFn(m, itlvdevidence, itlvdids, burnin, d_itlvd, nrats,
                                  [1, 2, 3, 4, 5, 6, 7, 8, 9, 16, 17, 24, 25, 32])
    else
        error("Unknown object function identifier $objfnname")
    end
end

# process arguments
if length(ARGS) ∉ [2,3]
    error("Expected two or three arguments")
end
const modelname = lowercase(ARGS[1])
const objfnname = lowercase(ARGS[2])
ratid = length(ARGS) >= 3 ? parse(Int, ARGS[3]) : 0  # ratid = 0 for all rats
m = createmodel(modelname)
o = objectivefn(objfnname, m, ratid)

# check if model has already been fit with parameter initialization objective fn
datafile(modelname, objfnname, ratid) = ratid == 0 ?
    "$(datapath)/fit_$(modelname)_$(objfnname).csv" :
    "$(datapath)/fit_$(modelname)_$(objfnname)_rat$(ratid).csv"
paraminifilename = datafile(modelname, paraminiobjfn, ratid)
if ratid != 0 && !isfile(paraminifilename)
    # default to collective rat fits if they don't exist for individual rats
    paraminifilename = datafile(modelname, paraminiobjfn, 0)
end
if isfile(paraminifilename)
    # loading initial parameters from previous model fit with $paraminiobjfn
    println("Loading initial parameters from $paraminifilename")
    d = CSV.read(paraminifilename)
    (maxllh, maxi) = findmax(d[:llh])
    if maxi != size(d,1) && maxllh > d[end,end]
        warn("Maximum llh $maxllh found in line $maxi instead of last line")
    end
    ϕini = Vector(d[maxi,1:end-1])
else
    # use initial parameters are determined by model
    ϕini = initialϕ(m)
end

# fit model
if ratid == 0
    println("Fitting $modelname model with objective $objfnname across all rats")
else
    println("Fitting $modelname model with objective $objfnname for rat $ratid")
end
d, ret = fitmodel(m, o, ϕini)
println("Optimizer stopped with return value $ret")
println("Minimum found at")
println(d[end,:])

if !(ret ∈ (:SUCCESS, :XTOL_REACHED, :STOPVAL_REACHED,
            :FTOL_REACHED, :MAXEVAL_REACHED, :MAXTIME_REACHED))
    println("Not writing to file, as optimization didn't complete successfully")
else
    # write to file
    outfile = datafile(modelname, objfnname, ratid)
    println("Writing data to $(outfile)")
    CSV.write(outfile, d)
end


module ModelHousehold

#
# This module provides a slightly abstracted view of a single Scottish household. The fields
# are predominantly based on the FRS, but could in principle come from other sources
# such as SHS, LCF and so on, or perhaps from a web form.
#
using Base: String
using Dates
import Base.show

using ScottishTaxBenefitModel
using .Definitions
using .Utils:
    has_non_z,
    todays_date,
    to_md_table

using .EquivalenceScales:
    EQ_P_Type,
    EQ_Person,
    EQScales,
    get_equivalence_scales,
    eq_dependent_child,
    eq_head,
    eq_other_adult,
    eq_spouse_of_head

import .EquivalenceScales.get_age
import .EquivalenceScales.eq_rel_to_hoh

import Base.isequal
import Base.hash

using .Uprating: uprate, UPRATE_MAPPINGS

export
    Household,
    BenefitUnit,
    BenefitUnits,
    BUAllocation,
    Person,
    People_Dict,

    adult_pids,
    age_then,
    between_ages,
    child_pids,
    default_bu_allocation,
    empl_status_in,
    ge_age,
    get_benefit_units,
    get_head,
    get_spouse,
    has_carer_member,
    has_children,
    has_disabled_member,
    has_income,
    interview_date,
    is_child,
    is_head,
    is_lone_parent,
    is_severe_disability,
    is_single,
    is_spouse,
    isless,
    le_age,
    make_benefit_unit,
    make_eq_scales!,
    num_adults,
    num_adults,
    num_carers,
    num_children,
    num_people,
    oldest_person,
    on_mt_benefits,
    pers_is_carer,
    pers_is_disabled,
    printpids,
    search,
    set_wage!,
    to_string,
    uprate!

mutable struct Person{RT<:Real}
    hid::BigInt # == sernum
    pid::BigInt # == unique id (year * 100000)+
    pno:: Int # person number in household
    is_hrp :: Bool
    default_benefit_unit:: Int
    is_standard_child :: Bool
    age:: Int

    sex::Sex
    ethnic_group::Ethnic_Group
    marital_status::Marital_Status
    highest_qualification::Qualification_Type

    sic::SIC_2007
    occupational_classification::Standard_Occupational_Classification
    public_or_private :: Employment_Sector
    principal_employment_type :: Employment_Type

    socio_economic_grouping::Socio_Economic_Group
    age_completed_full_time_education:: Int
    years_in_full_time_work:: Int
    employment_status::ILO_Employment
    actual_hours_worked::RT
    usual_hours_worked::RT

    age_started_first_job :: Int

    income::Incomes_Dict{RT}
    benefit_ratios :: Incomes_Dict{RT}

    jsa_type :: JSAType # FIXME change this name
    esa_type :: JSAType
    dla_self_care_type :: LowMiddleHigh
    dla_mobility_type :: LowMiddleHigh
    attendance_allowance_type :: LowMiddleHigh
    pip_daily_living_type :: PIPType
    pip_mobility_type ::  PIPType

    bereavement_type :: BereavementType
    had_children_when_bereaved :: Bool

    assets::Asset_Dict{RT}
    over_20_k_saving :: Bool
    pay_includes ::Included_In_Pay_Dict{Bool}

    # contracted_out_of_serps::Bool

    registered_blind::Bool
    registered_partially_sighted::Bool
    registered_deaf::Bool

    disabilities::Disability_Dict{Bool}

    health_status::Health_Status

    has_long_standing_illness :: Bool
    adls_are_reduced :: ADLS_Inhibited
    how_long_adls_reduced :: Illness_Length

    relationships::Relationship_Dict
    relationship_to_hoh :: Relationship
    is_informal_carer::Bool
    receives_informal_care_from_non_householder::Bool
    hours_of_care_received::RT
    hours_of_care_given::RT
    #
    # Childcare fields; assigned to children
    #
    hours_of_childcare :: RT
    cost_of_childcare :: RT
    childcare_type :: Child_Care_Type
    employer_provides_child_care :: Bool

    company_car_fuel_type :: Fuel_Type
    company_car_value :: RT
    company_car_contribution :: RT
    fuel_supplied :: RT
    onerand :: String
end

People_Dict = Dict{BigInt,Person{T}} where T<:Real
Pid_Array = Vector{BigInt}

mutable struct Household{RT<:Real}
    sequence:: Int # position in current generated dataset
    hid::BigInt
    data_year :: Int
    interview_year:: Int
    interview_month:: Int
    quarter:: Int
    tenure::Tenure_Type
    region::Standard_Region
    ct_band::CT_Band
    council_tax::RT
    water_and_sewerage ::RT
    mortgage_payment::RT
    mortgage_interest::RT
    years_outstanding_on_mortgage:: Int
    mortgage_outstanding::RT
    year_house_bought:: Int
    gross_rent::RT # rentg Gross rent including Housing Benefit  or rent Net amount of last rent payment
    rent_includes_water_and_sewerage::Bool
    other_housing_charges::RT
    gross_housing_costs::RT
    total_income::RT
    total_wealth::RT
    house_value::RT
    weight::RT
    council :: Symbol
    nhs_board :: Symbol
    bedrooms :: Int
    head_of_household :: BigInt
    people::People_Dict{RT}
    onerand :: String
    equivalence_scales :: EQScales{RT}
end

function to_string( hh :: Household ) :: String
    s = to_md_table( hh, exclude=[:people,:onerand])
    for pid in sort(collect(keys(hh.people)))
        s *= to_string( hh.people[pid] )
    end
    return s
end

function to_string( pers :: Person ) :: String
    s = """
    ### Person $(pers.pno)
    """
    s *= to_md_table( pers, depth=1,exclude=[:onerand])
    return s
end

"""
for equivalence scale implicit interface
"""
function get_age( p::Person ) :: Int
    return p.age;
end

"""
for the eq scale implicit interface
"""
function eq_rel_to_hoh( p :: Person ) :: EQ_P_Type
    if p.relationship_to_hoh == This_Person
        return eq_head
    elseif p.is_standard_child
        return eq_dependent_child
    elseif p.relationship_to_hoh in [Spouse,Cohabitee]
        return eq_spouse_of_head
    else
        return eq_other_adult
    end
end

function make_eq_scales!( hh :: Household{T} ) where T
    pl = collect(values(hh.people))
    # sz = size( pl )[1]
    # print( "num people $sz" )
    hh.equivalence_scales = get_equivalence_scales( T, pl )
end

function interview_date( hh :: Household ) :: Date
    return Date( hh.interview_year, hh.interview_month, 15 )
end

function uprate!( pid :: BigInt, year::Integer, quarter::Integer, person :: Person )
    for (t,inc) in person.income
            person.income[t] = uprate( inc, year, quarter, UPRATE_MAPPINGS[t])
    end
    for (a,ass) in person.assets
            person.assets[a] = uprate( ass, year, quarter, UPRATE_MAPPINGS[a])
    end
    person.cost_of_childcare = uprate( person.cost_of_childcare, year, quarter, upr_earnings )
end

function uprate!( hh :: Household )

    hh.water_and_sewerage  = uprate( hh.water_and_sewerage , hh.interview_year, hh.quarter, upr_housing_rents )
    hh.mortgage_payment = uprate( hh.mortgage_payment, hh.interview_year, hh.quarter, upr_housing_oo )
    hh.mortgage_interest = uprate( hh.mortgage_interest, hh.interview_year, hh.quarter, upr_housing_oo )
    hh.mortgage_outstanding = uprate( hh.mortgage_outstanding, hh.interview_year, hh.quarter, upr_housing_oo )
    hh.gross_rent = uprate( hh.gross_rent, hh.interview_year, hh.quarter, upr_housing_rents )
    hh.other_housing_charges = uprate( hh.other_housing_charges, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.gross_housing_costs = uprate( hh.gross_housing_costs, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.total_income = uprate( hh.total_income, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.total_wealth = uprate( hh.total_wealth, hh.interview_year, hh.quarter, upr_nominal_gdp )
    hh.house_value = uprate( hh.house_value, hh.interview_year, hh.quarter, upr_housing_oo )
    for (pid,person) in hh.people
        uprate!( pid, hh.interview_year, hh.quarter, person )
    end

end

function oldest_person( people :: People_Dict ) :: NamedTuple
    oldest = ( age=-999, pid=BigInt(0))
    for person in people
        if person.age > oldest.age
            oldest.age = person.age
            oldest.pid = person.pid
        end
    end
    oldest
end

PeopleArray = Vector{Person}

struct BenefitUnit{T<:Real}
    people :: People_Dict{T}
    head :: BigInt
    spouse :: BigInt
    adults :: Pid_Array
    children :: Pid_Array
end

BenefitUnits = Vector{BenefitUnit}
BUAllocation = Vector{PeopleArray}

function num_children( bu :: BenefitUnit ) :: Integer
    size( bu.children )[1]
end

function num_children( hh :: Household ) :: Integer
    count( hh, is_child )
end

"""
For Budget constraints: set the wage and make all the needed
corrections to employment status, hours worked, pensions(? pensions not done yet)
"""
function set_wage!( pers :: Person, gross :: Real, wage :: Real )
    pers.income[wages] = gross
    h = gross/wage
    pers.usual_hours_worked = h
    pers.employment_status = if h < 5
        Unemployed
    elseif h < 30
        Part_time_Employee
    else
        Full_time_Employee
    end
    # println( "made usual_hours_worked = $h gross=$gross wage=$wage pers.employment_status=$(pers.employment_status)")
end

#
# This creates a array of references to each person in the houshold, broken into
# benefit units using the default FRS/EFS benefit unit number.
#
function default_bu_allocation( hh :: Household ) :: BUAllocation
    bua = BUAllocation()
    nbus = 0
    # how many bus
    for (pid,person) in hh.people
        nbus = max( nbus, person.default_benefit_unit )
    end
    ## create benefit units
    for buno in 1:nbus
        push!( bua, PeopleArray())
    end
    sz = size( bua )[1]
    ## allocate peobuallocationple to them
    for (pid,person) in hh.people
        pp = bua[person.default_benefit_unit]
        push!( pp, person )
    end
    ## sort people in each by pid
    for buno in 1:nbus
        sort!( bua[buno], lt=(left,right)->isless(right.age,left.age))
    end
    bua
end

"""

Create a benefit unit from an array of people
really only needed for some unit tests.

FIXME: don't use outside of tests! likely mess up children and spouses.

"""
function make_benefit_unit(
    people :: PeopleArray,
    head :: BigInt,
    spouse:: BigInt = -1 ) :: BenefitUnit
    npeople = size( people )[1]
    T = typeof( people[1].hours_of_care_received )
    pd = People_Dict{T}()
    children = Pid_Array()
    adults = Pid_Array()
    push!( adults, head )
    if spouse > 0
        push!( adults, spouse )
    end
    palloc = spouse <= 0 ? 2 : 3
    for n in palloc:npeople
        push!(children, people[n].pid)
        @assert people[n].age <= 21 # vague idiot check
    end
    for pers in people
        pd[pers.pid] = pers
    end
    return BenefitUnit( pd, head, spouse, adults, children )
end

function child_pids( hh :: Household ) :: Pid_Array
    chlds = Pid_Array()
    for (pid,pers) in hh.people
        if is_child( pers )
            push!( chlds, pid )
        end
    end
    return chlds
end

#
# This creates a array of references to each person in the houshold, broken into
# benefit units using the default FRS/EFS benefit unit number.
#
function allocate_to_bus( T::Type, bua :: BUAllocation ) :: BenefitUnits
    nbus = size(bua)[1]
    bus = BenefitUnits(undef, nbus)
    for i in 1:nbus
        people = People_Dict{T}()
        head_pid :: BigInt = -1
        spouse_pid :: BigInt = -1
        children = Pid_Array()
        adults = Pid_Array()
        npeople = size( bua[i])[1]
        for p in 1:npeople
            person = bua[i][p]
            people[person.pid] = person
            if p == 1
                head_pid = person.pid
                push!( adults, head_pid )
            else
                # println( "on bu $i person $p relationships $(person.relationships)")
                reltohead = person.relationships[head_pid]
                if reltohead in [Spouse,Cohabitee,Civil_Partner]
                    spouse_pid = person.pid
                    push!( adults, spouse_pid )
                    # FIXME we need to remove these checks if
                    # we're using a non-default allocation to bus
                    # @assert person.age >= 16
                else
                    # @assert person.age <= 19
                    push!( children, person.pid )
                end
            end
        end
        new_bu = BenefitUnit( people, head_pid, spouse_pid, adults, children )
        bus[i] = new_bu
    end
    bus
end

function get_benefit_units(
    hh :: Household{T},
    allocator :: Function=default_bu_allocation ) :: BenefitUnits where T
    allocate_to_bus( T, allocator(hh))
end

function num_people( bu :: BenefitUnit )::Integer
    length( bu.people )
end

function num_adults( bu :: BenefitUnit )::Integer
    size( bu.adults )[1]
end

function adult_pids( hh :: Household )::Pid_Array
    pids = Pid_Array()
    for (pid,pers) in hh.people
        if ! is_child( pers )
            push!( pids, pid )
        end
    end
    return pids
end

function num_adults( hh )
    return size( adult_pids(hh))[1]
end

function num_people( hh :: Household ) :: Integer
    length( hh.people )
end

function get_head( bu :: BenefitUnit )::Person
    bu.people[bu.head]
end

function get_head( hh :: Household ) :: Person
    return hh.people[hh.head_of_household]
end

function is_head( bu :: BenefitUnit, pers :: Person  ) :: Bool
    return bu.head == pers.pid
end

function is_spouse( unit, pers :: Person  ) :: Bool
    head = get_head(unit)
    p = partner_of( head )
    if p === nothing
        return false
    end
    return p == pers.pid
end

function partner_of( pers :: Person ) :: Union{Nothing,BigInt}
    for (pid,rel) in pers.relationships
        if rel in [Spouse,Cohabitee]
            return pid
        end
    end
    return nothing
end

function get_spouse( hh :: Household ) ::Union{Nothing,Person}
    head = get_head( hh )
    p = partner_of( head )
    if p !== nothing
        return hh.people[p]
    end
    return nothing
end

function get_spouse( bu :: BenefitUnit )::Union{Nothing,Person}
    if bu.spouse <= 0
        return nothing
    end
    bu.people[bu.spouse]
end

function has_children( bu :: BenefitUnit ) :: Bool
    num_children( bu ) > 0
end

"""
relies on `is_child` being sensible
"""
function has_children( hh :: Household )::Bool
    for (pid, pers ) in hh.people
        if is_child( pers )
            return true
        end
    end
    return false
end

function is_lone_parent( bu :: BenefitUnit ) :: Bool
    return bu.spouse < 0 && has_children( bu )
end

function is_child( pers :: Person )
    # FIXME crude
    pers.is_standard_child
    # below won't work for hh  since
    # wouldn't be child if in own BU & we can't check that with 1 person
    #if pers.age <= 15
    #    return true
    #end
    #(pers.age <= 19) && (pers.employment_status in [Student])
end

#
# fixme just count people???
#
function is_lone_parent( hh :: Household,
    allocator :: Function=default_bu_allocation ) :: Bool

    bus = get_benefit_units( hh, allocator )
    if size(bus)[1] > 1
        return false
    end
    return bus[1].spouse < 0 && size( bus[1].children )[1] > 0
end

function is_single( bu :: BenefitUnit ) :: Bool
    num_people( bu ) == 1
end

function is_single( hh :: Household ) :: Bool
    num_people( hh ) == 1
end

"""
FIXME we're going to do this solely on benefit receipt/is_informal_carer flag
for now until we get the regressions done
we use historic benefits here since this uses actual data
"""
function pers_is_carer( pers :: Person, params ... ) :: Bool
    if has_non_z( pers.income, carers_allowance )
        return true
    end
    # FIXME check/parameterise this
    if (pers.is_informal_carer) && (pers.hours_of_care_given > 0)
        return true
    end
    return false
end

const DISABLE_BENEFITS = [
    severe_disability_allowance,
    attendance_allowance,
    incapacity_benefit,
    dlaself_care,
    dlamobility,
    personal_independence_payment_daily_living,
    personal_independence_payment_mobility]

"""

"""
function is_severe_disability( pers :: Person )
    return pers_is_disabled( pers ) ||
           pers.dla_self_care_type in [mid, high] ||
           pers.dla_mobility_type == high ||
           pers.attendance_allowance_type == high ||
           pers.pip_daily_living_type == enhanced_pip ||
           pers.pip_mobility_type == enhanced_pip ||
           pers.adls_are_reduced == reduced_a_lot ||
           pers.health_status == Very_Bad

end

"""
FIXME we're going to do this solely on benefit receipt
for now until we get the regressions done
we use historic benefits here since this uses actual data
"""
function pers_is_disabled( pers :: Person, params ... ) :: Bool
    if pers.registered_blind
        return true
    end
    if pers.employment_status in [Permanently_sick_or_disabled]
        return true
    end
    if pers.adls_are_reduced == reduced_a_lot
        return true
    end
    for dis in [mobility,
        dexterity,
        learning,
        memory,
        mental_health]
        if haskey(pers.disabilities,dis) && pers.disabilities[dis]
            return true
        end
    end

    for k in DISABLE_BENEFITS
        if has_non_z( pers.income, k )
            return true
        end
    end # loop
    return false
end

function empl_status_in( pers :: Person, statuses ...)
    return pers.employment_status in statuses
end

le_age( pers :: Person, age ... ) = pers.age <= age[1]

ge_age( pers :: Person, age ... ) = pers.age >= age[1]

between_ages( pers :: Person, age ... ) = age[1] >= pers.age <= age[2]

has_income( pers::Person, which :: Incomes_Type ) = haskey( pers.income, which )

function has_income( pers::Person, which ... )::Bool
    for k in which
        if has_income( pers, k )
            return true
        end
    end
    return false
end

function search( people :: People_Dict, func :: Function, params... ) :: Bool
    for ( pid,pers ) in people
        if func( pers, params... )
            return true
        end
    end
    return false
end

function count( people :: People_Dict, func :: Function, params... ) :: Int
    n = 0
    for ( pid,pers ) in people
        if func( pers, params... )
            n += 1
        end
    end
    return n
end

function search( bu :: BenefitUnit, func :: Function, params... ) :: Bool
    return search( bu.people, func, params... )
end

function search( hh :: Household, func :: Function, params ... ) :: Bool
    return search( hh.people, func, params... )
end

function count( bu :: BenefitUnit, func :: Function, params... ) :: Integer
     return count( bu.people, func, params... )
end

function count( hh :: Household, func :: Function, params ... ) :: Integer
    return count( hh.people, func, params... )
end

function on_mt_benefits( pers :: Person ) :: Bool
    if has_income( pers, [income_support, working_tax_credit, child_tax_credit, housing_benefit, universal_credit]...)
        return true
    end
    if pers.jsa_type == income_related_jsa
        return true
    end
    if pers.esa_type == income_related_jsa
        return true
    end
    return false
end

function on_mt_benefits( bu :: BenefitUnit ) :: Bool
    return search( bu, on_mt_benefits )
end

function on_mt_benefits( hh :: Household ) :: Bool
    return search( hh, on_mt_benefits )
end

function has_disabled_member( bu :: BenefitUnit ) :: Bool
    search( bu.people, pers_is_disabled )
end

function has_disabled_member( hh :: Household ) :: Bool
    search( hh.people, pers_is_disabled )
end

function has_carer_member( bu :: BenefitUnit ) :: Bool
    search( bu.people, pers_is_carer )
end

function has_carer_member( hh :: Household ) :: Bool
    search( hh.people, pers_is_carer )
end

function num_carers( bu :: BenefitUnit ) :: Int
    count( bu.people, pers_is_carer )
end

function carers( hh :: Household ) :: Int
    count( hh.people, pers_is_carer )
end

# simple diagnostic prints for testing allocation

function printpids( pers :: Person )
    println( "pid: $(pers.pid) age $(pers.age) sex $(pers.sex) default_benefit_unit $(pers.default_benefit_unit) ")
end

function printpids( bu::BenefitUnit)
      head = get_head( bu )
      print( "HEAD: ")
      printpids( head )
      spouse = get_spouse( bu )
      if spouse !== nothing
            print( "SPOUSE: ")
            printpids( spouse )

      end
      for chno in bu.children
            print( "CHILD_$(chno):")
            child = bu.people[chno]
            printpids( child )
      end
end

function printpids( people :: People_Dict )
    ks = sort(collect(keys(people)))
    for k in ks
        print( "PERSON_$k")
        printpids( people[k] )
    end
end

function printpids( buas::BUAllocation )
    nbus = size( buas )[1]
    for i in 1:nbus
        npeople = size( buas[i] )[1]
        for p in 1:npeople
            print( "PERS_$(p): ")
            printpids( buas[i][p])
        end
    end
end

function age_then( pers :: Person, when :: Int ) :: Int
    Utils.age_then( pers.age, when );
end

function age_then( pers :: Person, when :: Date = todays_date() ) :: Int
    Utils.age_then( pers.age, Dates.year( when ))
end

end # module


import RayTracing: num_dims, num_cells

"""
    MoCProblem{Dim,NRegions,NGroups,elType,G<:TrackGenerator,Q<:Quadrature,M} <: TransportProblem{Dim,NRegions,NGroups}

Defines a problem that ought to be solved using the Method of Characteristics, with
dimension `Dim`, number of flat source regions `NRegions`, number of energy groups
`NGroups`, ray tracing
"""
struct MoCProblem{Dim,NRegions,NGroups,elType,G<:TrackGenerator,Q<:Quadrature,Xs<:XSs} <: TransportProblem{Dim,NRegions,NGroups}
    nφ::Int
    nψ::Int

    trackgenerator::G
    quadrature::Q

    xss::Vector{Xs}
    fsr_tag::Vector{Int8}  # cell/element/fsr to xs tag (tag is the index in the xss array)

    function MoCProblem{Dim,NRegions,NGroups,elType}(
        nφ, nψ, tg::G, quad::Q, xss::Vector{X}, fsr_tag
    ) where {Dim,NRegions,NGroups,elType,G,Q,X}
        return new{Dim,NRegions,NGroups,elType,G,Q,X}(nφ, nψ, tg, quad, xss, fsr_tag)
    end
end

function MoCProblem(tg::TrackGenerator, polar_quadrature::PolarQuadrature, xss::Vector{<:XSs})
    @unpack mesh, azimuthal_quadrature, n_total_tracks = tg

    Dim = num_dims(mesh)

    NGroups = ngroups(first(xss))
    if !all(g -> isequal(g, NGroups), ngroups.(xss))
        error("all `CrossSections` *must* have the same number of energy groups.")
    end
    elType = promote_type(eltype.(xss)...)

    NRegions = num_cells(mesh)
    n_polar_2 = npolar2(polar_quadrature)

    nφ = NRegions * NGroups
    nψ = n_total_tracks * 2 * n_polar_2 * NGroups # 2 is the number of directions

    quadrature = Quadrature(azimuthal_quadrature, polar_quadrature)

    face_labeling = get_face_labeling(tg.mesh.model)
    tag_to_idx = Dict{Int8,Int8}()
    for (i, xs) in enumerate(xss)
        # given a xs name, find its gridap tag
        gridap_tag = get_tag_from_name(face_labeling, xs.name)

        # map gridap tag to xs idx
        push!(tag_to_idx, gridap_tag => i)
    end

    cell_tag = get_face_tag(face_labeling, Dim)  # gridap cell/element/fsr tags
    fsr_tag = similar(cell_tag)                  # NeutronTransport cell/element/fsr tags
    for (i, gridap_tag) in enumerate(cell_tag)
        fsr_tag[i] = tag_to_idx[gridap_tag]
    end

    return MoCProblem{Dim,NRegions,NGroups,elType}(nφ, nψ, tg, quadrature, xss, fsr_tag)
end

dimension(::MoCProblem{Dim}) where{Dim} = Dim
nregions(::MoCProblem{Dim,NRegions}) where{Dim,NRegions} = NRegions
ngroups(::MoCProblem{Dim,NRegions,NGroups}) where{Dim,NRegions,NGroups} = NGroups
eltype(::MoCProblem{Dim,NRegions,NGroups,elType}) where{Dim,NRegions,NGroups,elType} = elType

function show(io::IO, prob::MoCProblem)
    println(io, "  Problem Dimension: ", dimension(prob))
    println(io, "  Number of regions: ", nregions(prob))
    print(io,   "  Number of energy groups: ", ngroups(prob))
end

@inline function getxs(prob::MoCProblem, region_idx::Integer)
    @unpack xss, fsr_tag = prob
    @inbounds xs_idx = getindex(fsr_tag, region_idx) # get material tag
    @inbounds xs = getindex(xss, xs_idx) # get material XSs
    return xs
end

# si quiero que los primeros NRegions sean para el 1er group, etc...
# macro region_index(i, g)
#     return :((g - 1) * NRegions + i)
# end

# si quiero que esten intercalados
macro region_index(i, g)
    ex = quote
        ($i - 1) * NGroups + $g
    end
    return esc(ex)
end
# puedo hacerlas como function con @inline pero necesitan `groups` o `nfsr`

macro angular_index(t, d, p, g)
    ex = quote
        ($t - 1) * 2 * n_polar_2 * NGroups + $d * n_polar_2 * NGroups + ($p - 1) * NGroups + $g
    end
    return esc(ex)
end

macro reduced_angular_index(p, g)
    ex = quote
        ($p - 1) * NGroups + $g
    end
    return esc(ex)
end


using ReverseDiff: SpecialInstruction, istracked, increment_deriv!, deriv, TrackedArray
using ReverseDiff: InstructionTape, TrackedReal, seed!, unseed!, reverse_pass!, record!
import ReverseDiff: special_reverse_exec!, track

# if a value can't be tracked, return the untracked value instead, silently
track(value, tape::InstructionTape) = value

mutable struct GFBackpropParamsState
    trace::GFTrace
    score::TrackedReal
    tape::InstructionTape
    visitor::AddressVisitor
    tracked_params::Dict{Symbol,Any}
end

function GFBackpropParamsState(trace::GFTrace, tape, gf::GenFunction)
    tracked_params = Dict{Symbol,Any}()
    for (name, value) in gf.params
        tracked_params[name] = track(value, tape)
    end
    score = track(0., tape)
    GFBackpropParamsState(trace, score, tape, AddressVisitor(), tracked_params)
end

get_args_change(state::GFBackpropParamsState) = nothing
get_addr_change(state::GFBackpropParamsState, addr) = nothing
set_ret_change!(state::GFBackpropParamsState, value) = begin end

function read_param(state::GFBackpropParamsState, name::Symbol)
    value = state.tracked_params[name]
    value
end

function addr(state::GFBackpropParamsState, dist::Distribution{T}, args, addr) where {T}
    visit!(state.visitor, addr)
    call::CallRecord = get_primitive_call(state.trace, addr)
    retval::T = call.retval
    state.score += logpdf(dist, retval, args...)
    retval
end

struct BackpropParamsRecord
    generator::Generator
    subtrace::Any
end

function addr(state::GFBackpropParamsState, gen::Generator{T}, args, addr) where {T}
    visit!(state.visitor, addr)
    subtrace = get_subtrace(state.trace, addr)
    call::CallRecord = get_call_record(subtrace) # use the return value recorded in the trace
    retval::T = call.retval
    retval_maybe_tracked = track(retval, state.tape)
    if accepts_output_grad(gen) && !istracked(retval_maybe_tracked)
        error("Could not track return value at address $addr on AD tape.")
    end
    # some of the args may be tracked (see special_reverse_exec!)
    # note: we still need to run backprop_params on gen, even if it does not
    # accept an output gradient, because it may make random choices.
    record!(state.tape, SpecialInstruction,
        BackpropParamsRecord(gen, subtrace), (args...,), retval_maybe_tracked)
    retval_maybe_tracked
end

function splice(state::GFBackpropParamsState, gf::GenFunction, args::Tuple)
    exec(gf, state, args)
end

function maybe_track(arg, has_argument_grad::Bool, tape)
    has_argument_grad ? track(arg, tape) : arg
end

function backprop_params(gf::GenFunction, trace::GFTrace, retval_grad)
    tape = InstructionTape()
    state = GFBackpropParamsState(trace, tape, gf)
    call = get_call_record(trace)
    args = call.args
    args_maybe_tracked = (map(maybe_track, args, gf.has_argument_grads, fill(tape, length(args)))...,)
    retval_maybe_tracked = exec(gf, state, args_maybe_tracked)
    if gf.accepts_output_grad
        if !istracked(retval_maybe_tracked)
            error("Output of $gf accepts gradient but return value was not tracked on AD tape.")
        end
        if retval_grad != nothing
            # it is not an error if retval_grad is nothing, it means the gradient is zero
            # this is a convenience that
            deriv!(retval_maybe_tracked, retval_grad)
        end
    end
    seed!(state.score)
    reverse_pass!(tape)

    # increment the gradient accumulators for static parameters
    for (name, tracked) in state.tracked_params
        gf.params_grad[name] += deriv(tracked)
    end

    # return gradients with respect to inputs
    # NOTE: if a value isn't tracked the gradient is nothing
    input_grads::Tuple = (map((arg, has_grad) -> has_grad ? deriv(arg) : nothing,
                             args_maybe_tracked, gf.has_argument_grads)...,)
    input_grads
end

@noinline function special_reverse_exec!(instruction::SpecialInstruction{BackpropParamsRecord})
    record = instruction.func
    gen = record.generator
    args_maybe_tracked = instruction.input
    retval_maybe_tracked = instruction.output
    if accepts_output_grad(gen)
        @assert istracked(retval_maybe_tracked)
        retval_grad = deriv(retval_maybe_tracked)
    else
        retval_grad = nothing
    end
    arg_grads = backprop_params(gen, record.subtrace, retval_grad)
    for (arg, grad, has_grad) in zip(args_maybe_tracked, arg_grads, has_argument_grads(gen))
        if has_grad && istracked(arg)
            increment_deriv!(arg, grad)
        end
    end
    nothing
end


function ConstantProcessNoise(arg0::RealMatrix)
    return ConstantProcessNoise((RealMatrix,), arg0)
end

function ConstantProcessNoise(arg0::RealMatrix, arg1::RealMatrix)
    return ConstantProcessNoise((RealMatrix, RealMatrix), arg0, arg1)
end

function get_process_noise_matrix(obj::ConstantProcessNoise, arg0::SpacecraftState, arg1::SpacecraftState)
    return jcall(obj, "getProcessNoiseMatrix", RealMatrix, (SpacecraftState, SpacecraftState), arg0, arg1)
end


const MKL_DOMAIN_PARDISO = @compat Int32(4)

type MKLPardisoSolver <: AbstractPardisoSolver
    pt::Vector{Int}
    iparm::Vector{Int32}
    mtype::MatrixType
    solver::Solver
    phase::Phase
    msglvl::MessageLevel
    maxfct::Int32
    mnum::Int32
    perm::Vector{Int32}
end

function MKLPardisoSolver()
    if !MKL_PARDISO_LOADED
        error("mkl library was not loaded, unable to create solver")
    end

    pt = zeros(Int, 64)
    iparm = zeros(Int32, 64)
    mtype = REAL_NONSYM
    solver = DIRECT_SOLVER
    phase = ANALYSIS_NUM_FACT_SOLVE_REFINE
    msglvl = MESSAGE_LEVEL_OFF
    mnum = 1
    maxfct = 1
    perm = Int32[]

    MKLPardisoSolver(pt, iparm, mtype, solver,
                      phase, msglvl, maxfct, mnum, perm)
end

show(io::IO, ps::MKLPardisoSolver) = print(io, string("$MKLPardisoSolver:\n",
                                  "\tMatrix type: $(MATRIX_STRING[get_matrixtype(ps)])\n",
                                  "\tPhase: $(PHASE_STRING[get_phase(ps)])"))

set_nprocs!(ps::MKLPardisoSolver, n::Integer) = ccall(set_nthreads, Void, (Ptr{Int32},Ptr{Int32}), &(@compat Int32(n)), &MKL_DOMAIN_PARDISO)
get_nprocs(ps::MKLPardisoSolver) = ccall(get_nthreads, Int32, (Ptr{Int32},), &MKL_DOMAIN_PARDISO)

valid_phases(ps::MKLPardisoSolver) = keys(MKL_PHASES)
phases(ps::MKLPardisoSolver) = MKL_PHASES

function get_matrix(ps::MKLPardisoSolver, A, T)
    mtype = get_matrixtype(ps)

    if isposornegdef(mtype)
        T == :C && return conj(tril(A))
        return tril(A)
    end

    if !issymmetric(mtype)
        T == :C && return conj(A)
        return A
    end

    if mtype == COMPLEX_SYM
        T == :C && return conj(tril(A))
        return tril(A)
    end

    error("Unhandled matrix type")
end

function ccall_pardisoinit(ps::MKLPardisoSolver)
    ERR = Ref{Int32}(0)
    ccall(mkl_init, Void,
          (Ptr{Int}, Ptr{Int32}, Ptr{Int32}),
          ps.pt, &ps.mtype, ps.iparm)
    check_error(ps, ERR[])
end

function ccall_pardiso{Tv}(ps::MKLPardisoSolver, N, AA::Vector{Tv}, IA, JA,
                                   NRHS, B::VecOrMat{Tv}, X::VecOrMat{Tv})
    ERR = Ref{Int32}(0)
    ccall(mkl_pardiso_f, Void,
          (Ptr{Int}, Ptr{Int32}, Ptr{Int32}, Ptr{Int32}, Ptr{Int32},
           Ptr{Int32}, Ptr{Tv}, Ptr{Int32}, Ptr{Int32}, Ptr{Int32},
           Ptr{Int32}, Ptr{Int32}, Ptr{Int32}, Ptr{Tv}, Ptr{Tv},
           Ptr{Int32}),
          ps.pt, &ps.maxfct, &ps.mnum, &ps.mtype, &ps.phase,
          &N, AA, IA, JA, ps.perm,
          &NRHS, ps.iparm, &ps.msglvl, B, X,
          ERR)
    check_error(ps, ERR[])
end

function check_error(ps::MKLPardisoSolver, err::Integer)
    err != -1  || throw(PardisoException("Input inconsistent."))
    err != -2  || throw(PardisoException("Not enough memory."))
    err != -3  || throw(PardisoException("Reordering problem."))
    err != -4  || throw(PardisoPosDefException("Zero pivot, numerical fact. or iterative refinement problem."))
    err != -5  || throw(PardisoException("Unclassified (internal) error."))
    err != -6  || throw(PardisoException("Preordering failed (matrix types 11, 13 only)."))
    err != -7  || throw(PardisoException("Diagonal matrix is singular"))
    err != -8  || throw(PardisoException("32-bit integer overflow problem."))
    err != -9  || throw(PardisoException("Not enough memory for OOC."))
    err != -10 || throw(PardisoException("Error opening OOC files."))
    err != -11 || throw(PardisoException("Read/write error with OOC files."))
    err != -12 || throw(PardisoException("pardiso_64 called from 32-bit library"))
    return
end


export BrickletPiezoSpeakerV2Beep
struct BrickletPiezoSpeakerV2Beep
    frequency::Integer
    volume::Integer
    duration::Integer
    duration_remaining::Integer
end

export BrickletPiezoSpeakerV2Alarm
struct BrickletPiezoSpeakerV2Alarm
    start_frequency::Integer
    end_frequency::Integer
    step_size::Integer
    step_delay::Integer
    volume::Integer
    duration::Integer
    duration_remaining::Integer
    current_frequency::Integer
end

export BrickletPiezoSpeakerV2SPITFPErrorCount
struct BrickletPiezoSpeakerV2SPITFPErrorCount
    error_count_ack_checksum::Integer
    error_count_message_checksum::Integer
    error_count_frame::Integer
    error_count_overflow::Integer
end

export BrickletPiezoSpeakerV2Identity
struct BrickletPiezoSpeakerV2Identity
    uid::String
    connected_uid::String
    position::Char
    hardware_version::Vector{Integer}
    firmware_version::Vector{Integer}
    device_identifier::Integer
end

export BrickletPiezoSpeakerV2
"""
Creates beep and alarm with configurable volume and frequency
"""
mutable struct BrickletPiezoSpeakerV2 <: TinkerforgeDevice
    ipcon::IPConnection
    deviceInternal::PyObject

    """
    Creates an object with the unique device ID *uid* and adds it to
    the IP Connection *ipcon*.
    """
    function BrickletPiezoSpeakerV2(uid::String, ipcon::IPConnection)
        package = pyimport("tinkerforge.bricklet_piezo_speaker_v2")
        deviceInternal = package.BrickletPiezoSpeakerV2(uid, ipcon.ipconInternal)

        return new(ipcon, deviceInternal)
    end
end

export set_beep
"""
    $(SIGNATURES)

Beeps with the given frequency and volume for the duration.

A duration of 0 stops the current beep if any is ongoing.
A duration of 4294967295 results in an infinite beep.
"""
function set_beep(device::BrickletPiezoSpeakerV2, frequency, volume, duration)
    device.deviceInternal.set_beep(frequency, volume, duration)
end

export get_beep
"""
    $(SIGNATURES)

Returns the last beep settings as set by :func:`Set Beep`. If a beep is currently
running it also returns the remaining duration of the beep.

If the frequency or volume is updated during a beep (with :func:`Update Frequency`
or :func:`Update Volume`) this function returns the updated value.
"""
function get_beep(device::BrickletPiezoSpeakerV2)
    return device.deviceInternal.get_beep()
end

export set_alarm
"""
    $(SIGNATURES)

Creates an alarm (a tone that goes back and force between two specified frequencies).

The following parameters can be set:

* Start Frequency: Start frequency of the alarm.
* End Frequency: End frequency of the alarm.
* Step Size: Size of one step of the sweep between the start/end frequencies.
* Step Delay: Delay between two steps (duration of time that one tone is used in a sweep).
* Duration: Duration of the alarm.

A duration of 0 stops the current alarm if any is ongoing.
A duration of 4294967295 results in an infinite alarm.

Below you can find two sets of example settings that you can try out. You can use
these as a starting point to find an alarm signal that suits your application.

Example 1: 10 seconds of loud annoying fast alarm

* Start Frequency = 800
* End Frequency = 2000
* Step Size = 10
* Step Delay = 1
* Volume = 10
* Duration = 10000

Example 2: 10 seconds of soft siren sound with slow build-up

* Start Frequency = 250
* End Frequency = 750
* Step Size = 1
* Step Delay = 5
* Volume = 0
* Duration = 10000

The following conditions must be met:

* Start Frequency: has to be smaller than end frequency
* End Frequency: has to be bigger than start frequency
* Step Size: has to be small enough to fit into the frequency range
* Step Delay: has to be small enough to fit into the duration
"""
function set_alarm(device::BrickletPiezoSpeakerV2, start_frequency, end_frequency, step_size, step_delay, volume, duration)
    device.deviceInternal.set_alarm(start_frequency, end_frequency, step_size, step_delay, volume, duration)
end

export get_alarm
"""
    $(SIGNATURES)

Returns the last alarm settings as set by :func:`Set Alarm`. If an alarm is currently
running it also returns the remaining duration of the alarm as well as the
current frequency of the alarm.

If the volume is updated during an alarm (with :func:`Update Volume`)
this function returns the updated value.
"""
function get_alarm(device::BrickletPiezoSpeakerV2)
    return device.deviceInternal.get_alarm()
end

export update_volume
"""
    $(SIGNATURES)

Updates the volume of an ongoing beep or alarm.
"""
function update_volume(device::BrickletPiezoSpeakerV2, volume)
    device.deviceInternal.update_volume(volume)
end

export update_frequency
"""
    $(SIGNATURES)

Updates the frequency of an ongoing beep.
"""
function update_frequency(device::BrickletPiezoSpeakerV2, frequency)
    device.deviceInternal.update_frequency(frequency)
end

export get_spitfp_error_count
"""
    $(SIGNATURES)

Returns the error count for the communication between Brick and Bricklet.

The errors are divided into

* ACK checksum errors,
* message checksum errors,
* framing errors and
* overflow errors.

The errors counts are for errors that occur on the Bricklet side. All
Bricks have a similar function that returns the errors on the Brick side.
"""
function get_spitfp_error_count(device::BrickletPiezoSpeakerV2)
    return device.deviceInternal.get_spitfp_error_count()
end

export set_bootloader_mode
"""
    $(SIGNATURES)

Sets the bootloader mode and returns the status after the requested
mode change was instigated.

You can change from bootloader mode to firmware mode and vice versa. A change
from bootloader mode to firmware mode will only take place if the entry function,
device identifier and CRC are present and correct.

This function is used by Brick Viewer during flashing. It should not be
necessary to call it in a normal user program.
"""
function set_bootloader_mode(device::BrickletPiezoSpeakerV2, mode)
    return device.deviceInternal.set_bootloader_mode(mode)
end

export get_bootloader_mode
"""
    $(SIGNATURES)

Returns the current bootloader mode, see :func:`Set Bootloader Mode`.
"""
function get_bootloader_mode(device::BrickletPiezoSpeakerV2)
    return device.deviceInternal.get_bootloader_mode()
end

export set_write_firmware_pointer
"""
    $(SIGNATURES)

Sets the firmware pointer for :func:`Write Firmware`. The pointer has
to be increased by chunks of size 64. The data is written to flash
every 4 chunks (which equals to one page of size 256).

This function is used by Brick Viewer during flashing. It should not be
necessary to call it in a normal user program.
"""
function set_write_firmware_pointer(device::BrickletPiezoSpeakerV2, pointer)
    device.deviceInternal.set_write_firmware_pointer(pointer)
end

export write_firmware
"""
    $(SIGNATURES)

Writes 64 Bytes of firmware at the position as written by
:func:`Set Write Firmware Pointer` before. The firmware is written
to flash every 4 chunks.

You can only write firmware in bootloader mode.

This function is used by Brick Viewer during flashing. It should not be
necessary to call it in a normal user program.
"""
function write_firmware(device::BrickletPiezoSpeakerV2, data)
    return device.deviceInternal.write_firmware(data)
end

export set_status_led_config
"""
    $(SIGNATURES)

Sets the status LED configuration. By default the LED shows
communication traffic between Brick and Bricklet, it flickers once
for every 10 received data packets.

You can also turn the LED permanently on/off or show a heartbeat.

If the Bricklet is in bootloader mode, the LED is will show heartbeat by default.
"""
function set_status_led_config(device::BrickletPiezoSpeakerV2, config)
    device.deviceInternal.set_status_led_config(config)
end

export get_status_led_config
"""
    $(SIGNATURES)

Returns the configuration as set by :func:`Set Status LED Config`
"""
function get_status_led_config(device::BrickletPiezoSpeakerV2)
    return device.deviceInternal.get_status_led_config()
end

export get_chip_temperature
"""
    $(SIGNATURES)

Returns the temperature as measured inside the microcontroller. The
value returned is not the ambient temperature!

The temperature is only proportional to the real temperature and it has bad
accuracy. Practically it is only useful as an indicator for
temperature changes.
"""
function get_chip_temperature(device::BrickletPiezoSpeakerV2)
    return device.deviceInternal.get_chip_temperature()
end

export reset
"""
    $(SIGNATURES)

Calling this function will reset the Bricklet. All configurations
will be lost.

After a reset you have to create new device objects,
calling functions on the existing ones will result in
undefined behavior!
"""
function reset(device::BrickletPiezoSpeakerV2)
    device.deviceInternal.reset()
end

export write_uid
"""
    $(SIGNATURES)

Writes a new UID into flash. If you want to set a new UID
you have to decode the Base58 encoded UID string into an
integer first.

We recommend that you use Brick Viewer to change the UID.
"""
function write_uid(device::BrickletPiezoSpeakerV2, uid)
    device.deviceInternal.write_uid(uid)
end

export read_uid
"""
    $(SIGNATURES)

Returns the current UID as an integer. Encode as
Base58 to get the usual string version.
"""
function read_uid(device::BrickletPiezoSpeakerV2)
    return device.deviceInternal.read_uid()
end

export get_identity
"""
    $(SIGNATURES)

Returns the UID, the UID where the Bricklet is connected to,
the position, the hardware and firmware version as well as the
device identifier.

The position can be 'a', 'b', 'c', 'd', 'e', 'f', 'g' or 'h' (Bricklet Port).
A Bricklet connected to an :ref:`Isolator Bricklet <isolator_bricklet>` is always at
position 'z'.

The device identifier numbers can be found :ref:`here <device_identifier>`.
|device_identifier_constant|
"""
function get_identity(device::BrickletPiezoSpeakerV2)
    return device.deviceInternal.get_identity()
end

export register_callback
"""
Registers the given *function* with the given *callback_id*.
"""
function register_callback(device::BrickletPiezoSpeakerV2, callback_id, function_)
    if isnothing(function_)
        device.registered_callbacks.pop(callback_id, None)
    else
        device.registered_callbacks[callback_id] = function_
    end
end


using GridapMHD: hunt

hunt(
  nc=(4,4),
  L=1.0,
  B=(0.,50.,0.),
  debug=false,
  vtk=true,
  nsums=50,
  path=".",
  nruns=1,
  verbose=true,
  title="warmup_hunt",
  solver=:petsc,
  kmap=1,
  petsc_options="-snes_monitor -ksp_error_if_not_converged true -ksp_converged_reason -ksp_type preonly -pc_type lu -pc_factor_mat_solver_type mumps"
)

hunt(
  nc=(4,4),
  np=(1,1),
  backend=:mpi,
  L=1.0,
  B=(0.,50.,0.),
  debug=false,
  vtk=true,
  nsums=50,
  path=".",
  nruns=1,
  verbose=true,
  title="warmup_hunt_mpi",
  solver=:petsc,
  kmap=1,
  petsc_options="-snes_monitor -ksp_error_if_not_converged true -ksp_converged_reason -ksp_type preonly -pc_type lu -pc_factor_mat_solver_type mumps"
)


module TestTraits

# using Revise
using MLJBase
import MLJBase
using Test

mutable struct Dummy <: Deterministic{Int} end

MLJBase.target_kind(::Type{Dummy}) = :numeric
MLJBase.target_quantity(::Type{Dummy}) = :univariate
MLJBase.inputs_can_be(::Type{Dummy}) = [:numeric, ]
MLJBase.is_pure_julia(::Type{Dummy}) = :yes
MLJBase.package_name(::Type{Dummy}) = "MultivariateStats"
MLJBase.package_uuid(::Type{Dummy}) = "6f286f6a-111f-5878-ab1e-185364afe411"
MLJBase.package_url(::Type{Dummy}) = "https://mickey.mouse.org"

d = Dict(:model_name => "TestTraits.Dummy",
         :is_pure_julia => :yes,
         :package_uuid  => "6f286f6a-111f-5878-ab1e-185364afe411",
         :package_name  => "MultivariateStats",
         :target_is     => Symbol[:deterministic, :numeric, :univariate],
         :inputs_can_be => Symbol[:numeric],
         :package_url   => "https://mickey.mouse.org")

@show info(Dummy)[:model_name]
# @test info(Dummy) == d
# @test info(Dummy()) == d

end
true


import RDates: WeekendCalendar, @rd_str, setcachedcalendar!, calendar
import Dates: Date

setcachedcalendar!(CALENDARS, "US/SETTLEMENT", WeekendCalendar() + RuleBasedCalendar([
    # New Year's Day
    PeriodicCalendarRule(rd"1y", rd"1JAN@WEEKEND[NEAR]"; cal_mgr=CALENDARS),
    # New Year's Day (handle case where day falls back into this year)
    PeriodicCalendarRule(rd"1y", rd"1y+1JAN@WEEKEND[NEAR]"; cal_mgr=CALENDARS),
    # # Martin Luther King Jr Birthday
    WindowCalendarRule(PeriodicCalendarRule(rd"1y", rd"1JAN+3rd MON"); from=Date(1983,1,1)),
    # # Washingthon's Birthday
    PeriodicCalendarRule(rd"1y", rd"1FEB+3rd MON"),
    # # Memorial Day
    PeriodicCalendarRule(rd"1y", rd"1MAY+Last MON"),
    # # Juneteenth
    WindowCalendarRule(PeriodicCalendarRule(rd"1y", rd"19JUN@WEEKEND[NEAR]"; cal_mgr=CALENDARS); from=Date(2021,1,1)),
    # # Independence Day
    PeriodicCalendarRule(rd"1y", rd"4JUL@WEEKEND[NEAR]"; cal_mgr=CALENDARS),
    # # Labor Day
    PeriodicCalendarRule(rd"1y", rd"1SEP+1st MON"),
    # # Columbus Day
    PeriodicCalendarRule(rd"1y", rd"1OCT+2nd MON"),
    # # Veteran's Day
    PeriodicCalendarRule(rd"1y", rd"11NOV@WEEKEND[NEAR]"; cal_mgr=CALENDARS),
    # # Thanksgiving Day
    PeriodicCalendarRule(rd"1y", rd"1NOV+4th THU"),
    # # Christmas
    PeriodicCalendarRule(rd"1y", rd"25DEC@WEEKEND[NEAR]"; cal_mgr=CALENDARS),
]))

setcachedcalendar!(CALENDARS, "NEW YORK", calendar(CALENDARS, "US/SETTLEMENT"))


struct Scene
    objects::Vector{Object}

    function Scene()
        new(Vector{Object}[])
    end
end

function hit(scene::Scene, ray::Ray, range::NamedTuple{(:min, :max), Tuple{T, T}} where T <: AbstractFloat)::Tuple{Bool, Union{Nothing, HitPoint}, Union{Nothing, Material}}
    hit_anything::Bool = false
    closest_so_far::AbstractFloat = range.max
    local hp::HitPoint
    local m::Material
    for object in scene.objects
        succeeded::Bool, t::AbstractFloat, hp0::Union{Nothing, HitPoint}, m0::Union{Nothing, Material} = hit(object, ray, (min=range.min, max=closest_so_far))
        if succeeded
            hit_anything = true
            closest_so_far = t;
            hp = hp0
            m = m0
        end
    end
    if hit_anything
        return true, hp, m
    else
        return false, nothing, nothing
    end
end

function addObject(scene::Scene, object::Object)
    push!(scene.objects, object)
end


# everything related to a DG semidiscretization in 2D using MPI,
# currently limited to Lobatto-Legendre nodes

# TODO: MPI dimension agnostic
# TODO: MPI, adapt to different real types (and AD!)
mutable struct MPICache
  mpi_neighbor_ranks::Vector{Int}
  mpi_neighbor_interfaces::Vector{Vector{Int}}
  mpi_send_buffers::Vector{Vector{Float64}}
  mpi_recv_buffers::Vector{Vector{Float64}}
  mpi_send_requests::Vector{MPI.Request}
  mpi_recv_requests::Vector{MPI.Request}
  n_elements_by_rank::OffsetArray{Int, 1, Array{Int, 1}}
  n_elements_global::Int
  first_element_global_id::Int
end

function MPICache()
  mpi_neighbor_ranks = Vector{Int}(undef, 0)
  mpi_neighbor_interfaces = Vector{Vector{Int}}(undef, 0)
  mpi_send_buffers = Vector{Vector{Float64}}(undef, 0)
  mpi_recv_buffers = Vector{Vector{Float64}}(undef, 0)
  mpi_send_requests = Vector{MPI.Request}(undef, 0)
  mpi_recv_requests = Vector{MPI.Request}(undef, 0)
  n_elements_by_rank = OffsetArray(Vector{Int}(undef, 0), 0:-1)
  n_elements_global = 0
  first_element_global_id = 0

  MPICache(mpi_neighbor_ranks, mpi_neighbor_interfaces,
           mpi_send_buffers, mpi_recv_buffers,
           mpi_send_requests, mpi_recv_requests,
           n_elements_by_rank, n_elements_global,
           first_element_global_id)
end

# TODO: MPI dimension agnostic
function start_mpi_receive!(mpi_cache::MPICache)

  for (index, d) in enumerate(mpi_cache.mpi_neighbor_ranks)
    mpi_cache.mpi_recv_requests[index] = MPI.Irecv!(
      mpi_cache.mpi_recv_buffers[index], d, d, mpi_comm())
  end

  return nothing
end

# TODO: MPI dimension agnostic
function start_mpi_send!(mpi_cache::MPICache, mesh, equations, dg, cache)
  data_size = nvariables(equations) * nnodes(dg)^(ndims(mesh) - 1)

  for d in 1:length(mpi_cache.mpi_neighbor_ranks)
    send_buffer = mpi_cache.mpi_send_buffers[d]

    for (index, interface) in enumerate(mpi_cache.mpi_neighbor_interfaces[d])
      first = (index - 1) * data_size + 1
      last =  (index - 1) * data_size + data_size

      if cache.mpi_interfaces.remote_sides[interface] == 1 # local element in positive direction
        @views send_buffer[first:last] .= vec(cache.mpi_interfaces.u[2, :, :, interface])
      else # local element in negative direction
        @views send_buffer[first:last] .= vec(cache.mpi_interfaces.u[1, :, :, interface])
      end
    end
  end

  # Start sending
  for (index, d) in enumerate(mpi_cache.mpi_neighbor_ranks)
    mpi_cache.mpi_send_requests[index] = MPI.Isend(
      mpi_cache.mpi_send_buffers[index], d, mpi_rank(), mpi_comm())
  end

  return nothing
end

# TODO: MPI dimension agnostic
function finish_mpi_send!(mpi_cache::MPICache)
  MPI.Waitall!(mpi_cache.mpi_send_requests)
end

# TODO: MPI dimension agnostic
function finish_mpi_receive!(mpi_cache::MPICache, mesh, equations, dg, cache)
  data_size = nvariables(equations) * nnodes(dg)^(ndims(mesh) - 1)

  # Start receiving and unpack received data until all communication is finished
  d, _ = MPI.Waitany!(mpi_cache.mpi_recv_requests)
  while d != 0
    recv_buffer = mpi_cache.mpi_recv_buffers[d]

    for (index, interface) in enumerate(mpi_cache.mpi_neighbor_interfaces[d])
      first = (index - 1) * data_size + 1
      last =  (index - 1) * data_size + data_size

      if cache.mpi_interfaces.remote_sides[interface] == 1 # local element in positive direction
        @views vec(cache.mpi_interfaces.u[1, :, :, interface]) .= recv_buffer[first:last]
      else # local element in negative direction
        @views vec(cache.mpi_interfaces.u[2, :, :, interface]) .= recv_buffer[first:last]
      end
    end

    d, _ = MPI.Waitany!(mpi_cache.mpi_recv_requests)
  end

  return nothing
end

# This method is called when a SemidiscretizationHyperbolic is constructed.
# It constructs the basic `cache` used throughout the simulation to compute
# the RHS etc.
function create_cache(mesh::ParallelTreeMesh{2}, equations,
                      dg::DG, RealT, uEltype)
  # Get cells for which an element needs to be created (i.e. all leaf cells)
  leaf_cell_ids = local_leaf_cells(mesh.tree)

  elements = init_elements(leaf_cell_ids, mesh, equations, dg.basis, RealT, uEltype)

  interfaces = init_interfaces(leaf_cell_ids, mesh, elements)

  mpi_interfaces = init_mpi_interfaces(leaf_cell_ids, mesh, elements)

  boundaries = init_boundaries(leaf_cell_ids, mesh, elements)

  mortars = init_mortars(leaf_cell_ids, mesh, elements, dg.mortar)

  mpi_cache = init_mpi_cache(mesh, elements, mpi_interfaces, nvariables(equations), nnodes(dg))

  cache = (; elements, interfaces, mpi_interfaces, boundaries, mortars,
             mpi_cache)

  # Add specialized parts of the cache required to compute the volume integral etc.
  cache = (;cache..., create_cache(mesh, equations, dg.volume_integral, dg, uEltype)...)
  cache = (;cache..., create_cache(mesh, equations, dg.mortar, uEltype)...)

  return cache
end

function init_mpi_cache(mesh, elements, mpi_interfaces, nvars, nnodes)
  mpi_cache = MPICache()

  init_mpi_cache!(mpi_cache, mesh, elements, mpi_interfaces, nvars, nnodes)
  return mpi_cache
end

function init_mpi_cache!(mpi_cache, mesh, elements, mpi_interfaces, nvars, nnodes)
  mpi_neighbor_ranks, mpi_neighbor_interfaces =
    init_mpi_neighbor_connectivity(elements, mpi_interfaces, mesh)

  mpi_send_buffers, mpi_recv_buffers, mpi_send_requests, mpi_recv_requests =
    init_mpi_data_structures(mpi_neighbor_interfaces, ndims(mesh), nvars, nnodes)

  # Determine local and total number of elements
  n_elements_by_rank = Vector{Int}(undef, mpi_nranks())
  n_elements_by_rank[mpi_rank() + 1] = nelements(elements)
  MPI.Allgather!(MPI.UBuffer(n_elements_by_rank, 1), mpi_comm())
  n_elements_by_rank = OffsetArray(n_elements_by_rank, 0:(mpi_nranks() - 1))
  n_elements_global = MPI.Allreduce(nelements(elements), +, mpi_comm())
  @assert n_elements_global == sum(n_elements_by_rank) "error in total number of elements"

  # Determine the global element id of the first element
  first_element_global_id = MPI.Exscan(nelements(elements), +, mpi_comm())
  if mpi_isroot()
    # With Exscan, the result on the first rank is undefined
    first_element_global_id = 1
  else
    # On all other ranks we need to add one, since Julia has one-based indices
    first_element_global_id += 1
  end
  # TODO reuse existing structures
  @pack! mpi_cache = mpi_neighbor_ranks, mpi_neighbor_interfaces,
                     mpi_send_buffers, mpi_recv_buffers,
                     mpi_send_requests, mpi_recv_requests,
                     n_elements_by_rank, n_elements_global,
                     first_element_global_id
end

# Initialize connectivity between MPI neighbor ranks
function init_mpi_neighbor_connectivity(elements, mpi_interfaces, mesh::TreeMesh2D)
  tree = mesh.tree

  # Determine neighbor ranks and sides for MPI interfaces
  neighbor_ranks = fill(-1, nmpiinterfaces(mpi_interfaces))
  # The global interface id is the smaller of the (globally unique) neighbor cell ids, multiplied by
  # number of directions (2 * ndims) plus direction minus one
  global_interface_ids = fill(-1, nmpiinterfaces(mpi_interfaces))
  for interface_id in 1:nmpiinterfaces(mpi_interfaces)
    orientation = mpi_interfaces.orientations[interface_id]
    remote_side = mpi_interfaces.remote_sides[interface_id]
    # Direction is from local cell to remote cell
    if orientation == 1 # MPI interface in x-direction
      if remote_side == 1 # remote cell on the "left" of MPI interface
        direction = 1
      else # remote cell on the "right" of MPI interface
        direction = 2
      end
    else # MPI interface in y-direction
      if remote_side == 1 # remote cell on the "left" of MPI interface
        direction = 3
      else # remote cell on the "right" of MPI interface
        direction = 4
      end
    end
    local_element_id = mpi_interfaces.local_element_ids[interface_id]
    local_cell_id = elements.cell_ids[local_element_id]
    remote_cell_id = tree.neighbor_ids[direction, local_cell_id]
    neighbor_ranks[interface_id] = tree.mpi_ranks[remote_cell_id]
    if local_cell_id < remote_cell_id
      global_interface_ids[interface_id] = 2 * ndims(tree) * local_cell_id + direction - 1
    else
      global_interface_ids[interface_id] = (2 * ndims(tree) * remote_cell_id +
                                            opposite_direction(direction) - 1)
    end
  end

  # Get sorted, unique neighbor ranks
  mpi_neighbor_ranks = unique(sort(neighbor_ranks))

  # Sort interfaces by global interface id
  p = sortperm(global_interface_ids)
  neighbor_ranks .= neighbor_ranks[p]
  interface_ids = collect(1:nmpiinterfaces(mpi_interfaces))[p]

  # For each neighbor rank, init connectivity data structures
  mpi_neighbor_interfaces = Vector{Vector{Int}}(undef, length(mpi_neighbor_ranks))
  for (index, d) in enumerate(mpi_neighbor_ranks)
    mpi_neighbor_interfaces[index] = interface_ids[findall(x->(x == d), neighbor_ranks)]
  end

  # Sanity check that we counted all interfaces exactly once
  @assert sum(length(v) for v in mpi_neighbor_interfaces) == nmpiinterfaces(mpi_interfaces)

  return mpi_neighbor_ranks, mpi_neighbor_interfaces
end

# TODO: MPI dimension agnostic
# Initialize MPI data structures
function init_mpi_data_structures(mpi_neighbor_interfaces, ndims, nvars, n_nodes)
  data_size = nvars * n_nodes^(ndims - 1)
  mpi_send_buffers = Vector{Vector{Float64}}(undef, length(mpi_neighbor_interfaces))
  mpi_recv_buffers = Vector{Vector{Float64}}(undef, length(mpi_neighbor_interfaces))
  for index in 1:length(mpi_neighbor_interfaces)
    mpi_send_buffers[index] = Vector{Float64}(undef, length(mpi_neighbor_interfaces[index]) * data_size)
    mpi_recv_buffers[index] = Vector{Float64}(undef, length(mpi_neighbor_interfaces[index]) * data_size)
  end

  mpi_send_requests = Vector{MPI.Request}(undef, length(mpi_neighbor_interfaces))
  mpi_recv_requests = Vector{MPI.Request}(undef, length(mpi_neighbor_interfaces))

  return mpi_send_buffers, mpi_recv_buffers, mpi_send_requests, mpi_recv_requests
end

function rhs!(du, u, t,
              mesh::ParallelTreeMesh{2}, equations,
              initial_condition, boundary_conditions, source_terms,
              dg::DG, cache)
  # Start to receive MPI data
  @timed timer() "start MPI receive" start_mpi_receive!(cache.mpi_cache)

  # Prolong solution to MPI interfaces
  @timed timer() "prolong2mpiinterfaces" prolong2mpiinterfaces!(
    cache, u, mesh, equations, dg)

  # Start to send MPI data
  @timed timer() "start MPI send" start_mpi_send!(
    cache.mpi_cache, mesh, equations, dg, cache)

  # Reset du
  @timed timer() "reset ∂u/∂t" du .= zero(eltype(du))

  # Calculate volume integral
  @timed timer() "volume integral" calc_volume_integral!(
    du, u, mesh,
    have_nonconservative_terms(equations), equations,
    dg.volume_integral, dg, cache)

  # Prolong solution to interfaces
  # TODO: Taal decide order of arguments, consistent vs. modified cache first?
  @timed timer() "prolong2interfaces" prolong2interfaces!(
    cache, u, mesh, equations, dg)

  # Calculate interface fluxes
  @timed timer() "interface flux" calc_interface_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg, cache)

  # Prolong solution to boundaries
  @timed timer() "prolong2boundaries" prolong2boundaries!(
    cache, u, mesh, equations, dg)

  # Calculate boundary fluxes
  @timed timer() "boundary flux" calc_boundary_flux!(
    cache, t, boundary_conditions, mesh, equations, dg)

  # Prolong solution to mortars
  @timed timer() "prolong2mortars" prolong2mortars!(
    cache, u, mesh, equations, dg.mortar, dg)

  # Calculate mortar fluxes
  @timed timer() "mortar flux" calc_mortar_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg.mortar, dg, cache)

  # Finish to receive MPI data
  @timed timer() "finish MPI receive" finish_mpi_receive!(
    cache.mpi_cache, mesh, equations, dg, cache)

  # Calculate MPI interface fluxes
  @timed timer() "MPI interface flux" calc_mpi_interface_flux!(
    cache.elements.surface_flux_values, mesh,
    have_nonconservative_terms(equations), equations,
    dg, cache)

  # Calculate surface integrals
  @timed timer() "surface integral" calc_surface_integral!(
    du, mesh, equations, dg, cache)

  # Apply Jacobian from mapping to reference element
  @timed timer() "Jacobian" apply_jacobian!(
    du, mesh, equations, dg, cache)

  # Calculate source terms
  @timed timer() "source terms" calc_sources!(
    du, u, t, source_terms, equations, dg, cache)

  # Finish to send MPI data
  @timed timer() "finish MPI send" finish_mpi_send!(cache.mpi_cache)

  return nothing
end

function prolong2mpiinterfaces!(cache, u,
                                mesh::ParallelTreeMesh{2},
                                equations, dg::DG)
  @unpack mpi_interfaces = cache

  @threaded for interface in eachmpiinterface(dg, cache)
    local_element = mpi_interfaces.local_element_ids[interface]

    if mpi_interfaces.orientations[interface] == 1 # interface in x-direction
      if mpi_interfaces.remote_sides[interface] == 1 # local element in positive direction
        for j in eachnode(dg), v in eachvariable(equations)
          mpi_interfaces.u[2, v, j, interface] = u[v,          1, j, local_element]
        end
      else # local element in negative direction
        for j in eachnode(dg), v in eachvariable(equations)
          mpi_interfaces.u[1, v, j, interface] = u[v, nnodes(dg), j, local_element]
        end
      end
    else # interface in y-direction
      if mpi_interfaces.remote_sides[interface] == 1 # local element in positive direction
        for i in eachnode(dg), v in eachvariable(equations)
          mpi_interfaces.u[2, v, i, interface] = u[v, i,          1, local_element]
        end
      else # local element in negative direction
        for i in eachnode(dg), v in eachvariable(equations)
          mpi_interfaces.u[1, v, i, interface] = u[v, i, nnodes(dg), local_element]
        end
      end
    end
  end

  return nothing
end

function calc_mpi_interface_flux!(surface_flux_values,
                                  mesh::ParallelTreeMesh{2},
                                  nonconservative_terms::Val{false}, equations,
                                  dg::DG, cache)
  @unpack surface_flux = dg
  @unpack u, local_element_ids, orientations, remote_sides = cache.mpi_interfaces

  @threaded for interface in eachmpiinterface(dg, cache)
    # Get local neighboring element
    element = local_element_ids[interface]

    # Determine interface direction with respect to element:
    if orientations[interface] == 1 # interface in x-direction
      if remote_sides[interface] == 1 # local element in positive direction
        direction = 1
      else # local element in negative direction
        direction = 2
      end
    else # interface in y-direction
      if remote_sides[interface] == 1 # local element in positive direction
        direction = 3
      else # local element in negative direction
        direction = 4
      end
    end

    for i in eachnode(dg)
      # Call pointwise Riemann solver
      u_ll, u_rr = get_surface_node_vars(u, equations, dg, i, interface)
      flux = surface_flux(u_ll, u_rr, orientations[interface], equations)

      # Copy flux to local element storage
      for v in eachvariable(equations)
        surface_flux_values[v, i, direction, element] = flux[v]
      end
    end
  end

  return nothing
end