module TrackerCompat

using ..Tracker: Tracker,
                 TrackedReal,
                 TrackedVector,
                 TrackedMatrix,
                 TrackedArray,
                 TrackedVecOrMat,
                 @grad,
                 track,
                 data,
                 param

import ..Bijectors
using ..Bijectors: Log, SimplexBijector, ADBijector,
    TrackerAD, Inverse, Stacked, Exp

import ChainRulesCore
import LogExpFunctions

using Compat: eachcol
using LinearAlgebra
using Distributions: LocationScale

Bijectors.maporbroadcast(f, x::TrackedArray...) = f.(x...)
function Bijectors.maporbroadcast(
    f,
    x1::TrackedArray{T, N},
    x::AbstractArray{<:TrackedReal}...,
) where {T, N}
    return f.(convert(Array{TrackedReal{T}, N}, x1), x...)
end

Bijectors._eps(::Type{<:TrackedReal{T}}) where {T} = Bijectors._eps(T)
function Base.minimum(d::LocationScale{<:TrackedReal})
    m = minimum(d.ρ)
    if isfinite(m)
        return d.μ + d.σ * m
    else
        return m
    end
end
function Base.maximum(d::LocationScale{<:TrackedReal})
    m = maximum(d.ρ)
    if isfinite(m)
        return d.μ + d.σ * m
    else
        return m
    end
end

# AD implementations
function Bijectors.jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::Real
)
    return data(Tracker.gradient(b, x)[1])
end
function Bijectors.jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::AbstractVector{<:Real}
)
    # We extract `data` so that we don't return a `Tracked` type
    return data(Tracker.jacobian(b, x))
end

# implementations for Shift bijector
function Bijectors._logabsdetjac_shift(a::TrackedReal, x::Real, ::Val{0})
    return tracker_shift_logabsdetjac(a, x, Val(0))
end
function Bijectors._logabsdetjac_shift(a::TrackedReal, x::AbstractVector{<:Real}, ::Val{0})
    return tracker_shift_logabsdetjac(a, x, Val(0))
end
function Bijectors._logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractVector{<:Real},
    ::Val{1}
)
    return tracker_shift_logabsdetjac(a, x, Val(1))
end
function Bijectors._logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractMatrix{<:Real},
    ::Val{1}
)
    return tracker_shift_logabsdetjac(a, x, Val(1))
end
function tracker_shift_logabsdetjac(a, x, ::Val{N}) where {N}
    return param(Bijectors._logabsdetjac_shift(data(a), data(x), Val(N)))
end

# Log bijector

@grad function Bijectors.logabsdetjac(b::Log{1}, x::AbstractVector)
    return -sum(log, data(x)), Δ -> (nothing, -Δ ./ data(x))
end
@grad function Bijectors.logabsdetjac(b::Log{1}, x::AbstractMatrix)
    return -vec(sum(log, data(x); dims = 1)), Δ -> (nothing, .- Δ' ./ data(x))
end
@grad function Bijectors.logabsdetjac(b::Log{2}, x::AbstractMatrix)
    return -sum(log, data(x)), Δ -> (nothing, -Δ ./ data(x))
end

# implementations for Scale bijector
# Adjoints for 0-dim and 1-dim `Scale` using `Real`
function Bijectors._logabsdetjac_scale(a::TrackedReal, x::Real, ::Val{0})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(0))
end
@grad function Bijectors._logabsdetjac_scale(a::Real, x::Real, ::Val{0})
    return Bijectors._logabsdetjac_scale(data(a), data(x), Val(0)), Δ -> (inv(data(a)) .* Δ, nothing, nothing)
end
# Need to treat `AbstractVector` and `AbstractMatrix` separately due to ambiguity errors
function Bijectors._logabsdetjac_scale(a::TrackedReal, x::AbstractVector, ::Val{0})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(0))
end
@grad function Bijectors._logabsdetjac_scale(a::Real, x::AbstractVector, ::Val{0})
    da = data(a)
    J = fill(inv.(da), length(x))
    return Bijectors._logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
function Bijectors._logabsdetjac_scale(a::TrackedReal, x::AbstractMatrix, ::Val{0})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(0))
end
@grad function Bijectors._logabsdetjac_scale(a::Real, x::AbstractMatrix, ::Val{0})
    da = data(a)
    J = fill(size(x, 1) / da, size(x, 2))
    return Bijectors._logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
# adjoints for 1-dim and 2-dim `Scale` using `AbstractVector`
function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(1))
end
@grad function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    # ∂ᵢ (∑ⱼ log|aⱼ|) = ∑ⱼ δᵢⱼ ∂ᵢ log|aⱼ|
    #                 = ∂ᵢ log |aᵢ|
    #                 = (1 / aᵢ) ∂ᵢ aᵢ
    #                 = (1 / aᵢ)
    da = data(a)
    J = inv.(da)
    return Bijectors._logabsdetjac_scale(da, data(x), Val(1)), Δ -> (J .* Δ, nothing, nothing)
end
function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    return track(Bijectors._logabsdetjac_scale, a, data(x), Val(1))
end
@grad function Bijectors._logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    da = data(a)
    Jᵀ = repeat(inv.(da), 1, size(x, 2))
    return Bijectors._logabsdetjac_scale(da, data(x), Val(1)), Δ -> (Jᵀ * Δ, nothing, nothing)
end
# TODO: implement analytical gradient for scaling a vector using a matrix
# function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     track(_logabsdetjac_scale, a, data(x), Val{1})
# end
# @grad function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     throw
# end
# implementations for Stacked bijector
function Bijectors.logabsdetjac(b::Stacked, x::TrackedMatrix{<:Real})
    return map(eachcol(x)) do c
        Bijectors.logabsdetjac(b, c)
    end
end
# TODO: implement custom adjoint since we can exploit block-diagonal nature of `Stacked`
function (sb::Stacked)(x::TrackedMatrix{<:Real})
    return Bijectors.eachcolmaphcat(sb, x)
end
# Simplex adjoints
function Bijectors._simplex_bijector(X::TrackedVecOrMat, b::SimplexBijector{1})
    return track(Bijectors._simplex_bijector, X, b)
end
function Bijectors._simplex_inv_bijector(Y::TrackedVecOrMat, b::SimplexBijector{1})
    return track(Bijectors._simplex_inv_bijector, Y, b)
end
@grad function Bijectors._simplex_bijector(X::AbstractVector, b::SimplexBijector{1})
    Xd = data(X)
    return Bijectors._simplex_bijector(Xd, b), Δ -> (Bijectors.simplex_link_jacobian(Xd)' * Δ, nothing)
end
@grad function Bijectors._simplex_inv_bijector(Y::AbstractVector, b::SimplexBijector{1})
    Yd = data(Y)
    return Bijectors._simplex_inv_bijector(Yd, b), Δ -> (Bijectors.simplex_invlink_jacobian(Yd)' * Δ, nothing)
end

@grad function Bijectors._simplex_bijector(X::AbstractMatrix, b::SimplexBijector{1})
    Xd = data(X)
    return Bijectors._simplex_bijector(Xd, b), Δ -> begin
        Bijectors.maphcat(eachcol(Xd), eachcol(Δ)) do c1, c2
            Bijectors.simplex_link_jacobian(c1)' * c2
        end, nothing
    end
end
@grad function Bijectors._simplex_inv_bijector(Y::AbstractMatrix, b::SimplexBijector{1})
    Yd = data(Y)
    return Bijectors._simplex_inv_bijector(Yd, b), Δ -> begin
        Bijectors.maphcat(eachcol(Yd), eachcol(Δ)) do c1, c2
            Bijectors.simplex_invlink_jacobian(c1)' * c2
        end, nothing
    end
end

Bijectors.replace_diag(::typeof(log), X::TrackedMatrix) = track(Bijectors.replace_diag, log, X)
@grad function Bijectors.replace_diag(::typeof(log), X)
    Xd = data(X)
    f(i, j) = i == j ? log(Xd[i, j]) : Xd[i, j]
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = i == j ? ∇[i, j]/Xd[i, j] : ∇[i, j]
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

Bijectors.replace_diag(::typeof(exp), X::TrackedMatrix) = track(Bijectors.replace_diag, exp, X)
@grad function Bijectors.replace_diag(::typeof(exp), X)
    Xd = data(X)
    f(i, j) = ifelse(i == j, exp(Xd[i, j]), Xd[i, j])
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = ifelse(i == j, ∇[i, j]*exp(Xd[i, j]), ∇[i, j])
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

Bijectors.logabsdetjac(b::SimplexBijector{1}, x::TrackedVecOrMat) = track(Bijectors.logabsdetjac, b, x)
@grad function Bijectors.logabsdetjac(b::SimplexBijector{1}, x::AbstractVector)
    xd = data(x)
    return Bijectors.logabsdetjac(b, xd), Δ -> begin
        (nothing, Bijectors.simplex_logabsdetjac_gradient(xd) * Δ)
    end
end
@grad function Bijectors.logabsdetjac(b::SimplexBijector{1}, x::AbstractMatrix)
    xd = data(x)
    return Bijectors.logabsdetjac(b, xd), Δ -> begin
        (nothing, Bijectors.maphcat(eachcol(xd), Δ) do c, g
            Bijectors.simplex_logabsdetjac_gradient(c) * g
        end)
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractVector)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:α_, :β, :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedVecOrMat)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedVector)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
]
    @eval begin
        function Bijectors._radial_transform($(header...))
            α = LogExpFunctions.log1pexp(α_)            # from A.2
            β_hat = -α + LogExpFunctions.log1pexp(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = vectorof(typeof(β_hat))
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = typeof(z_0)
            else
                T = TV = typeof(z)
            end
            Tr = promote_type(eltype(z), eltype(z_0))
            r::Tr = norm((z .- z_0)::TV)
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractMatrix)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:α_, :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
]
    @eval begin
        function Bijectors._radial_transform($(header...))
            α = LogExpFunctions.log1pexp(α_)            # from A.2
            β_hat = -α + LogExpFunctions.log1pexp(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = matrixof(TV)
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = matrixof(TV)
            else
                T = typeof(z)
                TV = vectorof(T)
            end
            r::TV = eachcolnorm(z .- z_0)
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end
eachcolnorm(X) = map(norm, eachcol(X))
eachcolnorm(X::TrackedMatrix) = track(eachcolnorm, X)
@grad function eachcolnorm(X)
    Xd = data(X)
    y = map(norm, eachcol(Xd))
    y, Δ -> begin
        (Xd .* (Δ ./ y)',)
    end
end

function matrixof(::Type{<:TrackedArray{T,1,Vector{T}}}) where {T<:Real}
    return TrackedArray{T,2,Matrix{T}}
end
function matrixof(::Type{TrackedReal{T}}) where {T<:Real}
    return TrackedArray{T,2,Matrix{T}}
end
function vectorof(::Type{<:TrackedArray{T,2,Matrix{T}}}) where {T<:Real}
    return TrackedArray{T,1,Vector{T}}
end
function vectorof(::Type{TrackedReal{T}}) where {T<:Real}
    return TrackedArray{T,1,Vector{T}}
end

(b::Exp{0})(x::TrackedVector) = exp.(x)::vectorof(float(eltype(x)))
(b::Exp{1})(x::TrackedVector) = exp.(x)::vectorof(float(eltype(x)))
(b::Exp{1})(x::TrackedMatrix) = exp.(x)::matrixof(float(eltype(x)))
(b::Exp{2})(x::TrackedMatrix) = exp.(x)::matrixof(float(eltype(x)))

(b::Log{0})(x::TrackedVector) = log.(x)::vectorof(float(eltype(x)))
(b::Log{1})(x::TrackedVector) = log.(x)::vectorof(float(eltype(x)))
(b::Log{1})(x::TrackedMatrix) = log.(x)::matrixof(float(eltype(x)))
(b::Log{2})(x::TrackedMatrix) = log.(x)::matrixof(float(eltype(x)))

Bijectors.logabsdetjac(b::Log{0}, x::TrackedVector) = .-log.(x)::vectorof(float(eltype(x)))
Bijectors.logabsdetjac(b::Log{1}, x::TrackedMatrix) = - vec(sum(log.(x); dims = 1))

Bijectors.getpd(X::TrackedMatrix) = track(Bijectors.getpd, X)
@grad function Bijectors.getpd(X::AbstractMatrix)
    Xd = data(X)
    return Bijectors.LowerTriangular(Xd) * Bijectors.LowerTriangular(Xd)', Δ -> begin
        Xl = Bijectors.LowerTriangular(Xd)
        return (Bijectors.LowerTriangular(Δ' * Xl + Δ * Xl),)
    end
end

Bijectors.lower(A::TrackedMatrix) = track(Bijectors.lower, A)
@grad function Bijectors.lower(A::AbstractMatrix)
    Ad = data(A)
    return Bijectors.lower(Ad), Δ -> (Bijectors.lower(Δ),)
end

Bijectors._inv_link_chol_lkj(y::TrackedMatrix) = track(Bijectors._inv_link_chol_lkj, y)
@grad function Bijectors._inv_link_chol_lkj(y_tracked)
    y = data(y_tracked)

    K = LinearAlgebra.checksquare(y)

    w = similar(y)

    z_mat = similar(y) # cache for adjoint
    tmp_mat = similar(y)

    @inbounds for j in 1:K
        w[1, j] = 1
        for i in 2:j
            z = tanh(y[i-1, j])
            tmp = w[i-1, j]

            z_mat[i, j] = z
            tmp_mat[i, j] = tmp

            w[i-1, j] = z * tmp
            w[i, j] = tmp * sqrt(1 - z^2)
        end
        for i in (j+1):K
            w[i, j] = 0
        end
    end

    function pullback_inv_link_chol_lkj(Δw)
        LinearAlgebra.checksquare(Δw)

        Δy = zero(y)

        @inbounds for j in 1:K
            Δtmp = Δw[j,j]
            for i in j:-1:2
                Δz = Δw[i-1, j] * tmp_mat[i, j] - Δtmp * tmp_mat[i, j] / sqrt(1 - z_mat[i, j]^2) * z_mat[i, j]
                Δy[i-1, j] = Δz / cosh(y[i-1, j])^2
                Δtmp = Δw[i-1, j] * z_mat[i, j] + Δtmp * sqrt(1 - z_mat[i, j]^2)
            end
        end

        return (Δy,)
    end

    return w, pullback_inv_link_chol_lkj
end

Bijectors._link_chol_lkj(w::TrackedMatrix) = track(Bijectors._link_chol_lkj, w)
@grad function Bijectors._link_chol_lkj(w_tracked)
    w = data(w_tracked)

    K = LinearAlgebra.checksquare(w)

    z = similar(w)

    @inbounds z[1, 1] = 0

    tmp_mat = similar(w) # cache for pullback.

    @inbounds for j=2:K
        z[1, j] = atanh(w[1, j])
        tmp = sqrt(1 - w[1, j]^2)
        tmp_mat[1, j] = tmp
        for i in 2:(j - 1)
            p = w[i, j] / tmp
            tmp *= sqrt(1 - p^2)
            tmp_mat[i, j] = tmp
            z[i, j] = atanh(p)
        end
        z[j, j] = 0
    end

    function pullback_link_chol_lkj(Δz)
        LinearAlgebra.checksquare(Δz)

        Δw = similar(w)

        @inbounds Δw[1,1] = zero(eltype(Δz))

        @inbounds for j=2:K
            Δw[j, j] = 0
            Δtmp = zero(eltype(Δz)) # Δtmp_mat[j-1,j]
            for i in (j-1):-1:2
                p = w[i, j] / tmp_mat[i-1, j]
                ftmp = sqrt(1 - p^2)
                d_ftmp_p = -p / ftmp
                d_p_tmp = -w[i,j] / tmp_mat[i-1, j]^2

                Δp = Δz[i,j] / (1-p^2) + Δtmp * tmp_mat[i-1, j] * d_ftmp_p
                Δw[i, j] = Δp / tmp_mat[i-1, j]
                Δtmp = Δp * d_p_tmp + Δtmp * ftmp # update to "previous" Δtmp
            end
            Δw[1, j] = Δz[1, j] / (1-w[1,j]^2) - Δtmp / sqrt(1 - w[1,j]^2) * w[1,j]
        end

        return (Δw,)
    end

    return z, pullback_link_chol_lkj
end

function Bijectors.find_alpha(wt_y::T, wt_u_hat::T, b::T) where {T<:TrackedReal}
    return track(Bijectors.find_alpha, wt_y, wt_u_hat, b)
end
@grad function Bijectors.find_alpha(wt_y::TrackedReal, wt_u_hat::TrackedReal, b::TrackedReal)
    α = Bijectors.find_alpha(data(wt_y), data(wt_u_hat), data(b))

    ∂wt_y = inv(1 + wt_u_hat * sech(α + b)^2)
    ∂wt_u_hat = - tanh(α + b) * ∂wt_y
    ∂b = ∂wt_y - 1
    find_alpha_pullback(Δ::Real) = (Δ * ∂wt_y, Δ * ∂wt_u_hat, Δ * ∂b)

    return α, find_alpha_pullback
end

# `OrderedBijector`
Bijectors._transform_ordered(y::Union{TrackedVector,TrackedMatrix}) = track(Bijectors._transform_ordered, y)
@grad function Bijectors._transform_ordered(y::AbstractVecOrMat)
    x, dx = ChainRulesCore.rrule(Bijectors._transform_ordered, data(y))
    return x, (wrap_chainrules_output ∘ Base.tail ∘ dx)
end

function Bijectors._transform_inverse_ordered(x::Union{TrackedVector,TrackedMatrix})
    return track(Bijectors._transform_inverse_ordered, x)
end
@grad function Bijectors._transform_inverse_ordered(x::AbstractVecOrMat)
    y, dy = ChainRulesCore.rrule(Bijectors._transform_inverse_ordered, data(x))
    return y, (wrap_chainrules_output ∘ Base.tail ∘ dy)
end

# NOTE: Probably doesn't work in complete generality.
wrap_chainrules_output(x) = x
wrap_chainrules_output(x::ChainRulesCore.AbstractZero) = nothing
wrap_chainrules_output(x::Tuple) = map(wrap_chainrules_output, x)

end


export fasta2matrix, letter2num, num2letter, matrix2fasta

##############################################################
##############################################################

"""
  function fasta2matrix(filename::AbstractString; max_gap_fraction = 1)

  Reads an aligned fasta file and returns a matrix with amino acids converted
  to numbers following the convention A --> 1, C --> 2, ... '-' --> 21.

    "filename": full path of the fasta file.
    "max_gap_fraction": to be removed.

"""

#--------------------------------------------------------

function fasta2matrix(filename::AbstractString; max_gap_fraction = 1)
    f = FastaReader(filename)

    max_gap_fraction = Float64(max_gap_fraction)

    # pass 1

    seqs = Int[]
    inds = Int[]
    fseqlen = 0

    for (name, seq) in f
        ngaps = 0
        if f.num_parsed == 1
            ls = length(seq)
            resize!(inds, ls)
            for i = 1:ls
                c = seq[i]
                if c != '.' && c == uppercase(c)
                    fseqlen += 1
                    inds[fseqlen] = i
                    c == '-' && (ngaps += 1)
                end
            end
        else
            ls = length(seq)
            ls == length(inds) || error("inputs are not aligned")
            tstfseqlen = 0
            for i = 1:ls
                c = seq[i]
                if c != '.' && c == uppercase(c)
                    tstfseqlen += 1
                    inds[tstfseqlen] == i || error("inconsistent inputs")
                    c == '-' && (ngaps += 1)
                end
            end
            tstfseqlen == fseqlen || error("inconsistent inputs")
        end
        ngaps / fseqlen <= max_gap_fraction && push!(seqs, f.num_parsed)
    end

    length(seqs) > 0 || error("Out of $(f.num_parsed) sequences, none passed the filter (max_gap_fraction=$max_gap_fraction)")

    # pass 2
    Z = Array{Int16}(undef, fseqlen, length(seqs))
    seqid = 1
    for (name, seq) in f
        seqs[end] < f.num_parsed && break
        seqs[seqid] == f.num_parsed || continue
        for i = 1:fseqlen
            c = seq[inds[i]]
            Z[i, seqid] = letter2num(c)
        end
        seqid += 1
    end
    @assert seqid == length(seqs) + 1

    close(f)

    size(Z)[2] == 1 && return Z[:, 1]
    return Int16.(Z')
end

##############################################################
"""
    letter2num(c::Union{Char,UInt8})

    Takes as input a char (representing an amino acid)
    and returns its conventional number representation.
    In this case with the convention "1 2 .. 21" <==> "A C .. -"
"""

#--------------------------------------------------------

let alphabet = [ 1, 21, 2, 3, 4, 5, 6, 7, 8, 21, 9, 10, 11, 12, 21, 13, 14, 15, 16, 17, 21, 18, 19, 21, 20]
               # A, B,  C, D, E, F, G, H, I, J,  K, L,  M,  N,  O,  P,  Q,  R,  S,  T,  U,  V,  W,  X,  Y
    global letter2num
    function letter2num(c::Union{Char,UInt8})
        i = UInt8(c) - 0x40
        1 <= i <= 25 && return alphabet[i]
        return 21
     end
end

##############################################################
"""
    num2letter(i::Integer)

    Takes as input an integer (representing an amino acid)
    and returns its conventional character representation.
    In this case with the convention "1 2 .. 21" <==> "A C .. -"
"""

#--------------------------------------------------------

let alphabet = ["A", "C", "D", "E", "F", "G", "H", "I",  "K", "L",  "M",  "N", "P",  "Q",  "R",
"S",  "T", "V",  "W",  "Y"]
    global num2letter
    function num2letter(i :: Integer)
        1 <= i <= 20 && return alphabet[i]
        return "-"
    end
end

##############################################################
"""
    vec2string(v::Array{:<Integer, 1})

    Takes as input a vector of integers (representing an amino acids)
    and returns a list of characters, the corresponding amino acids.
    In this case with the convention "1 2 .. 21" <==> "A C .. -"
"""

#--------------------------------------------------------

function vec2string(v::Array{<:Integer, 1})
    s = ""
    for i in v
        s = s*num2letter(i)
    end
    return s
end

##############################################################
"""
	string2vec(s::AbstractString)

    Takes as input a string (representing amino acids)
    and returns a vector of numbers, the corresponding number representation.
    In this case with the convention "1 2 .. 21" <==> "A C .. -"
"""

#--------------------------------------------------------

function string2vec(s::AbstractString)
    v = Vector{Int8}(undef, length(s))
    for (i, l) in enumerate(s)
        v[i] = letter2num(l)
    end
    return v
end

##############################################################
"""
	write_single_muts_MSA(path_wt::AbstractString, fitness_wt::Array{Union{Missing, Float64}},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")
	write_single_muts_MSA(wt::Array{<:Integer, 1}, fitness_wt::Array{Union{Missing, Float64}},
			dir_out::AbstractString; wt_name::AbstractString = "WT")
	write_single_muts_MSA(path_wt::AbstractString, fitness_wt::Array{Any, 2},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")
	write_single_muts_MSA(wt::Array{<:Integer, 1}, fitness_wt::Array{Any, 2},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")

    USE:
    Given a sequence, it returns all single mutations in a MSA fashon.

    INPUT:
    "fitness_wt": if it is a 1d vector is must contain missings for missing mutations
    			  and should cover only the non gapped positions

    			  if it is 2d vector the first column must show the position, the second the mutation
    			  and the third the fitness value

    OUTPUT: writes a .fasta and a .fit file in the "dir_out" folder

"""

#--------------------------------------------------------

function write_single_muts_MSA(path_wt::AbstractString, fitness_wt::Array{Union{Missing, Float64}},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")

	wt = fasta2matrix(path_wt)
	dir_out == "" && (  dir_out = (*).(split(path_wt, "/")[1:end-1])  )
	path_MSA_out = joinpath(dir_out, "single_muts_MSA_$(wt_name).fasta")
	path_fit_out = joinpath(dir_out, "single_muts_fitness_$(wt_name).fit")
	vec_fit = []
	k = 0
	FastaWriter(path_MSA_out, "w") do file
	    writeentry(file, "$k | $(wt_name)", vec2string(wt))
	    for i in [pos for (pos, amino) in enumerate(wt) if amino!= 21]
	        mutant = copy(wt)
	        for amino in 1:20
	            mutant[i] = amino
	            new_amino = num2letter(mutant[i])
	            old_amino = num2letter(wt[i])
	            if !ismissing(fitness_wt[(i-1)*20 + amino]) && (amino != wt[i])
	                k+=1
	                append!(vec_fit, fitness_wt[(i-1)*20 + amino])
	                writeentry(file, "$k | $(old_amino)$(i)$(new_amino)", vec2string(mutant))
	            end
	        end
	    end
	end

	pushfirst!(vec_fit, 0)
	writedlm(path_fit_out, vec_fit)
end

#--------------------------------------------------------

function write_single_muts_MSA(wt::Array{<:Integer, 1}, fitness_wt::Array{Union{Missing, Float64}},
			dir_out::AbstractString; wt_name::AbstractString = "WT")

	path_MSA_out = joinpath(dir_out, "single_muts_MSA_$(wt_name).fasta")
	path_fit_out = joinpath(dir_out, "single_muts_fitness_$(wt_name).fit")
	vec_fit = []
	k = 0
	FastaWriter(path_MSA_out, "w") do file
	    writeentry(file, "$k | $(wt_name)", vec2string(wt))
	    for i in [pos for (pos, amino) in enumerate(wt) if amino!= 21]
	        mutant = copy(wt)
	        for amino in 1:20
	            mutant[i] = amino
	            new_amino = num2letter(mutant[i])
	            old_amino = num2letter(wt[i])
	            if !ismissing(fitness_wt[(i-1)*20 + amino]) && (amino != wt[i])
	                k+=1
	                append!(vec_fit, fitness_wt[(i-1)*20 + amino])
	                writeentry(file, "$k | $(old_amino)$(i)$(new_amino)", vec2string(mutant))
	            end
	        end
	    end
	end

	pushfirst!(vec_fit, 0)
	writedlm(path_fit_out, vec_fit)
end

#--------------------------------------------------------

function write_single_muts_MSA(path_wt::AbstractString, fitness_wt::Array{Any, 2},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")

	wt = fasta2matrix(path_wt)
	dir_out == "" && (  dir_out = (*).(split(path_wt, "/")[1:end-1])  )
	path_MSA_out = joinpath(dir_out, "single_muts_MSA_$(wt_name).fasta")
	path_fit_out = joinpath(dir_out, "single_muts_fitness_$(wt_name).fit")
	vec_fit = []

    l = 0
	FastaWriter(path_MSA_out, "w") do file
	    writeentry(file, "0 | $(wt_name)", vec2string(wt))
        for k in 1:size(fitness_wt, 1)

	        mutant = copy(wt)
            pos = fitness_wt[k, 1]

            old_amino_lettr = num2letter(wt[pos])
            new_amino_lettr = fitness_wt[k, 2][1]
            new_amino_num = letter2num(new_amino_lettr)

            mutant[pos] = new_amino_num

            if (new_amino_num != wt[pos])
                l +=1
                append!(vec_fit, fitness_wt[k, 3])
                writeentry(file, "$l | $(old_amino_lettr)$(pos)$(new_amino_lettr)", vec2string(mutant))
            end
	    end
	end

	pushfirst!(vec_fit, 0)
	writedlm(path_fit_out, vec_fit)
end

#--------------------------------------------------------

function write_single_muts_MSA(wt::Array{<:Integer, 1}, fitness_wt::Array{Any, 2},
			dir_out::AbstractString = ""; wt_name::AbstractString = "WT")

	dir_out == "" && (  dir_out = (*).(split(path_wt, "/")[1:end-1])  )
	path_MSA_out = joinpath(dir_out, "single_muts_MSA_$(wt_name).fasta")
	path_fit_out = joinpath(dir_out, "single_muts_fitness_$(wt_name).fit")
	vec_fit = []

    l = 0
	FastaWriter(path_MSA_out, "w") do file
	    writeentry(file, "0 | $(wt_name)", vec2string(wt))
        for k in 1:size(fitness_wt, 1)

	        mutant = copy(wt)
            pos = fitness_wt[k, 1]

            old_amino_lettr = num2letter(wt[pos])
            new_amino_lettr = fitness_wt[k, 2][1]
            new_amino_num = letter2num(new_amino_lettr)

            mutant[pos] = new_amino_num

            if (new_amino_num != wt[pos])
                l +=1
                append!(vec_fit, fitness_wt[k, 3])
                writeentry(file, "$l | $(old_amino_lettr)$(pos)$(new_amino_lettr)", vec2string(mutant))
            end
	    end
	end

	pushfirst!(vec_fit, 0)
	writedlm(path_fit_out, vec_fit)
end

##############################################################
"""
    function matrix2fasta(path::AbstractString, MSA, desc_name::AbstractString = "")

    (path, MSA) --> write MSA to path

    USE:
    Given a vector or matrix of numbers representing an amino acid, write it in letter format to "path".

    INPUT:
    "path": path of output
    "MSA": matrix or vector of numbers (1 to 21 in the convention used here)
    "desc_name": name of the sequence (works only if NOT an MSA)

    OUTPUT:

"""

 #--------------------------------------------------------

function matrix2fasta(path::AbstractString, MSA, desc_name::AbstractString = "", start_index = 1)

    M = size(MSA, 1)

    single_seq_flag = (length(size(MSA)) == 1)

    if single_seq_flag == true
        FastaWriter(path, "w") do file
            desc_name == "" && (desc_name = "1")
            writeentry(file, desc_name, vec2string(MSA) )
        end
    else
	seq_index = start_index
        FastaWriter(path, "w") do file
            for i in 1:M
                seq = MSA[i, :]
                writeentry(file, "$(seq_index) $(desc_name)", vec2string(seq) )
		seq_index +=1
            end
        end
    end
end


function get_corrected_measurement(obj::KalmanEstimation)
    return jcall(obj, "getCorrectedMeasurement", EstimatedMeasurement, ())
end

function get_corrected_spacecraft_states(obj::KalmanEstimation)
    return jcall(obj, "getCorrectedSpacecraftStates", Vector{SpacecraftState}, ())
end

function get_current_date(obj::KalmanEstimation)
    return jcall(obj, "getCurrentDate", AbsoluteDate, ())
end

function get_current_measurement_number(obj::KalmanEstimation)
    return jcall(obj, "getCurrentMeasurementNumber", jint, ())
end

function get_estimated_measurements_parameters(obj::KalmanEstimation)
    return jcall(obj, "getEstimatedMeasurementsParameters", ParameterDriversList, ())
end

function get_estimated_orbital_parameters(obj::KalmanEstimation)
    return jcall(obj, "getEstimatedOrbitalParameters", ParameterDriversList, ())
end

function get_estimated_propagation_parameters(obj::KalmanEstimation)
    return jcall(obj, "getEstimatedPropagationParameters", ParameterDriversList, ())
end

function get_physical_estimated_covariance_matrix(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalEstimatedCovarianceMatrix", RealMatrix, ())
end

function get_physical_estimated_state(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalEstimatedState", RealVector, ())
end

function get_physical_innovation_covariance_matrix(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalInnovationCovarianceMatrix", RealMatrix, ())
end

function get_physical_kalman_gain(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalKalmanGain", RealMatrix, ())
end

function get_physical_measurement_jacobian(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalMeasurementJacobian", RealMatrix, ())
end

function get_physical_state_transition_matrix(obj::KalmanEstimation)
    return jcall(obj, "getPhysicalStateTransitionMatrix", RealMatrix, ())
end

function get_predicted_measurement(obj::KalmanEstimation)
    return jcall(obj, "getPredictedMeasurement", EstimatedMeasurement, ())
end

function get_predicted_spacecraft_states(obj::KalmanEstimation)
    return jcall(obj, "getPredictedSpacecraftStates", Vector{SpacecraftState}, ())
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


"""
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