using Compose
using Test

include("misc.jl")
@testset "SVG Correctness Tests" begin
    include("svg.jl")
end
include("immerse.jl")

# Run the examples
cd(joinpath(@__DIR__, "output"))

exampledir = joinpath(@__DIR__, "examples")
for ex in readdir(exampledir)
    endswith(ex, ".jl") || continue
    file = joinpath(exampledir, ex)
    expr = quote
        module $(Symbol(replace(ex, "."=>"_")))
        using Compose
        using Colors
        using Random
        Random.seed!(1) #Needed so that SVG uuid is reproducible
        include($file)
        end
    end
    expr.head = :toplevel
    eval(expr)
end


using Plots
using Oceananigans

using Statistics

"""
    make_vertical_slice_movie(model::Model, nc_writer::NetCDFOutputWriter,
                              var_name, Nt, Δt, var_offset=0, slice_idx=1)

Make a movie of a vertical slice produced by `model` with output being saved by
`nc_writer`. The variable name `var_name` can be either of "u", "v", "w", "T",
or "S". `Nt` is the number of model iterations (or time steps) taken and ``Δt`
is the time step. A plotting offset `var_offset` can be specified to be
subtracted from the data before plotting (useful for plotting e.g. small
temperature perturbations around T₀). A `slice_idx` can be specified to select
the index of the y-slice to be plotted (useful when plotting vertical slices
from a 3D model, it should be set to 1 for 2D xz-slice models).
"""
function make_vertical_slice_movie(model::Model, nc_writer::NetCDFOutputWriter, var_name, Nt, Δt, var_offset=0, slice_idx=1)
    freq = nc_writer.output_frequency
    N_frames = Int(Nt/freq)

    print("Producing movie... ($N_frames frames)\n")
    Plots.gr(dpi=150)

    animation = @animate for n in 0:N_frames
        print("\rframe = $n / $N_frames   ")
        var = read_output(nc_writer, var_name, freq*n)
        Plots.heatmap(model.grid.xC, model.grid.zC, rotl90(var[:, slice_idx, :]) .- var_offset,
                      color=:balance, clims=(-0.01, 0.01), title="t=$(freq*n*Δt) s ($(round(freq*n*Δt/86400; digits=2)) days)")
    end

    mp4(animation, nc_writer.filename_prefix * "$(round(Int, time())).mp4", fps=30)
end

"""
    make_horizontal_slice_movie(model::Model, nc_writer::NetCDFOutputWriter,
                                var_name, Nt, Δt, var_offset=0)

Make a movie of a horizontal slice produced by `model` with output being saved by
`nc_writer`. The variable name `var_name` can be either of "u", "v", "w", "T",
or "S". `Nt` is the number of model iterations (or time steps) taken and ``Δt`
is the time step. A plotting offset `var_offset` can be specified to be
subtracted from the data before plotting (useful for plotting e.g. small
temperature perturbations around T₀).
"""
function make_horizontal_slice_movie(model::Model, nc_writer::NetCDFOutputWriter, var_name, Nt, Δt, var_offset=0)
    freq = nc_writer.output_frequency
    N_frames = Int(Nt/freq)

    print("Producing movie... ($N_frames frames)\n")
    Plots.gr(dpi=150)

    animation = @animate for n in 0:N_frames
        print("\rframe = $n / $N_frames   ")
        var = read_output(nc_writer, var_name, freq*n)
        Plots.heatmap(model.grid.xC, model.grid.yC, var[:, :, 1] .- var_offset,
                      color=:balance, clims=(-0.01, 0.01),
                      title="t=$(freq*n*Δt) s ($(round(freq*n*Δt/86400; digits=2)) days)")
    end

    mp4(animation, nc_writer.filename_prefix * "$(round(Int, time())).mp4", fps=30)
end

"""
    make_vertical_profile_movie(model::Model, nc_writer::NetCDFOutputWriter,
                                var_name, Nt, Δt, var_offset=0)

Make a movie of a vertical profile produced by `model` with output being saved by
`nc_writer`. The variable name `var_name` can be either of "u", "v", "w", "T",
or "S". `Nt` is the number of model iterations (or time steps) taken and ``Δt`
is the time step. A plotting offset `var_offset` can be specified to be
subtracted from the data before plotting (useful for plotting e.g. small
temperature perturbations around T₀).
"""
function make_vertical_profile_movie(model::Model, nc_writer::NetCDFOutputWriter, var_name, Nt, Δt, var_offset=0)
    freq = nc_writer.output_frequency
    N_frames = Int(Nt/freq)

    print("Producing movie... ($N_frames frames)\n")
    Plots.gr(dpi=150)

    animation = @animate for n in 0:N_frames
        print("\rframe = $n / $N_frames   ")
        var = read_output(nc_writer, var_name, freq*n)
        Plots.plot(var[1, 1, :] .- var_offset, model.grid.zC,
                   title="t=$(freq*n*Δt) s ($(round(freq*n*Δt/86400; digits=2)) days)")
    end

    mp4(animation, nc_writer.filename_prefix * "$(round(Int, time())).mp4", fps=30)
end

using NetCDF

function make_avg_temperature_profile_movie()
    Nt, dt = 86400, 0.5
    freq = 3600
    N_frames = Int(Nt/freq)
    filename_prefix = "convection"
    var_offset = 273.15

    Nz, Lz = 128, 100
    dz = Lz/Nz
    zC = -dz/2:-dz:-Lz

    print("Producing movie... ($N_frames frames)\n")
    Plots.gr(dpi=150)

    animation = @animate for n in 0:N_frames
        print("\rframe = $n / $N_frames   ")

        filepath = filename_prefix * lpad(freq*n, 9, "0") * ".nc"
        field_data = ncread(filepath, "T")
        ncclose(filepath)

        T_profile = mean(field_data; dims=[1,2])

        Plots.plot(reshape(T_profile, Nz) .- var_offset, zC,
                   title="t=$(freq*n*dt) s ($(round(freq*n*dt/86400; digits=2)) days)")
    end

    mp4(animation, filename_prefix * "$(round(Int, time())).mp4", fps=30)
end


using JWebImageDemo
using FileIO
using Test

const TEST_IMAGE = joinpath(@__DIR__, "resources/green_python.jpg")
const NUM = 5

const TMP_PATH = "tmp"
isdir(TMP_PATH) || mkdir(TMP_PATH)

@testset "image permute test" begin
    img = FileIO.load(TEST_IMAGE)
    for i = 2:NUM
        @time img_out = JWebImageDemo.permute_image(img, i)
        # save(Stream{format"PNG"}(io), data...)
        @test size(img) != size(img_out) ||
            any(((i, o),) -> i != o, zip(iterate(img), iterate(img_out)))
        save(joinpath(TMP_PATH, "out_$i.jpg"), img_out)
    end
end


# Code derived from metabolicEP (https://github.com/anna-pa-m/Metabolic-EP)

function read_mat(filename::String; model_id = nothing)
    mat_dict = MAT.matread(filename)
    if isnothing(model_id)
        model_id, model_mat = first(mat_dict)
    else
        model_mat = mat_dict[model_id]
    end
    return MetNet(model_mat; reshape = true)
end


using VLConstraintBasedModelGenerationUtilities
using DataFrames

# setup path to sequence file -
path_to_protein_sequence_file = "/Users/jeffreyvarner/Desktop/julia_work/VLConstraintBasedModelGenerationUtilities.jl/test/data/P-MZ373340.fasta"

# load -
protein_table = build_protein_table(path_to_protein_sequence_file) |> check
seq = protein_table[1,:protein_sequence]
seq_name = protein_table[1,:id]

# build the translation table -
translation_dictionary = Dict{String,DataFrame}()
(number_of_proteins, _) = size(protein_table)
for protein_index = 1:number_of_proteins

    local_seq = protein_table[protein_index,:protein_sequence]
    local_seq_name = protein_table[protein_index,:id]
    translation_table = build_translation_reaction_table(local_seq_name, local_seq) |> check

    # grab -
    translation_dictionary[local_seq_name] = translation_table
end


# include("HTgnss.jl")

path_rinex = abspath("../data_test/")

function test_obs()

    files = vcat(
                glob("*.rnx", path_rinex),
                glob("*.*o", path_rinex),
                    )

    sort!(files)
    ifile = files[1]
    this_file = abspath(ifile)
    ohd = HTgnss.read_obsheader(this_file)
    oh = HTgnss.parse_obsheader(ohd)
    # @show oh
    @time iobs = HTgnss.load_obs(this_file);
end
function test_nav()
    files = vcat(
                # glob("*.*[nN]", path_rinex),
                # glob("*.*[gG]", path_rinex),
                glob("*.*[pP]", path_rinex),
                )
    sort!(files)
    ifile = files[1]
    this_file = abspath(ifile)

    # nhd = HTgnss.read_navheader(this_file)
    # nh  = HTgnss.parse_navheader(nhd)
    @time inav = HTgnss.load_nav(this_file);
end

function test_pos()

    oh = test_obs()
    nh = test_nav()
    ndata = nh.data.data
    ks = sort!(collect(keys(oh.data.obs)), by=x -> x.str)
    isat = ks[30]
    idata = oh.data.obs[isat]
    dt = idata.dts[2]
    ees = ndata[isat]
    # merge(dict1,dict2)   data::Dict{SatID,EphEpochStore}
    x = HTgnss.calculate_pos(ees, dt)
    @show ees.si, dt
    @test x[1] ≈ 1.764682259684179e7
    @test x[2] ≈ 3.452196727133614e6
    @test x[3] ≈ 1.935522316823225e7

    for (isat, idata) in oh.data.obs
        break
    end

    # keys(nh.data.data)
    # nh.data.data[HTgnss.SatID("G11")].eph

end
function test_2017251()
    #=
    using Glob
    using HTgnss
    using Dates =#
    p = "/home/t/RD/GNSS/STEC/-1_00_01/2017/251/"
    fs = glob("brdm*.17p", p)
    mnh, mnd = HTgnss.load_navs(fs)
    dt = DateTime("2017-09-06T23:59:42")
    sats = sort!(collect(keys(mnd)), by=x -> x.str)
    for (i, k) in enumerate(sats)
        if string(k) == "S36"
            @show i, k #  (111, S36)
        end
    end
    isat = sats[111]
    ees = mnd[isat]
    # merge(dict1,dict2)   data::Dict{SatID,EphEpochStore}
    x = HTgnss.calculate_pos(ees, dt)

end


# this is copied and adapted from `constraints.jl`
"""
    supports_objective( model::ModelLike,
		::Type{F})::Bool where {F<MOI.AbstractFunction}

Return a `Bool` indicating whether `model` generally supports adding an objective
of type `F`.
That is, `copy_to(model, src)` does not throw [`UnsupportedObjective`](@ref) when
`src` contains objectives of type `F`.
"""
function MOI.supports_constraint(
    ::MOI.ModelLike,
    ::Type{<:MOI.AbstractFunction}
)
    return false
end

"""
    struct UnsupportedObjective{F<:MOI.AbstractFunction}
        message::String # Human-friendly explanation why the attribute cannot be set
    end

	An error indicating that objectives of type `F` are not supported by
the model, i.e. that [`supports_constraint`](@ref) returns `false`.
"""
struct UnsupportedObjective{F<:MOI.AbstractFunction} <: MOI.UnsupportedError
    message::String # Human-friendly explanation why the attribute cannot be set
end
UnsupportedObjective{F}() where F = UnsupportedObjective{F}("")

function MOI.element_name(::UnsupportedObjective{F}) where {F}
    return "`$F` objective"
end

"""
    struct AddObjectiveNotAllowed{F<:AbstractFunction} <: NotAllowedError
        message::String # Human-friendly explanation why the attribute cannot be set
    end
An error indicating that objectives of type `F` are supported (see
[`supports_constraint`](@ref)) but cannot be added.
"""
struct AddObjectiveNotAllowed{F<:MOI.AbstractFunction} <:
       MOI.NotAllowedError
    message::String # Human-friendly explanation why the attribute cannot be set
end
AddObjectiveNotAllowed{F}() where {F} = AddConstraintNotAllowed{F}("")

function MOI.operation_name(::AddObjectiveNotAllowed{F}) where F
    return "Adding `$F` objectives."
end

# This should be implemented by multiobjective solvers, i.e., if
# `get( model, NumberOfObjectiveOutputs ) > 1`
# currently, it falls back to the old way of adding a single objective
function add_objective( model :: MOI.ModelLike, objf :: F ) where {F<:MOI.AbstractFunction}
	# this is a fallback to the old behavior:
	MOI.set( model, MOI.ObjectiveFunction{F}(), objf )

	# do some smart error handling as in src/constraints.jl
	return ObjectiveIndex{F}(1)
end

function add_objectives( model :: MOI.ModelLike, objfs :: Vector{<:MOI.AbstractFunction} )
	return add_objective.( model, ois, objfs )
end

# Default attributes

# overwrite this for a multi-objective solver:
MOI.get(model::MOI.ModelLike, ::MaxOutputs) = 1
MOI.get(model::MOI.ModelLike, ::NumberOfObjectives) = 1
MOI.get(model::MOI.ModelLike, ::NumberOfObjectiveOutputs) = 1

# default fallback to get number of outputs of a single objective
# using `output_dimension` which is provided for all MOI functions
function MOI.get( model :: MOI.ModelLike, :: NumberOfObjectiveOutputs, oi :: ObjectiveIndex)
	objf_fun = MOI.get( model, ObjectiveFunction(), oi )
	return MOI.output_dimension( objf_fun )
end

# default fallback to get number of outputs
function MOI.get( model :: MOI.ModelLike, :: NumberOfOutputs, oi :: ObjectiveIndex )
	if MOI.supports( model, ListOfObjectiveIndices() )
		output_counter = 0
		for oi in MOI.get( model, ListOfObjectiveIndices() )
			objf_fun = MOI.get( model, oi )
			output_counter += MOI.get( model, NumberOfObjectiveOutputs(), objf_fun )
		end
		return output_counter
	else
		throw(MOI.UnsupportedAttribute( NumberOfOutputs(),
				"Cannot determine number of outputs because `ListOfObjectivesIndices` is not supported."))
	end
end

# TODO `transform`


import Automa
import Automa.RegExp: @re_str
const re = Automa.RegExp

machine = (function ()
    # Primitives
    label      = re"[^\(\);,:]+"
    len        = re":[^\(\);,:]+"
    cladestart = re"\("
    cladestop  = re"\)"
    sep        = re","
    finish     = re";"

    # Composites
    metadata   = (label * len) | label | len
    nodechange = cladestart | cladestop | sep
    element    = nodechange * re.opt(metadata)
    tree       = re.opt(metadata) * re.rep(element) * finish

    # Actions
    label.actions[:enter] = [:mark]
    label.actions[:exit]  = [:label]
    len.actions[:enter] = [:mark]
    len.actions[:exit] = [:len]
    cladestart.actions[:enter] = [:push]
    cladestop.actions[:enter] = [:pop]
    sep.actions[:enter] = [:pop, :push]
    finish.actions[:enter] = [:finish]

    # Finally, compile the final FASTA pattern into a state machine.
    return Automa.compile(tree)
end)()

mutable struct Node
    label::Union{Nothing, String}
    len::Union{Nothing, Float64}
    parent::Union{Nothing, Node}
end

function Base.show(io::IO, node::Node)
    parent = node.parent === nothing ? nothing : "Node(\"$(node.parent.label)\")"
    print(io, "Node(\"$(node.label)\", $(node.len), $parent)")
end

@noinline function popstack!(stack::Vector{Node})
    length(stack) > 1 || error("Too many closing parenthesis in file")
    pop!(stack)
    return nothing
end

actions = Dict(
    :finish => :(finished = true),
    :mark => :(mark = p),
    :label => :(last(stack).label = String(data[mark:p-1])),
    :len => :(last(stack).len = Base.parse(Float64, String(data[mark+1:p-1]))),
    :pop => :(popstack!(stack)),
    :push => quote
        push!(stack, Node(nothing, nothing, last(stack)))
        push!(nodes, last(stack))
    end
)

context = Automa.CodeGenContext()
@eval function parse(data::Union{String,Vector{UInt8}})
    mark = 0
    nodes = [Node(nothing, nothing, nothing)]
    stack = copy(nodes)
    finished = false

    $(Automa.generate_init_code(context, machine))
    p_end = p_eof = lastindex(data)
    $(Automa.generate_exec_code(context, machine, actions))

    if (cs != 0) & !finished
        error("failed to parse on byte ", p)
    end

    # Final stack should only have the root node.
    if length(stack) != 1
        error("Not enough closing parentheses in file")
    end

    return nodes
end


@testset "SpatialElem" begin
    @test SI.check_hasmbr(SI.Rect{Float64, 2}, SpatialElem{Float64, 2, Int, Int})
    @test_throws ArgumentError SI.check_hasmbr(SI.Rect{Float64, 2}, Int)
    @test_throws ArgumentError SI.check_hasmbr(SI.Rect{Float64, 3}, SpatialElem{Float64, 2, Int, Int})
    @test_throws ArgumentError SI.check_hasmbr(SI.Rect{Float64, 3}, SpatialElem{Int, 2, Int, Int})

    @test SI.check_hasid(Int, SpatialElem{Float64, 2, Int, Int})
    @test_throws ArgumentError SI.check_hasid(String, SpatialElem{Float64, 2, Int, Int})
    @test_throws ArgumentError SI.check_hasid(Int, SpatialElem{Float64, 2, Nothing, Int})
end

@testset "RTree" begin
@testset "Basic Operations" begin
    tree_vars = [SI.RTreeStar, SI.RTreeLinear, SI.RTreeQuadratic]
    @testset "RTree{Float,2,Int32,Int}(variant=$tree_var)" for tree_var in tree_vars
        ambr = SI.Rect((0.0, 0.0), (0.0, 0.0))
        bmbr = SI.Rect((0.0, 1.0), (0.0, 1.0))
        cmbr = SI.Rect((0.5, 0.5), (0.5, 0.6))

        tree = RTree{Float64, 2}(Int32, Int, variant=tree_var)
        @test tree isa RTree{Float64, 2, SpatialElem{Float64, 2, Int32, Int}}
        @test SI.variant(tree) === tree_var

        @test eltype(tree) === SpatialElem{Float64, 2, Int32, Int}
        @test SI.regiontype(tree) === SI.Rect{Float64, 2}
        @test SI.dimtype(tree) === Float64
        @test ndims(tree) === 2
        @test length(tree) == 0
        @test SI.height(tree) == 1
        @test isempty(tree)
        @test isempty(tree, ambr)
        @test isequal(SI.mbr(tree.root), SI.empty(SI.regiontype(tree)))
        @test SI.check(tree)
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)), 1)
        @test SI.check(tree)
        @test iterate(tree) === nothing
        @test iterate(contained_in(tree, SI.empty(SI.mbrtype(tree)))) === nothing
        @test iterate(intersects_with(tree, SI.empty(SI.mbrtype(tree)))) === nothing
        @test iterate(contained_in(tree, cmbr)) === nothing
        @test iterate(intersects_with(tree, cmbr)) === nothing
        @test collect(tree) == eltype(tree)[]
        @test collect(contained_in(tree, cmbr)) == eltype(tree)[]
        @test collect(intersects_with(tree, cmbr)) == eltype(tree)[]
        @test typeof(collect(tree)) === Vector{eltype(tree)}

        @test insert!(tree, ambr, 1, 2) === tree
        @test length(tree) == 1
        @test !isempty(tree)
        @test !isempty(tree, ambr)
        @test SI.height(tree) == 1
        @test isequal(SI.mbr(tree.root), ambr)
        @test SI.check(tree)
        @test_throws KeyError delete!(tree, bmbr, 1)
        @test_throws KeyError delete!(tree, ambr, 2)
        @test SI.check(tree)
        @test collect(tree) == [SpatialElem(ambr, Int32(1), 2)]

        @test insert!(tree, bmbr, 2, 2) === tree
        @test length(tree) == 2
        @test SI.height(tree) == 1
        @test isequal(SI.mbr(tree.root), SI.combine(ambr, bmbr))
        @test SI.check(tree)
        @test length(collect(tree)) == 2
        @test_throws KeyError delete!(tree, bmbr, 1)
        @test_throws KeyError delete!(tree, ambr, 2)
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)), 1)
        @test delete!(tree, ambr, 1) === tree
        @test length(tree) == 1
        @test SI.check(tree)

        insert!(tree, ambr, 1, 2)
        insert!(tree, cmbr, 2, 3)
        @test length(tree) == 3
        @test_throws SpatialIndexException SI.check(tree) # duplicate ID=2 (b and c)
        @test delete!(tree, cmbr, 2) === tree
        @test length(tree) == 2
        # prepare tree with a(id=1), b(id=2), c(id=3)
        @test insert!(tree, cmbr, 3, 3) === tree
        @test length(tree) == 3
        @test SI.check(tree)
        @test length(collect(contained_in(tree, SI.Rect((0.0, 0.0), (1.0, 1.0))))) == 3
        @test length(collect(intersects_with(tree, SI.Rect((0.0, 0.0), (1.0, 1.0))))) == 3
        @test length(collect(contained_in(tree, SI.Rect((0.0, 0.0), (0.6, 0.6))))) == 2
        @test length(collect(intersects_with(tree, SI.Rect((0.0, 0.0), (0.6, 0.6))))) == 2
        @test length(collect(contained_in(tree, SI.Rect((0.0, 0.0), (0.55, 0.55))))) == 1 # a only
        @test length(collect(intersects_with(tree, SI.Rect((0.0, 0.0), (0.55, 0.55))))) == 2 # a and c

        tree2 = similar(tree)
        @test typeof(tree2) === typeof(tree)
        @test tree2 !== tree
        @test isempty(tree2)

        @testset "findfirst()" begin
            @test findfirst(tree, ambr, 2) === nothing
            @test_throws MethodError findfirst(tree, ambr, "1") # wrong key type
            @test_throws MethodError findfirst(tree, SI.empty(SI.Rect{Int,2}()))
            @test findfirst(tree, ambr, 1) == (tree.root, 2)
            @test findfirst(tree, ambr) == (tree.root, 2) # search without id ignores it
            @test findfirst(tree, bmbr, 1) === nothing
            @test findfirst(tree, bmbr, 2) == (tree.root, 1)
            @test findfirst(tree, bmbr) == (tree.root, 1)
            @test findfirst(tree, cmbr, 2) === nothing
            @test findfirst(tree, cmbr, 3) == (tree.root, 3)
            @test findfirst(tree, cmbr) == (tree.root, 3)
            @test findfirst(tree, SI.combine(ambr, cmbr)) === nothing
            @test findfirst(tree, SI.empty(SI.mbrtype(tree))) === nothing
        end
    end

    @testset "RTree{Int,3,String,Nothing}(variant=$tree_var) (no id)" for tree_var in tree_vars
        ambr = SI.Rect((0, 0, 0), (0, 0, 0))
        bmbr = SI.Rect((0, 1, 1), (0, 1, 1))
        tree = RTree{Int, 3}(String, variant=tree_var)
        @test tree isa RTree{Int, 3, SpatialElem{Int, 3, Nothing, String}}
        @test SI.variant(tree) === tree_var
        @test eltype(tree) === SpatialElem{Int, 3, Nothing, String}
        @test SI.regiontype(tree) === SI.Rect{Int, 3}
        @test SI.dimtype(tree) === Int
        @test ndims(tree) === 3
        @test length(tree) == 0
        @test SI.height(tree) == 1
        @test isempty(tree)
        @test isempty(tree, ambr)
        @test isequal(SI.mbr(tree.root), SI.empty(SI.regiontype(tree)))
        @test SI.check(tree)
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)), 1)
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)))
        @test SI.check(tree)

        @test_throws MethodError insert!(tree, ambr, 1, "2")
        @test insert!(tree, ambr, "2") === tree
        @test length(tree) == 1
        @test !isempty(tree)
        @test !isempty(tree, ambr)
        @test SI.height(tree) == 1
        @test isequal(SI.mbr(tree.root), ambr)
        @test SI.check(tree)
        @test_throws MethodError delete!(tree, bmbr, 1)
        @test_throws KeyError delete!(tree, bmbr)
        @test SI.check(tree)

        @test insert!(tree, bmbr, "3") === tree
        @test length(tree) == 2
        @test SI.height(tree) == 1
        @test isequal(SI.mbr(tree.root), SI.combine(ambr, bmbr))
        @test SI.check(tree)
        @test_throws MethodError delete!(tree, bmbr, 1)
        @test_throws MethodError delete!(tree, SI.empty(SI.Rect{Int, 2}))
        @test_throws MethodError delete!(tree, SI.empty(SI.Rect{Float64, 3}))
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)))
        @test delete!(tree, bmbr) === tree
        @test length(tree) == 1
        @test SI.check(tree)
        @test insert!(tree, bmbr, "4") === tree
        @test length(tree) == 2
        @test SI.check(tree)

        @testset "findfirst()" begin
            @test_throws MethodError findfirst(tree, ambr, 2) # no key
            @test_throws MethodError findfirst(tree, ambr, "1")
            @test_throws MethodError findfirst(tree, SI.empty(SI.Rect{Int,2}()))
            @test_throws MethodError findfirst(tree, SI.empty(SI.Rect{Float64,3}()))
            @test findfirst(tree, ambr) == (tree.root, 1)
            @test findfirst(tree, bmbr) == (tree.root, 2)
            @test findfirst(tree, SI.combine(ambr, bmbr)) === nothing
            @test findfirst(tree, SI.empty(SI.mbrtype(tree))) === nothing
        end
    end
end

@testset "1000 vertices" begin
    # generate random point cloud
    Random.seed!(32123)
    mbrs = Vector{SI.Rect{Float64, 3}}()
    for i in 1:1000
        w, h, d = 5 .* rand(3)
        x, y, z = 50 .* randn(3)
        rmbr = SI.Rect((x-w, y-h, z-d), (x+w, y+h, z+d))
        push!(mbrs, rmbr)
    end

    @testset "sequential inserts" begin
        tree = RTree{SI.dimtype(eltype(mbrs)), ndims(eltype(mbrs))}(Int, String, leaf_capacity = 20, branch_capacity = 20)
        @test tree isa RTree{Float64, 3, SpatialElem{Float64, 3, Int, String}}
        for (i, rmbr) in enumerate(mbrs)
            insert!(tree, rmbr, i, string(i))
            @test length(tree) == i
            @test tree.nelem_insertions == i
            @test SI.check(tree)
            #@debug "$i: len=$(length(tree)) height=$(SI.height(tree))"
        end
        #@show SI.height(tree)
        #@show tree.nnodes_perlevel

        @testset "findfirst()" begin
            # check that the elements can be found
            for i in 1:length(tree)
                node_ix = findfirst(tree.root, mbrs[i], i)
                @test node_ix !== nothing
                if node_ix !== nothing
                    node, ix = node_ix
                    elem = node[ix]
                    @test SI.id(elem) == i
                    @test elem.val == string(i)
                end
            end
        end

        @testset "iterate" begin
            all_elems = collect(tree)
            @test length(all_elems) == length(tree)
            @test eltype(all_elems) === eltype(tree)

            bound_mbr = SI.Rect((-40.0, -20.0, -20.0), (30.0, 50.0, 40.0))

            in_elems = collect(contained_in(tree, bound_mbr))
            @test eltype(in_elems) === eltype(tree)
            @test length(in_elems) == sum(br -> in(br, bound_mbr), mbrs)

            isect_elems = collect(intersects_with(tree, bound_mbr))
            @test eltype(isect_elems) === eltype(tree)
            @test length(isect_elems) == sum(br -> SI.intersects(br, bound_mbr), mbrs)
        end
    end

    @testset "load!() (OMT bulk load)" begin
        tree = RTree{SI.dimtype(eltype(mbrs)), ndims(eltype(mbrs))}(Int, String,
                                leaf_capacity = 20, branch_capacity = 20)
        @test tree === SI.load!(tree, enumerate(mbrs), method=:OMT,
                                convertel = x -> eltype(tree)(x[2], x[1], string(x[1])))
        @test SI.check(tree)
        @test length(tree) == length(mbrs)
        #@show SI.height(tree)
        #@show tree.nnodes_perlevel
        # cannot bulk-load into non-empty tree
        @test_throws ArgumentError SI.load!(tree, enumerate(mbrs), method=:OMT,
                                            convertel = x -> eltype(tree)(x[2], x[1], string(x[1])))
        tree2 = similar(tree)
        # can load from tree into another tree
        @test SI.load!(tree2, tree) == tree2
        @test length(tree2) == length(tree)

        @testset "findfirst()" begin
            # check that the elements can be found
            for i in 1:length(tree)
                node_ix = findfirst(tree.root, mbrs[i], i)
                @test node_ix !== nothing
                if node_ix !== nothing
                    node, ix = node_ix
                    elem = node[ix]
                    @test SI.id(elem) == i
                    @test elem.val == string(i)
                end
            end
        end
    end
end

@testset "subtract!()" begin
    @testset "simple" begin
        tree = RTree{Int, 2}(Int, String, leaf_capacity = 5, branch_capacity = 5)
        pts = [(0, 0), (1, 0), (2, 2), (2, 0), (0, 1), (1, 1), (-1, -1)]
        SI.load!(tree, enumerate(pts),
                 convertel = x -> eltype(tree)(SI.Rect(x[2], x[2]), x[1], string(x[1])))
        @test length(tree) == length(pts)
        @test SI.check(tree)
        rect = SI.Rect((0, 0), (1, 1))
        @test !isempty(tree, rect)
        @test isempty(tree, SI.Rect((-2, 0), (-1, 1)))
        SI.subtract!(tree, rect)
        @test SI.check(tree)
        @test length(tree) == 3
        @test tree.nelem_deletions == 4
        for i in [3, 4, 7] # check that the correct points stayed in the tree
            @test findfirst(tree.root, SI.Rect(pts[i], pts[i]), i) !== nothing
        end
    end

    @testset "from 1000 points" begin
        # generate random point cloud
        Random.seed!(32123)
        pts = [ntuple(_ -> 5. * randn(), 3) for _ in 1:1000]
        tree = RTree{Float64, 3}(Int, String, leaf_capacity = 10, branch_capacity = 10)
        SI.load!(tree, enumerate(pts),
                 convertel = x -> eltype(tree)(SI.Rect(x[2], x[2]), x[1], string(x[1])))
        @test length(tree) == length(pts)
        #@show tree.nelems tree.nnodes_perlevel

        rect = SI.Rect((-3.0, -4.0, -5.0), (5.0, 8.0, 3.0))
        n_in_rect = sum(pt -> in(SI.Point(pt), rect), pts)
        SI.subtract!(tree, rect)
        @test SI.check(tree)
        @test tree.nelem_deletions == n_in_rect
        #@show tree.nelems tree.nnodes_perlevel tree.nelem_deletions tree.nelem_insertions
    end

    # test how well the R-tree can stand the deletion of all but a single element for various tree sizes
    @testset "subtract!() removing all but single point" begin
        Random.seed!(32123)
        pts = [ntuple(_ -> rand(), 2) for _ in 1:200]
        reftree = RTree{Float64, 2}(Int, String, leaf_capacity = 5, branch_capacity = 5)
        loner = eltype(reftree)(SI.Rect((5.0, 5.0), (5.0, 5.0)), 0, "loner")
        @testset "removing $n points" for (n, pt) in enumerate(pts)
            insert!(reftree, eltype(reftree)(SI.Rect(pt, pt), n, string(n)))
            tree = deepcopy(reftree)
            insert!(tree, loner)
            @test length(tree) == n + 1
            SI.subtract!(tree, SI.Rect((0.0, 0.0), (1.0, 1.0)))
            @test length(tree) == 1
            @test SI.check(tree)
            @test SI.id(first(tree)) == 0
        end
    end

    @testset "subtract!() removing n<=N central nodes" begin
        Random.seed!(32123)
        pts = [ntuple(_ -> 2rand()-1, 3) for _ in 1:200]
        # the test implies that all pts have different distances to the origin
        sort!(pts, by = pt -> maximum(abs, pt))
        reftree = RTree{Float64, 3}(Int, String, leaf_capacity = 5, branch_capacity = 5)
        SI.load!(reftree, enumerate(pts),
                 convertel = x -> eltype(reftree)(SI.Rect(x[2], x[2]), x[1], string(x[1])))
        corembr = SI.Rect((0.0, 0.0, 0.0), (0.0, 0.0, 0.0))
        tree1 = deepcopy(reftree)
        @test length(tree1) == length(reftree)
        @testset "removing $n points" for (n, pt) in enumerate(pts)
            corembr = SI.combine(corembr, SI.Rect(pt, pt))

            # subtracting centrermbr from tree1 removes just 1 point
            @test length(tree1) == length(reftree) - n + 1
            tree1_to_remove = collect(contained_in(tree1, corembr))
            @test length(tree1_to_remove) == 1
            @test SI.id(first(tree1_to_remove)) == n

            SI.subtract!(tree1, corembr)
            @test length(tree1) == length(reftree) - n
            @test SI.check(tree1)

            # subtracting centrermbr from tree2 removes n points
            tree2 = deepcopy(reftree)
            SI.subtract!(tree2, corembr)
            @test length(tree2) == length(reftree) - n
            @test SI.check(tree2)
        end
    end
end

end


using LinearAlgebra
using Adapt: adapt

import CUDAdrv
import CUDAdrv: CuPtr, CU_NULL

@testset "Memory" begin
  CuArrays.alloc(0)

  @test (CuArrays.@allocated CuArray{Int32}(undef,1)) == 4

  ret, out = @grab_output CuArrays.@time CuArray{Int32}(undef, 1)
  @test isa(ret, CuArray{Int32})
  @test occursin("1 GPU allocation: 4 bytes", out)

  ret, out = @grab_output CuArrays.@time Base.unsafe_wrap(CuArray, CuPtr{Int32}(12345678), (2, 3))
  @test isa(ret, CuArray{Int32})
  @test !occursin("GPU allocation", out)
end

@testset "Array" begin
  xs = CuArray{Int}(undef, 2, 3)
  @test collect(CuArray([1 2; 3 4])) == [1 2; 3 4]
  @test collect(cu[1, 2, 3]) == [1, 2, 3]
  @test collect(cu([1, 2, 3])) == [1, 2, 3]
  @test testf(vec, rand(5,3))
  @test cu(1:3) === 1:3
  @test Base.elsize(xs) == sizeof(Int)
  @test CuArray{Int, 2}(xs) === xs

  @test_throws ArgumentError Base.unsafe_convert(Ptr{Int}, xs)
  @test_throws ArgumentError Base.unsafe_convert(Ptr{Float32}, xs)

  # Check that allowscalar works
  @test_throws ErrorException xs[1]
  @test_throws ErrorException xs[1] = 1

  # unsafe_wrap
  @test Base.unsafe_wrap(CuArray, CU_NULL, 1; own=false).pooled == false
  @test Base.unsafe_wrap(CuArray, CU_NULL, 2)                == CuArray{Nothing,1}(CU_NULL, (2,),  false)
  @test Base.unsafe_wrap(CuArray{Nothing}, CU_NULL, 2)       == CuArray{Nothing,1}(CU_NULL, (2,),  false)
  @test Base.unsafe_wrap(CuArray{Nothing,1}, CU_NULL, 2)     == CuArray{Nothing,1}(CU_NULL, (2,),  false)
  @test Base.unsafe_wrap(CuArray, CU_NULL, (1,2))            == CuArray{Nothing,2}(CU_NULL, (1,2), false)
  @test Base.unsafe_wrap(CuArray{Nothing}, CU_NULL, (1,2))   == CuArray{Nothing,2}(CU_NULL, (1,2), false)
  @test Base.unsafe_wrap(CuArray{Nothing,2}, CU_NULL, (1,2)) == CuArray{Nothing,2}(CU_NULL, (1,2), false)

  @test collect(CuArrays.zeros(2, 2)) == zeros(Float32, 2, 2)
  @test collect(CuArrays.ones(2, 2)) == ones(Float32, 2, 2)

  @test collect(CuArrays.fill(0, 2, 2)) == zeros(Float32, 2, 2)
  @test collect(CuArrays.fill(1, 2, 2)) == ones(Float32, 2, 2)
end

@testset "Adapt" begin
  A = rand(Float32, 3, 3)
  dA = CuArray(A)
  @test adapt(Array, dA) ≈ A
  @test adapt(CuArray, A) ≈ dA
end

@testset "Broadcast" begin
  @test testf((x)       -> fill!(x, 1),  rand(3,3))
  @test testf((x, y)    -> map(+, x, y), rand(2, 3), rand(2, 3))
  @test testf((x)       -> sin.(x),      rand(2, 3))
  @test testf((x)       -> log.(x) .+ 1, rand(2, 3))
  @test testf((x)       -> 2x,           rand(2, 3))
  @test testf((x)       -> x .^ 0,      rand(2, 3))
  @test testf((x)       -> x .^ 1,      rand(2, 3))
  @test testf((x)       -> x .^ 2,      rand(2, 3))
  @test testf((x)       -> x .^ 3,      rand(2, 3))
  @test testf((x)       -> x .^ 5,      rand(2, 3))
  @test testf((x)       -> (z = Int32(5); x .^ z),      rand(2, 3))
  @test testf((x)       -> (z = Float64(π); x .^ z),      rand(2, 3))
  @test testf((x)       -> (z = Float32(π); x .^ z),      rand(Float32, 2, 3))
  @test testf((x, y)    -> x .+ y,       rand(2, 3), rand(1, 3))
  @test testf((z, x, y) -> z .= x .+ y,  rand(2, 3), rand(2, 3), rand(2))
  @test (CuArray{Ptr{Cvoid}}(undef, 1) .= C_NULL) == CuArray([C_NULL])
  @test CuArray([1,2,3]) .+ CuArray([1.0,2.0,3.0]) == CuArray([2,4,6])

  @eval struct Whatever{T}
      x::Int
  end
  @test Array(Whatever{Int}.(CuArray([1]))) == Whatever{Int}.([1])
end

@testset "Cufunc" begin
  gelu1(x) = oftype(x, 0.5) * x * (1 + tanh(oftype(x, √(2/π))*(x + oftype(x, 0.044715) * x^3)))
  sig(x) = one(x) / (one(x) + exp(-x))
  f(x) = gelu1(log(x)) * sig(x) * tanh(x)
  g(x) = x^7 - 2 * x^f(x^2) + 3

  CuArrays.@cufunc gelu1(x) = oftype(x, 0.5) * x * (1 + tanh(oftype(x, √(2/π))*(x + oftype(x, 0.044715) * x^3)))
  CuArrays.@cufunc sig(x) = one(x) / (one(x) + exp(-x))
  CuArrays.@cufunc f(x) = gelu1(log(x)) * sig(x) * tanh(x)
  CuArrays.@cufunc g(x) = x^7 - 2 * x^f(x^2) + 3

  @test :gelu1 ∈ CuArrays.cufuncs()
  @test :sig ∈ CuArrays.cufuncs()
  @test :f ∈ CuArrays.cufuncs()
  @test :g ∈ CuArrays.cufuncs()
  @test testf((x)  -> gelu1.(x), rand(3,3))
  @test testf((x)  -> sig.(x),  rand(3,3))
  @test testf((x)  -> f.(x),    rand(3,3))
  @test testf((x)  -> g.(x),    rand(3,3))
end

# https://github.com/JuliaGPU/CUDAnative.jl/issues/223
@testset "Ref Broadcast" begin
  foobar(idx, A) = A[idx]
  @test CuArray([42]) == foobar.(CuArray([1]), Base.RefValue(CuArray([42])))
end

@testset "Broadcast Fix" begin
  @test testf(x -> log.(x), rand(3,3))
  @test testf((x,xs) -> log.(x.+xs), Ref(1), rand(3,3))

  if CuArrays.has_cudnn()
    using NNlib

    @test testf(x -> logσ.(x), rand(5))
  end
end

@testset "Reduce" begin
  @test testf(x -> sum(x, dims=1), rand(2, 3))
  @test testf(x -> sum(x, dims=2), rand(2, 3))
  @test testf(x -> sum(x -> x^2, x, dims=1), rand(2, 3))
  @test testf(x -> prod(x, dims=2), rand(2, 3))

  @test testf(x -> sum(x), rand(2, 3))
  @test testf(x -> prod(x), rand(2, 3))

  @test testf(x -> minimum(x), rand(2, 3))
  @test testf(x -> minimum(x, dims=2), rand(2, 3))
  @test testf(x -> minimum(x, dims=(2, 3)), rand(2, 3, 4))

  @test testf(x -> maximum(x), rand(2, 3))
  @test testf(x -> maximum(x, dims=2), rand(2, 3))
  @test testf(x -> maximum(x, dims=(2, 3)), rand(2, 3, 4))

  myreducer(x1, x2) = x1 + x2 # bypass optimisations for sum()
  @test testf(x -> reduce(myreducer, x, dims=(2, 3), init=0.0), rand(2, 3, 4))
  @test testf(x -> reduce(myreducer, x, init=0.0), rand(2, 3))
  @test testf(x -> reduce(myreducer, x, dims=2, init=42.0), rand(2, 3))

  ex = ErrorException("Please supply a neutral element for &. E.g: mapreduce(f, &, A; init = 1)")
  @test_throws ex mapreduce(t -> t > 0.5, &, cu(rand(2, 3)))
  @test testf(x -> mapreduce(t -> t > 0.5, &, x, init=true), rand(2, 3))

  ex = UndefKeywordError(:init)
  cub = map(t -> t > 0.5, cu(rand(2, 3)))
  @test_throws ex reduce(|, cub)
  @test testf(x -> reduce(|, x, init=false), map(t -> t > 0.5, cu(rand(2, 3))))
end

@testset "0D" begin
  x = CuArray{Float64}(undef)
  x .= 1
  @test collect(x)[] == 1
  if VERSION >= v"1.3-"
    # broken test that throws
    # https://github.com/JuliaGPU/GPUArrays.jl/issues/204
    @test_throws CUDAnative.InvalidIRError x /= 2
  else
    x /= 2
    @test collect(x)[] == 0.5
  end
end

@testset "SubArray" begin
  @test testf(rand(5)) do x
    y = x[2:4]
    y .= 1
    x
  end

  @test testf(rand(5)) do x
    y = view(x, 2:4)
    y .= 1
    x
  end

  @test testf(x->view(x, :, 1:4, 3), rand(Float32, 5, 4, 3))

  let x = cu(rand(Float32, 5, 4, 3))
    @test_throws BoundsError view(x, :, :, 1:10)

    # Contiguous views should return new CuArray
    @test view(x, :, 1, 2) isa CuVector{Float32}
    @test view(x, 1:4, 1, 2) isa CuVector{Float32}
    @test view(x, :, 1:4, 3) isa CuMatrix{Float32}
    @test view(x, :, :, 1) isa CuMatrix{Float32}
    @test view(x, :, :, :) isa CuArray{Float32,3}
    @test view(x, :) isa CuVector{Float32}
    @test view(x, 1:3) isa CuVector{Float32}

    # Non-contiguous views should fall back to base's SubArray
    @test view(x, 1:3, 1:3, 3) isa SubArray
    @test view(x, 1, :, 3) isa SubArray
    @test view(x, 1, 1:4, 3) isa SubArray
    @test view(x, :, 1, 1:3) isa SubArray
    @test view(x, :, 1:2:4, 1) isa SubArray
    @test view(x, 1:2:5, 1, 1) isa SubArray
  end

  # non-contiguous copyto!
  let x = CuArrays.rand(4, 4)
    y = view(x, 2:3, 2:3)

    # to gpu
    gpu = CuArray{eltype(y)}(undef, size(y))
    copyto!(gpu, y)
    @test Array(gpu) == Array(y)

    # to cpu
    cpu = Array{eltype(y)}(undef, size(y))
    copyto!(cpu, y)
    @test cpu == Array(y)
  end

  # bug in parentindices conversion
  let x = CuArray{Int}(undef, 1, 1)
    x[1,:] .= 42
    @test Array(x)[1,1] == 42
  end

  # bug in copyto!
  ## needless N type parameter
  @test testf((x,y)->copyto!(y, selectdim(x, 2, 1)), ones(2,2,2), zeros(2,2))
  ## inability to copyto! smaller destination
  @test testf((x,y)->copyto!(y, selectdim(x, 2, 1)), ones(2,2,2), zeros(3,3))
end

@testset "reshape" begin
  A = [1 2 3 4
       5 6 7 8]
  gA = reshape(CuArray(A),1,8)
  _A = reshape(A,1,8)
  _gA = Array(gA)
  @test all(_A .== _gA)
  A = [1,2,3,4]
  gA = reshape(CuArray(A),4)
end

@testset "$f! with diagonal $d" for (f, f!) in ((triu, triu!), (tril, tril!)),
                                          d in -2:2
  A = randn(10, 10)
  @test f(A, d) == Array(f!(CuArray(A), d))
end

@testset "Utilities" begin
  t = Base.@elapsed ret = CuArrays.@sync begin
    # TODO: do something that takes a while on the GPU
    #       (need to wrap clock64 in CUDAnative for that)
    42
  end
  @test t >= 0
  @test ret == 42
end

@testset "accumulate" begin
  for n in (0, 1, 2, 3, 10, 10_000, 16384, 16384+1) # small, large, odd & even, pow2 and not
    @test testf(x->accumulate(+, x), rand(n))
  end

  # multidimensional
  for (sizes, dims) in ((2,) => 2,
                        (3,4,5) => 2,
                        (1, 70, 50, 20) => 3)
    @test testf(x->accumulate(+, x; dims=dims), rand(Int, sizes))
  end

  # using initializer
  for (sizes, dims) in ((2,) => 2,
                        (3,4,5) => 2,
                        (1, 70, 50, 20) => 3)
    @test testf(x->accumulate(+, x; dims=dims, init=100.), rand(Int, sizes))
  end

  # in place
  @test testf(x->(accumulate!(+, x, copy(x)); x), rand(2))

  # specialized
  @test testf(cumsum, rand(2))
  @test testf(cumprod, rand(2))
end

@testset "logical indexing" begin
  @test CuArray{Int}(undef, 2)[CuArray{Bool}(undef, 2)] isa CuArray
  @test CuArray{Int}(undef, 2, 2)[CuArray{Bool}(undef, 2, 2)] isa CuArray
  @test CuArray{Int}(undef, 2, 2, 2)[CuArray{Bool}(undef, 2, 2, 2)] isa CuArray

  @test CuArray{Int}(undef, 2)[Array{Bool}(undef, 2)] isa CuArray
  @test CuArray{Int}(undef, 2, 2)[Array{Bool}(undef, 2, 2)] isa CuArray
  @test CuArray{Int}(undef, 2, 2, 2)[Array{Bool}(undef, 2, 2, 2)] isa CuArray

  @test testf((x,y)->x[y], rand(2), rand(Bool, 2))
  @test testf((x,y)->x[y], rand(2, 2), rand(Bool, 2, 2))
  @test testf((x,y)->x[y], rand(2, 2, 2), rand(Bool, 2, 2, 2))

  @test testf(x -> x[x .> 0.5], rand(2))
  @test testf(x -> x[x .> 0.5], rand(2,2))
  @test testf(x -> x[x .> 0.5], rand(2,2,2))

  @test testf(x -> filter(y->y .> 0.5, x), rand(2))
  @test testf(x -> filter(y->y .> 0.5, x), rand(2,2))
  @test testf(x -> filter(y->y .> 0.5, x), rand(2,2,2))
end

@testset "generic fallbacks" begin
    a = rand(Int8, 3, 3)
    b = rand(Int8, 3, 3)
    d_a = CuArray{Int8}(a)
    d_b = CuArray{Int8}(b)
    d_c = d_a*d_b
    @test collect(d_c) == a*b
    a = rand(Complex{Int8}, 3, 3)
    b = rand(Complex{Int8}, 3, 3)
    d_a = CuArray{Complex{Int8}}(a)
    d_b = CuArray{Complex{Int8}}(b)
    d_c = d_a'*d_b
    @test collect(d_c) == a'*b
    d_c = d_a*d_b'
    @test collect(d_c) == a*b'
    d_c = d_a'*d_b'
    @test collect(d_c) == a'*b'
    d_c = transpose(d_a)*d_b'
    @test collect(d_c) == transpose(a)*b'
    d_c = d_a'*transpose(d_b)
    @test collect(d_c) == a'*transpose(b)
    d_c = transpose(d_a)*d_b
    @test collect(d_c) == transpose(a)*b
    d_c = d_a*transpose(d_b)
    @test collect(d_c) == a*transpose(b)
    d_c = transpose(d_a)*transpose(d_b)
    @test collect(d_c) == transpose(a)*transpose(b)
    d_c = rmul!(copy(d_a), Complex{Int8}(2, 2))
    @test collect(d_c) == a*Complex{Int8}(2, 2)
    d_c = lmul!(Complex{Int8}(2, 2), copy(d_a))
    @test collect(d_c) == Complex{Int8}(2, 2)*a
end

@testset "reverse" begin
    # 1-d out-of-place
    @test testf(x->reverse(x), rand(1000))
    @test testf(x->reverse(x, 10), rand(1000))
    @test testf(x->reverse(x, 10, 90), rand(1000))

    # 1-d in-place
    @test testf(x->reverse!(x), rand(1000))
    @test testf(x->reverse!(x, 10), rand(1000))
    @test testf(x->reverse!(x, 10, 90), rand(1000))

    # n-d out-of-place
    for shape in ([1, 2, 4, 3], [4, 2], [5], [2^5, 2^5, 2^5]),
        dim in 1:length(shape)
      @test testf(x->reverse(x; dims=dim), rand(shape...))

      cpu = rand(shape...)
      gpu = CuArray(cpu)
      reverse!(gpu; dims=dim)
      @test Array(gpu) == reverse(cpu; dims=dim)
    end
end

@testset "permutedims" begin
    @test testf(x->permutedims(x, [1, 2]), rand(4, 4))

    inds = rand(1:100, 150, 150)
    @test testf(x->permutedims(view(x, inds, :), (3, 2, 1)), rand(100, 100))
end

@testset "findall" begin
    # 1D
    @test testf(x->findall(x), rand(Bool, 100))
    @test testf(x->findall(y->y>0.5, x), rand(100))

    # ND
    let x = rand(Bool, 10, 10)
      @test findall(x) == Array(findall(CuArray(x)))
    end
    let x = rand(10, 10)
      @test findall(y->y>0.5, x) == Array(findall(y->y>0.5, CuArray(x)))
    end
end

@testset "findfirst" begin
    # 1D
    @test testf(x->findfirst(x), rand(Bool, 100))
    @test testf(x->findfirst(y->y>0.5, x), rand(100))

    # ND
    let x = rand(Bool, 10, 10)
      @test findfirst(x) == findfirst(CuArray(x))
    end
    let x = rand(10, 10)
      @test findfirst(y->y>0.5, x) == findfirst(y->y>0.5, CuArray(x))
    end
end

@testset "findmin" begin
  let x = rand(Float32, 100)
      @test findmin(x) == findmin(CuArray(x))
  end
  let x = rand(Float32, 10, 10)
      @test findmin(x) == findmin(CuArray(x))
  end
end

@testset "findmax" begin
  let x = rand(Float32, 100)
      @test findmax(x) == findmax(CuArray(x))
  end
  let x = rand(Float32, 10, 10)
      @test findmax(x) == findmax(CuArray(x))
  end
end


# This file was generated, do not modify it. # hide
knn = machine(KnnPipe(), Xtrain, ytrain)
multi = machine(MnPipe(), Xtrain, ytrain)


################################################################################
# Particle (flat LC)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle")
dir_sim   = joinpath(module_dir(), "src/simulation/particle")
model = deepcopy(particle)
env = deepcopy(flat_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_lc/residual.jld2")
path_jac = joinpath(dir_sim, "flat_lc/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Particle (flat + nonlinear cone)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle")
dir_sim   = joinpath(module_dir(), "src/simulation/particle")
model = deepcopy(particle)
env = deepcopy(flat_3D_nc)
s = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_nc/residual.jld2")
path_jac = joinpath(dir_sim, "flat_nc/jacobians.jld2")

instantiate_base!(s.model, path_base)
instantiate_dynamics!(s.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(s.model, s.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(s, path_res, path_jac)

################################################################################
# Particle (quadratic)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle")
dir_sim   = joinpath(module_dir(), "src/simulation/particle")
model = deepcopy(particle)
env = deepcopy(quadratic_bowl_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "quadratic/residual.jld2")
path_jac = joinpath(dir_sim, "quadratic/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Particle 2D (flat + LC)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/particle_2D")
model = deepcopy(particle_2D)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_lc/residual.jld2")
path_jac = joinpath(dir_sim, "flat_lc/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Particle 2D (flat + nonlinear cone)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/particle_2D")
model = deepcopy(particle_2D)
env = deepcopy(flat_2D_nc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_nc/residual.jld2")
path_jac = joinpath(dir_sim, "flat_nc/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Particle 2D (slope)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/particle_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/particle_2D")
model = deepcopy(particle_2D)
env = deepcopy(slope1_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "slope/residual.jld2")
path_jac = joinpath(dir_sim, "slope/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper (2D)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_2D")
model = deepcopy(hopper_2D)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper (2D) (sinusoidal)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_2D")
model = deepcopy(hopper_2D)
env = deepcopy(sine2_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "sinusoidal/residual.jld2")
path_jac = joinpath(dir_sim, "sinusoidal/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper 2D (Piecewise)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_2D")
model = deepcopy(hopper_2D)
env = deepcopy(piecewise1_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "piecewise/residual.jld2")
path_jac = joinpath(dir_sim, "piecewise/jacobians.jld2")

instantiate_base!(sim.model, path_base)
expr_dyn = generate_dynamics_expressions(sim.model, derivs = true)
save_expressions(expr_dyn, path_dyn, overwrite=true)
instantiate_dynamics!(sim.model, path_dyn, derivs = true)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env, jacobians = :approx)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac, jacobians = :approx)

################################################################################
# Hopper (2D) (stairs)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_2D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_2D")
model = deepcopy(hopper_2D)
env = deepcopy(stairs3_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "stairs/residual.jld2")
path_jac = joinpath(dir_sim, "stairs/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper (3D)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_3D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_3D")
model = deepcopy(hopper_3D)
env = deepcopy(flat_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper (3D nonlinear cone)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_3D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_3D")
model = deepcopy(hopper_3D)
env = deepcopy(flat_3D_nc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat_nc/residual.jld2")
path_jac = joinpath(dir_sim, "flat_nc/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Hopper (3D) (sinusoidal)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/hopper_3D")
dir_sim   = joinpath(module_dir(), "src/simulation/hopper_3D")
model = deepcopy(hopper_3D)
env = deepcopy(sine2_3D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "sinusoidal/residual.jld2")
path_jac = joinpath(dir_sim, "sinusoidal/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Quadruped
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/quadruped")
model = deepcopy(quadruped)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Quadruped Payload
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/quadruped")
model = deepcopy(quadruped_payload)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics_payload/base.jld2")
path_dyn = joinpath(dir_model, "dynamics_payload/dynamics.jld2")
path_res = joinpath(dir_sim, "payload/residual.jld2")
path_jac = joinpath(dir_sim, "payload/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Quadruped Sinusoidal
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/quadruped")
model = deepcopy(quadruped)
env = deepcopy(sine1_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "sinusoidal/residual.jld2")
path_jac = joinpath(dir_sim, "sinusoidal/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env, jacobians = :full)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Quadruped (Piecewise)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/quadruped")
dir_sim   = joinpath(module_dir(), "src/simulation/quadruped")
model = deepcopy(quadruped)
env = deepcopy(piecewise1_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "piecewise/residual.jld2")
path_jac = joinpath(dir_sim, "piecewise/jacobians.jld2")

instantiate_base!(sim.model, path_base)
expr_dyn = generate_dynamics_expressions(sim.model, derivs = true)
save_expressions(expr_dyn, path_dyn, overwrite=true)
instantiate_dynamics!(sim.model, path_dyn, derivs = true)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env, jacobians = :approx)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac, jacobians = :approx)

################################################################################
# Flamingo (flat)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/flamingo")
dir_sim   = joinpath(module_dir(), "src/simulation/flamingo")
model = deepcopy(flamingo)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Flamingo (sinusoidal)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/flamingo")
dir_sim   = joinpath(module_dir(), "src/simulation/flamingo")
model = deepcopy(flamingo)
env = deepcopy(sine3_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "sinusoidal/residual.jld2")
path_jac = joinpath(dir_sim, "sinusoidal/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# Flamingo (smooth slope)
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/flamingo")
dir_sim   = joinpath(module_dir(), "src/simulation/flamingo")
model = deepcopy(flamingo)
env = deepcopy(slope_smooth_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "slope/residual.jld2")
path_jac = joinpath(dir_sim, "slope/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)

################################################################################
# PushBot
################################################################################
dir_model = joinpath(module_dir(), "src/dynamics/pushbot")
dir_sim   = joinpath(module_dir(), "src/simulation/pushbot")
model = deepcopy(pushbot)
env = deepcopy(flat_2D_lc)
sim = Simulation(model, env)

path_base = joinpath(dir_model, "dynamics/base.jld2")
path_dyn = joinpath(dir_model, "dynamics/dynamics.jld2")
path_res = joinpath(dir_sim, "flat/residual.jld2")
path_jac = joinpath(dir_sim, "flat/jacobians.jld2")

instantiate_base!(sim.model, path_base)
instantiate_dynamics!(sim.model, path_dyn)

expr_res, rz_sp, rθ_sp = generate_residual_expressions(sim.model, sim.env)
save_expressions(expr_res, path_res, overwrite=true)
@save path_jac rz_sp rθ_sp
instantiate_residual!(sim, path_res, path_jac)


using MPI, PETSc2

  PetscInitialize()
  b = PetscVec(10, VECMPI, MPI.COMM_WORLD)

  idx = PetscInt[1, 2, 3]
  vals = PetscScalar[1.0, 2.0, 3.0]

  set_values1!(b, idx, vals, ADD_VALUES)

  # don't assemble the vector
  PetscView(b)

  PetscDestroy(b)


# Exports

export Shape, PostOpObj, create, combine, unpack, repack, repack!

abstract type Shape end

mutable struct PostOpObj
    objs::Vector{Any}
    operations::Array{Function}
    function PostOpObj(obj, func)
        if ~isa(obj, Vector)
            obj = [obj]
        end
        if ~isa(func, Vector)
                func = [func]
        end
        new(obj, func)
    end
end

function Base.copy(x::T) where T<:Union{Shape, PostOpObj}
    T(copy.([getfield(x, fn) for fn in fieldnames(T)])...)
end

"""
    create(shape::T; resolution=nothing, rand_=0.0, type::Int64=1) where T <: Shape

Abstact function for creating **Shape** objects.

## Returns
    - X : Initial reference position
    - V : Initial velocity
    - Y : Initial position
    - volume : Volume per particle point
    - type: Type of particle point

"""
function create(shape::T; resolution=nothing, rand_=0.0, type::Int64=1) where T <: Shape
    error("Not implemented for type **$(typeof(shape))** yet.")
end

function mycat(a, b)
    if length(size(a))==1
        return vcat(a, b)
    else
        return hcat(a, b)
    end
end

function create(pobj::PostOpObj, args...; kwargs...)
    temp = nothing
    out = nothing
    for obj in pobj.objs
        temp = create(obj, args...; kwargs...)
        if isa(out, Nothing)
            out = temp
        else
            out = mycat.(out, temp)
        end
    end
    for func in pobj.operations
        out = func(out)
    end
    return out
end

function combine(obj1::T1, obj2::T2) where {T1<:Union{Shape, PostOpObj}, T2<:Union{Shape, PostOpObj}}
    PostOpObj([obj1, obj2], [])
end

function unpack(d::Dict)
    return d[:x], d[:v], d[:y], d[:volume], d[:type]
end

function repack(args...; keys_ = (:x, :v, :y, :volume, :type))
    d = Dict()
    for i in 1:5
        d[keys_[i]] = args[i]
    end
    d
end

function repack!(d::Dict, keys_, vals)
    if length(keys_)==length(vals)
        for i in 1:length(vals)
            d[keys_[i]] = vals[i]
        end
    else
        error("Lengths are not same.")
    end
    d
end

include("./cuboid.jl")
include("./sphere.jl")
include("./disk.jl")
include("./cylinder.jl")
include("./cone.jl")


# ---
# jupyter:
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# IJulia rocks! So does Plotly. Check it out

using Plotly
api_key = "" # visit https://plot.ly/api to generate an API username and password
username = ""

Plotly.signin(username, api_key)

# +
# Following data taken from http://julialang.org/ frontpage
benchmarks = ["fib", "parse_int", "quicksort3", "mandel", "pi_sum", "rand_mat_stat", "rand_mat_mul"]
platforms = ["Fortran", "Julia", "Python", "R", "Matlab", "Mathematica", "Javascript", "Go"]

data = {
        platforms[1] => [0.26, 5.03, 1.11, 0.86, 0.80, 0.64, 0.96],
        platforms[2] => [0.91, 1.60, 1.14, 0.85, 1.00, 1.66, 1.01],
        platforms[3] => [30.37, 13.95, 31.98, 14.19, 16.33, 13.52, 3.41 ],
        platforms[4] => [411.36, 59.40, 524.29, 106.97, 15.42, 10.84, 3.98  ],
        platforms[5] => [1992.00, 1463.16, 101.84, 64.58, 1.29, 6.61, 1.10   ],
        platforms[6] => [64.46, 29.54, 35.74, 6.07, 1.32, 4.52, 1.16 ],
        platforms[7] => [2.18, 2.43, 3.51, 3.49, 0.84, 3.28, 14.60],
        platforms[8] => [1.03, 4.79, 1.25, 2.36, 1.41, 8.12, 8.51]
        }

pdata = [ {"x"=>benchmarks,"y"=>data[k],"bardir"=>"h","type"=>"bar","name"=>k} for k = platforms ]

layout = {
          "title"=> "Julia benchmark comparison (smaller is better, C performance = 1.0)",
          "barmode"=> "group",
          "autosize"=> false,
          "width"=> 900,
          "height"=> 900,
          "titlefont"=>
          {
           "family"=> "Open Sans",
           "size"=> 18,
           "color"=> "rgb(84, 39, 143)"
           },
          "margin"=> {"l"=>160, "pad"=>0},
          "xaxis"=> {
                     "title"=> "Benchmark log-time",
                     "type"=> "log"
                     },
          "yaxis"=> {"title"=> "Benchmark Name"}
          }

response = Plotly.plot(pdata,["layout"=>layout])

# Embed in an iframe within IJulia
s = string("")
display("text/html", s)

# +
# checkout https://plot.ly/api/ for more Julia examples!
# But to show off some other Plotly features:
x = 1:1500
y1 = sin(2*pi*x/1500.) + rand(1500)-0.5
y2 = sin(2*pi*x/1500.)

fish = {"x"=>x,"y"=> y1,
	"type"=>"scatter","mode"=>"markers",
	"marker"=>{"color"=>"rgb(0, 0, 255)","opacity"=>0.5 } }

fit = {"x"=> x,"y"=> y2,
	"type"=>"scatter", "mode"=>"markers", "opacity"=>0.8,
	"marker"=>{"color"=>"rgb(255, 0, 0)"} }

layout = {"autosize"=> false,
    "width"=> 650, "height"=> 550,
    "title"=>"Fish School",
	"xaxis"=>{ "ticks"=> "",
        "gridcolor"=> "white",
        "zerolinecolor"=> "white",
        "linecolor"=> "white",
        "autorange"=> false,
        "range"=>[0,1500] },
	"yaxis"=>{ "ticks"=> "",
        "gridcolor"=> "white",
        "zerolinecolor"=> "white",
		"linecolor"=> "white",
        "autorange"=> false,
        "range"=>[-2.2,2.2] },
	"plot_bgcolor"=> "rgb(245,245,247)",
    "showlegend"=> false,
    "hovermode"=> "closest"}

response = Plotly.plot([fish, fit],["layout"=>layout])
s = string("")
display("text/html", s)


## Return in set strings of infinite sets
# Fallback
function JuMP.in_set_string(print_mode, set::AbstractInfiniteSet)::String
    return string(JuMP._math_symbol(print_mode, :in), " ", set)
end

# Extend to return of in set string for interval sets
function JuMP.in_set_string(print_mode, set::IntervalSet)::String
    if set.lower_bound != set.upper_bound
        return string(JuMP._math_symbol(print_mode, :in), " [",
                      JuMP._string_round(set.lower_bound), ", ",
                      JuMP._string_round(set.upper_bound), "]")
    else
        return string(JuMP._math_symbol(print_mode, :eq), " ",
                      JuMP._string_round(set.lower_bound))
    end
end

# Extend to return of in set string for distribution sets
function JuMP.in_set_string(print_mode, set::DistributionSet)::String
    # get distribution string
    d_string = string(set.distribution)
    # remove number type
    first_bracket = findfirst(isequal('{'), d_string)
    if first_bracket != nothing
        last_bracket = findall(isequal('}'), d_string)[end]
        d_string = d_string[1:first_bracket-1] * d_string[last_bracket+1:end]
    end
    # remove new line characters
    new_lines = findall(isequal('\n'), d_string)
    for i = 1:length(new_lines)
        if new_lines[1] == length(d_string)
            d_string = d_string[1:end-1]
        elseif d_string[new_lines[1]-1] == '(' || d_string[new_lines[1]+1] == ')'
            d_string = d_string[1:new_lines[1]-1] * d_string[new_lines[1]+1:end]
        else
            d_string = d_string[1:new_lines[1]-1] * ", " *
                       d_string[new_lines[1]+1:end]
        end
        new_lines = findall(isequal('\n'), d_string)
    end
    return string(JuMP._math_symbol(print_mode, :in), " ", d_string)
end

# Return contraint parameter bounds as a string
function bound_string(print_mode, bounds::ParameterBounds)::String
    string_list = ""
    for (pref, set) in bounds.intervals
        string_list *= string(JuMP.function_string(print_mode, pref), " ",
                              JuMP.in_set_string(print_mode, set), ", ")
    end
    return string_list[1:end-2]
end

# Return string of a bounded constraint
function JuMP.constraint_string(print_mode,
                                constraint_object::BoundedScalarConstraint)
    # get needed strings
    func_str = JuMP.function_string(print_mode, constraint_object)
    in_set_str = JuMP.in_set_string(print_mode, constraint_object)
    bound_str = bound_string(print_mode, constraint_object.bounds)
    # modify as needed
    if print_mode == JuMP.REPLMode
        lines = split(func_str, '\n')
        lines[1 + div(length(lines), 2)] *= " " * in_set_str * ", " *
                                       JuMP._math_symbol(print_mode, :for_all) *
                                       " " * bound_str
        return join(lines, '\n')
    else
        return func_str * " " * in_set_str * ", " *
               JuMP._math_symbol(print_mode, :for_all) * " " * bound_str
    end
end

# Return a string for a constraint reference
function JuMP.constraint_string(print_mode, ref::GeneralConstraintRef)
    return JuMP.constraint_string(print_mode, JuMP.name(ref),
           JuMP.constraint_object(ref))
end

# return list of constraint
function JuMP.constraints_string(print_mode, model::InfiniteModel)
    strings = String[]
    # Sort by creation order
    constraints = sort(collect(model.constrs), by = c -> c.first)
    # produce a string for each constraint
    for (index, constraint) in constraints
        push!(strings, JuMP.constraint_string(print_mode, constraint))
    end
    # produce the parameters that are used along with their sets
    params = sort(collect(model.params), by = c -> c.first)
    ignore_groups = [] # used to indicate which groups should no longer be used
    for (index, param) in params
        pref = ParameterRef(model, index)
        if is_used(pref) && !(group_id(pref) in ignore_groups)
            if isa(infinite_set(pref), DistributionSet{<:Distributions.MultivariateDistribution})
                # print multivariate distributions with variables grouped together
                push!(ignore_groups, group_id(pref))
                push!(strings, string(_root_name(pref), " ",
                                JuMP.in_set_string(print_mode, param.set)))
            else
                # print parameter individually
                push!(strings, string(JuMP.function_string(print_mode, pref), " ",
                                      JuMP.in_set_string(print_mode, param.set)))
            end
        end
    end
    return strings
end

# Return the objective string
function JuMP.objective_function_string(print_mode, model::InfiniteModel)
    return JuMP.function_string(print_mode, JuMP.objective_function(model))
end

# Show IntervalSets in REPLMode
function Base.show(io::IO, set::IntervalSet)
    print(io, "[", JuMP._string_round(set.lower_bound), ", ",
          JuMP._string_round(set.upper_bound), "]")
    return
end

# Show IntervalSets in IJuliaMode
function Base.show(io::IO, ::MIME"text/latex", set::IntervalSet)
    print(io, "[", JuMP._string_round(set.lower_bound), ", ",
          JuMP._string_round(set.upper_bound), "]")
end

# Show IntervalSets in REPLMode
function Base.show(io::IO, set::DistributionSet)
    print(io, set.distribution)
    return
end

# Show IntervalSets in IJuliaMode
function Base.show(io::IO, ::MIME"text/latex", set::DistributionSet)
    print(io, set.distribution)
end

# TODO make list following group_id and use ... when necessary --> Subdomain bounds: t in [0, 1], x[1] == 0, x[2] == 0, ... x[6] == 0, a in [2, 3]
# Show ParameterBounds in REPLMode
function Base.show(io::IO, bounds::ParameterBounds)
    print(io, "Subdomain bounds (", length(bounds), "): ",
          bound_string(JuMP.REPLMode, bounds))
    return
end

# Show ParameterBounds in IJuliaMode
function Base.show(io::IO, ::MIME"text/latex", bounds::ParameterBounds)
    print(io, "Subdomain bounds (", length(bounds), "): ",
          bound_string(JuMP.IJuliaMode, bounds))
end

# TODO show hold variables better

# Show constraint in REPLMode
function Base.show(io::IO, ref::GeneralConstraintRef)
    print(io, JuMP.constraint_string(JuMP.REPLMode, ref))
    return
end

# Show constraint in IJuliaMode
function Base.show(io::IO, ::MIME"text/latex", ref::GeneralConstraintRef)
    print(io, JuMP.constraint_string(JuMP.IJuliaMode, ref))
end

# Show the backend information associated with the optimizer model
function JuMP.show_backend_summary(io::IO, model::InfiniteModel)
    println(io, "Optimizer model backend information: ")
    JuMP.show_backend_summary(io, optimizer_model(model))
    return
end

# Show the objective function type
function JuMP.show_objective_function_summary(io::IO, model::InfiniteModel)
    println(io, "Objective function type: ",
            JuMP.objective_function_type(model))
    return
end

# Return "s" if n is greater than one
_plural(n) = (isone(n) ? "" : "s")

# Return constraint summary in JuMP like manner
function JuMP.show_constraints_summary(io::IO, model::InfiniteModel)
    for (F, S) in JuMP.list_of_constraint_types(model)
        n_constraints = JuMP.num_constraints(model, F, S)
        println(io, "`$F`-in-`$S`: $n_constraints constraint",
                _plural(n_constraints))
    end
    return
end

# Have the infinitemodel object show similar to Model
function Base.show(io::IO, model::InfiniteModel)
    # indicate is an infinite model
    println(io, "An InfiniteOpt Model")
    # show objective ssnse info
    sense = JuMP.objective_sense(model)
    if sense == MOI.MAX_SENSE
        print(io, "Maximization")
    elseif sense == MOI.MIN_SENSE
        print(io, "Minimization")
    else
        print(io, "Feasibility")
    end
    println(io, " problem with:")
    # show variable info
    println(io, "Variable", _plural(JuMP.num_variables(model)), ": ",
            JuMP.num_variables(model))
    # show obejctive function info
    if sense != MOI.FEASIBILITY_SENSE
        JuMP.show_objective_function_summary(io, model)
    end
    # show constraint info
    JuMP.show_constraints_summary(io, model)
    # show other info
    names_in_scope = sort(collect(keys(JuMP.object_dictionary(model))))
    if !isempty(names_in_scope)
        print(io, "Names registered in the model: ",
              join(string.(names_in_scope), ", "), "\n")
    end
    JuMP.show_backend_summary(io, model)
end


module TBIO

    using Formatting
    export readTB!, writeTB

    function readTB!(
        filename   :: AbstractString,
        txt_nchars :: Integer,
        arrs       :: AbstractArray;
        endianess  :: Symbol = :LITTLE,   # Small: 0x04030201,  Big: 0x01020304
        trim       :: Bool   = true,
    )

        if endianess != :LITTLE && endianess != :BIG
            throw(ErrorException("Unknown symbol: " * string(endianess)))
        end

        if txt_nchars < 0
            throw(ErrorException("txt_nchars cannot be negative."))
        end

        msg = nothing
        if isfile(filename)

            get_filesize = filesize(filename)
            expect_filesize = sum(length.(arrs)) * 8 + txt_nchars + 8
            if get_filesize == expect_filesize

                open(filename, "r") do io
                    msg = String(read(io, txt_nchars))

                    trim && (msg = strip(msg))

                    for i = 1:length(arrs)
                        read!(io, arrs[i])
                    end

                    if     endianess == :LITTLE && Base.ENDIAN_BOM == 0x01020304
                        for i = 1:length(arrs)
                            arrs[i][:] = ltoh.(arrs[i])
                        end
                    elseif endianess == :BIG && Base.ENDIAN_BOM == 0x04030201
                        for i = 1:length(arrs)
                            arrs[i][:] = ntoh.(arrs[i])
                        end
                    end

                    received_cs   = reinterpret(UInt64, read(io, 8))[1]
                    calculated_cs = calChecksum(arrs)

                    if received_cs != calculated_cs
                        msg = nothing
                        println(format("[readTB!] Checksum does not match. Received: {:X}, calculated: {:X}", received_cs, calculated_cs))
                    end

                end
            else
                println(format("[readTB!] Expecting filesize: {} bytes, but got {} bytes", expect_filesize, get_filesize))
            end
        else
            println(format("[readTB!] File {} does not exist.", filename))
        end

        return msg
    end

    function writeTB(
        filename   :: AbstractString,
        msg        :: AbstractString,
        txt_nchars :: Integer,
        arrs       :: AbstractArray{T};
        endianess  :: Symbol = :LITTLE,
    ) where T <: AbstractArray{Float64}

        if endianess != :LITTLE && endianess != :BIG
            throw(ErrorException("Unknown symbol: " * string(endianess)))
        end

        if length(msg) > txt_nchars
            throw(ErrorException("Message length exceeds txt_nchars."))
        end

        if txt_nchars < 0
            throw(ErrorException("txt_nchars cannot be negative."))
        end

        open(filename, "w") do io
            write(io, msg)

            for i = 1:(txt_nchars - length(msg))
                write(io, " ")
            end

            if     endianess == :LITTLE && Base.ENDIAN_BOM == 0x01020304
                for i = 1:length(arrs)
                    write(io, htol.(arrs[i]))
                end
                write(io, htol(calChecksum(arrs)))
            elseif endianess == :BIG && Base.ENDIAN_BOM == 0x04030201
                for i = 1:length(arrs)
                    write(io, hton.(arrs[i]))
                end
                write(io, hton(calChecksum(arrs)))
            else
                for i = 1:length(arrs)
                    write(io, arrs[i])
                end
                write(io, calChecksum(arrs))
            end

        end
    end

    # ror, rol :: bitwise circular shift to right / left
    # Source:
    # https://discourse.julialang.org/t/efficient-bit-rotate-functions-ror-rol-what-is-the-official-way-for-julia-v1/19062
    ror(x::UInt64, k::Int) = (x >>> (0x3f & k)) | (x << (0x3f & -k))
    rol(x::UInt64, k::Int) = (x >>> (0x3f & -k)) | (x << (0x3f & k))

    function calChecksum(
        arrs :: AbstractArray,
    )

        k = 0
        s = UInt64(0)
        for i = 1:length(arrs)
            arr = reinterpret(UInt64, arrs[i])

            for j = 1:length(arr)

                s = xor(s, rol(arr[j], k))
                #s = xor(s, arr[j])
                k = mod(k+1, 64)
            end

        end

#        println(format("{:X}", s))
        return s
    end
end


include(joinpath("..", "required.jl"))

## Conditional arguments
outcome = "bart"
subset_qc = true
# save_figures = true

# Subset to QC data (optional)
if (@isdefined subset_qc) && subset_qc == true
    @rput subset_qc
end

## Train & apply models (RF or BART)
model_script = joinpath("src", "others", "x_direct-models_$(outcome).R")
@rput model_script
@time begin
    R"""
    source(model_script)
    """
end # about 1 minute
@rget results

## Fix missing values
predictions = results[:predictions]

## Get comparison layers
# Load species names
glossary = CSV.read(joinpath("data", "proc", "glossary.csv"), DataFrame)
spenames = filter(:type => ==("species"), glossary).full_name

# Load distributions
raw_distributions = [
    geotiff(
        SimpleSDMPredictor, joinpath("data", "raster", "distributions_raw.tif"), i
    ) for i in eachindex(spenames)
]
sdm_distributions = [
    geotiff(
        SimpleSDMPredictor, joinpath("data", "raster", "distributions_bart.tif"), i
    ) for i in eachindex(spenames)
]

# Subset to Quebec
if (@isdefined subset_qc) && subset_qc == true
    coords_qc = (left=-80.0, right=-55.0, bottom=45.0, top=63.0)
    raw_distributions = [d[coords_qc] for d in raw_distributions]
    sdm_distributions = [d[coords_qc] for d in sdm_distributions]
end

## Arrange predictions as layers
richness_pred = SimpleSDMResponse(predictions, :richness, raw_distributions[1])
lcbd_pred = SimpleSDMResponse(predictions, :lcbd, raw_distributions[1])
replace!(richness_pred, missing => nothing)
replace!(lcbd_pred, missing => nothing)
pred = (richness=richness_pred, lcbd=lcbd_pred)

## Plot predictions
richness_plot = plot_layer(
    richness_pred;
    c=:cividis,
    title="Direct richness predictions ($(uppercase(outcome)))",
    colorbar_title="Predicted number of species",
    dpi=200,
)

## Plot differences

# Get comparisons richness & lcbd
Yraw = Ymatrix(raw_distributions)
Ysdm = Ymatrix(sdm_distributions)
raw, sdm = [
    (
        richness=richness(Y, raw_distributions[1]),
        lcbd=lcbd(Y, raw_distributions[1]),
    ) for Y in (Yraw, Ysdm)
]

# Custom calculation function
function difference(prediction, comparison; absolute=false)
    # Calculate difference
    diff_layer = prediction - comparison
    # Get absolute values (optional)
    if absolute
        replace!(x -> !isnothing(x) ? abs(x) : x, diff_layer.grid)
    end
    return diff_layer
end

# Calculate differences
diff_pred_raw = difference(pred.richness, raw.richness)
diff_pred_sdm = difference(pred.richness, sdm.richness)
diff_sdm_raw = difference(sdm.richness, raw.richness)
diffs = [diff_pred_raw, diff_pred_sdm, diff_sdm_raw]
map(x -> extrema(filter(!isnothing, x.grid)), diffs)
maxlim = 30

# Custom plot function
function difference_plot(layer, lim; title="")
    limrange = (-lim, lim)
    diff_map = plot_layer(
        layer;
        c=cgrad(:PuOr; rev=true),
        clim=limrange,
        title="Richness difference",
        colorbar_title="Difference",
    )
    diff_hist = histogram(
        [
            filter(x -> !isnothing(x) && x > 0, layer.grid),
            filter(x -> !isnothing(x) && x <= 0, layer.grid),
        ];
        bins=:rice,
        c=cgrad(:PuOr; rev=true),
        legend=:none,
        ylim=limrange, # xlabel = "Difference",
        title="Distribution of difference values",
        orientation=:horizontal,
    )
    diff_title = plot(; annotation=(0.5, 0.5, "$(title)"), framestyle=:none)
    l = @layout [t{0.01h}; a{0.6w} b{0.38w}]
    diff_plot = plot(diff_title, diff_map, diff_hist; size=(800, 400), dpi=200, layout=l)
    return diff_plot
end

# Plot differences
plot_pred_sdm = difference_plot(
    diff_pred_sdm, maxlim; title="Direct richness $(uppercase(outcome)) vs SDM richness"
)
plot_pred_raw = difference_plot(
    diff_pred_raw, maxlim; title="Direct richness $(uppercase(outcome)) vs raw richness"
)
plot_sdm_raw = difference_plot(diff_sdm_raw, maxlim; title="SDM richness vs raw richness")

## Extras

# Absolute richness difference
richness_absdiff = plot_layer(
    difference(pred.richness, sdm.richness; absolute=true);
    c=:inferno, # clim = (-Inf, Inf),
    title="Direct richness predictions difference ($(uppercase(outcome)))",
    colorbar_title="Difference from SDM-predicted richness (absolute)",
    dpi=200,
)

# LCBD predictions
lcbd_plot = plot_layer(
    lcbd_pred;
    c=:viridis,
    title="Direct LCBD predictions ($(uppercase(outcome)))",
    colorbar_title="LCBD scores",
    clim=(0, 1),
    dpi=200,
)

# LCBD difference
lcbd_diff = plot_layer(
    difference(pred.lcbd, sdm.lcbd; absolute=true);
    c=:inferno, # clim = (-Inf, Inf),
    title="Direct richness predictions difference ($(uppercase(outcome)))",
    colorbar_title="Difference in predicted LCBD (absolute)",
    dpi=200,
)


export
    TestMode, TrainMode,
    FluxOptApp, OptimOptApp, SciMLOptApp

abstract type Flows end
abstract type NormalizingFlows <: Flows end
abstract type ContinuousNormalizingFlows <: NormalizingFlows end
abstract type InfinitesimalContinuousNormalizingFlows <: ContinuousNormalizingFlows end

abstract type AbstractICNF{T} <: InfinitesimalContinuousNormalizingFlows where {T <: AbstractFloat} end
abstract type AbstractCondICNF{T} <: InfinitesimalContinuousNormalizingFlows where {T <: AbstractFloat} end

abstract type Mode end
struct TestMode <: Mode end
struct TrainMode <: Mode end

abstract type OptApp end
struct FluxOptApp <: OptApp end
struct OptimOptApp <: OptApp end
struct SciMLOptApp <: OptApp end

# MLJ interface

abstract type MLJICNF <: MLJModelInterface.Unsupervised end

# Distributions interface

abstract type ICNFDistribution <: ContinuousMultivariateDistribution end


function simulate_waveforms( mcevents::TypedTables.Table, sim::Simulation{T},
                             output_dir::AbstractString,
                             output_base_name::AbstractString = "generated_waveforms";
                             chunk_n_physics_events::Int = 1000,
                             Δt::RealQuantity = 4u"ns",
                             max_nsteps::Int = 1000,
                             diffusion::Bool = false,
                             self_repulsion::Bool = false,
                             number_of_carriers::Int = 1,
                             number_of_shells::Int = 1,
                             verbose = false) where {T <: SSDFloat}
    n_total_physics_events = length(mcevents)
    Δtime = T(to_internal_units(Δt))
    n_contacts = length(sim.detector.contacts)
    @info "Detector has $(n_contacts) contact(s)"
    contacts = sim.detector.contacts
    if !ispath(output_dir) mkpath(output_dir) end
    nfmt(i::Int) = format(i, zeropadding = true, width = length(digits(n_total_physics_events)))
    evt_ranges = chunked_ranges(n_total_physics_events, chunk_n_physics_events)
    @info "-> $(length(flatview(mcevents.edep))) energy depositions to simulate."

    for evtrange in evt_ranges
        ofn = joinpath(output_dir, "$(output_base_name)_evts_$(nfmt(first(evtrange)))-$(nfmt(last(evtrange))).h5")
        @info "Now simulating $(evtrange) and storing it in\n\t \"$ofn\""
        mcevents_sub = simulate_waveforms(mcevents[evtrange], sim; Δt, max_nsteps, diffusion, self_repulsion, number_of_carriers, number_of_shells, verbose)

        HDF5.h5open(ofn, "w") do output
            LegendHDF5IO.writedata(output, "generated_waveforms", mcevents_sub)
        end
    end
end


export TabularCFRPolicy

struct InfoStateNode{M<:AbstractVector{Bool}}
    strategy::Vector{Float64}
    cumulative_regret::Vector{Float64}
    cumulative_strategy::Vector{Float64}
    mask::M
end

function InfoStateNode(mask)
    n = sum(mask)
    InfoStateNode(
        fill(1 / n, n),
        zeros(n),
        zeros(n),
        mask
    )
end

#####
# TabularCFRPolicy
#####

mutable struct TabularCFRPolicy{S,T,R<:AbstractRNG} <: AbstractCFRPolicy
    nodes::Dict{S,InfoStateNode}
    behavior_policy::QBasedPolicy{TabularLearner{S,T},WeightedExplorer{true,R}}
    is_reset_neg_regrets::Bool
    is_linear_averaging::Bool
    weighted_averaging_delay::Int
    is_alternating_update::Bool
    rng::R
    n_iteration::Int
end

(p::TabularCFRPolicy)(env::AbstractEnv) = p.behavior_policy(env)

RLBase.prob(p::TabularCFRPolicy, env::AbstractEnv) = prob(p.behavior_policy, env)

"""
    TabularCFRPolicy(;kwargs...)

Some useful papers while implementing this algorithm:

- [An Introduction to Counterfactual Regret Minimization](http://modelai.gettysburg.edu/2013/cfr/cfr.pdf)
- [MONTE CARLO SAMPLING AND REGRET MINIMIZATION FOR EQUILIBRIUM COMPUTATION AND DECISION-MAKING IN LARGE EXTENSIVE FORM GAMES](http://mlanctot.info/files/papers/PhD_Thesis_MarcLanctot.pdf)
- [Solving Large Imperfect Information Games Using CFR⁺](https://arxiv.org/pdf/1407.5042.pdf)
- [Revisiting CFR⁺ and Alternating Updates](https://arxiv.org/pdf/1810.11542v1.pdf)
- [Solving Imperfect-Information Games via Discounted Regret Minimization](https://arxiv.org/pdf/1809.04040.pdf)

# Keyword Arguments

- `is_alternating_update=true`: If `true`, we update the players alternatively.
- `is_reset_neg_regrets=true`: Whether to use **regret matching⁺**.
- `is_linear_averaging=true`
- `weighted_averaging_delay=0`. The averaging delay in number of iterations. Only valid when `is_linear_averaging` is set to `true`.
- `state_type=String`, the data type of information set.
- `rng=Random.GLOBAL_RNG`
"""
function TabularCFRPolicy(;
    is_reset_neg_regrets = true,
    is_linear_averaging = true,
    weighted_averaging_delay = 0,
    is_alternating_update = true,
    state_type = String,
    rng = Random.GLOBAL_RNG,
    n_iteration = 1,
)
    TabularCFRPolicy(
        Dict{state_type,InfoStateNode}(),
        TabularRandomPolicy(;
            rng = rng,
            table = Dict{state_type,Vector{Float64}}(),
            is_normalized = true,
        ),
        is_reset_neg_regrets,
        is_linear_averaging,
        weighted_averaging_delay,
        is_alternating_update,
        rng,
        n_iteration,
    )
end

"Update the `behavior_policy`"
function RLBase.update!(p::TabularCFRPolicy)
    for (k, v) in p.nodes
        s = sum(v.cumulative_strategy)
        if s != 0
            m = v.mask
            strategy = zeros(length(m))
            strategy[m] .= v.cumulative_strategy ./ s
            update!(p.behavior_policy, k => strategy)
        else
            # The TabularLearner will return uniform distribution by default.
            # So we do nothing here.
        end
    end
end

"Run one interation"
function RLBase.update!(p::TabularCFRPolicy, env::AbstractEnv)
    w = p.is_linear_averaging ? max(p.n_iteration - p.weighted_averaging_delay, 0) : 1
    if p.is_alternating_update
        for x in players(env)
            if x != chance_player(env)
                cfr!(p.nodes, env, x, w)
                regret_matching!(p)
            end
        end
    else
        cfr!(p.nodes, env, nothing, w)
        regret_matching!(p)
    end
    p.n_iteration += 1
end

"""
Symbol meanings:

π: reach prob
π′: new reach prob
π₋ᵢ: opponents' reach prob
p: player to update. `nothing` means simultaneous update.
w: weight
v: counterfactual value **before weighted by opponent's reaching probability**
V: a vector containing the `v` after taking each action with current information set. Used to calculate the **regret value**
"""
function cfr!(nodes, env, p, w, π = Dict(x => 1.0 for x in players(env)))
    if is_terminated(env)
        reward(env, p)
    else
        if current_player(env) == chance_player(env)
            v = 0.0
            for (a, pₐ) in zip(action_space(env), prob(env))
                if pₐ > 0
                    π′ = copy(π)
                    π′[current_player(env)] *= pₐ
                    v += pₐ * cfr!(nodes, child(env, a), p, w, π′)
                end
            end
            v
        else
            v = 0.0
            legal_actions = legal_action_space(env)
            M = legal_action_space_mask(env)
            node = get!(nodes, state(env), InfoStateNode(M))

            is_update = isnothing(p) || p == current_player(env)
            V = is_update ? Vector{Float64}(undef, length(legal_actions)) : nothing

            for (i, a) in enumerate(legal_actions)
                σᵢ = node.strategy[i]
                π′ = copy(π)
                π′[current_player(env)] *= σᵢ

                vₐ = cfr!(nodes, child(env, a), p, w, π′)
                is_update && (V[i] = vₐ)
                v += σᵢ * vₐ
            end

            if is_update
                πᵢ = π[current_player(env)]
                π₋ᵢ = reduce(*, (prob for (p, prob) in π if p != current_player(env)))
                node.cumulative_regret .+= π₋ᵢ .* (V .- v)
                node.cumulative_strategy .+= w .* πᵢ .* node.strategy
            end
            v
        end
    end
end

function regret_matching!(p::TabularCFRPolicy)
    for node in values(p.nodes)
        regret_matching!(node; is_reset_neg_regrets = p.is_reset_neg_regrets)
    end
end

regret_matching!(node::InfoStateNode; kwargs...) =
    regret_matching!(node.strategy, node.cumulative_regret; kwargs...)

function regret_matching!(strategy, cumulative_regret; is_reset_neg_regrets = true)
    if is_reset_neg_regrets
        for i in 1:length(cumulative_regret)
            if cumulative_regret[i] < 0
                cumulative_regret[i] = 0
            end
        end
    end
    s = mapreduce(x -> max(0, x), +, cumulative_regret)
    if s > 0
        strategy .= max.(0.0, cumulative_regret) ./ s
    else
        fill!(strategy, 1 / length(strategy))
    end
end


using Revise
using CUDA
using StaticArrays
using StatsBase: sample
using BenchmarkTools

nbins = 10
items = Int32(1e6)
hist = zeros(Float32, nbins)
x = ones(Float32, items)
idx = Int64.(rand(1:nbins, items))

hist_gpu = CuArray(hist)
x_gpu = CuArray(x)
idx_gpu = CuArray(idx)

hist .- Array(hist_gpu)
sum(hist) - sum(Array(hist_gpu))

# base kernel
function kernel!(h::CuDeviceVector{T}, x::CuDeviceVector{T}, xid::CuDeviceVector{S}) where {T,S}

    nbins = size(h, 1)
    it = threadIdx().x
    ib = blockIdx().x
    id = blockDim().x

    shared = @cuDynamicSharedMem(T, nbins)
    fill!(shared, 0)
    fill!(h, 0)
    sync_threads()

    i_tot = size(x, 1)
    iter = 0
    while iter * id < i_tot
        i = it + id * iter
        if i <= size(xid, 1)
            @inbounds k = Base._to_linear_index(h, xid[i])
            @inbounds CUDA.atomic_add!(pointer(shared, k), x[i])
        end
        iter += 1
    end
    sync_threads()
    # loop to cover cases where nbins > nthreads
    for i in 1:(nbins - 1) ÷ id + 1
        bin_id = it + id * (i - 1)
        if bin_id <= nbins
            @inbounds CUDA.atomic_add!(pointer(h, bin_id), shared[bin_id])
        end
    end
    return nothing
end

# base approach - block built along the cols first, the rows (limit collisions)
function hist!(h::AbstractVector{T}, x::AbstractVector{T}, xid::AbstractVector{S}; MAX_THREADS=256) where {T,S}
    threads = min(MAX_THREADS, size(xid, 1))
    @cuda blocks = 1 threads = threads shmem = sizeof(T) * size(h,1) kernel!(h, x, xid)
    return
end

@btime CUDA.@sync hist!($hist_gpu, $x_gpu, $idx_gpu, MAX_THREADS=1024)