# TODO
# - find need allocations in attractive case
# - decimate allocations here
# - I guess maybe it's additional Diagonal's etc?

################################################################################
### Type and general utility
################################################################################

struct BlockDiagonal{T, N, AT <: AbstractMatrix{T}} <: AbstractMatrix{T}
    blocks::NTuple{N, AT}

    function BlockDiagonal(blocks::NTuple{N, AT}) where {T, N, AT <: AbstractMatrix{T}}
        n, m = size(first(blocks))
        n == m || throw(ErrorException("Each block must be square."))
        all(b -> size(b) == (n, n), blocks) || throw(ErrorException(
            "All Blocks must have the same size."
        ))
        # One matrix with zeros is faster than multiple 3x3 matrices on my
        # machine (avx2). This will most likely depend on AVX version...
        # if n < 4
        #     full = zeros(T, n*N, n*N)
        #     for i in 1:N
        #         @views copyto!(full[(i-1)*n+1 : i*n, (i-1)*n+1 : i*n], blocks[i])
        #     end
        #     # new{T, 1, AT}((full,))
        #     # Is this bad? - Yes, because the rest of the code doesn't know about Matrix returns
        #     return full
        # else
            return new{T, N, AT}(blocks)
        # end
    end
end

BlockDiagonal(blocks::AT...) where {T, AT <: AbstractMatrix{T}} = BlockDiagonal(blocks)
function BlockDiagonal{T, N, AT}(::UndefInitializer, n, m) where {T, N, AT}
    n == m || throw(ErrorException("Each block must be square."))
    (n, m) = divrem(n, N)
    m == 0 || throw(ErrorException(
        "The size of the Matrix is incompatible with the blocking ($n % $N != 0)"
    ))
    BlockDiagonal([AT(undef, n, n) for _ in 1:N]...)
end
function BlockDiagonal{T, N, AT}(::UniformScaling, n, m) where {T, N, AT}
    n == m || throw(ErrorException("Each block must be square."))
    (n, m) = divrem(n, N)
    m == 0 || throw(ErrorException(
        "The size of the Matrix is incompatible with the blocking ($n % $N != 0)"
    ))
    BlockDiagonal([AT(I, n, n) for _ in 1:N]...)
end

function Base.Matrix(B::BlockDiagonal{T, N}) where {T, N}
    n = size(B.blocks[1], 1)
    M = N * n
    output = zeros(T, M, M)
    for i in 1:N
        @views copyto!(output[(i-1)*n+1 : i*n, (i-1)*n+1 : i*n], B.blocks[i])
    end
    output
end

function Base.checkbounds(B::BlockDiagonal{T, N}, i, j) where {T, N}
    @inbounds n = size(B.blocks[1], 1)
    (1 ≤ i ≤ N*n) && (1 ≤ j ≤ N*n)
end

function Base.getindex(B::BlockDiagonal{T, N}, i, j) where {T, N}
    @boundscheck checkbounds(B, i, j)
    @inbounds n = size(B.blocks[1], 1)

    b, k = divrem(i-1, n)
    _b, l = divrem(j-1, n)

    if b != _b
        return zero(T)
    else
        @inbounds return B.blocks[b+1][k+1, l+1]
    end
end

function Base.copyto!(B1::BlockDiagonal{T, N}, B2::BlockDiagonal{T, N}) where {T, N}
    @inbounds for i in 1:N
        copyto!(B1.blocks[i], B2.blocks[i])
    end
    return B1
end
function Base.copyto!(B::BlockDiagonal{T, N}, ::UniformScaling) where {T, N}
    @inbounds for i in 1:N
        copyto!(B.blocks[i], I)
    end
    return B
end
function Base.copyto!(B::BlockDiagonal{T, N}, D::Diagonal{T}) where {T, N}
    @inbounds n = size(B.blocks[1], 1)
    @inbounds for i in 1:N
        b = B.blocks[i]
        offset = (i-1)*n
        @turbo for j in 1:n, k in 1:n
            b[j, k] = ifelse(j == k, D.diag[offset+j], zero(T))
        end
    end
    return B
end

# for Base.show
function Base.size(B::BlockDiagonal{T, N}) where {T, N}
    @inbounds n = size(B.blocks[1], 1)
    (n*N, n*N)
end

# Needed for tests
Base.copy(B::BlockDiagonal) = deepcopy(B)

# Needed for stack building (shouldn't be performance critical)
function Base.:(*)(x::T, B::BlockDiagonal{T, N, AT}) where {T<:Number, N, AT<:AbstractMatrix{T}}
    BlockDiagonal(map(block -> x * block, B.blocks)...)
end
Base.:(*)(B::BlockDiagonal{T, N, AT}, x::T) where {T<:Number, N, AT<:AbstractMatrix{T}} = x * B
function Base.:(*)(B1::BlockDiagonal{T, N, AT}, B2::BlockDiagonal{T, N, AT}) where {
        T<:Number, N, AT<:AbstractMatrix{T}
    }
    BlockDiagonal(map(*, B1.blocks, B2.blocks)...)
end

function Base.exp(B::BlockDiagonal)
    BlockDiagonal(map(block -> exp(block), B.blocks)...)
end

function Base.log(B::BlockDiagonal)
    BlockDiagonal(map(block -> log(block), B.blocks)...)
end

function LinearAlgebra.det(B::BlockDiagonal{T}) where {T}
    output = T(1)
    for b in B.blocks
        output *= det(b)
    end
    output
end

# I thought this would be needed for greens(k, l), but it's not?
function LinearAlgebra.transpose!(A::BlockDiagonal{T, N}, B::BlockDiagonal{T, N}) where {T, N}
    @inbounds for i in 1:N
        transpose!(A.blocks[i], B.blocks[i])
    end
    A
end
# This however is used for greens(k, l)
function LinearAlgebra.rmul!(B::BlockDiagonal{T, N}, f::Number) where {T, N}
    @inbounds for i in 1:N
        rmul!(B.blocks[i], f)
    end
    B
end

################################################################################
### AVX-powered matrix multiplications
################################################################################

function vmul!(C::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, B::BlockDiagonal{T, N}) where {T, N}
    @inbounds for i in 1:N
        vmul!(C.blocks[i], A.blocks[i], B.blocks[i])
    end
end
function vmul!(C::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, B::Diagonal{T}) where {T<:Real, N}
    # Assuming correct size
    @inbounds n = size(A.blocks[1], 1)
    @inbounds for i in 1:N
        c = C.blocks[i]
        a = A.blocks[i]
        @turbo for k in 1:n, l in 1:n
            c[k,l] = a[k,l] * B.diag[(i-1)*n + l]
        end
    end
end
function vmul!(C::BlockDiagonal{T, N}, A::Diagonal{T}, B::BlockDiagonal{T, N}) where {T<:Real, N}
    # Assuming correct size
    @inbounds n = size(C.blocks[1], 1)
    @inbounds for i in 1:N
        c = C.blocks[i]
        b = B.blocks[i]
        @turbo for k in 1:n, l in 1:n
            c[k,l] = A.diag[(i-1)*n + k] * b[k,l]
        end
    end
end
function vmul!(C::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, X::Adjoint{T}) where {T<:Real, N}
    B = X.parent
    @inbounds n = size(C.blocks[1], 1)
    @inbounds for i in 1:N
        a = A.blocks[i]
        b = B.blocks[i]
        c = C.blocks[i]
        @turbo for k in 1:n, l in 1:n
            Ckl = zero(eltype(c))
            for m in 1:n
                Ckl += a[k,m] * b[l, m]
            end
            c[k,l] = Ckl
        end
    end
end
function vmul!(C::BlockDiagonal{T, N}, X::Adjoint{T}, B::BlockDiagonal{T, N}) where {T<:Real, N}
    A = X.parent
    @inbounds n = size(C.blocks[1], 1)
    @inbounds for i in 1:N
        a = A.blocks[i]
        b = B.blocks[i]
        c = C.blocks[i]
        @turbo for k in 1:n, l in 1:n
            Ckl = zero(eltype(c))
            for m in 1:n
                Ckl += a[m,k] * b[m,l]
            end
            c[k,l] = Ckl
        end
    end
end
function vmul!(C::BD, X1::Adjoint{T, BD}, X2::Adjoint{T, BD}) where {T <: Real, N, BD <: BlockDiagonal{T, N}}
    A = X1.parent
    B = X2.parent
    @inbounds n = size(C.blocks[1], 1)
    @inbounds for i in 1:N
        a = A.blocks[i]
        b = B.blocks[i]
        c = C.blocks[i]
        @turbo for k in 1:n, l in 1:n
            Ckl = zero(eltype(c))
            for m in 1:n
                Ckl += a[m,k] * b[l, m]
            end
            c[k,l] = Ckl
        end
    end
end
function rvmul!(A::BlockDiagonal{T, N}, B::Diagonal) where {T, N}
    # Assuming correct size
    @inbounds n = size(A.blocks[1], 1)
    @inbounds for i in 1:N
        @views rvmul!(A.blocks[i], Diagonal(B.diag[(i-1)*n+1 : i*n]))
    end
end
function lvmul!(A::Diagonal, B::BlockDiagonal{T, N}) where {T, N}
    # Assuming correct size
    @inbounds n = size(B.blocks[1], 1)
    @inbounds for i in 1:N
        @views lvmul!(Diagonal(A.diag[(i-1)*n+1 : i*n]), B.blocks[i])
    end
end

# used in greens(k, l)
function rvadd!(A::BlockDiagonal{T, N}, B::BlockDiagonal{T, N}) where {T <: Real, N}
    @inbounds for i in 1:N
        a = A.blocks[i]
        b = B.blocks[i]
        @turbo for j in axes(a, 1), k in axes(a, 2)
            a[j, k] = a[j, k] + b[j, k]
        end
    end
end
# used in equal time greens
function rvadd!(B::BlockDiagonal{T, N}, D::Diagonal{T}) where {T<:Real, N}
    # Assuming correct size
    @inbounds n = size(B.blocks[1], 1)
    @inbounds for i in 1:N
        a = B.blocks[i]
        offset = (i-1) * n
        @turbo for j in 1:n
            a[j, j] = a[j, j] + D.diag[j+offset]
        end
    end
end
# used in CombinedGreensIterator
function vsub!(O::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, ::UniformScaling) where {T<:Real, N}
    @inbounds n = size(O.blocks[1], 1)
    T1 = one(T)
    @inbounds for i in 1:N
        a = A.blocks[i]
        o = O.blocks[i]
        @turbo for j in 1:n, k in 1:n
            o[j, k] = a[j, k]
        end
        @turbo for j in 1:n
            o[j, j] -= T1
        end
    end
end

function rdivp!(A::BD, T::BD, O::BD, pivot) where {ET<:Real, N, BD <: BlockDiagonal{ET, N}}
    # assume Diagonal is ±1!
    @inbounds begin
        n = size(A.blocks[1], 1)

        # Apply pivot
        for b in 1:N
            o = O.blocks[b]
            a = A.blocks[b]
            t = T.blocks[b]
            offset = (b-1)*n

            for j in 1:n
                p = pivot[j + offset]
                @turbo for i in 1:n
                    o[i, j] = a[i, p]
                end
            end

            # do the rdiv
            # @turbo will segfault on `k in 1:0`, so pull out first loop
            invt11 = 1.0 / t[1, 1]
            @turbo for i in 1:n
                a[i, 1] = o[i, 1] * invt11
            end
            for j in 2:n
                invtjj = 1.0 / t[j, j]
                @turbo for i in 1:n
                    x = o[i, j]
                    for k in 1:j-1
                        x -= a[i, k] * t[k, j]
                    end
                    a[i, j] = x * invtjj
                end
            end
        end
    end
    A
end

################################################################################
### Fallbacks
################################################################################

function vmul!(C::BlockDiagonal{T, N}, A::Diagonal{T}, B::BlockDiagonal{T, N}) where {T, N}
    # Assuming correct size
    @inbounds n = size(C.blocks[1], 1)
    @inbounds for i in 1:N
        @views vmul!(C.blocks[i], Diagonal(A.diag[(i-1)*n+1 : i*n]), B.blocks[i])
    end
end
function vmul!(C::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, B::Diagonal) where {T, N}
    # Assuming correct size
    @inbounds n = size(C.blocks[1], 1)
    @inbounds for i in 1:N
        @views vmul!(C.blocks[i], A.blocks[i], Diagonal(B.diag[(i-1)*n+1 : i*n]))
    end
end
function vmul!(C::BlockDiagonal{T, N}, A::BlockDiagonal{T, N}, X::Adjoint{T}) where {T, N}
    B = X.parent
    @inbounds for i in 1:N
        vmul!(C.blocks[i], A.blocks[i], adjoint(B.blocks[i]))
    end
end
function vmul!(C::BlockDiagonal{T, N}, X::Adjoint{T}, B::BlockDiagonal{T, N}) where {T, N}
    A = X.parent
    @inbounds for i in 1:N
        vmul!(C.blocks[i], adjoint(A.blocks[i]), B.blocks[i])
    end
end
function rvadd!(B::BlockDiagonal{T, N}, D::Diagonal{T}) where {T, N}
    # Assuming correct size
    @inbounds n = size(B.blocks[1], 1)
    @inbounds for i in 1:N
        @views rvadd!(B.blocks[i], Diagonal(D.diag[(i-1)*n+1 : i*n]))
    end
end
function rvadd!(A::BlockDiagonal{T, N}, B::BlockDiagonal{T, N}) where {T, N}
    @inbounds for i in 1:N
        a = A.blocks[i]
        b = B.blocks[i]
        rvadd!(a, b)
    end
end
function rdivp!(A::BD, T::BD, O::BD, pivot) where {ET, N, BD <: BlockDiagonal{ET, N}}
    # Assuming correct size
    @inbounds n = size(T.blocks[1], 1)
    @inbounds for i in 1:N
        @views rdivp!(A.blocks[i], T.blocks[i], O.blocks[i], pivot[(i-1)*n+1 : i*n])
    end
    A
end

################################################################################
### UDT
################################################################################

function udt_AVX_pivot!(
        U::BlockDiagonal{Float64, N, AT},
        D::AbstractArray{Float64, 1},
        input::BlockDiagonal{Float64, N, AT},
        pivot::AbstractArray{Int64, 1} = Vector(UnitRange(1:size(input, 1))),
        temp::AbstractArray{Float64, 1} = Vector{Float64}(undef, length(D)),
        apply_pivot::Val = Val(true)
    ) where {N, AT <: AbstractMatrix{Float64}}
    # Assuming correct size
    @inbounds n = size(input.blocks[1], 1)
    @inbounds for i in 1:N
        @views udt_AVX_pivot!(
            U.blocks[i], D[(i-1)*n+1 : i*n], input.blocks[i],
            pivot[(i-1)*n+1 : i*n], temp[(i-1)*n+1 : i*n], apply_pivot
        )
        # pivot[(i-1)*n+1 : i*n] .+= (i-1)*n
    end
end
function udt_AVX_pivot!(
        U::BlockDiagonal{ComplexF64, N, AT},
        D::AbstractArray{Float64, 1},
        input::BlockDiagonal{ComplexF64, N, AT},
        pivot::AbstractArray{Int64, 1} = Vector(UnitRange(1:size(input, 1))),
        temp::AbstractArray{ComplexF64, 1} = Vector{ComplexF64}(undef, length(D)),
        apply_pivot::Val = Val(true)
    ) where {N, AT <: AbstractMatrix{ComplexF64}}
    # Assuming correct size
    @inbounds n = size(input.blocks[1], 1)
    @inbounds for i in 1:N
        @views udt_AVX_pivot!(
            U.blocks[i], D[(i-1)*n+1 : i*n], input.blocks[i],
            pivot[(i-1)*n+1 : i*n], temp[(i-1)*n+1 : i*n], apply_pivot
        )
        # pivot[(i-1)*n+1 : i*n] .+= (i-1)*n
    end
end


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl
=#

insts = OrderedDict()
insts["minimal"] = [
    ((1, 1, 1, false, false, true),),
    ((1, 1, 1, false, true, false),),
    ((1, 1, 1, true, false, false),),
    ]
insts["fast"] = [
    ((3, 1, 3, false, false, true),),
    ((3, 1, 3, false, true, false),),
    ((3, 1, 3, true, false, false),),
    ((1, 2, 3, false, false, true),),
    ((1, 2, 3, false, true, false),),
    ((1, 2, 3, true, false, false),),
    ]
insts["various"] = [
    ((4, 2, 5, false, false, true),),
    ((4, 2, 5, false, true, false),),
    ((4, 2, 5, true, false, false),),
    ((4, 2, 10, false, true, false),),
    ((4, 2, 10, true, false, false),),
    ((4, 4, 5, false, true, false),),
    ((4, 3, 6, true, false, false),),
    ((4, 3, 6, false, true, false),),
    ]
return (NearestPolyMatJuMP, insts)


module ClampTest

using Test
using ForneyLab
using ForneyLab: outboundType, isApplicable, @ensureVariables, generateId, addNode!, associate!
using ForneyLab: SPClamp

# Integration helper
mutable struct MockNode <: FactorNode
    id::Symbol
    interfaces::Vector{Interface}
    i::Dict{Symbol,Interface}

    function MockNode(out; id=generateId(MockNode))
        @ensureVariables(out)
        self = new(id, Array{Interface}(undef, 1), Dict{Int,Interface}())
        addNode!(currentGraph(), self)
        self.i[:out] = self.interfaces[1] = associate!(Interface(self), out)

        return self
    end
end

@testset "@ensureVariables" begin
    g = FactorGraph()
    nd_c = MockNode(constant(1.0))
    @test isa(nd_c.i[:out].partner.node, Clamp)

    nd_v = MockNode(1.0)
    @test isa(nd_v.i[:out].partner.node, Clamp)
end

@testset "Clamp" begin
    g = FactorGraph()
    nd = Clamp(Variable(), 1.0)

    @test isa(nd, Clamp)
    @test nd.value == 1.0
end

@testset "constant" begin
    g = FactorGraph()
    var = constant(1.0, id=:my_constant)
    nd = g.nodes[:my_constant]

    @test isa(var, Variable)
    @test isa(nd, Clamp)
    @test nd.value == 1.0
end

@testset "placeholder" begin
    g = FactorGraph()

    # Standard placeholder
    var = Variable(id=:y)
    placeholder(var, :y)
    nd = g.nodes[:placeholder_y]

    @test isa(nd, Clamp)
    @test g.placeholders[nd] == (:y, 0)

    # Placeholder without explicit variable initialization
    g = FactorGraph()

    placeholder(:y)
    nd = g.nodes[:placeholder_y]

    @test length(g.variables) == 1
    @test isa(nd, Clamp)
    @test g.placeholders[nd] == (:y, 0)

    # Indexed placeholder
    g = FactorGraph()

    var_i = Variable(id=:y)
    placeholder(var_i, :y, index=1)
    nd_i = g.nodes[:placeholder_y]

    @test isa(nd_i, Clamp)
    @test g.placeholders[nd_i] == (:y, 1)
end

#-------------
# Update rules
#-------------

@testset "SPClamp" begin
    @test SPClamp{Univariate} <: SumProductRule{Clamp{Univariate}}
    @test outboundType(SPClamp{Univariate}) == Message{PointMass, Univariate}
    @test isApplicable(SPClamp{Univariate}, DataType[])

    @test SPClamp{Multivariate} <: SumProductRule{Clamp{Multivariate}}
    @test outboundType(SPClamp{Multivariate}) == Message{PointMass, Multivariate}
    @test isApplicable(SPClamp{Multivariate}, DataType[])

    @test SPClamp{MatrixVariate} <: SumProductRule{Clamp{MatrixVariate}}
    @test outboundType(SPClamp{MatrixVariate}) == Message{PointMass, MatrixVariate}
    @test isApplicable(SPClamp{MatrixVariate}, DataType[])
end

end #module


################################################################################
#  Copyright 2020, Tom Van Acker                                               #
################################################################################
# MultiStateSystems.jl                                                         #
# A Julia package to solve multi-state system models.                          #
# See http://github.com/timmyfaraday/MultiStateSystems.jl                      #
################################################################################

"""
Example considering a double bridge network.
"""

# load pkgs
using Unitful
using MultiStateSystems

# include the state-transition diagrams for the pipes
stdᵖ = solvedSTD(prob = [0.1, 0.9], flow = [0.0u"m^3/hr", 1.0u"m^3/hr"])

# initialize the network corresponding with the bridge network
ntwᵈᵇⁿ = Network()

# add the sources, components and users to the network
add_sources!(ntwᵈᵇⁿ, node = [1 ,2])
add_components!(ntwᵈᵇⁿ, edge = [(1,3), (1,4), (2,3), (2,4), (3,5), (4,5)],
                       std  = stdᵖ)
add_bidirectional_component!(ntwᵈᵇⁿ, edge = (3,4), std = stdᵖ)
add_user!(ntwᵈᵇⁿ, node = 5)

# return the network
return ntwᵈᵇⁿ


################################################################################
# Active Set
################################################################################

function active(game_con::GameConstraintValues, stamp::CStamp)
	probsize = game_con.probsize
    N = probsize.N
    p = probsize.p
    px = probsize.px

    a = BitArray([0])
    if stamp.con == :col
        if valid(stamp, N, p)
            i = stamp.i
            j = stamp.j
            j_ = findfirst(x -> all(x.con.x2 .== px[j]), game_con.state_conval[i])
            conval = game_con.state_conval[i][j_]
            l = findfirst(x -> x==stamp.k, conval.inds)
            a[1] = conval.active[l][1]
			# if conval.vals[l][1] >= 0. || conval.λ[l][1] != 0.
                # a[1] = 1
            # end
        end
    end
    return a
end

function active_vertical_mask!(ascore::ActiveSetCore, game_con::GameConstraintValues)
    probsize = game_con.probsize
    N = probsize.N
	n = probsize.n
	m = probsize.m
	p = probsize.p
	S = probsize.S

    ascore.vmask = Vector{Int}(1:S)
	stamp = CStamp()
	for k = 2:N
		for i = 1:p
			for j = i+1:p
				stampify!(stamp, :v, :col, i, j, k)
				if active(game_con, stamp)[1] == 1
					ascore.vmask = vcat(ascore.vmask, vertical_idx(ascore, stamp))
				end
			end
		end
	end
    return nothing
end

function active_horizontal_mask!(ascore::ActiveSetCore, game_con::GameConstraintValues)
    probsize = game_con.probsize
    N = probsize.N
	n = probsize.n
	m = probsize.m
	p = probsize.p
	S = probsize.S

    ascore.hmask = Vector{Int}(1:S)
	stamp = CStamp()
	for k = 2:N
		for i = 1:p
			for j ∈ setdiff(1:p,i)
				stampify!(stamp, :h, :col, i, j, k)
				if active(game_con, stamp)[1] == 1
					ascore.hmask = vcat(ascore.hmask, horizontal_idx(ascore, stamp))
				end
			end
		end
	end
    return nothing
end

################################################################################
# Residual
################################################################################

function get_collision_conval(game_con::GameConstraintValues, i::Int, j::Int)
	px = game_con.probsize.px
	j_ = findfirst(
		x -> typeof(x.con) <: CollisionConstraint && all(x.con.x2 .== px[j]),
		game_con.state_conval[i]
	)
	if j_ == nothing
		valid = false
		return nothing, valid
	else
		valid = true
		return game_con.state_conval[i][j_], valid
	end
end

function residual!(ascore::ActiveSetCore, prob::GameProblem{KN,n,m,T,SVd,SVx},
	pdtraj::PrimalDualTraj{KN,n,m,T,SVd}) where {KN,n,m,T,SVd,SVx}

	probsize = ascore.probsize
	N = probsize.N
	p = probsize.p
	px = probsize.px
	S = probsize.S

	# Initialization
	ascore.res .= 0.0

	residual!(prob, pdtraj)
	ascore.res[1:S] .= prob.core.res

	evaluate!(prob.game_con, pdtraj.pr)
	stamp = CStamp()
	for i = 1:p
		for j = i+1:p
			conval, v = get_collision_conval(prob.game_con, i, j)
			if v
				for (l,k) in enumerate(conval.inds)
					stampify!(stamp, :v, :col, i, j, k)
					valid(stamp,N,p) ? add2sub(ascore.res_sub[stamp], conval.vals[l]) : nothing
				end
			end
		end
	end
    return nothing
end

################################################################################
# Residual Jacobian
################################################################################

function residual_jacobian!(ascore::ActiveSetCore, prob::GameProblem{KN,n,m,T,SVd,SVx},
	pdtraj::PrimalDualTraj{KN,n,m,T,SVd}) where {KN,n,m,T,SVd,SVx}

	probsize = ascore.probsize
	N = probsize.N
	p = probsize.p
	S = probsize.S

	# Reset!
	sparse_zero!(ascore.jac)

	residual_jacobian!(prob, pdtraj)
	ascore.jac[1:S,1:S] .= prob.core.jac

    # Allocations
	vs = VStamp()
	hs = HStamp()
    cs = CStamp()

	jacobian!(prob.game_con, pdtraj.pr)
	for i = 1:p
		for j ∈ setdiff(1:p,i)
			conval, v = get_collision_conval(prob.game_con, i, j)
			if v
				for (l,k) in enumerate(conval.inds)
					stampify!(vs, :opt, i, :x, 1, k)
					stampify!(cs, :h, :col, i, j, k)
					if valid(cs,N,p) && valid(vs,N,p)
						add2sub(ascore.jac_sub[(vs,cs)], conval.jac[l]')
					end
					if valid(hs,N,p) && valid(cs,N,p)
						stampify!(hs, :x, 1, k)
						stampify!(cs, :v, :col, i, j, k)
						add2sub(ascore.jac_sub[(hs,cs)], conval.jac[l])
					end
				end
			end
		end
	end
    return nothing
end

function update_nullspace!(ascore::ActiveSetCore, prob::GameProblem{KN,n,m,T,SVd,SVx},
	pdtraj::PrimalDualTraj{KN,n,m,T,SVd}; atol::T=1e-20) where {KN,n,m,T,SVd,SVx}

	update_active_set!(prob.game_con, pdtraj.pr)
	active_vertical_mask!(ascore, prob.game_con)
	active_horizontal_mask!(ascore, prob.game_con)
	residual_jacobian!(ascore, prob, pdtraj)
	djac = Matrix(ascore.jac[ascore.vmask, ascore.hmask])
	reset!(ascore.null)
	add_matrix!(ascore.null, nullspace(djac, atol=1e-20), ascore.hmask)
	return nothing
end


##########################################################################
# Copyright 2017 Samuel Ridler.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##########################################################################

# Demand (call) model

# get demand mode for given time and priority
function getDemandMode!(demand::Demand, priority::Priority, startTime::Float)
	@assert(startTime != nullTime)
	@assert(priority != nullPriority)

	updateDemandToTime!(demand, startTime) # update demand.recentSetsStartTimesIndex for given start time
	demandSetIndex = demand.setsTimeOrder[demand.recentSetsStartTimesIndex]
	demandModeIndex = demand.modeLookup[demandSetIndex, Int(priority)]

	return demand.modes[demandModeIndex]
end

# update demand.recentSetsStartTimesIndex up to given demand start time
function updateDemandToTime!(demand::Demand, startTime::Float)
	# shorthand:
	setsStartTimes = demand.setsStartTimes
	i = demand.recentSetsStartTimesIndex
	n = length(setsStartTimes)

	@assert(setsStartTimes[i] <= startTime) # otherwise, have gone back in time?

	if i == n || startTime < setsStartTimes[i+1]
		return # do nothing
	end

	# use linear search to update position in setsStartTimes
	while i < n && setsStartTimes[i+1] <= startTime
		i += 1
	end

	demand.recentSetsStartTimesIndex = i
end

# Creates one set of points to represent the demand in demand.rasters.
# Each point will store one or more demand values, one from each raster.
# Where there are multiple points per raster cell, the demand is divided between the points.
# Points are not created in places where no rasters have demand.
# Requires all rasters to have the same dimensions (only z values can differ).
function createDemandPointsFromRasters(demand::Demand; numCellRows::Int = 1, numCellCols::Int = 1)
	@assert(demand.numRasters >= 1)
	@assert(numCellRows >= 1)
	@assert(numCellCols >= 1)

	# shorthand
	@unpack rasters, numRasters = demand
	@unpack nx, ny = rasters[1]

	# check that all rasters have same dimensions,
	# otherwise creating points that work for all rasters may be tricky
	for i = 2:numRasters
		@assert(rasters[i].nx == nx)
		@assert(rasters[i].ny == ny)
		@assert(all(j -> isapprox(rasters[i].x[j], rasters[1].x[j]; rtol = eps()), 1:nx))
		@assert(all(j -> isapprox(rasters[i].y[j], rasters[1].y[j]; rtol = eps()), 1:ny))
	end

	# create point(s) for each raster cell,
	# except where point is in cell where all rasters have z = 0
	# numCellRows is for y dimension, numCellCols for x
	numCellPoints = numCellRows * numCellCols # number of points per cell
	points = Vector{Point}()
	maxNumPoints = (nx * numCellCols) * (ny * numCellRows)
	rastersPointDemands = [Vector{Float}(undef, maxNumPoints) for i = 1:numRasters] # rastersPointDemands[i][j] gives demand at point j for raster i
	for i = 1:nx, j = 1:ny # raster cell indices
		if all(raster -> raster.z[i,j] == 0, rasters)
			continue
		end
		for cellRow = 1:numCellRows, cellCol = 1:numCellCols
			point = Point()
			point.index = length(points) + 1

			# set point location
			point.location.x = rasters[1].x[i] + rasters[1].dx * ((cellCol - 0.5)/numCellCols - 0.5)
			point.location.y = rasters[1].y[j] + rasters[1].dy * ((cellRow - 0.5)/numCellRows - 0.5)

			# set demand of point for each raster,
			# divide raster demand by number of points per cell
			for rasterIndex = 1:numRasters
				rastersPointDemands[rasterIndex][point.index] = rasters[rasterIndex].z[i,j] / numCellPoints
			end

			push!(points, point)
		end
	end

	# remove unneeded values from rastersPointDemands (rastersPointDemands[i][j] where j > length(points))
	numPoints = length(points) # shorthand
	for i = 1:numRasters
		rastersPointDemands[i] = rastersPointDemands[i][1:numPoints]
	end

	return points, rastersPointDemands
end

# Sets the nearest node data for the points.
# mutates: points (fields: nearestNodeIndex, nearestNodeDist)
function setPointsNearestNodes!(sim::Simulation, points::Vector{Point})
	# find and set nearest node index of each point
	for point in points
		@assert(point.location.x != nullX && point.location.y != nullY)
		(point.nearestNodeIndex, point.nearestNodeDist) = findNearestNode(sim.map, sim.grid, sim.net.fGraph.nodes, point.location)
	end
end

# Creates an instance of PointsCoverageMode for the given travelMode and coverTime.
function createPointsCoverageMode(sim::Simulation, travelMode::TravelMode, coverTime::Float)
	@assert(coverTime >= 0)
	@assert(!isempty(sim.demandCoverage.nodesPoints))

	pointsCoverageMode = PointsCoverageMode()
	pointsCoverageMode.coverTime = coverTime
	pointsCoverageMode.travelMode = travelMode
	pointsCoverageMode.points = sim.demandCoverage.points

	# shorthand:
	points = pointsCoverageMode.points
	numPoints = length(points)
	@unpack stations, numStations = sim
	numNodes = length(sim.net.fGraph.nodes)
	nodesPoints = sim.demandCoverage.nodesPoints

	# # determine point coverage
	# stationsCoverPoints = fill(false, numStations, numPoints)
	# for i = 1:numStations
		# station = stations[i] # shorthand

		# # get travel time from station to nearest node
		# (node1, dist1) = (station.nearestNodeIndex, station.nearestNodeDist)
		# time1 = offRoadTravelTime(travelMode, dist1) # time to reach nearest node

		# for j = 1:numPoints
			# # get shortest path travel time
			# point = points[j] # shorthand
			# (node2, dist2) = (point.nearestNodeIndex, point.nearestNodeDist)
			# pathTravelTime = shortestPathTravelTime(net, travelMode.index, node1, node2)
			# time2 = offRoadTravelTime(travelMode, dist2) # time to reach nearest node
			# if time1 + pathTravelTime + time2 <= coverTime
				# stationsCoverPoints[i,j] = true
			# end
		# end
	# end

	# determine point coverage
	stationsCoverPoints = fill(false, numStations, numPoints)
	for i = 1:numStations
		station = stations[i] # shorthand

		# get travel time from station to nearest node
		(node1, dist1) = (station.nearestNodeIndex, station.nearestNodeDist)
		time1 = offRoadTravelTime(travelMode, dist1) # time to reach nearest node from station

		for j = 1:numNodes
			pathTravelTime = shortestPathTravelTime(sim.net, travelMode.index, node1, j)
			if time1 + pathTravelTime <= coverTime
				for k in nodesPoints[j] # indices of points for which node j is the nearest node
					point = points[k] # shorthand
					@assert(point.nearestNodeIndex == j)
					time2 = offRoadTravelTime(travelMode, point.nearestNodeDist) # time to reach point from nearest node
					stationsCoverPoints[i,k] = (time1 + pathTravelTime + time2 <= coverTime)
				end
			end
		end
	end

	# group points with coverage by same set of stations
	stationSetsPoints = Dict{Vector{Bool},Vector{Int}}() # stationSetsPoints[stationSet] gives the points that belong to the stationSet
	for j = 1:numPoints
		stationSet = stationsCoverPoints[:,j] # stationSet[i] = if station i covers point j
		# will put points that are not covered into a set, though alternatively they could be ignored
		stationSets = get!(stationSetsPoints, stationSet, Int[])
		push!(stationSets, j)
	end
	pointSets = collect(values(stationSetsPoints)) # pointSets[i] has set of all point indices covered by the same unique set of stations
	stationSets = [findall(stationSet) for stationSet in keys(stationSetsPoints)] # stationSets[i] = unique set of stations for pointSets[i]
	# according to documentation, 'values' and 'keys' return elements in the same order, so pointSets[i] corresponds with stationSets[i]

	stationsCoverPointSets = Vector{Vector{Int}}(undef, numStations) # stationsCoverPointSets[i] = list of point set indices covered by station i
	for i = 1:numStations
		stationsCoverPointSets[i] = findall(stationSet -> in(i,stationSet), stationSets)
	end

	pointsCoverageMode.pointSets = pointSets
	pointsCoverageMode.stationSets = stationSets
	pointsCoverageMode.stationsCoverPointSets = stationsCoverPointSets

	return pointsCoverageMode # will still need to set pointsCoverageMode.index
end

function getPointsCoverageMode(demandCoverage::DemandCoverage, travelMode::TravelMode, coverTime::Float)
	dc = demandCoverage # shorthand
	@assert(travelMode.index <= length(dc.pointsCoverageModeLookup))
	@assert(haskey(dc.pointsCoverageModeLookup[travelMode.index], coverTime))
	pointsCoverageMode = dc.pointsCoverageModes[dc.pointsCoverageModeLookup[travelMode.index][coverTime]]
	return pointsCoverageMode
end

# mutates: sim.travel state
function getPointsCoverageMode!(sim::Simulation, demandPriority::Priority, currentTime::Float)
	@assert(demandPriority != nullPriority)
	@assert(currentTime != nullTime)
	travelPriority = sim.responseTravelPriorities[demandPriority]
	travelMode = getTravelMode!(sim.travel, travelPriority, currentTime)
	coverTime = sim.demandCoverage.coverTimes[demandPriority]
	pointsCoverageMode = getPointsCoverageMode(sim.demandCoverage, travelMode, coverTime)
	@assert(pointsCoverageMode != nothing)
	return pointsCoverageMode
end

# mutates: sim.travel and sim.demand states
function getPointSetsDemands!(sim::Simulation, demandPriority::Priority, currentTime::Float;
	pointsCoverageMode::Union{PointsCoverageMode,Nothing} = nothing)
	if pointsCoverageMode == nothing
		pointsCoverageMode = getPointsCoverageMode!(sim, demandPriority, currentTime)
	end
	demandMode = getDemandMode!(sim.demand, demandPriority, currentTime)
	pointSetsDemands = sim.demandCoverage.pointSetsDemands[pointsCoverageMode.index, demandMode.rasterIndex]
	return pointSetsDemands
end

"""
	function initDemand!(sim::Simulation, demand::Union{Demand,Nothing} = nothing;
		demandFilename::String = "")
Initialise demand data for `sim`.
Will initialise from the first item in this list: `demand`, `demandFilename` file, `sim.demand` (if set), `sim.inputFiles[\"demand\"].path` file.
Deletes old demand coverage data in `sim.demandCoverage`.

Mutates: `sim.demand`, `sim.demandCoverage`
"""
function initDemand!(sim::Simulation, demand::Union{Demand,Nothing} = nothing;
	demandFilename::String = "")
	if demand == nothing
		if demandFilename != ""
			demand = readDemandFile(demandFilename)
		elseif sim.demand.numSets == nullIndex
			if haskey(sim.inputFiles, "demand")
				demand = readDemandFile(sim.inputFiles["demand"].path)
			else
				error("no demand data given")
			end
		else
			demand = sim.demand
		end
	end
	demand.initialised = true
	sim.demand = demand

	# old demand coverage data may no longer be valid
	sim.demandCoverage = DemandCoverage(sim.demandCoverage) # keep some old params
end

# Initialises the points coverage modes data in sim.demandCoverage.
# mutates: sim.demandCoverage (fields: pointsCoverageModes, pointsCoverageModeLookup)
function initPointsCoverageModes!(sim::Simulation)
	# shorthand:
	dc = sim.demandCoverage
	travel = sim.travel

	# go through all travel sets and demand priorities, create points coverage mode for each call response travel priority and cover time
	# dc.pointsCoverageModes
	dc.pointsCoverageModeLookup = [Dict{Float,Int}() for i = 1:travel.numModes]
	for i = 1:travel.numSets, demandPriority in priorities
		# get travel mode and cover time
		travelPriority = sim.responseTravelPriorities[demandPriority]
		travelMode = travel.modes[travel.modeLookup[i,Int(travelPriority)]]
		coverTime = sim.demandCoverage.coverTimes[demandPriority]

		# if combination of travel mode and cover time is new, need to calculate points coverage
		if !haskey(dc.pointsCoverageModeLookup[travelMode.index], coverTime)
			pointsCoverageMode = createPointsCoverageMode(sim, travelMode, coverTime)
			pointsCoverageMode.index = length(dc.pointsCoverageModes) + 1
			push!(dc.pointsCoverageModes, pointsCoverageMode)
			dc.pointsCoverageModeLookup[travelMode.index][coverTime] = pointsCoverageMode.index
		end
	end
end

"""
	function initDemandCoverage!(sim::Simulation;
		coverTimes::Union{Dict{Priority,Float},Nothing} = nothing,
		rasterCellNumRows::Int = 1, rasterCellNumCols::Int = 1)
Initialises demand coverage data for `sim`.
If `coverTimes` is given, will use this (along with `rasterCellNumRows` and `rasterCellNumCols`), otherwise will use the previous parameter values set in `sim.demandCoverage`.

# Keyword arguments
- `coverTimes` is the target coverage time for each demand priority, e.g. `coverTimes = Dict([p => 8/60/24 for p in priorities])` sets a target time of 8 minutes (converted to days) for each priority. Can omit this if `sim.demandCoverage.coverTimes` is already set.
- `rasterCellNumRows` and `rasterCellNumCols` are used to set the number of demand points to represent demand per raster cell for `sim.demand.rasters`

Mutates: `sim.demandCoverage`
"""
function initDemandCoverage!(sim::Simulation;
	coverTimes::Union{Dict{Priority,Float},Nothing} = nothing,
	rasterCellNumRows::Int = 1, rasterCellNumCols::Int = 1)

	@assert(!sim.used)
	@assert(sim.travel.recentSetsStartTimesIndex == 1)
	sim.demand.initialised || initDemand!(sim)
	@assert(sim.demand.recentSetsStartTimesIndex == 1)

	# shorthand
	@unpack demand, travel = sim
	@unpack rasters, numRasters = demand
	numNodes = length(sim.net.fGraph.nodes)

	# set demand coverage params
	oldDemandCoverage = sim.demandCoverage
	if coverTimes != nothing
		sim.demandCoverage = DemandCoverage(coverTimes, rasterCellNumRows, rasterCellNumCols)
	else # use old params
		@assert(!isempty(oldDemandCoverage.coverTimes))
		sim.demandCoverage = DemandCoverage(oldDemandCoverage)
	end
	dc = sim.demandCoverage # shorthand
	demandPriorities = priorities
	@assert(all(p -> haskey(dc.coverTimes, p), demandPriorities), "coverTimes not set")
	@assert(dc.rasterCellNumRows >= 1)
	@assert(dc.rasterCellNumCols >= 1)

	# create demand points
	dc.points, dc.rastersPointDemands = createDemandPointsFromRasters(sim.demand; numCellRows = dc.rasterCellNumRows, numCellCols = dc.rasterCellNumCols)
	points = dc.points # shorthand
	setPointsNearestNodes!(sim, points)
	numPoints = length(points) # shorthand

	# set nodesPoints
	dc.nodesPoints = [Int[] for i = 1:numNodes]
	for point in points
		push!(dc.nodesPoints[point.nearestNodeIndex], point.index)
	end

	# set points coverage modes
	initPointsCoverageModes!(sim)
	numPointsCoverageModes = length(dc.pointsCoverageModes)

	# set pointSetsDemands
	# note that pointSetsDemands does not need to be set for all combinations of pointsCoverageMode.index and rasterIndex, as some combinations may not occur in simulation
	dc.pointSetsDemands = [Vector{Float}() for i = 1:numPointsCoverageModes, j = 1:demand.numRasters]
	times = vcat(travel.setsStartTimes, demand.setsStartTimes) |> unique |> sort
	for time in times, demandPriority in demandPriorities
		pointsCoverageMode = getPointsCoverageMode!(sim, demandPriority, time)
		demandMode = getDemandMode!(demand, demandPriority, time)
		rasterIndex = demandMode.rasterIndex # shorthand

		# if combination of pointsCoverageMode and rasterIndex is new, need to calculate point sets demands
		pointDemands = dc.rastersPointDemands[rasterIndex]
		if isempty(dc.pointSetsDemands[pointsCoverageMode.index, rasterIndex])
			pointSetsDemand = zeros(Float, length(pointsCoverageMode.pointSets))
			for (pointSetIndex, pointSet) in enumerate(pointsCoverageMode.pointSets)
				for pointIndex in pointSet
					pointSetsDemand[pointSetIndex] += pointDemands[pointIndex]
				end
			end
			dc.pointSetsDemands[pointsCoverageMode.index, rasterIndex] = pointSetsDemand
		end
	end
	# reset travel and demand state (previously changed by getTravelMode! and getDemandMode!)
	travel.recentSetsStartTimesIndex = 1
	demand.recentSetsStartTimesIndex = 1

	dc.initialised = true
end

# Calculate, for a given points coverage mode, and number of ambulances assigned to each
# station, how many ambulances cover each point set.
function calcPointSetsCoverCounts(pointsCoverageMode::PointsCoverageMode, stationsNumAmbs::Vector{Int})
	# stationsNumAmbs[i] gives the number of ambulances that can provide coverage and are assigned to station i

	@assert(all(n -> n >= 0, stationsNumAmbs))

	# count how many times each point set is covered
	numPointSets = length(pointsCoverageMode.pointSets) # shorthand
	pointSetsCoverCount = zeros(Int, numPointSets)
	for i = 1:numPointSets
		for j in pointsCoverageMode.stationSets[i]
			pointSetsCoverCount[i] += stationsNumAmbs[j]
		end
	end

	# # alternative code:
	# @assert(length(stationsNumAmbs) == length(pointsCoverageMode.stationsCoverPointSets))
	# pointSetsCoverCount = zeros(Int, length(pointsCoverageMode.pointSets))
	# for (i,stationNumAmbs) in enumerate(stationsNumAmbs)
		# for j in pointsCoverageMode.stationsCoverPointSets[i]
			# pointSetsCoverCount[j] += stationNumAmbs
		# end
	# end

	return pointSetsCoverCount
end

# Calculate, for each demand priority, how many ambulances cover each point set in
# pointsCoverageMode.pointSet, where pointsCoverageMode is for the demand priority and current time.
# mutates: sim.travel state
function calcPointSetsCoverCounts!(sim::Simulation, currentTime::Float, stationsNumAmbs::Vector{Int};
	demandPriorities::Vector{Priority} = priorities)
	# stationsNumAmbs[i] gives the number of ambulances that can provide coverage and are assigned to station i

	@assert(length(stationsNumAmbs) == sim.numStations)
	@assert(all(n -> n >= 0, stationsNumAmbs))

	demandsPointSetsCoverCounts = Dict{Priority,Vector{Int}}() # demandsPointSetsCoverCounts[p][i] gives the number of ambulances assigned to stations that cover point set i for demand priority p
	pointsCoverageModesUsed = Dict{PointsCoverageMode,Priority}() # keep track of which points coverage mode are used, and any demand priority that uses each
	for demandPriority in demandPriorities
		pointsCoverageMode = getPointsCoverageMode!(sim, demandPriority, currentTime)
		if haskey(pointsCoverageModesUsed, pointsCoverageMode)
			# have already calculated the point sets cover counts for this points coverage mode
			priority = pointsCoverageModesUsed[pointsCoverageMode]
			demandsPointSetsCoverCounts[demandPriority] = demandsPointSetsCoverCounts[priority] # should use copy(), so that vectors are independent?
		else
			demandsPointSetsCoverCounts[demandPriority] = calcPointSetsCoverCounts(pointsCoverageMode, stationsNumAmbs)
			pointsCoverageModesUsed[pointsCoverageMode] = demandPriority
		end
	end

	return demandsPointSetsCoverCounts
end

# Calculates, for each demand priority, how much of the demand is covered by 1,2,... ambulances
# mutates: sim.travel and sim.demand states
function calcDemandCoverCounts!(sim::Simulation, currentTime::Float, stationsNumAmbs::Vector{Int};
	demandPriorities::Vector{Priority} = priorities)
	# stationsNumAmbs[i] gives the number of ambulances that can provide coverage and are assigned to station i

	@assert(length(stationsNumAmbs) == sim.numStations)
	@assert(all(n -> n >= 0, stationsNumAmbs))

	# shorthand
	numAmbs = sum(stationsNumAmbs)

	demandsPointSetsCoverCounts = calcPointSetsCoverCounts!(sim, currentTime, stationsNumAmbs; demandPriorities = demandPriorities)

	demandCoverCounts = Dict{Priority,Vector{Float}}() # demandCoverCounts[p][i] gives the amount of demand covered by i ambulances for demand priority p
	coverCountsCalculated = [Dict{Vector{Int},Vector{Int}}() for i = 1:sim.demand.numRasters] # coverCountsCalculated[rasterIndex][pointSetsCoverCount] stores the coverCounts (before multiplying by demandMode.rasterMultiplier) for the combination of args
	for demandPriority in demandPriorities
		pointsCoverageMode = getPointsCoverageMode!(sim, demandPriority, currentTime)
		demandMode = getDemandMode!(sim.demand, demandPriority, currentTime)
		rasterIndex = demandMode.rasterIndex # shorthand

		pointSetsCoverCount = demandsPointSetsCoverCounts[demandPriority]
		pointSetsDemands = sim.demandCoverage.pointSetsDemands[pointsCoverageMode.index, rasterIndex]

		coverCounts = nothing # init
		if haskey(coverCountsCalculated[rasterIndex], pointSetsCoverCount)
			coverCounts = coverCountsCalculated[rasterIndex][pointSetsCoverCount]
		else
			coverCounts = zeros(Float, numAmbs) # coverCounts[i] is the demand covered by i ambulances
			for (i,pointSet) in enumerate(pointsCoverageMode.pointSets)
				if pointSetsCoverCount[i] > 0
					coverCounts[pointSetsCoverCount[i]] += pointSetsDemands[i]
				end
			end
			coverCountsCalculated[rasterIndex][pointSetsCoverCount] = coverCounts
		end
		demandCoverCounts[demandPriority] = coverCounts * demandMode.rasterMultiplier
	end

	return demandCoverCounts
end


#=
    Used from https://github.com/lanl-ansi/Juniper.jl
=#

function init_original_problem!(op::OriginalProblem, model::MOI.AbstractOptimizer)
    num_variables = length(model.variable_info)
    num_linear_le_constraints = length(model.linear_le_constraints)
    num_linear_ge_constraints = length(model.linear_ge_constraints)
    num_linear_eq_constraints = length(model.linear_eq_constraints)
    num_quadratic_le_constraints = length(model.quadratic_le_constraints)
    num_quadratic_ge_constraints = length(model.quadratic_ge_constraints)
    num_quadratic_eq_constraints = length(model.quadratic_eq_constraints)

    op.status = MOI.OPTIMIZE_NOT_CALLED
    op.relaxation_status = MOI.OPTIMIZE_NOT_CALLED
    op.has_nl_objective = model.nlp_data.has_objective
    op.nlp_evaluator = model.nlp_data.evaluator
    op.objective = model.objective

    op.objval = NaN
    op.best_bound = NaN
    op.solution = fill(NaN, num_variables)
    op.nsolutions = 0
    op.solutions = []
    op.num_disc_var = 0
    op.nintvars = 0
    op.nbinvars = 0
    op.nnodes = 1 # is set to one for the root node

    op.nlp_constr_type = []

    op.start_time = time()

    op.nl_solver = model.options.nl_solver
    nl_vec_opts = Vector{Pair}()
    if isa(op.nl_solver, MOI.OptimizerWithAttributes)
        for arg in model.options.nl_solver.params
            push!(nl_vec_opts, arg)
        end
    else
        op.nl_solver = JuMP.optimizer_with_attributes(op.nl_solver)
    end
    op.nl_solver_options = nl_vec_opts

    if model.options.mip_solver !== nothing
        op.mip_solver = model.options.mip_solver
        mip_vec_opts = Vector{Pair}()
        if isa(op.mip_solver, MOI.OptimizerWithAttributes)
            for arg in model.options.mip_solver.params
                push!(mip_vec_opts, arg)
            end
        else
            op.mip_solver = JuMP.optimizer_with_attributes(op.mip_solver)
        end
        op.mip_solver_options = mip_vec_opts
    end
    op.options = model.options
    if model.sense == MOI.MIN_SENSE
        op.obj_sense = :Min
    else
        op.obj_sense = :Max
    end
    op.l_var = info_array_of_variables(model.variable_info, :lower_bound)
    op.u_var = info_array_of_variables(model.variable_info, :upper_bound)
    integer_bool_arr = info_array_of_variables(model.variable_info, :is_integer)
    binary_bool_arr = info_array_of_variables(model.variable_info, :is_binary)
    primal_start_arr = info_array_of_variables(model.variable_info, :start)
    op.primal_start = primal_start_arr
    op.nintvars = sum(integer_bool_arr)
    op.nbinvars = sum(binary_bool_arr)
    op.num_disc_var = sum(integer_bool_arr)+sum(binary_bool_arr)
    op.num_var = length(model.variable_info)
    op.var_type = [:Cont for i in 1:op.num_var]
    op.var_type[integer_bool_arr .== true] .= :Int
    op.var_type[binary_bool_arr .== true] .= :Bin
    op.disc2var_idx = zeros(op.num_disc_var)
    op.var2disc_idx = zeros(op.num_var)
    int_i = 1
    for i=1:op.num_var
        if op.var_type[i] != :Cont
            op.disc2var_idx[int_i] = i
            op.var2disc_idx[i] = int_i
            int_i += 1
        end
    end
    op.num_l_constr = num_linear_le_constraints+num_linear_ge_constraints+num_linear_eq_constraints
    op.num_q_constr = num_quadratic_le_constraints+num_quadratic_ge_constraints+num_quadratic_eq_constraints
    op.num_nl_constr = length(model.nlp_data.constraint_bounds)
    op.num_constr = op.num_l_constr+op.num_q_constr+op.num_nl_constr
end

function create_original_model!(optimizer::MOI.AbstractOptimizer, op::OriginalProblem; fix_start=false)
    ps = op.options.log_levels

    op.model = Model()
    lb = op.l_var
    ub = op.u_var
    # all continuous we solve relaxation first
    @variable(op.model, lb[i] <= x[i=1:op.num_var] <= ub[i])

    for i=1:op.num_var
        JuMP.set_start_value(x[i], op.primal_start[i])
    end

    register_functions!(op.model, op.options.registered_functions)

    # TODO check whether it is supported
    if optimizer.nlp_data.has_objective
        obj_expr = MOI.objective_expr(optimizer.nlp_data.evaluator)
        obj_expr = expr_dereferencing(obj_expr, op.model)
        try
            JuMP.set_NL_objective(op.model, optimizer.sense, obj_expr)
        catch
            error("Have you registered a function? Then please register the function also for Juniper see: \n
            https://lanl-ansi.github.io/Juniper.jl/stable/options/#registered_functions%3A%3AUnion%7BNothing%2CVector%7BRegisteredFunction%7D%7D-%5Bnothing%5D-1")
        end
    elseif optimizer.objective !== nothing
        MOI.set(op.model, MOI.ObjectiveFunction{typeof(optimizer.objective)}(), optimizer.objective)
        MOI.set(op.model, MOI.ObjectiveSense(), optimizer.sense)
    end

    backend = JuMP.backend(op.model);
    llc = optimizer.linear_le_constraints
    lgc = optimizer.linear_ge_constraints
    lec = optimizer.linear_eq_constraints
    qlc = optimizer.quadratic_le_constraints
    qgc = optimizer.quadratic_ge_constraints
    qec = optimizer.quadratic_eq_constraints
    for constr_type in [llc, lgc, lec, qlc, qgc, qec]
        for constr in constr_type
            MOI.add_constraint(backend, constr[1], constr[2])
        end
    end
    for i in 1:op.num_nl_constr
        constr_expr = MOI.constraint_expr(optimizer.nlp_data.evaluator, i)
        constr_expr = expr_dereferencing(constr_expr, op.model)
        push!(op.nlp_constr_type, constr_expr.args[1])
        try
            JuMP.add_NL_constraint(op.model, constr_expr)
        catch
            error("Have you registered a function? Then please register the function also for Juniper see: \n
            https://lanl-ansi.github.io/Juniper.jl/stable/options/#registered_functions%3A%3AUnion%7BNothing%2CVector%7BRegisteredFunction%7D%7D-%5Bnothing%5D-1")
        end
    end

    op.x = x
end

function fix_primal_start!(op::OriginalProblem)
    lb = op.l_var
    ub = op.u_var
    x = op.x
    for i=1:op.num_var
        if op.var_type[i] != :Cont && lb[i] <= op.primal_start[i] <= ub[i]
            JuMP.set_lower_bound(op.x[i], op.primal_start[i])
            JuMP.set_upper_bound(op.x[i], op.primal_start[i])
        else
            JuMP.set_start_value(x[i], op.primal_start[i])
        end
    end
end

function unfix_primal_start!(op::OriginalProblem)
    lb = op.l_var
    ub = op.u_var
    x = op.x
    for i=1:op.num_var
        JuMP.set_lower_bound(op.x[i], lb[i])
        JuMP.set_upper_bound(op.x[i], ub[i])
        JuMP.set_start_value(x[i], op.primal_start[i])
    end
end

function solve_root_incumbent_model(op::OriginalProblem)
    status, backend = optimize_get_status_backend(op.model; solver=op.nl_solver)
    incumbent = nothing
    ps = op.options.log_levels
    if state_is_optimal(status; allow_almost=op.options.allow_almost_solved)
        # set incumbent
        objval = MOI.get(backend, MOI.ObjectiveValue())
        solution = JuMP.value.(op.x)
        if are_type_correct(solution, op.var_type, op.disc2var_idx, op.options.atol)
            if only_almost_solved(status) && op.options.allow_almost_solved_integral
                @warn "Start value incumbent only almost locally solved. Disable with `allow_almost_solved_integral=false`"
            end
        end
        #incumbent = Incumbent(objval, solution, only_almost_solved(status))
    end
    return incumbent
end


# =====================================
# Structs to work with waiting lists
# =====================================

"""
    StateInLocation{ST, M<:Integer}

Struct that associates a set with a hybrid automaton's location index,
usually an integer.
"""
struct StateInLocation{ST, M}
    X::ST     # set representation
    loc_id::M # discrete location index
end

state(s::StateInLocation) = s.X
location(s::StateInLocation) = s.loc_id

# convert to a struct with broader type
function Base.convert(::Type{StateInLocation{ST, M}}, s::StateInLocation{WT, M}) where {ST, WT<:ST, M}
    StateInLocation{ST, M}(s.X, s.loc_id)
end

# ===============================================================
# Abstract interface
# ===============================================================

"""
    AbstractWaitingList

Abstract supertype for all waiting list types.
"""
abstract type AbstractWaitingList end

# iterator interface
# array(w) should return the array of StateInLocation elements
@inline Base.getindex(w::AbstractWaitingList, i::Int) = getindex(array(w), i)
@inline Base.getindex(w::AbstractWaitingList, i::Number) = getindex(array(w), convert(Int, i))
@inline Base.getindex(w::AbstractWaitingList, I::AbstractVector) = getindex(array(w), I)
@inline Base.isempty(w::AbstractWaitingList) = isempty(array(w))
@inline Base.iterate(w::AbstractWaitingList) = iterate(array(w))
@inline Base.iterate(w::AbstractWaitingList, state) = iterate(array(w), state)
@inline Base.length(w::AbstractWaitingList) = length(array(w))
@inline Base.size(w::AbstractWaitingList) = (length(array(w)),)
@inline Base.first(w::AbstractWaitingList) = getindex(w, 1)
@inline Base.last(w::AbstractWaitingList) = getindex(w, lastindex(w))
@inline Base.firstindex(w::AbstractWaitingList) = 1
@inline Base.lastindex(w::AbstractWaitingList) = length(array(w))
@inline Base.eachindex(w::AbstractWaitingList) = eachindex(array(w))

# ===============================================================
# Waiting list with all sets of the same type
# ===============================================================

"""
    WaitingList{TN, ST, M, QT<:StateInLocation{ST, M}} <: AbstractWaitingList

Iterable container representing a list of pairs `(set, mode)` of a hybrid automaton.

### Fields

- `times`  -- vector with a time interval associated to each state
- `array`  -- vector of `StateInLocation`

### Notes

A `WaitingList` is a list of pairs ``(set_i, loc_i)` for `i in 1..k` where
`times` is a vector with a time interval associated to each state.

This waiting list allows for a unique set representation (`ST`) for all elements
of the list.
"""
struct WaitingList{TN, ST, M, QT<:StateInLocation{ST, M}} <: AbstractWaitingList
    times::Vector{TN}
    array::Vector{QT}

    function WaitingList(times::Vector{TN}, array::Vector{QT}) where {TN, ST, M, QT<:StateInLocation{ST, M}}
        @assert length(times) == length(array) || throw(ArgumentError("the lengths of the time " *
        "stamps and the array of sets should match, but they are $(length(times)) and $(length(array)) respectively"))

        return new{TN, ST, M, QT}(times, array)
    end
end

# constructor of empty waiting list
WaitingList{TN, ST, M}() where {TN, ST, M} = WaitingList(Vector{TN}(), Vector{StateInLocation{ST, M}}())
WaitingList{TN, ST, M, QT}() where {TN, ST, M, QT<:StateInLocation{ST, M}} = WaitingList(Vector{TN}(), Vector{QT}())

# constructor without time-stamp
WaitingList(array::Vector{QT}) where {ST, M, QT<:StateInLocation{ST, M}} = WaitingList{TimeInterval}(array)
function WaitingList{TN}(array::Vector{QT}) where {TN, ST, M, QT<:StateInLocation{ST, M}}
    times = fill(zeroI, length(array))
    WaitingList(times, array)
end

# getter functions
@inline array(w::WaitingList) = w.array
@inline times(w::WaitingList) = w.times
setrep(w::WaitingList{TN, ST}) where {TN, ST} = ST
locrep(w::WaitingList{TN, ST, M}) where {TN, ST, M} = M
tspan(w::WaitingList, i) = w.times[i]

# we let elem::ET be possibly different than QT, for waiting lists which mixed set representations
@inline function Base.push!(w::WaitingList{TN, ST, M, QT}, elem::ET) where {TN, ST, M, QT, ET}
    push!(w.times, zeroI)
    push!(w.array, elem)
    return w
end
@inline function Base.push!(w::WaitingList{TN, ST, M, QT}, Δt::TN, elem::ET) where {TN, ST, M, QT, ET}
    push!(w.times, Δt)
    push!(w.array, elem)
    return w
end

@inline Base.pop!(w::WaitingList) = (pop!(w.times), pop!(w.array))

# containment check
function Base.:⊆(s::StateInLocation, w::WaitingList)
    contained = false
    q = location(s)
    X = state(s)
    for elem in w
        if (location(elem) == q) && (X ⊆ state(elem))
            contained = true
            break
        end
    end
    return contained
end

# conversion from vector-of-tuples to waiting list
function Base.convert(::Type{TW}, Q::Vector{Tuple{M, ST}}) where {TW<:WaitingList, M<:Integer, ST<:AdmissibleSet}
    waiting_list = TW()
    for Qi in Q # no intersection chcks
        q, Xi = Qi
        push!(waiting_list, StateInLocation(Xi, q))
    end
    return waiting_list
end

# "duck-typing" conversion from vector-of-tuples to waiting list
function Base.convert(::Type{TW}, Q::Vector{Tuple{ST, M}}) where {TW<:WaitingList, ST<:AdmissibleSet, M<:Integer}
    waiting_list = TW()
    for Qi in Q # no intersection chcks
        Xi, q = Qi
        push!(waiting_list, StateInLocation(Xi, q))
    end
    return waiting_list
end

# ===============================================================
# Waiting list allowing storage of sets of different types
# ===============================================================

# non-strictly typed waiting list, useful for cases in which the set reprsentation
# to be added in the list is not known a priori
struct MixedWaitingList{TN, QT<:StateInLocation} <: AbstractWaitingList
    times::Vector{TN}
    array::Vector{QT}

    function MixedWaitingList(times::Vector{TN}, array::Vector{QT}) where {TN, QT<:StateInLocation}
        @assert length(times) == length(array) || throw(ArgumentError("the lengths of the time " *
        "stamps and the array of sets should match, but they are $(length(times)) and $(length(array)) respectively"))

        return new{TN, QT}(times, array)
    end
end

# constructor of empty waiting list
MixedWaitingList{TN, QT}() where {TN, QT} = MixedWaitingList(Vector{TN}(), Vector{StateInLocation}())


# Paper_Experiment_Functions.jl
#= ---------------- #
Set of datatype and functions for efficient testing.
# ---------------- =#

module TestFunctions

using BayesianSVM
using Distributions
using KernelFunctions

export TestingModel
export DefaultParameters, BSVMParameters
export CreateModel, TrainModel, RunTests, ProcessResults, PrintResults, WriteResults
export ComputePrediction, ComputePredictionAccuracy

#Datatype for containing the model, its results and its parameters
type TestingModel
  MethodName::String #Name of the method
  DatasetName::String #Name of the dataset
  ExperimentType::String #Type of experiment
  MethodType::String #Type of method used ("SVM","BSVM","ECM","GPC")
  Param::Dict{String,Any} #Some paramters to run the method
  Results::Dict{String,Any} #Saved results
  Model::Any
  TestingModel(methname,dataset,exp,methtype) = new(methname,dataset,exp,methtype,Dict{String,Any}(),Dict{String,Any}())
  TestingModel(methname,dataset,exp,methtype,params) = new(methname,dataset,exp,methtype,params,Dict{String,Any}())
end

#Create a default dictionary
function DefaultParameters()
  param = Dict{String,Any}()
  param["ϵ"]= 1e-8 #Convergence criteria
  param["BatchSize"] = 10 #Number of points used for stochasticity
  param["Kernel"] = "rbf" # Kernel function
  param["Θ"] = 1.0 # Hyperparameter for the kernel function
  param["γ"] = 1.0 #Variance of introduced noise
  param["M"] = 32 #Number of inducing points
  param["Window"] = 5 #Number of points used to check convergence (smoothing effect)
  param["Verbose"] = 0 #Verbose
  return param
end

#Create a default parameters dictionary for BSVM
function BSVMParameters(;Stochastic=true,NonLinear=true,Sparse=true,ALR=true,Autotuning=false,main_param=DefaultParameters())
  param = Dict{String,Any}()
  param["Stochastic"] = Stochastic #Is the method stochastic
  param["Sparse"] = Sparse #Is the method using inducing points
  param["NonLinear"] = NonLinear #Is the method using kernels
  param["Intercept"] = NonLinear ? false : true;
  param["ALR"] = ALR #Is the method using adpative learning rate (in case of the stochastic case)
  param["AutoTuning"] = Autotuning #Is hyperoptimization performed
  param["ATFrequency"] = 1 #How even autotuning is performed
  param["κ"] = 1.0;  param["τ"] = 100; #Parameters for learning rate of autotuning
  param["κ_s"] = 1.0;  param["τ_s"] = 100; #Parameters for learning rate of Stochastic gradient descent when ALR is not used
  param["ϵ"] = main_param["ϵ"]; param["Window"] = main_param["Window"]; #Convergence criteria (checking parameters norm variation on a window)
  param["Kernels"] = [Kernel(main_param["Kernel"],1.0,params=main_param["Θ"])] #Kernel creation (standardized for now)
  param["Verbose"] = main_param["Verbose"] ? 3 : 0 #Verbose
  param["BatchSize"] = main_param["BatchSize"] #Number of points used for stochasticity
  param["M"] = main_param["M"] #Number of inducing points
  param["γ"] = main_param["γ"] #Variance of introduced noise
  return param
end

#Create a model given the parameters passed in p
function CreateModel(tm::TestingModel,X_training,y_training) #tm testing_model, p parameters
  tm.Model = BSVM(X_training,y_training;Stochastic=tm.Param["Stochastic"],batchSize=tm.Param["BatchSize"],Sparse=tm.Param["Sparse"],m=tm.Param["M"],NonLinear=tm.Param["NonLinear"],
  kernels=tm.Param["Kernels"],Autotuning=tm.Param["AutoTuning"],autotuningfrequency=tm.Param["ATFrequency"],AdaptativeLearningRate=tm.Param["ALR"],Intercept=tm.Param["Intercept"],κ_s=tm.Param["κ_s"],τ_s = tm.Param["τ_s"],ϵ=tm.Param["ϵ"],γ=tm.Param["γ"],
  κ_Θ=tm.Param["κ"],τ_Θ=tm.Param["τ"],smoothingWindow=tm.Param["Window"],VerboseLevel=tm.Param["Verbose"])
end

#Train the model on trainin set (X,y) for #iterations
function TrainModel(tm::TestingModel,iterations)
  time_training = 0;
  tm.Model.nEpochs = iterations
  time_training = @elapsed tm.Model.Train()
  return time_training;
end

#Run tests accordingly to the arguments and save them
function RunTests(tm::TestingModel,X_test,y_test;accuracy::Bool=false,brierscore::Bool=false,logscore::Bool=false)
  if accuracy
    push!(tm.Results["accuracy"],TestAccuracy(y_test,ComputePrediction(tm,X_test)))
  end
  y_predic_acc = 0
  if brierscore
    y_predic_acc = ComputePredictionAccuracy(tm::TestingModel, X_test)
    push!(tm.Results["brierscore"],TestBrierScore(y_test,y_predic_acc))
  end
  if logscore
    if y_predic_acc == 0
      y_predic_acc = ComputePredictionAccuracy(tm::TestingModel,X_test)
    end
    push!(tm.Results["logscore"],TestLogScore(y_test,y_predic_acc))
  end
end

#Compute the mean and the standard deviation and assemble in one result
function ProcessResults(tm::TestingModel,writing_order)
  all_results = Array{Float64,1}()
  names = Array{String,1}()
  for name in writing_order
    result = [mean(tm.Results[name]), std(tm.Results[name])]
    all_results = vcat(all_results,result)
    names = vcat(names,name)
  end
  if haskey(tm.Results,"allresults")
    tm.Results["allresults"] = vcat(tm.Results["allresults"],all_results')
  else
    tm.Results["allresults"] = all_results'
  end
  if !haskey(tm.Results,"names")
    tm.Results["names"] = names
  end
end

function PrintResults(results,method_name,writing_order)
  i = 1
  for category in writing_order
    println("$category : $(results[i*2-1]) ± $(results[i*2])")
    i+=1
  end
end

function WriteResults(tm::TestingModel,location)
  fold = String(location*"/"*tm.ExperimentType*"Experiment_"*tm.DatasetName*"Dataset")
  if !isdir(fold); mkdir(fold); end;
  writedlm(String(fold*"/Results_"*tm.MethodName*".txt"),tm.Results["allresults"])
end

#Return predicted labels (-1,1) for test set X_test
function ComputePrediction(tm::TestingModel, X_test)
  return sign.(tm.Model.Predict(X_test))
end

#Return prediction certainty for class 1 on test set X_test
function ComputePredictionAccuracy(tm::TestingModel, X_test)
  return tm.Model.PredictProba(X_test)
end

#Return Accuracy on test set
function TestAccuracy(y_test, y_predic)
  return 1-sum(1-y_test.*y_predic)/(2*length(y_test))
end
#Return Brier Score
function TestBrierScore(y_test, y_predic)
  return sum(((y_test+1)./2 - y_predic).^2)/length(y_test)
end
#Return Log Score
function TestLogScore(y_test, y_predic)
  return sum((y_test+1)./2.*log(y_predic)+(1-(y_test+1)./2).*log(1-y_predic))/length(y_test)
end
#Return ROC
function TestROC(y_test,y_predic)
    nt = length(y_test)
    truepositive = zeros(npoints); falsepositive = zeros(npoints)
    truenegative = zeros(npoints); falsenegative = zeros(npoints)
    thresh = collect(linspace(0,1,npoints))
    for i in 1:npoints
      for j in 1:nt
        truepositive[i] += (yp[j]>=thresh[i] && y_test[j]>=0.9) ? 1 : 0;
        truenegative[i] += (yp[j]<=thresh[i] && y_test[j]<=-0.9) ? 1 : 0;
        falsepositive[i] += (yp[j]>=thresh[i] && y_test[j]<=-0.9) ? 1 : 0;
        falsenegative[i] += (yp[j]<=thresh[i] && y_test[j]>=0.9) ? 1 : 0;
      end
    end
    return (truepositive./(truepositive+falsenegative),falsepositive./(truenegative+falsepositive))
end

end #end of module


using PyCall
using Cambrian
import Distances
import Random
import Formatting
import Base.GC

export GymEnv, fitness, new_gen!, gym_envs

import RDatasets

"Gym BERL environment"
mutable struct GymEnv <: BERLenv
    name::String
    memory::Array # Array of storable data
    gym
    env_name
    gen
    n_steps::Int64
end

environments["gym"]=GymEnv
gym_envs = ["Acrobot-v1", "CartPole-v1", "MountainCar-v0", "MountainCarContinuous-v0", "Pendulum-v0"]

function GymEnv(cfg::Dict)
    env_name = get!(cfg, "gym_env", "MountainCar-v0") # Default mountain car
    cfg["env"]="Gym-$env_name"

    gym = pyimport("gym")
    game = gym.make(env_name)
    cfg["n_out"] = game.action_space.n # length(env.action_space.sample())
    cfg["n_in"] = length(game.observation_space.sample())

    gen = get!(cfg, "seed", 0)
    Random.seed!(gen)

    GymEnv(cfg["env"], [], gym, env_name, gen, 0)
end

"Returns the fitness of one Cambrian Individual in this environment"
function fitness(indiv::Cambrian.Individual, env::GymEnv)
    game = env.gym.make(env.env_name)
    game.seed(env.gen)
    obs = game.reset()
    total_reward = 0.0
    done = false

    max_obs = Float64(max(-minimum(game.observation_space.low),
                          maximum(game.observation_space.high)))

    while ~done
        action_vect = process(indiv, obs ./ max_obs)
        # println(action_vect)
        action = argmax(action_vect)-1
        obs, reward, done, _ = game.step(action)
        total_reward += reward
        env.n_steps += 1
    end

    game.close()
    game = nothing
    Base.GC.gc()
    [total_reward]
end

"Sets up the env for the next generation"
function new_gen!(env::GymEnv, cfg::Dict=nothing)
    env.gen += 1
end


transD_GP.SurfaceRegression.plot_surface_posterior(opt, optlog10λ, cmapmean="RdBu", property="moho depth",
            property_units="km", cmappdf="viridis")
## plot shapefile
ax = gcf().axes
using Shapefile, ArchGDAL, GeoInterface, GeoFormatTypes
t = Shapefile.Table("/home/547/ar0754/Downloads/STE_2016_AUST.shp")
for iaxis in 1:2
    for row in t
        if (row.STE_NAME16 == "Queensland") || (row.STE_NAME16 == "Northern Territory")
            xy = GeoInterface.coordinates(Shapefile.shape(row))
            for (i, points) in enumerate(xy)
                XY = permutedims(reverse(reduce(hcat, points[1]), dims=1))
                reprojXY = Vector{Array{Float64,1}}(undef, size(XY,1))
                for i in 1:size(XY, 1)
                        reprojXY[i] = XY[i,:]
                end
                XY = permutedims(reduce(hcat, ArchGDAL.reproject(reprojXY, EPSG(4283), EPSG(4462))))
                ax[iaxis].plot(XY[:,1], XY[:,2], "-k", linewidth=0.5)
            end
        end
    end
end


# Universal properties
struct Map
    interstitialDirections::Vector{Vector{Int64}}
end
function Map(nDimension::Int64)
    Map(GenerateIntersitialDirections(nDimension))
end

struct Box
    map::Map
    nDimension::Int64
    halfLength::Vector{Float64}
    length::Vector{Int64}
    boundary::Matrix{Int64}
end
function Box(boundary::Matrix{Int64})
    nDimension = size(boundary)[1]
    length = boundary[:,2] - boundary[:,1]
    halfLength = length/2.0
    map = Map(nDimension)
    Box(map, nDimension, halfLength, length, boundary)
end

mutable struct Thermo
    time::Float64
    lastTime::Float64
    iterTime::Int64
    Temp::Float64
end
function Thermo(initTemp::Float64)
    Thermo(0.0, 0.0, 0, initTemp)
end

EventContainer{UniverseEvent}() = EventContainer{UniverseEvent}(0,[],[])

function CellContainer(box::Box, cellCutoff::Float64)
    nDimension = box.nDimension
    dCellIndexs = DCellIndexs(nDimension)
    lengths = Vector{Int64}(undef, nDimension)
    for d in 1:nDimension
        lo = fld(box.boundary[d,1],cellCutoff)
        hi = fld(box.boundary[d,2],cellCutoff)-1
        lengths[d] = hi-lo
    end
    cells = Array{Cell}(undef, Tuple(lengths))
    for index = 1:length(cells)
        cells[index] = Cell(Vector{DefectObj}())
    end
    CellContainer(cells, cellCutoff, dCellIndexs)
end

# Global type Universe
# Contain all the objects and universal properties
mutable struct Universe
    objs::Vector{DefectObj}
    cellContainer::CellContainer
    universeEventHolders::Vector{UniverseEventHolder}
    maxID::Int64
    fre::Float64
    box::Box
    thermo::Thermo
    dumpFileNames::Vector{String}
    cache::Dict{String, Union{Int64, Float64}}
    stop::Bool
end
function Universe(boundary, initTemp)
    box = Box(boundary)
    cellContainer = CellContainer(box, cellCutoff)
    thermo = Thermo(initTemp)
    universeEventContainer = EventContainer{UniverseEvent}()
    universeEventHolder = UniverseEventHolder(universeEventContainer)
    Universe(Vector{DefectObj}(), cellContainer, [universeEventHolder],
            0, 0, box, thermo,
            Vector{String}(), Dict{String, Union{Int64, Float64}}(), false)
end

# Methods of Universe
function Base.getindex(universe::Universe, n::Int64)
    universe.objs[n]
end
function Base.push!(universe::Universe, obj::DefectObj)
    universe.maxID += 1
    obj.id = universe.maxID
   # if obj.id == 3224
   #     push!(debug, obj)
   # end
    push!(universe.objs, obj)
    universe.fre += obj.eventContainer.fre
end
function Base.push!(universe::Universe, obj::UniverseEventHolder)
    push!(universe.universeEventHolders, obj)
    universe.fre += obj.eventContainer.fre
end
function Base.delete!(universe::Universe, obj::DefectObj)
    deleteat!(universe.objs, findfirst(member->member==obj,universe.objs))
    universe.fre -= obj.eventContainer.fre
end
function Base.iterate(universe::Universe)
    iterate(universe.objs)
end
function Base.iterate(universe::Universe, n::Int64)
    iterate(universe.objs, n)
end
function Base.length(universe::Universe)
    length(universe.objs)
end


############################################################################
### models w.o. meanstructure
############################################################################

model_ml = Sem(
    specification = spec,
    data = dat,
    diff = semdiff
)

model_ls_sym = Sem(
    specification = spec,
    data = dat,
    imply = RAMSymbolic,
    loss = SemWLS,
    start_val = start_simple,
    diff = semdiff
)

model_ml_sym = Sem(
    specification = spec,
    data = dat,
    imply = RAMSymbolic,
    diff = semdiff
)

model_ridge = Sem(
    specification = spec,
    data = dat,
    loss = (SemML, SemRidge),
    α_ridge = .001,
    which_ridge = 16:20,
    diff = semdiff
)

model_constant = Sem(
    specification = spec,
    data = dat,
    loss = (SemML, SemConstant),
    constant_loss = 3.465,
    diff = semdiff
)

model_ml_weighted = Sem(
    specification = partable,
    data = dat,
    loss_weights = (n_obs(model_ml),),
    diff = semdiff
)

############################################################################
### test gradients
############################################################################

models = [model_ml, model_ls_sym, model_ridge, model_constant, model_ml_sym, model_ml_weighted]
names = ["ml", "ls_sym", "ridge", "constant", "ml_sym", "ml_weighted"]

for (model, name) in zip(models, names)
    try
        @testset "$(name)_gradient" begin
            @test test_gradient(model, start_test; rtol = 1e-9)
        end
    catch
    end
end

############################################################################
### test solution
############################################################################

models = [model_ml, model_ls_sym, model_ml_sym, model_constant]
names = ["ml", "ls_sym", "ml_sym", "constant"]
solution_names = Symbol.("parameter_estimates_".*["ml", "ls", "ml", "ml"])

for (model, name, solution_name) in zip(models, names, solution_names)
    try
        @testset "$(name)_solution" begin
            solution = sem_fit(model)
            update_estimate!(partable, solution)
            @test compare_estimates(partable, solution_lav[solution_name]; atol = 1e-2)
        end
    catch
    end
end

@testset "ridge_solution" begin
    solution_ridge = sem_fit(model_ridge)
    solution_ml = sem_fit(model_ml)
    # solution_ridge_id = sem_fit(model_ridge_id)
    @test abs(solution_ridge.minimum - solution_ml.minimum) < 1
end

# test constant objective value
@testset "constant_objective_and_gradient" begin
    @test (objective!(model_constant, start_test) - 3.465) ≈ objective!(model_ml, start_test)
    @test gradient!(model_constant, start_test) ≈ gradient!(model_ml, start_test)
end

@testset "ml_solution_weighted" begin
    solution_ml = sem_fit(model_ml)
    solution_ml_weighted = sem_fit(model_ml_weighted)
    @test isapprox(solution(solution_ml), solution(solution_ml_weighted), rtol = 1e-3)
    @test isapprox(n_obs(model_ml)*StructuralEquationModels.minimum(solution_ml), StructuralEquationModels.minimum(solution_ml_weighted), rtol = 1e-6)
end

############################################################################
### test fit assessment
############################################################################

@testset "fitmeasures/se_ml" begin
    solution_ml = sem_fit(model_ml)
    @test all(test_fitmeasures(fit_measures(solution_ml), solution_lav[:fitmeasures_ml]; atol = 1e-3))

    update_partable!(partable, identifier(model_ml), se_hessian(solution_ml), :se)
    @test compare_estimates(partable, solution_lav[:parameter_estimates_ml]; atol = 1e-3, col = :se, lav_col = :se)
end

@testset "fitmeasures/se_ls" begin
    solution_ls = sem_fit(model_ls_sym)
    fm = fit_measures(solution_ls)
    @test all(test_fitmeasures(fm, solution_lav[:fitmeasures_ls]; atol = 1e-3, fitmeasure_names = fitmeasure_names_ls))
    @test (fm[:AIC] === missing) & (fm[:BIC] === missing) & (fm[:minus2ll] === missing)

    update_partable!(partable, identifier(model_ls_sym), se_hessian(solution_ls), :se)
    @test compare_estimates(partable, solution_lav[:parameter_estimates_ls]; atol = 1e-2, col = :se, lav_col = :se)
end

############################################################################
### test hessians
############################################################################

if semdiff == SemDiffOptim
    using Optim, LineSearches

    model_ls = Sem(
        specification = spec,
        data = dat,
        imply = RAMSymbolic,
        loss = SemWLS,
        hessian = true,
        algorithm = Newton(
            ;linesearch = BackTracking(order=3),
            alphaguess = InitialHagerZhang())
    )

    model_ml = Sem(
        specification = spec,
        data = dat,
        imply = RAMSymbolic,
        hessian = true,
        algorithm = Newton()
    )

    @testset "ml_hessians" begin
        @test test_hessian(model_ml, start_test; atol = 1e-4)
    end

    @testset "ls_hessians" begin
        @test test_hessian(model_ls, start_test; atol = 1e-4)
    end

    @testset "ml_solution_hessian" begin
        solution = sem_fit(model_ml)
        update_estimate!(partable, solution)
        @test compare_estimates(partable, solution_lav[:parameter_estimates_ml]; atol = 1e-3)
    end

    @testset "ls_solution_hessian" begin
        solution = sem_fit(model_ls)
        update_estimate!(partable, solution)
        @test compare_estimates(partable, solution_lav[:parameter_estimates_ls]; atol = 1e-3)
    end

end

############################################################################
### meanstructure
############################################################################

# models
model_ls = Sem(
    specification = spec_mean,
    data = dat,
    imply = RAMSymbolic,
    loss = SemWLS,
    meanstructure = true,
    diff = semdiff
)

model_ml = Sem(
    specification = spec_mean,
    data = dat,
    meanstructure = true,
    diff = semdiff
)

model_ml_sym = Sem(
    specification = spec_mean,
    data = dat,
    imply = RAMSymbolic,
    meanstructure = true,
    start_val = start_test_mean,
    diff = semdiff
)

############################################################################
### test gradients
############################################################################

models = [model_ml, model_ls, model_ml_sym]
names = ["ml", "ls_sym", "ml_sym"]

for (model, name) in zip(models, names)
    try
        @testset "$(name)_gradient_mean" begin
            @test test_gradient(model, start_test_mean; rtol = 1e-9)
        end
    catch
    end
end

############################################################################
### test solution
############################################################################

solution_names = Symbol.("parameter_estimates_".*["ml", "ls", "ml"].*"_mean")

for (model, name, solution_name) in zip(models, names, solution_names)
    try
        @testset "$(name)_solution_mean" begin
            solution = sem_fit(model)
            update_estimate!(partable_mean, solution)
            @test compare_estimates(partable_mean, solution_lav[solution_name]; atol = 1e-2)
        end
    catch
    end
end

############################################################################
### test fit assessment
############################################################################

@testset "fitmeasures/se_ml_mean" begin
    solution_ml = sem_fit(model_ml)
    @test all(test_fitmeasures(fit_measures(solution_ml), solution_lav[:fitmeasures_ml_mean]; atol = 1e-3))

    update_partable!(partable_mean, identifier(model_ml), se_hessian(solution_ml), :se)
    @test compare_estimates(partable_mean, solution_lav[:parameter_estimates_ml_mean]; atol = 1e-3, col = :se, lav_col = :se)
end

@testset "fitmeasures/se_ls_mean" begin
    solution_ls = sem_fit(model_ls)
    fm = fit_measures(solution_ls)
    @test all(test_fitmeasures(fm, solution_lav[:fitmeasures_ls_mean]; atol = 1e-3, fitmeasure_names = fitmeasure_names_ls))
    @test (fm[:AIC] === missing) & (fm[:BIC] === missing) & (fm[:minus2ll] === missing)

    update_partable!(partable_mean, identifier(model_ls), se_hessian(solution_ls), :se)
    @test compare_estimates(partable_mean, solution_lav[:parameter_estimates_ls_mean]; atol = 1e-2, col = :se, lav_col = :se)
end

############################################################################
### fiml
############################################################################

# models
model_ml = Sem(
    specification = spec_mean,
    data = dat_missing,
    observed = SemObsMissing,
    loss = SemFIML,
    diff = semdiff
)

model_ml_sym = Sem(
    specification = spec_mean,
    data = dat_missing,
    observed = SemObsMissing,
    imply = RAMSymbolic,
    loss = SemFIML,
    start_val = start_test_mean,
    diff = semdiff
)

############################################################################
### test gradients
############################################################################

@testset "fiml_gradient" begin
    @test test_gradient(model_ml, start_test_mean; atol = 1e-6)
end

@testset "fiml_gradient_symbolic" begin
    @test test_gradient(model_ml_sym, start_test_mean; atol = 1e-6)
end

############################################################################
### test solution
############################################################################

@testset "fiml_solution" begin
    solution = sem_fit(model_ml)
    update_estimate!(partable_mean, solution)
    @test compare_estimates(partable_mean, solution_lav[:parameter_estimates_fiml]; atol = 1e-2)
end

@testset "fiml_solution_symbolic" begin
    solution = sem_fit(model_ml_sym)
    update_estimate!(partable_mean, solution)
    @test compare_estimates(partable_mean, solution_lav[:parameter_estimates_fiml]; atol = 1e-2)
end

############################################################################
### test fit measures
############################################################################

@testset "fitmeasures/se_fiml" begin
    solution_ml = sem_fit(model_ml)
    @test all(test_fitmeasures(fit_measures(solution_ml), solution_lav[:fitmeasures_fiml]; atol = 1e-3))

    update_partable!(partable_mean, identifier(model_ml), se_hessian(solution_ml), :se)
    @test compare_estimates(partable_mean, solution_lav[:parameter_estimates_fiml]; atol = 1e-3, col = :se, lav_col = :se)
end


module SfMaterials

using Unitful, ..SfMatter

aluminium = Material(2.7u"g/cm^3", 270u"MPa") # yield strength for 6061 alloy, https://en.wikipedia.org/wiki/6061_aluminium_alloy
iron = Material(7.874u"g/cm^3")
steel = Material(7700u"kg/m^3", 350u"MPa") # 1020 steel, http://www.matweb.com/search/datasheet.aspx?bassnum=MS0001
titanium = Material(4.506u"g/cm^3", 140u"MPa") # http://www.matweb.com/search/datasheet.aspx?bassnum=METi00&ckck=1
titanium_alloy = Material(4.51u"g/cm^3", 830u"MPa") # 6% Al, 4% V, https://en.wikipedia.org/wiki/Yield_(engineering)
tungsten = Material(19.25u"g/cm^3", 550u"MPa") # annealed, https://en.wikipedia.org/wiki/Yield_(engineering)
water = Material(997u"kg/m^3", 0u"Pa")
ice = Material(0.9168u"g/cm^3", 0u"Pa")

end


type LRNDescriptor
  ptr::Ptr{Void}
end

function LRNDescriptor(n::Float64, alpha::Float64, beta::Float64, k::Float64)
  p = Ptr{Void}[0]
  cudnnCreateLRNDescriptor(p)
  cudnnSetLRNDescriptor(p[1], n, alpha, beta, k)
  desc = new(p[1])
  finalizer(desc, cudnnDestroyLRNDescriptor)
  desc
end

Base.unsafe_convert(::Type{Ptr{Void}}, desc::LRNDescriptor) = desc.ptr

function crosschanel()
  cudnnLRNCrossChannelForward()
end


# This file is part of Kpax3. License is MIT.

mutable struct KWeight
  c::Vector{Float64}
  w::Matrix{Float64}
  z::Matrix{Float64}
end

mutable struct MCMCSupport
  m::Int
  n::Int

  u::Vector{Int}
  t::Vector{Float64}

  lp::Array{Float64, 3}
  lq::Array{Float64, 3}
  lr::Array{Float64, 3}

  vi::Int
  ni::Vector{Float64}
  ui::Vector{Int}
  wi::KWeight
  lpi::Matrix{Float64}

  vj::Int
  nj::Vector{Float64}
  uj::Vector{Int}
  wj::KWeight
  lpj::Matrix{Float64}

  tmp::Vector{Float64}

  cl::Vector{Int}
  k::Int

  lograR::Float64

  logmlik::Float64
  logmlikcandidate::Float64
end

function MCMCSupport(state::State,
                     priorC::AminoAcidPriorCol)
  n = length(state.R)
  (maxclust, m) = size(state.C)

  u = Int[a for a in 1:n]
  t = zeros(Float64, n)

  lp = zeros(Float64, 4, maxclust, m)
  lq = zeros(Float64, 4, maxclust, m)
  lr = zeros(Float64, 3, maxclust, m)

  g = 0
  for b in 1:m, l in 1:state.k
    g = state.cl[l]

    lp[1, g, b] = logmarglik(state.n1s[g, b], state.v[g], priorC.A[1, b],
                             priorC.B[1, b])
    lp[2, g, b] = logmarglik(state.n1s[g, b], state.v[g], priorC.A[2, b],
                             priorC.B[2, b])
    lp[3, g, b] = logmarglik(state.n1s[g, b], state.v[g], priorC.A[3, b],
                             priorC.B[3, b])
    lp[4, g, b] = logmarglik(state.n1s[g, b], state.v[g], priorC.A[4, b],
                             priorC.B[4, b])
  end

  vi = 0
  ni = zeros(Float64, m)
  ui = zeros(Int, n)
  wi = KWeight(zeros(Float64, m), zeros(Float64, 4, m), zeros(Float64, 4, m))

  lpi = zeros(Float64, 4, m)

  vj = 0
  nj = zeros(Float64, m)
  uj = zeros(Int, n)
  wj = KWeight(zeros(Float64, m), zeros(Float64, 4, m), zeros(Float64, 4, m))

  lpj = zeros(Float64, 4, m)

  tmp = zeros(Float64, 4)

  cl = zeros(Int, n)
  k = 0

  logmlik = logmarglikelihood(state.cl, state.k, lp, priorC)

  MCMCSupport(m, n, u, t, lp, lq, lr, vi, ni, ui, wi, lpi, vj, nj, uj, wj, lpj,
              tmp, cl, k, 0.0, logmlik, 0.0)
end

function resizesupport!(support::MCMCSupport,
                        maxclust::Int)
  if size(support.lp, 2) < maxclust
    # we don't want to allocate new resources too often, so allocate double the
    # previous size. this should guarantee a logarithmic number of allocations
    newsize = min(support.n, max(maxclust, 2 * size(support.lp, 2)))

    lp = zeros(Float64, 4, newsize, support.m)

    for b in 1:support.m, g in 1:size(support.lp, 2)
      lp[1, g, b] = support.lp[1, g, b]
      lp[2, g, b] = support.lp[2, g, b]
      lp[3, g, b] = support.lp[3, g, b]
      lp[4, g, b] = support.lp[4, g, b]
    end

    support.lp = lp
    support.lq = zeros(Float64, 4, newsize, support.m)
    support.lr = zeros(Float64, 3, newsize, support.m)
  end

  nothing
end

mutable struct KOffspring
  R::Vector{Int}
  v::Vector{Int}
end

mutable struct GASupport
  m::Int
  n::Int

  oi::KOffspring
  oj::KOffspring
end

function GASupport(m::Int,
                   n::Int)
  oi = KOffspring(zeros(Int, n), zeros(Int, n))
  oj = KOffspring(zeros(Int, n), zeros(Int, n))
  GASupport(m, n, oi, oj)
end


module PhysicsCodesLabXD

# The structure of the package is the following:

# 1. Define tensor category objects: simple objects, dual, braiding structure, tensor product

end


function grid_init_do_not_profile(inputs::Inputs, mype::Int32)::SimulationData

    ## Structure to hold all allocated simuluation data arrays
    simulation_data = SimulationData()

    ## Keep track of how much data we're allocating
    nbytes::UInt64 = 0

    ## Set the initial seed value
    seed::UInt64 = 42

    ##############################
    ## Initialize Nuclide Grids
    ##############################

    if mype == 0
        println("Initializing nuclide grids...")
    end

    ## First, we need to initialize our nuclide grid. This comes in the form
    ## of a flattened 2D array that hold all the information we need to define
    ## the cross sections for all isotopes in the simulation.
    ## The grid is composed of "NuclideGridPoint" structures, which hold the
    ## energy level of the grid point and all associated XS data at that level.
    ## An array of structures (AOS) is used instead of
    ## a structure of arrays, as the grid points themselves are accessed in
    ## a random order, but all cross section interaction channels and the
    ## energy level are read whenever the gridpoint is accessed, meaning the
    ## AOS is more cache efficient.

    ## Initialize Nuclide Grid
    simulation_data.length_nuclide_grid = inputs.n_isotopes * inputs.n_gridpoints

    resize!(simulation_data.nuclide_grid, simulation_data.length_nuclide_grid)

    nbytes += size(simulation_data.nuclide_grid)[1] * sizeof(NuclideGridPoint)
    println("nbytes: ", nbytes)

    return simulation_data

end