using Statistics
using ..Transformers.Basic: gather

using ZygoteRules: @adjoint, pullback

using AbstractTrees
import AbstractTrees: children

const ACT2FN = (gelu = gelu, relu = relu, swish = swish, gelu_new = gelu, mish = mish)

abstract type THModule end
abstract type HGFPreTrainedModel <: THModule end

# torch.nn.LayerNorm

struct FakeTHLayerNorm{T<:AbstractArray} <: THModule
  eps::Float32
  weight::T
  bias::T
end

Functors.functor(::Type{<:FakeTHLayerNorm}, layernorm) = (weight = layernorm.weight, bias = layernorm.bias), y -> FakeTHLayerNorm(layernorm.eps, y...)

(ln::FakeTHLayerNorm)(x) = ln.weight .* Flux.normalise(x, dims=1, ϵ=ln.eps) .+ ln.bias

function load_state!(layer::FakeTHLayerNorm, state)
  for k in keys(state)
    if k == :gamma
      nk = :weight
    elseif k == :beta
      nk = :bias
    else
      nk = k
    end
    load_state!(getfield(layer, nk), getfield(state, k))
  end
end

# torch.nn.Linear

struct FakeTHLinear{W<:AbstractArray, B<:Union{Nothing, AbstractArray}} <: THModule
  weight::W
  bias::B
end

FakeTHLinear(w) = FakeTHLinear(w, nothing)

_has_bias(::FakeTHLinear{W, Nothing}) where {W<:AbstractArray}= false
_has_bias(::FakeTHLinear) = true

function Functors.functor(::Type{<:FakeTHLinear}, linear)
  (_has_bias(linear) ? (weight = linear.weight, bias = linear.bias) : (weight = linear.weight,)),
  y -> FakeTHLinear(y...)
end

(l::FakeTHLinear)(x::AbstractMatrix) = _has_bias(l) ? l.weight * x .+ l.bias : l.weight * x
function (l::FakeTHLinear)(x::AbstractArray)
  old_size = size(x)
  new_size = Base.setindex(old_size, size(l.weight, 1), 1)

  new_x = reshape(x, old_size[1], :)
  y = l(new_x)
  return reshape(y, new_size)
end

# torch.nn.Embedding

struct FakeTHEmbedding{T<:AbstractArray} <: THModule
  pad_idx::Int
  weight::T
end

function FakeTHEmbedding(pad_idx, weight)
  @assert pad_idx <= size(weight, 2)
  if !iszero(pad_idx)
    @view(weight[:, pad_idx]) .= 0
  end

  FakeTHEmbedding(pad_idx, weight)
end

Functors.functor(::Type{<:FakeTHEmbedding}, embedding) = (weight = embedding.weight,), y -> FakeTHEmbedding(embedding.pad_idx, y...)

#TODO: refactor basic/embeds/gather.jl
_padded_gather(w, x, padded_idx) = gather(w, x)
@adjoint function _padded_gather(w, x, padded_idx)
  y, back = pullback(gather, w, x)
  return y, Δ -> begin
    Δ′, _ = back(Δ)
    if !iszero(padded_idx)
      @view(Δ′[:, padded_idx]) .= 0
    end
    return (Δ′, nothing, nothing)
  end
end

(e::FakeTHEmbedding)(x) = _padded_gather(e.weight, x, e.pad_idx)

function load_state!(layer::FakeTHEmbedding, state)
  load_state!(layer.weight, state.weight')
end

function get_state_dict(state, prefix, embedding::FakeTHEmbedding)
  param = Functors.functor(embedding)[1]
  k = :weight
  cprefix = isnothing(prefix) ? String(k) : join((prefix, k), '.')
  state[cprefix] = param.weight'
end

# torch.nn.ModuleList

struct FakeTHModuleList{N, T<:Tuple} <: THModule
  _modules::T
  FakeTHModuleList(ms) = new{length(ms), typeof(ms)}(ms)
end

FakeTHModuleList(ms...) = FakeTHModuleList(ms)
FakeTHModuleList(ms::Vector) = FakeTHModuleList(Tuple(ms))

Functors.functor(::Type{<:FakeTHModuleList}, modulelist) = modulelist._modules, y -> FakeTHModuleList(y)

Base.iterate(modulelist::FakeTHModuleList) = iterate(modulelist._modules)
Base.iterate(modulelist::FakeTHModuleList, i...) = iterate(modulelist._modules, i...)
Base.length(::FakeTHModuleList{N}) where N = N
Base.getindex(modulelist::FakeTHModuleList, i) = modulelist._modules[i]

function get_state_dict(state, prefix, modulelist::FakeTHModuleList)
  param = Functors.functor(modulelist)[1]
  for (i, layer) in enumerate(modulelist)
    cprefix = join((prefix, i-1), '.')
    get_state_dict(state, cprefix, layer)
  end
end

function load_state!(layer::FakeTHModuleList, state)
  for (layerᵢ, stateᵢ) in zip(layer, state)
    load_state!(layerᵢ, stateᵢ)
  end
end

# displaying layer in tree structure

function Base.show(io::IO, x::M; depth=10) where {M<:THModule}
  print_tree(_printnode, io, x, depth)
  io
end

children(m::THModule) = Tuple(pairs(first(Functors.functor(m))))
function children(::THModule, x::Pair{Symbol})
  v = x[2]
  v isa AbstractArray ? () : children(v)
end
children(::THModule, x::Pair{Int}) = children(x[2])

_printnode(io, x) = print(io, Base.typename(typeof(x)))
function _printnode(io::IO, p::Pair)
  print(io, p[1])
  print(io, " => ")
  v = p[2]
  if v isa AbstractArray
    summary(io, p[2])
  else
    print(io, Base.typename(typeof(v)))
  end
end

# transformers.modeling_utils.Conv1D

struct FakeHGFConv1D{W<:AbstractArray, B<:AbstractArray} <: THModule
  weight::W
  bias::B
end

@functor FakeHGFConv1D

(c::FakeHGFConv1D)(x::AbstractMatrix) = c.weight' * x .+ c.bias

function (c::FakeHGFConv1D)(x::AbstractArray)
  old_size = size(x)
  new_size = Base.setindex(old_size, size(c.weight, 2), 1)

  new_x = reshape(x, old_size[1], :)
  y = c(new_x)
  return reshape(y, new_size)
end

# transformers.modeling_utils.SequenceSummary

abstract type AbstractSummaryType end
struct LastSummary <: AbstractSummaryType end
struct FirstSummary <: AbstractSummaryType end
struct MeanSummary <: AbstractSummaryType end
struct CLSindexSummary <: AbstractSummaryType end

struct FakeHGFSequenceSummary{T<:AbstractSummaryType, S, F} <: THModule
  summary::S
  activation::F

  FakeHGFSequenceSummary{T}(s, a) where {T} = new{T, typeof(s), typeof(a)}(s, a)
end

SummaryType(s::Symbol) = (LastSummary, FirstSummary, MeanSummary, CLSindexSummary)[
  findfirst(isequal(s), (:last, :first, :mean, :cls_index))
]

summarytype(s::FakeHGFSequenceSummary{T}) where {T<:AbstractSummaryType} = T

Functors.functor(::Type{<:FakeHGFSequenceSummary}, ss) = (summary = ss.summary, activation = ss.activation), y -> FakeHGFSequenceSummary{summarytype(ss)}(y...)

summary_input(s::FakeHGFSequenceSummary{LastSummary}, x) = x[:, end, :, :]
summary_input(s::FakeHGFSequenceSummary{FirstSummary}, x) = x[:, 1, :, :]
summary_input(s::FakeHGFSequenceSummary{MeanSummary}, x) = mean(x, dims=2)

summary_input(s::FakeHGFSequenceSummary, x, _) = summary_input(s, x)
summary_input(s::FakeHGFSequenceSummary{CLSindexSummary}, x, ::Nothing) = x[:, end, :, :]
summary_input(s::FakeHGFSequenceSummary{CLSindexSummary}, x, cls_index) = gather(x, cls_index)

function (s::FakeHGFSequenceSummary)(hidden_states, cls_index)
  output = summary_input(s, hidden_states, cls_index)
  output = s.summary(output)
  output = s.activation(output)
  return output
end


module ParameterJuMP

using SparseArrays

import MutableArithmetics
const MA = MutableArithmetics

using JuMP
export index

export ParameterRef, Param, all_parameters, parametrized_dual_objective_value

# types
# ------------------------------------------------------------------------------
struct ParameterRef <: AbstractVariableRef
    ind::Int64 # local reference
    model::Model
end

# Reference to a constraint in which the parameter has coefficient coef
struct ParametrizedConstraintRef{C}
    cref::ConstraintRef{Model,C}
    coef::Float64
end

const CtrRef{F,S} = ConstraintRef{Model,MOI.ConstraintIndex{F,S},ScalarShape}
const SAF = MOI.ScalarAffineFunction{Float64}

mutable struct _ParameterData
    inds::Vector{Int64}
    next_ind::Int64
    current_values::Vector{Float64}
    future_values::Vector{Float64}
    # Whether the JuMP model is in sync with the value of the parameters
    sync::Bool
    # constraints where it is a RHS
    constraints_map::Dict{Int64,Vector{ParametrizedConstraintRef}}
    parameters_map_saf_in_eq::Dict{
        CtrRef{SAF,MOI.EqualTo{Float64}},
        GenericAffExpr{Float64,ParameterRef},
    }
    parameters_map_saf_in_le::Dict{
        CtrRef{SAF,MOI.LessThan{Float64}},
        GenericAffExpr{Float64,ParameterRef},
    }
    parameters_map_saf_in_ge::Dict{
        CtrRef{SAF,MOI.GreaterThan{Float64}},
        GenericAffExpr{Float64,ParameterRef},
    }
    names::Dict{ParameterRef,String}
    lazy::Bool
    no_duals::Bool
    dual_values::Vector{Float64}

    function _ParameterData()
        return new(
            Int64[],
            1,
            Float64[],
            Float64[],
            true,
            Dict{Int64,Vector{ConstraintRef}}(),
            Dict{CtrRef{SAF,MOI.EqualTo{Float64}},GenericAffExpr{Float64,ParameterRef}}(),
            Dict{CtrRef{SAF,MOI.LessThan{Float64}},GenericAffExpr{Float64,ParameterRef}}(),
            Dict{CtrRef{SAF,MOI.GreaterThan{Float64}},GenericAffExpr{Float64,ParameterRef}}(),
            Dict{ParameterRef,String}(),
            false,
            false,
            Float64[],
        )
    end
end

no_duals(data::_ParameterData) = data.no_duals
no_duals(model::Model) = no_duals(_getparamdata(model))

set_no_duals(model::Model) = set_no_duals(_getparamdata(model))
function set_no_duals(data::_ParameterData)
    if isempty(data.current_values)
        data.no_duals = true
    elseif no_duals(data)
        @warn "No duals mode is already activated"
    else
        error("Parameter JuMP's no duals mode can only be activated in empty models.")
    end
end

"""
    index(p::ParameterRef)::Int64

Return the internal index of the parameter `p`.
"""
JuMP.index(p::ParameterRef) = p.ind

lazy_duals(data::_ParameterData) = data.lazy
lazy_duals(model::Model) = lazy_duals(_getparamdata(model))

set_lazy_duals(model::Model) = set_lazy_duals(_getparamdata(model))
function set_lazy_duals(data::_ParameterData)
    if isempty(data.current_values)
        data.lazy = true
    elseif lazy_duals(data)
        @warn "Lazy mode is already activated"
    else
        error("Parameter JuMP's lazy mode can only be activated in empty models.")
    end
end

set_not_lazy_duals(model::Model) = set_not_lazy_duals(_getparamdata(model))
function set_not_lazy_duals(data::_ParameterData)
    if isempty(data.current_values)
        data.lazy = false
    elseif !lazy_duals(data)
        @warn "Lazy mode is already de-activated"
    else
        error("Parameter JuMP's lazy mode can only be de-activated in empty models.")
    end
end

function _get_param_dict(data::_ParameterData, ::Type{MOI.EqualTo{Float64}})
    return data.parameters_map_saf_in_eq
end

function _get_param_dict(data::_ParameterData, ::Type{MOI.LessThan{Float64}})
    return data.parameters_map_saf_in_le
end

function _get_param_dict(data::_ParameterData, ::Type{MOI.GreaterThan{Float64}})
    return data.parameters_map_saf_in_ge
end

_getmodel(p::ParameterRef) = p.model

function _getparamdata(p::ParameterRef)::_ParameterData
    return _getparamdata(_getmodel(p))::_ParameterData
end

function _getparamdata(model::Model)::_ParameterData
    params = get(model.ext, :ParameterJuMP, nothing)
    if params !== nothing
        return params
    end
    return enable_parameters(model)
end

"""
    is_sync(model::Model)

Test if the JuMP Model is updated to the latest value of all parameters.
"""
is_sync(data::_ParameterData) = data.sync
is_sync(model::Model) = is_sync(_getparamdata(model))

function _add_parameter(model::Model, val::Real)
    params = _getparamdata(model)::_ParameterData
    ind = params.next_ind
    params.next_ind += 1
    push!(params.inds, ind)
    push!(params.current_values, 0.0)
    push!(params.future_values, val)
    if !iszero(val)
        params.sync = false
    end
    push!(params.dual_values, NaN)
    params.constraints_map[ind] = ParametrizedConstraintRef[]
    return ParameterRef(ind, model)
end

_addsizehint!(vec::Vector, len::Integer) = sizehint!(vec, len+length(vec))

"""
    all_parameters(model::AbstractModel)::Vector{ParameterRef}

Returns a list of all parameters currently in the model. The parameters are
ordered by creation time.
"""
function all_parameters(model::AbstractModel)
    data = _getparamdata(model)
    return ParameterRef[ParameterRef(ind, model) for ind in data.inds]
end

# solve
# ------------------------------------------------------------------------------

"""
    enable_parameters(m::Model)

Enables a `Model` to handle `Parameters`.
"""
function enable_parameters(m::Model)
    if haskey(m.ext, :params)
        error("Model already has parameter enabled")
    end
    set_optimize_hook(m, parameter_optimizehook)
    return _initialize_parameter_data(m)
end

function _initialize_parameter_data(model::Model)
    return model.ext[:ParameterJuMP] = _ParameterData()
end

# The constraint has been deleted. We cannot keep `dict` in sync with deletions
# as we return a `ConstraintRef`, not a custom type when the user create a
# parametrized constraint.
function _update_dicts_with_deletion(data::_ParameterData, ::Type{S}) where {S}
    dict = _get_param_dict(data, S)
    to_delete = eltype(keys(dict))[]
    for cref in keys(dict)
        if !is_valid(cref.model, cref)
            push!(to_delete, cref)
        end
    end
    for cref in to_delete
        delete!(dict, cref)
    end
end

function _update_dicts_with_deletion(data::_ParameterData)
    _update_dicts_with_deletion(data, MOI.EqualTo{Float64})
    _update_dicts_with_deletion(data, MOI.LessThan{Float64})
    _update_dicts_with_deletion(data, MOI.GreaterThan{Float64})
    return
end

function parameter_optimizehook(m::Model; kwargs...)
    data = _getparamdata(m)::_ParameterData
    _update_dicts_with_deletion(data)
    sync(data)
    optimize!(m::Model, ignore_optimize_hook = true, kwargs...)
    if !no_duals(data) && has_duals(m)
        _update_duals(data)
    end
    return
end

"""
    sync(model::Model)

Forces the model to update its constraints to the new values of the parameters.
"""
function sync(data::_ParameterData)
    if !is_sync(data)
        _update_constraints(data)
        data.current_values .= data.future_values
        data.sync = true
    end
    return
end
sync(model::Model) = sync(_getparamdata(model))

function parametrized_dual_objective_value(model::AbstractModel)
    params = all_parameters(model)
    linear = [param => dual(param) for param in params]
    obj = dual_objective_value(model)
    constant = obj - sum(value(term.first) * term.second for term in linear)
    return GenericAffExpr(constant, linear)
end

include("constraints.jl")
include("variable_interface.jl")
include("operators.jl")
include("mutable_arithmetics.jl")
include("macros.jl")
include("print.jl")

include("deprecations.jl")

end


module AttitudesWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const AngularCoordinates = @jimport org.orekit.utils.AngularCoordinates
const AngularDerivativesFilter = @jimport org.orekit.utils.AngularDerivativesFilter
const Attitude = @jimport org.orekit.attitudes.Attitude
const AttitudeProvider = @jimport org.orekit.attitudes.AttitudeProvider
const AttitudeProviderModifier = @jimport org.orekit.attitudes.AttitudeProviderModifier
const AttitudesSequence = @jimport org.orekit.attitudes.AttitudesSequence
const AttitudesSequence_SwitchHandler = @jimport org.orekit.attitudes.AttitudesSequence$SwitchHandler
const BodyCenterPointing = @jimport org.orekit.attitudes.BodyCenterPointing
const BodyShape = @jimport org.orekit.bodies.BodyShape
const CelestialBodyPointed = @jimport org.orekit.attitudes.CelestialBodyPointed
const Class = @jimport java.lang.Class
const Collection = @jimport java.util.Collection
const Ellipsoid = @jimport org.orekit.bodies.Ellipsoid
const EventDetector = @jimport org.orekit.propagation.events.EventDetector
const Field = @jimport org.hipparchus.Field
const FieldAbsoluteDate = @jimport org.orekit.time.FieldAbsoluteDate
const FieldAngularCoordinates = @jimport org.orekit.utils.FieldAngularCoordinates
const FieldAttitude = @jimport org.orekit.attitudes.FieldAttitude
const FieldPVCoordinatesProvider = @jimport org.orekit.utils.FieldPVCoordinatesProvider
const FieldPropagator = @jimport org.orekit.propagation.FieldPropagator
const FieldRotation = @jimport org.hipparchus.geometry.euclidean.threed.FieldRotation
const FieldTimeInterpolable = @jimport org.orekit.time.FieldTimeInterpolable
const FieldVector3D = @jimport org.hipparchus.geometry.euclidean.threed.FieldVector3D
const FixedRate = @jimport org.orekit.attitudes.FixedRate
const Frame = @jimport org.orekit.frames.Frame
const GeodeticPoint = @jimport org.orekit.bodies.GeodeticPoint
const GroundPointing = @jimport org.orekit.attitudes.GroundPointing
const InertialProvider = @jimport org.orekit.attitudes.InertialProvider
const JString = @jimport java.lang.String
const LOFType = @jimport org.orekit.frames.LOFType
const List = @jimport java.util.List
const LofOffset = @jimport org.orekit.attitudes.LofOffset
const LofOffsetPointing = @jimport org.orekit.attitudes.LofOffsetPointing
const NadirPointing = @jimport org.orekit.attitudes.NadirPointing
const Object = @jimport java.lang.Object
const PVCoordinatesProvider = @jimport org.orekit.utils.PVCoordinatesProvider
const Propagator = @jimport org.orekit.propagation.Propagator
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const Rotation = @jimport org.hipparchus.geometry.euclidean.threed.Rotation
const RotationOrder = @jimport org.hipparchus.geometry.euclidean.threed.RotationOrder
const SpinStabilized = @jimport org.orekit.attitudes.SpinStabilized
const Stream = @jimport java.util.stream.Stream
const TabulatedLofOffset = @jimport org.orekit.attitudes.TabulatedLofOffset
const TabulatedProvider = @jimport org.orekit.attitudes.TabulatedProvider
const TargetPointing = @jimport org.orekit.attitudes.TargetPointing
const TimeInterpolable = @jimport org.orekit.time.TimeInterpolable
const TimeStampedAngularCoordinates = @jimport org.orekit.utils.TimeStampedAngularCoordinates
const TimeStampedFieldAngularCoordinates = @jimport org.orekit.utils.TimeStampedFieldAngularCoordinates
const TimeStampedFieldPVCoordinates = @jimport org.orekit.utils.TimeStampedFieldPVCoordinates
const TimeStampedPVCoordinates = @jimport org.orekit.utils.TimeStampedPVCoordinates
const Vector3D = @jimport org.hipparchus.geometry.euclidean.threed.Vector3D
const YawCompensation = @jimport org.orekit.attitudes.YawCompensation
const YawSteering = @jimport org.orekit.attitudes.YawSteering

export Attitude
export AttitudeProvider
export AttitudeProviderModifier
export AttitudesSequence
export BodyCenterPointing
export CelestialBodyPointed
export FieldAttitude
export FixedRate
export GroundPointing
export InertialProvider
export LofOffset
export LofOffsetPointing
export NadirPointing
export SpinStabilized
export TabulatedLofOffset
export TabulatedProvider
export TargetPointing
export YawCompensation
export YawSteering
export add_switching_condition
export equals
export get_attitude
export get_base_state
export get_body_frame
export get_class
export get_date
export get_orientation
export get_reference_attitude
export get_reference_frame
export get_rotation
export get_rotation_acceleration
export get_spin
export get_table
export get_target_pv
export get_underlying_attitude_provider
export get_yaw_angle
export hash_code
export interpolate
export notify
export notify_all
export register_switch_events
export reset_active_provider
export shifted_by
export to_attitude
export to_string
export wait
export with_reference_frame

include("attitude.jl")
include("attitude_provider.jl")
include("attitude_provider_modifier.jl")
include("attitudes_sequence.jl")
include("body_center_pointing.jl")
include("celestial_body_pointed.jl")
include("field_attitude.jl")
include("fixed_rate.jl")
include("ground_pointing.jl")
include("inertial_provider.jl")
include("lof_offset.jl")
include("lof_offset_pointing.jl")
include("nadir_pointing.jl")
include("spin_stabilized.jl")
include("tabulated_lof_offset.jl")
include("tabulated_provider.jl")
include("target_pointing.jl")
include("yaw_compensation.jl")
include("yaw_steering.jl")

end


################
#@require Sundials begin
import Sundials

# note that using a Jacobian is messy in Sundials

# Monkey-patch Sundials v0.2
# TODO: remove this after a new Sundials.jl version gets tagged
eval(Sundials, quote
macro checkflag(ex)
    # Insert a check that the given function call returns 0,
    # throw an error otherwise. Only apply directly to function calls.
    @assert Base.Meta.isexpr(ex, :call)
    fname = ex.args[1]
    quote
        flag = $(esc(ex))
        if flag != 0
            error($(string(fname, " failed with error code = ")), flag)
        end
        flag
    end
end
end)

extraoutputs = 10^6
const sundialssolvers = Dict{Any,Solver}()

# CVode wrapper to high level interface: Sundials.cvode
#####

function wrapper_CVODE_simple{N,T}(tr::TestRun{N,T})
    # Wraps the specific MyPkg.mysolver such that it works within
    # IVPTestSuite setup.
    tc = tr.tc
    so = tr.solver
    ###
    # 0) Wrap tc.fn!, tc.jac!, tc.mass! if necessary
    # not necessary
    ###
    # 1) Make call signature

    tsteps = collect(linspace(tc.tspan[1], tc.tspan[2], extraoutputs)) # TODO: fix by setting mxstep higher
    args = (tc.fn!, tc.ic, tsteps)
    kwargs  = ((:reltol, tr.reltol), (:abstol, tr.abstol))

    ###
    # 2) Call solver, if it does not succeed throw an error (if that
    # is not done anyway)
    out = so.solverfn(args...; kwargs...)
    # (probably no need to modify this section)

    ###
    # 3) Transform output to conform to standard:
    # tend -- end time reached
    # yend -- solution at tend
    # stats -- statistics, if available: (steps_total,steps_accepted, fn_evals, jac_evals, linear_solves)
    #                         otherwise  (-1, -1, -1, -1, -1)
    tend = tc.tspan[2]
    yend = squeeze(out[end,:],1)
    stats = (-1, -1, -1, -1, -1)
    return tend, yend, stats
end
daeindex = 0
adaptive = true
cvode_simple = Solver{:im}(Sundials.cvode, Sundials, wrapper_CVODE_simple, stiff, adaptive, daeindex, explicit_eq)
sundialssolvers[Sundials.cvode] = cvode_simple

# CVode wrapper to low level interface: Sundials.CVode
#####

# Note that this function cannot be nested inside wrapper_CVODE
function cvodefun(t::Float64, y::Sundials.N_Vector, yp::Sundials.N_Vector, fn!::Function)
    y = Sundials.asarray(y)
    yp = Sundials.asarray(yp)
    fn!(t, y, yp)
    return Int32(0)
end

function wrapper_CVODE{N,T}(tr::TestRun{N,T})
    # Adapted from the Sundials.jl wrapper

    # Wraps the specific MyPkg.mysolver such that it works within
    # IVPTestSuite setup.
    tc = tr.tc
    so = tr.solver
    ###
    # 0) Wrap tc.fn!, tc.jac!, tc.mass! if necessary
    # not necessary

    ###
    # 1) Make call signature

    mem = Sundials.CVodeCreate(Sundials.CV_BDF, Sundials.CV_NEWTON)
    if mem == C_NULL
        error("Failed to allocate CVODE solver object")
    end

    flag = Sundials.@checkflag Sundials.CVodeInit(mem,
                              cfunction(cvodefun, Int32, (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Ref{Function})),
                              tc.tspan[1], Sundials.nvector(tc.ic))
    flag = Sundials.@checkflag Sundials.CVodeSetUserData(mem, tc.fn!)
    flag = Sundials.@checkflag Sundials.CVodeSStolerances(mem, tr.reltol, tr.abstol)
    flag = Sundials.@checkflag Sundials.CVodeSetMaxNumSteps(mem, -1)
    flag = Sundials.@checkflag Sundials.CVDense(mem, tc.dof)
    y = copy(tc.ic)
    tout = [0.0]

    ###
    # 2) Call solver, if it does not succeed throw an error (if that
    # is not done anyway)
    for k in 2:length(tc.tspan)
        flag = so.solverfn(mem, tc.tspan[k], y, tout, Sundials.CV_NORMAL) # solverfn==CVode
    end

    ###
    # 3) Transform output to conform to standard:
    # tend -- end time reached
    # yend -- solution at tend
    # stats -- statistics, if available: (steps_total,steps_accepted, fn_evals, jac_evals, linear_solves)
    #                         otherwise  (-1, -1, -1, -1, -1)
    tend = tout[1]
    yend = y
    stats = (-1, -1, -1, -1, -1) # TODO
    return tend, yend, stats
end
daeindex = 0
adaptive = true
cvode = Solver{:im}(Sundials.CVode, Sundials, wrapper_CVODE, stiff, adaptive, daeindex, explicit_eq)
sundialssolvers[Sundials.CVode] = cvode

# IDA wrapper to high-level interface: Sundials.idasol
#####

function wrapper_IDA_simple{N,T}(tr::TestRun{N,T})
    # Wraps the specific MyPkg.mysolver such that it works within
    # IVPTestSuite setup.
    tc = tr.tc
    so = tr.solver
    ###
    # 0) Wrap tc.fn!, tc.jac!, tc.mass! if necessary

    if hasmass(tc)
        function residual!(t, y, dydt, res)
            m = tc.mass!()
            tc.mass!(t,y,m)
            tc.fn!(t,y,res)
            res[:] = res-m*dydt
        end
    else
        function residual!(t, y, dydt, res)
            tc.fn!(t,y,res)
            res[:] = res-dydt
        end
    end

    ###
    # 1) Make call signature
    ic_dydt = tc.fn!()
    tc.fn!(tc.tspan[1], tc.ic, ic_dydt)
    tsteps = collect(linspace(tc.tspan[1], tc.tspan[2], extraoutputs)) # TODO: fix by setting mxstep higher
    args = (residual!, tc.ic, ic_dydt, tsteps)
    kwargs  = ((:reltol, tr.reltol), (:abstol, tr.abstol))

    ###
    # 2) Call solver, if it does not succeed throw an error (if that
    # is not done anyway)
    y,dydt = so.solverfn(args...; kwargs...)
    # (probably no need to modify this section)

    ###
    # 3) Transform output to conform to standard:
    # tend -- end time reached
    # yend -- solution at tend
    # stats -- statistics, if available: (steps_total,steps_accepted, fn_evals, jac_evals, linear_solves)
    #                         otherwise  (-1, -1, -1, -1, -1)
    tend = tc.tspan[2]
    yend = squeeze(y[end,:],1)
    stats = (-1, -1, -1, -1, -1)
    return tend, yend, stats
end
daeindex = 1
adaptive = true
ida_simple = Solver{:im}(Sundials.idasol, Sundials, wrapper_IDA_simple, stiff, adaptive, daeindex, implicit_eq)
sundialssolvers[Sundials.idasol] = ida_simple

# IDA wrapper to low-level interface: Sundials.ISASolve
#####

function idasolfun(t::Float64, y::Sundials.N_Vector, yp::Sundials.N_Vector, r::Sundials.N_Vector, userfun::Function)
    y = Sundials.asarray(y)
    yp = Sundials.asarray(yp)
    r = Sundials.asarray(r)
    userfun(t, y, yp, r)
    return Int32(0)   # indicates normal return
end

function wrapper_IDA{N,T}(tr::TestRun{N,T})
    # Adapted from the Sundials.jl wrapper

    # Wraps the specific MyPkg.mysolver such that it works within
    # IVPTestSuite setup.
    tc = tr.tc
    so = tr.solver
    ###
    # 0) Wrap tc.fn!, tc.jac!, tc.mass! if necessary
    if hasmass(tc)
        function residual!(t, y, dydt, res)
            m = tc.mass!()
            tc.mass!(t,y,m)
            tc.fn!(t,y,res)
            res[:] = res-m*dydt
        end
    else
        function residual!(t, y, dydt, res)
            tc.fn!(t,y,res)
            res[:] = res-dydt
        end
    end

    ###
    # 1) Make call signature
    ic_dydt = tc.fn!()
    tc.fn!(tc.tspan[1], tc.ic, ic_dydt)

    mem = Sundials.IDACreate()
    if mem == C_NULL
        error("Failed to allocate CVODE solver object")
    end

    flag = Sundials.@checkflag Sundials.IDAInit(mem, cfunction(idasolfun, Int32, (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Sundials.N_Vector, Ref{Function})),
                            tc.tspan[1], Sundials.nvector(tc.ic), Sundials.nvector(ic_dydt))
    flag = Sundials.@checkflag Sundials.IDASetUserData(mem, residual!)
    flag = Sundials.@checkflag Sundials.IDASStolerances(mem, tr.reltol, tr.abstol)
    flag = Sundials.@checkflag Sundials.IDADense(mem, tc.dof)
    flag = Sundials.@checkflag Sundials.IDASetMaxNumSteps(mem, -1)

    rtest = zeros(T, tc.dof)
    residual!(tc.tspan[1], tc.ic, ic_dydt, rtest)
    if any(abs(rtest) .>= tr.reltol)
        error("Inconsistent ICs!  This shouldn't happen...")
    end
    yres = zeros(T,length(tc.tspan), length(tc.ic))
    ypres = zeros(T,length(tc.tspan), length(tc.ic))
    yres[1,:] = tc.ic
    ypres[1,:] = ic_dydt
    y = copy(tc.ic)
    yp = copy(ic_dydt)
    tout = [0.0]

    ###
    # 2) Call solver, if it does not succeed throw an error (if that
    # is not done anyway)
    for k in 2:length(tc.tspan)
        flag = so.solverfn(mem, tc.tspan[k], tout, y, yp, Sundials.IDA_NORMAL) # solverfn==IDASolve
    end

    ###
    # 3) Transform output to conform to standard:
    # tend  -- end time reached
    # yend  -- solution at tend
    # stats -- statistics, if available: (steps_total,steps_accepted, fn_evals, jac_evals, linear_solves)
    #                         otherwise  (-1, -1, -1, -1, -1)
    tend = tout[1]
    yend = y
    stats = (-1, -1, -1, -1, -1) # TODO
    return tend, yend, stats
end
daeindex = 1
adaptive = true
ida = Solver{:im}(Sundials.IDASolve, Sundials, wrapper_IDA, stiff, adaptive, daeindex, implicit_eq)

sundialssolvers[Sundials.IDASolve] = ida

merge!(allsolvers, sundialssolvers)


using MacroTools
using MacroTools: @q, combinedef

named(arg) =
    isexpr(arg, :(::)) && length(arg.args) == 1 ? :($(gensym())::$(arg.args[1])) : arg

typeless(x) = MacroTools.postwalk(x -> isexpr(x, :(::), :kw) ? x.args[1] : x, x)
isvararg(x) = isexpr(x, :(::)) && namify(x.args[2]) == :Vararg

for n = 0:3
    gradtuple = Symbol(:gradtuple, n)
    @eval begin
        $gradtuple(x::Tuple) = ($(ntuple(_ -> :nothing, n)...), x...)
        $gradtuple(x::Nothing) = nothing
        $gradtuple(x) = error("Gradient $x should be a tuple")
    end
end

abstract type AContext end
function adjoint end
function _forward end

function gradm(ex, mut = false)
    @capture(
        shortdef(ex),
        (name_(args__) = body_) | (name_(args__) where {Ts__} = body_),
    ) || error("Need a function definition")
    kw = length(args) > 1 && isexpr(args[1], :parameters) ? esc(popfirst!(args)) : nothing
    isclosure = isexpr(name, :(::)) && length(name.args) > 1
    f, T =
        isexpr(name, :(::)) ?
        (length(name.args) == 1 ? (esc(gensym()), esc(name.args[1])) : esc.(name.args)) :
        (esc(gensym()), :(Core.Typeof($(esc(name)))))
    kT = :(Core.kwftype($T))
    Ts == nothing && (Ts = [])
    args = named.(args)
    argnames = Any[typeless(arg) for arg in args]
    !isempty(args) && isvararg(args[end]) && (argnames[end] = :($(argnames[end])...,))
    args = esc.(args)
    argnames = esc.(argnames)
    Ts = esc.(Ts)
    cx = :($(esc(:__context__))::AContext)
    fargs = kw == nothing ? [cx, :($f::$T), args...] : [kw, cx, :($f::$T), args...]
    gradtuple = isclosure ? gradtuple0 : gradtuple1
    gradtuplekw = isclosure ? gradtuple2 : gradtuple3
    adj = @q @inline ZygoteRules.adjoint($(fargs...)) where {$(Ts...)} = $(esc(body))
    quote
        $adj
        @inline function ZygoteRules._forward($cx, $f::$T, $(args...)) where {$(Ts...)}
            y, _back = adjoint(__context__, $f, $(argnames...))
            $(mut ? nothing : :(back(::Nothing) = nothing))
            back(Δ) = $gradtuple(_back(Δ))
            return y, back
        end
        @inline function ZygoteRules._forward(
            $cx,
            ::$kT,
            kw,
            $f::$T,
            $(args...),
        ) where {$(Ts...)}
            y, _back = adjoint(__context__, $f, $(argnames...); kw...)
            $(mut ? nothing : :(back(::Nothing) = nothing))
            back(Δ) = $gradtuplekw(_back(Δ))
            return y, back
        end
        nothing
    end
end

macro adjoint(ex)
    gradm(ex)
end

macro adjoint!(ex)
    gradm(ex, true)
end


@testset "Flux reports" begin

@testset "function nodeFluxReport" begin
    @test isempty( nodeFluxReport( mpSim, [-6.0, -12.0], :in, "active",
        "A junior", "Career", "Blort" ) )
    @test isempty( nodeFluxReport( mpSim, [12.0, 8.0, 0.0], :bang, "active",
        "A junior", "Career", "Blort" ) )
    @test isempty( nodeFluxReport( mpSim, [12.0, 8.0, 0.0], :in, "Blort" ) )
    @test isempty( nodeFluxReport( mpSim, -6, :in, "active",
        "A junior", "Career", "Blort" ) )

    report = nodeFluxReport( mpSim, 12, :in, "active", "A junior", "Career",
        "Blort" )
    @test !haskey( report, "Blort" ) && all( haskey.( Ref( report ), ["active",
        "Career", "A junior"] ) )
    @test all( report["active"][:, "external => active"] .== 30 )
    @test all( report["A junior"][:, "other => A junior"] .== 10 )
    @test all( report["Career"][:, "other => Career"] .==
        vcat( 0, 0, 24, 24, fill( 16, 22 ) ) )

    report = nodeFluxReport( mpSim, 12, :out, "active", "A junior", "Career",
        "Blort" )
    @test all( report["active"][:, "active => external"] .==
        vcat( 0, 0, 6, 6, fill( 14, 6 ), 38, 38, fill( 30, 14 ) ) )
    @test all( report["A junior"][:, "A junior => other"] .==
        vcat( 0, 0, fill( 10, 24 ) ) )
    @test all( report["Career"][:, "Career => other"] .==
        vcat( fill( 0, 10 ), 24, 24, fill( 16, 14 ) ) )

    report = nodeFluxReport( mpSim, 12, :within, "active", "A junior", "Career",
        "Blort" )
    @test all( report["active"][:, "within active"] .== [0, 0,
        24, 24, 16, 26, 26, 21, 16, 16, 26, 26, 21, 16, 16, 26, 26, 21, 16, 16,
        26, 26, 21, 16, 16, 26] )
end  # @testset "function nodeFluxReport"

@testset "function transitionFluxReport" begin
    @test isempty( transitionFluxReport( mpSim, [-12.0, -8.0], "Promotion",
        ("A junior", "A senior"), ("Reserve", "Reserve junior", "B senior") ) )
    @test isempty( transitionFluxReport( mpSim, [12.0, 24.0], "Blook",
        ("Foo", "Bar"), ("A senior", "B junior"), ("Blook", "Foo", "Bar"),
        ("Reserve", "A junior", "A senior") ) )
    @test isempty( transitionFluxReport( mpSim, -12.0, "Promotion",
        ("A junior", "A senior"), ("Reserve", "Reserve junior", "B senior") ) )
    report = transitionFluxReport( mpSim, 12, "attrition", "Promotion", "EW",
        "B-" )
    @test all( report[:, "attrition"] .== 0 )
    @test all( report[:, "Promotion"] .== [0, 0, 16, 16, 16, 26, 26, 21, 16,
        16, 26, 26, 21, 16, 16, 26, 26, 21, 16, 16, 26, 26, 21, 16, 16, 26] )
    @test all( report[:, "EW"] .== 30 )
    @test all( report[:, "B-"] .==
        vcat( 0, 0, 6, 6, fill( 14, 22 ) ) )

    report = transitionFluxReport( mpSim, 12, ("external", "A junior"),
        ("B senior", ""), ("Reserve junior", "A senior") )
    @test all( report[:, "external => A junior"] .== 10 )
    @test all( report[:, "B senior => external"] .==
        vcat( fill( 0, 10 ), 7, 7, 8, 8, 8, 3, 3, 8, 8, 8, 3, 3, 8, 8, 8, 3 ) )
    @test all( report[:, "Reserve junior => A senior"] .==
        vcat( 0, 0, 4, 4, fill( 0, 22 ) ) )

    report = transitionFluxReport( mpSim, 12, ("attrition", "Master", "out"),
        ("PE", "A senior", ""), ("EW", "OUT", "A junior"),
        ("Reserve", "Reserve junior", "B senior") )
    @test all( report[:, "attrition: Master => external"] .== 0 )
    @test all( report[:, "PE: A senior => external"] .==
        vcat( fill( 0, 10 ), 7, 7, 3, 8, 8, 3, 3, 3, 8, 8, 3, 3, 3, 8, 8, 3 ) )
    @test all( report[:, "EW: external => A junior"] .== 10 )
    @test all( report[:, "Reserve: Reserve junior => B senior"] .==
        vcat( 0, 0, 4, 4, fill( 0, 22 ) ) )
end  # @testset "function transitionFluxReport"

end  # @testset "Flux reports"


abstract type AbstractSimulation{T <: SSDFloat} end

"""
    mutable struct Simulation{T <: SSDFloat} <: AbstractSimulation{T}

Collection of all parts of a Simulation of a Solid State Detector.
"""
mutable struct Simulation{T <: SSDFloat} <: AbstractSimulation{T}
    detector::Union{SolidStateDetector{T}, Missing}
    ρ::Union{ChargeDensity{T}, Missing}
    ρ_fix::Union{ChargeDensity{T}, Missing}
    ϵ::Union{DielectricDistribution{T}, Missing}
    point_types::Union{PointTypes{T}, Missing}
    electric_potential::Union{ElectricPotential{T}, Missing}
    weighting_potentials::Vector{Any}
    electric_field::Union{ElectricField{T}, Missing}

    charge_drift_model::Union{<:AbstractChargeDriftModel{T}, Missing}

    electron_drift_field::Union{ElectricField{T}, Missing}
    hole_drift_field::Union{ElectricField{T}, Missing}
end

function Simulation{T}() where {T <: SSDFloat}
    Simulation{T}(
        SolidStateDetector{T}(),
        missing,
        missing,
        missing,
        missing,
        missing,
        [missing],
        missing,
        ElectricFieldChargeDriftModel{T}(),
        missing,
        missing
    )
end

function get_precision_type(sim::Simulation{T}) where {T}
    typeof(sim).parameters[1]
end

function NamedTuple(sim::Simulation{T}) where {T <: SSDFloat}
    wps_strings = AbstractString[]
    for contact in sim.detector.contacts
        if !ismissing(sim.weighting_potentials[contact.id])
            push!(wps_strings, "WeightingPotential_$(contact.id)")
        end
    end
    wps_syms = Symbol.(wps_strings)
    nt = (
        detector_json_string = NamedTuple(sim.detector.config_dict),
        electric_potential = NamedTuple(sim.electric_potential),
        ρ = NamedTuple(sim.ρ),
        ρ_fix = NamedTuple(sim.ρ_fix),
        ϵ = NamedTuple(sim.ϵ),
        point_types = NamedTuple(sim.point_types),
        electric_field = NamedTuple(sim.electric_field),
        electron_drift_field = NamedTuple(sim.electron_drift_field),
        hole_drift_field = NamedTuple(sim.hole_drift_field)
    )
    if length(wps_strings) > 0
        nt = merge(nt, (weighting_potentials = NamedTuple{Tuple(wps_syms)}(NamedTuple.( skipmissing(sim.weighting_potentials))),))
    end
    return nt
end
Base.convert(T::Type{NamedTuple}, x::Simulation) = T(x)

function Simulation(nt::NamedTuple)
    ep = ElectricPotential(nt.electric_potential)
    T = eltype(ep.data)
    det = SolidStateDetector{T}( Dict(nt.detector_json_string) )
    sim = Simulation( det )
    if !ismissing(nt.electric_potential) sim.electric_potential = ep end
    if !ismissing(nt.ρ) sim.ρ = ChargeDensity(nt.ρ) end
    if !ismissing(nt.ρ_fix) sim.ρ_fix = ChargeDensity(nt.ρ_fix) end
    if !ismissing(nt.ϵ) sim.ϵ = DielectricDistribution(nt.ϵ) end
    if !ismissing(nt.point_types) sim.point_types = PointTypes(nt.point_types) end
    if !ismissing(nt.electric_field) sim.electric_field = ElectricField(nt.electric_field) end
    sim.weighting_potentials = [missing for contact in sim.detector.contacts]
    if !ismissing(nt.weighting_potentials)
        for contact in sim.detector.contacts
            if !ismissing(values(nt.weighting_potentials[contact.id])[1])
                sim.weighting_potentials[contact.id] = WeightingPotential(nt.weighting_potentials[contact.id])
            end
        end
    end
    if !ismissing(nt.electron_drift_field) sim.electron_drift_field = ElectricField(nt.electron_drift_field) end
    if !ismissing(nt.hole_drift_field) sim.hole_drift_field = ElectricField(nt.hole_drift_field) end
    sim.charge_drift_model = ADLChargeDriftModel(T=T)
    @info "I/O of charge drift model not yet supported. Loading default: ADLChargeDriftModel"
    return sim
end
Base.convert(T::Type{Simulation}, x::NamedTuple) = T(x)

function println(io::IO, sim::Simulation{T}) where {T <: SSDFloat}
    println(typeof(sim), " - Coordinate system: ", get_coordinate_system(sim.detector))
    println("  Detector: $(sim.detector.name)")
    println("  Electric potential: ", !ismissing(sim.electric_potential) ? size(sim.electric_potential) : missing)
    println("  Charge density: ", !ismissing(sim.ρ) ? size(sim.ρ) : missing)
    println("  Fix Charge density: ", !ismissing(sim.ρ_fix) ? size(sim.ρ_fix) : missing)
    println("  Dielectric distribution: ", !ismissing(sim.ϵ) ? size(sim.ϵ) : missing)
    println("  Point types: ", !ismissing(sim.point_types) ? size(sim.point_types) : missing)
    println("  Electric field: ", !ismissing(sim.electric_field) ? size(sim.electric_field) : missing)
    println("  Weighting potentials: ")
    for contact in sim.detector.contacts
        print("    Contact $(contact.id): ")
        println(!ismissing(sim.weighting_potentials[contact.id]) ? size(sim.weighting_potentials[contact.id]) : missing)
    end
    println("  Charge drift model: ", !ismissing(sim.electric_field) ? typeof(sim.charge_drift_model) : missing)
    println("  Electron drift field: ", !ismissing(sim.electron_drift_field) ? size(sim.electron_drift_field) : missing)
    println("  Hole drift field: ", !ismissing(sim.hole_drift_field) ? size(sim.hole_drift_field) : missing)
end

function print(io::IO, sim::Simulation{T}) where {T <: SSDFloat}
    print(io, "Simulation{$T} - ", "$(sim.detector.name)")
end

function show(io::IO, sim::Simulation{T}) where {T <: SSDFloat} println(io, sim) end

function show(io::IO, ::MIME"text/plain", sim::Simulation{T}) where {T <: SSDFloat}
    show(io, sim)
end

function Simulation(detector::SolidStateDetector{T})::Simulation{T} where {T <: SSDFloat}
    sim::Simulation{T} = Simulation{T}()
    sim.detector = detector
    sim.weighting_potentials = Missing[ missing for i in 1:length(sim.detector.contacts)]
    return sim
end

function Simulation{T}(config_file::AbstractString)::Simulation{T} where{T <: SSDFloat}
    return Simulation( SolidStateDetector{T}(config_file) )
end
function Simulation(config_file::AbstractString)::Simulation{Float32}
    return Simulation{Float32}( config_file )
end

# Functions
"""
    function apply_initial_state!(sim::Simulation{T}, ::Type{ElectricPotential}, grid::Grid{T} = Grid(sim.detector))::Nothing

Applies the initial state of the electric potential calculation.
It overwrites `sim.electric_potential`, `sim.ρ`, `sim.ρ_fix`, `sim.ϵ` and `sim.point_types`.
"""
function apply_initial_state!(sim::Simulation{T}, ::Type{ElectricPotential}, grid::Grid{T} = Grid(sim.detector))::Nothing where {T <: SSDFloat}
    fssrb::PotentialSimulationSetupRB{T, 3, 4, get_coordinate_system(sim.detector)} =
        PotentialSimulationSetupRB(sim.detector, grid);

    sim.ρ = ChargeDensity(ChargeDensityArray(fssrb), grid)
    sim.ρ_fix = ChargeDensity(FixedChargeDensityArray(fssrb), grid)
    sim.ϵ = DielectricDistribution(DielektrikumDistributionArray(fssrb), grid)
    sim.point_types = PointTypes(PointTypeArray(fssrb), grid)
    sim.electric_potential = ElectricPotential(ElectricPotentialArray(fssrb), grid)
    nothing
end

"""
    function apply_initial_state!(sim::Simulation{T}, ::Type{WeightingPotential}, contact_id::Int, grid::Grid{T} = Grid(sim.detector))::Nothing

Applies the initial state of the weighting potential calculation for the contact with the id `contact_id`.
It overwrites `sim.weighting_potentials[contact_id]`.
"""
function apply_initial_state!(sim::Simulation{T}, ::Type{WeightingPotential}, contact_id::Int, grid::Grid{T} = Grid(sim.detector))::Nothing where {T <: SSDFloat}
    fssrb::PotentialSimulationSetupRB{T, 3, 4, get_coordinate_system(sim.detector)} =
        PotentialSimulationSetupRB(sim.detector, grid, weighting_potential_contact_id = contact_id);

    sim.weighting_potentials[contact_id] = WeightingPotential(ElectricPotentialArray(fssrb), grid)
    nothing
end

"""
    function update_till_convergence!( sim::Simulation{T} ::Type{ElectricPotential}, convergence_limit::Real; kwargs...)::T

Takes the current state of `sim.electric_potential` and updates it until it has converged.
"""
function update_till_convergence!( sim::Simulation{T},
                                   ::Type{ElectricPotential},
                                   convergence_limit::Real = 1e-7;
                                   n_iterations_between_checks::Int = 500,
                                   max_n_iterations::Int = -1,
                                   depletion_handling::Bool = false,
                                   use_nthreads::Int = Base.Threads.nthreads(),
                                   sor_consts::Union{Missing, T, NTuple{2, T}} = missing
                                    )::T where {T <: SSDFloat}
    CS = get_coordinate_system(sim.detector)
    if ismissing(sor_consts)
        sor_consts = CS == :cylindrical ? (T(1.4), T(1.85)) : T(1.4)
    elseif length(sor_consts) == 1 && CS == :cylindrical
        sor_consts = (T(sor_consts), T(sor_consts))
    elseif length(sor_consts) > 1 && CS == :cartesian
        sor_consts = T(sor_consts[1])
    end
    only_2d::Bool = length(sim.electric_potential.grid.axes[2]) == 1

    fssrb::PotentialSimulationSetupRB{T, 3, 4, get_coordinate_system(sim.electric_potential.grid)} =
        PotentialSimulationSetupRB(sim.detector, sim.electric_potential.grid, sim.electric_potential.data, sor_consts = T.(sor_consts))

    cf::T = _update_till_convergence!( fssrb, T(convergence_limit);
                                       only2d = Val{only_2d}(),
                                       depletion_handling = Val{depletion_handling}(),
                                       is_weighting_potential = Val{false}(),
                                       use_nthreads = use_nthreads,
                                       n_iterations_between_checks = n_iterations_between_checks,
                                       max_n_iterations = max_n_iterations )

    grid::Grid = Grid(fssrb)
    sim.ρ = ChargeDensity(ChargeDensityArray(fssrb), grid)
    sim.ρ_fix = ChargeDensity(FixedChargeDensityArray(fssrb), grid)
    sim.ϵ = DielectricDistribution(DielektrikumDistributionArray(fssrb), grid)
    sim.electric_potential = ElectricPotential(ElectricPotentialArray(fssrb), grid)
    sim.point_types = PointTypes(PointTypeArray(fssrb), grid)

    if depletion_handling == false
        update_again::Bool = false # With SOR-Constant = 1
        if fssrb.bulk_is_ptype
            @inbounds for i in eachindex(sim.electric_potential.data)
                if sim.electric_potential.data[i] < fssrb.minimum_applied_potential
                    sim.electric_potential.data[i] = fssrb.minimum_applied_potential
                    update_again = true
                end
            end
        else # ntype
            @inbounds for i in eachindex(sim.electric_potential.data)
                if sim.electric_potential.data[i] > fssrb.maximum_applied_potential
                    sim.electric_potential.data[i] = fssrb.maximum_applied_potential
                    update_again = true
                end
            end
        end
        if update_again
            fssrb.sor_const[:] .= T(1)
            cf = _update_till_convergence!( fssrb, T(convergence_limit);
                                            only2d = Val{only_2d}(),
                                            depletion_handling = Val{depletion_handling}(),
                                            is_weighting_potential = Val{false}(),
                                            use_nthreads = use_nthreads,
                                            n_iterations_between_checks = n_iterations_between_checks,
                                            max_n_iterations = max_n_iterations )
            sim.electric_potential = ElectricPotential(ElectricPotentialArray(fssrb), grid)

            if fssrb.bulk_is_ptype
                @inbounds for i in eachindex(sim.electric_potential.data)
                    if sim.electric_potential.data[i] < fssrb.minimum_applied_potential
                        @warn """Detector seems not to be fully depleted at a bias voltage of $(fssrb.bias_voltage) V.
                            At least one grid point has a smaller potential value ($(sim.electric_potential.data[i]) V)
                            than the minimum applied potential ($(fssrb.minimum_applied_potential) V). This should not be.
                            However, small overshoots could be due to numerical precision."""
                        break
                    end
                end
            else # ntype
                @inbounds for i in eachindex(sim.electric_potential.data)
                    if sim.electric_potential.data[i] > fssrb.maximum_applied_potential
                        @warn """Detector seems not to be not fully depleted at a bias voltage of $(fssrb.bias_voltage) V.
                            At least one grid point has a higher potential value ($(sim.electric_potential.data[i]) V)
                            than the maximum applied potential ($(fssrb.maximum_applied_potential) V). This should not be.
                            However, small overshoots could be due to numerical precision."""
                        break
                    end
                end
            end
        end
    end

    cf
end

"""
    function update_till_convergence!( sim::Simulation{T} ::Type{WeightingPotential}, contact_id::Int, convergence_limit::Real; kwargs...)::T

Takes the current state of `sim.weighting_potentials[contact_id]` and updates it until it has converged.
"""
function update_till_convergence!( sim::Simulation{T},
                                   ::Type{WeightingPotential},
                                   contact_id::Int,
                                   convergence_limit::Real;
                                   n_iterations_between_checks::Int = 500,
                                   max_n_iterations::Int = -1,
                                   depletion_handling::Bool = false,
                                   use_nthreads::Int = Base.Threads.nthreads(),
                                   sor_consts::Union{Missing, T, NTuple{2, T}} = missing
                                    )::T where {T <: SSDFloat}
    CS = get_coordinate_system(sim.detector)
    if ismissing(sor_consts)
        sor_consts = CS == :cylindrical ? (T(1.4), T(1.85)) : T(1.4)
    elseif length(sor_consts) == 1 && CS == :cylindrical
        sor_consts = (T(sor_consts), T(sor_consts))
    elseif length(sor_consts) > 1 && CS == :cartesian
        sor_consts = T(sor_consts[1])
    end

    only_2d::Bool = length(sim.weighting_potentials[contact_id].grid.axes[2]) == 1
    fssrb::PotentialSimulationSetupRB{T, 3, 4, get_coordinate_system(sim.weighting_potentials[contact_id].grid)} =
        PotentialSimulationSetupRB(sim.detector, sim.weighting_potentials[contact_id].grid, sim.weighting_potentials[contact_id].data,
                sor_consts = T.(sor_consts), weighting_potential_contact_id = contact_id)

    cf::T = _update_till_convergence!( fssrb, T(convergence_limit);
                                       only2d = Val{only_2d}(),
                                       depletion_handling = Val{depletion_handling}(),
                                       is_weighting_potential = Val{true}(),
                                       use_nthreads = use_nthreads,
                                       n_iterations_between_checks = n_iterations_between_checks,
                                       max_n_iterations = max_n_iterations )

    sim.weighting_potentials[contact_id] = WeightingPotential(ElectricPotentialArray(fssrb), sim.weighting_potentials[contact_id].grid)

    cf
end

"""
    function refine!(sim::Simulation{T}, ::Type{ElectricPotential}, max_diffs::Tuple{<:Real,<:Real,<:Real}, minimum_distances::Tuple{<:Real,<:Real,<:Real})

Takes the current state of `sim.electric_potential` and refines it with respect to the input arguments
`max_diffs` and `minimum_distances`.
"""
function refine!(sim::Simulation{T}, ::Type{ElectricPotential},
                    max_diffs::Tuple{<:Real,<:Real,<:Real} = (T(0), T(0), T(0)),
                    minimum_distances::Tuple{<:Real,<:Real,<:Real} = (T(0), T(0), T(0));
                    update_other_fields::Bool = false) where {T <: SSDFloat}
    sim.electric_potential = refine(sim.electric_potential, max_diffs, minimum_distances)

    if update_other_fields
        fssrb::PotentialSimulationSetupRB{T, 3, 4, get_coordinate_system(sim.electric_potential.grid)} =
            PotentialSimulationSetupRB(sim.detector, sim.electric_potential.grid, sim.electric_potential.data)

        sim.ρ = ChargeDensity(ChargeDensityArray(fssrb), sim.electric_potential.grid)
        sim.ρ_fix = ChargeDensity(FixedChargeDensityArray(fssrb), sim.electric_potential.grid)
        sim.ϵ = DielectricDistribution(DielektrikumDistributionArray(fssrb), sim.electric_potential.grid)
        sim.point_types = PointTypes(PointTypeArray(fssrb), sim.electric_potential.grid)
    end
    nothing
end
"""
    function refine!(sim::Simulation{T}, ::Type{WeightingPotential}, max_diffs::Tuple{<:Real,<:Real,<:Real}, minimum_distances::Tuple{<:Real,<:Real,<:Real})

Takes the current state of `sim.weighting_potentials[contact_id]` and refines it with respect to the input arguments
`max_diffs` and `minimum_distances`.
"""
function refine!(sim::Simulation{T}, ::Type{WeightingPotential}, contact_id::Int,
                    max_diffs::Tuple{<:Real,<:Real,<:Real} = (T(0), T(0), T(0)),
                    minimum_distances::Tuple{<:Real,<:Real,<:Real} = (T(0), T(0), T(0))) where {T <: SSDFloat}
    sim.weighting_potentials[contact_id] = refine(weighting_potentials[contact_id], max_diffs, minimum_distances)
    nothing
end

function _calculate_potential!( sim::Simulation{T}, potential_type::UnionAll, contact_id::Union{Missing, Int} = missing;
        init_grid_size::Union{Missing, NTuple{3, Int}} = missing,
        init_grid_spacing::Union{Missing, Tuple{<:Real,<:Real,<:Real}} = missing,
        grid::Union{Missing, Grid{T}} = missing,
        convergence_limit::Real = 1e-7,
        max_refinements::Int = 3,
        refinement_limits::Union{Missing, Tuple{<:Real,<:Real,<:Real}} = missing,
        min_grid_spacing::Union{Missing, Tuple{<:Real,<:Real,<:Real}} = missing,
        depletion_handling::Bool = false,
        use_nthreads::Int = Base.Threads.nthreads(),
        sor_consts::Union{Missing, <:Real, Tuple{<:Real,<:Real}} = missing,
        max_n_iterations::Int = 50000,
        verbose::Bool = true,
    )::Nothing where {T <: SSDFloat}

    begin # preperations
        convergence_limit::T = T(convergence_limit)
        isEP::Bool = potential_type == ElectricPotential
        isWP::Bool = !isEP
        if isWP depletion_handling = false end
        CS = get_coordinate_system(sim.detector)
        if ismissing(grid)
            grid = Grid(sim.detector, init_grid_size = init_grid_size, init_grid_spacing = init_grid_spacing, for_weighting_potential = isWP)
        end
        if ismissing(sor_consts)
            sor_consts = CS == :cylindrical ? (T(1.4), T(1.85)) : T(1.4)
        elseif length(sor_consts) == 1 && CS == :cylindrical
            sor_consts = (T(sor_consts), T(sor_consts))
        elseif length(sor_consts) > 1 && CS == :cartesian
            sor_consts = T(sor_consts[1])
        end
        if ismissing(refinement_limits)
            refinement_limits::NTuple{3, T} = CS == :cylindrical ? (T(1e-5), T(1e-5) / (0.25 * grid.axes[1][end]), T(1e-5)) : (T(1e-5), T(1e-5), T(1e-5))
        end
        if ismissing(min_grid_spacing)
            min_grid_spacing::NTuple{3, T} = CS == :cylindrical ? (T(1e-5), T(1e-5) / (0.25 * grid.axes[1][end]), T(1e-5)) : (T(1e-5), T(1e-5), T(1e-5))
        end
        n_iterations_between_checks::Int = div(10^7, length(grid))
        if n_iterations_between_checks < 20 n_iterations_between_checks = 20 end
        if use_nthreads > Base.Threads.nthreads()
            use_nthreads = Base.Threads.nthreads();
            @warn "`use_nthreads` was set to `1`. The environment variable `JULIA_NUM_THREADS` must be set appropriately before the julia session is started."
        end
        refine::Bool = max_refinements > 0 ? true : false
        only_2d::Bool = length(grid.axes[2]) == 1 ? true : false
        if CS == :cylindrical
            cyclic::T = grid.axes[2].interval.right - grid.axes[2].interval.left
            n_φ_sym::Int = only_2d ? 1 : Int(round(T(2π) / cyclic, digits = 3))
            n_φ_sym_info_txt = if only_2d
                "φ symmetry: Detector is φ-symmetric -> 2D computation."
            else
                "φ symmetry: calculating just 1/$(n_φ_sym) in φ of the detector."
            end
        end
        contact_potentials::Vector{T} = [contact.potential for contact in sim.detector.contacts]
        bias_voltage::T = (length(contact_potentials) > 0) ? (maximum(contact_potentials) - minimum(contact_potentials)) : T(0)
        if isWP bias_voltage = T(1) end
        if verbose
            println(
                "$(isEP ? "Electric" : "Weighting") Potential Calculation\n",
                if isEP "Bias voltage: $(bias_voltage) V\n" else "" end,
                if CS == :cylindrical "$n_φ_sym_info_txt\n" else "" end,
                "Precision: $T\n",
                "Convergence limit: $convergence_limit => $(round(abs(bias_voltage * convergence_limit), sigdigits=2)) V\n",
                "Threads: $use_nthreads\n",
                "Coordinate system: $(CS)\n",
                "Initial grid dimension: $(size(grid))\n",
                "Refine? -> $refine\n",
                "Refinement parameters:\n",
                "\tmaximum number of refinements: $(max_refinements)\n",
                "\tminimum grid spacing:\n",
                "\t\tr: $(min_grid_spacing[1]) m\n",
                "\t\tφ: $(min_grid_spacing[2]) rad\n",
                "\t\tz: $(min_grid_spacing[3]) m\n",
                "\tRefinement limits:\n",
                "\t\tr: $(refinement_limits[1]) -> $(round(abs(bias_voltage * refinement_limits[1]), sigdigits=2)) V\n",
                "\t\tφ: $(refinement_limits[2]) -> $(round(abs(bias_voltage * refinement_limits[2]), sigdigits=2)) V\n",
                "\t\tz: $(refinement_limits[3]) -> $(round(abs(bias_voltage * refinement_limits[3]), sigdigits=2)) V\n",
                ""
            )
        end
    end
    if isEP
        apply_initial_state!(sim, potential_type, grid)
        update_till_convergence!( sim, potential_type, convergence_limit,
                                  n_iterations_between_checks = n_iterations_between_checks,
                                  max_n_iterations = max_n_iterations,
                                  depletion_handling = depletion_handling,
                                  use_nthreads = use_nthreads,
                                  sor_consts = sor_consts )
    else
        apply_initial_state!(sim, potential_type, contact_id, grid)
        update_till_convergence!( sim, potential_type, contact_id, convergence_limit,
                                    n_iterations_between_checks = n_iterations_between_checks,
                                    max_n_iterations = max_n_iterations,
                                    depletion_handling = depletion_handling,
                                    use_nthreads = use_nthreads,
                                    sor_consts = sor_consts )
    end

    refinement_counter::Int = 1
    if refine
        max_diffs::NTuple{3, T} = abs.(refinement_limits .* bias_voltage)
        refine_at_inds::NTuple{3, Vector{Int}} = if isEP
            _get_refinement_inds(sim.electric_potential.data, sim.electric_potential.grid, max_diffs, min_grid_spacing)
        else
            _get_refinement_inds(sim.weighting_potentials[contact_id].data, sim.weighting_potentials[contact_id].grid, max_diffs, min_grid_spacing)
        end
        while any(!isempty, refine_at_inds[1]) && refinement_counter <= max_refinements
            if isEP
                sim.electric_potential = ElectricPotential(add_points_and_interpolate(
                        sim.electric_potential.data, sim.electric_potential.grid, refine_at_inds...)...)
                if verbose @info "New Grid Size = $(size(sim.electric_potential.grid))" end
            else
                sim.weighting_potentials[contact_id] = WeightingPotential(add_points_and_interpolate(
                    sim.weighting_potentials[contact_id].data, sim.weighting_potentials[contact_id].grid, refine_at_inds...)...)
                if verbose @info "New Grid Size = $(size(sim.weighting_potentials[contact_id].grid))" end
            end
            n_iterations_between_checks = div(10^7, length(isEP ? sim.electric_potential.grid : sim.weighting_potentials[contact_id].grid))
            if n_iterations_between_checks < 20 n_iterations_between_checks = 20 end

            if isEP
                update_till_convergence!( sim, potential_type, convergence_limit,
                    n_iterations_between_checks = n_iterations_between_checks,
                    max_n_iterations = max_n_iterations,
                    depletion_handling = depletion_handling,
                    use_nthreads = use_nthreads,
                    sor_consts = sor_consts )
            else
                update_till_convergence!( sim, potential_type, contact_id, convergence_limit,
                                    n_iterations_between_checks = n_iterations_between_checks,
                                    max_n_iterations = max_n_iterations,
                                    depletion_handling = depletion_handling,
                                    use_nthreads = use_nthreads,
                                    sor_consts = sor_consts )
            end

            refine_at_inds = if isEP
                _get_refinement_inds(sim.electric_potential.data, sim.electric_potential.grid, max_diffs, min_grid_spacing)
            else
                _get_refinement_inds(sim.weighting_potentials[contact_id].data, sim.weighting_potentials[contact_id].grid, max_diffs, min_grid_spacing)
            end
            refinement_counter += 1
        end
    end

    nothing
end

"""
    calculate_weighting_potential!(sim::Simulation{T}, contact_id::Int; kwargs...)::Nothing

Compute the weighting potential for the contact with id `contact_id`
for the given Simulation `sim` on an adaptive grid through successive over relaxation.

There are serveral `<keyword arguments>` which can be used to tune the computation:

# Keywords
- `convergence_limit::Real`: `convergence_limit` times the bias voltage sets the convergence limit of the relaxation.
    The convergence value is the absolute maximum difference of the potential between two iterations of all grid points.
    Default of `convergence_limit` is `2e-6` (times bias voltage).
- `max_refinements::Int`: Number of maximum refinements. Default is `2`. Set it to `0` to switch off refinement.
- `refinement_limits::Tuple{<:Real, <:Real, <:Real}`: Tuple of refinement limits for each dimension
    (in case of cylindrical coordinates the order is `r`, `φ`, `z`).
    A refinement limit (e.g. `refinement_limits[1]`) times the bias voltage of the detector `det` is the
    maximum allowed voltage difference between two neighbouring grid points in the respective dimension.
    When the difference is larger, new points are created inbetween. Default is `[1e-5, 1e-5, 1e-5]`.
- `init_grid_spacing::Tuple{<:Real, <:Real, <:Real}`: Tuple of the initial distances between two grid points for each dimension.
    For normal coordinates the unit is meter. For angular coordinates, the unit is radiance.
    It prevents the refinement to make the grid to fine. Default is `[0.005, 10.0, 0.005]``.
- `min_grid_spacing::Tuple{<:Real, <:Real, <:Real}`: Tuple of the mimimum allowed distance between two grid points for each dimension.
    For normal coordinates the unit is meter. For angular coordinates, the unit is radiance.
    It prevents the refinement to make the grid to fine. Default is [`1e-6`, `1e-6`, `1e-6`].
- `grid::Grid{T, N, S}`: Initial grid used to start the simulation. Default is `Grid(detector, init_grid_spacing=init_grid_spacing)`.
- `depletion_handling::Bool`: Enables the handling of undepleted regions. Default is false.
- `use_nthreads::Int`: Number of threads to use in the computation. Default is `Base.Threads.nthreads()`.
    The environment variable `JULIA_NUM_THREADS` must be set appropriately before the Julia session was
    started (e.g. `export JULIA_NUM_THREADS=8` in case of bash).
- `sor_consts::Union{<:Real, NTuple{2, <:Real}}`: Two element tuple in case of cylindrical coordinates.
    First element contains the SOR constant for `r` = 0.
    Second contains the constant at the outer most grid point in `r`. A linear scaling is applied in between.
    First element should be smaller than the second one and both should be ∈ [1.0, 2.0]. Default is [1.4, 1.85].
    In case of cartesian coordinates only one value is taken.
- `max_n_iterations::Int`: Set the maximum number of iterations which are performed after each grid refinement.
    Default is `10000`. If set to `-1` there will be no limit.
- `verbose::Bool=true`: Boolean whether info output is produced or not.
"""
function calculate_weighting_potential!(sim::Simulation{T}, contact_id::Int, args...; n_points_in_φ::Union{Missing, Int} = missing, kwargs...)::Nothing where {T <: SSDFloat}
    # S = get_coordinate_system(sim.detector)
    # periodicity::T = get_periodicity(sim.detector.world.intervals[2])
    # if S == :cylindrical && periodicity == T(0)
    #     if ismissing(n_points_in_φ)
    #         @info "\tIn weighing potential calculation: Keyword `n_points_in_φ` not set.\n\t\tDefault is `n_points_in_φ = 36`. 2D field will be extended to 36 points in φ."
    #         n_points_in_φ = 36
    #     else
    #         if !(n_points_in_φ > 1 && iseven(n_points_in_φ))
    #             @info "\tIn weighing potential calculation: Keyword `n_points_in_φ` is $(n_points_in_φ) but must be even and larger than 1.\n\t\t`n_points_in_φ` is now set to 36. 2D field will be extended to 36 points in φ."
    #             n_points_in_φ = 36
    #         end
    #     end
    # end
    # wps = calculate_weighting_potential(sim.detector, contact_id, args...; kwargs...)
    # if S == :cylindrical && size(wps.potential, 2) == 1 && !ismissing(n_points_in_φ)
    #     wp = WeightingPotential(wps, n_points_in_φ = n_points_in_φ)
    # else
    #     wp = WeightingPotential(wps)
    # end
    _calculate_potential!(sim, WeightingPotential, contact_id, args...; kwargs...)
    nothing
end

"""
    calculate_electric_potential!(sim::Simulation{T}; kwargs...)::Nothing

Compute the electric potential for the given Simulation `sim` on an adaptive grid
through successive over relaxation.

There are serveral `<keyword arguments>` which can be used to tune the computation:

# Keywords
- `convergence_limit::Real`: `convergence_limit` times the bias voltage sets the convergence limit of the relaxation.
    The convergence value is the absolute maximum difference of the potential between two iterations of all grid points.
    Default of `convergence_limit` is `2e-6` (times bias voltage).
- `max_refinements::Int`: Number of maximum refinements. Default is `2`. Set it to `0` to switch off refinement.
- `refinement_limits::Tuple{<:Real, <:Real, <:Real}`: Tuple of refinement limits for each dimension
    (in case of cylindrical coordinates the order is `r`, `φ`, `z`).
    A refinement limit (e.g. `refinement_limits[1]`) times the bias voltage of the detector `det` is the
    maximum allowed voltage difference between two neighbouring grid points in the respective dimension.
    When the difference is larger, new points are created inbetween. Default is `[1e-5, 1e-5, 1e-5]`.
- `init_grid_spacing::Tuple{<:Real, <:Real, <:Real}`: Tuple of the initial distances between two grid points for each dimension.
    For normal coordinates the unit is meter. For angular coordinates, the unit is radiance.
    It prevents the refinement to make the grid to fine. Default is `[0.005, 10.0, 0.005]``.
- `min_grid_spacing::Tuple{<:Real, <:Real, <:Real}`: Tuple of the mimimum allowed distance between two grid points for each dimension.
    For normal coordinates the unit is meter. For angular coordinates, the unit is radiance.
    It prevents the refinement to make the grid to fine. Default is [`1e-6`, `1e-6`, `1e-6`].
- `grid::Grid{T, N, S}`: Initial grid used to start the simulation. Default is `Grid(detector, init_grid_spacing=init_grid_spacing)`.
- `depletion_handling::Bool`: Enables the handling of undepleted regions. Default is false.
- `use_nthreads::Int`: Number of threads to use in the computation. Default is `Base.Threads.nthreads()`.
    The environment variable `JULIA_NUM_THREADS` must be set appropriately before the Julia session was
    started (e.g. `export JULIA_NUM_THREADS=8` in case of bash).
- `sor_consts::Union{<:Real, NTuple{2, <:Real}}`: Two element tuple in case of cylindrical coordinates.
    First element contains the SOR constant for `r` = 0.
    Second contains the constant at the outer most grid point in `r`. A linear scaling is applied in between.
    First element should be smaller than the second one and both should be ∈ [1.0, 2.0]. Default is [1.4, 1.85].
    In case of cartesian coordinates only one value is taken.
- `max_n_iterations::Int`: Set the maximum number of iterations which are performed after each grid refinement.
    Default is `10000`. If set to `-1` there will be no limit.
- `verbose::Bool=true`: Boolean whether info output is produced or not.
"""
function calculate_electric_potential!(sim::Simulation{T}, args...; kwargs...)::Nothing where {T <: SSDFloat}
    _calculate_potential!(sim, ElectricPotential, args...; kwargs...)
    nothing
end

"""
    calculate_electric_field!(sim::Simulation{T}, args...; n_points_in_φ::Union{Missing, Int} = missing, kwargs...)::Nothing

ToDo...
"""
function calculate_electric_field!(sim::Simulation{T}, args...; n_points_in_φ::Union{Missing, Int} = missing, kwargs...)::Nothing where {T <: SSDFloat}
    S = get_coordinate_system(sim.detector)
    periodicity::T = get_periodicity(sim.detector.world.intervals[2])
    e_pot, point_types = if S == :cylindrical && periodicity == T(0) # 2D, only one point in φ
        if ismissing(n_points_in_φ)
            @info "\tIn electric field calculation: Keyword `n_points_in_φ` not set.\n\t\tDefault is `n_points_in_φ = 36`. 2D field will be extended to 36 points in φ."
            n_points_in_φ = 36
        else
            if !(n_points_in_φ > 1 && iseven(n_points_in_φ))
                @info "\tIn electric field calculation: Keyword `n_points_in_φ` is $(n_points_in_φ) but must be even and larger than 1.\n\t\t`n_points_in_φ` is now set to 36. 2D field will be extended to 36 points in φ."
                n_points_in_φ = 36
            end
        end
        get_2π_potential(sim.electric_potential, n_points_in_φ = n_points_in_φ),
        get_2π_potential(sim.point_types,  n_points_in_φ = n_points_in_φ);
    elseif S == :cylindrical
        get_2π_potential(sim.electric_potential),
        get_2π_potential(sim.point_types)
    else
        sim.electric_potential,
        sim.point_types
    end
    sim.electric_field = get_electric_field_from_potential(e_pot, point_types);
    nothing
end

function set_charge_drift_model!(sim::Simulation{T}, charge_drift_model::AbstractChargeDriftModel{T})::Nothing where {T <: SSDFloat}
    sim.charge_drift_model = charge_drift_model
    nothing
end

function calculate_drift_fields!(sim::Simulation{T};
    use_nthreads::Int = Base.Threads.nthreads())::Nothing where {T <: SSDFloat}
    sim.electron_drift_field = ElectricField(get_electron_drift_field(sim.electric_field.data, sim.charge_drift_model, use_nthreads = use_nthreads), sim.electric_field.grid)
    sim.hole_drift_field = ElectricField(get_hole_drift_field(sim.electric_field.data, sim.charge_drift_model, use_nthreads = use_nthreads), sim.electric_field.grid)
    nothing
end
@deprecate apply_charge_drift_model!(args...; kwargs...) calculate_drift_fields!(args...; kwargs...)

function get_interpolated_drift_field(ef::ElectricField)
    get_interpolated_drift_field(ef.data, ef.grid)
end

function drift_charges( sim::Simulation{T}, starting_positions::Vector{CartesianPoint{T}};
                        Δt::RealQuantity = 5u"ns", max_nsteps::Int = 1000, verbose::Bool = true )::Vector{EHDriftPath{T}} where {T <: SSDFloat}
    return _drift_charges(   sim.detector, sim.point_types.grid, sim.point_types, starting_positions,
                             get_interpolated_drift_field(sim.electron_drift_field), get_interpolated_drift_field(sim.hole_drift_field),
                             Δt, max_nsteps = max_nsteps, verbose = verbose)::Vector{EHDriftPath{T}}
end

function get_signal(sim::Simulation{T}, drift_paths::Vector{EHDriftPath{T}}, energy_depositions::Vector{T}, contact_id::Int; Δt::TT = T(5) * u"ns") where {T <: SSDFloat, TT}
    dt::T = to_internal_units(internal_time_unit, Δt)
    wp::Interpolations.Extrapolation{T, 3} = interpolated_scalarfield(sim.weighting_potentials[contact_id])
    timestamps = _common_timestamps( drift_paths, dt )
    signal::Vector{T} = zeros(T, length(timestamps))
    add_signal!(signal, timestamps, drift_paths, energy_depositions, wp, Val(get_coordinate_system(sim.detector)))
    return RDWaveform( range(zero(T) * unit(Δt), step = T(ustrip(Δt)) * unit(Δt), length = length(signal)), signal )
end

"""
    function simulate!(sim::Simulation{T};  max_refinements::Int = 1, verbose::Bool = false,
                                        depletion_handling::Bool = false, convergence_limit::Real = 1e-5 ) where {T <: SSDFloat}

ToDo...
"""
function simulate!(sim::Simulation{T};  max_refinements::Int = 1, verbose::Bool = false,
                                        depletion_handling::Bool = false, convergence_limit::Real = 1e-7 ) where {T <: SSDFloat}
    calculate_electric_potential!(  sim,
                                    max_refinements = max_refinements,
                                    verbose = verbose,
                                    init_grid_size = (10,10,10),
                                    depletion_handling = depletion_handling,
                                    convergence_limit = convergence_limit )
    for contact in sim.detector.contacts
        calculate_weighting_potential!(sim, contact.id, max_refinements = max_refinements,
                init_grid_size = (10,10,10),
                verbose = verbose, convergence_limit = convergence_limit)
    end
    calculate_electric_field!(sim)
    set_charge_drift_model!(sim, sim.charge_drift_model)
    calculate_drift_fields!(sim)
    @info "Detector simulation done"
end

function _get_abs_bias_voltage(det::SolidStateDetector{T}) where {T <: SSDFloat}
    potentials::Vector{T} = map(c -> c.potential, det.contacts)
    return (maximum(potentials) - minimum(potentials)) * u"V"
end

calculate_stored_energy(sim::Simulation) =
    calculate_stored_energy(sim.electric_field, sim.ϵ)

function calculate_stored_energy(ef::ElectricField{T,3,S}, ϵ::DielectricDistribution{T,3,S}) where {T <: SSDFloat, S}
    W::T = 0

    cylindric::Bool = S == :cylindrical
    cartesian::Bool = !cylindric
    r0_handling::Bool = typeof(ef.grid.axes[1]).parameters[2] == :r0

    ax1::Vector{T} = collect(ef.grid[1])
    ax2::Vector{T} = collect(ef.grid[2])
    ax3::Vector{T} = collect(ef.grid[3])
    mp1::Vector{T} = midpoints(get_extended_ticks(ef.grid[1]))
    mp2::Vector{T} = midpoints(get_extended_ticks(ef.grid[2]))
    mp3::Vector{T} = midpoints(get_extended_ticks(ef.grid[3]))
    Δax1::Vector{T} = diff(ax1)
    Δax2::Vector{T} = diff(ax2)
    Δax3::Vector{T} = diff(ax3)
    Δmp1::Vector{T} = diff(mp1)
    Δmp2::Vector{T} = diff(mp2)
    Δmp3::Vector{T} = diff(mp3)

    w1r::Vector{T} = inv.(Δmp1) .* (mp1[2:end] .- ax1)
    w1l::Vector{T} = inv.(Δmp1) .* (ax1 - mp1[1:end-1])
    w2r::Vector{T} = inv.(Δmp2) .* (mp2[2:end] .- ax2)
    w2l::Vector{T} = inv.(Δmp2) .* (ax2 - mp2[1:end-1])
    w3r::Vector{T} = inv.(Δmp3) .* (mp3[2:end] .- ax3)
    w3l::Vector{T} = inv.(Δmp3) .* (ax3 - mp3[1:end-1])

    if cylindric
        mp1[1] = 0
        mp1[end] = ax1[end]
        Δmp1 = ((mp1[2:end].^2) .- (mp1[1:end-1].^2)) ./ 2
    end
    V::T = 0
    for i3 in 1:size(ϵ, 3)-1
        _Δmp3::T = Δmp3[i3]
        if (i3 == 1 || i3 == size(ϵ, 3)-1) _Δmp3 /= 2 end
        for i2 in 1:size(ϵ, 2)-1
            _Δmp2::T = Δmp2[i2]
            if (cartesian && (i2 == 1 || i2 == size(ϵ, 2)-1)) _Δmp2 /= 2 end
            for i1 in 1:size(ϵ, 1)-1
                _Δmp1::T = Δmp1[i1]
                if (cartesian && (i1 == 1 || i1 == size(ϵ, 1)-1)) _Δmp1 /= 2 end
                ev::SArray{Tuple{3},Float32,1,3} = ef.data[i1, i2, i3]
                dV::T = _Δmp3 * _Δmp2 * _Δmp1
                _ϵ::T = sum([
                    ϵ[i1, i2, i3]             * w1l[i1] * w2l[i2] * w3l[i3],
                    ϵ[i1 + 1, i2, i3]         * w1r[i1] * w2l[i2] * w3l[i3],
                    ϵ[i1, i2 + 1, i3]         * w1l[i1] * w2r[i2] * w3l[i3],
                    ϵ[i1, i2, i3 + 1]         * w1l[i1] * w2l[i2] * w3r[i3],
                    ϵ[i1 + 1, i2 + 1, i3]     * w1r[i1] * w2r[i2] * w3l[i3],
                    ϵ[i1 + 1, i2, i3 + 1]     * w1r[i1] * w2l[i2] * w3r[i3],
                    ϵ[i1, i2 + 1, i3 + 1]     * w1l[i1] * w2r[i2] * w3r[i3],
                    ϵ[i1 + 1, i2 + 1, i3 + 1] * w1r[i1] * w2r[i2] * w3r[i3]
                ])
                V += dV
                W += sum(ev.^2) * dV * _ϵ
            end
        end
    end
    E =  W * ϵ0 / 2 * u"J"
    if cylindric && (size(ϵ, 2) - 1 != size(ef, 2))
        E *= size(ef, 2) / (size(ϵ, 2) - 1)
    end
    return E
end

export calculate_capacitance
"""
    calculate_capacitance(sim::Simulation{T})::T where {T <: SSDFloat}

Calculates the capacitance of an detector in Farad.
"""
function calculate_capacitance(sim::Simulation{T}) where {T <: SSDFloat}
    W = calculate_stored_energy(sim)
    return uconvert(u"pF", 2 * W / (_get_abs_bias_voltage(sim.detector)^2))
end


# This file establishes the default initial conditions, boundary conditions and sources
# for the baroclinicwave_problem experiment, following
#
# Ullrich, P. A., Melvin, T., Jablonowski, C., and Staniforth, A.:
# A proposed baroclinic wave test case for deep- and shallow atmosphere
# dynamical cores, Q. J. Roy. Meteor. Soc., 140, 1590-1602, doi:10.1002/qj.2241, 2014.

# Override default CLIMAParameters for consistency with literature on this case
CLIMAParameters.Planet.press_triple(::EarthParameterSet) = 610.78

struct BaroclinicWaveProblem{BC, ISP, ISA, WP, BS, MP} <: AbstractAtmosProblem
    boundarycondition::BC
    init_state_prognostic::ISP
    init_state_auxiliary::ISA
    perturbation::WP
    base_state::BS
    moisture_profile::MP
end
function BaroclinicWaveProblem(;
    boundarycondition = (AtmosBC(), AtmosBC()),
    perturbation = nothing,
    base_state = nothing,
    moisture_profile = nothing,
)
    # Set up defaults
    if isnothing(perturbation)
        perturbation = DeterministicPerturbation()
    end
    if isnothing(base_state)
        base_state = BCWaveBaseState()
    end
    if isnothing(moisture_profile)
        moisture_profile = MoistLowTropicsMoistureProfile()
    end

    problem = (
        boundarycondition,
        init_gcm_experiment!,
        (_...) -> nothing,
        perturbation,
        base_state,
        moisture_profile,
    )
    return BaroclinicWaveProblem{typeof.(problem)...}(problem...)
end

problem_name(::BaroclinicWaveProblem) = "BaroclinicWave"

setup_source(::BaroclinicWaveProblem) = (Gravity(), Coriolis())


# wrappers of the low-level CUBLAS functionality

import CUDAdrv: CuStream

# Note: we don't implement padded storage dimensions
function _mkplan(xtype, xdims, region)
    nrank = length(region)
    sz = [xdims[i] for i in region]
    csz = copy(sz)
    csz[1] = div(sz[1],2) + 1
    batch = prod(xdims) ÷ prod(sz)

    pp = Ref{cufftHandle_t}()
    if (nrank == 1) && (batch == 1)
        cufftPlan1d(pp, sz[1], xtype, 1)
    elseif (nrank == 2) && (batch == 1)
        cufftPlan2d(pp, sz[2], sz[1], xtype)
    elseif (nrank == 3) && (batch == 1)
        cufftPlan3d(pp, sz[3], sz[2], sz[1], xtype)
    else
        rsz = (length(sz) > 1) ? rsz = reverse(sz) : sz
        if ((region...,) == ((1:nrank)...,))
            # handle simple case ... simply! (for robustness)
           cufftPlanMany(pp, nrank, Cint[rsz...], C_NULL, 1, 1, C_NULL, 1, 1,
                         xtype, batch)
        else
            if nrank==1 || all(diff(collect(region)) .== 1)
                # _stride: successive elements in innermost dimension
                # _dist: distance between first elements of batches
                if region[1] == 1
                    istride = 1
                    idist = prod(sz)
                    cdist = prod(csz)
                else
                    if region[end] != length(xdims)
                        throw(ArgumentError("batching dims must be sequential"))
                    end
                    istride = prod(xdims[1:region[1]-1])
                    idist = 1
                    cdist = 1
                end
                inembed = Cint[rsz...]
                cnembed = (length(csz) > 1) ? Cint[reverse(csz)...] : Cint[csz[1]]
                ostride = istride
                if xtype == CUFFT_R2C || xtype == CUFFT_D2Z
                    odist = cdist
                    onembed = cnembed
                else
                    odist = idist
                    onembed = inembed
                end
                if xtype == CUFFT_C2R || xtype == CUFFT_Z2D
                    idist = cdist
                    inembed = cnembed
                end
            else
                if any(diff(collect(region)) .< 1)
                    throw(ArgumentError("region must be an increasing sequence"))
                end
                cdims = collect(xdims)
                cdims[region[1]] = div(cdims[region[1]],2)+1

                if region[1] == 1
                    istride = 1
                    ii=1
                    while (ii < nrank) && (region[ii] == region[ii+1]-1)
                        ii += 1
                    end
                    idist = prod(xdims[1:ii])
                    cdist = prod(cdims[1:ii])
                    ngaps = 0
                else
                    istride = prod(xdims[1:region[1]-1])
                    idist = 1
                    cdist = 1
                    ngaps = 1
                end
                nem = ones(Int,nrank)
                cem = ones(Int,nrank)
                id = 1
                for ii=1:nrank-1
                    if region[ii+1] > region[ii]+1
                        ngaps += 1
                    end
                    while id < region[ii+1]
                        nem[ii] *= xdims[id]
                        cem[ii] *= cdims[id]
                        id += 1
                    end
                    @assert nem[ii] >= sz[ii]
                end
                if region[end] < length(xdims)
                    ngaps += 1
                end
                # CUFFT represents batches by a single stride (_dist)
                # so we must verify that region is consistent with this:
                if ngaps > 1
                    throw(ArgumentError("batch regions must be sequential"))
                end

                inembed = Cint[reverse(nem)...]
                cnembed = Cint[reverse(cem)...]
                ostride = istride
                if xtype == CUFFT_R2C || xtype == CUFFT_D2Z
                    odist = cdist
                    onembed = cnembed
                else
                    odist = idist
                    onembed = inembed
                end
                if xtype == CUFFT_C2R || xtype == CUFFT_Z2D
                    idist = cdist
                    inembed = cnembed
                end
            end
            cufftPlanMany(pp, nrank, Cint[rsz...],
                          inembed, istride, idist, onembed, ostride, odist,
                          xtype, batch)
        end
    end
    pp[]
end

# this is used implicitly in the unsafe_execute methods below:
unsafe_convert(::Type{cufftHandle_t}, p::CuFFTPlan) = p.plan

convert(::Type{cufftHandle_t}, p::CuFFTPlan) = p.plan

destroy_plan(plan::CuFFTPlan) = cufftDestroy(plan)

set_stream(plan::CuFFTPlan, stream::CuStream) = cufftSetStream(plan, stream)

function assert_applicable(p::CuFFTPlan{T,K}, X::CuArray{T}) where {T,K}
    (size(X) == p.sz) ||
        throw(ArgumentError("CuFFT plan applied to wrong-size input"))
end

function assert_applicable(p::CuFFTPlan{T,K}, X::CuArray{T}, Y::CuArray{Ty}) where {T,K,Ty}
    assert_applicable(p, X)
    (size(Y) == p.osz) ||
        throw(ArgumentError("CuFFT plan applied to wrong-size output"))
    # type errors should be impossible by dispatch, but just in case:
    if p.xtype ∈ [CUFFT_C2R, CUFFT_Z2D]
        (Ty == real(T)) ||
            throw(ArgumentError("Type mismatch for argument Y"))
    elseif p.xtype ∈ [CUFFT_R2C, CUFFT_D2Z]
        (Ty == complex(T)) ||
            throw(ArgumentError("Type mismatch for argument Y"))
    else
        (Ty == T) ||
            throw(ArgumentError("Type mismatch for argument Y"))
    end
end

function unsafe_execute!(plan::cCuFFTPlan{cufftComplex,K,true,N},
                         x::CuArray{cufftComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_C2C
    cufftExecC2C(plan, x, x, K)
end
function unsafe_execute!(plan::rCuFFTPlan{cufftComplex,K,true,N},
                         x::CuArray{cufftComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_C2R
    cufftExecC2R(plan, x, x)
end

function unsafe_execute!(plan::cCuFFTPlan{cufftComplex,K,false,N},
                         x::CuArray{cufftComplex,N}, y::CuArray{cufftComplex}
                         ) where {K,N}
    @assert plan.xtype == CUFFT_C2C
    cufftExecC2C(plan, x, y, K)
end
function unsafe_execute!(plan::rCuFFTPlan{cufftComplex,K,false,N},
                         x::CuArray{cufftComplex,N}, y::CuArray{cufftReal}
                         ) where {K,N}
    @assert plan.xtype == CUFFT_C2R
    cufftExecC2R(plan, x, y)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftReal,K,false,N},
                         x::CuArray{cufftReal,N}, y::CuArray{cufftComplex,N}
                         ) where {K,N}
    @assert plan.xtype == CUFFT_R2C
    cufftExecR2C(plan, x, y)
end

# double prec.
function unsafe_execute!(plan::cCuFFTPlan{cufftDoubleComplex,K,true,N},
                         x::CuArray{cufftDoubleComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_Z2Z
    cufftExecZ2Z(plan, x, x, K)
end
function unsafe_execute!(plan::rCuFFTPlan{cufftDoubleComplex,K,true,N},
                         x::CuArray{cufftDoubleComplex,N}) where {K,N}
    @assert plan.xtype == CUFFT_Z2D
    cufftExecZ2D(plan, x, x)
end

function unsafe_execute!(plan::cCuFFTPlan{cufftDoubleComplex,K,false,N},
                         x::CuArray{cufftDoubleComplex,N}, y::CuArray{cufftDoubleComplex}
                         ) where {K,N}
    @assert plan.xtype == CUFFT_Z2Z
    cufftExecZ2Z(plan, x, y, K)
end
function unsafe_execute!(plan::rCuFFTPlan{cufftDoubleComplex,K,false,N},
                         x::CuArray{cufftDoubleComplex,N}, y::CuArray{cufftDoubleReal}
                         ) where {K,N}
    @assert plan.xtype == CUFFT_Z2D
    cufftExecZ2D(plan, x, y)
end

function unsafe_execute!(plan::rCuFFTPlan{cufftDoubleReal,K,false,N},
                         x::CuArray{cufftDoubleReal,N}, y::CuArray{cufftDoubleComplex,N}
                         ) where {K,N}
    @assert plan.xtype == CUFFT_D2Z
    cufftExecD2Z(plan, x, y)
end


#=

  "PANDA: Phylogenetic ANalyses of DiversificAtion" module package

=#

module PANDA

__precompile__(true) #put back to true when checking done

#=
 Submodules
=#

# Utilities
include("Utils.jl")

# ESSE: Environmental and State Dependent Diversification
include("ESSE.jl")

# TRIBE: Trait and Range Interspecific Biogeographic Evolution
include("TRIBE.jl")

# ClaDS: Cladogenetic XX
include("ClaDS.jl")

#=
 Exported functions
=#

using .ESSE: esse, simulate_sse
export esse, simulate_sse

using .TRIBE: tribe, simulate_tribe
export tribe, simulate_tribe

using .ClaDS: Tree, tip_labels, n_tips, sample_tips,
    sim_ClaDS2_ntips, plot_ClaDS, infer_ClaDS,
    CladsOutput, print_CladsOutput, plot_CladsOutput, tip_rate,
    load_tree, save_ClaDS_in_R
export Tree, tip_labels, n_tips, sample_tips,
    sim_ClaDS2_ntips, plot_ClaDS, infer_ClaDS,
    CladsOutput, print_CladsOutput, plot_CladsOutput, tip_rate,
    load_tree, save_ClaDS_in_R
end


# © Przemyslaw Szufel, 2018, run under Julia 1.0.0

function do_item(modelC :: ModelConfig,rng::MersenneTwister,worker_grid::Matrix{Int},worker_locations::Dict{Int,Tuple{Int,Int}},p2p::Bool) :: String
    GRID_SIZE=modelC.GRID_SIZE
    status=zeros(Float64,(GRID_SIZE))
    borders=Matrix{BorderData}(undef,GRID_SIZE...)
    leaveLocData=Matrix{LeaveLocData}(undef,GRID_SIZE...)
    dim_i=modelC.dim_i
    dim_j=modelC.dim_j

    avg_asyn_bord = 0.
    max_asyn_bord = 0.

    start = time()
    if p2p
        @sync for ix in modelC.slaves
            (i,j) = worker_locations[ix]
            @async status[i,j] = remotecall_fetch(step_setborder_nodes_p2p!,worker_grid[i,j])
        end
        avg_asyn_bord = mean(status)
        max_asyn_bord = maximum(status)
    else
        @sync for ix in modelC.slaves
            (i,j) = worker_locations[ix]
            @async borders[i,j] = remotecall_fetch(get_border,worker_grid[i,j])
        end

        @sync for ix in modelC.slaves
            (i,j) = worker_locations[ix]
            j_left = j-1
            if j_left < 1; j_left = GRID_SIZE[2]; end
            j_right = j+1
            if j_right > GRID_SIZE[2]; j_right = 1; end
            i_top = i-1
            if i_top < 1; i_top = GRID_SIZE[1]; end
            i_bottom = i+1
            if i_bottom > GRID_SIZE[1]; i_bottom = 1; end

            sendto(worker_grid[i,j],border_data=BorderData( #left right top bottom
                    borders[i,j_left].right,  borders[i,j_right].left,
                    borders[i_top,j].bottom,  borders[i_bottom,j].top),
                corner_data=CornerData( #topleft topright bottomleft bottomright
                    borders[i_top,j_left].bottom[modelC.dim_j],  borders[i_top,j_right].bottom[1],
                    borders[i_bottom,j_left].top[modelC.dim_j],  borders[i_bottom,j_right].top[1])
                    )
            @async status[i,j] = remotecall_fetch(step_setborder_nodes!,worker_grid[i,j])
        end
    end

    @sync for ix in modelC.slaves
        (i,j) = worker_locations[ix]
        @async begin
            leaveLocData[i,j] = remotecall_fetch(step_leave_location!,worker_grid[i,j])
            status[i,j] = leaveLocData[i,j].elapsed_time
        end
    end
    max_leave_loc_time = maximum(status)
    avg_leave_loc_time = mean(status)

    reloc_start_time = time()

    add_new_citizens=Matrix{Vector{Int}}(undef,GRID_SIZE...)

    # all units to be located on the board of each color type
    relocate_counts = zeros(Int,length(modelC.probs_w)-1)
    free_cells_count=zeros(Int,GRID_SIZE)

    for ix in modelC.slaves
        (i,j) = worker_locations[ix]
        relocate_counts += leaveLocData[i,j].relocate_counts
        add_new_citizens[i,j]=zeros(Int,length(modelC.probs_w)-1)
        free_cells_count[i,j]=leaveLocData[i,j].free_cells_count
    end

    while sum(relocate_counts) > 0
        for t in 1:length(relocate_counts)
            if relocate_counts[t] == 0
                continue
            end
            relocated = 0
            total_free_cells_count = sum(free_cells_count)
            for ix in shuffle(rng,modelC.slaves)
                (i,j) = worker_locations[ix]
                ρ = free_cells_count[i,j]/total_free_cells_count
                μ = ρ * relocate_counts[t]
                if μ < 2.0
                    μ = 2.0
                end
                σ = √(μ * (1-ρ))
                val = round(randn(rng)*σ+μ)
                t_rand = max(0, Int(min(val,relocate_counts[t]-relocated,free_cells_count[i,j]))  )
                free_cells_count[i,j] -= t_rand
                add_new_citizens[i,j][t] += t_rand
                relocated += t_rand
            end
            relocate_counts[t] -= relocated
        end
    end
    reloc_time = time() - reloc_start_time

    @sync for ix in modelC.slaves
        (i,j) = worker_locations[ix]

        sendto(worker_grid[i,j],new_citizens=add_new_citizens[i,j])
        @async status[i,j] = remotecall_fetch(step_set_new_location!,worker_grid[i,j])
    end
    max_new_loc_time = maximum(status)
    avg_new_loc_time = mean(status)
    elapsed_time = time()-start
    return string("$GRID_SIZE x $dim_i x $dim_j t: $(@sprintf("%.3f", elapsed_time  )) = loc $(@sprintf("%.3f", reloc_time  ))  avg/max_asyn_bord $(@sprintf("%.3f", avg_asyn_bord  ))/$(@sprintf("%.3f", max_asyn_bord  )) avg/max_leave $(@sprintf("%.3f", avg_leave_loc_time))/$(@sprintf("%.3f", max_leave_loc_time)) avg/max_new $(@sprintf("%.3f", avg_new_loc_time  ))/$(@sprintf("%.3f", max_new_loc_time  )) ")

end
function do_item() :: String
    return do_item(modelC,rng,worker_grid,worker_locations,p2p)
end


## ScalarAffineFunction-in-Set
## ScalarQuadraticFunction-in-Set

function _build_constraint(
    model::Optimizer,
    f::F,
    s::MOI.GreaterThan{T},
) where {
    T <: Real,
    F <: Union{MOI.ScalarAffineFunction{T}, MOI.ScalarQuadraticFunction{T}},
}
    return cpo_java_ge(model.inner, _parse(model, f), s.lower)
end

function _build_constraint(
    model::Optimizer,
    f::F,
    s::MOI.LessThan{T},
) where {
    T <: Real,
    F <: Union{MOI.ScalarAffineFunction{T}, MOI.ScalarQuadraticFunction{T}},
}
    return cpo_java_le(model.inner, _parse(model, f), s.upper)
end

function _build_constraint(
    model::Optimizer,
    f::F,
    s::MOI.EqualTo{T},
) where {
    T <: Real,
    F <: Union{MOI.ScalarAffineFunction{T}, MOI.ScalarQuadraticFunction{T}},
}
    return cpo_java_eq(model.inner, _parse(model, f), s.value)
end

function _build_constraint(
    model::Optimizer,
    f::F,
    s::MOI.EqualTo{T},
) where {
    T <: Integer,
    F <: Union{MOI.ScalarAffineFunction{T}, MOI.ScalarQuadraticFunction{T}},
}
    return cpo_java_eq(model.inner, _parse(model, f), Int32(s.value))
end

function _build_constraint(
    model::Optimizer,
    f::F,
    s::MOI.Interval{T},
) where {
    T <: Real,
    F <: Union{MOI.ScalarAffineFunction{T}, MOI.ScalarQuadraticFunction{T}},
}
    if s.lower == Inf
        return _build_constraint(model, f, MOI.LessThan(s.upper))
    elseif s.upper == Inf
        return _build_constraint(model, f, MOI.GreaterThan(s.lower))
    elseif s.lower == s.upper
        return _build_constraint(model, f, MOI.EqualTo(s.lower))
    end

    return cpo_java_range(model.inner, s.lower, _parse(model, f), s.upper)
end

# No vector of constraints, there is no more efficient way to do it.

function MOI.delete(
    model::Optimizer,
    c::MOI.ConstraintIndex{F, S},
) where {
    T <: Real,
    F <: Union{MOI.ScalarAffineFunction{T}, MOI.ScalarQuadraticFunction{T}},
    S <:
    Union{MOI.GreaterThan{T}, MOI.LessThan{T}, MOI.EqualTo{T}, MOI.Interval{T}},
}
    cpo_java_remove(model.inner, _info(model, c).constraint)
    delete!(model.constraint_info, c)
    return
end

# TODO: function MOI.set(model::Optimizer, ::MOI.ConstraintSet, c::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, S}, s::S) where {S}

## VectorOfVariables-in-SOS{I|II}
# Not available. Bridge it?

## VectorOfVariables-in-SecondOrderCone

function _build_constraint(
    model::Optimizer,
    f::MOI.VectorOfVariables,
    s::MOI.SecondOrderCone,
)
    if length(f.variables) != s.dimension
        error("Dimension of $(s) does not match number of terms in $(f)")
    end

    # First, check the lower bound on t.
    t_info = _info(model, f.variables[1])
    if !_has_lb(model, t_info.index) || _get_lb(model, t_info.index) < 0.0
        if _get_lb(model, t_info.index) < 0.0
            t_info.old_lb = _get_lb(model, t_info.index)
        end

        t_info.n_socs += 1
        _set_lb(model, t_info.index, 0.0)
    end

    # Then, add the quadratic constraint.
    expr = _parse(model, f, s)
    return cpo_java_gt(model, expr, 0)
end

function MOI.delete(
    model::Optimizer,
    c::MOI.ConstraintIndex{MOI.VectorOfVariables, MOI.SecondOrderCone},
)
    # Remove the constraint.
    cpo_java_remove(model, _info(model, c).constraint)

    # Maybe restore the old bound on t.
    t_info = _info(model, f.variables[1])
    t_info.n_socs -= 1

    if t_info.n_socs == 0
        _set_lb(model, t_info.index, t_info.old_lb, typeof(t_info.old_lb))
        t_info.old_lb = nothing
    end

    return
end


ur2() = uppercase(randstring(2))

# Acceptable type codes in :misc
OK = UInt8[0x00, 0x01,
      0x10, 0x11, 0x12, 0x13, 0x14,
      0x20, 0x21, 0x22, 0x23, 0x24,
      0x30, 0x31, 0x32,
      0x50, 0x51, 0x52, 0x53, 0x54,
      0x60, 0x61, 0x62, 0x63, 0x64,
      0x70, 0x71, 0x72,
      0x80, 0x81,
      0x90, 0x91, 0x92, 0x93, 0x94,
      0xa0, 0xa1, 0xa2, 0xa3, 0xa4,
      0xb0, 0xb1, 0xb2,
      0xd0, 0xd1, 0xd2, 0xd3, 0xd4,
      0xe0, 0xe1, 0xe2, 0xe3, 0xe4,
      0xf0, 0xf1, 0xf2]

# Fill :misc with garbage
function pop_rand_dict!(D::Dict{String,Any}, N::Int)
  for n = 1:N
    t = code2typ(rand(OK))
    k = randstring(rand(2:12))
    if Bool(t <: Real) == true
      D[k] = rand(t)
    elseif Bool(t <: Complex) == true
      D[k] = rand(Complex{real(t)})
    elseif Bool(t <: Array) == true
      y = eltype(t)
      if Bool(y <: Number) == true
        D[k] = rand(y, rand(1:1000))
      end
    end
  end
  if haskey(D, "hc")
    delete!(D, "hc")
  end
  return D
end

"""
   (i,c,u) = getyp2codes(b::Char, s=false::Bool)

Using band code `b`, generate quasi-sane random instrument char code (`i`)
and channel char code (`c`), plus unit string `u`. if s=true, use only seismic
data codes.
"""
function getyp2codes(b::Char, s=false::Bool)
  if s
    # Neglecting gravimeters ('G') and mass position sensors ('M')
    i = rand(['H','L','N'])
    if rand() > 0.2
      c = rand(['Z','N','E'])
    else
      c = rand(['A','B','C','1','2','3','U','V','W'])
    end
    if Base.in(i, ['H','L'])
      u = rand(["m", "m/s"])
    else
      u = "m/s2"
    end
  else
    i = rand(['A','B','D','F','G','I','J','K','M','O','P','Q','R','S','T','U','V','W','Z'])
    if i == 'A' # tiltmeter
      c = rand(['N','E'])
      u = "rad"
    elseif i == 'B' # creep meter
      c = '_'
      u = "m"
      # C is calibration input
    elseif i == 'D' # pressure (barometer, infrasound, hydrophone ∈ 'D')
      c = rand(['O','I','D','F','H','U'])
      u = "Pa"
      # E is an electronic test point
    elseif i == 'F' # magnetometer
      c = rand(['Z','N','E'])
      u = "T"
    elseif i == 'G' # tiltmeter
        c = rand(['A','B','C','1','2','3','U','V','W'])
        u = "m/s2"
    elseif i == 'I' # humidity
      c = rand(['O','I','D'])
      u = "%"
    elseif i == 'J' # rotational seismometer
      c = rand(['Z','N','E','A','B','C','T','R','1','2','3','U','V','W'])
      u = rand(["rad", "rad/s", "rad/s2"])
    elseif i == 'K' # thermal (thermometer or radiometer)
      c = rand(['O', 'I', 'D'])
      u = rand(["Cel","K"])
    elseif i == 'M' # mass position sensor
        c = rand(['A','B','C','1','2','3','U','V','W'])
        u = "m"
    elseif i == 'O' # current gauge
      c = '_'
      u = "m/s"
    elseif i == 'P' # very short-period geophone
      c = rand(['Z','N','E'])
      u = rand(["m", "m/s", "m/s2"])
    elseif i == 'Q' # voltmeter
      c = '_'
      u = "V"
    elseif i == 'R' # rain gauge
      c = '_'
      u = rand(["m", "m/s"])
    elseif i == 'S' # strain gauge
      c = rand(['Z','N','E'])
      u = "m/m"
    elseif i == 'T' # tide gauge
      c = 'Z'
      u = "m"
    elseif i == 'U' # bolometer
      c = '_'
      u = "%{cloud_cover}"
    elseif i == 'V' # volumetric strainmeter
      c = '_'
      u = "m3/m3"
    elseif i == 'W' # wind speed ('S') or direction ('D')
      c = rand(['S','D'])
      u = c == 'S' ? "m/s" : "{direction_vector}"

    # X, Y are instrument-specific...excluded

    elseif i == 'Z' # synthesized beam or stack
      c = rand(['I','C','F','O'])
      u = rand(["m", "m/s", "m/s2"])
    end
  end
  return i,c,u
end