text:
using Compat
using Distances
using PersistenceDiagrams
using Ripserer
using SparseArrays
using StaticArrays
using Test

include("../testdatasets.jl")
include("interfacetest.jl")
using Ripserer: nv, edges

# Julia 1.0 does not allow these to be defined inside @testset.

# This test is mostly supposed to test that pre-finding apparent pairs with
# `find_apparent_pairs` first works correctly. On the side, it also tests that rips
# filtrations that only overload `adjacency_matrix` and `threshold` work fine.
struct ApparentPairsRips{I,T,R<:Rips{I,T}} <: Ripserer.AbstractRipsFiltration{I,T}
    rips::R
    apparent::Ref{Int}
    normal::Ref{Int}
end

function ApparentPairsRips(data; kwargs...)
    return ApparentPairsRips(Rips(data; kwargs...), Ref(0), Ref(0))
end

Ripserer.adjacency_matrix(rw::ApparentPairsRips) = Ripserer.adjacency_matrix(rw.rips)
Ripserer.threshold(rw::ApparentPairsRips) = Ripserer.threshold(rw.rips)

# This works fine but is slower than the original algorithm (even if the pair finding code
# was optimized) as apparent pairs are found in `initialize_coboundary!` anyway. Doing this
# in parallel or on the GPU is an option, however.  See https://arxiv.org/abs/2003.07989
function Ripserer.find_apparent_pairs(rw::ApparentPairsRips, columns, _)
    S = eltype(columns)
    C = Ripserer.simplex_type(rw, dim(S) + 1)
    cols_left = S[]
    apparent = Tuple{S,C}[]
    for σ in columns
        τ = minimum(Ripserer.coboundary(rw, σ)) # This is broken if coboundary is empty.
        σ′ = maximum(Ripserer.boundary(rw, τ))
        if σ′ == σ
            rw.apparent[] += 1
            push!(apparent, (σ, τ))
        else
            rw.normal[] += 1
            push!(cols_left, σ)
        end
    end
    return cols_left, apparent
end

@testset "Rips with `find_apparent_pairs` overloaded." begin
    for m in (2, 3, 5), t in (2, 1)
        @testset "projective plane with modulus=$m, threshold=$t" begin
            appa = ApparentPairsRips(projective_plane; threshold=t)
            res_rips = ripserer(Rips(projective_plane; threshold=t); modulus=m)
            res_appa = ripserer(appa; modulus=m)
            @test res_rips == res_appa
            # The nv - 1 edges were cleared by the twist algorithm.
            @test appa.apparent[] + appa.normal[] + nv(appa) - 1 == length(edges(appa))
        end
    end
    for m in (2, 3, 5), t in (8, 4)
        @testset "cycle with modulus=$m, threshold=$t" begin
            appa = ApparentPairsRips(cycle; threshold=t)
            res_rips = ripserer(Rips(cycle; threshold=t); modulus=m, dim_max=3)
            res_appa = ripserer(appa; modulus=m, dim_max=3)
            @test res_rips == res_appa
            @test appa.apparent[] > 0
        end
    end

    @testset "Interface" begin
        test_filtration(ApparentPairsRips, projective_plane; modulus=2)
        test_filtration(ApparentPairsRips, projective_plane; modulus=3)
    end
end

# This test checks that apparent pairs can produce correct intervals. On the side, it checks
# `AbstractCustomFiltration`s.
struct ApparentPairsCustom <: Ripserer.AbstractCustomFiltration{Int,Int}
    found::Ref{Bool}
end
ApparentPairsCustom() = ApparentPairsCustom(Ref(false))

function Ripserer.simplex_dicts(::ApparentPairsCustom)
    return [Dict(1 => 0, 2 => 0, 3 => 0), Dict(1 => 1, 2 => 1, 3 => 2), Dict(1 => 3)]
end
Ripserer.adjacency_matrix(::ApparentPairsCustom) = sparse(Bool[0 1 1; 1 0 1; 1 1 0])
function Ripserer.find_apparent_pairs(a::ApparentPairsCustom, columns, _)
    σ = columns[1]
    a.found[] = true
    return typeof(σ)[], [(σ, only(Ripserer.coboundary(a, σ)))]
end

@testset "Apparent pairs can produce valid intervals." begin
    @testset "cohomology" begin
        appa = ApparentPairsCustom()
        @test !appa.found[]
        d0, d1 = ripserer(appa; reps=true)
        @test appa.found[]
        @test d0 == [(0, 1), (0, 1), (0, Inf)]
        @test d1 == [(2, 3)]
        @test d1[1].birth_simplex == Simplex{1}(3, 2)
        @test d1[1].death_simplex == Simplex{2}(1, 3)
        @test simplex(only(d1[1].representative)) == Simplex{1}(3, 2)
    end
end

# The main idea of this test is to test postprocess_diagrams
struct ReversedResult <: Ripserer.AbstractFiltration{Int,Int} end

function Ripserer.unsafe_simplex(::Type{Simplex{0,Int,Int}}, ::ReversedResult, (v,))
    return Simplex{0}(v, 0)
end
function Ripserer.unsafe_simplex(
    ::Type{Simplex{D,Int,Int}}, ::ReversedResult, vertices
) where {D}
    return Simplex{D}(index(vertices), 1)
end
Ripserer.nv(::ReversedResult) = 10
Ripserer.simplex_type(::Type{ReversedResult}, D) = Simplex{D,Int,Int}
Ripserer.edges(::ReversedResult) = Simplex{1}.(10:-1:1, 1)
function Ripserer.postprocess_diagram(::ReversedResult, diagram)
    return reverse!(diagram)
end

@testset "New filtration" begin
    d0, d1 = ripserer(ReversedResult())
    @test d0 == [fill((0.0, Inf), 6); fill((0.0, 1.0), 4)]
    @test d1 == []
end


text:
# This file investigates integration error with respect to free variables.

using FractalTools
using GeometryBasics
using Makie

# Construct interpolation data
f(x, y) = [
    x^2 + y^2 + 1,
    x^2 - y^2
    ]
ngon = Triangle(
    Point(BigFloat(0.), BigFloat(0.)),
    Point(BigFloat(1.), BigFloat(0.)),
    Point(BigFloat(0.5), BigFloat(1.)))
npts = 100
pts = getdata(f, ngon, npts)

# Compute errors
fval = 35 / 48
freevars = 0.001 : 0.001 : 0.025
mse = map(freevars) do freevar
    ival = integrate(pts, HInterp2D(fill(freevar, 2, 2)))[1]
    abs(fval - ival)
end

# Plot mse
fig = Figure()
ax = fig[1, 1] = Axis(fig, xlabel="Free Variable", ylabel="Integration Error", title="2D Hidden Integration Error")
stem!(ax, freevars, mse, color=:black)
save(joinpath(@__DIR__, "hinteg2d_error.png"), fig)
display(fig)

ext: jl


text:
abstract type AbstractExtendedKalmanFilter <: AbstractKalmanFilter end
@with_kw struct ExtendedKalmanFilter{KF <: KalmanFilter, F, G} <: AbstractExtendedKalmanFilter
    kf::KF
    dynamics::F
    measurement::G
end

function Base.getproperty(ekf::EKF, s::Symbol) where EKF <: AbstractExtendedKalmanFilter
    s ∈ fieldnames(EKF) && return getfield(ekf, s)
    return getproperty(ekf.kf, s)
end

function Base.propertynames(ekf::EKF, private::Bool=false) where EKF <: AbstractExtendedKalmanFilter
    return (fieldnames(EKF)..., propertynames(ekf.kf, private)...)
end

function predict!(kf::AbstractExtendedKalmanFilter, u, t::Integer = index(kf))
    @unpack x,R,R1 = kf
    A = ForwardDiff.jacobian(x->kf.dynamics(x,u,t), x)
    x .= kf.dynamics(x, u, t)
    R .= symmetrize(A*R*A') + R1
    kf.t[] += 1
end

function correct!(kf::AbstractExtendedKalmanFilter, u, y, t::Integer = index(kf))
    @unpack x,R,R2 = kf
    C = ForwardDiff.jacobian(x->kf.measurement(x,u,t), x)
    e  = y .- kf.measurement(x,u,t)
    S   = symmetrize(C*R*C') + R2
    Sᵪ  = cholesky(S)
    K   = (R*C')/Sᵪ
    x .+= vec(K*e)
    R  .= symmetrize((I - K*C)*R) # WARNING against I .- A
    ll = logpdf(MvNormal(PDMat(S, Sᵪ)), e)[]# - 1/2*logdet(S) # logdet is included in logpdf
    ll, e
end

function smooth(kf::AbstractExtendedKalmanFilter, u::AbstractVector, y::AbstractVector)
    reset!(kf)
    T            = length(y)
    x,xt,R,Rt,ll = forward_trajectory(kf, u, y)
    xT           = similar(xt)
    RT           = similar(Rt)
    xT[end]      = xt[end]      |> copy
    RT[end]      = Rt[end]      |> copy
    for t = T-1:-1:1
        A = ForwardDiff.jacobian(x->kf.dynamics(x,u[t+1],t+1), xt[t+1])
        C     = Rt[t]*A/R[t+1]
        xT[t] = xt[t] .+ C*(xT[t+1] .- x[t+1])
        # xT[t][end] = clamp(xT[t][end], 0.01, 7.499)
        RT[t] = Rt[t] .+ symmetrize(C*(RT[t+1] .- R[t+1])*C')
    end
    xT,RT,ll
end

sample_state(kf::AbstractExtendedKalmanFilter) = rand(kf.d0)
sample_state(kf::AbstractExtendedKalmanFilter, x, u, t) = kf.dynamics(x, u, t) .+ rand(MvNormal(kf.R1))
sample_measurement(kf::AbstractExtendedKalmanFilter, x, u, t) = kf.measurement(x, u, t) .+ rand(MvNormal(kf.R2))


text:
using CsminWel
using Base.Test

# write your own tests here
f(x) = sum(x.^2)
function grad(x)
    2*x, false
end

function g!(x, stor)
    stor[:] = 2*x
end

res1 = optimize(f, g!, [.1, .1], Csminwel())
res2 = optimize(f, [.1, .1], Csminwel())
res3 = optimize(f, [.1, .1], Csminwel(), Optim.Options(autodiff=true))

@test_approx_eq_eps Optim.minimum(res2)  Optim.minimum(res1) 1e-09
@test_approx_eq_eps Optim.minimum(res2)  Optim.minimum(res3) 1e-09

@test_approx_eq_eps Optim.minimizer(res2)  Optim.minimizer(res1) 1e-09
@test_approx_eq_eps Optim.minimizer(res2)  Optim.minimizer(res3) 1e-09

#=
Maximizing loglikelihood logistic models
=#
using StatsFuns
srand(1)
x = [ones(200) randn(200,4)]
y = [rand() < 0.5 ? 1. : 0. for j in 1:200]

function loglik(beta)
    xb = x*beta
    sum(-y.*xb + log1pexp.(xb))
end

function dloglik(beta)
    xb = x*beta
    px = logistic.(xb)
    -x'*(y.-px)
end

function fg!(beta, stor)
    stor[:] = dloglik(beta)
end

res1 = optimize(loglik, fg!, zeros(5), BFGS())
res2 = optimize(loglik, fg!, zeros(5), Csminwel())
res3 = optimize(loglik, zeros(5), Csminwel(), Optim.Options(autodiff=true))
res4 = optimize(loglik, zeros(5), Csminwel(), Optim.Options(autodiff=false))

@test_approx_eq_eps Optim.minimum(res2)  Optim.minimum(res1) 1e-06
@test_approx_eq_eps Optim.minimum(res2)  Optim.minimum(res3) 1e-06
@test_approx_eq_eps Optim.minimum(res4)  Optim.minimum(res3) 1e-06

@test_approx_eq_eps Optim.minimizer(res2)  Optim.minimizer(res1) 1e-06
@test_approx_eq_eps Optim.minimizer(res2)  Optim.minimizer(res3) 1e-06
@test_approx_eq_eps Optim.minimizer(res3)  Optim.minimizer(res4) 1e-06


text:
# Load Julia packages (libraries)

using Pkg, DrWatson

using MonteCarloMeasurements
using Random
using StanSample
using StatisticalRethinking

Random.seed!(12345)

stan12_3 = "
data{
    int N;
    int y[N];
}
parameters{
    real ap;
    real al;
}
model{
    real p;
    real lambda;
    al ~ normal( 1 , 0.5 );
    ap ~ normal( -1.5 , 1 );
    lambda = al;
    lambda = exp(lambda);
    p = ap;
    p = inv_logit(p);
    for ( i in 1:N ) {
        if ( y[i]==0 )
            target += log_mix( p , 0 , poisson_lpmf(0|lambda) );
        if ( y[i] > 0 )
            target += log1m( p ) + poisson_lpmf(y[i] | lambda );
    }
}";

# Define the SampleModel.

m12_3s = SampleModel("m12.3s",  stan12_3);

# Input data for cmdstan

prob_drink = 0.2 # 20% of days
rate_work = 1    # average 1 manuscript per day

# sample one year of production
N = 365
drink = rand(Binomial(1, prob_drink ), N)

# simulate manuscripts completed
y = (1 .- drink) .* rand(Poisson(rate_work ), N)

## R code 12.8
#histogram( y , xlab="manuscripts completed" , lwd=4 )

zeros_drink = sum(drink)
zeros_work = sum(y==0 && drink==0)
zeros_total = sum(y==0)

m12_3_data = Dict("N" => length(y), "y" => y);

# Sample using cmdstan's sample option

rc12_3s = stan_sample(m12_3s, data=m12_3_data);

# Describe the draws

if success(rc12_3s)
  part12_3s = read_samples(m12_3s, :particles)
  part12_3s |> display
end


text:
using DiffEqSensitivity, ParameterizedFunctions
using OrdinaryDiffEq, Calculus, Test

f = @ode_def AAA begin
    dx = -a
    dy = b
end a b

cb = ContinuousCallback((u,t,i) -> u[1], (integrator)->(println("Stopped.");integrator.p[2]=0.0))
p = [2.0, 1.0]

function test_f(p)
  prob = ODEProblem(f,eltype(p).([1.0,0.0]),eltype(p).((0.0,1.0)),copy(p))
  solve(prob,Tsit5(),abstol=1e-14,reltol=1e-14,callback=cb,save_everystep=false)[end]
end
p = [2.0, 1.0]
findiff = Calculus.finite_difference_jacobian(test_f,p)
findiff

using ForwardDiff
ad = ForwardDiff.jacobian(test_f,p)
ad

@test ad ≈ findiff


text:
struct NodeIterator
    iter
end

Iterators.IteratorSize(::NodeIterator) = Iterators.SizeUnknown()
Base.iterate(x::NodeIterator, args...) = iterate(x.iter, args...)

function eachnode(tree)
    return NodeIterator(node_iterator(tree))
end

node_iterator(node::AbstractTreeViewNode) = Iterators.flatten(((node,), child_node_iterator(node)))
child_node_iterator(node::AbstractTreeViewNode) = Iterators.flatten(node_iterator(x) for x in children(node))

node_iterator(tree::TreeViewRoot) = Iterators.flatten(node_iterator(x) for x in tree.nodes)
node_iterator(treeview::TreeView) = node_iterator(treeview.content)

function find_by_id(tree::TreeViewRoot, id::String)
    for node in eachnode(tree)
        if tree_id(node)==id
            return node
        end
    end
    return nothing
end

Base.getindex(tree::TreeViewRoot, id::String) = find_by_id(tree, id)
Base.getindex(treeview::TreeView, id::String) = find_by_id(treeview.content, id)

Base.getindex(tree::TreeViewRoot, d::Dict{String, String}) = find_by_id(tree, d["id"])
Base.getindex(treeview::TreeView, d::Dict{String, String}) = find_by_id(treeview.content, d["id"])


text:
using teste
using Test

@testset "teste.jl" begin
    # Write your own tests here.
end


text:
# Julia wrapper for header: api.h
# Automatically generated using Clang.jl

function ts_parser_new()
    ccall((:ts_parser_new, libtreesitter), Ptr{TSParser}, ())
end

function ts_parser_delete(parser)
    ccall((:ts_parser_delete, libtreesitter), Cvoid, (Ptr{TSParser},), parser)
end

function ts_parser_set_language(self, language)
    ccall(
        (:ts_parser_set_language, libtreesitter),
        Bool,
        (Ptr{TSParser}, Ptr{TSLanguage}),
        self,
        language,
    )
end

function ts_parser_language(self)
    ccall(
        (:ts_parser_language, libtreesitter),
        Ptr{TSLanguage},
        (Ptr{TSParser},),
        self,
    )
end

function ts_parser_set_included_ranges(self, ranges, length)
    ccall(
        (:ts_parser_set_included_ranges, libtreesitter),
        Bool,
        (Ptr{TSParser}, Ptr{TSRange}, UInt32),
        self,
        ranges,
        length,
    )
end

function ts_parser_included_ranges(self, length)
    ccall(
        (:ts_parser_included_ranges, libtreesitter),
        Ptr{TSRange},
        (Ptr{TSParser}, Ptr{UInt32}),
        self,
        length,
    )
end

function ts_parser_parse(self, old_tree, input)
    ccall(
        (:ts_parser_parse, libtreesitter),
        Ptr{TSTree},
        (Ptr{TSParser}, Ptr{TSTree}, TSInput),
        self,
        old_tree,
        input,
    )
end

function ts_parser_parse_string(self, old_tree, string, length)
    ccall(
        (:ts_parser_parse_string, libtreesitter),
        Ptr{TSTree},
        (Ptr{TSParser}, Ptr{TSTree}, Cstring, UInt32),
        self,
        old_tree,
        string,
        length,
    )
end

function ts_parser_parse_string_encoding(
    self,
    old_tree,
    string,
    length,
    encoding,
)
    ccall(
        (:ts_parser_parse_string_encoding, libtreesitter),
        Ptr{TSTree},
        (Ptr{TSParser}, Ptr{TSTree}, Cstring, UInt32, TSInputEncoding),
        self,
        old_tree,
        string,
        length,
        encoding,
    )
end

function ts_parser_reset(self)
    ccall((:ts_parser_reset, libtreesitter), Cvoid, (Ptr{TSParser},), self)
end

function ts_parser_set_timeout_micros(self, timeout)
    ccall(
        (:ts_parser_set_timeout_micros, libtreesitter),
        Cvoid,
        (Ptr{TSParser}, UInt64),
        self,
        timeout,
    )
end

function ts_parser_timeout_micros(self)
    ccall(
        (:ts_parser_timeout_micros, libtreesitter),
        UInt64,
        (Ptr{TSParser},),
        self,
    )
end

function ts_parser_set_cancellation_flag(self, flag)
    ccall(
        (:ts_parser_set_cancellation_flag, libtreesitter),
        Cvoid,
        (Ptr{TSParser}, Ptr{Csize_t}),
        self,
        flag,
    )
end

function ts_parser_cancellation_flag(self)
    ccall(
        (:ts_parser_cancellation_flag, libtreesitter),
        Ptr{Csize_t},
        (Ptr{TSParser},),
        self,
    )
end

function ts_parser_set_logger(self, logger)
    ccall(
        (:ts_parser_set_logger, libtreesitter),
        Cvoid,
        (Ptr{TSParser}, TSLogger),
        self,
        logger,
    )
end

function ts_parser_logger(self)
    ccall((:ts_parser_logger, libtreesitter), TSLogger, (Ptr{TSParser},), self)
end

function ts_parser_print_dot_graphs(self, file)
    ccall(
        (:ts_parser_print_dot_graphs, libtreesitter),
        Cvoid,
        (Ptr{TSParser}, Cint),
        self,
        file,
    )
end

function ts_tree_copy(self)
    ccall((:ts_tree_copy, libtreesitter), Ptr{TSTree}, (Ptr{TSTree},), self)
end

function ts_tree_delete(self)
    ccall((:ts_tree_delete, libtreesitter), Cvoid, (Ptr{TSTree},), self)
end

function ts_tree_root_node(self)
    ccall((:ts_tree_root_node, libtreesitter), TSNode, (Ptr{TSTree},), self)
end

function ts_tree_language(arg1)
    ccall(
        (:ts_tree_language, libtreesitter),
        Ptr{TSLanguage},
        (Ptr{TSTree},),
        arg1,
    )
end

function ts_tree_edit(self, edit)
    ccall(
        (:ts_tree_edit, libtreesitter),
        Cvoid,
        (Ptr{TSTree}, Ptr{TSInputEdit}),
        self,
        edit,
    )
end

function ts_tree_get_changed_ranges(old_tree, new_tree, length)
    ccall(
        (:ts_tree_get_changed_ranges, libtreesitter),
        Ptr{TSRange},
        (Ptr{TSTree}, Ptr{TSTree}, Ptr{UInt32}),
        old_tree,
        new_tree,
        length,
    )
end

function ts_tree_print_dot_graph(arg1, arg2)
    ccall(
        (:ts_tree_print_dot_graph, libtreesitter),
        Cvoid,
        (Ptr{TSTree}, Ptr{Libc.FILE}),
        arg1,
        arg2,
    )
end

function ts_node_type(arg1)
    ccall((:ts_node_type, libtreesitter), Cstring, (TSNode,), arg1)
end

function ts_node_symbol(arg1)
    ccall((:ts_node_symbol, libtreesitter), TSSymbol, (TSNode,), arg1)
end

function ts_node_start_byte(arg1)
    ccall((:ts_node_start_byte, libtreesitter), UInt32, (TSNode,), arg1)
end

function ts_node_start_point(arg1)
    ccall((:ts_node_start_point, libtreesitter), TSPoint, (TSNode,), arg1)
end

function ts_node_end_byte(arg1)
    ccall((:ts_node_end_byte, libtreesitter), UInt32, (TSNode,), arg1)
end

function ts_node_end_point(arg1)
    ccall((:ts_node_end_point, libtreesitter), TSPoint, (TSNode,), arg1)
end

function ts_node_string(arg1)
    ccall((:ts_node_string, libtreesitter), Cstring, (TSNode,), arg1)
end

function ts_node_is_null(arg1)
    ccall((:ts_node_is_null, libtreesitter), Bool, (TSNode,), arg1)
end

function ts_node_is_named(arg1)
    ccall((:ts_node_is_named, libtreesitter), Bool, (TSNode,), arg1)
end

function ts_node_is_missing(arg1)
    ccall((:ts_node_is_missing, libtreesitter), Bool, (TSNode,), arg1)
end

function ts_node_is_extra(arg1)
    ccall((:ts_node_is_extra, libtreesitter), Bool, (TSNode,), arg1)
end

function ts_node_has_changes(arg1)
    ccall((:ts_node_has_changes, libtreesitter), Bool, (TSNode,), arg1)
end

function ts_node_has_error(arg1)
    ccall((:ts_node_has_error, libtreesitter), Bool, (TSNode,), arg1)
end

function ts_node_parent(arg1)
    ccall((:ts_node_parent, libtreesitter), TSNode, (TSNode,), arg1)
end

function ts_node_child(arg1, arg2)
    ccall((:ts_node_child, libtreesitter), TSNode, (TSNode, UInt32), arg1, arg2)
end

function ts_node_child_count(arg1)
    ccall((:ts_node_child_count, libtreesitter), UInt32, (TSNode,), arg1)
end

function ts_node_named_child(arg1, arg2)
    ccall(
        (:ts_node_named_child, libtreesitter),
        TSNode,
        (TSNode, UInt32),
        arg1,
        arg2,
    )
end

function ts_node_named_child_count(arg1)
    ccall((:ts_node_named_child_count, libtreesitter), UInt32, (TSNode,), arg1)
end

function ts_node_child_by_field_name(self, field_name, field_name_length)
    ccall(
        (:ts_node_child_by_field_name, libtreesitter),
        TSNode,
        (TSNode, Cstring, UInt32),
        self,
        field_name,
        field_name_length,
    )
end

function ts_node_child_by_field_id(arg1, arg2)
    ccall(
        (:ts_node_child_by_field_id, libtreesitter),
        TSNode,
        (TSNode, TSFieldId),
        arg1,
        arg2,
    )
end

function ts_node_next_sibling(arg1)
    ccall((:ts_node_next_sibling, libtreesitter), TSNode, (TSNode,), arg1)
end

function ts_node_prev_sibling(arg1)
    ccall((:ts_node_prev_sibling, libtreesitter), TSNode, (TSNode,), arg1)
end

function ts_node_next_named_sibling(arg1)
    ccall((:ts_node_next_named_sibling, libtreesitter), TSNode, (TSNode,), arg1)
end

function ts_node_prev_named_sibling(arg1)
    ccall((:ts_node_prev_named_sibling, libtreesitter), TSNode, (TSNode,), arg1)
end

function ts_node_first_child_for_byte(arg1, arg2)
    ccall(
        (:ts_node_first_child_for_byte, libtreesitter),
        TSNode,
        (TSNode, UInt32),
        arg1,
        arg2,
    )
end

function ts_node_first_named_child_for_byte(arg1, arg2)
    ccall(
        (:ts_node_first_named_child_for_byte, libtreesitter),
        TSNode,
        (TSNode, UInt32),
        arg1,
        arg2,
    )
end

function ts_node_descendant_for_byte_range(arg1, arg2, arg3)
    ccall(
        (:ts_node_descendant_for_byte_range, libtreesitter),
        TSNode,
        (TSNode, UInt32, UInt32),
        arg1,
        arg2,
        arg3,
    )
end

function ts_node_descendant_for_point_range(arg1, arg2, arg3)
    ccall(
        (:ts_node_descendant_for_point_range, libtreesitter),
        TSNode,
        (TSNode, TSPoint, TSPoint),
        arg1,
        arg2,
        arg3,
    )
end

function ts_node_named_descendant_for_byte_range(arg1, arg2, arg3)
    ccall(
        (:ts_node_named_descendant_for_byte_range, libtreesitter),
        TSNode,
        (TSNode, UInt32, UInt32),
        arg1,
        arg2,
        arg3,
    )
end

function ts_node_named_descendant_for_point_range(arg1, arg2, arg3)
    ccall(
        (:ts_node_named_descendant_for_point_range, libtreesitter),
        TSNode,
        (TSNode, TSPoint, TSPoint),
        arg1,
        arg2,
        arg3,
    )
end

function ts_node_edit(arg1, arg2)
    ccall(
        (:ts_node_edit, libtreesitter),
        Cvoid,
        (Ptr{TSNode}, Ptr{TSInputEdit}),
        arg1,
        arg2,
    )
end

function ts_node_eq(arg1, arg2)
    ccall((:ts_node_eq, libtreesitter), Bool, (TSNode, TSNode), arg1, arg2)
end

function ts_tree_cursor_new(arg1)
    ccall((:ts_tree_cursor_new, libtreesitter), TSTreeCursor, (TSNode,), arg1)
end

function ts_tree_cursor_delete(arg1)
    ccall(
        (:ts_tree_cursor_delete, libtreesitter),
        Cvoid,
        (Ptr{TSTreeCursor},),
        arg1,
    )
end

function ts_tree_cursor_reset(arg1, arg2)
    ccall(
        (:ts_tree_cursor_reset, libtreesitter),
        Cvoid,
        (Ptr{TSTreeCursor}, TSNode),
        arg1,
        arg2,
    )
end

function ts_tree_cursor_current_node(arg1)
    ccall(
        (:ts_tree_cursor_current_node, libtreesitter),
        TSNode,
        (Ptr{TSTreeCursor},),
        arg1,
    )
end

function ts_tree_cursor_current_field_name(arg1)
    ccall(
        (:ts_tree_cursor_current_field_name, libtreesitter),
        Cstring,
        (Ptr{TSTreeCursor},),
        arg1,
    )
end

function ts_tree_cursor_current_field_id(arg1)
    ccall(
        (:ts_tree_cursor_current_field_id, libtreesitter),
        TSFieldId,
        (Ptr{TSTreeCursor},),
        arg1,
    )
end

function ts_tree_cursor_goto_parent(arg1)
    ccall(
        (:ts_tree_cursor_goto_parent, libtreesitter),
        Bool,
        (Ptr{TSTreeCursor},),
        arg1,
    )
end

function ts_tree_cursor_goto_next_sibling(arg1)
    ccall(
        (:ts_tree_cursor_goto_next_sibling, libtreesitter),
        Bool,
        (Ptr{TSTreeCursor},),
        arg1,
    )
end

function ts_tree_cursor_goto_first_child(arg1)
    ccall(
        (:ts_tree_cursor_goto_first_child, libtreesitter),
        Bool,
        (Ptr{TSTreeCursor},),
        arg1,
    )
end

function ts_tree_cursor_goto_first_child_for_byte(arg1, arg2)
    ccall(
        (:ts_tree_cursor_goto_first_child_for_byte, libtreesitter),
        Int64,
        (Ptr{TSTreeCursor}, UInt32),
        arg1,
        arg2,
    )
end

function ts_tree_cursor_copy(arg1)
    ccall(
        (:ts_tree_cursor_copy, libtreesitter),
        TSTreeCursor,
        (Ptr{TSTreeCursor},),
        arg1,
    )
end

function ts_query_new(language, source, source_len, error_offset, error_type)
    ccall(
        (:ts_query_new, libtreesitter),
        Ptr{TSQuery},
        (Ptr{TSLanguage}, Cstring, UInt32, Ptr{UInt32}, Ptr{TSQueryError}),
        language,
        source,
        source_len,
        error_offset,
        error_type,
    )
end

function ts_query_delete(arg1)
    ccall((:ts_query_delete, libtreesitter), Cvoid, (Ptr{TSQuery},), arg1)
end

function ts_query_pattern_count(arg1)
    ccall(
        (:ts_query_pattern_count, libtreesitter),
        UInt32,
        (Ptr{TSQuery},),
        arg1,
    )
end

function ts_query_capture_count(arg1)
    ccall(
        (:ts_query_capture_count, libtreesitter),
        UInt32,
        (Ptr{TSQuery},),
        arg1,
    )
end

function ts_query_string_count(arg1)
    ccall(
        (:ts_query_string_count, libtreesitter),
        UInt32,
        (Ptr{TSQuery},),
        arg1,
    )
end

function ts_query_start_byte_for_pattern(arg1, arg2)
    ccall(
        (:ts_query_start_byte_for_pattern, libtreesitter),
        UInt32,
        (Ptr{TSQuery}, UInt32),
        arg1,
        arg2,
    )
end

function ts_query_predicates_for_pattern(self, pattern_index, length)
    ccall(
        (:ts_query_predicates_for_pattern, libtreesitter),
        Ptr{TSQueryPredicateStep},
        (Ptr{TSQuery}, UInt32, Ptr{UInt32}),
        self,
        pattern_index,
        length,
    )
end

function ts_query_capture_name_for_id(arg1, id, length)
    ccall(
        (:ts_query_capture_name_for_id, libtreesitter),
        Cstring,
        (Ptr{TSQuery}, UInt32, Ptr{UInt32}),
        arg1,
        id,
        length,
    )
end

function ts_query_string_value_for_id(arg1, id, length)
    ccall(
        (:ts_query_string_value_for_id, libtreesitter),
        Cstring,
        (Ptr{TSQuery}, UInt32, Ptr{UInt32}),
        arg1,
        id,
        length,
    )
end

function ts_query_disable_capture(arg1, arg2, arg3)
    ccall(
        (:ts_query_disable_capture, libtreesitter),
        Cvoid,
        (Ptr{TSQuery}, Cstring, UInt32),
        arg1,
        arg2,
        arg3,
    )
end

function ts_query_disable_pattern(arg1, arg2)
    ccall(
        (:ts_query_disable_pattern, libtreesitter),
        Cvoid,
        (Ptr{TSQuery}, UInt32),
        arg1,
        arg2,
    )
end

function ts_query_cursor_new()
    ccall((:ts_query_cursor_new, libtreesitter), Ptr{TSQueryCursor}, ())
end

function ts_query_cursor_delete(arg1)
    ccall(
        (:ts_query_cursor_delete, libtreesitter),
        Cvoid,
        (Ptr{TSQueryCursor},),
        arg1,
    )
end

function ts_query_cursor_exec(arg1, arg2, arg3)
    ccall(
        (:ts_query_cursor_exec, libtreesitter),
        Cvoid,
        (Ptr{TSQueryCursor}, Ptr{TSQuery}, TSNode),
        arg1,
        arg2,
        arg3,
    )
end

function ts_query_cursor_set_byte_range(arg1, arg2, arg3)
    ccall(
        (:ts_query_cursor_set_byte_range, libtreesitter),
        Cvoid,
        (Ptr{TSQueryCursor}, UInt32, UInt32),
        arg1,
        arg2,
        arg3,
    )
end

function ts_query_cursor_set_point_range(arg1, arg2, arg3)
    ccall(
        (:ts_query_cursor_set_point_range, libtreesitter),
        Cvoid,
        (Ptr{TSQueryCursor}, TSPoint, TSPoint),
        arg1,
        arg2,
        arg3,
    )
end

function ts_query_cursor_next_match(arg1, match)
    ccall(
        (:ts_query_cursor_next_match, libtreesitter),
        Bool,
        (Ptr{TSQueryCursor}, Ptr{TSQueryMatch}),
        arg1,
        match,
    )
end

function ts_query_cursor_remove_match(arg1, id)
    ccall(
        (:ts_query_cursor_remove_match, libtreesitter),
        Cvoid,
        (Ptr{TSQueryCursor}, UInt32),
        arg1,
        id,
    )
end

function ts_query_cursor_next_capture(arg1, match, capture_index)
    ccall(
        (:ts_query_cursor_next_capture, libtreesitter),
        Bool,
        (Ptr{TSQueryCursor}, Ptr{TSQueryMatch}, Ptr{UInt32}),
        arg1,
        match,
        capture_index,
    )
end

function ts_language_symbol_count(arg1)
    ccall(
        (:ts_language_symbol_count, libtreesitter),
        UInt32,
        (Ptr{TSLanguage},),
        arg1,
    )
end

function ts_language_symbol_name(arg1, arg2)
    ccall(
        (:ts_language_symbol_name, libtreesitter),
        Cstring,
        (Ptr{TSLanguage}, TSSymbol),
        arg1,
        arg2,
    )
end

function ts_language_symbol_for_name(self, string, length, is_named)
    ccall(
        (:ts_language_symbol_for_name, libtreesitter),
        TSSymbol,
        (Ptr{TSLanguage}, Cstring, UInt32, Bool),
        self,
        string,
        length,
        is_named,
    )
end

function ts_language_field_count(arg1)
    ccall(
        (:ts_language_field_count, libtreesitter),
        UInt32,
        (Ptr{TSLanguage},),
        arg1,
    )
end

function ts_language_field_name_for_id(arg1, arg2)
    ccall(
        (:ts_language_field_name_for_id, libtreesitter),
        Cstring,
        (Ptr{TSLanguage}, TSFieldId),
        arg1,
        arg2,
    )
end

function ts_language_field_id_for_name(arg1, arg2, arg3)
    ccall(
        (:ts_language_field_id_for_name, libtreesitter),
        TSFieldId,
        (Ptr{TSLanguage}, Cstring, UInt32),
        arg1,
        arg2,
        arg3,
    )
end

function ts_language_symbol_type(arg1, arg2)
    ccall(
        (:ts_language_symbol_type, libtreesitter),
        TSSymbolType,
        (Ptr{TSLanguage}, TSSymbol),
        arg1,
        arg2,
    )
end

function ts_language_version(arg1)
    ccall(
        (:ts_language_version, libtreesitter),
        UInt32,
        (Ptr{TSLanguage},),
        arg1,
    )
end


text:
const makie_linetype = Dict{Symbol, Any}(
    :auto => nothing,
    :solid => nothing,
    :dash => :dash,
    :dot => :dot,
    :dashdot => :dashdot,
    :dashdotdot => [0.4, 0.2, 0.1, 0.2, 0.4]
)

function makie_color(c)
    if color === :match
        return AbstractPlotting.Colors.colorant"blue"
    end
    convert(AbstractPlotting.RGBA, c)
end
makie_seriestype_map = Dict{Symbol, Type}(
    :path => AbstractPlotting.Lines,
    :path3d => AbstractPlotting.Lines,
    :scatter => AbstractPlotting.Scatter,
    :linesegments => AbstractPlotting.LineSegments,
    :heatmap => AbstractPlotting.Heatmap,
    :image => AbstractPlotting.Image,
    :spy => AbstractPlotting.Spy,
    :surface => AbstractPlotting.Surface,
    :shape => AbstractPlotting.Poly,
    :contour => AbstractPlotting.Contour,
    :curves => Bezier,
    :bar => AbstractPlotting.BarPlot,
    # TODO: line, contour,
)


text:
############################################################################################
# Assign abstract types
"""
$(TYPEDEF)
Abstract super type for Baytes Models.
"""
abstract type ModelName end
"""
$(TYPEDEF)
Default modelname of Baytes.Model struct.
"""
struct BaseModel <: ModelName end

############################################################################################
#!NOTE: For M<:Union{P,ModelName} where {P}, keep A so option using of other PPLs is easily configurable.
"""
$(TYPEDEF)
Baytes Model struct.

Contains information about current Model value, name, and information, see also [`ParameterInfo`](@ref).

# Fields
$(TYPEDFIELDS)
"""
mutable struct ModelWrapper{
    M<:Union{P,ModelName} where {P},A<:NamedTuple,B<:ParameterInfo
} <: BaytesCore.AbstractModelWrapper
    "Current Model values as NamedTuple - works with Nested Tuples."
    val::A
    "Information about parameter distributions, transformations and constraints, see [`ParameterInfo`](@ref)."
    info::B
    "Model id, per default BaseModel. Useful for dispatching ModelWrapper struct."
    id::M
    function ModelWrapper(
        val::A, info::B, id::M
    ) where {M<:Union{P,ModelName} where {P},A<:NamedTuple,B<:ParameterInfo}
        return new{M,A,B}(val, info, id)
    end
end
# Convenient Constructor
function ModelWrapper(
    id::M, parameter::A, flattendefault::F=FlattenDefault()
) where {M<:Union{P,ModelName} where {P},A<:NamedTuple,F<:FlattenDefault}
    ## Check if all values in val are of type Param
    ArgCheck.@argcheck _checkparams(parameter) "All values in (nested) NamedTuple have to be of Type Param."
    ## Split between values and constraints
    val = _get_val(parameter)
    constraint = _get_constraint(parameter)
    ## Create ParameterInfo struct
    paraminfo = ParameterInfo(constraint, val, flattendefault)
    ## Return ModelWrapper
    return ModelWrapper(val, paraminfo, id)
end
ModelWrapper(parameter::A, flattendefault::F=FlattenDefault()) where {A<:NamedTuple,F<:FlattenDefault} =
    ModelWrapper(BaseModel(), parameter, flattendefault)

############################################################################################
# Basic functions for Model struct
length(model::ModelWrapper) = model.info.reconstruct.unflatten.strict._unflatten.sz[end]
paramnames(model::ModelWrapper) = keys(model.val)

############################################################################################
# A bunch of functions that can be used/extended for target model in Sampling process
"""
$(SIGNATURES)
Simulate data given Model parameter.

# Examples
```julia
```

"""
function simulate(model::ModelWrapper) end

############################################################################################
# Dispatch Model struct for .Core functions
"""
$(SIGNATURES)
Fill 'model' values with NamedTuple 'θ'.

# Examples
```julia
```

"""
function fill(model::ModelWrapper, θ::NamedTuple)
    return merge(model.val, θ)
end

"""
$(SIGNATURES)
Inplace version of [`fill`](@ref).

# Examples
```julia
```

"""
function fill!(model::ModelWrapper, θ::NamedTuple)
    model.val = merge(model.val, θ)
    return nothing
end

function subset(model::ModelWrapper, sym)
    return subset(model.val, sym)
end

#########################################
"""
$(SIGNATURES)
Unconstrain 'model' values and return as NamedTuple.

# Examples
```julia
```

"""
function unconstrain(model::ModelWrapper)
    return unconstrain(model.info.transform, model.val)
end

"""
$(SIGNATURES)
Flatten 'model' values and return as vector.

# Examples
```julia
```

"""
function flatten(model::ModelWrapper)
    return flatten(model.info.reconstruct, model.val)
end

"""
$(SIGNATURES)
Flatten and unconstrain 'model' values and return as vector.

# Examples
```julia
```

"""
function unconstrain_flatten(model::ModelWrapper)
    return flatten(model.info.reconstruct, unconstrain(model))
end

#########################################
"""
$(SIGNATURES)
Unlatten Vector 'θ' given constraints from 'model' and return as NamedTuple.

# Examples
```julia
```

"""
function unflatten(model::ModelWrapper, θ::AbstractVector{T}) where {T<:Real}
    return unflatten(model.info.reconstruct, θ)
end

"""
$(SIGNATURES)
Inplace version of [`unflatten`](@ref).

# Examples
```julia
```

"""
function unflatten!(model::ModelWrapper, θ::AbstractVector{T}) where {T<:Real}
    model.val = unflatten(model, θ)
    return nothing
end

"""
$(SIGNATURES)
Constrain and Unflatten vector 'θᵤ' given 'model' constraints.

# Examples
```julia
```

"""
function unflatten_constrain(model::ModelWrapper, θᵤ::AbstractVector{T}) where {T<:Real}
    return constrain(model.info.transform, unflatten(model, θᵤ))
end

"""
$(SIGNATURES)
Inplace version of [`unflatten_constrain`](@ref).

# Examples
```julia
```

"""
function unflatten_constrain!(model::ModelWrapper, θᵤ::AbstractVector{T}) where {T<:Real}
    model.val = unflatten_constrain(model, θᵤ)
    return nothing
end

#########################################
"""
$(SIGNATURES)
Sample from 'model' prior and return as NamedTuple.

# Examples
```julia
```

"""
function sample(_rng::Random.AbstractRNG, model::ModelWrapper)
    return sample_constraint(_rng, model.info.constraint, model.val)
end
sample(model::ModelWrapper) = sample(Random.GLOBAL_RNG, model)

"""
$(SIGNATURES)
Inplace version of [`sample`](@ref).

# Examples
```julia
```

"""
function sample!(_rng::Random.AbstractRNG, model::ModelWrapper)
    ArgCheck.@argcheck _checkprior(model.info.constraint) "For inplace sample version, all constraints need to be a Distribution."
    model.val = sample(_rng, model)
    return nothing
end
sample!(model::ModelWrapper) = sample!(Random.GLOBAL_RNG, model)

#########################################
"""
$(SIGNATURES)
Evaluate Log density of 'model' prior given current 'model' values.

# Examples
```julia
```

"""
function log_prior(model::ModelWrapper)
    return log_prior(model.info.constraint, model.val)
end

"""
$(SIGNATURES)
Evaluate Log density and eventual Jacobian adjustments of 'model' prior given current 'model' values.

# Examples
```julia
```

"""
function log_prior_with_transform(model::ModelWrapper)
    return log_prior_with_transform(model.info.constraint, model.val)
end

"""
$(SIGNATURES)
Evaluate eventual Jacobian adjustments from transformations at 'model' values.

# Examples
```julia
```

"""
function log_abs_det_jac(model::ModelWrapper)
    return log_abs_det_jac(model.info.transform.unconstrain, model.val)
end

############################################################################################
export
    ModelName,
    BaseModel,
    ModelWrapper,
    simulate,
    length,
    paramnames,
    fill,
    fill!,
    subset,
    flatten,
    unconstrain_flatten,
    unflatten,
    unflatten!,
    unconstrain,
    unflatten_constrain,
    unflatten_constrain!,
    sample,
    sample!,
    log_prior,
    log_prior_with_transform,
    log_abs_det_jac


text:
# Julia wrapper for header: /usr/local/include/libavutil/frame.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

export
    av_frame_alloc,
    av_frame_free,
    av_frame_ref,
    av_frame_clone,
    av_frame_unref,
    av_frame_move_ref,
    av_frame_get_buffer,
    av_frame_is_writable,
    av_frame_make_writable,
    av_frame_copy_props,
    av_frame_get_plane_buffer,
    av_frame_new_side_data,
    av_frame_get_side_data

function av_frame_alloc()
    ccall((:av_frame_alloc,libavutil),Ptr{AVFrame},())
end

function av_frame_free(frame)
    ccall((:av_frame_free,libavutil),Void,(Ptr{Ptr{AVFrame}},),frame)
end

function av_frame_ref(dst,src)
    ccall((:av_frame_ref,libavutil),Cint,(Ptr{AVFrame},Ptr{AVFrame}),dst,src)
end

function av_frame_clone(src)
    ccall((:av_frame_clone,libavutil),Ptr{AVFrame},(Ptr{AVFrame},),src)
end

function av_frame_unref(frame)
    ccall((:av_frame_unref,libavutil),Void,(Ptr{AVFrame},),frame)
end

function av_frame_move_ref(dst,src)
    ccall((:av_frame_move_ref,libavutil),Void,(Ptr{AVFrame},Ptr{AVFrame}),dst,src)
end

function av_frame_get_buffer(frame,align::Integer)
    ccall((:av_frame_get_buffer,libavutil),Cint,(Ptr{AVFrame},Cint),frame,align)
end

function av_frame_is_writable(frame)
    ccall((:av_frame_is_writable,libavutil),Cint,(Ptr{AVFrame},),frame)
end

function av_frame_make_writable(frame)
    ccall((:av_frame_make_writable,libavutil),Cint,(Ptr{AVFrame},),frame)
end

function av_frame_copy_props(dst,src)
    ccall((:av_frame_copy_props,libavutil),Cint,(Ptr{AVFrame},Ptr{AVFrame}),dst,src)
end

function av_frame_get_plane_buffer(frame,plane::Integer)
    ccall((:av_frame_get_plane_buffer,libavutil),Ptr{AVBufferRef},(Ptr{AVFrame},Cint),frame,plane)
end

function av_frame_new_side_data(frame,_type::AVFrameSideDataType,size::Integer)
    ccall((:av_frame_new_side_data,libavutil),Ptr{AVFrameSideData},(Ptr{AVFrame},AVFrameSideDataType,Cint),frame,_type,size)
end

function av_frame_get_side_data(frame,_type::AVFrameSideDataType)
    ccall((:av_frame_get_side_data,libavutil),Ptr{AVFrameSideData},(Ptr{AVFrame},AVFrameSideDataType),frame,_type)
end


text:
using StaticArrays
using ClimateMachine.VariableTemplates
using ClimateMachine.BalanceLaws:
    BalanceLaw,
    Prognostic,
    Auxiliary,
    Gradient,
    GradientFlux,
    GradientLaplacian,
    Hyperdiffusive
import ClimateMachine.BalanceLaws:
    vars_state,
    flux_first_order!,
    flux_second_order!,
    source!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    nodal_init_state_auxiliary!,
    init_state_prognostic!,
    boundary_conditions,
    boundary_state!,
    wavespeed,
    transform_post_gradient_laplacian!

using ClimateMachine.Mesh.Geometry: LocalGeometry
using ClimateMachine.DGMethods.NumericalFluxes:
    NumericalFluxFirstOrder, NumericalFluxSecondOrder

abstract type HyperDiffusionProblem end
struct HyperDiffusion{dim, P} <: BalanceLaw
    problem::P
    function HyperDiffusion{dim}(
        problem::P,
    ) where {dim, P <: HyperDiffusionProblem}
        new{dim, P}(problem)
    end
end

vars_state(::HyperDiffusion, ::Auxiliary, FT) = @vars(D::SMatrix{3, 3, FT, 9})
#
# Density is only state
vars_state(::HyperDiffusion, ::Prognostic, FT) = @vars(ρ::FT)

# Take the gradient of density
vars_state(::HyperDiffusion, ::Gradient, FT) = @vars(ρ::FT)
# Take the gradient of laplacian of density
vars_state(::HyperDiffusion, ::GradientLaplacian, FT) = @vars(ρ::FT)

vars_state(::HyperDiffusion, ::GradientFlux, FT) = @vars()
# The hyperdiffusion DG auxiliary variable: D ∇ Δρ
vars_state(::HyperDiffusion, ::Hyperdiffusive, FT) = @vars(σ::SVector{3, FT})

function flux_first_order!(m::HyperDiffusion, _...) end

"""
    flux_second_order!(m::HyperDiffusion, flux::Grad, state::Vars,
                     auxDG::Vars, auxHDG::Vars, aux::Vars, t::Real)

Computes diffusive flux `F` in:

```
∂ρ
-- = - ∇ • (σ) = - ∇ • F
∂t
```
Where

 - `σ` is hyperdiffusion DG auxiliary variable (`σ = D ∇ Δρ` with D being the hyperdiffusion tensor)
"""
function flux_second_order!(
    m::HyperDiffusion,
    flux::Grad,
    state::Vars,
    auxDG::Vars,
    auxHDG::Vars,
    aux::Vars,
    t::Real,
)
    σ = auxHDG.σ
    flux.ρ += σ
end

"""
    compute_gradient_argument!(m::HyperDiffusion, transform::Vars, state::Vars,
                   aux::Vars, t::Real)

Set the variable to take the gradient of (`ρ` in this case)
"""
function compute_gradient_argument!(
    m::HyperDiffusion,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform.ρ = state.ρ
end

compute_gradient_flux!(m::HyperDiffusion, _...) = nothing
function transform_post_gradient_laplacian!(
    m::HyperDiffusion,
    auxHDG::Vars,
    gradvars::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    ∇Δρ = gradvars.ρ
    D = aux.D
    auxHDG.σ = D * ∇Δρ
end

"""
    source!(m::HyperDiffusion, _...)

There is no source in the hyperdiffusion model
"""
source!(m::HyperDiffusion, _...) = nothing

function init_state_prognostic!(
    m::HyperDiffusion,
    state::Vars,
    aux::Vars,
    localgeo,
    t::Real,
)
    initial_condition!(m.problem, state, aux, localgeo, t)
end

boundary_conditions(::HyperDiffusion) = ()
boundary_state!(nf, ::HyperDiffusion, _...) = nothing


text:
function read_file(file, p, cache, particle_cache)
    h, blocks = open(file_summary, file)
    SDFFile(file, h, blocks, p, cache, particle_cache)
end

function Base.read(sdf::SDFFile, entry::Symbol)
    open(sdf.name) do f
        read(f, getindex(sdf.blocks, entry))
    end
end

function Base.read(sdf::SDFFile, entries...)
    open(sdf.name) do f
        asyncmap(i->read(f, getindex(sdf.blocks,i)), entries)
    end
end

function read_expensive(sdf, ids)
    open(sdf.name) do f
        asyncmap(i->make_grid(
            Variable(),
            sdf.blocks[i],
            nothing,
            f;
            cache=sdf.particle_cache[]
        ), ids)
    end
end

function Base.getindex(sdf::SDFFile, idx::Symbol)
    open(sdf.name) do f
        read_entry(f, sdf, idx)
    end
end

function expensive_grids(ids, idx)
    # Values corresponding to "grid/[species]") are expensive to read
    cid = [ids...]
    for (i, id) in enumerate(ids)
        if isnothing(id) && occursin("grid/", string(idx[i]))
            cid[i] = idx[i]
        end
    end
    s_id = string.(cid)
    expensive = map(i->occursin("grid/", i), s_id)
    !reduce(|, expensive) && return nothing, expensive
    @debug "Found expensive to read mesh entries"
    map(Symbol, unique(s_id[expensive])), expensive
end

# The data for the particles is the most expensive to read since there are
# a lot of particles. Since the ScalarVariables have a grid, that grid
# might be the same for more variables and should only be read once
function Base.getindex(sdf::SDFFile, idx::Vararg{Symbol, N}) where N
    mesh_ids = get_mesh_id.((sdf,), idx)
    @debug "Reading entries for $idx with mesh ids $mesh_ids"
    expensive_ids, is_expensive = expensive_grids(mesh_ids, idx)
    if isnothing(expensive_ids)
        simple_read(sdf, idx)
    else
        @debug "Expensive idxs: $is_expensive"
        @debug "Reading data without expensive grids"
        partial_data = selective_read(sdf, idx, expensive_ids)
        @debug "Reading expensive grids: $expensive_ids"
        grids = read_expensive(sdf, expensive_ids)
        grid_map = (; zip(expensive_ids, grids)...)
        complete_data = ()
        for (i, d) in enumerate(partial_data)
            @debug "Is $(idx[i]) expensive? $(is_expensive[i])"
            if is_expensive[i]
                id = mesh_ids[i]
                @debug "Expensive grid $id at $i"
                if isnothing(id)
                    grid_id = idx[i]
                    @debug "Storing grid entry $grid_id"
                    grid = getproperty(grid_map, grid_id)
                    complete_data = push!!(complete_data, grid)
                else
                    @debug "Storing $id"
                    grid = getproperty(grid_map, id)
                    data = d.data
                    data_block = d.block

                    f = store_entry(data_block, data, grid)
                    complete_data = push!!(complete_data, f)
                end
            else
                complete_data = push!!(complete_data, d)
            end
        end
        complete_data
    end
end

function selective_read(sdf, idx, expensive_ids)
    open(sdf.name) do f
        asyncmap(i->read_selected(f, sdf.blocks, i, expensive_ids), idx)
    end
end

function simple_read(sdf, idx)
    open(sdf.name) do f
        asyncmap(i->read_entry(f, sdf, i), idx)
    end
end


text:
# Test the importance sampler
# The test is based on running a reference implementation of a model with the same seed as the importance sampler.
# Both predictions and logevidence are compared for equality.

using Turing
using Distributions
using Test
using Random

function logsum(xs :: Vector{Float64})
  largest = maximum(xs)
  ys = map(x -> exp.(x - largest), xs)
  result = log(sum(ys)) + largest
  return result
end

function reference(n :: Int)
  logweights = zeros(Float64, n)
  samples = Array{Dict{Symbol,Any}}(undef, n)
  for i = 1:n
    samples[i] = reference()
    logweights[i] = samples[i][:logweight]
  end
  logevidence = logsum(logweights) - log(n)
  results = Dict{Symbol,Any}()
  results[:logevidence] = logevidence
  results[:logweights] = logweights
  results[:samples] = samples
  return results
end

function reference()
  x = rand(Normal(4,5))
  y = rand(Normal(x,1))
  log_lik = logpdf(Normal(x,2), 3) + logpdf(Normal(y,2), 1.5)
  d = Dict()
  d[:logweight] = log_lik
  d[:a] = x
  d[:b] = y
  return d
end

let n = 10
  @model normal() = begin
    a ~ Normal(4,5)
    3 ~ Normal(a,2)
    b ~ Normal(a,1)
    1.5 ~ Normal(b,2)
    a, b
  end

  alg = IS(n)
  seed = 0

  _f = normal();
  for i=1:100
    Random.seed!(seed)
    exact = reference(n)
    Random.seed!(seed)
    tested = sample(_f, alg)
    for i = 1:n
      @test exact[:samples][i][:a] == tested[:samples][i][:a]
      @test exact[:samples][i][:b] == tested[:samples][i][:b]
      @test exact[:logweights][i]  == tested[:logweights][i]
    end
    @test exact[:logevidence] == tested[:logevidence]
  end
end


text:
data_dir = withenv("DATADEPS_ALWAY_ACCEPT"=>"true") do
    datadep"TUDataset"
end

@testset "TUDataset - PROTEINS" begin
    data  = TUDataset("PROTEINS")

    @test data.num_nodes == 43471
    @test data.num_edges == 162088
    @test data.num_graphs == 1113

    @test length(data.source) == data.num_edges
    @test length(data.target) == data.num_edges

    @test size(data.node_attributes) == (1, data.num_nodes)
    @test data.edge_attributes === nothing
    @test data.graph_attributes === nothing

    @test size(data.node_labels) == (data.num_nodes,)
    @test data.edge_labels === nothing
    @test size(data.graph_labels) == (data.num_graphs,)

    @test length(data.graph_indicator) == data.num_nodes
    @test all(sort(unique(data.graph_indicator)) .== 1:data.num_graphs)

    @testset "indexing" begin
        d1, d2 = data[5], data[6]
        d12 = data[5:6]

        @test d12.num_nodes == d1.num_nodes + d2.num_nodes
        @test d12.num_edges == d1.num_edges + d2.num_edges
        @test d12.num_graphs == d1.num_graphs + d2.num_graphs == 2

        @test length(d12.source) == d12.num_edges
        @test length(d12.target) == d12.num_edges
        @test length(d1.source) == d1.num_edges
        @test length(d1.target) == d1.num_edges
        @test length(d2.source) == d2.num_edges
        @test length(d2.target) == d2.num_edges

        @test length(d12.graph_indicator) == d12.num_nodes
        @test length(d1.graph_indicator) == d1.num_nodes
        @test length(d2.graph_indicator) == d2.num_nodes

        @test all(d1.graph_indicator .== 1)
        @test all(d2.graph_indicator .== 1)
        @test sort(unique(d12.graph_indicator)) == [1,2]
    end
end

@testset "TUDataset - QM9" begin
    data  = TUDataset("QM9")

    @test data.num_nodes == 2333625
    @test data.num_edges == 4823498
    @test data.num_graphs == 129433

    @test length(data.source) == data.num_edges
    @test length(data.target) == data.num_edges

    @test size(data.node_attributes) == (16, data.num_nodes)
    @test size(data.edge_attributes) == (4, data.num_edges)
    @test size(data.graph_attributes) == (19, data.num_graphs)

    @test data.node_labels === nothing
    @test data.edge_labels === nothing
    @test data.graph_labels === nothing

    @test length(data.graph_indicator) == data.num_nodes
    @test all(sort(unique(data.graph_indicator)) .== 1:data.num_graphs)

    @testset "indexing" begin
        d1, d2 = data[5], data[6]
        d12 = data[5:6]

        @test d12.num_nodes == d1.num_nodes + d2.num_nodes
        @test d12.num_edges == d1.num_edges + d2.num_edges
        @test d12.num_graphs == d1.num_graphs + d2.num_graphs == 2

        @test length(d12.source) == d12.num_edges
        @test length(d12.target) == d12.num_edges
        @test length(d1.source) == d1.num_edges
        @test length(d1.target) == d1.num_edges
        @test length(d2.source) == d2.num_edges
        @test length(d2.target) == d2.num_edges

        @test length(d12.graph_indicator) == d12.num_nodes
        @test length(d1.graph_indicator) == d1.num_nodes
        @test length(d2.graph_indicator) == d2.num_nodes

        @test all(d1.graph_indicator .== 1)
        @test all(d2.graph_indicator .== 1)
        @test sort(unique(d12.graph_indicator)) == [1,2]

        @test size(d12.node_attributes) == (16, d12.num_nodes)
        @test size(d12.edge_attributes) == (4, d12.num_edges)
        @test size(d12.graph_attributes) == (19, d12.num_graphs)
        @test size(d1.node_attributes) == (16, d1.num_nodes)
        @test size(d1.edge_attributes) == (4, d1.num_edges)
        @test size(d1.graph_attributes) == (19, d1.num_graphs)
        @test size(d2.node_attributes) == (16, d2.num_nodes)
        @test size(d2.edge_attributes) == (4, d2.num_edges)
        @test size(d2.graph_attributes) == (19, d2.num_graphs)
    end
end


text:
using .NumericalFluxes: numerical_volume_flux_first_order!

# {{{ FIXME: remove this after we've figure out how to pass through to kernel
const _ξ1x1, _ξ2x1, _ξ3x1 = Grids._ξ1x1, Grids._ξ2x1, Grids._ξ3x1
const _ξ1x2, _ξ2x2, _ξ3x2 = Grids._ξ1x2, Grids._ξ2x2, Grids._ξ3x2
const _ξ1x3, _ξ2x3, _ξ3x3 = Grids._ξ1x3, Grids._ξ2x3, Grids._ξ3x3
const _M = Grids._M

const _n1, _n2, _n3 = Grids._n1, Grids._n2, Grids._n3
const _sM, _vMI = Grids._sM, Grids._vMI
# }}}

@doc """
    esdg_volume_tendency!(
        balance_law::BalanceLaw,
        ::Val{dim},
        ::Val{polyorder},
        tendency,
        state_prognostic,
        state_auxiliary,
        vgeo,
        D,
        α,
        β,
    ) where {dim, polyorder}

Computes the reference element horizontal tendency using a two-point flux
approximation of all derivatives.
"""
@kernel function esdg_volume_tendency!(
    balance_law::BalanceLaw,
    ::Val{dir},
    ::Val{info},
    volume_numerical_flux_first_order,
    tendency,
    state_prognostic,
    state_auxiliary,
    vgeo,
    D,
    α,
    β,
    add_source = false,
) where {dir, info}
    @uniform begin
        dim = info.dim

        FT = eltype(state_prognostic)
        num_state = number_states(balance_law, Prognostic(), FT)
        num_aux = number_states(balance_law, Auxiliary(), FT)

        local_H = MArray{Tuple{3, num_state}, FT}(undef)

        state_2 = MArray{Tuple{num_state}, FT}(undef)
        aux_2 = MArray{Tuple{num_aux}, FT}(undef)

        local_source = MArray{Tuple{num_state}, FT}(undef)

        @inbounds Nq1 = info.Nq[1]
        @inbounds Nq2 = info.Nq[2]
        Nq3 = info.Nqk

        if dir == 1
            _ξdx1, _ξdx2, _ξdx3 = _ξ1x1, _ξ1x2, _ξ1x3
            Nqd = Nq1
        elseif dir == 2
            _ξdx1, _ξdx2, _ξdx3 = _ξ2x1, _ξ2x2, _ξ2x3
            Nqd = Nq2
        elseif dir == 3
            _ξdx1, _ξdx2, _ξdx3 = _ξ3x1, _ξ3x2, _ξ3x3
            Nqd = Nq3
        end
    end

    e = @index(Group, Linear)
    i, j, k = @index(Local, NTuple)

    state_1 = @private FT (num_state,)
    aux_1 = @private FT (num_aux,)
    local_tendency = @private FT (num_state,)
    local_MI = @private FT (1,)
    shared_G = @localmem FT (Nq1 * Nq2 * Nq3, 3)
    shared_D = @localmem FT (Nqd, Nqd)

    @inbounds begin
        ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))

        # generalization for different polynomial orders
        @unroll for l in 1:Nqd
            if dir == 1
                id = i
                ild = l + Nq1 * ((j - 1) + Nq2 * (k - 1))
            elseif dir == 2
                id = j
                ild = i + Nq1 * ((l - 1) + Nq2 * (k - 1))
            elseif dir == 3
                id = k
                ild = i + Nq1 * ((j - 1) + Nq2 * (l - 1))
            end
            shared_D[id, l] = D[id, l]
        end

        M = vgeo[ijk, _M, e]
        shared_G[ijk, 1] = M * vgeo[ijk, _ξdx1, e]
        shared_G[ijk, 2] = M * vgeo[ijk, _ξdx2, e]
        if dim == 3
            shared_G[ijk, 3] = M * vgeo[ijk, _ξdx3, e]
        end

        # Build ode scaling into mass matrix (so it doesn't show up later)
        local_MI[1] = α / M

        # Get the volume tendency (scaling by β done below)
        @unroll for s in 1:num_state
            local_tendency[s] = β != 0 ? tendency[ijk, s, e] : -zero(FT)
        end

        # Scale β into the volume tendency
        @unroll for s in 1:num_state
            local_tendency[s] *= β
        end

        # Compute the volume tendency
        # ∑_{i,j}^{d} ( G_ij (Q_i ∘ H_j) - (H_j ∘ Q_i^T) G_ij)
        #  =
        # ( G_11 (Q_1 ∘ H_1) - (H_1 ∘ Q_1^T) G_11) 1 +
        # ( G_12 (Q_1 ∘ H_2) - (H_2 ∘ Q_1^T) G_12) 1 +
        # ( G_13 (Q_1 ∘ H_3) - (H_3 ∘ Q_1^T) G_13) 1 +
        # ( G_21 (Q_2 ∘ H_1) - (H_1 ∘ Q_2^T) G_21) 1 +
        # ( G_22 (Q_2 ∘ H_2) - (H_2 ∘ Q_2^T) G_22) 1 +
        # ( G_23 (Q_2 ∘ H_3) - (H_3 ∘ Q_2^T) G_23) 1 +
        # ( G_31 (Q_3 ∘ H_1) - (H_1 ∘ Q_3^T) G_31) 1 +
        # ( G_32 (Q_3 ∘ H_2) - (H_2 ∘ Q_3^T) G_32) 1 +
        # ( G_33 (Q_3 ∘ H_3) - (H_3 ∘ Q_3^T) G_33) 1
        @unroll for s in 1:num_state
            state_1[s] = state_prognostic[ijk, s, e]
        end
        @unroll for s in 1:num_aux
            aux_1[s] = state_auxiliary[ijk, s, e]
        end

        if add_source
            fill!(local_source, -zero(eltype(local_source)))
            source!(
                balance_law,
                Vars{vars_state(balance_law, Prognostic(), FT)}(local_source),
                Vars{vars_state(balance_law, Prognostic(), FT)}(state_1),
                Vars{vars_state(balance_law, Auxiliary(), FT)}(aux_1),
            )

            @unroll for s in 1:num_state
                local_tendency[s] += α * local_source[s]
            end
        end

        @synchronize

        ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))

        # Note: unrolling this loop makes things slower
        @views for l in 1:Nqd
            if dir == 1
                id = i
                ild = l + Nq1 * ((j - 1) + Nq2 * (k - 1))
            elseif dir == 2
                id = j
                ild = i + Nq1 * ((l - 1) + Nq2 * (k - 1))
            elseif dir == 3
                id = k
                ild = i + Nq1 * ((j - 1) + Nq2 * (l - 1))
            end

            # Compute derivatives wrt ξd
            # ( G_31 (Q_3 ∘ H_1) - (H_1 ∘ Q_3^T) G_31) 1 +
            # ( G_32 (Q_3 ∘ H_2) - (H_2 ∘ Q_3^T) G_32) 1 +
            # ( G_33 (Q_3 ∘ H_3) - (H_3 ∘ Q_3^T) G_33) 1 +
            @unroll for s in 1:num_state
                state_2[s] = state_prognostic[ild, s, e]
            end
            @unroll for s in 1:num_aux
                aux_2[s] = state_auxiliary[ild, s, e]
            end
            fill!(local_H, -zero(FT))
            numerical_volume_flux_first_order!(
                volume_numerical_flux_first_order,
                balance_law,
                local_H,
                state_1[:],
                aux_1[:],
                state_2,
                aux_2,
            )
            @unroll for s in 1:num_state
                # G_21 (Q_2 ∘ H_1) 1 +
                # G_22 (Q_2 ∘ H_2) 1 +
                # G_23 (Q_2 ∘ H_3) 1
                local_tendency[s] -=
                    local_MI[1] *
                    shared_D[id, l] *
                    (
                        shared_G[ijk, 1] * local_H[1, s] +
                        shared_G[ijk, 2] * local_H[2, s] +
                        (
                            dim == 3 ? shared_G[ijk, 3] * local_H[3, s] :
                            -zero(FT)
                        )
                    )
                #  (H_1 ∘ Q_2^T) G_21 1 +
                #  (H_2 ∘ Q_2^T) G_22 1 +
                #  (H_3 ∘ Q_2^T) G_23 1
                local_tendency[s] +=
                    local_MI[1] *
                    (
                        local_H[1, s] * shared_G[ild, 1] +
                        local_H[2, s] * shared_G[ild, 2] +
                        (
                            dim == 3 ? local_H[3, s] * shared_G[ild, 3] :
                            -zero(FT)
                        )
                    ) *
                    shared_D[l, id]
            end
        end

        @unroll for s in 1:num_state
            tendency[ijk, s, e] = local_tendency[s]
        end
    end
end


text:
Random.rand(r::JaxRNGKey, dims::Integer...) = rand(r, Float64, Dims(dims))
Random.rand(r::JaxRNGKey, T::Type, dims::Dims)  = uniform(T, r, dims, minval=zero(T), maxval=one(T))
Random.rand(r::JaxRNGKey, T::Type{<:Integer}, dims::Dims)  = uniform(T, r, dims, minval=typemin(T), maxval=typemax(T))

Random.randn(r::JaxRNGKey, dims::Integer...) = rand(r, Float64, Dims(dims))
Random.randn(r::JaxRNGKey, T::Type, dims::Dims)  = normal(T, r, dims)

#=
Random.rand(r::JaxRNGKey, X, d::Integer, dims::Integer...) = rand(r, X, Dims((d, dims...)))

Random.rand(r::JaxRNGKey, ::Type{X}, dims::Dims) where {X} = rand!(r, Array{X}(undef, dims), X)

Random.rand(r::JaxRNGKey, ::Type{X}, d::Integer, dims::Integer...) where {X} = rand(r, X, Dims((d, dims...)))
=#


text:
using LazyInverses
using LinearAlgebra
using BenchmarkTools

n = 1024
x = randn(n)
A = randn(n, n)
@time B = inv(A)
@time C = inverse(A)
println(A \ x ≈ C * x)
println(A * x ≈ C \ x)


text:
# using Plots
# using StatPlots
# using DataFrames, RDatasets

gr()

function linescatter()
    x=1:4
    y=[10, 15, 13, 17]
    plot(x, y)
end
linescatter()

function multiple_scatter_traces()
    x=1:4
    y1 = [10, 15, 13, 17]
    y2 = [16, 5, 11, 9]
    y3 = [12, 9, 15, 12]
    y4 = [5, 10, 8, 12]
    plot(y1, line = (:scatter, 1), label="marker")
    plot!(y2, marker =(:circle, 4), label="line+marker")
    plot!(y3, label="line")
    plot!(y4, line = (:dash, 2), marker =(:circle, 4), label="dash")
end
multiple_scatter_traces()

function multiple_scatter_traces_v2()
    x=1:4
    y1 = [10, 15, 13, 17]
    y2 = [16, 5, 11, 9]
    y3 = [12, 9, 15, 12]
    y4 = [5, 10, 8, 12]
    y = [y1 y2 y3 y4]
    plot(y, line=(2, [:solid :solid :dash :dash]), marker=(10, [:diamond :circle :circle :circle]))
end
multiple_scatter_traces_v2()

function data_labels()
    x = 1:5
    y1 = [1, 6, 3, 6, 1]
    y2 = [4, 1, 7, 1, 4]

    plot(y1, line = (:scatter, 1),
         series_annotations=[text("A-1", :bottom, 12, :blue, "Raleway, sans-serif"),
                             "A-2","A-3","A-4","A-5"],
         label="Team A",
         title="Data Labels on the plot")
    plot!(y2, line = (:scatter, 1),
         series_annotations=[text("B-1", :top, 12, :red, "Times"),
                                  "B-2","B-3","B-4","B-5"],
         label="Team B")
end
data_labels()

# pyplot() -- to fill to one
function area1()
    y1=[0, 2, 3, 5] #, fill="tozeroy")
    y2=[3, 5, 1, 7] #, fill="tonexty")
    plot(y1, label="lines", w=3, fill=0, α=1)
    plot!(y2, label="lines", w=3, fill=0, α=0.6)
end
area1()

gr()
function plot_matrices()
    z = rand(10,10)
    plot(spy(z), heatmap(z), contour(z), surface(z))
end
plot_matrices()

plotlyjs()
function advanced_layouts()
    l = @layout [  a{0.3w} [grid(3,3)
                             b{0.2h} ]]
    plot(
        rand(10,11),
        layout = l, legend = false, seriestype = [:bar :scatter :path],
        title = reshape(["($i)" for i=1:11], 1, 11), titleloc = :center, titlefont = font(8)
    )
end
advanced_layouts()

plotlyjs()
function grouped_bar_example()
    x=["giraffes", "orangutans", "monkeys"]
    y1=[20, 14, 23]
    y2=[12, 18, 29]

    bar(x, y1, label= "SF Zoo")
    bar!(x, y2, label= "LA Zoo")
end
grouped_bar_example()

pyplot()
function stacked_bar_example()
    x=["giraffes", "orangutans", "monkeys"]
    y1=[20, 14, 23]
    y2=[12, 18, 29]

    bar(x, y1, bar_position=:stack, label= "SF Zoo")
    bar!(x, y2, bar_position=:stack, label= "LA Zoo")
end
stacked_bar_example()

function two_hists()
    x0 = randn(500)
    x1 = x0+1

    histogram(x0, opacity=0.75)
    histogram!(x1, opacity=0.75)
end

two_hists()

function box_plot()
    x1 = ["day 1", "day 1", "day 1", "day 1", "day 1", "day 1",
          "day 2", "day 2", "day 2", "day 2", "day 2", "day 2"]
    x2 = ["day 1B", "day 1B", "day 1B", "day 1B", "day 1B", "day 1B",
          "day 2B", "day 2B", "day 2B", "day 2B", "day 2B", "day 2B"]
    x3 = ["day 1C", "day 1C", "day 1C", "day 1C", "day 1C", "day 1C",
      "day 2C", "day 2C", "day 2C", "day 2C", "day 2C", "day 2C"]

    y1=[0.2, 0.2, 0.6, 1.0, 0.5, 0.4, 0.2, 0.7, 0.9, 0.1, 0.5, 0.3]
    y2=[0.6, 0.7, 0.3, 0.6, 0.0, 0.5, 0.7, 0.9, 0.5, 0.8, 0.7, 0.2]
    y3=[0.1, 0.3, 0.1, 0.9, 0.6, 0.6, 0.9, 1.0, 0.3, 0.6, 0.8, 0.5]

    boxplot(x1, y1, label="kale", color="#3D9970")
    boxplot!(x2, y2, label="radishes", color="#FF4136")
    boxplot!(x3, y3, label="carrot", color="#FF851B")
#
end
box_plot()

function dataframe_scatter()
    gr()
    iris = dataset("datasets", "iris")#
     # display(head(iris))
    @df iris scatter(:SepalLength, :SepalWidth, group=:Species,
        title = "Iris Sepal length vs width",
        xlabel = "Length", ylabel = "Width",
        m=(0.5, [:cross :hex :star7], 12))
end

dataframe_scatter()


text:
pushfirst!(LOAD_PATH, joinpath(@__DIR__, ".."))

using Documenter, GCMMesh

makedocs(
    sitename = "GCMMesh.jl",
    doctest = false,
    strict = false,
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true",
        mathengine = MathJax(
            Dict(
                :TeX =>
                    Dict(:equationNumbers => Dict(:autoNumber => "AMS"), :Macros => Dict()),
            ),
        ),
    ),
    clean = true,
    modules = GCMMesh,
    pages = Any[
        "Home"=>"index.md",
        "Graphs"=>"Graphs.md",
        "Mesh"=>"Mesh.md",
        "Partition"=>"Partition.md",
    ],
)

deploydocs(
    repo = "github.com/CliMA/GCMMesh.git",
    target = "build",
    push_preview = true,
    devbranch = "main",
    forcepush = true,
)

max_stars_repo_name: CliMA/GCMMesh

max_issues_repo_name: CliMA/GCMMesh

max_forks_repo_name: CliMA/GCMMesh


text:
"""
        load_timeseries_data_provided(region::String="GER_1"; T::Int=24, years::Array{Int,1}=[2016], att::Array{String,1}=Array{String,1}())
- Adding the information in the `*.csv` file at `data_path` to the data dictionary
The `*.csv` files shall have the following structure and must have the same length:

|`Timestamp`  |`Year`  |[column names...] |
|-------------|--------|------------------|
|[iterator]   |[year]  |[values]          |

The first column should be called `Timestamp` if it contains a time iterator
The other columns can specify the single timeseries like specific geolocation.
for regions:
- `"GER_1"`: Germany 1 node
- `"GER_18"`: Germany 18 nodes
- `"CA_1"`: California 1 node
- `"CA_14"`: California 14 nodes
- `"TX_1"`: Texas 1 node
"""
function load_timeseries_data_provided(region::String="GER_1";
                              T::Int=24,
                              years::Array{Int,1}=[2016],
                              att::Array{String,1}=Array{String,1}())
    # Check for existance of the region in data
    region in readdir(normpath(joinpath(@__DIR__,"..","..","data"))) || throw(@error "The region $region is not found. The provided regions are: GER_1: Germany 1 node, GER_18: Germany 18 nodes, CA_1: California 1 node, CA_14: California 14 nodes, TX_1: Texas 1 node")
    # Generate the data path based on application and region
    data_path=normpath(joinpath(@__DIR__,"..","..","data",region,"TS"))
    return load_timeseries_data(data_path; region=region, T=T, years=years, att=att)
end

"""
    load_cep_data_techs(data_path::String)
load the `techs.csv` in the folder `data_path` with the following columns:
- `tech`: defines the sets `tech`
- `categ`: the category of this technology (is it storage, transmission or generation)
- `sector`: sector of the technology (electricity or heat)
- `eff`: efficiency of this technologies conversion [-]
- `time_series`: time_series name for availability
- `lifetime`: product lifetime [a]
- `financial_lifetime`: financial time to break even [a]
- `discount_rate`: discount rate for technology [a]
- `annuityfactor`: annuity factor, important for cap-costs [-]
returns `techs::OptVariable`    techs[tech] - OptDataCEPTech
"""
function load_cep_data_techs(data_path::String)
    tab=CSV.read(joinpath(data_path,"techs.csv");strict=true)
    #Check existance of necessary column
    :tech in names(tab) || throw(@error "No column called `tech` in `techs.csv`")
    #Build empty OptVariable
    techs=OptVariable{OptDataCEPTech}(undef, unique(tab[:tech]); type="fv", axes_names=["tech"])
    # loop through all axes
    for tech in axes(techs,"tech")
        #name
        name=tech
        #categ
        categ=tab[(:tech,tech),:categ][1]
        #sector::String
        sector=tab[(:tech,tech),:sector][1]
        #eff::Number
        eff=tab[(:tech,tech),:eff][1]
        #time_series::String
        time_series=tab[(:tech,tech),:time_series][1]
        #lifetime::Number
        lifetime=tab[(:tech,tech),:lifetime][1]
        #financial_lifetime::Number
        financial_lifetime=tab[(:tech,tech),:financial_lifetime][1]
        #discount_rate::Number
        discount_rate=tab[(:tech,tech),:discount_rate][1]
        # The time for the cap-investion to be paid back is the minimum of the max. financial lifetime and the lifetime of the product (If it's just good for 5 years, you'll have to rebuy one after 5 years)
        # annuityfactor = (1+i)^y*i/((1+i)^y-1) , i-discount_rate and y-payoff years
        annuityfactor=round((1+discount_rate)^(min(financial_lifetime,lifetime)) *discount_rate/ ((1+discount_rate) ^(min(financial_lifetime,lifetime))-1); sigdigits=9)
        # Add single data entry
        techs[tech]=OptDataCEPTech(name,categ,sector,eff,time_series,lifetime,financial_lifetime,discount_rate,annuityfactor)
    end
    return techs
end

"""
    load_cep_data_nodes(data_path::String, techs::OptVariable)
load the `nodes.csv` in the folder `data_path` with the following columns:
- `nodes` defines the set `nodes`
- `infrastruct` to indicate that this row either constains `ex` `power_ex` existing capacity [MW or MWh] or `lim` `power_lim` capacity limit [MW or MWh]
- `region`
- `lat`
- `lon` hold geolocation information in [°,°]
- each `tech` and the actual capacity [MW or MWh]
- ...
returns `nodes::OptVariable`    nodes[tech, node] - OptDataCEPNode
"""
function load_cep_data_nodes(data_path::String,
                             techs::OptVariable)
    tab=CSV.read(joinpath(data_path,"nodes.csv");strict=true)
    # Check exisistance of columns
    check_column(tab,[:node, :infrastruct])
    #Create empty OptVariable
    nodes=OptVariable{OptDataCEPNode}(undef, axes(techs,"tech"), unique(tab[:node]); type="fv", axes_names=["tech", "node"])
    for tech in axes(nodes,"tech")
        for node in axes(nodes,"node")
            #name
            name=node
            #value
            power_ex=tab[(:node,node)][(:infrastruct,"ex"),Symbol(tech)][1]
            data=tab[(:node,node)][(:infrastruct,"lim"),Symbol(tech)]
            power_lim= isempty(data) ? Inf : data[1]
            #region
            region=tab[(:node,node),:region][1]
            #lat and lon
            latlon=LatLon(tab[(:node,node),:lat][1],tab[(:node,node),:lon][1])
            nodes[tech,node]=OptDataCEPNode(name,power_ex, power_lim, region, latlon)
        end
    end
    return nodes
end

"""
    load_cep_data_lines(data_path::String, techs::OptVariable)
load the `lines.csv` in the folder `data_path` with the following columns:
- `tech` for each of the row
- `line` defines the set `lines`
- `node_start` Node where line starts
- `node_end` Node where line ends
- `reactance`
- `resistance` [Ω]
- `power_ex`: existing power limit [MW]
- `power_lim`: limit power limit [MW]
- `circuits` [-]
- `voltage` [V]
- `length` [km]
- `eff` [-]
returns `lines::OptVarible`     lines[tech, line] - OptDataCEPLine
"""
function load_cep_data_lines(data_path::String,
                             techs::OptVariable)
    if isfile(joinpath(data_path,"lines.csv"))
        tab=CSV.read(joinpath(data_path,"lines.csv");strict=true)
        #Check existance of necessary column
        check_column(tab, [:line])

        #Create empty OptVariable
        lines=OptVariable{OptDataCEPLine}(undef, unique(tab[:tech]), unique(tab[:line]); type="fv", axes_names=["tech", "line"])
        for tech in axes(lines,"tech")
            for line in axes(lines,"line")
                #name
                name=line
                #node_start
                node_start=tab[(:tech,tech)][(:line,line),:node_start][1]
                #node_end
                node_end=tab[(:tech,tech)][(:line,line),:node_end][1]
                #reactance
                reactance=tab[(:tech,tech)][(:line,line),:reactance][1]
                #resistance
                resistance=tab[(:tech,tech)][(:line,line),:resistance][1]
                #power
                power_ex=tab[(:tech,tech)][(:line,line),:power_ex][1]
                #power
                power_lim=tab[(:tech,tech)][(:line,line),:power_lim][1]
                #circuits
                circuits=tab[(:tech,tech)][(:line,line),:circuits][1]
                #voltage
                voltage=tab[(:tech,tech)][(:line,line),:voltage][1]
                #length
                length=tab[(:tech,tech)][(:line,line),:length][1]
                #eff calculate the efficiency provided as eff/km in techs
                #η=1-l_{line}⋅(1-η_{tech}) [-]
                eff=1-length*(1-techs[tech].eff)
                lines[tech,line]=OptDataCEPLine(name,node_start,node_end,reactance,resistance,power_ex,power_lim,circuits,voltage,length,eff)
            end
        end
        return lines
    else
        return lines=OptVariable{OptDataCEPLine}(undef, Array{String,1}(), Array{String,1}(); type="fv", axes_names=["tech", "line"])
    end
end

"""
    get_region_data(nodes::OptVariable,tab::DataFrame,tech::String,node::String,account::String)
Return the name of the region `region` or `"all"` that data is provided for in the `tab`
"""
function get_location_data(nodes::OptVariable,
                            tab::DataFrame,
                            tech::String,
                            node::String,
                            account::String)
    #determine region for this technology and node based on infromation in nodes
    region=nodes[tech,node].region
    #determine regions provided for this tech and this account in the data
    locations_data=unique(tab[(:tech,tech)][(:account,account),:location])
    #check if either specific `node`, `region` or a value for `all` regions is given
    if node in locations_data
        return node
    elseif region in locations_data
        return region
    elseif "all" in locations_data
        return "all"
    else
        return @error "region $region not provided in $(repr(tab))"
    end
end

"""
        load_cep_data_costs(data_path::String,techs::OptVariable, nodes::OptVariable)
load the `costs.csv` in the folder `data_path` with the following columns:
- `tech`
- `location` - in which location is this value valid? Either `"all"`, specific region, or node
-  `year` - in which year is this value valid?
- `account` - `cap` - total capacity costs for the entire lifetime, `fix` - yearly fixed (maintanance) costs, `var` - variable costs per MW
- `EUR` or `USD` or your preffered currency - the monetary cost
- `CO2` - the environmental cost as an impact category
- ... other impact categories
returns `costs::OptVariable`    costs[tech,node,year,account,impact] - Number
"""
function load_cep_data_costs(data_path::String,
                            techs::OptVariable,
                            nodes::OptVariable)
    tab=CSV.read(joinpath(data_path,"costs.csv");strict=true)
    check_column(tab,[:tech, :location, :year, :account])
    impacts=String.(names(tab)[findfirst(names(tab).==Symbol("|"))+1:end])
    #Create empty OptVariable
    costs=OptVariable{Number}(undef, axes(techs,"tech"), axes(nodes,"node"), unique(tab[:year]), ["cap_fix", "var"], impacts; type="fv", axes_names=["tech", "node", "year", "account", "impact"])
    for tech in axes(costs,"tech")
        for node in axes(costs,"node")
            for year in axes(costs,"year")
                for impact in axes(costs, "impact")
                    #Addition of capacity costs and fix maintanance cost - For numerical benefit in solving
                    account="cap_fix"
                        cap_location=get_location_data(nodes,tab,tech,node,"cap")
                        total_cap_cost=tab[(:tech,tech)][(:location,cap_location)][(:account,"cap")][(:year,year),Symbol(impact)][1]
                        #First impact shall always be currency - Currency of capacity cost is annulized with annuityfactor
                        if impact==axes(costs,"impact")[1]
                            annulized_cap_cost=round(total_cap_cost*techs[tech].annuityfactor;sigdigits=9)
                        else #Emissions of capacity cost are annulized with total lifetime
                            annulized_cap_cost=round(total_cap_cost/techs[tech].lifetime;sigdigits=9)
                        end
                        fix_location=get_location_data(nodes,tab,tech,node,"fix")
                        fix_cost=tab[(:tech,tech)][(:location,fix_location)][(:account,"fix")][(:year,year),Symbol(impact)][1]
                        costs[tech,node,year,account,impact]=annulized_cap_cost+fix_cost
                    #Variable cost is seperate
                    account="var"
                        var_location=get_location_data(nodes,tab,tech,node,account)
                        var_cost=tab[(:tech,tech)][(:location,var_location)][(:account,account)][(:year,year),Symbol(impact)][1]
                        costs[tech,node,year,account,impact]=var_cost
                end
            end
        end
    end
    return costs
end

"""
    load_cep_data_provided(region::String)
Loading from .csv files in a the folder `../CEP/data/{region}/`
Follow instructions preparing your own data:
-`region::String`: name of state or region data belongs to
-`costs::OptVariable`: `costs[tech,node,year,account,impact] - annulized costs [USD in USD/MW_el, CO2 in kg-CO₂-eq./MW_el]`
-`techs::OptVariable`: `techs[tech] - OptDataCEPTech`
-`nodes::OptVariable`: `nodes[tech,node] - OptDataCEPNode`
-`lines::OptVarible`: `lines[tech,line] - OptDataCEPLine`
for regions:
- `"GER_1"`: Germany 1 node
- `"GER_18"`: Germany 18 nodes
- `"CA_1"`: California 1 node
- `"CA_14"`: California 14 nodes
- `"TX_1"`: Texas 1 node
"""
function load_cep_data_provided(region::String)
  data_path=normpath(joinpath(dirname(@__FILE__),"..","..","data",region))
  return load_cep_data(data_path;region=region)
end

"""
    load_cep_data(data_path::String)
Loading from .csv files in a the folder `/data_path/`
Follow instructions for the CSV-Files:
-`region::String`: name of state or region data belongs to
-`costs::OptVariable`: costs[tech,node,year,account,impact] - annulized costs [USD in USD/MW_el, CO2 in kg-CO₂-eq./MW_el]`
-`techs::OptVariable`: techs[tech] - OptDataCEPTech -
-`nodes::OptVariable`: nodes[tech,node] - OptDataCEPNode
-`lines::OptVarible`: lines[tech,line] - OptDataCEPLine
"""
function load_cep_data(data_path::String;region::String="none")
  techs=load_cep_data_techs(data_path)
  nodes=load_cep_data_nodes(data_path, techs)
  lines=load_cep_data_lines(data_path, techs)
  costs=load_cep_data_costs(data_path, techs, nodes)
  return OptDataCEP(region,costs,techs,nodes,lines)
end #load_pricedata

#= Interpolation
"""
    get_number_interpolation(numbers_data::Array{Number,1}, number::Number)
find the neighboring values to do an interpolation of `number` in `numbers_data`
if `number` has no higher or lower neighbor, return the closest neighbor twice
"""
function get_number_interpolation(numbers_data::Array, number::Number)
    #Find numbers being greater and lower than the current numb
    numbers_g=numbers_data[numbers_data.>=number]
    numbers_l=numbers_data[numbers_data.<=number]
    #Create numbers for interpolation
    number_int=Tuple()
    number_int[1]= isempty(numbers_l) ? minimum(numbers_g) : maximum(numbers_l)
    number_int[2]= isempty(numbers_g) ? maximum(numbers_l) : minimum(numbers_g)
    return number_int
end

function get_interpolation(tab::DataFrame,col_val_ind::Tuple{Symbol,Number},colon_ind::Symbol)
    col_int=col_val_ind[1]
    val_int=col_val_ind[2]
    neighbors_int=get_number_interpolation(tab[col_int],val_int)
    if neighbors_int[1]==neighbors_int[2]
        return tab[(col_int,neighbors_int[1]),colon_ind][1]
    else
        #interpolation
        res_1=tab[(col_int,neighbors_int[1]),colon_ind]
        res_2=tab[(col_int,neighbors_int[2]),colon_ind]
        return res_1+(res_2-res_1)*(val_int-neighbors_int[1])/(neighbors_int[2]-neighbors_int[1])
    end
end
=#


text:
@testset "gradient" begin

    vars = @variables x, y

    @testset "Single variable" begin
        f( (x, ) ) = x + 1
        g = ReversePropagation.gradient( f, (x,))
        @test g(3) == (4, (1,))

        f2( (x, ) ) = x^2 + x
        g2 = ReversePropagation.gradient( f2, (x,) )
        @test g2(3) == (12, (7,))
    end

    @testset "Two variables" begin
        h( (a, b) ) = (a + b) * (a + b)
        vars = @variables x, y
        ∇h = ReversePropagation.gradient( h, (x, y) )
        @test ∇h( (1, 2) ) == (9, (6, 6))

        h2( (a, b) ) = a^2 + b^2
        vars = @variables x, y
        ∇h2 = ReversePropagation.gradient( h2, (x, y) )
        @test ∇h2( (1, 2) ) == (5, (2, 4))
    end

end


text:
using Flux, CUDA
using Plots
using OrdinaryDiffEq
using BenchmarkTools

function truth_ODE!(du, u, p, t)
    A = cu([-0.1 2.; -2. -0.1])
    du .= A * u
end

u₀ = [2., 0.] |> gpu
trange = 0:0.1:1
tspan = (trange[1], trange[end])

truth_prob = ODEProblem(truth_ODE!, u₀, tspan, saveat=trange)
truth_timeseries = solve(truth_prob, Tsit5())

@btime solve(truth_prob, Tsit5())

function NDE(u, p, t)
    A =
end

m = Dense(10,5) |> gpu

x = rand(10) |> gpu

using CUDA
using BenchmarkTools

W = cu(rand(2, 5)) # a 2×5 CuArray
b = cu(rand(2))

predict(x) = W*x .+ b
loss(x, y) = sum((predict(x) .- y).^2)

x, y = cu(rand(5)), cu(rand(2)) # Dummy data
@btime loss(cu(rand(5)), cu(rand(2))) # ~ 3

using JLD2

file = jldopen("D:\\University Matters\\Massachusetts Institute of Technology\\CLiMA Project\\OceanParameterizations.jl\\training_output\\parameter_optimisation_modified_pacanowski_philander.jld2", "r")


text:
function solve!(solver::EulerMaruyama!, Y::VSamplePath, u::T, W, P::ProcessOrCoefficients) where {T}
    N = length(W)
    N != length(Y) && error("Y and W differ in length.")

    tt = Y.tt
    tt[:] = W.tt
    yy = Y.yy
    y::T = copy(u)

    size(Y.yy) != (length(y), N) && error("Starting point has wrong length.")
    #assert(size(W.yy) == (length(y), N))
    tmp1 = copy(y)
    tmp2 = copy(y)
    dw = W.yy[.., 1]
    dw2 = copy(dw)
    for i in 1:N-1
        t¯ = tt[i]
        dt = tt[i+1] - t¯
        for k in eachindex(tmp1)
            @inbounds yy[k, i] = y[k]
        end
        for k in eachindex(dw)
            @inbounds dw[k] = W.yy[k, i+1] - W.yy[k, i]
        end
        _b!((i,t¯), y, tmp1, P)
        σ!(t¯, y, dw, tmp2, P)
        for k in eachindex(y)
            @inbounds y[k] = y[k] + tmp1[k]*dt + tmp2[k]
        end
    end
    yy[.., N] = endpoint(y, P)
    Y
end

function Bridge.solve!(solver::EulerMaruyama!, Y::SamplePath, u, W::SamplePath, P::Bridge.ProcessOrCoefficients)
    N = length(W)
    N != length(Y) && error("Y and W differ in length.")

    tt = Y.tt
    tt[:] = W.tt
    yy = Y.yy
    y = copy(u)

    tmp1 = copy(y)
    tmp2 = copy(y)
    dw = copy(W.yy[1])
    for i in 1:N-1
        t¯ = tt[i]
        dt = tt[i+1] - t¯
        copyto!(yy[i], y)
        if dw isa Number
            dw = W.yy[i+1] - W.yy[i]
        else
            for k in eachindex(dw)
                dw[k] = W.yy[i+1][k] - W.yy[i][k]
            end
        end

        _b!((i,t¯), y, tmp1, P)
        σ!(t¯, y, dw, tmp2, P)

        for k in eachindex(y)
            y[k] = y[k] + tmp1[k]*dt + tmp2[k]
        end
    end
    copyto!(yy[end], endpoint(y, P))
    Y
end


text:
"""
    AlgRemez.jl is a wrapper for AlgRemez written in c++.
    Please see
    https://github.com/maddyscientist/AlgRemez

"""
module AlgRemez
    using AlgRemez_jll
# Write your package code here.
    const exe = algremez(x -> x)

    struct AlgRemez_coeffs
        α0::Float64
        α::Array{Float64,1}
        β::Array{Float64,1}
        n::Int64
    end

    function fittedfunction(coeff::AlgRemez_coeffs)
        function func(x)
            value = coeff.α0
            for i=1:coeff.n
                value += coeff.α[i]/(x + coeff.β[i])
            end
            return value
        end
        x -> func(x)
    end

    """
Output: coeff_plus::AlgRemez_coeffs,coeff_minus::AlgRemez_coeffs
struct AlgRemez_coeffs
    α0::Float64
    α::Array{Float64,1}
    β::Array{Float64,1}
    n::Int64
end
f(x) = x^(y/z) = coeff_plus.α0 + sum_i^n coeff_plus.α[i]/(x + coeff_plus.β[i])
f(x) = x^(-y/z) = coeff_minus.α0 + sum_i^n coeff_minus.α[i]/(x + coeff_minus.β[i])

The function to be approximated f(x) = x^(y/z) and f(x) = x^(-y/z),
with degree (n,n) over the spectral range [lambda_low,lambda_high],
using precision digits of precision in the arithmetic.
The parameters y and z must be positive, the approximation to f(x) = x^(-y/z) is simply
the inverse of the approximation to f(x) = x^(y/z).
The default value of precision is 42.
    """
    function calc_coefficients(y,z,n,lambda_low,lambda_high;precision=42)
        run(`$exe $y $z $n $n $lambda_low $lambda_high $precision`)
        datas = readlines("approx.dat")
        icount = 3
        αplus0 = parse(Float64,split(datas[icount],"=")[2])
        #println(αplus0)
        αplus = zeros(Float64,n)
        βplus = zeros(Float64,n)
        for i=1:n
            icount += 1
            u = split(datas[icount],",")
            αplus[i] = parse(Float64,split(u[1],"=")[2])
            βplus[i] = parse(Float64,split(u[2],"=")[2])
        end
        coeff_plus = AlgRemez_coeffs(αplus0,αplus,βplus,n)
        #println(αplus)
        #println(βplus)

        icount += 4
        αminus0 = parse(Float64,split(datas[icount],"=")[2])
        #println(αminus0)
        αminus = zeros(Float64,n)
        βminus = zeros(Float64,n)
        for i=1:n
            icount += 1
            u = split(datas[icount],",")
            αminus[i] = parse(Float64,split(u[1],"=")[2])
            βminus[i] = parse(Float64,split(u[2],"=")[2])
        end
        coeff_minus = AlgRemez_coeffs(αminus0,αminus,βminus,n)
        #println(αminus)
        #println(βminus)
        #for data in datas
        #    println(data)
        #end

        rm("approx.dat")
        rm("error.dat")

        return coeff_plus,coeff_minus
    end

end


text:
#==
Example taken from
Indraneel Das and John E Dennis.
Normal-boundary intersection: A new method for generating the pareto surface in
nonlinear multicriteria optimization problems.
SIAM Journal on Optimization, 8(3):631–657, 1998.
==#

using MultiJuMP, JuMP
using Ipopt

m = multi_model(Ipopt.Optimizer)
@variable(m, x[i=1:5])
@NLexpression(m, f1, sum(x[i]^2 for i=1:5))
@NLexpression(m, f2, 3x[1]+2x[2]-x[3]/3+0.01*(x[4]-x[5])^3)
@NLconstraint(m, x[1]+2x[2]-x[3]-0.5x[4]+x[5]==2)
@NLconstraint(m, 4x[1]-2x[2]+0.8x[3]+0.6x[4]+0.5x[5]^2 == 0)
@NLconstraint(m, sum(x[i]^2 for i=1:5) <= 10)

obj1 = SingleObjective(f1)
obj2 = SingleObjective(f2)

md = get_multidata(m)
md.objectives = [obj1, obj2]
md.pointsperdim = 20
optimize!(m, method = NBI())

using Plots
plot(md)


text:
export RewardTransformedEnv

"""
    RewardTransformedEnv(env, f)

Apply `f` on `reward(env)`.
"""
struct RewardTransformedEnv{F,E<:AbstractEnv} <: AbstractEnvWrapper
    env::E
    reward_mapping::F
end

RewardTransformedEnv(env; reward_mapping=identity) =
    RewardTransformedEnv(env, reward_mapping)

RLBase.reward(env::RewardTransformedEnv, args...; kwargs...) =
    env.reward_mapping(reward(env.env, args...; kwargs...))


text:
using Oceananigans.Advection

#####
##### Momentum advection operators
#####

"""
    u∇u(i, j, k, grid, U)

Calculate the advection of momentum in the x-direction U·∇u

    1/Vᵘ * [δxᶠᵃᵃ(ℑxᶜᵃᵃ(Ax * u) * ℑxᶜᵃᵃ(u)) + δy_fca(ℑxᶠᵃᵃ(Ay * v) * ℑyᵃᶠᵃ(u)) + δz_fac(ℑxᶠᵃᵃ(Az * w) * ℑzᵃᵃᶠ(u))]

which will end up at the location `fcc`.
"""
@inline function div_ũu(i, j, k, grid, advection, U)
    return 1/Vᵃᵃᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, momentum_flux_uu, advection, U.u)    +
                                    δyᵃᶜᵃ(i, j, k, grid, momentum_flux_uv, advection, U.u, U.v) +
                                    δzᵃᵃᶜ(i, j, k, grid, momentum_flux_uw, advection, U.u, U.w))
end

"""
    u∇v(i, j, k, grid, U)

Calculates the advection of momentum in the y-direction U·∇v

    1/Vʸ * [δx_cfa(ℑyᵃᶠᵃ(Ax * u) * ℑxᶠᵃᵃ(v)) + δyᵃᶠᵃ(ℑyᵃᶜᵃ(Ay * v) * ℑyᵃᶜᵃ(v)) + δz_afc(ℑxᶠᵃᵃ(Az * w) * ℑzᵃᵃᶠ(w))]

which will end up at the location `cfc`.
"""
@inline function div_ũv(i, j, k, grid, advection, U)
    return 1/Vᵃᵃᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, momentum_flux_vu, advection, U.u, U.v) +
                                    δyᵃᶠᵃ(i, j, k, grid, momentum_flux_vv, advection, U.v)    +
                                    δzᵃᵃᶜ(i, j, k, grid, momentum_flux_vw, advection, U.v, U.w))
end

"""
    u∇w(i, j, k, grid, U)

Calculates the advection of momentum in the z-direction U·∇w

    1/Vʷ * [δx_caf(ℑzᵃᵃᶠ(Ax * u) * ℑxᶠᵃᵃ(w)) + δy_acf(ℑzᵃᵃᶠ(Ay * v) * ℑyᵃᶠᵃ(w)) + δzᵃᵃᶠ(ℑzᵃᵃᶜ(Az * w) * ℑzᵃᵃᶜ(w))]

which will end up at the location `ccf`.
"""
@inline function div_ũw(i, j, k, grid, advection, U)
    return 1/Vᵃᵃᶠ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, momentum_flux_wu, advection, U.u, U.w) +
                                    δyᵃᶜᵃ(i, j, k, grid, momentum_flux_wv, advection, U.v, U.w) +
                                    δzᵃᵃᶠ(i, j, k, grid, momentum_flux_ww, advection, U.w))
end


text:
# kernels
function copy(C, A; kwargs...)
    for i in eachindex(C, A)
        @inbounds C[i] = A[i]
    end
    return nothing
end

function scale(B, C, s; kwargs...)
    for i in eachindex(C)
        @inbounds B[i] = s * C[i]
    end
    return nothing
end

function add(C, A, B; kwargs...)
    for i in eachindex(C)
        @inbounds C[i] = A[i] + B[i]
    end
    return nothing
end

function triad(A, B, C, s; kwargs...)
    for i in eachindex(C)
        @inbounds A[i] = B[i] + s * C[i]
    end
    return nothing
end


text:
function has_flashes(levels :: Array{Array{Int64, 1}, 1}) :: Bool
    has_flashes = false
    for i in 1:length(levels)
        for j in 1:length(levels[1])
            if levels[i][j] == 10
                has_flashes = true
                break
            end
        end
    end
    return has_flashes
end

function get_adjacents(levels :: Array{Array{Int64, 1}, 1}, i :: Int64, j :: Int64) :: Array{Tuple{Int, Int}}
    up_left = Union{Nothing, Tuple{Int, Int}}
    up_right = Union{Nothing, Tuple{Int, Int}}

    down = Union{Nothing, Tuple{Int, Int}}
    down_left = Union{Nothing, Tuple{Int, Int}}
    down_right = Union{Nothing, Tuple{Int, Int}}

    left = Union{Nothing, Tuple{Int, Int}}
    right = Union{Nothing, Tuple{Int, Int}}

    if i <= 1
        up = Nothing
    else
        up = (i - 1, j)
    end

    if i <= 1 || j <= 1
        up_left = Nothing
    else
        up_left = (i - 1, j - 1)
    end

    if i <= 1 || j >= 10
        up_right = Nothing
    else
        up_right = (i - 1, j + 1)
    end

    if i >= 10
        down = Nothing
    else
        down = (i + 1, j)
    end

    if i >= 10 || j <= 1
        down_left = Nothing
    else
        down_left = (i + 1, j - 1)
    end

    if i >= 10 || j >= 10
        down_right = Nothing
    else
        down_right = (i + 1, j + 1)
    end

    if j <= 1
        left = Nothing
    else
        left = (i, j - 1)
    end

    if j >= 10
        right = Nothing
    else
        right = (i, j + 1)
    end

    result = [up, up_left, up_right, down, down_left, down_right, left, right]
    return filter(x -> x != Nothing, result)
end

function step(levels :: Array{Array{Int64, 1}, 1}) :: Int64
    for i in 1:length(levels)
        for j in 1:length(levels[1])
            levels[i][j] += 1
        end
    end

    while has_flashes(levels)
        for i in 1:length(levels)
            for j in 1:length(levels[1])
                if levels[i][j] == 10
                    levels[i][j] += 1
                    for adj in get_adjacents(levels, i, j)
                        if levels[adj[1]][adj[2]] < 10
                            levels[adj[1]][adj[2]] += 1
                        end
                    end
                end
            end
        end
    end

    flashes = 0
    for i in 1:length(levels)
        for j in 1:length(levels[1])
            if levels[i][j] > 9
                levels[i][j] = 0
                flashes += 1
            end
        end
    end

    return flashes
end

function parse_file(path :: String) :: Array{Array{Int, 1}, 1}
    levels = Array{Array{Int64, 1}, 1}()
    open(path) do file
        for line in eachline(file)
            line_numbers = Int[]
            if length(line) > 0
                for i = firstindex(line):lastindex(line)
                    push!(line_numbers, Int(line[i]) - Int('0'))
                end
                push!(levels, line_numbers)
            else
            end
        end
    end
    return levels
end

function part1(path :: String)
    levels = parse_file(path)

    total_flashes = 0
    for i in 1:100
        flashes = step(levels)
        total_flashes += flashes
    end

    println("Part 1: ", total_flashes)
end

function part2(path :: String)
    levels = parse_file(path)

    flashes = 0
    i = 0
    while flashes != 100
        flashes = step(levels)
        i += 1
    end

    println("Part 2: ", i)
end

for arg in ARGS
    part1(arg)
    part2(arg)
end