text:
using Random

include("input/fMultiSinGen.jl")

function gensignalARX(θ_true, τ_true; M1=2, M2=2, fMin=0.8, fMax=1.0, fs=1.0, T=100)

    # Orders
    M = M1+1+M2

    # Parameters
    θ_true = θ_scale .*(rand(M,) .- 0.5)

    # Input frequency and amplitude
    input, _ = fMultiSinGen(T, 1, 1, fMin=fMin, fMax=fMax, fs=fs, type_signal="odd")

    # Observation array
    output = zeros(T,)
    errors = zeros(T,)

    for k = 1:T

        # Input
        input[k] = A[k]*cos(ω*k)

        # Generate noise
        errors[k] = sqrt(inv(τ_true))*randn(1)[1]

        # Output
        if k < M
            output[k] = input[k] + errors[k]
        else
            # Update history vectors
            x_kmin1 = output[k-1:-1:k-M1]
            z_kmin1 = input[k-1:-1:k-M2]

            # Compute output
            output[k] = θ_true'*[x_kmin1; input[k]; z_kmin1] + errors[k]
        end
    end

    return input, output
end

function gensignalNARX(ϕ, θ_scale, τ_true; M1=2, M2=2, degree=1, fMin=0.8, fMax=1.0, fs=1.0, T=100)

    # Orders
    M = (M1+1+M2)*degree + 1

    # Parameters
    θ_true = θ_scale .*(rand(M,) .- 0.5)
    θ_true[end] = 0.

    # Input frequency and amplitude
    input, _ = fMultiSinGen(T, 1, 1, fMin=fMin, fMax=fMax, fs=fs, type_signal="odd")

    # Observation array
    output = zeros(T,)
    errors = zeros(T,)

    for k = 1:T

        # Input
        input[k] = A[k]*cos(ω*k)

        # Generate noise
        errors[k] = sqrt(inv(τ_true))*randn(1)[1]

        # Output
        if k < M
            output[k] = input[k] + errors[k]
        else
            # Update history vectors
            x_kmin1 = output[k-1:-1:k-M1]
            z_kmin1 = input[k-1:-1:k-M2]

            # Compute output
            output[k] = θ_true'*ϕ([x_kmin1; input[k]; z_kmin1]) + errors[k]
        end
    end

    return input, output
end

function gensignalNARMAX(ϕ, θ_scale, τ_true; M1=2, M2=2, M3=2, degree=1, fMin=0.8, fMax=1.0, fs=1.0, uStd=1., T=100, num_periods=1, points_period=1)

    # Orders
    M = (M1+1+M2+M3)*degree + 1

    # Parameters
    θ_true = θ_scale .*(rand(M,) .- 0.5)
    θ_true[end] = 0.0

    # Input frequency and amplitude
    input, _ = fMultiSinGen(points_period, num_periods, 1, fMin=fMin, fMax=fMax, fs=fs, type_signal="full", uStd=uStd)

    # Observation array
    output = zeros(T,)
    errors = zeros(T,)

    for k = 1:T

        # Generate noise
        errors[k] = sqrt(inv(τ_true))*randn(1)[1]

        # Output
        if k < (maximum([M1,M2,M3])+1)
            output[k] = input[k] + errors[k]
        else
            # Update history vectors
            x_kmin1 = output[k-1:-1:k-M1]
            z_kmin1 = input[k-1:-1:k-M2]
            r_kmin1 = errors[k-1:-1:k-M3]

            # Compute output
            output[k] = θ_true'*ϕ([x_kmin1; input[k]; z_kmin1; r_kmin1]) + errors[k]
        end
    end

    return input, output
end

function split_data(T, split_index; start_index=10)

    # Check inputs
    if split_index > T
        error("Split index "*string(split_index)*"must be lower than end time T="*string(T))
    end

    # Select training set
    ix_trn = Array{Int64,1}(start_index:split_index)

    # Select validation set
    ix_val = Array{Int64,1}(split_index+1:T)

    return ix_trn, ix_val
end


text:
# Device pointer with address space information

#
# Address spaces
#

export
    AS, addrspace

abstract type AddressSpace end

module AS

using CUDAnative
import CUDAnative: AddressSpace

struct Generic  <: AddressSpace end
struct Global   <: AddressSpace end
struct Shared   <: AddressSpace end
struct Constant <: AddressSpace end
struct Local    <: AddressSpace end

end

#
# Device pointer
#

struct DevicePtr{T,A}
    ptr::Ptr{T}

    # inner constructors, fully parameterized
    DevicePtr{T,A}(ptr::Ptr{T}) where {T,A<:AddressSpace} = new(ptr)
end

# outer constructors, partially parameterized
DevicePtr{T}(ptr::Ptr{T}) where {T} = DevicePtr{T,AS.Generic}(ptr)

# outer constructors, non-parameterized
DevicePtr(ptr::Ptr{T})              where {T} = DevicePtr{T,AS.Generic}(ptr)

## getters

Base.pointer(p::DevicePtr) = p.ptr

Base.eltype(::Type{<:DevicePtr{T}}) where {T} = T

addrspace(x::DevicePtr) = addrspace(typeof(x))
addrspace(::Type{DevicePtr{T,A}}) where {T,A} = A

## conversions

# between regular and device pointers
## simple conversions disallowed
Base.convert(::Type{Ptr{T}}, p::DevicePtr{T})        where {T} = throw(InexactError(:convert, Ptr{T}, p))
Base.convert(::Type{<:DevicePtr{T}}, p::Ptr{T})      where {T} = throw(InexactError(:convert, DevicePtr{T}, p))
## unsafe ones are allowed
Base.unsafe_convert(::Type{Ptr{T}}, p::DevicePtr{T}) where {T} = pointer(p)

# defer conversions to DevicePtr to unsafe_convert
Base.cconvert(::Type{<:DevicePtr}, x) = x

# between device pointers
Base.convert(::Type{<:DevicePtr}, p::DevicePtr)                         = throw(InexactError(:convert, DevicePtr, p))
Base.convert(::Type{DevicePtr{T,A}}, p::DevicePtr{T,A})   where {T,A}   = p
Base.unsafe_convert(::Type{DevicePtr{T,A}}, p::DevicePtr) where {T,A}   = DevicePtr{T,A}(reinterpret(Ptr{T}, pointer(p)))
## identical addrspaces
Base.convert(::Type{DevicePtr{T,A}}, p::DevicePtr{U,A}) where {T,U,A} = Base.unsafe_convert(DevicePtr{T,A}, p)
## convert to & from generic
Base.convert(::Type{DevicePtr{T,AS.Generic}}, p::DevicePtr)               where {T}     = Base.unsafe_convert(DevicePtr{T,AS.Generic}, p)
Base.convert(::Type{DevicePtr{T,A}}, p::DevicePtr{U,AS.Generic})          where {T,U,A} = Base.unsafe_convert(DevicePtr{T,A}, p)
Base.convert(::Type{DevicePtr{T,AS.Generic}}, p::DevicePtr{T,AS.Generic}) where {T}     = p  # avoid ambiguities
## unspecified, preserve source addrspace
Base.convert(::Type{DevicePtr{T}}, p::DevicePtr{U,A}) where {T,U,A} = Base.unsafe_convert(DevicePtr{T,A}, p)

## limited pointer arithmetic & comparison

Base.:(==)(a::DevicePtr, b::DevicePtr) = pointer(a) == pointer(b) && addrspace(a) == addrspace(b)

Base.isless(x::DevicePtr, y::DevicePtr) = Base.isless(pointer(x), pointer(y))
Base.:(-)(x::DevicePtr, y::DevicePtr)   = pointer(x) - pointer(y)

Base.:(+)(x::DevicePtr{T,A}, y::Integer) where {T,A} = DevicePtr{T,A}(pointer(x) + y)
Base.:(-)(x::DevicePtr{T,A}, y::Integer) where {T,A} = DevicePtr{T,A}(pointer(x) - y)
Base.:(+)(x::Integer, y::DevicePtr) = y + x

## memory operations

Base.convert(::Type{Int}, ::Type{AS.Generic})  = 0
Base.convert(::Type{Int}, ::Type{AS.Global})   = 1
Base.convert(::Type{Int}, ::Type{AS.Shared})   = 3
Base.convert(::Type{Int}, ::Type{AS.Constant}) = 4
Base.convert(::Type{Int}, ::Type{AS.Local})    = 5

@generated function Base.unsafe_load(p::DevicePtr{T,A}, i::Integer=1,
                                     ::Val{align}=Val(1)) where {T,A,align}
    eltyp = convert(LLVMType, T)

    T_int = convert(LLVMType, Int)
    T_ptr = convert(LLVMType, Ptr{T})

    T_actual_ptr = LLVM.PointerType(eltyp)

    # create a function
    param_types = [T_ptr, T_int]
    llvm_f, _ = create_function(eltyp, param_types)

    # generate IR
    Builder(JuliaContext()) do builder
        entry = BasicBlock(llvm_f, "entry", JuliaContext())
        position!(builder, entry)

        ptr = inttoptr!(builder, parameters(llvm_f)[1], T_actual_ptr)

        ptr = gep!(builder, ptr, [parameters(llvm_f)[2]])
        ptr_with_as = addrspacecast!(builder, ptr, LLVM.PointerType(eltyp, convert(Int, A)))
        val = load!(builder, ptr_with_as)
        alignment!(val, align)
        ret!(builder, val)
    end

    call_function(llvm_f, T, Tuple{Ptr{T}, Int}, :((pointer(p), Int(i-one(i)))))
end

@generated function Base.unsafe_store!(p::DevicePtr{T,A}, x, i::Integer=1,
                                       ::Val{align}=Val(1)) where {T,A,align}
    eltyp = convert(LLVMType, T)

    T_int = convert(LLVMType, Int)
    T_ptr = convert(LLVMType, Ptr{T})

    T_actual_ptr = LLVM.PointerType(eltyp)

    # create a function
    param_types = [T_ptr, eltyp, T_int]
    llvm_f, _ = create_function(LLVM.VoidType(JuliaContext()), param_types)

    # generate IR
    Builder(JuliaContext()) do builder
        entry = BasicBlock(llvm_f, "entry", JuliaContext())
        position!(builder, entry)

        ptr = inttoptr!(builder, parameters(llvm_f)[1], T_actual_ptr)

        ptr = gep!(builder, ptr, [parameters(llvm_f)[3]])
        ptr_with_as = addrspacecast!(builder, ptr, LLVM.PointerType(eltyp, convert(Int, A)))
        val = parameters(llvm_f)[2]
        inst = store!(builder, val, ptr_with_as)
        alignment!(inst, align)
        ret!(builder)
    end

    call_function(llvm_f, Cvoid, Tuple{Ptr{T}, T, Int}, :((pointer(p), convert(T,x), Int(i-one(i)))))
end

## loading through the texture cache

# NOTE: CUDA 8.0 supports more caching modifiers, but those aren't supported by LLVM yet

# operand types supported by llvm.nvvm.ldg.global
const CachedLoadOperands = Union{UInt8, UInt16, UInt32, UInt64,
                                 Int8, Int16, Int32, Int64,
                                 Float32, Float64}

# containing DevicePtr types
const CachedLoadPointers = Union{Tuple(DevicePtr{T,AS.Global}
                                 for T in Base.uniontypes(CachedLoadOperands))...}

@generated function unsafe_cached_load(p::DevicePtr{T,AS.Global}, i::Integer=1,
                                       ::Val{align}=Val(1)) where
                                      {T<:CachedLoadOperands,align}
    # NOTE: we can't `ccall(..., llvmcall)`, because
    #       1) Julia passes pointer arguments as plain integers
    #       2) we need to addrspacecast the pointer argument

    eltyp = convert(LLVMType, T)

    T_int = convert(LLVMType, Int)
    T_int32 = LLVM.Int32Type(JuliaContext())
    T_ptr = convert(LLVMType, Ptr{T})

    T_actual_ptr = LLVM.PointerType(eltyp)
    T_actual_ptr_as = LLVM.PointerType(eltyp, convert(Int, AS.Global))

    # create a function
    param_types = [T_ptr, T_int]
    llvm_f, _ = create_function(eltyp, param_types)

    # create the intrinsic
    intrinsic_name = let
        class = if isa(eltyp, LLVM.IntegerType)
            :i
        elseif isa(eltyp, LLVM.FloatingPointType)
            :f
        else
            error("Cannot handle $eltyp argument to unsafe_cached_load")
        end
        width = sizeof(T)*8
        typ = Symbol(class, width)
        "llvm.nvvm.ldg.global.$class.$typ.p1$typ"
    end
    mod = LLVM.parent(llvm_f)
    intrinsic_typ = LLVM.FunctionType(eltyp, [T_actual_ptr_as, T_int32])
    intrinsic = LLVM.Function(mod, intrinsic_name, intrinsic_typ)

    # generate IR
    Builder(JuliaContext()) do builder
        entry = BasicBlock(llvm_f, "entry", JuliaContext())
        position!(builder, entry)

        ptr = inttoptr!(builder, parameters(llvm_f)[1], T_actual_ptr)

        ptr = gep!(builder, ptr, [parameters(llvm_f)[2]])
        ptr_with_as = addrspacecast!(builder, ptr, T_actual_ptr_as)
        val = call!(builder, intrinsic,
                    [ptr_with_as, ConstantInt(Int32(align), JuliaContext())])
        ret!(builder, val)
    end

    call_function(llvm_f, T, Tuple{Ptr{T}, Int}, :((pointer(p), Int(i-one(i)))))
end

@inline function unsafe_cached_load(p::DevicePtr{T,AS.Global}, args...) where {T}
    recurse_pointer_invocation(unsafe_cached_load, p, CachedLoadPointers, args...)
end

ext: jl


text:
using SMC, Base.Test

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

srand(125)

## Linear Gaussian

dx, dy = 5, 3
A = randn(dx,dx); A /= 0.9norm(A)
B = randn(dy,dx); B /= 1.1norm(B)
Q = randn(dx,dx); Q *= Q'; Q += 0.5*eye(dx); Q += Q'; Q /= 5
R = randn(dy,dy); R *= R'; R += 0.5*eye(dy); R += R'; R /= 20

lg  = LinearGaussian(A,B,Q,R)
hmm = HMM(lg)
x0  = randn(dx)

### LinearGaussian -- testing the generation

srand(123)
(states, observations) = generate(lg, x0, 3)

srand(123)
noisex = chol(Q)' * randn(dx, 3)
noisey = chol(R)' * randn(dy, 3)
state1 = x0
state2 = A*state1+noisex[:,2] # we do not use the first noise
state3 = A*state2+noisex[:,3] # since x0 is given
obs1   = B*state1+noisey[:,1]
obs2   = B*state2+noisey[:,2]
obs3   = B*state3+noisey[:,3]

@test state1==states[:,1] &&
      state2==states[:,2] &&
      state3==states[:,3]
@test obs1==observations[:,1] &&
      obs2==observations[:,2] &&
      obs3==observations[:,3]

@test isapprox( hmm.transmean(0,state1), A*state1)
@test isapprox( hmm.obsmean(0,state1), B*state1)
@test isapprox( hmm.transloglik(0,state1,state2),
        (-norm(chol(Q)'\(state2-A*state1))^2/2) )
@test isapprox( hmm.obsloglik(0,obs2,state2),
        (-norm(chol(R)'\(  obs2-B*state2))^2/2) )


text:
img = SquareImage(rows=4, cols=4, l=2, calibration_value=1, background=0)
@testset "Testing FBP algorithms" begin
    @testset "Testing FBP FFT algorithm" begin
        sinog = radon(img, Radon())
        tomog = iradon(sinog, FBP())
        @test tomog ≈ [
            [-0.642772  0.374742  0.374742  -0.642772];
            [0.033042  1.05056   1.05056    0.0330419];
            [0.033042  1.05056   1.05056    0.0330422];
            [-0.642772  0.374742  0.374742  -0.642772]
        ]
    end
    @testset "Testing FBP FFT square algorithm" begin
        sinog = radon(img, RadonSquare())
        tomog = iradon(sinog, FBPFFTSquare())
        @test tomog ≈ [
            [0.177991   0.55048  0.554279  -0.260271];
            [0.0837881  1.22652  1.21234   -0.185625];
            [-0.148054   1.21234  1.22652   -0.132705];
            [-0.0358718  0.55428  0.55048    0.140421]
        ]
    end
end


text:
using Base.Test
using Klara

fnames = (
  :value,
  :loglikelihood,
  :logprior,
  :logtarget,
  :gradloglikelihood,
  :gradlogprior,
  :gradlogtarget,
  :tensorloglikelihood,
  :tensorlogprior,
  :tensorlogtarget,
  :dtensorloglikelihood,
  :dtensorlogprior,
  :dtensorlogtarget
)

println("    Testing BasicContMuvParameterState constructors and methods...")

v = Float64[1., 1.5]
s = BasicContMuvParameterState(v)

@test eltype(s) == Float64
@test s.value == v
for i in 2:4
  @test isnan(getfield(s, fnames[i]))
end
for i in 5:7
  @test length(getfield(s, fnames[i])) == 0
end
for i in 8:10
  @test size(getfield(s, fnames[i])) == (0, 0)
end
for i in 11:13
  @test size(getfield(s, fnames[i])) == (0, 0, 0)
end
@test s.diagnosticvalues == []
@test s.size == length(v)
@test s.diagnostickeys == Symbol[]

@test diagnostics(s) == Dict{Symbol, Any}()

v = Float16[0.24, 5.5, -6.3]
dv = [false]
ssize = length(v)
s = BasicContMuvParameterState(v, [:gradlogtarget], [:accept], nothing, nothing, dv)

@test eltype(s) == Float16
@test s.value == v
@test length(s.gradlogtarget) == ssize
for i in 2:4
  @test isnan(getfield(s, fnames[i]))
end
for i in 5:6
  @test length(getfield(s, fnames[i])) == 0
end
for i in 8:10
  @test size(getfield(s, fnames[i])) == (0, 0)
end
for i in 11:13
  @test size(getfield(s, fnames[i])) == (0, 0, 0)
end
@test s.diagnosticvalues == dv
@test s.size == ssize
@test s.diagnostickeys == [:accept]

@test diagnostics(s) == Dict(:accept => dv[1])

ssize = 3
s = BasicContMuvParameterState(ssize)

@test eltype(s) == Float64
@test isa(s.value, Vector{Float64})
@test length(s.value) == ssize
for i in 2:4
  @test isnan(getfield(s, fnames[i]))
end
for i in 5:7
  @test length(getfield(s, fnames[i])) == 0
end
for i in 8:10
  @test size(getfield(s, fnames[i])) == (0, 0)
end
for i in 11:13
  @test size(getfield(s, fnames[i])) == (0, 0, 0)
end
@test s.diagnosticvalues == []
@test s.size == ssize
@test s.diagnostickeys == Symbol[]

@test diagnostics(s) == Dict{Symbol, Any}()

dv = [true]
ssize = 5
s = BasicContMuvParameterState(ssize, [:tensorlogtarget], [:accept], BigFloat, nothing, nothing, dv)

@test eltype(s) == BigFloat
@test isa(s.value, Vector{BigFloat})
@test length(s.value) == ssize
@test size(s.tensorlogtarget) == (ssize, ssize)
for i in 2:4
  @test isnan(getfield(s, fnames[i]))
end
for i in 5:7
  @test length(getfield(s, fnames[i])) == 0
end
for i in 8:9
  @test size(getfield(s, fnames[i])) == (0, 0)
end
for i in 11:13
  @test size(getfield(s, fnames[i])) == (0, 0, 0)
end
@test s.diagnosticvalues == dv
@test s.size == ssize
@test s.diagnostickeys == [:accept]

@test diagnostics(s) == Dict(:accept => dv[1])

z = BasicContMuvParameterState(Float64[-0.12, 12.15], [:gradloglikelihood], [:accept], nothing, nothing, [false])
w = deepcopy(z)
@test isequal(z, w)


text:
# Test of breadth-first-search

using Graphs
using Test

g = simple_adjlist(6)

add_edge!(g, 1, 2)
add_edge!(g, 1, 3)
add_edge!(g, 1, 6)
add_edge!(g, 2, 4)
add_edge!(g, 2, 5)
add_edge!(g, 3, 5)
add_edge!(g, 3, 6)
add_edge!(g, 5, 1)

@assert num_vertices(g) == 6
@assert num_edges(g) == 8

vs1 = visited_vertices(g, BreadthFirst(), 1)
ds1 = gdistances(g, 1)
@test vs1 == [1, 2, 3, 6, 4, 5]
@test ds1 == [0, 1, 1, 2, 2, 1]

vs2 = visited_vertices(g, BreadthFirst(), [1, 3])
ds2 = gdistances(g, [1, 3])
@test vs2 == [1, 3, 2, 6, 5, 4]
@test ds2 == [0, 1, 0, 2, 1, 1]


text:
"""
    function write(output;
                   header=missing,
                   data=missing,
                   delim=',',
                   quotes=missing,
                   quotetypes=AbstractString)

Write a dataset to disk or IO

# Arguments
- `output`
    - the path on disk or IO where you want to write to
- `header`
    - the column names for the data to `output`
    - default: `header=missing`
        - no header is written
- `data`
    - the dataset to write to `output`
    - default: `data=missing`
        - no data is written
- `delim`
    - the delimiter to seperate fields by
    - default: `delim=','`
        - for CSV files
    - frequently used:
        - `delim='\\t'`
        - `delim=' '`
        - `delim='|'`
- `quotes`
    - the quoting character to use when writing fields
    - default: `quotes=missing`
        - fields are not quoted by default, and fields are written using julia's
          default string-printing mechanisms
- `quotetypes::Type`
    - when quoting fields, quote only columns where `coltype <: quotetypes`
        - columns of type `Union{<:quotetypes, Missing}` will also be quoted
    - default: `quotetypes=AbsractString`
        - only the header and fields where `coltype <: AbsractString` will be quoted
    - frequently used:
        - `quotetypes=Any`
            - quote every field in the dataset
"""
function write(fullpath::Union{String, IO};
               header::Union{Vector{String}, Missing}=missing,
               data::Union{Vector{<:Any}, Missing}=missing,
               delim::Union{Char, String}=',',
               quotes::Union{Char, Missing}=missing,
               quotetypes::Type=AbstractString)
    if ismissing(header) && ismissing(data)
        throw(ArgumentError("no header or data provided"))
    elseif !ismissing(data)
        @assert length(unique(length.(data))) == 1
        if !ismissing(header)
            @assert length(header) == length(data)
        end
    end
    if isa(fullpath, IO)
        if iswritable(fullpath)
            f = fullpath
        else
            throw(ArgumentError("""
                                Provided IO is not writable
                                """))
        end
    else
        f = open(fullpath, "w")
    end
    if !ismissing(header)
        if !ismissing(quotes)
            for i in eachindex(header)
                header[i] = string(quotes, header[i], quotes)
            end
        end
        Base.write(f, join(header, delim) * "\n")
    end
    if !ismissing(data)
        numcols = length(data)
        eltypes = eltype.(data)
        for row in 1:length(data[1])
            rowvalues = [string(data[col][row]) for col in 1:numcols]
            if !ismissing(quotes)
                for i in findall(e -> e <: quotetypes || (e != Missing && e <: Union{quotetypes, Missing}), eltypes)
                    rowvalues[i] = string(quotes, rowvalues[i], quotes)
                end
            end
            Base.write(f, join(rowvalues, delim) * "\n")
        end
    end
    close(f)
end

"""
    function write(output,
                   df;
                   delim=',',
                   quotes=missing,
                   quotetypes=AbstractString)

Write a DataFrame to disk or IO
"""
function write(fullpath, df::DataFrame; kwargs...)
    write(fullpath; header = string.(names(df)), data = DataFrames.columns(df), kwargs...)
end


text:
const PressureArray{x} = Array{T, x} where T <: Union{Unitful.Pressure, Number}
const Volt = typeof(1.0u"V")
const Meter = typeof(1.0u"m")
const Sec = typeof(1.0u"s")

abstract type UltrasoundMetric end

abstract type ISPPA <: UltrasoundMetric end      # intensity_sppa
abstract type ISPTA <: UltrasoundMetric end      # intensity_spta
abstract type MI <: UltrasoundMetric end         # mechanical_index

struct Metric{T, D}
    val::AbstractArray
end

get_metric_display_name(a::Metric{ISPPA, T}) where T = "Intensity_SPPA"
get_metric_display_name(a::Metric{ISPTA, T}) where T = "Intensity_SPTA"
get_metric_display_name(a::Metric{MI, T}) where T = "Mechanical Index"

@kwdef struct ScanParameters
    medium::Any
    excitation::Any
    f0::Float64
    hydrophone_id::Symbol
    preamp_id::Union{Symbol, Nothing} = nothing
    calibration_factor::typeof(1.0u"Pa/V") = volt_to_pressure(f0, hydrophone_id, preamp_id)
end

function Base.show(io::IO, scan_setup::ScanParameters)
    println(io, "medium: ", scan_setup.medium)
    println(io, "excitation: ", scan_setup.excitation)
    println(io, "f0: ", scan_setup.f0)
    println(io, "hydrophone_id: ", scan_setup.hydrophone_id)
    println(io, "preamp_id: ", scan_setup.preamp_id)
    println(io, "calibration_factor: ", scan_setup.calibration_factor)
end

struct ScanMetric
    params::ScanParameters
    pressure::PressureArray
    isppa::Metric{ISPPA}
    isppa_max::Tuple{Unitful.Quantity, Vector{Meter}}
    ispta::Metric{ISPTA}
    ispta_max::Tuple{Unitful.Quantity, Vector{Meter}}
    mechanical_index::Metric{MI}
    mechanical_index_max::Tuple{Float64, Vector{Meter}}
end

"""
Input:
       xyz: The handle of desired xyz stage
     scope: Handle of desired scope
   channel: This input decides which channel the scope will read from.

```
ps = initialize(PositionerSystem)
scope = initialize(Scope350MHz)
hydrophone = IntensityScan(xyz = ps, scope = scope, channel = 1, precapture_delay=0u"s")
```
"""
@kwdef struct IntensityScan
    xyz::T where T <: ThorlabsLTStage.PositionerSystem
    scope::TcpInstruments.Instr{T} where T <: Oscilloscope
    channel::Int64
    post_move_delay::typeof(1.0u"s") = 0u"s"
    precapture_delay::typeof(1.0u"s")
    trigger_function::Function = () -> nothing
end

function Base.show(io::IO, scan_setup::IntensityScan)
    println(io, "positioning_system: ", scan_setup.xyz)
    println(io, "scope: ", scan_setup.scope)
    println(io, "scope_channel: ", scan_setup.channel)
    println(io, "post_move_delay: ", scan_setup.post_move_delay)
    println(io, "precapture_delay: ", scan_setup.precapture_delay)
    println(io, "trigger_function: ", scan_setup.trigger_function)
end

const version = v"1"

struct ScanData
    version::VersionNumber
    axes::String
    scope_info::TcpInstruments.ScopeInfo
    time::Vector{Sec}
    waveform::Array{Volt, N} where {N}
    coordinates::Array{Meter, N} where {N}
    metrics::Union{Nothing, ScanMetric}
end

function ScanData(
    axes::AbstractString,
    scope_info::TcpInstruments.ScopeInfo,
    time::AbstractArray,
    num_points_per_waveform,
    num_points_per_axis::AbstractArray
)
    waveform_dimensions = Tuple([num_points_per_waveform; num_points_per_axis])
    waveform = u"V" * zeros(waveform_dimensions)

    num_coordinates_per_point = 3       # x, y, z location
    coordinates_dimensions = Tuple([num_coordinates_per_point; num_points_per_axis])
    coordinates = u"m" * zeros(coordinates_dimensions)

    return ScanData(version, axes, scope_info, time, waveform, coordinates, nothing)
end

function Base.show(io::IO, scan::ScanData)
    println(io, "version: ", scan.version)
    println(io, "axes: ", scan.axes)
    println(io, "scope_info: ", scan.scope_info)
    println(io, "time: ", size(scan.time))
    println(io, "waveform: ", size(scan.waveform))
    println(io, "coordinates: ", size(scan.coordinates))
    metrics_fieldnames = fieldnames(typeof(scan.metrics))
    if isempty(metrics_fieldnames)
        println(io, "metrics: [none]")
    else
        println(io, "metrics: ")
        for fieldname in metrics_fieldnames
            println(io, "  ", fieldname)
        end
    end
end


text:
####### general utilities

## macro for argument checking from Distributions.jl
macro check_args(D, cond)
    quote
        if !($(esc(cond)))
            throw(ArgumentError(string(
                $(string(D)), ": the condition ", $(string(cond)), " is not satisfied.")))
        end
    end
end

# set all keys of a dictionary to zero
function reset!( d::Dict; z0=missing )
    for k in keys(d)
        if ismissing(z0)
            d[k] = zero(typeof(d[k]))
        else
            d[k] = z0
        end
    end
end

function increment!( d::Dict{S, T}, k::S, i::T) where {T<:Real, S<:AbstractString}
    if haskey(d, k)
        d[k] += i
    else
        d[k] = i
    end
end

function increment!( d::Dict{S, Dict{S, T}}, k1::S, k2::AbstractString, i::T) where {T<:Real, S<:AbstractString}
    if !haskey(d, k1)
        d[k1] = Dict{AbstractString, Real}()
    end
    if !haskey(d[k1], k2)
        d[k1][k2] = i
    else
        d[k1][k2] += i
    end
end

"""
    linesearch()

 Finds the minimum of a 1D function. Particular algorithms make assumptions about the form of the function.

## Input Arguments
* `algorithm::String`: algorithm to use. Current choices are `Grid`, `Golden` ...
* `func::Function`: function to optimize
* `x0::AbstractFloat`: [optional] starting point (only used for some algorithms)
* `range::Array{AbstractFloat,1}`: [optional] search range (only used for some algorithms)
* `ε::AbstractFloat`:  [default = 1.0e-3] error tolerance.

The input function should take one floating point parameter and map it
to a floating point output. At present there are not facilities to
pass additional parameters to the function, but this can be
accomplished by creating an anonymous function.

## Output arguments
* `x_star`: estimate of the location of the maximum
* `y_star`: estimate of the value at the maximum
* `error_estimate`: en estimate of the maximum possible error, which should be ≦ ε
* `n_evaluations`: number of function evaluations

## Examples
```jldoctest
julia> r = [0.0, 2.0]
julia> ε = 1.0e-6
julia> (x_star1, y_star1, error_estimate1, n_evaluations2) = linesearch( "Golden", x->sin(x); range=r, ε=ε)
(1.5707963972987837, 0.9999999999999976, 3.322605284239444e-7, 33)
```
"""
function linesearch( algorithm::String, func::Function; x0::AbstractFloat=0.0, range::Array{T,1}, ε::AbstractFloat=1.0e-3) where {T<:Real}
    if algorithm=="Grid"
        # dumb-arse grid search
        x = Float64(range[1]) : ε : Float64(range[2])
        y = func.(x)
        (y_star, i) = findmax(y)
        x_star = x[i]
        error_estimate = ε
        n_evaluations = length(x)
    elseif algorithm=="Golden"
        # Golden section search
        γ = (sqrt(5)-1)/2
        a = Float64( range[1] )
        b = Float64( range[2] )
        ell = b - a
        no_of_iterations=ceil( log(ε/ell)/log(γ) )
        p = b - γ*(b-a)
        q = a + b - p
        fp = -func(p) # default search is a minimisation
        fq = -func(q)
        n_evaluations = 2
        for i = 1 : no_of_iterations
            ell = b - a
            if fp < fq
                a = a
                b = q
                if i < no_of_iterations
                    q = p
                    p = a + b - q
                    fq = fp
                    fp = -func(p)
                    n_evaluations = n_evaluations + 1
                end
            else
                a = p
                b = b
                if i < no_of_iterations
                    p = q
                    q = a + b - q
                    fp = fq
                    fq = -func(q)
                    n_evaluations = n_evaluations + 1
                end
            end
        end
        x_star = (a + b) / 2
        y_star = func(x_star)
        n_evaluations = n_evaluations + 1
        error_estimate = (b - a)/2.0
    else
        error("algorithm isn't implemented")
    end
    return (x_star, y_star, error_estimate, n_evaluations)
end


text:
# # Poisson And Laplace Equation Solvers On The Sphere
#
# Compute scalar potential, and divergent component, from vertically integrated transport.

using SparseArrays, Statistics

"""
    MaskWetPoints(TrspCon)

Mask land points with NaN.
```
(TrspCon, mskWet, mskDry)=MaskWetPoints(TrspCon)
```
"""
function MaskWetPoints(TrspCon)
    mskWet=1.0 .+ 0.0 * TrspCon
    mskDry=1.0 * isnan.(mskWet)
    mskDry=mask(mskDry,NaN,0.0)
    #
    tmp1=fill(1.0,mskWet); tmp2=exchange(tmp1);
    for I=1:size(tmp1.f,1)
        tmp3=mskWet[I]; tmp4=tmp2[I];
        tmp4=tmp4[2:end-1,1:end-2]+tmp4[2:end-1,3:end]+tmp4[1:end-2,2:end-1]+tmp4[3:end,2:end-1];
        !isempty(findall(isnan.(tmp4) .& (!isnan).(tmp3))) ? println("warning: modified mask") : nothing
        tmp3[findall(isnan.(tmp4))] .= NaN
        mskWet[I]=tmp3
    end
    #
    TrspCon=mask(TrspCon,0.0)*mskWet;
    return TrspCon, mskWet, mskDry
end

"""
    MapWetPoints(mskWet)

Mapping from global array to global ocean vector.
```
(Kvec,Lvec,Kmap,Lmap)=MapWetPoints(mskWet)
```
"""
function MapWetPoints(mskWet)
    tmp1=write(mskWet)[:]
    kk=findall((!isnan).(tmp1))
    nn=length(kk); s0=size(tmp1); s1=mskWet.grid.ioSize;
    Kvec=fill(0.0,s0...); Kvec[kk]=kk; Kmap=read(reshape(Kvec,s1...),mskWet) #global array indices
    Lvec=fill(0.0,s0...); Lvec[kk]=1:nn; Lmap=read(reshape(Lvec,s1...),mskWet) #global vector indices
    return Kvec,Lvec,Kmap,Lmap
end

"""
    SeedWetPoints(tmp::MeshArray,Kmap::MeshArray,Lmap::MeshArray,I...)

Seed a subset of grid points.
```
(FLDones,FLDkkFROM)=SeedWetPoints(tmp::MeshArray,Kmap::MeshArray,Lmap::MeshArray,I...)
```
"""
function SeedWetPoints(tmp::MeshArray,Kmap::MeshArray,Lmap::MeshArray,I...)
    aa=I[1]
    ii=I[2]
    jj=I[3]

    #1) seed points (FLDones) and neighborhood of influence (FLDkkFROM)
    FLDones=fill(0.0,tmp)
    FLDones[aa][ii:3:end,jj:3:end].=1.0
    FLDones[aa][findall(Kmap[aa].==0.0)].=0.0

    FLDkkFROMtmp=fill(0.0,tmp)
    FLDkkFROMtmp[aa][ii:3:end,jj:3:end]=Kmap[aa][ii:3:end,jj:3:end]
    FLDkkFROMtmp[aa][findall(isnan.(tmp[aa]))].=0.0

    FLDkkFROM=exchange(FLDkkFROMtmp)
    FLDkkFROM=mask(FLDkkFROM,0.0)

    for bb in 1:tmp.grid.nFaces
        tmp1=FLDkkFROM[bb]
        tmp2=zeros(size(tmp1) .- 2)
        for ii2 in 1:3
            for jj2 in 1:3
                tmp2=tmp2+tmp1[ii2:end-3+ii2,jj2:end-3+jj2]
            end
        end
        FLDkkFROM[bb]=tmp2
    end

    return FLDones,FLDkkFROM
end

"""
    MatrixForPoisson(TrspCon,mskWet,mskDry,Kvec,Lvec,Kmap,Lmap)

Assemble sparse matrix using mskWet, Kvec, Lvec directly and Kmap, Lmap via SeedWetPoints
```
A=MatrixForPoisson(TrspCon,mskWet,mskDry,Kvec,Lvec,Kmap,Lmap)
```
"""
function MatrixForPoisson(TrspCon,mskWet,mskDry,Kvec,Lvec,Kmap,Lmap)
    I=Array{Int}(undef,0)
    J=Array{Int}(undef,0)
    V=Array{Float64}(undef,0)

    for aa=1:TrspCon.grid.nFaces;
        for ii=1:3; for jj=1:3;
            #1) compute effect of each point on neighboring target point:
            (FLDones,FLDkkFROM)=SeedWetPoints(TrspCon,Kmap,Lmap,aa,ii,jj);
            (tmpU,tmpV)=gradient(FLDones,(;),false)
            dFLDdt=convergence(tmpU,tmpV);
            #2) mask `dFLDdt` since we use a **Neumann** boundary condition.
            #Extrapolation uses a **Dirichlet** boundary condition, so mskFreeze should not be applied then.
            isa(FLDkkFROM,MeshArray) ? FLDkkFROM=write(FLDkkFROM)[:] : nothing;
            #3.1) For wet points -- add contributions in matrix:
            dFLDdtWet=write(dFLDdt.*mskWet)[:];
            tmp1=findall( (dFLDdtWet .!= 0.0) .* (!isnan).(dFLDdtWet));
            tmpV=dFLDdtWet[tmp1]; tmpJ=FLDkkFROM[tmp1]; tmpI=Kvec[tmp1];
            I=[I;Lvec[Int.(tmpI)]]; J=[J;Lvec[Int.(tmpJ)]]; V=[V;tmpV];
            size(tmpV)
            #3.2) For dry points -- This part reflects the `Neumann` boundary condition:
            dFLDdtDry=write(dFLDdt.*mskDry)[:];
            tmp1=findall( (dFLDdtDry .!= 0.0) .* (!isnan).(dFLDdtDry));
            tmpV=dFLDdtDry[tmp1]; tmpIJ=FLDkkFROM[tmp1];
            I=[I;Lvec[Int.(tmpIJ)]]; J=[J;Lvec[Int.(tmpIJ)]]; V=[V;tmpV];
        end; end;
    end;

    nn=sum((!isnan).(mskWet))
    A=sparse(I,J,V,nn,nn)
end

"""
    ScalarPotential(TrspCon)

Scalar potential inversion.
```
TrspPot=ScalarPotential(TrspCon)
```
"""
function ScalarPotential(TrspCon)
    (TrspCon, mskWet, mskDry)=MaskWetPoints(TrspCon)
    (Kvec,Lvec,Kmap,Lmap)=MapWetPoints(mskWet)

    A=MatrixForPoisson(TrspCon,mskWet,mskDry,Kvec,Lvec,Kmap,Lmap)
    yy=write(TrspCon)[:]
    yy=yy[findall(Kvec .!= 0)]

    xx=A\yy
    xx=xx .- median(xx)

    TrspPot=0.0 * write(TrspCon)[:]
    TrspPot[findall(Kvec .!= 0)]=xx
    TrspPot=read(TrspPot,TrspCon)

    return TrspPot
end

"""
    VectorPotential(TrspX,TrspY,Γ,method::Int=1)

Vector potential inversion.
```
TrspPot=ScalarPotential(TrspCon)
```
"""
function VectorPotential(TrspX::MeshArray,TrspY::MeshArray,Γ::NamedTuple,method::Int=1)

    # 1)  streamfunction face by face:

    (fldU,fldV)=exch_UV(TrspX,TrspY);
    fldU=mask(fldU,0.0); fldV=mask(fldV,0.0);

    psi=similar(fldV)
    for I in eachindex(fldV)
        tmp2=cumsum(fldV[I],dims=1)
        psi[I]=[zeros(1,size(tmp2,2));tmp2]
    end

    # 2a)  reset one land value per face to zero (method 1)

    if method==1
        mskW=mask(1.0 .+ 0.0 * mask(view(Γ.hFacW,:,1),NaN,0.0),0.0)
        mskS=mask(1.0 .+ 0.0 * mask(view(Γ.hFacS,:,1),NaN,0.0),0.0)
        (mskW,mskS)=exch_UV(mskW,mskS); mskW=abs.(mskW); mskS=abs.(mskS)

        for iF=1:TrspX.grid.nFaces
            tmp2=mskS[iF]
            tmp3a=[ones(1,size(tmp2,2));tmp2]
            tmp3b=[tmp2;ones(1,size(tmp2,2))]
            tmpB=tmp3a.*tmp3b
            #
            tmpA=psi[iF];
            I=findall(tmpB.==0)
            ii=I[1][1]; jj=I[1][2]
            tmpA[:,jj]=tmpA[:,jj] .- tmpB[ii,jj]
            #
            for kk=1:jj-1
                tmpE=tmpA[:,kk+1]-tmpA[:,kk]
                tmpE+=fldU[iF][:,kk]
                tmpE=median(tmpE[findall((!isnan).(tmpE))])
                tmpA[:,kk]=tmpA[:,kk] .- tmpE
            end
            #
            for kk=jj+1:size(tmpA,2)
                tmpE=tmpA[:,kk]-tmpA[:,kk-1]
                tmpE+=fldU[iF][:,kk-1]
                tmpE=median(tmpE[findall((!isnan).(tmpE))])
                tmpA[:,kk]=tmpA[:,kk] .- tmpE
            end;
            #
            psi[iF]=tmpA
        end
    end

    # 2b)  subtract divergent flow line by line (method 2)

    if method==2
        for I in eachindex(fldV)
            tmp2=diff(psi[I],dims=2)+fldU[I]
            tmp3=cumsum(mean(tmp2,dims=1),dims=2)
            psi[I]=psi[I]-ones(size(psi[I],1),1)*[0.0 tmp3]
        end
    end

    # 3) match edge values:

    if fldU.grid.nFaces>1
        TMP1=similar(psi)
        for I in eachindex(TMP1); TMP1[I] = fill(I,size(psi[I])); end
        TMP2=exchange(TMP1) #this is a trick

    for I in 1:TrspX.grid.nFaces-1
        tmp2=exchange(psi) #this is a trick
        tmp3=tmp2[I+1]; tmp3[3:end-2,3:end-2].=NaN #mask out interior points
        TMP3=TMP2[I+1]; tmp3[findall(TMP3.>I+1)].=NaN #mask out edges points coming from unadjusted faces
        tmp3[findall(TMP3.==0)].=NaN
        tmp3[:,1]=tmp3[:,1]-tmp3[:,2]; tmp3[:,end]=tmp3[:,end]-tmp3[:,end-1] #compare edge points
        tmp3[1,:]=tmp3[1,:]-tmp3[2,:]; tmp3[end,:]=tmp3[end,:]-tmp3[end-1,:] #compare edge points
        tmp3[2:end-1,2:end-1] .= NaN #mask out remaining interior points
        psi[I+1]=psi[I+1] .+ median(tmp3[findall((!isnan).(tmp3))]) #adjust the face data
    end
    end

    # 4) put streamfunction at cell center

    for I in 1:TrspX.grid.nFaces
        tmp3=psi[I]
        tmp3=(tmp3[:,1:end-1]+tmp3[:,2:end])/2
        tmp3=(tmp3[1:end-1,:]+tmp3[2:end,:])/2
        psi[I]=tmp3
    end

    # 5) reset one land point to zero

    #all values:
    tmp1=write(psi)

    #all land points:
    tmp2=write(view(Γ.hFacC,:,1))
    tmp2=findall( (tmp2 .== 0.0) .& (!isnan).(tmp1))

    #closest to Boston:
    tmp_lon=write(Γ.XC)[tmp2]
    tmp_lat=write(Γ.YC)[tmp2]
    tmp_dis=(tmp_lat .- 42.3601).^2 + (tmp_lon .- 71.0589).^2
    tmp2=tmp2[findall(tmp_dis .== minimum(tmp_dis))]

    #set that point to zero
    psi=psi .- tmp1[tmp2[1]];

    return psi
end


text:
# LibTOAST.jl: interface to the TOAST++ library
# Copyright (C) 2019 Samuel Powell

# Imports
import SparseArrays: sparse

# Export types
export SystemMatrix, sparse

mutable struct SystemMatrix

  ptr::Cxx.CxxCore.CppPtr
  mesh::Mesh

  function SystemMatrix(mesh)

    nnd = nodecount(mesh)
    rowptr = mesh.rowptr
    colidx = mesh.colidx
    matptr = @cxxnew RCompRowMatrix(nnd, nnd, pointer(rowptr), pointer(colidx))

    sysmat = new(matptr, mesh)

    finalizer(_sysmat_delete, sysmat)

    return sysmat

  end

end

ncols(S::SystemMatrix) = @cxx S.ptr->nCols()
nrows(S::SystemMatrix) = @cxx S.ptr->nRows()
valptr(S::SystemMatrix) = @cxx S.ptr->ValPtr()

# Delete underlying pointer to a Toast++ system matrix
# TODO: Check if Cxx.jl does this for me
_sysmat_delete(sysmat::SystemMatrix) = finalize(sysmat.ptr) #icxx"""delete $(sysmat.ptr);"""

function sparse(sysmat::SystemMatrix)
  nnd = nodecount(sysmat.mesh)
  rowptr = sysmat.mesh.rowptr
  colidx = sysmat.mesh.colidx
  _CSR0_to_CSC1(nnd, nnd, rowptr, colidx, valptr(sysmat))
end


text:
module TruncatedNormal
    using Random, Statistics
    using ChainRulesCore, SpecialFunctions

    include("tnmean.jl")
    include("tnmom2.jl")
    include("tnvar.jl")
    include("tnmom1c.jl")
    include("tnmom2c.jl")
    include("util.jl")

    include("truncnorm.jl")
    include("rejection.jl")
end


text:
function snap_1_d(filename::String, data::Dict{Any,Any})

    f = open(filename)

    seek(f, 264)

    N = sum(data["Header"]["npart"])
    skipsize = read(f, Int32)
    bit_size = skipsize/(3*N)

    if Int(bit_size) == 4

    elseif Int(bit_size) == 8

    else
        println("read error! neither 32 nor 64 bits data!")
        return -1
    end

    # read positions
    for i in 1:length(data["Header"]["PartTypes"])

        if data["Header"]["npart"][i] != 0

            data[data["Header"]["PartTypes"][i]] = Dict()
            N = Int64(data["Header"]["npart"][i])
            #dummy = read(f, Float32, (3,N))

            if Int(bit_size) == 4
                data[data["Header"]["PartTypes"][i]]["POS"] = copy(transpose(read!(f, Array{Float32,2}(undef,(3,N)))))
            else
                data[data["Header"]["PartTypes"][i]]["POS"] = copy(transpose(read!(f, Array{Float64,2}(undef,(3,N)))))
            end

        end

    end

    p = position(f)

    # skip identifiers
    seek(f, p+8)

    # read Velocities
    for i in 1:length(data["Header"]["PartTypes"])

        if data["Header"]["npart"][i] != 0

            N = Int64(data["Header"]["npart"][i])

            if Int(bit_size) == 4
                data[data["Header"]["PartTypes"][i]]["VEL"] = copy(transpose(read!(f, Array{Float32,2}(undef,(3,N)))))
            else
                data[data["Header"]["PartTypes"][i]]["VEL"] = copy(transpose(read!(f, Array{Float64,2}(undef,(3,N)))))
            end

        end

    end

    #dummy = 0
    #gc()

    p = position(f)

    # skip identifiers
    seek(f, p+8)

    # Read IDs
    for i in 1:length(data["Header"]["PartTypes"])

        if data["Header"]["npart"][i] != 0

            N = Int64(data["Header"]["npart"][i])
            data[data["Header"]["PartTypes"][i]]["ID"] = copy(transpose(read!(f, Array{UInt32,2}(undef,(1,N)))))

        end

    end

    p = position(f)

    # skip identifiers
    seek(f, p+8)

    # Read IDs
    for i in 1:length(data["Header"]["PartTypes"])

        if data["Header"]["npart"][i] != 0

            if data["Header"]["massarr"][i] == 0

                N = Int64(data["Header"]["npart"][i])

                if Int(bit_size) == 4
                    data[data["Header"]["PartTypes"][i]]["MASS"] = copy(transpose(read!(f, Array{Float32,2}(undef,(1,N)))))
                else
                    data[data["Header"]["PartTypes"][i]]["MASS"] = copy(transpose(read!(f, Array{Float64,2}(undef,(1,N)))))
                end

            else

                N = Int64(data["Header"]["npart"][i])
                data[data["Header"]["PartTypes"][i]]["MASS"] = zeros(N)
                data[data["Header"]["PartTypes"][i]]["MASS"] .= data["Header"]["massarr"][i]

            end

        end

    end

    if data["Header"]["npart"][1] != 0

        # skip identifiers
        p = position(f)
        seek(f, p+8)

        N = Int64(data["Header"]["npart"][1])

        # read U
        if Int(bit_size) == 4
            data["PartType0"]["InternalEnergy"] = copy(transpose(read!(f, Array{Float32,2}(undef,(1,N)))))
        else
            data["PartType0"]["InternalEnergy"] = copy(transpose(read!(f, Array{Float64,2}(undef,(1,N)))))
        end

        # skip identifiers
        p = position(f)
        seek(f, p+8)

        # Read Density
        if Int(bit_size) == 4
            data["PartType0"]["Density"] = copy(transpose(read!(f, Array{Float32,2}(undef,(1,N)))))
        else
            data["PartType0"]["Density"] = copy(transpose(read!(f, Array{Float64,2}(undef,(1,N)))))
        end

        # skip identifiers
        p = position(f)
        seek(f, p+8)

        # Read SmoothingLength
        if Int(bit_size) == 4
            data["PartType0"]["SmoothingLength"] = copy(transpose(read!(f, Array{Float32,2}(undef,(1,N)))))
        else
            data["PartType0"]["SmoothingLength"] = copy(transpose(read!(f, Array{Float64,2}(undef,(1,N)))))
        end

    end

    close(f)

    return data

end

max_stars_repo_name: LudwigBoess/GadJet

max_issues_repo_name: LudwigBoess/GadJet

max_forks_repo_name: LudwigBoess/GadJet


text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

module TrajSaveMCBest

export MCBestStudy, MCBestStudyResults

using AdaptiveStressTesting
using SISLES: GenerativeModel

using CPUTime
using RLESUtils, Obj2Dict, RunCases, Observers, Loggers, Obj2DataFrames

using ..Config_ACASX_GM
using ..ConfigAST
using ..ConfigMCBest
using ..DefineSave
using ..TrajSaveCommon
using ..DefineLog
using ..SaveTypes
using ..SaveHelpers
using ..PostProcess

import ..DefineSave.trajSave
import ..TrajSaveCommon: get_study_params, get_study_results

type MCBestStudy
  fileroot::String
  studytype::String  #timed or samples
  trial::Int64
end
function MCBestStudy(;
                     fileroot::AbstractString="trajSaveMCBest",
                     studytype::AbstractString="samples",
                     trial::Int64=0);
  MCBestStudy(fileroot, studytype, trial)
end

type MCBestStudyResults
  rewards::Vector{Float64}  #vector of all the rewards
end

MCBestStudyResults() = MCBestStudyResults(0, Float64[])

function trajSave(study_params::MCBestStudy,
                  cases::Cases = Cases(Case());
                  outdir::AbstractString = "./",
                  postproc::PostProcessing=StandardPostProc(),
                  print_rate::Int64=1000)

  println("Starting MCBest Study...")
  pmap(case -> begin
         starttime_us = CPUtime_us()
         startnow = string(now())

         sim_params = extract_params!(defineSimParams(), case, "sim_params")
         ast_params = extract_params!(defineASTParams(), case, "ast_params")
         mcbest_params = extract_params!(defineMCBestParams(), case, "mcbest_params")
         study_params = extract_params!(study_params, case, "study_params")

         sim = defineSim(sim_params)
         ast = defineAST(sim, ast_params)

         if study_params.studytype == "timed"
           results = sample_timed(ast, mcbest_params.maxtime_s; print_rate=print_rate)
         elseif study_params.studytype == "samples"
           results = sample(ast, mcbest_params.n; print_rate=print_rate)
         end
         rewards = map(x -> x[1], results)

         study_results = MCBestStudyResults(rewards)
         index = indmax(rewards)
         reward, action_seq = results[index]

         compute_info = ComputeInfo(startnow,
                                    string(now()),
                                    gethostname(),
                                    (CPUtime_us() - starttime_us) / 1e6)

         fileroot_ = "$(study_params.fileroot)_$(sim.string_id)"
         outfileroot = joinpath(outdir, fileroot_)
         log = trajLoggedPlay(ast, reward, action_seq, compute_info, sim_params, ast_params)

         set_run_type!(log, "MCBest")
         set_mcbest_params!(log, mcbest_params)
         set_study_params!(log, study_params)
         set_study_results!(log, study_results)

         outfile = trajSave(outfileroot, log)

         #callback for postprocessing
         postprocess(outfile, postproc)

         reward
       end,

       cases)
end

function get_study_params(d::TrajLog, ::Type{Val{:MCBest}})
    study = MCBestStudy()
    Obj2DataFrames.set!(study, ObjDataFrame(d["study_params"]))
end
function get_study_results(d::TrajLog, ::Type{Val{:MCBest}})
    result = MCBestStudyResults()
    Obj2DataFrames.set!(result, ObjDataFrame(d["study_results"]))
end

end #module

max_stars_repo_name: sisl/RLESCAS

max_issues_repo_name: sisl/RLESCAS

max_forks_repo_name: sisl/RLESCAS


text:
module InteratomicPotentials

using AtomsBase
using LinearAlgebra
using StaticArrays
using Unitful
using UnitfulAtomic

using Distances
using NearestNeighbors

include("unit_convention.jl")
include("constants.jl")
include("nnlist.jl")
include("api.jl")
include("types.jl")

end


text:
# from Float64 to Posit8,16,32
Posit8(x::Float64) = ccall((:convertDoubleToP8, SoftPositPath), Posit8, (Float64,),x)
Posit16_old(x::Float64) = ccall((:convertDoubleToP16, SoftPositPath), Posit16, (Float64,),x)
Posit32(x::Float64) = ccall((:convertDoubleToP32, SoftPositPath), Posit32, (Float64,),x)

# from Float32/16 to Posit8,16,32
Posit8(x::T where {T <: Float16or32}) = Posit8(Float64(x))
Posit32(x::T where {T <: Float16or32}) = Posit32(Float64(x))

Posit16(x::Float16) = Posit16(Float32(x))
Posit16(x::Float64) = Posit16(Float32(x))

function Posit16(x::Float32)
    ui = reinterpret(UInt32,x)

    # REGIME AND EXPONENT BITS
    # exponent without the exponent bias 127
    e = reinterpret(Int32,(ui & 0x7f80_0000) >> 23) - Int32(127)
    abs_e = abs(e)                          # exponent without sign
    signbit_e = signbit(e)                  # sign of exponent
    n_regimebits = (abs_e >> 1) + 1         # number of regime bits
    exponent_bit = abs_e & 0x1              # exponent bit from exponent odd = 3,5,...?

    # combine regime bit and exponent and then arithmetic bitshift them in for e.g. 111110_e_....
    regime_exponent = reinterpret(Int32,0x8000_0000 | (exponent_bit << 29)) >> n_regimebits

    # for x < 1 use two's complement to the regime and exponent bits to flip them correctly
    regime_exponent = signbit_e ? -regime_exponent : regime_exponent
    regime_exponent &= 0x7fff_ffff          # remove any sign that results from arith bitshift

    # MANTISSA - isolate mantissa bits and shift in position
    mantissa = (ui & 0x007f_ffff) >> (n_regimebits-6+signbit_e*(exponent_bit-1))
    # add u/2 = 0x0000_7fff or 0x0000_8000 for tie to even
    u_half = 0x0000_7fff + ((mantissa >> 16) & 0x0000_0001)

    # combine regime, exponent and mantissa, round to nearest, tie to even
    p32 = (regime_exponent | mantissa) + u_half
    p16 = ((p32 >> 16) % UInt16) - (15 < n_regimebits < 64)    # after +u_half round down via >>

    # check for sign bit and apply two's complement for negative numbers
    p16 = signbit(x) ? -p16 : p16
    return reinterpret(Posit16,p16)
end

# legacy
Posit16_old(x::Float32) = Posit16_old(Float64(x))
Posit16_old(x::Float16) = Posit16_old(Float64(x))

# from Float64 to Posit_1 (convertDoubleToPX1 currently not available)
# Posit8_1(x::Float64) = ccall((:convertDoubleToPX1, SoftPositPath), Posit8_1, (Float64,Int64),x,8)
# Posit16_1(x::Float64) = ccall((:convertDoubleToPX1, SoftPositPath), Posit16_1, (Float64,Int64),x,16)
# Posit24_1(x::Float64) = ccall((:convertDoubleToPX1, SoftPositPath), Posit24_1, (Float64,Int64),x,24)

# use detour
Posit8_1(x::Float64) = Posit8_1(Posit32(x))
Posit16_1(x::Float64) = Posit16_1(Posit32(x))
Posit24_1(x::Float64) = Posit24_1(Posit32(x))

# from Float32/16 to Posit_1
Posit8_1(x::T where {T <: Float16or32}) = Posit8_1(Float64(x))
Posit16_1(x::T where {T <: Float16or32}) = Posit16_1(Float64(x))
Posit24_1(x::T where {T <: Float16or32}) = Posit24_1(Float64(x))

# from Float64 to Posit_2
Posit8_2(x::Float64) = ccall((:convertDoubleToPX2, SoftPositPath), Posit8_2, (Float64,Int64),x,8)
Posit16_2(x::Float64) = ccall((:convertDoubleToPX2, SoftPositPath), Posit16_2, (Float64,Int64),x,16)
Posit24_2(x::Float64) = ccall((:convertDoubleToPX2, SoftPositPath), Posit24_2, (Float64,Int64),x,24)

# from Float32/16 to Posit_2
Posit8_2(x::T where {T <: Float16or32}) = Posit8_2(Float64(x))
Posit16_2(x::T where {T <: Float16or32}) = Posit16_2(Float64(x))
Posit24_2(x::T where {T <: Float16or32}) = Posit24_2(Float64(x))


text:
primitive type ARGB 32 end

ARGB(c::UInt32) = reinterpret(ARGB, c)
ARGB(c) = ARGB(UInt32(c))
ARGB(α::UInt8, red::UInt8, green::UInt8, blue::UInt8) =
    ARGB(UInt32(α) << 24 + UInt32(red) << 16 +
         UInt32(green) << 8 + UInt32(blue))
ARGB(α, red, green, blue) = ARGB(UInt8(α), UInt8(red),
                                 UInt8(green), UInt8(blue))

function ARGB(c::AbstractString)
   if !occursin(r"^#[0-9a-fA-F]{8}$", c)
       throw(DomainError("wrong color string: $c"))
   end
   ARGB(parse(UInt32, c[2:end], base=16))
end

macro ARGB_str(s) ARGB(s) end

α(c::ARGB)::UInt8 = (UInt32(c) >> 24) & 0x000000FF
red(c::ARGB)::UInt8 = (UInt32(c) >> 16) & 0x000000FF
green(c::ARGB)::UInt8 = (UInt32(c) >> 8) & 0x000000FF
blue(c::ARGB)::UInt8 = UInt32(c) & 0x000000FF

Base.UInt32(c::ARGB) = reinterpret(UInt32, c)
convert(UInt32, c::ARGB) = UInt32(c)
convert(ARGB, c::UInt32) = ARGB(c)
Base.String(c::ARGB) = "#" * lpad(string(UInt32(c), base=16), 8, "0")
convert(String, c::ARGB) = String(c)
convert(ARGB, c::AbstractString) = ARGB(c)

max_stars_repo_path: Chapter06/04. Defining primitive types/argb.jl


text:
function _precompile_()
    ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
    precompile(Core.Inference.indexed_next, (Tuple{Void, TypeConstructor}, Int32, Int32,))
    precompile(Core.Inference.indexed_next, (Tuple{Void, Union}, Int32, Int32,))
    precompile(Core.Inference.start, (Tuple{Void, Union},))
    precompile(Core.Inference.gensym_increment, (Float64, Int32,))
    precompile(Core.Inference.eltype, (Type{Array{Expr, 1}},))
    precompile(Core.Inference.start, (Tuple{Void, TypeConstructor},))
end

max_stars_repo_name: waTeim/gpsFlow

max_issues_repo_name: waTeim/flying-monkey

max_forks_repo_name: waTeim/flying-monkey


text:
function distributeZ2(D; p = 0.5)
    D0 = ceil(Int, p*D)
    D1 = D - D0
    return [(0, D0), (1, D1)]
end

distributeU1(D; p = 0.25) = distributeU1_poisson(D; p = p)

function distributeU1_exponential(D; p = 0.25)
    λ = (1-p)/(1+p)
    D0 = ceil(Int, p*D)
    if isodd(D-D0)
        D0 = D0 == 1 ? 2 : D0-1
    end
    sectors = [(0, D0)]
    Drem = D - D0
    n = 1
    while Drem > 0
        pn = p * λ^n
        Dn = ceil(Int, pn*D)
        sectors = push!(sectors, (n, Dn), (-n, Dn))
        Drem -= 2*Dn
        n += 1
    end
    return sort!(sectors, by = first)
end

function distributeU1_poisson(D; p = 0.25)
    λ = log((1/p+1)/2)
    D0 = ceil(Int, p*D)
    if isodd(D-D0)
        D0 = D0 == 1 ? 2 : D0-1
    end
    sectors = [(0, D0)]
    Drem = D - D0
    n = 1
    while Drem > 0
        pn = p * λ^n / factorial(n)
        Dn = ceil(Int, pn*D)
        sectors = push!(sectors, (n, Dn), (-n, Dn))
        Drem -= 2*Dn
        n += 1
    end
    return sort!(sectors, by = first)
end

module Timers
    struct Timer{F,D<:Tuple}
        f::F
        argsref::Base.RefValue{D}
        Timer(f, args...) = new{typeof(f), typeof(args)}(f, Ref(args))
    end

    @noinline donothing(arg) = arg
    function (t::Timer)(; inner = 1, outer = 1)
        args = t.argsref[]
        f = t.f
        f(args...) # run once to compile
        times = zeros(Float64, (outer,))
        gctimes = zeros(Float64, (outer,))
        @inbounds for i = 1:outer
            if inner == 1
                gcstart = Base.gc_num()
                start = Base.time_ns()
                donothing(f(args...))
                stop = Base.time_ns()
                gcstop = Base.gc_num()
            else
                gcstart = Base.gc_num()
                start = Base.time_ns()
                for _ = 1:inner
                    donothing(f(args...))
                end
                stop = Base.time_ns()
                gcstop = Base.gc_num()
            end
            times[i] = (stop - start)/(1e9)/inner
            gctimes[i] = Base.GC_Diff(gcstop, gcstart).total_time/1e9/inner
        end
        return times, gctimes
    end
end


text:
# # Turbulent mixing of a three layer boundary layer driven by constant surface fluxes

# This script runs a simulation of a turbulent oceanic boundary layer with an initial
# three-layer temperature stratification. Turbulent mixing is driven by constant fluxes
# of momentum and heat at the surface.
#
# This script is set up to be configurable on the command line --- a useful property
# when launching multiple jobs at on a cluster.

using Pkg
using Statistics
using Printf
using Logging

using ArgParse
using JLD2
using NCDatasets
using GeoData
using CairoMakie

using LESbrary
using Oceanostics
using Oceananigans
using Oceananigans.Grids
using Oceananigans.BuoyancyModels
using Oceananigans.BoundaryConditions
using Oceananigans.Fields
using Oceananigans.Advection
using Oceananigans.OutputWriters
using Oceananigans.Utils
using Oceananigans.Units

using Oceananigans.Grids: Face, Center
using Oceananigans.Fields: PressureField
using Oceanostics.FlowDiagnostics: richardson_number_ccf!
using Oceanostics.TurbulentKineticEnergyTerms: TurbulentKineticEnergy, ShearProduction_z

using LESbrary.Utils: SimulationProgressMessenger, fit_cubic, poly
using LESbrary.NearSurfaceTurbulenceModels: SurfaceEnhancedModelConstant
using LESbrary.TurbulenceStatistics: first_through_second_order, turbulent_kinetic_energy_budget,
                                     subfilter_momentum_fluxes, subfilter_tracer_fluxes,
                                     ViscousDissipation

Logging.global_logger(OceananigansLogger())

# To start, we ensure that all packages in the LESbrary environment are installed:

Pkg.instantiate()

# Next, we parse the command line arguments

"Returns a dictionary of command line arguments."
function parse_command_line_arguments()
    settings = ArgParseSettings()

    @add_arg_table! settings begin
        "--size"
            help = "The number of grid points in x, y, and z."
            nargs = 3
            default = [32, 32, 32]
            arg_type = Int

        "--extent"
            help = "The length of the x, y, and z dimensions."
            nargs = 3
            default = [512meters, 512meters, 256meters]
            arg_type = Float64

        "--buoyancy-flux"
            help = """The surface buoyancy flux in units of m² s⁻³.
                      A positive buoyancy flux implies cooling.

                      Note:
                          buoyancy-flux = + 1e-7 corresponds to cooling at 208 W / m²
                          buoyancy-flux = + 1e-8 corresponds to cooling at 21 W / m²
                          buoyancy-flux = - 1e-7 corresponds to heating at 208 W / m²"""
            default = 1e-8
            arg_type = Float64

        "--momentum-flux"
            help = """The surface x-momentum flux divided by density in units of m² s⁻².
                      A negative flux drives currents in the positive x-direction.

                      Note:
                        momentum-flux = - 1e-4 corresponds to U₁₀ = +6 m/s, roughly speaking
                        momentum-flux = - 1e-3 corresponds to U₁₀ = +19 m/s, roughly speaking"""
            default = -1e-4
            arg_type = Float64

        "--coriolis"
            help = "The Coriolis parameter, calculated as f=2Ω*sinϕ, where Ω is the rotation rate of the Earth in rad/s and ϕ is the latitude."
            default = 1e-4
            arg_type = Float64

        "--surface-temperature"
            help = """The temperature at the surface in ᵒC."""
            default = 20.0
            arg_type = Float64

        "--surface-layer-depth"
            help = "The depth of the surface layer in units of m."
            default = 48.0
            arg_type = Float64

        "--thermocline"
            help = """Two choices for the thermocline structure:
                        * linear: a thermocline with a linear buoyancy structure (constant stratification)
                        * cubic: a thermocline with a fitted cubic structure"""
            default = "linear"
            arg_type = String

        "--thermocline-width"
            help = "The width of the thermocline in units of m."
            default = 24.0
            arg_type = Float64

        "--surface-layer-buoyancy-gradient"
            help = "The buoyancy gradient in the surface layer in units s⁻²."
            default = 2e-6
            arg_type = Float64

        "--thermocline-buoyancy-gradient"
            help = "The buoyancy gradient in the thermocline in units s⁻²."
            default = 1e-5
            arg_type = Float64

        "--deep-buoyancy-gradient"
            help = "The buoyancy gradient below the thermocline in units s⁻²."
            default = 2e-6
            arg_type = Float64

        "--hours"
            help = "Number of hours to run the simulation for"
            default = 0.1
            arg_type = Float64

        "--time-averaged-statistics"
            help = "Compute and output time-averaged statistics."
            action = :store_true

        "--animation"
            help = "Make an animation of the horizontal and vertical velocity when the simulation completes."
            action = :store_true

        "--pickup"
            help = "Whether or not to pick the simulation up from latest checkpoint"
            default = false
            arg_type = Bool

        "--name"
            help = "A name to add to the end of the output filename, e.g. weak_wind_strong_cooling."
            default = ""
            arg_type = String
    end

    return parse_args(settings)
end

# # Setup
#
# We start by parsing the arguments received on the command line, and prescribing
# a few additional output-related arguments.

@info "Parsing command line arguments..."

args = parse_command_line_arguments()

Nx, Ny, Nz = args["size"]
Lx, Ly, Lz = args["extent"]
stop_hours = args["hours"]
f = args["coriolis"]
name = args["name"]

snapshot_time_interval = 10minutes
averages_time_interval = 3hours
averages_time_window = 15minutes

slice_depth = 8.0

## Determine filepath prefix

Qᵇ = args["buoyancy-flux"]
Qᵘ = args["momentum-flux"]

thermocline_type = args["thermocline"]

prefix = @sprintf("three_layer_constant_fluxes_%s_hr%d_Qu%.1e_Qb%.1e_f%.1e_Nh%d_Nz%d_",
                  thermocline_type, stop_hours, abs(Qᵘ), Qᵇ, f, Nx, Nz)
data_directory = joinpath(@__DIR__, "..", "data", prefix * name) # save data in /data/prefix

# Copy this file into the directory with data
mkpath(data_directory)
cp(@__FILE__, joinpath(data_directory, basename(@__FILE__)), force=true)

# Save command line arguments used to an executable bash script
open(joinpath(data_directory, "run_three_layer_constant_fluxes.sh"), "w") do io
    write(io, "#!/bin/sh\n")
    write(io, "julia " * basename(@__FILE__) * " " * join(ARGS, " ") * "\n")
end

# Domain
#
# We use a three-dimensional domain that's twice as wide as it is deep.
# We choose this aspect ratio so that the horizontal scale is 4x larger
# than the boundary layer depth when the boundary layer penetrates half
# the domain depth.

@info "Mapping grid..."

grid = RegularRectilinearGrid(size=(Nx, Ny, Nz), x=(0, Lx), y=(0, Ly), z=(-Lz, 0))

# Buoyancy and boundary conditions

@info "Enforcing boundary conditions..."

surface_layer_depth = args["surface-layer-depth"]
thermocline_width = args["thermocline-width"]

N²_surface_layer = args["surface-layer-buoyancy-gradient"]
N²_thermocline = args["thermocline-buoyancy-gradient"]
N²_deep = args["deep-buoyancy-gradient"]

buoyancy = SeawaterBuoyancy(equation_of_state=LinearEquationOfState(α=2e-4), constant_salinity=35.0)

α = buoyancy.equation_of_state.α
g = buoyancy.gravitational_acceleration

Qᶿ = Qᵇ / (α * g)
dθdz_surface_layer = N²_surface_layer / (α * g)
dθdz_thermocline   = N²_thermocline   / (α * g)
dθdz_deep          = N²_deep          / (α * g)

θ_bcs = TracerBoundaryConditions(grid, top = BoundaryCondition(Flux, Qᶿ),
                                       bottom = BoundaryCondition(Gradient, dθdz_deep))

u_bcs = UVelocityBoundaryConditions(grid, top = BoundaryCondition(Flux, Qᵘ))

# Tracer forcing

@info "Forcing and sponging tracers..."

# # Initial condition and sponge layer

## Fiddle with indices to get a correct discrete profile
k_transition = searchsortedfirst(grid.zC, -surface_layer_depth)
k_deep = searchsortedfirst(grid.zC, -(surface_layer_depth + thermocline_width))

z_transition = grid.zC[k_transition]
z_deep = grid.zC[k_deep]

θ_surface = args["surface-temperature"]
θ_transition = θ_surface + z_transition * dθdz_surface_layer
θ_deep = θ_transition + (z_deep - z_transition) * dθdz_thermocline

@inline passive_tracer_forcing(x, y, z, t, p) = p.μ⁺ * exp(-(z - p.z₀)^2 / (2 * p.λ^2)) - p.μ⁻

λ = 4.0
μ⁺ = 1 / 6hour
μ₀ = √(2π) * λ / grid.Lz * μ⁺ / 2
μ∞ = √(2π) * λ / grid.Lz * μ⁺

c₀_forcing = Forcing(passive_tracer_forcing, parameters=(z₀=  0.0, λ=λ, μ⁺=μ⁺, μ⁻=μ₀))
c₁_forcing = Forcing(passive_tracer_forcing, parameters=(z₀=-48.0, λ=λ, μ⁺=μ⁺, μ⁻=μ∞))
c₂_forcing = Forcing(passive_tracer_forcing, parameters=(z₀=-96.0, λ=λ, μ⁺=μ⁺, μ⁻=μ∞))

# Sponge layer for u, v, w, and T
gaussian_mask = GaussianMask{:z}(center=-grid.Lz, width=grid.Lz/10)

u_sponge = v_sponge = w_sponge = Relaxation(rate=4/hour, mask=gaussian_mask)

T_sponge = Relaxation(rate = 4/hour,
                      target = LinearTarget{:z}(intercept = θ_deep - z_deep*dθdz_deep, gradient = dθdz_deep),
                      mask = gaussian_mask)

# # LES Model

# Wall-aware AMD model constant which is 'enhanced' near the upper boundary.
# Necessary to obtain a smooth temperature distribution.

@info "Building the wall model..."

Cᴬᴹᴰ = SurfaceEnhancedModelConstant(grid.Δz, C₀ = 1/12, enhancement = 7, decay_scale = 4 * grid.Δz)

# # Instantiate Oceananigans.IncompressibleModel

@info "Framing the model..."

model = IncompressibleModel(
           architecture = GPU(),
            timestepper = :RungeKutta3,
              advection = WENO5(),
                   grid = grid,
                tracers = (:T, :c₀, :c₁, :c₂),
               buoyancy = buoyancy,
               coriolis = FPlane(f=f),
                closure = AnisotropicMinimumDissipation(),
    boundary_conditions = (T=θ_bcs, u=u_bcs),
                forcing = (u=u_sponge, v=v_sponge, w=w_sponge, T=T_sponge,
                           c₀=c₀_forcing, c₁=c₁_forcing, c₂=c₂_forcing)
)

# # Set Initial condition

@info "Setting initial conditions..."

## Noise with 8 m decay scale
Ξ(z) = rand() * exp(z / 8)

function thermocline_structure_function(thermocline_type, z_transition, θ_transition, z_deep, θ_deep, dθdz_surface_layer, dθdz_thermocline, dθdz_deep)
    if thermocline_type == "linear"
        return z -> θ_transition + dθdz_thermocline * (z - z_transition)

    elseif thermocline_type == "cubic"
        p1 = (z_transition, θ_transition)
        p2 = (z_deep, θ_deep)
        coeffs = fit_cubic(p1, p2, dθdz_surface_layer, dθdz_deep)
        return z -> poly(z, coeffs)

    else
        @error "Invalid thermocline type: $thermocline"
    end
end

θ_thermocline = thermocline_structure_function(thermocline_type, z_transition, θ_transition, z_deep, θ_deep,
                                               dθdz_surface_layer, dθdz_thermocline, dθdz_deep)

"""
    initial_temperature(x, y, z)

Returns a three-layer initial temperature distribution. The average temperature varies in z
and is augmented by three-dimensional, surface-concentrated random noise.
"""
function initial_temperature(x, y, z)

    noise = 1e-6 * Ξ(z) * dθdz_surface_layer * grid.Lz

    if z_transition < z <= 0
        return θ_surface + dθdz_surface_layer * z + noise

    elseif z_deep < z <= z_transition
        return θ_thermocline(z) + noise

    else
        return θ_deep + dθdz_deep * (z - z_deep) + noise

    end
end

set!(model, T = initial_temperature)

# # Prepare the simulation

@info "Conjuring the simulation..."

# Adaptive time-stepping
wizard = TimeStepWizard(cfl=0.8, Δt=1.0, max_change=1.1, min_Δt=0.01, max_Δt=30.0)

stop_time = stop_hours * hour

simulation = Simulation(model,
                    Δt = wizard,
             stop_time = stop_time,
    iteration_interval = 10,
              progress = SimulationProgressMessenger(wizard)
)

# # Prepare Output

@info "Strapping on checkpointer..."

pickup = args["pickup"]
force = pickup ? false : true

simulation.output_writers[:checkpointer] =
    Checkpointer(model, schedule = TimeInterval(stop_time/3), prefix = prefix * "_checkpointer", dir = data_directory)

@info "Squeezing out statistics..."

# Prepare turbulence statistics
k_xy_slice = searchsortedfirst(grid.zF[:], -slice_depth)

b = BuoyancyField(model)
p = PressureField(model)

ccc_scratch = Field(Center, Center, Center, model.architecture, model.grid)
ccf_scratch = Field(Center, Center, Face, model.architecture, model.grid)
fcf_scratch = Field(Face, Center, Face, model.architecture, model.grid)
cff_scratch = Field(Center, Face, Face, model.architecture, model.grid)

primitive_statistics = first_through_second_order(model, b=b, p=p, w_scratch=ccf_scratch, c_scratch=ccc_scratch)

subfilter_flux_statistics = merge(
    subfilter_momentum_fluxes(model, uz_scratch=ccf_scratch, vz_scratch=cff_scratch, c_scratch=ccc_scratch),
    subfilter_tracer_fluxes(model, w_scratch=ccf_scratch),
)

U = primitive_statistics[:u]
V = primitive_statistics[:v]
B = primitive_statistics[:b]

e = TurbulentKineticEnergy(model, U=U, V=V)
shear_production = ShearProduction_z(model, U=U, V=V)
dissipation = ViscousDissipation(model)

tke_budget_statistics = turbulent_kinetic_energy_budget(model, b=b, p=p, U=U, V=V, e=e,
                                                        shear_production=shear_production, dissipation=dissipation)

# FIXME: This 3D kernel actually wastes a lot of computation since we just need a 1D kernel.
# See: https://github.com/CliMA/LESbrary.jl/issues/114
Ri_kcf = KernelComputedField(Center, Center, Face, richardson_number_ccf!, model,
                               computed_dependencies=(U, V, B), parameters=(dUdz_bg=0, dVdz_bg=0, N2_bg=0))

Ri = AveragedField(Ri_kcf, dims=(1, 2))

dynamics_statistics = Dict(:Ri => Ri)

fields_to_output = merge(model.velocities, model.tracers, (e=e, ϵ=dissipation))

statistics_to_output = merge(primitive_statistics, subfilter_flux_statistics, tke_budget_statistics, dynamics_statistics)

@info "Garnishing output writers..."

# Code credit: https://discourse.julialang.org/t/collecting-all-output-from-shell-commands/15592
function execute(cmd::Cmd)
    out, err = Pipe(), Pipe()

    process = run(pipeline(ignorestatus(cmd), stdout=out, stderr=err))
    close(out.in)
    close(err.in)

    return (stdout = out |> read |> String, stderr = err |> read |> String, code = process.exitcode)
end

global_attributes = (
    LESbrary_jl_commit_SHA1 = execute(`git rev-parse HEAD`).stdout |> strip,
    name = name,
    thermocline_type = thermocline_type,
    buoyancy_flux = Qᵇ,
    momentum_flux = Qᵘ,
    temperature_flux = Qᶿ,
    coriolis_parameter = f,
    thermal_expansion_coefficient = α,
    gravitational_acceleration = g,
    boundary_condition_θ_top = Qᶿ,
    boundary_condition_θ_bottom = dθdz_deep,
    boundary_condition_u_top = Qᵘ,
    boundary_condition_u_bottom = 0.0,
    surface_layer_depth = surface_layer_depth,
    thermocline_width = thermocline_width,
    N²_surface_layer = N²_surface_layer,
    N²_thermocline = N²_thermocline,
    N²_deep = N²_deep,
    dθdz_surface_layer = dθdz_surface_layer,
    dθdz_thermocline = dθdz_thermocline,
    dθdz_deep = dθdz_deep,
    θ_surface = θ_surface,
    θ_transition = θ_transition,
    θ_deep = θ_deep,
    z_transition = z_transition,
    z_deep = z_deep,
    k_transition = k_transition,
    k_deep = k_deep
)

function init_save_some_metadata!(file, model)
    for (name, value) in pairs(global_attributes)
        file["parameters/$(string(name))"] = value
    end
    return nothing
end

## Add JLD2 output writers

name_with_halos(name, with_halos) = with_halos ? "$(name)_with_halos" : name

for with_halos in [true, false]

    simulation.output_writers[Symbol(name_with_halos("xy_jld2", with_halos))] =
        JLD2OutputWriter(model, fields_to_output,
                                  dir = data_directory,
                               prefix = name_with_halos("xy_slice", with_halos),
                             schedule = TimeInterval(snapshot_time_interval),
                         field_slicer = FieldSlicer(k=k_xy_slice, with_halos=with_halos),
                                force = force,
                                 init = init_save_some_metadata!)

    simulation.output_writers[Symbol(name_with_halos("xz_jld2", with_halos))] =
        JLD2OutputWriter(model, fields_to_output,
                                  dir = data_directory,
                               prefix = name_with_halos("xz_slice", with_halos),
                             schedule = TimeInterval(snapshot_time_interval),
                         field_slicer = FieldSlicer(j=1, with_halos=with_halos),
                                force = force,
                                 init = init_save_some_metadata!)

    simulation.output_writers[Symbol(name_with_halos("yz_jld2", with_halos))] =
        JLD2OutputWriter(model, fields_to_output,
                                  dir = data_directory,
                               prefix = name_with_halos("yz_slice", with_halos),
                             schedule = TimeInterval(snapshot_time_interval),
                         field_slicer = FieldSlicer(i=1, with_halos=with_halos),
                                force = force,
                                 init = init_save_some_metadata!)

    simulation.output_writers[Symbol(name_with_halos("stats_jld2", with_halos))] =
        JLD2OutputWriter(model, statistics_to_output,
                                  dir = data_directory,
                               prefix = name_with_halos("instantaneous_statistics", with_halos),
                             schedule = TimeInterval(snapshot_time_interval),
                         field_slicer = FieldSlicer(with_halos=with_halos),
                                force = force,
                                 init = init_save_some_metadata!)

    if args["time-averaged-statistics"]
        simulation.output_writers[Symbol(name_with_halos("averaged_stats_jld2", with_halos))] =
            JLD2OutputWriter(model, statistics_to_output,
                                      dir = data_directory,
                                   prefix = name_with_halos("time_averaged_statistics", with_halos),
                                 schedule = AveragedTimeInterval(averages_time_interval,
                                                                 window = averages_time_window),
                             field_slicer = FieldSlicer(with_halos=with_halos),
                                    force = force,
                                     init = init_save_some_metadata!)
    end
end

## Add NetCDF output writers

statistics_to_output = Dict(string(k) => v for (k, v) in statistics_to_output)

simulation.output_writers[:xy_nc] =
    NetCDFOutputWriter(model, fields_to_output,
                     mode = "c",
                 filepath = joinpath(data_directory, "xy_slice.nc"),
                 schedule = TimeInterval(snapshot_time_interval),
             field_slicer = FieldSlicer(k=k_xy_slice),
        global_attributes = global_attributes)

simulation.output_writers[:xz_nc] =
    NetCDFOutputWriter(model, fields_to_output,
                     mode = "c",
                 filepath = joinpath(data_directory, "xz_slice.nc"),
                 schedule = TimeInterval(snapshot_time_interval),
             field_slicer = FieldSlicer(j=1),
        global_attributes = global_attributes)

simulation.output_writers[:yz_nc] =
    NetCDFOutputWriter(model, fields_to_output,
                     mode = "c",
                 filepath = joinpath(data_directory, "yz_slice.nc"),
                 schedule = TimeInterval(snapshot_time_interval),
             field_slicer = FieldSlicer(i=1),
        global_attributes = global_attributes)

simulation.output_writers[:stats_nc] =
    NetCDFOutputWriter(model, statistics_to_output,
                     mode = "c",
                 filepath = joinpath(data_directory, "instantaneous_statistics.nc"),
                 schedule = TimeInterval(snapshot_time_interval),
        global_attributes = global_attributes)

if args["time-averaged-statistics"]
    simulation.output_writers[:averaged_stats_nc] =
        NetCDFOutputWriter(model, statistics_to_output,
                         mode = "c",
                     filepath = joinpath(data_directory, "time_averaged_statistics.nc"),
                     schedule = AveragedTimeInterval(averages_time_interval, window = averages_time_window),
            global_attributes = global_attributes)
end

# # Run

@info "Teaching simulation to run!..."

run!(simulation)

# # Load and plot turbulence statistics

make_animation = args["animation"]

function squeeze(A)
    singleton_dims = tuple((d for d in 1:ndims(A) if size(A, d) == 1)...)
    return dropdims(A, dims=singleton_dims)
end

if make_animation
    ds_xy = NCDstack(joinpath(data_directory, "xy_slice.nc"))
    ds_xz = NCDstack(joinpath(data_directory, "xz_slice.nc"))

    times = ds_xy[:time]
    Nt = length(times)

    xc = ds_xy[:xC].data
    xf = ds_xy[:xF].data
    yc = ds_xy[:yC].data
    yf = ds_xy[:yF].data
    zc = ds_xz[:zC].data
    zf = ds_xz[:zF].data

    fig = Figure(resolution=(3000, 750))

    u_max = max(maximum(abs, ds_xy[:u]), maximum(abs, ds_xz[:u]))
    v_max = max(maximum(abs, ds_xy[:v]), maximum(abs, ds_xz[:v]))
    w_max = max(maximum(abs, ds_xy[:w]), maximum(abs, ds_xz[:w]))
    U_max = max(u_max, v_max, w_max)
    U_lims = 0.5 .* (-U_max, +U_max)

    frame = Node(1)

    plot_title = @lift "LESbrary.jl three-layer constant fluxes $name: time = $(prettytime(times[$frame]))"

    u_xy = @lift ds_xy[:u][Ti=$frame].data |> squeeze
    v_xy = @lift ds_xy[:v][Ti=$frame].data |> squeeze
    w_xy = @lift ds_xy[:w][Ti=$frame].data |> squeeze
    T_xy = @lift ds_xy[:T][Ti=$frame].data |> squeeze
    c₀_xy = @lift ds_xy[:c₀][Ti=$frame].data |> squeeze
    c₁_xy = @lift ds_xy[:c₁][Ti=$frame].data |> squeeze
    c₂_xy = @lift ds_xy[:c₂][Ti=$frame].data |> squeeze
    e_xy = @lift ds_xy[:e][Ti=$frame].data |> squeeze .|> log10
    ϵ_xy = @lift ds_xy[:ϵ][Ti=$frame].data |> squeeze .|> log10

    u_xz = @lift ds_xz[:u][Ti=$frame].data |> squeeze
    v_xz = @lift ds_xz[:v][Ti=$frame].data |> squeeze
    w_xz = @lift ds_xz[:w][Ti=$frame].data |> squeeze
    T_xz = @lift ds_xz[:T][Ti=$frame].data |> squeeze
    c₀_xz = @lift ds_xz[:c₀][Ti=$frame].data |> squeeze
    c₁_xz = @lift ds_xz[:c₁][Ti=$frame].data |> squeeze
    c₂_xz = @lift ds_xz[:c₂][Ti=$frame].data |> squeeze
    e_xz = @lift ds_xz[:e][Ti=$frame].data |> squeeze .|> log10
    ϵ_xz = @lift ds_xz[:ϵ][Ti=$frame].data |> squeeze .|> log10

    ax_u_xy = fig[1, 1] = Axis(fig, title="u-velocity")
    hm_u_xy = heatmap!(ax_u_xy, xf, yc, u_xy, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_u_xy)

    ax_v_xy = fig[1, 2] = Axis(fig, title="v-velocity")
    hm_v_xy = heatmap!(ax_v_xy, xc, yf, v_xy, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_v_xy)

    ax_w_xy = fig[1, 3] = Axis(fig, title="w-velocity")
    hm_w_xy = heatmap!(ax_w_xy, xc, yc, w_xy, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_w_xy)

    ax_T_xy = fig[1, 4] = Axis(fig, title="temperature")
    hm_T_xy = heatmap!(ax_T_xy, xc, yc, T_xy, colormap=:thermal, colorrange=extrema(ds_xy[:T]))
    hidedecorations!(ax_T_xy)

    ax_c₀_xy = fig[1, 5] = Axis(fig, title="tracer 0")
    hm_c₀_xy = heatmap!(ax_c₀_xy, xc, yc, c₀_xy, colormap=:ice, colorrange=extrema(ds_xy[:c₀]))
    hidedecorations!(ax_c₀_xy)

    ax_c₁_xy = fig[1, 6] = Axis(fig, title="tracer 1")
    hm_c₁_xy = heatmap!(ax_c₁_xy, xc, yc, c₁_xy, colormap=:ice, colorrange=extrema(ds_xy[:c₁]))
    hidedecorations!(ax_c₁_xy)

    ax_c₂_xy = fig[1, 7] = Axis(fig, title="tracer 2")
    hm_c₂_xy = heatmap!(ax_c₂_xy, xc, yc, c₂_xy, colormap=:ice, colorrange=extrema(ds_xy[:c₂]))
    hidedecorations!(ax_c₂_xy)

    ax_e_xy = fig[1, 8] = Axis(fig, title="log TKE")
    hm_e_xy = heatmap!(ax_e_xy, xc, yc, e_xy, colormap=:deep, colorrange=(-5, -2))
    hidedecorations!(ax_e_xy)

    ax_ϵ_xy = fig[1, 9] = Axis(fig, title="log TKE dissipation")
    hm_ϵ_xy = heatmap!(ax_ϵ_xy, xc, yc, ϵ_xy, colormap=:dense, colorrange=(-10, -5))
    hidedecorations!(ax_ϵ_xy)

    ax_u_xz = fig[2, 1] = Axis(fig)
    hm_u_xz = heatmap!(ax_u_xz, xf, zc, u_xz, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_u_xz)

    ax_v_xz = fig[2, 2] = Axis(fig)
    hm_v_xz = heatmap!(ax_v_xz, xc, zc, v_xz, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_v_xz)

    ax_w_xz = fig[2, 3] = Axis(fig)
    hm_w_xz = heatmap!(ax_w_xz, xc, zf, w_xz, colormap=:balance, colorrange=U_lims)
    hidedecorations!(ax_w_xz)

    ax_T_xz = fig[2, 4] = Axis(fig)
    hm_T_xz = heatmap!(ax_T_xz, xc, zc, T_xz, colormap=:thermal, colorrange=extrema(ds_xz[:T]))
    hidedecorations!(ax_T_xz)

    ax_c₀_xz = fig[2, 5] = Axis(fig)
    hm_c₀_xz = heatmap!(ax_c₀_xz, xc, zc, c₀_xz, colormap=:ice, colorrange=extrema(ds_xz[:c₀]))
    hidedecorations!(ax_c₀_xz)

    ax_c₁_xz = fig[2, 6] = Axis(fig)
    hm_c₁_xz = heatmap!(ax_c₁_xz, xc, zc, c₁_xz, colormap=:ice, colorrange=extrema(ds_xz[:c₁]))
    hidedecorations!(ax_c₁_xz)

    ax_c₂_xz = fig[2, 7] = Axis(fig)
    hm_c₂_xz = heatmap!(ax_c₂_xz, xc, zc, c₂_xz, colormap=:ice, colorrange=extrema(ds_xz[:c₂]))
    hidedecorations!(ax_c₂_xz)

    ax_e_xz = fig[2, 8] = Axis(fig)
    hm_e_xz = heatmap!(ax_e_xz, xc, zc, e_xz, colormap=:deep, colorrange=(-8, -2))
    hidedecorations!(ax_e_xz)

    ax_ϵ_xz = fig[2, 9] = Axis(fig)
    hm_ϵ_xz = heatmap!(ax_ϵ_xz, xc, zc, ϵ_xz, colormap=:dense, colorrange=(-15, -5))
    hidedecorations!(ax_ϵ_xz)

    supertitle = fig[0, :] = Label(fig, plot_title, textsize=30)

    filepath = joinpath(data_directory, "xy_xz_movie.mp4")
    record(fig, filepath, 1:Nt, framerate=15) do n
        @info "Animating xy/xz movie frame $n/$Nt..."
        frame[] = n
    end

    @info "Movie saved: $filepath"
end

if make_animation
    ds = NCDstack(joinpath(data_directory, "instantaneous_statistics.nc"))

    times = ds[:time]
    Nt = length(times)

    zc = ds[:zC].data
    zf = ds[:zF].data

    fig = Figure(resolution=(3000, 2000))

    frame = Node(1)

    plot_title = @lift "LESbrary.jl three-layer constant fluxes $name: time = $(prettytime(times[$frame]))"

    u = @lift ds[:u][Ti=$frame]
    v = @lift ds[:v][Ti=$frame]
    T = @lift ds[:T][Ti=$frame]
    c₀ = @lift ds[:c₀][Ti=$frame]
    c₁ = @lift ds[:c₁][Ti=$frame]
    c₂ = @lift ds[:c₂][Ti=$frame]
    uu = @lift ds[:uu][Ti=$frame]
    vv = @lift ds[:vv][Ti=$frame]
    ww = @lift ds[:ww][Ti=$frame]
    uv = @lift ds[:uv][Ti=$frame]
    wu = @lift ds[:wu][Ti=$frame]
    wv = @lift ds[:wv][Ti=$frame]
    uT = @lift ds[:uT][Ti=$frame]
    vT = @lift ds[:vT][Ti=$frame]
    wT = @lift ds[:wT][Ti=$frame]
    Ri = @lift ds[:Ri][Ti=$frame]
    νₑ_∂z_u = @lift ds[:νₑ_∂z_u][Ti=$frame]
    νₑ_∂z_v = @lift ds[:νₑ_∂z_v][Ti=$frame]
    νₑ_∂z_w = @lift ds[:νₑ_∂z_w][Ti=$frame]
    κₑ_∂z_T = @lift ds[:κₑ_∂z_T][Ti=$frame]
    κₑ_∂z_c₀ = @lift ds[:κₑ_∂z_c₀][Ti=$frame]
    κₑ_∂z_c₁ = @lift ds[:κₑ_∂z_c₁][Ti=$frame]
    κₑ_∂z_c₂ = @lift ds[:κₑ_∂z_c₂][Ti=$frame]
    e = @lift ds[:e][Ti=$frame]
    tke_dissipation = @lift ds[:tke_dissipation][Ti=$frame]
    tke_advective_flux = @lift ds[:tke_advective_flux][Ti=$frame]
    tke_buoyancy_flux = @lift ds[:tke_buoyancy_flux][Ti=$frame]
    tke_pressure_flux = @lift ds[:tke_pressure_flux][Ti=$frame]
    tke_shear_production = @lift ds[:tke_shear_production][Ti=$frame]

    colors = ["dodgerblue2", "crimson", "forestgreen"]

    ax_uv = fig[1, 1] = Axis(fig, xlabel="m/s", ylabel="z (m)")
    line_u = lines!(ax_uv, u, zc, label="U", linewidth=3, color=colors[1])
    line_v = lines!(ax_uv, v, zc, label="V", linewidth=3, color=colors[2])
    axislegend(ax_uv, position=:rb, framevisible=false)
    xlims!(ax_uv, extrema([extrema(ds[:u])..., extrema(ds[:v])...]))
    ylims!(ax_uv, extrema(zf))

    ax_T = fig[1, 2] = Axis(fig, xlabel="°C")
    line_T = lines!(ax_T, T, zc, linewidth=3, color=colors[1])
    xlims!(ax_T, extrema(ds[:T]))
    ylims!(ax_T, extrema(zf))
    hideydecorations!(ax_T, grid=false)

    ax_c = fig[1, 3] = Axis(fig, xlabel="c")
    line_c₀ = lines!(ax_c, c₀, zc, linewidth=3, label="c₀", color=colors[1])
    line_c₁ = lines!(ax_c, c₁, zc, linewidth=3, label="c₁", color=colors[2])
    line_c₂ = lines!(ax_c, c₂, zc, linewidth=3, label="c₂", color=colors[3])
    axislegend(ax_c, position=:rb, framevisible=false)
    xlims!(ax_c, extrema([extrema(ds[:c₀])..., extrema(ds[:c₁])..., extrema(ds[:c₂])...]))
    ylims!(ax_c, extrema(zf))
    hideydecorations!(ax_c, grid=false)

    ax_ke = fig[1, 4] = Axis(fig, xlabel="m²/s²")
    line_uu = lines!(ax_ke, uu, zc, linewidth=3, label="uu", color=colors[1])
    line_vv = lines!(ax_ke, vv, zc, linewidth=3, label="vv", color=colors[2])
    line_ww = lines!(ax_ke, ww, zf, linewidth=3, label="ww", color=colors[3])
    axislegend(ax_ke, position=:rb, framevisible=false)
    xlims!(ax_ke, extrema([extrema(ds[:c₀])..., extrema(ds[:c₁])..., extrema(ds[:c₂])...]))
    ylims!(ax_ke, extrema(zf))
    hideydecorations!(ax_ke, grid=false)

    ax_UΣ = fig[1, 5] = Axis(fig, xlabel="m²/s²")
    line_uv = lines!(ax_UΣ, uv, zc, linewidth=3, label="uv", color=colors[1])
    line_wu = lines!(ax_UΣ, wu, zf, linewidth=3, label="wu", color=colors[2])
    line_wv = lines!(ax_UΣ, wv, zf, linewidth=3, label="wv", color=colors[3])
    axislegend(ax_UΣ, position=:rb, framevisible=false)
    xlims!(ax_UΣ, extrema([extrema(ds[:uv])..., extrema(ds[:wu])..., extrema(ds[:wv])...]))
    ylims!(ax_UΣ, extrema(zf))
    hideydecorations!(ax_UΣ, grid=false)

    ax_UT = fig[1, 6] = Axis(fig, xlabel="m·K/s")
    line_uT = lines!(ax_UT, uT, zc, linewidth=3, label="uT", color=colors[1])
    line_vT = lines!(ax_UT, vT, zc, linewidth=3, label="vT", color=colors[2])
    line_wT = lines!(ax_UT, wT, zf, linewidth=3, label="wT", color=colors[3])
    axislegend(ax_UT, position=:rb, framevisible=false)
    xlims!(ax_UT, extrema([extrema(ds[:uT])..., extrema(ds[:vT])..., extrema(ds[:wT])...]))
    ylims!(ax_UT, extrema(zf))
    hideydecorations!(ax_UT, grid=false)

    ax_νₑ = fig[1, 7] = Axis(fig, xlabel="m²/s²")
    line_νₑ_∂z_u = lines!(ax_νₑ, νₑ_∂z_u, zf, linewidth=3, label="νₑ ∂z(u)", color=colors[1])
    line_νₑ_∂z_v = lines!(ax_νₑ, νₑ_∂z_v, zf, linewidth=3, label="νₑ ∂z(v)", color=colors[2])
    line_νₑ_∂z_w = lines!(ax_νₑ, νₑ_∂z_w, zc, linewidth=3, label="νₑ ∂z(w)", color=colors[3])
    axislegend(ax_νₑ, position=:rb, framevisible=false)
    xlims!(ax_νₑ, extrema([extrema(ds[:νₑ_∂z_u])..., extrema(ds[:νₑ_∂z_v])..., extrema(ds[:νₑ_∂z_w])...]))
    ylims!(ax_νₑ, extrema(zf))
    hideydecorations!(ax_νₑ, grid=false)

    ax_κₑ∂zT = fig[1, 8] = Axis(fig, xlabel="m⋅K/s")
    line_κₑ∂zT = lines!(ax_κₑ∂zT, κₑ_∂z_T, zf, label="κₑ ∂z(T)", linewidth=3, color=colors[1])
    axislegend(ax_κₑ∂zT, position=:rb, framevisible=false)
    xlims!(ax_κₑ∂zT, extrema(ds[:κₑ_∂z_T]))
    ylims!(ax_κₑ∂zT, extrema(zf))
    hideydecorations!(ax_κₑ∂zT, grid=false)

    ax_κₑ∂zC = fig[2, 1] = Axis(fig, xlabel="m²/s²")
    line_κₑ_∂z_c₀ = lines!(ax_κₑ∂zC, κₑ_∂z_c₀, zf, linewidth=3, label="κₑ ∂z(c₀)", color=colors[1])
    line_κₑ_∂z_c₁ = lines!(ax_κₑ∂zC, κₑ_∂z_c₂, zf, linewidth=3, label="κₑ ∂z(c₁)", color=colors[2])
    line_κₑ_∂z_c₂ = lines!(ax_κₑ∂zC, κₑ_∂z_c₂, zf, linewidth=3, label="κₑ ∂z(c₂)", color=colors[3])
    axislegend(ax_κₑ∂zC, position=:rb, framevisible=false)
    xlims!(ax_κₑ∂zC, extrema([extrema(ds[:κₑ_∂z_c₀])..., extrema(ds[:κₑ_∂z_c₂])..., extrema(ds[:κₑ_∂z_c₂])...]))
    ylims!(ax_κₑ∂zC, extrema(zf))

    ax_Ri = fig[2, 2] = Axis(fig, xlabel="Ri")
    line_Ri = lines!(ax_Ri, Ri, zf, linewidth=3, color=colors[1])
    xlims!(ax_Ri, (-0.5, 4))
    ylims!(ax_Ri, extrema(zf))
    hideydecorations!(ax_Ri, grid=false)

    ax_e = fig[2, 3] = Axis(fig, xlabel="TKE")
    line_e = lines!(ax_e, e, zc, linewidth=3, color=colors[1])
    xlims!(ax_e, extrema(ds[:e]))
    ylims!(ax_e, extrema(zf))
    hideydecorations!(ax_e, grid=false)

    ax_ϵ = fig[2, 4] = Axis(fig, xlabel="TKE dissipation")
    line_ϵ = lines!(ax_ϵ, tke_dissipation, zc, linewidth=3, color=colors[1])
    xlims!(ax_ϵ, extrema(ds[:tke_dissipation]))
    ylims!(ax_ϵ, extrema(zf))
    hideydecorations!(ax_ϵ, grid=false)

    ax_tke_a = fig[2, 5] = Axis(fig, xlabel="TKE advective flux")
    line_tke_a = lines!(ax_tke_a, tke_advective_flux, zf, linewidth=3, color=colors[1])
    xlims!(ax_tke_a, extrema(ds[:tke_advective_flux]))
    ylims!(ax_tke_a, extrema(zf))
    hideydecorations!(ax_tke_a, grid=false)

    ax_tke_b = fig[2, 6] = Axis(fig, xlabel="TKE buoyancy flux")
    line_tke_b = lines!(ax_tke_b, tke_buoyancy_flux, zc, linewidth=3, color=colors[1])
    xlims!(ax_tke_b, extrema(ds[:tke_buoyancy_flux]))
    ylims!(ax_tke_b, extrema(zf))
    hideydecorations!(ax_tke_b, grid=false)

    ax_tke_p = fig[2, 7] = Axis(fig, xlabel="TKE pressure flux")
    line_tke_p = lines!(ax_tke_p, tke_pressure_flux, zf, linewidth=3, color=colors[1])
    xlims!(ax_tke_p, extrema(ds[:tke_pressure_flux]))
    ylims!(ax_tke_p, extrema(zf))
    hideydecorations!(ax_tke_p, grid=false)

    ax_tke_sp = fig[2, 8] = Axis(fig, xlabel="TKE shear production")
    line_tke_sp = lines!(ax_tke_sp, tke_shear_production, zc, linewidth=3, color=colors[1])
    xlims!(ax_tke_sp, extrema(ds[:tke_shear_production]))
    ylims!(ax_tke_sp, extrema(zf))
    hideydecorations!(ax_tke_sp, grid=false)

    supertitle = fig[0, :] = Label(fig, plot_title, textsize=30)

    filepath = joinpath(data_directory, "instantaneous_statistics.mp4")
    record(fig, filepath, 1:Nt, framerate=15) do n
        @info "Animating instantaneous statistics movie frame $n/$Nt..."
        frame[] = n
    end

    @info "Movie saved: $filepath"
end


text:
"""
    struct Particles <: AbstractParticles
        type::Symbol
        props::Dict{Symbol,Any}
    end

Particles of a certain type (typically something like `:dm` or `:gas` in the cosmological
context) with their properties.

The following naming convention is to be used for specific properties:
- `:id`: vector of IDs
- `:pos`: ``2×N`` or ``3×N`` matrix with positions
- `:vel`: ``2×N`` or ``3×N`` matrix with velocities
- `:mass`: vector of masses or a scalar if all particles have the same mass
- `:temp`: vector of temperatures
"""
struct Particles <: AbstractParticles
    type::Symbol
    props::Dict{Symbol,Any}
end

"""
    Particles(type[, pairs::Pair...])

Create a [`Particles`](@ref) object with the given `type` and pairs of `Symbol` to the property
values. These are passed to the underlying `Dict`.
If no pairs are passed to the method, an empty `Dict` is created.
"""
Particles(type) = Particles(type, Dict{Symbol,Any}())
Particles(type, pairs::Pair...) = Particles(type, Dict{Symbol,Any}(pairs...))

Base.copy(p::Particles) = Particles(p.type, copy(p.props))
Base.empty(p::Particles) = Particles(p.type, empty(p.props))
Base.:(==)(p1::Particles, p2::Particles) = p1.type == p2.type && isequal(p1.props, p2.props)
particle_name(::Particles) = "Particles"

function Base.propertynames(p::Particles; private=false)
    if private
        return [fieldnames(Particles) |> collect; keys(p) |> collect]
    else
        return [:type; keys(p) |> collect]
    end
end

function Base.show(io::IO, mime::MIME"text/plain", p::Particles)
    printstyled(io, String(p.type); bold=true)
    print(io, ": ")
    show_properties(io, mime, p)
end


text:
  const DATADIR = joinpath(dirname(pathof(@__MODULE__)), "..", "data")

  # US Equities from A to E (so far)
  const AAPL = readtimearray(joinpath(DATADIR, "AAPL.csv"), meta = "AAPL")
  const BA   = readtimearray(joinpath(DATADIR, "BA.csv"),   meta = "BA")
  const CAT  = readtimearray(joinpath(DATADIR, "CAT.csv"),  meta = "CAT")
  const DELL = readtimearray(joinpath(DATADIR, "DELL.csv"), meta = "DELL")
  const EBAY = readtimearray(joinpath(DATADIR, "EBAY.csv"), meta = "EBAY")
  const F    = readtimearray(joinpath(DATADIR, "F.csv"),    meta = "F")
  const GE   = readtimearray(joinpath(DATADIR, "GE.csv"),   meta = "GE")
  const TX   = readtimearray(joinpath(DATADIR, "TX.csv"),   meta = "TX")

  # smaller datasets used for testing time-related packages
  const dates = Date(2000,1,1):Day(1):Date(2001,12,31)
  const cl        = AAPL[:Close][dates]
  const op        = AAPL[:Open][dates]
  const ohlc      = AAPL[:Open, :High, :Low, :Close][dates]
  const ohlcv     = AAPL[:Open, :High, :Low, :Close, :Volume][dates]
  const datetime1 = readtimearray(joinpath(DATADIR, "datetime1.csv"))
  const datetime2 = readtimearray(joinpath(DATADIR, "datetime2.csv"),
                                  format="yyyy-mm-dd HH:MM:SS")
  const mdata     = TimeArray(timestamp(cl), values(cl), colnames(cl), "Apple") # data with meta field occupied


text:
"""
Tests that _may_ destroy your Julia setup!
"""
module DestructiveTests

using JuliaManager
using Test

macro test_nothrow(ex)
    esc(:($Test.@test $ex isa $Any))
end

if Sys.iswindows()
    pathsep = ";"
else
    pathsep = ":"
end

if Sys.which("jlm") === nothing
    destdir = joinpath(homedir(), ".julia", "bin")

    ENV["PATH"] = destdir * pathsep * ENV["PATH"]

    @test JuliaManager.install_cli(destdir) === nothing
    @assert Sys.which("jlm") !== nothing  # abort test

    # At this point `install_cli` should be a no-op
    @test JuliaManager.install_cli(destdir) === nothing
    @test JuliaManager.install_cli(destdir; upgrade=true) === nothing
end

@show Sys.which("jlm")
@test Sys.which("jlm") !== nothing

@info "jlm --help"
@test_nothrow run(`jlm --help`)

@info "jlm create-default-sysimage"
@time @test_nothrow run(`jlm create-default-sysimage`)

if Sys.which("tox") === nothing
    run(`python3 -m pip install --user tox`)
end

jlmdir = joinpath(@__DIR__, "..", "jlm")
@info "cd $jlmdir && tox"
cd(jlmdir) do
    @time @test_nothrow run(`tox`)
end

end  # module