text:
using AbstractTrees

# Language
#

mutable struct Language
    name::Symbol
    ptr::Ptr{API.TSLanguage}
    Language(name::Symbol) = new(name, API.lang_ptr(name))
end
Base.show(io::IO, l::Language) = print(io, "Language(", repr(l.name), ")")

#
# Parser
#

mutable struct Parser
    language::Language
    ptr::Ptr{API.TSParser}
    function Parser(lang::Language)
        parser = new(lang, API.ts_parser_new())
        finalizer(p -> API.ts_parser_delete(p.ptr), parser)
        set_language!(parser, lang)
        return parser
    end
    Parser(name::Symbol) = Parser(Language(name))
end
Base.show(io::IO, p::Parser) = print(io, "Parser(", p.language, ")")

function set_language!(parser::Parser, language::Language)
    API.ts_parser_set_language(parser.ptr, language.ptr)
    parser.language = language
    return parser
end

Base.parse(p::Parser, text::AbstractString) = Tree(API.ts_parser_parse_string(p.ptr, C_NULL, text, sizeof(text)))

#
# Tree
#

mutable struct Tree
    ptr::Ptr{API.TSTree}
    function Tree(ptr::Ptr{API.TSTree})
        tree = new(ptr)
        finalizer(t -> API.ts_tree_delete(t.ptr), tree)
        return tree
    end
end
Base.show(io::IO, t::Tree) = show(io, AbstractTrees.print_tree(root(t)))

root(t::Tree) = Node(API.ts_tree_root_node(t.ptr))

traverse(f, tree::Tree, iter=children) = traverse(f, root(tree), iter)

#
# Node
#

struct Node{T}
    ptr::API.TSNode
    Node(ptr::API.TSNode) = new{Symbol(node_type(ptr))}(ptr)
end
Base.show(io::IO, n::Node{T}) where T = print(io, "$(typeof(n))")

node_string(n::Node{T}) where T = unsafe_string(API.ts_node_string(n.ptr))
node_symbol(n::Node{T}) where T = API.ts_node_symbol(n.ptr)
node_type(n::Node{T}) where T = unsafe_string(API.ts_node_type(n.ptr))
node_type(n::API.TSNode) = unsafe_string(API.ts_node_type(n))

is_null(n::Node{T}) where T = API.ts_node_is_null(n.ptr)
is_named(n::Node{T}) where T = API.ts_node_is_named(n.ptr)
is_missing(n::Node{T}) where T = API.ts_node_is_missing(n.ptr)
is_extra(n::Node{T}) where T = API.ts_node_is_extra(n.ptr)
is_leaf(n::Node{T}) where T = iszero(count_nodes(n))

count_nodes(n::Node{T}) where T = Int(API.ts_node_child_count(n.ptr))
count_named_nodes(n::Node{T}) where T = Int(API.ts_node_named_child_count(n.ptr))

child(n::Node{T}, nth::Integer) where T = Node(API.ts_node_child(n.ptr, nth-1))
named_child(n::Node{T}, nth::Integer) where T = Node(API.ts_node_named_child(n.ptr, nth-1))

children(n::Node{T}) where T = (child(n, ind) for ind = 1:count_nodes(n))
AbstractTrees.children(n::Node{T}) where T = collect(children(n))
named_children(n::Node{T}) where T = (named_child(n, ind) for ind = 1:count_named_nodes(n))

function traverse(f, n::Node{T}, iter=children) where T
    f(n, true)
    for child in iter(n)
        traverse(f, child, iter)
    end
    f(n, false)
    return nothing
end

Base.:(==)(left::Node{T}, right::Node{T}) where T = API.ts_node_eq(left.ptr, right.ptr)

byte_range(n::Node{T}) where T = (Int(API.ts_node_start_byte(n.ptr)) + 1, Int(API.ts_node_end_byte(n.ptr)))

slice(src::AbstractString, n::Node{T}) where T = slice(src, byte_range(n))
slice(src::AbstractString, (from, to)) = SubString(src, from, thisind(src, to))

child(n::Node{T}, name::AbstractString) where T = Node(API.ts_node_child_by_field_name(n.ptr, String(name), sizeof(name)))

#
# Query
#

struct QueryException <: Exception
    msg::String
end

mutable struct Query
    language::Language
    ptr::Ptr{API.TSQuery}

    function Query(language::Language, source)
        source_text = load_source(language, source)
        error_offset_ref = Ref{UInt32}()
        error_type_ref = Ref{API.TSQueryError}()
        ptr = API.ts_query_new(language.ptr, String(source_text), sizeof(source_text), error_offset_ref, error_type_ref)
        if ptr === C_NULL
            type =
                error_type_ref[] === API.TSQueryErrorSyntax ? "syntax" :
                error_type_ref[] === API.TSQueryErrorNodeType ? "node type" :
                error_type_ref[] === API.TSQueryErrorField ? "field" :
                error_type_ref[] === API.TSQueryErrorCapture ? "capture" :
                error("unknown query error type: $(error_type_ref[])")
            offset = error_offset_ref[] + 1
            throw(QueryException("'$type' error starting at index $offset"))
        else
            query = new(language, ptr)
            finalizer(q -> API.ts_query_delete(q.ptr), query)
            return query
        end
    end
    Query(language::Symbol, source) = Query(Language(language), source)

    function load_source(lang::Language, files)
        queries = API.lang_queries(lang.name)
        out = IOBuffer()
        for file in files
            println(out, get(queries, file, ""))
        end
        return String(take!(out))
    end
    load_source(::Language, source::AbstractString) = source
end
Base.show(io::IO, q::Query) = print(io, "Query(", q.language, ")")

macro query_cmd(body, language=error("no language provided in query"))
    return :(Query($(esc(Meta.quot(Symbol(language)))), $(esc(body))))
end

pattern_count(q::Query) = Int(API.ts_query_pattern_count(q.ptr))
capture_count(q::Query) = Int(API.ts_query_capture_count(q.ptr))
string_count(q::Query) = Int(API.ts_query_string_count(q.ptr))

mutable struct QueryCursor
    ptr::Ptr{API.TSQueryCursor}

    function QueryCursor()
        ptr = API.ts_query_cursor_new()
        cursor = new(ptr)
        finalizer(c -> API.ts_query_cursor_delete(c.ptr), cursor)
        return cursor
    end
end
Base.show(io::IO, ::QueryCursor) = print(io, "QueryCursor()")

exec(c::QueryCursor, q::Query, n::Node) = (API.ts_query_cursor_exec(c.ptr, q.ptr, n.ptr); c)
exec(c::QueryCursor, q::Query, t::Tree) = (exec(c, q, root(t)); c)

Base.eachmatch(query::Query, tree::Tree) = exec(QueryCursor(), query, tree)

function Base.iterate(cursor::QueryCursor, state=nothing)
    result = next_match(cursor)
    return result === nothing ? nothing : (result, nothing)
end

mutable struct QueryMatch
    obj::API.TSQueryMatch
end

capture_count(qm::QueryMatch) = Int(qm.obj.capture_count)

function next_match(cursor::QueryCursor)
    match_ref = Ref{API.TSQueryMatch}()
    success = API.ts_query_cursor_next_match(cursor.ptr, match_ref)
    return success ? QueryMatch(match_ref[]) : nothing
end

struct QueryCapture
    node::Node
    id::UInt32
end

function captures(qm::QueryMatch)
    fn = function (ith)
        ptr = unsafe_load(qm.obj.captures, ith)
        node = Node(ptr.node)
        return QueryCapture(node, ptr.index)
    end
    return (fn(i) for i = 1:capture_count(qm))
end

function capture_name(q::Query, qc::QueryCapture)
    unsafe_string(API.ts_query_capture_name_for_id(q.ptr, qc.id, Ref{UInt32}()))
end

query_string(q, id) = unsafe_string(API.ts_query_string_value_for_id(q.ptr, id, Ref{UInt32}()))

function predicate(q::Query, m::QueryMatch, source::AbstractString)
    len = Ref{UInt32}()
    ptr = API.ts_query_predicates_for_pattern(q.ptr, m.obj.pattern_index, len)
    if len[] > 0
        func, args = "", []
        for i in 1:len[]
            step = unsafe_load(ptr, i)
            if step.type === API.TSQueryPredicateStepTypeCapture
                # Iterate over the captures rather than allocating a dict.
                for address in 1:capture_count(m)
                    capture = unsafe_load(m.obj.captures, address)
                    if capture.index == step.value_id
                        node = Node(capture.node)
                        str = slice(source, node)
                        push!(args, str)
                        break
                    end
                end
            elseif step.type === API.TSQueryPredicateStepTypeString
                # Either a literal argument name, or predicate name.
                str = query_string(q, step.value_id)
                func == "" ? (func = str) : push!(args, str)
            elseif step.type === API.TSQueryPredicateStepTypeDone
                # This marks the end of an individual predicate. Check whether
                # we actually have enough arguments to call the predicate and
                # then continue.
                if length(args) != 2
                    @warn "incorrect number of arguments to '$func', expected 2"
                    return false
                end
                # Handle the following predicates. Source: tree-sitter rust lib.
                #
                #   - eq?
                #   - not-eq?
                #   - is?
                #   - is-not?
                #   - match?
                #   - set!
                #
                result = if func == "eq?"
                    left, right = args
                    left == right
                elseif func == "not-eq?"
                    left, right = args
                    left != right
                elseif func == "is?"
                    @warn "'$func' not implemented" # TODO
                    false
                elseif func == "is-not?"
                    @warn "'$func' not implemented" # TODO
                    false
                elseif func == "match?"
                    arg_1, arg_2 = args
                    occursin(Regex(arg_2), arg_1)
                elseif func == "set!"
                    @warn "'$func' not implemented" # TODO
                    false
                else
                    # Invalid predicate functions will fail with a warning:
                    @warn "unknown predicate function '$func'"
                    false
                end
                if result
                    # Success, reset for the next predicate.
                    func = ""
                    empty!(args)
                else
                    # Failed to match the current predicate, so we can bale here.
                    return false
                end
            else
                # We really shouldn't get here if all is well.
                error("unreachable reached")
            end
        end
    end
    return true
end

function each_capture(tree::Tree, query::Query, source::AbstractString)
    return (c for m in eachmatch(query, tree) for c in captures(m) if predicate(query, m, source))
end

function tokens(parser::Parser, query::Query, source::AbstractString)
    out = Tuple{String,String}[]
    tree = parse(parser, source)
    for m in eachmatch(query, tree)
        if predicate(query, m, source)
            for c in captures(m)
                id = capture_name(query, c)
                text = slice(source, c.node)
                push!(out, (text, id))
            end
        end
    end
    return out
end


text:
# Removed in CUDA 11.0

@checked function cublasHgemm(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb,
                              beta, C, ldc)
    initialize_context()
    ccall((:cublasDgemm_v2, libcublas()), cublasStatus_t,
                   (cublasHandle_t, cublasOperation_t, cublasOperation_t, Cint, Cint,
                    Cint, RefOrCuRef{Float16}, CuPtr{Float16}, Cint, CuPtr{Float16}, Cint,
                    RefOrCuRef{Float16}, CuPtr{Float16}, Cint),
                   handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc)
end

ext: jl


text:
# MethodError (no method matching foo(::String))

function foo(x::Int64)
    42
end

function foo(x::Bool)
    1729
end

function cant_know_which_one_at_this_point(x)
    foo(x)
end

cant_know_which_one_at_this_point("oups")

max_stars_repo_name: gabriel-doriath-dohler/pjulia

max_issues_repo_name: gabriel-doriath-dohler/pjulia

max_forks_repo_name: gabriel-doriath-dohler/pjulia


text:
"""
`BufferedOutputStream{T}` provides buffered writing to a sink of type `T`.

Any type `T` wrapped in a `BufferedOutputStream` must implement:

    BufferedStreams.writebytes(sink::T, buffer::Vector{UInt8}, n::Int, eof::Bool)

This function should:

* write `n` bytes from `buffer`, starting at the first position
* `eof` is `true` when there will be no more write operations to `sink`, to
  facilitate flushing or closing the sink, etc.
* `writebytes` should return the number of bytes written. This must be
  `n` or 0. A return value of 0 indicates data should not be evicted from
  the buffer.

The buffer passed to this function never reallocated by `BufferedOutputStream`,
so it's safe to retain a reference to it to, for example, report some bytes as
written but do so lazily or asynchronously.
"""
mutable struct BufferedOutputStream{T} <: IO
    sink::T
    buffer::Vector{UInt8}

    # Position of the next unused byte in buffer;
    # `position ≤ 0` indicates that the stream is closed.
    position::Int
end

function BufferedOutputStream(sink::T, bufsize::Integer=default_buffer_size) where T
    if bufsize ≤ 0
        throw(ArgumentError("buffer size must be positive"))
    end
    return BufferedOutputStream{T}(sink, Vector{UInt8}(undef, bufsize), 1)
end

function Base.show(io::IO, stream::BufferedOutputStream{T}) where T
    bufsize = length(stream.buffer)
    filled = stream.position
    if isopen(stream)
        print(io,
            summary(stream), "(<",
            _datasize(bufsize), " buffer, ",
            round(Int, filled / bufsize * 100), "% filled>)")
    else
        print(io, summary(stream), "(<closed>)")
    end
end

"""
Flush all accumulated data from the buffer.
"""
function flushbuffer!(stream::BufferedOutputStream, eof::Bool = false)
    buffered = stream.position - 1
    written = writebytes(stream.sink, stream.buffer, buffered, eof)
    if written != buffered
        error("BufferedOutputStream sink failed to write all data")
    end
    stream.position = 1
    return
end

function checkopen(stream::BufferedOutputStream)
    if !isopen(stream)
        error("buffered output stream is already closed")
    end
end

@inline function Base.write(stream::BufferedOutputStream, b::UInt8)
    checkopen(stream)
    position = stream.position
    buffer = stream.buffer
    if position > length(buffer)
        flushbuffer!(stream)
        position = stream.position
        buffer = stream.buffer
        if position > length(buffer)
            throw(EOFError)
        end
    end
    @inbounds buffer[position] = b
    stream.position = position + 1
    return 1
end

function Base.write(stream::BufferedOutputStream, data::Vector{UInt8})
    checkopen(stream)
    # TODO: find a way to write large vectors directly to the sink bypassing the buffer
    #append!(stream, data, 1, length(data))
    n_avail = lastindex(stream.buffer) - stream.position + 1
    n = min(length(data), n_avail)
    copyto!(stream.buffer, stream.position, data, 1, n)
    stream.position += n
    written = n
    while written < length(data)
        flushbuffer!(stream)
        n_avail = lastindex(stream.buffer) - stream.position + 1
        @assert n_avail > 0
        n = min(lastindex(data) - written, n_avail)
        copyto!(stream.buffer, stream.position, data, written + 1, n)
        stream.position += n
        written += n
    end
    return written
end

# TODO: This is too slow. I think this pointer/pointer_to_array trick may
# allocate, so we should try to avoid it everywhere, but especially here.
"""
Write part of a byte array.
"""
function Base.append!(stream::BufferedOutputStream, data::Vector{UInt8},
                      start::Int, stop::Int)
    buffer = stream.buffer
    position = stream.position
    writelen = stop - start + 1
    buflen = length(buffer)
    while true
        if position > buflen
            stream.position = position
            flushbuffer!(stream)
            position = stream.position
            buffer = stream.buffer
            buflen = length(buffer)
        end

        num_chunk_bytes = min(stop - start + 1, buflen - position + 1)
        copyto!(buffer, position, data, start, num_chunk_bytes)
        start += num_chunk_bytes
        position += num_chunk_bytes
        if start > stop
            break
        end
    end
    stream.position = position
    return writelen
end

function Base.flush(stream::BufferedOutputStream)
    flushbuffer!(stream)
    if applicable(flush, stream.sink)
        flush(stream.sink)
    end
    return
end

function Base.isopen(stream::BufferedOutputStream)
    return stream.position > 0
end

function Base.close(stream::BufferedOutputStream)
    if !isopen(stream)
        return
    end
    flush(stream)
    if applicable(close, stream.sink)
        close(stream.sink)
    end
    empty!(stream.buffer)
    stream.position = 0
    return
end

function Base.eof(stream::BufferedOutputStream)
    return true
end

function Base.pointer(stream::BufferedOutputStream, index::Integer = 1)
    return pointer(stream.buffer, stream.position + index - 1)
end

function available_bytes(stream::BufferedOutputStream)
    return max(lastindex(stream.buffer) - stream.position + 1, 0)
end


text:
### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# ╔═╡ 4655bdc2-1ed7-11ec-022d-7df0b8fdb907
using Pkg; Pkg.develop(path=".."); Pkg.activate("..")

# ╔═╡ 661cc8fe-44f3-422a-91f5-96956b748847
begin
	using SqState
	using Plots
	using QuantumStateBase
	using Flux
	using LinearAlgebra
end

# ╔═╡ 510b566f-4fd9-4d96-bfc5-b54385b713ce
using Fetch, DataDeps, MAT

# ╔═╡ c55d7af1-0cd3-4b0a-bbd5-0048f15ea851
# m = SqState.get_model("model_q2args_sqth_x")

# ╔═╡ 8f82e42f-8c2c-4bd0-ab20-95b0930670dc
# args = (1.72,  4.17646, 0.4);
args = SqState.rand_arg_sqth((0.23, 0.231), (0, 2π), (0, 0.3))
# args = (1., π/3, 0.4, 0.07, 0.8, 0.2)

# ╔═╡ 9588e66c-9c19-4257-8fd2-12455e0c0f0d
relu(x) = x<0 ? zero(typeof(x)) : x

# ╔═╡ f0cfd6f3-1e55-46cf-bd58-937e52cb7daa
begin
	dim = 100
	state = SqState.construct_state_sqth(args..., 1000)
	d = Float32.(rand(state, 4096, IsGaussian))
	argŝ = relu.(m(reshape(d, 2, :, 1)))
	statê = SqState.construct_state_sqth(argŝ..., 1000)
	ρ = statê.𝛒
	argŝ
end

# ╔═╡ 57e75729-4bf4-4e0c-9ef2-bdaf6fb4e22b
function fidelity(ρ1, ρ2)
	sqrt_ρ1 = sqrt(ρ1)
	return tr(sqrt(sqrt_ρ1*ρ2*sqrt_ρ1))^2
end

# ╔═╡ 6739ac26-7590-46a8-a75a-2160a1d7c11d
fidelity(ρ, state.𝛒)

# ╔═╡ 00c09086-1eef-4b3f-b764-fe7201e1a354
scatter(d[1, :], d[2, :])

# ╔═╡ be08410c-4fce-405c-9b27-ababf152402e
begin
	d̂ = rand(statê, 4096, IsGaussian)
	scatter(d̂[1, :], d̂[2, :])
end

# ╔═╡ 32c40e60-0989-44f6-b518-9f6203aa6209
function calc_f(; r=(0, 1.8), θ=(0, 2π), n̄=(0, 0.3), dim=700)
	args = SqState.rand_arg_sqth(r, θ, n̄)
	state = SqState.construct_state_sqth(args..., dim)
	argŝ = relu.(m(reshape(Float32.(rand(state, 4096, IsGaussian)), 2, :, 1)))
	statê = SqState.construct_state_sqth(argŝ..., dim)

	fidelity(statê.𝛒, state.𝛒)
end

# ╔═╡ 53da20d3-2d64-4cab-ae4d-0615f750df67
# 5db -> 0.575
# 10db -> 1.151
# 15db -> 1.727

# ╔═╡ bc685a2b-7807-4dd9-b6cd-4a9a1448c09b
sum(calc_f(r=(0., 0.575), dim=100) for _ in 1:100) / 100

# ╔═╡ a4f54b38-ecaa-42b2-bd57-16e9dc52aea0
sum(calc_f(r=(1.15, 1.151), dim=500) for _ in 1:50) / 50

# ╔═╡ 8083f0ef-ad7f-49fe-a248-5931c56ce998
begin
	mse_x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
	mse_y = [0.997, 0.9922, 0.9839, 0.9770, 0.962,0.9576, 0.9537, 0.9451,0.939, 0.9382]
	cnn_x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
	cnn_y = [0.999, 0.999, 0.997, 0.997, 0.997, 0.997, 0.996, 0.995, 0.996, 0.992]
	fno_x = [2, 5, 10, 11]
	fno_y = [0.9989662782956508, 0.999431614249688, 0.9996389623744389, 0.9973824329127635]
end

# ╔═╡ 48e7622f-92dd-4a83-8d52-2d0905e87d18
begin
	plotly()
	plot(title="Fidelity of different model", ylabel="Fidelity", xlabel="Squeezing lavel (dB)")
	plot!(mse_x, mse_y, lw=2, label="MLE")
	plot!(cnn_x, cnn_y, lw=2, label="CNN")
	plot!(fno_x, fno_y, lw=2, label="FNO")
	plot!(legend=:bottomright,)
end

# ╔═╡ c7dd44e4-2f9d-43ba-ac59-d959ccf44c79
ENV["DATADEPS_ALWAYS_ACCEPT"] = true

# ╔═╡ da114fef-f840-47cb-8b7a-710573151b2e
register(DataDep(
		"QSTDemo",
		"""Quantum state tomography demo dataset""",
		"https://drive.google.com/file/d/1Z6g2ZEhMUhqSEQFebVrilQ2gACRJkTVU/view?usp=sharing",
		fetch_method=gdownload,
		post_fetch_method=unpack
	))

# ╔═╡ 720a99ab-ba50-4fb5-a4c1-499546b73591
function get_data(data_name::String, field::String="data_sq")
    data_file = matopen(data_name)
    data = read(data_file, field)
    close(data_file)

    return data
end

# ╔═╡ 970797a8-f845-48a7-b865-e16255137d95
function sample(data::Matrix, n::Integer)
    data_indices = sort!(rand(1:size(data, 1), n))

	θs = data[data_indices, 2]
	θs .-= minimum(θs)
	θs ./= maximum(θs)
	θs .*= 2π

    return vcat(# 1: x; 2: θ
		reshape(θs, 1, :, 1),
		reshape(data[data_indices, 1], 1, :, 1)
	)
end

# ╔═╡ bc1b37e7-893e-46d2-a3d6-b0e13db987ec
function infer(
    data_name::String;
    n_sample=10, fix_θ=true,
    wf=WignerFunction(LinRange(-3, 3, 101), LinRange(-3, 3, 101), dim=100),
    m=SqState.get_model("model_q2args_sqth_x")
)
    data = get_data(data_name)

    argv = zeros(Float32, 3)
    for _ in 1:n_sample
        argv += relu.(m(sample(data, 4096)))
    end
    argv ./= n_sample

    r, θ, n̄ = argv
    θ = fix_θ ? zero(typeof(θ)) : θ

    state = SqState.construct_state_sqth(r, θ, n̄, wf.m_dim)

    return argv, 𝛒(state), wf(state).𝐰_surface
end

# ╔═╡ a55b5296-1596-4441-ba39-1eff14d5ce54
readdir(datadep"QSTDemo")

# ╔═╡ 983f2809-65f4-40d0-a596-ad8ef54348f9
p = datadep"QSTDemo"

# ╔═╡ f76e6855-b72b-4fbc-8ba4-1f53c16d5746
_, _, w = infer(joinpath(p, "SQ10_5mW.mat"), m=m)

# ╔═╡ f9db36e9-29bc-4d96-9549-beb3049bf356
lim = maximum(abs.(w))

# ╔═╡ 1e8002ec-3371-4cd5-a9b1-93aeb1c29393
surface(w, color=:coolwarm, clim=(-lim, lim), border=:none, title="5mW", camera=(40, 10))

# ╔═╡ Cell order:
# ╠═4655bdc2-1ed7-11ec-022d-7df0b8fdb907
# ╠═661cc8fe-44f3-422a-91f5-96956b748847
# ╠═c55d7af1-0cd3-4b0a-bbd5-0048f15ea851
# ╠═8f82e42f-8c2c-4bd0-ab20-95b0930670dc
# ╠═9588e66c-9c19-4257-8fd2-12455e0c0f0d
# ╠═f0cfd6f3-1e55-46cf-bd58-937e52cb7daa
# ╠═57e75729-4bf4-4e0c-9ef2-bdaf6fb4e22b
# ╠═6739ac26-7590-46a8-a75a-2160a1d7c11d
# ╠═00c09086-1eef-4b3f-b764-fe7201e1a354
# ╠═be08410c-4fce-405c-9b27-ababf152402e
# ╠═32c40e60-0989-44f6-b518-9f6203aa6209
# ╠═53da20d3-2d64-4cab-ae4d-0615f750df67
# ╠═bc685a2b-7807-4dd9-b6cd-4a9a1448c09b
# ╠═a4f54b38-ecaa-42b2-bd57-16e9dc52aea0
# ╠═8083f0ef-ad7f-49fe-a248-5931c56ce998
# ╠═48e7622f-92dd-4a83-8d52-2d0905e87d18
# ╠═510b566f-4fd9-4d96-bfc5-b54385b713ce
# ╠═c7dd44e4-2f9d-43ba-ac59-d959ccf44c79
# ╠═da114fef-f840-47cb-8b7a-710573151b2e
# ╠═720a99ab-ba50-4fb5-a4c1-499546b73591
# ╠═970797a8-f845-48a7-b865-e16255137d95
# ╠═bc1b37e7-893e-46d2-a3d6-b0e13db987ec
# ╠═a55b5296-1596-4441-ba39-1eff14d5ce54
# ╠═983f2809-65f4-40d0-a596-ad8ef54348f9
# ╠═f76e6855-b72b-4fbc-8ba4-1f53c16d5746
# ╠═f9db36e9-29bc-4d96-9549-beb3049bf356
# ╠═1e8002ec-3371-4cd5-a9b1-93aeb1c29393


text:
@testset "PigEnv" begin
    env = PigEnv()
    RLBase.test_interfaces!(env)
    RLBase.test_runnable!(env)
end


text:
using Cbc
using Test
using LinearAlgebra

# using Gurobi
# solver = GurobiSolver(OutputFlag=0)

# using CPLEX
# solver = CplexSolver(CPX_PARAM_SCRIND=0)

import JuMP
import MathOptInterface
const MOI = MathOptInterface

using PiecewiseLinearOpt

# TODO: Add :SOS2 to this list, once Cbc supports it
methods_1D = (:CC,:MC,:Logarithmic,:LogarithmicIB,:ZigZag,:ZigZagInteger,:GeneralizedCelaya,:SymmetricCelaya,:Incremental,:DisaggLogarithmic)
# TODO: Add :SOS2 to this list, once Cbc supports it
# TODO: Add :MC to this list, Cbc (but not Gurobi) gives a different answer below, only for :MC (maybe a bug in Cbc?)
methods_2D = (:CC,:Logarithmic,:LogarithmicIB,:ZigZag,:ZigZagInteger,:GeneralizedCelaya,:SymmetricCelaya,:DisaggLogarithmic)
patterns_2D = (:Upper,:Lower,:BestFit,:UnionJack,:K1,:Random) # :OptimalTriangleSelection and :Stencil not supported currently

println("\nunivariate tests")
@testset "1D: $method" for method in methods_1D
    model = JuMP.Model(JuMP.with_optimizer(Cbc.Optimizer))
    JuMP.@variable(model, x)
    z = piecewiselinear(model, x, range(1,stop=2π, length=8), sin, method=method)
    JuMP.@objective(model, Max, z)

    JuMP.optimize!(model)

    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.value(x) ≈ 1.75474 rtol=1e-4
    @test JuMP.value(z) ≈ 0.98313 rtol=1e-4

    JuMP.@constraint(model, x ≤ 1.5z)

    JuMP.optimize!(model)

    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.value(x) ≈ 1.36495 rtol=1e-4
    @test JuMP.value(z) ≈ 0.90997 rtol=1e-4
    @test JuMP.objective_value(model) ≈ 0.90997 rtol=1e-4
    @test JuMP.objective_value(model) ≈ JuMP.value(z) rtol=1e-4
end

println("\nbivariate tests")
@testset "2D: $method, $pattern" for method in methods_2D, pattern in patterns_2D
    model = JuMP.Model(JuMP.with_optimizer(Cbc.Optimizer))
    JuMP.@variable(model, x[1:2])
    d = range(0,stop=1,length=8)
    f = (x1,x2) -> 2*(x1-1/3)^2 + 3*(x2-4/7)^4
    z = piecewiselinear(model, x[1], x[2], BivariatePWLFunction(d, d, f, pattern=pattern), method=method)
    JuMP.@objective(model, Min, z)

    JuMP.optimize!(model)

    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.value(x[1]) ≈ 0.285714 rtol=1e-4
    @test JuMP.value(x[2]) ≈ 0.571429 rtol=1e-4
    @test JuMP.value(z) ≈ 0.004535 rtol=1e-3
    @test JuMP.objective_value(model) ≈ 0.004535 rtol=1e-3
    @test JuMP.objective_value(model) ≈ JuMP.value(z) rtol=1e-3

    JuMP.@constraint(model, x[1] ≥ 0.6)

    JuMP.optimize!(model)

    @test JuMP.termination_status(model) == MOI.OPTIMAL
    @test JuMP.value(x[1]) ≈ 0.6 rtol=1e-4
    @test JuMP.value(x[2]) ≈ 0.571428 rtol=1e-4
    @test JuMP.value(z) ≈ 0.148753 rtol=1e-4
    @test JuMP.objective_value(model) ≈ 0.148753 rtol=1e-3
    @test JuMP.objective_value(model) ≈ JuMP.value(z) rtol=1e-3
end
#
# println("\nbivariate optimal IB scheme tests")
# @testset "2D: optimal IB, UnionJack" begin
#     model = JuMP.Model(JuMP.with_optimizer(Cbc.Optimizer))
#     JuMP.@variable(model, x)
#     JuMP.@variable(model, y)
#     d = range(0,stop=1,length=3)
#     f = (x,y) -> 2*(x-1/3)^2 + 3*(y-4/7)^4
#     z = piecewiselinear(model, x, y, BivariatePWLFunction(d, d, f, pattern=:UnionJack), method=:OptimalIB, subsolver=solver)
#     JuMP.@objective(model, Min, z)
#
#     JuMP.optimize!(model)
#
#     @test JuMP.termination_status(model) == MOI.OPTIMAL
#     @test JuMP.value(x) ≈ 0.5 rtol=1e-4
#     @test JuMP.value(y) ≈ 0.5 rtol=1e-4
#     @test JuMP.value(z) ≈ 0.055634 rtol=1e-3
#     @test getobjectivevalue(model) ≈ 0.055634 rtol=1e-3
#     @test getobjectivevalue(model) ≈ JuMP.value(z) rtol=1e-3
#
#     JuMP.@constraint(model, x ≥ 0.6)
#
#     JuMP.optimize!(model)
#
#     @test JuMP.termination_status(model) == MOI.OPTIMAL
#     @test JuMP.value(x) ≈ 0.6 rtol=1e-4
#     @test JuMP.value(y) ≈ 0.5 rtol=1e-4
#     @test JuMP.value(z) ≈ 0.222300 rtol=1e-3
#     @test JuMP.objective_value(model) ≈ 0.222300 rtol=1e-3
#     @test JuMP.objective_value(model) ≈ JuMP.value(z) rtol=1e-3
# end


text:
using VoronoiDelaunay
using GeometricalPredicates
import GeometricalPredicates
import GeometricalPredicates: area  # needed so that area() can be called in runtests.jl
import VoronoiDelaunay.isexternal, VoronoiDelaunay.locate, VoronoiDelaunay.findindex
import Base: getindex, ==, -, sizeof, endof

if Pkg.installed("Distributions") == nothing
    @eval const haveDistributions = false
else
    @eval const haveDistributions = true
    using Distributions
end

export VoronoiCell, VoronoiCellsA, voronoicells, voronoicellsnogrid, findindex, locate, invoronoicell, area, avoronoicellsnogrid
export getcellindex,isexternal, nverts, nedges, getgenerator, scale, iscale, get_

export VoronoiCellIndex, isvalid, getinvalidcellindex

# For VoronoiCellsA
export cellstogrid, poissontesselation, poissonvoronoicells, poissonvoronoicellsnogrid,
approxpoissontesselation, approxpoissonvoronoicells, approxpoissonvoronoicellsnogrid,
       ngrid, findindex0
export slocate, sfindindex0, sfindindex, area, sarea, smaxcoord, smincoord, sgetgenerator
export getareascale, getscale, getshift, getcells, random_cell, random_cell_index

####

const min_coord = GeometricalPredicates.min_coord + eps(Float64)
const max_coord = GeometricalPredicates.max_coord - eps(Float64)

#### Some functions that only operate on GeometricalPredicates

# Euclidean distance
distance(p1::Point2D,p2::Point2D) = (getx(p2)-getx(p1))^2 + (gety(p2)-gety(p1))^2

# Maybe faster than Euclidean, and is sufficient for use below
distanceL1(p1::Point2D,p2::Point2D) = abs(getx(p2)-getx(p1)) + abs(gety(p2)-gety(p1))

# Removed use of this. Maybe remove this, as well
-(p1::Point2D,p2::Point2D) = Point2D(getx(p1)-getx(p2),gety(p1)-gety(p2))

# sign shows which side of line between P(x1,y1) P(x2,y2) the point

lineside(p1,p2,p) = (gety(p)-gety(p1))*(getx(p2)-getx(p1)) - (getx(p)-getx(p1))*(gety(p2)-gety(p1))

# true if p is inside poly, ie. if lineside gives the same sign for all edges in Voronoi cell
"""
    inconvexpolygon(poly::Array{Point2D,1}, p0::Point2D)

return `true` if `p0` is inside the polygon `poly`.
"""
function inconvexpolygon(poly::Array{Point2D,1}, p0::Point2D)
    is_inconvexpolygon = true
    n = length(poly)
    p1 = poly[n]
    p2 = poly[1]
    s = lineside(p1,p2,p0)
    @inbounds for i in 1:length(poly)-1
        p1 = poly[i]
        p2 = poly[i+1]
        s1::Float64 = lineside(p1,p2,p0)
        if sign(s1) != sign(s)
            is_inconvexpolygon = false
            break
        end
    end
    is_inconvexpolygon
end

function random_coordinate()
    width = max_coord - min_coord
    min_coord+rand()*width
end

function random_coordinate(cutoff)
    width = (max_coord - min_coord) - 2*cutoff
    min_coord + cutoff + rand()*width
end

random_point() = Point2D(random_coordinate(),random_coordinate())
random_point(cutoff) = Point2D(random_coordinate(cutoff),random_coordinate(cutoff))

#### Some functions that only operate on VoronoiDelaunay objects.

# Get vertex by integer index (1,2,3) are (a,b,c)
function get_vertex(trig::VoronoiDelaunay.DelaunayTriangle, iv::Int)
    if iv == 1
        geta(trig)
    elseif iv == 2
        getb(trig)
    elseif iv == 3
        getc(trig)
    else
        error("Vertex index must be 1,2, or 3. Got $iv.")
    end
end

# Get the neigboring triangle sharing side opposite vertex a,b,c via indices 1,2,3.
# Note regarding DelaunayTriangle:
# tr._neighbour_a is the neighboring triangle sharing the edge in tr that is opposite point a
# tr._neighbour_b, tr._neighbour_c contain the vertex a, while tr._neighbour_a does not.
function get_neighbor(trigs,trig,iv::Int)
    if iv == 1
        trigs[trig._neighbour_a]
    elseif iv == 2
        trigs[trig._neighbour_b]
    elseif iv == 3
        trigs[trig._neighbour_c]
    else
        error("Neighbor index must be 1,2, or 3. Got $iv.")
    end
end

# Same as above, but return index of triangle in array of Delaunay triangles,
# rather than the triangle itself
function get_neighbor_index(trigs,trig,iv::Int)
    if iv == 1
        trig._neighbour_a
    elseif iv == 2
        trig._neighbour_b
    elseif iv == 3
        trig._neighbour_c
    else
        error("Neighbor index must be 1,2, or 3. Got $iv.")
    end
end

# Return 1,2, or 3 if one of the three vertices a,b, or c
# in triangle tr is equal to point p0. error otherwise.
function match_vertex(tr, p0)
    eps0 = eps(min_coord) # same as for max_coord
    pa = geta(tr)
    if distanceL1(pa,p0) <= eps0
        return 1
    else
        pb = getb(tr)
        if distanceL1(pb,p0) <= eps0
            return 2
        else
            pc = getc(tr)
            if distanceL1(pc,p0) <= eps0
                return 3
            end
        end
    end
    error("match_vertex: Given triangle does not contain given vertex")
end

# return complement of {v1,v2} in {1,2,3}
# i.e. if input is 3,2, we return 1.
function match_other_vertex(v1,v2)
    v1 == v2 &&  error("Can't find other vertex: v1=$v1 v2=$v2")
    if v1 == 1
        if v2 == 2
            return 3
        else
            return 2
        end
    elseif v1 == 2
        if v2 == 1
            return 3
        else
            return 1
        end
    elseif v1 == 3    # wastes time, but safer
        if v2 == 1
            return 2
        else
            return 1
        end
    else
        error("Can't find other vertex: v1=$v1 v2=$v2")
    end
end

# Find the next DelaunayTriangle triangle when building a VoronoiCell
# trigs -- array of DelaunayTriangles
# tr -- current triangle
# iv_gen -- index (recall (1,2,3) <--> (a,b,c)) of generator point in the triangle tr
#  for the cell we are building. Each triangle has the generator as a vertex, but
#  the index may vary with the triangles.
# iv_opp -- index of the vertex opposite the side that is shared with the
#  next triangle.
# iv_other -- index of remaining vertex
# return:
# tr2  -- next triangle in the cell
# itr2  -- index into the array of all DelaunayTriangle's of tr2
# iv_gen2 -- index (1,2,3) in tr2 of vertex coincident with the generator.
# iv_opp2 -- index in tr2 of vertex opposite the edge in tr2 shared with
#  the next-next triangle.
# iv_other2 -- remaining vertex in tr2
@inline function next_cell_triangle(trigs,tr,iv_gen::Int,iv_opp::Int, iv_other::Int)
    eps0 = eps(min_coord)
    tr2 = get_neighbor(trigs,tr,iv_opp)
    itr2 = get_neighbor_index(trigs,tr,iv_opp)
    iv_gen2 = match_vertex(tr2,get_vertex(tr,iv_gen))
    iv_opp2 = match_vertex(tr2,get_vertex(tr,iv_other))
    iv_other2 = match_other_vertex(iv_gen2, iv_opp2)
    return (tr2, itr2, iv_gen2, iv_opp2, iv_other2)
end

#### Begin code specific to VoronoiCells

"""
    VoronoiCell

The VoronoiCell type. It contains a generator point and an array of the vertices of
the cell. The field `_generator` is one of the points in the point process. There is exactly one such point in each
cell.
"""
immutable VoronoiCell
    _generator::Point2D
    _verts::Array{Point2D,1}
    # neigbhors ?
end

VoronoiCell(generator) = VoronoiCell(generator, Array(Point2D, 0))
nverts(c::VoronoiCell) = length(c._verts)
nedges(c::VoronoiCell) = length(c._verts) + 1
getgenerator(c::VoronoiCell) = c._generator

# Area of irregular polygon. Don't make use of convex property.
"""
    area(c::VoronoiCell)

return the area of `c::VoronoiCell`.
"""
function area(c::VoronoiCell)
    vs::Array{Point2D,1} = c._verts
    A::Float64 = 0.0
    n::Int = length(vs)
   @inbounds @simd for i in 1:n-1
        A += getx(vs[i])*gety(vs[i+1]) -  getx(vs[i+1])*gety(vs[i])
    end
    A += getx(vs[n])*gety(vs[1]) -  getx(vs[1])*gety(vs[n])
    0.5 * abs(A)
end

# is any vertex on boundary of cell outside of allowed range
function isexternal(c::VoronoiCell)
    vs::Array{Point2D,1} = c._verts
    n::Int = length(vs)
    found_external::Bool = false
    @inbounds for i in 1:n
        xc::Float64 = getx(vs[i])
        if xc < min_coord || xc > max_coord
            found_external = true
            break
        end
        yc::Float64 = gety(vs[i])
        if yc < min_coord || yc > max_coord
            found_external = true
            break
        end
    end
    found_external
end

invoronoicell(c::VoronoiCell, p::Point2D) = inconvexpolygon(c._verts,p)

# Find the index in an array of cells of the (first) cell containing point p,
# or zero if no cell contains p.
# Use linear search. We do not use this function on type VoronoiCellsA below.
# It is too expensive to search all cells. Searching for one point in each of 10^6
# cells takes on the order of 1 day cpu time.
# Below, there is a method for this function for VoronoiCellsA.
function findindex(cells::Array{VoronoiCell,1}, p::Point2D)
    ifound::Int = 0
    @inbounds for i in 1:length(cells)
        if invoronoicell(cells[i],p)
            ifound = i
            break
        end
    end
    ifound
end

# could do this with iterator, I suppose, but that constructs all cells every time
# return the cell rather than index. Do not check for error (index == 0 )
locate(cells::Array{VoronoiCell,1}, p::Point2D) = cells[findindex(cells,p)]

function find_cell(trigs,tr, iv_gen, iv_opp, visited)
    cell::VoronoiCell = VoronoiCell(get_vertex(tr, iv_gen))
    iv_other = match_other_vertex(iv_gen,iv_opp)
    tr0 = tr
    is_visited::Bool = false
    while true
        (tr2, itr2, iv_gen, iv_opp, iv_other) = next_cell_triangle(trigs, tr, iv_gen, iv_opp, iv_other)
        if visited[itr2] && tr2 != tr0
            is_visited = true
            break
        end
        push!(cell._verts, circumcenter(tr))
        if tr2 == tr0 break end
        tr = tr2
    end
    (is_visited,cell)
end

# Using  produce, collect is much slower than just returning an array with avoronoicellsnogrid

"""
    avoronoicellsnogrid(t::DelaunayTessellation2D)

return an iterator of the cells `a::Array{VoronoiCell,1}` in the tesselation `t`.
This function uses `Task` and is much slower than `avoronoicellsnogrid(t::DelaunayTessellation2D)`, which
simply returns an array.
"""
function voronoicellsnogrid(t::DelaunayTessellation2D)
    visited::Array{Bool,1} = zeros(Bool, t._last_trig_index)
    visited[1] = true
    function voronoicelliterator()
        j::Int = 0
	@inbounds for ix::Int in 2:t._last_trig_index
	    visited[ix] && continue
	    const tr = t._trigs[ix]
	    visited[ix] = true
            isexternal(tr) && continue
            for iv_gen in 1:3

                iv_opp = mod1(iv_gen+1,3) # pick one of the other two vertices
                ix2 = get_neighbor_index(t._trigs,tr,iv_opp)
                visited[ix2] && continue

                iv_opp = mod1(iv_gen+2,3) # pick the remaining vertex
                ix2 = get_neighbor_index(t._trigs,tr,iv_opp)
                visited[ix2] && continue

                (is_visited, cell::VoronoiCell) = find_cell(t._trigs, tr, iv_gen, iv_opp, visited)
                is_visited && continue
                j += 1
#                j > 10000 && return
                isexternal(cell) && continue # just checking each triangle does not seem to work.
                produce(cell)
            end
        end
    end
    Task(voronoicelliterator)
end

"""
    avoronoicellsnogrid(t::DelaunayTessellation2D)

return an array of the cells `a::Array{VoronoiCell,1}` in the tesselation `t`.
This function is faster than `voronoicellsnogrid(t::DelaunayTessellation2D)`, which
uses lightweight threads.
"""
function avoronoicellsnogrid(t::DelaunayTessellation2D)
    visited = zeros(Bool, t._last_trig_index)
    visited[1] = true
    cells = Array(VoronoiCell,0)
    j = 0
    @inbounds for ix in 2:t._last_trig_index
	visited[ix] && continue
	const tr = t._trigs[ix]
	visited[ix] = true
        isexternal(tr) && continue
        @inbounds for iv_gen in 1:3

            iv_opp = mod1(iv_gen+1,3) # pick one of the other two vertices
            ix2 = get_neighbor_index(t._trigs,tr,iv_opp)
            visited[ix2] && continue

            iv_opp = mod1(iv_gen+2,3) # pick the remaining vertex
            ix2 = get_neighbor_index(t._trigs,tr,iv_opp)
            visited[ix2] && continue

            (is_visited, cell) = find_cell(t._trigs, tr, iv_gen, iv_opp, visited)
            is_visited && continue
            j += 1
            #                j > 10000 && return
            isexternal(cell) && continue # just checking each triangle does not seem to work.
            push!(cells,cell)
        end
    end
    cells
end

#### VoronoiCellsA

# Voronoi cells arranged in efficiently searchable array
# structure is _ngrid by _ngrid array of arrays of indices into the linear array.
# elements in the 2d array correspond to square regions covering area of point process.
#  of all VoronoiCell's.
# _cells -- linear array of all cells (currently in (almost) no order)
# _grid -- 2d array of arrays of indices into _cells.
# _scale, _shift, _areascale. For convenience of user. See scaled versions of functions below.
immutable VoronoiCellsA
    _ngrid::Int
    _cells::Array{VoronoiCell,1}
    _grid::Array{Array{Int,1},2}
    _scale::Float64
    _shift::Float64
    _areascale::Float64
end

# index into grid and array structure VoronoiCellsA
immutable VoronoiCellIndex
    _ix::Int
    _iy::Int
    _ind::Int
end

# Pick random points on the shifted unit square and return the cell there.
# This weights cells by their areas.
function random_cell(c::VoronoiCellsA)
    ntries::Int = 0
    local idx::VoronoiCellIndex
    while true
        pt = random_point()
        ntries += 1
        idx = findindex0(c,pt)
        if isvalid(idx) break end
        if ntries > 1000 error("VoronoiCells: can't find a random cell") end
    end
    c[idx]
end

function random_cell_index(c::VoronoiCellsA)
    ntries::Int = 0
    local idx::VoronoiCellIndex
    while true
        pt = random_point()
        ntries += 1
        idx = findindex0(c,pt)
        if isvalid(idx) break end
        if ntries > 1000 error("VoronoiCells: can't find a random cell") end
    end
    idx
end

function random_cell(c::VoronoiCellsA, cutoff)
    ntries::Int = 0
    local idx::VoronoiCellIndex
    while true
        pt = random_point(cutoff)
        ntries += 1
        idx = findindex0(c,pt)
        if isvalid(idx) break end
        if ntries > 1000 error("VoronoiCells: can't find a random cell") end
    end
    c[idx]
end

function random_cell_index(c::VoronoiCellsA, cutoff)
    ntries::Int = 0
    local idx::VoronoiCellIndex
    while true
        pt = random_point(cutoff)
        ntries += 1
        idx = findindex0(c,pt)
        if isvalid(idx) break end
        if ntries > 1000 error("VoronoiCells: can't find a random cell") end
    end
    idx
end

"""
isvalid(idx::VoronoiCellIndex)

Return `true` if `idx` is valid cell index.

## Example

idx = findindex(cells,x,y)
isvalid(idx)

return true if a cell was found containing `(x,y)`, otherwise `false`.
"""
Base.isvalid(iv::VoronoiCellIndex) = iv._ind != 0
splat(iv::VoronoiCellIndex) = (iv._ix,iv._iy,iv._ind)
==(iv1::VoronoiCellIndex, iv2::VoronoiCellIndex) = splat(iv1) == splat(iv2)

"""
    getinvalidcellindex()

return an invalid cell index. This is useful for initializing an index
that will later be checked for validity.
"""
getinvalidcellindex() = VoronoiCellIndex(0,0,0)

getareascale(gcells::VoronoiCellsA) = gcells._areascale
getscale(gcells::VoronoiCellsA) = gcells._scale
getshift(gcells::VoronoiCellsA) = gcells._shift
getcells(gcells::VoronoiCellsA) = gcells._cells
# Note that c[i], c[i,j], c[i,j,k] return very different things
getindex(c::VoronoiCellsA, i::Int) = c._cells[i]
getindex(c::VoronoiCellsA, i::Int, j::Int) = c._grid[i,j]
getindex(c::VoronoiCellsA, i::Int, j::Int, k::Int) =  c._cells[c._grid[i,j][k]]
getindex(c::VoronoiCellsA, iv::VoronoiCellIndex) = getindex(c,splat(iv)...)
endof(c::VoronoiCellsA) = endof(c._cells)
# return integer index into big 1d array of all cells
getcellindex(c::VoronoiCellsA, i::Int, j::Int, k::Int) =  c._grid[i,j][k]
getcellindex(c::VoronoiCellsA, iv::VoronoiCellIndex) = getcellindex(c, splat(iv)...)
Base.length(c::VoronoiCellsA) = length(c._cells)
ngrid(c::VoronoiCellsA) = c._ngrid

# Some of these look broken to me... iscale must take two args
# inverse of scaling of coordinates
iscale(gc::VoronoiCellsA, x) = (x / gc._scale) + gc._shift
iscale(gc::VoronoiCellsA, p::Point2D) = Point2D(iscale(gc,getx(p)),iscale(gc,gety(p)))

Base.scale(gc::VoronoiCellsA, p::Point2D) = Point2D(scale(gc,getx(p)),scale(gc,gety(p)))
Base.scale(gc::VoronoiCellsA, x) = (x - gc._shift) * gc._scale

# for initializing member _grid above
function make_grid_array(ngrid::Int)
    gridcells = Array(Array{Int,1},ngrid,ngrid)
    for i in 1:ngrid
        for j in 1:ngrid
        @inbounds gridcells[i,j] = Array(Int,0)
        end
    end
    gridcells
end

# Find the element in the 2d grid containing the point (x,y)
# Assume input point satisfies 1.0 <= x,y <= 2.0 (actually a bit different)
# Divide this region into ngrid x ngrid grid and return index.
# of box that gp is in.
# This probably throws away carefully preserved precision.
# Should we use width of area rather than 1.0 ?
function find_grid_element(pt::Point2D, ngrid::Int)
    ix::Int = round(Int64, (getx(pt) - 1.0) * ngrid) + 1
    iy::Int = round(Int64, (gety(pt) - 1.0) * ngrid) + 1
    ix > ngrid ? ix = ngrid : nothing
    iy > ngrid ? iy = ngrid : nothing
    ix < 1 ? ix = 1 : nothing
    iy < 1 ? iy = 1 : nothing
    return (ix,iy)  # return indices
end

#find_grid_element(p::Point2D, ngrid::Int) = find_grid_element(getx(p),gety(p),ngrid)
find_grid_element(cell::VoronoiCell, ngrid::Int) = find_grid_element(cell._generator, ngrid)

function make_grid_cells(cells::Array{VoronoiCell,1}, ngrid::Int)
    gridcells = make_grid_array(ngrid)
    n::Int = length(cells)
    @inbounds for i in 1:n
        cell = cells[i]
        (ix::Int,iy::Int) = find_grid_element(cell,ngrid)
        push!(gridcells[ix,iy], i)
    end
    gridcells
end

# BUG!! area scale is wrong here, because cells have been thrown away
# Assign indices of cells in big linear array to arrays in each
# grid element and return new object.
# Cells are assigned to the grid element the generator of the cell is in.
function cellstogrid(cells::Array{VoronoiCell,1}, ngrid::Int)
    gridcells = make_grid_cells(cells,ngrid)
    n = length(cells)
    gshift::Float64 = 1.5
    gscale::Float64 = sqrt(n)
    gareascale::Float64 = convert(Float64,n)
    VoronoiCellsA(ngrid,cells,gridcells, gscale, gshift, gareascale)
end

function cellstogrid(cells::Array{VoronoiCell,1}, ngrid::Int, gscale, gshift, gareascale)
    gridcells = make_grid_cells(cells,ngrid)
    VoronoiCellsA(ngrid,cells,gridcells, gscale, gshift, gareascale)
end

# Version 0.5.0-dev+3385 (2016-04-02 23:53 UTC) throws an error when
# trying to use 'collect' as in commented out code above. So
# we make an explicit loop to create the cell files.
function voronoicells(t::DelaunayTessellation2D, ndiv)
    cells = avoronoicellsnogrid(t)
    ngrid = round(Int,sqrt(length(cells))/ndiv)
    cellstogrid(cells, ngrid)
end

function voronoicells(t::DelaunayTessellation2D)
     voronoicells(t,10) # ndiv defaults to 10
end

# find index in short array of indices into big array of indices of cells
# such that the cell contains point (x,y). If no such index exists in the
# short array, return zero.
# cells -- array of all cells
# indarray -- array of indices into cells corresponding one grid element
# x,y -- point in cell we are searching for
function findindexA(cells::Array{VoronoiCell,1}, indarray::Array{Int,1}, p::Point2D)
    ifound = 0
    @inbounds for i in 1:length(indarray)
        if invoronoicell(cells[indarray[i]],p)
            ifound = i
            break
        end
    end
    ifound
end

# same as findindexA, but gc contains all structures and we identify index (i,j)
# of grid element corresponding to short array of indices.
# Find cell containing (x,y) in bin (i,j) and return index into gc._cells
findindexingrid(gc::VoronoiCellsA, i, j, pt) = findindexA(gc._cells,gc._grid[i,j],pt)

function findindex00(gridcells::VoronoiCellsA, pt::Point2D)
    (ix::Int,iy::Int) = find_grid_element(pt,size(gridcells._grid,1))
    ind::Int = findindexA(gridcells._cells, gridcells._grid[ix,iy], pt)
    return(ix,iy,ind)
end

macro maybe_return_from_find_index()
    esc(quote
        ind = findindexingrid(grc,ix0,iy0,p)
        ind != 0 && return VoronoiCellIndex(ix0,iy0,ind)
        end)
end

# Find index of cell in gc containing point p.
# using the grid makes the search efficient.
# Rounding errors cause point to not be found in a computed grid bin for about 2/3 percent
# of randomly chosen points. In these cases, we look for the point in the neighboring bins.
# Test shows that this works for all random points (no misses found in 10^7 or more trials)
# Each case caught below occurs in 10^6 trials.
function findindex0(grc::VoronoiCellsA, p::Point2D)
    (ix::Int,iy::Int, ind::Int) = findindex00(grc,p)
    ind != 0 && return VoronoiCellIndex(ix,iy,ind)
    if ix > 1
        ix0 = ix-1
        iy0 = iy
        @maybe_return_from_find_index
        if iy > 1
            iy0 = iy-1
            @maybe_return_from_find_index
        end
        if iy < ngrid(grc)
            iy0 = iy+1
            @maybe_return_from_find_index
        end
    end
    if ix < ngrid(grc)
        ix0 = ix+1
        iy0 = iy
        @maybe_return_from_find_index
        if iy > 1
            iy0 = iy-1
            @maybe_return_from_find_index
        end
        if iy < ngrid(grc)
            iy0 = iy+1
            @maybe_return_from_find_index
        end
    end
    if iy > 1
        iy0 = iy-1
        ix0 = ix
        @maybe_return_from_find_index
    end
    if iy < ngrid(grc)
        iy0 = iy+1
        ix0 = ix
        @maybe_return_from_find_index
    end
    return VoronoiCellIndex(ix,iy,ind)
end

# User gives last returned index as hint if new p is close to p for last call
# i.e. assume we are in the same cell. We also assume that the same indices
# ix,iy will be computed. If hint is wrong, we do the usual search.
function findindex0(gridcells::VoronoiCellsA, hint::Int, p::Point2D)
    (ix::Int,iy::Int) = find_grid_element(p,size(gridcells._grid,1))
    length(gridcells[ix,iy]) >= hint && invoronoicell(gridcells[ix,iy,hint],p) && return VoronoiCellIndex(ix,iy,hint)
    findindex0(gridcells,p)
end

findindex(gridcells::VoronoiCellsA, x,y) = findindex0(gridcells,Point2D(x,y))
findindex(gridcells::VoronoiCellsA, p::Point2D) = findindex0(gridcells,p)

findindex(gridcells::VoronoiCellsA, hint::VoronoiCellIndex, p::Point2D) =
    findindex0(gridcells,hint._ind, p)
findindex(gridcells::VoronoiCellsA, hint::VoronoiCellIndex, x,y) = findindex0(gridcells,hint._ind, Point2D(x,y))

# Return false if there is no complete cell containing p
function isexternal(gridcells::VoronoiCellsA, p::Point2D)
    idx = findindex0(gridcells,p)
    idx._ind == 0 && return true
    return false
end

isexternal(gridcells::VoronoiCellsA, x, y) = isexternal(gridcells,Point2D(x,y))

# Careful with these, they don't check for ind == 0
function locate(gridcells::VoronoiCellsA, p::Point2D)
    (ix,iy,ind) = findindex(gridcells,p)
    cellind = gridcells._grid[ix,iy][ind]
    gridcells._cells[cellind]
end

locate(gridcells::VoronoiCellsA,x,y) =  locate(gridcells, Point2D(x,y))

# Return just the tesselation structure.

# Generate a sample of Poisson point process in the plane and return Delaunay tesselation.
# If Distributions is installed, sample the number of generator points npts in the unit square, assuming that
# the mean number is n. Otherwise, we take npts to be n. In any case approxpoissontesselation
# always does the latter.

# We should generate all these with a macro, of course!
if haveDistributions

"""
    poissontesselation(n::Int)

return a tesselation `t::DelaunayTessellation` of a sample of the
Poisson point process in the region of the plane used by `GeometricalPredicates`.
This depends on `Distributions.jl`.
"""
@eval function poissontesselation{T<:Real}(n::T)
    width = max_coord - min_coord
    npts = rand(Poisson(n))
    a = Point2D[Point(min_coord+rand()*width, min_coord+rand()*width) for i in 1:npts]  # this is very fast
    tess = DelaunayTessellation()   # this is 5 or 6 times faster than making the cell structure
    push!(tess,a)
    tess
end
else
 @eval   function poissontesselation(n::Int)
    width = max_coord - min_coord
    a = Point2D[Point(min_coord+rand()*width, min_coord+rand()*width) for i in 1:n]  # this is very fast
    tess = DelaunayTessellation()   # this is 5 or 6 times faster than making the cell structure
    push!(tess,a)
    tess
   end
end

# Not an exact sample of the Poisson point process in the plane
"""
    approxpoissontesselation(n::Int)

return a tesselation `t::DelaunayTessellation` of an approximation of a sample of the
Poisson point process in the region of the plane used by `GeometricalPredicates`.
Exactly `n` generator points will be in the sample.
"""
function approxpoissontesselation(n::Int)
    width = max_coord - min_coord
    a = Point2D[Point(min_coord+rand()*width, min_coord+rand()*width) for i in 1:n]  # this is very fast
    tess = DelaunayTessellation()   # this is 5 or 6 times faster than making the cell structure
    push!(tess,a)
    tess
end

# Create tesselation of poisson point process sample.
# Return only the efficient cell structure.
# ndiv^2 is the average number of generator points per square in the grid
function poissonvoronoicells(n, ndiv)
    tess = poissontesselation(n)
    gcells = voronoicells(tess,ndiv)
end

# Same as above, but calculate standard grid sized from number of cells
function poissonvoronoicells(n)
    tess = poissontesselation(n)
    gcells = voronoicells(tess)
end

# Generate and return single array of all cells. Do not create grid
# structure. Cannot search efficiently for random points in these cells.
# See below about not using 'collect'
poissonvoronoicellsnogrid(n::Int) = avoronoicellsnogrid(poissontesselation(n))

function approxpoissonvoronoicells(n, ndiv)
    tess = approxpoissontesselation(n)
    gcells = voronoicells(tess,ndiv)
end

function approxpoissonvoronoicells(n)
    tess = approxpoissontesselation(n)
    gcells = voronoicells(tess)
end

approxpoissonvoronoicellsnogrid(n::Int) = avoronoicellsnogrid(approxpoissontesselation(n))

####

# Scaled versions of some functions. The tesselated region, approximately
# 1.0 <= x,y, 2.0  is scaled and shifted to coordinates convenient for the
# user.
slocate(gcells::VoronoiCellsA,x,y) = locate(gcells,iscale(gcells,x),iscale(gcells,y))
slocate(gcells::VoronoiCellsA, p::Point2D) = locate(gridcells, iscale(gcells,p))
sfindindex0(gcells::VoronoiCellsA, x,y) = findindex0(gcells,iscale(gcells,x),iscale(gcells,y))
sfindindex0(gcells::VoronoiCellsA, p::Point2D) = findindex0(gcells, iscale(gcells,p))
sfindindex0(gcells::VoronoiCellsA, hint, x,y) = findindex0(gcells,hint,iscale(gcells,x),iscale(gcells,y))
sfindindex0(gcells::VoronoiCellsA, hint, p::Point2D) = findindex0(gcells, hint, iscale(gcells,p))
sfindindex(gcells::VoronoiCellsA, x,y) = findindex(gcells,iscale(gcells,x),iscale(gcells,y))
sfindindex(gcells::VoronoiCellsA, p::Point2D) = findindex(gcells, iscale(gcells,p))
sfindindex(gridcells::VoronoiCellsA, hint::VoronoiCellIndex, x,y) = findindex(gridcells, hint, iscale(gridcells,x),iscale(gridcells,y))

sarea(gcells::VoronoiCellsA, i::Int) = area(gcells[i]) * getareascale(gcells)
sarea(gcells::VoronoiCellsA, i::Int, j::Int, k::Int) = area(gcells[i,j,k]) * getareascale(gcells)
sarea(gcells::VoronoiCellsA, idx::VoronoiCellIndex) = sarea(gcells, splat(idx)...)

area(gcells::VoronoiCellsA, i::Int, j::Int, k::Int) = area(gcells[i,j,k])
area(gcells::VoronoiCellsA, idx::VoronoiCellIndex) = area(gcells, splat(idx)...)

smaxcoord(gcells) = scale(gcells,max_coord)
smincoord(gcells) = scale(gcells,min_coord)
sgetgenerator(gcells::VoronoiCellsA, c::VoronoiCell) = scale(gcells,c._generator)

#####

function sizeof(c::VoronoiCell)
    s1::Int = sizeof(c._generator)
    s2::Int = sizeof(c._verts)
    s3::Int = 8  # for address of c._verts
    s1 + s2 + s3
end

function sizeof(cells::VoronoiCellsA)
    s1::Int = sizeof(cells._cells)
    s2::Int = sizeof(cells._grid)
    s3::Int = 32 # for four simple members
    s1 + s2 + s3
end

##### Following not exported

# This should be about ndiv^2.
function mean_num_points_per_grid_square(cells::VoronoiCellsA)
    sum::Int = 0
    @inbounds for i in 1:length(cells._grid)
        sum += length(cells._grid[i])
    end
    sum / length(cells._grid)
end

function number_of_generators_in_grid(cells::VoronoiCellsA)
    sum::Int = 0
    @inbounds for i in 1:length(cells._grid)
        sum += length(cells._grid[i])
    end
    sum
end


text:
# ---
# title: 594. Longest Harmonious Subsequence
# id: problem594
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Hash Table
# link: <https://leetcode.com/problems/longest-harmonious-subsequence/description/>
# hidden: true
# ---
#
# We define a harmonious array as an array where the difference between its
# maximum value and its minimum value is **exactly** `1`.
#
# Given an integer array `nums`, return _the length of its longest harmonious
# subsequence among all its possible subsequences_.
#
# A **subsequence** of array is a sequence that can be derived from the array by
# deleting some or no elements without changing the order of the remaining
# elements.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,3,2,2,5,2,3,7]
#     Output: 5
#     Explanation: The longest harmonious subsequence is [3,2,2,2,3].
#
#
# **Example 2:**
#
#
#
#     Input: nums = [1,2,3,4]
#     Output: 2
#
#
# **Example 3:**
#
#
#
#     Input: nums = [1,1,1,1]
#     Output: 0
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 2 * 104`
#   * `-109 <= nums[i] <= 109`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
println("INFO: Loading packages")

using FFTW
using Parameters
using DifferentialEquations
using Plots
using Statistics
using PyCall

include("./parameters.jl")
include("./derivative.jl")
include("./initial.jl")
include("./util.jl")
include("./postprocessing.jl")

np = pyimport("numpy")

num_of_samples = 20

Re_range = range(10.0, stop = 69.0, length = num_of_samples)

struct SolverParameters
    k2cut
    N
    parameters
end

function simplified_model!(du, u, p, t)
    # unpack parameters
    k2cut = p.k2cut
    N = p.N
    parameters = p.parameters

    # alias for U
    # refactor: use lowercase `u` throughout
    U = u

    # update derivative
    du[:] = dUdt_simplified_model(U, k2cut, N, parameters)
end

T_MAX = 25.0
radius = 0.6600
tspan = (0.0, T_MAX)
ScaleH = 1
ScaleU = 1
no_t = 50
t_range = LinRange(0.0, T_MAX, no_t)
DIMENSION = 3

Threads.@threads for sample_val in 1:num_of_samples

    global parameters = GeneralParams(Re = Re_range[sample_val], Ka = 1000.)

    N = 4 * parameters.N2 * (parameters.N1 + 1);
    global kx, kz, k2cut = aliasing_filter(radius, parameters)

    u0 = create_initial_condition(kx, kz, k2cut, N, parameters)

    p = SolverParameters(k2cut, N, parameters)

    prob = ODEProblem(simplified_model!, u0, tspan, p)

    println("INFO: Solving equation")
    @time sol = solve(prob, Tsit5(), maxiters = 10000)

    println("INFO: Postprocessing results")

    @unpack M1, M2 = parameters

    # generate plots using Nusselt scaling
    ScaleX = 2 * pi / parameters.k0x
    ScaleZ = 2 * pi / parameters.k0z

    x = (0:1 / (M1 - 1):1) * ScaleX
    z = (0:1 / (M2 - 1):1) * ScaleZ

    lll = zeros(Float64, no_t, 64, 64)
    for i in 1:no_t
        t = t_range[i]
        U = sol(t)
        datah, dataq, datap = reconstruct_U(U, parameters)
        surf = real(datah) * ScaleH
        lll[i,:,:] = surf
    end

    nplll = np.array(lll)

    if sol.retcode == :MaxIters
        println("Unsuccesful iteration.")
    else
        println("Iteration:", sample_val)
        np.save("data/zvals_Re" * string(Re_range[sample_val]) * ".npy", nplll)
    end
end


text:
# This file contains default implementations for the `MOI.copy_to` function that can be used by a model.

"""
    automatic_copy_to(dest::MOI.ModelLike, src::MOI.ModelLike;
                      copy_names::Bool=true,
                      filter_constraints::Union{Nothing, Function}=nothing)

Use [`Utilities.supports_default_copy_to`](@ref) and
[`Utilities.supports_allocate_load`](@ref) to automatically choose between
[`Utilities.default_copy_to`](@ref) or [`Utilities.allocate_load`](@ref) to
apply the copy operation.

If the `filter_constraints` arguments is given, only the constraints for which
this function returns `true` will be copied. This function is given a
constraint index as argument.
"""
function automatic_copy_to(
    dest::MOI.ModelLike,
    src::MOI.ModelLike;
    copy_names::Bool = true,
    filter_constraints::Union{Nothing,Function} = nothing,
)
    if supports_default_copy_to(dest, copy_names)
        default_copy_to(dest, src, copy_names, filter_constraints)
    elseif supports_allocate_load(dest, copy_names)
        allocate_load(dest, src, copy_names, filter_constraints)
    else
        error(
            "Model $(typeof(dest)) does not support copy",
            copy_names ? " with names" : "",
            ".",
        )
    end
end

"""
    supports_default_copy_to(model::ModelLike, copy_names::Bool)

Return a `Bool` indicating whether the model `model` supports
[`default_copy_to(model, src, copy_names=copy_names)`](@ref) if all the
attributes set to `src` and constraints added to `src` are supported by `model`.

This function can be used to determine whether a model can be loaded into
`model` incrementally or whether it should be cached and copied at once instead.
This is used by JuMP to determine whether to add a cache or not in two
situations:
1. A first cache can be used to store the model as entered by the user as well
   as the names of variables and constraints. This cache is created if this
   function returns `false` when `copy_names` is `true`.
2. If bridges are used, then a second cache can be used to store the bridged
   model with unnamed variables and constraints. This cache is created if this
   function returns `false` when `copy_names` is `false`.

## Examples

If [`MathOptInterface.set`](@ref), [`MathOptInterface.add_variable`](@ref) and
[`MathOptInterface.add_constraint`](@ref) are implemented for a model of type
`MyModel` and names are supported, then [`MathOptInterface.copy_to`](@ref) can
be implemented as
```julia
MOI.Utilities.supports_default_copy_to(model::MyModel, copy_names::Bool) = true
function MOI.copy_to(dest::MyModel, src::MOI.ModelLike; kws...)
    return MOI.Utilities.automatic_copy_to(dest, src; kws...)
end
```
The [`Utilities.automatic_copy_to`](@ref) function automatically redirects to
[`Utilities.default_copy_to`](@ref).

If names are not supported, simply change the first line by
```julia
MOI.supports_default_copy_to(model::MyModel, copy_names::Bool) = !copy_names
```
The [`Utilities.default_copy_to`](@ref) function automatically throws an helpful
error in case `copy_to` is called with `copy_names` equal to `true`.
"""
supports_default_copy_to(model::MOI.ModelLike, copy_names::Bool) = false

"""
    _index_to_variable(i::Int)

Simply returns `MOI.VariableIndex(i)`. This is necessary to pass a function that
creates the `VariableIndex` from an integer. If we pass `MOI.VariableIndex`
julia understands is as a `DataType` and not a function, this leads to type
instability issues.
"""
_index_to_variable(i) = MOI.VariableIndex(i)
const DenseVariableDict{V} = CleverDicts.CleverDict{
    MOI.VariableIndex,
    V,
    typeof(MOI.index_value),
    typeof(_index_to_variable),
}
function dense_variable_dict(::Type{V}, n) where {V}
    return CleverDicts.CleverDict{MOI.VariableIndex,V}(
        MOI.index_value,
        _index_to_variable,
        n,
    )
end

struct IndexMap <: AbstractDict{MOI.Index,MOI.Index}
    # just keeping the union to make it more backward compatible
    # we should remove as soon as possible.
    varmap::Union{
        DenseVariableDict{MOI.VariableIndex},
        Dict{MOI.VariableIndex,MOI.VariableIndex},
    }
    conmap::DoubleDicts.MainIndexDoubleDict
end

function IndexMap(n = 0)
    return IndexMap(
        dense_variable_dict(MOI.VariableIndex, n),
        DoubleDicts.IndexDoubleDict(),
    )
end

"""
    index_map_for_variable_indices(variables)

This function does not add variables to the IndexMap.
It simply initializes the IndexMap with a proper data struture.
If the variable indices are contiguous and start from 1, then
an optimized data structure with pre allocated memory is initialized.
Otherwise the data structure will start empty and will try to
keep using performant structure for as long as possible.
"""
function index_map_for_variable_indices(variables)
    n = length(variables)
    if all(i -> variables[i] == MOI.VariableIndex(i), 1:n)
        return IndexMap(n)
    else
        return IndexMap()
    end
end

Base.getindex(idxmap::IndexMap, vi::MOI.VariableIndex) = idxmap.varmap[vi]
function Base.getindex(
    idxmap::IndexMap,
    ci::MOI.ConstraintIndex{F,S},
) where {F,S}
    return idxmap.conmap[ci]::MOI.ConstraintIndex{F,S}
end

function Base.setindex!(
    idxmap::IndexMap,
    vi1::MOI.VariableIndex,
    vi2::MOI.VariableIndex,
)
    return Base.setindex!(idxmap.varmap, vi1, vi2)
end
function Base.setindex!(
    idxmap::IndexMap,
    ci1::MOI.ConstraintIndex{F,S},
    ci2::MOI.ConstraintIndex{F,S},
) where {F,S}
    return Base.setindex!(idxmap.conmap, ci1, ci2)
end

function Base.delete!(idxmap::IndexMap, vi::MOI.VariableIndex)
    return delete!(idxmap.varmap, vi)
end
function Base.delete!(idxmap::IndexMap, ci::MOI.ConstraintIndex)
    return delete!(idxmap.conmap, ci)
end

function Base.haskey(idxmap::IndexMap, ci::MOI.ConstraintIndex)
    return haskey(idxmap.conmap, ci)
end
Base.haskey(idxmap::IndexMap, vi::MOI.VariableIndex) = haskey(idxmap.varmap, vi)

function Base.keys(idxmap::IndexMap)
    return Iterators.flatten((keys(idxmap.varmap), keys(idxmap.conmap)))
end

Base.length(idxmap::IndexMap) = length(idxmap.varmap) + length(idxmap.conmap)
function Base.iterate(idxmap::IndexMap, args...)
    return iterate(
        Base.Iterators.flatten((idxmap.varmap, idxmap.conmap)),
        args...,
    )
end

"""
    pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, pass_attr::Function=MOI.set)

Pass the model attributes from the model `src` to the model `dest` using `canpassattr` to check if the attribute can be passed and `pass_attr` to pass the attribute. Does not copy `Name` if `copy_names` is `false`.

    pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, vis_src::Vector{MOI.VariableIndex}, pass_attr::Function=MOI.set)

Pass the variable attributes from the model `src` to the model `dest` using `canpassattr` to check if the attribute can be passed and `pass_attr` to pass the attribute. Does not copy `VariableName` if `copy_names` is `false`.

    pass_attributes(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool, idxmap::IndexMap, cis_src::Vector{MOI.ConstraintIndex{F, S}}, pass_attr::Function=MOI.set) where {F, S}

Pass the constraint attributes of `F`-in-`S` constraints from the model `src` to the model `dest` using `canpassattr` to check if the attribute can be passed and `pass_attr` to pass the attribute. Does not copy `ConstraintName` if `copy_names` is `false`.
"""
function pass_attributes end

function pass_attributes(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    copy_names::Bool,
    idxmap::IndexMap,
    pass_attr::Function = MOI.set,
)
    # Copy model attributes
    attrs = MOI.get(src, MOI.ListOfModelAttributesSet())
    if !copy_names
        attrs = filter(attr -> !(attr isa MOI.Name), attrs)
    end
    return _pass_attributes(
        dest,
        src,
        idxmap,
        attrs,
        tuple(),
        tuple(),
        tuple(),
        pass_attr,
    )
end
function pass_attributes(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    copy_names::Bool,
    idxmap::IndexMap,
    vis_src::Vector{VI},
    pass_attr::Function = MOI.set,
)
    # Copy variable attributes
    attrs = MOI.get(src, MOI.ListOfVariableAttributesSet())
    if !copy_names
        attrs = filter(attr -> !(attr isa MOI.VariableName), attrs)
    end
    if !isempty(attrs) # If `attrs` is empty, we can spare the computation of `vis_dest`
        vis_dest = map(vi -> idxmap[vi], vis_src)
        _pass_attributes(
            dest,
            src,
            idxmap,
            attrs,
            (VI,),
            (vis_src,),
            (vis_dest,),
            pass_attr,
        )
    end
end
function pass_attributes(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    copy_names::Bool,
    idxmap::IndexMap,
    cis_src::Vector{CI{F,S}},
    pass_attr::Function = MOI.set,
) where {F,S}
    # Copy constraint attributes
    attrs = MOI.get(src, MOI.ListOfConstraintAttributesSet{F,S}())
    if !copy_names
        attrs = filter(attr -> !(attr isa MOI.ConstraintName), attrs)
    end
    if !isempty(attrs) # If `attrs` is empty, we can spare the computation of `cis_dest`
        cis_dest = map(ci -> idxmap[ci], cis_src)
        _pass_attributes(
            dest,
            src,
            idxmap,
            attrs,
            (CI{F,S},),
            (cis_src,),
            (cis_dest,),
            pass_attr,
        )
    end
end

function _pass_attributes(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    idxmap::IndexMap,
    attrs,
    supports_args,
    get_args,
    set_args,
    pass_attr!::Function = MOI.set,
)
    for attr in attrs
        @assert MOI.is_copyable(attr)
        if attr isa MOI.VariablePrimalStart ||
           attr isa MOI.ConstraintPrimalStart ||
           attr isa MOI.ConstraintDualStart
            # As starting values are simply *hints* for the optimization, not
            # supporting them gives a warning, not an error
            if !MOI.supports(dest, attr, supports_args...)
                @warn(
                    "$attr is not supported by $(typeof(dest)). This " *
                    "information will be discarded."
                )
                continue
            end
        end
        value = MOI.get(src, attr, get_args...)
        if value !== nothing
            mapped_value = map_indices(idxmap, value)
            pass_attr!(dest, attr, set_args..., mapped_value)
        end
    end
end

"""
    copy_vector_of_variables(dest::MOI.ModelLike, src::MOI.ModelLike,
                             idxmap::IndexMap,
                             S::Type{<:MOI.AbstractSet})

Copy the constraints of type `MOI.VectorOfVariables`-in-`S` from the model `src`
to the model `dest` and fill `idxmap` accordingly. The copy is only done when
the variables to be copied are not already keys of `idxmap`. It returns a list
of the constraints copied and not copied.
"""
function copy_vector_of_variables(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    idxmap::IndexMap,
    S::Type{<:MOI.AbstractSet},
    copy_constrained_variables::Function = MOI.add_constrained_variables,
)
    added = MOI.ConstraintIndex{MOI.VectorOfVariables,S}[]
    not_added = MOI.ConstraintIndex{MOI.VectorOfVariables,S}[]
    cis_src =
        MOI.get(src, MOI.ListOfConstraintIndices{MOI.VectorOfVariables,S}())
    fs_src =
        MOI.get(
            src,
            MOI.ConstraintFunction(),
            cis_src,
        )::Vector{MOI.VectorOfVariables}
    for (ci_src, f_src) in zip(cis_src, fs_src)
        if all(vi -> !haskey(idxmap, vi), f_src.variables) &&
           allunique(f_src.variables)
            set = MOI.get(src, MOI.ConstraintSet(), ci_src)::S
            vis_dest, ci_dest = copy_constrained_variables(dest, set)
            idxmap[ci_src] = ci_dest
            for (vi_src, vi_dest) in zip(f_src.variables, vis_dest)
                idxmap[vi_src] = vi_dest
            end
            push!(added, ci_src)
        else
            push!(not_added, ci_src)
        end
    end
    return added, not_added
end

"""
    copy_single_variable(dest::MOI.ModelLike, src::MOI.ModelLike,
                          idxmap::IndexMap,
                          S::Type{<:MOI.AbstractSet})

Copy the constraints of type `MOI.SingleVariable`-in-`S` from the model `src` to
the model `dest` and fill `idxmap` accordingly. The copy is only done when the
variables to be copied are not already keys of `idxmap`. It returns a list of
the constraints not copied.
"""
function copy_single_variable(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    idxmap::IndexMap,
    S::Type{<:MOI.AbstractSet},
    copy_constrained_variable::Function = MOI.add_constrained_variable,
)
    added = MOI.ConstraintIndex{MOI.SingleVariable,S}[]
    not_added = MOI.ConstraintIndex{MOI.SingleVariable,S}[]
    cis_src = MOI.get(src, MOI.ListOfConstraintIndices{MOI.SingleVariable,S}())
    fs_src =
        MOI.get(
            src,
            MOI.ConstraintFunction(),
            cis_src,
        )::Vector{MOI.SingleVariable}
    for (ci_src, f_src) in zip(cis_src, fs_src)
        if !haskey(idxmap, f_src.variable)
            set = MOI.get(src, MOI.ConstraintSet(), ci_src)::S
            vi_dest, ci_dest = copy_constrained_variable(dest, set)
            idxmap[ci_src] = ci_dest
            idxmap[f_src.variable] = vi_dest
            push!(added, ci_src)
        else
            push!(not_added, ci_src)
        end
    end
    return added, not_added
end

"""
    copy_constraints(dest::MOI.ModelLike, src::MOI.ModelLike,
                     idxmap::IndexMap,
                     cis_src::Vector{<:MOI.ConstraintIndex},
                     filter_constraints::Union{Nothing, Function}=nothing)

Copy the constraints `cis_src` from the model `src` to the model `dest` and fill
`idxmap` accordingly. Note that the attributes are not copied; call
[`pass_attributes`] to copy the constraint attributes.

If the `filter_constraints` arguments is given, only the constraints for which
this function returns `true` will be copied. This function is given a
constraint index as argument.
"""
function copy_constraints(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    idxmap::IndexMap,
    cis_src::Vector{<:MOI.ConstraintIndex},
    filter_constraints::Union{Nothing,Function} = nothing,
)
    # Retrieve the constraints to copy.
    f_src = MOI.get(src, MOI.ConstraintFunction(), cis_src)

    # Filter this set of constraints if needed before.
    if filter_constraints !== nothing
        filter!(filter_constraints, cis_src)
    end

    # Copy the constraints into the new model and build the map.
    f_dest = map_indices.(Ref(idxmap), f_src)
    s = MOI.get(src, MOI.ConstraintSet(), cis_src)
    cis_dest = MOI.add_constraints(dest, f_dest, s)
    for (ci_src, ci_dest) in zip(cis_src, cis_dest)
        idxmap.conmap[ci_src] = ci_dest
    end
end

function pass_constraints(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    copy_names::Bool,
    idxmap::IndexMap,
    single_variable_types,
    single_variable_indices,
    vector_of_variables_types,
    vector_of_variables_indices,
    pass_cons = copy_constraints,
    pass_attr = MOI.set;
    filter_constraints::Union{Nothing,Function} = nothing,
)
    # copy_constraints can also take a filter_constraints argument; however, filtering
    # is performed within this function (because it also calls MOI.set on the constraints).
    # Don't pass this argument to copy_constraints/pass_cons to avoid a double filtering.
    for (S, cis_src) in zip(single_variable_types, single_variable_indices)
        if filter_constraints !== nothing
            filter!(filter_constraints, cis_src)
        end
        # do the rest in `pass_cons` which is type stable
        pass_cons(dest, src, idxmap, cis_src)

        cis_src =
            MOI.get(src, MOI.ListOfConstraintIndices{MOI.SingleVariable,S}())
        if filter_constraints !== nothing
            filter!(filter_constraints, cis_src)
        end
        pass_attributes(dest, src, copy_names, idxmap, cis_src, pass_attr)
    end

    for (S, cis_src) in
        zip(vector_of_variables_types, vector_of_variables_indices)
        if filter_constraints !== nothing
            filter!(filter_constraints, cis_src)
        end
        # do the rest in `pass_cons` which is type stable
        pass_cons(dest, src, idxmap, cis_src)

        cis_src =
            MOI.get(src, MOI.ListOfConstraintIndices{MOI.VectorOfVariables,S}())
        if filter_constraints !== nothing
            filter!(filter_constraints, cis_src)
        end
        pass_attributes(dest, src, copy_names, idxmap, cis_src, pass_attr)
    end

    nonvariable_constraint_types = [
        (F, S)
        for
        (F, S) in MOI.get(src, MOI.ListOfConstraints()) if
        F != MOI.SingleVariable && F != MOI.VectorOfVariables
    ]
    for (F, S) in nonvariable_constraint_types
        cis_src = MOI.get(src, MOI.ListOfConstraintIndices{F,S}())
        if filter_constraints !== nothing
            filter!(filter_constraints, cis_src)
        end
        # do the rest in `pass_cons` which is type stable
        pass_cons(dest, src, idxmap, cis_src)
        pass_attributes(dest, src, copy_names, idxmap, cis_src, pass_attr)
    end
end

function copy_free_variables(
    dest::MOI.ModelLike,
    idxmap::IndexMap,
    vis_src,
    copy_variables::F,
) where {F<:Function}
    if length(vis_src) != length(keys(idxmap.varmap))
        vars = copy_variables(dest, length(vis_src) - length(idxmap.varmap))
        i = 1
        for vi in vis_src
            if !haskey(idxmap.varmap, vi)
                idxmap.varmap[vi] = vars[i]
                i += 1
            end
        end
        @assert i == length(vars) + 1
        @assert length(vis_src) == length(idxmap.varmap)
    end
end

function default_copy_to(dest::MOI.ModelLike, src::MOI.ModelLike)
    Base.depwarn(
        "default_copy_to(dest, src) is deprecated, use default_copy_to(dest, src, true) instead or default_copy_to(dest, src, false) if you do not want to copy names.",
        :default_copy_to,
    )
    return default_copy_to(dest, src, true)
end

function sorted_variable_sets_by_cost(dest::MOI.ModelLike, src::MOI.ModelLike)
    constraint_types = MOI.get(src, MOI.ListOfConstraints())
    single_or_vector_variables_types = [
        (F, S)
        for
        (F, S) in constraint_types if
        F === MOI.SingleVariable || F === MOI.VectorOfVariables
    ]
    sort!(
        single_or_vector_variables_types;
        by = ((F, S),) ->
        # We give priority for sets such that there is a big cost reduction
        # constraining the variable on creation.
            (
                MOI.get(dest, MOI.VariableBridgingCost{S}()) -
                MOI.get(dest, MOI.ConstraintBridgingCost{F,S}()),
                # In case of ties, we give priority to vector sets.
                # See issue #987
                F === MOI.SingleVariable,
            ),
    )
    return single_or_vector_variables_types
end

function try_constrain_variables_on_creation(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    idxmap,
    copy_constrained_variables::F1,
    copy_constrained_variable::F2,
) where {F1<:Function,F2<:Function}
    single_or_vector_variables_types = sorted_variable_sets_by_cost(dest, src)
    vector_of_variables_types = Type{<:MOI.AbstractVectorSet}[]
    vector_of_variables_added =
        Vector{<:MOI.ConstraintIndex{MOI.VectorOfVariables}}[]
    vector_of_variables_not_added =
        Vector{<:MOI.ConstraintIndex{MOI.VectorOfVariables}}[]
    single_variable_types = Type{<:MOI.AbstractScalarSet}[]
    single_variable_added = Vector{<:MOI.ConstraintIndex{MOI.SingleVariable}}[]
    single_variable_not_added =
        Vector{<:MOI.ConstraintIndex{MOI.SingleVariable}}[]
    for (F, S) in single_or_vector_variables_types
        if F === MOI.VectorOfVariables
            added, not_added = copy_vector_of_variables(
                dest,
                src,
                idxmap,
                S,
                copy_constrained_variables,
            )
            push!(vector_of_variables_added, added)
            push!(vector_of_variables_not_added, not_added)
            push!(vector_of_variables_types, S)
        elseif F === MOI.SingleVariable
            added, not_added = copy_single_variable(
                dest,
                src,
                idxmap,
                S,
                copy_constrained_variable,
            )
            push!(single_variable_added, added)
            push!(single_variable_not_added, not_added)
            push!(single_variable_types, S)
        end
    end
    return vector_of_variables_types,
    vector_of_variables_added,
    vector_of_variables_not_added,
    single_variable_types,
    single_variable_added,
    single_variable_not_added
end

"""
    default_copy_to(dest::MOI.ModelLike, src::MOI.ModelLike, copy_names::Bool,
                    filter_constraints::Union{Nothing, Function}=nothing)

Implements `MOI.copy_to(dest, src)` by adding the variables and then the
constraints and attributes incrementally. The function
[`supports_default_copy_to`](@ref) can be used to check whether `dest` supports
the copying a model incrementally.

If the `filter_constraints` arguments is given, only the constraints for which
this function returns `true` will be copied. This function is given a
constraint index as argument.
"""
function default_copy_to(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    copy_names::Bool,
    filter_constraints::Union{Nothing,Function} = nothing,
)
    MOI.empty!(dest)

    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    idxmap = index_map_for_variable_indices(vis_src)

    # The `NLPBlock` assumes that the order of variables does not change (#849)
    if MOI.NLPBlock() in MOI.get(src, MOI.ListOfModelAttributesSet())
        constraint_types = MOI.get(src, MOI.ListOfConstraints())
        single_variable_types =
            [S for (F, S) in constraint_types if F == MOI.SingleVariable]
        vector_of_variables_types =
            [S for (F, S) in constraint_types if F == MOI.VectorOfVariables]
        vector_of_variables_not_added = [
            MOI.get(
                src,
                MOI.ListOfConstraintIndices{MOI.VectorOfVariables,S}(),
            ) for S in vector_of_variables_types
        ]
        single_variable_not_added = [
            MOI.get(src, MOI.ListOfConstraintIndices{MOI.SingleVariable,S}()) for S in single_variable_types
        ]
    else
        vector_of_variables_types,
        _,
        vector_of_variables_not_added,
        single_variable_types,
        _,
        single_variable_not_added = try_constrain_variables_on_creation(
            dest,
            src,
            idxmap,
            MOI.add_constrained_variables,
            MOI.add_constrained_variable,
        )
    end

    copy_free_variables(dest, idxmap, vis_src, MOI.add_variables)

    # Copy variable attributes
    pass_attributes(dest, src, copy_names, idxmap, vis_src)

    # Copy model attributes
    pass_attributes(dest, src, copy_names, idxmap)

    # Copy constraints
    pass_constraints(
        dest,
        src,
        copy_names,
        idxmap,
        single_variable_types,
        single_variable_not_added,
        vector_of_variables_types,
        vector_of_variables_not_added,
        filter_constraints = filter_constraints,
    )

    return idxmap
end

# Allocate-Load Interface: 2-pass copy of a MathOptInterface model
# Some solver wrappers (e.g. SCS, ECOS, SDOI) do not supporting copying an optimization model using `MOI.add_constraints`, `MOI.add_variables` and `MOI.set`
# as they first need to figure out some information about a model before being able to pass the problem data to the solver.
#
# During the first pass (called allocate) : the model collects the relevant information about the problem so that
# on the second pass (called load), the constraints can be loaded directly to the solver (in case of SDOI) or written directly into the matrix of constraints (in case of SCS and ECOS).

# To support `MOI.copy_to` using this 2-pass mechanism, implement the allocate-load interface defined below and do:
# MOI.copy_to(dest::ModelType, src::MOI.ModelLike) = MOIU.allocate_load(dest, src)
# In the implementation of the allocate-load interface, it can be assumed that the different functions will the called in the following order:
# 1) `allocate_variables`
# 2) `allocate` and `allocate_constraint`
# 3) `load_variables`
# 4) `load` and `load_constraint`
# The interface is not meant to be used to create new constraints with `allocate_constraint` followed by `load_constraint` after a solve, it is only meant for being used in this order to implement `MOI.copy_to`.

# TODO deprecate this in MOI v0.7
"""
    needs_allocate_load(model::MOI.ModelLike)::Bool
Return a `Bool` indicating whether `model` does not support `add_variables`/`add_constraint`/`set` but supports `allocate_variables`/`allocate_constraint`/`allocate`/`load_variables`/`load_constraint`/`load`.
That is, the allocate-load interface need to be used to copy an model to `model`.
"""
needs_allocate_load(::MOI.ModelLike) = false

"""
    supports_allocate_load(model::MOI.ModelLike, copy_names::Bool)::Bool

Return a `Bool` indicating whether `model` supports
[`allocate_load(model, src, copy_names=copy_names)`](@ref) if all the
attributes set to `src` and constraints added to `src` are supported by `model`.
"""
supports_allocate_load(::MOI.ModelLike, copy_names::Bool) = false

"""
    allocate_variables(model::MOI.ModelLike, nvars::Integer)

Creates `nvars` variables and returns a vector of `nvars` variable indices.
"""
function allocate_variables end

"""
    allocate_constrained_variable(model::MOI.ModelLike,
                                  set::MOI.AbstractScalarSet)

Returns a tuple with the variable index and the index for the constraint to be
used in `load_allocate_constraint`.
"""
function allocate_constrained_variable(
    model::MOI.ModelLike,
    set::MOI.AbstractScalarSet,
)
    variables = allocate_variables(model, 1)
    variable = variables[1]
    func = MOI.SingleVariable(variable)
    constraint = allocate_constraint(model, func, set)
    return variable, constraint
end

"""
    allocate_constrained_variables(model::MOI.ModelLike,
                                   set::MOI.AbstractVectorSet)

Returns a tuple with the variable indices and the index for the constraint to be
used in `load_allocate_constraint`.
"""
function allocate_constrained_variables(
    model::MOI.ModelLike,
    set::MOI.AbstractVectorSet,
)
    variables = allocate_variables(model, MOI.dimension(set))
    func = MOI.VectorOfVariables(variables)
    constraint = allocate_constraint(model, func, set)
    return variables, constraint
end

const ALLOCATE_LOAD_NOT_IMPLEMENTED = ErrorException(
    "The Allocate-Load interface is" * " not implemented by the model",
)

"""
    allocate(model::ModelLike, attr::ModelLikeAttribute, value)
    allocate(model::ModelLike, attr::AbstractVariableAttribute, v::VariableIndex, value)
    allocate(model::ModelLike, attr::AbstractConstraintAttribute, c::ConstraintIndex, value)

Informs `model` that `load` will be called with the same arguments after `load_variables` is called.
"""
function allocate(model::MOI.ModelLike, args...)
    return MOI.throw_set_error_fallback(
        model,
        args...;
        error_if_supported = ALLOCATE_LOAD_NOT_IMPLEMENTED,
    )
end

function allocate(
    model::MOI.ModelLike,
    attr::Union{MOI.AbstractVariableAttribute,MOI.AbstractConstraintAttribute},
    indices::Vector,
    values::Vector,
)
    for (index, value) in zip(indices, values)
        allocate(model, attr, index, value)
    end
end

"""
    allocate_constraint(model::MOI.ModelLike, f::MOI.AbstractFunction, s::MOI.AbstractSet)

Returns the index for the constraint to be used in `load_constraint` that will be called after `load_variables` is called.
"""
function allocate_constraint(
    model::MOI.ModelLike,
    func::MOI.AbstractFunction,
    set::MOI.AbstractSet,
)
    return MOI.throw_add_constraint_error_fallback(
        model,
        func,
        set;
        error_if_supported = ALLOCATE_LOAD_NOT_IMPLEMENTED,
    )
end

"""
    load_variables(model::MOI.ModelLike, nvars::Integer)

Prepares `model` for [`load`](@ref) and [`load_constraint`](@ref).
"""
function load_variables end

"""
    load_constrained_variable(
        model::MOI.ModelLike, vi::MOI.VariableIndex,
        ci::MOI.ConstraintIndex{MOI.SingleVariable},
        set::MOI.AbstractScalarSet)

Load the constrained variable `vi` to set `set` to `model`.
"""
function load_constrained_variable(
    model::MOI.ModelLike,
    vi::MOI.VariableIndex,
    ci::MOI.ConstraintIndex{MOI.SingleVariable},
    set::MOI.AbstractScalarSet,
)
    func = MOI.SingleVariable(vi)
    return load_constraint(model, ci, func, set)
end

"""
    load_constrained_variables(
        model::MOI.ModelLike, vi::MOI.VariableIndex,
        ci::MOI.ConstraintIndex{MOI.VectorOfVariables},
        set::MOI.AbstractVectorSet)

Load the constrained variable `vi` to set `set` to `model`.
"""
function load_constrained_variables(
    model::MOI.ModelLike,
    vis::Vector{MOI.VariableIndex},
    ci::MOI.ConstraintIndex{MOI.VectorOfVariables},
    set::MOI.AbstractVectorSet,
)
    func = MOI.VectorOfVariables(vis)
    return load_constraint(model, ci, func, set)
end

"""
    load(model::ModelLike, attr::ModelLikeAttribute, value)
    load(model::ModelLike, attr::AbstractVariableAttribute, v::VariableIndex, value)
    load(model::ModelLike, attr::AbstractConstraintAttribute, c::ConstraintIndex, value)

This has the same effect that `set` with the same arguments except that `allocate` should be called first before `load_variables`.
"""
function load(model::MOI.ModelLike, args...)
    return MOI.throw_set_error_fallback(
        model,
        args...;
        error_if_supported = ALLOCATE_LOAD_NOT_IMPLEMENTED,
    )
end

function load(
    model::MOI.ModelLike,
    attr::Union{MOI.AbstractVariableAttribute,MOI.AbstractConstraintAttribute},
    indices::Vector,
    values::Vector,
)
    for (index, value) in zip(indices, values)
        load(model, attr, index, value)
    end
end

"""
    load_single_variable(
        dest::MOI.ModelLike, src::MOI.ModelLike, idxmap::IndexMap,
        cis_src::Vector{MOI.ConstraintIndex{MOI.SingleVariable, S}}) where S

Load the constraints in `cis_src` from the model `src` into the model `dest`
using `load_constrained_variable`.
"""
function load_single_variable(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    idxmap::IndexMap,
    cis_src::Vector{MOI.ConstraintIndex{MOI.SingleVariable,S}},
) where {S}
    fs_src =
        MOI.get(
            src,
            MOI.ConstraintFunction(),
            cis_src,
        )::Vector{MOI.SingleVariable}
    sets = MOI.get(src, MOI.ConstraintSet(), cis_src)::Vector{S}
    for (ci_src, f_src, set) in zip(cis_src, fs_src, sets)
        vi_dest = idxmap[f_src.variable]
        ci_dest = idxmap[ci_src]
        load_constrained_variable(dest, vi_dest, ci_dest, set)
    end
end

"""
    load_vector_of_variables(
        dest::MOI.ModelLike, src::MOI.ModelLike, idxmap::IndexMap,
        cis_src::Vector{MOI.ConstraintIndex{MOI.VectorOfVariables, S}}) where S

Load the constraints in `cis_src` from the model `src` into the model `dest`
using `load_constrained_variable`.
"""
function load_vector_of_variables(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    idxmap::IndexMap,
    cis_src::Vector{MOI.ConstraintIndex{MOI.VectorOfVariables,S}},
) where {S}
    fs_src =
        MOI.get(
            src,
            MOI.ConstraintFunction(),
            cis_src,
        )::Vector{MOI.VectorOfVariables}
    sets = MOI.get(src, MOI.ConstraintSet(), cis_src)::Vector{S}
    for (ci_src, f_src, set) in zip(cis_src, fs_src, sets)
        vis_dest = [idxmap[vi] for vi in f_src.variables]
        ci_dest = idxmap[ci_src]
        load_constrained_variables(dest, vis_dest, ci_dest, set)
    end
end

"""
    load_constraint(model::MOI.ModelLike, ci::MOI.ConstraintIndex, f::MOI.AbstractFunction, s::MOI.AbstractSet)

Sets the constraint function and set for the constraint of index `ci`.
"""
function load_constraint(
    model::MOI.ModelLike,
    ::MOI.ConstraintIndex,
    func::MOI.AbstractFunction,
    set::MOI.AbstractSet,
)
    return MOI.throw_add_constraint_error_fallback(
        model,
        func,
        set;
        error_if_supported = ALLOCATE_LOAD_NOT_IMPLEMENTED,
    )
end

function allocate_constraints(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    idxmap::IndexMap,
    cis_src::Vector{<:MOI.ConstraintIndex},
)
    for ci_src in cis_src
        f_src = MOI.get(src, MOI.ConstraintFunction(), ci_src)
        s = MOI.get(src, MOI.ConstraintSet(), ci_src)
        f_dest = map_indices(idxmap, f_src)
        ci_dest = allocate_constraint(dest, f_dest, s)
        idxmap.conmap[ci_src] = ci_dest
    end
end

function load_constraints(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    idxmap::IndexMap,
    cis_src::Vector{<:MOI.ConstraintIndex},
)
    for ci_src in cis_src
        ci_dest = idxmap[ci_src]
        f_src = MOI.get(src, MOI.ConstraintFunction(), ci_src)
        f_dest = map_indices(idxmap, f_src)
        s = MOI.get(src, MOI.ConstraintSet(), ci_src)
        load_constraint(dest, ci_dest, f_dest, s)
    end
end

"""
    allocate_load(dest::MOI.ModelLike, src::MOI.ModelLike,
                  filter_constraints::Union{Nothing, Function}=nothing
                  )

Implements `MOI.copy_to(dest, src)` using the Allocate-Load API. The function
[`supports_allocate_load`](@ref) can be used to check whether `dest` supports
the Allocate-Load API.

If the `filter_constraints` arguments is given, only the constraints for which
this function returns `true` will be copied. This function is given a
constraint index as argument.
"""
function allocate_load(
    dest::MOI.ModelLike,
    src::MOI.ModelLike,
    copy_names::Bool,
    filter_constraints::Union{Nothing,Function} = nothing,
)
    MOI.empty!(dest)

    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    idxmap = index_map_for_variable_indices(vis_src)
    constraint_types = MOI.get(src, MOI.ListOfConstraints())
    single_variable_types =
        [S for (F, S) in constraint_types if F === MOI.SingleVariable]
    vector_of_variables_types =
        [S for (F, S) in constraint_types if F === MOI.VectorOfVariables]

    # Allocate variables
    vector_of_variables_types,
    vector_of_variables_allocated,
    vector_of_variables_not_allocated,
    single_variable_type,
    single_variable_allocated,
    single_variable_not_allocated = try_constrain_variables_on_creation(
        dest,
        src,
        idxmap,
        allocate_constrained_variables,
        allocate_constrained_variable,
    )

    copy_free_variables(dest, idxmap, vis_src, allocate_variables)

    # Allocate variable attributes
    pass_attributes(dest, src, copy_names, idxmap, vis_src, allocate)

    # Allocate model attributes
    pass_attributes(dest, src, copy_names, idxmap, allocate)

    # Allocate constraints
    pass_constraints(
        dest,
        src,
        copy_names,
        idxmap,
        single_variable_types,
        single_variable_not_allocated,
        vector_of_variables_types,
        vector_of_variables_not_allocated,
        allocate_constraints,
        allocate,
        filter_constraints = filter_constraints,
    )

    # Load variables
    load_variables(dest, length(vis_src))
    for cis_src in vector_of_variables_allocated
        load_vector_of_variables(dest, src, idxmap, cis_src)
    end
    for cis_src in single_variable_allocated
        load_single_variable(dest, src, idxmap, cis_src)
    end

    # Load variable attributes
    pass_attributes(dest, src, copy_names, idxmap, vis_src, load)

    # Load model attributes
    pass_attributes(dest, src, copy_names, idxmap, load)

    # Load constraints
    pass_constraints(
        dest,
        src,
        copy_names,
        idxmap,
        single_variable_types,
        single_variable_not_allocated,
        vector_of_variables_types,
        vector_of_variables_not_allocated,
        load_constraints,
        load,
        filter_constraints = filter_constraints,
    )

    return idxmap
end


text:
using Test
using OpenQuantumSystems
using Random, SparseArrays, LinearAlgebra, StableRNGs

import OrdinaryDiffEq

@testset "liouville" begin

    Random.seed!(StableRNG(0), 1)

    D(op1::Array, op2::Array) = abs(norm(op1 - op2))
    D(x1::StateVector, x2::StateVector) = norm(x2 - x1)
    D(op1::AbstractOperator, op2::AbstractOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))
    D(op1::AbstractSuperOperator, op2::AbstractSuperOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))

    mode1 = Mode(0.2, 1.0)
    mode2 = Mode(0.3, 2.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1], 3, [2.0, 200.0])
    mol2 = Molecule([mode2], 3, [3.0, 300.0])
    aggCore = AggregateCore([mol1, mol2])
    aggCore.coupling[2, 3] = 50
    aggCore.coupling[3, 2] = 50
    agg = setupAggregate(aggCore)
    aggTools = agg.tools

    Ham = agg.operators.Ham
    Ham_0 = agg.operators.Ham_0
    basis = agg.tools.basis

    ket0 = randstate(basis)
    W0 = dm(ket0)
    # tests have to be quick enough
    tspan = [0.0:0.1:1.0;]

    _, rho_int_t = LvN_sI(
        W0,
        tspan,
        agg;
        reltol = 1e-10,
        abstol = 1e-10,
        alg = OrdinaryDiffEq.Tsit5(),
    )
    for t_i = 1:length(tspan)
        t = tspan[t_i]
        U_op = evolutionOperator(Ham, t)
        W = U_op * W0 * U_op'
        U_0_op = evolutionOperator(Ham_0, t)
        W_int = U_0_op' * W * U_0_op
        rho_int = trace_bath(W_int, aggCore, aggTools)
        @test 1e-7 > D(rho_int, rho_int_t[t_i])
        # println(t_i, " ", D(W_int, W_int_t[t_i]))
    end

    _, rho_t = LvN_sS(
        W0,
        tspan,
        agg;
        reltol = 1e-10,
        abstol = 1e-10,
        alg = OrdinaryDiffEq.Tsit5(),
    )
    for t_i = 1:length(tspan)
        t = tspan[t_i]
        U_op = evolutionOperator(Ham, t)
        W = U_op * W0 * U_op'
        rho = trace_bath(W, aggCore, aggTools)
        @test 1e-7 > D(rho, rho_t[t_i])
        # println(t_i, " ", D(W_int, W_int_t[t_i]))
    end

    _, W_int_t = LvN_SI(
        W0,
        tspan,
        agg;
        reltol = 1e-10,
        abstol = 1e-10,
        alg = OrdinaryDiffEq.Tsit5(),
    )
    for t_i = 1:length(tspan)
        t = tspan[t_i]
        U_op = evolutionOperator(Ham, t)
        W = U_op * W0 * U_op'
        U_0_op = evolutionOperator(Ham_0, t)
        W_int = U_0_op' * W * U_0_op
        @test 1e-7 > D(W_int, W_int_t[t_i])
        # println(t_i, " ", D(W_int, W_int_t[t_i]))
    end

    _, W_t = LvN_SS(
        W0,
        tspan,
        agg;
        reltol = 1e-10,
        abstol = 1e-10,
        alg = OrdinaryDiffEq.Tsit5(),
    )
    for t_i = 1:length(tspan)
        t = tspan[t_i]
        U_op = evolutionOperator(Ham, t)
        W = U_op * W0 * U_op'
        @test 1e-7 > D(W, W_t[t_i])
        # println(t_i, " ", D(rho.data, rho_t[t_i].data))
    end

end # testset

max_stars_repo_name: detrin/OpenQuantumSystems

max_issues_repo_name: detrin/OpenQuantumSystems

max_forks_repo_name: detrin/OpenQuantumSystems