text:
using .NumericalFluxes:
    CentralNumericalFluxHigherOrder, CentralNumericalFluxDivergence

"""
    SpaceDiscretization

Supertype for spatial discretizations.

Must have the following properties:

    - `grid`
    - `balance_law`
    - `state_auxiliary`
"""
abstract type SpaceDiscretization end

struct DGFVModel{BL, G, FVR, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD, GF, TF} <:
       SpaceDiscretization
    balance_law::BL
    grid::G
    fv_reconstruction::FVR
    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
    gradient_filter::GF
    tendency_filter::TF
    check_for_crashes::Bool
end

function DGFVModel(
    balance_law,
    grid,
    fv_reconstruction,
    numerical_flux_first_order,
    numerical_flux_second_order,
    numerical_flux_gradient;
    fill_nan = false,
    check_for_crashes = false,
    state_auxiliary = create_state(
        balance_law,
        grid,
        Auxiliary(),
        fill_nan = fill_nan,
    ),
    state_gradient_flux = create_state(balance_law, grid, GradientFlux()),
    states_higher_order = (
        create_state(balance_law, grid, GradientLaplacian()),
        create_state(balance_law, grid, Hyperdiffusive()),
    ),
    direction = EveryDirection(),
    diffusion_direction = direction,
    modeldata = nothing,
    gradient_filter = nothing,
    tendency_filter = nothing,
)
    # Make sure we are FVM in the vertical
    @assert polynomialorders(grid)[end] == 0
    @assert isstacked(grid.topology)
    state_auxiliary =
        init_state(state_auxiliary, balance_law, grid, direction, Auxiliary())
    DGFVModel(
        balance_law,
        grid,
        fv_reconstruction,
        numerical_flux_first_order,
        numerical_flux_second_order,
        numerical_flux_gradient,
        state_auxiliary,
        state_gradient_flux,
        states_higher_order,
        direction,
        diffusion_direction,
        modeldata,
        gradient_filter,
        tendency_filter,
        check_for_crashes,
    )
end

struct DGModel{BL, G, NFND, NFD, GNF, AS, DS, HDS, D, DD, MD, GF, TF} <:
       SpaceDiscretization
    balance_law::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
    gradient_filter::GF
    tendency_filter::TF
    check_for_crashes::Bool
end

function DGModel(
    balance_law,
    grid,
    numerical_flux_first_order,
    numerical_flux_second_order,
    numerical_flux_gradient;
    fill_nan = false,
    check_for_crashes = false,
    state_auxiliary = create_state(
        balance_law,
        grid,
        Auxiliary(),
        fill_nan = fill_nan,
    ),
    state_gradient_flux = create_state(balance_law, grid, GradientFlux()),
    states_higher_order = (
        create_state(balance_law, grid, GradientLaplacian()),
        create_state(balance_law, grid, Hyperdiffusive()),
    ),
    direction = EveryDirection(),
    diffusion_direction = direction,
    modeldata = nothing,
    gradient_filter = nothing,
    tendency_filter = nothing,
)
    state_auxiliary =
        init_state(state_auxiliary, balance_law, grid, direction, Auxiliary())
    DGModel(
        balance_law,
        grid,
        numerical_flux_first_order,
        numerical_flux_second_order,
        numerical_flux_gradient,
        state_auxiliary,
        state_gradient_flux,
        states_higher_order,
        direction,
        diffusion_direction,
        modeldata,
        gradient_filter,
        tendency_filter,
        check_for_crashes,
    )
end

# Include the remainder model for composing DG models and balance laws
include("remainder.jl")

basic_grid_info(spacedisc::SpaceDiscretization) =
    basic_grid_info(spacedisc.grid)
function basic_grid_info(grid)
    dim = dimensionality(grid)
    # Tuple of polynomial degrees (N₁, N₂, N₃)
    N = polynomialorders(grid)
    @assert dim == 2 || N[1] == N[2]

    # Number of quadrature point in each direction (Nq₁, Nq₂, Nq₃)
    Nq = N .+ 1

    # Number of quadrature point in the horizontal direction Nq₁ * Nq₂
    Nqh = Nq[1] * Nq[2]

    # Number of quadrature points in the vertical Nq₃
    Nqk = dim == 2 ? 1 : Nq[dim]

    Np = dofs_per_element(grid)
    Nfp_v, Nfp_h = div.(Np, (Nq[1], Nq[end]))

    topology_info = basic_topology_info(grid.topology)

    ninteriorelem = length(grid.interiorelems)
    nexteriorelem = length(grid.exteriorelems)

    nface = 2 * dim

    grid_info = (
        dim = dim,
        N = N,
        Nq = Nq,
        Nqh = Nqh,
        Nqk = Nqk,
        Nfp_v = Nfp_v,
        Nfp_h = Nfp_h,
        Np = Np,
        nface = nface,
        ninteriorelem = ninteriorelem,
        nexteriorelem = nexteriorelem,
    )

    return merge(grid_info, topology_info)
end

function basic_launch_info(spacedisc::SpaceDiscretization)
    device = array_device(spacedisc.state_auxiliary)
    grid_info = basic_grid_info(spacedisc.grid)
    return merge(grid_info, (device = device,))
end

function (spacedisc::SpaceDiscretization)(
    tendency,
    state_prognostic,
    param,
    t;
    increment = false,
)
    # TODO deprecate increment argument
    spacedisc(tendency, state_prognostic, param, t, true, increment)
end

"""
    (dgfvm::DGFVModel)(tendency, state_prognostic, _, t, α, β)

Uses spectral element discontinuous Galerkin in the horizontal and finite volume
in the vertical to compute the tendency.

Computes the tendency terms compatible with `IncrementODEProblem`

    tendency .= α .* dQdt(state_prognostic, p, t) .+ β .* tendency

The 4-argument form will just compute

    tendency .= dQdt(state_prognostic, p, t)
"""
function (dgfvm::DGFVModel)(tendency, state_prognostic, _, t, α, β)
    device = array_device(state_prognostic)

    FT = eltype(state_prognostic)
    num_state_prognostic = number_states(dgfvm.balance_law, Prognostic())
    num_state_gradient_flux = number_states(dgfvm.balance_law, GradientFlux())
    @assert 0 == number_states(dgfvm.balance_law, Hyperdiffusive())
    num_state_tendency = size(tendency, 2)

    if num_state_prognostic < num_state_tendency && β != 1
        # If we don't operate on the full state, then we need to scale here instead of volume_tendency!
        tendency .*= β
        β = β != 0 # if β==0 then we can avoid the memory load in volume_tendency!
    end

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

    update_auxiliary_state!(
        dgfvm,
        dgfvm.balance_law,
        state_prognostic,
        t,
        dgfvm.grid.topology.realelems,
    )

    exchange_state_gradient_flux = NoneEvent()
    exchange_state_prognostic = NoneEvent()

    comp_stream = Event(device)

    if communicate
        exchange_state_prognostic = MPIStateArrays.begin_ghost_exchange!(
            state_prognostic;
            dependencies = comp_stream,
        )
    end

    if num_state_gradient_flux > 0
        ########################
        # Gradient Computation #
        ########################

        comp_stream = launch_volume_gradients!(
            dgfvm,
            state_prognostic,
            t;
            dependencies = comp_stream,
        )

        comp_stream = launch_interface_gradients!(
            dgfvm,
            state_prognostic,
            t;
            surface = :interior,
            dependencies = comp_stream,
        )

        if communicate
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dgfvm.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dgfvm.check_for_crashes,
            )
            update_auxiliary_state!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end

        comp_stream = launch_interface_gradients!(
            dgfvm,
            state_prognostic,
            t;
            surface = :exterior,
            dependencies = (comp_stream, exchange_state_prognostic),
        )

        if dgfvm.gradient_filter !== nothing
            comp_stream = Filters.apply_async!(
                dgfvm.state_gradient_flux,
                1:num_state_gradient_flux,
                dgfvm.grid,
                dgfvm.gradient_filter;
                dependencies = comp_stream,
            )
        end

        if communicate
            if num_state_gradient_flux > 0
                exchange_state_gradient_flux =
                    MPIStateArrays.begin_ghost_exchange!(
                        dgfvm.state_gradient_flux,
                        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.
            checked_wait(device, comp_stream, nothing, dgfvm.check_for_crashes)
            update_auxiliary_state_gradient!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.realelems,
            )
            comp_stream = Event(device)
        end
    end

    ###################
    # RHS Computation #
    ###################
    comp_stream = launch_volume_tendency!(
        dgfvm,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        dependencies = comp_stream,
    )

    comp_stream = launch_interface_tendency!(
        dgfvm,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :interior,
        dependencies = comp_stream,
    )

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

            # Update_aux_diffusive may start asynchronous work on the
            # compute device and we synchronize those here through a device
            # event.
            checked_wait(
                device,
                exchange_state_gradient_flux,
                nothing,
                dgfvm.check_for_crashes,
            )
            update_auxiliary_state_gradient!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.ghostelems,
            )
            exchange_state_gradient_flux = Event(device)
        else
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dgfvm.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dgfvm.check_for_crashes,
            )
            update_auxiliary_state!(
                dgfvm,
                dgfvm.balance_law,
                state_prognostic,
                t,
                dgfvm.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end
    end

    comp_stream = launch_interface_tendency!(
        dgfvm,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :exterior,
        dependencies = (
            comp_stream,
            exchange_state_prognostic,
            exchange_state_gradient_flux,
            # XXX: This is disabled until FVM with hyperdiffusion for DG is implemented: exchange_Qhypervisc_grad,
        ),
    )

    if dgfvm.tendency_filter !== nothing
        comp_stream = Filters.apply_async!(
            tendency,
            1:num_state_tendency,
            dgfvm.grid,
            dgfvm.tendency_filter;
            dependencies = comp_stream,
        )
    end

    # 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.
    checked_wait(device, comp_stream, nothing, dgfvm.check_for_crashes)
end

"""
    (dg::DGModel)(tendency, state_prognostic, _, t, α, β)

Uses spectral element discontinuous Galerkin in all direction to compute the
tendency.

Computes the tendency terms compatible with `IncrementODEProblem`

    tendency .= α .* dQdt(state_prognostic, p, t) .+ β .* tendency

The 4-argument form will just compute

    tendency .= dQdt(state_prognostic, p, t)
"""
function (dg::DGModel)(tendency, state_prognostic, _, t, α, β)

    device = array_device(state_prognostic)
    Qhypervisc_grad, Qhypervisc_div = dg.states_higher_order

    FT = eltype(state_prognostic)
    num_state_prognostic = number_states(dg.balance_law, Prognostic())
    num_state_gradient_flux = number_states(dg.balance_law, GradientFlux())
    nhyperviscstate = number_states(dg.balance_law, Hyperdiffusive())
    num_state_tendency = size(tendency, 2)

    @assert num_state_prognostic ≤ num_state_tendency

    if num_state_prognostic < num_state_tendency && β != 1
        # If we don't operate on the full state, then we need to scale here instead of volume_tendency!
        tendency .*= β
        β = β != 0 # if β==0 then we can avoid the memory load in volume_tendency!
    end

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

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

    exchange_state_prognostic = NoneEvent()
    exchange_state_gradient_flux = NoneEvent()
    exchange_Qhypervisc_grad = NoneEvent()
    exchange_Qhypervisc_div = NoneEvent()

    comp_stream = Event(device)

    if communicate
        exchange_state_prognostic = MPIStateArrays.begin_ghost_exchange!(
            state_prognostic;
            dependencies = comp_stream,
        )
    end

    if num_state_gradient_flux > 0 || nhyperviscstate > 0
        ########################
        # Gradient Computation #
        ########################

        comp_stream = launch_volume_gradients!(
            dg,
            state_prognostic,
            t;
            dependencies = comp_stream,
        )

        comp_stream = launch_interface_gradients!(
            dg,
            state_prognostic,
            t;
            surface = :interior,
            dependencies = comp_stream,
        )

        if communicate
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dg.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dg.check_for_crashes,
            )
            update_auxiliary_state!(
                dg,
                dg.balance_law,
                state_prognostic,
                t,
                dg.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end

        comp_stream = launch_interface_gradients!(
            dg,
            state_prognostic,
            t;
            surface = :exterior,
            dependencies = (comp_stream, exchange_state_prognostic),
        )

        if dg.gradient_filter !== nothing
            comp_stream = Filters.apply_async!(
                dg.state_gradient_flux,
                1:num_state_gradient_flux,
                dg.grid,
                dg.gradient_filter;
                dependencies = comp_stream,
            )
        end

        if communicate
            if num_state_gradient_flux > 0
                exchange_state_gradient_flux =
                    MPIStateArrays.begin_ghost_exchange!(
                        dg.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.
            checked_wait(device, comp_stream, nothing, dg.check_for_crashes)
            update_auxiliary_state_gradient!(
                dg,
                dg.balance_law,
                state_prognostic,
                t,
                dg.grid.topology.realelems,
            )
            comp_stream = Event(device)
        end
    end

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

        comp_stream = launch_volume_divergence_of_gradients!(
            dg,
            state_prognostic,
            t;
            dependencies = comp_stream,
        )

        comp_stream = launch_interface_divergence_of_gradients!(
            dg,
            state_prognostic,
            t;
            surface = :interior,
            dependencies = comp_stream,
        )

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

        comp_stream = launch_interface_divergence_of_gradients!(
            dg,
            state_prognostic,
            t;
            surface = :exterior,
            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 = launch_volume_gradients_of_laplacians!(
            dg,
            state_prognostic,
            t,
            dependencies = (comp_stream,),
        )

        comp_stream = launch_interface_gradients_of_laplacians!(
            dg,
            state_prognostic,
            t;
            surface = :interior,
            dependencies = (comp_stream,),
        )

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

        comp_stream = launch_interface_gradients_of_laplacians!(
            dg,
            state_prognostic,
            t;
            surface = :exterior,
            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 = launch_volume_tendency!(
        dg,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        dependencies = (comp_stream,),
    )

    comp_stream = launch_interface_tendency!(
        dg,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :interior,
        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!(
                        dg.state_gradient_flux;
                        dependencies = exchange_state_gradient_flux,
                        check_for_crashes = dg.check_for_crashes,
                    )

                # Update_aux_diffusive may start asynchronous work on the
                # compute device and we synchronize those here through a device
                # event.
                checked_wait(
                    device,
                    exchange_state_gradient_flux,
                    nothing,
                    dg.check_for_crashes,
                )
                update_auxiliary_state_gradient!(
                    dg,
                    dg.balance_law,
                    state_prognostic,
                    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,
                    check_for_crashes = dg.check_for_crashes,
                )
            end
        else
            exchange_state_prognostic = MPIStateArrays.end_ghost_exchange!(
                state_prognostic;
                dependencies = exchange_state_prognostic,
                check_for_crashes = dg.check_for_crashes,
            )

            # Update_aux may start asynchronous work on the compute device and
            # we synchronize those here through a device event.
            checked_wait(
                device,
                exchange_state_prognostic,
                nothing,
                dg.check_for_crashes,
            )
            update_auxiliary_state!(
                dg,
                dg.balance_law,
                state_prognostic,
                t,
                dg.grid.topology.ghostelems,
            )
            exchange_state_prognostic = Event(device)
        end
    end

    comp_stream = launch_interface_tendency!(
        dg,
        tendency,
        state_prognostic,
        t,
        α,
        β;
        surface = :exterior,
        dependencies = (
            comp_stream,
            exchange_state_prognostic,
            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.
    if dg.tendency_filter !== nothing
        comp_stream = Filters.apply_async!(
            tendency,
            1:num_state_tendency,
            dg.grid,
            dg.tendency_filter;
            dependencies = comp_stream,
        )
    end
    checked_wait(device, comp_stream, nothing, dg.check_for_crashes)
end

function init_ode_state(
    spacedisc::SpaceDiscretization,
    args...;
    init_on_cpu = false,
    fill_nan = false,
)
    grid = spacedisc.grid
    balance_law = spacedisc.balance_law
    state_auxiliary = spacedisc.state_auxiliary

    device = arraytype(grid) <: Array ? CPU() : CUDADevice()

    state_prognostic =
        create_state(balance_law, grid, Prognostic(), fill_nan = fill_nan)

    topology = grid.topology
    Np = dofs_per_element(grid)

    dim = dimensionality(grid)
    N = polynomialorders(grid)
    nrealelem = length(topology.realelems)

    if !init_on_cpu
        event = Event(device)
        event = kernel_init_state_prognostic!(device, min(Np, 1024))(
            balance_law,
            Val(dim),
            Val(N),
            state_prognostic.data,
            state_auxiliary.data,
            grid.vgeo,
            topology.realelems,
            args...;
            ndrange = Np * nrealelem,
            dependencies = (event,),
        )
        wait(device, event)
    else
        h_state_prognostic = similar(state_prognostic, Array)
        h_state_auxiliary = similar(state_auxiliary, Array)
        h_state_auxiliary .= state_auxiliary
        event = kernel_init_state_prognostic!(CPU(), Np)(
            balance_law,
            Val(dim),
            Val(N),
            h_state_prognostic.data,
            h_state_auxiliary.data,
            Array(grid.vgeo),
            topology.realelems,
            args...;
            ndrange = Np * nrealelem,
        )
        wait(event)
        state_prognostic .= h_state_prognostic
        state_auxiliary .= h_state_auxiliary
    end

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

    return state_prognostic
end

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

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

    device = arraytype(dg.grid) <: Array ? CPU() : CUDADevice()
    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, direction)
    state_auxiliary = create_state(bl, grid, Auxiliary())
    state_auxiliary =
        init_state(state_auxiliary, bl, grid, direction, Auxiliary())
    state_auxiliary .= aux_data
    return state_auxiliary
end

# TODO: this should really be a separate function
"""
    init_state_auxiliary!(
        bl::BalanceLaw,
        f!,
        statearray_auxiliary,
        grid,
        direction;
        state_temporary = nothing
    )

Apply `f!(bl, state_auxiliary, tmp, geom)` at each node, storing the result in
`statearray_auxiliary`, where `tmp` are the values at the corresponding node in
`state_temporary` and `geom` contains the geometry information.
"""
function init_state_auxiliary!(
    balance_law,
    init_f!,
    state_auxiliary,
    grid,
    direction;
    state_temporary = nothing,
)
    topology = grid.topology
    dim = dimensionality(grid)
    Np = dofs_per_element(grid)
    N = polynomialorders(grid)
    vgeo = grid.vgeo
    device = array_device(state_auxiliary)
    nrealelem = length(topology.realelems)

    event = Event(device)
    event = kernel_nodal_init_state_auxiliary!(
        device,
        min(Np, 1024),
        Np * nrealelem,
    )(
        balance_law,
        Val(dim),
        Val(N),
        init_f!,
        state_auxiliary.data,
        isnothing(state_temporary) ? nothing : state_temporary.data,
        Val(isnothing(state_temporary) ? @vars() : vars(state_temporary)),
        vgeo,
        topology.realelems,
        dependencies = (event,),
    )

    event = MPIStateArrays.begin_ghost_exchange!(
        state_auxiliary;
        dependencies = event,
    )
    event = MPIStateArrays.end_ghost_exchange!(
        state_auxiliary;
        dependencies = event,
    )
    wait(device, event)
end

update_auxiliary_state_gradient!(::SpaceDiscretization, _...) = false

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

    device = array_device(state_prognostic)

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorders(grid)
    Nq = N .+ 1
    Nqj = dim == 2 ? 1 : Nq[2]

    FT = eltype(state_prognostic)

    # Compute 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[1], Nqj))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        state_prognostic.data,
        state_auxiliary.data,
        grid.vgeo,
        # Only need the vertical Imat since this kernel is vertically oriented
        grid.Imat[dim],
        horzelems;
        ndrange = (length(horzelems) * Nq[1], Nqj),
        dependencies = (event,),
    )
    checked_wait(device, event, nothing, dg.check_for_crashes)
end

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

    device = array_device(state_auxiliary)

    grid = dg.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorders(grid)
    Nq = N .+ 1
    Nqj = dim == 2 ? 1 : Nq[2]

    FT = eltype(state_auxiliary)

    # Compute 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[1], Nqj))(
        m,
        Val(dim),
        Val(N),
        Val(nvertelem),
        state_prognostic.data,
        state_auxiliary.data,
        horzelems;
        ndrange = (length(horzelems) * Nq[1], Nqj),
        dependencies = (event,),
    )
    checked_wait(device, event, nothing, dg.check_for_crashes)
end

# By default, we call update_auxiliary_state!, given
# nodal_update_auxiliary_state!, defined for the
# particular balance_law:

# TODO: this should really be a separate function
function update_auxiliary_state!(
    f!,
    spacedisc::SpaceDiscretization,
    m::BalanceLaw,
    state_prognostic::MPIStateArray,
    t::Real,
    elems::UnitRange = spacedisc.grid.topology.realelems;
    diffusive = false,
)
    device = array_device(state_prognostic)

    grid = spacedisc.grid
    topology = grid.topology

    dim = dimensionality(grid)
    N = polynomialorders(grid)
    nelem = length(elems)

    Np = dofs_per_element(grid)

    knl_nodal_update_auxiliary_state! =
        kernel_nodal_update_auxiliary_state!(device, min(Np, 1024))
    ### Update state_auxiliary variables
    event = Event(device)
    if diffusive
        event = knl_nodal_update_auxiliary_state!(
            m,
            Val(dim),
            Val(N),
            f!,
            state_prognostic.data,
            spacedisc.state_auxiliary.data,
            spacedisc.state_gradient_flux.data,
            t,
            elems,
            grid.activedofs;
            ndrange = Np * nelem,
            dependencies = (event,),
        )
    else
        event = knl_nodal_update_auxiliary_state!(
            m,
            Val(dim),
            Val(N),
            f!,
            state_prognostic.data,
            spacedisc.state_auxiliary.data,
            t,
            elems,
            grid.activedofs;
            ndrange = Np * nelem,
            dependencies = (event,),
        )
    end
    checked_wait(device, event, nothing, spacedisc.check_for_crashes)
end

"""
    courant(local_courant::Function, dg::DGModel, m::BalanceLaw,
            state_prognostic::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_prognostic::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::SpaceDiscretization,
    m::BalanceLaw,
    state_prognostic::MPIStateArray,
    Δt,
    simtime,
    direction = EveryDirection(),
)
    grid = dg.grid
    topology = grid.topology
    nrealelem = length(topology.realelems)

    if nrealelem > 0
        N = polynomialorders(grid)
        dim = dimensionality(grid)
        Nq = N .+ 1
        Nqk = dim == 2 ? 1 : Nq[dim]
        device = array_device(grid.vgeo)
        pointwise_courant = similar(grid.vgeo, prod(Nq), nrealelem)
        event = Event(device)
        event = Grids.kernel_min_neighbor_distance!(
            device,
            min(Nq[1] * Nq[2] * Nqk, 1024),
        )(
            Val(N),
            Val(dim),
            direction,
            pointwise_courant,
            grid.vgeo,
            topology.realelems;
            ndrange = (nrealelem * Nq[1] * Nq[2] * Nqk),
            dependencies = (event,),
        )
        event = kernel_local_courant!(device, min(Nq[1] * Nq[2] * Nqk, 1024))(
            m,
            Val(dim),
            Val(N),
            pointwise_courant,
            local_courant,
            state_prognostic.data,
            dg.state_auxiliary.data,
            dg.state_gradient_flux.data,
            topology.realelems,
            Δt,
            simtime,
            direction;
            ndrange = (nrealelem * Nq[1] * Nq[2] * Nqk),
            dependencies = (event,),
        )
        checked_wait(device, event, nothing, dg.check_for_crashes)

        rank_courant_max = maximum(pointwise_courant)
    else
        rank_courant_max = typemin(eltype(state_prognostic))
    end

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

function MPIStateArrays.MPIStateArray(dg::SpaceDiscretization)
    balance_law = dg.balance_law
    grid = dg.grid

    state_prognostic = create_state(balance_law, grid, Prognostic())

    return state_prognostic
end

"""
    auxiliary_field_gradient!(::BalanceLaw, ∇state::MPIStateArray,
                               vars_out, state::MPIStateArray, vars_in, grid;
                               direction = EveryDirection())

Take the gradient of the variables `vars_in` located in the array `state`
and stores it in the variables `vars_out` of `∇state`. This function computes
element wise gradient without accounting for numerical fluxes and hence
its primary purpose is to take the gradient of continuous reference fields.

## Examples
```julia
FT = eltype(state_auxiliary)
grad_Φ = similar(state_auxiliary, vars=@vars(∇Φ::SVector{3, FT}))
auxiliary_field_gradient!(
    model,
    grad_Φ,
    ("∇Φ",),
    state_auxiliary,
    ("orientation.Φ",),
    grid,
)
```
"""
function auxiliary_field_gradient!(
    m::BalanceLaw,
    ∇state::MPIStateArray,
    vars_out,
    state::MPIStateArray,
    vars_in,
    grid,
    direction = EveryDirection(),
)
    topology = grid.topology
    nrealelem = length(topology.realelems)
    N = polynomialorders(grid)
    dim = dimensionality(grid)
    Nq = N .+ 1
    Nqk = dim == 2 ? 1 : Nq[dim]
    device = array_device(state)

    I = varsindices(vars(state), vars_in)
    O = varsindices(vars(∇state), vars_out)

    event = Event(device)

    if direction isa EveryDirection || direction isa HorizontalDirection
        # We assume N₁ = N₂, so the same polyorder, quadrature weights,
        # and differentiation operators are used
        horizontal_polyorder = N[1]
        horizontal_D = grid.D[1]
        horizontal_ω = grid.ω[1]
        event = dgsem_auxiliary_field_gradient!(device, (Nq[1], Nq[2]))(
            m,
            Val(dim),
            Val(N),
            HorizontalDirection(),
            ∇state.data,
            state.data,
            grid.vgeo,
            horizontal_D,
            horizontal_ω,
            Val(I),
            Val(O),
            false,
            ndrange = (nrealelem * Nq[1], Nq[2]),
            dependencies = (event,),
        )
    end

    if direction isa EveryDirection || direction isa VerticalDirection
        vertical_polyorder = N[dim]
        if vertical_polyorder > 0
            vertical_D = grid.D[dim]
            vertical_ω = grid.ω[dim]
            event = dgsem_auxiliary_field_gradient!(device, (Nq[1], Nq[2]))(
                m,
                Val(dim),
                Val(N),
                VerticalDirection(),
                ∇state.data,
                state.data,
                grid.vgeo,
                vertical_D,
                vertical_ω,
                Val(I),
                Val(O),
                # If we are computing in every direction, we need to
                # increment after we compute the horizontal values
                (direction isa EveryDirection);
                ndrange = (nrealelem * Nq[1], Nq[2]),
                dependencies = (event,),
            )
        else
            info = basic_grid_info(grid)
            event = vert_fvm_auxiliary_field_gradient!(device, info.Nfp_h)(
                m,
                Val(info),
                ∇state.data,
                state.data,
                grid.vgeo,
                grid.sgeo,
                grid.vmap⁻,
                grid.vmap⁺,
                grid.elemtobndy,
                Val(I),
                Val(O),
                # If we are computing in every direction, we need to
                # increment after we compute the horizontal values
                (direction isa EveryDirection);
                ndrange = (nrealelem * info.Nfp_h),
                dependencies = (event,),
            )
        end
    end
    wait(device, event)
end

function hyperdiff_indexmap(balance_law, ::Type{FT}) where {FT}
    ns_hyperdiff = number_states(balance_law, Hyperdiffusive())
    if ns_hyperdiff > 0
        return varsindices(
            vars_state(balance_law, Gradient(), FT),
            fieldnames(vars_state(balance_law, GradientLaplacian(), FT)),
        )
    else
        return nothing
    end
end

"""
    launch_volume_gradients!(spacedisc, state_prognostic, t; dependencies)

Launches horizontal and vertical kernels for computing the volume gradients.
"""
function launch_volume_gradients!(spacedisc, state_prognostic, t; dependencies)
    FT = eltype(state_prognostic)
    # XXX: This is until FVM with hyperdiffusion for DG is implemented
    if spacedisc isa DGFVModel
        @assert 0 == number_states(spacedisc.balance_law, Hyperdiffusive())
        Qhypervisc_grad_data = nothing
    elseif spacedisc isa DGModel
        Qhypervisc_grad_data = spacedisc.states_higher_order[1].data
    end

    # Workgroup is determined by the number of quadrature points
    # in the horizontal direction. For each horizontal quadrature
    # point, we operate on a stack of quadrature in the vertical
    # direction. (Iteration space is in the horizontal)
    info = basic_launch_info(spacedisc)

    # We assume (in 3-D) that both x and y directions
    # are discretized using the same polynomial order, Nq[1] == Nq[2].
    # In 2-D, the workgroup spans the entire set of quadrature points:
    # Nq[1] * Nq[2]
    workgroup = (info.Nq[1], info.Nq[2])
    ndrange = (info.Nq[1] * info.nrealelem, info.Nq[2])
    comp_stream = dependencies

    # If the model direction is EveryDirection, we need to perform
    # both horizontal AND vertical kernel calls; otherwise, we only
    # call the kernel corresponding to the model direction `spacedisc.diffusion_direction`
    if spacedisc.diffusion_direction isa EveryDirection ||
       spacedisc.diffusion_direction isa HorizontalDirection

        # We assume N₁ = N₂, so the same polyorder, quadrature weights,
        # and differentiation operators are used
        horizontal_polyorder = info.N[1]
        horizontal_D = spacedisc.grid.D[1]
        comp_stream = volume_gradients!(info.device, workgroup)(
            spacedisc.balance_law,
            Val(info),
            HorizontalDirection(),
            state_prognostic.data,
            spacedisc.state_gradient_flux.data,
            Qhypervisc_grad_data,
            spacedisc.state_auxiliary.data,
            spacedisc.grid.vgeo,
            t,
            horizontal_D,
            Val(hyperdiff_indexmap(spacedisc.balance_law, FT)),
            spacedisc.grid.topology.realelems,
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end

    # Now we call the kernel corresponding to the vertical direction
    if spacedisc isa DGModel && (
        spacedisc.diffusion_direction isa EveryDirection ||
        spacedisc.diffusion_direction isa VerticalDirection
    )

        # Vertical polynomial degree and differentiation matrix
        vertical_polyorder = info.N[info.dim]
        vertical_D = spacedisc.grid.D[info.dim]
        comp_stream = volume_gradients!(info.device, workgroup)(
            spacedisc.balance_law,
            Val(info),
            VerticalDirection(),
            state_prognostic.data,
            spacedisc.state_gradient_flux.data,
            Qhypervisc_grad_data,
            spacedisc.state_auxiliary.data,
            spacedisc.grid.vgeo,
            t,
            vertical_D,
            Val(hyperdiff_indexmap(spacedisc.balance_law, FT)),
            spacedisc.grid.topology.realelems,
            # If we are computing the volume gradient in every direction, we
            # need to increment into the appropriate fields _after_ the
            # horizontal computation.
            !(spacedisc.diffusion_direction isa VerticalDirection),
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end
    return comp_stream
end

"""
    launch_interface_gradients!(spacedisc, state_prognostic, t; surface::Symbol, dependencies)

Launches horizontal and vertical kernels for computing the interface gradients.
The argument `surface` is either `:interior` or `:exterior`, which denotes whether
we are computing interface gradients on boundaries which are interior (exterior resp.)
to the _parallel_ boundary.
"""
function launch_interface_gradients!(
    spacedisc,
    state_prognostic,
    t;
    surface::Symbol,
    dependencies,
)
    @assert surface === :interior || surface === :exterior
    # XXX: This is until FVM with DG hyperdiffusion is implemented
    if spacedisc isa DGFVModel
        @assert 0 == number_states(spacedisc.balance_law, Hyperdiffusive())
        Qhypervisc_grad_data = nothing
    elseif spacedisc isa DGModel
        Qhypervisc_grad_data = spacedisc.states_higher_order[1].data
    end

    FT = eltype(state_prognostic)

    info = basic_launch_info(spacedisc)
    comp_stream = dependencies

    # If the model direction is EveryDirection, we need to perform
    # both horizontal AND vertical kernel calls; otherwise, we only
    # call the kernel corresponding to the model direction `spacedisc.diffusion_direction`
    if spacedisc.diffusion_direction isa EveryDirection ||
       spacedisc.diffusion_direction isa HorizontalDirection

        workgroup = info.Nfp_v
        if surface === :interior
            elems = spacedisc.grid.interiorelems
            ndrange = workgroup * info.ninteriorelem
        else
            elems = spacedisc.grid.exteriorelems
            ndrange = workgroup * info.nexteriorelem
        end

        # Hoirzontal polynomial order (assumes same for both horizontal directions)
        horizontal_polyorder = info.N[1]

        comp_stream = dgsem_interface_gradients!(info.device, workgroup)(
            spacedisc.balance_law,
            Val(info),
            HorizontalDirection(),
            spacedisc.numerical_flux_gradient,
            state_prognostic.data,
            spacedisc.state_gradient_flux.data,
            Qhypervisc_grad_data,
            spacedisc.state_auxiliary.data,
            spacedisc.grid.vgeo,
            spacedisc.grid.sgeo,
            t,
            spacedisc.grid.vmap⁻,
            spacedisc.grid.vmap⁺,
            spacedisc.grid.elemtobndy,
            Val(hyperdiff_indexmap(spacedisc.balance_law, FT)),
            elems;
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end

    # Vertical interface kernel call
    if spacedisc.diffusion_direction isa EveryDirection ||
       spacedisc.diffusion_direction isa VerticalDirection

        workgroup = info.Nfp_h
        if surface === :interior
            elems = spacedisc.grid.interiorelems
            ndrange = workgroup * info.ninteriorelem
        else
            elems = spacedisc.grid.exteriorelems
            ndrange = workgroup * info.nexteriorelem
        end

        # Vertical polynomial degree
        vertical_polyorder = info.N[info.dim]

        if spacedisc isa DGModel
            comp_stream = dgsem_interface_gradients!(info.device, workgroup)(
                spacedisc.balance_law,
                Val(info),
                VerticalDirection(),
                spacedisc.numerical_flux_gradient,
                state_prognostic.data,
                spacedisc.state_gradient_flux.data,
                Qhypervisc_grad_data,
                spacedisc.state_auxiliary.data,
                spacedisc.grid.vgeo,
                spacedisc.grid.sgeo,
                t,
                spacedisc.grid.vmap⁻,
                spacedisc.grid.vmap⁺,
                spacedisc.grid.elemtobndy,
                Val(hyperdiff_indexmap(spacedisc.balance_law, FT)),
                elems;
                ndrange = ndrange,
                dependencies = comp_stream,
            )
        elseif spacedisc isa DGFVModel
            # Make sure FVM in the vertical
            @assert info.N[info.dim] == 0

            # The FVM will only work on stacked grids!
            @assert isstacked(spacedisc.grid.topology)
            nvertelem = spacedisc.grid.topology.stacksize
            periodicstack = spacedisc.grid.topology.periodicstack

            # 1 thread per degree freedom per element
            comp_stream = vert_fvm_interface_gradients!(info.device, workgroup)(
                spacedisc.balance_law,
                Val(info),
                Val(nvertelem),
                Val(periodicstack),
                VerticalDirection(),
                state_prognostic.data,
                spacedisc.state_gradient_flux.data,
                spacedisc.state_auxiliary.data,
                spacedisc.grid.vgeo,
                spacedisc.grid.sgeo,
                t,
                spacedisc.grid.elemtobndy,
                elems,
                # If we are computing in every direction, we need to
                # increment after we compute the horizontal values
                spacedisc.direction isa EveryDirection,
                ndrange = ndrange,
                dependencies = comp_stream,
            )
        else
            error("unknown spatial discretization: $(typeof(spacedisc))")
        end
    end
    return comp_stream
end

"""
    launch_volume_divergence_of_gradients!(dg, state_prognostic, t; dependencies)

Launches horizontal and vertical volume kernels for computing the divergence of gradients.
"""
function launch_volume_divergence_of_gradients!(
    dg,
    state_prognostic,
    t;
    dependencies,
)
    Qhypervisc_grad, Qhypervisc_div = dg.states_higher_order

    info = basic_launch_info(dg)
    workgroup = (info.Nq[1], info.Nq[2])
    ndrange = (info.nrealelem * info.Nq[1], info.Nq[2])
    comp_stream = dependencies

    # If the model direction is EveryDirection, we need to perform
    # both horizontal AND vertical kernel calls; otherwise, we only
    # call the kernel corresponding to the model direction `dg.diffusion_direction`
    if dg.diffusion_direction isa EveryDirection ||
       dg.diffusion_direction isa HorizontalDirection

        # Horizontal polynomial order and differentiation matrix
        horizontal_polyorder = info.N[1]
        horizontal_D = dg.grid.D[1]

        comp_stream = volume_divergence_of_gradients!(info.device, workgroup)(
            dg.balance_law,
            Val(info),
            HorizontalDirection(),
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            dg.grid.vgeo,
            horizontal_D,
            dg.grid.topology.realelems;
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end

    # And now the vertical kernel call
    if dg.diffusion_direction isa EveryDirection ||
       dg.diffusion_direction isa VerticalDirection

        # Vertical polynomial order and differentiation matrix
        vertical_polyorder = info.N[info.dim]
        vertical_D = dg.grid.D[info.dim]

        comp_stream = volume_divergence_of_gradients!(info.device, workgroup)(
            dg.balance_law,
            Val(info),
            VerticalDirection(),
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            dg.grid.vgeo,
            vertical_D,
            dg.grid.topology.realelems,
            # If we are computing the volume gradient in every direction, we
            # need to increment into the appropriate fields _after_ the
            # horizontal computation.
            !(dg.diffusion_direction isa VerticalDirection);
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end
    return comp_stream
end

"""
    launch_interface_divergence_of_gradients!(dg, state_prognostic, t; surface::Symbol, dependencies)

Launches horizontal and vertical interface kernels for computing the divergence of gradients.
The argument `surface` is either `:interior` or `:exterior`, which denotes whether
we are computing values on boundaries which are interior (exterior resp.)
to the _parallel_ boundary.
"""
function launch_interface_divergence_of_gradients!(
    dg,
    state_prognostic,
    t;
    surface::Symbol,
    dependencies,
)
    @assert surface === :interior || surface === :exterior
    Qhypervisc_grad, Qhypervisc_div = dg.states_higher_order

    info = basic_launch_info(dg)
    comp_stream = dependencies

    # If the model direction is EveryDirection, we need to perform
    # both horizontal AND vertical kernel calls; otherwise, we only
    # call the kernel corresponding to the model direction `dg.diffusion_direction`
    if dg.diffusion_direction isa EveryDirection ||
       dg.diffusion_direction isa HorizontalDirection

        workgroup = info.Nfp_v
        if surface === :interior
            elems = dg.grid.interiorelems
            ndrange = info.Nfp_v * info.ninteriorelem
        else
            elems = dg.grid.exteriorelems
            ndrange = info.Nfp_v * info.nexteriorelem
        end

        # Hoirzontal polynomial order (assumes same for both horizontal directions)
        horizontal_polyorder = info.N[1]

        comp_stream =
            interface_divergence_of_gradients!(info.device, workgroup)(
                dg.balance_law,
                Val(info),
                HorizontalDirection(),
                CentralNumericalFluxDivergence(),
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                dg.state_auxiliary.data,
                dg.grid.vgeo,
                dg.grid.sgeo,
                dg.grid.vmap⁻,
                dg.grid.vmap⁺,
                dg.grid.elemtobndy,
                t,
                elems;
                ndrange = ndrange,
                dependencies = comp_stream,
            )
    end

    # Vertical kernel call
    if dg.diffusion_direction isa EveryDirection ||
       dg.diffusion_direction isa VerticalDirection

        workgroup = info.Nfp_h
        if surface === :interior
            elems = dg.grid.interiorelems
            ndrange = info.Nfp_h * info.ninteriorelem
        else
            elems = dg.grid.exteriorelems
            ndrange = info.Nfp_h * info.nexteriorelem
        end

        # Vertical polynomial degree
        vertical_polyorder = info.N[info.dim]

        comp_stream =
            interface_divergence_of_gradients!(info.device, workgroup)(
                dg.balance_law,
                Val(info),
                VerticalDirection(),
                CentralNumericalFluxDivergence(),
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                dg.state_auxiliary.data,
                dg.grid.vgeo,
                dg.grid.sgeo,
                dg.grid.vmap⁻,
                dg.grid.vmap⁺,
                dg.grid.elemtobndy,
                t,
                elems;
                ndrange = ndrange,
                dependencies = comp_stream,
            )
    end

    return comp_stream
end

"""
    launch_volume_gradients_of_laplacians!(dg, state_prognostic, t; dependencies)

Launches horizontal and vertical volume kernels for computing the DG gradient of
a second-order DG gradient (Laplacian).
"""
function launch_volume_gradients_of_laplacians!(
    dg,
    state_prognostic,
    t;
    dependencies,
)
    Qhypervisc_grad, Qhypervisc_div = dg.states_higher_order

    info = basic_launch_info(dg)
    workgroup = (info.Nq[1], info.Nq[2])
    ndrange = (info.nrealelem * info.Nq[1], info.Nq[2])
    comp_stream = dependencies

    # If the model direction is EveryDirection, we need to perform
    # both horizontal AND vertical kernel calls; otherwise, we only
    # call the kernel corresponding to the model direction `dg.diffusion_direction`
    if dg.diffusion_direction isa EveryDirection ||
       dg.diffusion_direction isa HorizontalDirection

        # Horizontal polynomial degree
        horizontal_polyorder = info.N[1]
        # Horizontal quadrature weights and differentiation matrix
        horizontal_ω = dg.grid.ω[1]
        horizontal_D = dg.grid.D[1]

        comp_stream = volume_gradients_of_laplacians!(info.device, workgroup)(
            dg.balance_law,
            Val(info),
            HorizontalDirection(),
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            state_prognostic.data,
            dg.state_auxiliary.data,
            dg.grid.vgeo,
            horizontal_ω,
            horizontal_D,
            dg.grid.topology.realelems,
            t;
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end

    # Vertical kernel call
    if dg.diffusion_direction isa EveryDirection ||
       dg.diffusion_direction isa VerticalDirection

        # Vertical polynomial degree
        vertical_polyorder = info.N[info.dim]
        # Vertical quadrature weights and differentiation matrix
        vertical_ω = dg.grid.ω[info.dim]
        vertical_D = dg.grid.D[info.dim]

        comp_stream = volume_gradients_of_laplacians!(info.device, workgroup)(
            dg.balance_law,
            Val(info),
            VerticalDirection(),
            Qhypervisc_grad.data,
            Qhypervisc_div.data,
            state_prognostic.data,
            dg.state_auxiliary.data,
            dg.grid.vgeo,
            vertical_ω,
            vertical_D,
            dg.grid.topology.realelems,
            t,
            # If we are computing the volume gradient in every direction, we
            # need to increment into the appropriate fields _after_ the
            # horizontal computation.
            !(dg.diffusion_direction isa VerticalDirection);
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end

    return comp_stream
end

"""
    launch_interface_gradients_of_laplacians!(dg, state_prognostic, t; surface::Symbol, dependencies)

Launches horizontal and vertical interface kernels for computing the gradients of Laplacians
(second-order gradients). The argument `surface` is either `:interior` or `:exterior`,
which denotes whether we are computing values on boundaries which are interior (exterior resp.)
to the _parallel_ boundary.
"""
function launch_interface_gradients_of_laplacians!(
    dg,
    state_prognostic,
    t;
    surface::Symbol,
    dependencies,
)
    @assert surface === :interior || surface === :exterior
    Qhypervisc_grad, Qhypervisc_div = dg.states_higher_order
    comp_stream = dependencies
    info = basic_launch_info(dg)

    # If the model direction is EveryDirection, we need to perform
    # both horizontal AND vertical kernel calls; otherwise, we only
    # call the kernel corresponding to the model direction `dg.diffusion_direction`
    if dg.diffusion_direction isa EveryDirection ||
       dg.diffusion_direction isa HorizontalDirection

        workgroup = info.Nfp_v
        if surface === :interior
            elems = dg.grid.interiorelems
            ndrange = info.Nfp_v * info.ninteriorelem
        else
            elems = dg.grid.exteriorelems
            ndrange = info.Nfp_v * info.nexteriorelem
        end

        # Hoirzontal polynomial order (assumes same for both horizontal directions)
        horizontal_polyorder = info.N[1]

        comp_stream =
            interface_gradients_of_laplacians!(info.device, workgroup)(
                dg.balance_law,
                Val(info),
                HorizontalDirection(),
                CentralNumericalFluxHigherOrder(),
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                state_prognostic.data,
                dg.state_auxiliary.data,
                dg.grid.vgeo,
                dg.grid.sgeo,
                dg.grid.vmap⁻,
                dg.grid.vmap⁺,
                dg.grid.elemtobndy,
                elems,
                t;
                ndrange = ndrange,
                dependencies = comp_stream,
            )
    end

    # Vertical kernel call
    if dg.diffusion_direction isa EveryDirection ||
       dg.diffusion_direction isa VerticalDirection

        workgroup = info.Nfp_h
        if surface === :interior
            elems = dg.grid.interiorelems
            ndrange = info.Nfp_h * info.ninteriorelem
        else
            elems = dg.grid.exteriorelems
            ndrange = info.Nfp_h * info.nexteriorelem
        end

        # Vertical polynomial degree
        vertical_polyorder = info.N[info.dim]

        comp_stream =
            interface_gradients_of_laplacians!(info.device, workgroup)(
                dg.balance_law,
                Val(info),
                VerticalDirection(),
                CentralNumericalFluxHigherOrder(),
                Qhypervisc_grad.data,
                Qhypervisc_div.data,
                state_prognostic.data,
                dg.state_auxiliary.data,
                dg.grid.vgeo,
                dg.grid.sgeo,
                dg.grid.vmap⁻,
                dg.grid.vmap⁺,
                dg.grid.elemtobndy,
                elems,
                t;
                ndrange = ndrange,
                dependencies = comp_stream,
            )
    end

    return comp_stream
end

"""
    launch_volume_tendency!(spacedisc, state_prognostic, t; dependencies)

Launches horizontal and vertical volume kernels for computing tendencies (sources, sinks, etc).
"""
function launch_volume_tendency!(
    spacedisc,
    tendency,
    state_prognostic,
    t,
    α,
    β;
    dependencies,
)
    # XXX: This is until FVM with hyperdiffusion is implemented
    if spacedisc isa DGFVModel
        @assert 0 == number_states(spacedisc.balance_law, Hyperdiffusive())
        Qhypervisc_grad_data = nothing
    elseif spacedisc isa DGModel
        Qhypervisc_grad_data = spacedisc.states_higher_order[1].data
    end
    grad_flux_data = spacedisc.state_gradient_flux.data

    # Workgroup is determined by the number of quadrature points
    # in the horizontal direction. For each horizontal quadrature
    # point, we operate on a stack of quadrature in the vertical
    # direction. (Iteration space is in the horizontal)
    info = basic_launch_info(spacedisc)

    # We assume (in 3-D) that both x and y directions
    # are discretized using the same polynomial order, Nq[1] == Nq[2].
    # In 2-D, the workgroup spans the entire set of quadrature points:
    # Nq[1] * Nq[2]
    workgroup = (info.Nq[1], info.Nq[2])
    ndrange = (info.Nq[1] * info.nrealelem, info.Nq[2])
    comp_stream = dependencies

    # If the model direction is EveryDirection, we need to perform
    # both horizontal AND vertical kernel calls; otherwise, we only
    # call the kernel corresponding to the model direction
    # `spacedisc.diffusion_direction`
    if spacedisc.direction isa EveryDirection ||
       spacedisc.direction isa HorizontalDirection

        # Horizontal polynomial degree
        horizontal_polyorder = info.N[1]
        # Horizontal quadrature weights and differentiation matrix
        horizontal_ω = spacedisc.grid.ω[1]
        horizontal_D = spacedisc.grid.D[1]

        comp_stream = volume_tendency!(info.device, workgroup)(
            spacedisc.balance_law,
            Val(info),
            spacedisc.direction,
            HorizontalDirection(),
            tendency.data,
            state_prognostic.data,
            grad_flux_data,
            Qhypervisc_grad_data,
            spacedisc.state_auxiliary.data,
            spacedisc.grid.vgeo,
            t,
            horizontal_ω,
            horizontal_D,
            spacedisc.grid.topology.realelems,
            α,
            β,
            # If the model direction is horizontal or FV in the vertical,
            # we want to be sure to add sources
            spacedisc.direction isa HorizontalDirection ||
            spacedisc isa DGFVModel,
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end

    # Vertical kernel
    if spacedisc isa DGModel && (
        spacedisc.direction isa EveryDirection ||
        spacedisc.direction isa VerticalDirection
    )

        # Vertical polynomial degree
        vertical_polyorder = info.N[info.dim]
        # Vertical quadrature weights and differentiation matrix
        vertical_ω = spacedisc.grid.ω[info.dim]
        vertical_D = spacedisc.grid.D[info.dim]

        comp_stream = volume_tendency!(info.device, workgroup)(
            spacedisc.balance_law,
            Val(info),
            spacedisc.direction,
            VerticalDirection(),
            tendency.data,
            state_prognostic.data,
            grad_flux_data,
            Qhypervisc_grad_data,
            spacedisc.state_auxiliary.data,
            spacedisc.grid.vgeo,
            t,
            vertical_ω,
            vertical_D,
            spacedisc.grid.topology.realelems,
            α,
            # If we are computing the volume gradient in every direction, we
            # need to increment into the appropriate fields _after_ the
            # horizontal computation.
            spacedisc.direction isa EveryDirection ? true : β,
            # Boolean to add source. In the case of EveryDirection, we always add the sources
            # in the vertical kernel. Here, we make the assumption that we're either computing
            # in every direction, or _just_ the vertical direction.
            true;
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end

    return comp_stream
end

"""
    launch_interface_tendency!(spacedisc, state_prognostic, t; surface::Symbol, dependencies)

Launches horizontal and vertical interface kernels for computing tendencies (sources, sinks, etc).
The argument `surface` is either `:interior` or `:exterior`, which denotes whether we are computing
values on boundaries which are interior (exterior resp.) to the _parallel_ boundary.
"""
function launch_interface_tendency!(
    spacedisc,
    tendency,
    state_prognostic,
    t,
    α,
    β;
    surface::Symbol,
    dependencies,
)
    @assert surface === :interior || surface === :exterior
    # XXX: This is until FVM with diffusion is implemented
    if spacedisc isa DGFVModel
        @assert 0 == number_states(spacedisc.balance_law, Hyperdiffusive())
        Qhypervisc_grad_data = nothing
    elseif spacedisc isa DGModel
        Qhypervisc_grad_data = spacedisc.states_higher_order[1].data
    end
    grad_flux_data = spacedisc.state_gradient_flux.data
    numerical_flux_second_order = spacedisc.numerical_flux_second_order

    info = basic_launch_info(spacedisc)
    comp_stream = dependencies

    # If the model direction is EveryDirection, we need to perform
    # both horizontal AND vertical kernel calls; otherwise, we only
    # call the kernel corresponding to the model direction
    # `spacedisc.diffusion_direction`
    if spacedisc.direction isa EveryDirection ||
       spacedisc.direction isa HorizontalDirection

        workgroup = info.Nfp_v
        if surface === :interior
            elems = spacedisc.grid.interiorelems
            ndrange = workgroup * info.ninteriorelem
        else
            elems = spacedisc.grid.exteriorelems
            ndrange = workgroup * info.nexteriorelem
        end

        # Hoirzontal polynomial order (assumes same for both horizontal
        # directions)
        horizontal_polyorder = info.N[1]

        comp_stream = dgsem_interface_tendency!(info.device, workgroup)(
            spacedisc.balance_law,
            Val(info),
            HorizontalDirection(),
            spacedisc.numerical_flux_first_order,
            numerical_flux_second_order,
            tendency.data,
            state_prognostic.data,
            grad_flux_data,
            Qhypervisc_grad_data,
            spacedisc.state_auxiliary.data,
            spacedisc.grid.vgeo,
            spacedisc.grid.sgeo,
            t,
            spacedisc.grid.vmap⁻,
            spacedisc.grid.vmap⁺,
            spacedisc.grid.elemtobndy,
            elems,
            α;
            ndrange = ndrange,
            dependencies = comp_stream,
        )
    end

    # Vertical kernel call
    if spacedisc.direction isa EveryDirection ||
       spacedisc.direction isa VerticalDirection
        elems =
            surface === :interior ? elems = spacedisc.grid.interiorelems :
            spacedisc.grid.exteriorelems

        if spacedisc isa DGModel
            workgroup = info.Nfp_h
            ndrange = workgroup * length(elems)

            # Vertical polynomial degree
            vertical_polyorder = info.N[info.dim]

            comp_stream = dgsem_interface_tendency!(info.device, workgroup)(
                spacedisc.balance_law,
                Val(info),
                VerticalDirection(),
                spacedisc.numerical_flux_first_order,
                numerical_flux_second_order,
                tendency.data,
                state_prognostic.data,
                grad_flux_data,
                Qhypervisc_grad_data,
                spacedisc.state_auxiliary.data,
                spacedisc.grid.vgeo,
                spacedisc.grid.sgeo,
                t,
                spacedisc.grid.vmap⁻,
                spacedisc.grid.vmap⁺,
                spacedisc.grid.elemtobndy,
                elems,
                α;
                ndrange = ndrange,
                dependencies = comp_stream,
            )
        elseif spacedisc isa DGFVModel
            # Make sure FVM in the vertical
            @assert info.N[info.dim] == 0

            # The FVM will only work on stacked grids!
            @assert isstacked(spacedisc.grid.topology)

            # Figute out the stacking of the mesh
            nvertelem = spacedisc.grid.topology.stacksize
            nhorzelem = div(length(elems), nvertelem)
            periodicstack = spacedisc.grid.topology.periodicstack

            # 2-D workgroup
            workgroup = info.Nfp_h
            ndrange = workgroup * nhorzelem

            # XXX: This will need to be updated to diffusion
            comp_stream = vert_fvm_interface_tendency!(info.device, workgroup)(
                spacedisc.balance_law,
                Val(info),
                Val(nvertelem),
                Val(periodicstack),
                VerticalDirection(),
                spacedisc.fv_reconstruction,
                spacedisc.numerical_flux_first_order,
                numerical_flux_second_order,
                tendency.data,
                state_prognostic.data,
                grad_flux_data,
                spacedisc.state_auxiliary.data,
                spacedisc.grid.vgeo,
                spacedisc.grid.sgeo,
                t,
                spacedisc.grid.elemtobndy,
                elems,
                α,
                β,
                # If we are computing in every direction, we need to
                # increment after we compute the horizontal values
                spacedisc.direction isa EveryDirection,
                # If we are computing in vertical direction, we need to
                # add sources here
                spacedisc.direction isa VerticalDirection,
                ndrange = ndrange,
                dependencies = comp_stream,
            )
        else
            error("unknown spatial discretization: $(typeof(spacedisc))")
        end
    end

    return comp_stream
end

function fvm_balance!(
    balance_func,
    m::BalanceLaw,
    state_auxiliary::MPIStateArray,
    grid,
)
    device = array_device(state_auxiliary)

    dim = dimensionality(grid)
    N = polynomialorders(grid)
    Nq = N .+ 1
    Nqj = dim == 2 ? 1 : Nq[2]

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

    event = Event(device)
    event = kernel_fvm_balance!(device, (Nq[1], Nqj))(
        balance_func,
        m,
        Val(nvertelem),
        state_auxiliary.data,
        grid.vgeo,
        horzelems;
        ndrange = (length(horzelems) * Nq[1], Nqj),
        dependencies = (event,),
    )
    wait(device, event)
end


text:
export AbstractHook,
    ComposedHook,
    EmptyHook,
    StepsPerEpisode,
    RewardsPerEpisode,
    TotalRewardPerEpisode,
    TotalBatchRewardPerEpisode,
    CumulativeReward,
    TimePerStep,
    DoEveryNEpisode,
    DoEveryNStep

"""
A hook is called at different stage duiring a [`run`](@ref) to allow users to inject customized runtime logic.
By default, a `AbstractHook` will do nothing. One can override the behavior by implementing the following methods:

- `(hook::YourHook)(::PreActStage, agent, env, obs, action)`, note that there's an extra argument of `action`.
- `(hook::YourHook)(::PostActStage, agent, env, obs)`
- `(hook::YourHook)(::PreEpisodeStage, agent, env, obs)`
- `(hook::YourHook)(::PostEpisodeStage, agent, env, obs)`
"""
abstract type AbstractHook end

(hook::AbstractHook)(args...) = nothing

# https://github.com/JuliaLang/julia/issues/14919
# function (f::Function)(stage::T, args...;kw...) where T<: AbstractStage end

#####
# ComposedHook
#####

"""
    ComposedHook(hooks::AbstractHook...)

Compose different hooks into a single hook.
"""
struct ComposedHook{T<:Tuple} <: AbstractHook
    hooks::T
    ComposedHook(hooks...) = new{typeof(hooks)}(hooks)
end

function (hook::ComposedHook)(stage::AbstractStage, args...; kw...)
    for h in hook.hooks
        h(stage, args...; kw...)
    end
end

Base.getindex(hook::ComposedHook, inds...) = getindex(hook.hooks, inds...)

#####
# EmptyHook
#####

"""
Do nothing
"""
struct EmptyHook <: AbstractHook end

const EMPTY_HOOK = EmptyHook()

#####
# StepsPerEpisode
#####

"""
    StepsPerEpisode(; steps = Int[], count = 0, tag = "TRAINING")

Store steps of each episode in the field of `steps`.
"""
Base.@kwdef mutable struct StepsPerEpisode <: AbstractHook
    steps::Vector{Int} = Int[]
    count::Int = 0
    tag::String = "TRAINING"
end

function (hook::StepsPerEpisode)(::PostActStage, args...)
    hook.count += 1
end

function (hook::StepsPerEpisode)(
    ::Union{PostEpisodeStage,PostExperimentStage},
    agent,
    env,
    obs,
)
    push!(hook.steps, hook.count)
    hook.count = 0
    @debug hook.tag STEPS_PER_EPISODE = hook.steps[end]
end

#####
# RewardsPerEpisode
#####

"""
    RewardsPerEpisode(; rewards = Vector{Vector{Float64}}(), tag = "TRAINING")

Store each reward of each step in every episode in the field of `rewards`.
"""
Base.@kwdef mutable struct RewardsPerEpisode <: AbstractHook
    rewards::Vector{Vector{Float64}} = Vector{Vector{Float64}}()
    tag::String = "TRAINING"
end

function (hook::RewardsPerEpisode)(::PreEpisodeStage, agent, env, obs)
    push!(hook.rewards, [])
end

function (hook::RewardsPerEpisode)(::PostActStage, agent, env, obs)
    push!(hook.rewards[end], get_reward(obs))
end

function (hook::RewardsPerEpisode)(::PostEpisodeStage, agent, env, obs)
    @debug hook.tag REWARDS_PER_EPISODE = hook.rewards[end]
end

#####
# TotalRewardPerEpisode
#####

"""
    TotalRewardPerEpisode(; rewards = Float64[], reward = 0.0, tag = "TRAINING")

Store the total rewards of each episode in the field of `rewards`.
"""
Base.@kwdef mutable struct TotalRewardPerEpisode <: AbstractHook
    rewards::Vector{Float64} = Float64[]
    reward::Float64 = 0.0
    tag::String = "TRAINING"
end

function (hook::TotalRewardPerEpisode)(::PostActStage, agent, env, obs)
    hook.reward += get_reward(obs)
end

function (hook::TotalRewardPerEpisode)(
    ::Union{PostEpisodeStage,PostExperimentStage},
    agent,
    env,
    obs,
)
    push!(hook.rewards, hook.reward)
    hook.reward = 0
    @debug hook.tag REWARD_PER_EPISODE = hook.rewards[end]
end

#####
# TotalBatchRewardPerEpisode
#####
struct TotalBatchRewardPerEpisode <: AbstractHook
    rewards::Vector{Vector{Float64}}
    reward::Vector{Float64}
    tag::String
end

"""
    TotalBatchRewardPerEpisode(batch_size::Int;tag="TRAINING")

Similar to [`TotalRewardPerEpisode`](@ref), but will record total rewards per episode in [`BatchObs`](@ref).
"""
function TotalBatchRewardPerEpisode(batch_size::Int; tag = "TRAINING")
    TotalBatchRewardPerEpisode([Float64[] for _ in 1:batch_size], zeros(batch_size), tag)
end

function (hook::TotalBatchRewardPerEpisode)(::PostActStage, agent, env, obs::BatchObs)
    for i in 1:length(obs)
        hook.reward[i] += get_reward(obs[i])
        if get_terminal(obs[i])
            push!(hook.rewards[i], hook.reward[i])
            hook.reward[i] = 0.0
        end
    end
end

#####
# CumulativeReward
#####

"""
    CumulativeReward(rewards::Vector{Float64} = [0.0], tag::String = "TRAINING")

Store cumulative rewards since the beginning to the field of `rewards`.
"""
Base.@kwdef struct CumulativeReward <: AbstractHook
    rewards::Vector{Float64} = [0.0]
    tag::String = "TRAINING"
end

function (hook::CumulativeReward)(::PostActStage, agent, env, obs)
    push!(hook.rewards, get_reward(obs) + hook.rewards[end])
    @debug hook.tag CUMULATIVE_REWARD = hook.rewards[end]
end

#####
# TimePerStep
#####

"""
    TimePerStep(;max_steps=100)
    TimePerStep(times::CircularArrayBuffer{Float64}, t::UInt64)

Store time cost of the latest `max_steps` in the `times` field.
"""
mutable struct TimePerStep <: AbstractHook
    times::CircularArrayBuffer{Float64,1}
    t::UInt64
end

TimePerStep(; max_steps = 100) =
    TimePerStep(CircularArrayBuffer{Float64}(max_steps), time_ns())

function (hook::TimePerStep)(::PostActStage, agent, env, obs)
    push!(hook.times, (time_ns() - hook.t) / 1e9)
    hook.t = time_ns()
end

"""
    DoEveryNStep(f; n=1, t=0)

Execute `f(agent, env, obs)` every `n` step.
`t` is a counter of steps.
"""
Base.@kwdef mutable struct DoEveryNStep{F} <: AbstractHook
    f::F
    n::Int = 1
    t::Int = 0
end

DoEveryNStep(f, n = 1, t = 0) = DoEveryNStep(f, n, t)

function (hook::DoEveryNStep)(::PostActStage, agent, env, obs)
    hook.t += 1
    if hook.t % hook.n == 0
        hook.f(hook.t, agent, env, obs)
    end
end

"""
    DoEveryNEpisode(f; n=1, t=0)

Execute `f(agent, env, obs)` every `n` episode.
`t` is a counter of steps.
"""
Base.@kwdef mutable struct DoEveryNEpisode{F} <: AbstractHook
    f::F
    n::Int = 1
    t::Int = 0
end

DoEveryNEpisode(f, n = 1, t = 0) = DoEveryNEpisode(f, n, t)

function (hook::DoEveryNEpisode)(::PostEpisodeStage, agent, env, obs)
    hook.t += 1
    if hook.t % hook.n == 0
        hook.f(hook.t, agent, env, obs)
    end
end

ext: jl


text:
struct HamiltonianGSEnergyProblem{B, SM} <: HermitianMatrixProblem where {B<:Basis,
                                                 SM}
    HilbSpace::B # 0
    H::SM
    ρss
end

HamiltonianGSEnergyProblem(args...) = HamiltonianGSEnergyProblem(STD_REAL_PREC, args...)
HamiltonianGSEnergyProblem(T::Type{<:Number}, gl::GraphOperator; operators=true) = begin
    if operators
        return HamiltonianGSEnergyProblem(basis(gl), to_linear_operator(gl), 0.0)
    else
        return HamiltonianGSEnergyProblem(T, SparseOperator(gl))
    end
end
HamiltonianGSEnergyProblem(T::Type{<:Number}, Ham::SparseOperator) =
    HamiltonianGSEnergyProblem(Ham.basis_l, data(Ham), 0.0)

basis(prob::HamiltonianGSEnergyProblem) = prob.HilbSpace

function compute_Cloc(prob::HamiltonianGSEnergyProblem{B,SM}, net::KetNet, σ::State,
                      lnψ=net(σ), σp=deepcopy(σ)) where {B,SM<:SparseMatrixCSC}
    H = prob.H

    #### Now compute E(S) = Σₛ⟨s|Hψ⟩/⟨s|ψ⟩
    C_loc = zero(Complex{real(out_type(net))})
    # Iterate through all elements in row i_σ of the matrix computing
    # ⟨i_σ|H|ψ⟩ = Σ_{i_σp} ⟨i_σ|H|i_σp⟩⟨i_σp|ψ⟩
    # NOTE: H is CSC, but I would like a CSR matrix. Since it is hermitian I
    # can simply take the conjugate of the elements in the columns
    i_σ = index(σ)
    for row_id = H.colptr[i_σ]:(H.colptr[i_σ+1]-1)
      # Find nonzero elements s by doing <i_sp|ℒdagℒ|i_σ>
      i_σp = H.rowval[row_id]
      # BackConvert to int
      set_index!(σp, i_σp)
      # Compute the log(ψ(σ)/ψ(σ')), by only computing differences.
      log_ratio = logψ(net, σp) - lnψ
      # Conj because I am taking the transpose... see the note above.
      C_loc += conj(H.nzval[row_id]) * exp(log_ratio)
    end

    return C_loc
end

function compute_Cloc(prob::HamiltonianGSEnergyProblem{B,SM}, net::KetNet, σ::State,
                      lnψ=net(σ), σp=deepcopy(σ)) where {B,SM<:AbsLinearOperator}
    H = prob.H

    #### Now compute E(S) = Σₛ⟨s|H|ψ⟩/⟨s|ψ⟩
    C_loc = zero(Complex{real(out_type(net))})
    for op=operators(H)
        r = local_index(σ, sites(op))
        for (mel, changes)=op.op_conns[r]
            set_index!(σp, index(σ))
            for (site,val)=changes
                setat!(σp, site, val)
            end

            log_ratio = logψ(net, σp) - lnψ
            C_loc += mel * exp(log_ratio)
        end
    end

    return C_loc
end

# pretty printing
Base.show(io::IO, p::HamiltonianGSEnergyProblem) = print(io,
    """
    HamiltonianGSEnergyProblem: target minimum ground state energy
        - space : $(basis(p))
        - using operators : $(p.H isa AbsLinearOperator)""")


text:
# General vector utilities that may have application outside of MkCell.

"Calculate angle between two vectors (in radians)"
ang(x,y) = acos(dot(x,y)/(norm(x)*norm(y)))

"""Get lattice a, b, c from cell"""
cell_abc(cell) = norm(cell[1,:]), norm(cell[2,:]), norm(cell[3,:])

"""Get cell angles α, β, γ (in degrees) from a cell."""
cell_angles(cell) = tuple((ang(cell[p[1],:], cell[p[2],:])*180.0/π for p in [(2,3), (1,3), (1,2)])...)


text:
@testset "simple workflow" begin
    env = StateOverriddenEnv(CartPoleEnv{Float32}(), deepcopy)
    policy = RandomPolicy(action_space(env))
    N_EPISODE = 10_000
    hook = TotalRewardPerEpisode()
    run(policy, env, StopAfterEpisode(N_EPISODE), hook)

    @test isapprox(sum(hook[]) / N_EPISODE, 21; atol = 2)
end


text:
# Axis-aligned multivariate normal (i.e., independent entries, i.e., diagonal covariance matrix) using conjugate prior.
module MVNaaC

module MVNaaCmodel # submodule for component family definitions
export Theta, Data, log_marginal, new_theta, Theta_clear!, Theta_adjoin!, Theta_remove!,
       Hyperparameters, construct_hyperparameters, update_hyperparameters!

using Statistics
using SpecialFunctions
lgamma_(x) = logabsgamma(x)[1]

const Data = Array{Float64,1}

mutable struct Theta
    n::Int64                 # number of data points assigned to this cluster
    sum_x::Array{Float64,1}  # sum of the data points x assigned to this cluster
    sum_xx::Array{Float64,1} # sum of x.*x for the data points assigned to this cluster
    Theta(d) = (p=new(); p.n=0; p.sum_x=zeros(d); p.sum_xx=zeros(d); p)
end
new_theta(H) = Theta(H.d)
Theta_clear!(p) = (p.sum_x[:] .= 0.; p.sum_xx[:] .= 0.; p.n = 0)
Theta_adjoin!(p,x) = (for i=1:length(x); p.sum_x[i] += x[i]; p.sum_xx[i] += x[i]*x[i]; end; p.n += 1)
Theta_remove!(p,x) = (for i=1:length(x); p.sum_x[i] -= x[i]; p.sum_xx[i] -= x[i]*x[i]; end; p.n -= 1)

# In each dimension independently,
# X_1,...,X_n ~ Normal(mu,1/lambda) with Normal(mu|m,1/(c*lambda))Gamma(lambda|a,b) prior on mean=mu, precision=lambda.
function log_marginal(p,H)
    n = p.n
    LB = 0.0
    for i=1:H.d; LB += log(H.b + 0.5*p.sum_xx[i] - 0.5*p.sum_x[i]*p.sum_x[i]/n + 0.5*H.c*n*(p.sum_x[i]/n - H.m)^2/(H.c+n)); end
    return H.d*(H.constant - 0.5*n*log(2*pi) - 0.5*log(H.c+n) + H.log_Ga[n]) - (H.a+0.5*n)*LB
end

function log_marginal(x,p,H)
    Theta_adjoin!(p,x)
    result = log_marginal(p,H)
    Theta_remove!(p,x)
    return result
end

mutable struct Hyperparameters
    d::Int64    # dimension
    m::Float64  # prior mean of mu's
    c::Float64  # prior precision multiplier for mu's
    a::Float64  # prior shape of lambda's
    b::Float64  # prior rate of lambda's
    constant::Float64
    log_Ga::Array{Float64,1}
end

function construct_hyperparameters(options)
    x = options.x
    n = length(x)
    d = length(x[1])
    mu = mean(x)
    v = mean([xi.*xi for xi in x]) .- mu.*mu  # sample variance
    @assert(all(abs.(mu) .< 1e-10) && all(abs.(v .- 1.0) .< 1e-10), "Data must be normalized to zero mean, unit variance.")
    m = 0.0
    c = 1.0
    a = 1.0
    b = 1.0
    log_Ga = lgamma_.(a .+ 0.5*(1:n+1))
    constant = 0.5*log(c) + a*log(b) - lgamma_(a)
    return Hyperparameters(d,m,c,a,b,constant,log_Ga)
end

function update_hyperparameters!(H,theta,list,t,x,z)
    #error("update_hyperparameters is not yet implemented.")
end

end # module MVNaaCmodel
using .MVNaaCmodel

# Include generic code
include("generic.jl")

# Include core sampler code
include("coreConjugate.jl")

end # module MVNaaC


text:
""" Check https://arxiv.org/pdf/1707.03429.pdf for grammar specification
"""
module QASM
using RBNF
using PrettyPrint

struct QASMLang end

second((a, b)) = b
second(vec::V) where V <: AbstractArray = vec[2]

RBNF.@parser QASMLang begin
    # define ignorances
    ignore{space}

    @grammar
    # define grammars
    mainprogram := ["OPENQASM", ver=nnreal, ';', prog=program]
    program     = statement{*}
    statement   = (decl | gate | opaque | qop | ifstmt | barrier)
    # stmts
    ifstmt      := ["if", '(', l=id, "==", r=nninteger, ')', body=qop]
    opaque      := ["opaque", id=id, ['(', [arglist1=idlist].?, ')'].? , arglist2=idlist, ';']
    barrier     := ["barrier", value=mixedlist]
    decl        := [regtype="qreg" | "creg", id=id, '[', int=nninteger, ']', ';']

    # gate
    gate        := [decl=gatedecl, [goplist=goplist].?, '}']
    gatedecl    := ["gate", id=id, ['(', [arglist1=idlist].?, ')'].?, arglist2=idlist, '{']

    goplist     = (uop |barrier_ids){*}
    barrier_ids := ["barrier", ids=idlist, ';']
    # qop
    qop         = (uop | measure | reset)
    reset       := ["reset", arg=argument, ';']
    measure     := ["measure", arg1=argument, "->", arg2=argument, ';']

    uop         = (iduop | u | cx)
    iduop      := [op=id, ['(', [lst1=explist].?, ')'].?, lst2=mixedlist, ';']
    u          := ['U', '(', exprs=explist, ')', arg=argument, ';']
    cx         := ["CX", arg1=argument, ',', arg2=argument, ';']

    idlist     = @direct_recur begin
        init = id
        prefix = (recur, (',', id) % second)
    end

    mixeditem   := [id=id, ['[', arg=nninteger, ']'].?]
    mixedlist   = @direct_recur begin
        init = mixeditem
        prefix = (recur, (',', mixeditem) % second)
    end

    argument   := [id=id, ['[', (arg=nninteger), ']'].?]

    explist    = @direct_recur begin
        init = nnexp
        prefix = (recur,  (',', nnexp) % second)
    end

    atom       = (nnreal | nninteger | "pi" | id | fnexp) | (['(', nnexp, ')'] % second) | neg
    fnexp      := [fn=fn, '(', arg=nnexp, ')']
    neg        := ['-', value=nnexp]
    nnexp        = @direct_recur begin
        init = atom
        prefix = (recur, binop, atom)
    end
    fn         = ("sin" | "cos" | "tan" | "exp" | "ln" | "sqrt")
    binop      = ('+' | '-' | '*' | '/')

    # define tokens
    @token
    id        := r"\G[a-z]{1}[A-Za-z0-9_]*"
    nnreal    := r"\G([0-9]+\.[0-9]*|[0-9]*\.[0.9]+)([eE][-+]?[0-9]+)?"
    nninteger := r"\G([1-9]+[0-9]*|0)"
    space     := r"\G\s+"
end

src1 = """
OPENQASM 2.0;

gate cu1(lambda) a,b
{
    U(0,0,theta/2) a;
    CX a,b;
    U(0,0,-theta/2) b;
    CX a,b;
    U(0,0,theta/2) b;
}

qreg q[3];
qreg a[2];
creg c[3];
creg syn[2];
cu1(pi/2) q[0],q[1];
"""
ast, ctx = RBNF.runparser(mainprogram, RBNF.runlexer(QASMLang, src1))
pprint(ast)
end


text:
using RCall, DataFrames

"""Load characters data"""
function initialize(label_a::Int, label_b::Int)::Tuple{DataFrames.DataFrame,DataFrames.DataFrame}
    @rput label_a
    @rput label_b
    R"""
    load('./mnist_train.rda')
    mnist = subset(mnist_train, mnist_train[,1] %in% c(label_a, label_b))
    mnist_means = aggregate(mnist[,-1], by = list(mnist[,1]), FUN = mean)
    """
    @rget mnist
    @rget mnist_means
    return mnist, mnist_means
end

"""Single pixel classifier based on mean pixel value"""
function baseClassifier_j(mnist_means::DataFrames.DataFrame, pixel::Int, j::Int, correct::Int)
    if abs(pixel - mnist_means[1,j]) <= abs(pixel - mnist_means[2,j])
        label = mnist_means[1,1]
    else
        label = mnist_means[2,1]
    end
    return label, correct == label
end

"""Return which label is the most common in a vector"""
function findMostCommon(vector::Vector{Int}, a::Int, b::Int)
    count_a = 0; count_b = 0
    for el in vector
        if el == a
            count_a += 1
        else
            count_b += 1
        end
    end
    if count_a > count_b
        return a
    else
        return b
    end
end

function main(label_a::Int, label_b::Int, ϵ::Float64, γ::Float64)
    mnist, mnist_means = initialize(label_a, label_b)

    d = size(mnist, 2) - 1
    n = size(mnist, 1)
    w = ones(n)
    η = γ
    T = convert(Int, ceil(2/γ^2 * log(1/ϵ)))
    info("T = ", T)
    corrects = falses(n)
    labels = Vector{Int}(n)
    labels_per_classifier = Matrix{Int}(n, T)
    correct_final_label = mnist[:X5]
    final_label = Vector{Int64}(n);
    hits = falses(n)
    percent_hits = zeros(T)
    count_base_classifiers = zeros(d-1)

    for t in 1:T
        j = 1
        S = 0.0
        while S < 0.5 + γ
            p = w / sum(w)
            S = 0.0
            for i in 1:n
                label, correct_label = baseClassifier_j(mnist_means, Int64(mnist[i, j + 1]), j + 1, Int64(mnist[i, 1]))
                labels[i] = label
                corrects[i] = correct_label
            end

            S = sum(p .* corrects)

            j += 1
            if j > d
                warn("No valid weak learner found for iteration ", t)
            end
        end
        count_base_classifiers[j] = count_base_classifiers[j] + 1
        labels_per_classifier[:, t] = labels
        # Found weak learner corresponding to j-th baseClassifier
        # Update weights
        for i in 1:n
            w[i] = w[i] * exp(η * !corrects[i]) # Increase weights for those that were wrongly classified
        end
        p = w / sum(w)
        for i in 1:n
            final_label[i] = findMostCommon(labels_per_classifier[i,1:t], label_a, label_b)
            hits[i] = final_label[i] == correct_final_label[i]
        end

        percent_hits[t] = sum(hits)/length(hits)

        if t % 10 == 0
            info("t = ", t, " of ", T)
            info("Correct ratio: ", percent_hits[t])
        end
        if percent_hits[t] > 1-ϵ
          return percent_hits, count_base_classifiers
        end
    end
    return percentual_hits, count_base_classifiers
end

max_stars_repo_name: raphaelsaavedra/Boosting

max_issues_repo_name: raphaelsaavedra/Boosting

max_forks_repo_name: raphaelsaavedra/Boosting


text:
function _create_planar_push!(vis, model::PlanarPush;
        i = 1,
        r = 0.1,
        r_pusher = 0.025,
        tl = 1.0,
        box_color = Colors.RGBA(0.0, 0.0, 0.0, tl),
        pusher_color = Colors.RGBA(0.5, 0.5, 0.5, tl))

    r_box = r - r_pusher

    setobject!(vis["box_$i"], GeometryBasics.Rect(Vec(-1.0 * r_box,
		-1.0 * r_box,
		-1.0 * r_box),
		Vec(2.0 * r_box, 2.0 * r_box, 2.0 * r_box)),
		MeshPhongMaterial(color = box_color))

    setobject!(vis["pusher_$i"],
        Cylinder(Point(0.0, 0.0, 0.0), Point(0.0, 0.0, r_box), r_pusher),
        MeshPhongMaterial(color = pusher_color))
end

function _set_planar_push!(vis, model::PlanarPush, q;
    i = 1)
    settransform!(vis["box_$i"],
		compose(Translation(q[1], q[2], 0.01 * i), LinearMap(RotZ(q[3]))))
    settransform!(vis["pusher_$i"], Translation(q[4], q[5], 0.01 * i))
end

function visualize!(vis, model::PlanarPush, q;
        i = 1,
        r = 0.1,
        r_pusher = 0.025,
        tl = 1.0,
        box_color = Colors.RGBA(0.0, 0.0, 0.0, tl),
        pusher_color = Colors.RGBA(0.5, 0.5, 0.5, tl),
        Δt = 0.1)

	default_background!(vis)

    _create_planar_push!(vis, model,
        i = i,
        r = r,
        r_pusher = r_pusher,
        tl = tl,
        box_color = box_color,
        pusher_color = pusher_color)

    anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

	T = length(q)
    for t = 1:T-1
        MeshCat.atframe(anim, t) do
            _set_planar_push!(vis, model, q[t])
        end
    end

	settransform!(vis["/Cameras/default"],
    compose(Translation(0.0, 0.0, 50.0), LinearMap(RotZ(0.5 * pi) * RotY(-pi/2.5))))
    setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", 50)

    MeshCat.setanimation!(vis, anim)
end


text:
# Apply Dirichlet boundary condition
struct Dirichlet{T}  #<: Constraint
    prescribed_dofs::Vector{Int}
    values::Vector{T}
end

function Dirichlet(dh::DofHandler{2}, u::TrialFunction, edgeset::String,f::Union{Function,Real})
    element = getelement(getfunctionspace(u))
    edgeset = getedgeset(dh.mesh, edgeset)
    prescribed_dofs = Vector{Int}()
    values = Vector{Float64}()
    for edge_idx in edgeset
        cell_idx = edge_idx[1]  #cell idx
        cell = getcells(dh.mesh)[cell_idx]
        cell_topology = gettopology(dh.mesh,cell, element)
        edge_lidx = edge_idx[2]  # edge idx
        l_dof = Int[]
        offset::Int = dh.cell_dofs_offset[cell_idx] - 1 + field_offset(dh, u, cell_idx)
        for j = 1:2 #Add vertex dofs
            local_offset = reference_edge_vertices(dh.mesh,cell)[edge_lidx][j]
            if !(dh.cell_dofs[offset+local_offset] ∈ prescribed_dofs)
                push!(prescribed_dofs, dh.cell_dofs[offset+local_offset])
                push!(l_dof, local_offset)
            end
        end
        dofs_per_edge = Int(cell_topology[1]/getnedges(dh.mesh,cell))
        for j in 1:dofs_per_edge  #Add edge dofs
            local_offset::Int = cell_topology[0] + dofs_per_edge*(edge_lidx-1) + j
            if !(dh.cell_dofs[offset+local_offset] ∈ prescribed_dofs)
                push!(prescribed_dofs, dh.cell_dofs[offset+local_offset])
                push!(l_dof, local_offset)
            end
        end
        _push_values!(values, cell_idx, dh.mesh, l_dof, element, f)
    end
    #now put all in order
    p = sortperm(prescribed_dofs)
    return Dirichlet(prescribed_dofs[p], values[p])
end

function Dirichlet(dh::DofHandler{dim}, u::TrialFunction, facetset::String,f::Union{Function,Real}) where {dim}
    element = getelement(getfunctionspace(u))
    facetset = getfacetset(dh.mesh, facetset)
    prescribed_dofs = Vector{Int}()
    values = Vector{Float64}()
    for facet_idx in facetset
        cell_idx = facet_idx[1]  #cell idx
        cell = getcells(dh.mesh)[cell_idx]
        cell_topology = gettopology(dh.mesh,cell, element)
        facet_lidx = facet_idx[2]  # facet idx
        l_dof = Int[]
        offset::Int = dh.cell_dofs_offset[cell_idx] - 1 + field_offset(dh, u, cell_idx)
        for j = 1:getnfacetvertices(dh.mesh, cell_idx, facet_lidx) #Add vertex dofs
            local_offset = reference_facet_vertices(dh.mesh,cell)[facet_lidx][j]
            if !(dh.cell_dofs[offset+local_offset] ∈ prescribed_dofs)
                push!(prescribed_dofs, dh.cell_dofs[offset+local_offset])
                push!(l_dof, local_offset)
            end
        end
        if dim > 1
            dofs_per_edge = Int(cell_topology[1]/getnedges(dh.mesh,cell))
            nedges = dim == 2 ? 1 : getnedges(dh.mesh, cell)
            for k in 1:nedges
                edge_lidx = dim == 2 ? facet_lidx : k
                for j in 1:dofs_per_edge  #Add edge dofs
                    local_offset::Int = cell_topology[0] + dofs_per_edge*(edge_lidx-1) + j
                    if !(dh.cell_dofs[offset+local_offset] ∈ prescribed_dofs)
                        push!(prescribed_dofs, dh.cell_dofs[offset+local_offset])
                        push!(l_dof, local_offset)
                    end
                end
            end
        end
        if dim > 2
            dofs_per_face = Int(cell_topology[2]/getnfaces(dh.mesh,cell))
            for j in 1:dofs_per_face  #Add face dofs
                local_offset::Int = cell_topology[0] + cell_topology[1] + dofs_per_face*(facet_lidx-1) + j
                if !(dh.cell_dofs[offset+local_offset] ∈ prescribed_dofs)
                    push!(prescribed_dofs, dh.cell_dofs[offset+local_offset])
                    push!(l_dof, local_offset)
                end
            end
        end
        _push_values!(values, cell_idx, dh.mesh, l_dof, element, f)
    end
    #now put all in order
    p = sortperm(prescribed_dofs)
    return Dirichlet(prescribed_dofs[p], values[p])
end

function _push_values!(values::Vector, cell::Int, mesh,l_dof::Vector{Int}, felem::AbstractGalerkinElement, f::Function)
    for i in l_dof
        vals = f(spatial_nodal_coordinate(mesh,cell,felem,i))
        push!(values,vals)
    end
end

function _push_values!(values::Vector, cell::Int, mesh,l_dof::Vector{Int}, felem::AbstractGalerkinElement, vals::Real)
    for i in l_dof
        push!(values,vals)
    end
end

@enum(ApplyStrategy, APPLY_TRANSPOSE, APPLY_INPLACE)

function apply!(KK::Union{SparseMatrixCSC,Symmetric}, f::AbstractVector, dirichlet::Dirichlet;
                strategy::ApplyStrategy=APPLY_TRANSPOSE)
    K = isa(KK, Symmetric) ? KK.data : KK
    @assert length(f) == 0 || length(f) == size(K, 1)
    @boundscheck checkbounds(K, dirichlet.prescribed_dofs, dirichlet.prescribed_dofs)
    @boundscheck length(f) == 0 || checkbounds(f, dirichlet.prescribed_dofs)

    m = meandiag(K) # Use the mean of the diagonal here to not ruin things for iterative solver
    @inbounds for i in 1:length(dirichlet.values)
        d = dirichlet.prescribed_dofs[i]
        v = dirichlet.values[i]

        if v != 0
            for j in nzrange(K, d)
                f[K.rowval[j]] -= v * K.nzval[j]
            end
        end
    end
    zero_out_columns!(K, dirichlet.prescribed_dofs)
    if strategy == APPLY_TRANSPOSE
        K′ = copy(K)
        transpose!(K′, K)
        zero_out_columns!(K′, dirichlet.prescribed_dofs)
        transpose!(K, K′)
    elseif strategy == APPLY_INPLACE
        K[dirichlet.prescribed_dofs, :] = 0
    else
        error("Unknown apply strategy")
    end
    @inbounds for i in 1:length(dirichlet.values)
        d = dirichlet.prescribed_dofs[i]
        v = dirichlet.values[i]
        K[d, d] = m
        if length(f) != 0
            f[d] = v * m
        end
    end
end

# columns need to be stored entries, this is not checked
function zero_out_columns!(K, dofs::Vector{Int}) # can be removed in 0.7 with #24711 merged
    #@debug assert(issorted(dofs))
    for col in dofs
        r = nzrange(K, col)
        K.nzval[r] .= 0.0
    end
end

function meandiag(K::AbstractMatrix)
    z = zero(eltype(K))
    for i in 1:size(K, 1)
        z += abs(K[i, i])
    end
    return z / size(K, 1)
end

max_stars_repo_name: Paulms/jFEMToools

max_issues_repo_name: Paulms/jFEMToools

max_forks_repo_name: Paulms/jFEMToools


text:
module PlotsUnitTests

using Test
using UnitfulAstrodynamics
using UnitfulAstrodynamics.TwoBody.Systems: Earth
@testset "OrbitPlots" begin

    # Twobody Orbit
    r̅ = [0.0, 11681, 0.0] * u"km"
    v̅ = [5.134, 4.226, 2.787] * u"km/s"
    orbit = CartesianOrbit(r̅, v̅, Earth, 0.0u"s")

    # Does plot run?
    fig = orbitplot(propagate(orbit, 1u"s"))

    @test true

end

end