text:
#######
# Demo setup

cool(x) = x + 1
cool(x, y) = x + y + 1

# a rule we define so we can test rules
dummy_identity(x) = x
@scalar_rule(dummy_identity(x), true)

nice(x) = 1
@scalar_rule(nice(x), ZeroTangent())

sum_two(x, y) = x + y
@scalar_rule(sum_two(x, y), (true, true))

complex_times(x) = (1 + 2im) * x
@scalar_rule(complex_times(x), 1 + 2im)

# Tests that aim to ensure that the API for frules doesn't regress and make these things
# hard to implement.

varargs_function(x...) = sum(x)
function ChainRulesCore.frule(dargs, ::typeof(varargs_function), x...)
    Δx = Base.tail(dargs)
    return sum(x), sum(Δx)
end

mixed_vararg(x, y, z...) = x + y + sum(z)
function ChainRulesCore.frule(
    dargs::Tuple{Any, Any, Any, Vararg},
    ::typeof(mixed_vararg), x, y, z...,
)
    Δx = dargs[2]
    Δy = dargs[3]
    Δz = dargs[4:end]
    return mixed_vararg(x, y, z...), Δx + Δy + sum(Δz)
end

type_constraints(x::Int, y::Float64) = x + y
function ChainRulesCore.frule(
    (_, Δx, Δy)::Tuple{Any, Int, Float64},
    ::typeof(type_constraints), x::Int, y::Float64,
)
    return type_constraints(x, y), Δx + Δy
end

mixed_vararg_type_constaint(x::Float64, y::Real, z::Vararg{Float64}) = x + y + sum(z)
function ChainRulesCore.frule(
    dargs::Tuple{Any, Float64, Real, Vararg{Float64}},
    ::typeof(mixed_vararg_type_constaint), x::Float64, y::Real, z::Vararg{Float64},
)
    Δx = dargs[2]
    Δy = dargs[3]
    Δz = dargs[4:end]
    return x + y + sum(z), Δx + Δy + sum(Δz)
end

ChainRulesCore.frule(dargs, ::typeof(Core._apply), f, x...) = frule(dargs[2:end], f, x...)

#######

_second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))

@testset "frule and rrule" begin
    dself = ZeroTangent()
    @test frule((dself, 1), cool, 1) === nothing
    @test frule((dself, 1), cool, 1; iscool=true) === nothing
    @test rrule(cool, 1) === nothing
    @test rrule(cool, 1; iscool=true) === nothing

    # add some methods:
    ChainRulesCore.@scalar_rule(Main.cool(x), one(x))
    @test hasmethod(rrule, Tuple{typeof(cool),Number})
    ChainRulesCore.@scalar_rule(Main.cool(x::String), "wow such dfdx")
    @test hasmethod(rrule, Tuple{typeof(cool),String})
    # Ensure those are the *only* methods that have been defined
    cool_methods = Set(m.sig for m in methods(rrule) if _second(m.sig) == typeof(cool))
    only_methods = Set([Tuple{typeof(rrule),typeof(cool),Number},
                        Tuple{typeof(rrule),typeof(cool),String}])
    @test cool_methods == only_methods

    frx, cool_pushforward = frule((dself, 1), cool, 1)
    @test frx === 2
    @test cool_pushforward === 1
    rrx, cool_pullback = rrule(cool, 1)
    self, rr1 = cool_pullback(1)
    @test self === NoTangent()
    @test rrx === 2
    @test rr1 == 1.0   # projection may make this ===

    frx, nice_pushforward = frule((dself, 1), nice, 1)
    @test nice_pushforward === ZeroTangent()
    rrx, nice_pullback = rrule(nice, 1)
    @test (NoTangent(), ZeroTangent()) === nice_pullback(1)

    # Test that these run. Do not care about numerical correctness.
    @test frule((nothing, 1.0, 1.0, 1.0), varargs_function, 0.5, 0.5, 0.5) == (1.5, 3.0)

    @test frule((nothing, 1.0, 2.0, 3.0, 4.0), mixed_vararg, 1.0, 2.0, 3.0, 4.0) == (10.0, 10.0)

    @test frule((nothing, 3, 2.0), type_constraints, 5, 4.0) == (9.0, 5.0)
    @test frule((nothing, 3.0, 2.0im), type_constraints, 5, 4.0) == nothing

    @test(frule(
        (nothing, 3.0, 2.0, 1.0, 0.0),
        mixed_vararg_type_constaint, 3.0, 2.0, 1.0, 0.0,
    ) == (6.0, 6.0))

    # violates type constraints, thus an frule should not be found.
    @test frule(
        (nothing, 3, 2.0, 1.0, 5.0),
        mixed_vararg_type_constaint, 3, 2.0, 1.0, 0,
    ) == nothing

    @test frule((nothing, nothing, 5.0), Core._apply, dummy_identity, 4.0) == (4.0, 5.0)

    @testset "broadcasting true" begin
        sx = @SVector [1, 2]
        sy = @SVector [3, 4]

        # Test that @scalar_rule and `true` play nice together, w.r.t broadcasting
        @inferred frule((ZeroTangent(), sx, sy), sum_two, 1, 2)
    end

    @testset "complex numbers" begin
        x, ẋ, Ω̄ = randn(ComplexF64, 3)
        Ω = complex_times(x)

        # forwards
        Ω_fwd, Ω̇ = frule((nothing, ẋ), complex_times, x)
        @test Ω_fwd == Ω
        @test Ω̇ ≈ (1 + 2im) * ẋ

        # reverse
        Ω_rev, back = rrule(complex_times, x)
        @test Ω_rev == Ω
        ∂self, ∂x = back(Ω̄)
        @test ∂self == NoTangent()
        @test ∂x ≈ (1 - 2im) * Ω̄

        # real argument, complex output
        xr = rand()
        Ωr = complex_times(xr)
        Ωr_rev, backr = rrule(complex_times, xr)
        ∂selfr, ∂xr = backr(Ω̄)
        @test_skip ∂xr isa Float64  # to be made true with projection
        @test_skip ∂xr ≈ real(∂x)
    end

    @testset "@opt_out" begin
        first_oa(x, y) = x
        @scalar_rule(first_oa(x, y), (1, 0))
        @opt_out ChainRulesCore.rrule(::typeof(first_oa), x::T, y::T) where T<:Float32
        @opt_out(
            ChainRulesCore.frule(::Any, ::typeof(first_oa), x::T, y::T) where T<:Float32
        )

        @testset "rrule" begin
            @test rrule(first_oa, 3.0, 4.0)[2](1) == (NoTangent(), 1, 0)
            @test rrule(first_oa, 3f0, 4f0) === nothing

            @test !isempty(Iterators.filter(methods(ChainRulesCore.no_rrule)) do m
                m.sig <:Tuple{Any, typeof(first_oa), T, T} where T<:Float32
            end)
        end

        @testset "frule" begin
            @test frule((NoTangent(), 1,0), first_oa, 3.0, 4.0) == (3.0, 1)
            @test frule((NoTangent(), 1,0), first_oa, 3f0, 4f0) === nothing

            @test !isempty(Iterators.filter(methods(ChainRulesCore.no_frule)) do m
                m.sig <:Tuple{Any, Any, typeof(first_oa), T, T} where T<:Float32
            end)
        end
    end
end


text:
using Primes
issemiprime(n::Integer) = sum(values(factor(n))) == 2
@show filter(issemiprime, 1:100)

ext: jl

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
import MPI, HDF5
import Test: @testset, @test, @test_throws

include("../../src/main.jl")

@testset "main" begin
    @test main(["hello"]) == 0
    if HDF5.has_parallel()
        @test_exists_and_rm("hello.h5")
    else
        @test_exists_and_rm("hello_0.h5")
    end
    filename = "testcase-temp-config.yaml"
    @test main(["-c", filename, "generate-config"]) == 0
    @test_exists_and_rm(filename)
end


text:
"""
    initialize_sediment_model(config::Config)

Initial part of the sediment model concept. Reads the input settings and data as defined in the
Config object. Will return a Model that is ready to run.
"""
function initialize_sediment_model(config::Config)

    # unpack the paths to the NetCDF files
    tomldir = dirname(config)
    static_path = joinpath(tomldir, config.input.path_static)

    reader = prepare_reader(config)
    clock = Clock(config, reader)
    Δt = clock.Δt

    reinit = get(config.model, "reinit", true)::Bool

    do_river = get(config.model, "runrivermodel", false)::Bool

    nc = NCDataset(static_path)
    dims = dimnames(nc[param(config, "input.subcatchment")])

    subcatch_2d = ncread(nc, param(config, "input.subcatchment"); allow_missing = true)
    # indices based on catchment
    inds, rev_inds = active_indices(subcatch_2d, missing)
    n = length(inds)
    modelsize_2d = size(subcatch_2d)

    river_2d = ncread(nc, param(config, "input.river_location"); type = Bool, fill = false)
    river = river_2d[inds]
    riverwidth_2d =
        ncread(nc, param(config, "input.lateral.river.width"); type = Float, fill = 0)
    riverwidth = riverwidth_2d[inds]
    riverlength_2d =
        ncread(nc, param(config, "input.lateral.river.length"); type = Float, fill = 0)
    riverlength = riverlength_2d[inds]

    inds_riv, rev_inds_riv = active_indices(river_2d, 0)
    nriv = length(inds_riv)

    # Needed to update the forcing
    reservoir = ()
    lake = ()

    # read x, y coordinates and calculate cell length [m]
    y_nc = read_y_axis(nc)
    x_nc = read_x_axis(nc)
    y = permutedims(repeat(y_nc, outer = (1, length(x_nc))))[inds]
    cellength = abs(mean(diff(x_nc)))

    sizeinmetres = get(config.model, "sizeinmetres", false)::Bool
    xl, yl = cell_lengths(y, cellength, sizeinmetres)
    riverfrac = river_fraction(river, riverlength, riverwidth, xl, yl)

    eros = initialize_landsed(nc, config, river, riverfrac, xl, yl, inds)

    ldd_2d = ncread(nc, param(config, "input.ldd"); allow_missing = true)
    ldd = ldd_2d[inds]

    # # lateral part sediment in overland flow
    rivcell = float(river)
    ols = OverlandFlowSediment{Float}(
        n = n,
        rivcell = rivcell,
        soilloss = fill(mv, n),
        erosclay = fill(mv, n),
        erossilt = fill(mv, n),
        erossand = fill(mv, n),
        erossagg = fill(mv, n),
        eroslagg = fill(mv, n),
        TCsed = fill(mv, n),
        TCclay = fill(mv, n),
        TCsilt = fill(mv, n),
        TCsand = fill(mv, n),
        TCsagg = fill(mv, n),
        TClagg = fill(mv, n),
        olsed = fill(mv, n),
        olclay = fill(mv, n),
        olsilt = fill(mv, n),
        olsand = fill(mv, n),
        olsagg = fill(mv, n),
        ollagg = fill(mv, n),
        inlandsed = fill(mv, n),
        inlandclay = fill(mv, n),
        inlandsilt = fill(mv, n),
        inlandsand = fill(mv, n),
        inlandsagg = fill(mv, n),
        inlandlagg = fill(mv, n),
    )

    graph = flowgraph(ldd, inds, pcr_dir)

    # River processes

    # the indices of the river cells in the land(+river) cell vector
    βₗ = ncread(nc, param(config, "input.lateral.land.slope"); sel = inds, type = Float)
    clamp!(βₗ, 0.00001, Inf)
    riverlength = riverlength_2d[inds_riv]
    riverwidth = riverwidth_2d[inds_riv]
    ldd_riv = ldd_2d[inds_riv]
    graph_riv = flowgraph(ldd_riv, inds_riv, pcr_dir)

    index_river = filter(i -> !isequal(river[i], 0), 1:n)
    frac_toriver = fraction_runoff_toriver(graph, ldd, index_river, βₗ, n)

    rs = initialize_riversed(nc, config, riverwidth, riverlength, inds_riv)

    modelmap = (vertical = eros, lateral = (land = ols, river = rs))
    indices_reverse = (
        land = rev_inds,
        river = rev_inds_riv,
        reservoir = isempty(reservoir) ? nothing : reservoir.reverse_indices,
        lake = isempty(lake) ? nothing : lake.reverse_indices,
    )
    writer = prepare_writer(config, reader, modelmap, indices_reverse, x_nc, y_nc, nc)
    close(nc)

    # for each domain save the directed acyclic graph, the traversion order,
    # and the indices that map it back to the two dimensional grid
    land = (
        graph = graph,
        order = topological_sort_by_dfs(graph),
        indices = inds,
        reverse_indices = rev_inds,
    )
    river = (
        graph = graph_riv,
        order = topological_sort_by_dfs(graph_riv),
        indices = inds_riv,
        reverse_indices = rev_inds_riv,
    )

    model = Model(
        config,
        (; land, river, reservoir, lake, index_river, frac_toriver),
        (land = ols, river = rs),
        eros,
        clock,
        reader,
        writer,
    )

    # read and set states in model object if reinit=false
    if reinit == false
        instate_path = joinpath(tomldir, config.state.path_input)
        state_ncnames = ncnames(config.state)
        set_states(instate_path, model, state_ncnames; type = Float)
    end

    return model
end

function update(model::Model{N,L,V,R,W}) where {N,L,V<:LandSediment,R,W}
    @unpack lateral, vertical, network, clock, config = model

    update_until_ols(vertical, config)
    update_until_oltransport(vertical, config)

    lateral.land.soilloss .= vertical.soilloss
    lateral.land.erosclay .= vertical.erosclay
    lateral.land.erossilt .= vertical.erossilt
    lateral.land.erossand .= vertical.erossand
    lateral.land.erossagg .= vertical.erossagg
    lateral.land.eroslagg .= vertical.eroslagg

    lateral.land.TCsed .= vertical.TCsed
    lateral.land.TCclay .= vertical.TCclay
    lateral.land.TCsilt .= vertical.TCsilt
    lateral.land.TCsand .= vertical.TCsand
    lateral.land.TCsagg .= vertical.TCsagg
    lateral.land.TClagg .= vertical.TClagg

    update(lateral.land, network.land, config)

    do_river = get(config.model, "runrivermodel", false)::Bool

    if do_river
        inds_riv = network.index_river
        lateral.river.inlandclay .= lateral.land.inlandclay[inds_riv]
        lateral.river.inlandsilt .= lateral.land.inlandsilt[inds_riv]
        lateral.river.inlandsand .= lateral.land.inlandsand[inds_riv]
        lateral.river.inlandsagg .= lateral.land.inlandsagg[inds_riv]
        lateral.river.inlandlagg .= lateral.land.inlandlagg[inds_riv]

        update(lateral.river, network.river, config)
    end

    write_output(model)

    # update the clock
    advance!(clock)

    return model
end


text:
using SparseArrays, DiffEqOperators, LinearAlgebra, Random,
      Test, BandedMatrices, FillArrays

import DiffEqOperators: BoundaryPaddedVector

function second_derivative_stencil(N)
  A = zeros(N,N+2)
  for i in 1:N, j in 1:N+2
      (j-i==0 || j-i==2) && (A[i,j]=1)
      j-i==1 && (A[i,j]=-2)
  end
  A
end

# Analytic solutions to higher order operators.
# Do not modify unless you are completely certain of the changes.
function fourth_deriv_approx_stencil(N)
    A = zeros(N,N+2)
    A[1,1:8] = [3.5 -56/3 42.5 -54.0 251/6 -20.0 5.5 -2/3]
    A[2,1:8] = [2/3 -11/6 0.0 31/6 -22/3 4.5 -4/3 1/6]

    A[N-1,N-5:end] = reverse([2/3 -11/6 0.0 31/6 -22/3 4.5 -4/3 1/6], dims=2)
    A[N,N-5:end] = reverse([3.5 -56/3 42.5 -54.0 251/6 -20.0 5.5 -2/3], dims=2)

    for i in 3:N-2
        A[i,i-2:i+4] = [-1/6 2.0 -13/2 28/3 -13/2 2.0 -1/6]
    end
    return A
end

function second_deriv_fourth_approx_stencil(N)
    A = zeros(N,N+2)
    A[1,1:6] = [5/6 -1.25 -1/3 7/6 -0.5 5/60]
    A[N,N-3:end] = reverse([5/6 -1.25 -1/3 7/6 -0.5 5/60], dims=2)
    for i in 2:N-1
        A[i,i-1:i+3] = [-1/12 4/3 -5/2 4/3 -1/12]
    end
    return A
end

# The following test-sets test for the correctness of the convolutions against
# AbstractVector and BoundaryPaddedArray in one dimension.

@testset "L.boundary_point_count is zero" begin
    N = 20
    L = CenteredDifference(2, 2, 1.0, N)
    cl = rand()
    cr = rand()
    u = rand(N)
    Qu = BoundaryPaddedVector(cl, cr, u)
    arrayQu = [cl;u;cr]

    # Check that BoundaryPaddedVector is constructed correctly
    @test Qu == arrayQu

    # Test for correctness of DerivativeOperator*AbstractVector
    @test second_derivative_stencil(N)*arrayQu ≈ L*arrayQu

    # Test for correctness of DerivativeOperator*BoundaryPaddedVector
    @test second_derivative_stencil(N)*arrayQu ≈ L*Qu
end

@testset "Fourth Derivative and Approximation Order Test" begin
    N = 20
    L = CenteredDifference(4, 4, 1.0, N)
    cl = rand()
    cr = rand()
    u = rand(N)
    Qu = BoundaryPaddedVector(cl, cr, u)
    arrayQu = [cl;u;cr]

    # Check that BoundaryPaddedVector is constructed correctly
    @test Qu == arrayQu

    # Test for correctness of DerivativeOperator*AbstractVector
    @test fourth_deriv_approx_stencil(N)*arrayQu ≈ L*arrayQu

    # Test for correctness of DerivativeOperator*BoundaryPaddedVector
    @test fourth_deriv_approx_stencil(N)*arrayQu ≈ L*Qu
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

struct VirtualNetworksApi <: SwaggerApi
    client::Swagger.Client
end

"""

Creates or updates a virtual network in the specified resource group.
Param: resourceGroupName::String (required)
Param: virtualNetworkName::String (required)
Param: parameters::VirtualNetwork (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: VirtualNetwork
"""
function virtualNetworksCreateOrUpdate(_api::VirtualNetworksApi, resourceGroupName::String, virtualNetworkName::String, parameters, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "PUT", VirtualNetwork, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/virtualNetworks/{virtualNetworkName}", ["azure_auth"], parameters)
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "virtualNetworkName", virtualNetworkName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Deletes the specified virtual network.
Param: resourceGroupName::String (required)
Param: virtualNetworkName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: Nothing
"""
function virtualNetworksDelete(_api::VirtualNetworksApi, resourceGroupName::String, virtualNetworkName::String, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "DELETE", Nothing, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/virtualNetworks/{virtualNetworkName}", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "virtualNetworkName", virtualNetworkName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets the specified virtual network by resource group.
Param: resourceGroupName::String (required)
Param: virtualNetworkName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Param: expand::String
Return: VirtualNetwork
"""
function virtualNetworksGet(_api::VirtualNetworksApi, resourceGroupName::String, virtualNetworkName::String, api_version::String, subscriptionId::String; expand=nothing, _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", VirtualNetwork, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/virtualNetworks/{virtualNetworkName}", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "virtualNetworkName", virtualNetworkName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_param(_ctx.query, "\$expand", expand)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets all virtual networks in a resource group.
Param: resourceGroupName::String (required)
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: VirtualNetworkListResult
"""
function virtualNetworksList(_api::VirtualNetworksApi, resourceGroupName::String, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", VirtualNetworkListResult, "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Network/virtualNetworks", ["azure_auth"])
    Swagger.set_param(_ctx.path, "resourceGroupName", resourceGroupName)  # type String
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

"""

Gets all virtual networks in a subscription.
Param: api_version::String (required)
Param: subscriptionId::String (required)
Return: VirtualNetworkListResult
"""
function virtualNetworksListAll(_api::VirtualNetworksApi, api_version::String, subscriptionId::String; _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", VirtualNetworkListResult, "/subscriptions/{subscriptionId}/providers/Microsoft.Network/virtualNetworks", ["azure_auth"])
    Swagger.set_param(_ctx.path, "subscriptionId", subscriptionId)  # type String
    Swagger.set_param(_ctx.query, "api-version", api_version)  # type String
    Swagger.set_header_accept(_ctx, ["application/json", "text/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/json", "text/json"] : [_mediaType])
    Swagger.exec(_ctx)
end

export virtualNetworksCreateOrUpdate, virtualNetworksDelete, virtualNetworksGet, virtualNetworksList, virtualNetworksListAll


text:
#' # Modeling and inversion with JUDI
#' ---
#' title: Overview of JUDI modeling and inversion usage
#' author: Mathias Louboutin, Philipp Witte
#' date: April 2022
#' ---

#' This example script is written using [Weave.jl](https://github.com/JunoLab/Weave.jl) and can be converted to different format for documentation and usage
#' This example is converted to a markdown file for the documentation.

#' # Import JUDI, Linear algebra utilities and Plotting
using JUDI, PyPlot, LinearAlgebra

#+ echo = false; results = "hidden"
close("all")

#' # Create a JUDI model structure
#' In JUDI, a `Model` structure contains the grid information (origin, spacing, number of gridpoints)
#' and the physical parameters. The squared slowness is always required as the base physical parameter for propagation. In addition,
#' JUDI supports additional physical representations. First we accept `density` that can either be a direct input `Model(n, d, o, m, rho)` or
#' an optional keyword argument `Model(n,d,o,m;rho=rho)`. Second, we also provide VTI/TTI kernels parametrized by the THomsen parameters that can be input as keyword arguments
#' `Model(n,d,o,m; rho=rho, epsilon=epsilon;delta=delta,theta=theta,phi=phi)`. Because the thomsen parameters are optional the propagator wil lonloy use the ones provided.
#' For example `Model(n,d,o,m; rho=rho, epsilon=epsilon;delta=delta)` will infer a VTI propagation

#' ## Create discrete parameters
# Set up model structure
n = (120, 100)   # (x,y,z) or (x,z)
d = (10., 10.)
o = (0., 0.)

# Velocity [km/s]
v = ones(Float32,n) .+ 0.5f0
v0 = ones(Float32,n) .+ 0.5f0
v[:,Int(round(end/2)):end] .= 3.5f0
rho = (v0 .+ .5f0) ./ 2

# Slowness squared [s^2/km^2]
m = (1f0 ./ v).^2
m0 = (1f0 ./ v0).^2
dm = vec(m0 - m)

# Setup model structure
nsrc = 2	# number of sources
model = Model(n, d, o, m)
model0 = Model(n, d, o, m0)

#' # Create acquisition geometry
#' In this simple usage example, we create a simple acquisiton by hand. In practice the acquisition geometry will be defined by the dataset
#' beeing inverted. We show in a spearate tutorial how to use [SegyIO.jl](https://github.com/slimgroup/SegyIO.jl) to handle SEGY seismic datasets in JUDI.

#' ## Create source and receivers positions at the surface
# Set up receiver geometry
nxrec = 120
xrec = range(0f0, stop=(n[1]-1)*d[1], length=nxrec)
yrec = 0f0 # WE have to set the y coordiante to zero (or any number) for 2D modeling
zrec = range(d[1], stop=d[1], length=nxrec)

# receiver sampling and recording time
timeD = 1250f0   # receiver recording time [ms]
dtD = 2f0    # receiver sampling interval [ms]

# Set up receiver structure
recGeometry = Geometry(xrec, yrec, zrec; dt=dtD, t=timeD, nsrc=nsrc)

#' The source geometry is a but different. Because we want to create a survey with `nsrc` shot records, we need
#' to convert the vector of sources postions `[s0, s1, ... sn]` into an array of array [[s0], [s1], ...] so that
#' JUDI understands that this is a set of indepednet `nsrc`

xsrc = convertToCell(range(0f0, stop=(n[1]-1)*d[1], length=nsrc))
ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
zsrc = convertToCell(range(d[1], stop=d[1], length=nsrc))

# Set up source geometry (cell array with source locations for each shot)
xsrc = convertToCell(range(0f0, stop=(n[1]-1)*d[1], length=nsrc))
ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
zsrc = convertToCell(range(d[1], stop=d[1], length=nsrc))

# Set up source structure
srcGeometry = Geometry(xsrc, ysrc, zsrc; dt=dtD, t=timeD)

#' # Source judiVector
#' Finally, with the geometry defined, we can create a source wavelet (a simple Ricker wavelet here) a our first `judiVector`
#' In JUDI, a `judiVector` is the core structure that represent a acquisition-geometry based dataset. This structure encapsulate
#' the physical locations (trace coordinates) and corrsponding data trace in a source-based structure. for a given `judiVector` `d` then
#' `d[1]` will be the shot record for the first source, or in the case of the source term, the first source wavelet and its positon.

# setup wavelet
f0 = 0.01f0     # kHz
wavelet = ricker_wavelet(timeD, dtD, f0)
q = judiVector(srcGeometry, wavelet)

#' # Modeling
#' With our survey and subsurface model setup, we can now model and image seismic data. We first define a few options. In this tutorial
#' we will choose to compute gradients/images subsampling the forward wavefield every two time steps `subsampling_factor=2` and we fix the computational
#' time step to be `1ms` wiuth `dt_comp=1.0` know to satisfy the CFL condition for this simple example. In practice, when `dt_comp` isn't provided, JUDI will compute the CFL
#' condition for the propagation.

# Setup options
opt = Options(subsampling_factor=2, dt_comp=1.0)

#' Linear Operators
#' The core idea behind JUDI is to abstract seismic inverse problems in term of linear algebra. In its simplest form, seismic inversion can be formulated as
#' ```math
#' \underset{\mathbf{m}}{\text{argmin}} \ \ \phi(\mathbf{m}) = \frac{1}{2} ||\mathbf{P}_r \mathbf{F}(\mathbf{m}) \mathbf{P}_s^{\top} \mathbf{q} - \mathbf{d} ||_2^2 \\
#' \text{   } \\
#' \nabla_{\mathbf{m}} \phi(\mathbf{m}) = \mathbf{J}(\mathbf{m}, \mathbf{q})^{\top} (\mathbf{P}_r \mathbf{F}(\mathbf{m}) \mathbf{P}_s^{\top} \mathbf{q} - \mathbf{d})
#' ```
#'
#' where $\mathbf{P}_r$ is the receiver projection (measurment operator) and $\mathbf{P}_s^{\top}$ is the source injection operator (adjoint of measurment at the source location).
#' Therefore, we bastracted these operation to be able to define these operators

# Setup operators
Pr = judiProjection(recGeometry)
F = judiModeling(model; options=opt)
F0 = judiModeling(model0; options=opt)
Ps = judiProjection(srcGeometry)
J = judiJacobian(Pr*F0*adjoint(Ps), q)

#' # Model and image data

#' We first model synthetic data using our defined source and true model
# Nonlinear modeling
dobs = Pr*F*adjoint(Ps)*q

#' Plot the shot record
fig = figure()
imshow(dobs.data[1], vmin=-1, vmax=1, cmap="PuOr", extent=[xrec[1], xrec[end], timeD/1000, 0], aspect="auto")
xlabel("Receiver position (m)")
ylabel("Time (s)")
title("Synthetic data")
display(fig)

#' Because we have abstracted the linear algebra, we can solve the adjoint wave-equation as well
#' where the data becomes the source. This adjoint solve will be part of the imaging procedure.
# # Adjoint
qad = Ps*adjoint(F)*adjoint(Pr)*dobs

#' We can easily now test the adjointness of our operator with the standard dot test. Because we
#' intend to conserve our linear algebra abstraction, `judiVector` implements all the necessary linear
#' algebra functions such as dot product or norm to be used directly.
# <x, F'y>
dot1 = dot(q, qad)
# <F x, y>
dot2 = dot(dobs, dobs)
# Compare
@show dot1, dot2, (dot2 - dot2)/(dot1 + dot2)

#' # Inversion
#' Our main goal is to provide an unversion framework for seismic inversion. To this end, as shown earlier,
#' users can easily define the Jacobian operator and compute an RTM image (i.e FWI gradient) with a simple matrix-vector product.
#' Once again, we provide both the Jacobian and its adjoint and we can compute Born linearized data.

# Linearized modeling J*dm
dD = J*dm
# Adjoint jacobian, RTM image
rtm = adjoint(J)*dD

#' We show the linearized data.
fig = figure()
imshow(dD.data[1], vmin=-1, vmax=1, cmap="PuOr", extent=[xrec[1], xrec[end], timeD/1000, 0], aspect="auto")
xlabel("Receiver position (m)")
ylabel("Time (s)")
title("Linearized data")
display(fig)

#' And the RTM image
fig = figure()
imshow(rtm', vmin=-1e2, vmax=1e2, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
xlabel("Lateral position(m)")
ylabel("Depth (m)")
title("RTM image")
display(fig)

#' ## Inversion utility functions
#' We currently introduced the lineaar operators that allow to write seismic modeling and inversion in a high-level, linear algebra way. These linear operators allow the script to closely follow the mathematics and to be readable and understandable.
#'
#' However, these come with overhead. In particular, consider the following compuation on the FWI gradient:
#'
#' ```julia
#' d_syn = F*q
#' r = judiJacobian(F, q)' * (d_syn - d_obs)
#' ```
#'
#' In this two lines, the forward modeling is performed twice: once to compute `d_syn` then once again to compute the Jacobian adjoint. In order to avoid this overhead for practical inversion, we provide utility function that directly comput the gradient and objective function (L2- misfit) of FWI, LSRTM and TWRI with minimum overhead.

#' FWI misfit and gradient
# evaluate FWI objective function
f, g = fwi_objective(model0, q, dobs; options=opt)

#' Plot gradient
fig = figure()
imshow(g', vmin=-1e2, vmax=1e2, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
xlabel("Lateral position(m)")
ylabel("Depth (m)")
title("FWI gradient")
display(fig)

#' LSRTM misfit and gradient
# evaluate LSRTM objective function
fj, gj = lsrtm_objective(model0, q, dD, dm; options=opt)
fjn, gjn = lsrtm_objective(model0, q, dobs, dm; nlind=true, options=opt)

#' Plot gradients
fig = figure()
imshow(gj', vmin=-1, vmax=1, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
xlabel("Lateral position(m)")
ylabel("Depth (m)")
title("LSRTM gradient")
display(fig)

fig = figure()
imshow(gjn', vmin=-1, vmax=1, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
xlabel("Lateral position(m)")
ylabel("Depth (m)")
title("LSRTM gradient with background data substracted")
display(fig)

#' By extension, lsrtm_objective is the same as fwi_objecive when `dm` is zero
#' And with computing of the residual. Small noise can be seen in the difference
#' due to floating point roundoff errors with openMP, but running with
#' OMP_NUM_THREADS=1 (no parllelism) produces the exact (difference == 0) same result
#' gjn2 == g
fjn2, gjn2 = lsrtm_objective(model0, q, dobs, 0f0.*dm; nlind=true, options=opt)
fig = figure()

#' Plot gradient
imshow(gjn2', vmin=-1e2, vmax=1e2, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
xlabel("Lateral position(m)")
ylabel("Depth (m)")
title("LSRTM gradient with zero perturbation")
display(fig)

#' # TWRI
#' Finally, JUDI implements TWRI, an augmented method to tackle cycle skipping. Once again we provide a computationnally efficient wrapper function that returns the objective value and necessary gradients
f, gm, gy = twri_objective(model0, q, dobs, nothing; options=opt, optionswri=TWRIOptions(params=:all))
# With on-the-fly DFT, experimental
f, gmf = twri_objective(model0, q, dobs, nothing; options=Options(frequencies=[[.009, .011], [.008, .012]]), optionswri=TWRIOptions(params=:m))

#' Plot gradients
fig = figure()
imshow(gm', vmin=-1, vmax=1, cmap="Greys", extent=[0, (n[1]-1)*d[1], (n[2]-1)*d[2], 0 ], aspect="auto")
xlabel("Lateral position(m)")
ylabel("Depth (m)")
title("TWRI gradient w.r.t m")
display(fig)

fig = figure()
imshow(gy.data[1], vmin=-1e2, vmax=1e2, cmap="PuOr", extent=[xrec[1], xrec[end], timeD/1000, 0], aspect="auto")
xlabel("Receiver position (m)")
ylabel("Time (s)")
title("TWRI gradient w.r.t y")
display(fig)


text:
import .FVReconstructions: FVConstant, FVLinear
using ..BalanceLaws: Primitive
import ..BalanceLaws:
    prognostic_to_primitive!,
    primitive_to_prognostic!,
    construct_face_auxiliary_state!
import .FVReconstructions: width
import StaticArrays: SUnitRange

@doc """
    function vert_fvm_interface_tendency!(
        balance_law::BalanceLaw,
        ::Val{info},
        ::Val{nvertelem},
        ::Val{periodicstack},
        ::VerticalDirection,
        numerical_flux_first_order,
        tendency,
        state_prognostic,
        state_auxiliary,
        vgeo,
        sgeo,
        t,
        vmap⁻,
        vmap⁺,
        elemtobndy,
        elems,
        α,
    )

Compute kernel for evaluating the interface tendencies using vertical FVM
reconstructions with a DG method of the form:

∫ₑ ψ⋅ ∂q/∂t dx - ∫ₑ ∇ψ⋅(Fⁱⁿᵛ + Fᵛⁱˢᶜ) dx + ∮ₑ n̂ ψ⋅(Fⁱⁿᵛ⋆ + Fᵛⁱˢᶜ⋆) dS,

or equivalently in matrix form:

dQ/dt = M⁻¹(MS + DᵀM(Fⁱⁿᵛ + Fᵛⁱˢᶜ) + ∑ᶠ LᵀMf(Fⁱⁿᵛ⋆ + Fᵛⁱˢᶜ⋆)).

This kernel computes the surface terms: M⁻¹ ∑ᶠ LᵀMf(Fⁱⁿᵛ⋆ + Fᵛⁱˢᶜ⋆)), where M
is the mass matrix, Mf is the face mass matrix, L is an interpolator from
volume to face, and Fⁱⁿᵛ⋆, Fᵛⁱˢᶜ⋆ are the numerical fluxes for the inviscid
and viscous fluxes, respectively.

A finite volume reconstruction is used to construction `Fⁱⁿᵛ⋆`
""" vert_fvm_interface_tendency!
@kernel function vert_fvm_interface_tendency!(
    balance_law::BalanceLaw,
    ::Val{info},
    ::Val{nvertelem},
    ::Val{periodicstack},
    ::VerticalDirection,
    reconstruction!,
    numerical_flux_first_order,
    numerical_flux_second_order,
    tendency,
    state_prognostic,
    state_gradient_flux,
    state_auxiliary,
    vgeo,
    sgeo,
    t,
    elemtobndy,
    elems,
    α,
    β,
    increment,
    add_source,
) where {info, nvertelem, periodicstack}
    @uniform begin
        dim = info.dim
        FT = eltype(state_prognostic)
        num_state_prognostic = number_states(balance_law, Prognostic())
        num_state_primitive = number_states(balance_law, Primitive())
        num_state_auxiliary = number_states(balance_law, Auxiliary())
        num_state_gradient_flux = number_states(balance_law, GradientFlux())
        num_state_hyperdiffusive = number_states(balance_law, Hyperdiffusive())
        @assert num_state_hyperdiffusive == 0

        nface = info.nface
        Np = info.Np
        Nqk = info.Nqk # Can only be 1 for the FVM method!
        @assert Nqk == 1

        # We only have the vertical faces
        faces = (nface - 1):nface

        stencil_width = width(reconstruction!)

        # If this fails the `@nif` below needs to change
        @assert stencil_width < 4

        # In the case of stencil_width = 0 we still need two values to evaluate
        # the fluxes, so the minimum stencil diameter is 2
        stencil_diameter = max(2, 2stencil_width + 1)

        # Value in the stencil that corresponds to the top face with respect to
        # face being updated
        stencil_center = max(stencil_width, 1) + 1

        # 1 → element i, face i - 1/2 (bottom face)
        # 2 → element i, face i + 1/2 (top face)
        local_state_face_prognostic = ntuple(Val(2)) do _
            MArray{Tuple{num_state_prognostic}, FT}(undef)
        end

        local_cell_weights = MArray{Tuple{stencil_diameter}, FT}(undef)

        # Two mass matrix inverse corresponding to +/- cells
        vMI = MArray{Tuple{2}, FT}(undef)

        # Storing the value below element when walking up the stack
        # cell i-1, face i - 1/2
        local_state_face_prognostic_neighbor =
            MArray{Tuple{num_state_prognostic}, FT}(undef)

        local_state_face_auxiliary_neighbor =
            MArray{Tuple{num_state_auxiliary}, FT}(undef)

        local_state_face_primitive = ntuple(Val(2)) do _
            MArray{Tuple{num_state_primitive}, FT}(undef)
        end

        local_state_face_auxiliary = ntuple(Val(2)) do _
            MArray{Tuple{num_state_auxiliary}, FT}(undef)
        end

        # Storage for all the values in the stencil
        local_state_prognostic = ntuple(Val(stencil_diameter)) do _
            MArray{Tuple{num_state_prognostic}, FT}(undef)
        end

        # Need to wrap in SVector so we can use views later
        local_state_primitive = SVector(ntuple(Val(stencil_diameter)) do _
            MArray{Tuple{num_state_primitive}, FT}(undef)
        end...)

        local_state_auxiliary = ntuple(Val(stencil_diameter)) do _
            MArray{Tuple{num_state_auxiliary}, FT}(undef)
        end

        # FIXME: These two arrays could be smaller
        # (only 2 elements not stencil_diameter)
        local_state_gradient_flux = ntuple(Val(stencil_diameter)) do _
            MArray{Tuple{num_state_gradient_flux}, FT}(undef)
        end

        local_state_hyperdiffusive = ntuple(Val(stencil_diameter)) do _
            MArray{Tuple{num_state_hyperdiffusive}, FT}(undef)
        end

        # Storage for the numerical flux
        local_flux = MArray{Tuple{num_state_prognostic}, FT}(undef)

        # Storage for the extra boundary points for some BCs
        local_state_prognostic_bottom1 =
            MArray{Tuple{num_state_prognostic}, FT}(undef)
        local_state_gradient_flux_bottom1 =
            MArray{Tuple{num_state_gradient_flux}, FT}(undef)
        local_state_auxiliary_bottom1 =
            MArray{Tuple{num_state_auxiliary}, FT}(undef)

        # Storage for the tendency
        local_tendency = MArray{Tuple{num_state_prognostic}, FT}(undef)
        local_source = MArray{Tuple{num_state_prognostic}, FT}(undef)

        # The remainder model needs to know which direction of face the model is
        # being evaluated for. In this case we only have `VerticalDirection()`
        # faces
        face_direction = (VerticalDirection(),)
    end

    # Optimization ideas:
    #  - More than 1 thread per stack
    #  - shift pointers not data
    #  - Don't keep wide stencil for all variables / data
    #  - Don't load periodic data when domain is not periodic

    # Get the horizontal group IDs
    grp_H = @index(Group, Linear)

    # Determine the index for the element at the bottom of the stack
    eHI = (grp_H - 1) * nvertelem + 1

    # Compute bottom stack element index minus one (so we can add vert element
    # number directly)
    eH = elems[eHI] - 1

    # Which degree of freedom do we handle in the element
    n = @index(Local, Linear)

    # Loads the data for a given element
    function load_data!(
        local_state_prognostic,
        local_state_auxiliary,
        local_state_gradient_flux,
        e,
    )
        @unroll for s in 1:num_state_prognostic
            @inbounds local_state_prognostic[s] = state_prognostic[n, s, e]
        end

        @unroll for s in 1:num_state_auxiliary
            @inbounds local_state_auxiliary[s] = state_auxiliary[n, s, e]
        end

        @unroll for s in 1:num_state_gradient_flux
            @inbounds local_state_gradient_flux[s] =
                state_gradient_flux[n, s, e]
        end
    end

    # To update the first element we either need to apply a BCS on the bottom
    # face in the nonperiodic case, or we need to reconstruct the periodic value
    # for the bottom face
    @inbounds begin
        # If periodic, we are doing the reconstruction in element periodically
        # below the first element, otherwise we are reconstructing the first
        # element
        eV = periodicstack ? nvertelem : 1

        # Figure out the data we need
        els = ntuple(Val(stencil_diameter)) do k
            eH + mod1(eV - 1 + k - (stencil_center - 1), nvertelem)
        end

        # Load all the stencil data
        @unroll for k in 1:stencil_diameter
            load_data!(
                local_state_prognostic[k],
                local_state_auxiliary[k],
                local_state_gradient_flux[k],
                els[k],
            )
            # If local cell weights are NOT _M we need to load _vMI out of sgeo
            local_cell_weights[k] = 2 * vgeo[n, _JcV, els[k]]
        end

        # Transform all the data into primitive variables
        @unroll for k in 1:stencil_diameter
            prognostic_to_primitive!(
                balance_law,
                local_state_primitive[k],
                local_state_prognostic[k],
                local_state_auxiliary[k],
            )
        end
        vMI[2] = sgeo[_vMI, n, faces[2], els[stencil_center]]

        # If we are periodic we reconstruct the top and bottom values for eV
        # then start with eV update in loop below
        if periodicstack
            # Reconstruct the top and bottom values
            rng1, rng2 = stencil_center .+ (-stencil_width, stencil_width)
            rng = SUnitRange(rng1, rng2)
            reconstruction!(
                local_state_face_primitive[1],
                local_state_face_primitive[2],
                local_state_primitive[rng],
                local_cell_weights[rng],
            )

            construct_face_auxiliary_state!(
                balance_law,
                local_state_face_auxiliary[1],
                local_state_auxiliary[stencil_center],
                -local_cell_weights[stencil_center],
            )
            construct_face_auxiliary_state!(
                balance_law,
                local_state_face_auxiliary[2],
                local_state_auxiliary[stencil_center],
                local_cell_weights[stencil_center],
            )

            # Transform the values back to prognostic state
            @unroll for f in 1:2
                primitive_to_prognostic!(
                    balance_law,
                    local_state_face_prognostic[f],
                    local_state_face_primitive[f],
                    # Use the cell auxiliary data
                    local_state_face_auxiliary[f],
                )
            end

            # Initialize local tendency
            @unroll for s in 1:num_state_prognostic
                local_tendency[s] = -zero(FT)
            end
        else

            bctag = elemtobndy[faces[1], eH + eV]

            sM = sgeo[_sM, n, faces[1], eH + eV]
            normal = SVector(
                sgeo[_n1, n, faces[1], eH + eV],
                sgeo[_n2, n, faces[1], eH + eV],
                sgeo[_n3, n, faces[1], eH + eV],
            )

            # Reconstruction using only eVs cell value
            rng = SUnitRange(stencil_center, stencil_center)
            reconstruction!(
                local_state_face_primitive[1],
                local_state_face_primitive[2],
                local_state_primitive[rng],
                local_cell_weights[rng],
            )

            construct_face_auxiliary_state!(
                balance_law,
                local_state_face_auxiliary[1],
                local_state_auxiliary[stencil_center],
                -local_cell_weights[stencil_center],
            )
            construct_face_auxiliary_state!(
                balance_law,
                local_state_face_auxiliary[2],
                local_state_auxiliary[stencil_center],
                local_cell_weights[stencil_center],
            )

            # Transform the values back to prognostic state
            @unroll for k in 1:2
                primitive_to_prognostic!(
                    balance_law,
                    local_state_face_prognostic[k],
                    local_state_face_primitive[k],
                    local_state_face_auxiliary[k],
                )
            end

            # Fill ghost cell data
            fill!(local_flux, -zero(FT))
            local_state_face_prognostic_neighbor .=
                local_state_face_prognostic[1]
            local_state_face_auxiliary_neighbor .= local_state_face_auxiliary[1]

            numerical_boundary_flux_first_order!(
                numerical_flux_first_order,
                bctag,
                balance_law,
                local_flux,
                normal,
                # current cell
                local_state_face_prognostic[1],
                local_state_face_auxiliary[1],
                # ghost cell
                local_state_face_prognostic_neighbor,
                local_state_face_auxiliary_neighbor,
                t,
                face_direction,
                # use current cell states as bottom states
                local_state_face_prognostic[1],
                local_state_face_auxiliary[1],
            )

            # Fill / reset ghost cell data
            local_state_prognostic[stencil_center - 1] .=
                local_state_prognostic[stencil_center]
            local_state_gradient_flux[stencil_center - 1] .=
                local_state_gradient_flux[stencil_center]
            local_state_hyperdiffusive[stencil_center - 1] .=
                local_state_hyperdiffusive[stencil_center]
            local_state_auxiliary[stencil_center - 1] .=
                local_state_auxiliary[stencil_center]

            numerical_boundary_flux_second_order!(
                numerical_flux_second_order,
                bctag,
                balance_law,
                local_flux,
                normal,
                # current cell
                local_state_prognostic[stencil_center],
                local_state_gradient_flux[stencil_center],
                local_state_hyperdiffusive[stencil_center],
                local_state_auxiliary[stencil_center],
                # ghost cell
                local_state_prognostic[stencil_center - 1],
                local_state_gradient_flux[stencil_center - 1],
                local_state_hyperdiffusive[stencil_center - 1],
                local_state_auxiliary[stencil_center - 1],
                t,
                # use current cell states as bottom states
                local_state_prognostic[stencil_center],
                local_state_gradient_flux[stencil_center],
                local_state_auxiliary[stencil_center],
            )

            # Compute boundary flux and add it in bottom element of the mesh
            @unroll for s in 1:num_state_prognostic
                local_tendency[s] = -α * sM * vMI[2] * local_flux[s]
            end
        end

        # The rest of the elements in the stack
        # Compute flux and update for face between elements eV and eV - 1
        #    top face of eV - 1
        #    bottom face of eV
        # For the reconstruction arrays `stencil_center - 1` corresponds to `eV
        # - 1` and `stencil_center` corresponds to `eV`
        #
        # Loop for periodic case has to go beyond the top element so we compute
        # the flux through the top face of vertical element `nvertelem` and
        # bottom face of vertical element 1
        for eV_up in (periodicstack ? (1:nvertelem) : (2:nvertelem))
            # mod1 handles periodicity
            eV_dn = mod1(eV_up - 1, nvertelem)

            # Shift data in storage in order to load new upper element for
            # reconstruction
            # FIXME: shift pointers not data?
            @unroll for k in 1:(stencil_diameter - 1)
                local_state_prognostic[k] .= local_state_prognostic[k + 1]
                local_state_primitive[k] .= local_state_primitive[k + 1]
                local_state_auxiliary[k] .= local_state_auxiliary[k + 1]
                local_state_gradient_flux[k] .= local_state_gradient_flux[k + 1]
                local_cell_weights[k] = local_cell_weights[k + 1]
            end

            # Update volume mass inverse as we move up the stack of elements
            vMI[1] = vMI[2]

            # Load surface metrics for the face we will update (bottom face of
            # `eV_up`)
            sM = sgeo[_sM, n, faces[1], eH + eV_up]
            normal = SVector(
                sgeo[_n1, n, faces[1], eH + eV_up],
                sgeo[_n2, n, faces[1], eH + eV_up],
                sgeo[_n3, n, faces[1], eH + eV_up],
            )

            # Reconstruction for eV_dn was computed in last time through the
            # loop, so we need to store the upper reconstructed values to
            # compute flux for this face
            local_state_face_prognostic_neighbor .=
                local_state_face_prognostic[2]
            local_state_face_auxiliary_neighbor .= local_state_face_auxiliary[2]
            # Next data we need to load (assume periodic, mod1, for now  will
            # mask out below as needed for boundary conditions)
            eV_load = mod1(eV_up + stencil_width, nvertelem)

            # Get element number
            e_load = eH + eV_load

            # Load the next cell into the end of the element arrays
            load_data!(
                local_state_prognostic[stencil_diameter],
                local_state_auxiliary[stencil_diameter],
                local_state_gradient_flux[stencil_diameter],
                e_load,
            )

            # Get local volume mass matrix inverse
            local_cell_weights[stencil_diameter] = 2vgeo[n, _JcV, e_load]
            vMI[2] = sgeo[_vMI, n, faces[2], eH + eV_up]

            # Tranform the prognostic data to primitive data
            prognostic_to_primitive!(
                balance_law,
                local_state_primitive[stencil_diameter],
                local_state_prognostic[stencil_diameter],
                local_state_auxiliary[stencil_diameter],
            )

            # Do the reconstruction! for this cell and compute the values at the
            # bottom (1) and top (2) faces of element `eV_up`
            if periodicstack ||
               stencil_width < eV_up < nvertelem - stencil_width + 1
                # If we are in the interior or periodic just use the
                # reconstruction
                rng1, rng2 = stencil_center .+ (-stencil_width, stencil_width)
                rng = SUnitRange(rng1, rng2)
                reconstruction!(
                    local_state_face_primitive[1],
                    local_state_face_primitive[2],
                    local_state_primitive[rng],
                    local_cell_weights[rng],
                )
            elseif eV_up <= stencil_width
                # Bottom of the element stack requires reconstruct using a
                # subset of the elements
                # Values around stencil center that we need for this
                # reconstruction
                Base.Cartesian.@nif 4 w -> (eV_up == w) w -> begin
                    rng1, rng2 = stencil_center .+ (1 - w, w - 1)
                    rng = SUnitRange(rng1, rng2)
                    reconstruction!(
                        local_state_face_primitive[1],
                        local_state_face_primitive[2],
                        local_state_primitive[rng],
                        local_cell_weights[rng],
                    )
                end w -> throw(BoundsError(local_state_primitive, w))
            elseif eV_up >= nvertelem - stencil_width + 1
                # Top of the element stack requires reconstruct using a
                # subset of the elements
                Base.Cartesian.@nif 4 w -> (w == (nvertelem - eV_up + 1)) w ->
                    begin
                        rng1, rng2 = stencil_center .+ (1 - w, w - 1)
                        rng = SUnitRange(rng1, rng2)
                        reconstruction!(
                            local_state_face_primitive[1],
                            local_state_face_primitive[2],
                            local_state_primitive[rng],
                            local_cell_weights[rng],
                        )
                    end w -> throw(BoundsError(local_state_primitive, w))
            end

            construct_face_auxiliary_state!(
                balance_law,
                local_state_face_auxiliary[1],
                local_state_auxiliary[stencil_center],
                -local_cell_weights[stencil_center],
            )
            construct_face_auxiliary_state!(
                balance_law,
                local_state_face_auxiliary[2],
                local_state_auxiliary[stencil_center],
                local_cell_weights[stencil_center],
            )

            # Transform reconstructed primitive values to prognostic
            @unroll for k in 1:2
                primitive_to_prognostic!(
                    balance_law,
                    local_state_face_prognostic[k],
                    local_state_face_primitive[k],
                    local_state_face_auxiliary[k],
                )
            end

            # Compute the flux for the bottom face of the element we are
            # considering
            fill!(local_flux, -zero(FT))
            numerical_flux_first_order!(
                numerical_flux_first_order,
                balance_law,
                local_flux,
                normal,
                local_state_face_prognostic[1],
                local_state_face_auxiliary[1],
                local_state_face_prognostic_neighbor,
                local_state_face_auxiliary_neighbor,
                t,
                face_direction,
            )

            numerical_flux_second_order!(
                numerical_flux_second_order,
                balance_law,
                local_flux,
                normal,
                local_state_prognostic[stencil_center],
                local_state_gradient_flux[stencil_center],
                local_state_hyperdiffusive[stencil_center],
                local_state_auxiliary[stencil_center],
                local_state_prognostic[stencil_center - 1],
                local_state_gradient_flux[stencil_center - 1],
                local_state_hyperdiffusive[stencil_center - 1],
                local_state_auxiliary[stencil_center - 1],
                t,
            )

            if add_source
                fill!(local_source, -zero(eltype(local_source)))
                source_arr!(
                    balance_law,
                    local_source,
                    local_state_prognostic[stencil_center - 1],
                    local_state_gradient_flux[stencil_center - 1],
                    local_state_auxiliary[stencil_center - 1],
                    t,
                    (VerticalDirection(),),
                )

                @unroll for s in 1:num_state_prognostic
                    local_tendency[s] += local_source[s]
                end
            end

            # Update the bottom element:
            # numerical flux is computed with respect to the top element, so
            # `+=` is used to reverse the flux
            @unroll for s in 1:num_state_prognostic
                local_tendency[s] += α * sM * vMI[1] * local_flux[s]

                if increment
                    tendency[n, s, eH + eV_dn] += local_tendency[s]
                else
                    if β != 0
                        T = local_tendency[s] + β * tendency[n, s, eH + eV_dn]
                    else
                        T = local_tendency[s]
                    end
                    tendency[n, s, eH + eV_dn] = T
                end

                # Store contribution to the top element tendency
                local_tendency[s] = -α * sM * vMI[2] * local_flux[s]
            end

            # Update top element of the stack
            if eV_up == nvertelem
                if add_source
                    fill!(local_source, -zero(eltype(local_source)))
                    source_arr!(
                        balance_law,
                        local_source,
                        local_state_prognostic[stencil_center],
                        local_state_gradient_flux[stencil_center],
                        local_state_auxiliary[stencil_center],
                        t,
                        (VerticalDirection(),),
                    )

                    @unroll for s in 1:num_state_prognostic
                        local_tendency[s] += local_source[s]
                    end
                end
                # If periodic just add in the tendency that we just computed
                if periodicstack
                    @unroll for s in 1:num_state_prognostic
                        if increment
                            tendency[n, s, eH + eV_up] += local_tendency[s]
                        else
                            # This tendency has already been updated in the first element (eV_up=1)
                            tendency[n, s, eH + eV_up] += local_tendency[s]
                        end
                    end
                else
                    # Load surface metrics for the face we will update
                    # (top face of `eV_up`)
                    bctag = elemtobndy[faces[2], eH + eV_up]
                    sM = sgeo[_sM, n, faces[2], eH + eV_up]
                    normal = SVector(
                        sgeo[_n1, n, faces[2], eH + eV_up],
                        sgeo[_n2, n, faces[2], eH + eV_up],
                        sgeo[_n3, n, faces[2], eH + eV_up],
                    )

                    # Since we are at the last element to update, we can safely use
                    # `stencil_center - 1` to store the ghost data

                    fill!(local_flux, -zero(FT))

                    # Fill ghost cell data
                    # Use the top reconstruction (since handling top face)
                    local_state_face_prognostic_neighbor .=
                        local_state_face_prognostic[2]
                    local_state_face_auxiliary_neighbor .=
                        local_state_face_auxiliary[2]
                    numerical_boundary_flux_first_order!(
                        numerical_flux_first_order,
                        bctag,
                        balance_law,
                        local_flux,
                        normal,
                        # Use the top reconstruction (since handling top face)
                        # current state
                        local_state_face_prognostic[2],
                        local_state_face_auxiliary[2],
                        # ghost state
                        local_state_face_prognostic_neighbor,
                        local_state_face_auxiliary_neighbor,
                        t,
                        face_direction,
                        # use current cell states as bottom states
                        local_state_face_prognostic[2],
                        local_state_face_auxiliary[2],
                    )

                    # Fill / reset ghost cell data
                    local_state_prognostic[stencil_center - 1] .=
                        local_state_prognostic[stencil_center]
                    local_state_gradient_flux[stencil_center - 1] .=
                        local_state_gradient_flux[stencil_center]
                    local_state_hyperdiffusive[stencil_center - 1] .=
                        local_state_hyperdiffusive[stencil_center]
                    local_state_auxiliary[stencil_center - 1] .=
                        local_state_auxiliary[stencil_center]
                    numerical_boundary_flux_second_order!(
                        numerical_flux_second_order,
                        bctag,
                        balance_law,
                        local_flux,
                        normal,
                        # current state
                        local_state_prognostic[stencil_center],
                        local_state_gradient_flux[stencil_center],
                        local_state_hyperdiffusive[stencil_center],
                        local_state_auxiliary[stencil_center],
                        # ghost state
                        local_state_prognostic[stencil_center - 1],
                        local_state_gradient_flux[stencil_center - 1],
                        local_state_hyperdiffusive[stencil_center - 1],
                        local_state_auxiliary[stencil_center - 1],
                        t,
                        # use current cell states as bottom states
                        local_state_prognostic[stencil_center],
                        local_state_gradient_flux[stencil_center],
                        local_state_auxiliary[stencil_center],
                    )

                    @unroll for s in 1:num_state_prognostic
                        local_tendency[s] -= α * sM * vMI[2] * local_flux[s]
                        if increment
                            tendency[n, s, eH + eV_up] += local_tendency[s]
                        else
                            if β != 0
                                T =
                                    local_tendency[s] +
                                    β * tendency[n, s, eH + eV_up]
                            else
                                T = local_tendency[s]
                            end
                            tendency[n, s, eH + eV_up] = T
                        end
                    end
                end
            end
        end
    end
end

@kernel function vert_fvm_interface_gradients!(
    balance_law::BalanceLaw,
    ::Val{info},
    ::Val{nvertelem},
    ::Val{periodicstack},
    ::VerticalDirection,
    state_prognostic,
    state_gradient_flux,
    state_auxiliary,
    vgeo,
    sgeo,
    t,
    elemtobndy,
    elems,
    increment,
) where {info, nvertelem, periodicstack}
    @uniform begin

        dim = info.dim
        FT = eltype(state_prognostic)
        num_state_prognostic = number_states(balance_law, Prognostic())
        ngradstate = number_states(balance_law, Gradient())
        num_state_gradient_flux = number_states(balance_law, GradientFlux())
        num_state_auxiliary = number_states(balance_law, Auxiliary())
        nface = info.nface
        Np = info.Np
        faces = (nface - 1, nface)

        # Storage for the prognostic state for e-1, e, e+1
        local_state_prognostic =
            ntuple(_ -> MArray{Tuple{num_state_prognostic}, FT}(undef), Val(3))

        # Storage for the auxiliary state for e-1, e, e+1
        local_state_auxiliary =
            ntuple(_ -> MArray{Tuple{num_state_auxiliary}, FT}(undef), Val(3))

        # Storage for the transform state for e-1, e, e+1
        # e.g., the state we take the gradient of
        l_grad_arg = ntuple(_ -> MArray{Tuple{ngradstate}, FT}(undef), Val(3))

        # Storage for the contribution to the gradient from these faces
        l_nG = MArray{Tuple{3, ngradstate}, FT}(undef)

        l_nG_bc = MArray{Tuple{3, ngradstate}, FT}(undef)

        # Storage for the state gradient flux locally
        local_state_gradient_flux =
            MArray{Tuple{num_state_gradient_flux}, FT}(undef)

        local_state_prognostic_bottom1 =
            MArray{Tuple{num_state_prognostic}, FT}(undef)
        local_state_auxiliary_bottom1 =
            MArray{Tuple{num_state_auxiliary}, FT}(undef)

        # XXX: will revisit this later for FVM
        fill!(local_state_prognostic_bottom1, NaN)
        fill!(local_state_auxiliary_bottom1, NaN)
    end

    # Element index
    eI = @index(Group, Linear)
    # Index of a quadrature point on a face
    n = @index(Local, Linear)

    @inbounds begin
        e = elems[eI]
        eV = mod1(e, nvertelem)

        # Figure out the element above and below e
        e_dn, bc_dn = if eV > 1
            e - 1, 0
        elseif periodicstack
            e + nvertelem - 1, 0
        else
            e, elemtobndy[faces[1], e]
        end

        e_up, bc_up = if eV < nvertelem
            e + 1, 0
        elseif periodicstack
            e - nvertelem + 1, 0
        else
            e, elemtobndy[faces[2], e]
        end

        bctag = (bc_dn, bc_up)

        els = (e_dn, e, e_up)

        # Load the normal vectors and surface mass matrix on the faces
        normal_vector = ntuple(
            k -> SVector(
                sgeo[_n1, n, faces[k], e],
                sgeo[_n2, n, faces[k], e],
                sgeo[_n3, n, faces[k], e],
            ),
            Val(2),
        )
        sM = ntuple(k -> sgeo[_sM, n, faces[k], e], Val(2))

        # Volume mass same on both faces
        vMI = sgeo[_vMI, n, faces[1], e]

        # Get the mass matrix for each of the elements
        M = ntuple(k -> vgeo[n, _M, els[k]], Val(3))

        # Load prognostic data
        @unroll for k in 1:3
            @unroll for s in 1:num_state_prognostic
                local_state_prognostic[k][s] = state_prognostic[n, s, els[k]]
            end
        end

        # Load auxiliary data
        @unroll for k in 1:3
            @unroll for s in 1:num_state_auxiliary
                local_state_auxiliary[k][s] = state_auxiliary[n, s, els[k]]
            end
        end

        # Transform to the gradient argument (i.e., the values we take the
        # gradient of)
        @unroll for k in 1:3
            fill!(l_grad_arg[k], -zero(eltype(l_grad_arg[k])))
            compute_gradient_argument_arr!(
                balance_law,
                l_grad_arg[k],
                local_state_prognostic[k],
                local_state_auxiliary[k],
                t,
            )
        end

        # Compute the surface integral contribution from these two faces:
        #   M⁻¹ Lᵀ Mf n̂ G*
        # Since this is FVM we do not subtract the interior state
        fill!(l_nG, -zero(eltype(l_nG)))
        @unroll for f in 1:2
            if bctag[f] == 0
                @unroll for s in 1:ngradstate
                    # Interpolate to the face (using the mass matrix for the
                    # interpolation weights) -- "the gradient numerical flux"
                    G =
                        (
                            M[f] * l_grad_arg[f + 1][s] +
                            M[f + 1] * l_grad_arg[f][s]
                        ) / (M[f] + M[f + 1])

                    # Compute the surface integral for this component and face
                    # multiplied by the normal to get the rotation to the
                    # physical space
                    @unroll for i in 1:3
                        l_nG[i, s] += vMI * sM[f] * normal_vector[f][i] * G
                    end
                end
            else
                # Computes G* incorporating boundary conditions
                numerical_boundary_flux_gradient!(
                    CentralNumericalFluxGradient(),
                    bctag[f],
                    balance_law,
                    l_nG_bc,
                    normal_vector[f],
                    l_grad_arg[2],
                    local_state_prognostic[2],
                    local_state_auxiliary[2],
                    l_grad_arg[2f - 1],
                    local_state_prognostic[2f - 1],
                    local_state_auxiliary[2f - 1],
                    t,
                    local_state_prognostic_bottom1,
                    local_state_auxiliary_bottom1,
                )

                @unroll for s in 1:ngradstate
                    @unroll for i in 1:3
                        l_nG[i, s] += vMI * sM[f] * l_nG_bc[i, s]
                    end
                end
            end
        end

        # Applies linear transformation of gradients to the diffusive variables
        # for storage
        compute_gradient_flux_arr!(
            balance_law,
            local_state_gradient_flux,
            l_nG,
            local_state_prognostic[2],
            local_state_auxiliary[2],
            t,
        )

        # This is the surface integral evaluated discretely
        # M^(-1) Mf G*
        @unroll for s in 1:num_state_gradient_flux
            if increment
                state_gradient_flux[n, s, e] += local_state_gradient_flux[s]
            else
                state_gradient_flux[n, s, e] = local_state_gradient_flux[s]
            end
        end
    end
end

@kernel function vert_fvm_auxiliary_field_gradient!(
    balance_law::BalanceLaw,
    ::Val{info},
    ∇state,
    state,
    vgeo,
    sgeo,
    vmap⁻,
    vmap⁺,
    elemtobndy,
    ::Val{I},
    ::Val{O},
    increment,
) where {info, I, O}
    @uniform begin
        FT = eltype(state)
        ngradstate = length(I)

        dim = info.dim
        nface = info.nface
        Np = info.Np

        # We only have the vertical faces
        faces = (nface - 1):nface

        local_state_bottom = fill!(MArray{Tuple{ngradstate}, FT}(undef), NaN)
        local_state = fill!(MArray{Tuple{ngradstate}, FT}(undef), NaN)
        local_state_top = fill!(MArray{Tuple{ngradstate}, FT}(undef), NaN)
    end

    e = @index(Group, Linear)
    n = @index(Local, Linear)

    @inbounds begin
        face_bottom = faces[1]
        face_top = faces[2]

        bctag_bottom = elemtobndy[face_bottom, e]
        bctag_top = elemtobndy[face_top, e]

        # TODO: exploit structured grid
        id = vmap⁻[n, face_bottom, e]
        id_bottom = vmap⁺[n, face_bottom, e]
        id_top = vmap⁺[n, face_top, e]

        e_bottom = ((id_bottom - 1) ÷ Np) + 1
        e_top = ((id_top - 1) ÷ Np) + 1

        vid = ((id - 1) % Np) + 1
        vid_bottom = ((id_bottom - 1) % Np) + 1
        vid_top = ((id_top - 1) % Np) + 1

        if bctag_bottom != 0
            e_bottom = e
            vid_bottom = vid
        end
        if bctag_top != 0
            e_top = e
            vid_top = vid
        end

        dzh_bottom = vgeo[vid_bottom, _JcV, e_bottom]
        dzh = vgeo[vid, _JcV, e]
        dzh_top = vgeo[vid_top, _JcV, e_top]

        if dim == 2
            ξvx1 = vgeo[vid, _ξ2x1, e]
            ξvx2 = vgeo[vid, _ξ2x2, e]
            ξvx3 = vgeo[vid, _ξ2x3, e]
        elseif dim == 3
            ξvx1 = vgeo[vid, _ξ3x1, e]
            ξvx2 = vgeo[vid, _ξ3x2, e]
            ξvx3 = vgeo[vid, _ξ3x3, e]
        end

        @unroll for s in 1:ngradstate
            local_state_bottom[s] = state[vid_bottom, I[s], e_bottom]
            local_state_top[s] = state[vid_top, I[s], e_top]
        end

        # only need the middle state near the boundaries
        if bctag_bottom != 0 || bctag_top != 0
            @unroll for s in 1:ngradstate
                local_state[s] = state[vid, I[s], e]
            end
        end

        # extrapolation at the boundaries equivalent to one-sided differencing
        if bctag_bottom != 0
            @unroll for s in 1:ngradstate
                local_state_bottom[s] = 2 * local_state[s] - local_state_top[s]
            end
            dzh_bottom = dzh_top
        end

        if bctag_top != 0
            @unroll for s in 1:ngradstate
                local_state_top[s] = 2 * local_state[s] - local_state_bottom[s]
            end
            dzh_top = dzh_bottom
        end

        @unroll for s in 1:ngradstate
            dz = dzh_top + 2dzh + dzh_bottom
            ∇s_v = (local_state_top[s] - local_state_bottom[s]) / dz
            # rotate back to Cartesian
            if increment
                ∇state[vid, O[3 * (s - 1) + 1], e] += ξvx1 * dzh * ∇s_v
                ∇state[vid, O[3 * (s - 1) + 2], e] += ξvx2 * dzh * ∇s_v
                ∇state[vid, O[3 * (s - 1) + 3], e] += ξvx3 * dzh * ∇s_v
            else
                ∇state[vid, O[3 * (s - 1) + 1], e] = ξvx1 * dzh * ∇s_v
                ∇state[vid, O[3 * (s - 1) + 2], e] = ξvx2 * dzh * ∇s_v
                ∇state[vid, O[3 * (s - 1) + 3], e] = ξvx3 * dzh * ∇s_v
            end
        end
    end
end

"""
    kernel_fvm_balance!(f!, balance_law::BalanceLaw, ::Val{nvertelem}, state_auxiliary, vgeo, elems)

    pᵢ(ρᵢ, Tᵢ) =  pᵢ₋₁ - g ρᵢ Δzᵢ/2 - g ρᵢ₋₁ Δzᵢ₋₁/2
    for i = 2:nvertelem
        pᵢ(ρᵢ, Tᵢ) =  pᵢ₋₁ - g ρᵢ Δzᵢ/2 - g ρᵢ₋₁ Δzᵢ₋₁/2
    end

 - `f!`: update function
 - `balance_law`: atmosphere model
 - `state_auxiliary`: auxiliary variables, ρ is updated
 - `vgeo`: 2*vgeo[ijk ,_JcV , e] is Δz
 - `elems`: horizontal element list
"""
@kernel function kernel_fvm_balance!(
    f!,
    balance_law::BalanceLaw,
    ::Val{nvertelem},
    state_auxiliary,
    vgeo,
    elems,
) where {nvertelem}
    @uniform begin
        FT = eltype(state_auxiliary)
        num_state_auxiliary = number_states(balance_law, Auxiliary())
        # ρᵢ pᵢ
        local_state_auxiliary_top =
            MArray{Tuple{num_state_auxiliary}, FT}(undef)
        # ρᵢ₋₁ pᵢ₋₁
        local_state_auxiliary = MArray{Tuple{num_state_auxiliary}, FT}(undef)

        Δz = MArray{Tuple{2}, FT}(undef)
    end

    _eH = @index(Group, Linear)
    n = @index(Local, Linear)

    @inbounds begin
        eH = elems[_eH]

        # handle bottom element
        eV = 1
        e = eV + (eH - 1) * nvertelem
        @unroll for s in 1:num_state_auxiliary
            local_state_auxiliary[s] = state_auxiliary[n, s, e]
        end
        # Δzᵢ₋₁
        Δz[1] = 2 * vgeo[n, _JcV, e]

        # Loop up the stack of elements
        for eV in 2:nvertelem
            e = eV + (eH - 1) * nvertelem

            @unroll for s in 1:num_state_auxiliary
                local_state_auxiliary_top[s] = state_auxiliary[n, s, e]
            end
            # Δzᵢ
            Δz[2] = 2 * vgeo[n, _JcV, e]

            f!(
                balance_law,
                # ρᵢ pᵢ
                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                    local_state_auxiliary_top,
                ),
                # ρᵢ₋₁ pᵢ₋₁
                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                    local_state_auxiliary,
                ),
                Δz,
            )

            # update to the global array
            @unroll for s in 1:num_state_auxiliary
                state_auxiliary[n, s, e] = local_state_auxiliary_top[s]
            end

            # (ρᵢ pᵢ Δzᵢ) -> (ρᵢ₋₁ pᵢ₋₁ Δzᵢ₋₁)
            @unroll for s in 1:num_state_auxiliary
                local_state_auxiliary[s] = local_state_auxiliary_top[s]
            end
            Δz[1] = Δz[2]
        end
    end
end


text:
# Michelle Medina
# July 7, 2020
# Tests Hierarchical Basis Functions

# Testing Level Constructors
level0 = Level([1.0,3.0])
level1 = Level([2.0])
level2 = Level([3.0,4.0])

# Testing Hierarchical Basis Contructor
bases = HierarchicalBasis([level0,level1,level2])

# Testing Nodal Basis Constructor
xVal = collect(range(0.0, stop=1.0, length = 9))
y = exp.(xVal)
nodal = calcNodal(x -> exp(x), xVal)
@test nodal.values[3] == 1.2840254166877414

# Tests getting maximum level
maxLevelBase = maxlevel(bases)
@test maxLevelBase == 2

maxLevelNodal = maxlevel(nodal)
@test maxLevelNodal == 3
#
# Tests getting values at index of a Hierarchical and Nodal Basis
indexBase = getindex(bases, 1, 0)
# can also name them this way
bases[1, 0]

indexNodal = getindex(nodal, 2)
@test nodal[2] == 1.2840254166877414

# Tests setting index of a basis
setindex!(bases, Float64(pi), 1, 0)
println(bases)

nodal[0] = cos(pi/6)
@show nodal

# Tests going from a nodal to a hierarchical basis
newValues = Nodal_2_H(nodal)

# Tests getting zeros
zerosFromH = zeros(bases)
zerosFromN = zeros(nodal)

@test maxlevel(nodal) == maxlevel(newValues)

# Tests evaluating at a particular x value
valueTest = evaluate(bases, 0.5)

# Tests how coefficients fall off with level
l = 12
xYLOPHONE = collect(range(0.0, stop=1.0, length = 2^l +1))
nodalTest = calcNodal(x -> exp(x), xYLOPHONE)
Values = Nodal_2_H(nodalTest)
MaxCoeff = getMaximum(Values)

using PyPlot
semilogy(MaxCoeff, "g o")
xlabel("Levels")
ylabel("Maximum Coefficient")
title("Max Coefficient vs Level")
show()

# Testing analytical and numeric calculations (up to a certain level)
9 points --> exp(x) [Nodal Basis] --> Hierarchical basis (level 3)
100 points between [0, 1]
y_analytic = exp(x) where x has 100 points
y_numeric_upto_level_l = evaluate(HierarchichalBasis(up to level l)) on these 100 points
error_upto_level_l = y_numeric_upto_level_l - y_analytic
plot(upto_level_l, error_upto_level_l)

xValue = collect(range(0.0, stop=1.0, length = 5))
nodalTest = calcNodal(x -> exp(x), xValue)
newVal = Nodal_2_H(nodalTest)
newVal2 = Nodal_2_H_new(nodalTest)
@show newVal
@show newVal2

# Testing error of going up to a certain level
using LinearAlgebra
xNew = collect(range(0.0, stop=1.0, length = 2000))
y2 = exp.(xNew)
E = []
for l in 0:10
	y1 = evaluate_upto(newVal, l, xNew)
	error_upto_level_l = norm(y1 - y2)
	append!(E, error_upto_level_l)
end

using PyPlot
semilogy(E)
xlabel("Levels")
ylabel("Interpolation Error")
title("Interpolation Error vs Levels")
show()


text:
using Trixi

trixi_include(default_example())
trixi_include(joinpath(examples_dir(), "tree_2d_dgsem", "elixir_euler_ec.jl"))
trixi_include(joinpath(examples_dir(), "tree_2d_dgsem", "elixir_eulermulti_ec.jl"), tspan=(0.0, 0.1))
trixi_include(joinpath(examples_dir(), "tree_2d_dgsem", "elixir_advection_amr_visualization.jl"))
trixi_include(joinpath(examples_dir(), "tree_2d_dgsem", "elixir_euler_blast_wave_amr.jl"), tspan=(0.0, 1.0))
trixi_include(joinpath(examples_dir(), "tree_1d_dgsem", "elixir_euler_positivity.jl"), tspan=(0.0, 1.0))
trixi_include(joinpath(examples_dir(), "tree_3d_dgsem", "elixir_euler_shockcapturing_amr.jl"), tspan=(0.0, 0.1))


text:
module Tracking
    using
        DocStringExtensions,
        GNSSSignals,
        StaticArrays,
        TrackingLoopFilters,
        StructArrays,
        LoopVectorization
    using Unitful: upreferred, Hz, dBHz, ms
    import Base.zero, Base.length, Base.resize!

    export
        get_early,
        get_prompt,
        get_late,
        get_correlator,
        get_accumulators,
        get_num_accumulators,
        get_early_index,
        get_prompt_index,
        get_late_index,
        get_accumulator,
        get_carrier_doppler,
        get_carrier_phase,
        get_code_doppler,
        get_code_phase,
        get_correlator_sample_shifts,
        get_early_late_sample_spacing,
        get_early_late_index_shift,
        get_secondary_code_or_bit_found,
        get_correlator_carrier_phase,
        get_correlator_carrier_frequency,
        get_state,
        get_system,
        get_cn0,
        track,
        TrackingState,
        NumAnts,
        NumAccumulators,
        MomentsCN0Estimator,
        AbstractCN0Estimator,
        get_bits,
        get_num_bits,
        EarlyPromptLateCorrelator,
        #VeryEarlyPromptLateCorrelator,
        SecondaryCodeOrBitDetector,
        GainControlledSignal

    struct NumAnts{x}
    end

    NumAnts(x) = NumAnts{x}()

    struct NumAccumulators{x}
    end

    NumAccumulators(x) = NumAccumulators{x}()

    include("code_replica.jl")
    include("carrier_replica.jl")
    include("downconvert.jl")
    include("cn0_estimation.jl")
    include("discriminators.jl")
    include("bit_buffer.jl")
    include("correlator.jl")
    include("secondary_code_or_bit_detector.jl")
    include("tracking_state.jl")
    include("tracking_results.jl")
    include("tracking_loop.jl")
    include("gpsl1.jl")
    include("gpsl5.jl")
    include("galileo_e1b.jl")
    include("boc.jl")
end


text:
module Flatbuf

using ..FlatBuffers

include("Schema.jl")
include("File.jl")
include("Message.jl")

end # module


text:
# Handles generating random seeds and associated helper functions.
import Random: seed!

hash_uint32(seed::Float64) = hash_uint32(hash(seed))
hash_uint32(seed::Union{UInt32, UInt64}) = UInt32(hash(seed) & 0x00000000FFFFFFFF)

function set_seed!(mdp::ASTMDP)
    # Use initial seed from mdp.params.seed
    seed = isnothing(mdp.current_seed) ? mdp.params.seed : mdp.current_seed
    mdp.current_seed = hash_uint32(seed)
end

set_global_seed(seed::UInt32) = Random.seed!(seed)
set_global_seed(a::ASTSeedAction) = set_global_seed(a.seed)


text:
#####
##### LSMR implementation copied from IterativeSolvers.jl
##### modified to make it non-allocating and to add multi-threading
##### https://github.com/JuliaLinearAlgebra/IterativeSolvers.jl/blob/v0.9.2/src/lsmr.jl
#####

struct LSMRWorkspace{TA, Tx, Tb, Tw}
    A::TA
    x::Tx
    b::Tb
    v::Tw
    h::Tw
    hbar::Tw
    Av::Tb
    Atu::Tw
end

function LSMRWorkspace(A)
    T = eltype(A)

    x = tfill!(Vector{T}(undef, size(A, 2)), zero(T))
    b = tfill!(Vector{T}(undef, size(A, 1)), zero(T))

    v = similar(x, T)
    h = similar(x, T)
    h̄ = similar(x, T)
    Atu = similar(x, T)
    Av = similar(b)

    return LSMRWorkspace(A, x, b, v, h, h̄, Av, Atu)
end

function LSMRWorkspace(A, b::AbstractArray)
    T = typeof(one(eltype(b)) / one(eltype(A)))
    x = similar(b, T, size(A, 2))
    return LSMRWorkspace(x, A, b)
end

function LSMRWorkspace(x::AbstractArray, A, b::AbstractArray)
    T = typeof(one(eltype(b)) / one(eltype(A)))

    u = isa(b, Array) ? tcopy(b) : copy(b)

    v = similar(x, T)
    h = similar(x, T)
    h̄ = similar(x, T)
    Atu = similar(x, T)
    Av = similar(b)

    return LSMRWorkspace(A, x, u, v, h, h̄, Av, Atu)
end

#== Not exported. Don't generate docs
"""
    lsmr(A, b; kwargs...) -> x

Same as [`lsmr!`](@ref), but allocates a solution vector `x` initialized with
zeros.
"""
==#
lsmr(A, b; kwargs...) = lsmr!(LSMRWorkspace(A, b); kwargs...)

#== Not exported. Don't generate docs
"""
    lsmr!(x, A, b; kwargs...) -> x

Minimizes ``\\|Ax - b\\|^2 + \\|λx\\|^2`` in the Euclidean norm. If multiple
solutions exists the minimum norm solution is returned.

The method is based on the Golub-Kahan bidiagonalization process. It is
algebraically equivalent to applying MINRES to the normal equations
``(A^*A + λ^2I)x = A^*b``, but has better numerical properties, especially if
``A`` is ill-conditioned.

### Arguments
- `x`: solution array
- `A`: linear operator
- `b`: right-hand side

### Keywords
- `lambda::Real = 0`: regularization parameter
- `atol::Real = 1e-6`: stopping tolerance
- `btol::Real = 1e-6`: stopping tolerance
- `conlim::Real = 1e8`: stopping tolerance
- `maxit::Int = 1 + min(size(A, 1), size(A, 2))`: maximum number of iterations
- `verbose::Bool = false`: print convergence information

### Returns
- `x`: approximated solution

### Notes
- `atol`, `btol`: LSMR continues iterations until a certain backward error
    estimate is smaller than some quantity depending on ATOL and BTOL.
    Let RES = B - A*X be the residual vector for the current approximate
    solution X.  If A*X = B seems to be consistent, LSMR terminates when
    NORM(RES) <= ATOL*NORM(A)*NORM(X) + BTOL*NORM(B).  Otherwise, LSMR
    terminates when NORM(A'*RES) <= ATOL*NORM(A)*NORM(RES).  If both tolerances
    are 1.0e-6 (say), the final NORM(RES) should be accurate to about 6 digits.
    (The final X will usually have fewer correct digits, depending on cond(A)
    and the size of LAMBDA.) If ATOL or BTOL is [], a default value of 1.0e-6
    will be used.  Ideally, they should be estimates of the relative error in
    the entries of A and B respectively.  For example, if the entries of A
    have 7 correct digits, set ATOL = 1e-7. This prevents the algorithm
    from doing unnecessary work beyond the uncertainty of the input data.
- `conlim`: LSMR terminates if an estimate of cond(A) exceeds CONLIM. For
    compatible systems Ax = b, conlim could be as large as 1.0e+12 (say).
    For least-squares problems, conlim should be less than 1.0e+8. If CONLIM
    is [], the default value is CONLIM = 1e+8. Maximum precision can be
    obtained by setting ATOL = BTOL = CONLIM = 0, but the number of iterations
    may then be excessive.

### References
    D. C.-L. Fong and M. A. Saunders,
    LSMR: An iterative algorithm for sparse least-squares problems,
    SIAM J. Sci. Comput., submitted 1 June 2010.
    See http://www.stanford.edu/~clfong/lsmr.html.
"""
==#
lsmr!(x, A, b; kwargs...) = lsmr!(LSMRWorkspace(x, A, b); kwargs...)

function lsmr!(
    L::LSMRWorkspace;
    atol::Real = 1e-6,
    btol::Real = 1e-6,
    conlim::Real = 1e8,
    maxit::Int = 1 + min(size(L.A, 1), size(L.A, 2)),
    lambda::Real = 0,
    verbose::Bool = false
)
    A, x, b = L.A, L.x, L.b
    v, h, hbar, Av, Atu = L.v, L.h, L.hbar, L.Av, L.Atu

    T = typeof(one(eltype(b)) / one(eltype(A)))
    Tr = real(T)

    At = adjoint(A)

    # form the first vectors u and v (satisfy  β*u = b,  α*v = A'u)
    u = b
    β = norm(u)

    if !iszero(β)
        u = _scal!(inv(β), u)
    end

    mul!(v, At, u)
    α = norm(v)

    if !iszero(α)
        v = _scal!(inv(α), v)
    end

    # Initialize variables for 1st iteration.
    λ = convert(Tr, lambda)

    ζbar = α * β
    αbar = α
    ρ    = one(Tr)
    ρbar = one(Tr)
    cbar = one(Tr)
    sbar = zero(Tr)

    if isa(x, Array)
        tfill!(x, 0)
    else
        fill!(x, 0)
    end

    if isa(h, Array) && isa(v, Array)
        _tcopyto!(h, v)
    else
        copyto!(h, v)
    end

    if isa(hbar, Array)
        tfill!(hbar, 0)
    else
        fill!(hbar, 0)
    end

    # Initialize variables for estimation of ||r||.
    βdd = β
    βd = zero(Tr)
    ρdold = one(Tr)
    τtildeold = zero(Tr)
    θtilde  = zero(Tr)
    ζ = zero(Tr)
    d = zero(Tr)

    # Initialize variables for estimation of ||A|| and cond(A).
    normA2  = abs2(α)
    maxrbar = zero(Tr)
    minrbar = typemax(Tr)

    # Items for use in stopping rules.
    i = 0
    istop = 0
    normb = β
    normr = β

    ϵa = convert(Tr, atol)
    ϵb = convert(Tr, btol)
    ϵc = conlim > 0 ? convert(Tr, inv(conlim)) : zero(Tr)

    # Exit if b = 0 or A'b = 0.
    normAr = α * β
    if iszero(normAr)
        if verbose
            @printf("The exact solution is x = 0")
        end
        return x
    end

    if verbose
        @printf("\nLSMR            Least-squares solution of  Ax = b\n")
        @printf("The matrix A has %d rows and %d cols\n", size(A, 1), size(A, 2))
        @printf("lambda = %16.10e\n", lambda)
        @printf("atol   = %8.2e\n", atol)
        @printf("btol   = %8.2e\n", btol)
        @printf("conlim = %8.2e\n", conlim)
        @printf("\n%7s %11s\t%8s %11s %5s\t%6s\t%8s\n",
            "i", "||r||", "||A'r||", "compatible", "LS", "||A||", "cond A"
        )
    end

    while i < maxit
        i += 1

        # u = A*v - α*u
        # β = norm(u)
        mul!(Av, A, v)
        u, β = _xpby_norm!(Av, -α, u)

        if !iszero(β)
            u = _scal!(inv(β), u)

            # v = A'*u - β*v
            # α = norm(v)
            mul!(Atu, At, u)
            v, α = _xpby_norm!(Atu, -β, v)

            if !iszero(α)
                v = _scal!(inv(α), v)
            end
        end

        # Construct rotation Qhat_{k,2k+1}.
        αhat = hypot(αbar, λ)
        chat = αbar / αhat
        shat = λ / αhat

        # Use a plane rotation (Q_i) to turn B_i to R_i.
        ρold = ρ
        ρ    = hypot(αhat, β)
        c    = αhat / ρ
        s    = β / ρ
        θnew = s * α
        αbar = c * α

        # Use a plane rotation (Qbar_i) to turn R_i^T to R_i^bar.
        ρbarold = ρbar
        ζold = ζ
        θbar = sbar * ρ
        ρtmp = cbar * ρ
        ρbar = hypot(ρtmp, θnew)
        cbar = ρtmp / ρbar
        sbar = θnew / ρbar
        ζ    = cbar * ζbar
        ζbar = -sbar * ζbar

        # Update h, h_hat, x.
        hbar, x, h, normx =
            _update_hbar_x_h!(hbar, x, h, v, θbar, ρ, ρold, ρbarold, ζ, ρbar, θnew)

        # Estimate of ||r||

        # Apply rotation Qhat_{k,2k+1}.
        βacute =  chat * βdd
        βcheck = -shat * βdd

        # Apply rotation Q_{k,k+1}.
        βhat =  c * βacute
        βdd  = -s * βacute

        # Apply rotation Qtilde_{k-1}.
        θtildeold = θtilde
        ρtildeold = hypot(ρdold, θbar)
        ctildeold = ρdold / ρtildeold
        stildeold = θbar / ρtildeold
        θtilde    = stildeold * ρbar
        ρdold     = ctildeold * ρbar
        βd        = -stildeold*βd + ctildeold*βhat

        τtildeold = (ζold - θtildeold * τtildeold) / ρtildeold
        τd = (ζ - θtilde * τtildeold) / ρdold

        d += abs2(βcheck)
        normr = sqrt(d + abs2(βd - τd) + abs2(βdd))

        # Estimate ||A||.
        normA2 += abs2(β)
        normA   = sqrt(normA2)
        normA2 += abs2(α)

        # Estimate cond(A).
        maxrbar = max(maxrbar, ρbarold)
        if i > 1
            minrbar = min(minrbar, ρbarold)
        end
        condA = max(maxrbar, ρtmp) / min(minrbar, ρtmp)

        # Test for convergence

        # Compute norms for convergence testing.
        normAr = abs(ζbar)

        # Now use these norms to estimate certain other quantities,
        # some of which will be small near a solution.
        test1 = normr / normb
        test2 = normAr / (normA * normr)
        test3 = inv(condA)

        t1 = test1 / (1 + normA * normx / normb)
        ϵr = ϵb + ϵa * normA * normx / normb

        # Allow for tolerances set by the user.
        if     test1 <= ϵr      istop = 1
        elseif test2 <= ϵa      istop = 2
        elseif test3 <= ϵc      istop = 3
        # The following tests guard against extremely small values of
        # atol, btol or ctol.  (The user may have set any or all of
        # the parameters atol, btol, conlim  to 0.)
        # The effect is equivalent to the normAl tests using
        # atol = eps,  btol = eps,  conlim = 1/eps.
        elseif 1 + t1    <= 1   istop = 4
        elseif 1 + test2 <= 1   istop = 5
        elseif 1 + test3 <= 1   istop = 6
        elseif i >= maxit       istop = 7
        end

        if verbose
            if size(A, 2) <= 50 ||
                maxit <= 50     ||
                i <= 10         ||
                i >= maxit-10   ||
                mod(i, 25) == 0 ||
                test3 <= 1.1*ϵc ||
                test2 <= 1.1*ϵa ||
                test1 <= 1.1*ϵr ||
                istop != 0
                @printf("[%4d/%d]", i, maxit)
                @printf("  %10.3e  %10.3e", normr, normAr)
                @printf("  %8.1e  %8.1e", test1, test2)
                @printf("  %8.1e  %8.1e\n", normA, condA)
            end
        end

        if istop > 0
            break
        end
    end

    if verbose
        @printf("\n")
    end

    return x
end

@inline function _xpby_norm!(X, b, Y)
    @. Y = muladd(b, Y, X)
    return Y, norm(Y)
end

@inline function _xpby_norm!(X::Array, b::Real, Y::Array{T}) where {T}
    Tr = real(T)
    bT = convert(Tr, b)
    @inbounds @batch threadlocal=zero(T)::T for I in eachindex(Y)
        Y[I] = muladd(bT, Y[I], X[I])
        threadlocal = muladd(conj(Y[I]), Y[I], threadlocal)
    end
    normY = sqrt(sum(real, threadlocal::Vector{T}))
    return Y, normY
end

@inline function _scal!(a, X)
    X .*= a
    return X
end

@inline function _scal!(a, X::Array{T}) where {T}
    aT = convert(real(T), a)
    @inbounds @batch for I in eachindex(X)
        X[I] *= aT
    end
    return X
end

@inline function _update_hbar_x_h!(h̄, x, h, v, θbar, ρ, ρold, ρbarold, ζ, ρbar, θnew)
    δ = -θbar * ρ / (ρold * ρbarold)
    σ = ζ / (ρ * ρbar)
    ν = -θnew / ρ

    @. h̄ = muladd(δ, h̄, h)
    @. x = muladd(σ, h̄, x)
    @. h = muladd(ν, h, v)

    return h̄, x, h, norm(x)
end

@inline function _update_hbar_x_h!(
    h̄::Array,
    x::Array{T},
    h::Array,
    v::Array,
    θbar::Real,
    ρ::Real,
    ρold::Real,
    ρbarold::Real,
    ζ::Real,
    ρbar::Real,
    θnew::Real
) where {T}
    δ = -θbar * ρ / (ρold * ρbarold)
    σ = ζ / (ρ * ρbar)
    ν = -θnew / ρ

    @inbounds @batch threadlocal=zero(T)::T for I in eachindex(x)
        h̄[I] = muladd(δ, h̄[I], h[I])
        x[I] = muladd(σ, h̄[I], x[I])
        h[I] = muladd(ν, h[I], v[I])
        threadlocal = muladd(conj(x[I]), x[I], threadlocal)
    end

    normx = sqrt(sum(real, threadlocal::Vector{T}))

    return h̄, x, h, normx
end


text:
using TimeZones: TZData

# Note: These tests are only meant to be run in a CI environment as they will modify the
# build time zones in the COMPILED_DIR.

@testset "build process" begin
    # Clean out deps directories for a clean re-build
    rm(TZData.COMPILED_DIR, recursive=true)
    for file in readdir(TZData.TZ_SOURCE_DIR)
        file == "utc" && continue
        rm(joinpath(TZData.TZ_SOURCE_DIR, file))
    end

    @test !isdir(TZData.COMPILED_DIR)
    @test length(readdir(TZData.TZ_SOURCE_DIR)) == 1

    # Using a version we already have avoids triggering a download
    TimeZones.build(TZDATA_VERSION)

    @test isdir(TZData.COMPILED_DIR)
    @test length(readdir(TZData.COMPILED_DIR)) > 0
    @test readdir(TZData.TZ_SOURCE_DIR) == sort(TZData.REGIONS)

    # Compile tz source files with an extended max_year. An example from the FAQ.
    warsaw = TimeZone("Europe/Warsaw")

    @test last(warsaw.transitions).utc_datetime == DateTime(2037, 10, 25, 1)
    @test warsaw.cutoff == DateTime(2038, 3, 28, 1)
    @test_throws TimeZones.UnhandledTimeError ZonedDateTime(DateTime(2039), warsaw)

    TZData.compile(max_year=2200)
    new_warsaw = TimeZone("Europe/Warsaw")

    @test warsaw !== new_warsaw
    @test last(new_warsaw.transitions).utc_datetime == DateTime(2200, 10, 26, 1)
    @test new_warsaw.cutoff == DateTime(2201, 3, 29, 1)
    ZonedDateTime(2100, new_warsaw)  # Test this doesn't throw an exception

    @test_throws TimeZones.UnhandledTimeError ZonedDateTime(2100, warsaw)

    # Using the tz string macro which runs at parse time means that the resulting TimeZone
    # will not reflect changes made from compile or new builds during runtime.
    @test tz"Africa/Windhoek" != TimeZone("Africa/Windhoek")
    @test tz"Africa/Windhoek".cutoff != DateTime(2201, 4, 5)
    @test TimeZone("Africa/Windhoek").cutoff == DateTime(2201, 4, 5)
end


text:
function euclidean(v1::Array{T1,1}, v2::Array{T2,1})::Float64 where {T1 <: Number,T2 <: Number}
    (v1 .- v2).^2.0 |> sum |> sqrt
end

function euclidean(v1::Array{T,1})::Float64 where T <: Number
    v2 = zeros(Float64, length(v1))
    return euclidean(v1, v2)
end

function normalize(v1::Array{T,1})::Array{Float64,1} where T <: Number
    return v1 ./ euclidean(v1)
end

function normalize(data::DataFrame)::DataFrame
    df = similar(data)
    _, p = size(df)
    for i in 1:p
        df[:, i] = normalize(data[:, i])
    end
    return df
end

function apply_columns(fs::Array{Function,1}, data::Union{DataFrame,Array{T,2}} where T <: Number)
    _, m = size(data)
    return [fs[i](data[:,i]) for i in 1:m]
end

function apply_columns(f::Function, data::Union{DataFrame,Array{T,2}} where T <: Number)
    return [f(c) for c in eachcol(data)]
end

function colmins(data::DataFrame)
    return apply_columns(minimum, data)
end

function colmaxs(data::DataFrame)
    return apply_columns(maximum, data)
end

function colsums(data)::Array{Number,1}
    return apply_columns(sum, data)
end

function colmeans(data)::Array{Number,1}
    return apply_columns(mean, data)
end

function apply_rows(f::Function, data::Union{DataFrame,Array{T,2}} where T <: Number)
    return [f(c) for c in eachrow(data)]
end

function rowmins(data::DataFrame)::Array{Float64,1}
    return apply_rows(minimum, data)
end

function rowmaxs(data::DataFrame)::Array{Float64,1}
    return apply_rows(maximum, data)
end

function rowsums(data)::Array{Float64,1}
    return apply_rows(sum, data)
end

function rowmeans(data)::Array{Float64,1}
    return apply_rows(mean, data)
end

function unitize(v::Array{T,1})::Array{Float64,1} where {T <: Number}
    return v ./ sum(v)
end

function Base.:*(w::Array{Float64,1}, data::DataFrame)::DataFrame
    newdf = copy(data)
    _, p = size(newdf)
    for i in 1:p
        newdf[!, i] = convert.(Float64, newdf[!,i])
        newdf[:, i] = w[i] .* data[:, i]
    end
    return newdf
end

function Base.:-(r1::DataFrameRow, r2::DataFrameRow)::Array{Number,1}
    # v1 = convert(Array{Float64,1}, r1)
    v1 = Vector{Float64}(r1)
    # v2 = convert(Array{Float64,1}, r2)
    v2 = Vector{Float64}(r2)
    return v1 .- v2
end

function Base.:-(r1::Array{T,1}, r2::DataFrameRow)::Array{T,1} where T <: Number
    # v2 = convert(Array{Float64,1}, r2)
    v2 = Vector{Float64}(r2)
    return r1 .- v2
end

function Base.:-(r1::DataFrameRow, r2::Array{T,1})::Array{T,1} where T <: Number
    v1 = convert(Array{Float64,1}, r1)
    return v1 .- r2
end

function makeDecisionMatrix(mat::Array{T,2}; names::Union{Nothing,Array{String,1}}=nothing)::DataFrame where {T <: Number}
    _, m = size(mat)
    df = DataFrame()
    for i in 1:m
        if names isa Nothing
            name = string("Crt", i)
        else
            name = names[i]
        end
        df[:,Symbol(name)] = convert(Array{Float64,1}, mat[:, i])
    end
    return df
end

function Base.minimum(df::DataFrame)
    df |> Matrix |> minimum
end

function reverseminmax(fns::Array{Function,1})::Array{Function,1}
    newfs = map(x -> if x == minimum maximum else minimum end, fns)
    return newfs
end

function makeminmax(fns::Array{K,1} where K)::Array{Function,1}
    return convert(Array{Function,1}, fns)
end

max_stars_repo_name: bahadirfyildirim/JMcDM

max_issues_repo_name: bahadirfyildirim/JMcDM

max_forks_repo_name: bahadirfyildirim/JMcDM


text:
"""
# Counting Sort
OVERVIEW:
Counting Sort is a sorting algorithm that sort elements within a specific range.
The sorting technique is to count the existing element and stored its occurrence time in a new list, then only print it out.

STEPS:
Assume the input as --> x=[-3, 1, -5, 0, -3]
  minimum = -5
- STEP 1: Create a list size within the range, in this case is -5 --> 1 which have range of 7 (-5, -4, -3, -2, -1, 0, 1), so list with size 7 and assign all to 0 is created
- STEP 2: Count the occurances of element in the list
          First number = -3 it is the third number in the range, so count[3]+=1
          Final view:
          index : ( 1,  2,  3,  4,  5, 6, 7)
          range : (-5, -4, -3, -2, -1, 0, 1)
          count : [ 1,  0,  2,  0,  0, 1, 1] <-- the list will store this occurences
- STEP 3: Make the count list accumulate the occurances
          The final count is (1, 1, 3, 3, 3, 4, 5)
- STEP 4: Assign the elements in x into correct possition by creating a new list (will call 'output' in this sample)
          the 1st element in 'x' is -3, it is third in range, so it will call the index of 3 in 'count', which is 3 and assign the -3 in to 3rd position in 'output',
          then the third element in range will deduct by 1, so the next repeated element will get the correct position, new 'count' --> [1, 1, 2, 3, 3, 4, 5]

          the 2nd element in 'x' is  1, it is last  in range, so it will call the index of 7 in 'count', which is 5 and assign the  1 in to 5th position in 'output',
          new 'count' --> [1, 1, 2, 3, 3, 4, 4]
          ......
          ......
          *If you want the order of original array to have the same order as the output array use can change this to decremental for loop
- STEP 5: Assign the 'ouput' list back to 'x'

FINAL RESULT -->  [-5, -3, -3, 0, 1]
"""
function counting_sort!(
    arr::Vector{T},
    l::Int = 1,
    r::Int = length(arr),
) where {T}
    if l >= r
        return nothing
    end
    max = maximum(arr)
    min = minimum(arr)
    range = max - min + 1
    count = Vector{T}(undef, range)
    output = Vector{T}(undef, r)

    for i in 1:range
        count[i] = 0
    end

    # Store count of the item that appear in the 'arr' (STEP 2)
    for i in 1:r
        count[arr[i]-min+1] += 1
    end

    # Calculate cumulative frequency in 'count' (STEP 3)
    for i in 2:length(count)
        count[i] += count[i-1]
    end

    # Build the 'output' by assigning the item into correct position (STEP 4)
    for i in 1:r
        index = arr[i] - min + 1
        output[count[index]] = arr[i]
        count[index] -= 1
    end

    # Copy the 'output' to 'arr', so that 'arr' contain sorted item (STEP 5)
    for i in 1:r
        arr[i] = output[i]
    end
end

max_stars_repo_name: AugustoCL/Julia

max_issues_repo_name: AugustoCL/Julia

max_forks_repo_name: AugustoCL/Julia