module RingArrays

using VirtualArrays

import Base: size, getindex, checkbounds, display, RefValue, showerror
export RingArray, size, checkbounds, display
export getindex, load_block, RingArrayBoundsError
export showerror, OverwriteError, RingArrayFullError, RingArrayBoundsError

type RingArray{T, N} <: AbstractArray{T, N}
    next_write::Int
    max_blocks::Int
    blocks::Array{AbstractArray{T, N}, 1}
    num_users::Array{Int, 1}
    block_size::Tuple
    block_length::Int
    range::UnitRange{Int}
    data_length::Int

    function RingArray(;max_blocks::Int=10, block_size::Tuple=(10,), data_length::Int=100)
        if length(block_size) != N
            throw(DimensionMismatch("block size $(block_size) does not match dimension of RingArray $(N)"))
        end
        return new(1, max_blocks,
            Array{AbstractArray{T, N}, 1}(max_blocks),
            zeros(Int, max_blocks), block_size,
            block_size[1], 1:0, data_length)
    end
end

type OverwriteError <: Exception
    ring::RingArray
end

type RingArrayBoundsError <: Exception
    ring::RingArray
    i
end

type RingArrayFullError <: Exception
    ring::RingArray
end

function showerror(io::IO, err::OverwriteError)
    next_write = err.ring.next_write
    num_users = err.ring.num_users[err.ring.next_write]
    print(io, "OverwriteError: Cannot overwrite block $(next_write) since it has $(num_users) views")
end

function showerror(io::IO, err::RingArrayBoundsError)
    range = tuple(err.ring.range, err.ring.block_size[2:end]...)
    print(io, "RingArrayBoundsError: Cannot index $(err.i), outside of range $range")
end

function showerror(io::IO, err::RingArrayFullError)
    print(io, "RingArrayFullError: Cannot load another block, max data length is $(err.ring.data_length)")
end

function display(ring::RingArray)
    display(ring.blocks)
end

function size{T, N}(ring::RingArray{T, N})
    if N <= 1
        return tuple(ring.data_length,)
    else
        return tuple(ring.data_length, ring.block_size[2:end]...)
    end
end

checkbounds(ring::RingArray) = true # because of warnings

function checkbounds{T, N}(ring::RingArray{T, N}, indexes::UnitRange{Int}...)
    num_ranges = length(indexes)
    if num_ranges >= N
        if ring.range.start > indexes[1].start || ring.range.stop < indexes[1].stop
            throw(RingArrayBoundsError(ring, indexes))
        end
        for i in 2:N
            if 1 > indexes[i].start || ring.block_size[i] < indexes[i].stop
                throw(RingArrayBoundsError(ring, indexes))
            end
        end
    elseif num_ranges == 1
        index = indexes[1]
        if ring.range.start > fix_zero_index(index.start, ring.data_length)  ||
            ring.range.stop < fix_zero_index(index.stop, ring.data_length)
            throw(RingArrayBoundsError(ring, indexes))
        end
        if divide(index.start, ring.data_length) != divide(index.stop, ring.data_length)
            throw(RingArrayBoundsError(ring, indexes))
        end
        if get_last_index(ring) < index.stop
            throw(RingArrayBoundsError(ring, indexes))
        end
    else
        if ring.range.start > indexes[1].start || ring.range.stop < indexes[1].stop
            throw(RingArrayBoundsError(ring, indexes))
        end
        for i in 2:num_ranges - 1
            if 1 > indexes[i].start || ring.block_size[i] < indexes[i].stop
                throw(RingArrayBoundsError(ring, indexes))
            end
        end
        largest_index_allowed = 1
        for i in num_ranges:N
            largest_index_allowed *= size(ring)[i]
        end
        if 1 > indexes[num_ranges].start || largest_index_allowed < indexes[num_ranges].stop
            throw(RingArrayBoundsError(ring, indexes))
        end
    end
    return true
end

function checkbounds{T, N}(ring::RingArray{T, N}, indexes::Int...)

    if ring.range.start > indexes[1] || ring.range.stop < indexes[1]
        throw(RingArrayBoundsError(ring, indexes))
    else
        for i in 2:N
            if 1 > indexes[i] || ring.block_size[i] < indexes[i]
                throw(RingArrayBoundsError(ring, indexes))
            end
        end
    end
    return true
end

getindex(ring::RingArray) = nothing # Warnings told me to make this

function getindex(ring::RingArray, i::Int...)
    i = expand_index(ring, i...)
    checkbounds(ring, i...)
    result = 0

    block_index = fix_zero_index(divide(i[1], ring.block_length), ring.max_blocks)
    index = collect([fix_zero_index(i[1], ring.block_length), i[2:end]...])

    return ring.blocks[block_index][index...]
end

function getindex(ring::RingArray, i::UnitRange...)
    add_users(ring, i...)
    return get_view(ring, i...)
end

function load_block{T, N}(ring::RingArray{T ,N}, block::AbstractArray{T, N})
    can_load_block!(ring, block)

    ring.blocks[ring.next_write] = block
    ring.next_write = fix_zero_index(ring.next_write + 1, ring.max_blocks)

    if ring.range.stop < ring.max_blocks * ring.block_length
        ring.range = ring.range.start:ring.range.stop + ring.block_length
    else
        ring.range = ring.range.start + ring.block_length:ring.range.stop + ring.block_length
    end
    nothing
end

# Util functions

function can_load_block!(ring::RingArray, block::AbstractArray)
    if ring.num_users[ring.next_write] > 0
        gc()
        if ring.num_users[ring.next_write] > 0
            throw(OverwriteError(ring))
        end
    end

    if ring.range.stop >= ring.data_length
        throw(RingArrayFullError(ring))
    end

    check_dimensions(ring, block)
end

function expand_index{T, N}(ring::RingArray{T, N}, i::Int...)
    result = collect(i)
    len_needed = N
    len_have = length(i)

    for at in len_have:len_needed - 1
        last_value = result[end]
        if at == 1
            result[end] = fix_zero_index(last_value, ring.data_length)
            push!(result, divide(last_value, ring.data_length))
        else
            result[end] = fix_zero_index(last_value, ring.block_size[at])
            push!(result, divide(last_value, ring.block_size[at]))
        end
    end

    return result
end

function fix_zero_index(value::Int, s::Int)
    value = value % s
    if value == 0
        return s
    else
        return value
    end
end

function divide(value::Int, s::Int)
    if value % s == 0
        return div(value, s)
    else
        return div(value, s) + 1
    end
end

function add_users(ring::RingArray, i::UnitRange...)

    first_block = divide(i[1].start, ring.block_length)
    last_block = divide(i[1].stop, ring.block_length)

    for i in first_block:last_block
        block_index = fix_zero_index(i, ring.max_blocks)
        ring.num_users[block_index] += 1
    end

end

function remove_users(ring::RingArray, i::UnitRange...)

    first_block = divide(i[1].start, ring.block_length)
    last_block = divide(i[1].stop, ring.block_length)

    for i in first_block:last_block
        block_index = fix_zero_index(i, ring.max_blocks)
        ring.num_users[block_index] -= 1
    end

end

function get_view(ring::RingArray, i::UnitRange...)

    view = virtual_vcat(sub(ring, i...))

    first_block = divide(i[1].start, ring.block_length)
    last_block = divide(i[1].stop, ring.block_length)

    function when_done(view)
        remove_users(ring, i...)
    end

    finalizer(view, when_done)

    return view
end

function check_dimensions(ring::RingArray, block::AbstractArray)
    if ring.block_size != size(block)
        throw(DimensionMismatch("block size $(size(block)) does not match what RingArray expects $(ring.block_size)"))
    end
end

function get_last_index(ring::RingArray)
    last_index = 1
    for i in size(ring)
        last_index *= i
    end
    return last_index
end

end


for dir in (:x, :y, :z)
    func! = Symbol(dir, :ticks!)
    attr1 = Symbol(dir, :axis)
    attr2 = Symbol(:set_, dir, :ticklabels)
    types = (AbstractRange, NTuple{N} where N)
    for X in types
        @eval begin
            function $(func!)(ax, locs::$X)
                ax.$(attr1).set_major_locator(plt.matplotlib.ticker.MaxNLocator(length(locs)))
                ax.$(attr1).set_major_locator(plt.matplotlib.ticker.FixedLocator(locs))
                return nothing
            end
            function $(func!)(ax, ticks::NTuple{N,S} where {N,S<:AbstractString}, locs::$X)
                $(func!)(ax, locs); ax.$(attr2)(ticks)
                return nothing
            end
        end
        for Y in types
            @eval begin
                function $(func!)(ax, ticks::$Y, locs::$X)
                    $(func!)(ax, locs); ax.$(attr2)(tick2label(ticks))
                    return nothing
                end
            end
        end
    end
end

function tick2label(ticks::AbstractRange{Int})
    labels = Vector{String}(undef, length(ticks))

    @inbounds @simd for i in eachindex(ticks)
        labels[i] = string(ticks[i])
    end

    return labels
end

function tick2label(ticks::AbstractRange{T}) where T<:Real
    labels = Vector{String}(undef, length(ticks))

    @inbounds @simd for i in eachindex(ticks)
        labels[i] = @sprintf("%.1f", ticks[i])
    end

    return labels
end

function tick2label(ticks::NTuple{N,Int}) where N
    if @generated
        e = Expr(:vect)
        a = Vector{Any}(undef, N)

        @inbounds for i in eachindex(a)
            a[i] = :(string(ticks[$i]))
        end

        e.args = a

        return quote
            $(Expr(:meta, :inline))
            @inbounds return $e
        end
    else
        labels = Vector{String}(undef, N)

        @inbounds @simd for i in eachindex(ticks)
            labels[i] = string(ticks[i])
        end

        return labels
    end
end

function tick2label(ticks::NTuple{N,T}) where {N,T<:Real}
    if @generated
        e = Expr(:vect)
        a = Vector{Any}(undef, N)

        @inbounds for i in eachindex(a)
            a[i] = :(@sprintf("%.1f", ticks[$i]))
        end

        e.args = a

        return quote
            $(Expr(:meta, :inline))
            @inbounds return $e
        end
    else
        labels = Vector{String}(undef, N)

        @inbounds @simd for i in eachindex(ticks)
            labels[i] = @sprintf("%.1f", ticks[i])
        end

        return labels
    end
end


import LowLevelJPEG: @code_str

@testset "Constructors" begin
    @test CodeBits() == CodeBits(0, 0x1)
    @test CodeBits(0x0006000D).data === 0x6000D
    @test CodeBits(6, 0b001101).data === 0x6000D
    @test CodeBits(0, 0b001101).data === 0x00000
end

@testset "Traits" begin
    @test LowLevelJPEG.word(code"001101") === UInt16(0b001101)
    @test LowLevelJPEG.isvalid(code"1000100010001000")
    @test !LowLevelJPEG.isvalid(CodeBits(0x10000000))
    @test length(code"") == 0
    @test length(code"001101") == 6
    @test size(code"001101") == (6,)

    @test code"01"[1] == false
end

@testset "Display" begin
    @test shown_string(code"") == "code\"\""
    @test shown_string(code"001101") == "code\"001101\""
    @test shown_string(CodeBits(0x11FFFF)) == "code\"?\""
end

@testset "Functions" begin
    @test LowLevelJPEG.inc(code"001101") === code"001110"
    @test LowLevelJPEG.shl(code"001101") === code"0011010"

    @test BitArray(code"001101") == BitArray((0, 0, 1, 1, 0, 1))
end


function get_return_type(test::TransferEntropyCausalityTest)
    T = typeof(1.0)

    if test.ηs isa Number
        return T
    else test.ηs isa Vector
        return Vector{T}
    end
end

function get_return_type(test::CrossMappingTest)
    T = typeof(1.0)

    return Vector{T}
end

function get_return_type(test::ConvergentCrossMappingTest)
    T = typeof(1.0)

    return Vector{Vector{T}}
end

function get_return_type(test::JointDistanceDistributionTest)
    T = typeof(1.0)

    return Vector{T}
end

function get_return_type(test::JointDistanceDistributionTTest)
    return OneSampleTTest
end

function get_return_type(test::SMeasureTest)
    T = typeof(1.0)
    return Vector{T}
end

function get_return_type(test::AbstractPredictiveAsymmetryTest)
    T = typeof(1.0)
    return Vector{T}
end

export get_return_type


# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder

# Collection of sources required to build MPCBuilder
name = "SymEngine"
version = v"0.4.0"
sources = [
    "https://github.com/symengine/symengine/releases/download/v$version/symengine-$version.tar.gz" =>
    "dd755901a9e2a49e53ba3bbe3f565f94265af05299e57a7b592186dd35916a1b",
]

# Bash recipe for building across all platforms

script = """cd symengine-$version
"""*raw"""
if [[ ${target} == x86_64-* ]]; then
    export CXXFLAGS=-march=core2
fi
cmake -DCMAKE_INSTALL_PREFIX=$prefix -DCMAKE_TOOLCHAIN_FILE=/opt/$target/$target.toolchain -DBUILD_TESTS=no -DBUILD_BENCHMARKS=no  -DBUILD_SHARED_LIBS=yes -DWITH_MPC=yes -DBUILD_FOR_DISTRIBUTION=yes -DWITH_SYMENGINE_THREAD_SAFE=yes .
make -j
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()

# The products that we will ensure are always built
products(prefix) = [
    LibraryProduct(prefix, "libsymengine", :libsymengine)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    "https://github.com/JuliaMath/GMPBuilder/releases/download/v6.1.2-2/build_GMP.v6.1.2.jl"
    "https://github.com/JuliaMath/MPFRBuilder/releases/download/v4.0.1-3/build_MPFR.v4.0.1.jl"
    "https://github.com/isuruf/MPCBuilder/releases/download/v1.1.0-2/build_MPC.v1.1.0.jl"
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)


# work around for julia issue #28948, Gurobi.jl issue #152
if VERSION ≤ v"0.7-"
    sparse_transpose(A) = sparse(transpose(A))
else
    sparse_transpose(A) = SparseMatrixCSC(transpose(A))
end

## Add Linear constraints

# add single constraint

function add_constr!(model::Model, inds::IVec, coeffs::FVec, rel::Cchar, rhs::Float64)
    length(inds) == length(coeffs) || error("Inconsistent argument dimensions.")
    ret = @grb_ccall(addconstr, Cint, (
        Ptr{Cvoid},    # model
        Cint,         # numnz
        Ptr{Cint},    # cind
        Ptr{Float64}, # cvals
        Cchar,        # sense
        Float64,      # rhs
        Ptr{UInt8}    # name
        ),
        model, length(inds), inds .- Cint(1), coeffs, rel, rhs, C_NULL)
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
    nothing
end

function add_constr!(model::Model, inds::Vector, coeffs::Vector, rel::GChars, rhs::Real)
    add_constr!(model, ivec(inds), fvec(coeffs), cchar(rel), Float64(rhs))
end

function add_constr!(model::Model, coeffs::Vector, rel::GChars, rhs::Real)
    inds = findall(x->x!=0, coeffs)
    vals = coeffs[inds]
    add_constr!(model, inds, vals, rel, rhs)
end

# add multiple constraints

function add_constrs!(model::Model, cbegins::IVec, inds::IVec, coeffs::FVec, rel::CVec, rhs::FVec)
    m = length(cbegins)
    nnz = length(inds)
    (m == length(rel) == length(rhs) && nnz == length(coeffs)) || error("Inconsistent argument dimensions.")

    if m > 0
        ret = @grb_ccall(addconstrs, Cint, (
            Ptr{Cvoid},    # model
            Cint,         # num constraints
            Cint,         # num non-zeros
            Ptr{Cint},    # cbeg
            Ptr{Cint},    # cind
            Ptr{Float64}, # cval
            Ptr{Cchar},   # sense
            Ptr{Float64}, # rhs
            Ptr{UInt8}    # names
            ),
            model, m, nnz, cbegins .- Cint(1), inds .- Cint(1), coeffs,
            rel, rhs, C_NULL)

        if ret != 0
            throw(GurobiError(model.env, ret))
        end
    end
    nothing
end

function add_constrs!(model::Model, cbeg::Vector, inds::Vector, coeffs::Vector, rel::GCharOrVec, rhs::Vector)
    add_constrs!(model, ivec(cbeg), ivec(inds), fvec(coeffs), cvecx(rel, length(cbeg)), fvec(rhs))
end

function add_constrs_t!(model::Model, At::SparseMatrixCSC{Float64}, rel::GCharOrVec, b::Vector)
    n, m = size(At)
    (m == length(b) && n == num_vars(model)) || error("Incompatible argument dimensions.")
    add_constrs!(model, At.colptr[1:At.n], At.rowval, At.nzval, rel, b)
end

function add_constrs_t!(model::Model, At::Matrix{Float64}, rel::GCharOrVec, b::Vector)
    n, m = size(At)
    (m == length(b) && n == num_vars(model)) || error("Incompatible argument dimensions.")
    add_constrs_t!(model, sparse(At), rel, b)
end

function add_constrs!(model::Model, A::CoeffMat, rel::GCharOrVec, b::Vector{Float64})
    m, n = size(A)
    (m == length(b) && n == num_vars(model)) || error("Incompatible argument dimensions.")
    add_constrs_t!(model, sparse_transpose(A), rel, b)
end

# add single range constraint

function add_rangeconstr!(model::Model, inds::IVec, coeffs::FVec, lb::Float64, ub::Float64)
    ret = @grb_ccall(addrangeconstr, Cint, (
        Ptr{Cvoid},    # model
        Cint,         # numnz
        Ptr{Cint},    # cind
        Ptr{Float64}, # cvals
        Float64,      # lower
        Float64,      # upper
        Ptr{UInt8}    # name
        ),
        model, length(inds), inds .- Cint(1), coeffs, lb, ub, C_NULL)
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
    nothing
end

function add_rangeconstr!(model::Model, inds::Vector, coeffs::Vector, lb::Real, ub::Real)
    add_rangeconstr!(model, ivec(inds), fvec(coeffs), Float64(lb), Float64(ub))
end

function add_rangeconstr!(model::Model, coeffs::Vector, lb::Real, ub::Real)
    inds = find(coeffs)
    vals = coeffs[inds]
    add_rangeconstr!(model, inds, vals, lb, ub)
end

# add multiple range constraints

function add_rangeconstrs!(model::Model, cbegins::IVec, inds::IVec, coeffs::FVec, lb::FVec, ub::FVec)
    m = length(cbegins)
    nnz = length(inds)
    (m == length(lb) == length(ub) && nnz == length(coeffs)) || error("Incompatible argument dimensions.")

    if m > 0
        ret = @grb_ccall(addrangeconstrs, Cint, (
            Ptr{Cvoid},    # model
            Cint,         # num constraints
            Cint,         # num non-zeros
            Ptr{Cint},    # cbeg
            Ptr{Cint},    # cind
            Ptr{Float64}, # cval
            Ptr{Float64}, # lower
            Ptr{Float64}, # upper
            Ptr{UInt8}    # names
            ),
            model, m, nnz, cbegins .- Cint(1), inds .- Cint(1), coeffs, lb, ub, C_NULL)

        if ret != 0
            throw(GurobiError(model.env, ret))
        end
    end
    nothing
end

function add_rangeconstrs!(model::Model, cbeg::Vector, inds::Vector, coeffs::Vector, lb::Vector, ub::Vector)
    add_rangeconstrs!(model, ivec(cbeg), ivec(inds), fvec(coeffs), fvec(lb), fvec(ub))
end

function add_rangeconstrs_t!(model::Model, At::SparseMatrixCSC{Float64}, lb::Vector, ub::Vector)
    add_rangeconstrs!(model, At.colptr[1:At.n], At.rowval, At.nzval, lb, ub)
end

function add_rangeconstrs_t!(model::Model, At::Matrix{Float64}, lb::Vector, ub::Vector)
    add_rangeconstrs_t!(model, sparse(At), lb, ub)
end

function add_rangeconstrs!(model::Model, A::CoeffMat, lb::Vector, ub::Vector)
    m, n = size(A)
    (m == length(lb) == length(ub) && n == num_vars(model)) || error("Incompatible argument dimensions.")

    add_rangeconstrs_t!(model, sparse_transpose(A), lb, ub)
end

function get_constrmatrix(model::Model)
    get_constrs(model::Model, 1, num_constrs(model))
end
function get_constrs(model::Model, start::Integer, len::Integer)

    m = num_constrs(model)
    @assert start >= 1
    @assert len >= 0
    @assert start + len <= m + 1
    n = num_vars(model)
    numnzP = Ref{Cint}()
    cbeg = Array{Cint}(undef, len+1)

    ret = @grb_ccall(getconstrs, Cint, (
        Ptr{Cvoid},
        Ptr{Cint},
        Ptr{Cint},
        Ptr{Cint},
        Ptr{Cdouble},
        Cint, # start
        Cint #len
        ),
        model, numnzP, cbeg, C_NULL, C_NULL, Cint(start-1), Cint(len))

    nnz = numnzP[]
    cind = Array{Cint}(undef, nnz)
    cval = Array{Cdouble}(undef, nnz)
    ret = @grb_ccall(getconstrs, Cint, (
                     Ptr{Cvoid},
                     Ptr{Cint},
                     Ptr{Cint},
                     Ptr{Cint},
                     Ptr{Cdouble},
                     Cint, # start
                     Cint #len
                     ),
                     model, numnzP, cbeg, cind, cval, Cint(start-1), Cint(len))
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
    cbeg[end] = nnz
    I = Array{Int64}(undef, nnz)
    J = Array{Int64}(undef, nnz)
    V = Array{Float64}(undef, nnz)
    for i in 1:length(cbeg)-1
        for j in (cbeg[i]+1):cbeg[i+1]
            I[j] = i
            J[j] = cind[j]+1
            V[j] = cval[j]
        end
    end
    return sparse(I, J, V, m, n)
end

const GRB_SOS_TYPE1 = convert(Cint, 1)
const GRB_SOS_TYPE2 = convert(Cint, 2)

function add_sos!(model::Model, sostype::Symbol, idx::Vector{Int}, weight::Vector{Cdouble})
    ((nelem = length(idx)) == length(weight)) || error("Index and weight vectors of unequal length")
    (sostype == :SOS1) ? (typ = GRB_SOS_TYPE1) : ( (sostype == :SOS2) ? (typ = GRB_SOS_TYPE2) : error("Invalid SOS constraint type") )
    ret = @grb_ccall(addsos, Cint, (
                     Ptr{Cvoid},
                     Cint,
                     Cint,
                     Ptr{Cint},
                     Ptr{Cint},
                     Ptr{Cint},
                     Ptr{Cdouble}
                     ),
                     model, convert(Cint, 1), convert(Cint, nelem), Cint[typ], Cint[0], convert(Vector{Cint}, idx.-1), weight)
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
end
function del_sos!(model::Model, idx::Vector{Cint})
    numdel = length(idx)
    ret = @grb_ccall(delsos, Cint, (
                     Ptr{Cvoid},
                     Cint,
                     Ptr{Cint}),
                     model, convert(Cint,numdel), ivec(idx.-1) )
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
end

get_sos_matrix(model::Model) = get_sos(model::Model, 1, num_sos(model))
function get_sos(model::Model, start::Integer, len::Integer)
    m = num_sos(model)
    sostype = Array{Cint}(undef, m)
    @assert m > 0
    @assert start <= m
    @assert len <= m
    n = num_vars(model)
    numnzP = Ref{Cint}()
    cbeg = Array{Cint}(undef, len+1)

    ret = @grb_ccall(getsos, Cint, (
        Ptr{Cvoid},
        Ptr{Cint},
        Ptr{Cint},
        Ptr{Cint},
        Ptr{Cint},
        Ptr{Cdouble},
        Cint, # start
        Cint #len
        ),
        model, numnzP, sostype, cbeg, C_NULL, C_NULL, Cint(start-1), Cint(len))

    nnz = numnzP[]

    cind = Array{Cint}(undef, nnz)
    cval = Array{Cdouble}(undef, nnz)
    ret = @grb_ccall(getsos, Cint, (
                     Ptr{Cvoid},
                     Ptr{Cint},
                     Ptr{Cint},
                     Ptr{Cint},
                     Ptr{Cint},
                     Ptr{Cdouble},
                     Cint, # start
                     Cint #len
                     ),
                     model, numnzP, sostype, cbeg, cind, cval, Cint(start-1), Cint(len))
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
    cbeg[end] = nnz
    I = Array{Int64}(undef, nnz)
    J = Array{Int64}(undef, nnz)
    V = Array{Float64}(undef, nnz)
    for i in 1:length(cbeg)-1
        for j in (cbeg[i]+1):cbeg[i+1]
            I[j] = i
            J[j] = cind[j]+1
            V[j] = cval[j]
        end
    end
    return sparse(I, J, V, m, n), sostype#map(x-> x==Cint(1) ? :SOS1 : :SOS2 ,sostype)
end

del_constrs!(model::Model, idx::T) where {T<:Real} = del_constrs!(model, Cint[idx])
del_constrs!(model::Model, idx::Vector{T}) where {T<:Real} = del_constrs!(model, convert(Vector{Cint},idx))
function del_constrs!(model::Model, idx::Vector{Cint})
    numdel = length(idx)
    ret = @grb_ccall(delconstrs, Cint, (
                     Ptr{Cvoid},
                     Cint,
                     Ptr{Cint}),
                     model, convert(Cint,numdel), idx .- Cint(1))
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
end

function chg_coeffs!(model::Model, cidx::Real, vidx::Real, val::Real)
    ret = @grb_ccall(chgcoeffs, Cint, (
                     Ptr{Cvoid},
                     Cint,
                     Ref{Cint},
                     Ref{Cint},
                     Ref{Float64}),
                     model, 1, cidx - 1, vidx - 1, val)
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
end

chg_coeffs!(model::Model, cidx::AbstractVector{<:Real}, vidx::AbstractVector{<:Real}, val::AbstractVector{<:Real}) =
    chg_coeffs!(model, ivec(cidx), ivec(vidx), fvec(val))

function chg_coeffs!(model::Model, cidx::IVec, vidx::IVec, val::FVec)
    (length(cidx) == length(vidx) == length(val)) || error("Inconsistent argument dimensions.")

    numchgs = length(cidx)
    ret = @grb_ccall(chgcoeffs, Cint, (
                     Ptr{Cvoid},
                     Cint,
                     Ptr{Cint},
                     Ptr{Cint},
                     Ptr{Float64}),
                     model, convert(Cint,numchgs), cidx .- Cint(1), vidx .- Cint(1), val)
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
end

function getcoeff!(val::FVec, model::Model, cidx::Integer, vidx::Integer)
    Base.depwarn("getcoeff!(val, model, cidx, vidx) is deprecated. Instead you can retrieve a coefficient without allocating a vector by doing `val = getcoeff(model, cidx, vidx)`", :grb_getcoeff)
    @assert length(val) == 1
    ret = @grb_ccall(getcoeff, Cint, (
        Ptr{Cvoid},
        Cint,
        Cint,
        Ptr{Float64}),
        model, cidx-Cint(1), vidx-Cint(1), val)
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
end

function getcoeff(model::Model, cidx::Integer, vidx::Integer)
    out = Ref{Float64}()
    ret = @grb_ccall(getcoeff, Cint, (
        Ptr{Cvoid},
        Cint,
        Cint,
        Ref{Float64}),
        model, cidx - 1, vidx - 1, out)
    if ret != 0
        throw(GurobiError(model.env, ret))
    end
    out[]
end


# Sample postprocessing script for virtual satellite trackings.
# Simultaneous virtual satellite trackings at multiple locations.
# Outputs are stored in binary format for sharing within Julia.
#
# Usage:
#   julia -t nthreads demo_pointsextracting.jl
# or
#   JULIA_NUM_THREADS=nthreads julia demo_pointsextracting.jl
#
# Hongyang Zhou, hyzhou@umich.edu

using Glob, Vlasiator
using Vlasiator: RE # Earth radius, [m]
using JLD2: jldsave

function extract_vars(files, locations)
   nfiles = length(files)
   nlocs = length(locations)
   # variables to be extracted
   t   = zeros(Float32, nfiles)
   rho = zeros(Float32, nlocs, nfiles)
   v   = zeros(Float32, 3, nlocs, nfiles)
   p   = zeros(Float32, nlocs, nfiles)
   b   = zeros(Float32, 3, nlocs, nfiles)
   e   = zeros(Float32, 3, nlocs, nfiles)

   ids = Vector{UInt64}(undef, nlocs)
   load(files[1]) do meta
      for iloc in eachindex(locations)
         ids[iloc] = getcell(meta, locations[iloc])
      end
   end

   # Extract data from each frame
   Threads.@threads for i = eachindex(files)
      meta = load(files[i])
      t[i] = meta.time
      rho[:,i] = readvariable(meta, "proton/vg_rho", ids)
      v[:,:,i] = readvariable(meta, "proton/vg_v",   ids)
      p[  :,i] = readvariable(meta, "vg_pressure",   ids)
      b[:,:,i] = readvariable(meta, "vg_b_vol",      ids)
      e[:,:,i] = readvariable(meta, "vg_e_vol",      ids)
   end

   # Save into binary file
   jldsave("satellites.jld2"; locations, t, rho, v, p, b, e)
end

#####

files = glob("bulk*.vlsv", "./")

# virtual satellite locations
locations = [[7RE, 0, 0], [9RE, 0, 0], [11RE, 0, 0], [12RE, 0, 0], [13RE, 0, 0],
   [14RE, 0, 0], [15RE, 0, 0], [16RE, 0, 0], [17RE, 0, 0], [29.3RE, 0, 0]]

println("Number of files: $(length(files))")
println("Number of virtual satellites: $(length(locations))")
println("Running with $(Threads.nthreads()) threads...")

@time extract_vars(files, locations)

println("Virtual satellite extraction done!")


"""
A `PointCloud` is a container for points sharing a common set of per-point
attributes.  Points within the cloud can be accessed by index, and vectors of
attributes may be accessed by name.  There is one required attribute
`position`, which is used to build a datastructure for fast spatial lookup of
k-nearest neighbours, or points within a fixed radius.

### Constructors

```julia
PointCloud(positions::Vector{FixedSizeArrays.Vec{3, AbstractFloat}})
PointCloud(positions::Matrix{AbstractFloat})
```

### Example

```julia
using FixedSizeArrays
using PointClouds

# Create positions of 3D Vec coordinates
positions = [Vec(i,-i,1.0) for i = 1:10]

# Create a PointCloud from positions
cloud = PointCloud(positions)

# Create new intensity attributes
cloud[:intensity] = collect(1:10)

# Create capture time for each point sample
cloud[:time] = [time() for i = 1:10]

# Find low intensity points
low_intensity_cloud = cloud[cloud[:intensity] .< 5]

# Find vector of time stamps per point
gps_time = low_intensity_cloud[:time]

# Find points within 5 units of [1,1,1]
nearby = cloud[inrange(cloud, [1,1,1], 5.0)]
```
"""

type PointCloud{Dim,T,SIndex}
    positions::Vector{SVector{Dim,T}}
    spatial_index::SIndex
    attributes::Dict{Symbol,Vector}
end

# Create a `PointCloud` for FixedSizeArrays Vec using a KDTree for spatial
# indexing based on `positions`.
# TODO should be able to change the spatial index
function PointCloud{Dim, T <: AbstractFloat}(positions::Vector{SVector{Dim, T}})
    PointCloud(positions, KDTree(positions), Dict{Symbol,Vector}())
end

# Create a `PointCloud` from an 3xN array of points, using a KDTree for spatial
# indexing based on `positions`.
PointCloud{T <: AbstractFloat}(points::Matrix{T}) = PointCloud(convert_positions(points))

"""
    convert_positions{T <: Real}(points::Matrix{T}) -> positions::Vector{FixedSizeArrays.Vec{3, Float64}}

Convert a standard julia array of points into a FixedSizeArray
"""
function convert_positions{T}(points::Matrix{T})
    ndims = size(points, 1)
    npoints = size(points, 2)
    if isbits(T)
        data = reinterpret(SVector{ndims, T}, points, (npoints, ))
    else
        data = SVector{ndims, T}[SVector{ndims, T}(points[:, i]) for i in 1:npoints]
    end
    return data
end

function show(io::IO, cloud::PointCloud)
    print(io, "PointCloud(N=$(length(cloud)), attributes=$(keys(cloud.attributes)))")
end

#------------------------
# Container functionality acting on columns
keys(cloud::PointCloud) = keys(cloud.attributes)
haskey(cloud::PointCloud, attrname::Symbol) = haskey(cloud.attributes, attrname)

function getindex(cloud::PointCloud, attrname::Symbol)
    haskey(cloud, attrname) || error("Point cloud $cloud has no attribute $attrname")
    cloud.attributes[attrname]
end

function setindex!(cloud::PointCloud, value::Vector, attrname::Symbol)
    if length(value) != length(cloud)
        error("length($attrname) = $(length(value)) not equal to number of points = $(length(cloud))")
    end
    cloud.attributes[attrname] = value
end

# Access to columns with "blessed" type for efficiency.
# TODO: Figure out how we can remove these, so that eg, cloud[:normal] is the
# definitive way to access normals.
# Note that the `positions()` is used rather than `position()`, to avoid a
# clash with `Base.position()` which means IO stream position... ugh.

"Return vector of point cloud positions"
positions(cloud::PointCloud)   = cloud.positions

"Return array of point cloud normals as fixed size vectors"
function normals{Dim,T,SIndex}(cloud::PointCloud{Dim,T,SIndex})
    cloud.attributes[:normal]::Vector{Vec{Dim,Float32}}
end

#--------------------------
# Container functionality acting on rows

length(cloud::PointCloud) = length(cloud.positions)
endof(cloud::PointCloud) = length(cloud.positions)

# Subset of a point cloud
function getindex{Dim,T,SIndex}(cloud::PointCloud{Dim,T,SIndex}, row_inds::AbstractVector)
    pos = positions(cloud)[row_inds]
    tree = KDTree(pos)  # TODO NEED TO GET SPATIAL INDEX FROM TYPE
    attrs = Dict{Symbol,Vector}()
    for (k,v) in cloud.attributes
        attrs[k] = v[row_inds]
    end
    PointCloud{Dim,T,SIndex}(pos, tree, attrs)
end

# Concatenate point clouds
function vcat{Dim,T,SIndex}(cloud1::PointCloud{Dim,T,SIndex}, clouds::PointCloud{Dim,T,SIndex}...)
    # Workaround for a weird compiler bug, which causes the deepcopy type assert to fail.
    # See https://github.com/JuliaLang/julia/issues/19041
    #attrs = deepcopy(cloud1.attributes)
    attrs = Dict{Symbol,Vector}()
    for (k,v) in cloud1.attributes
        attrs[k] = deepcopy(v)
    end
    pos = positions(cloud1)
    ks = Set(keys(attrs))
    for cloud in clouds
        pos = [pos; positions(cloud)]
        ki = Set(keys(cloud))
        ki == ks || error("Cannot join clouds with inconsistent attributes")
        for (k,v) in attrs
            append!(v, cloud.attributes[k])
        end
    end
    spatial_index = KDTree(pos)  # TODO NEED TO GET SPATIAL INDEX FROM TYPE
    PointCloud{Dim,T,SIndex}(pos, spatial_index, attrs)
end

"""
Group points into `M = unique(scanid)` point clouds, each containing points
with the same scan id.  Returns a dictionary mapping the unique scanid to each
cloud.
"""
function split_cloud(allpoints, scanid; min_points::Int=10000)
    unique_scans = unique(scanid)
    scans = Dict{eltype(unique_scans), typeof(allpoints)}()
    for id in unique_scans
        scans[id] = allpoints[id .== scanid]
    end
    scans
end

#--------------------------
# Spatial index lookup (knn and ball queries)
knn(cloud::PointCloud, points, k) = knn(cloud.spatial_index, points, k)
inrange(cloud::PointCloud, points, radius) = inrange(cloud.spatial_index, points, radius)

#------------------------
# Utilities for adding columns

"""
    pca_normals!(normals, query_points, cloud; neighbours::Int=10)

Compute normals of a point cloud via PCA

For each point in the 3xN array `query_points`, compute an approximate surface
normal using PCA on the local neighbourhood, using the closest `neighbours`
points.  `cloud` should support spatial lookup via the `knn()` function.
Normals are returned in the 3xN array `normals`.

TODO: Need a generalized version of this for the PCA covariance computation.
"""
function pca_normals!{V<:AbstractVector}(normals::Vector{V}, query_points, cloud; neighbours::Int=10)
    position = positions(cloud)
    for i = 1:length(query_points)
        @inbounds (inds, _) = knn(cloud, query_points[i], neighbours)
        isempty(inds) && continue
        cov_tmp = zeros(SMatrix{3,3,eltype(V)})
        for j = 1:length(inds)
            @inbounds p = position[inds[j]] - query_points[i]
            cov_tmp += p * p'
        end
        # FIXME: Avoid converting to Matrix when eig(SMatrix) is solid
        # See https://github.com/JuliaArrays/StaticArrays.jl/issues/80
        evals, evecs = eig(Symmetric(Matrix(cov_tmp)))
        # Select evec of smallest eval
        normals[i] = evecs[:,1]
    end
    return normals
end

"""
Add normals to point cloud in Float32 precision, compute using PCA.

See pca_normals!() for details.
"""
function add_normals!{Dim}(cloud::PointCloud{Dim}; kwargs...)
    # Use Float32 for estimated normals - heaps of precision at half the memory
    query_points = positions(cloud)
    normals = similar(query_points, SVector{Dim,Float32})
    pca_normals!(normals, query_points, cloud; kwargs...)
    cloud[:normals] = normals
    return cloud
end


function foo_reorder(args...; node::Int64 = Int[], n::Int64 = Int[], m::Int64 = Int[],
    e1::Array{Int64, 1} = Int[], e2::Array{Int64, 1} = Int[],
    neworder::Array{Int64, 1} = Int[], L::Array{Int64, 1} = Int[], pos::Array{Int64, 1} = Int[])
    i  = node - n
    for j in 1:pos[i]
        k = L[i + m * j]
        neworder[iii] = k + 1
        if e2[k] > n
            foo_reorder(node = e2[k], n = n, m = m, e1 = e1, e2 = e2, neworder = neworder, L = L, pos = pos)
        end
        iii = iii + 1
    end
end


abstract type Scalarness end
struct Scalar <: Scalarness end
struct NotScalar <: Scalarness end

isscalar(::Type{Any}) = NotScalar() # if we don't know the type we can't really know if scalar or not
isscalar(::Type{<:AbstractString}) = NotScalar() # We consider strings to be nonscalar
isscalar(::Type{<:Number}) = Scalar() # We consider Numbers to be scalar
isscalar(::Type{Char}) = Scalar() # We consider Character to be scalar
isscalar(::Type{T}) where T = hasmethod(iterate, (T,)) ? NotScalar() : Scalar()

"""
    apply_to_interior_levels(f, iter, lvls)

f is called on all levels up to `max_depth`
that have children and who's children (probably) have children.

`f` must be a function taking the current level number it's first argument,
and the level content as it's second.
That level is replaced with the return value of `f`
"""
function apply_to_interior_levels(ff, iter, lvls, max_depth)
    function inner(level_iter::T, cur_level) where T
        # This is a static Holy Trait dispatch
        this_scalar = isscalar(T)
        children_scalar = isscalar(eltype(T))
        inner(this_scalar, children_scalar, level_iter, cur_level)
    end

    # scalar, any children (Fall back if didn't manage to stop earlier)
    function inner(::Scalar,::Any, childs, cur_level)
        childs # This is just a scalare element
    end

    # not scalar, children scalar
    function inner(::NotScalar,::Scalar, childs, cur_level)
        # Only ops we care about are ones that act on nonscalar children
        childs
    end

    # not scalar, children not scalar (probably)
    function inner(::NotScalar, ::NotScalar, childs, cur_level)
        if cur_level <= max_depth
            # Only expand down if not yet at deepset level we act on
            # (This saves on allocations and time
            childs = apply(childs) do _childs
                # `apply` needs to be wrapped around all operations, including imap/map
                imap(_childs) do child
                    inner(child, cur_level+1)
                end
            end
        end

        ff(cur_level, childs)
    end

    inner(iter, 1)
end

"""
    apply_at_level(ff, iter, lvls)

Apply the given function `ff` at all levels specified in `lvls`
"""
function apply_at_level(ff, iter, lvls)
    max_depth = maximum(lvls)
    apply_to_interior_levels(iter, lvls, max_depth) do cur_level, childs
        ret = if cur_level ∈ lvls
            # If on the level where we act, then act
            apply(ff, childs)
        else
            childs
        end
        ret
    end
end

"""
    apply_at_level(iter, lvl_ops::Associative{<:Integer})

Applies a op
`ops` is a list of functions the same length as `lvls`.
`apply_pel_level` applies the given op, at the correponding level.

If only a single level and a single op is provider,
then this is identical to `apply_at_level`.
"""
function apply_at_level(iter, lvl_ops::AbstractDict{<:Integer})
    max_depth = maximum(keys(lvl_ops))
    apply_to_interior_levels(iter, lvl_ops, max_depth) do cur_level, childs
        # if level is a key then  do the thing, else no change
        if haskey(lvl_ops, cur_level)
            apply(lvl_ops[cur_level], childs)
        else
            childs
        end
    end
end


function BackgroundVelocityFields(bg, grid, clock)
    u = :u ∈ keys(bg) ? regularize_background_field(Face, Center, Center, bg[:u], grid, clock) : ZeroField()
    v = :v ∈ keys(bg) ? regularize_background_field(Center, Face, Center, bg[:v], grid, clock) : ZeroField()
    w = :w ∈ keys(bg) ? regularize_background_field(Center, Center, Face, bg[:w], grid, clock) : ZeroField()

    return (u=u, v=v, w=w)
end

function BackgroundTracerFields(bg, tracer_names, grid, clock)
    tracer_fields =
        Tuple(c ∈ keys(bg) ?
              regularize_background_field(Center, Center, Center, getindex(bg, c), grid, clock) :
              ZeroField()
              for c in tracer_names)

    return NamedTuple{tracer_names}(tracer_fields)
end

#####
##### Convenience for model constructor
#####

function BackgroundFields(background_fields, tracer_names, grid, clock)
    velocities = BackgroundVelocityFields(background_fields, grid, clock)
    tracers = BackgroundTracerFields(background_fields, tracer_names, grid, clock)
    return (velocities=velocities, tracers=tracers)
end

"""
    BackgroundField{F, P}

Temporary container for storing information about BackgroundFields.
"""
struct BackgroundField{F, P}
    func:: F
    parameters :: P
end

"""
    BackgroundField(func; parameters=nothing)

Returns a `BackgroundField` to be passed to `IncompressibleModel` for use
as a background velocity or tracer field.

If `parameters` is not provided, `func` must be callable with the signature

```julia
func(x, y, z, t)
```

If `parameters` is provided, `func` must be callable with the signature

```julia
func(x, y, z, t, parameters)
```
"""
BackgroundField(func; parameters=nothing) = BackgroundField(func, parameters)

regularize_background_field(X, Y, Z, f::BackgroundField{<:Function}, grid, clock) =
    FunctionField{X, Y, Z}(f.func, grid; clock=clock, parameters=f.parameters)

regularize_background_field(X, Y, Z, func::Function, grid, clock) =
    FunctionField{X, Y, Z}(func, grid; clock=clock)

function regularize_background_field(X, Y, Z, field::AbstractField, grid, clock)
    if location(field) != (X, Y, Z)
        throw(ArgumentError("Cannot use field at $(location(field)) as a background field at $((X, Y, Z))"))
    end

    return field
end

Base.show(io::IO, field::BackgroundField{F, P}) where {F, P} =
    print(io, "BackgroundField{$F, $P}", '\n',
          "├── func: $(short_show(field.func))", '\n',
          "└── parameters: $(field.parameters)")


using LinearAlgebra, NLPModels, Stopping

include("backls.jl")
include("uncons.jl")

#https://juliasmoothoptimizers.github.io/SolverBenchmark.jl/latest/tutorial/
#In this tutorial we illustrate the main uses of SolverBenchmark.
using DataFrames, Printf, SolverBenchmark
#CUTEst is a collection of test problems
using CUTEst

problems_unconstrained = CUTEst.select(contype="unc")
n = length(problems_unconstrained) #240
#problems_boundconstrained = CUTEst.select(contype="bounds")
#n = length(problems_boundconstrained) #124

printstyled("Benchmark solvers: \n", color = :green)

n = min(n, 3)

#Names of 3 solvers:
names = [:armijo, :wolfe, :armijo_wolfe]
p1 = PrmUn(); p2 = PrmUn(ls_func = wolfe); p3 = PrmUn(ls_func = armijo_wolfe)
paramDict = Dict(:armijo => p1, :wolfe => p2, :armijo_wolfe => p3)
#Initialization of the DataFrame for n problems.
stats = Dict(name => DataFrame(:id => 1:n,
         :name => [@sprintf("prob%s", problems_unconstrained[i]) for i = 1:n],
         :nvar => zeros(Int64, n),
         :status => [:Unknown for i = 1:n],
         :f => NaN*ones(n),
         :t => NaN*ones(n),
         :iter => zeros(Int64, n),
         :eval_f => zeros(Int64, n),
         :eval_g => zeros(Int64, n),
         :eval_H => zeros(Int64, n),
         :score => NaN*ones(n)) for name in names)

for i=1:n
  nlp_cutest = CUTEst.CUTEstModel(problems_unconstrained[i])
  @show i, problems_unconstrained[i], nlp_cutest.meta.nvar
  #update the stopping with the new problem
  stop_nlp = NLPStopping(nlp_cutest,
                         NLPAtX(nlp_cutest.meta.x0),
                         max_iter = 20,
                         optimality_check = unconstrained_check)

  for name in names

    #solve the problem
    global_newton(stop_nlp, paramDict[name])

    #update the stats from the Stopping
    stats[name].nvar[i] = nlp_cutest.meta.nvar
    stats[name].status[i] = status(stop_nlp)
    stats[name].f[i] = stop_nlp.current_state.fx
    stats[name].t[i] = stop_nlp.current_state.current_time - stop_nlp.meta.start_time
    stats[name].iter[i] = stop_nlp.meta.nb_of_stop
    stats[name].score[i] = unconstrained_check(nlp_cutest, stop_nlp.current_state)
    stats[name].eval_f[i] = getfield(stop_nlp.current_state.evals, :neval_obj)
    stats[name].eval_g[i] = getfield(stop_nlp.current_state.evals, :neval_grad)
    stats[name].eval_H[i] = getfield(stop_nlp.current_state.evals, :neval_hess)

    #reinitialize the Stopping and the nlp
    reinit!(stop_nlp, rstate = true, x = nlp_cutest.meta.x0)
    reset!(stop_nlp.pb)
  end

  #finalize nlp
  finalize(nlp_cutest)
end #end of main loop

for name in names
@show stats[name]
end

#You can export the table in Latex
#latex_table(stdout, stats[:armijo])

#or run a performance profile:
#using Plots
#pyplot()
#cost(df) = (df.status != :Optimal) * Inf + df.t
#p = performance_profile(stats, cost)
#Plots.svg(p, "profile2")

#or a profile wall:
#solved(df) = (def.status .== :Optimal)
#costs = [df -> .!sovled(df) * Inf + df.t, df -> .!sovled(df) * Inf + df.iter]
#costnames = ["Time", "Iterations"]
#p = profile_solvers(stats, costs, costnames)
#Plots.svg(p, "profile3")

printstyled("The End.", color = :green)


#####
##### Script to replicate the mysterious OOM errors
#####

# NUM_FILTERS=64 julia --project scripts/profile/debug_oom.jl
# NUM_FILTERS=128 julia --project scripts/profile/debug_oom.jl

using AlphaZero
using Setfield

if !haskey(ENV, "NUM_FILTERS")
  error("You should set the NUM_FILTERS env variable.")
end

num_filters = parse(Int, ENV["NUM_FILTERS"])

println("Number of filters: $num_filters")

exp = Examples.experiments["connect-four"]

exp = @set exp.netparams.num_filters = num_filters

Scripts.test_grad_updates(exp, num_games=5000)


module QuOptimalControl

include("problems.jl")
include("cost_functions.jl")
include("tools.jl")
include("grape_tools.jl")
include("timeevolution.jl")
include("algorithms.jl")
include("dCRAB.jl")
include("GRAPE.jl")
include("solve.jl")
include("visualisation.jl")

export C1, C2, C3, C4, C5, C6, C7, fom_func
export init_GRAPE, init_ensemble
export commutator, eig_factors, expm_exact_gradient, trace_matmul, save, load
export pw_evolve,
    pw_evolve_save,
    pw_ham_save!,
    pw_gen_save!,
    pw_prop_save!,
    evolve_func,
    evolve_func!
export IntegratorType, Piecewise, Continuous
export algorithm,
    gradientBased, gradientFree, SolutionResult, GRAPE, ADGRAPE, EnsembleSolutionResult
export solve
export grad_func
export Problem,
    SystemType,
    ClosedSystem,
    OpenSystem,
    StateTransfer,
    UnitaryGate,
    CoherenceTransfer,
    EnsembleProblem,
    Experiment,
    ExperimentInterface
export visualise_expt_val, visualise_pulse

end


using Test
using Random
using ClimaCoupler, Dates, Unitful

@testset "Coupler Interface" begin
    Random.seed!(26)

    coupler = CouplerState()

    data = rand(10, 10)
    date = DateTime(2021)
    coupler_add_field!(coupler, :test1, data, nothing, date, u"kg")
    coupler_add_field!(coupler, :test2, data, nothing, date, u"km/hr")
    coupler_add_field!(coupler, :test3, data, nothing, date, u"°C")

    @show coupler

    @testset "coupler_get" begin
        @test data === coupler_get(coupler, :test1, nothing, date, u"kg")
        # unit conversion
        @test data .* 1000 == coupler_get(coupler, :test1, nothing, date, u"g")
        @test data .* (5 / 18) == coupler_get(coupler, :test2, nothing, date, u"m/s")
        @test ustrip.(u"°F", data * u"°C") == coupler_get(coupler, :test3, nothing, date, u"°F")

        # key not in coupler dict
        @test_throws KeyError coupler_get(coupler, :idontexist, nothing, date, u"kg")
        # retreival at inconsistent datetime
        @test_throws ErrorException coupler_get(coupler, :test1, nothing, DateTime(2022), u"kg")
        # incompatible units
        @test_throws Unitful.DimensionError coupler_get(coupler, :test1, nothing, date, u"°C")
    end

    @testset "coupler_put!" begin
        newdata = rand(10, 10)
        newdate = DateTime(2022)
        coupler_put!(coupler, :test1, newdata, nothing, newdate, u"kg")

        @test newdata == coupler_get(coupler, :test1, nothing, newdate, u"kg")
        # coupler_put! is in-place; original data array has been modified
        @test data === coupler_get(coupler, :test1, nothing, newdate, u"kg")
        # unit conversion; data is converted from g -> kg when put in coupler
        coupler_put!(coupler, :test1, newdata, nothing, newdate, u"g")
        @test newdata ≈ 1000 * data

        # coupler_put! must be to a previously add_fielded field
        @test_throws KeyError coupler_put!(coupler, :idontexist, newdata, nothing, newdate, u"kg")
        # incoming data must match dimensions of add_fielded field
        @test_throws DimensionMismatch coupler_put!(coupler, :test1, rand(10, 5), nothing, newdate, u"kg")
        # incompatible units
        @test_throws Unitful.DimensionError coupler_put!(coupler, :test1, newdata, nothing, newdate, u"J/m")
        # coupler_put! updates coupler timestamp
        @test_throws ErrorException coupler_get(coupler, :test1, nothing, date, u"kg")
    end

    @testset "grid interpolation" begin
        # placeholder
    end
end


module ReadIR
using JSON2
using MLStyle
using CamJulia
using CamJulia.CAM

export load_aeson, aeson_to_ir

function load_aeson(path)
    open(path) do f
        let s = read(f, String)
            JSON2.read(s, NamedTuple)
        end
    end
end

function aeson_to_ir(x::NamedTuple)
    f = @match x.tag begin
        "ComLet"    => com_let
        "ComLetrec" => com_letrec
        "ComIf"     => com_if
        "ComWhile"  => com_while
        "ComMutate" => com_mut
        "ComFun"    => com_fun
        "ComApp"    => com_app
        "ComVar"    => com_var
        "ComBlock"  => com_block
        "ComTuple"  => com_tuple
        "ComProj"   => com_proj
        "ComSymbol" => com_sym
        "ComInt"    => com_int
        "ComBigInt" => com_bigint
        "ComDouble" => com_double
        "ComCh"     => com_ch
        "ComBool"   => com_bool
        "ComStr"    => com_str
        "ComNil"    => nothing
        "ComInternal" => com_internal
        a => throw(a)
    end
    f === nothing ? NilConst() : f(x.contents)
end

@active Load(x) begin
    aeson_to_ir(x)
end

list_conv(xs) = map(aeson_to_ir, xs)

com_let = @λ [name, Load(value), Load(body)] ->
          Let(name, value, body)

com_letrec = @λ [
        [[fst, Load(snd)] for (fst, snd) in bindings],
        Load(body)
    ] -> LetRec(bindings, body)

com_if = @λ [Load(cond), Load(iftrue), Load(iffasle)] ->
            If(cond, iftrue, iffasle)

com_while = @λ [Load(cond), Load(body)] ->
               While(cond, body)

com_mut = @λ [varname, Load(value)] -> Mutate(varname, value)

com_fun = @λ [args, Load(body)] -> Fun(args, body)

com_app = @λ [Load(fn), [Load(arg) for arg in args]] -> App(fn, args)

com_var = Var

com_block  = @λ [Load(expr) for expr in blocks] -> Block(blocks)

com_tuple = @λ [Load(elt) for elt in elts] -> Join(elts)

com_proj = @λ [Load(major), Load(ith)] -> Proj(major, ith)

com_bigint = BigIntConst ∘ IdrisIntegerT

com_sym = SymConst

com_int = IntConst

com_str = StrConst

com_double = DoubleConst

com_ch =  @λ s -> ChConst(s[1])

com_bool = BoolConst

com_internal = Internal

end


# test_replacement_scan.jl

function RangeReplacementScan(info)
    table_name = DuckDB.get_table_name(info)
    number = tryparse(Int64, table_name)
    if number === nothing
        return
    end
    DuckDB.set_function_name(info, "range")
    DuckDB.add_function_parameter(info, DuckDB.create_value(number))
    return
end

@testset "Test replacement scans" begin
    con = DBInterface.connect(DuckDB.DB)

    # add a replacement scan that turns any number provided as a table name into range(X)
    DuckDB.add_replacement_scan!(con, RangeReplacementScan, nothing)

    df = DataFrame(DBInterface.execute(con, "SELECT * FROM \"2\" tbl(a)"))
    @test df.a == [0, 1]

    # this still fails
    @test_throws DuckDB.QueryException DBInterface.execute(con, "SELECT * FROM nonexistant")

    DBInterface.close!(con)
end


# Copyright (c) 2019 Guido Kraemer
# Copyright (c) 2019 Simon Castano
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.

"""
    TxOut

Each output spends a certain number of satoshis, placing them under control of
anyone who can satisfy the provided pubkey script.
A `TxOut` is composed of
- `value::UInt64`, number of satoshis to spend. May be zero; the sum of all
outputs may not exceed the sum of satoshis previously spent to the outpoints
provided in the input section. (Exception: coinbase transactions spend the
block subsidy and collected transaction fees.)
- `scriptpubkey::Script` which defines the conditions which must be
satisfied to spend this output.

"""
struct TxOut
    value           :: UInt64
    scriptpubkey    :: Script
end

"""
    TxOut(io::IOBuffer)

Parse an `IOBuffer` to a `TxOut`
"""
function TxOut(io::IOBuffer)
    value = read(io, UInt64)
    scriptpubkey = Script(io)

    TxOut(value, scriptpubkey)
end

function Base.show(io::IO, output::TxOut)
    println(io, "Transaction output: ₿" * string(output.value/100000000))
    println(io, "ScriptPubKey:")
    println(io, output.scriptpubkey)
end

"""
    serialize(tx::TxOut) -> Vector{UInt8}

Returns the byte serialization of the transaction output
"""
function serialize(tx::TxOut)
    result = bytes(tx.value, len=8, little_endian=true)
    append!(result, serialize(tx.scriptpubkey))
    return result
end


using WordCloud
using Random

background = loadmask(pkgdir(WordCloud)*"/res/butterfly.png")
istrans = c->maximum(c[1:3])*(c[4]/255)<128
mask = WordCloud.imagemask(background, istrans)
showmask(background, mask, highlight=(1,0,0,0.7))
#md# `showmask` might be helpful to find a proper `istrans` function
words = [randstring(1) for i in 1:600]
weights = randexp(length(words)) .+ 1

wc = wordcloud(
    words, weights,
    mask = background,
    colors = "LimeGreen",
    angles = -30,
    density = 0.45,
    transparent = istrans,
    spacing = 1,
) |> generate!;
#md# ## average style
recolor!(wc, style=:average)
avgimg = paint(wc, background=loadmask(background, color=0.99))
#md# ## clipping style
recolor!(wc, style=:clipping)
clipimg = paint(wc, background=loadmask(background, color=0.99))
#md# ## blending style
recolor!(wc, style=:reset)
recolor!(wc, style=:blending, alpha=0.5) #blending with origin color - LimeGreen
blendimg = paint(wc, background=loadmask(background, color=0.99))
#md# ## mix style
#md# styles can also be mixed
# setcolors!(wc, :, "LimeGreen")
recolor!(wc, style=:reset)
recolor!(wc, 1:3:length(words), style=:average) #vector index is ok
recolor!(wc, 2:3:length(words), style=:clipping)
recolor!(wc, 3:3:length(words), style=:blending)
setcolors!(wc, 200:250, "black")
recolor!(wc, 200:250, style=:reset)
setcolors!(wc, 1, "black")
recolor!(wc, 1, style=:reset) #single index is ok
mixstyleimg = paint(wc, background=loadmask(background, color=0.99))
#md#
h, w = size(avgimg)
lw = 21
lc = eltype(avgimg)(parsecolor(0.1))
vbar = zeros(eltype(avgimg), (h, lw))
hbar = zeros(eltype(avgimg), (lw, 2w+lw))
vbar[:, lw÷2+1] .= lc
hbar[lw÷2+1, :] .= lc
image = [avgimg vbar clipimg; hbar; blendimg vbar mixstyleimg]
println("results are saved to recolor.png")
WordCloud.save("recolor.png", image)
wc
#eval# runexample(:recolor)
#md# ![recolor](recolor.png)


module PairwiseAlignment

export
    pairalign,
    # types
    AlignmentResult,
    # scoring systems
    AbstractSubstitutionMatrix,
    SubstitutionMatrix,
    AffineGapScoreModel,
    CostModel,
    AbstractSubstitutionMatrix,
    # alignment types
    GlobalAlignment,
    LocalAlignment,
    SemiGlobalAlignment,
    # distances
    EditDistance,
    HammingDistance,
    LevenshteinDistance,
    # gapped sequence
    GappedSequence,
    n_chars,
    n_gaps,
    leading_chars,
    leading_gaps,
    trailing_chars,
    trailing_gaps,
    counts,
    reversed_counts,
    push_chars!,
    push_gaps!,
    gapmap

# alignment types
# ---------------

# alignment
immutable GlobalAlignment end
immutable LocalAlignment end
immutable SemiGlobalAlignment end

# distances
immutable EditDistance end
immutable HammingDistance end
immutable LevenshteinDistance end

include("gapped_sequence.jl")
include("model.jl")
include("result.jl")
include("pairalign.jl")
include("algorithm/common.jl")
include("algorithm/affinegap_global_align.jl")
include("algorithm/affinegap_local_align.jl")
include("algorithm/affinegap_semiglobal_align.jl")
include("algorithm/affinegap_banded_global_align.jl")
include("algorithm/edit_distance.jl")
include("algorithm/hamming_distance.jl")

end # module


let
    cost_model = SumOfMakeSpans()
    for (i, f) in enumerate(pctapf_test_problems())
        solver = NBSSolver(
            assignment_model = TaskGraphsMILPSolver(GreedyAssignment()),
            path_planner = PIBTPlanner{NTuple{3,Float64}}()
            )
        let
            pc_tapf = f(solver;cost_function=cost_model,verbose=false)
            base_search_env = pc_tapf.env
            prob = formulate_assignment_problem(solver.assignment_model,pc_tapf)
            sched, _ = solve_assignment_problem!(solver.assignment_model,prob,pc_tapf)

            pibt_planner = route_planner(solver)
            search_env = construct_search_env(solver,deepcopy(sched),base_search_env)
            pc_mapf = PC_MAPF(search_env)
            set_iteration_limit!(pibt_planner,50)
            set_verbosity!(pibt_planner,0)
            solution, valid_flag = pibt!(pibt_planner,deepcopy(pc_mapf))
            # @show i, f, get_cost(solution)
            # @show convert_to_vertex_lists(solution.route_plan)
            @test valid_flag

            reset_solver!(solver)
            solution, cost = solve!(pibt_planner,pc_mapf)
            @test cost != typemax(cost_type(pc_mapf))
        end
    end
end
let
    solver = NBSSolver(
        assignment_model = TaskGraphsMILPSolver(GreedyAssignment()),
        path_planner = PIBTPlanner{NTuple{3,Float64}}(partial=true)
        )
    set_iteration_limit!(solver.path_planner,5)
    rpctapf = replanning_problem_1(solver)
    replan_model = MergeAndBalance()

    base_env = deepcopy(rpctapf.env)
    env = replan!(solver,replan_model,base_env,rpctapf.requests[1])
    pc_tapf = PC_TAPF(env)
    reset_solver!(solver)
    solution, cost = solve!(solver,pc_tapf)
    fix_precutoff_nodes!(solution)
    reset_solver!(solver)
    solution2, cost = solve!(solver,PC_TAPF(solution))

    for v in solution2.cache.active_set
        node = get_node_from_vtx(solution2.schedule,v)
        @test isa(node,GO)
        @test get_destination_location_id(node) == LocationID(-1)
    end

    # prob = formulate_assignment_problem(solver.assignment_model,pc_tapf)
    # sched, _ = solve_assignment_problem!(solver.assignment_model,prob,pc_tapf)
    #
    # pibt_planner = route_planner(solver)
    # pc_mapf = PC_MAPF(construct_search_env(solver,deepcopy(sched),pc_tapf.env))
    # set_iteration_limit!(pibt_planner,5)
    # set_verbosity!(pibt_planner,0)
    # solution, valid_flag = pibt!(pibt_planner,pc_mapf);
    #
    # remap_object_ids!(rpctapf.requests[2].schedule,solution.schedule)
    # env = replan!(solver,replan_model,solution,rpctapf.requests[2])
    # # update_planning_cache!(solver,env)
    # reset_solver!(solver)
    # solution2, cost = solve!(solver,PC_TAPF(env))

end

# function show_cache_state(env,cache=env.cache)
#     for (name,c) in zip(["closed","active"],[cache.closed_set, cache.active_set])
#        vtxs = map(v->string(get_node_from_vtx(env.schedule,v)), collect(c))
#        println("***",name,"***")
#        for node in vtxs
#            println(node)
#        end
#        println("\n")
#    end
#    println("*** FIXED ***")
#    for v in vertices(env.schedule)
#        path_spec = get_path_spec(env.schedule,v)
#        pre = path_spec.fixed ? " - FIXED" : ""
#        println(string(get_node_from_vtx(env.schedule,v)), pre)
#    end
# end


"""
    struct SkeletonTriangulation{Dc,Dp,B} <: Triangulation{Dc,Dp}
      plus::B
      minus::B
    end
"""
struct SkeletonTriangulation{Dc,Dp,B} <: Triangulation{Dc,Dp}
  plus::B
  minus::B
  function SkeletonTriangulation(plus::B,minus::B) where B<:Triangulation
    Dc = num_cell_dims(plus)
    Dp = num_point_dims(plus)
    #@assert Dc + 1 == Dp
    new{Dc,Dp,B}(plus,minus)
  end
end

function Base.getproperty(x::SkeletonTriangulation, sym::Symbol)
  if sym == :⁺
    x.plus
  elseif sym == :⁻
    x.minus
  else
    getfield(x, sym)
  end
end

function Base.propertynames(x::SkeletonTriangulation, private::Bool=false)
  (fieldnames(typeof(x))...,:⁺,:⁻)
end

#have_compatible_domains(a::SkeletonTriangulation,b::Triangulation) = a.plus===b || a.minus===b
#have_compatible_domains(a::Triangulation,b::SkeletonTriangulation) = have_compatible_domains(b,a)
#have_compatible_domains(a::SkeletonTriangulation,b::SkeletonTriangulation) = a===b

"""
    SkeletonTriangulation(model::DiscreteModel,face_to_mask::Vector{Bool})
    SkeletonTriangulation(model::DiscreteModel)
"""
function SkeletonTriangulation(model::DiscreteModel,face_to_mask::AbstractVector{Bool})
  left_cell_around = 1
  plus = BoundaryTriangulation(model,face_to_mask,left_cell_around)
  right_cell_around = 2
  minus = BoundaryTriangulation(model,face_to_mask,right_cell_around)
  SkeletonTriangulation(plus,minus)
end

function SkeletonTriangulation(model::DiscreteModel)
  topo = get_grid_topology(model)
  D = num_cell_dims(model)
  face_to_mask = collect(Bool, .!get_isboundary_face(topo,D-1))
  SkeletonTriangulation(model,face_to_mask)
end

function SkeletonTriangulation(
  model::DiscreteModel,
  reffe::ReferenceFE,
  face_own_dofs::Vector{Vector{Int}})

  ncells = num_cells(model)
  cell_to_reffe = Fill(1,ncells)
  SkeletonTriangulation(model,[reffe],cell_to_reffe,[face_own_dofs])
end

function SkeletonTriangulation(
  model::DiscreteModel,
  reffes::Vector{<:ReferenceFE},
  cell_to_reffe::AbstractVector,
  reffe_to_face_own_dofs::Vector{Vector{Vector{Int}}})

  reffe_to_polytope = map(get_polytope,reffes)
  reffe_to_faces = map(get_faces,reffe_to_polytope)
  reffe_to_range = map(p->get_dimrange(p,num_dims(p)-1),reffe_to_polytope)

  reffe_to_lfacet_to_mask = _compute_reffe_to_facet_to_mask(
    reffe_to_faces,reffe_to_range,reffe_to_face_own_dofs)

  topo = get_grid_topology(model)

  D = num_cell_dims(model)
  nfacets = num_facets(model)
  facet_to_isboundary = get_isboundary_face(topo,D-1)
  cell_to_facets = get_faces(topo,D,D-1)
  facet_to_mask = _compute_disc_facet_mask(
    facet_to_isboundary,cell_to_facets,cell_to_reffe,reffe_to_lfacet_to_mask)

  SkeletonTriangulation(model,facet_to_mask)

end

function _compute_disc_facet_mask(
  facet_to_isboundary, cell_to_facets::Table, cell_to_reffe, reffe_to_lfacet_to_mask)
  nfacets = length(facet_to_isboundary)
  facet_to_mask = Vector{Bool}(undef,nfacets)
  ncells = length(cell_to_reffe)
  for cell in 1:ncells
    pini = cell_to_facets.ptrs[cell]
    pend = cell_to_facets.ptrs[cell+1]-1
    reffe = cell_to_reffe[cell]
    lfacet_to_mask = reffe_to_lfacet_to_mask[reffe]
    for (lfacet,p) in enumerate(pini:pend)
      facet = cell_to_facets.data[p]
      mask = lfacet_to_mask[lfacet]
      isinterior = ! facet_to_isboundary[facet]
      facet_to_mask[facet] = mask && isinterior
    end
  end
  facet_to_mask
end

function _compute_reffe_to_facet_to_mask(
  reffe_to_faces,reffe_to_range,reffe_to_face_own_dofs)

  reffe_to_facet_to_mask = Vector{Bool}[]
  nreffes = length(reffe_to_faces)
  for reffe in 1:nreffes
    range = reffe_to_range[reffe]
    facet_to_faces = reffe_to_faces[reffe][range]
    facet_to_own_dofs = reffe_to_face_own_dofs[reffe][range]
    face_to_own_dofs = reffe_to_face_own_dofs[reffe]
    nfacets = length(facet_to_faces)
    facet_to_mask = fill(false,nfacets)
    for facet in 1:nfacets
      n = length(facet_to_own_dofs[facet])
      faces = facet_to_faces[facet]
      for face in faces
        n += length(face_to_own_dofs[face])
      end
      facet_to_mask[facet] = n == 0
    end
    push!(reffe_to_facet_to_mask,facet_to_mask)
  end

  reffe_to_facet_to_mask
end

const IN = -1
const OUT = 1

"""
"""
function InterfaceTriangulation(model::DiscreteModel,cell_to_is_in::Vector{Bool})
  cell_to_inout = fill(Int8(OUT),length(cell_to_is_in))
  cell_to_inout[cell_to_is_in] .= IN
  InterfaceTriangulation(model,cell_to_inout)
end

function InterfaceTriangulation(model::DiscreteModel,cells_in,cells_out)
  cell_to_inout = fill(Int8(OUT),num_cells(model))
  cell_to_inout[cells_in] .= IN
  cell_to_inout[cells_out] .= OUT
  InterfaceTriangulation(model,cell_to_inout)
end

function InterfaceTriangulation(model::DiscreteModel,cell_to_inout::AbstractVector{<:Integer})

  D = num_cell_dims(model)
  facet_grid = Grid(ReferenceFE{D-1},model)
  cell_grid = Grid(ReferenceFE{D},model)

  topo = get_grid_topology(model)
  facet_to_cells = Table(get_faces(topo,D-1,D))
  ifacet_to_facet, facet_to_lcell_left, facet_to_lcell_right = _find_interface_facets(
    cell_to_inout, facet_to_cells)

  ifacet_trian = RestrictedTriangulation(facet_grid,ifacet_to_facet)

  glue_left = FaceToCellGlue(topo,cell_grid,ifacet_trian,ifacet_to_facet,facet_to_lcell_left)
  glue_right = FaceToCellGlue(topo,cell_grid,ifacet_trian,ifacet_to_facet,facet_to_lcell_right)

  plus = BoundaryTriangulation(ifacet_trian,cell_grid,glue_left)
  minus = BoundaryTriangulation(ifacet_trian,cell_grid,glue_right)

  SkeletonTriangulation(plus,minus)
end

function _find_interface_facets( cell_to_inout, facet_to_cells::Table)

  nifacets = 0
  for facet in 1:length(facet_to_cells)
    a = facet_to_cells.ptrs[facet]
    b = facet_to_cells.ptrs[facet+1]
    if b-a == 2
      cell1 = facet_to_cells.data[a]
      cell2 = facet_to_cells.data[a+1]
      inout_1 = cell_to_inout[cell1]
      inout_2 = cell_to_inout[cell2]
      if (inout_1== IN && inout_2==OUT) || (inout_1== OUT && inout_2==IN)
        nifacets += 1
      end
    end
  end

  T = eltype(eltype(facet_to_cells))
  ifacet_to_facet = zeros(T,nifacets)
  nfacets = length(facet_to_cells)
  facet_to_lcell_left = fill(Int8(1),nfacets)
  facet_to_lcell_right = fill(Int8(2),nfacets)

  nifacets = 0
  for facet in 1:length(facet_to_cells)
    a = facet_to_cells.ptrs[facet]
    b = facet_to_cells.ptrs[facet+1]
    if b-a == 2
      cell1 = facet_to_cells.data[a]
      cell2 = facet_to_cells.data[a+1]
      inout_1 = cell_to_inout[cell1]
      inout_2 = cell_to_inout[cell2]
      if (inout_1== IN && inout_2==OUT) || (inout_1== OUT && inout_2==IN)
        nifacets += 1
        ifacet_to_facet[nifacets] = facet
        if inout_1 == IN
          facet_to_lcell_left[facet] = 1
          facet_to_lcell_right[facet] = 2
        else
          facet_to_lcell_left[facet] = 2
          facet_to_lcell_right[facet] = 1
        end
      end
    end
  end

   ifacet_to_facet, facet_to_lcell_left, facet_to_lcell_right
 end

# Triangulation interface

function get_cell_coordinates(trian::SkeletonTriangulation)
  get_cell_coordinates(trian.plus)
end

function get_node_coordinates(trian::SkeletonTriangulation)
  get_node_coordinates(trian.plus)
end

function get_cell_node_ids(trian::SkeletonTriangulation)
  get_cell_node_ids(trian.plus)
end

function get_reffes(trian::SkeletonTriangulation)
    get_reffes(trian.plus)
end

function get_cell_type(trian::SkeletonTriangulation)
  get_cell_type(trian.plus)
end

function get_cell_map(trian::SkeletonTriangulation)
  get_cell_map(trian.plus)
end

function get_facet_normal(trian::SkeletonTriangulation)
  plus = get_facet_normal(trian.plus)
  minus = get_facet_normal(trian.minus)
  SkeletonPair(plus,minus)
end

TriangulationStyle(::Type{<:SkeletonTriangulation}) = SubTriangulation()

function get_background_triangulation(trian::SkeletonTriangulation)
  get_background_triangulation(trian.plus)
end

function get_cell_to_bgcell(trian::SkeletonTriangulation)
  plus = get_cell_to_bgcell(trian.plus)
  minus = get_cell_to_bgcell(trian.minus)
  SkeletonPair(plus,minus)
end

function get_cell_ref_map(trian::SkeletonTriangulation)
  plus = get_cell_ref_map(trian.plus)
  minus = get_cell_ref_map(trian.minus)
  SkeletonPair(plus,minus)
end

function RestrictedTriangulation(
  oldtrian::SkeletonTriangulation,cell_to_oldcell::AbstractVector{<:Integer})
  plus = RestrictedTriangulation(oldtrian.plus,cell_to_oldcell)
  minus = RestrictedTriangulation(oldtrian.minus,cell_to_oldcell)
  SkeletonTriangulation(plus,minus)
end


#
#   Tool-Module for calculating the coupling to a beam
#
#   V: 2019-04-05
#
#   Stefan Knirck
#

# This old export doesnt make a whole lot of sense. Can the func
#export matching_ratio, unnormalized_match, antenna
export unnormalized_match

function unnormalized_match(Eout, Eantenna; dx=0.01, dy=0.01)
    return sum(Eout.*conj(Eantenna)).*dx.*dy
end


function acos(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "acos", Dfp, (Dfp,), arg0)
end

function asin(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "asin", Dfp, (Dfp,), arg0)
end

function atan(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "atan", Dfp, (Dfp,), arg0)
end

function cos(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "cos", Dfp, (Dfp,), arg0)
end

function equals(obj::Object, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function exp(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "exp", Dfp, (Dfp,), arg0)
end

function hash_code(obj::Object)
    return jcall(obj, "hashCode", jint, ())
end

function log(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "log", Dfp, (Dfp,), arg0)
end

function pow(::Type{DfpMath}, arg0::Dfp, arg1::Dfp)
    return jcall(DfpMath, "pow", Dfp, (Dfp, Dfp), arg0, arg1)
end

function pow(::Type{DfpMath}, arg0::Dfp, arg1::jint)
    return jcall(DfpMath, "pow", Dfp, (Dfp, jint), arg0, arg1)
end

function sin(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "sin", Dfp, (Dfp,), arg0)
end

function tan(::Type{DfpMath}, arg0::Dfp)
    return jcall(DfpMath, "tan", Dfp, (Dfp,), arg0)
end


# Recipes for beamforming grids, vespagrams and ARFs

using RecipesBase: RecipesBase, @recipe, @series

#=
    Beamforming grids
=#
@recipe function f(bf::BeamformGrid;
        phases=[], model="iasp91", powscale=:linear,
        s_km=false,
        dazimuth=30,
        dslowness=_default_dslowness(bf, s_km),
    )
    # Slowness grid
    sx, sy = if s_km
        s_per_km.(bf.sx), s_per_km.(bf.sy)
    else
        bf.sx, bf.sy
    end

    # Plotting defaults
    xlims --> extrema(sx)
    ylims --> extrema(sy)
    aspect_ratio --> :equal
    framestyle --> :box
    title --> "Normalised beam power ($(powscale))"
    fontfamily --> "Helvetica"

    # Power heatmap
    @series begin
        seriescolor --> :lightrainbow
        seriestype := :heatmap
        power = if powscale == :linear
            bf.power'
        elseif powscale == :log10
            log10.(bf.power')
        elseif powscale == :dB
            DSP.pow2db.(bf.power')
        else
            @error("unrecognised plot powscale '$powscale'")
        end
        sx, sy, power
    end

    # Backazimuth and slowness lines
    @series begin
        linecolor --> :white
        linewidth --> 1
        linestyle --> :dash
        label := ""
        x, y = _guidelines(sx, sy, dslowness, dazimuth)
        x, y
    end

    # 1D phase annotations, if event information is present
    phases = phases isa Union{Tuple,AbstractArray} ? phases : (phases,)
    phases_x, phases_y = Float64[], Float64[]
    phase_annotations = String[]
    if !isempty(phases)
        ! _hastraces(bf) && error("beamforming object does not contain traces, " *
                                 "which are needed to add phase arrivals")
        evt = bf.traces[1].evt
        β = Geodesics.azimuth(bf.lon, bf.lat, evt.lon, evt.lat, true)
        for phase in phases
            # FIXME: Find a better way to conditionally calculate this depening
            # on whether SeisTau is loaded
            try
                arr = Main.SeisTau.travel_time(first(bf.traces), phase, model=model)
                isempty(arr) && continue
                dtdd = first(arr).dtdd
                s_km && (dtdd = s_per_km(dtdd))
                push!(phases_x, sind(β)*dtdd)
                push!(phases_y, cosd(β)*dtdd)
                push!(phase_annotations, phase)
            catch err
                if err isa UndefVarError && err.var == :SeisTau
                    @warn("SeisTau is not loaded; no phase arrivals plotted.  " *
                          "To add phase arrivals, do `using SeisTau`, then " *
                          "try this plotting command again.")
                else
                    rethrow(err)
                end
            end
        end
    end

    # Phase location
    @series begin
        seriestype := :scatter
        markercolor --> :white
        markersize --> 4
        markershape --> :square
        markerstrokecolor --> :black
        markerstrokewidth --> 2
        label := ""
        phases_x, phases_y
    end

    # Phase label
    @series begin
        seriestype := :scatter
        primary := false
        markersize := 0
        label := ""
        # FIXME: Find a better way to set font properties
        series_annotations := Main.Plots.series_annotations(phase_annotations,
            Main.Plots.font("Helvetica", 9, :left, :bottom))
        phases_x, phases_y
    end
end

#=
    ARFs
=#
@recipe function f(arf::AbstractArrayResponse;
        powscale=:linear, s_km=false,
        dslowness=_default_dslowness(arf, s_km), # s_km ? 0.2 : 2,
        dazimuth=30
    )
    sx, sy = s_km ? (arf.sx, arf.sy) : s_per_degree.((arf.sx, arf.sy))
    xlabel, ylabel = _axis_labels(s_km)
    xlims --> extrema(sx)
    ylims --> extrema(sy)
    aspect_ratio --> :equal
    framestyle --> :box
    fontfamily --> "Helvetica"
    xguide --> xlabel
    yguide --> ylabel
    f1, f2 = round.(extrema(arf.freqs), sigdigits=2)
    freqlabel = length(arf.freqs) == 1 ? "$f1" : "$f1 - $f2"
    title --> "Array response ($freqlabel Hz)"
    colorbar_title --> "Beam power ($powscale)"
    # Power
    @series begin
        seriescolor --> :lightrainbow
        seriestype := :heatmap
        sx, sy, _scaled_power(arf, powscale, norm=true)
    end
    # Azimuth and slowness contours
    @series begin
        primary := false
        linecolor --> :white
        linewidth --> 1
        linestyle --> :dash
        x, y = _guidelines(sx, sy, dslowness, dazimuth)
        x, y
    end
end

#=
    Vespagrams
=#
@recipe function f(vespa::VespaGrid; phases=[], model="iasp91", normalise=true)
    xlims --> extrema(vespa.time)
    ylims --> extrema(vespa.slow)
    framestyle --> :box
    fontfamily --> "Helvetica"
    maxamp = maximum(abs, vespa.power)
    clims --> (normalise ? (-1, 1) : (-maxamp, maxamp))
    grid = normalise ? vespa.power' ./ maxamp : vespa.power'
    xguide --> "Time / s"
    yguide --> "Slowness / s/°"
    plot_title = (normalise ? "Normalised v" : "V") * "espagram " *
        (vespa.envelope ? "envelope" : "amplitude")
    title --> (vespa.envelope ? "Vespagram envelope" : "Vespagram amplitude")
    @series begin
        seriestype := :heatmap
        seriescolor --> (vespa.envelope ? :lightrainbow : :RdBu)
        vespa.time, vespa.slow, grid
    end
    # 1D phase annotations, if event information is present
    phases = phases isa Union{Tuple,AbstractArray} ? phases : (phases,)
    phases_x, phases_y = Float64[], Float64[]
    phase_annotations = String[]
    if !isempty(phases)
        ! _hastraces(vespa) && error("beamforming object does not contain traces, " *
                                     "which are needed to add phase arrivals")
        evt = vespa.traces[1].evt
        for phase in phases
            # FIXME: Find a better way to conditionally calculate this depening
            # on whether SeisTau is loaded
            try
                arr = Main.SeisTau.travel_time(first(vespa.traces), phase, model=model)
                isempty(arr) && continue
                dtdd = first(arr).dtdd
                tt = first(arr).time
                push!(phases_x, tt)
                push!(phases_y, dtdd)
                push!(phase_annotations, phase)
            catch err
                if err isa UndefVarError && err.var == :SeisTau
                    @warn("SeisTau is not loaded; no phase arrivals plotted.  " *
                          "To add phase arrivals, do `using SeisTau`, then " *
                          "try this plotting command again.")
                else
                    rethrow(err)
                end
            end
        end
    end

    # Phase location
    @series begin
        seriestype := :scatter
        markercolor --> :white
        markersize --> 4
        markershape --> :square
        markerstrokecolor --> :black
        markerstrokewidth --> 2
        label := ""
        phases_x, phases_y
    end

    # Phase label
    @series begin
        seriestype := :scatter
        primary := false
        markersize := 0
        label := ""
        # FIXME: Find a better way to set font properties
        series_annotations := Main.Plots.series_annotations(phase_annotations,
            Main.Plots.font("Helvetica", 9, :left, :bottom))
        phases_x, phases_y
    end
end

"""
Return a scaled beam power correctly transposed for plotting.

`powscale` can be:
- `:linear`: No transformation applied.
- `:log10`: Base-10 log of power is taken.
- `:dB`: Power converted to dB = 10log₁₀(power).

If `norm` is `true`, power is scaled such that the maximum absolute power is
unity.
"""
function _scaled_power(bf::Union{BeamformGrid,AbstractArrayResponse}, powscale; norm=true)
    scale = norm ? maximum(abs, bf.power) : one(eltype(bf.power))
    if powscale == :linear
        bf.power'./scale
    elseif powscale == :log10
        log10.(bf.power'./scale)
    elseif powscale == :dB
        DSP.pow2db.(bf.power'./scale)
    else
        throw(ArgumentError("unrecognised plot powscale '$powscale'"))
    end
end

"""
Create paths for lines of constant slowness and constant azimuth on a grid of
`sx` and `sy` slowness points.

Constant-azimuth lines are placed `dazimuth`° apart, whilst constant-slowness
lines are placed `dslowness` apart; units must match those of `sx` and `sy`.
"""
function _guidelines(sx, sy, dslowness, dazimuth)
    x = Vector{Float64}[]
    y = similar(x)
    pmax = √2*maximum(abs, [sx; sy])
    azimuths = range(0, stop=360, step=dazimuth)
    slownesses = range(0, stop=pmax, step=dslowness)
    # Constant azimuth lines
    for azimuth in azimuths
        push!(x, [0, sind(azimuth)*pmax])
        push!(y, [0, cosd(azimuth)*pmax])
    end
    # Constant slowness circles
    for slowness in slownesses
        azis = 0:360
        push!(x, sind.(azis)*slowness)
        push!(y, cosd.(azis)*slowness)
    end
    x, y
end

"Return a pair of axis labels with units s/km if `s_km` is `true`, and s/° otherwise."
function _axis_labels(s_km)
    unit = s_km ? "s/km" : "s/°"
    "East slowness ($unit)", "North slowness ($unit)"
end

"""
Pick a nice round number to give about 5 isoslowness contours for a beam power
plot.
"""
function _default_dslowness(grid::Union{AbstractArrayResponse,BeamformGrid}, s_km)
    maxslow = maximum(abs, [grid.sx; grid.sy])
    !s_km && (maxslow = s_per_degree(maxslow))
    dslow = maxslow/5
    n = floor(Int, log10(dslow))
    first_digit = floor(Int, dslow*10.0^-n)
    if first_digit == 0
        1/10.0^-n
    elseif first_digit in 1:2
        first_digit/10.0^-n
    elseif first_digit in 3:9
        5/10.0^-n
    end
end

"""
    plot(beamform_grid; kwargs...) -> ::Plot

Create a plot of a [`BeamformGrid`](@ref), as created by [`beamform`](@ref)
or [`crosscorrelation_beamform`](@ref).

# Keyword arguments:
- `phases = []`: Vector of phases to add to plot.  Before using this option,
  you must have imported `SeisTau`.
- `model = "iasp91"`: Model to use in finding predicted arrival slownesses.
- `powscale = :linear`: Scaling of beam power.  Choices are: `:linear`,
  `:log10` and `:dB`.
- `dazimuth = 30`: Spacing of constant-azimuth lines on plot in °.
- `dslowness = 2`: Spacing of constant-slowness lines on plot,
  in s/° if `s_km` is `false` (the default), or s/km if `s_km` is `true`.
- `s_km = false`: Plot grid in s/km rather than s/° if `true`.

---

    plot(vespagram_grid; phases=[], model="iasp91") -> ::Plot

Create a plot of a vespagram computed by [`vespagram`](@ref).

Optionally provide a `Vector` of phases to plot.  Before using this option,
you must have imported `SeisTau`.  Specify the model for predicted
slownesses and travel times with `model`.

---

    plot(arf; kwargs...) -> ::Plot

Create a plot of an array response function as computed by [`array_response`](@ref)
or [`crosscorrelation_array_response`](@ref).

# Keyword arguments
- `powscale = :linear`: Scaling of array response power.  Choices are:
  `:linear`, `:log10` and `:dB`.
- `dazimuth = 30`: Spacing of constant-azimuth lines on plot in °.
- `dslowness = 2`: Spacing of constant-slowness lines on plot,
  in s/° if `s_km` is `false` (the default), or s/km if `s_km` is `true`.
- `s_km = false`: Plot grid in s/km rather than s/° if `true`.

""" plot


"""
    InvertedList{I<:Unsigned, U<:Unsigned}

Basic structure which corresponds to the points found within
a Voronoi cell. The fields `idxs` contains the indices of the
points while `codes` contains quantized vector data.
"""
struct InvertedList{I<:Unsigned, U<:Unsigned}
    idxs::Vector{I}
    codes::Vector{Vector{U}}
end

Base.show(io::IO, ivlist::InvertedList{I,U}) where {I,U} = begin
    n = length(ivlist.idxs)
    print(io, "InvertedList{$I,$U}, $n vectors")
end

"""
Simple alias for `Vector{InvertedList{I<:Unsigned, U<:Unsigned}}`.
"""
const InvertedIndex{I,U} = Vector{InvertedList{I,U}}

"""
    IVFADCIndex{U<:Unsigned, I<:Unsigned, Dc<:Distances.PreMetric, Dr<:Distances.PreMetric, T<:AbstractFloat, Q<:AbstractCoarseQuantizer{Dc,T}}

The inverse file system object. It allows for approximate nearest
neighbor search into the contained vectors.

# Fields
  * `coarse_quantizer::AbstractCoarseQuantizer{Dc,T}` contains the coarse vectors
  * `residual_quantizer::QuantizedArrays.OrthogonalQuantizer{U,Dr,T,2}`
is employed to quantize vectors when adding to the index
  * `inverse_index::InvertedIndex{I,U}` is the actual inverse index employed
to perform the search.
"""
struct IVFADCIndex{U<:Unsigned,
                   I<:Unsigned,
                   Dc<:Distances.PreMetric,
                   Dr<:Distances.PreMetric,
                   T<:AbstractFloat,
                   Q<:AbstractCoarseQuantizer{Dc,T}}
    coarse_quantizer::Q
    residual_quantizer::QuantizedArrays.OrthogonalQuantizer{U,Dr,T,2}
    inverse_index::InvertedIndex{I,U}
end

"""
    length(ivfadc::IVFADCIndex)

Returns the number of vectors indexed by `ivfadc`.
"""
Base.length(ivfadc::IVFADCIndex) =
    mapreduce(ivlist->length(ivlist.idxs), +, ivfadc.inverse_index)

"""
    size(ivfadc::IVFADCIndex)

Returns a tuple with the dimensionality and number of the vectors indexed by `ivfadc`.
"""
Base.size(ivfadc::IVFADCIndex) = (size(ivfadc.coarse_quantizer, 1), length(ivfadc))
Base.size(ivfadc::IVFADCIndex, i::Int) = size(ivfadc)[i]

Base.show(io::IO, ivfadc::IVFADCIndex{U,I,Dc,Dr,T,Q}) where {U,I,Dc,Dr,T,Q} = begin
    nvars, nvectors = size(ivfadc)
    idxsize = sizeof(I)
    m = length(ivfadc.residual_quantizer.codebooks)
    compsize = sizeof(U)
    codesize = m * compsize + idxsize
    cqstr = ifelse(Q <: HNSWQuantizer, "HNSW", "naive")
    print(io, "IVFADCIndex, $cqstr coarse quantizer, $codesize-byte encoding ($idxsize + $compsize×$m), $nvectors $T vectors")
end

"""
    IVFADCIndex(data [;kwargs])

Main constructor for building an inverse file system for billion-scale ANN search.

# Arguments
  * `Matrix{T<:AbstractFloat}` input data

# Keyword arguments
  * `kc::Int=DEFAULT_COARSE_K` number of clusters (Voronoi cells) to employ
in the coarse quantization step
  * `k::Int=DEFAULT_QUANTIZATION_K` number of residual quantization levels to use
  * `m::Int=DEFAULT_QUANTIZATION_M` number of residual quantizers to use
  * `coarse_quantizer::Symbol=DEFAULT_COARSE_QUANTIZER` coarse quantizer
  * `coarse_distance=DEFAULT_COARSE_DISTANCE` coarse quantization distance
  * `quantization_distance=DEFAULT_QUANTIZATION_DISTANCE` residual quantization distance
  * `quantization_method=DEFAULT_QUANTIZATION_METHOD` residual quantization method
  * `coarse_maxiter=DEFAULT_COARSE_MAXITER` number of clustering iterations for obtaining
the coarse vectors
  * `quantization_maxiter=DEFAULT_QUANTIZATION_MAXITER` number of clustering iterations for
residual quantization
  * `index_type=UInt32` type for the indexes of the vectors in the inverted list
"""
function IVFADCIndex(data::Matrix{T};
                     kc::Int=DEFAULT_COARSE_K,
                     k::Int=DEFAULT_QUANTIZATION_K,
                     m::Int=DEFAULT_QUANTIZATION_M,
                     coarse_quantizer::Symbol=DEFAULT_COARSE_QUANTIZER,
                     coarse_distance::Distances.PreMetric=DEFAULT_COARSE_DISTANCE,
                     quantization_distance::Distances.PreMetric=DEFAULT_QUANTIZATION_DISTANCE,
                     quantization_method::Symbol=DEFAULT_QUANTIZATION_METHOD,
                     coarse_maxiter::Int=DEFAULT_COARSE_MAXITER,
                     quantization_maxiter::Int=DEFAULT_QUANTIZATION_MAXITER,
                     index_type::Type{I}=UInt32
                    ) where {I<:Unsigned, T<:AbstractFloat}
    # Checks
    nrows, nvectors = size(data)
    bits_required = ceil(Int, log2(nvectors))
    @assert kc >= 2 "Number of coarse clusters has to be >= 2"
    @assert k <= nvectors "Number of quantization levels  has to be <= $nvectors"
    @assert m >= 1 && m <= nrows "Number of codebooks has to be between 1 and $nrows"
    @assert coarse_quantizer in [:naive, :hnsw] "Coarse quantizer can be :naive or :hnsw only"
    @assert coarse_maxiter > 0 "Number of clustering iterations has to be > 0"
    @assert quantization_maxiter > 0 "Number of clustering iterations has to be > 0"
    @assert QuantizedArrays.TYPE_TO_BITS[index_type] >=
        bits_required "$nvectors vectors require at least $bits_required index bits"

    # Run kmeans, build coarse quantizer
    @debug "Clustering..."
    cmodel = kmeans(data,
                    kc,
                    maxiter=coarse_maxiter,
                    distance=coarse_distance,
                    init=:kmpp,
                    display=:none)

    # Calculate residuals
    @debug "Residual calculation..."
    residuals = _build_residuals(cmodel, data)

    # Build residuals quantizer
    @debug "Building quantizer..."
    rq = build_quantizer(residuals,
                         k=k,
                         m=m,
                         method=quantization_method,
                         distance=quantization_distance,
                         maxiter=quantization_maxiter)

    # Quantize residuals for each cluster
    @debug "Building inverted index..."
    ii = _build_inverted_index(rq, cmodel, residuals, index_type=index_type)

    # Return
    @debug "Building coarse quantizer..."
    if coarse_quantizer == :naive
        cq = NaiveQuantizer{typeof(coarse_distance),
                eltype(cmodel.centers)}(cmodel.centers)
    else
        hnsw = HierarchicalNSW([cmodel.centers[:,i] for i in 1:kc],
                                metric=coarse_distance)
        add_to_graph!(hnsw)
        cq = HNSWQuantizer(hnsw)
    end
    return IVFADCIndex(cq, rq, ii)
end

function _build_residuals(km::KmeansResult, data::Matrix{T}) where {T}
    residuals = similar(data)
    for cluster in 1:nclusters(km)
        bidxs = km.assignments .== cluster
        @. residuals[:, bidxs] = data[:, bidxs] - km.centers[:, cluster]
    end
    return residuals
end

function _build_inverted_index(rq::QuantizedArrays.OrthogonalQuantizer{U,D,T,2},
                               km::KmeansResult,
                               data::Matrix{T};
                               index_type::Type{I}=UInt32
                              ) where {U,I<:Unsigned,D,T}
    n = nclusters(km)
    invindex = InvertedIndex{I,U}(undef, n)
    for cluster in 1:n
        idxs = findall(isequal(cluster), km.assignments)
        qdata = QuantizedArrays.quantize_data(rq, data[:, idxs])
        ivlist = InvertedList{I,U}(
                    idxs .- one(I),
                    [qdata[:, j] for j in 1:length(idxs)])
        invindex[cluster] = ivlist
    end
    return invindex
end

"""
    knn_search(ivfadc, point, k[; w=1])

Searches at most `k` closest neighbors of `point` in the index `ivfadc`;
the neighbors will be searched for in the points contained in the closest
`w` clusters.
"""
function knn_search(ivfadc::IVFADCIndex{U,I,Dc,Dr,T},
                    point::Vector{T},
                    k::Int;
                    w::Int=1
                   ) where {U,I,Dc,Dr,T}
    # Checks and initializations
    @assert k >= 1 "Number of neighbors must be k >= 1"
    @assert w >= 1 "Number of clusters to search in must be w >= 1"
    cq = ivfadc.coarse_quantizer
    nclusters = size(cq, 2)
    rq_cbooks = ivfadc.residual_quantizer.codebooks
    m, n = length(rq_cbooks), length(point)
    w = min(w, nclusters)

    # Find the 'w' closest coarse vectors and calculate residuals
    closest_clusters, coarse_distances = coarse_search(cq, point, w)
    residuals = _closest_cluster_residuals(cq, point, closest_clusters)

    # Calculate distances between point and the vectors
    # from the clusters using the residual distances.
    distances = Vector{T}()
    neighbors = SortedMultiDict{T,I}()
    maxdist = zero(T)
    difftables = Vector{LittleDict{U,T}}(undef, m)
    @inbounds for j = eachindex(closest_clusters)
        dc = coarse_distances[j]
        # Calculate all residual distances
        # (between the vector and the codebooks of the residual quantizer)
        for i = eachindex(difftables)
            rr = QuantizedArrays.rowrange(n, m, i)
            diffs = colwise(Dc(), rq_cbooks[i].vectors, residuals[rr, j])
            difftables[i] = LittleDict{U,T}(rq_cbooks[i].codes, diffs)
        end
        # Loop through the iverted list and calculate
        # the actual (quantized) distances through lookup;
        # use SortedMultiDict as a maxheap for `k` neighbors
        ivlist = ivfadc.inverse_index[closest_clusters[j]]
        for i = eachindex(ivlist.idxs)
            d = dc
            @simd for ii = eachindex(ivlist.codes[i])
                # add difference vector for the ii'th code of the i'th subquantizer
                d += difftables[ii][ivlist.codes[i][ii]]
            end
            if length(neighbors) < k
                push!(neighbors, d => ivlist.idxs[i])
                maxdist, _ = last(neighbors)
            elseif maxdist > d
                delete!((neighbors, lastindex(neighbors)))  # delete the max key
                push!(neighbors, d => ivlist.idxs[i])
                maxdist, _ = last(neighbors)
            end
        end
    end
    return collect(values(neighbors)), collect(keys(neighbors))
end

function knn_search(ivfadc::IVFADCIndex{U,I,Dc,Dr,T},
                    points::Vector{Vector{T}},
                    k::Int;
                    w::Int=1
                   ) where {U,I,Dc,Dr,T}
    n = length(points)
    idxs = Vector{Vector{I}}(undef, n)
    dists = Vector{Vector{T}}(undef, n)
    for i = eachindex(points)  # Threads.@threads for i = ... for multi-threading
        idxs[i], dists[i] = knn_search(ivfadc, points[i], k, w=w)
    end
    return idxs, dists
end


using Luxor

include("environments/collect_points.jl")
include("tools/play.jl")

function kernel_test_initialize(environments, env_seed)

    tx = threadIdx().x
    bx = blockIdx().x

    input = @cuDynamicSharedMem(Float32, environments.number_inputs)
    offset = sizeof(input)

    sync_threads()

    initialize(tx, bx, input, environments, offset, env_seed)

    return
end

function main()

    config_environment = OrderedDict()
    config_environment["maze_columns"] = 5
    config_environment["maze_rows"] = 5
    config_environment["maze_cell_size"] = 80
    config_environment["agent_radius"] = 12
    config_environment["point_radius"] = 8
    config_environment["agent_movement_range"] = 10.0
    config_environment["use_sensors"] = true
    config_environment["reward_per_collected_positive_point"] = 500.0
    config_environment["reward_per_collected_negative_point"] = -700.0
    config_environment["number_time_steps"] = 1000

    env_seed = rand(1:1000)
    number_individuals = 100

    environments = CollectPoints(config_environment, number_individuals)

    shared_memory = get_memory_requirements(environments) + sizeof(Float32) * environments.number_inputs

    CUDA.@cuda threads = 10 blocks = number_individuals shmem = shared_memory kernel_test_initialize(environments, env_seed)

    CUDA.synchronize()

    println(environments.agents_positions[:,1])

    width = environments.maze_cell_size * environments.maze_columns
    height = environments.maze_cell_size * environments.maze_rows
    number_iterations = 100

    @play width height number_iterations begin
        render(environments)
    end

end

main()


"""
Constraint handler solving the semi-subproblem (fixed y, binary z).
"""
mutable struct SemiSubHdlr <: SCIP.AbstractConstraintHandler
    scip::SCIP.Optimizer
    solver
    finaltime::Float64 # seconds
    counter::Int
    primal_bound::Float64
    best_solution::Union{CandSol, Nothing}

    function SemiSubHdlr(scip, solver, finaltime)
        return new(scip, solver, finaltime, 0, Inf, nothing)
    end
end

"""
Constraint data, referencing the topology and binary edges variables.
"""
mutable struct SemiSubCons <: SCIP.AbstractConstraint{SemiSubHdlr}
    inst::Instance
    topo::Topology
    y::Array{JuMP.VariableRef}
    q::Array{JuMP.VariableRef}
end

function solve_semi_sub(ch::SemiSubHdlr, cons::SemiSubCons,
                        sol::Ptr{SCIP.SCIP_SOL}=C_NULL)
    ch.counter += 1

    # preparation
    inst = cons.inst
    topo = cons.topo
    ndiams = length(inst.diameters)
    narcs = length(topo.arcs)

    # fetch solution values
    ysol = SCIP.sol_values(ch.scip, cons.y, sol)
    qsol = SCIP.sol_values(ch.scip, cons.q, sol)
    @debug "  cand. sol:$(findall(!iszero, qsol))"

    # complete solution candidate
    zcand = fill(false, narcs, ndiams)
    zcand[ysol .> 0.5, 1] .= true  # pretend to use smallest diameter
    cand = CandSol(zcand, qsol, qsol.^2)

    # solve subproblem (from scratch)
    submodel, candarcs, z = make_semisub(inst, topo, cand, ch.solver)
    # TODO: solver.writemodels && writeLP(submodel, "sub_$(counter).lp", genericnames=false)
    settimelimit!(submodel, ch.solver, ch.finaltime - time())
    JuMP.optimize!(submodel)
    substatus = JuMP.termination_status(submodel)
    if substatus == MOI.OPTIMAL
        # have found improving solution?
        newobj = JuMP.objective_value(submodel)
        if newobj < ch.primal_bound
            ch.primal_bound = newobj
            znew = fill(false, narcs, ndiams)
            znew[candarcs, :] = JuMP.value.(z) .> 0.5
            ch.best_solution = CandSol(znew, qsol, qsol.^2)
            @debug "  new sol: $(ch.primal_bound)"
            return true # found improvement
        end
    elseif substatus != MOI.INFEASIBLE
        error("Unexpected status: $(:substatus)")
    end

    # TODO: what can we do about dual/primal stopping criterion?
    # TODO: add cut on objective var?

    return false # did not find improvement
end

function SCIP.check(ch::SemiSubHdlr,
                    constraints::Array{Ptr{SCIP.SCIP_CONS}},
                    sol::Ptr{SCIP.SCIP_SOL},
                    checkintegrality::Bool,
                    checklprows::Bool,
                    printreason::Bool,
                    completely::Bool)
    @assert length(constraints) == 1
    cons::SemiSubCons = SCIP.user_constraint(constraints[1])

    # Always solve semi-subproblem (to find improving solutions)
    # TODO: cache candidates to avoid re-solving?
    improvement = solve_semi_sub(ch, cons, sol)

    # Never accept the solution candidate.
    return SCIP.SCIP_INFEASIBLE
end

function enforce_semi_sub(ch::SemiSubHdlr, cons::SemiSubCons)
    # Always solve semi-subproblem (to find improving solutions)
    # TODO: cache candidates to avoid re-solving?
    improvement = solve_semi_sub(ch, cons)

    # Always add no-good to cut off this candidate.
    ysol = SCIP.sol_values(ch.scip, cons.y)
    active = ysol .> 0.5
    coefs = 2.0 * active .- 1.0
    ci = @cb_constraint(ch.scip, coefs ⋅ cons.y ≤ sum(active) - 1)

    return SCIP.SCIP_CONSADDED
end

function SCIP.enforce_lp_sol(ch::SemiSubHdlr,
                             constraints::Array{Ptr{SCIP.SCIP_CONS}},
                             nusefulconss::Cint,
                             solinfeasible::Bool)
    @assert length(constraints) == 1
    cons::SemiSubCons = SCIP.user_constraint(constraints[1])
    return enforce_semi_sub(ch, cons)
end

function SCIP.enforce_pseudo_sol(ch::SemiSubHdlr,
                                 constraints::Array{Ptr{SCIP.SCIP_CONS}},
                                 nusefulconss::Cint,
                                 solinfeasible::Bool,
                                 objinfeasible::Bool)
    @assert length(constraints) == 1
    cons::SemiSubCons = SCIP.user_constraint(constraints[1])
    return enforce_semi_sub(ch, cons)
end

function SCIP.lock(ch::SemiSubHdlr,
                   constraint::Ptr{SCIP.SCIP_CONS},
                   locktype::SCIP.SCIP_LOCKTYPE,
                   nlockspos::Cint,
                   nlocksneg::Cint)
    cons::SemiSubCons = SCIP.user_constraint(constraint)
    for y in cons.y
        # TODO: understand why lock is called during SCIPfree, after the
        # constraint should have been deleted already. Does it mean we should
        # implement CONSTRANS?
        var_::Ptr{SCIP.SCIP_VAR} = SCIP.var(ch.scip, JuMP.index(y))
        var_ != C_NULL || continue  # avoid segfault!

        SCIP.@SC SCIP.SCIPaddVarLocksType(
            ch.scip, var_, locktype, nlockspos + nlocksneg, nlockspos + nlocksneg)
    end
end