struct _Undefined end
struct _Defined end

"""
    rasterize(x; kw...)
    rasterize(points, values; kw...)

Rasterize the points and values in a Tables.jl compatible object,
or separate `points` and `values` vetors/iterators, which must be
the same length.

If a GeoInterface `AbstractGeometry` or nested `Vector`s of `Tuple`/`Vector`
points is passed in, a `fill` keyword is also required to provide the value
that the array will be filled with. If `fill` is a function, it will be
applied to the existing value present in the array.

# Arguments

- `x`: a Tables.jl compatible object containing points and values columns,
    an GeoInterface.jl `AbstractGeometry`, or a nested `Vector` of `Vectors`.
- `points`: A `Vector` or nested `Vectors` holding `Vector` or `Tuple` of `Real`
- `values` A `Vector` of values to be written to a `Raster`, or a Vector of `NamedTuple`
    to write to a `RasterStack`.

# Keywords

These are detected automatically from `A` and `data` where possible.

- `to`: a `Raster`, `RasterStack` of `Tuple` of `Dimension` to use as a to.
- `order`: A `Tuple` of pairs `Dim => Symbol` for the keys in the data that match
    the dimension, or for the order of dimensions in ppoint values, like `(X, Y)`.
- `atol`: an absolute tolerance for rasterizing to dimensions with `Points` sampling.
- `filename`: a filename to write to directly, useful for large files.
- `suffix`: a string or value to append to the filename.
    A tuple of `suffix` will be applied to stack layers. `keys(st)` are the default.

## Geometry keywords

These can be used when a `GeoInterface.AbstractGeometry` is passed in.

- `fill`: the value to fill a polygon with, if `data` is a polygon.
- `shape`: Force `data` to be treated as `:polygon`, `:line` or `:point`.

And specifically for `shape=:polygon`:

- `boundary`: include pixels where the `:center` is inside the polygon, where
    the line `:touches` the pixel, or that are completely `:inside` inside the polygon.

## Table keywords

- `name`: A `Symbol` to return a `Raster` from a single column,
    or `Tuple` of `Symbol` to return a `RasterStack` from multiple columns.

# Example

Rasterize a shapefile for China and plot, with a border.

```jldoctest
using Rasters, Plots, Dates, Shapefile, Downloads
using Rasters.LookupArrays

# Download a borders shapefile
shapefile_url = "https://github.com/nvkelso/natural-earth-vector/raw/master/10m_cultural/ne_10m_admin_0_countries.shp"
shapefile_name = "boundary_lines.shp"
isfile(shapefile_name) || Downloads.download(shapefile_url, shapefile_name)

# Loade the shapes for china
china_border = Shapefile.Handle(shapefile_name).shapes[10]

# Define dims for the china area
dms = Y(Projected(15.0:0.1:55.0; order=ForwardOrdered(), span=Regular(0.1), sampling=Intervals(Start()), crs=EPSG(4326))),
      X(Projected(70.0:0.1:140; order=ForwardOrdered(), span=Regular(0.1), sampling=Intervals(Start()), crs=EPSG(4326)))

# Rasterize the border polygon
china = rasterize(china_border;
    to=dms, missingval=-9999, fill=1,
    order=(X, Y), shape=:polygon,
    boundary=:touches,
)

# And plot
p = plot(china; color=:spring)
plot!(p, china_border; fillalpha=0, linewidth=0.6)
savefig("build/china_rasterized.png")

# output
```

![rasterize](china_rasterized.pngfill)

$EXPERIMENTAL
"""
function rasterize(data; to, order=nothing, name=nothing, kw...)
    if Tables.istable(data)
        points, values, order, name = _table_to_points_values(to, data; order, name)
        return _rasterize(to, points, values; order, name, kw...)
    else
        order = isnothing(order) ? DEFAULT_POINT_ORDER : order
        order = dims(to, order)
        return _rasterize(to, data; order, init=_Undefined(), kw...)
    end
end
rasterize(points, values; to, kw...)  = _rasterize(to, points, values; kw...)

function _rasterize(to::DimTuple, points;
    filename=nothing, suffix=nothing, metadata=NoMetadata(), keys=nothing, name=keys, parent=nothing,
    fill, eltype=typeof(fill), missingval=_writeable_missing(filename, eltype), kw...
)
    A = _alloc_rasterize(filename, eltype, to; missingval, suffix, parent) do a
        rasterize!(a, points; fill, missingval, name, kw...)
    end
    return A
end
function _rasterize(to::DimTuple, points, vals;
    missingval=nothing, filename=nothing, suffix=nothing,
    metadata=NoMetadata(), keys=nothing, name=keys, parent=nothing, kw...
)
    firstval = first(vals)

    if firstval isa Union{NamedTuple,Tuple,AbstractArray}
        keys = _filter_name(name, vals)
        # Rasterize mutiple values to a stack
        layers = map(keys, values(firstval)) do key, val
            missingval = isnothing(missingval) ? _writeable_missing(filename, typeof(val)) : missingval
            _alloc_rasterize(filename, typeof(val), to; name, metadata, missingval, suffix=key, parent) do a
                a .= missingval
            end
        end |> NamedTuple{keys}
        st = RasterStack(layers, to)
        return rasterize!(st, points, vals; kw...)
    else
        name = first(cleankeys(name))
        # Rasterize to an array
        missingval = isnothing(missingval) ? _writeable_missing(filename, typeof(firstval)) : missingval
        A = _alloc_rasterize(filename, typeof(firstval), to; name, metadata, missingval, suffix, parent) do a
            a .= missingval
        end
        return rasterize!(A, points, vals; kw...)
    end
end
function _rasterize(to::AbstractRaster, args...;
    missingval=missingval(to), keys=name(to), name=keys, kw...
)
    _rasterize(dims(to), args...; name, missingval, kw...)
end
function _rasterize(to::AbstractRasterStack, args...; keys=keys(to), name=keys, kw...)
    return _rasterize(dims(to), args...; name, kw...)
end

function _alloc_rasterize(f, filename, T, to; missingval, suffix=nothing, kw...)
    T = promote_type(typeof(missingval), T)
    A = create(filename, T, to; suffix, missingval, kw...)
    open(f, A; write=true)
    return A
end

"""
    rasterize!(x, data; order, name, atol)
    rasterize!(x, points, vals; order, atol)

Rasterize the points and vals in `data`, or the `points` and `vals` objects,
into the [`Raster`](@ref) or [`RasterStack`](@ref) `x`.

# Arguments

- `x`: a `Raster` or `RasterStack` to rasterize to.
- `data`: a Tables.jl compatible object containing points and values or a
    polygon - an GeoInterface.jl `AbstractGeometry`, or a nested `Vector` of `Vectors`.
- `points`: A `Vector` or nested `Vector` holding `Vector` or `Tuple` of `Real`
- `vals` A `Vector` of values to be written when `x` is a `Raster`, or a Vector of
    `NamedTuple` to write when `x` is a `RasterStack`.

# Keywords

These are detected automatically from `A` and `data` where possible.

- `order`: A `Tuple` of pairs `Dim => Symbol` for the keys in the data that match
    the dimension, or for the order of dimensions in points, like `(X, Y)`.
- `atol`: an absolute tolerance for rasterizing to dimensions with `Points` sampling.
- `filename`: a filename to write to directly, useful for large files.
- `suffix`: a string or value to append to the filename.
    A tuple of `suffix` will be applied to stack layers. `keys(st)` are the default.

## Geometry keywords

These can be used when a `GeoInterface.AbstractGeometry` is passed in.

- `fill`: the value to fill a polygon with, if `data` is a polygon.
    `fill can also be a `Function` of the existing value.
- `shape`: Force `data` to be treated as `:polygon`, `:line` or `:point`.

And specifically for `shape=:polygon`:

- `boundary`: include pixels where the `:center` is inside the polygon, where
    the line `:touches` the pixel, or that are completely `:inside` inside the polygon.

## Table keywords

- `name`: A `Symbol` to return a `Raster` from a single column,
    or `Tuple` of `Symbol` to return a `RasterStack` from multiple columns.

# Example

```jldoctest
using Rasters, Plots, Dates, Shapefile, GeoInterface, Downloads
using Rasters.LookupArrays

# Download a borders shapefile
shapefile_url = "https://github.com/nvkelso/natural-earth-vector/raw/master/10m_cultural/ne_10m_admin_0_countries.shp"
shapefile_name = "indonesia_border.shp"
isfile(shapefile_name) || Downloads.download(shapefile_url, shapefile_name)

# Load the shapes for denmark
indonesia_border = Shapefile.Handle(shapefile_name).shapes[1]

# Make an empty EPSG 4326 projected Raster of the area of Indonesia
dimz = Y(-15.0:0.1:10.9; mode=Projected(; sampling=Intervals(Start()), crs=EPSG(4326))),
       X(90.0:0.1:145; mode=Projected(; sampling=Intervals(Start()), crs=EPSG(4326)))

A = Raster(zeros(UInt16, dimz); missingval=0)

# Rasterize each island with a different number
for (i, shp) in enumerate(coordinates(indonesia_border))
    rasterize!(A, shp; fill=i, order=(X, Y), shape=:polygon, boundary=:touches)
end

# And plot
p = plot(A; color=:spring)
plot!(p, indonesia_border; fillalpha=0, linewidth=0.7)
savefig("build/indonesia_rasterized.png")

# output

```

![rasterize](indonesia_rasterized.png)

$EXPERIMENTAL
"""
function rasterize!(x::RasterStackOrArray, data;
    order=nothing, name=nothing, missingval=missingval(x), kw...
)
    if Tables.istable(data)
        _rasterize_table!(x, data; name, kw...)
    else
        order = isnothing(order) ? DEFAULT_POINT_ORDER : order
        buffer = Raster(falses(commondims(x, order)))
        _rasterize_geometry!(x, data; order, buffer, init=_Defined(), missingval, kw...)
    end
end
function rasterize!(A::AbstractRaster, geom::GI.AbstractGeometry, vals; shape=:point, kw...)
    rasterize!(A, _flat_nodes(geom), vals; kw...)
end
function rasterize!(st::AbstractRasterStack, geom::GI.AbstractGeometry, vals; shape=:point, kw...)
    rasterize!(st, _flat_nodes(geom), vals; kw...)
end
function rasterize!(st::AbstractRasterStack, points, vals; keys=keys(st), name=keys, kw...)
    keys = _filter_name(name, vals)
    return _rasterize!(st[name], points, vals; name, kw...)
end
function rasterize!(A::AbstractRaster, points, vals; name=nothing, kw...)
    _rasterize!(A, points, vals; name, kw...)
end

function _rasterize!(x::RasterStackOrArray, points, vals;
    order=DEFAULT_POINT_ORDER, atol=nothing, name=nothing
)
    check_points(points, order)
    isdisk(first(x)) && _warn_disk()
    ordered_dims = dims(x, order)
    _without_mapped_crs(x) do x
        map(points, vals) do p, v
            any(map(ismissing, p)) && return nothing
            selectors = map((d, x) -> _at_or_contains(d, x, atol), ordered_dims, p)
            all(map(DD.hasselection, ordered_dims, selectors)) || return nothing
            _set_at_selection!(x, selectors, v)
            return nothing
        end
    end
    return x
end

function _set_at_selection!(st::AbstractRasterStack, selectors, vals)
    map(values(st), values(vals)) do A, val
        _set_at_selection!(A, selectors, val)
    end
end
function _set_at_selection!(A::AbstractRaster, selectors, val::Union{NamedTuple,Tuple})
    _set_at_selection!(A, selectors, first(val))
end
function _set_at_selection!(A::AbstractRaster, selectors, val)
    if length(selectors) == length(dims(A))
        A[selectors...] = val
    else
        A[selectors...] .= val
    end
end

function check_points(points, order)
    lp = length(first(points))
    lo = length(order)
    if lp !== lo
        throw(ArgumentError("length of points in `points` ($lp) is different to the length of `order` ($lo)"))
    end
end

# _rasterize_table!
function _rasterize_table!(x::RasterStackOrArray, data; order=nothing, name=nothing, kw...)
    points, vals, order, name = _table_to_points_values(x, data; order, name)
    return rasterize!(x, points, vals; order, name, kw...)
end

# _rasterize_geometry!
function _rasterize_geometry!(x::RasterStackOrArray, geom;
    buffer, fill, order, init, missingval, kw...
)
    boolmask!(buffer, geom; order, kw...)
    return _fill!(x, buffer, fill, init, missingval)
end

_fill(B, fill, missingval) where T = broadcast_dims(x -> x ? fill : missingval, B)

function _fill!(st::AbstractRasterStack, B, fill, args...)
    map((a, f) -> _fill!(a, B, f, args...), st, fill)
    return st
end
# If the array is initialised, we can use the existing values
function _fill!(A::AbstractRaster{T}, B, fill, init::_Defined, missingval) where T
    broadcast_dims!(A, A, B) do a, b
        val = b ? (fill isa Function ? fill(a) : fill) : a
        convert(T, val) # In case we are writing to disk
    end
end
# If the array is not yet initialised, we have to fill with fill and misssingval
function _fill!(A::AbstractRaster{T}, B, fill, init::_Undefined, missingval) where T
    fill = convert(T, fill) # In case we are writing to disk
    missingval = convert(T, missingval)
    broadcast_dims!(A, B) do b
        b ? fill : missingval
    end
end

function _at_or_contains(d, v, atol)
    selector = sampling(d) isa Intervals ? Contains(v) : At(v; atol=atol)
    DD.basetypeof(d)(selector)
end

function _filter_name(name, vals)
    firstval = first(vals)
    name = if isnothing(name)
        if firstval isa NamedTuple
            keys(firstval)
        else
            stackkeys
        end
    else
        if firstval isa NamedTuple
            name == keys(firstval) || throw(ArgumentError("`name` keyword does not match point names"))
        elseif firstval isa Union{Tuple,Array}
            length(name) == length(firstval) || throw(ArgumentError("`name` keyword does not match length of point"))
        end
        name
    end
end


for T in subtypes(AbstractPointwise)
    @eval begin
        function (ρ::$T)(node::AbstractNode)
            return Node(ρ(node.data), node.ranges)
        end
    end
end


using Test
using Random
using NCDatasets
using OrderedCollections
using NCRegressionTests
Random.seed!(123)

@testset "NCRegressionTests - perfect match" begin

    sz = (123, 145)
    data = randn(MersenneTwister(152), sz)

    filenameA = tempname()
    filenameB = tempname()

    NCDataset(filenameA, "c") do ds
        defDim(ds, "lon", sz[1])
        defDim(ds, "lat", sz[2])
        v = defVar(ds, "var", Float64, ("lon", "lat"))
        v[:, :] = data
        v = defVar(ds, "var2", Float64, ("lon", "lat"))
        v[:, :] = 2 * data
    end
    NCDataset(filenameB, "c") do ds
        defDim(ds, "lon", sz[1])
        defDim(ds, "lat", sz[2])
        v = defVar(ds, "var", Float64, ("lon", "lat"))
        v[:, :] = data
        v = defVar(ds, "var2", Float64, ("lon", "lat"))
        v[:, :] = 2 * data
    end

    reference_mse = OrderedCollections.OrderedDict("var" => 0, "var2" => 0)

    computed_mse = NCRegressionTests.compute_mse(;
        job_name = "NCTest",
        reference_mse,
        ds_filename_computed = filenameA,
        ds_filename_reference = filenameB,
    )
    @test computed_mse["var"] == 0
    @test computed_mse["var2"] == 0

    NCRegressionTests.test_mse(computed_mse, reference_mse)

end

@testset "NCRegressionTests - Mismatch" begin

    sz = (123, 145)
    data = randn(MersenneTwister(152), sz)

    filenameA = tempname()
    filenameB = tempname()

    NCDataset(filenameA, "c") do ds
        defDim(ds, "lon", sz[1])
        defDim(ds, "lat", sz[2])
        v = defVar(ds, "var", Float64, ("lon", "lat"))
        v[:, :] = data .+ 1
        v = defVar(ds, "var2", Float64, ("lon", "lat"))
        v[:, :] = 2 * data .+ 2
    end
    NCDataset(filenameB, "c") do ds
        defDim(ds, "lon", sz[1])
        defDim(ds, "lat", sz[2])
        v = defVar(ds, "var", Float64, ("lon", "lat"))
        v[:, :] = data
        v = defVar(ds, "var2", Float64, ("lon", "lat"))
        v[:, :] = 2 * data
    end

    reference_mse = OrderedCollections.OrderedDict("var" => 2.8633275236008583e+04, "var2" => 0)

    computed_mse = NCRegressionTests.compute_mse(;
        job_name = "NCTest",
        reference_mse,
        ds_filename_computed = filenameA,
        ds_filename_reference = filenameB,
    )

    @test computed_mse["var"] ≈ 2.8633275236008583e+04
    @test computed_mse["var2"] ≈ 2.8633275236008583e+04

    reference_mse = OrderedCollections.OrderedDict("var" => 2.8633275236008583e+04, "var2" => 0, "no_key" => 1)

    err = ErrorException("No key no_key for mse computation.")
    @test_throws err NCRegressionTests.compute_mse(;
        job_name = "NCTest",
        reference_mse,
        ds_filename_computed = filenameA,
        ds_filename_reference = filenameB,
    )

end


using Test
using MPI

# TODO: enable CUDA tests once OpenMPI has full support
ArrayType = Array

MPI.Init()

comm = MPI.COMM_WORLD

const rank = MPI.Comm_rank(comm)
const N = MPI.Comm_size(comm)

buf = ArrayType{Int}(fill(Int(rank),N))
received = ArrayType{Int}(fill(-1,N))

# Test Get using fence
win = MPI.Win_create(buf, comm)
MPI.Win_fence(0, win)
MPI.Get!(received, (rank+1)%N, win)
MPI.Win_fence(0, win)

@test received == ArrayType{Int}(fill(Int(rank+1)%N, size(received)))

# Test locked window
if rank != 0
  MPI.Win_lock(MPI.LOCK_EXCLUSIVE, 0, 0, win)
  received[1] = rank
  MPI.Put!(view(received,1:1), 0, rank, win)
  MPI.Win_unlock(0, win)
else
  buf[1] = 0
end

MPI.Win_fence(0, win)

if rank == 0
  @test buf == collect(0:N-1)
end

MPI.Barrier(comm)

if rank == 1
  MPI.Win_lock(MPI.LOCK_EXCLUSIVE, 1, 0, win)
  fill!(buf,3)
  MPI.Win_unlock(1,win)
end

MPI.Barrier(comm)

if rank == 0
  MPI.Win_lock(MPI.LOCK_EXCLUSIVE, 0, 0, win)
  fill!(buf, 2)
  MPI.Win_unlock(0,win)
  MPI.Win_lock(MPI.LOCK_EXCLUSIVE, 1, 0, win)
  result = similar(buf)
  MPI.Get_accumulate!(buf, result, 1, 0, MPI.SUM, win)
  MPI.Win_unlock(1,win)
  @test all(result .== 3)
end

MPI.Barrier(comm)

if rank == 1
  MPI.Win_lock(MPI.LOCK_EXCLUSIVE, 1, 0, win)
  @test all(buf .== 5)
  fill!(buf,-2)
  MPI.Win_unlock(1,win)
  MPI.Win_lock(MPI.LOCK_EXCLUSIVE, 0, 0, win)
  MPI.Accumulate!(buf, 0, 0, MPI.SUM, win)
  MPI.Win_unlock(0,win)
  MPI.Win_lock(MPI.LOCK_EXCLUSIVE, 1, 0, win)
  fill!(buf,1)
  MPI.Win_unlock(1,win)
end

MPI.Barrier(comm)

if rank == 0
  MPI.Win_lock(MPI.LOCK_EXCLUSIVE, 0, 0, win)
  @test all(buf .== 0)
  MPI.Win_unlock(0,win)
end

MPI.Barrier(comm)
MPI.free(win)
MPI.Barrier(comm)

# Try a dynamic window
win = MPI.Win_create_dynamic(comm)
MPI.Win_attach!(win, buf)

address_buf = Vector{Cptrdiff_t}(undef, 1)
address_win = MPI.Win_create(address_buf, comm)

# Address for dynamic window is global instead of relative to the start of the buffer, so we need to share it
MPI.Win_lock(MPI.LOCK_EXCLUSIVE, rank, 0, address_win)
address_buf[1] = MPI.Get_address(buf)
MPI.Win_unlock(rank, address_win)

MPI.Barrier(comm)

# Test atomic operation
if rank == 0
  received = Ref(0)
  to_send = Ref(0)
  for r in 0:N-1
    address = Ref{Cptrdiff_t}(0)
    MPI.Win_lock(MPI.LOCK_EXCLUSIVE, r, 0, address_win)
    MPI.Get!(address, r, address_win)
    MPI.Win_flush(r, address_win)
    to_send[] = r+5
    MPI.Win_lock(MPI.LOCK_EXCLUSIVE, r, 0, win)
    MPI.Fetch_and_op!(to_send, received, r, address[]+r*sizeof(Int), MPI.REPLACE, win)
    MPI.Win_flush(r, win)
    @test received[] == r
    MPI.Win_unlock(r, win)
    MPI.Win_unlock(r, address_win)
  end
end

MPI.Barrier(comm)

@test buf[rank+1] == rank+5

MPI.Barrier(comm)
MPI.Win_detach!(win, buf)
MPI.free(win)
MPI.free(address_win)

MPI.Finalize()
@test MPI.Finalized()


using CUDA
using CUDA: cufunc, curand_rng
using CUDA.CUSPARSE: CuSparseMatrix, CuSparseMatrixCSR, mv!, switch2csr
using CUDA.CUSOLVER: CuQR

const CuFlatS0{P,T,M<:CuArray} = FlatS0{P,T,M}

# a function version of @cuda which can be referenced before CUDA is
# loaded as long as it exists by run-time (unlike the macro @cuda which must
# exist at compile-time)
function cuda(f, args...; threads=256)
    @cuda threads=threads f(args...)
end

is_gpu_backed(f::FlatField) = fieldinfo(f).M <: CuArray
global_rng_for(::Type{<:CuArray}) = curand_rng()
seed_for_storage!(::Type{<:CuArray}, seed=nothing) =
    Random.seed!(global_rng_for(CuArray), seed)

### broadcasting
preprocess(dest::F, bc::Broadcasted) where {F<:CuFlatS0} =
    Broadcasted{Nothing}(cufunc(bc.f), preprocess_args(dest, bc.args), map(OneTo,content_size(F)))
preprocess(dest::F, arg) where {M,F<:CuFlatS0{<:Any,<:Any,M}} =
    adapt(M,broadcastable(F, arg))
function copyto!(dest::F, bc::Broadcasted{Nothing}) where {F<:CuFlatS0}
    bc′ = preprocess(dest, bc)
    copyto!(firstfield(dest), bc′)
    return dest
end
BroadcastStyle(::FlatS0Style{F,Array}, ::FlatS0Style{F,CuArray}) where {P,F<:FlatS0{P}} =
    FlatS0Style{basetype(F){P},CuArray}()

### misc
# the generic versions of these trigger scalar indexing of CUDA, so provide
# specialized versions:

function copyto!(dst::F, src::F) where {F<:CuFlatS0}
    copyto!(firstfield(dst),firstfield(src))
    dst
end
pinv(D::Diagonal{T,<:CuFlatS0}) where {T} = Diagonal(@. ifelse(isfinite(inv(D.diag)), inv(D.diag), $zero(T)))
inv(D::Diagonal{T,<:CuFlatS0}) where {T} = any(Array((D.diag.==0)[:])) ? throw(SingularException(-1)) : Diagonal(inv.(D.diag))
fill!(f::CuFlatS0, x) = (fill!(firstfield(f),x); f)
==(a::CuFlatS0, b::CuFlatS0) = (==)(firstfield.(promote(a,b))...)
≈(a::CuFlatS0, b::CuFlatS0) = (≈)(firstfield.(promote(a,b))...)
≈(a::Diagonal{<:Any,<:CuFlatS0}, b::Diagonal{<:Any,<:CuFlatS0}) = basis(diag(a)) == basis(diag(b)) && diag(a) ≈ diag(b)
sum(f::CuFlatS0; dims=:) = (dims == :) ? sum(firstfield(f)) : (1 in dims) ? error("Sum over invalid dims of CuFlatS0.") : f

# these only work for Reals in CUDA
# with these definitions, they work for Complex as well
CUDA.isfinite(x::Complex) = Base.isfinite(x)
CUDA.sqrt(x::Complex) = CUDA.sqrt(CUDA.abs(x)) * CUDA.exp(im*CUDA.angle(x)/2)
CUDA.culiteral_pow(::typeof(^), x::Complex, ::Val{2}) = x * x

# this makes cu(::SparseMatrixCSC) return a CuSparseMatrixCSR rather than a
# dense CuArray
adapt_structure(::Type{<:CuArray}, L::SparseMatrixCSC) = CuSparseMatrixCSR(L)

# CUDA somehow missing this one
# see https://github.com/JuliaGPU/CuArrays.jl/issues/103
# and https://github.com/JuliaGPU/CuArrays.jl/pull/580
ldiv!(qr::CuQR, x::CuVector) = qr.R \ (CuMatrix(qr.Q)' * x)

# some Random API which CUDA doesn't implement yet
Random.randn(rng::CUDA.CURAND.RNG, T::Random.BitFloatType) =
    adapt(Array,randn!(rng, CuVector{T}(undef,1)))[1]

# perhaps minor type-piracy, but this lets us simulate into a CuArray using the
# CPU random number generator
Random.randn!(rng::MersenneTwister, A::CuArray{T}) where {T} =
    (A .= adapt(CuArray{T}, randn!(rng, adapt(Array{T},A))))

# CUDA makes some copies here as a workaround for JuliaGPU/CuArrays.jl#345 &
# NVIDIA/cuFFT#2714055 but it doesn't appear to be needed in the R2C case, and
# in the C2R case we pre-allocate the memory only once (via memoization) as well
# as doing the copy asynchronously
import CUDA.CUFFT: unsafe_execute!
using CUDA.CUFFT: rCuFFTPlan, cufftReal, cufftComplex, CUFFT_R2C, cufftExecR2C, cufftExecC2R, CUFFT_C2R, unsafe_copyto!, CuDefaultStream, pointer

plan_buffer(x) = plan_buffer(eltype(x),size(x))
@memoize plan_buffer(T,dims) = CuArray{T}(undef,dims...)

function unsafe_execute!(plan::rCuFFTPlan{cufftReal,K,false,N},
                            x::CuArray{cufftReal,N}, y::CuArray{cufftComplex,N}
                            ) where {K,N}
    @assert plan.xtype == CUFFT_R2C
    cufftExecR2C(plan, x, y)
end
function unsafe_execute!(plan::rCuFFTPlan{cufftComplex,K,false,N},
                            x::CuArray{cufftComplex,N}, y::CuArray{cufftReal}
                            ) where {K,N}
    @assert plan.xtype == CUFFT_C2R
    cufftExecC2R(plan, unsafe_copyto!(pointer(plan_buffer(x)),pointer(x),length(x),async=true,stream=CuDefaultStream()), y)
end

# monkey-patched version of https://github.com/JuliaGPU/CUDA.jl/pull/436
# until it hits a release
using CUDA.CURAND: curandSetPseudoRandomGeneratorSeed, curandSetGeneratorOffset,
    CURAND_STATUS_ALLOCATION_FAILED, CURAND_STATUS_PREEXISTING_FAILURE, CURAND_STATUS_SUCCESS,
    unsafe_curandGenerateSeeds, throw_api_error, @retry_reclaim, RNG

function Random.seed!(rng::RNG, seed=Base.rand(UInt64), offset=0)
    curandSetPseudoRandomGeneratorSeed(rng, seed)
    curandSetGeneratorOffset(rng, offset)
    res = @retry_reclaim err->isequal(err, CURAND_STATUS_ALLOCATION_FAILED) ||
                              isequal(err, CURAND_STATUS_PREEXISTING_FAILURE) begin
        unsafe_curandGenerateSeeds(rng)
    end
    if res != CURAND_STATUS_SUCCESS
        throw_api_error(res)
    end
    return
end

gc = () -> (GC.gc(true); CUDA.reclaim())

"""
    assign_GPU_workers()

Assuming you submitted a SLURM job and got several GPUs, possibly across several
nodes, this assigns each Julia worker process a unique GPU using `CUDA.device!`.
Assumes the SLURM variables `SLURM_STEP_GPUS` and `GPU_DEVICE_ORDINAL` are
defined on the workers.
"""
function assign_GPU_workers()
    @everywhere @eval using CUDA, Distributed
    topo = @eval Main pmap(workers()) do _
        hostname = gethostname()
        virtgpus = parse.(Int,split(ENV["GPU_DEVICE_ORDINAL"],","))
        if "SLURM_STEP_GPUS" in keys(ENV)
            physgpus = parse.(Int,split(ENV["SLURM_STEP_GPUS"],","))
        else
            @warn "SLURM_STEP_GPUS not defined, assign_GPU_workers may fail."
            # SLURM_STEP_GPUS seems not correctly set on all systems. this
            # will work if you requested a full node's worth of GPUs at least
            physgpus = virtgpus
        end
        if Set(virtgpus)!=Set(deviceid.(devices()))
            @warn "Virtual GPUs not same as CUDA.devices(), using latter"
            virtgpus = deviceid.(devices())
        end
        (i=myid(), hostname=hostname, virtgpus=virtgpus, physgpus=physgpus)
    end
    claimed = Set()
    assignments = Dict(map(topo) do (i,hostname,virtgpus,physgpus)
        for (virtgpu,physgpu) in zip(virtgpus,physgpus)
            if !((hostname,physgpu) in claimed)
                push!(claimed,(hostname,physgpu))
                return i => virtgpu
            end
        end
    end)
    @everywhere workers() device!($assignments[myid()])
    return topo, claimed, assignments
end


export judiModeling, judiDataModeling, judiDataSourceModeling, judiPointSourceModeling, judiJacobian

const adjoint_map = Dict(:forward => :adjoint, :adjoint => :forward, :born => :adjoint_born, :adjoint_born => :born)

# Base abstract types
abstract type judiPropagator{D, O} <: joAbstractLinearOperator{D, D} end
abstract type judiComposedPropagator{D, O} <: judiPropagator{D, O} end
abstract type judiAbstractJacobian{D, O, FT} <: judiComposedPropagator{D, O} end

############################################################################################################################
# Concrete types
struct judiModeling{D, O} <: judiPropagator{D, O}
    m::AbstractSize
    n::AbstractSize
    model::Model
    options::JUDIOptions
end

struct judiPointSourceModeling{D, O} <: judiComposedPropagator{D, O}
    m::AbstractSize
    n::AbstractSize
    F::judiModeling{D, O}
    qInjection::AdjointProjection{D}
    function judiPointSourceModeling{D, O}(F::judiModeling{D, O}, qInjection::AdjointProjection{D}) where {D, O}
        # instantiate un-initialized sizes
        ts = time_space_src(get_nsrc(qInjection), get_nt(qInjection), F.model.n)
        merge!(F.n, ts)
        merge!(F.m, ts)
        update_size(qInjection, F)
        new(F.m, qInjection.n, F, qInjection)
    end
end

struct judiDataModeling{D, O} <: judiComposedPropagator{D, O}
    m::AbstractSize
    n::AbstractSize
    rInterpolation::Projection{D}
    F::judiModeling{D, O}
    function judiDataModeling{D, O}(rInterpolation::Projection{D}, F::judiModeling{D, O}) where {D, O}
        ts = time_space_src(get_nsrc(rInterpolation), get_nt(rInterpolation), F.model.n)
        merge!(F.n, ts)
        merge!(F.m, ts)
        update_size(rInterpolation, F)
        new(rInterpolation.m, F.n, rInterpolation, F)
    end
end

struct judiDataSourceModeling{D, O} <: judiComposedPropagator{D, O}
    m::AbstractSize
    n::AbstractSize
    rInterpolation::Projection{D}
    F::judiModeling{D, O}
    qInjection::AdjointProjection{D}
    function judiDataSourceModeling{D, O}(rInterpolation::Projection{D}, F::judiModeling{D, O}, qInjection::AdjointProjection{D}) where {D, O}
        ts = time_space_src(get_nsrc(rInterpolation), get_nt(rInterpolation), F.model.n)
        merge!(F.n, ts)
        merge!(F.m, ts)
        update_size(rInterpolation, F)
        update_size(qInjection, F)
        new(rInterpolation.m, qInjection.n, rInterpolation, F, qInjection)
    end
end

struct judiJacobian{D, O, FT} <: judiAbstractJacobian{D, O, FT}
    m::AbstractSize
    n::AbstractSize
    F::FT
    q::judiMultiSourceVector
end

struct LazyScal
    s::Number
    P::judiPropagator
end

############################################################################################################################
# Base and JOLI compat
_get_property(J::judiPropagator, ::Val{:fop}) = x -> J*x
_get_property(J::judiPropagator, ::Val{:fop_T}) = x -> J'*x
_get_property(J::judiPropagator, ::Val{:name}) = "$(typeof(J))"
_get_property(J::judiPropagator, ::Val{s}) where {s} = getfield(J, s)

_get_property(C::judiComposedPropagator, ::Val{:model}) = C.F.model
_get_property(C::judiComposedPropagator, ::Val{:options}) = C.F.options

getproperty(J::judiPropagator, s::Symbol) = _get_property(J, Val{s}())

display(P::judiPropagator{D, O}) where {D, O} = println("JUDI $(String(O)){$D} propagator $(repr(P.n)) -> $(repr(P.m))")
show(io::Union{IOBuffer, IOContext}, P::judiPropagator{D, O}) where {D, O} = print(io, "JUDI $(String(O)){$D} propagator $(repr(P.n)) -> $(repr(P.m))")

############################################################################################################################
# Constructors
"""
    judiModeling(model; options=Options())
    judiModeling(model, src_geometry, rec_geometry; options=Options())

Create seismic modeling operator for a velocity model given as a `Model` structure.
The function also takes the source and receiver geometries
as additional input arguments, which creates a combined operator `judiProjection*judiModeling*judiProjection'`.

Example
=======
`Pr` and `Ps` are projection operatos of type `judiProjection` and
`q` is a data vector of type `judiVector`:
    F = judiModeling(model)
    dobs = Pr*F*Ps'*q
    F = judiModeling(model, q.geometry, rec_geometry)
    dobs = F*q
"""
function judiModeling(model::Model; options=Options())
    D = eltype(model.m)
    m = time_space(model.n)
    return judiModeling{D, :forward}(m, m, model, options)
end

judiModeling(model::Model, src_geom::Geometry, rec_geom::Geometry; options=Options()) =
    judiProjection(rec_geom) * judiModeling(model; options=options) * adjoint(judiProjection(src_geom))

"""
    judiJacobian(F,q)

Create a linearized modeling operator from the non-linear modeling operator `F` and
the source `q`. `q` can be either a judiVector (point source Jacobian) or a `judiWeight` for
extended source modeling.
`F` is a full modeling operator including source/receiver projections.

Examples
========
1) `F` is a modeling operator without source/receiver projections:
    J = judiJacobian(Pr*F*Ps',q)
2) `F` is the combined operator `Pr*F*Ps'`:
    J = judiJacobian(F,q)
"""
function judiJacobian(F::judiComposedPropagator{D, O}, q::judiMultiSourceVector; options=nothing) where {D, O}
    update!(F.F.options, options)
    return judiJacobian{D, :born, typeof(F)}(F.m, space(F.model.n), F, q)
end

# Backward compat with giving weights as array. Not recommened
function judiJacobian(F::judiComposedPropagator{D, O}, q::Array{D, N}; options=nothing) where {D, O, N}
    @warn "judiWeights is recommended for judiJacobian(F, weights)"
    nsrc = try length(F.qInjection.data) catch; length(F.rInterpolation.data) end
    update!(F.F.options, options)
    return judiJacobian(F, judiWeights(reshape(q, F.model.n); nsrc=nsrc))
end

############################################################################################################################
# Linear algebra (conj/adj/...)
# Adjoints
conj(F::judiPropagator) = F
transpose(F::judiPropagator) = adjoint(F)

adjoint(s::Symbol) = adjoint_map[s]
adjoint(F::judiModeling{D, O}) where {D, O} = judiModeling{D, adjoint(O)}(F.n, F.m, F.model, F.options)
adjoint(F::judiDataModeling{D, O}) where {D, O} = judiPointSourceModeling{D, adjoint(O)}(adjoint(F.F), adjoint(F.rInterpolation))
adjoint(F::judiPointSourceModeling{D, O}) where {D, O}= judiDataModeling{D, adjoint(O)}(adjoint(F.qInjection), adjoint(F.F))
adjoint(F::judiDataSourceModeling{D, O}) where {D, O} = judiDataSourceModeling{D, adjoint(O)}(adjoint(F.qInjection), adjoint(F.F), adjoint(F.rInterpolation))
adjoint(J::judiJacobian{D, O, FT}) where {D, O, FT} = judiJacobian{D, adjoint(O), FT}(J.n, J.m, J.F, J.q)
adjoint(L::LazyScal) = LazyScal(L.s, adjoint(L.P))

# Composition
*(F::judiModeling{D, O}, P::AdjointProjection{D}) where {D, O} = judiPointSourceModeling{D, O}(F, P)
*(P::Projection{D}, F::judiModeling{D, O}) where {D, O} = judiDataModeling{D, O}(P, F)
*(P::Projection{D}, F::judiPointSourceModeling{D, O}) where {D, O} = judiDataSourceModeling{D, O}(P, F.F, F.qInjection)
*(F::judiDataModeling{D, O}, P::AdjointProjection{D}) where {D, O} = judiDataSourceModeling{D, O}(F.rInterpolation, F.F, P)
*(s::Number, P::judiPropagator) = LazyScal(s, P)
\(P::judiPropagator, s::Number) = LazyScal(1/s, P)
*(L::LazyScal, x::SourceType{T}) where {T<:Number} = L.s * (L.P * x)
*(L::LazyScal, x::Array{T, N}) where {T<:Number, N} = L.s * (L.P * x)

# Propagation via linear algebra `*`
*(F::judiPropagator{T, O}, q::SourceType{T}) where {T<:Number, O} = multi_src_propagate(F, q)
*(F::judiAbstractJacobian{T, O, FT}, q::dmType{T}) where {T<:Number, O, FT} = multi_src_propagate(F, q)

mul!(out::SourceType{T}, F::judiPropagator{T, O}, q::SourceType{T}) where {T<:Number, O} = begin y = F*q; copyto!(out, y) end
mul!(out::SourceType{T}, F::judiPropagator{T, O}, q::Vector{T}) where {T<:Number, O} = begin y = F*q; copyto!(out, y) end
mul!(out::SourceType{T}, F::joLinearFunction{T, T}, q::SourceType{T}) where {T<:Number} = begin y = F*q; copyto!(out, y) end
mul!(out::SourceType{T}, F::judiAbstractJacobian{T, :born, FT}, q::Vector{T}) where {T<:Number, FT} = begin y = F*q[:]; copyto!(out, y) end
mul!(out::SourceType{T}, F::judiAbstractJacobian{T, :born, FT}, q::Array{T, N}) where {T<:Number, FT, N} = begin y = F*q[:]; copyto!(out, y) end
mul!(out::Array{T, 2}, F::judiAbstractJacobian{T, :adjoint_born, FT}, q::SourceType{T}) where {T<:Number, FT} =  begin y = F*q; copyto!(out, y) end
mul!(out::Array{T, 3}, F::judiAbstractJacobian{T, :adjoint_born, FT}, q::SourceType{T}) where {T<:Number, FT} =  begin y = F*q; copyto!(out, y) end

############################################################################################################################
# Propagation input
process_input_data(::judiPropagator, data::judiMultiSourceVector) = data
process_input_data(F::judiModeling, q::Vector) = process_input_data(q, F.model)
process_input_data(F::judiPointSourceModeling, q::Vector) = process_input_data(q, F.qInjection.geometry)
process_input_data(F::judiDataSourceModeling, q::Vector) = process_input_data(q, F.qInjection.data)
process_input_data(F::judiDataModeling, q::Vector) = process_input_data(q, F.model)
process_input_data(J::judiJacobian{D, :adjoint_born, FT}, q::Vector) where {D, FT} =
    process_input_data(q, J.F.qInjection.data)
process_input_data(::judiJacobian{D, :born, FT}, q::dmType{D}) where {D<:Number, FT} = q

make_input(::judiModeling, q::SourceType) = (nothing, make_input(q), nothing, nothing, nothing)
make_input(::judiModeling, rhs::judiRHS) = (make_src(rhs)..., nothing, nothing, nothing)
make_input(F::judiPointSourceModeling, q::SourceType{T}) where {T} = (make_src(q, F.qInjection)..., nothing, nothing, nothing)
make_input(F::judiPointSourceModeling, q::Matrix{T}) where {T} = (F.qInjection.data[1], q, nothing, nothing, nothing)
make_input(F::judiDataModeling, q::SourceType{T}) where {T} = (nothing, make_input(q), F.rInterpolation.geometry, nothing, nothing)
make_input(F::judiDataModeling{T, O}, q::LazyAdd{T}) where {T, O} = (make_src(q)..., F.rInterpolation.geometry, nothing, nothing)
make_input(F::judiDataModeling, rhs::judiRHS) = (make_src(rhs)..., F.rInterpolation.geometry, nothing, nothing)

make_input(F::judiDataSourceModeling, q::SourceType{T}) where {T} = (make_src(q, F.qInjection)..., F.rInterpolation.data[1], nothing, nothing)
make_input(F::judiDataSourceModeling, q::Matrix{T}) where {T} = (F.qInjection.data[1], q, F.rInterpolation.data[1], nothing, nothing)

function make_input(J::judiJacobian{D, :born, FT}, q::dmType{D}) where {D<:Number, FT}
    srcGeom, srcData = make_src(J.q, J.F.qInjection)
    return srcGeom, srcData, J.F.rInterpolation.data[1], nothing, reshape(q, J.model.n)
end

function make_input(J::judiJacobian{D, :adjoint_born, FT}, q::SourceType{D}) where {D, FT}
    srcGeom, srcData = make_src(J.q, J.F.qInjection)
    recGeom, recData = make_src(q, J.F.rInterpolation)
    return srcGeom, srcData, recGeom, recData, nothing
end

############################################################################################################################
# Size update based on linear operator

update_size(w::judiProjection, F::judiPropagator) = set_space_size!(w.n, F.model.n)
update_size(w::jAdjoint{<:judiProjection}, F::judiPropagator) = set_space_size!(w.op.n, F.model.n)

update_size(w::judiWavelet, F::judiPropagator) = set_space_size!(w.m, F.model.n)
update_size(w::jAdjoint{<:judiWavelet}, F::judiPropagator) = set_space_size!(w.op.m, F.model.n)
############################################################################################################################
# indexing
getindex(F::judiModeling{D, O}, i) where {D, O} = judiModeling{D, O}(F.m[i], F.n[i], F.model, F.options[i])
getindex(F::judiDataModeling{D, O}, i) where {D, O} = judiDataModeling{D, O}(F.rInterpolation[i], F.F[i])
getindex(F::judiPointSourceModeling{D, O}, i) where {D, O}= judiPointSourceModeling{D, O}(F.F[i], F.qInjection[i])
getindex(F::judiDataSourceModeling{D, O}, i) where {D, O} = judiDataSourceModeling{D, O}(F.rInterpolation[i], F.F[i], F.qInjection[i])
getindex(J::judiJacobian{D, O, FT}, i) where {D, O, FT} = judiJacobian{D, O, FT}(J.m[i], J.n[i], J.F[i], J.q[i])

############################################################################################################################
# Comparisons
isequal(P1::judiPropagator, P2::judiPropagator) = P1 == P2
==(F1::judiModeling{D, O1}, F2::judiModeling{D, O2}) where {D, O1, O2} = (O1 == O2 && F1.model == F2.model && F1.options == F2.options)
==(F1::judiDataModeling, F2::judiDataModeling) = (F1.F == F2.F && F1.rInterpolation == F2.rInterpolation)
==(F1::judiPointSourceModeling, F2::judiPointSourceModeling) = (F1.F == F2.F && F1.qInjection == F2.qInjection)
==(F1::judiDataSourceModeling, F2::judiDataSourceModeling) = (F1.F == F2.F && F1.qInjection == F2.qInjection && F1.rInterpolation == F2.rInterpolation)
==(F1::judiJacobian{D, O1, FT1}, F2::judiJacobian{D, O2, FT2}) where {D, O1, O2, FT1, FT2} = (O1 == O2 && FT1 == FT2 && F1.F == F2.F && F1.q == F2.q)


# Fallbacks!
NLᵁ(m, i) = 0
NLⱽ(m, i) = 0
NLᵀ(m, i) = 0
NLˢ(m, i) = 0
NLᵉ(m, i) = 0

∂z_NLᵁ(m, i) = 0
∂z_NLⱽ(m, i) = 0
∂z_NLᵀ(m, i) = 0
∂z_NLˢ(m, i) = 0
∂z_NLᵉ(m, i) = 0

#####
##### Counter gradient flux model proposed by Large et al (1994)
#####

Base.@kwdef struct CounterGradientFlux{T} <: AbstractParameters
    Cᴺ :: T = 1.0 # Mass flux proportionality constant
end

const CGModel = Model{L, K, W, <:CounterGradientFlux} where {L, K, W}

#∂z_NLᵀ(m::CGModel, i) where K = ∂NL∂z(m.nonlocal_flux.Cᴺ, m.state.Qθ, m)
#∂z_NLˢ(m::CGModel, i) where K = ∂NL∂z(m.nonlocal_flux.Cᴺ, m.state.Qs, m)

#####
##### Diagnostic plume model
#####

abstract type AbstractDiagnosticPlumeModel <: AbstractParameters end

Base.@kwdef struct WitekDiagnosticPlumeModel{T} <: AbstractDiagnosticPlumeModel
     Ca :: T = 0.1
    Cbw :: T = 2.86
     Ce :: T = 0.1
    Cew :: T = 0.572
     CQ :: T = 1.0
end

const ModelWithPlumes = Model{L, K, W, <:AbstractDiagnosticPlumeModel} where {L, K, W}

instantiate_plume(::AbstractDiagnosticPlumeModel, grid) =
    (T=CellField(grid), S=CellField(grid), W²=FaceField(grid), e=nothing)

#####
##### Entrainment
#####

"""
    dynamic_entrainment(Ce, ΔB̆, W̆²)

Returns

``C^e * min(0, ΔB̆ / W̆²) ``,

which is always less than zero.
"""
@inline dynamic_entrainment(Ce, ΔB̆, W̆²::T) where T =
    ifelse(W̆²==0, zero(T), Ce * min(zero(T), ΔB̆ / W̆²))

@inline function tracer_entrainment(m::ModelWithPlumes, i)
    ΔB̆ = plume_buoyancy_excess(i, m.grid, m)
    W̆² = oncell(m.state.plume.W², i)
    return dynamic_entrainment(m.nonlocal_flux.Ce, ΔB̆, W̆²)
end

@inline function momentum_entrainment(m::ModelWithPlumes, i)
    ΔB̆ = onface(i, m.grid, plume_buoyancy_excess, m)
    W̆² = @inbounds m.state.plume.W²[i]
    return dynamic_entrainment(m.nonlocal_flux.Ce, ΔB̆, W̆²)
end

#####
##### Plume boundary conditions
#####

function set_tracer_plume_bc!(ϕ̆, Φ, Qϕ, CQ, model)
    n = ϕ̆.grid.N

    v★ = maxsqrt(model.solution.e, n) # turbulent velocity scale at the surface

    # Surface layer model: √e Δϕ̆ = - C Qϕ, where Δϕ̆ is plume excess.
    # Also, plumes can't exist if v★=0. This makes no sense, but our best option atm.
    @inbounds ϕ̆[n] = ifelse(v★==0, Φ[n], Φ[n] - CQ * Qϕ / v★)

    return nothing
end

function set_vertical_momentum_plume_bc!(W)
    @inbounds W[W.grid.N+1] = 0
    return nothing
end

#####
##### Plume equations
#####

#####
##### Buoyancy
#####

@inline plume_buoyancy_excess(i, grid, T̆, S̆, T, S, α, β, g) =
    @inbounds g * (α*(T̆[i] - T[i]) - β*(S̆[i] - S[i]))

@inline plume_buoyancy_excess(i, grid, model) =
    plume_buoyancy_excess(i, grid, model.state.plume.T, model.state.plume.S,
                                   model.solution.T, model.solution.S,
                                   model.constants.α, model.constants.β, model.constants.g)

function clip_infinite!(ϕ)
    for i in eachindex(ϕ)
        @inbounds begin
            ϕᵢ = ϕ[i]
            if !(isfinite(ϕᵢ))
                 ϕ[i] = 0
            end
        end
    end
    return nothing
end

function clip_positive!(ϕ)
    for i in eachindex(ϕ)
        @inbounds begin
            ϕᵢ = ϕ[i]
            if ϕᵢ > 0
                ϕ[i] = 0
            end
        end
    end
    return nothing
end

function clip_negative!(ϕ)
    for i in eachindex(ϕ)
        @inbounds begin
            ϕᵢ = ϕ[i]
            if ϕᵢ < 0
                ϕ[i] = 0
            end
        end
    end
    return nothing
end

function update_nonlocal_flux!(model::ModelWithPlumes) where K

    set_tracer_plume_bc!(model.state.plume.T, model.solution.T,
                         model.state.Qθ, model.nonlocal_flux.CQ, model)

    set_tracer_plume_bc!(model.state.plume.S, model.solution.S,
                         model.state.Qs, model.nonlocal_flux.CQ, model)

    set_vertical_momentum_plume_bc!(model.state.plume.W²)

    integrate_plume_equations!(model.state.plume.T, model.state.plume.S, model.state.plume.W²,
                               model.solution.T, model.solution.S, model.grid, model)

    clip_infinite!(model.state.plume.W²)
    clip_negative!(model.state.plume.W²)
    clip_infinite!(model.state.plume.T)

    return nothing
end

function integrate_plume_equations!(T̆, S̆, W̆², T, S, grid, model)

    n = grid.N

    # Vertical momentum at the nᵗʰ cell interface, approximating excess buoyancy at
    # interface with excess at top cell center:
    ΔB̆ᵢ₊₁ = plume_buoyancy_excess(n, grid, model)
    @inbounds W̆²[n] = -model.nonlocal_flux.Cbw * Δf(grid, n) * ΔB̆ᵢ₊₁

    # Integrate from surface cell `N-1` downwards
    for i in n-1 : -1 : 1
        if W̆²[i+1] <= 0 # plume is stopped
            @inbounds W̆²[i] = 0
            @inbounds T̆[i] = T[i]
            @inbounds S̆[i] = S[i]
        else # plume still lives

            # Integrate temperature and salinity
            @inbounds T̆[i] = T̆[i+1] + Δc(grid, i+1) * tracer_entrainment(model, i) * (T̆[i+1] - T[i+1])
            @inbounds S̆[i] = S̆[i+1] + Δc(grid, i+1) * tracer_entrainment(model, i) * (S̆[i+1] - S[i+1])

            # Integrate vertical momentum
            ΔB̆ᵢ₊₁ = onface(i+1, grid, plume_buoyancy_excess, model)

            @inbounds W̆²[i] = W̆²[i+1] - Δf(grid, i+1) * (
                                model.nonlocal_flux.Cbw * ΔB̆ᵢ₊₁
                              - model.nonlocal_flux.Cew * momentum_entrainment(model, i) * W̆²[i+1])
        end
    end

    return nothing
end

#####
##### Mass flux contribution to tracer budgets
#####

maxzero(ϕ::T) where T = max(zero(T), ϕ)

@inline M(m, i) = @inbounds -m.nonlocal_flux.Ca * maxsqrt(m.state.plume.W², i)

"Returns the non-local flux of temperature at cell centers."
@inline function NLᵀ(m::ModelWithPlumes, i)
    Mᵢ = oncell(M, m, i)
    ΔTᵢ = @inbounds m.state.plume.T[i] - m.solution.T[i]
    return Mᵢ * ΔTᵢ
end

"Returns the non-local flux of salinity at cell centers."
@inline function NLˢ(m::ModelWithPlumes, i)
    Mᵢ = oncell(M, m, i)
    ΔSᵢ = @inbounds m.state.plume.S[i] - m.solution.S[i]
    return Mᵢ * ΔSᵢ
end

"Returns the temperature transport for a plume model, with mass flux on cell interfaces."
@inline ∂z_NLᵀ(m::ModelWithPlumes, i) = 1 / (2 * Δc(m.grid, i+1)) * (NLᵀ(m, i+1) - NLᵀ(m, i))

"Returns the salinity transport for a plume model, with mass flux on cell interfaces."
@inline ∂z_NLˢ(m::ModelWithPlumes, i) = 1 / (2 * Δc(m.grid, i+1)) * (NLˢ(m, i+1) - NLˢ(m, i))


abstract type AbstractBoundaryCondition end

struct DefaultBC <: AbstractBoundaryCondition end

Base.@kwdef struct BulkFormulaTemperature{𝒯, 𝒰, 𝒱} <: AbstractBoundaryCondition
    drag_coef_temperature::𝒯
    drag_coef_moisture::𝒰
    surface_temperature::𝒱
end

Base.@kwdef struct CoupledPrimarySlabOceanBC{𝒯, 𝒰} <: AbstractBoundaryCondition
    drag_coef_temperature::𝒯
    drag_coef_moisture::𝒰
end

abstract type TemperatureBC end
struct Insulating <: TemperatureBC end
struct CoupledSecondaryAtmosModelBC <: TemperatureBC end

function numerical_boundary_flux_first_order!(
    numerical_flux::NumericalFluxFirstOrder,
    ::DefaultBC,
    balance_law::DryAtmosModel,
    fluxᵀn::Vars{S},
    n̂::SVector,
    state⁻::Vars{S},
    aux⁻::Vars{A},
    state⁺::Vars{S},
    aux⁺::Vars{A},
    t,
    direction,
    state1⁻::Vars{S},
    aux1⁻::Vars{A},
) where {S, A}
    state⁺.ρ = state⁻.ρ
    state⁺.ρe = state⁻.ρe
    state⁺.ρq = state⁻.ρq

    # project and reflect for impenetrable condition, but
    # leave tangential component untouched
    ρu⁻ = state⁻.ρu
    state⁺.ρu = ρu⁻ - 2n̂ ⋅ ρu⁻ .* SVector(n̂)

    numerical_flux_first_order!(numerical_flux, balance_law, fluxᵀn, n̂, state⁻, aux⁻, state⁺, aux⁺, t, direction)
end

function numerical_boundary_flux_first_order!(
    numerical_flux::NumericalFluxFirstOrder,
    bctype::BulkFormulaTemperature,
    model::DryAtmosModel,
    fluxᵀn::Vars{S},
    n̂::SVector,
    state⁻::Vars{S},
    aux⁻::Vars{A},
    state⁺::Vars{S},
    aux⁺::Vars{A},
    t,
    direction,
    state1⁻::Vars{S},
    aux1⁻::Vars{A},
) where {S, A}
    # Impenetrable free-slip condition to reflect and project momentum
    # at the boundary
    numerical_boundary_flux_first_order!(
        numerical_flux,
        DefaultBC(),
        model,
        fluxᵀn,
        n̂,
        state⁻,
        aux⁻,
        state⁺,
        aux⁺,
        t,
        direction,
        state1⁻,
        aux1⁻,
    )

    E, H = calc_ocean_sfc_fluxes(model.physics, bctype, state⁻, aux⁻)
    LH_v0 = model.physics.parameters.LH_v0

    fluxᵀn.ρ -= E / LH_v0
    fluxᵀn.ρe -= E + H
    fluxᵀn.ρq -= E / LH_v0
end

function numerical_boundary_flux_first_order!(
    numerical_flux::NumericalFluxFirstOrder,
    bctype::CoupledPrimarySlabOceanBC,
    model::DryAtmosModel,
    fluxᵀn::Vars{S},
    n̂::SVector,
    state⁻::Vars{S},
    aux⁻::Vars{A},
    state⁺::Vars{S},
    aux⁺::Vars{A},
    t,
    direction,
    state1⁻::Vars{S},
    aux1⁻::Vars{A},
) where {S, A}
    # Impenetrable free-slip condition to reflect and project momentum
    # at the boundary

    numerical_boundary_flux_first_order!(
        numerical_flux,
        DefaultBC(),
        model,
        fluxᵀn,
        n̂,
        state⁻,
        aux⁻,
        state⁺,
        aux⁺,
        t,
        direction,
        state1⁻,
        aux1⁻,
    )

    # The following will be moved the the second-order kernel (as a Neumann BC) once available
    LH_v0 = model.physics.parameters.LH_v0
    E, H = calc_ocean_sfc_fluxes(model.physics, bctype, state⁻, aux⁻) #[W/m^2]

    fluxᵀn.ρ -= E / LH_v0
    fluxᵀn.ρe -= E + H
    fluxᵀn.ρq -= E / LH_v0

end


"""
`module  JAC.PairAnnihilation1Photon`  ... a submodel of JAC that contains all methods for computing positron-bound-electron pair annihilation
                                           (PEPA) with two-photon emission cross sections and rates; e^+ + |i(N)> --> |f(N-1)> + photon + photon;
                                           it is using JAC, JAC.ManyElectron.
"""
module PairAnnihilation2Photon

    using JAC, JAC.ManyElectron

    """
    `struct  PairAnnihilation2Photon.Settings`  ... defines a type for the details and parameters in computing positron-bound-electron pair
             annihilation (PEPA) with single-photon emission cross sections and rates; e^+ + |i(N)> --> |f(N-1)> + photon1 + photon2.

        + multipoles              ::Array{EmMultipole}           ... Specifies the multipoles of the radiation field that are to be included.
        + gauges                  ::Array{UseGauge}              ... Specifies the gauges to be included into the computations.
        + positronEnergies        ::Array{Float64,1}             ... List of positron energies.
        + printBeforeComputation  ::Bool                         ... True, if all energies and lines are printed before their evaluation.
        + selectLines             ::Bool                         ... True, if lines are selected individually for the computations.
        + selectedLines           ::Array{Tuple{Int64,Int64},1}  ... List of lines, given by tupels (inital-level, final-level).
    """
    struct Settings
        multipoles                ::Array{EmMultipole}
        gauges                    ::Array{UseGauge}
        positronEnergies          ::Array{Float64,1}
        printBeforeComputation    ::Bool
        selectLines               ::Bool
        selectedLines             ::Array{Tuple{Int64,Int64},1}
    end

    """
    `JAC.PairAnnihilation2Photon.Settings()`  ... constructor for the default values of pair-annihilation photon line computations
    """
    function Settings()
        Settings(EmMultipole[], UseGauge[], Float64[], false, false, Tuple{Int64,Int64}[])
    end

    # `Base.show(io::IO, settings::PairAnnihilation2Photon.Settings)`
    #		... prepares a proper printout of the variable settings::PairAnnihilation2Photon.Settings.
    function Base.show(io::IO, settings::PairAnnihilation2Photon.Settings)
        println(io, "multipoles:               $(settings.multipoles)  ")
        println(io, "gauges:                   $(settings.gauges)  ")
        println(io, "positronEnergies:         $(settings.positronEnergies)  ")
        println(io, "printBeforeComputation:   $(settings.printBeforeComputation)  ")
        println(io, "selectLines:              $(settings.selectLines)  ")
        println(io, "selectedLines:            $(settings.selectedLines)  ")
    end

    """
    `struct  PairAnnihilation2Photon.Channel`  ... defines a type for a positron-bound-electron pair annihilation (PEPA) with two-photon
                                                   emission channel that specifies all quantum numbers, phases and amplitudes.

        + multipole1     ::EmMultipole          ... Multipole of the emitted photon1.
        + multipole2     ::EmMultipole          ... Multipole of the emitted photon2.
        + gauge          ::EmGauge              ... Gauge for dealing with the (coupled) radiation field.
        + kappa          ::Int64                ... partial-wave of the incoming free positron
        + symmetry       ::LevelSymmetry        ... total angular momentum and parity of the scattering state
        + phase          ::Float64              ... phase of the partial wave
        + amplitude      ::Complex{Float64}     ... PairAnnihilation2PhotonChannel amplitude associated with the given channel.
   """
    struct  Channel
        multipole1       ::EmMultipole
        multipole2       ::EmMultipole
        gauge            ::EmGauge
        kappa            ::Int64
        symmetry         ::LevelSymmetry
        phase            ::Float64
        amplitude        ::Complex{Float64}
    end

    """
    `struct  PairAnnihilation2Photon.Line`  ... defines a type for a positron-bound-electron pair-annihilation (two-photon) line that may
                                                include the definition of channels.

        + initialLevel   ::Level                 ... initial-(state) level
        + finalLevel     ::Level                 ... final-(state) level
        + positronEnergy ::Float64                ... Energy of the (incoming free) positron.
        + photon1Energy  ::Float64                ... Energy of the emitted photon1.
        + photon2Energy  ::Float64                ... Energy of the emitted photon2.
        + crossSection   ::EmProperty             ... Cross section for this pair-annihilation (photon) line.
        + hasChannels    ::Bool                   ... Determines whether the individual (sub-) channels are defined in terms of their
                                                      free-positron energy, kappa, multipole, etc., or not.
        + channels       ::Array{PairAnnihilation2Photon.Channel,1}  ... List of PairAnnihilation2Photon.Channels of this line.
    """
    struct  Line
        initialLevel     ::Level
        finalLevel       ::Level
        positronEnergy   ::Float64
        photon1Energy    ::Float64
        photon2Energy    ::Float64
        crossSection     ::EmProperty
        hasChannels      ::Bool
        channels         ::Array{PairAnnihilation2Photon.Channel,1}
    end

    """
    `JAC.PairAnnihilation2Photon.Line()`  ... 'empty' constructor for a pair-annihilation (photon) line between a specified initial and
                                              final level.
    """
    function Line()
        Line(Level(), Level(), 0., 0., 0., EmProperty(0., 0.), false, PairAnnihilation2Photon[] )
    end

    # `Base.show(io::IO, line::PairAnnihilation2Photon.Line)`  ... prepares a proper printout of the variable line::PairAnnihilation2Photon.Line.
    function Base.show(io::IO, line::PairAnnihilation1Photon.Line)
        println(io, "initialLevel:      $(line.initialLevel)  ")
        println(io, "finalLevel:        $(line.finalLevel)  ")
        println(io, "positronEnergy:    $(line.positronEnergy)  ")
        println(io, "photon1Energy:     $(line.photon1Energy)  ")
        println(io, "photon2Energy:     $(line.photon2Energy)  ")
        println(io, "crossSection:      $(line.crossSection)  ")
        println(io, "hasChannels:       $(line.hasChannels)  ")
        println(io, "channels:          $(line.channels)  ")
    end

end # module


# initialisation functions for GeMM

"""
    createpop(cellsize)

Create a new, random individual and replicates it a certain number of times
(depending on metabolic variables) to create a new population of organisms.
Returns an array of individuals.
"""
function createpop(cellsize::Float64)
    # start with an archetypical individual
    archetype = createind(true)
    population = [archetype]
    # calculate the population size
    popsize = 0
    if occursin("metabolic", setting("popsize")) || occursin("single", setting("popsize"))
        # population size determined by adult size
        popsize = round(setting("fertility") * archetype.traits["repsize"] ^ (-1 / 4) *
                        exp(-act / (boltz * 298.0)))
    elseif occursin("bodysize", setting("popsize"))
        # population size up to 25% of the maximum possible in this cell
        quarterpopsize = Integer(ceil(cellsize / archetype.traits["repsize"]) / 4)
        popsize = rand(1:quarterpopsize) + 1
    elseif occursin("minimal", setting("popsize")) || popsize == 0
        popsize = 2 #Takes two to tangle ;-)
    else
        simlog("Invalid value for `popsize`: $(setting("popsize"))", 'e')
    end
    # clone & mutate the archetype N times
    locivar = rand()
    for i in 2:popsize
        #TODO avoid `deepcopy` here
        ind = deepcopy(archetype) # XXX this is potentially very slow!
        ind.id = rand(UInt64)
        varyalleles!(ind.genome, locivar)
        ind.traits = gettraitdict(ind.genome, setting("traitnames"))
        if !(setting("indsize") == "adult" || setting("indsize") == "seed")
            #XXX sizes uniformally distributed?
            ind.size = ind.traits["seedsize"] + rand() * (ind.traits["repsize"] - ind.traits["seedsize"])
        end
        push!(population, ind)
    end
    population
end

"""
    createind(marked=false)

Create an individual organism of a new species with a random genome.
"""
function createind(marked::Bool = false)
    id = rand(UInt64)
    lineage = randstring(6)
    ngenes = rand(1:setting("maxloci")) * length(setting("traitnames"))
    setting("degpleiotropy") == 0 && (ngenes = length(setting("traitnames")))
    ngenes < 1 && (simlog("0 genes would be created", 'e')) #shouldn't happen
    traits = createtraits()
    genes = creategenes(ngenes, traits)
    if setting("linkage") == "none"
        nchrms = length(genes)
    elseif setting("linkage") == "full"
        nchrms = 1
    else
        linkage = rand(1:length(genes))
        nchrms = Integer(round(ngenes/linkage))
    end
    chromosomes = createchrms(nchrms, genes, lineage)
    varyalleles!(chromosomes, rand())
    traitdict = gettraitdict(chromosomes, setting("traitnames"))
    if setting("indsize") == "adult"
        indsize = traitdict["repsize"]
    elseif setting("indsize") == "seed"
        indsize = traitdict["seedsize"]
    else
        indsize = traitdict["seedsize"] + rand() * (traitdict["repsize"] - traitdict["seedsize"])
    end
    Individual(lineage, chromosomes, traitdict, marked, 1.0, 1.0, indsize, hermaphrodite, 0, id)
end

"""
    genesis(cellsize)

Create a new community, composed of random new species populations, for a patch.
Returns an array of individuals.
"""
function genesis(cellsize::Float64)
    community = Individual[]
    totalmass = 0.0
    while true
        population = createpop(cellsize)
        popsize = length(population)
        # Check the cell capacity
        popmass = sum(map(x -> x.size, population))
        if totalmass + popmass > cellsize * setting("overfill")
            # stop loop if cell is full
            if totalmass >= cellsize * 0.9 || occursin("single", setting("popsize"))
                #make sure the cell is full enough
                @simlog("Cell is now $(round((totalmass/cellsize)*100))% full.", 'd')
                break
            else
                continue
            end
        end
        totalmass += popmass
        append!(community, population)
        occursin("single", setting("popsize")) && break
    end
    @simlog("Patch initialized with $(length(community)) individuals.", 'd')
    community
end

"""
    createpatch(patchentry)

Create a new patch and set its parameters by parsing one line of the map file.
"""
function createpatch(patchentry::Array{String,1})
    size(patchentry,1) < 3 && simlog("please check your map file for incomplete or faulty entries. \n
Each line must contain patch information with at least \n
\t - a unique integer ID, \n
\t - an integer x coordinate, \n
\t - an integer y coordinate, \n
separated by a whitespace character (<ID> <x> <y>).", 'e')
    # create the basic patch
    id = parse(Int, patchentry[1])
    xcord = parse(Int, patchentry[2])
    ycord = parse(Int, patchentry[3])
    defaultcapacity = setting("cellsize")
    # XXX the 'global' here is a hack so that I can use eval() later on
    # (eval() always works on the global scope)
    global newpatch = Patch(id, (xcord, ycord), defaultcapacity)
    @simlog("Creating patch $id at $xcord/$ycord", 'd')
    # parse other parameter options
    for p in patchentry[4:end]
        varval = split(p, '=')
        var = varval[1]
        if !(var in map(string, fieldnames(Patch)))
            simlog("Unrecognized patch parameter $var.", 'w')
            continue
        elseif length(varval) < 2
            # if no value is specified, assume 'true'
            # (unless the variable is prepended with an exclamation mark, then take 'false')
            if var[1] == '!'
                var = var[2:end]
                val = false
            else
                val = true
            end
        else
            val = Meta.parse(varval[2])
        end
        # check for correct type and modify the new patch
        vartype = typeof(eval(Meta.parse("newpatch."*var)))
        if !isa(val, vartype)
            try
                val = convert(vartype, val)
            catch
                simlog("Invalid patch parameter type $var: $val", 'w')
                continue
            end
        end
        eval(Meta.parse("newpatch."*string(var)*" = $val"))
    end
    newpatch
end

"""
    createworld(maptable)

Use a parsed map file (as returned by `readmapfile`) to create the world.
Initialises each patch with its parameters and a new community, then returns
an array of patches.
"""
function createworld(maptable::Array{Array{String,1},1})
    simlog("Creating world...")
    world = Array{Patch}(undef, length(maptable))
    for entry in eachindex(maptable)
        newpatch = createpatch(maptable[entry])
        if newpatch.initpop
            if setting("mode") == "zosterops"
                append!(newpatch.community, zgenesis(newpatch))
            elseif setting("indsize") != "seed"
                append!(newpatch.community, genesis(newpatch.capacity))
            else
                append!(newpatch.seedbank, genesis(newpatch.capacity))
            end
        end
        world[entry] = newpatch
    end
    global newpatch = nothing # remove variable used in `createpatch()`
    (setting("mode") == "zosterops") && findneighbours!(world)
    world
end

"""
    updateworld(world, maptable, cellsize)

Reinitialise the world from another parsed map file. Works analogously to
`createworld`. Intended for use in scenarios where the model world changes
during a run (e.g. through global warming or island ontogeny).
If the ID of an old patch matches the ID of a new patch, its community and
seed bank are transferred to the new patch.
"""
function updateworld(oldworld::Array{Patch,1}, maptable::Array{Array{String,1},1})
    simlog("Updating world...")
    newworld = Array{Patch}(undef, length(maptable))
    for entry in eachindex(maptable)
        newpatch = createpatch(maptable[entry])
        updates = findfirst(p -> p.id==newpatch.id, oldworld)
        if !isnothing(updates) # did this patch exist in the old world?
            newpatch.community = oldworld[updates].community
            newpatch.seedbank = oldworld[updates].seedbank
        end
        newworld[entry] = newpatch
    end
    global newpatch = nothing # remove variable used in `createpatch()`
    (setting("mode") == "zosterops") && findneighbours!(newworld)
    return newworld
end


using Random
using CUDAdrv
using CUDAnative
using CuArrays
using StaticArrays
using GPUifyLoops # for @unroll macro

# Base.@propagate_inbounds function Base.getindex(v::MArray, i::Int)
#     @boundscheck checkbounds(v,i)
#     T = eltype(v)
#
#     @assert isbitstype(T) "Only isbits is supported on the device"
#     return GC.@preserve v begin
#         ptr  = pointer_from_objref(v)
#         dptr = reinterpret(CUDAnative.DevicePtr{T, CUDAnative.AS.Local}, ptr)
#         unsafe_load(dptr, i)
#     end
# end
#
# Base.@propagate_inbounds function Base.setindex!(v::MArray, val, i::Int)
#     @boundscheck checkbounds(v,i)
#     T = eltype(v)
#
#     @assert isbitstype(T) "Only isbits is supported on the device"
#     GC.@preserve v begin
#     ptr  = pointer_from_objref(v)
#         dptr = reinterpret(CUDAnative.DevicePtr{T, CUDAnative.AS.Local}, ptr)
#         unsafe_store!(dptr, convert(T, val), i)
#     end
#
#     return val
# end

for (f, T) in Base.Iterators.product((:add, :mul, :sub), (Float32, Float64))
    name = Symbol("$(f)_float_contract")
    if T === Float32
        llvmt = "float"
    elseif T === Float64
        llvmt = "double"
    end

    ir = """
    %x = f$f contract $llvmt %0, %1
    ret $llvmt %x
    """
    @eval begin
        # the @pure is necessary so that we can constant propagate.
        Base.@pure function $name(a::$T, b::$T)
            @Base._inline_meta
            Base.llvmcall($ir, $T, Tuple{$T, $T}, a, b)
        end
    end
end

# note the order of the fields below is also assumed in the code.
const _nstate = 5
const _ρ, _U, _V, _W, _E = 1:_nstate

const _nvgeo = 14
const _ξx, _ηx, _ζx, _ξy, _ηy, _ζy, _ξz, _ηz, _ζz, _MJ, _MJI,
       _x, _y, _z = 1:_nvgeo

Base.@irrational gdm1 0.4 BigFloat(0.4)

function volumerhs!(::Val{N}, rhs, Q, vgeo, D, nelem, t) where {N}
  Nq = N + 1

  s_D = @cuStaticSharedMem eltype(D) (Nq, Nq)
  s_F = @cuStaticSharedMem eltype(Q) (Nq, Nq, Nq, _nstate)
  s_G = @cuStaticSharedMem eltype(Q) (Nq, Nq, Nq, _nstate)
  s_H = @cuStaticSharedMem eltype(Q) (Nq, Nq, Nq, _nstate)

  e = blockIdx().x
  k = threadIdx().z
  j = threadIdx().y
  i = threadIdx().x

  @inbounds begin
    # fetch D into shared
    k == 1 && (s_D[i, j] = D[i, j])

    # Load values will need into registers
    MJ = vgeo[i, j, k, _MJ, e]
    ξx, ξy, ξz = vgeo[i,j,k,_ξx,e], vgeo[i,j,k,_ξy,e], vgeo[i,j,k,_ξz,e]
    ηx, ηy, ηz = vgeo[i,j,k,_ηx,e], vgeo[i,j,k,_ηy,e], vgeo[i,j,k,_ηz,e]
    ζx, ζy, ζz = vgeo[i,j,k,_ζx,e], vgeo[i,j,k,_ζy,e], vgeo[i,j,k,_ζz,e]
    x, y, z = vgeo[i,j,k,_x,e], vgeo[i,j,k,_y,e], vgeo[i,j,k,_z,e]

    U, V, W = Q[i, j, k, _U, e], Q[i, j, k, _V, e], Q[i, j, k, _W, e]
    ρ, E = Q[i, j, k, _ρ, e], Q[i, j, k, _E, e]

    P = gdm1*(E - (U^2 + V^2 + W^2)/(2*ρ))

    ρinv = 1 / ρ

    fluxρ_x = U
    fluxU_x = ρinv * U * U + P
    fluxV_x = ρinv * U * V
    fluxW_x = ρinv * U * W
    fluxE_x = ρinv * U * (E + P)

    fluxρ_y = V
    fluxU_y = ρinv * V * U
    fluxV_y = ρinv * V * V + P
    fluxW_y = ρinv * V * W
    fluxE_y = ρinv * V * (E + P)

    fluxρ_z = W
    fluxU_z = ρinv * W * U
    fluxV_z = ρinv * W * V
    fluxW_z = ρinv * W * W + P
    fluxE_z = ρinv * W * (E + P)

    s_F[i, j, k, _ρ] = MJ * (ξx * fluxρ_x + ξy * fluxρ_y + ξz * fluxρ_z)
    s_F[i, j, k, _U] = MJ * (ξx * fluxU_x + ξy * fluxU_y + ξz * fluxU_z)
    s_F[i, j, k, _V] = MJ * (ξx * fluxV_x + ξy * fluxV_y + ξz * fluxV_z)
    s_F[i, j, k, _W] = MJ * (ξx * fluxW_x + ξy * fluxW_y + ξz * fluxW_z)
    s_F[i, j, k, _E] = MJ * (ξx * fluxE_x + ξy * fluxE_y + ξz * fluxE_z)

    s_G[i, j, k, _ρ] = MJ * (ηx * fluxρ_x + ηy * fluxρ_y + ηz * fluxρ_z)
    s_G[i, j, k, _U] = MJ * (ηx * fluxU_x + ηy * fluxU_y + ηz * fluxU_z)
    s_G[i, j, k, _V] = MJ * (ηx * fluxV_x + ηy * fluxV_y + ηz * fluxV_z)
    s_G[i, j, k, _W] = MJ * (ηx * fluxW_x + ηy * fluxW_y + ηz * fluxW_z)
    s_G[i, j, k, _E] = MJ * (ηx * fluxE_x + ηy * fluxE_y + ηz * fluxE_z)

    s_H[i, j, k, _ρ] = MJ * (ζx * fluxρ_x + ζy * fluxρ_y + ζz * fluxρ_z)
    s_H[i, j, k, _U] = MJ * (ζx * fluxU_x + ζy * fluxU_y + ζz * fluxU_z)
    s_H[i, j, k, _V] = MJ * (ζx * fluxV_x + ζy * fluxV_y + ζz * fluxV_z)
    s_H[i, j, k, _W] = MJ * (ζx * fluxW_x + ζy * fluxW_y + ζz * fluxW_z)
    s_H[i, j, k, _E] = MJ * (ζx * fluxE_x + ζy * fluxE_y + ζz * fluxE_z)

    r_rhsρ = pi*(-CUDAnative.sin(pi*t)*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 2*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y) +
                 CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 3*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*z) +
                 2*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 3*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 3*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z))
    r_rhsU = pi*(-10*CUDAnative.sin(pi*t)*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 15*CUDAnative.sin(pi*t)*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*z) -
                 15*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 10*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 5*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 30*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*z) -
                 3*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 15*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y) +
                 9*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 15*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 6*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y) +
                 18*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 2*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*z))*CUDAnative.cos(pi*y)/5
    r_rhsV = pi*(-10*CUDAnative.sin(pi*t)*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 15*CUDAnative.sin(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 3*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 9*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 10*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 5*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 6*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y) -
                 18*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 15*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y) +
                 15*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 15*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 2*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*z) +
                 30*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z))*CUDAnative.sin(pi*x)/5
    r_rhsW = pi*(-10*CUDAnative.sin(pi*t)*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 15*CUDAnative.sin(pi*t) -
                 15*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.sin(pi*(x + y))*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*z) +
                 15*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*z) +
                 36*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 18*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 4*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y) +
                 18*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 4*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y) +
                 30*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*z)*CUDAnative.cos(pi*(x + y)) -
                 2*CUDAnative.cos(pi*t))*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*y)/5
    r_rhsE = pi*(3*(1 - CUDAnative.cos(2*pi*x))*
                 (1 - CUDAnative.cos(4*pi*z))*(CUDAnative.cos(2*pi*t) + 1)*
                 (CUDAnative.cos(2*pi*y) + 1)/512 - 5*CUDAnative.sin(pi*t)*
                 CUDAnative.sin(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 4*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 8*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y) -
                 2*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 9*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 18*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 4*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 3*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 8*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 6*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 14*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 9*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 7*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*z)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y) -
                 18*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 7*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) -
                 700*CUDAnative.sin(pi*x)*CUDAnative.sin(pi*y)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*z) +
                 14*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 700*CUDAnative.sin(pi*x)*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z) +
                 700*CUDAnative.cos(pi*t)*CUDAnative.cos(pi*x)*CUDAnative.cos(pi*y)*CUDAnative.cos(pi*z))/5

    sync_threads()

    # loop of ξ-grid lines
    @unroll for n = 1:Nq
      Dni = s_D[n, i]
      Dnj = s_D[n, j]
      Dnk = s_D[n, k]

      r_rhsρ += Dni * s_F[n, j, k, _ρ]
      r_rhsρ += Dnj * s_G[i, n, k, _ρ]
      r_rhsρ += Dnk * s_H[i, j, n, _ρ]

      r_rhsU += Dni * s_F[n, j, k, _U]
      r_rhsU += Dnj * s_G[i, n, k, _U]
      r_rhsU += Dnk * s_H[i, j, n, _U]

      r_rhsV += Dni * s_F[n, j, k, _V]
      r_rhsV += Dnj * s_G[i, n, k, _V]
      r_rhsV += Dnk * s_H[i, j, n, _V]

      r_rhsW += Dni * s_F[n, j, k, _W]
      r_rhsW += Dnj * s_G[i, n, k, _W]
      r_rhsW += Dnk * s_H[i, j, n, _W]

      r_rhsE += Dni * s_F[n, j, k, _E]
      r_rhsE += Dnj * s_G[i, n, k, _E]
      r_rhsE += Dnk * s_H[i, j, n, _E]
    end

    MJI = vgeo[i, j, k, _MJI, e]

    rhs[i, j, k, _U, e] += MJI * r_rhsU
    rhs[i, j, k, _V, e] += MJI * r_rhsV
    rhs[i, j, k, _W, e] += MJI * r_rhsW
    rhs[i, j, k, _ρ, e] += MJI * r_rhsρ
    rhs[i, j, k, _E, e] += MJI * r_rhsE
  end

  nothing
end

function volumerhs_trigpi!(::Val{N}, rhs, Q, vgeo, D, nelem, t) where {N}
  Nq = N + 1

  s_D = @cuStaticSharedMem eltype(D) (Nq, Nq)
  s_F = @cuStaticSharedMem eltype(Q) (Nq, Nq, Nq, _nstate)
  s_G = @cuStaticSharedMem eltype(Q) (Nq, Nq, Nq, _nstate)
  s_H = @cuStaticSharedMem eltype(Q) (Nq, Nq, Nq, _nstate)

  e = blockIdx().x
  k = threadIdx().z
  j = threadIdx().y
  i = threadIdx().x

  @inbounds begin
    # fetch D into shared
    k == 1 && (s_D[i, j] = D[i, j])

    # Load values will need into registers
    MJ = vgeo[i, j, k, _MJ, e]
    ξx, ξy, ξz = vgeo[i,j,k,_ξx,e], vgeo[i,j,k,_ξy,e], vgeo[i,j,k,_ξz,e]
    ηx, ηy, ηz = vgeo[i,j,k,_ηx,e], vgeo[i,j,k,_ηy,e], vgeo[i,j,k,_ηz,e]
    ζx, ζy, ζz = vgeo[i,j,k,_ζx,e], vgeo[i,j,k,_ζy,e], vgeo[i,j,k,_ζz,e]
    x, y, z = vgeo[i,j,k,_x,e], vgeo[i,j,k,_y,e], vgeo[i,j,k,_z,e]

    U, V, W = Q[i, j, k, _U, e], Q[i, j, k, _V, e], Q[i, j, k, _W, e]
    ρ, E = Q[i, j, k, _ρ, e], Q[i, j, k, _E, e]

    P = gdm1*(E - (U^2 + V^2 + W^2)/(2*ρ))

    ρinv = 1 / ρ

    fluxρ_x = U
    fluxU_x = ρinv * U * U + P
    fluxV_x = ρinv * U * V
    fluxW_x = ρinv * U * W
    fluxE_x = ρinv * U * (E + P)

    fluxρ_y = V
    fluxU_y = ρinv * V * U
    fluxV_y = ρinv * V * V + P
    fluxW_y = ρinv * V * W
    fluxE_y = ρinv * V * (E + P)

    fluxρ_z = W
    fluxU_z = ρinv * W * U
    fluxV_z = ρinv * W * V
    fluxW_z = ρinv * W * W + P
    fluxE_z = ρinv * W * (E + P)

    s_F[i, j, k, _ρ] = MJ * (ξx * fluxρ_x + ξy * fluxρ_y + ξz * fluxρ_z)
    s_F[i, j, k, _U] = MJ * (ξx * fluxU_x + ξy * fluxU_y + ξz * fluxU_z)
    s_F[i, j, k, _V] = MJ * (ξx * fluxV_x + ξy * fluxV_y + ξz * fluxV_z)
    s_F[i, j, k, _W] = MJ * (ξx * fluxW_x + ξy * fluxW_y + ξz * fluxW_z)
    s_F[i, j, k, _E] = MJ * (ξx * fluxE_x + ξy * fluxE_y + ξz * fluxE_z)

    s_G[i, j, k, _ρ] = MJ * (ηx * fluxρ_x + ηy * fluxρ_y + ηz * fluxρ_z)
    s_G[i, j, k, _U] = MJ * (ηx * fluxU_x + ηy * fluxU_y + ηz * fluxU_z)
    s_G[i, j, k, _V] = MJ * (ηx * fluxV_x + ηy * fluxV_y + ηz * fluxV_z)
    s_G[i, j, k, _W] = MJ * (ηx * fluxW_x + ηy * fluxW_y + ηz * fluxW_z)
    s_G[i, j, k, _E] = MJ * (ηx * fluxE_x + ηy * fluxE_y + ηz * fluxE_z)

    s_H[i, j, k, _ρ] = MJ * (ζx * fluxρ_x + ζy * fluxρ_y + ζz * fluxρ_z)
    s_H[i, j, k, _U] = MJ * (ζx * fluxU_x + ζy * fluxU_y + ζz * fluxU_z)
    s_H[i, j, k, _V] = MJ * (ζx * fluxV_x + ζy * fluxV_y + ζz * fluxV_z)
    s_H[i, j, k, _W] = MJ * (ζx * fluxW_x + ζy * fluxW_y + ζz * fluxW_z)
    s_H[i, j, k, _E] = MJ * (ζx * fluxE_x + ζy * fluxE_y + ζz * fluxE_z)

    r_rhsρ = pi*(-CUDAnative.sinpi(t)*CUDAnative.sinpi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 2*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y) +
                 CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 3*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(z) +
                 2*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 3*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 3*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z))
    r_rhsU = pi*(-10*CUDAnative.sinpi(t)*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 15*CUDAnative.sinpi(t)*CUDAnative.sinpi(x)*CUDAnative.cospi(z) -
                 15*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 10*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 5*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 30*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(z) -
                 3*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 15*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y) +
                 9*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 15*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 6*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y) +
                 18*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 2*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(z))*CUDAnative.cospi(y)/5
    r_rhsV = pi*(-10*CUDAnative.sinpi(t)*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 15*CUDAnative.sinpi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 3*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 9*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 10*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 5*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 6*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y) -
                 18*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 15*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y) +
                 15*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 15*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 2*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(z) +
                 30*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z))*CUDAnative.sinpi(x)/5
    r_rhsW = pi*(-10*CUDAnative.sinpi(t)*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 15*CUDAnative.sinpi(t) -
                 15*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.sinpi((x + y))*CUDAnative.cospi(t)*CUDAnative.cospi(z) +
                 15*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(z) +
                 36*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 18*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 4*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y) +
                 18*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 4*CUDAnative.cospi(t)*CUDAnative.cospi(y) +
                 30*CUDAnative.cospi(t)*CUDAnative.cospi(z)*CUDAnative.cospi((x + y)) -
                 2*CUDAnative.cospi(t))*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(y)/5
    r_rhsE = pi*(3*(1 - CUDAnative.cos(2*pi*x))*
                 (1 - CUDAnative.cos(4*pi*z))*(CUDAnative.cos(2*pi*t) + 1)*
                 (CUDAnative.cos(2*pi*y) + 1)/512 - 5*CUDAnative.sinpi(t)*
                 CUDAnative.sinpi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 4*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 8*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y) -
                 2*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 9*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 18*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 4*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 3*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 8*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 6*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 14*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 9*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 7*CUDAnative.sinpi(x)*CUDAnative.sinpi(z)*CUDAnative.cospi(t)*CUDAnative.cospi(y) -
                 18*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 7*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) -
                 700*CUDAnative.sinpi(x)*CUDAnative.sinpi(y)*CUDAnative.cospi(t)*CUDAnative.cospi(z) +
                 14*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 700*CUDAnative.sinpi(x)*CUDAnative.cospi(t)*CUDAnative.cospi(y)*CUDAnative.cospi(z) +
                 700*CUDAnative.cospi(t)*CUDAnative.cospi(x)*CUDAnative.cospi(y)*CUDAnative.cospi(z))/5

    sync_threads()

    # loop of ξ-grid lines
    @unroll for n = 1:Nq
      Dni = s_D[n, i]
      Dnj = s_D[n, j]
      Dnk = s_D[n, k]

      r_rhsρ += Dni * s_F[n, j, k, _ρ]
      r_rhsρ += Dnj * s_G[i, n, k, _ρ]
      r_rhsρ += Dnk * s_H[i, j, n, _ρ]

      r_rhsU += Dni * s_F[n, j, k, _U]
      r_rhsU += Dnj * s_G[i, n, k, _U]
      r_rhsU += Dnk * s_H[i, j, n, _U]

      r_rhsV += Dni * s_F[n, j, k, _V]
      r_rhsV += Dnj * s_G[i, n, k, _V]
      r_rhsV += Dnk * s_H[i, j, n, _V]

      r_rhsW += Dni * s_F[n, j, k, _W]
      r_rhsW += Dnj * s_G[i, n, k, _W]
      r_rhsW += Dnk * s_H[i, j, n, _W]

      r_rhsE += Dni * s_F[n, j, k, _E]
      r_rhsE += Dnj * s_G[i, n, k, _E]
      r_rhsE += Dnk * s_H[i, j, n, _E]
    end

    MJI = vgeo[i, j, k, _MJI, e]

    rhs[i, j, k, _U, e] += MJI * r_rhsU
    rhs[i, j, k, _V, e] += MJI * r_rhsV
    rhs[i, j, k, _W, e] += MJI * r_rhsW
    rhs[i, j, k, _ρ, e] += MJI * r_rhsρ
    rhs[i, j, k, _E, e] += MJI * r_rhsE
  end

  nothing
end

function main()
    N = 4
    nelem = 4000
    ntrials = 1
    DFloat = Float32

    rnd = MersenneTwister(0)

    Nq = N + 1
    Q = 1 .+ CuArray(rand(rnd, DFloat, Nq, Nq, Nq, _nstate, nelem))
    Q[:, :, :, _E, :] .+= 20
    vgeo = CuArray(rand(rnd, DFloat, Nq, Nq, Nq, _nvgeo, nelem))

    # Make sure the entries of the mass matrix satisfy the inverse relation
    vgeo[:, :, :, _MJ, :] .+= 3
    vgeo[:, :, :, _MJI, :] .= 1 ./ vgeo[:, :, :, _MJ, :]

    D = CuArray(rand(rnd, DFloat, Nq, Nq))

    rhs = CuArray(zeros(DFloat, Nq, Nq, Nq, _nstate, nelem))

    t = DFloat(1)
    @cuda(threads=(N+1, N+1, N+1), blocks=nelem, maxregs=255,
          volumerhs!(Val(N), rhs, Q, vgeo, D, nelem, t))

    CUDAdrv.@profile  begin
        @cuda(threads=(N+1, N+1, N+1), blocks=nelem, maxregs=255,
              volumerhs!(Val(N), rhs, Q, vgeo, D, nelem, t))
    end

    @cuda(threads=(N+1, N+1, N+1), blocks=nelem, maxregs=255,
          volumerhs_trigpi!(Val(N), rhs, Q, vgeo, D, nelem, t))

    CUDAdrv.@profile  begin
        @cuda(threads=(N+1, N+1, N+1), blocks=nelem, maxregs=255,
              volumerhs_trigpi!(Val(N), rhs, Q, vgeo, D, nelem, t))
    end

    #=
    CUDAnative.@device_code dir="jlir" begin
        @cuda(threads=(N+1, N+1, N+1), blocks=nelem, maxregs=255,
              volumerhs!(Val(N), rhs, Q, vgeo, D, nelem, t))
    end
    =#

    nothing
end

main()