text:
using Statistics
using NCDatasets
using Oceananigans.BoundaryConditions: BoundaryFunction, PBC, FBC, ZFBC
using Oceananigans.Diagnostics
using Oceananigans.Fields

function run_thermal_bubble_netcdf_tests(arch)
    Nx, Ny, Nz = 16, 16, 16
    Lx, Ly, Lz = 100, 100, 100

    grid = RegularCartesianGrid(size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz))
    closure = IsotropicDiffusivity(ν=4e-2, κ=4e-2)
    model = IncompressibleModel(architecture=arch, grid=grid, closure=closure)
    simulation = Simulation(model, Δt=6, stop_iteration=10)

    # Add a cube-shaped warm temperature anomaly that takes up the middle 50%
    # of the domain volume.
    i1, i2 = round(Int, Nx/4), round(Int, 3Nx/4)
    j1, j2 = round(Int, Ny/4), round(Int, 3Ny/4)
    k1, k2 = round(Int, Nz/4), round(Int, 3Nz/4)
    CUDA.@allowscalar model.tracers.T.data[i1:i2, j1:j2, k1:k2] .+= 0.01

    outputs = Dict(
        "v" => model.velocities.v,
        "u" => model.velocities.u,
        "w" => model.velocities.w,
        "T" => model.tracers.T,
        "S" => model.tracers.S
    )
    nc_filename = "test_dump_$(typeof(arch)).nc"
    nc_writer = NetCDFOutputWriter(model, outputs, filename=nc_filename, iteration_interval=10, verbose=true)
    push!(simulation.output_writers, nc_writer)

    xC_slice = 1:10
    xF_slice = 2:11
    yC_slice = 10:15
    yF_slice = 1
    zC_slice = 10
    zF_slice = 9:11

    nc_sliced_filename = "test_dump_sliced_$(typeof(arch)).nc"
    nc_sliced_writer =
        NetCDFOutputWriter(model, outputs, filename=nc_sliced_filename, iteration_interval=10, verbose=true,
                           xC=xC_slice, xF=xF_slice, yC=yC_slice,
                           yF=yF_slice, zC=zC_slice, zF=zF_slice)

    push!(simulation.output_writers, nc_sliced_writer)

    run!(simulation)

    ds3 = Dataset(nc_filename)

    @test haskey(ds3.attrib, "date") && !isnothing(ds3.attrib["date"])
    @test haskey(ds3.attrib, "Julia") && !isnothing(ds3.attrib["Julia"])
    @test haskey(ds3.attrib, "Oceananigans") && !isnothing(ds3.attrib["Oceananigans"])

    @test length(ds3["xC"]) == Nx
    @test length(ds3["yC"]) == Ny
    @test length(ds3["zC"]) == Nz
    @test length(ds3["xF"]) == Nx
    @test length(ds3["yF"]) == Ny
    @test length(ds3["zF"]) == Nz+1  # z is Bounded

    @test ds3["xC"][1] == grid.xC[1]
    @test ds3["xF"][1] == grid.xF[1]
    @test ds3["yC"][1] == grid.yC[1]
    @test ds3["yF"][1] == grid.yF[1]
    @test ds3["zC"][1] == grid.zC[1]
    @test ds3["zF"][1] == grid.zF[1]

    @test ds3["xC"][end] == grid.xC[Nx]
    @test ds3["xF"][end] == grid.xF[Nx]
    @test ds3["yC"][end] == grid.yC[Ny]
    @test ds3["yF"][end] == grid.yF[Ny]
    @test ds3["zC"][end] == grid.zC[Nz]
    @test ds3["zF"][end] == grid.zF[Nz+1]  # z is Bounded

    u = ds3["u"][:, :, :, end]
    v = ds3["v"][:, :, :, end]
    w = ds3["w"][:, :, :, end]
    T = ds3["T"][:, :, :, end]
    S = ds3["S"][:, :, :, end]

    close(ds3)

    @test all(u .≈ Array(interiorparent(model.velocities.u)))
    @test all(v .≈ Array(interiorparent(model.velocities.v)))
    @test all(w .≈ Array(interiorparent(model.velocities.w)))
    @test all(T .≈ Array(interiorparent(model.tracers.T)))
    @test all(S .≈ Array(interiorparent(model.tracers.S)))

    ds2 = Dataset(nc_sliced_filename)

    @test haskey(ds2.attrib, "date") && !isnothing(ds2.attrib["date"])
    @test haskey(ds2.attrib, "Julia") && !isnothing(ds2.attrib["Julia"])
    @test haskey(ds2.attrib, "Oceananigans") && !isnothing(ds2.attrib["Oceananigans"])

    @test length(ds2["xC"]) == length(xC_slice)
    @test length(ds2["xF"]) == length(xF_slice)
    @test length(ds2["yC"]) == length(yC_slice)
    @test length(ds2["yF"]) == length(yF_slice)
    @test length(ds2["zC"]) == length(zC_slice)
    @test length(ds2["zF"]) == length(zF_slice)

    @test ds2["xC"][1] == grid.xC[xC_slice[1]]
    @test ds2["xF"][1] == grid.xF[xF_slice[1]]
    @test ds2["yC"][1] == grid.yC[yC_slice[1]]
    @test ds2["yF"][1] == grid.yF[yF_slice]
    @test ds2["zC"][1] == grid.zC[zC_slice]
    @test ds2["zF"][1] == grid.zF[zF_slice[1]]

    @test ds2["xC"][end] == grid.xC[xC_slice[end]]
    @test ds2["xF"][end] == grid.xF[xF_slice[end]]
    @test ds2["yC"][end] == grid.yC[yC_slice[end]]
    @test ds2["yF"][end] == grid.yF[yF_slice]
    @test ds2["zC"][end] == grid.zC[zC_slice]
    @test ds2["zF"][end] == grid.zF[zF_slice[end]]

    u_sliced = ds2["u"][:, :, :, end]
    v_sliced = ds2["v"][:, :, :, end]
    w_sliced = ds2["w"][:, :, :, end]
    T_sliced = ds2["T"][:, :, :, end]
    S_sliced = ds2["S"][:, :, :, end]

    close(ds2)

    @test all(u_sliced .≈ Array(interiorparent(model.velocities.u))[xF_slice, yC_slice, zC_slice])
    @test all(v_sliced .≈ Array(interiorparent(model.velocities.v))[xC_slice, yF_slice, zC_slice])
    @test all(w_sliced .≈ Array(interiorparent(model.velocities.w))[xC_slice, yC_slice, zF_slice])
    @test all(T_sliced .≈ Array(interiorparent(model.tracers.T))[xC_slice, yC_slice, zC_slice])
    @test all(S_sliced .≈ Array(interiorparent(model.tracers.S))[xC_slice, yC_slice, zC_slice])
end

function run_thermal_bubble_netcdf_tests_with_halos(arch)
    Nx, Ny, Nz = 16, 16, 16
    Lx, Ly, Lz = 100, 100, 100

    grid = RegularCartesianGrid(size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz))
    closure = IsotropicDiffusivity(ν=4e-2, κ=4e-2)
    model = IncompressibleModel(architecture=arch, grid=grid, closure=closure)
    simulation = Simulation(model, Δt=6, stop_iteration=10)

    # Add a cube-shaped warm temperature anomaly that takes up the middle 50%
    # of the domain volume.
    i1, i2 = round(Int, Nx/4), round(Int, 3Nx/4)
    j1, j2 = round(Int, Ny/4), round(Int, 3Ny/4)
    k1, k2 = round(Int, Nz/4), round(Int, 3Nz/4)
    CUDA.@allowscalar model.tracers.T.data[i1:i2, j1:j2, k1:k2] .+= 0.01

    outputs = Dict(
        "v" => model.velocities.v,
        "u" => model.velocities.u,
        "w" => model.velocities.w,
        "T" => model.tracers.T,
        "S" => model.tracers.S
    )
    nc_filename = "test_dump_with_halos_$(typeof(arch)).nc"
    nc_writer = NetCDFOutputWriter(model, outputs, filename=nc_filename, iteration_interval=10, with_halos=true)
    push!(simulation.output_writers, nc_writer)

    run!(simulation)

    ds = Dataset(nc_filename)

    @test haskey(ds.attrib, "date") && !isnothing(ds.attrib["date"])
    @test haskey(ds.attrib, "Julia") && !isnothing(ds.attrib["Julia"])
    @test haskey(ds.attrib, "Oceananigans") && !isnothing(ds.attrib["Oceananigans"])

    Hx, Hy, Hz = grid.Hx, grid.Hy, grid.Hz
    @test length(ds["xC"]) == Nx+2Hx
    @test length(ds["yC"]) == Ny+2Hy
    @test length(ds["zC"]) == Nz+2Hz
    @test length(ds["xF"]) == Nx+2Hx
    @test length(ds["yF"]) == Ny+2Hy
    @test length(ds["zF"]) == Nz+2Hz+1  # z is Bounded

    @test ds["xC"][1] == grid.xC[1-Hx]
    @test ds["xF"][1] == grid.xF[1-Hx]
    @test ds["yC"][1] == grid.yC[1-Hy]
    @test ds["yF"][1] == grid.yF[1-Hy]
    @test ds["zC"][1] == grid.zC[1-Hz]
    @test ds["zF"][1] == grid.zF[1-Hz]

    @test ds["xC"][end] == grid.xC[Nx+Hx]
    @test ds["xF"][end] == grid.xF[Nx+Hx]
    @test ds["yC"][end] == grid.yC[Ny+Hy]
    @test ds["yF"][end] == grid.yF[Ny+Hy]
    @test ds["zC"][end] == grid.zC[Nz+Hz]
    @test ds["zF"][end] == grid.zF[Nz+Hz+1]  # z is Bounded

    u = ds["u"][:, :, :, end]
    v = ds["v"][:, :, :, end]
    w = ds["w"][:, :, :, end]
    T = ds["T"][:, :, :, end]
    S = ds["S"][:, :, :, end]

    close(ds)

    @test all(u .≈ Array(model.velocities.u.data.parent))
    @test all(v .≈ Array(model.velocities.v.data.parent))
    @test all(w .≈ Array(model.velocities.w.data.parent))
    @test all(T .≈ Array(model.tracers.T.data.parent))
    @test all(S .≈ Array(model.tracers.S.data.parent))
end

function run_netcdf_function_output_tests(arch)
    N = 16
    L = 1
    Δt = 1.25
    iters = 3

    model = IncompressibleModel(grid=RegularCartesianGrid(size=(N, N, N), extent=(L, 2L, 3L)))
    simulation = Simulation(model, Δt=Δt, stop_iteration=iters)
    grid = model.grid

    # Define scalar, vector, and 2D slice outputs
    f(model) = model.clock.time^2

    g(model) = model.clock.time .* exp.(znodes(Cell, grid))

    h(model) = model.clock.time .* (   sin.(xnodes(Cell, grid, reshape=true)[:, :, 1])
                                    .* cos.(ynodes(Face, grid, reshape=true)[:, :, 1]))

    outputs = Dict("scalar" => f,  "profile" => g,       "slice" => h)
       dims = Dict("scalar" => (), "profile" => ("zC",), "slice" => ("xC", "yC"))

    output_attributes = Dict(
        "scalar"  => Dict("longname" => "Some scalar", "units" => "bananas"),
        "profile" => Dict("longname" => "Some vertical profile", "units" => "watermelons"),
        "slice"   => Dict("longname" => "Some slice", "units" => "mushrooms")
    )

    global_attributes = Dict("location" => "Bay of Fundy", "onions" => 7)

    nc_filename = "test_function_outputs_$(typeof(arch)).nc"
    simulation.output_writers[:food] =
        NetCDFOutputWriter(model, outputs;
            iteration_interval=1, filename=nc_filename, dimensions=dims, verbose=true,
            global_attributes=global_attributes, output_attributes=output_attributes)

    run!(simulation)

    ds = Dataset(nc_filename, "r")

    @test haskey(ds.attrib, "date") && !isnothing(ds.attrib["date"])
    @test haskey(ds.attrib, "Julia") && !isnothing(ds.attrib["Julia"])
    @test haskey(ds.attrib, "Oceananigans") && !isnothing(ds.attrib["Oceananigans"])

    @test length(ds["xC"]) == N
    @test length(ds["yC"]) == N
    @test length(ds["zC"]) == N
    @test length(ds["xF"]) == N
    @test length(ds["yF"]) == N
    @test length(ds["zF"]) == N+1  # z is Bounded

    @test ds["xC"][1] == grid.xC[1]
    @test ds["xF"][1] == grid.xF[1]
    @test ds["yC"][1] == grid.yC[1]
    @test ds["yF"][1] == grid.yF[1]
    @test ds["zC"][1] == grid.zC[1]
    @test ds["zF"][1] == grid.zF[1]

    @test ds["xC"][end] == grid.xC[N]
    @test ds["yC"][end] == grid.yC[N]
    @test ds["zC"][end] == grid.zC[N]
    @test ds["xF"][end] == grid.xF[N]
    @test ds["yF"][end] == grid.yF[N]
    @test ds["zF"][end] == grid.zF[N+1]  # z is Bounded

    @test ds.attrib["location"] == "Bay of Fundy"
    @test ds.attrib["onions"] == 7

    @test length(ds["time"]) == iters+1
    @test ds["time"][:] == [n*Δt for n in 0:iters]

    @test length(ds["scalar"]) == iters+1
    @test ds["scalar"].attrib["longname"] == "Some scalar"
    @test ds["scalar"].attrib["units"] == "bananas"
    @test ds["scalar"][:] == [(n*Δt)^2 for n in 0:iters]
    @test dimnames(ds["scalar"]) == ("time",)

    @test ds["profile"].attrib["longname"] == "Some vertical profile"
    @test ds["profile"].attrib["units"] == "watermelons"
    @test size(ds["profile"]) == (N, iters+1)
    @test dimnames(ds["profile"]) == ("zC", "time")

    for n in 0:iters
        @test ds["profile"][:, n+1] == n*Δt .* exp.(znodes(Cell, grid))
    end

    @test ds["slice"].attrib["longname"] == "Some slice"
    @test ds["slice"].attrib["units"] == "mushrooms"
    @test size(ds["slice"]) == (N, N, iters+1)
    @test dimnames(ds["slice"]) == ("xC", "yC", "time")

    for n in 0:iters
        @test ds["slice"][:, :, n+1] == n*Δt .* (   sin.(xnodes(Cell, grid, reshape=true)[:, :, 1])
                                                 .* cos.(ynodes(Face, grid, reshape=true)[:, :, 1]))
    end

    close(simulation.output_writers[:food])

    #####
    ##### Take 1 more time step and test that appending to a NetCDF file works
    #####

    iters += 1
    simulation = Simulation(model, Δt=Δt, stop_iteration=iters)

    simulation.output_writers[:food] =
        NetCDFOutputWriter(model, outputs;
            iteration_interval=1, filename=nc_filename, mode="a", dimensions=dims, verbose=true,
            global_attributes=global_attributes, output_attributes=output_attributes)

    run!(simulation)

    ds = Dataset(nc_filename, "r")

    @test length(ds["time"]) == iters+1
    @test length(ds["scalar"]) == iters+1
    @test size(ds["profile"]) == (N, iters+1)
    @test size(ds["slice"]) == (N, N, iters+1)

    @test ds["time"][:] == [n*Δt for n in 0:iters]
    @test ds["scalar"][:] == [(n*Δt)^2 for n in 0:iters]

    for n in 0:iters
        @test ds["profile"][:, n+1] == n*Δt .* exp.(znodes(Cell, grid))
        @test ds["slice"][:, :, n+1] == n*Δt .* (   sin.(xnodes(Cell, grid, reshape=true)[:, :, 1])
                                                 .* cos.(ynodes(Face, grid, reshape=true)[:, :, 1]))
    end

    close(simulation.output_writers[:food])

    return nothing
end

function run_jld2_file_splitting_tests(arch)
    model = IncompressibleModel(grid=RegularCartesianGrid(size=(16, 16, 16), extent=(1, 1, 1)))
    simulation = Simulation(model, Δt=1, stop_iteration=10)

    function fake_bc_init(file, model)
        file["boundary_conditions/fake"] = π
    end

    ow = JLD2OutputWriter(model, (u=model.velocities.u,); dir=".", prefix="test", iteration_interval=1,
                          init=fake_bc_init, including=[:grid],
                          field_slicer=nothing, array_type=Array{Float64},
                          max_filesize=200KiB, force=true)

    push!(simulation.output_writers, ow)

    # 531 KiB of output will be written which should get split into 3 files.
    run!(simulation)

    # Test that files has been split according to size as expected.
    @test filesize("test_part1.jld2") > 200KiB
    @test filesize("test_part2.jld2") > 200KiB
    @test filesize("test_part3.jld2") < 200KiB

    for n in string.(1:3)
        filename = "test_part" * n * ".jld2"
        jldopen(filename, "r") do file
            # Test to make sure all files contain structs from `including`.
            @test file["grid/Nx"] == 16

            # Test to make sure all files contain info from `init` function.
            @test file["boundary_conditions/fake"] == π
        end

        # Leave test directory clean.
        rm(filename)
    end
end

"""
Run two coarse rising thermal bubble simulations and make sure that when
restarting from a checkpoint, the restarted simulation matches the non-restarted
simulation to machine precision.
"""
function run_thermal_bubble_checkpointer_tests(arch)
    Nx, Ny, Nz = 16, 16, 16
    Lx, Ly, Lz = 100, 100, 100
    Δt = 6

    grid = RegularCartesianGrid(size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz))
    closure = IsotropicDiffusivity(ν=4e-2, κ=4e-2)
    true_model = IncompressibleModel(architecture=arch, grid=grid, closure=closure)

    # Add a cube-shaped warm temperature anomaly that takes up the middle 50%
    # of the domain volume.
    i1, i2 = round(Int, Nx/4), round(Int, 3Nx/4)
    j1, j2 = round(Int, Ny/4), round(Int, 3Ny/4)
    k1, k2 = round(Int, Nz/4), round(Int, 3Nz/4)
    CUDA.@allowscalar true_model.tracers.T.data[i1:i2, j1:j2, k1:k2] .+= 0.01

    checkpointed_model = deepcopy(true_model)

    true_simulation = Simulation(true_model, Δt=Δt, stop_iteration=9)
    run!(true_simulation)

    checkpointed_simulation = Simulation(checkpointed_model, Δt=Δt, stop_iteration=5)
    checkpointer = Checkpointer(checkpointed_model, iteration_interval=5, force=true)
    push!(checkpointed_simulation.output_writers, checkpointer)

    # Checkpoint should be saved as "checkpoint5.jld" after the 5th iteration.
    run!(checkpointed_simulation)

    # Remove all knowledge of the checkpointed model.
    checkpointed_model = nothing

    # model_kwargs = Dict{Symbol, Any}(:boundary_conditions => SolutionBoundaryConditions(grid))
    restored_model = restore_from_checkpoint("checkpoint_iteration5.jld2")

    for n in 1:4
        time_step!(restored_model, Δt, euler=false)
    end

    rm("checkpoint_iteration0.jld2", force=true)
    rm("checkpoint_iteration5.jld2", force=true)

    # Now the true_model and restored_model should be identical.
    CUDA.@allowscalar begin
        @test all(restored_model.velocities.u.data     .≈ true_model.velocities.u.data)
        @test all(restored_model.velocities.v.data     .≈ true_model.velocities.v.data)
        @test all(restored_model.velocities.w.data     .≈ true_model.velocities.w.data)
        @test all(restored_model.tracers.T.data        .≈ true_model.tracers.T.data)
        @test all(restored_model.tracers.S.data        .≈ true_model.tracers.S.data)
        @test all(restored_model.timestepper.Gⁿ.u.data .≈ true_model.timestepper.Gⁿ.u.data)
        @test all(restored_model.timestepper.Gⁿ.v.data .≈ true_model.timestepper.Gⁿ.v.data)
        @test all(restored_model.timestepper.Gⁿ.w.data .≈ true_model.timestepper.Gⁿ.w.data)
        @test all(restored_model.timestepper.Gⁿ.T.data .≈ true_model.timestepper.Gⁿ.T.data)
        @test all(restored_model.timestepper.Gⁿ.S.data .≈ true_model.timestepper.Gⁿ.S.data)
    end

    return nothing
end

function run_checkpoint_with_function_bcs_tests(arch)
    grid = RegularCartesianGrid(size=(16, 16, 16), extent=(1, 1, 1))

    @inline some_flux(x, y, t) = 2x + exp(y)
    some_flux_bf = BoundaryFunction{:z, Cell, Cell}(some_flux)
    top_u_bc = top_T_bc = FluxBoundaryCondition(some_flux_bf)
    u_bcs = UVelocityBoundaryConditions(grid, top=top_u_bc)
    T_bcs = TracerBoundaryConditions(grid, top=top_T_bc)

    model = IncompressibleModel(architecture=arch, grid=grid, boundary_conditions=(u=u_bcs, T=T_bcs))
    set!(model, u=π/2, v=ℯ, T=Base.MathConstants.γ, S=Base.MathConstants.φ)

    checkpointer = Checkpointer(model)
    write_output(model, checkpointer)
    model = nothing

    restored_model = restore_from_checkpoint("checkpoint_iteration0.jld2")
    @test  ismissing(restored_model.velocities.u.boundary_conditions)
    @test !ismissing(restored_model.velocities.v.boundary_conditions)
    @test !ismissing(restored_model.velocities.w.boundary_conditions)
    @test  ismissing(restored_model.tracers.T.boundary_conditions)
    @test !ismissing(restored_model.tracers.S.boundary_conditions)

    CUDA.@allowscalar begin
        @test all(interior(restored_model.velocities.u) .≈ π/2)
        @test all(interior(restored_model.velocities.v) .≈ ℯ)
        @test all(interior(restored_model.velocities.w) .== 0)
        @test all(interior(restored_model.tracers.T) .≈ Base.MathConstants.γ)
        @test all(interior(restored_model.tracers.S) .≈ Base.MathConstants.φ)
    end
    restored_model = nothing

    properly_restored_model = restore_from_checkpoint("checkpoint_iteration0.jld2",
                                                      boundary_conditions=(u=u_bcs, T=T_bcs))

    CUDA.@allowscalar begin
        @test all(interior(properly_restored_model.velocities.u) .≈ π/2)
        @test all(interior(properly_restored_model.velocities.v) .≈ ℯ)
        @test all(interior(properly_restored_model.velocities.w) .== 0)
        @test all(interior(properly_restored_model.tracers.T) .≈ Base.MathConstants.γ)
        @test all(interior(properly_restored_model.tracers.S) .≈ Base.MathConstants.φ)
    end

    @test !ismissing(properly_restored_model.velocities.u.boundary_conditions)
    @test !ismissing(properly_restored_model.velocities.v.boundary_conditions)
    @test !ismissing(properly_restored_model.velocities.w.boundary_conditions)
    @test !ismissing(properly_restored_model.tracers.T.boundary_conditions)
    @test !ismissing(properly_restored_model.tracers.S.boundary_conditions)

    u, v, w = properly_restored_model.velocities
    T, S = properly_restored_model.tracers

    @test u.boundary_conditions.x.left  isa PBC
    @test u.boundary_conditions.x.right isa PBC
    @test u.boundary_conditions.y.left  isa PBC
    @test u.boundary_conditions.y.right isa PBC
    @test u.boundary_conditions.z.left  isa ZFBC
    @test u.boundary_conditions.z.right isa FBC
    @test u.boundary_conditions.z.right.condition isa BoundaryFunction
    @test u.boundary_conditions.z.right.condition.func(1, 2, 3) == some_flux(1, 2, 3)

    @test T.boundary_conditions.x.left  isa PBC
    @test T.boundary_conditions.x.right isa PBC
    @test T.boundary_conditions.y.left  isa PBC
    @test T.boundary_conditions.y.right isa PBC
    @test T.boundary_conditions.z.left  isa ZFBC
    @test T.boundary_conditions.z.right isa FBC
    @test T.boundary_conditions.z.right.condition isa BoundaryFunction
    @test T.boundary_conditions.z.right.condition.func(1, 2, 3) == some_flux(1, 2, 3)

    # Test that the restored model can be time stepped
    time_step!(properly_restored_model, 1)
    @test properly_restored_model isa IncompressibleModel

    return nothing
end

function run_cross_architecture_checkpointer_tests(arch1, arch2)
    grid = RegularCartesianGrid(size=(16, 16, 16), extent=(1, 1, 1))
    model = IncompressibleModel(architecture=arch1, grid=grid)
    set!(model, u=π/2, v=ℯ, T=Base.MathConstants.γ, S=Base.MathConstants.φ)

    checkpointer = Checkpointer(model)
    write_output(model, checkpointer)
    model = nothing

    restored_model = restore_from_checkpoint("checkpoint_iteration0.jld2", architecture=arch2)

    @test restored_model.architecture == arch2

    ArrayType = array_type(restored_model.architecture)
    CUDA.@allowscalar begin
        @test restored_model.velocities.u.data.parent isa ArrayType
        @test restored_model.velocities.v.data.parent isa ArrayType
        @test restored_model.velocities.w.data.parent isa ArrayType
        @test restored_model.tracers.T.data.parent isa ArrayType
        @test restored_model.tracers.S.data.parent isa ArrayType

        @test all(interior(restored_model.velocities.u) .≈ π/2)
        @test all(interior(restored_model.velocities.v) .≈ ℯ)
        @test all(interior(restored_model.velocities.w) .== 0)
        @test all(interior(restored_model.tracers.T) .≈ Base.MathConstants.γ)
        @test all(interior(restored_model.tracers.S) .≈ Base.MathConstants.φ)
    end

    # Test that the restored model can be time stepped
    time_step!(restored_model, 1)
    @test restored_model isa IncompressibleModel

    return nothing
end

function instantiate_windowed_time_average(model)

    set!(model, u = (x, y, z) -> rand())

    u, v, w = model.velocities

    u₀ = similar(interior(u))
    u₀ .= interior(u)

    wta = WindowedTimeAverage(model.velocities.u, time_window=1.0, time_interval=10.0)

    return all(wta(model) .== u₀)
end

function time_step_with_windowed_time_average(model)
    model.clock.iteration = 0
    model.clock.time = 0.0

    set!(model, u=0, v=0, w=0, T=0, S=0)

    wta = WindowedTimeAverage(model.velocities.u, time_window=2.0, time_interval=4.0)

    simulation = Simulation(model, Δt=1.0, stop_time=4.0)
    simulation.diagnostics[:u_avg] = wta
    run!(simulation)

    return all(wta(model) .== interior(model.velocities.u))
end

function dependencies_added_correctly!(model, windowed_time_average, output_writer)

    model.clock.iteration = 0
    model.clock.time = 0.0

    simulation = Simulation(model, Δt=1.0, stop_iteration=1)
    push!(simulation.output_writers, output_writer)
    run!(simulation)

    return windowed_time_average ∈ values(simulation.diagnostics)
end

function jld2_field_output(model)

    model.clock.iteration = 0
    model.clock.time = 0.0

    set!(model, u = (x, y, z) -> rand(),
                v = (x, y, z) -> rand(),
                w = (x, y, z) -> rand(),
                T = 0,
                S = 0)

    simulation = Simulation(model, Δt=1.0, stop_iteration=1)

    simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocities,
                                                              time_interval = 1.0,
                                                                        dir = ".",
                                                                     prefix = "test",
                                                                      force = true)

    u₀ = data(model.velocities.u)[3, 3, 3]
    v₀ = data(model.velocities.v)[3, 3, 3]
    w₀ = data(model.velocities.w)[3, 3, 3]

    run!(simulation)

    file = jldopen("test.jld2")

    # Data is saved without halos by default
    u₁ = file["timeseries/u/0"][3, 3, 3]
    v₁ = file["timeseries/v/0"][3, 3, 3]
    w₁ = file["timeseries/w/0"][3, 3, 3]

    close(file)

    rm("test.jld2")

    FT = typeof(u₁)

    return FT(u₀) == u₁ && FT(v₀) == v₁ && FT(w₀) == w₁
end

function jld2_sliced_field_output(model)

    model.clock.iteration = 0
    model.clock.time = 0.0

    set!(model, u = (x, y, z) -> rand(),
                v = (x, y, z) -> rand(),
                w = (x, y, z) -> rand())

    simulation = Simulation(model, Δt=1.0, stop_iteration=1)

    simulation.output_writers[:velocities] =
        JLD2OutputWriter(model, model.velocities,
                                time_interval = 1.0,
                                 field_slicer = FieldSlicer(i=1:2, j=1:3, k=:),
                                          dir = ".",
                                       prefix = "test",
                                        force = true)

    run!(simulation)

    file = jldopen("test.jld2")

    u₁ = file["timeseries/u/0"]
    v₁ = file["timeseries/v/0"]
    w₁ = file["timeseries/w/0"]

    close(file)

    rm("test.jld2")

    return size(u₁) == (2, 3, 4) && size(v₁) == (2, 3, 4) && size(w₁) == (2, 3, 5)
end

function jld2_time_averaging_of_horizontal_averages(model)

    model.clock.iteration = 0
    model.clock.time = 0.0

    set!(model, u = (x, y, z) -> 1,
                v = (x, y, z) -> 2,
                w = (x, y, z) -> 0,
                T = (x, y, z) -> 4)

    simulation = Simulation(model, Δt=1.0, stop_iteration=5)

    u, v, w = model.velocities
    T, S = model.tracers

    average_fluxes = (wu = AveragedField(w * u, dims=(1, 2)),
                      uv = AveragedField(u * v, dims=(1, 2)),
                      wT = AveragedField(w * T, dims=(1, 2)))

    simulation.output_writers[:velocities] = JLD2OutputWriter(model, average_fluxes,
                                                                      time_interval = 4.0,
                                                              time_averaging_window = 2.0,
                                                                                dir = ".",
                                                                             prefix = "test",
                                                                              force = true)

    run!(simulation)

    file = jldopen("test.jld2")

    # Data is saved with halos by default
    wu = file["timeseries/wu/4"][1, 1, 3]
    uv = file["timeseries/uv/4"][1, 1, 3]
    wT = file["timeseries/wT/4"][1, 1, 3]

    close(file)

    rm("test.jld2")

    FT = eltype(model.grid)

    return wu == zero(FT) && wT == zero(FT) && uv == FT(2)
end

@testset "Output writers" begin
    @info "Testing output writers..."

    for arch in archs
         @testset "NetCDF [$(typeof(arch))]" begin
             @info "  Testing NetCDF output writer [$(typeof(arch))]..."
             run_thermal_bubble_netcdf_tests(arch)
             run_thermal_bubble_netcdf_tests_with_halos(arch)
             run_netcdf_function_output_tests(arch)
         end

        @testset "JLD2 [$(typeof(arch))]" begin
            @info "  Testing JLD2 output writer [$(typeof(arch))]..."
            run_jld2_file_splitting_tests(arch)
        end

        @testset "Checkpointer [$(typeof(arch))]" begin
            @info "  Testing Checkpointer [$(typeof(arch))]..."
            run_thermal_bubble_checkpointer_tests(arch)
            run_checkpoint_with_function_bcs_tests(arch)
            @hascuda run_cross_architecture_checkpointer_tests(CPU(), GPU())
            @hascuda run_cross_architecture_checkpointer_tests(GPU(), CPU())
        end

        grid = RegularCartesianGrid(size=(4, 4, 4), extent=(1, 1, 1))
        model = IncompressibleModel(architecture=arch, grid=grid)

        @testset "WindowedTimeAverage and FieldSlicer [$(typeof(arch))]" begin
            @info "  Testing WindowedTimeAverage and FieldSlicer [$(typeof(arch))]"

            #####
            ##### Field slicing and field output
            #####

            @test FieldSlicer() isa FieldSlicer
            @test instantiate_windowed_time_average(model)
            @test jld2_field_output(model)

            #####
            ##### Dependency-adding
            #####

            windowed_time_average = WindowedTimeAverage(model.velocities.u, time_window=2.0, time_interval=4.0)

            output = Dict("time_average" => windowed_time_average)
            attributes = Dict("time_average" => Dict("longname" => "A time average",  "units" => "arbitrary"))
            dimensions = Dict("time_average" => ("xF", "yC", "zC"))

            # JLD2 dependencies test
            jld2_output_writer = JLD2OutputWriter(model, output, time_interval=4.0, dir=".", prefix="test", force=true)

            @test dependencies_added_correctly!(model, windowed_time_average, jld2_output_writer)

            # NetCDF dependency test
            netcdf_output_writer = NetCDFOutputWriter(model, output,
                                                          time_interval = 4.0,
                                                               filename = "test.nc",
                                                      output_attributes = attributes,
                                                             dimensions = dimensions)

            @test dependencies_added_correctly!(model, windowed_time_average, netcdf_output_writer)

            #####
            ##### Time-stepping and Simulations.run! with WindowedTimeAverage
            #####

            @test time_step_with_windowed_time_average(model)

            model.clock.iteration = model.clock.time = 0
            simulation = Simulation(model, Δt=1.0, stop_iteration=0)

            jld2_output_writer = JLD2OutputWriter(model, model.velocities,
                                                          time_interval = π,
                                                  time_averaging_window = 1.0,
                                                                 prefix = "test",
                                                                  force = true)

            outputs_are_time_averaged = Tuple(typeof(out) <: WindowedTimeAverage for out in jld2_output_writer.outputs)

            @test all(outputs_are_time_averaged)

            # Test the collection does *not* start when a simulation is initialized
            # when time_interval ≠ time_averaging_window
            simulation.output_writers[:jld2] = jld2_output_writer

            run!(simulation)

            windowed_time_average = simulation.output_writers[:jld2].outputs.u

            @test !(windowed_time_average.collecting)

            # Test that time-averaging is finalized prior to output even when averaging over
            # time_window is not fully realized. For this, step forward to a time at which
            # collection should start. Note that time_interval = π and time_window = 1.0.
            simulation.Δt = 1.5
            simulation.stop_iteration = 2
            run!(simulation) # model.clock.time = 3.0, just before output but after average-collection.

            @test windowed_time_average.collecting

            # Step forward such that time_window is not reached, but
            # output will occur.
            simulation.Δt = π - 3 + 0.01 # ≈ 0.15 < 1.0
            simulation.stop_iteration = 3
            run!(simulation) # model.clock.time ≈ 3.15, after output

            @test windowed_time_average.previous_interval_stop_time ==
                model.clock.time - rem(model.clock.time, windowed_time_average.time_interval)

            # Test the collection does start when a simulation is initialized and
            # time_interval = time_averaging_window
            model.clock.iteration = model.clock.time = 0
            simulation.output_writers[:jld2] = JLD2OutputWriter(model, model.velocities,
                                                                        time_interval = π,
                                                                time_averaging_window = π,
                                                                               prefix = "test",
                                                                                force = true)

            run!(simulation)

            windowed_time_average = simulation.output_writers[:jld2].outputs.u

            @test windowed_time_average.collecting

            @test jld2_time_averaging_of_horizontal_averages(model)
        end
    end
end


text:
######## GAS DEFINITIONS #######

export AbstractLiquid, Liquid

abstract type AbstractLiquid <: AbstractFluid end

# default values
const DefaultLiquidRefTemperature = Temperature(15.6u"°C")
const DefaultLiquidRefDensity = Density(999)
const DefaultLiquidViscosity = Viscosity(1.12e-3)
const DefaultSurfaceTension = SurfaceTension(7.34e-2)
const DefaultVaporPressure = VaporPressure(1.77e3)
const DefaultBulkModulus = BulkModulus(2.15e9)

struct Liquid <: AbstractLiquid
    Tref :: Temperature
    ρref :: Density
    specwt :: SpecificWeight
    μ :: Viscosity
    ν :: KinematicViscosity
    σ :: SurfaceTension
    pv :: VaporPressure
    Ev :: BulkModulus
end

function Liquid(;
            Tref::Temperature = DefaultLiquidRefTemperature,
            ρref::Density = DefaultLiquidRefDensity,
            μ::Viscosity = DefaultLiquidViscosity,
            σ::SurfaceTension = DefaultSurfaceTension,
            pv::VaporPressure = DefaultVaporPressure,
            Ev::BulkModulus = DefaultBulkModulus
            )

    return Liquid(Tref,ρref,SpecificWeight(ρref*1u"ge"),μ,KinematicViscosity(μ/ρref),σ,pv,Ev)
end

Density(g::Liquid) = g.ρref
Viscosity(g::Liquid) = g.μ
KinematicViscosity(g::Liquid) = g.ν
SpecificWeight(g::Liquid) = g.specwt
SurfaceTension(g::Liquid) = g.σ
VaporPressure(g::Liquid) = g.pv
BulkModulus(g::Liquid) = g.Ev

function Base.show(io::IO, m::MIME"text/plain", g::Liquid)
    println(io,"Liquid with")
    println(io,"   Density = $(value(Density(g)))")
    println(io,"   Viscosity = $(value(Viscosity(g)))")
    println(io,"   Kinematic viscosity = $(value(KinematicViscosity(g)))")
    println(io,"   Specific weight = $(value(SpecificWeight(g)))")
    println(io,"   Surface tension = $(value(SurfaceTension(g)))")
    println(io,"   Bulk modulus = $(value(BulkModulus(g)))")
    println(io,"   Vapor pressure = $(value(VaporPressure(g)))")
    println(io,"   at reference temp $(value(g.Tref))")
end

ext: jl

max_stars_repo_name: UCLAMAEThreads/ThermofluidQuantities

max_issues_repo_name: UCLAMAEThreads/ThermofluidQuantities

max_forks_repo_name: UCLAMAEThreads/ThermofluidQuantities


text:
"""
    uniformprobvec(n)

Return a uniform probability distribution vector of size `n`.
"""
function uniformprobvec(n)
    return ones(n) ./ n
end

"""
    uniformtransmat(n)

Return a stochastic matrix of size `n` with uniform transition probability distributions.
"""
function uniformtransmat(n)
    return ones(n, n) ./ n
end

"""
    randprobvec(n)

Return a random probability distribution vector of size `n`.
"""
function randprobvec(n)
    p = rand(n)
    return p ./ sum(p)
end

"""
    randtransmat(n)

Return a stochastic matrix of size `n` with random transition probability distributions.
"""
function randtransmat(n)
    P = rand(n, n)
    return P ./ sum(P, dims = 2)
end


text:
mutable struct ListNode{T}
    data::T
    prev::ListNode{T}
    next::ListNode{T}
    function ListNode{T}() where T
        node = new{T}()
        node.next = node
        node.prev = node
        return node
    end
    function ListNode{T}(data) where T
        node = new{T}(data)
        return node
    end
end

mutable struct MutableLinkedList{T}
    len::Int
    node::ListNode{T}
    function MutableLinkedList{T}() where T
        l = new{T}()
        l.len = 0
        l.node = ListNode{T}()
        l.node.next = l.node
        l.node.prev = l.node
        return l
    end
end

MutableLinkedList() = MutableLinkedList{Any}()

function MutableLinkedList{T}(elts...) where T
    l = MutableLinkedList{T}()
    for elt in elts
        push!(l, elt)
    end
    return l
end

Base.iterate(l::MutableLinkedList) = l.len == 0 ? nothing : (l.node.next.data, l.node.next.next)
Base.iterate(l::MutableLinkedList, n::ListNode) = n === l.node ? nothing : (n.data, n.next)

Base.isempty(l::MutableLinkedList) = l.len == 0
Base.length(l::MutableLinkedList) = l.len
Base.collect(l::MutableLinkedList{T}) where T = T[x for x in l]
Base.eltype(::Type{<:MutableLinkedList{T}}) where T = T
Base.lastindex(l::MutableLinkedList) = l.len

function Base.first(l::MutableLinkedList)
    isempty(l) && throw(ArgumentError("List is empty"))
    return l.node.next.data
end

function Base.last(l::MutableLinkedList)
    isempty(l) && throw(ArgumentError("List is empty"))
    return l.node.prev.data
end

Base.:(==)(l1::MutableLinkedList{T}, l2::MutableLinkedList{S}) where {T,S} = false

function Base.:(==)(l1::MutableLinkedList{T}, l2::MutableLinkedList{T}) where T
    length(l1) == length(l2) || return false
    for (i, j) in zip(l1, l2)
        i == j || return false
    end
    return true
end

function Base.map(f::Base.Callable, l::MutableLinkedList{T}) where T
    if isempty(l) && f isa Function
        S = Core.Compiler.return_type(f, Tuple{T})
        return MutableLinkedList{S}()
    elseif isempty(l) && f isa Type
        return MutableLinkedList{f}()
    else
        S = typeof(f(first(l)))
        l2 = MutableLinkedList{S}()
        for h in l
            el = f(h)
            if el isa S
                push!(l2, el)
            else
                R = typejoin(S, typeof(el))
                l2 = MutableLinkedList{R}(collect(l2)...)
                push!(l2, el)
            end
        end
        return l2
    end
end

function Base.filter(f::Function, l::MutableLinkedList{T}) where T
    l2 = MutableLinkedList{T}()
    for h in l
        if f(h)
            push!(l2, h)
        end
    end
    return l2
end

function Base.reverse(l::MutableLinkedList{T}) where T
    l2 = MutableLinkedList{T}()
    for h in l
        pushfirst!(l2, h)
    end
    return l2
end

function Base.copy(l::MutableLinkedList{T}) where T
    l2 = MutableLinkedList{T}()
    for h in l
        push!(l2, h)
    end
    return l2
end

function Base.getindex(l::MutableLinkedList, idx::Int)
    @boundscheck 0 < idx <= l.len || throw(BoundsError(l, idx))
    node = l.node
    for i in 1:idx
        node = node.next
    end
    return node.data
end

function Base.getindex(l::MutableLinkedList{T}, r::UnitRange) where T
    @boundscheck 0 < first(r) < last(r) <= l.len || throw(BoundsError(l, r))
    l2 = MutableLinkedList{T}()
    node = l.node
    for i in 1:first(r)
        node = node.next
    end
    len = length(r)
    for j in 1:len
        push!(l2, node.data)
        node = node.next
    end
    l2.len = len
    return l2
end

function Base.setindex!(l::MutableLinkedList{T}, data, idx::Int) where T
    @boundscheck 0 < idx <= l.len || throw(BoundsError(l, idx))
    node = l.node
    for i in 1:idx
        node = node.next
    end
    node.data = convert(T, data)
    return l
end

function Base.append!(l1::MutableLinkedList{T}, l2::MutableLinkedList{T}) where T
    l1.node.prev.next = l2.node.next # l1's last's next is now l2's first
    l2.node.prev.next = l1.node # l2's last's next is now l1.node
    l1.len += length(l2)
    return l1
end

function Base.append!(l::MutableLinkedList, elts...)
    for elt in elts
        push!(l, elt)
    end
    return l
end

function Base.delete!(l::MutableLinkedList, idx::Int)
    @boundscheck 0 < idx <= l.len || throw(BoundsError(l, idx))
    node = l.node
    for i = 1:idx
        node = node.next
    end
    prev = node.prev
    next = node.next
    prev.next = next
    next.prev = prev
    l.len -= 1
    return l
end

function Base.delete!(l::MutableLinkedList, r::UnitRange)
    @boundscheck 0 < first(r) < last(r) <= l.len || throw(BoundsError(l, r))
    node = l.node
    for i in 1:first(r)
        node = node.next
    end
    prev = node.prev
    len = length(r)
    for j in 1:len
        node = node.next
    end
    next = node
    prev.next = next
    next.prev = prev
    l.len -= len
    return l
end

function Base.push!(l::MutableLinkedList{T}, data) where T
    oldlast = l.node.prev
    node = ListNode{T}(data)
    node.next = l.node
    node.prev = oldlast
    l.node.prev = node
    oldlast.next = node
    l.len += 1
    return l
end

function Base.pushfirst!(l::MutableLinkedList{T}, data) where T
    oldfirst = l.node.next
    node = ListNode{T}(data)
    node.prev = l.node
    node.next = oldfirst
    l.node.next = node
    oldfirst.prev = node
    l.len += 1
    return l
end

function Base.pop!(l::MutableLinkedList)
    isempty(l) && throw(ArgumentError("List must be non-empty"))
    last = l.node.prev.prev
    data = l.node.prev.data
    last.next = l.node
    l.node.prev = last
    l.len -= 1
    return data
end

function Base.popfirst!(l::MutableLinkedList)
    isempty(l) && throw(ArgumentError("List must be non-empty"))
    first = l.node.next.next
    data = l.node.next.data
    first.prev = l.node
    l.node.next = first
    l.len -= 1
    return data
end

function Base.show(io::IO, node::ListNode)
    x = node.data
    print(io, "$(typeof(node))($x)")
end

function Base.show(io::IO, l::MutableLinkedList)
    print(io, typeof(l), '(')
    join(io, l, ", ")
    print(io, ')')
end


text:
# Get all distributions in Distributions.jl
distributions = map(subtypes(Distribution)) do t
    Symbol(t)
end
push!(distributions, :Mixture)
push!(distributions, :Product)

# Desugar "sugared" tilde expressions.
function _sugar(expr)
    MacroTools.postwalk(expr) do s
        if @capture(s, {addr_} ~ fn_(args__))
            if Symbol("Distributions.$fn") in distributions || fn in distributions
                k = Expr(:call, :trace, addr, Expr(:call, fn, args...))
            else
                k = Expr(:call, :trace, addr, fn, args...)
            end

        elseif @capture(s, val_ ~ fn_(args__))
            val isa Expr && error("Raw value assignment ~ for choice requires that value be a variable name (e.g. x, y, z, ...).")
            addr = QuoteNode(val)
            if Symbol("Distributions.$fn") in distributions || fn in distributions
                k = Expr(:(=), val, Expr(:call, :trace, addr, Expr(:call, fn, args...)))
            else
                k = Expr(:(=), val, Expr(:call, :trace, addr, fn, args...))
            end

        elseif @capture(s, val_ <- fn_(args__))
            k = quote $val = deep($fn, $(args...)) end

        else
            # Fallthrough.
            k = s
        end
        (unblock ∘ rmlines)(k)
    end
end

# Core Jaynes parser.
function _jaynes(def, opt)

    # Matches longdef function definitions.
    R = Any
    if @capture(def, function decl_(args__) body__ end) || @capture(def, function decl_(args__)::R_ body__ end)
        argtypes = Expr(:tuple, map(args) do a
            if a isa Expr
                a.head == :(::) ? a.args[2] : :Any
            else
                :Any
            end
        end...)
        trans = quote
            $def
            argtypes = $argtypes
            JFunction($opt(),
                      $decl,
                      argtypes,
                      tuple([false for _ in argtypes]...),
                      false,
                      $R)
        end
        trans

        # Matches thunks.
    elseif @capture(def, () -> body__)
        trans = quote
            JFunction($opt(),
                      $def,
                      (Tuple{}, ),
                      (false, ),
                      false,
                      Any)
        end

        # Matches lambdas with formals.
    elseif @capture(def, (args_) -> body__)
        if args isa Symbol
            argtypes = [Any]
        elseif args.head == :(::)
            argtypes = [eval(args.args[2])]
        else
            argtypes = map(args.args) do a
                if a isa Expr
                    a.head == :(::) ? eval(a.args[2]) : Any
                else
                    Any
                end
            end
        end
        trans = quote
            JFunction($opt(),
                      $def,
                      tuple($argtypes...),
                      tuple([false for _ in $argtypes]...),
                      false,
                      Any)
        end
    else
        error("ParseError (@jaynes): requires a longdef function definition or an anonymous function definition.")
    end

    MacroTools.postwalk(unblock ∘ rmlines, trans)
end

# @jaynes macro.
macro jaynes(expr, opt)
    def = _sugar(expr)
    trans = _jaynes(def, opt)
    esc(trans)
end

macro jaynes(expr)
    def = _sugar(expr)
    trans = _jaynes(def, :DefaultPipeline)
    esc(trans)
end


text:
# ----------------------------------------------------------
# ---------  Struct definitions and constructors  ----------
# ----------------------------------------------------------

# Default value of the FunctionOperatorComposite's plan field
# This variable is also used to check if FunctionOperatorComposite's plan field has the default value
noplan = () -> nothing

"""
Supertype for FunctionOperator and FunctionOperatorComposite
"""
abstract type FunOp end

# Returns the number of arguments
nargs(f::Function) = first(methods(f)).nargs
checkTwoInputs(f::Function) = (nargs(f) == 3)

mutable struct Counter
    num::Threads.Atomic{Int}
end
const counter = Counter(Threads.Atomic{Int}(1))
getNextNum() = Threads.atomic_add!(counter.num, 1)

"""
FunctionOperator is an operator that maps from a multidimensional space to another multidimensional space. The mapping is defined by a function (`forw`), and optionally the reverse mapping can also be defined (`backw`). The input the mapping must be subtype of AbstractArray.

The following constructors are available:
 - Positional constructor #1: `FunctionOperator{eltype}(forw, inDims, outDims)`
 - Positional constructor #2: `FunctionOperator{eltype}(forw, backw, inDims, outDims)`
 - Positional constructor #3: `FunctionOperator{eltype}(name, forw, inDims, outDims)`
 - Positional constructor #4: `FunctionOperator{eltype}(name, forw, backw, inDims, outDims)`
 - Keyword constructor: `FunctionOperator{eltype}(;kwargs...)`
where `eltype` is the type enforced on elements of input array.

Arguments
 - `name::String` (Optional but strongly recommended) The operator is referenced later in error messages by this string. **Warning!** It is also used to check equality of (composite) FunctionOperators. Default value: `OpX` where X is a number incremented in each constructor-call.
 - `forw::Function` Function defining the mapping. Must accept one or two arguments. In case of two arguments, the first argument is a preallocated buffer to write the result into (to speed up code by avoiding repeated allocations). In case of both one and two arguments, the return value must be the result of the mapping.
 - `backw::Function` (Optional) Same as backw, but defines the backward mapping
 - `inDims::Tuple{Vararg{Int}}` Size of input array
 - `outDims::Tuple{Vararg{Int}}` Size of output array
"""
@with_kw_noshow struct FunctionOperator{T <: Number} <: FunOp
    name::String = "Op$(getNextNum())"
    forw::Function
    backw::Function = (nargs(forw) == 2 ?
        (x) -> error("backward function not implemented for "*name) :
        (buffer, x) -> error("backward function not implemented for "*name))
    adjoint::Bool = false # adjoint operator creates a new object where this field is negated
    ismutating::Bool = checkTwoInputs(forw) # true if forw has two arguments
    scaling::Bool = false # true if created from LinearAlgebra.UniformScaling object
    getScale::Function = () -> nothing # This is used only if scaling field is true
    inDims::Tuple{Vararg{Int}}
    outDims::Tuple{Vararg{Int}}
    @assert (2 <= nargs(forw) <= 3)  "forw can only accept either one or two inputs"
    @assert (nargs(forw) == nargs(backw)) "forw and backw must accept the same number of inputs!"
end

# Constructor with positional arguments without default valued fields
FunctionOperator{T}(forw::Function,
        inDims::Tuple{Vararg{Int}}, outDims::Tuple{Vararg{Int}}) where {T} =
    FunctionOperator{T}(forw = forw, inDims = inDims, outDims = outDims)

# Constructor with positional arguments without backw
FunctionOperator{T}(name::String, forw::Function,
        inDims::Tuple{Vararg{Int}}, outDims::Tuple{Vararg{Int}}) where {T} =
    FunctionOperator{T}(name = name, forw = forw, inDims = inDims, outDims = outDims)

# Constructor with positional arguments without name
FunctionOperator{T}(forw::Function, backw::Function,
        inDims::Tuple{Vararg{Int}}, outDims::Tuple{Vararg{Int}}) where {T} =
    FunctionOperator{T}(forw = forw, backw = backw, inDims=inDims, outDims=outDims)

# Constructor with positional arguments with all public fields
FunctionOperator{T}(name::String, forw::Function, backw::Function,
        inDims::Tuple{Vararg{Int}}, outDims::Tuple{Vararg{Int}}) where {T} =
    FunctionOperator{T}(name = name, forw = forw, backw = backw, inDims = inDims, outDims = outDims)

# Named array for nicer string representation of computed plan.
# The name is always "bufferX" where X is a number which is incremented at each allocation
mutable struct Buffer
    buffer::Array
    name::String
    number::Int64
    available::Bool
end

# Structure holding a combination of FunctionOperators
# It is mutable because plan, plan_string and plan_buffer fields are changed
# after applying the operator to an array
@with_kw_noshow mutable struct FunctionOperatorComposite{T <: Number} <: FunOp
    name::String
    left::FunOp
    right::FunOp
    operator::Symbol
    adjoint::Bool = false
    ismutating::Bool = true
    inDims::Tuple{Vararg{Int}} = right.inDims
    outDims::Tuple{Vararg{Int}} = left.outDims
    plan_function::Function = noplan
    plan_string::String = "no plan"
end

# Get the value of the name field and append ' if it is adjoint
getName(FO::FunctionOperator) = FO.adjoint ? FO.name*"'" : FO.name
getName(FO::FunctionOperatorComposite) = FO.adjoint ? "($(FO.name))'" : FO.name

# Constructor that combines two FunOp objects
function FunctionOperatorComposite(FO1::FunOp, FO2::FunOp, op::Symbol)
    name = getName(FO1) * " $op " * getName(FO2)
    name = op == :* ? name : "($name)"
    FunctionOperatorComposite{eltype(FO1)}(name = name, left = FO1, right = FO2, operator = op)
end


text:
using BenchmarkTools

for i in 1:5
    sleep(1)
end

@time for i in 1:5
    sleep(1)
end

@time for i in 1:5
    @async sleep(1)
end

@time @sync for i in 1:5
    @async sleep(1)
end

@time for i in 1:50
    sin.(rand(1000, 1000))
end

@time @sync for i in 1:50
    @async sin.(rand(1000, 1000))
end

@time @sync for i in 1:5
    @async ccall(("sleep", :libc), Cint, (Cint, ), 1)
end

# ## Task Lifecycle

t=Task(()->println("Hello from tasks"))

schedule(t)

istaskstarted(t)

schedule(t)

istaskdone(t)

istaskstarted(t)

current_task()

t == current_task()

# ## Task local storage

task_local_storage("x", 1)

task_local_storage("x") == 1

# # Communicating between tasks

c = Channel{Int}(10)

function producer(c::Channel)
   put!(c, "start")
   for n=1:4
      put!(c, 2n)
   end
   put!(c, "stop")
end

chnl = Channel(producer)

take!(chnl)
take!(chnl)
take!(chnl)
take!(chnl)
take!(chnl)
take!(chnl)

# ## Task Iteration

chnl = Channel(producer)
for i in chnl
    @show i
end

function consume(c)
   println("Starting Channel iteration")
   for i in c
       println("Got $i from Channel")
   end
   println("Channel iteration is complete")
end

chnl = Channel(1)

@async consume(chnl)

put!(chnl, 1)

put!(chnl, 2)

close(chnl)

using Pkg
Pkg.add("Mux")

using Mux
@app basicapp = (
         Mux.defaults,
         page("/", respond("
Hello World!
")),
         Mux.notfound())

serve(basicapp; reuseaddr=true)


text:
# Starting functions example for Learning Julia

# Functions are defined with the function keyword and are usually
# lowercase names, optionally with underscores if they are hard to read
function myFunc(a, b)
    println("This is a function")
    a + b # Julia automatically returns the last line of statement
          # so this last line is equivalent to return a + b
end

result1 = myFunc(10, 15)
println(result1)

# function arguments can have default values
function foo(a, b, z = 10)
    return (a + b) * z
end
println(foo(2, 3))
println(foo(2, 3, 5))

# you can also use keyword arguments - define them after a semicolon
function bar(a, b; multiplier = 10) # after ;, arguments must be called explicitly with the name
    return (a + b) * multiplier
end
println(bar(multiplier = 5, 4, 5))

# The Julia shorthand way of defining a function
anotherFunc(x, y) = (a = x - 1 ; 2a + y) # last argument evaluated is returned
result2 = anotherFunc(3, 4)
println(result2)

# use the ... notation for variable arguments
function summit(args...)
    sum = 0

    # process each argument
    for a in args
        sum += a
    end

    return sum
end
println(summit(1, 5, 10))
println(summit(2, 4, 6, 8))

max_stars_repo_name: scaotravis/Julia-Course

max_issues_repo_name: scaotravis/Julia-Course

max_forks_repo_name: scaotravis/Julia-Course


text:
include("simplexoperations.jl")

"""
    `nontrivial_simplexpair()`

Genereate a set of non-trivially intersecting simplices (i.e. they don't intersect along boundaries or vertices only).
"""
function nontrivial_simplexpair()

end


text:
using FNVHash
using Test

@testset "FNVHash.jl" begin

@testset "type $T" for T in [UInt32, UInt64, UInt128]
    @test fnv1(T, "test") isa T
    @test fnv1a(T, "test") isa T
end

@testset "empty $T" for T in [UInt32, UInt64, UInt128]
    @test fnv1(T, "") == fnv_offset_basis(T)
    @test fnv1a(T, "") == fnv_offset_basis(T)
end

@testset "vector/string equivalence $T" for T in [UInt32, UInt64, UInt128]
    data = UInt8['a', 'b', 'c']
    s = String(copy(data))
    @test fnv1(T, data) == fnv1(T, s)
    @test fnv1a(T, data) == fnv1a(T, s)
end

@testset "changing $T" for T in [UInt32, UInt64, UInt128]
    @test fnv1(T, "test") != fnv1(T, "tesu")
    @test fnv1a(T, "test") != fnv1a(T, "tesu")
end

end


text:
# # I/O related functions
#
# The following function are used to read datastructures from KM3NeT
# dataformats like HDF5 or custom binary types.

"""
    function read_compound(dset::HDF5.HDF5Dataset, T::DataType)

Read an HDF5Compund structure from an HDF5 dataset.
"""
function read_compound(dset::HDF5.HDF5Dataset, T::DataType)
    filetype = HDF5.datatype(dset) # packed layout on disk
    memtype_id = HDF5.h5t_get_native_type(filetype.id) # padded layout in memory
    @assert sizeof(T) == HDF5.h5t_get_size(memtype_id) "Type sizes don't match!"
    out = Vector{T}(undef, length(dset))
    HDF5.h5d_read(dset.id, memtype_id, HDF5.H5S_ALL, HDF5.H5S_ALL, HDF5.H5P_DEFAULT, out)
    HDF5.h5t_close(memtype_id)
    out
end

function read_compound(filename::AbstractString,
                       h5loc::AbstractString,
                       T::DataType)
    fobj = HDF5.h5open(filename)
    data = read_compound(fobj, h5loc, T)
    close(fobj)
    data
end

function read_compound(fobj::HDF5.HDF5File,
                       h5loc::AbstractString,
                       T::DataType)
    data = read_compound(fobj[h5loc], T)
    data
end

function read_hits(fobj::HDF5.HDF5File, idx::Int, n_hits::Int)
    hits = Vector{Hit}()
    channel_id = fobj["hits/channel_id"][idx+1:idx+n_hits]
    dom_id = fobj["hits/dom_id"][idx+1:idx+n_hits]
    t = fobj["hits/time"][idx+1:idx+n_hits]
    tot = fobj["hits/tot"][idx+1:idx+n_hits]
    triggered = fobj["hits/triggered"][idx+1:idx+n_hits]
    for i ∈ 1:n_hits
        hit =  Hit(channel_id[i], dom_id[i], t[i], tot[i], triggered[i])
        push!(hits, hit)
    end
    return hits
end

function read_hits(fobj::HDF5.HDF5File, group_id::Int)
    hit_indices = read_indices(fobj, "/hits")
    idx = hit_indices[group_id+1][1]
    n_hits = hit_indices[group_id+1][2]
    hits = read_hits(fobj, idx, n_hits)::Vector{Hit}
    hits
end

function read_hits(filename::AbstractString, group_id::Int)
    f = h5open(filename, "r")
    hits = read_hits(f, group_id)
    close(f)
    return hits
end

function read_hits(filename::AbstractString,
                    group_ids::Union{Array{T}, UnitRange{T}}) where {T<:Integer}
    f = h5open(filename, "r")
    hit_indices = read_indices(f, "/hits")

    hits_collection = Dict{Int, Vector{SnapshotHit}}()
    for group_id ∈ group_ids
        idx = hit_indices[group_id+1][1]
        n_hits = hit_indices[group_id+1][2]
        hits = read_hits(f, idx, n_hits)::Vector{SnapshotHit}
        hits_collection[group_id] = hits
    end
    close(f)
    return hits_collection
end

function read_hits(f::DAQEventFile, group_id)
    idx = f._hit_indices[group_id+1][1]
    n_hits = f._hit_indices[group_id+1][2]
    read_hits(f._fobj, idx, n_hits)::Vector{SnapshotHit}
end

function read_mchits(fobj::HDF5.HDF5File, idx::Int, n_hits::Int)
    hits = Vector{MCHit}()
    a = fobj["mc_hits/a"][idx+1:idx+n_hits]
    origin = fobj["mc_hits/origin"][idx+1:idx+n_hits]
    pmt_id = fobj["mc_hits/pmt_id"][idx+1:idx+n_hits]
    t = fobj["mc_hits/time"][idx+1:idx+n_hits]
    for i ∈ 1:n_hits
        hit =  MCHit(a[i], origin[i], pmt_id[i], t[i])
        push!(hits, hit)
    end
    return hits
end

function read_mchits(filename::AbstractString, group_id::Int)
    f = h5open(filename, "r")
    hit_indices = read_indices(f, "/mc_hits")
    idx = hit_indices[group_id+1][1]
    n_hits = hit_indices[group_id+1][2]
    hits = read_mchits(f, idx, n_hits)::Vector{MCHit}
    close(f)
    return hits
end

function read_mchits(filename::AbstractString,
                    group_ids::Union{Array{T}, UnitRange{T}}) where {T<:Integer}
    f = h5open(filename, "r")
    hit_indices = read_indices(f, "/mc_hits")

    hits_collection = Dict{Int, Vector{MCHit}}()
    for group_id ∈ group_ids
        idx = hit_indices[group_id+1][1]
        n_hits = hit_indices[group_id+1][2]
        hits = read_hits(f, idx, n_hits)::Vector{MCHit}
        hits_collection[group_id] = hits
    end
    close(f)
    return hits_collection
end

function read_mctracks(filename::AbstractString)
    mc_tracks = Dict{Int64}{Vector{MCTrack}}()
    for track in NeRCA.read_compound(filename, "/mc_tracks", MCTrack)
        group_id = track.group_id
        if !haskey(mc_tracks, group_id)
            mc_tracks[group_id] = Vector{MCTrack}()
        end
        push!(mc_tracks[group_id], track)
    end
    mc_tracks
end

function read_event_info(f::DAQEventFile, group_id)
    f._event_infos[group_id + 1]
end

function read_indices(filename::AbstractString, from::AbstractString)
    f = h5open(filename, "r")
    indices = read_indices(f, from)
    close(f)
    return indices
end

function read_indices(fobj::HDF5.HDF5File, from::AbstractString)
    idc = read(fobj, from * "/_indices")
    indices = [i.data for i ∈ idc]::Vector{Tuple{Int64,Int64}}
    return indices
end

function read_calibration(filename::AbstractString)
    lines = readlines(filename)
    filter!(e->!startswith(e, "#") && !isempty(strip(e)), lines)

    if 'v' ∈ first(lines)
        det_id, version = map(x->parse(Int,x), split(first(lines), 'v'))
        n_doms = parse(Int, lines[4])
        idx = 5
    else
        det_id, n_doms = map(x->parse(Int,x), split(first(lines)))
        version = 1
        idx = 2
    end

    pos = Dict{Int32,Vector{NeRCA.Position}}()
    dir = Dict{Int32,Vector{NeRCA.Direction}}()
    t0s = Dict{Int32,Vector{Float64}}()
    dus = Dict{Int32,UInt8}()
    floors = Dict{Int32,UInt8}()
    omkeys = Dict{Int32,OMKey}()

    max_z = 0.0
    for dom ∈ 1:n_doms
        dom_id, du, floor, n_pmts = map(x->parse(Int,x), split(lines[idx]))
        pos[dom_id] = Vector{NeRCA.Position}()
        dir[dom_id] = Vector{NeRCA.Direction}()
        t0s[dom_id] = Vector{Float64}()
        dus[dom_id] = du
        floors[dom_id] = floor

        for pmt in 1:n_pmts
            l = split(lines[idx+pmt])
            pmt_id = parse(Int,first(l))
            x, y, z, dx, dy, dz = map(x->parse(Float64, x), l[2:7])
            max_z = max(max_z, z)
            t0 = parse(Float64,l[8])
            push!(pos[dom_id], Position(x, y, z))
            push!(dir[dom_id], Direction(dx, dy, dz))
            push!(t0s[dom_id], t0)
            omkeys[pmt_id] = OMKey(dom_id, pmt-1)
        end
        idx += n_pmts + 1
    end
    n_dus = length(unique(values(dus)))
    Calibration(det_id, pos, dir, t0s, dus, floors, omkeys, max_z, n_dus)
end

# Triggers
is_3dmuon(e::DAQEvent) = Bool(e.trigger_mask & 16 > 0)
is_3dshower(e::DAQEvent) = Bool(e.trigger_mask & 2 > 0)
is_mxshower(e::DAQEvent) = Bool(e.trigger_mask & 4 > 0)


text:
## Mimi UI
using VegaLite
using FilePaths

global app = nothing

#include functions and modules
include("buildspecs.jl")

"""
    explore(m::Model; title = "Electron")

Produce a UI to explore the parameters and variables of Model `m` in a Window with title `title`.
"""
function explore(m::Model; title = "Electron")

    if m.mi === nothing
        error("A model must be run before it can be plotted")
    end

    #start Electron app
    if app === nothing
        global app = Application()
    end

    #load main html file
    mainpath = replace(joinpath(@__DIR__, "assets", "main.html"), "\\" => "/")

    #window options
    windowopts = Dict("title" => title, "width" => 1000, "height" => 700)
    w = Window(app, URI(joinpath(@__PATH__, "assets", "main.html")), options = windowopts)

    #set async block to process messages
    @async for msg in msgchannel(w)

        spec = _spec_for_item(m, Symbol(msg["comp_name"]), Symbol(msg["item_name"]))
        specJSON = JSON.json(spec)

        run(w, "display($specJSON)")
    end

    #refresh variable list
    menulist = menu_item_list(m)
    menulistJSON = JSON.json(menulist)

    result = run(w, "refresh($menulistJSON)")

    return w

end

"""
    explore(m::Model, comp_name::Symbol, datum_name::Symbol)

Plot a specific `datum_name` (a `variable` or `parameter`) of Model `m`.
"""
function explore(m::Model, comp_name::Symbol, datum_name::Symbol)

    if m.mi === nothing
        error("A model must be run before it can be plotted")
    end

    spec = Mimi._spec_for_item(m, comp_name, datum_name)["VLspec"]
    spec === nothing && error("Spec cannot be built.")

    return VegaLite.VLSpec{:plot}(spec)
end


text:
# This file was generated, do not modify it. # hide
@show krb.knn_weight
@show krb.knn_model.K
@show krb.ridge_model.lambda

max_stars_repo_name: tlienart/MLJTutorials

max_issues_repo_name: tlienart/MLJTutorials

max_forks_repo_name: tlienart/MLJTutorials


text:
using DifferentialEquations
using Printf
using Plots
using LaTeXStrings
gr()
# define the set of odes
function stellar_structure!(du, u, params, m)
    a, b, n, s, nu, lambda, A = params
    p, x, t, l = u
    du[1] = - (m / x^4)
    du[2] = (t^b / (x^2 * p^a))
    du[3] = - (p^(a * n) * l) / (x^4 * t^(3 + s + b * n))
    du[4] = A * p^(a * lambda) * t^(nu - (b * lambda))
end

# We define the condition to find the event where p becomes less than some tolerace set by user.
condition(u, m, integrator) =  u[1] - 1e-4
affect!(integrator) = terminate!(integrator)
cb = ContinuousCallback(condition, affect!)

# Initial conditions, at m = 0, as per eq 1.2
p0 = 1.0
x0 = 1e-4
t0 = 1.0
l0 = 0.0

u0 = [p0, x0, t0, l0]

# parameters
a = 1.0
b = 1.0
n = 0.0
s = 0.0
nu = 4.0
lambda = 1.0
A = 0.535

params = [a, b, n, s, nu, lambda, A]
mspan = (0, 20.0)
problem1 = ODEProblem(stellar_structure!, u0, mspan, params)
solution1 = solve(problem1, alg_hints=[:stiff], callback=cb)
println(solution1.t[end], "\t", solution1.u[end][1])
plot(solution1, vars=(0, 1), lw = 2, xlabel = L"m", label = L"p")
plot!(solution1, vars=(0, 2), lw = 2, xlabel = L"m", label = L"x")
plot!(solution1, vars=(0, 3), lw = 2, xlabel = L"m", label = L"t")
plot!(solution1, vars=(0, 4), lw = 2, xlabel = L"m", label = L"l")
plot!(title = L"A=%$A", titlefontsize = 10, xlims = [0, 12], ylims = [0, 5])
savefig("p_vs_m_A_$A.pdf")