const GAEVF = JuMP.GenericAffExpr{Float64, VariableRef}
const GQEVF = JuMP.GenericQuadExpr{Float64, VariableRef}

function moi_tests(
    op_problem::OperationsProblem,
    params::Bool,
    vars::Int,
    interval::Int,
    lessthan::Int,
    greaterthan::Int,
    equalto::Int,
    binary::Bool,
)

    JuMPmodel = op_problem.psi_container.JuMPmodel
    @test (:params in keys(JuMPmodel.ext)) == params
    @test JuMP.num_variables(JuMPmodel) == vars
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.Interval{Float64}) == interval
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.LessThan{Float64}) == lessthan
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.GreaterThan{Float64}) == greaterthan
    @test JuMP.num_constraints(JuMPmodel, GAEVF, MOI.EqualTo{Float64}) == equalto
    @test ((JuMP.VariableRef, MOI.ZeroOne) in JuMP.list_of_constraint_types(JuMPmodel)) ==
          binary

    return

end

function psi_constraint_test(
    op_problem::OperationsProblem,
    constraint_names::Vector{Symbol},
)
    for con in constraint_names
        @test !isnothing(get(op_problem.psi_container.constraints, con, nothing))
    end
    return
end

function psi_checkbinvar_test(
    op_problem::OperationsProblem,
    bin_variable_names::Vector{Symbol},
)
    for variable in bin_variable_names
        for v in PSI.get_variable(op_problem.psi_container, variable)
            @test JuMP.is_binary(v)
        end
    end
    return
end

function psi_checkobjfun_test(op_problem::OperationsProblem, exp_type)
    @test JuMP.objective_function_type(op_problem.psi_container.JuMPmodel) == exp_type
    return
end

function moi_lbvalue_test(op_problem::OperationsProblem, con_name::Symbol, value::Number)
    for con in op_problem.psi_container.constraints[con_name]
        @test JuMP.constraint_object(con).set.lower == value
    end
    return
end

function psi_checksolve_test(op_problem::OperationsProblem, status)
    JuMP.optimize!(op_problem.psi_container.JuMPmodel)
    @test termination_status(op_problem.psi_container.JuMPmodel) in status
end

function psi_checksolve_test(
    op_problem::OperationsProblem,
    status,
    expected_result,
    tol = 0.0,
)
    res = solve!(op_problem)
    @test termination_status(op_problem.psi_container.JuMPmodel) in status
    @test isapprox(get_total_cost(res)[:OBJECTIVE_FUNCTION], expected_result, atol = tol)
end


# Unicode iteration and decoding for StringView, copied from the corresponding String functions in Base

@inline between(b::T, lo::T, hi::T) where {T<:Integer} = (lo ≤ b) & (b ≤ hi)

@inline function Base.iterate(s::StringView, i::Int=firstindex(s))
    (i % UInt) - 1 < ncodeunits(s) || return nothing
    b = @inbounds codeunit(s, i)
    u = UInt32(b) << 24
    between(b, 0x80, 0xf7) || return reinterpret(Char, u), i+1
    return iterate_continued(s, i, u)
end

function iterate_continued(s::StringView, i::Int, u::UInt32)
    u < 0xc0000000 && (i += 1; @goto ret)
    n = ncodeunits(s)
    # first continuation byte
    (i += 1) > n && @goto ret
    @inbounds b = codeunit(s, i)
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b) << 16
    # second continuation byte
    ((i += 1) > n) | (u < 0xe0000000) && @goto ret
    @inbounds b = codeunit(s, i)
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b) << 8
    # third continuation byte
    ((i += 1) > n) | (u < 0xf0000000) && @goto ret
    @inbounds b = codeunit(s, i)
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b); i += 1
@label ret
    return reinterpret(Char, u), i
end

Base.@propagate_inbounds function Base.getindex(s::StringView, i::Int)
    b = codeunit(s, i)
    u = UInt32(b) << 24
    between(b, 0x80, 0xf7) || return reinterpret(Char, u)
    return getindex_continued(s, i, u)
end

function getindex_continued(s::StringView, i::Int, u::UInt32)
    if u < 0xc0000000
        # called from `getindex` which checks bounds
        @inbounds isvalid(s, i) && @goto ret
        Base.string_index_err(s, i)
    end
    n = ncodeunits(s)

    (i += 1) > n && @goto ret
    @inbounds b = codeunit(s, i) # cont byte 1
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b) << 16

    ((i += 1) > n) | (u < 0xe0000000) && @goto ret
    @inbounds b = codeunit(s, i) # cont byte 2
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b) << 8

    ((i += 1) > n) | (u < 0xf0000000) && @goto ret
    @inbounds b = codeunit(s, i) # cont byte 3
    b & 0xc0 == 0x80 || @goto ret
    u |= UInt32(b)
@label ret
    return reinterpret(Char, u)
end

Base.getindex(s::StringView, r::UnitRange{<:Integer}) = s[Int(first(r)):Int(last(r))]

@inline function Base.getindex(s::StringView, r::UnitRange{Int})
    isempty(r) && return ""
    i, j = first(r), last(r)
    @boundscheck begin
        checkbounds(s, r)
        @inbounds isvalid(s, i) || string_index_err(s, i)
        @inbounds isvalid(s, j) || string_index_err(s, j)
    end
    j = nextind(s, j) - 1
    return StringView(@view s.data[i:j])
end

Base.length(s::StringView) = length_continued(s, 1, ncodeunits(s), ncodeunits(s))

@inline function Base.length(s::StringView, i::Int, j::Int)
    @boundscheck begin
        0 < i ≤ ncodeunits(s)+1 || throw(BoundsError(s, i))
        0 ≤ j < ncodeunits(s)+1 || throw(BoundsError(s, j))
    end
    j < i && return 0
    @inbounds i, k = thisind(s, i), i
    c = j - i + (i == k)
    length_continued(s, i, j, c)
end

@inline function length_continued(s::StringView, i::Int, n::Int, c::Int)
    i < n || return c
    @inbounds b = codeunit(s, i)
    @inbounds while true
        while true
            (i += 1) ≤ n || return c
            0xc0 ≤ b ≤ 0xf7 && break
            b = codeunit(s, i)
        end
        l = b
        b = codeunit(s, i) # cont byte 1
        c -= (x = b & 0xc0 == 0x80)
        x & (l ≥ 0xe0) || continue

        (i += 1) ≤ n || return c
        b = codeunit(s, i) # cont byte 2
        c -= (x = b & 0xc0 == 0x80)
        x & (l ≥ 0xf0) || continue

        (i += 1) ≤ n || return c
        b = codeunit(s, i) # cont byte 3
        c -= (b & 0xc0 == 0x80)
    end
end


using StaticArrays: @SVector, normalize
using Accessors: Accessors, setproperties
export position
export direction

position(o) = o.position
direction(o) = o.direction
set_position(o, pos)  = setproperties(o, (position=pos,))
set_direction(o, dir) = setproperties(o, (direction=dir,))
set_position_direction(o, pos, dir) = setproperties(o, (position=pos, direction=dir))

Accessors.set(o, ::typeof(position), val) = set_position(o, val)
Accessors.set(o, ::typeof(direction), val) = set_direction(o, val)

function direction(p::AbstractParticle)
    @SVector [p.u, p.v, p.w]
end

has_z(::Type{<:AbstractParticle}) = true
has_z(::Type{<:EGSParticle})      = false

function position(p::P, z=nothing) where {P <: AbstractParticle}
    if (z == nothing)
        if has_z(P)
            return @SVector[p.x, p.y, p.z]
        else
            return @SVector[p.x, p.y]
        end
    else
        if has_z(P)
            @argcheck p.z ≈ z
            z = p.z
        end
        return @SVector[p.x, p.y, z]
    end
end

function set_direction(p::AbstractParticle, dir)
    u,v,w = dir
    setproperties(p, (u=u,v=v,w=w))
end

@inline function _xyz(pos, z)
    len = length(pos)
    if len == 3
        return pos
    elseif len == 2
        x,y = pos
        return x,y,z
    else
        throw(ArgumentError("Wrong length $(len))"))
    end
end

function set_position(p::AbstractParticle, pos)
    x,y,z = _xyz(pos, nothing)
    setproperties(p, (x=x,y=y,z=z))
end

function set_position(p::CompressedIAEAParticle, pos)
    set_position(IAEAParticle(p), pos)
end
function set_direction(p::CompressedIAEAParticle, dir)
    set_direction(IAEAParticle(p), dir)
end
function set_position_direction(p::CompressedIAEAParticle, pos, dir)
    set_position_direction(IAEAParticle(p), pos, dir)
end

function set_position_direction(p::AbstractParticle, pos, dir)
    x,y,z = _xyz(pos, nothing)
    u,v,w = dir
    setproperties(p, (x=x,y=y,z=z,u=u,v=v,w=w))
end


function MOI.add_constraint(
   model::Optimizer,
   func::MOI.ScalarAffineFunction{Float64},
   set::MOI.LessThan{Float64}
)
   check_inbounds(model, func)
   push!(model.linear_le_constraints, ConstraintInfo(func, set))
   n = length(model.linear_le_constraints)
   return MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}(n)
end

function MOI.add_constraint(
   model::Optimizer,
   func::MOI.ScalarAffineFunction{Float64},
   set::MOI.GreaterThan{Float64}
)
   check_inbounds(model, func)
   push!(model.linear_ge_constraints, ConstraintInfo(func, set))
   n = length(model.linear_ge_constraints)
   return MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}}(n)
end

function MOI.add_constraint(
   model::Optimizer,
   func::MOI.ScalarAffineFunction{Float64},
   set::MOI.EqualTo{Float64}
)
   check_inbounds(model, func)
   push!(model.linear_eq_constraints, ConstraintInfo(func, set))
   n = length(model.linear_eq_constraints)
   return MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}(n)
end

function MOI.add_constraint(
   model::Optimizer,
   func::MOI.ScalarQuadraticFunction{Float64},
   set::MOI.LessThan{Float64}
)
   check_inbounds(model, func)
   push!(model.quadratic_le_constraints, ConstraintInfo(func, set))
   n = length(model.quadratic_le_constraints)
   return MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}}(n)
end

function MOI.add_constraint(
   model::Optimizer,
   func::MOI.ScalarQuadraticFunction{Float64},
   set::MOI.GreaterThan{Float64}
)
   check_inbounds(model, func)
   push!(model.quadratic_ge_constraints, ConstraintInfo(func, set))
   n = length(model.quadratic_ge_constraints)
   return MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.GreaterThan{Float64}}(n)
end

function MOI.add_constraint(
   model::Optimizer,
   func::MOI.ScalarQuadraticFunction{Float64},
   set::MOI.EqualTo{Float64}
)
   check_inbounds(model, func)
   push!(model.quadratic_eq_constraints, ConstraintInfo(func, set))
   n = length(model.quadratic_eq_constraints)
   return MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.EqualTo{Float64}}(n)
end

function MOI.add_constraint(
   model::Optimizer,
   func::MOI.VectorAffineFunction{Float64},
   set::MOI.Complements
)
   check_inbounds(model, func)
   push!(model.complementarity_constraints, ConstraintInfo(func, set))
   n = length(model.complementarity_constraints)
   return MathOptInterface.ConstraintIndex{MOI.VectorAffineFunction{Float64}, MOI.Complements}(n)
end

function MOI.get(
   model::Optimizer,
   ::MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}
)
   return length(model.linear_le_constraints)
end

function MOI.get(
   model::Optimizer,
   ::MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}}
)
   return length(model.linear_ge_constraints)
end

function MOI.get(
   model::Optimizer,
   ::MOI.NumberOfConstraints{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}
)
   return length(model.linear_eq_constraints)
end

function MOI.get(
   model::Optimizer,
   ::MOI.NumberOfConstraints{MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}}
)
   return length(model.quadratic_le_constraints)
end

function MOI.get(
   model::Optimizer,
   ::MOI.NumberOfConstraints{MOI.ScalarQuadraticFunction{Float64}, MOI.GreaterThan{Float64}}
)
   return length(model.quadratic_ge_constraints)
end

function MOI.get(
   model::Optimizer,
   ::MOI.NumberOfConstraints{MOI.ScalarQuadraticFunction{Float64}, MOI.EqualTo{Float64}}
)
   return length(model.quadratic_eq_constraints)
end

function MOI.get(
   model::Optimizer,
   ::MOI.NumberOfConstraints{MOI.VectorAffineFunction{Float64}, MOI.Complements}
)
   return length(model.complementarity_constraints)
end

function MOI.set(
   model::Optimizer,
   attr::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.SingleVariable},
   value
)
   error("Constraint names for variable bound constraints not supported.")
   return
end

function MOI.set(
   model::Optimizer,
   attr::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}},
   value
)
   if ! (1 <= ci.value <= length(model.linear_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   model.linear_le_constraints[ci.value].name = value
   return
end

function MOI.set(
   model::Optimizer,
   attr::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}},
   value
)
   if ! (1 <= ci.value <= length(model.linear_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   model.linear_ge_constraints[ci.value].name = value
   return
end

function MOI.set(
   model::Optimizer,
   attr::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}},
   value
)
   if ! (1 <= ci.value <= length(model.linear_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   model.linear_eq_constraints[ci.value].name = value
   return
end

function MOI.set(
   model::Optimizer,
   attr::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}},
   value
)
   if ! (1 <= ci.value <= length(model.quadratic_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   model.quadratic_le_constraints[ci.value].name = value
   return
end

function MOI.set(
   model::Optimizer,
   attr::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.GreaterThan{Float64}},
   value
)
   if ! (1 <= ci.value <= length(model.quadratic_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   model.quadratic_ge_constraints[ci.value].name = value
   return
end

function MOI.set(
   model::Optimizer,
   attr::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.EqualTo{Float64}},
   value
)
   if ! (1 <= ci.value <= length(model.quadratic_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   model.quadratic_eq_constraints[ci.value].name = value
   return
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}
)
   if ! (1 <= ci.value <= length(model.linear_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.linear_le_constraints[ci.value].name
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}}
)
   if ! (1 <= ci.value <= length(model.linear_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.linear_ge_constraints[ci.value].name
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}
)
   if ! (1 <= ci.value <= length(model.linear_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.linear_eq_constraints[ci.value].name
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}}
)
   if ! (1 <= ci.value <= length(model.quadratic_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.quadratic_le_constraints[ci.value].name
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.GreaterThan{Float64}}
)
   if ! (1 <= ci.value <= length(model.quadratic_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.quadratic_ge_constraints[ci.value].name
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.EqualTo{Float64}}
)
   if ! (1 <= ci.value <= length(model.quadratic_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.quadratic_eq_constraints[ci.value].name
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintSet,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}
)
   if ! (1 <= ci.value <= length(model.linear_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return MOI.LessThan(model.linear_le_constraints[ci.value].set.upper)
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintSet,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}}
)
   if ! (1 <= ci.value <= length(model.linear_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return MOI.GreaterThan(model.linear_ge_constraints[ci.value].set.lower)
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintSet,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}
)
   if ! (1 <= ci.value <= length(model.linear_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return MOI.EqualTo(model.linear_eq_constraints[ci.value].set.value)
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintSet,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}}
)
   if ! (1 <= ci.value <= length(model.quadratic_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return MOI.LessThan(model.quadratic_le_constraints[ci.value].set.upper)
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintSet,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.GreaterThan{Float64}}
)
   if ! (1 <= ci.value <= length(model.quadratic_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return MOI.GreaterThan(model.quadratic_ge_constraints[ci.value].set.lower)
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintSet,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.EqualTo{Float64}}
)
   if ! (1 <= ci.value <= length(model.quadratic_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return MOI.EqualTo(model.quadratic_eq_constraints[ci.value].set.value)
end

function MOI.get(
   model::Optimizer,
   ::MOI.ConstraintName,
   ci::MOI.ConstraintIndex{MOI.VectorAffineFunction{Float64}, MOI.Complements}
)
   if ! (1 <= ci.value <= length(model.complementarity_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.complementarity_constraints[ci.value].name
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintPrimal,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if ! (1 <= ci.value <= length(model.linear_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.sol.equ["eq$(ci.value + offset_linear_le(model))"].level[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintPrimal,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if ! (1 <= ci.value <= length(model.linear_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.sol.equ["eq$(ci.value + offset_linear_ge(model))"].level[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintPrimal,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if ! (1 <= ci.value <= length(model.linear_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.sol.equ["eq$(ci.value + offset_linear_eq(model))"].level[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintPrimal,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if ! (1 <= ci.value <= length(model.quadratic_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.sol.equ["eq$(ci.value + offset_quadratic_le(model))"].level[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintPrimal,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.GreaterThan{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if ! (1 <= ci.value <= length(model.quadratic_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.sol.equ["eq$(ci.value + offset_quadratic_ge(model))"].level[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintPrimal,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.EqualTo{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if ! (1 <= ci.value <= length(model.quadratic_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.sol.equ["eq$(ci.value + offset_quadratic_eq(model))"].level[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintDual,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if !(1 <= ci.value <= length(model.linear_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return -model.sol.equ["eq$(ci.value + offset_linear_le(model))"].dual[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintDual,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if !(1 <= ci.value <= length(model.linear_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.sol.equ["eq$(ci.value + offset_linear_ge(model))"].dual[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintDual,
   ci::MOI.ConstraintIndex{MOI.ScalarAffineFunction{Float64}, MOI.EqualTo{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if !(1 <= ci.value <= length(model.linear_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return model.sol.equ["eq$(ci.value + offset_linear_eq(model))"].dual[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintDual,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.LessThan{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if !(1 <= ci.value <= length(model.quadratic_le_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return -model.sol.equ["eq$(ci.value + offset_quadratic_le(model))"].dual[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintDual,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.GreaterThan{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if !(1 <= ci.value <= length(model.quadratic_ge_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return -model.sol.equ["eq$(ci.value + offset_quadratic_ge(model))"].dual[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.ConstraintDual,
   ci::MOI.ConstraintIndex{MOI.ScalarQuadraticFunction{Float64}, MOI.EqualTo{Float64}}
)
   MOI.check_result_index_bounds(model, attr)
   if !(1 <= ci.value <= length(model.quadratic_eq_constraints))
      error("Invalid constraint index ", ci.value)
   end
   return -model.sol.equ["eq$(ci.value + offset_quadratic_eq(model))"].dual[1]
end

function MOI.get(
   model::Optimizer,
   attr::MOI.NLPBlockDual
)
   MOI.check_result_index_bounds(model, attr)
   values = zeros(model.m_nonlin)
   for i = offset_nonlin(model)+1:offset_nonlin(model)+model.m_nonlin
      values[i] = -model.sol.equ["eq$i"].dual[1]
   end
   return values
end


"""
    rotAngMom(eph,jd0,time,target,unit)

Compute angular momentum due to rotation and first derivative of target at
epoch jd0+time.
To get the best precision for the interpolation, the time is split in two
floating-point numbers. The argument jd0 should be an integer and time should
be a fraction of the day. But you may call this function with time=0 and jd0,
the desired time, if you don't take care about precision.

# Arguments
- `eph`: ephemeris
- `jd0::Float64`: jd0+time must be equal to the Julian Day for the time coordinate corresponding to the ephemeris (usually TDB or TCB)
- `time::Float64`: jd0+time must be equal to the Julian Day for the time coordinate corresponding to the ephemeris (usually TDB or TCB)
- `target::Integer`: The body whose angular momentum is required. The numbering system depends on the parameter unit.
- `unit::Integer` : The units of the result. This integer is a sum of some unit constants (unit*) and/or the constant useNaifId. If the unit contains useNaifId, the NAIF identification numbering system is used for the target and the center. If the unit does not contain useNaifId, the old number system is used for the target and the center (see the list in the documentation of function compute). The angles are expressed in radians if unit contains unitRad.

"""
function rotAngMom(eph::Ephem,jd0::Float64,time::Float64,
   target::Integer,unit::Integer)
    @_checkPointer eph.data "Ephemeris is not properly initialized!"
    result = Array{Float64,1}(undef,6)
    stat = unsafe_rotAngMom!(result,eph,jd0,time,target,unit)
    @_checkStatus stat "Unable to compute ephemeris"
    return result
end

"""
    unsafe_rotAngMom!(result,eph,jd0,time,target,unit)

In place version of the rotAngMom function. Does not perform any checks!
Compute angular momentum due to rotation and first derivative of target at
epoch jd0+time.
To get the best precision for the interpolation, the time is split in two
floating-point numbers. The argument jd0 should be an integer and time should
be a fraction of the day. But you may call this function with time=0 and jd0,
the desired time, if you don't take care about precision.

# Arguments
- `result`: container for result. It is not checked if it is sufficiently large enough!
- `eph`: ephemeris
- `jd0::Float64`: jd0+time must be equal to the Julian Day for the time coordinate corresponding to the ephemeris (usually TDB or TCB)
- `time::Float64`: jd0+time must be equal to the Julian Day for the time coordinate corresponding to the ephemeris (usually TDB or TCB)
- `target::Integer`: The body whose angular momentum is required. The numbering system depends on the parameter unit.
- `unit::Integer` : The units of the result. This integer is a sum of some unit constants (unit*) and/or the constant useNaifId. If the unit contains useNaifId, the NAIF identification numbering system is used for the target and the center. If the unit does not contain useNaifId, the old number system is used for the target and the center (see the list in the documentation of function compute). The angles are expressed in radians if unit contains unitRad.

# Return:
- status integer from CALCEPH: 0 if an error occured

"""
function unsafe_rotAngMom!(result,eph::Ephem,jd0::Float64,time::Float64,
   target::Integer,unit::Integer)
    stat = ccall((:calceph_rotangmom_unit, libcalceph), Cint,
    (Ptr{Cvoid},Cdouble,Cdouble,Cint,Cint,Ref{Cdouble}),
    eph.data,jd0,time,target,unit,result)
    return stat
end

"""
    rotAngMom(eph,jd0,time,target,unit,order)

Compute angular momentum due to rotation and derivatives up to order of target
at epoch jd0+time.
To get the best precision for the interpolation, the time is split in two
floating-point numbers. The argument jd0 should be an integer and time should
be a fraction of the day. But you may call this function with time=0 and jd0,
the desired time, if you don't take care about precision.

# Arguments
- `eph`: ephemeris
- `jd0::Float64`: jd0+time must be equal to the Julian Day for the time coordinate corresponding to the ephemeris (usually TDB or TCB)
- `time::Float64`: jd0+time must be equal to the Julian Day for the time coordinate corresponding to the ephemeris (usually TDB or TCB)
- `target::Integer`: The body whose angular momentum is required. The numbering system depends on the parameter unit.
- `unit::Integer` : The units of the result. This integer is a sum of some unit constants (unit*) and/or the constant useNaifId. If the unit contains useNaifId, the NAIF identification numbering system is used for the target and the center. If the unit does not contain useNaifId, the old number system is used for the target and the center (see the list in the documentation of function compute).
- `order::Integer` : The order of derivatives
    * 0: only the angles are computed.
    * 1: only the angles and 1st derivatives are computed.
    * 2: only the angles, the 1st derivatives and 2nd derivatives are computed.
    * 3: the angles, the 1st derivatives, 2nd derivatives and 3rd derivatives are computed.

"""
function rotAngMom(eph::Ephem,jd0::Float64,time::Float64,
   target::Integer,unit::Integer,order::Integer)
    @_checkPointer eph.data "Ephemeris is not properly initialized!"
    @_checkOrder order
    result = Array{Float64,1}(undef,3+3order)
    stat = unsafe_rotAngMom!(result,eph,jd0,time,target,unit,order)
    @_checkStatus stat "Unable to compute ephemeris"
    return result
end

"""
    unsafe_rotAngMom!(result,eph,jd0,time,target,unit,order)

In place version of the rotAngMom function. Does not perform any checks!
Compute angular momentum due to rotation and derivatives up to order of target
at epoch jd0+time.
To get the best precision for the interpolation, the time is split in two
floating-point numbers. The argument jd0 should be an integer and time should
be a fraction of the day. But you may call this function with time=0 and jd0,
the desired time, if you don't take care about precision.

# Arguments
- `result`: container for result. It is not checked if it is sufficiently large enough!
- `eph`: ephemeris
- `jd0::Float64`: jd0+time must be equal to the Julian Day for the time coordinate corresponding to the ephemeris (usually TDB or TCB)
- `time::Float64`: jd0+time must be equal to the Julian Day for the time coordinate corresponding to the ephemeris (usually TDB or TCB)
- `target::Integer`: The body whose angular momentum is required. The numbering system depends on the parameter unit.
- `unit::Integer` : The units of the result. This integer is a sum of some unit constants (unit*) and/or the constant useNaifId. If the unit contains useNaifId, the NAIF identification numbering system is used for the target and the center. If the unit does not contain useNaifId, the old number system is used for the target and the center (see the list in the documentation of function compute).
- `order::Integer` : The order of derivatives
    * 0: only the angles are computed.
    * 1: only the angles and 1st derivatives are computed.
    * 2: only the angles, the 1st derivatives and 2nd derivatives are computed.
    * 3: the angles, the 1st derivatives, 2nd derivatives and 3rd derivatives are computed.

    # Return:
    - status integer from CALCEPH: 0 if an error occured

"""
function unsafe_rotAngMom!(result,eph::Ephem,jd0::Float64,time::Float64,
   target::Integer,unit::Integer,order::Integer)
    stat = ccall((:calceph_rotangmom_order, libcalceph), Cint,
    (Ptr{Cvoid},Cdouble,Cdouble,Cint,Cint,Cint,Ref{Cdouble}),
    eph.data,jd0,time,target,unit,order,result)
    return stat
end


config[@__DIR__] = (
    directory=@__DIR__,
    case="GaAs",
    n_repeat=40,
    kwargs_scf=(randomise_guess=true, random_ampl=1e-3, ),
    firstmixkey="RestaGaAs",
    # firstα=0.1,
    #
    construct_localiser=(lattice, atoms) -> identity,
    mixings=Dict(
        "Kerker"     => (α, localiser) -> KerkerMixing(α=α, kTF=1.0),
        "None"       => (α, localiser) -> SimpleMixing(α=α),
        "RestaGaAs"  => (α, localiser) -> RestaMixing( α=α, kTF=1.0, εr=14),
        "RestaSiO2"  => (α, localiser) -> RestaMixing( α=α, kTF=1.0, εr=2.4),
    ),
)


using ClimateMachine
using ClimateMachine.ConfigTypes
using ClimateMachine.Mesh.Topologies:
    StackedCubedSphereTopology, equiangular_cubed_sphere_warp, grid1d
using ClimateMachine.Mesh.Grids:
    DiscontinuousSpectralElementGrid,
    VerticalDirection,
    HorizontalDirection,
    EveryDirection,
    min_node_distance
using ClimateMachine.Mesh.Filters
using ClimateMachine.DGMethods: DGModel, init_ode_state, remainder_DGModel
using ClimateMachine.DGMethods.NumericalFluxes:
    RusanovNumericalFlux,
    CentralNumericalFluxGradient,
    CentralNumericalFluxSecondOrder
using ClimateMachine.ODESolvers
using ClimateMachine.SystemSolvers
using ClimateMachine.VTK: writevtk, writepvtu
using ClimateMachine.GenericCallbacks:
    EveryXWallTimeSeconds, EveryXSimulationSteps
using ClimateMachine.Thermodynamics:
    air_density, soundspeed_air, internal_energy, PhaseDry_pT, PhasePartition
using ClimateMachine.TemperatureProfiles: IsothermalProfile
using ClimateMachine.Atmos:
    AtmosPhysics,
    AtmosModel,
    DryModel,
    NoPrecipitation,
    NoRadiation,
    NTracers,
    ConstantDynamicViscosity,
    vars_state,
    Gravity,
    HydrostaticState,
    AtmosAcousticGravityLinearModel,
    AtmosAcousticLinearModel,
    parameter_set
using ClimateMachine.Orientations:
    SphericalOrientation, gravitational_potential, altitude, latitude, longitude
using ClimateMachine.VariableTemplates: flattenednames

using CLIMAParameters
using CLIMAParameters.Planet: planet_radius
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

using MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test

const output_vtk = false

const ntracers = 1

function main()
    ClimateMachine.init()
    ArrayType = ClimateMachine.array_type()

    mpicomm = MPI.COMM_WORLD

    polynomialorder = 5
    numelem_horz = 10
    numelem_vert = 5

    timeend = 60 * 60
    # timeend = 33 * 60 * 60 # Full simulation
    outputtime = 60 * 60

    expected_result = Dict()
    expected_result[Float64, true] = 9.5073337869322578e+13
    expected_result[Float64, false] = 9.5073455070673781e+13

    @testset "acoustic wave" begin
        for FT in (Float64,)# Float32)
            for explicit in (true, false)
                result = test_run(
                    mpicomm,
                    polynomialorder,
                    numelem_horz,
                    numelem_vert,
                    timeend,
                    outputtime,
                    ArrayType,
                    FT,
                    explicit,
                )
                @test result ≈ expected_result[FT, explicit]
            end
        end
    end
end

function test_run(
    mpicomm,
    polynomialorder,
    numelem_horz,
    numelem_vert,
    timeend,
    outputtime,
    ArrayType,
    FT,
    explicit_solve,
)

    setup = AcousticWaveSetup{FT}()

    _planet_radius::FT = planet_radius(param_set)
    vert_range = grid1d(
        _planet_radius,
        FT(_planet_radius + setup.domain_height),
        nelem = numelem_vert,
    )
    topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)

    grid = DiscontinuousSpectralElementGrid(
        topology,
        FloatType = FT,
        DeviceArray = ArrayType,
        polynomialorder = polynomialorder,
        meshwarp = equiangular_cubed_sphere_warp,
    )
    hmnd = min_node_distance(grid, HorizontalDirection())
    vmnd = min_node_distance(grid, VerticalDirection())

    T_profile = IsothermalProfile(param_set, setup.T_ref)
    δ_χ = @SVector [FT(ii) for ii in 1:ntracers]

    fullphysics = AtmosPhysics{FT}(
        param_set;
        ref_state = HydrostaticState(T_profile),
        turbulence = ConstantDynamicViscosity(FT(0)),
        moisture = DryModel(),
        tracers = NTracers{length(δ_χ), FT}(δ_χ),
    )
    fullmodel = AtmosModel{FT}(
        AtmosLESConfigType,
        fullphysics;
        orientation = SphericalOrientation(),
        init_state_prognostic = setup,
        source = (Gravity(),),
    )
    dg = DGModel(
        fullmodel,
        grid,
        RusanovNumericalFlux(),
        CentralNumericalFluxSecondOrder(),
        CentralNumericalFluxGradient(),
    )
    Q = init_ode_state(dg, FT(0))

    # The linear model which contains the fast modes
    # acousticmodel = AtmosAcousticLinearModel(fullmodel)
    acousticmodel = AtmosAcousticGravityLinearModel(fullmodel)

    vacoustic_dg = DGModel(
        acousticmodel,
        grid,
        RusanovNumericalFlux(),
        CentralNumericalFluxSecondOrder(),
        CentralNumericalFluxGradient();
        direction = VerticalDirection(),
        state_auxiliary = dg.state_auxiliary,
    )

    # Advection model is the difference between the fullmodel and acousticmodel.
    # This will be handled with explicit substepping (time step in between the
    # vertical and horizontally acoustic models)
    rem_dg = remainder_DGModel(dg, (vacoustic_dg,))

    # determine the time step for the model components
    acoustic_speed = soundspeed_air(param_set, FT(setup.T_ref))
    advection_speed = 1 # What's a reasonable number here?

    vacoustic_dt = vmnd / acoustic_speed
    remainder_dt = min(min(hmnd, vmnd) / advection_speed, hmnd / acoustic_speed)
    odesolver = if explicit_solve
        remainder_dt = 5vacoustic_dt

        nsteps_output = ceil(outputtime / remainder_dt)
        remainder_dt = outputtime / nsteps_output
        nsteps = ceil(Int, timeend / remainder_dt)
        @assert nsteps * remainder_dt ≈ timeend

        vacoustic_solver =
            LSRK54CarpenterKennedy(vacoustic_dg, Q; dt = vacoustic_dt)
        rem_solver = MRIGARKERK45aSandu(
            rem_dg,
            vacoustic_solver,
            Q;
            dt = remainder_dt,
        )

        rem_solver
    else
        vacoustic_dt = 200vacoustic_dt
        element_size = (setup.domain_height / numelem_vert)

        nsteps_output = ceil(outputtime / vacoustic_dt)
        vacoustic_dt = outputtime / nsteps_output
        nsteps = ceil(Int, timeend / vacoustic_dt)
        @assert nsteps * vacoustic_dt ≈ timeend

        rem_solver = LSRK54CarpenterKennedy(rem_dg, Q; dt = remainder_dt)
        vacoustic_solver = MRIGARKESDIRK24LSA(
            vacoustic_dg,
            LinearBackwardEulerSolver(ManyColumnLU(); isadjustable = false),
            rem_solver,
            Q;
            dt = vacoustic_dt,
        )

        vacoustic_solver
    end

    filterorder = 18
    filter = ExponentialFilter(grid, 0, filterorder)
    cbfilter = EveryXSimulationSteps(1) do
        Filters.apply!(Q, :, grid, filter, direction = VerticalDirection())
        nothing
    end

    eng0 = norm(Q)
    @info @sprintf(
        """Starting
           ArrayType       = %s
           FT              = %s
           polynomialorder = %d
           numelem_horz    = %d
           numelem_vert    = %d
           acoustic dt     = %.16e
           remainder_dt    = %.16e
           norm(Q₀)        = %.16e
           """,
        "$ArrayType",
        "$FT",
        polynomialorder,
        numelem_horz,
        numelem_vert,
        vacoustic_dt,
        remainder_dt,
        eng0
    )
    # Set up the information callback
    starttime = Ref(now())
    cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)
        if s
            starttime[] = now()
        else
            energy = norm(Q)
            runtime = Dates.format(
                convert(DateTime, now() - starttime[]),
                dateformat"HH:MM:SS",
            )
            @info @sprintf """Update
                              simtime = %.16e
                              runtime = %s
                              norm(Q) = %.16e
                              """ gettime(odesolver) runtime energy
        end
    end
    callbacks = (cbinfo, cbfilter)

    if output_vtk
        # create vtk dir
        vtkdir =
            "vtk_acousticwave" *
            "_poly$(polynomialorder)_horz$(numelem_horz)_vert$(numelem_vert)" *
            "_$(ArrayType)_$(FT)" *
            "_$(explicit_solve ? "Explicit_Explicit" : "Implicit_Explicit")"
        mkpath(vtkdir)

        vtkstep = 0
        # output initial step
        do_output(mpicomm, vtkdir, vtkstep, dg, Q, fullmodel)

        # setup the output callback
        cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do
            vtkstep += 1
            Qe = init_ode_state(dg, gettime(odesolver))
            do_output(mpicomm, vtkdir, vtkstep, dg, Q, fullmodel)
        end
        callbacks = (callbacks..., cbvtk)
    end

    solve!(
        Q,
        odesolver;
        numberofsteps = nsteps,
        adjustfinalstep = false,
        callbacks = callbacks,
    )

    # final statistics
    engf = norm(Q)
    @info @sprintf """Finished
    norm(Q)                 = %.16e
    norm(Q) / norm(Q₀)      = %.16e
    norm(Q) - norm(Q₀)      = %.16e
    """ engf engf / eng0 engf - eng0
    engf
end

Base.@kwdef struct AcousticWaveSetup{FT}
    domain_height::FT = 10e3
    T_ref::FT = 300
    α::FT = 3
    γ::FT = 100
    nv::Int = 1
end

function (setup::AcousticWaveSetup)(problem, bl, state, aux, localgeo, t)
    # callable to set initial conditions
    FT = eltype(state)
    param_set = parameter_set(bl)

    λ = longitude(bl, aux)
    φ = latitude(bl, aux)
    z = altitude(bl, aux)

    β = min(FT(1), setup.α * acos(cos(φ) * cos(λ)))
    f = (1 + cos(FT(π) * β)) / 2
    g = sin(setup.nv * FT(π) * z / setup.domain_height)
    Δp = setup.γ * f * g
    p = aux.ref_state.p + Δp

    ts = PhaseDry_pT(param_set, p, setup.T_ref)
    q_pt = PhasePartition(ts)
    e_pot = gravitational_potential(bl.orientation, aux)
    e_int = internal_energy(ts)

    state.ρ = air_density(ts)
    state.ρu = SVector{3, FT}(0, 0, 0)
    state.energy.ρe = state.ρ * (e_int + e_pot)

    state.tracers.ρχ = @SVector [FT(ii) for ii in 1:ntracers]
    nothing
end

function do_output(
    mpicomm,
    vtkdir,
    vtkstep,
    dg,
    Q,
    model,
    testname = "acousticwave",
)
    ## name of the file that this MPI rank will write
    filename = @sprintf(
        "%s/%s_mpirank%04d_step%04d",
        vtkdir,
        testname,
        MPI.Comm_rank(mpicomm),
        vtkstep
    )

    statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))
    auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))
    writevtk(filename, Q, dg, statenames, dg.state_auxiliary, auxnames)

    ## Generate the pvtu file for these vtk files
    if MPI.Comm_rank(mpicomm) == 0
        ## name of the pvtu file
        pvtuprefix = @sprintf("%s/%s_step%04d", vtkdir, testname, vtkstep)

        ## name of each of the ranks vtk files
        prefixes = ntuple(MPI.Comm_size(mpicomm)) do i
            @sprintf("%s_mpirank%04d_step%04d", testname, i - 1, vtkstep)
        end

        writepvtu(pvtuprefix, prefixes, (statenames..., auxnames...), eltype(Q))

        @info "Done writing VTK: $pvtuprefix"
    end
end

main()


"""
Runs the semi-distributed bucket model for a given area with one precipitation and temperature input.

$(SIGNATURES)

The function returns the modeled discharge for each timestep, the modelled snow extend for each elevation and timestep and the Waterbalance.
"""
function run_model(Area, Evaporation_Mean::Array{Float64,1}, Precipitation::Array{Float64,2}, Temp::Array{Float64,2},
                bare_input::HRU_Input, forest_input::HRU_Input, grass_input::HRU_Input, rip_input::HRU_Input,
                bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64,
                bare_parameters::Parameters, forest_parameters::Parameters, grass_parameters::Parameters, rip_parameters::Parameters, slow_parameters::Slow_Paramters, Total_Elevationbands, Elevation_Percentage)
    # the function takes as input the parameters of each HRU, the inital storage values of each HRU, the inital value of the slow storage
    # KS, ratio riparian, all inputs

    # define the maximum time
    tmax::Int128 = length(Precipitation[:,1])

    Discharge::Array{Float64,1} = zeros(tmax)
    WBtotal::Array{Float64,1} = zeros(tmax)
    Snow_Extend::Array{Float64,2} = zeros(tmax, Total_Elevationbands)

    # store the initial storage values, Does not take into account GW storage!!
    # Assumption_ GW storage is 0 at start
    #TO DO: average values over area
    # as form now only works if storage input is 0 at start
    Initial_Storage_bare::Float64 = bare_storage.Fast + sum(bare_storage.Interception) + sum(bare_storage.Snow) + bare_storage.Soil
    Initial_Storage_forest::Float64 = forest_storage.Fast + sum(forest_storage.Interception) + sum(forest_storage.Snow) + forest_storage.Soil
    Initial_Storage_grass::Float64 = grass_storage.Fast + sum(grass_storage.Interception) + sum(grass_storage.Snow) + grass_storage.Soil
    Initial_Storage_rip::Float64 = rip_storage.Fast + sum(rip_storage.Interception) + sum(rip_storage.Snow) + rip_storage.Soil
    Initial_Storage::Float64 = Initial_Storage_bare + Initial_Storage_forest + Initial_Storage_grass + Initial_Storage_rip + Slowstorage

    for t in 1:tmax
        # at each timestep new temp, precipitation and Epot values have to be delivered
        # areas don't change
        # riparian discharge from former timestep has to be used
        # gives the current precipitation, evaporation and temperature
        Evaporation_Mean_Current = Evaporation_Mean[t]
        Precipitation_Current = Precipitation[t, :]
        Temperature_Current = Temp[t, :]

        bare_input::HRU_Input = input_timestep(bare_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        forest_input::HRU_Input = input_timestep(forest_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        grass_input::HRU_Input = input_timestep(grass_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        rip_input::HRU_Input = input_timestep(rip_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)

        Riparian_Discharge::Float64, Total_Discharge::Float64, Total_Interception_Evaporation::Float64, Total_Soil_Evaporation::Float64, bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64, WB::Float64, Total_Prec::Float64 = allHRU(bare_input, forest_input, grass_input, rip_input,
                                                                                bare_storage, forest_storage, grass_storage, rip_storage,
                                                                                bare_parameters, forest_parameters, grass_parameters, rip_parameters,
                                                                                Slowstorage, slow_parameters)
        # give new riparian discharge as input for next timestep
        rip_input.Riparian_Discharge = Riparian_Discharge
        # save the fluxes of the current timestep
        Discharge[t]::Float64 = Total_Discharge/1000 * Area / (3600 * 24) # mm convert it to meter and than * area / seconds in one day
        # OPTIONAL: store the total amount of snow in each elevation (considering whole catchment)
        Snow_Elevations, Snow_Extend[t,:]::Array{Float64,1} = snowperelevation(bare_input, forest_input, grass_input, rip_input, bare_storage, forest_storage, grass_storage, rip_storage, Total_Elevationbands)
        #Bare_Snow[t,:]::Array{Float64,1} = bare_storage.Snow
        Snow_Extend[t,:] = Snow_Extend[t,:] ./ Elevation_Percentage

        WBtotal[t]::Float64 = WB
    end

    # Check Water Balance
    Waterbalance = sum(WBtotal)::Float64
    @assert Waterbalance <= 10^(-10)
    return Discharge::Array{Float64,1}, Snow_Extend::Array{Float64,2}, Waterbalance::Float64
end

"""
Runs the semi-distributed bucket model for a given area with one precipitation and temperature input.

$(SIGNATURES)

The function returns the modeled discharge for each timestep, the modelled snow extend for each elevation and timestep and the Waterbalance.
"""
function run_model_glacier(Area, Evaporation_Mean::Array{Float64,1},Glacier::Array{Float64,2}, Precipitation::Array{Float64,2}, Temp::Array{Float64,2},
                bare_input::HRU_Input, forest_input::HRU_Input, grass_input::HRU_Input, rip_input::HRU_Input,
                bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64,
                bare_parameters::Parameters, forest_parameters::Parameters, grass_parameters::Parameters, rip_parameters::Parameters, slow_parameters::Slow_Paramters, Total_Elevationbands, Elevation_Percentage)
    # the function takes as input the parameters of each HRU, the inital storage values of each HRU, the inital value of the slow storage
    # KS, ratio riparian, all inputs

    # define the maximum time
    tmax::Int128 = length(Precipitation[:,1])

    Discharge::Array{Float64,1} = zeros(tmax)
    WBtotal::Array{Float64,1} = zeros(tmax)
    Snow_Extend::Array{Float64,2} = zeros(tmax, Total_Elevationbands)

    # store the initial storage values, Does not take into account GW storage!!
    # Assumption_ GW storage is 0 at start
    #TO DO: average values over area
    # as form now only works if storage input is 0 at start
    Initial_Storage_bare::Float64 = bare_storage.Fast + sum(bare_storage.Interception) + sum(bare_storage.Snow) + bare_storage.Soil
    Initial_Storage_forest::Float64 = forest_storage.Fast + sum(forest_storage.Interception) + sum(forest_storage.Snow) + forest_storage.Soil
    Initial_Storage_grass::Float64 = grass_storage.Fast + sum(grass_storage.Interception) + sum(grass_storage.Snow) + grass_storage.Soil
    Initial_Storage_rip::Float64 = rip_storage.Fast + sum(rip_storage.Interception) + sum(rip_storage.Snow) + rip_storage.Soil
    Initial_Storage::Float64 = Initial_Storage_bare + Initial_Storage_forest + Initial_Storage_grass + Initial_Storage_rip + Slowstorage

    for t in 1:tmax
        # at each timestep new temp, precipitation and Epot values have to be delivered
        # areas don't change
        # riparian discharge from former timestep has to be used
        # gives the current precipitation, evaporation and temperature
        Evaporation_Mean_Current = Evaporation_Mean[t]
        Precipitation_Current = Precipitation[t, :]
        Temperature_Current = Temp[t, :]
        Glacier_Current = Glacier[:,t]

        bare_input::HRU_Input = input_timestep_glacier(bare_input, Evaporation_Mean_Current, Glacier_Current, Precipitation_Current, Temperature_Current)
        forest_input::HRU_Input = input_timestep(forest_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        grass_input::HRU_Input = input_timestep(grass_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        rip_input::HRU_Input = input_timestep(rip_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)

        Riparian_Discharge::Float64, Total_Discharge::Float64, Total_Interception_Evaporation::Float64, Total_Soil_Evaporation::Float64, bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64, WB::Float64, Total_Prec::Float64 = allHRU(bare_input, forest_input, grass_input, rip_input,
                                                                                bare_storage, forest_storage, grass_storage, rip_storage,
                                                                                bare_parameters, forest_parameters, grass_parameters, rip_parameters,
                                                                                Slowstorage, slow_parameters)
        # give new riparian discharge as input for next timestep
        rip_input.Riparian_Discharge = Riparian_Discharge
        # save the fluxes of the current timestep
        Discharge[t]::Float64 = Total_Discharge/1000 * Area / (3600 * 24) # mm convert it to meter and than * area / seconds in one day
        # OPTIONAL: store the total amount of snow in each elevation (considering whole catchment)
        Snow_Elevations, Snow_Extend[t,:]::Array{Float64,1} = snowperelevation(bare_input, forest_input, grass_input, rip_input, bare_storage, forest_storage, grass_storage, rip_storage, Total_Elevationbands)
        #Bare_Snow[t,:]::Array{Float64,1} = bare_storage.Snow
        Snow_Extend[t,:] = Snow_Extend[t,:] ./ Elevation_Percentage

        WBtotal[t]::Float64 = WB
    end

    # Check Water Balance
    Waterbalance = sum(WBtotal)::Float64
    @assert Waterbalance <= 10^(-10)
    return Discharge::Array{Float64,1}, Snow_Extend::Array{Float64,2}, Waterbalance::Float64
end

"""
Runs the semi-distributed bucket model for a given area with one precipitation and temperature input.

$(SIGNATURES)

The function returns the modeled discharge for each timestep, the modelled snow extend for each elevation and timestep and the Waterbalance and the snowstorage.
"""
function run_model_future(Area, Evaporation_Mean::Array{Float64,1}, Precipitation::Array{Float64,2}, Temp::Array{Float64,2},
                bare_input::HRU_Input, forest_input::HRU_Input, grass_input::HRU_Input, rip_input::HRU_Input,
                bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64,
                bare_parameters::Parameters, forest_parameters::Parameters, grass_parameters::Parameters, rip_parameters::Parameters, slow_parameters::Slow_Paramters, Total_Elevationbands, Elevation_Percentage)
    # the function takes as input the parameters of each HRU, the inital storage values of each HRU, the inital value of the slow storage
    # KS, ratio riparian, all inputs

    # define the maximum time
    tmax::Int128 = length(Precipitation[:,1])

    Discharge::Array{Float64,1} = zeros(tmax)
    Total_Melt::Array{Float64,1} = zeros(tmax)
    WBtotal::Array{Float64,1} = zeros(tmax)
    Snow_Extend::Array{Float64,2} = zeros(tmax, Total_Elevationbands)
    Snowstorage::Array{Float64,1} = zeros(tmax)

    # store the initial storage values, Does not take into account GW storage!!
    # Assumption_ GW storage is 0 at start
    #TO DO: average values over area
    # as form now only works if storage input is 0 at start
    Initial_Storage_bare::Float64 = bare_storage.Fast + sum(bare_storage.Interception) + sum(bare_storage.Snow) + bare_storage.Soil
    Initial_Storage_forest::Float64 = forest_storage.Fast + sum(forest_storage.Interception) + sum(forest_storage.Snow) + forest_storage.Soil
    Initial_Storage_grass::Float64 = grass_storage.Fast + sum(grass_storage.Interception) + sum(grass_storage.Snow) + grass_storage.Soil
    Initial_Storage_rip::Float64 = rip_storage.Fast + sum(rip_storage.Interception) + sum(rip_storage.Snow) + rip_storage.Soil
    Initial_Storage::Float64 = Initial_Storage_bare + Initial_Storage_forest + Initial_Storage_grass + Initial_Storage_rip + Slowstorage

    for t in 1:tmax
        # at each timestep new temp, precipitation and Epot values have to be delivered
        # areas don't change
        # riparian discharge from former timestep has to be used
        # gives the current precipitation, evaporation and temperature
        Evaporation_Mean_Current = Evaporation_Mean[t]
        Precipitation_Current = Precipitation[t, :]
        Temperature_Current = Temp[t, :]

        bare_input::HRU_Input = input_timestep(bare_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        forest_input::HRU_Input = input_timestep(forest_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        grass_input::HRU_Input = input_timestep(grass_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        rip_input::HRU_Input = input_timestep(rip_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)

        Riparian_Discharge::Float64, Total_Discharge::Float64, Total_Interception_Evaporation::Float64, Total_Soil_Evaporation::Float64, bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64, WB::Float64, Total_Prec::Float64, Melt::Float64 = allHRU_future(bare_input, forest_input, grass_input, rip_input,
                                                                                bare_storage, forest_storage, grass_storage, rip_storage,
                                                                                bare_parameters, forest_parameters, grass_parameters, rip_parameters,
                                                                                Slowstorage, slow_parameters)
        # give new riparian discharge as input for next timestep
        rip_input.Riparian_Discharge = Riparian_Discharge
        # save the fluxes of the current timestep
        Discharge[t]::Float64 = Total_Discharge/1000 * Area / (3600 * 24) # mm convert it to meter and than * area / seconds in one day
        # OPTIONAL: store the total amount of snow in each elevation (considering whole catchment)
        Snow_Elevations, Snow_Extend[t,:]::Array{Float64,1} = snowperelevation(bare_input, forest_input, grass_input, rip_input, bare_storage, forest_storage, grass_storage, rip_storage, Total_Elevationbands)
        #Bare_Snow[t,:]::Array{Float64,1} = bare_storage.Snow
        Snow_Extend[t,:] = Snow_Extend[t,:] ./ Elevation_Percentage
        Total_Melt[t] = Melt

        Bare_Interceptionstorage::Float64, Bare_Snowstorage::Float64 = Storage_Total(bare_storage, bare_input)
        Forest_Interceptionstorage::Float64, Forest_Snowstorage::Float64 = Storage_Total(forest_storage, forest_input)
        Grass_Interceptionstorage::Float64, Grass_Snowstorage::Float64 = Storage_Total(grass_storage, grass_input)
        Rip_Interceptionstorage::Float64, Rip_Snowstorage::Float64 = Storage_Total(rip_storage, rip_input)

        Snowstorage[t] = Bare_Snowstorage * bare_input.Area_HRU + Forest_Snowstorage * forest_input.Area_HRU + Grass_Snowstorage * grass_input.Area_HRU + Rip_Snowstorage * rip_input.Area_HRU

        WBtotal[t]::Float64 = WB
    end

    # Check Water Balance
    Waterbalance = sum(WBtotal)::Float64
    @assert Waterbalance <= 10^(-10)
    return Discharge::Array{Float64,1}, Snow_Extend::Array{Float64,2}, Snowstorage::Array{Float64,1}, Waterbalance::Float64, Total_Melt::Array{Float64,1}
end

"""
Runs the semi-distributed bucket model for a given area with one precipitation and temperature input.

$(SIGNATURES)

The function returns the modeled discharge for each timestep, the modelled snow extend for each elevation and timestep and the Waterbalance and the snowstorage.
"""
function run_model_glacier_future(Area, Evaporation_Mean::Array{Float64,1}, Glacier::Array{Float64,2}, Precipitation::Array{Float64,2}, Temp::Array{Float64,2},
                bare_input::HRU_Input, forest_input::HRU_Input, grass_input::HRU_Input, rip_input::HRU_Input,
                bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64,
                bare_parameters::Parameters, forest_parameters::Parameters, grass_parameters::Parameters, rip_parameters::Parameters, slow_parameters::Slow_Paramters, Total_Elevationbands, Elevation_Percentage)
    # the function takes as input the parameters of each HRU, the inital storage values of each HRU, the inital value of the slow storage
    # KS, ratio riparian, all inputs

    # define the maximum time
    tmax::Int128 = length(Precipitation[:,1])

    Discharge::Array{Float64,1} = zeros(tmax)
    Total_Melt::Array{Float64,1} = zeros(tmax)
    WBtotal::Array{Float64,1} = zeros(tmax)
    Snow_Extend::Array{Float64,2} = zeros(tmax, Total_Elevationbands)
    Snowstorage::Array{Float64,1} = zeros(tmax)

    # store the initial storage values, Does not take into account GW storage!!
    # Assumption_ GW storage is 0 at start
    #TO DO: average values over area
    # as form now only works if storage input is 0 at start
    Initial_Storage_bare::Float64 = bare_storage.Fast + sum(bare_storage.Interception) + sum(bare_storage.Snow) + bare_storage.Soil
    Initial_Storage_forest::Float64 = forest_storage.Fast + sum(forest_storage.Interception) + sum(forest_storage.Snow) + forest_storage.Soil
    Initial_Storage_grass::Float64 = grass_storage.Fast + sum(grass_storage.Interception) + sum(grass_storage.Snow) + grass_storage.Soil
    Initial_Storage_rip::Float64 = rip_storage.Fast + sum(rip_storage.Interception) + sum(rip_storage.Snow) + rip_storage.Soil
    Initial_Storage::Float64 = Initial_Storage_bare + Initial_Storage_forest + Initial_Storage_grass + Initial_Storage_rip + Slowstorage

    for t in 1:tmax
        # at each timestep new temp, precipitation and Epot values have to be delivered
        # areas don't change
        # riparian discharge from former timestep has to be used
        # gives the current precipitation, evaporation and temperature
        Evaporation_Mean_Current = Evaporation_Mean[t]
        Precipitation_Current = Precipitation[t, :]
        Temperature_Current = Temp[t, :]
        Glacier_Current = Glacier[:,t]

        bare_input::HRU_Input = input_timestep_glacier(bare_input, Evaporation_Mean_Current, Glacier_Current, Precipitation_Current, Temperature_Current)
        forest_input::HRU_Input = input_timestep(forest_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        grass_input::HRU_Input = input_timestep(grass_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        rip_input::HRU_Input = input_timestep(rip_input, Evaporation_Mean_Current, Precipitation_Current, Temperature_Current)
        Riparian_Discharge::Float64, Total_Discharge::Float64, Total_Interception_Evaporation::Float64, Total_Soil_Evaporation::Float64, bare_storage::Storages, forest_storage::Storages, grass_storage::Storages, rip_storage::Storages, Slowstorage::Float64, WB::Float64, Total_Prec::Float64, Melt::Float64 = allHRU_future(bare_input, forest_input, grass_input, rip_input,
                                                                                bare_storage, forest_storage, grass_storage, rip_storage,
                                                                                bare_parameters, forest_parameters, grass_parameters, rip_parameters,
                                                                                Slowstorage, slow_parameters)
        # give new riparian discharge as input for next timestep
        rip_input.Riparian_Discharge = Riparian_Discharge
        # save the fluxes of the current timestep
        Discharge[t]::Float64 = Total_Discharge/1000 * Area / (3600 * 24) # mm convert it to meter and than * area / seconds in one day
        # OPTIONAL: store the total amount of snow in each elevation (considering whole catchment)
        Snow_Elevations, Snow_Extend[t,:]::Array{Float64,1} = snowperelevation(bare_input, forest_input, grass_input, rip_input, bare_storage, forest_storage, grass_storage, rip_storage, Total_Elevationbands)
        #Bare_Snow[t,:]::Array{Float64,1} = bare_storage.Snow
        Snow_Extend[t,:] = Snow_Extend[t,:] ./ Elevation_Percentage
        Total_Melt[t] = Melt

        Bare_Interceptionstorage::Float64, Bare_Snowstorage::Float64 = Storage_Total(bare_storage, bare_input)
        Forest_Interceptionstorage::Float64, Forest_Snowstorage::Float64 = Storage_Total(forest_storage, forest_input)
        Grass_Interceptionstorage::Float64, Grass_Snowstorage::Float64 = Storage_Total(grass_storage, grass_input)
        Rip_Interceptionstorage::Float64, Rip_Snowstorage::Float64 = Storage_Total(rip_storage, rip_input)

        Snowstorage[t] = Bare_Snowstorage * bare_input.Area_HRU + Forest_Snowstorage * forest_input.Area_HRU + Grass_Snowstorage * grass_input.Area_HRU + Rip_Snowstorage * rip_input.Area_HRU

        WBtotal[t]::Float64 = WB
    end

    # Check Water Balance
    Waterbalance = sum(WBtotal)::Float64
    #@assert Waterbalance <= 10^(-10)
    return Discharge::Array{Float64,1}, Snow_Extend::Array{Float64,2}, Snowstorage::Array{Float64,1}, Waterbalance::Float64, Total_Melt::Array{Float64,1}
end


using .NumericalFluxes:
    CentralNumericalFluxHigherOrder, CentralNumericalFluxDivergence

struct DGModel{BL, G, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD}
    balancelaw::BL
    grid::G
    numerical_flux_first_order::NFND
    numerical_flux_second_order::NFD
    numerical_flux_gradient::GNF
    state_auxiliary::AS
    state_gradient_flux::DS
    states_higher_order::HDS
    direction::D
    diffusion_direction::DD
    modeldata::MD
end
function DGModel(
    balancelaw,
    grid,
    numerical_flux_first_order,
    numerical_flux_second_order,
    numerical_flux_gradient;
    state_auxiliary = create_auxiliary_state(balancelaw, grid),
    state_gradient_flux = create_gradient_state(balancelaw, grid),
    states_higher_order = create_higher_order_states(balancelaw, grid),
    direction = EveryDirection(),
    diffusion_direction = direction,
    modeldata = nothing,
)
    DGModel(
        balancelaw,
        grid,
        numerical_flux_first_order,
        numerical_flux_second_order,
        numerical_flux_gradient,
        state_auxiliary,
        state_gradient_flux,
        states_higher_order,
        direction,
        diffusion_direction,
        modeldata,
    )
end

function (dg::DGModel)(
    tendency,
    state_conservative,
    ::Nothing,
    t;
    increment = false,
)

    balance_law = dg.balancelaw
    device = typeof(state_conservative.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq
    Nfp = Nq * Nqk
    nrealelem = length(topology.realelems)

    state_gradient_flux = dg.state_gradient_flux
    Qhypervisc_grad, Qhypervisc_div = dg.states_higher_order
    state_auxiliary = dg.state_auxiliary

    FT = eltype(state_conservative)
    num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
    nhyperviscstate = num_hyperdiffusive(balance_law, FT)

    Np = dofs_per_element(grid)

    workgroups_volume = (Nq, Nq, Nqk)
    ndrange_volume = (nrealelem * Nq, Nq, Nqk)
    workgroups_surface = Nfp
    ndrange_interior_surface = Nfp * length(grid.interiorelems)
    ndrange_exterior_surface = Nfp * length(grid.exteriorelems)

    communicate =
        !(isstacked(topology) && typeof(dg.direction) <: VerticalDirection)

    update_auxiliary_state!(
        dg,
        balance_law,
        state_conservative,
        t,
        dg.grid.topology.realelems,
    )

    if nhyperviscstate > 0
        hypervisc_indexmap = create_hypervisc_indexmap(balance_law)
    else
        hypervisc_indexmap = nothing
    end

    exchange_state_conservative = NoneEvent()
    exchange_state_gradient_flux = NoneEvent()
    exchange_Qhypervisc_grad = NoneEvent()
    exchange_Qhypervisc_div = NoneEvent()

    comp_stream = Event(device)

    ########################
    # Gradient Computation #
    ########################
    if communicate
        exchange_state_conservative = MPIStateArrays.begin_ghost_exchange!(
            state_conservative;
            dependencies = comp_stream,
        )
    end

    if num_state_gradient_flux > 0 || nhyperviscstate > 0

        comp_stream = volume_gradients!(device, (Nq, Nq))(
            balance_law,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            state_conservative.data,
            state_gradient_flux.data,
            Qhypervisc_grad.data,
            state_auxiliary.data,
            grid.vgeo,
            t,
            grid.D,
            hypervisc_indexmap,
            topology.realelems,
            ndrange = (Nq * nrealelem, Nq),
            dependencies = (comp_stream,),
        )

        comp_stream = interface_gradients!(device, workgroups_surface)(
            balance_law,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            dg.numerical_flux_gradient,
            state_conservative.data,
            state_gradient_flux.data,
            Qhypervisc_grad.data,
            state_auxiliary.data,
            grid.vgeo,
            grid.sgeo,
            t,
            grid.vmap⁻,
            grid.vmap⁺,
            grid.elemtobndy,
            hypervisc_indexmap,
            grid.interiorelems;
            ndrange = ndrange_interior_surface,
            dependencies = (comp_stream,),
        )

        if communicate
            exchange_state_conservative = MPIStateArrays.end_ghost_exchange!(
                state_conservative;
                dependencies = exchange_state_conservative,
            )

            # update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            wait(device, exchange_state_conservative)
            update_auxiliary_state!(
                dg,
                balance_law,
                state_conservative,
                t,
                dg.grid.topology.ghostelems,
            )
            exchange_state_conservative = Event(device)
        end

        comp_stream = interface_gradients!(device, workgroups_surface)(
            balance_law,
            Val(dim),
            Val(N),
            dg.diffusion_direction,
            dg.numerical_flux_gradient,
            state_conservative.data,
            state_gradient_flux.data,
            Qhypervisc_grad.data,
            state_auxiliary.data,
            grid.vgeo,
            grid.sgeo,
            t,
            grid.vmap⁻,
            grid.vmap⁺,
            grid.elemtobndy,
            hypervisc_indexmap,
            grid.exteriorelems;
            ndrange = ndrange_exterior_surface,
            dependencies = (comp_stream, exchange_state_conservative),
        )

        if communicate
            if num_state_gradient_flux > 0
                exchange_state_gradient_flux =
                    MPIStateArrays.begin_ghost_exchange!(
                        state_gradient_flux,
                        dependencies = comp_stream,
                    )
            end
            if nhyperviscstate > 0
                exchange_Qhypervisc_grad = MPIStateArrays.begin_ghost_exchange!(
                    Qhypervisc_grad,
                    dependencies = comp_stream,
                )
            end
        end

        if num_state_gradient_flux > 0
            # update_aux_diffusive may start asynchronous work on the compute device
            # and we synchronize those here through a device event.
            wait(device, comp_stream)
            update_auxiliary_state_gradient!(
                dg,
                balance_law,
                state_conservative,
                t,
                dg.grid.topology.realelems,
            )
            comp_stream = Event(device)
        end
    end

    if nhyperviscstate > 0
        #########################
        # Laplacian Computation #
        #########################

        comp_stream =
            volume_divergence_of_gradients!(device, workgroups_volume)(
                balance_law,
                Val(dim),
                Val(N),
                dg.diffusion_direction,
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                grid.vgeo,
                grid.D,
                topology.realelems;
                ndrange = ndrange_volume,
                dependencies = (comp_stream,),
            )

        comp_stream =
            interface_divergence_of_gradients!(device, workgroups_surface)(
                balance_law,
                Val(dim),
                Val(N),
                dg.diffusion_direction,
                CentralNumericalFluxDivergence(),
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                grid.vgeo,
                grid.sgeo,
                grid.vmap⁻,
                grid.vmap⁺,
                grid.elemtobndy,
                grid.interiorelems;
                ndrange = ndrange_interior_surface,
                dependencies = (comp_stream,),
            )

        if communicate
            exchange_Qhypervisc_grad = MPIStateArrays.end_ghost_exchange!(
                Qhypervisc_grad,
                dependencies = exchange_Qhypervisc_grad,
            )
        end

        comp_stream =
            interface_divergence_of_gradients!(device, workgroups_surface)(
                balance_law,
                Val(dim),
                Val(N),
                dg.diffusion_direction,
                CentralNumericalFluxDivergence(),
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                grid.vgeo,
                grid.sgeo,
                grid.vmap⁻,
                grid.vmap⁺,
                grid.elemtobndy,
                grid.exteriorelems;
                ndrange = ndrange_exterior_surface,
                dependencies = (comp_stream, exchange_Qhypervisc_grad),
            )

        if communicate
            exchange_Qhypervisc_div = MPIStateArrays.begin_ghost_exchange!(
                Qhypervisc_div,
                dependencies = comp_stream,
            )
        end

        ####################################
        # Hyperdiffusive terms computation #
        ####################################

        comp_stream =
            volume_gradients_of_laplacians!(device, workgroups_volume)(
                balance_law,
                Val(dim),
                Val(N),
                dg.diffusion_direction,
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                state_conservative.data,
                state_auxiliary.data,
                grid.vgeo,
                grid.ω,
                grid.D,
                topology.realelems,
                t;
                ndrange = ndrange_volume,
                dependencies = (comp_stream,),
            )

        comp_stream =
            interface_gradients_of_laplacians!(device, workgroups_surface)(
                balance_law,
                Val(dim),
                Val(N),
                dg.diffusion_direction,
                CentralNumericalFluxHigherOrder(),
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                state_conservative.data,
                state_auxiliary.data,
                grid.vgeo,
                grid.sgeo,
                grid.vmap⁻,
                grid.vmap⁺,
                grid.elemtobndy,
                grid.interiorelems,
                t;
                ndrange = ndrange_interior_surface,
                dependencies = (comp_stream,),
            )

        if communicate
            exchange_Qhypervisc_div = MPIStateArrays.end_ghost_exchange!(
                Qhypervisc_div,
                dependencies = exchange_Qhypervisc_div,
            )
        end

        comp_stream =
            interface_gradients_of_laplacians!(device, workgroups_surface)(
                balance_law,
                Val(dim),
                Val(N),
                dg.diffusion_direction,
                CentralNumericalFluxHigherOrder(),
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                state_conservative.data,
                state_auxiliary.data,
                grid.vgeo,
                grid.sgeo,
                grid.vmap⁻,
                grid.vmap⁺,
                grid.elemtobndy,
                grid.exteriorelems,
                t;
                ndrange = ndrange_exterior_surface,
                dependencies = (comp_stream, exchange_Qhypervisc_div),
            )

        if communicate
            exchange_Qhypervisc_grad = MPIStateArrays.begin_ghost_exchange!(
                Qhypervisc_grad,
                dependencies = comp_stream,
            )
        end
    end

    ###################
    # RHS Computation #
    ###################
    comp_stream = volume_tendency!(device, (Nq, Nq))(
        balance_law,
        Val(dim),
        Val(N),
        dg.direction,
        tendency.data,
        state_conservative.data,
        state_gradient_flux.data,
        Qhypervisc_grad.data,
        state_auxiliary.data,
        grid.vgeo,
        t,
        grid.ω,
        grid.D,
        topology.realelems,
        increment;
        ndrange = (nrealelem * Nq, Nq),
        dependencies = (comp_stream,),
    )

    comp_stream = interface_tendency!(device, workgroups_surface)(
        balance_law,
        Val(dim),
        Val(N),
        dg.direction,
        dg.numerical_flux_first_order,
        dg.numerical_flux_second_order,
        tendency.data,
        state_conservative.data,
        state_gradient_flux.data,
        Qhypervisc_grad.data,
        state_auxiliary.data,
        grid.vgeo,
        grid.sgeo,
        t,
        grid.vmap⁻,
        grid.vmap⁺,
        grid.elemtobndy,
        grid.interiorelems;
        ndrange = ndrange_interior_surface,
        dependencies = (comp_stream,),
    )

    if communicate
        if num_state_gradient_flux > 0 || nhyperviscstate > 0
            if num_state_gradient_flux > 0
                exchange_state_gradient_flux =
                    MPIStateArrays.end_ghost_exchange!(
                        state_gradient_flux;
                        dependencies = exchange_state_gradient_flux,
                    )

                # update_aux_diffusive may start asynchronous work on the
                # compute device and we synchronize those here through a device
                # event.
                wait(device, exchange_state_gradient_flux)
                update_auxiliary_state_gradient!(
                    dg,
                    balance_law,
                    state_conservative,
                    t,
                    dg.grid.topology.ghostelems,
                )
                exchange_state_gradient_flux = Event(device)
            end
            if nhyperviscstate > 0
                exchange_Qhypervisc_grad = MPIStateArrays.end_ghost_exchange!(
                    Qhypervisc_grad;
                    dependencies = exchange_Qhypervisc_grad,
                )
            end
        else
            exchange_state_conservative = MPIStateArrays.end_ghost_exchange!(
                state_conservative;
                dependencies = exchange_state_conservative,
            )

            # update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            wait(device, exchange_state_conservative)
            update_auxiliary_state!(
                dg,
                balance_law,
                state_conservative,
                t,
                dg.grid.topology.ghostelems,
            )
            exchange_state_conservative = Event(device)
        end
    end

    comp_stream = interface_tendency!(device, workgroups_surface)(
        balance_law,
        Val(dim),
        Val(N),
        dg.direction,
        dg.numerical_flux_first_order,
        dg.numerical_flux_second_order,
        tendency.data,
        state_conservative.data,
        state_gradient_flux.data,
        Qhypervisc_grad.data,
        state_auxiliary.data,
        grid.vgeo,
        grid.sgeo,
        t,
        grid.vmap⁻,
        grid.vmap⁺,
        grid.elemtobndy,
        grid.exteriorelems;
        ndrange = ndrange_exterior_surface,
        dependencies = (
            comp_stream,
            exchange_state_conservative,
            exchange_state_gradient_flux,
            exchange_Qhypervisc_grad,
        ),
    )

    # The synchronization here through a device event prevents CuArray based and
    # other default stream kernels from launching before the work scheduled in
    # this function is finished.
    wait(device, comp_stream)
end

function init_ode_state(dg::DGModel, args...; init_on_cpu = false)
    device = arraytype(dg.grid) <: Array ? CPU() : CUDA()

    balance_law = dg.balancelaw
    grid = dg.grid

    state_conservative = create_conservative_state(balance_law, grid)

    topology = grid.topology
    Np = dofs_per_element(grid)

    state_auxiliary = dg.state_auxiliary
    dim = dimensionality(grid)
    N = polynomialorder(grid)
    nrealelem = length(topology.realelems)

    if !init_on_cpu
        event = Event(device)
        event = kernel_init_state_conservative!(device, min(Np, 1024))(
            balance_law,
            Val(dim),
            Val(N),
            state_conservative.data,
            state_auxiliary.data,
            grid.vgeo,
            topology.realelems,
            args...;
            ndrange = Np * nrealelem,
            dependencies = (event,),
        )
        wait(device, event)
    else
        h_state_conservative = similar(state_conservative, Array)
        h_state_auxiliary = similar(state_auxiliary, Array)
        h_state_auxiliary .= state_auxiliary
        event = kernel_init_state_conservative!(CPU(), Np)(
            balance_law,
            Val(dim),
            Val(N),
            h_state_conservative.data,
            h_state_auxiliary.data,
            Array(grid.vgeo),
            topology.realelems,
            args...;
            ndrange = Np * nrealelem,
        )
        wait(event) # XXX: This could be `wait(device, event)` once KA supports that.
        state_conservative .= h_state_conservative
    end

    event = Event(device)
    event = MPIStateArrays.begin_ghost_exchange!(
        state_conservative;
        dependencies = event,
    )
    event = MPIStateArrays.end_ghost_exchange!(
        state_conservative;
        dependencies = event,
    )
    wait(device, event)

    return state_conservative
end

function restart_ode_state(dg::DGModel, state_data; init_on_cpu = false)
    bl = dg.balancelaw
    grid = dg.grid

    state = create_state(bl, grid)
    state .= state_data

    device = arraytype(dg.grid) <: Array ? CPU() : CUDA()
    event = Event(device)
    event = MPIStateArrays.begin_ghost_exchange!(state; dependencies = event)
    event = MPIStateArrays.end_ghost_exchange!(state; dependencies = event)
    wait(device, event)

    return state
end

function restart_auxiliary_state(bl, grid, aux_data)
    state_auxiliary = create_auxiliary_state(bl, grid)
    state_auxiliary .= aux_data
    return state_auxiliary
end

# fallback
function update_auxiliary_state!(
    dg::DGModel,
    balance_law::BalanceLaw,
    state_conservative::MPIStateArray,
    t::Real,
    elems::UnitRange,
)
    return false
end

function update_auxiliary_state_gradient!(
    dg::DGModel,
    balance_law::BalanceLaw,
    state_conservative::MPIStateArray,
    t::Real,
    elems::UnitRange,
)
    return false
end

function indefinite_stack_integral!(
    dg::DGModel,
    m::BalanceLaw,
    state_conservative::MPIStateArray,
    state_auxiliary::MPIStateArray,
    t::Real,
    elems::UnitRange = dg.grid.topology.elems,
)

    device = typeof(state_conservative.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq

    FT = eltype(state_conservative)

    # do integrals
    nelem = length(elems)
    nvertelem = topology.stacksize
    horzelems = fld1(first(elems), nvertelem):fld1(last(elems), nvertelem)

    event = Event(device)
    event = kernel_indefinite_stack_integral!(device, (Nq, Nqk))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        state_conservative.data,
        state_auxiliary.data,
        grid.vgeo,
        grid.Imat,
        horzelems;
        ndrange = (length(horzelems) * Nq, Nqk),
        dependencies = (event,),
    )
    wait(device, event)
end

function reverse_indefinite_stack_integral!(
    dg::DGModel,
    m::BalanceLaw,
    state_conservative::MPIStateArray,
    state_auxiliary::MPIStateArray,
    t::Real,
    elems::UnitRange = dg.grid.topology.elems,
)

    device = typeof(state_auxiliary.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq

    FT = eltype(state_auxiliary)

    # do integrals
    nelem = length(elems)
    nvertelem = topology.stacksize
    horzelems = fld1(first(elems), nvertelem):fld1(last(elems), nvertelem)

    event = Event(device)
    event = kernel_reverse_indefinite_stack_integral!(device, (Nq, Nqk))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        state_conservative.data,
        state_auxiliary.data,
        horzelems;
        ndrange = (length(horzelems) * Nq, Nqk),
        dependencies = (event,),
    )
    wait(device, event)
end

function nodal_update_auxiliary_state!(
    f!,
    dg::DGModel,
    m::BalanceLaw,
    state_conservative::MPIStateArray,
    t::Real,
    elems::UnitRange = dg.grid.topology.realelems;
    diffusive = false,
)
    device = typeof(state_conservative.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    nelem = length(elems)

    Np = dofs_per_element(grid)

    nodal_update_auxiliary_state! =
        kernel_nodal_update_auxiliary_state!(device, min(Np, 1024))
    ### update state_auxiliary variables
    event = Event(device)
    if diffusive
        event = nodal_update_auxiliary_state!(
            m,
            Val(dim),
            Val(N),
            f!,
            state_conservative.data,
            dg.state_auxiliary.data,
            dg.state_gradient_flux.data,
            t,
            elems,
            grid.activedofs;
            ndrange = Np * nelem,
            dependencies = (event,),
        )
    else
        event = nodal_update_auxiliary_state!(
            m,
            Val(dim),
            Val(N),
            f!,
            state_conservative.data,
            dg.state_auxiliary.data,
            t,
            elems,
            grid.activedofs;
            ndrange = Np * nelem,
            dependencies = (event,),
        )
    end
    wait(device, event)
end

"""
    courant(local_courant::Function, dg::DGModel, m::BalanceLaw,
            state_conservative::MPIStateArray, direction=EveryDirection())
Returns the maximum of the evaluation of the function `local_courant`
pointwise throughout the domain.  The function `local_courant` is given an
approximation of the local node distance `Δx`.  The `direction` controls which
reference directions are considered when computing the minimum node distance
`Δx`.
An example `local_courant` function is
    function local_courant(m::AtmosModel, state_conservative::Vars, state_auxiliary::Vars,
                           diffusive::Vars, Δx)
      return Δt * cmax / Δx
    end
where `Δt` is the time step size and `cmax` is the maximum flow speed in the
model.
"""
function courant(
    local_courant::Function,
    dg::DGModel,
    m::BalanceLaw,
    state_conservative::MPIStateArray,
    Δt,
    simtime,
    direction = EveryDirection(),
)
    grid = dg.grid
    topology = grid.topology
    nrealelem = length(topology.realelems)

    if nrealelem > 0
        N = polynomialorder(grid)
        dim = dimensionality(grid)
        Nq = N + 1
        Nqk = dim == 2 ? 1 : Nq
        device = grid.vgeo isa Array ? CPU() : CUDA()
        pointwise_courant = similar(grid.vgeo, Nq^dim, nrealelem)
        event = Event(device)
        event = Grids.kernel_min_neighbor_distance!(device, (Nq, Nq, Nqk))(
            Val(N),
            Val(dim),
            direction,
            pointwise_courant,
            grid.vgeo,
            topology.realelems;
            ndrange = (nrealelem * Nq, Nq, Nqk),
            dependencies = (event,),
        )
        event = kernel_local_courant!(device, Nq * Nq * Nqk)(
            m,
            Val(dim),
            Val(N),
            pointwise_courant,
            local_courant,
            state_conservative.data,
            dg.state_auxiliary.data,
            dg.state_gradient_flux.data,
            topology.realelems,
            Δt,
            simtime,
            direction;
            ndrange = nrealelem * Nq * Nq * Nqk,
            dependencies = (event,),
        )
        wait(device, event)
        rank_courant_max = maximum(pointwise_courant)
    else
        rank_courant_max = typemin(eltype(state_conservative))
    end

    MPI.Allreduce(rank_courant_max, max, topology.mpicomm)
end

function copy_stack_field_down!(
    dg::DGModel,
    m::BalanceLaw,
    state_auxiliary::MPIStateArray,
    fldin,
    fldout,
    elems = topology.elems,
)

    device = typeof(state_auxiliary.data) <: Array ? CPU() : CUDA()

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorder(grid)
    Nq = N + 1
    Nqk = dim == 2 ? 1 : Nq

    # do integrals
    nelem = length(elems)
    nvertelem = topology.stacksize
    horzelems = fld1(first(elems), nvertelem):fld1(last(elems), nvertelem)

    event = Event(device)
    event = kernel_copy_stack_field_down!(device, (Nq, Nqk))(
        Val(dim),
        Val(N),
        Val(nvertelem),
        state_auxiliary.data,
        horzelems,
        Val(fldin),
        Val(fldout);
        ndrange = (length(horzelems) * Nq, Nqk),
        dependencies = (event,),
    )
    wait(device, event)
end

function MPIStateArrays.MPIStateArray(dg::DGModel)
    balance_law = dg.balancelaw
    grid = dg.grid

    state_conservative = create_conservative_state(balance_law, grid)

    return state_conservative
end

function create_hypervisc_indexmap(balance_law::BalanceLaw)
    # helper function
    _getvars(v, ::Type) = v
    function _getvars(v::Vars, ::Type{T}) where {T <: NamedTuple}
        fields = getproperty.(Ref(v), fieldnames(T))
        collect(Iterators.Flatten(_getvars.(fields, fieldtypes(T))))
    end

    gradvars = vars_state_gradient(balance_law, Int)
    gradlapvars = vars_gradient_laplacian(balance_law, Int)
    indices = Vars{gradvars}(1:varsize(gradvars))
    SVector{varsize(gradlapvars)}(_getvars(indices, gradlapvars))
end


mutable struct AtomicSpeciesVars
    # species files path
    sppath::String
    # species filenames
    spfname::Vector{String}
    # species name
    spname::Vector{String}
    # species symbol
    spsymb::Vector{String}
    # species nuclear charge
    spzn::Vector{Float64}
    # ptnucl is .true. if the nuclei are to be treated as point charges, if .false.
    # the nuclei have a finite spherical distribution
    ptnucl::Bool
    # nuclear radius
    rnucl::Vector{Float64}
    # nuclear volume
    volnucl::Vector{Float64}
    # number of radial mesh points to nuclear radius
    nrnucl::Vector{Int64}
    # number of coarse radial mesh points to nuclear radius
    nrcnucl::Vector{Int64}
    # nuclear Coulomb potential
    vcln::Array{Float64,2}
    # species electronic charge
    spze::Vector{Float64}
    # species mass
    spmass::Vector{Float64}
    # smallest radial point for each species
    rminsp::Vector{Float64}
    # effective infinity for species
    rmaxsp::Vector{Float64}
    # number of radial points to effective infinity for each species
    nrsp::Vector{Int64}
    # maximum nrsp over all the species
    nrspmax::Int64
    # maximum allowed states for each species
    maxstsp::Int64 # parameter: 40
    # number of states for each species
    nstsp::Vector{Int64}
    # maximum nstsp over all the species
    nstspmax::Int64
    # core-valence cut-off energy for species file generation
    ecvcut::Float64
    # semi-core-valence cut-off energy for species file generation
    esccut::Float64
    # state principle quantum number for each species
    nsp::Array{Int64,2} # (maxstsp,maxspecies)
    # state l value for each species
    lsp::Array{Int64,2} # (maxstsp,maxspecies)
    # state k value for each species
    ksp::Array{Int64,2} #(maxstsp,maxspecies)
    # spcore is .true. if species state is core
    spcore::Array{Bool,2} #(maxstsp,maxspecies)
    # total number of core states
    nstcr::Int64
    # state eigenvalue for each species
    evalsp::Array{Float64,2} # (maxstsp,maxspecies)
    # state occupancy for each species
    occsp::Array{Float64,2} # (maxstsp,maxspecies)
    # species radial mesh to effective infinity
    rsp::Array{Float64,2}
    # r^l on radial mesh to muffin-tin radius
    rlsp::Array{Float64,3}
    # species charge density
    rhosp::Array{Float64,2}
    # species self-consistent potential
    vrsp::Array{Float64,2}
    # exchange-correlation type for atomic species (the converged ground-state of
    # the crystal does not depend on this choice)
    xctsp::Tuple{Int64,Int64,Int64}
end

function AtomicSpeciesVars( Nspecies::Int64 )

    sppath = ""
    spfname = Vector{String}(undef, Nspecies)
    spname = Vector{String}(undef, Nspecies)
    spsymb = Vector{String}(undef, Nspecies)

    spzn = zeros(Float64, Nspecies)

    ptnucl = true

    rnucl = zeros(Float64, Nspecies)

    volnucl = zeros(Float64, Nspecies)

    nrnucl = zeros(Int64, Nspecies)
    nrcnucl = zeros(Int64, Nspecies)

    vcln = zeros(Float64,1,1) # XXX will be properly setup later

    spze = zeros(Float64, Nspecies)
    spmass = zeros(Float64, Nspecies)

    rminsp = zeros(Float64, Nspecies)
    rmaxsp = zeros(Float64, Nspecies)
    nrsp = zeros(Int64, Nspecies)
    nrspmax = 0
    maxstsp = 40
    nstsp = zeros(Int64, Nspecies)
    nstspmax = 0

    ecvcut = 0.0
    esccut = 0.0

    nsp = zeros(Int64,maxstsp,Nspecies)
    lsp = zeros(Int64,maxstsp,Nspecies)
    ksp = zeros(Int64,maxstsp,Nspecies)
    spcore = zeros(Bool,maxstsp,Nspecies)
    nstcr = 0
    evalsp = zeros(Float64,maxstsp,Nspecies)
    occsp  = zeros(Float64,maxstsp,Nspecies)

    # XXX Will be setup properly later
    rsp = zeros(Float64,1,1)
    rlsp  = zeros(Float64,1,1,1)
    rhosp = zeros(Float64,1,1)
    vrsp  = zeros(Float64,1,1)
    xctsp  = (3,0,0)

    return AtomicSpeciesVars(
        sppath, spfname, spname, spsymb, spzn,
        ptnucl, rnucl, volnucl, nrnucl, nrcnucl,
        vcln, spze, spmass, rminsp, rmaxsp, nrsp,
        nrspmax, maxstsp, nstsp, nstspmax, ecvcut, esccut,
        nsp, lsp, ksp, spcore, nstcr, evalsp, occsp, rsp, rlsp, rhosp, vrsp, xctsp
    )
end

# Setting up vcln
function init_nuclear_pot!( atsp_vars::AtomicSpeciesVars )
    # spherical harmonic for l=m=0
    y00=0.28209479177387814347

    rsp = atsp_vars.rsp
    nrsp = atsp_vars.nrsp
    nspecies = size(nrsp,1)
    ptnucl = atsp_vars.ptnucl
    nrspmax = atsp_vars.nrspmax
    spzn = atsp_vars.spzn

    # determine the nuclear Coulomb potential
    atsp_vars.vcln = zeros(Float64,nrspmax,nspecies)
    vcln = atsp_vars.vcln
    t1 = 1.0/y00
    for is in 1:nspecies
        nr = nrsp[is]
        @views potnucl!(ptnucl, nr, rsp[:,is], spzn[is], vcln[:,is])
        vcln[1:nr,is] = t1*vcln[1:nr,is]
    end
    return
end