## Functions to handle input files.

using Chemfiles

include("clustering.jl")

"""
    parse_cif(file_path)

Parse a CIF file and return a dictionary where each identifier (without the
starting '_' character) is linked to its value.
Values are either a string or a vector of string (if defined in a loop).
"""
function parse_cif(file)
    all_data = Dict{String, Union{String, Vector{String}}}()
    inloop = false
    loopisspecified = false
    loopspec = String[]
    loop_n = 0

    l = read(file, String)
    i, j, x = nextword(l, 0)
    while i != 0

        if inloop
            if !loopisspecified
                if l[i] != '_' # This indicates the start of the values
                    loop_n = length(loopspec)
                    loopisspecified = true
                else # The identifier is part of the loop specification
                    push!(loopspec, l[i+1:j])
                    all_data[l[i+1:j]] = String[]
                    i, j, x = nextword(l, x); continue
                end
            end

            # From this point, the loop has been specified
            @assert loopisspecified
            if l[i] != '_' && l[i:j] != "loop_"
                for k in 1:loop_n
                    push!(all_data[loopspec[k]], l[i:j])
                    i, j, x = nextword(l, x)
                end
                continue
            end
            if l[i:j] == "loop_"
                loopisspecified = false
                loopspec = String[]
                i, j, x = nextword(l, x); continue
            end

            # This point can only be reached if we just quitted a loop
            inloop = false
        end

        @assert !inloop
        if l[i] == '_' # Simple identifier definition
            next_i, next_j, next_x = nextword(l, x)
            @assert next_i != 0
            all_data[l[i+1:j]] = l[next_i:next_j]
            x = next_x
        else
            if l[i:j] == "loop_"
                inloop = true
                loopisspecified = false
                loopspec = String[]
            elseif j-i > 4 && l[i:i+4] == "data_"
                @assert !haskey(all_data, "data")
                all_data["data"] = l[i+5:j]
            else
                k::Int = findprev(isequal('\n'), l, i)
                n::Int = count("\n", l[1:k])
                error("Unkown word \"$(l[i:j])\" at line $(n+1), position $(i-k):$(j-k)")
            end
        end

        i, j, x = nextword(l, x)
    end

    return all_data
end

function parsestrip(s)
    s = s[end] == ')' ? s[1:prevind(s, findlast('(', s))] : s
    return parse(BigFloat, s)
end

"""
    CIF(file_path::AbstractString)

Make a CIF object out of the parsed file.
"""
CIF(file::AbstractString) = CIF(parse_cif(file))
function CIF(parsed::Dict{String, Union{Vector{String},String}})
    natoms = length(parsed["atom_site_label"])
    equivalentpositions = pop!(parsed,
        haskey(parsed, "symmetry_equiv_pos_as_xyz") ?
            "symmetry_equiv_pos_as_xyz" : "space_group_symop_operation_xyz")
    refid = find_refid(equivalentpositions)
    cell = Cell(Symbol(haskey(parsed, "symmetry_cell_setting") ?
                            pop!(parsed, "symmetry_cell_setting") : ""),
                pop!(parsed, "symmetry_space_group_name_H-M"),
                haskey(parsed, "symmetry_Int_Tables_number") ?
                    parse(Int, pop!(parsed, "symmetry_Int_Tables_number")) : 0,
                parsestrip(pop!(parsed, "cell_length_a")),
                parsestrip(pop!(parsed, "cell_length_b")),
                parsestrip(pop!(parsed, "cell_length_c")),
                parsestrip(pop!(parsed, "cell_angle_alpha")),
                parsestrip(pop!(parsed, "cell_angle_beta")),
                parsestrip(pop!(parsed, "cell_angle_gamma")),
                parse.(EquivalentPosition, equivalentpositions, Ref(refid)))

    haskey(parsed, "symmetry_equiv_pos_site_id") && pop!(parsed, "symmetry_equiv_pos_site_id")
    removed_identity = false
    for i in eachindex(cell.equivalents)
        eq = cell.equivalents[i]
        if isone(eq.mat) && iszero(eq.ofs)
            deleteat!(cell.equivalents, i)
            removed_identity = true
            break
        end
    end
    @assert removed_identity

    labels =  pop!(parsed, "atom_site_label")
    _symbols = haskey(parsed, "atom_site_type_symbol") ?
                    pop!(parsed, "atom_site_type_symbol") : copy(labels)
    symbols = String[]
    @inbounds for x in _symbols
        i = findfirst(!isletter, x)
        push!(symbols, isnothing(i) ? x : x[1:i-1])
    end
    pos_x = pop!(parsed, "atom_site_fract_x")
    pos_y = pop!(parsed, "atom_site_fract_y")
    pos_z = pop!(parsed, "atom_site_fract_z")

    types = Symbol[]
    pos = Matrix{Float64}(undef, 3, natoms)
    correspondence = Dict{String, Int}()
    for i in 1:natoms
        @assert !haskey(correspondence, labels[i])
        correspondence[labels[i]] = i
        push!(types, Symbol(symbols[i]))
        pos[:,i] = parsestrip.([pos_x[i], pos_y[i], pos_z[i]])
        pos[:,i] .-= floor.(Int, pos[:,i])
    end

    invids = sortperm(types)
    types = types[invids]
    ids = invperm(invids)
    bonds = falses(natoms, natoms)
    if haskey(parsed, "geom_bond_atom_site_label_1") &&
       haskey(parsed, "geom_bond_atom_site_label_2")
        bond_a = pop!(parsed, "geom_bond_atom_site_label_1")
        bond_b = pop!(parsed, "geom_bond_atom_site_label_2")
        for i in 1:length(bond_a)
            x = correspondence[bond_a[i]]
            y = correspondence[bond_b[i]]
            bonds[x,y] = bonds[y,x] = 1
        end
    end

    return CIF(parsed, cell, ids, types, pos, bonds)
end

"""
    parse_arc(file)

Parse a .arc Systre archive such as the one used by the RCSR.
Return a pair `(flag, pairs)`.

`flag` is set if the archive corresponds to one generated by a compatible release
of CrystalNets. If unset, the genomes of the archive may not be the same as those
computed by CrystalNets for the same nets.
`pairs` is a `Dict{String,String}` whose entries have the form `genome => id` where `id`
is the name of the net and `genome` is the topological genome corresponding to this net
(given as a string of whitespace-separated values parseable by `PeriodicGraph`).
"""
function parse_arc(file)
    pairs = Dict{String,String}()
    curr_key = ""
    counter = 1
    firstline = readline(file)
    flag = startswith(firstline, "Made by CrystalNets.jl v")
    if flag
        flag = CRYSTAL_NETS_VERSION == VersionNumber(firstline[25:end])
    end
    for l in eachline(file)
        if length(l) > 3 && l[1:3] == "key"
            @assert isempty(curr_key)
            i = 4
            while isspace(l[i])
                i += 1
            end
            curr_key = l[i:end]
            @assert !haskey(pairs, curr_key)
        elseif length(l) > 2 && l[1:2] == "id"
            @assert !isempty(curr_key)
            i = 3
            while isspace(l[i])
                i += 1
            end
            pairs[curr_key] = l[i:end]
            curr_key = ""
        end
    end
    return flag, pairs
end

"""
    parse_arcs(file)

Parse a folder containing .arc Systre archives such as the one used by the RCSR.
Return a pair `(flag, pairs)` with the same convention than `parse_arc`
"""
function parse_arcs(path)
    combine(x, y) = x * ", " * y
    dict = Dict{String,String}()
    flag = true
    for f in readdir(path; sort=true)
        _flag, _dict = parse_arc(path*f)
        flag &= _flag
        mergewith!(combine, dict, _dict)
    end
    return flag, dict
end

function parse_atom_name(name::AbstractString)
    firstsep = findfirst(x -> ispunct(x) || isspace(x) || isnumeric(x), name)
    if firstsep isa Nothing
        return String(name)
    else
        firstsep::Int
        if firstsep != 1 && !any(isuppercase, @view name[nextind(name, firstsep):end])
            return String(name[1:prevind(firstsep)])
        else
            return String(name)
        end
    end
end

function parse_atom(name)
    atom = Chemfiles.Atom(name)
    set_type!(atom, parse_atom_name(name))
    return atom
end

function set_unique_bond_type!(frame::Frame, types, bond_length, bonded_atoms::Tuple{Symbol, Symbol}, onlykeep, tol)
    @ifwarn @info "To avoid guessing bonds, use a file format that contains the bonds."
    n = Int(size(frame))
    pos = Chemfiles.positions(frame)
    mat = Chemfiles.matrix(UnitCell(frame))'
    indices = [i for i in 1:n if types[i] ∈ onlykeep]
    #=@inbounds=# for _i in 1:length(indices)
        i = indices[_i]
        for _j in _i+1:length(indices)
            j = indices[_j]
            if minmax(types[i], types[j]) == bonded_atoms
                bonded = abs2(periodic_distance(pos[:,i], pos[:,j], mat) - bond_length) <= tol
                if bonded
                    Chemfiles.add_bond!(frame, i-1, j-1)
                end
            end
        end
    end
    nothing
end

function try_guess_bonds!(frame::Frame)
    # n = Int(size(frame))
    # types = Vector{Symbol}(undef, n)
    # for i in 1:n
    #     types[i] = Symbol(parse_atom_name(Chemfiles.type(Chemfiles.Atom(frame, i-1))))
    # end
    # unique_types = unique!(sort(types))
    # if unique_types == [:C] || unique_types == [:C, :H]
    #     if !(NOWARN::Bool)
    #         @warn "Guessing bonds. The structure seems to be made of only carbons (and possibly hydrogens): using an interatomic distance of 1.54±0.3 Å to assign edges between C atoms."
    #     end
    #     set_unique_bond_type!(frame, types, 1.54, (:C, :C), (:C,), 0.3)
    # elseif unique_types == [:O, :Si] || unique_types == [:Si]
    #     if !(NOWARN::Bool)
    #         @warn "Guessing bonds. The structure seems to be a zeolite: using an interatomic distance of 3.1±0.2 Å to assign edges between Si atoms."
    #     end
    #     set_unique_bond_type!(frame, types, 1.63, (:O, :Si), (:O, :Si), 0.1)
    # else
        @ifwarn begin
            @warn "Guessing bonds through Chemfiles. This may take a while for big structures and may be inexact."
            @info "To avoid guessing bonds, use a file format that contains the bonds."
        end
        guess_bonds!(frame)
    # end
end

function attribute_residues(topology, assert_use_existing_residues)
    m = Int(count_residues(topology))
    n = Int(size(topology))
    atoms_in_residues = m == 0 ? 0 : sum(length(atoms(Residue(topology, i))) for i in 0:(m-1))
    @assert atoms_in_residues <= n

    if atoms_in_residues < n
        if assert_use_existing_residues
            throw(ArgumentError("""
            Cannot use existing residues as vertices because not all atoms have an associated residue.
            To fix this, either assign a residue to each atom or provide another way to detect the vertices.
            """))
        end
        @ifwarn begin
            if atoms_in_residues > 0
                @warn "Some but not all atoms have an associated residue, so we cannot rely on existing residues"
            end
        end
        attributions = Int[]
    else
        attributions = zeros(Int, n)
        for i_residue in 1:m
            for atom in atoms(Residue(topology, i_residue-1))
                attributions[atom+1] = i_residue
            end
        end
    end
    return attributions
end

@static if !isdefined(Chemfiles, :atoms) # up to 0.9.3 included
    function atoms(residue::Residue)
        count = size(residue)
        result = Array{UInt64}(undef, count)
        Chemfiles.__check(Chemfiles.lib.chfl_residue_atoms(Chemfiles.__const_ptr(residue), pointer(result), count))
        return result
    end
end

function check_collision(pos, mat)
    _n, n = size(pos)
    @assert _n == 3
    invmat = inv(mat)
    toremove = Int[]
    for i in 1:n, j in (i+1):n
        if periodic_distance(invmat*pos[:,i], invmat*pos[:,j], mat) < 0.55
            push!(toremove, i)
            push!(toremove, j)
        end
    end
    tokeep = collect(1:n)
    if !isempty(toremove)
        @ifwarn @warn "This file contains multiple colliding atoms. All colliding atoms will be removed."
        unique!(sort!(toremove))
    end
    return toremove
end

hasHneighbor(types, graph, i) = any(x -> types[x.v] === :H, neighbors(graph,i))
function check_valence(types, graph)
    n = length(types)

    ## Small atoms valence check
    invalidatoms = Set{Symbol}()
    for i in 1:n
        d = degree(graph, i)
        t = types[i]
        if t === :O && d > 2 && (d > 3 || !hasHneighbor(types, graph, i))
            push!(invalidatoms, :O)
        elseif t === :C && d > 4
            push!(invalidatoms, :C)
        elseif t === :N && d > 3 && (d > 4 || !hasHneighbor(types, graph, i))
            push!(invalidatoms, :N)
        end
    end
    if !isempty(invalidatoms)
        s = String.(collect(invalidatoms))
        @ifwarn @warn "Found $(join(s, ',', " and ")) with too many bonds."
        return true
    end
    return false
end

"""
    @enum BondingMode

Selection mode for the detection of bonds. The choices are:
-   `InputBonds`: use the input bonds. Fail if those are not specified.
-   `ChemfilesBonds`: use chemfiles built-in bond detection mechanism.
-   `AutoBonds`: if the input specifies bonds, use them unless they look suspicious (too or too
    large according to a heuristic). Otherwise, fall back to `ChemfilesBonds`.
"""
@enum BondingMode begin
    InputBonds
    ChemfilesBonds
    AutoBonds
end

function sanity_checks!(pos, types, graph, mat, bondingmode)
    ## Bond length check
    removeedges = PeriodicEdge3D[]
    for e in edges(graph)
        s, d = e.src, e.dst.v
        bondlength = norm(pos[:,d] .+ (mat * e.dst.ofs) .- pos[:,s])
        if bondlength > 3
            @ifwarn @warn "Suspiciously large bond found."
            return true
        elseif bondlength < 0.85 && types[s] !== :H && types[d] !== :H
            push!(removeedges, e)
        end
    end
    if bondingmode == AutoBonds
        @ifwarn begin
            if !isempty(removeedges)
                @warn "Suspiciously small bond found. Such bonds are probably spurious and will be deleted."
                @info "To force retaining these bonds, use --bond-detect=input or --bond-detect=chemfiles"
            end
        end
        for e in removeedges
            rem_edge!(graph, e)
        end
    end
    return false
end

"""
       parse_chemfiles(path)

Parse a file given in any reckognised chemical format and extract the topological
information.
Such format can be .cif or any file format reckognised by Chemfiles.jl that
contains all the necessary topological information.
"""
function parse_chemfile(_path, exportto=tempdir(), bondingmode::BondingMode=AutoBonds, assert_use_existing_residues=false ; ignore_atoms=[])
    # Separate the cases unhandled by Chemfiles from the others
    path = expanduser(_path)
    cell = Cell()
    frame::Frame = if lowercase(last(splitext(path))) == ".cif"
        framecif = Frame()
        cif = expand_symmetry(CIF(path))
        a, b, c, α, β, γ = cell_parameters(cif.cell)
        set_cell!(framecif, UnitCell(a, b, c, α, β, γ))
        n = length(cif.ids)
        ignored = 0
        for i in 1:n
            typ = cif.types[cif.ids[i]]
            if typ ∈ ignore_atoms
                ignored += 1
                continue
            end
            pos = cif.cell.mat * cif.pos[:,i]
            atom = parse_atom(string(typ))
            add_atom!(framecif, atom, Vector{Float64}(pos))
        end
        n -= ignored
        if iszero(cif.bonds)
            try_guess_bonds!(framecif)
        else
            for i in 1:n, j in (i+1):n
                if cif.bonds[i,j]
                    add_bond!(framecif, i-1, j-1)
                end
            end
        end
        attributions = Int[]

        framecif
    else # The rest is handled by Chemfiles

        frameelse = read(Trajectory(path))
        for i in Int(size(frameelse))-1:-1:0
            if Symbol(type(Chemfiles.Atom(frameelse, i))) ∈ ignore_atoms
                Chemfiles.remove_atom!(frameelse, i)
            end
        end

        toremove = reverse(check_collision(positions(frameelse), matrix(UnitCell(frame))))
        for j in toremove
            Chemfiles.remove_atom!(frameelse, j-1)
        end

        topology = Topology(frameelse)
        n = Int(size(topology))

        if bonds_count(topology) == 0
            try_guess_bonds!(frameelse)
            topology = Topology(frameelse) # safer but useless since the underlying pointer is the same
        end

        attributions = attribute_residues(topology, assert_use_existing_residues)

        frameelse
    end

    cell = Cell(Cell(), SMatrix{3,3,BigFloat}(matrix(UnitCell(frame)))')
    if !all(isfinite, cell.mat) || iszero(det(cell.mat))
        @ifwarn @warn "Suspicious unit cell of matrix $(Float64.(cell.mat)). Is the input really periodic? Using a cubic unit cell instead"
        cell = Cell()
    end

    types = [Symbol(type(Chemfiles.Atom(frame, i))) for i in 0:(n-1)]
    poss = copy(positions(frame))

    adjacency = zeros(Bool, n, n)
    for (a,b) in eachcol(bonds(Topology(frame)))
        adjacency[a+1,b+1] = true
        adjacency[b+1,a+1] = true
    end
    mat = Float64.(cell.mat)
    graph = PeriodicGraph3D(n, edges_from_bonds(adjacency, mat, poss))

    if bondingmode != InputBonds
        bad_valence = check_valence(types, graph)
        recompute_bonds = bad_valence || sanity_checks!(poss, types, graph, mat, bondingmode)
        if recompute_bonds
            @ifwarn begin
                @warn "Disregarding all bonds from the input file."
                @info "To force retaining the initial bonds, use --bond-detect=input"
            end
            try_guess_bonds!(frame)
            topology = Topology(frame)
            adjacency = zeros(Bool, n, n)
            for (a,b) in eachcol(bonds(topology))
                adjacency[a+1,b+1] = true
                adjacency[b+1,a+1] = true
            end
            graph = PeriodicGraph3D(n, edges_from_bonds(adjacency, cell.mat, poss))
            @ifwarn if check_valence(types, graph)
                @warn "Remaining atoms with invalid valence. Proceeding anyway."
            end
            recompute_bonds = sanity_checks!(poss, types, graph, cell.mat, bondingmode)
            @ifwarn if recompute_bonds
                 @warn "Remaining bonds of suspicious lengths. Proceeding anyway."
            end
            attributions = attribute_residues(topology, assert_use_existing_residues)
        end
    end

    if isempty(attributions)
        crystalnothing = Crystal{Nothing}(cell, types, nothing, poss, graph)
        ifexport(crystalnothing, splitext(splitdir(path)[2])[1], exportto)
        return crystalnothing
    else
        crystalclusters = Crystal{Clusters}(cell, types, regroup_sbus(graph, attributions), poss, graph)
        ifexport(crystalclusters, splitext(splitdir(path)[2])[1], exportto)
        return crystalclusters
    end
end


export datachunk
export split_into_n, split_tod_mpi, get_chunk_properties

import Healpix
import CorrNoise
using Random
using FITSIO

try
    import MPI
catch
end

"""
This structure holds a number of parameters relative to a certain chunk of data.

Field              | Type           | Meaning
:----------------- |:-------------- |:------------------------------------------------------------
`pol_number`       | Int            | ID number of the polarimeter
`first_idx`        | Int            | Index of the first element of the chunk
`last_idx`         | Int            | Index of the last element of the chunk
`num_of_elements`  | Int            | Total number of elements of the chunk

# Example
given the following array of data: [1, 10, 20, 3, 4, 5, 7]
measured by the polarimeter number 2

the chunk: [20, 3, 4]
will correspond to the following structure:
datachunk(2, 3, 5, 3)
"""
struct datachunk
    pol_number ::Int
    first_idx::Int
    last_idx::Int
    num_of_elements::Int
end

"""
    function split_into_n(length, num_of_segments)
        Given the `length` of the array, it convenientely
        splits it into `num_of_segments` sections of as similar length as possible.
        It returns an array containing the number of elements of each section.

        # Example
        julia> split_into_n(20, 3)
        3-element Array{Int64,1}:
        6
        7
        7
"""
function split_into_n(length, num_of_segments)
    @assert num_of_segments >0
    @assert length >= num_of_segments
    start_pos = zeros(Int, num_of_segments+1)

    for i in 1:num_of_segments+1
        start_pos[i] =  floor(((i-1)*length/num_of_segments))
    end

    return start_pos[2:end]-start_pos[1:end-1]
end

"""
    function split_tod_mpi(total_time, baseline_length_s, baselines_per_process, num_of_MPI_proc)

        This function can be used to split the TOD production of many polarimeters among MPI processes.

        It requires in input:

        -the total time (in seconds) of the simulated observation
        -the length (in seconds) of each 1/f noise baseline
        -the array containing the number of 1/f baselines to simulate for each process.
            It can be obtained by using the function `plit_into_n` in the following way:

            split_into_n(num_of_polarimeters*baselines_per_pol, num_of_MPI_proc)

            where baselines_per_pol = Int64(total_time/baseline_length_s) is the number of baselines of each polarimeter

        -the number of MPI processes used

        It returns an array of arrays of `datachunk` instances, of length == num_of_MPI_proc
        where each element tells the chunk of data that each process should simulate (see Example)

        # Example
        ```julia-repl
        julia> num_of_polarimeters = 4
        julia> num_of_MPI_proc = 3
        julia> total_time = 50
        julia> baseline_length_s = 10
        julia> baselines_per_pol = Int64(total_time/baseline_length_s)
        5

        julia> baselines_per_process = split_into_n(20, 3)
        3-element Array{Int64,1}:
        6
        7
        7

        julia> chunks = split_tod_mpi(total_time, baseline_length_s, baselines_per_process, num_of_MPI_proc)
        3-element Array{Any,1}:
        Any[datachunk(1, 1, 5, 5), datachunk(2, 1, 1, 1)]
        Any[datachunk(2, 2, 5, 4), datachunk(3, 1, 3, 3)]
        Any[datachunk(3, 4, 5, 2), datachunk(4, 1, 5, 5)]
        ```

        which means:

        - process number 0 should simulate:
        polarimeter number 1 from baseline 1 to baseline 5,  total number of baselines = 5
        polarimeter number 2 from baseline 1 to baseline 1 , total number of baselines = 1

        - process number 1 should simulate:
        polarimeter number 2 from baseline 2 to baseline 5,  total number of baselines = 4
        polarimeter number 3 from baseline 1 to baseline 3 , total number of baselines = 3

        - process number 2 should simulate:
        polarimeter number 3 from baseline 4 to baseline 5,  total number of baselines = 2
        polarimeter number 4 from baseline 1 to baseline 5,  total number of baselines = 5

"""
function split_tod_mpi(total_time, baseline_length_s, baselines_per_process, num_of_MPI_proc)

    duration = Int64(total_time/baseline_length_s)

    #initialization
    detector_num = 1
    sample_idx = 0
    samples_in_det = duration
    result = []

    for rank_num in 0:(num_of_MPI_proc-1)  #loop on MPI processes

        samples_for_this_process = baselines_per_process[rank_num+1]
        samples_left = samples_for_this_process
        data_this_rank = []

        while samples_left > 0  #loop on detectors

            #if the current detector has more samples than needed to fill the current MPI process
            if samples_in_det > samples_left

                first_idx = sample_idx+1
                last_idx = sample_idx+samples_left
                data = datachunk(detector_num, first_idx, last_idx, samples_left)
                data_this_rank = append!(data_this_rank, [data])

                sample_idx = sample_idx + samples_left
                samples_in_det = samples_in_det - samples_left
                samples_left = 0

            #if the current detector has not enough samples to provide the current MPI process
            #with the required number of samples. In this case we need to increase "detector_num" before the next iteration
            else

                first_idx = sample_idx+1
                last_idx = sample_idx+samples_in_det
                data = datachunk(detector_num, first_idx, last_idx, samples_in_det)
                data_this_rank = append!(data_this_rank, [data])

                samples_left = samples_left - samples_in_det
                samples_in_det = 0
            end

            if samples_in_det == 0
                detector_num +=1
                sample_idx = 0
                samples_in_det = duration
            end

        end

        result = append!(result, [data_this_rank])

    end

    return result

end

"""
    function get_chunk_properties(chunks, baseline_length_s, fsamp_hz, rank)

        Given:
        - the data chunks (which can be obtained by using the function `split_tod_mpi`)
        - the length (in seconds) of each 1/f noise baseline
        - the sampling frequency (in Hz)
        - the number of current MPI rank

        this function extracts useful information to perform the TOD simulation in the current rank.

        It returns a tuple containing 5 arrays:
        - the ID number of the polarimeters that the current rank will simulate
        - the start time of the acquisition portion for each polarimeter
        - the stop time of the acquisition portion for each polarimeter
        - the number of 1/f baselines for each polarimeter
        - the total number of samples for each polarimeter

        # Example
        ```julia-repl
        julia> baseline_length_s = 10
        julia> fsamp_hz = 10
        julia> rank = 1

        julia> chunks = [[datachunk(1, 1, 5, 5), datachunk(2, 1, 1, 1)], [datachunk(2, 2, 5, 4), datachunk(3, 1, 3, 3)], [datachunk(3, 4, 5, 2), datachunk(4, 1, 5, 5)]]

        get_chunk_properties(chunks, baseline_length_s, fsamp_hz, rank)
        ([2, 3], [10.0, 0.0], [50.0, 30.0], [4, 3], [400, 300])

        which means that rank 1 will simulate:
        - polarimeter number 2 from 10 s (from the starting of the acquisition) to 50 s,
          with a total of 4 1/f baselines and 400 samples.
        - polarimeter number 3 from 0 s (from the starting of the acquisition) to 30 s,
          with a total of 3 1/f baselines and 300 samples.

    """
    function get_chunk_properties(chunks, baseline_length_s, fsamp_hz, rank)

        this_rank_chunk = chunks[rank+1]
        first_time, last_time = [Array{Float64}(undef, length(this_rank_chunk)) for i in (1:2)]
        detector_number, num_of_baselines, baseline_len, num_of_samples = [Array{Int64}(undef, length(this_rank_chunk)) for i in (1:4)]

        for i in 1:length(this_rank_chunk)
            detector_number[i] = this_rank_chunk[i].pol_number
            first_time[i] = (this_rank_chunk[i].first_idx-1)*baseline_length_s
            last_time[i] = this_rank_chunk[i].last_idx*baseline_length_s -0.99*(1/fsamp_hz)
            num_of_baselines[i] =  this_rank_chunk[i].num_of_elements
            num_of_samples[i] = num_of_baselines[i]*baseline_length_s*fsamp_hz
        end
        return (detector_number, first_time, last_time, num_of_baselines, num_of_samples)
    end


##############################################
######### Define all processing ##############
##############################################

#This is an abstract function
"""
It will generally be called as
       new_timed_events = process_event(time, state, event)
It will generate 0 or more new timed events based on the current event
"""
function process_event end

function process_event(time::Float64, state::State, es_event::EndSimEvent)::Vector{TimedEvent}
    println("Ending simulation at time $time.")
    return []
end

function process_event(time::Float64, state::State, ls_event::LogStateEvent)::Vector{TimedEvent}
    println("Logging state at time $time.")
    println(state)
    return []
end;

"""
Process an arrival event adding a new job to the system

Returns all events created in the process
"""
function process_event(time::Float64, state::State, arrival_event::ExternalArrivalEvent)::Vector{TimedEvent}
    new_timed_events = TimedEvent[]

    q = get_entry_queue(state)
    if !is_full(q, state)
        new_job(q, time, state)
        #if this is the only job on the server engage service
        num_in_queue(q, state) == 1 && push!(new_timed_events,
                                    TimedEvent(EndOfServiceAtQueueEvent(q), time + next_service_time(state, q)))
    else
        #Finds new queue using overflow matrix
        next_q = get_next_queue(q, state, :overflow)
        if next_q > 0
            transit_time = time + travel_time(state)
            new_transit(time, transit_time, state)
            push!(new_timed_events, TimedEvent(InTransitEvent(next_q), transit_time))
        else
            failed_arrival(time, state)
        end
    end

    #prepare next arrival
    push!(new_timed_events, TimedEvent(ExternalArrivalEvent(), time + next_arrival_time(state)))
    return new_timed_events
end

"""
Process an end of service event by sending the next job in the specified queue either to transit to its next destination
or out of the sysytem.

Returns all events created in the process
"""
function process_event(time::Float64, state::State, eos_event::EndOfServiceAtQueueEvent)::Vector{TimedEvent}
    q = eos_event.q
    new_timed_events = TimedEvent[]

    #if another customer in the queue then start a new service
    if num_in_queue(q, state) > 1
        push!(new_timed_events, TimedEvent(EndOfServiceAtQueueEvent(q), time + next_service_time(state, q)))
    end

    next_q = get_next_queue(q, state, :routing)
    if next_q > 0
        transit_time = time + travel_time(state)
        queue_to_transit(q, time, transit_time, state)
        push!(new_timed_events, TimedEvent(InTransitEvent(next_q), transit_time))
    else
        pop_queue(q, time, state)
    end
    return new_timed_events
end

"""
Process a transit event by removing the next job from transit and attemping to add it to its destination queue

Returns all events created in the process
"""
function process_event(time::Float64, state::State, transit_event::InTransitEvent)::Vector{TimedEvent}
    new_timed_events = TimedEvent[]

    q = transit_event.q
    if !is_full(q, state)
        transit_to_queue(q, time, state)
        #if this is the only job on the server engage service
        num_in_queue(q, state) == 1 && push!(new_timed_events,
                                    TimedEvent(EndOfServiceAtQueueEvent(q), time + next_service_time(state, q)))
    else
        #Finds new queue using overflow matrix
        next_q = get_next_queue(q, state, :overflow)
        if next_q > 0
            transit_time = time + travel_time(state)
            update_transit(time, transit_time, state)
            push!(new_timed_events, TimedEvent(InTransitEvent(next_q), transit_time))
        else
            pop_transit(time, state)
        end
    end
    return new_timed_events
end

""" simulate(scenario, [job_tracking, max_time, log_times, call_back])

Runs a simulation of the given scenario until the max_time is reached.
The simulation can be run with, no job tracking (:none) as default, only tracking job entry and exit times (:times) or
tracking full job history (:full).

An overview of the system state will be printed at the provided log_times.
The call_back function will be called each time an event occurs.
"""
function simulate(scenario::NetworkParameters;
                    job_tracking::Symbol = :none,
                    max_time::Real = 10.0,
                    log_times::Vector{Float64} = Float64[],
                    call_back = (time,state) -> clear_left(state))
    # Initializes the the state depeding on simulation mode
    if job_tracking == :none
        state = NetworkState(scenario)
    elseif job_tracking == :times
        state = TrackedState(scenario)
    elseif job_tracking == :full
        state = FullTrackedState(scenario)
    else
        error("$mode is not a valid simulation mode")
    end

    init_timed_event = TimedEvent(ExternalArrivalEvent(), 0.0)

    #The event queue
    priority_queue = BinaryMinHeap{TimedEvent}()

    #Put the standard events in the queue
    push!(priority_queue, init_timed_event)
    push!(priority_queue, TimedEvent(EndSimEvent(), max_time))
    for lt in log_times
        push!(priority_queue,TimedEvent(LogStateEvent(), lt))
    end

    time = 0.0
    call_back(time, state)

    #The main discrete event simulation loop
    while true
        #Get the next event
        timed_event = pop!(priority_queue)

        #advance the time
        time = timed_event.time

        #Act on the event
        new_timed_events = process_event(time, state, timed_event.event)

        #if the event was an end of simulation then stop
        isa(timed_event.event, EndSimEvent) && break

        #The event may spawn 0 or more events which we put in the priority queue
        for new_event in new_timed_events
            push!(priority_queue, new_event)
        end

        call_back(time, state)
    end
end


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# VectorOfDecisions #
# ========================== #
struct VectorDecisionConstraintBridge{T, S <: MOI.AbstractVectorSet} <: MOIB.Constraint.AbstractBridge
    constraint::CI{MOI.VectorOfVariables, S}
    decisions::VectorOfDecisions
end

function MOIB.Constraint.bridge_constraint(::Type{VectorDecisionConstraintBridge{T,S}},
                                           model,
                                           f::VectorOfDecisions,
                                           set::S) where {T, S <: MOI.AbstractVectorSet}
    # Perform the bridge mapping manually
    mapped_variables = map(f.decisions) do decision
        bridged = MOIB.bridged_function(
            model,
            AffineDecisionFunction{T}(SingleDecision(decision)))
        bridged.decision_part.terms[1].variable_index
    end
    # Add the bridged constraint
    constraint = MOI.add_constraint(model,
                                    MOI.VectorOfVariables(mapped_variables),
                                    set)
    # Save the constraint indices and the decision to allow modifications
    return VectorDecisionConstraintBridge{T,S}(constraint, f)
end

function MOI.supports_constraint(::Type{<:VectorDecisionConstraintBridge{T}},
                                 ::Type{VectorOfDecisions},
                                 ::Type{<:MOI.AbstractVectorSet}) where T
    return true
end
function MOIB.added_constrained_variable_types(::Type{<:VectorDecisionConstraintBridge})
    return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{<:VectorDecisionConstraintBridge{T,S}}) where {T,S}
    return [(MOI.VectorOfVariables, S), (MOI.ScalarAffineFunction{T}, MOI.EqualTo{T})]
end
function MOIB.Constraint.concrete_bridge_type(::Type{<:VectorDecisionConstraintBridge{T}},
                              ::Type{VectorOfDecisions},
                              S::Type{<:MOI.AbstractVectorSet}) where T
    return VectorDecisionConstraintBridge{T,S}
end

MOI.get(b::VectorDecisionConstraintBridge{T,S}, ::MOI.NumberOfConstraints{MOI.VectorOfVariables, S}) where {T,S} = 1
MOI.get(b::VectorDecisionConstraintBridge{T}, ::MOI.ListOfConstraintIndices{MOI.VectorOfVariables, S}) where {T,S} = [b.constraint]

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 bridge::VectorDecisionConstraintBridge)
    return bridge.decisions
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 bridge::VectorDecisionConstraintBridge)
    return MOI.get(model, MOI.ConstraintSet(), bridge.constraint)
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintPrimal,
                 bridge::VectorDecisionConstraintBridge)
    return MOI.get(model, attr, bridge.constraint)
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintDual,
                 bridge::VectorDecisionConstraintBridge)
    return MOI.get(model, attr, bridge.constraint)
end

function MOI.get(model::MOI.ModelLike, ::DecisionIndex,
                 bridge::VectorDecisionConstraintBridge{T}) where T
    return bridge.constraint
end

function MOI.delete(model::MOI.ModelLike, bridge::VectorDecisionConstraintBridge)
    MOI.delete(model, bridge.constraint)
    return nothing
end

function MOI.set(model::MOI.ModelLike, ::MOI.ConstraintSet,
                 bridge::VectorDecisionConstraintBridge{T,S}, change::S) where {T,S}
    MOI.set(model, MOI.ConstraintSet(), bridge.constraint, change)
    return nothing
end


#///////////////////////////////////////
#// File Name: RiskSensitiveSAC.jl
#// Author: Haruki Nishimura (hnishimura@stanford.edu)
#// Date Created: 2020/12/04
#// Description: Julia package for Risk Sensitive Sequential Action
#// Control, with public version of Trajectron++
#///////////////////////////////////////

module RiskSensitiveSAC

import Base: ==, isequal, isapprox
import PlotUtils: plot_color
import RobotOS: Time, Duration
import StatsBase: fit, Histogram
import StatsFuns: logsumexp

using CUDA
using DataStructures
using Distributions
using ForwardDiff
using LinearAlgebra
using Plots
pyplot();
using Printf
using ProgressMeter
using PyCall
using Random

# Modules for Buffered Input Cell and Related Functions (and constrained SAC optimizations)
import Convex: Variable, norm, quadform, minimize, dot, solve!
using GeometricalPredicates
using LazySets
using Polyhedra
using SCS
using ECOS
using VoronoiCells

function __init__()
    @info "Number of Julia Thread(s): $(Threads.nthreads())";
    @info "CUDA Device: $(name(CuDevice(0)))"
    @info "Python executable used by PyCall: $(pyimport("sys").executable)";

    # Append the code directory to python path
    # (see https://github.com/JuliaPy/PyCall.jl#troubleshooting)
    file_path = @__FILE__;
    pushfirst!(PyVector(pyimport("sys")."path"),
               py"'/'.join($file_path.split('/')[:-1]) + '/../Trajectron-plus-plus/trajectron'");

    # Define python functions for Trajectron bridge
    py"""
    import os
    import time
    import json
    import torch
    import dill
    import random
    import pathlib
    import pandas as pd
    # import evaluation
    import numpy as np
    import visualization as vis
    from easydict import EasyDict
    from model.online.online_trajectron import OnlineTrajectron
    from model.model_registrar import ModelRegistrar
    from environment import Environment, Scene, Node

    # For CrowdNav Baseline
    import configparser
    import crowd_nav
    import gym
    from crowd_nav.policy.policy_factory import policy_factory
    from crowd_sim.envs.utils.robot import Robot
    from crowd_sim.envs.utils.state import ObservableState

    def load_hyperparams(model_dir, conf_file):
        '''
        param model_dir: directory where the learned model resides
        param conf_file: config file name
        '''
        # Load hyperparameters from json
        config_file = os.path.join(model_dir, conf_file)
        if not os.path.exists(config_file):
            raise ValueError('Config json not found!')
        with open(config_file, 'r') as conf_json:
            hyperparams = json.load(conf_json)
        return hyperparams

    def load_online_scene(data_dir, data_name, maximum_history_length, state_def,
                         scene_idx, incl_robot_node, init_timestep=2,
                         ado_id_removed=None):

        '''
        param data_dir: directory to look in for evaluation data
        param data_name: name of the evaluation data
        param maximum_history_length: maximum history length considered in the model
        param state_def: state definition
        param scene_idx: scene_idx (starting from 0)
        param incl_robot_node: whether to include the robot node in the scene
        param init_timestep: init_timestep (starting form 0). You need to have at least acceleration, so you want 2 timesteps of prior data, e.g. [0, 1],
                             so that you can immediately start incremental inference from the 3rd timestep onwards.
        param ado_id_removed: ado_id in string to be removed from online_scene. Note this ado will not be removed from eval_scene
        '''
        log = []
        eval_data_path = os.path.join(data_dir, data_name)
        with open(eval_data_path, 'rb') as f:
            eval_env = dill.load(f, encoding='latin1')
        log.append('Loaded evaluation data from %s' % (eval_data_path,))
        if eval_env.robot_type is None and incl_robot_node:
             eval_env.robot_type = eval_env.NodeType[0]  # TODO: Make more general, allow the user to specify?

             x = [np.nan]
             y = [np.nan]
             vx = [np.nan]
             vy = [np.nan]
             ax = [np.nan]
             ay = [np.nan]
             data_dict = {('position', 'x'): x,
                          ('position', 'y'): y,
                          ('velocity', 'x'): vx,
                          ('velocity', 'y'): vy,
                          ('acceleration', 'x'): ax,
                          ('acceleration', 'y'): ay}
             data_columns = pd.MultiIndex.from_product([['position', 'velocity', 'acceleration'], ['x', 'y']])
             node_data = pd.DataFrame(data_dict, columns=data_columns)
             robot_node = Node(eval_env.robot_type, 'ROBOT', node_data, is_robot=True, first_timestep=init_timestep)
             for scene in eval_env.scenes:
                 #scene.add_robot_from_nodes(eval_env.robot_type)
                 scene.robot = robot_node
                 scene.nodes.append(scene.robot)

        assert 0 <= scene_idx < len(eval_env.scenes), "scene_idx must be between 0 and {}!".format(len(eval_env.scenes) - 1)
        eval_scene = eval_env.scenes[scene_idx]

        online_scene = Scene(timesteps=init_timestep,
                             map=eval_scene.map,
                             dt=eval_scene.dt)
        online_scene.nodes = eval_scene.get_nodes_clipped_at_time(timesteps=np.arange(init_timestep + 1 - maximum_history_length,
                                                                                init_timestep),
                                                                  state=state_def)
        online_scene.robot = eval_scene.robot
        node_to_remove = None
        if ado_id_removed is not None:
            for node in online_scene.nodes:
                if str(node) == ado_id_removed:
                    node_to_remove = node
        if node_to_remove is not None:
            online_scene.nodes.remove(node_to_remove)

        return eval_env, eval_scene, online_scene, log

    def initialize_model(args, model_dir, conf_file):
        '''
        param args: easydict object for parsing arguments
        param model_dir: directory where the learned model resides
        param confi_file: config file name
        '''
        log = []
        # Load hyperparameters from json
        hyperparams = load_hyperparams(model_dir, conf_file)

        if args.seed is not None:
            random.seed(args.seed)
            np.random.seed(args.seed)
            torch.manual_seed(args.seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed_all(args.seed)

        # Add hyperparams from arguments
        hyperparams['dynamic_edges'] = args.dynamic_edges
        hyperparams['edge_state_combine_method'] = args.edge_state_combine_method
        hyperparams['edge_influence_combine_method'] = args.edge_influence_combine_method
        hyperparams['edge_addition_filter'] = args.edge_addition_filter
        hyperparams['edge_removal_filter'] = args.edge_removal_filter
        hyperparams['offline_scene_graph'] = args.offline_scene_graph
        hyperparams['incl_robot_node'] = args.incl_robot_node
        hyperparams['edge_encoding'] = not args.no_edge_encoding
        hyperparams['use_map_encoding'] = args.map_encoding

        iter_num = 100
        model_registrar = ModelRegistrar(model_dir, args.eval_device)
        model_registrar.load_models(iter_num=iter_num)
        log.append('Loaded Trajectron model from %s' % (os.path.join(model_dir, 'model_registrar-{}.pt'.format(iter_num))))
        trajectron = OnlineTrajectron(model_registrar, hyperparams,
                                  args.eval_device)
        return trajectron, hyperparams, log

    def configure_rl_robot(model_dir, env_config_path, policy_config_path,
                           policy_name='sarl'):
        env_config_file = os.path.join(model_dir, os.path.basename(env_config_path))
        policy_config_file = os.path.join(model_dir, os.path.basename(policy_config_path))
        model_weights = os.path.join(model_dir, 'rl_model.pth')
        device = torch.device('cpu')

        # configure policy
        policy = policy_factory[policy_name]()
        policy_config = configparser.RawConfigParser()
        policy_config.read(policy_config_file)
        policy.configure(policy_config)
        if policy.trainable:
            policy.get_model().load_state_dict(torch.load(model_weights))
        phase = 'test'
        policy.set_phase(phase)
        policy.set_device(device)
        if hasattr(policy, 'query_env'):
            policy.query_env = False

        # configure environment
        env_config = configparser.RawConfigParser()
        env_config.read(env_config_file)

        robot = Robot(env_config, 'robot')
        robot.set_policy(policy)

        # additional setup (from crowd_sim.CrowdSim.reset)
        time_step = env_config.getfloat('env', 'time_step')
        robot.time_step = time_step
        robot.policy.time_step = time_step
        if robot.sensor == 'RGB':
            raise NotImplementedError

        return robot
    """
end

# Type Definition
export
    Parameter,
    State,
    WayPoint2D,
    Trajectory2D,
    get_position,
    RobotState,
    get_velocity,
    WorldState
include("type_definition.jl")

# State Transition
export
    transition,
    transition_jacobian,
    transition_control_coeff
    #sample_future_positions
include("state_transition.jl")

# Cost
export
    instant_position_cost,
    instant_control_cost,
    instant_collision_cost,
    terminal_position_cost,
    terminal_collision_cost,
    instant_position_cost_gradient,
    instant_collision_cost_gradient,
    terminal_position_cost_gradient,
    terminal_collision_cost_gradient,
    CostParameter
include("cost.jl")

# Cost (CUDA kernels)
export
    kernel_instant_position_cost!,
    kernel_instant_control_cost!,
    kernel_instant_collision_cost!,
    kernel_terminal_position_cost!,
    kernel_terminal_collision_cost!,
    kernel_instant_position_cost_gradient!,
    kernel_instant_collision_cost_gradient!,
    kernel_terminal_position_cost_gradient!,
    kernel_terminal_collision_cost_gradient!
include("cost_gpu.jl")

# Scene Loader
export
    fetch_ado_positions!,
    convert_nodes_to_str,
    SceneLoader,
    TrajectronSceneParameter,
    TrajectronSceneLoader,
    reduce_to_positions,
    get_trajectory_for_ado,
    SyntheticSceneParameter,
    SyntheticSceneLoader
include("scene_loader.jl")

# Predictor
export
    initialize_scene_graph!,
    sample_future_ado_positions!,
    Predictor,
    TrajectronPredictorParameter,
    TrajectronPredictor,
    OraclePredictorParameter,
    OraclePredictor,
    GaussianPredictorParameter,
    GaussianPredictor
include("predictor.jl")

# Forward-Backward Simulation
export
    get_measurement_schedule,
    get_target_pos_array,
    process_u_arrays,
    simulate_forward,
    sample_ado_positions,
    process_ap_dict,
    compute_costs,
    integrate_costs,
    choose_best_nominal_control,
    compute_cost_gradients,
    sum_cost_gradients,
    simulate_backward,
    simulate,
    # evaluate_risk,
    SimulationParameter,
    SimulationCostResult,
    SimulationCostGradResult,
    SimulationResult
include("forward_backward_simulation.jl")

# Risk Sensitive SAC Controller
export
    control!,
    adjust_old_prediction!,
    schedule_prediction!,
    schedule_control_update!,
    convert_to_schedule,
    get_nominal_u_arrays,
    get_robot_present_and_future,
    get_control_coeffs,
    solve_multi_qcqp,
    get_control_schedule,
    determine_control_time,
    sac_control_update,
    ControlParameter,
    ControlSchedule,
    RSSACController
include("rs_sac_controller.jl")

# Buffered Input Cells and Related Functions (benchmark)
export
    vertex_direction,
    sort_vertices,
    scale_and_shift,
    find_extrema,
    retract_cell,
    get_input_positions,
    compute_buffered_voronoi_cells,
    compute_buffered_input_cell,
    project,
    lq_tracking,
    BICControlParameter,
    BICController,
    bic_control_update,
    bic_parameter_setup
include("bic.jl")

# CrowdNav RL Controller (benchmark)
export
    CrowdNavControlParameter,
    CrowdNavController,
    get_action!,
    crowdnav_control_update!,
    control!,
    schedule_control_update!
include("crowd_nav_controller.jl")

# Evaluation Functions
export
    evaluate,
    EvaluationResult,
    BICEvaluationResult
include("evaluation.jl")

# Helper Functions
export
    get_nominal_trajectory,
    init_condition_setup,
    controller_setup,
    display_log,
    visualize!,
    make_gif,
    fetch_stats_filtered,
    plot_histogram
include("utils.jl")

end # module


# VectorOfDecisions #
# ========================== #
struct VectorDecisionConstraintBridge{T, S <: MOI.AbstractVectorSet} <: MOIB.Constraint.AbstractBridge
    constraint::CI{MOI.VectorOfVariables, S}
    fixing_constraints::Vector{FixingConstraint{T}}
    decisions::VectorOfDecisions
end

function MOIB.Constraint.bridge_constraint(::Type{VectorDecisionConstraintBridge{T,S}},
                                           model,
                                           f::VectorOfDecisions,
                                           set::S) where {T, S <: MOI.AbstractVectorSet}
    # Perform the bridge mapping manually
    mapped_variables = map(f.decisions) do decision
        bridged = MOIB.bridged_function(
            model,
            AffineDecisionFunction{T}(SingleDecision(decision)))
        bridged.decision_part.terms[1].variable_index
    end
    # Add the bridged constraint
    constraint = MOI.add_constraint(model,
                                    MOI.VectorOfVariables(mapped_variables),
                                    set)
    # Save the constraint indices and the decision to allow modifications
    return VectorDecisionConstraintBridge{T,S}(constraint, fill(FixingConstraint{T}(0), MOI.dimension(set)), f)
end

function MOI.supports_constraint(::Type{<:VectorDecisionConstraintBridge{T}},
                                 ::Type{VectorOfDecisions},
                                 ::Type{<:MOI.AbstractVectorSet}) where T
    return true
end
function MOIB.added_constrained_variable_types(::Type{<:VectorDecisionConstraintBridge})
    return Tuple{DataType}[]
end
function MOIB.added_constraint_types(::Type{<:VectorDecisionConstraintBridge{T,S}}) where {T,S}
    return [(MOI.VectorOfVariables, S), (MOI.ScalarAffineFunction{T}, MOI.EqualTo{T})]
end
function MOIB.Constraint.concrete_bridge_type(::Type{<:VectorDecisionConstraintBridge{T}},
                              ::Type{VectorOfDecisions},
                              S::Type{<:MOI.AbstractVectorSet}) where T
    return VectorDecisionConstraintBridge{T,S}
end

MOI.get(b::VectorDecisionConstraintBridge{T,S}, ::MOI.NumberOfConstraints{MOI.VectorOfVariables, S}) where {T,S} = 1
MOI.get(b::VectorDecisionConstraintBridge{T}, ::MOI.NumberOfConstraints{MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}}) where T =
    count(fc -> fc.value != 0, b.fixing_constraints)
MOI.get(b::VectorDecisionConstraintBridge{T}, ::MOI.ListOfConstraintIndices{MOI.VectorOfVariables, S}) where {T,S} = [b.constraint]
MOI.get(b::VectorDecisionConstraintBridge{T}, ::MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{T}, MOI.EqualTo{T}}) where {T} =
    filter(fc -> fc.value != 0, b.fixing_constraints)

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintFunction,
                 bridge::VectorDecisionConstraintBridge)
    return bridge.decisions
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintSet,
                 bridge::VectorDecisionConstraintBridge)
    return MOI.get(model, MOI.ConstraintSet(), bridge.constraint)
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintPrimal,
                 bridge::VectorDecisionConstraintBridge)
    return MOI.get(model, attr, bridge.constraint)
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintDual,
                 bridge::VectorDecisionConstraintBridge)
    return MOI.get(model, attr, bridge.constraint)
end

function MOI.delete(model::MOI.ModelLike, bridge::VectorDecisionConstraintBridge)
    MOI.delete(model, bridge.constraint)
    return nothing
end

function MOI.set(model::MOI.ModelLike, ::MOI.ConstraintSet,
                 bridge::VectorDecisionConstraintBridge{T,S}, change::S) where {T,S}
    MOI.set(model, MOI.ConstraintSet(), bridge.constraint, change)
    return nothing
end

function MOI.modify(model::MOI.ModelLike, bridge::VectorDecisionConstraintBridge{T,S}, change::DecisionStateChange) where {T,S}
    i = findfirst(d -> d.decision == change.decision, bridge.decisions.decisions)
    # Switch on state transition
    if change.new_state == NotTaken
        if bridge.fixing_constraints[i].value != 0
            # Remove the fixing constraint
            MOI.delete(model, bridge.fixing_constraints[i])
            bridge.fixing_constraints[i] = FixingConstraint{T}(0)
        end
    end
    if change.new_state == Taken
        if bridge.fixing_constraints[i].value != 0
            # Remove any existing fixing constraint
            MOI.delete(model, bridge.fixing_constraints[i])
        end
        # Perform the bridge mapping manually
        aff = MOIB.bridged_function(model, AffineDecisionFunction{T}(bridge.decisions.decisions[i]))
        f = MOI.ScalarAffineFunction{T}(MOI.SingleVariable(aff.decision_part.terms[1].variable_index))
        # Get the decision value
        set = MOI.EqualTo(aff.decision_part.constant)
        # Add a fixing constraint to ensure that fixed decision is feasible.
        bridge.fixing_constraints[i] = MOI.add_constraint(model, f, set)
    end
    return nothing
end


function polar_kernel(n::Int, nlines::Int, line_start::Int, scale::Float64,
                     u_curr::CuDeviceArray{Float64,1}, v_curr::CuDeviceArray{Float64,1},
                     l_curr::CuDeviceArray{Float64,1}, rho::CuDeviceArray{Float64,1},
                     shift_lines::Int, param::CuDeviceArray{Float64,2},
                     _YffR::CuDeviceArray{Float64,1}, _YffI::CuDeviceArray{Float64,1},
                     _YftR::CuDeviceArray{Float64,1}, _YftI::CuDeviceArray{Float64,1},
                     _YttR::CuDeviceArray{Float64,1}, _YttI::CuDeviceArray{Float64,1},
                     _YtfR::CuDeviceArray{Float64,1}, _YtfI::CuDeviceArray{Float64,1},
                     frBound::CuDeviceArray{Float64,1}, toBound::CuDeviceArray{Float64,1})

    tx = threadIdx().x
    ty = threadIdx().y
    I = blockIdx().x
    id_line = I + shift_lines

    x = @cuDynamicSharedMem(Float64, n)
    xl = @cuDynamicSharedMem(Float64, n, n*sizeof(Float64))
    xu = @cuDynamicSharedMem(Float64, n, (2*n)*sizeof(Float64))

    @inbounds begin
        YffR = _YffR[id_line]; YffI = _YffI[id_line]
        YftR = _YftR[id_line]; YftI = _YftI[id_line]
        YttR = _YttR[id_line]; YttI = _YttI[id_line]
        YtfR = _YtfR[id_line]; YtfI = _YtfI[id_line]

        pij_idx = line_start + 8*(I-1)

        xl[1] = sqrt(frBound[2*(id_line-1)+1])
        xu[1] = sqrt(frBound[2*id_line])
        xl[2] = sqrt(toBound[2*(id_line-1)+1])
        xu[2] = sqrt(toBound[2*id_line])
        xl[3] = -2*pi
        xu[3] = 2*pi
        xl[4] = -2*pi
        xu[4] = 2*pi

        x[1] = min(xu[1], max(xl[1], sqrt(u_curr[pij_idx+4])))
        x[2] = min(xu[2], max(xl[2], sqrt(u_curr[pij_idx+5])))
        x[3] = min(xu[3], max(xl[3], u_curr[pij_idx+6]))
        x[4] = min(xu[4], max(xl[4], u_curr[pij_idx+7]))

        param[1,id_line] = l_curr[pij_idx]
        param[2,id_line] = l_curr[pij_idx+1]
        param[3,id_line] = l_curr[pij_idx+2]
        param[4,id_line] = l_curr[pij_idx+3]
        param[5,id_line] = l_curr[pij_idx+4]
        param[6,id_line] = l_curr[pij_idx+5]
        param[7,id_line] = l_curr[pij_idx+6]
        param[8,id_line] = l_curr[pij_idx+7]
        param[9,id_line] = rho[pij_idx]
        param[10,id_line] = rho[pij_idx+1]
        param[11,id_line] = rho[pij_idx+2]
        param[12,id_line] = rho[pij_idx+3]
        param[13,id_line] = rho[pij_idx+4]
        param[14,id_line] = rho[pij_idx+5]
        param[15,id_line] = rho[pij_idx+6]
        param[16,id_line] = rho[pij_idx+7]
        param[17,id_line] = v_curr[pij_idx]
        param[18,id_line] = v_curr[pij_idx+1]
        param[19,id_line] = v_curr[pij_idx+2]
        param[20,id_line] = v_curr[pij_idx+3]
        param[21,id_line] = v_curr[pij_idx+4]
        param[22,id_line] = v_curr[pij_idx+5]
        param[23,id_line] = v_curr[pij_idx+6]
        param[24,id_line] = v_curr[pij_idx+7]

        CUDA.sync_threads()

        status, minor_iter = tron_kernel(n, shift_lines, 500, 200, 1e-6, scale, true, x, xl, xu,
                                         param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)

        vi_vj_cos = x[1]*x[2]*cos(x[3] - x[4])
        vi_vj_sin = x[1]*x[2]*sin(x[3] - x[4])

        u_curr[pij_idx] = YffR*x[1]^2 + YftR*vi_vj_cos + YftI*vi_vj_sin
        u_curr[pij_idx+1] = -YffI*x[1]^2 - YftI*vi_vj_cos + YftR*vi_vj_sin
        u_curr[pij_idx+2] = YttR*x[2]^2 + YtfR*vi_vj_cos - YtfI*vi_vj_sin
        u_curr[pij_idx+3] = -YttI*x[2]^2 - YtfI*vi_vj_cos - YtfR*vi_vj_sin
        u_curr[pij_idx+4] = x[1]^2
        u_curr[pij_idx+5] = x[2]^2
        u_curr[pij_idx+6] = x[3]
        u_curr[pij_idx+7] = x[4]
    end

    return
end

function polar_kernel_cpu(n::Int, nline::Int, line_start::Int, scale::Float64,
                          u_curr::AbstractVector{Float64}, v_curr::AbstractVector{Float64},
                          l_curr::AbstractVector{Float64}, rho::AbstractVector{Float64},
                          shift::Int,
                          param::Array{Float64},
                          YffR::Array{Float64}, YffI::Array{Float64},
                          YftR::Array{Float64}, YftI::Array{Float64},
                          YttR::Array{Float64}, YttI::Array{Float64},
                          YtfR::Array{Float64}, YtfI::Array{Float64},
                          frBound::Array{Float64}, toBound::Array{Float64})
    avg_minor_it = 0

    x = zeros(n)
    xl = zeros(n)
    xu = zeros(n)

    xl[3] = -2*pi
    xu[3] = 2*pi
    xl[4] = -2*pi
    xu[4] = 2*pi

    @inbounds for I=1:nline
        pij_idx = line_start + 8*(I-1)
        id_line = shift + I

        xl[1] = sqrt(frBound[2*(id_line-1)+1])
        xu[1] = sqrt(frBound[2*id_line])
        xl[2] = sqrt(toBound[2*(id_line-1)+1])
        xu[2] = sqrt(toBound[2*id_line])

        x[1] = min(xu[1], max(xl[1], sqrt(u_curr[pij_idx+4])))
        x[2] = min(xu[2], max(xl[2], sqrt(u_curr[pij_idx+5])))
        x[3] = min(xu[3], max(xl[3], u_curr[pij_idx+6]))
        x[4] = min(xu[4], max(xl[4], u_curr[pij_idx+7]))

        param[1,id_line] = l_curr[pij_idx]
        param[2,id_line] = l_curr[pij_idx+1]
        param[3,id_line] = l_curr[pij_idx+2]
        param[4,id_line] = l_curr[pij_idx+3]
        param[5,id_line] = l_curr[pij_idx+4]
        param[6,id_line] = l_curr[pij_idx+5]
        param[7,id_line] = l_curr[pij_idx+6]
        param[8,id_line] = l_curr[pij_idx+7]
        param[9,id_line] = rho[pij_idx]
        param[10,id_line] = rho[pij_idx+1]
        param[11,id_line] = rho[pij_idx+2]
        param[12,id_line] = rho[pij_idx+3]
        param[13,id_line] = rho[pij_idx+4]
        param[14,id_line] = rho[pij_idx+5]
        param[15,id_line] = rho[pij_idx+6]
        param[16,id_line] = rho[pij_idx+7]
        param[17,id_line] = v_curr[pij_idx]
        param[18,id_line] = v_curr[pij_idx+1]
        param[19,id_line] = v_curr[pij_idx+2]
        param[20,id_line] = v_curr[pij_idx+3]
        param[21,id_line] = v_curr[pij_idx+4]
        param[22,id_line] = v_curr[pij_idx+5]
        param[23,id_line] = v_curr[pij_idx+6]
        param[24,id_line] = v_curr[pij_idx+7]

        function eval_f_cb(x)
            f = eval_f_polar_kernel_cpu(id_line, scale, x, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return f
        end

        function eval_g_cb(x, g)
            eval_grad_f_polar_kernel_cpu(id_line, scale, x, g, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        function eval_h_cb(x, mode, rows, cols, _scale, lambda, values)
            eval_h_polar_kernel_cpu(id_line, x, mode, scale, rows, cols, lambda, values,
                                    param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        nele_hess = 10
        tron = ExaTron.createProblem(4, xl, xu, nele_hess, eval_f_cb, eval_g_cb, eval_h_cb;
                                     :tol => 1e-6, :matrix_type => :Dense, :max_minor => 200,
                                     :frtol => 1e-12)

        tron.x .= x
        status = ExaTron.solveProblem(tron)
        x .= tron.x
        avg_minor_it += tron.minor_iter

        cos_ij = cos(x[3] - x[4])
        sin_ij = sin(x[3] - x[4])
        vi_vj_cos = x[1]*x[2]*cos_ij
        vi_vj_sin = x[1]*x[2]*sin_ij

        u_curr[pij_idx] = YffR[id_line]*x[1]^2 + YftR[id_line]*vi_vj_cos + YftI[id_line]*vi_vj_sin
        u_curr[pij_idx+1] = -YffI[id_line]*x[1]^2 - YftI[id_line]*vi_vj_cos + YftR[id_line]*vi_vj_sin
        u_curr[pij_idx+2] = YttR[id_line]*x[2]^2 + YtfR[id_line]*vi_vj_cos - YtfI[id_line]*vi_vj_sin
        u_curr[pij_idx+3] = -YttI[id_line]*x[2]^2 - YtfI[id_line]*vi_vj_cos - YtfR[id_line]*vi_vj_sin
        u_curr[pij_idx+4] = x[1]^2
        u_curr[pij_idx+5] = x[2]^2
        u_curr[pij_idx+6] = x[3]
        u_curr[pij_idx+7] = x[4]
    end

    return 0, avg_minor_it / nline
end

function polar_kernel_two_level(
    n::Int, nline::Int, line_start::Int, bus_start::Int, scale::Float64,
    u::CuDeviceArray{Float64}, xbar::CuDeviceArray{Float64},
    z::CuDeviceArray{Float64}, l::CuDeviceArray{Float64}, rho::CuDeviceArray{Float64},
    shift_lines::Int, param::CuDeviceArray{Float64},
    _YffR::CuDeviceArray{Float64}, _YffI::CuDeviceArray{Float64},
    _YftR::CuDeviceArray{Float64}, _YftI::CuDeviceArray{Float64},
    _YttR::CuDeviceArray{Float64}, _YttI::CuDeviceArray{Float64},
    _YtfR::CuDeviceArray{Float64}, _YtfI::CuDeviceArray{Float64},
    frBound::CuDeviceArray{Float64}, toBound::CuDeviceArray{Float64},
    brBusIdx::CuDeviceArray{Int}
)
    I = blockIdx().x
    id_line = I + shift_lines

    x = @cuDynamicSharedMem(Float64, n)
    xl = @cuDynamicSharedMem(Float64, n, n*sizeof(Float64))
    xu = @cuDynamicSharedMem(Float64, n, (2*n)*sizeof(Float64))

    @inbounds begin
        YffR = _YffR[id_line]; YffI = _YffI[id_line]
        YftR = _YftR[id_line]; YftI = _YftI[id_line]
        YttR = _YttR[id_line]; YttI = _YttI[id_line]
        YtfR = _YtfR[id_line]; YtfI = _YtfI[id_line]

        pij_idx = line_start + 8*(I-1)
        xbar_pij_idx = line_start + 4*(I-1)

        xl[1] = sqrt(frBound[2*(id_line-1)+1])
        xu[1] = sqrt(frBound[2*id_line])
        xl[2] = sqrt(toBound[2*(id_line-1)+1])
        xu[2] = sqrt(toBound[2*id_line])
        xl[3] = -2*pi
        xu[3] = 2*pi
        xl[4] = -2*pi
        xu[4] = 2*pi

        x[1] = min(xu[1], max(xl[1], sqrt(u[pij_idx+4])))
        x[2] = min(xu[2], max(xl[2], sqrt(u[pij_idx+5])))
        x[3] = min(xu[3], max(xl[3], u[pij_idx+6]))
        x[4] = min(xu[4], max(xl[4], u[pij_idx+7]))

        fr_idx = bus_start + 2*(brBusIdx[2*(id_line-1)+1] - 1)
        to_idx = bus_start + 2*(brBusIdx[2*id_line] - 1)

        param[1,id_line] = l[pij_idx]
        param[2,id_line] = l[pij_idx+1]
        param[3,id_line] = l[pij_idx+2]
        param[4,id_line] = l[pij_idx+3]
        param[5,id_line] = l[pij_idx+4]
        param[6,id_line] = l[pij_idx+5]
        param[7,id_line] = l[pij_idx+6]
        param[8,id_line] = l[pij_idx+7]
        param[9,id_line] = rho[pij_idx]
        param[10,id_line] = rho[pij_idx+1]
        param[11,id_line] = rho[pij_idx+2]
        param[12,id_line] = rho[pij_idx+3]
        param[13,id_line] = rho[pij_idx+4]
        param[14,id_line] = rho[pij_idx+5]
        param[15,id_line] = rho[pij_idx+6]
        param[16,id_line] = rho[pij_idx+7]
        param[17,id_line] = xbar[xbar_pij_idx] - z[pij_idx]
        param[18,id_line] = xbar[xbar_pij_idx+1] - z[pij_idx+1]
        param[19,id_line] = xbar[xbar_pij_idx+2] - z[pij_idx+2]
        param[20,id_line] = xbar[xbar_pij_idx+3] - z[pij_idx+3]
        param[21,id_line] = xbar[fr_idx] - z[pij_idx+4]
        param[22,id_line] = xbar[to_idx] - z[pij_idx+5]
        param[23,id_line] = xbar[fr_idx+1] - z[pij_idx+6]
        param[24,id_line] = xbar[to_idx+1] - z[pij_idx+7]

        CUDA.sync_threads()

        status, minor_iter = tron_kernel(n, shift_lines, 500, 200, 1e-6, scale, true, x, xl, xu,
                                         param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)

        vi_vj_cos = x[1]*x[2]*cos(x[3] - x[4])
        vi_vj_sin = x[1]*x[2]*sin(x[3] - x[4])

        u[pij_idx] = YffR*x[1]^2 + YftR*vi_vj_cos + YftI*vi_vj_sin
        u[pij_idx+1] = -YffI*x[1]^2 - YftI*vi_vj_cos + YftR*vi_vj_sin
        u[pij_idx+2] = YttR*x[2]^2 + YtfR*vi_vj_cos - YtfI*vi_vj_sin
        u[pij_idx+3] = -YttI*x[2]^2 - YtfI*vi_vj_cos - YtfR*vi_vj_sin
        u[pij_idx+4] = x[1]^2
        u[pij_idx+5] = x[2]^2
        u[pij_idx+6] = x[3]
        u[pij_idx+7] = x[4]
    end

    return
end

function polar_kernel_two_level_cpu(
    n::Int, nline::Int, line_start::Int, bus_start::Int, scale::Float64,
    u::AbstractVector{Float64}, xbar::AbstractVector{Float64},
    z::AbstractVector{Float64}, l::AbstractVector{Float64}, rho::AbstractVector{Float64},
    shift::Int,
    param::Array{Float64},
    YffR::Array{Float64}, YffI::Array{Float64},
    YftR::Array{Float64}, YftI::Array{Float64},
    YttR::Array{Float64}, YttI::Array{Float64},
    YtfR::Array{Float64}, YtfI::Array{Float64},
    frBound::Array{Float64}, toBound::Array{Float64},
    brBusIdx::Array{Int}
)
    avg_minor_it = 0

    x = zeros(n)
    xl = zeros(n)
    xu = zeros(n)

    xl[3] = -2*pi
    xu[3] = 2*pi
    xl[4] = -2*pi
    xu[4] = 2*pi

    @inbounds for I=1:nline
        pij_idx = line_start + 8*(I-1)
        xbar_pij_idx = line_start + 4*(I-1)
        id_line = shift + I

        xl[1] = sqrt(frBound[2*(id_line-1)+1])
        xu[1] = sqrt(frBound[2*id_line])
        xl[2] = sqrt(toBound[2*(id_line-1)+1])
        xu[2] = sqrt(toBound[2*id_line])

        x[1] = min(xu[1], max(xl[1], sqrt(u[pij_idx+4])))
        x[2] = min(xu[2], max(xl[2], sqrt(u[pij_idx+5])))
        x[3] = min(xu[3], max(xl[3], u[pij_idx+6]))
        x[4] = min(xu[4], max(xl[4], u[pij_idx+7]))

        fr_idx = bus_start + 2*(brBusIdx[2*(id_line-1)+1] - 1)
        to_idx = bus_start + 2*(brBusIdx[2*id_line] - 1)

        param[1,id_line] = l[pij_idx]
        param[2,id_line] = l[pij_idx+1]
        param[3,id_line] = l[pij_idx+2]
        param[4,id_line] = l[pij_idx+3]
        param[5,id_line] = l[pij_idx+4]
        param[6,id_line] = l[pij_idx+5]
        param[7,id_line] = l[pij_idx+6]
        param[8,id_line] = l[pij_idx+7]
        param[9,id_line] = rho[pij_idx]
        param[10,id_line] = rho[pij_idx+1]
        param[11,id_line] = rho[pij_idx+2]
        param[12,id_line] = rho[pij_idx+3]
        param[13,id_line] = rho[pij_idx+4]
        param[14,id_line] = rho[pij_idx+5]
        param[15,id_line] = rho[pij_idx+6]
        param[16,id_line] = rho[pij_idx+7]
        param[17,id_line] = xbar[xbar_pij_idx] - z[pij_idx]
        param[18,id_line] = xbar[xbar_pij_idx+1] - z[pij_idx+1]
        param[19,id_line] = xbar[xbar_pij_idx+2] - z[pij_idx+2]
        param[20,id_line] = xbar[xbar_pij_idx+3] - z[pij_idx+3]
        param[21,id_line] = xbar[fr_idx] - z[pij_idx+4]
        param[22,id_line] = xbar[to_idx] - z[pij_idx+5]
        param[23,id_line] = xbar[fr_idx+1] - z[pij_idx+6]
        param[24,id_line] = xbar[to_idx+1] - z[pij_idx+7]

        function eval_f_cb(x)
            f = eval_f_polar_kernel_cpu(id_line, scale, x, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return f
        end

        function eval_g_cb(x, g)
            eval_grad_f_polar_kernel_cpu(id_line, scale, x, g, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        function eval_h_cb(x, mode, rows, cols, _scale, lambda, values)
            eval_h_polar_kernel_cpu(id_line, x, mode, scale, rows, cols, lambda, values,
                                    param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        nele_hess = 10
        tron = ExaTron.createProblem(4, xl, xu, nele_hess, eval_f_cb, eval_g_cb, eval_h_cb;
                                     :tol => 1e-6, :matrix_type => :Dense, :max_minor => 200,
                                     :frtol => 1e-12)

        tron.x .= x
        status = ExaTron.solveProblem(tron)
        x .= tron.x
        avg_minor_it += tron.minor_iter

        vi_vj_cos = x[1]*x[2]*cos(x[3] - x[4])
        vi_vj_sin = x[1]*x[2]*sin(x[3] - x[4])

        u[pij_idx] = YffR[id_line]*x[1]^2 + YftR[id_line]*vi_vj_cos + YftI[id_line]*vi_vj_sin
        u[pij_idx+1] = -YffI[id_line]*x[1]^2 - YftI[id_line]*vi_vj_cos + YftR[id_line]*vi_vj_sin
        u[pij_idx+2] = YttR[id_line]*x[2]^2 + YtfR[id_line]*vi_vj_cos - YtfI[id_line]*vi_vj_sin
        u[pij_idx+3] = -YttI[id_line]*x[2]^2 - YtfI[id_line]*vi_vj_cos - YtfR[id_line]*vi_vj_sin
        u[pij_idx+4] = x[1]^2
        u[pij_idx+5] = x[2]^2
        u[pij_idx+6] = x[3]
        u[pij_idx+7] = x[4]
    end

    return 0, avg_minor_it / nline
end


module OBJFileIO

export Shapeobject, loadOBJFile, createCylinderVertices

############### Begin Type declaration for Shape Object ################
type Shapeobject
    #A new type declaration. This is a shape object that is composed of
    #An array of vertices and an array of faces
    vertices::Array{Float64, 2}
    faces::Array{Int64, 2}

    #Constructor that creates an object from an obj file
    function Shapeobject(filename::String)
        #Create instance of the object
        shapeOBJ = new()
        #Populate the object fields using the loadOBJFile method
        shapeOBJ.vertices, shapeOBJ.faces = loadOBJFile(filename)
        #Return object
        return shapeOBJ
    end #end constructor

    #Constructor that creates an object from an obj file
    function Shapeobject(cylDiamter::Float64, cylLength::Float64,
                         numVertices::Int64, filename::String)
        #Create instance of the object
        shapeOBJ = new()
        #Populate the object.faces field using the loadOBJFile method
        x, shapeOBJ.faces = loadOBJFile(filename)
        #Populate the object.vertices field using the createCylinderVertices method
        shapeOBJ.vertices = createCylinderVertices(cylDiamter, cylLength, numVertices)
        #Return object
        return shapeOBJ
    end #end constructor

end # end of type declaration

############### Begin Functions ###############

#FUNCTION
#SUMMARY:
#This function loads an obj file and returns an array of vertices and an array of faces
#
#INPUTS:
#    objFile::String -
#    Path to the obj file containing the object that you want
#
#OUTPUTS:
#    Tuple of arrays
#
#    objectVertices::Array{Float64,2} -
#    array containing the vertices of the object
#
#    objectFaces::Array{Float64,2} -
#    array containing the faces of the object
function loadOBJFile(objFile::String)
    #Initiate variables for the object vertices and faces
    objectVertices = zeros(Float64,1,3)
    objectFaces = zeros(Int64,1,3)

    #Initialise variables to temporarily store the data for each line.
    lineVertex = zeros(Float64,1,3)
    lineFace = zeros(Int64,1,3)

    #Initialise counter for the vertex and face line number
    vertexLineCounter = 1
    faceLineCounter = 1

    #Open OBJ file for reading
    objFileContent = open(objFile)

    #Loop through each line of the file and look for vertex lines and face lines
    #Those lines begin with a 'v' or and 'f' respectively.
    #Foor those lines store the values in the repsective objectvertices or object faces
    #array.
    for line in eachline(objFileContent)
        if line[1] == 'v'
            vertexLine = split(line)
            #Initialise counter for the vertex line number
            for i = 1:length(vertexLine)
                if i != 1
                    if vertexLineCounter == 1
                        objectVertices[vertexLineCounter,i-1] = parsefloat(vertexLine[i])
                    else
                        lineVertex[i-1] = parsefloat(vertexLine[i])
                    end
                end
            end
            if vertexLineCounter != 1
                #Concatenate the vertex from the parsed line to the object vertices array
                objectVertices = vcat(objectVertices,lineVertex)
            end
            #Increment counter
            vertexLineCounter += 1

        elseif line[1] == 'f'
            faceLine = split(line)
            for i = 1:length(faceLine)
                if i != 1
                    if faceLineCounter == 1
                        objectFaces[faceLineCounter,i-1] = parseint(faceLine[i])
                    else
                        lineFace[i-1] = parsefloat(faceLine[i])
                    end
                end
            end
            if faceLineCounter != 1
                #Concatenate the face from the parsed line to the object vertices array
                objectFaces = vcat(objectFaces,lineFace)
            end
            #Increment counter
            faceLineCounter += 1
        end
    end

    #Close the file now we are done with it
    close(objFileContent)

    #return Object vertices and faces
    return objectVertices, objectFaces
end #end loadOBJFile function

#FUNCTION
#SUMMARY:
#Takes in parameters and returns a cylinder object from those parameters
#
#INPUTS:
#    cylDiameter::Float64 -
#    Diameter of cylinder
#
#    cylLength::Float64 -
#    Axial length of the cylinder
#
#    numVertices::Int64 -
#    number of vertices around each end of cylinder i.e. this is half of the total number
#    of vertices in the cylinder object
#
#OUTPUTS:
#    vertices::Array{Float64,2}
#    An 2*numVertices x 3 array containing the coordinates of each vertex of the cylinder
function createCylinderVertices(cylDiamter::Float64, cylLength::Float64, numVertices::Int64)
    #Calculate radius
    cylRadius = cylDiamter/2

    #Create xCoordinates
    cylBase = cylLength/2
    xCoordBase = -cylBase
    xCoordTop = cylBase

    #Calculate angular step around circle
    angleToVertex = -2π/numVertices

    #Preallocate vertices array
    vertices = zeros(Float64, 2 * numVertices, 3)

    #Loop through each vertex
    for vertex = 1:numVertices
        #Calculate points around the circle
        yCoord = cylRadius * cos((vertex - 1) * angleToVertex)
        ZCoord = cylRadius * sin((vertex - 1) * angleToVertex)

        #Add points to the vertices array for the Base
        vertices[2 * vertex - 1, 1] = xCoordBase
        vertices[2 * vertex - 1, 2] = yCoord
        vertices[2 * vertex - 1, 3] = ZCoord

        #Add points to the vertices array for the Top
        vertices[2 * vertex, 1] = xCoordTop
        vertices[2 * vertex, 2] = yCoord
        vertices[2 * vertex, 3] = ZCoord
    end

    #return vertices
    return vertices
end #end createCylinderVertices method

end #end module


abstract type CropMethod end
struct RedMax <: CropMethod end

function isredmax(p::T) where T <: ColorTypes.AbstractRGB
    pmax = max(p.r, p.g, p.b)
    # disallow masked and overexposed pixels, assume plants reflect better than dirt
    pmax == 1 && return false
    return p.r == pmax
end

gapfraction(pixs::AbstractArray, thresh::RedMax) = sum(isredmax(p) for p in pixs) / length(pixs)


function cumulative_probability(obj::RealDistribution, arg0::jdouble)
    return jcall(obj, "cumulativeProbability", jdouble, (jdouble,), arg0)
end

function density(obj::RealDistribution, arg0::jdouble)
    return jcall(obj, "density", jdouble, (jdouble,), arg0)
end

function get_numerical_mean(obj::RealDistribution)
    return jcall(obj, "getNumericalMean", jdouble, ())
end

function get_numerical_variance(obj::RealDistribution)
    return jcall(obj, "getNumericalVariance", jdouble, ())
end

function get_support_lower_bound(obj::RealDistribution)
    return jcall(obj, "getSupportLowerBound", jdouble, ())
end

function get_support_upper_bound(obj::RealDistribution)
    return jcall(obj, "getSupportUpperBound", jdouble, ())
end

function inverse_cumulative_probability(obj::RealDistribution, arg0::jdouble)
    return jcall(obj, "inverseCumulativeProbability", jdouble, (jdouble,), arg0)
end

function is_support_connected(obj::RealDistribution)
    return jcall(obj, "isSupportConnected", jboolean, ())
end

function log_density(obj::RealDistribution, arg0::jdouble)
    return jcall(obj, "logDensity", jdouble, (jdouble,), arg0)
end

function probability(obj::RealDistribution, arg0::jdouble, arg1::jdouble)
    return jcall(obj, "probability", jdouble, (jdouble, jdouble), arg0, arg1)
end


"""implementation of Simultaneous Sampling-and-Search Planning (SSSP)"""
module LibSSSP
export SSSP

import Printf: @sprintf, @printf
import Base: @kwdef
import ...MRMP: AbsState, Node, now, elapsed_sec, gen_uniform_sampling, get_mid_status, dist
import DataStructures: PriorityQueue, enqueue!, dequeue!
import ..Solvers: gen_g_func, get_distance_tables, get_distance_table

@kwdef mutable struct SuperNode{State<:AbsState}
    Q::Vector{Node{State}}  # set of search nodes
    next::Int64  # next agent, 0 -> fixed agents
    id::String = get_Q_id(Q, next)
    parent_id::Union{Nothing,String} = nothing  # parent node
    g::Float64 = 0.0  # g-value
    h::Float64 = 0.0  # h-value
    f::Float64 = g + h  # f-value
    depth::Int64 = 1  # depth
end

"""
    SSSP(
        config_init::Vector{State},
        config_goal::Vector{State},
        connect::Function,
        collide::Function,
        check_goal::Function;
        g_func::Function = gen_g_func(greedy = true),
        steering_depth::Int64 = 2,
        num_vertex_expansion::Int64 = 10,
        init_min_dist_thread::Float64 = 0.1,
        decreasing_rate_min_dist_thread::Float64 = 0.99,
        epsilon::Union{Float64,Nothing} = nothing,
        TIME_LIMIT::Union{Nothing,Real} = 30,
        VERBOSE::Int64 = 0,

        use_random_h_func::Bool = false,         # for ablation study
        no_roadmap_at_beginning::Bool = false,   # for ablation study

        no_fast_collision_check::Bool = false,   # use 'slow' collision checker
    )::Tuple{
        Union{Nothing,Vector{Vector{Node{State}}}},  # solution
        Vector{Vector{Node{State}}},  # roadmap
    } where {State<:AbsState}

implementation of SSSP
"""
function SSSP(
    config_init::Vector{State},
    config_goal::Vector{State},
    connect::Function,
    collide::Function,
    check_goal::Function;
    g_func::Function = gen_g_func(greedy = true),
    steering_depth::Int64 = 2,
    num_vertex_expansion::Int64 = 10,
    init_min_dist_thread::Float64 = 0.1,
    decreasing_rate_min_dist_thread::Float64 = 0.99,
    epsilon::Union{Float64,Nothing} = nothing,
    TIME_LIMIT::Union{Nothing,Real} = 30,
    VERBOSE::Int64 = 0,

    # for ablation study
    use_random_h_func::Bool = false,
    no_roadmap_at_beginning::Bool = false,

    # just for fairness of experiments
    no_fast_collision_check::Bool = false,
)::Tuple{
    Union{Nothing,Vector{Vector{Node{State}}}},  # solution
    Vector{Vector{Node{State}}},  # roadmap
} where {State<:AbsState}
    """get initial solution / refinement"""

    t_s = now()
    elapsed() = elapsed_sec(t_s)
    timeover() = TIME_LIMIT != nothing && elapsed() > TIME_LIMIT

    # number of agents
    N = length(config_init)

    # define sampler
    sampler = gen_uniform_sampling(config_init[1])

    # used in steering
    conn(q_from::State, q_to::State, i::Int64) = begin
        (isnothing(epsilon) || dist(q_from, q_to) <= epsilon) && connect(q_from, q_to, i)
    end

    # get initial roadmap by RRT-connect
    roadmaps = (
        no_roadmap_at_beginning ?
        map(
            i -> begin
                v_init = Node{State}(config_init[i], 1, [])
                v_goal = Node{State}(config_goal[i], 2, [])
                conn(v_init.q, v_goal.q, i) && push!(v_init.neighbors, v_goal.id)
                conn(v_goal.q, v_init.q, i) && push!(v_goal.neighbors, v_init.id)
                [v_init, v_goal]
            end,
            1:N,
        ) :
        gen_RRT_connect_roadmaps(
            config_init,
            config_goal,
            connect;
            steering_depth = steering_depth,
            epsilon = epsilon,
            TIME_LIMIT = (isnothing(TIME_LIMIT) ? nothing : TIME_LIMIT - elapsed()),
        )
    )
    if isnothing(roadmaps)
        VERBOSE > 0 &&
            @info @sprintf("\t%6.4f sec: failed to construct initial roadmaps\n", elapsed())
        return (nothing, map(i -> Vector{Node{State}}(), 1:N))
    end

    VERBOSE > 0 && !no_roadmap_at_beginning && @info ("\tdone, setup initial roadmaps")

    # setup distance tables
    distance_tables = get_distance_tables(roadmaps)

    # verbose
    print_progress! =
        (S::SuperNode{State}, loop_cnt::Int64; force::Bool = false) -> begin
            (VERBOSE < 1 || (!force && (loop_cnt % 100 != 0))) && return
            @printf(
                "\r\t\t%6.4f sec, explored node: %08d, f: %.4f, depth: %04d",
                elapsed(),
                loop_cnt,
                S.f,
                S.depth
            )
        end

    # setup heuristic function
    h_func(Q::Vector{Node{State}}) =
        (use_random_h_func ? rand() : sum(map(i -> distance_tables[i][Q[i].id], 1:N)) / N)

    # initial configuration
    Q_init = [roadmaps[i][1] for i = 1:N]

    # initial search node
    S_init = SuperNode(Q = Q_init, next = 1, id = get_Q_id(Q_init, 0), h = h_func(Q_init))

    k = 0
    while !timeover()
        k += 1

        # open list
        OPEN = PriorityQueue{SuperNode{State},Float64}()

        # discovered list to avoid duplication
        VISITED = Dict{String,SuperNode{State}}()

        # threshold of space-filling metric
        min_dist_thread = init_min_dist_thread * (decreasing_rate_min_dist_thread^(k - 1))

        # setup initail node
        enqueue!(OPEN, S_init, S_init.f)
        VISITED[S_init.id] = S_init

        loop_cnt = 0
        while !isempty(OPEN) && !timeover()
            loop_cnt += 1

            # pop
            S = dequeue!(OPEN)

            # check goal
            if check_goal(S.Q)
                print_progress!(S, loop_cnt, force = true)
                VERBOSE > 0 && @info @sprintf("\n\t%6.4f sec: found solution\n", elapsed())
                return (backtrack(S, VISITED), roadmaps)
            end

            # initial search or update for refine agents
            i = S.next
            j = mod1(S.next + 1, N)

            v = S.Q[i]

            # explore new states
            expand!(
                (q_from::State, q_to::State) -> conn(q_from, q_to, i),
                sampler,
                v,
                roadmaps[i],
                min_dist_thread,
                num_vertex_expansion,
                steering_depth,
            ) && (distance_tables[i] = get_distance_table(roadmaps[i]))

            # expand search node
            for p_id in vcat(v.neighbors, v.id)

                # create new configuration
                p = roadmaps[i][p_id]
                Q = copy(S.Q)
                Q[i] = p

                # check duplication and collision
                Q_id = get_Q_id(Q, j)
                haskey(VISITED, Q_id) && continue
                !no_fast_collision_check && collide(S.Q, p.q, i) && continue
                no_fast_collision_check && collide(S.Q, Q) && continue

                # create new search node
                S_new = SuperNode(
                    Q = Q,
                    next = j,
                    id = Q_id,
                    parent_id = S.id,
                    h = h_func(Q),
                    g = S.g + g_func(S.Q, Q),
                    depth = S.depth + 1,
                )

                # insert
                enqueue!(OPEN, S_new, S_new.f)
                VISITED[S_new.id] = S_new
            end
            print_progress!(S, loop_cnt, force = isempty(OPEN))
        end

        VERBOSE > 1 && println()
    end

    VERBOSE > 0 && @info @sprintf("\t%6.4f sec: failed to find solution\n", elapsed())
    return (nothing, roadmaps)
end

"""implementation of vertex expansion"""
function expand!(
    connect::Function,
    sampler::Function,
    v_from::Node{State},
    roadmap::Vector{Node{State}},
    min_dist_thread::Float64,
    num_vertex_expansion::Int64,
    steering_depth::Int64,
)::Bool where {State<:AbsState}

    updated = false
    for _ = 1:num_vertex_expansion
        # steering
        q_new = steering(connect, sampler(), v_from.q, steering_depth)
        # check space-filling metric
        if minimum(v -> dist(v.q, q_new), roadmap) > min_dist_thread
            # add vertex and edges
            u = Node(
                q_new,
                length(roadmap) + 1,
                map(v -> v.id, filter(v -> connect(q_new, v.q), roadmap)),
            )
            push!(roadmap, u)
            # update neighbors
            foreach(
                v -> push!(v.neighbors, u.id),
                filter(v -> connect(v.q, q_new), roadmap),
            )
            updated = true
        end
    end
    return updated
end

function steering(
    connect::Function,
    q_rand::State,
    q_near::State,
    steering_depth::Int64,
)::State where {State<:AbsState}

    q_h = q_rand  # probably unsafe -> eventually safe
    if !connect(q_near, q_h)
        # binary search
        q_l = q_near  # safe
        for _ = 1:steering_depth
            q = get_mid_status(q_l, q_h)
            if connect(q_near, q)
                q_l = q
            else
                q_h = q
            end
        end
        q_h = q_l
    end
    return q_h
end

"""generate initial roadmaps via RRT-Connect"""
function gen_RRT_connect_roadmaps(
    config_init::Vector{State},
    config_goal::Vector{State},
    connect::Function;
    steering_depth::Int64 = 2,
    epsilon::Union{Float64,Nothing} = 0.2,
    TIME_LIMIT::Union{Nothing,Real} = 30,
)::Union{Nothing,Vector{Vector{Node{State}}}} where {State<:AbsState}

    # utilities for timeout
    t_s = now()
    elapsed() = elapsed_sec(t_s)
    timeover() = TIME_LIMIT != nothing && elapsed() > TIME_LIMIT

    N = length(config_init)
    roadmaps = Vector{Vector{Node{State}}}()
    for i = 1:N
        rmp = gen_RRT_connect_roadmap(
            i,
            config_init[i],
            config_goal[i],
            connect;
            steering_depth = steering_depth,
            epsilon = epsilon,
            TIME_LIMIT = (isnothing(TIME_LIMIT) ? nothing : TIME_LIMIT - elapsed()),
        )
        isnothing(rmp) && return nothing
        push!(roadmaps, rmp)
    end
    return roadmaps
end

function gen_RRT_connect_roadmap(
    i::Int64,
    q_init::State,
    q_goal::State,
    connect::Function;
    steering_depth::Int64 = 2,
    epsilon::Union{Float64,Nothing} = 0.2,
    TIME_LIMIT::Union{Nothing,Real} = 30,
)::Union{Nothing,Vector{Node{State}}} where {State<:AbsState}

    # utilities for timeout
    t_s = now()
    elapsed() = elapsed_sec(t_s)
    timeover() = TIME_LIMIT != nothing && elapsed() > TIME_LIMIT

    # define sampler
    sampler = gen_uniform_sampling(q_init)

    # re-define connect function
    conn(q_from::State, q_to::State)::Bool = begin
        # check distance
        !(isnothing(epsilon) || dist(q_from, q_to) <= epsilon) && return false

        !connect(q_from, q_to, i) && return false

        return true
    end

    conn(q::State)::Bool = connect(q, i)

    # special case
    if conn(q_init, q_goal)
        v_init = Node(q_init, 1, [2])
        v_goal = Node(q_goal, 2, Vector{Int64}())
        conn(q_goal, q_init) && push!(v_goal.neighbors, v_init.id)
        return [v_init, v_goal]
    end

    # store all samples
    V1 = [q_init]
    V2 = [q_goal]

    V = V1

    # main loop
    iter = 0
    while !timeover()
        iter += 1
        from_start = (iter % 2 == 1)

        (flg, q_new) = extend!(conn, sampler(), V1, from_start, steering_depth)
        if flg != :trapped &&
           extend!(conn, q_new, V2, !from_start, steering_depth)[1] == :reached
            # fix tree
            if from_start
                V1, V2 = V2, V1
            end
            # create roadmap, insert vertices
            roadmap = [Node{State}(q_init, 1, []), Node{State}(q_goal, 2, [])]
            foreach(q -> push!(roadmap, Node{State}(q, length(roadmap) + 1, [])), V1[2:end])
            foreach(q -> push!(roadmap, Node{State}(q, length(roadmap) + 1, [])), V2[2:end])
            # create roadmap, insert edges
            for (k, v_from) in enumerate(roadmap)
                for (l, v_to) in enumerate(roadmap)
                    if k != l && conn(v_from.q, v_to.q)
                        push!(v_from.neighbors, v_to.id)
                    end
                end
                timeover() && return nothing
            end
            return roadmap
        end

        # swap tree
        V1, V2 = V2, V1
    end

    return nothing
end

"""used in RRT-Connect"""
function extend!(
    connect::Function,
    q_rand::State,  # sampled point
    V::Vector{State},  # vertices
    from_start::Bool,  # true -> tree rooted C_init, false -> tree rooted C_goal
    steering_depth::Int64,
)::Tuple{Symbol,State} where {State<:AbsState}

    # find nearest existing vertex
    q_near =
        (from_start ? argmin(q -> dist(q, q_rand), V) : argmin(q -> dist(q_rand, q), V))
    flg = :reached

    q_l = q_near  # safe sample
    q_h = q_rand  # probably unsafe sample -> eventually safe

    # steering, binary search
    if from_start ? !connect(q_near, q_h) : !(connect(q_h) && connect(q_h, q_near))
        flg = :trapped
        for _ = 1:steering_depth
            q = from_start ? get_mid_status(q_l, q_h) : get_mid_status(q_h, q_l)
            if from_start ? connect(q_near, q) : (connect(q) && connect(q, q_near))
                flg = :advanced
                q_l = q
            else
                q_h = q
            end
        end
        q_h = q_l
    end

    flg != :trapped && push!(V, q_h)
    return (flg, q_h)
end

"""generate id of search nodes"""
function get_Q_id(Q::Vector{Node{State}}, next::Int64)::String where {State<:AbsState}
    return @sprintf("%s_%d", join([v.id for v in Q], "-"), next)
end

"""obtain solution from search nodes by backtracking"""
function backtrack(
    S_fin::SuperNode{State},
    VISITED::Dict{String,SuperNode{State}},
)::Vector{Vector{Node{State}}} where {State<:AbsState}

    S = S_fin
    solution = Vector{Vector{Node{State}}}()
    while S.parent_id != nothing
        pushfirst!(solution, S.Q)
        S = VISITED[S.parent_id]
    end
    pushfirst!(solution, S.Q)
    return solution
end

end


# we also need this before defining a module
using DotOverload

@dotted module PlotTest

using DotOverload
using PyCall
using PyPlot

DotOverload.getMember(t::PyObject, k) = t[k]
DotOverload.setMember!(t::PyObject, k, v) = t[k] = v
DotOverload.getMember(t::Figure, k) = t[k]
DotOverload.setMember!(t::Figure, k, v) = t[k] = v

plt.ioff()

fig, ax = plt.subplots()
line, = ax.plot([],[])

plt.ion()

function init()
  fig.show()
  fig.canvas.draw()
  plt.axis([0, 2*π, -2, 2])
end

function doStuff()
  x = linspace(0,2*pi,1000);
  y = sin(3*x + 4*cos(2*x))

  line.set_xdata(x)
  line.set_ydata(y)

  ax.draw_artist(line)
  fig.canvas.update()
end

end #module

# usage:
# using PlotTest
# PlotTest.init()
# PlotTest.doStuff()


using Firefly
using Test
using MCMCChains

@testset "Firefly.jl" begin
    @testset "Distributions" begin
        include("distributions.jl")
    end
    @testset "findpeak" begin
        X = ones(100)
        @test findpeak(X) ≈ 1 rtol = 1e-2
        @test findpeaks(X, 1) ≈ [1] rtol = 1e-2

        X = append!(ones(100), 2 .* ones(80))
        @test findpeak(X) ≈ 1 rtol = 1e-2
        @test findpeaks(X, 1) ≈ [1] rtol = 1e-2
        @test findpeaks(X, 2) ≈ [1, 2] rtol = 1e-2

        # interface
        c = Chains(reshape(X, (180, 1, 1)), ["x"])
        @test findpeak(c) == findpeak(X)
        @test findpeaks(c) == findpeaks(X)
        @test findpeak(c) == findpeak(c, :x)
        @test findpeaks(c) == findpeaks(c, :x)
        @test findpeaks(c, 1) == findpeaks(c, :x, 1)
        @test findpeaks(c, 1) == findpeaks(c, :x)[1:1]

        c2 = Chains(ones(100, 2, 1), ["x", "y"])
        @test_throws MethodError findpeak(c2)
        @test_throws MethodError findpeaks(c2)
        @test_throws MethodError findpeaks(ones(100, 100))
        @test_throws MethodError findpeaks(ones(100), 2.4)
        @test_throws MethodError findpeaks(c, 2.4)
    end
end


# Copyright (C) 2020 Argonne National Laboratory
# Written by Alinson Santos Xavier <axavier@anl.gov>

using RELOG

@testset "build_graph" begin
    basedir = dirname(@__FILE__)
    instance = RELOG.parsefile("$basedir/../../instances/s1.json")
    graph = RELOG.build_graph(instance)
    process_node_by_location_name =
        Dict(n.location.location_name => n for n in graph.process_nodes)

    @test length(graph.plant_shipping_nodes) == 8
    @test length(graph.collection_shipping_nodes) == 10
    @test length(graph.process_nodes) == 6

    node = graph.collection_shipping_nodes[1]
    @test node.location.name == "C1"
    @test length(node.incoming_arcs) == 0
    @test length(node.outgoing_arcs) == 2
    @test node.outgoing_arcs[1].source.location.name == "C1"
    @test node.outgoing_arcs[1].dest.location.plant_name == "F1"
    @test node.outgoing_arcs[1].dest.location.location_name == "L1"
    @test node.outgoing_arcs[1].values["distance"] == 1095.62

    node = process_node_by_location_name["L1"]
    @test node.location.plant_name == "F1"
    @test node.location.location_name == "L1"
    @test length(node.incoming_arcs) == 10
    @test length(node.outgoing_arcs) == 2

    node = process_node_by_location_name["L3"]
    @test node.location.plant_name == "F2"
    @test node.location.location_name == "L3"
    @test length(node.incoming_arcs) == 2
    @test length(node.outgoing_arcs) == 2

    @test length(graph.arcs) == 38
end