module myTesting
using SpinLattice
using PyCall, RecursiveArrayTools
unshift!(PyVector(pyimport("sys")["path"]), "")
@pyimport mkl
mkl.set_num_threads(1)
@pyimport madness as oop
@pyimport scipy.sparse as pysparse
@pyimport scipy.sparse.linalg as slinalg

L1d_1 = ((15,), (-0.41,-0.41,0.82), Interaction(nearest_neighbours), (0.,0.,0.))
L1d_2 = ((15,), (1.0,1.0,1.0), Interaction(nearest_neighbours), (0.,0.,0.))
L1d_3 = ((15,), (0.707,0.707,0.0), Interaction(nearest_neighbours), (0.,0.,0.))
L1d_4 = ((15,), (0.518, 0.830, 0.207), Interaction(nearest_neighbours), (0.,0.,0.))
L1d_5 = ((15,), (0.518, 0.830, 0.207), Interaction(nearest_neighbours), (0.207,0.,0.))
L1d_6 = ((15,), (0.518, 0.830, 0.207), Interaction(nearest_neighbours), (0.207,0.,0.100))
L1d_7 = ((15,), (0.518, 0.830, 0.207), Interaction(nearest_neighbours), (0.207,0.,0.207))

L2d_1 = ((9,9), (-0.41,-0.41,0.82), Interaction(nearest_neighbours), (0.,0.,0.))
L2d_2 = ((9,9), (1.0,1.0,1.0), Interaction(nearest_neighbours), (0.,0.,0.))
L2d_3 = ((9,9), (0.707,0.707,0.0), Interaction(nearest_neighbours), (0.,0.,0.))
L2d_4 = ((9,9), (0.518, 0.830, 0.207), Interaction(nearest_neighbours), (0.,0.,0.))
L2d_5 = ((9,9), (0.518, 0.830, 0.207), Interaction(nearest_neighbours), (0.207,0.,0.))
L2d_6 = ((9,9), (0.518, 0.830, 0.207), Interaction(nearest_neighbours), (0.207,0.,0.100))
L2d_7 = ((9,9), (0.518, 0.830, 0.207), Interaction(nearest_neighbours), (0.207,0.,0.207))

L3d_1 = ((9,9,9), (-0.41,-0.41,0.82), Interaction(nearest_neighbours), (0.,0.,0.))
Lfid0x0x1 = ((9,9,9), (-0.5,-0.5,1.0), Interaction(MDI(0,0,1)), (0.,0.,0.), [(1,1,i) for i in 1:9], [(1,1,i) for i in 1:9])
Lfid0x1x1 = ((9,9,9), (-0.5,-0.5,1.0), Interaction(MDI(0,1,1)), (0.,0.,0.), [(i,1,1) for i in 1:9], [(i,1,1) for i in 1:9])
Lfid1x1x1 = ((9,9,9), (-0.5,-0.5,1.0), Interaction(MDI(1,1,1)), (0.,0.,0.), [(i,i,i) for i in 1:9], [(i,i,i) for i in 1:9])

L1d = [L1d_1, L1d_2, L1d_3, L1d_4, L1d_5, L1d_6, L1d_7]
L2d = [L2d_1, L2d_2, L2d_3, L2d_4, L2d_5, L2d_6, L2d_7]
L3d = [L3d_1, Lfid0x0x1, Lfid0x1x1, Lfid1x1x1]

d1block = (5,)
d2block = (3,3)
d3block = (:cross, sort([(2,2,2),(2,2,1),(2,1,2),(1,2,2),(2,2,3),(2,3,2),(3,2,2)]), [(2,2,2)])
d3block_py = ([Tuple(reverse(collect(ind) .- 1)) for ind in d3block[2]], [Tuple(reverse(collect(ind) .- 1)) for ind in d3block[3]])
get_block(::Type{Val{1}}) = d1block
get_block(::Type{Val{2}}) = d2block
get_block(::Type{Val{3}}) = d3block
get_block_py(x::Type{Val{D}}) where D = get_block(x)
get_block_py(::Type{Val{3}}) = d3block_py

jlmat2pymat(S::Union{SparseMatrixCSC, SpinLattice.SharedSparseMatrixCSC}) = pysparse.csc_matrix((S.nzval, S.rowval .- 1, S.colptr .- 1), shape=size(S))
function py_initialize(M::Exact, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{typeof(nearest_neighbours)}, hs::NTuple{3,Float64}; kwargs...) where {D}
    sp = oop.UniformLattice(reverse(dims), hs, Js; kwargs...)
    return sp
end
function py_initialize(M::Exact, dims::NTuple{2,Int64}, Js::NTuple{3,Float64}, I::Interaction{typeof(nearest_neighbours)}, hs::NTuple{3,Float64}; kwargs...)
    sp = oop.UniformLattice(reverse((4,4)), hs, Js; kwargs...)
    return sp
end
function py_initialize(M::Clustered, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{typeof(nearest_neighbours)}, hs::NTuple{3,Float64}; kwargs...) where {D}
    dims_r = reverse(dims)
    block_dims = reverse(get_block_py(Val{D}))
    blocks = Tuple(div.(collect(dims_r), collect(block_dims)))
    sp = oop.ClusteredLattice(block_dims, blocks, hs, Js; kwargs...)
    return sp
end
function py_initialize(M::PureClassical, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{typeof(nearest_neighbours)}, hs::NTuple{3,Float64}; kwargs...) where {D}
    sp = oop.ClassicalLattice(reverse(dims), hs, Js; kwargs...)
    return sp
end
function py_initialize(M::Hybrid, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{typeof(nearest_neighbours)}, hs::NTuple{3,Float64}; kwargs...) where {D}
    dims_r = reverse(dims)
    block_dims = reverse(get_block_py(Val{D}))
    sp = oop.HybridLattice(block_dims, dims_r, hs, Js; kwargs...)
    return sp
end
function py_initialize(M::Hybrid, dims::NTuple{3,Int64}, Js::NTuple{3,Float64}, I::Interaction{typeof(nearest_neighbours)}, hs::NTuple{3,Float64}; kwargs...)
    q_spins, middle_spins = d3block_py
    q_spins_r = [reverse(x) for x in q_spins]
    sp = oop.HybridLatticeIR(q_spins_r, reverse(dims), middle_spins, hs, Js; kwargs...)
    return sp
end
function py_initialize(M::Hybrid, dims::NTuple{3,Int64}, Js::NTuple{3,Float64}, I::Interaction{MDI{3}}, hs::NTuple{3,Float64}; kwargs...)
    q_spins, middle_spins = d3block_py
    q_spins_r = [reverse(x) for x in q_spins]
    sp = oop.FIDLattice(3, reverse(dims), q_spins_r, middle_spins, direction = reverse(I.func.dir); kwargs...)
    return sp
end
function py_initialize(M::Hybrid, dims::NTuple{3,Int64}, Js::NTuple{3,Float64}, I::Interaction{MDI{3}}, hs::NTuple{3,Float64}, q_spins::Vector{NTuple{3,Int64}}, middle_spins::Vector{NTuple{3,Int64}}; kwargs...)
    q_spins_r = [Tuple(reverse(collect(ind) .- 1)) for ind in q_spins]
    middle_spins_r = [Tuple(reverse(collect(ind) .- 1)) for ind in middle_spins]
    sp = oop.FIDLattice(3, reverse(dims), q_spins_r, middle_spins_r, direction = reverse(I.func.dir); kwargs...)
    return sp
end
function py_initialize(M::PureClassical, dims::NTuple{3,Int64}, Js::NTuple{3,Float64}, I::Interaction{MDI{3}}, hs::NTuple{3,Float64}; kwargs...)
    sp = oop.FIDLatticeClassical(3, reverse(dims), direction = reverse(I.func.dir); kwargs...)
    return sp
end
@inline function py_initialize(M::PureClassical, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, q_spins::Vector{NTuple{3,Int64}}, middle_spins::Vector{NTuple{3,Int64}}; kwargs...) where {D,F}
    py_initialize(M::PureClassical, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}; kwargs...)
end

@inline default_links(L::Lattice, M::Model) where Model<:Union{Exact, PureClassical} = (:all,:all)
@inline default_links(L::Lattice{D}, M::Model) where {D, Model<:Union{Clustered,Hybrid}} = (:all,:central)
@inline default_links(L::Lattice{3}, M::Hybrid) = (:all, SpinArray(L, [(2,2,2)], :central))

function julia_initialize(M::Model, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, B::Bool=true) where {D, Model<:Union{Exact,PureClassical}} where F
    L = Lattice(dims, Js, I, hs, Int64, B)
    A = build_Approximation(M, L)
    RHS = build_RHS_function(A)
    rules, OSET = set_CorrFuncs(A, default_links(L,M))
    return A, RHS, OSET, rules
end
@inline function julia_initialize(M::PureClassical, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, q_spins::Vector{NTuple{3,Int64}}, middle_spins::Vector{NTuple{3,Int64}}, B::Bool=true) where {D, F}
    julia_initialize(M, dims, Js, I, hs, Int64, B)
end
function julia_initialize(M::Exact, dims::NTuple{2,Int64}, Js::NTuple{3,Float64}, I::Interaction{typeof(nearest_neighbours)}, hs::NTuple{3,Float64}, B::Bool=true)
    L = Lattice((4,4), Js, I, hs, Int64, B)
    A = build_Approximation(M, L)
    RHS = build_RHS_function(A)
    rules, OSET = set_CorrFuncs(A, default_links(L,M))
    return A, RHS, OSET, rules
end

function julia_initialize(M::Model, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, B::Bool=true) where {D, Model<:Union{Clustered,Hybrid}, F}
    L = Lattice(dims, Js, I, hs, Int64, B)
    A = build_Approximation(M, L, get_block(Val{D}))
    RHS = build_RHS_function(A)
    rules, OSET = set_CorrFuncs(A, default_links(L,M))
    return A, RHS, OSET, rules
end

function julia_initialize(M::Hybrid, dims::NTuple{3,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, B::Bool=true) where F
    L = Lattice(dims, Js, I, hs, Int64, B)
    name, q_spins, middle_spins = d3block
    sarray = SpinArray(L, q_spins, name)
    A = build_Approximation(M, L, sarray)
    RHS = build_RHS_function(A)
    rules, OSET = set_CorrFuncs(A, default_links(L,M))
    return A, RHS, OSET, rules
end
function julia_initialize(M::Hybrid, dims::NTuple{3,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, q_spins::Vector{NTuple{3,Int64}}, middle_spins::Vector{NTuple{3,Int64}}, B::Bool=true) where F
    L = Lattice(dims, Js, I, hs, Int64, B)
    name = :column
    sarray = SpinArray(L, q_spins, name)
    A = build_Approximation(M, L, sarray)
    RHS = build_RHS_function(A)
    rules, OSET = set_CorrFuncs(A, (:all, SpinArray(L, middle_spins, :central)))
    return A, RHS, OSET, rules
end

function test_julia_init()
    Models = [Exact(), PureClassical(), Clustered(), Hybrid()]
    for model in Models
        for lat in L1d
            try
                julia_initialize(model, lat...)
                print("all green\n\n")
            catch
                print("$model   $lat\n\n")
            end
        end
    end
    for model in Models
        for lat in L2d
            try
                julia_initialize(model, lat...)
                print("all green\n\n")
            catch
                print("$model   $lat\n\n")
            end
        end
    end
    for model in Models[[2,4]]
        for lat in L3d
            try
                julia_initialize(model, lat...)
                print("all green\n\n")
            catch
                print("$model   $lat\n\n")
            end
        end
    end
end
function test_py_init()
    Models = [Exact(), PureClassical(), Clustered(), Hybrid()]
    for model in Models
        for lat in L1d
            try
                py_initialize(model, lat...)
                print("all green\n\n")
            catch
                print("$model   $lat\n\n")
            end
        end
    end
    for model in Models
        for lat in L2d
            try
                py_initialize(model, lat...)
                print("all green\n\n")
            catch
                print("$model   $lat\n\n")
            end
        end
    end
    for model in Models[[2,4]]
        for lat in L3d
            try
                py_initialize(model, lat...)
                print("all green\n\n")
            catch
                print("$model   $lat\n\n")
            end
        end
    end
end
function test_py_init_3d()
    Models = [Exact(), PureClassical(), Clustered(), Hybrid()]
    for model in Models[[2,4]]
        for lat in L3d
            try
                py_initialize(model, lat...)
                print("all green\n\n")
            catch
                print("$model   $lat\n\n")
            end
        end
    end
end

const toler = 1e-14

function convert_state(A::ExactApprox, sp)
    v = Vector{Vector{Complex{Float64}}}()
    push!(v, copy(sp[:psi]))
    for i = 1:size(sp[:psi_aid],1)
        push!(v, sp[:psi_aid][i,:])
    end
    return VectorOfArray(v)
end

function test_operators(M::Exact, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, B::Bool=true) where {D, F}
    A,RHS,OSET, rules = julia_initialize(M, dims, Js, I, hs, B)
    sp = py_initialize(M, dims, Js, I, hs; cb=B)
    H_j = jlmat2pymat(RHS.H)
    Ms_j = Dict()
    for Obs in OSET.Observables
        axis = axis_dict[Obs.axis]-1
        Ms_j[axis] = jlmat2pymat(Obs.QO.O)
    end
    sp[:initState]()
    u0 = copy(sp[:psi])
    du = copy(u0)
    tj = -time()
    RHS(0.0,u0,du)
    tj += time()
    tp = -time()
    du2 = sp[:rhs](sp[:psi])
    tp += time()
    print("$((dims, Js, I, hs))\n")
    @printf("Check Hamiltonian:   %e\n", slinalg.norm(sp[:H]-H_j)/slinalg.norm(sp[:H]))
    @printf("Check tot magnetiz: %e   %e   %e\n", [slinalg.norm(sp[:Ms][i]-Ms_j[i])/slinalg.norm(sp[:Ms][i]) for i in 0:2]...)
    @printf("Check RHS: %e\n", norm(du.*(2.0^-7)-du2)/norm(du2))
    @printf("Check time: %10e", tj/tp)
    print("\n\n")

end

function convert_state(A::ClusteredApprox, sp)
    cluster_num = size(sp[:psi],1)
    u0 = VectorOfArray([copy(sp[:psi][i,:]) for i in 1:cluster_num])
    return u0
end

function test_operators(M::Clustered, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}) where {D, F}
    d = Dict([(:all,"Global"), (:central,"Local")])
    A,RHS,OSET, rules = julia_initialize(M, dims, Js, I, hs)
    sp = py_initialize(M, dims, Js, I, hs)
    H_j = jlmat2pymat(RHS.H)
    ops_j = Dict()
    for Obs in OSET.Observables
        index = (string(Obs.axis), d[Obs.name])
        ops_j[index] = jlmat2pymat(Obs.QO.O)
    end
    num_of_espins = RHS.num_of_espins
    taus_j = Matrix(num_of_espins, 3)
    for espin in 1:num_of_espins for sigma in 1:3
        taus_j[espin, sigma] = jlmat2pymat(RHS.edgeSpinOperators[sigma, espin])
    end end
    sp[:initState]()
    cluster_num = size(sp[:psi],1)
    u0 = VectorOfArray([copy(sp[:psi][i,:]) for i in 1:cluster_num])
    du = VectorOfArray([copy(sp[:psi][i,:]) for i in 1:cluster_num])
    tj = -time()
    RHS(0.0,u0,du)
    tj += time()
    tp = -time()
    du2 = sp[:rhs](sp[:psi])
    tp += time()

    print("$((dims, Js, I, hs))\n")
    @printf("Check Hamiltonian:   %e\n", slinalg.norm(sp[:H]-H_j)/slinalg.norm(sp[:H]))
    print("Check spin operators:\n")
    for espin in 1:num_of_espins
        for sigma in 1:3
            @printf("%10e ", slinalg.norm(sp[:taus][espin, sigma]-taus_j[espin, sigma])/slinalg.norm(sp[:taus][espin, sigma]))
        end
        print("\n")
    end
    ###################
    links = sp[:links]
    noe = num_of_espins
    cl_num = cluster_num
    len = noe*cl_num
    IMp = zeros(len, len)
    for link in links
        p1 = sub2ind((noe,cl_num), link[1][2]+1, link[1][1]+1)
        p2 = sub2ind((noe,cl_num), link[2][2]+1, link[2][1]+1)
        IMp[p1,p2] += 1.0
    end
    IMp .= .-IMp.*0.25.*sqrt(sp[:N])
    @printf("Check intrablock matrix: %10e\n", norm(IMp-RHS.IM)/norm(IMp))
    ####################
    print("Check observables:\n")
    print("Local: ")
    for axis in ["x","y","x"]
        index = (axis,"Local")
        @printf("%10e ", slinalg.norm(sp[:ops][index] - ops_j[index])/slinalg.norm(sp[:ops][index]))
    end
    print("\n")
    print("Global: ")
    for axis in ["x","y","x"]
        index = (axis,"Global")
        @printf("%10e ", slinalg.norm(sp[:ops][index] - ops_j[index])/slinalg.norm(sp[:ops][index]))
    end
    print("\n")
    ###################
    print("Check rhs: ")
    for i in 1:cluster_num
        @printf("%10e ", norm(du[:,i].*2.0^-7-du2[i,:])/norm(du2[i,:]))
    end
    print("\n")
    @printf("Check time: %10e", tj/tp)
    print("\n\n")
end

function convert_state(A::HybridApprox, sp)
    u0 = ArrayPartition(copy(sp[:psi]), VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3]))
    return u0
end

function test_operators(M::Hybrid, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}) where {D, F}
    d = Dict([(:all,"Global"), (:central,"Local")])
    A,RHS,OSET, rules = julia_initialize(M, dims, Js, I, hs)
    sp = py_initialize(M, dims, Js, I, hs)
    H_j = jlmat2pymat(RHS.H)
    ops_j = Dict()
    for Obs in OSET.Observables
        index = (axis_dict[Obs.axis]-1, d[Obs.name])
        ops_j[index] = (jlmat2pymat(Obs.QO.O),Obs)
    end
    num_of_espins = size(RHS.edgeSpinOperators,1)
    taus_j = Matrix(num_of_espins, 3)
    for espin in 1:num_of_espins for sigma in 1:3
        taus_j[espin, sigma] = jlmat2pymat(RHS.edgeSpinOperators[espin, sigma])
    end end
    sp[:initState]()
    sp[:updateCorrelator]()
    u0 = ArrayPartition(copy(sp[:psi]), VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3]))
    du = ArrayPartition(copy(sp[:psi]), VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3]))
    tj = -time()
    RHS(0.0,u0,du)
    tj += time()
    tp = -time()
    du2 = sp[:rhs](sp[:psi], sp[:clState])
    tp += time()

    print("$((dims, Js, I, hs))\n")
    @printf("Check Hamiltonian:   %e\n", slinalg.norm(sp[:H]-H_j)/slinalg.norm(sp[:H]))
    print("Check spin operators:\n")
    for espin in 1:num_of_espins
        for sigma in 1:3
            @printf("%10e ", slinalg.norm(sp[:taus][espin, sigma]-taus_j[espin, sigma])/slinalg.norm(sp[:taus][espin, sigma]))
        end
        print("\n")
    end
    ###################
    print("Check observable ops:\n")
    print("Local: ")
    for axis in [0,1,2]
        index = (axis,"Local")
        @printf("%10e ", slinalg.norm(sp[:ops][index] - ops_j[index][1])/slinalg.norm(sp[:ops][index]))
    end
    print("\n")
    print("Global: ")
    for axis in [0,1,2]
        index = (axis,"Global")
        @printf("%10e ", slinalg.norm(sp[:ops][index] - ops_j[index][1])/slinalg.norm(sp[:ops][index]))
    end
    print("\n")
    ###################
    print("Check observable vals:\n")
    print("Local: ")
    for axis in [0,1,2]
        index = (axis,"Local")
        @printf("%10e ", norm(sp[:means][index][end] - ops_j[index][2](u0))/norm(sp[:means][index][end]))
    end
    print("\n")
    print("Global: ")
    for axis in [0,1,2]
        index = (axis,"Global")
        @printf("%10e ", norm(sp[:means][index][end] - ops_j[index][2](u0))/norm(sp[:means][index][end]))
    end
    print("\n")
    ###################
    IM_j = jlmat2pymat(RHS.IM)
    q2cl_j = jlmat2pymat(RHS.q2cl)
    cl2q_j = jlmat2pymat(RHS.cl2q)
    q2cl_p = sp[:Q2CL]*0.5*sqrt(get_Dh(A))
    cl2q_p = sp[:CL2Q]*-0.5
    @printf("Check IM: %10e\n", slinalg.norm(sp[:FM]-IM_j)/slinalg.norm(sp[:FM]))
    @printf("Check Q2CL: %10e\n", slinalg.norm(q2cl_p-q2cl_j)/slinalg.norm(q2cl_p))
    @printf("Check CL2Q: %10e\n", slinalg.norm(cl2q_p-cl2q_j)/slinalg.norm(cl2q_p))
    ###################
    print("Check rhs:\n")
    @printf("%10e\n", norm(du2[1] - du.x[1].*2.0^-7)/norm(du2[1]))
    for sigma in 1:3
        @printf("%10e ", norm(du2[2][sigma,:]-du.x[2][:,sigma].*2.0^-7)/norm(du2[2][sigma,:]))
    end
    print("\n")
    @printf("Check time: %10e", tj/tp)
    print("\n\n")
end

function test_operators(M::Hybrid, dims::NTuple{3,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, q_spins::Vector{NTuple{3,Int64}}, middle_spins::Vector{NTuple{3,Int64}}) where {F}
    d = Dict([(:all,"Global"), (:central,"Local")])
    A,RHS,OSET, rules = julia_initialize(M, dims, Js, I, hs, q_spins, middle_spins)
    sp = py_initialize(M, dims, Js, I, hs, q_spins, middle_spins)
    H_j = jlmat2pymat(RHS.H)
    ops_j = Dict()
    for Obs in OSET.Observables
        if Obs.axis == :x
            index = d[Obs.name]
            ops_j[index] = (jlmat2pymat(Obs.QO.O),Obs)
        end
    end
    num_of_espins = size(RHS.edgeSpinOperators,1)
    taus_j = Matrix(num_of_espins, 3)
    for espin in 1:num_of_espins for sigma in 1:3
        taus_j[espin, sigma] = jlmat2pymat(RHS.edgeSpinOperators[espin, sigma])
    end end
    sp[:initState]()
    sp[:updateCorrelator]()
    u0 = ArrayPartition(copy(sp[:psi]), VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3]))
    du = ArrayPartition(copy(sp[:psi]), VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3]))
    tj = -time()
    RHS(0.0,u0,du)
    tj += time()
    tp = -time()
    du2 = sp[:rhs](sp[:psi], sp[:clState])
    tp += time()

    print("$((dims, Js, I, hs))\n")
    @printf("Check Hamiltonian:   %e\n", slinalg.norm(sp[:H]-H_j)/slinalg.norm(sp[:H]))
    print("Check spin operators:\n")
    for espin in 1:num_of_espins
        for sigma in 1:3
            @printf("%10e ", slinalg.norm(sp[:taus][espin, sigma]-taus_j[espin, sigma])/slinalg.norm(sp[:taus][espin, sigma]))
        end
        print("\n")
    end
    ###################
    print("Check observable ops:\n")
    index = "Local"
    @printf("Local: %10e\n", slinalg.norm(sp[:ops][index] - ops_j[index][1])/slinalg.norm(sp[:ops][index]))
    index = "Global"
    @printf("Global: %10e\n", slinalg.norm(sp[:ops][index] - ops_j[index][1])/slinalg.norm(sp[:ops][index]))
    ###################
    print("Check observable vals:\n")
    index = "Local"
    @printf("Local: %10e\n", norm(sp[:means][index][end] - ops_j[index][2](u0))/norm(sp[:means][index][end]))
    index = "Global"
    @printf("Global: %10e\n", norm(sp[:means][index][end] - ops_j[index][2](u0))/norm(sp[:means][index][end]))
    ###################
    IM_j = RHS.IM
    q2cl_j = RHS.q2cl
    cl2q_j = RHS.cl2q
    q2cl_p = sp[:Q2CL]
    cl2q_p = -sp[:CL2Q]
    @printf("Check IM: %10e\n", norm(sp[:FM]-IM_j)/norm(sp[:FM]))
    @printf("Check Q2CL: %10e\n", norm(q2cl_p-q2cl_j)/norm(q2cl_p))
    @printf("Check CL2Q: %10e\n", norm(cl2q_p-cl2q_j)/norm(cl2q_p))
    ###################
    print("Check rhs:\n")
    @printf("%10e\n", norm(du2[1] - du.x[1].*2.0^-7)/norm(du2[1]))
    for sigma in 1:3
        @printf("%10e ", norm(du2[2][sigma,:]-du.x[2][:,sigma].*2.0^-7)/norm(du2[2][sigma,:]))
    end
    print("\n")
    @printf("Check time: %10e", tj/tp)
    print("\n\n")
end

function convert_state(A::PureClassicalApprox, sp)
    u0 = VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3])
    return u0
end

function test_operators(M::PureClassical, dims::NTuple{D,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, B::Bool=true) where {D, F}
    A,RHS,OSET, rules = julia_initialize(M, dims, Js, I, hs, B)
    sp = py_initialize(M, dims, Js, I, hs; cb=B)
    ops_j = Dict()
    for Obs in OSET.Observables
        index = axis_dict[Obs.axis]-1
        ops_j[index] = Obs
    end
    sp[:initState]()
    sp[:updateCorrelator]()
    u0 = VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3])
    du = VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3])
    tj = -time()
    RHS(0.0,u0,du)
    tj += time()
    tp = -time()
    du2 = sp[:rhs](sp[:clState])
    tp += time()

    print("$((dims, Js, I, hs))\n")
    ###################
    print("Check observable vals: ")
    for axis in [0,1,2]
        @printf("%10e ", norm(sp[:means][axis][end] - ops_j[axis](u0))/norm(sp[:means][axis][end]))
    end
    print("\n")
    ###################
    IM_j = jlmat2pymat(RHS.IM)
    @printf("Check IM: %10e\n", slinalg.norm(sp[:FM]-IM_j)/slinalg.norm(sp[:FM]))
    ###################
    print("Check rhs: ")
    for sigma in 1:3
        @printf("%10e ", norm(du2[sigma,:]-du[:,sigma].*2.0^-7)/norm(du2[sigma,:]))
    end
    print("\n")
    @printf("Check time: %10e", tj/tp)
    print("\n\n")
end

function test_operators(M::PureClassical, dims::NTuple{3,Int64}, Js::NTuple{3,Float64}, I::Interaction{F}, hs::NTuple{3,Float64}, q_spins, middle_spins) where {F}
    A,RHS,OSET, rules = julia_initialize(M, dims, Js, I, hs)
    sp = py_initialize(M, dims, Js, I, hs)
    ops_j = OSET.Observables[1]
    for Obs in OSET.Observables
        if Obs.axis == :x
            ops_j = Obs
        end
    end
    sp[:initState]()
    sp[:updateCorrelator]()
    u0 = VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3])
    du = VectorOfArray([copy(sp[:clState][i,:]) for i in 1:3])
    tj = -time()
    RHS(0.0,u0,du)
    tj += time()
    tp = -time()
    du2 = sp[:rhs](sp[:clState])
    tp += time()

    print("$((dims, Js, I, hs))\n")
    ###################
    @printf("Check observable val: %10e\n", norm(sp[:means][end] - ops_j(u0))/norm(sp[:means][end]))
    ###################
    IM_j = RHS.IM
    @printf("Check IM: %10e\n", norm(sp[:FM]-IM_j)/norm(sp[:FM]))
    ###################
    print("Check rhs: ")
    for sigma in 1:3
        @printf("%10e ", norm(du2[sigma,:]-du[:,sigma].*2.0^-7)/norm(du2[sigma,:]))
    end
    print("\n")
    @printf("Check time: %10e", tj/tp)
    print("\n\n")
end

@inline store_state(A::ExactApprox, sp) = (deepcopy(sp[:psi]),deepcopy(sp[:psi_aid]))
@inline store_state(A::ClusteredApprox, sp) = (deepcopy(sp[:psi]),)
@inline store_state(A::HybridApprox, sp) = (deepcopy(sp[:psi]), deepcopy(sp[:clState]))
@inline store_state(A::PureClassicalApprox, sp) = (deepcopy(sp[:clState]),)

function test_propagation(M::AbstractModel, args...; Tmax::Float64 = 10.0, tstep::Float64 = 2.0^-7, delimiter = 10)
    A, ARHS, OSET, rules = julia_initialize(M, args...)
    sp = py_initialize(M, args...; :Tmax => Tmax, :tstep => tstep)
    sp[:initState]()
    u0 = convert_state(A, sp)
    state = store_state(A, sp)
    sp[:propagate](state...)
    ####
    l = length(OSET.Observables)
    tspan = (0.0,Tmax*delimiter)
    t_means = collect(0.0:tstep:(Tmax*delimiter))
    t_cors = collect(0.0:tstep:Tmax)
    saved_values = SpinLattice.SavedValues(t_means, [zeros(l) for i in eachindex(t_means)])
    prob = SpinLattice.ODEProblem(ARHS, u0, tspan)
    cb = SpinLattice.SavingCallback(OSET, saved_values, u0; saveat = t_means)
    #cb = SavingCallback(OSET, saved_values; save_everystep=true)
    cf_vals = CFVals(rules, t_cors)
    LP = LProblem(A, cb, prob, rules, cf_vals, Tmax, tstep, delimiter)
    ####
    cf_vals, integrator = simulate(LP, 1, SpinLattice.RK4(); adaptive = false, dt = LP.tstep)

    return LP, cf_vals, integrator, sp
end

function test_clustered()
    for L in L1d[1:4]
        test_operators(Clustered(), L...)
        gc()
    end
    for L in L2d[1:4]
        test_operators(Clustered(), L...)
        gc()
    end
end

function test_all()
    Models = [Exact(), PureClassical(), Clustered(), Hybrid()]
    for model in Models
        for L in L1d[1:4]
            print(model,"\n")
            test_operators(model, L...)
            gc()
        end
    end
    for model in Models
        for L in L2d[1:4]
            print(model,"\n")
            test_operators(model, L...)
            gc()
        end
    end
    for model in Models[2:4]
        for L in L3d[2:4]
            print(model,"\n")
            test_operators(model, L...)
            gc()
        end
    end
end

function test_boundary()
    Models = [Exact(), PureClassical()]
    cbs = [true, false]
    for model in Models
        for L in L1d[1:4]
            for cb in cbs
                print(model, " cb=", cb,"\n")
                test_operators(model, L...,cb)
                gc()
            end
        end
    end
end

function test_3d()
    Models = [Exact(), PureClassical(), Clustered(), Hybrid()]
    for model in (Models[4],)
        for L in L3d[2:4]
            print(model,"\n")
            test_operators(model, L...)
            gc()
        end
    end
end

end


function build_wall!(vis::Visualizer, walls::AbstractVector{<:AbstractWall}; name::Symbol=:Env, α=0.2)
    for i in eachindex(walls)
        build_wall!(vis, walls[i], name=name, id=i, α=α)
    end
    return nothing
end

function build_cylinder!(vis::Visualizer, walls::AbstractVector{<:AbstractWall}; radius=0.0, name::Symbol=:Env, α=0.2)
    for i in eachindex(walls)
        build_cylinder!(vis, walls[i], radius=radius, name=name, id=i, α=α)
    end
    return nothing
end

function build_wall!(vis::Visualizer, wall::Wall3D; name::Symbol=:Env, id::Int, α=0.2)
    p1 = wall.p1
    p2 = wall.p2
    p3 = wall.p3
    p4 = p3 + (p1 - p2)
    l1 = norm(p1 - p2)
    l2 = norm(p3 - p2)
    v1 = (p1 - p2)/l1
    v2 = (p3 - p2)/l2
    orange_mat, blue_mat, black_mat = get_material(;α=α)

    setobject!(vis[name]["wall$id"][:slab],
        MeshCat.HyperRectangle(MeshCat.Vec(0.0, 0.0, 0.0), MeshCat.Vec(l1, l2, 0.0)),
        black_mat)
    for j = 1:3
        setobject!(vis[name]["wall$id"]["p$j"],
            GeometryBasics.Sphere(MeshCat.Point3f0(0.0), 0.05),
            black_mat)
    end

    rt = MeshCat.LinearMap([v1 v2 cross(v1, v2)])
    settransform!(vis[name]["wall$id"][:slab], MeshCat.compose(MeshCat.Translation(p2...), rt))
    settransform!(vis[name]["wall$id"]["p1"], MeshCat.Translation(p1...))
    settransform!(vis[name]["wall$id"]["p2"], MeshCat.Translation(p2...))
    settransform!(vis[name]["wall$id"]["p3"], MeshCat.Translation(p3...))
    return nothing
end

function build_cylinder!(vis::Visualizer, wall::CylinderWall; radius=0.0, name::Symbol=:Env, id::Int, α=0.2)
    p = wall.p
    v = wall.v
    v_ = Int.([v == :x, v == :y, v == :z])
    l = convert(Float32, wall.l)
    r = convert(Float32, wall.r-radius)
    orange_mat, blue_mat, black_mat = get_material(;α=α)

    setobject!(vis[name]["cylinder$id"][:body],
        MeshCat.Cylinder(MeshCat.Point3f0(0.0), MeshCat.Point3f0((l*v_)...), r),
        black_mat)

    # MeshCat.HyperRectangle(MeshCat.Vec(0.0, 0.0, 0.0),
    #     MeshCat.Vec(l1, l2, 0.0)), MeshPhongMaterial(color = MeshCat.RGBA(0.5, 0.5, 0.5, α)))
    setobject!(vis[name]["cylinder$id"]["p"], GeometryBasics.Sphere(MeshCat.Point3f0(0.0), 0.05), black_mat)

    # rt = MeshCat.LinearMap([v1 v2 cross(v1, v2)])
    # settransform!(vis[name]["cylinder$id"][:body], MeshCat.compose(MeshCat.Translation(p...), rt))
    settransform!(vis[name]["cylinder$id"][:body], MeshCat.Translation(p...))
    settransform!(vis[name]["cylinder$id"]["p"], MeshCat.Translation(p...))
    return nothing
end


export SumProductNetwork
export SPNNode, Node, Leaf, SumNode, ProductNode
export FiniteSumNode, FiniteProductNode
export IndicatorNode, UnivariateNode
export MultivariateNode

export root, leaves

# Abstract definition of an SumProductNetwork node.
abstract type SPNNode end
abstract type Node <: SPNNode end
abstract type SumNode <: Node end
abstract type ProductNode <: Node end
abstract type Leaf <: SPNNode end

header(node::SPNNode) = "$(summary(node))($(node.id))"

function Base.show(io::IO, node::SPNNode)
    println(io, header(node))
    if hasweights(node)
        println(io, "\tweights = $(weights(node))")
        println(io, "\tnormalized = $(isnormalized(node))")
    end
    if hasscope(node)
        println(io, "\tscope = $(scope(node))")
    else
        println(io, "\tNo scope set!")
    end
    if hasobs(node)
        println(io, "\tassigns = $(obs(node))")
    end
end

function Base.show(io::IO, node::Leaf)
    println(io, header(node))
    println(io, "\tscope = $(scope(node))")
    println(io, "\tparameters = $(params(node))")
end

struct SumProductNetwork
    root::Node
    nodes::Vector{<:SPNNode}
    leaves::Vector{<:SPNNode}
    idx::Dict{Symbol,Int}
    topological_order::Vector{Int}
    layers::Vector{AbstractVector{SPNNode}}
    info::Dict{Symbol,<:Real}
end

function SumProductNetwork(root::Node)
    nodes = getOrderedNodes(root)
    leaves = filter(n -> isa(n, Leaf), nodes)
    idx = Dict(n.id => indx for (indx, n) in enumerate(nodes))
    toporder = collect(1:length(nodes))

    maxdepth = depth(root)
    nodedepth = map(n -> depth(n), nodes)
    layers = Vector{Vector{SPNNode}}(undef, maxdepth+1)
    for d in 0:maxdepth
        layers[d+1] = nodes[findall(nodedepth .== d)]
    end

    return SumProductNetwork(root, nodes, leaves, idx, toporder, layers, Dict{Symbol, Real}())
end

Base.keys(spn::SumProductNetwork) = keys(spn.idx)
Base.values(spn::SumProductNetwork) = spn.nodes
Base.getindex(spn::SumProductNetwork, i...) = getindex(spn.nodes, spn.idx[i...])
Base.setindex!(spn::SumProductNetwork, v, i...) = setindex!(spn.nodes, v, spn.idx[i...])
Base.length(spn::SumProductNetwork) = length(spn.nodes)
function Base.show(io::IO, spn::SumProductNetwork)
    println(io, summary(spn))
    println(io, "\t#nodes = $(length(spn))")
    println(io, "\t#leaves = $(length(spn.leaves))")
    println(io, "\tdepth = $(length(spn.layers))")
end
leaves(spn::SumProductNetwork) = spn.leaves
root(spn::SumProductNetwork) = spn.root

"""
   FiniteSumNode <: SumNode

A sum node computes a convex combination of its weight and the pdf's of its children.

## Usage:

```julia
node = FiniteSumNode{Float64}(;D=4, N=1)
add!(node, ..., log(0.5))
add!(node, ..., log(0.5))
logpdf(node, rand(4))
```

"""
mutable struct FiniteSumNode{T<:Real} <: SumNode
    id::Symbol
    parents::Vector{<:Node}
    children::Vector{<:SPNNode}
    logweights::Vector{T}
    scopeVec::BitArray{1}
    obsVec::BitArray{1}
end

function FiniteSumNode{T}(;D=1, N=1, parents::Vector{<:Node}=Node[]) where {T<:Real}
    return FiniteSumNode{T}(gensym(:sum), parents, SPNNode[], T[], falses(D), falses(N))
end
function FiniteSumNode(;D=1, N=1, parents::Vector{<:Node}=Node[])
    return FiniteSumNode{Float64}(;D=D, N=N, parents=parents)
end

eltype(::Type{FiniteSumNode{T}}) where {T<:Real} = T
eltype(n::SPNNode) = eltype(typeof(n))

"""
   FiniteProductNode <: ProductNode

A product node computes a product of the pdf's of its children.

## Usage:

```julia
node = FiniteProductNode(;D=4, N=1)
add!(node, ...)
add!(node, ...)
logpdf(node, rand(4))
```

"""
mutable struct FiniteProductNode <: ProductNode
    id::Symbol
    parents::Vector{<:Node}
    children::Vector{<:SPNNode}
    scopeVec::BitArray{1}
    obsVec::BitArray{1}
end

function FiniteProductNode(;D=1, N=1, parents::Vector{<:Node} = Node[])
    return FiniteProductNode(gensym(:prod), parents, SPNNode[], falses(D), falses(N))
end

"""
   IndicatorNode <: Leaf

An indicator node evaluates an indicator function.

## Usage:

```julia
value = 1
dimension = 1
node = IndicatorNode(value, dimension)
logpdf(node, [1, 2, 1]) # == 0.0
logpdf(node, [2, 2, 1]) # == -Inf
```

"""
mutable struct IndicatorNode <: Leaf
    id::Symbol
    value::Int
    scope::Int
    parents::Vector{<:Node}
end

function IndicatorNode(value::Int, dim::Int; parents::Vector{<:Node} = Node[])
    return IndicatorNode(gensym(:indicator), value, dim, parents)
end

function Base.isequal(n1::IndicatorNode, n2::IndicatorNode)
    return (n1.value == n2.value) && (n1.scope == n2.scope)
end
params(n::IndicatorNode) = (n.value)

"""
   UnivariateNode <: Leaf

A univariate node evaluates the logpdf of its univariate distribution function.

## Usage:

```julia
distribution = Normal()
dimension = 1
node = UnivariateNode(distribution, dimension)
logpdf(node, rand(2)) # == logpdf(Normal(), rand(2))
```

"""
mutable struct UnivariateNode <: Leaf
    id::Symbol
    parents::Vector{<:Node}
    dist::UnivariateDistribution
    scope::Int
end

function UnivariateNode(distribution::T, dim::Int; parents::Vector{<:Node} = Node[]) where {T<:UnivariateDistribution}
    return UnivariateNode(gensym(:univ), parents, distribution, dim)
end
params(n::UnivariateNode) = Distributions.params(n.dist)

"""
   MultivariateNode <: Leaf

A multivariate node evaluates the logpdf of its multivariate distribution function.

## Usage:

```julia
distribution = MvNormal()
dimensions = [1, 2]
node = UnivariateNode(distribution, dimensions)
logpdf(node, rand(2)) # == logpdf(MvNormal(), rand(2))
```

"""
mutable struct MultivariateNode <: Leaf
    id::Symbol
    parents::Vector{<:Node}
    dist::MultivariateDistribution
    scope::Vector{Int}
end

function MultivariateNode(distribution::T, dims::Vector{Int}; parents::Vector{<:Node} = Node[]) where {T<:MultivariateDistribution}
    return MultivariateNode(gensym(:multiv), parents, distribution, dims)
end
params(n::MultivariateNode) = Distributions.params(n.dist)


# SPDX-License-Identifier: MIT

"""
    abstract type AbstractModel

Abstract type for all (Bayesian) models.
"""
abstract type AbstractModel <: BMObject end

"""
    abstract type AbstractLatentVariableModel

Abstract type for model with latent variable (mixture, pca, ...).
"""
abstract type AbstractLatentVariableModel <: AbstractModel end

"""
    loglikelihood(model::AbstractModel, X)

Return the log-likelihood of `model` of each data sample of `X`.
"""
loglikelihood

"""
    posterior(model::AbstractLatentVariableModel, X)

Return the posterior over the latent variables of `model`.
"""
posterior

"""
    predict(model::AbstractLatentVariableModel, X)

Return the most likely latent variables of `model` given `X`.
"""
predict

"""
    todict(model)

Return a dictionary of the state of the model.
"""
function todict(m::AbstractModel)
    d = Dict()
    names = fieldnames(typeof(m))
    d[:length] = length(names)
    for (i,name) in enumerate(names)
        prop = getproperty(m, name)
        d[Symbol("prop_$(i)_type")] = typeof(prop)
        d[Symbol("prop_$i")] = todict(prop)
    end
    d
end

"""
    fromdict(T, dict)

Return a model of type `T` from the state stored in `dict`.
"""
function fromdict(T::Type{<:AbstractModel}, d::AbstractDict)
    params = []
    for i in 1:d[:length]
        ptype = d[Symbol("prop_$(i)_type")]
        pstate = d[Symbol("prop_$(i)")]
        push!(params, fromdict(ptype, pstate))
    end
    T(params...)
end

"""
    ModelList{N,T}(m1, m2, m3, ...)

Store a list of models for an object's attribute. `N` is the number of
models in the list and `T` is the type of the models. The list is
immutable.
"""
const ModelList{N,T<:AbstractModel} = NTuple{N,T}


# RELOG: Reverse Logistics Optimization
# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
# Released under the modified BSD license. See COPYING.md for more details.

using DataStructures
using JSON
using JSONSchema
using Printf
using Statistics

function parsefile(path::String)::Instance
    return RELOG.parse(JSON.parsefile(path))
end

function parse(json)::Instance
    basedir = dirname(@__FILE__)
    json_schema = JSON.parsefile("$basedir/../schemas/input.json")
    validate(json, Schema(json_schema))

    T = json["parameters"]["time horizon (years)"]
    json_schema["definitions"]["TimeSeries"]["minItems"] = T
    json_schema["definitions"]["TimeSeries"]["maxItems"] = T
    validate(json, Schema(json_schema))

    building_period = [1]
    if "building period (years)" in keys(json)
        building_period = json["building period (years)"]
    end

    plants = Plant[]
    products = Product[]
    collection_centers = CollectionCenter[]
    prod_name_to_product = Dict{String,Product}()

    # Create products
    for (product_name, product_dict) in json["products"]
        cost = product_dict["transportation cost (\$/km/tonne)"]
        energy = zeros(T)
        emissions = Dict()

        if "transportation energy (J/km/tonne)" in keys(product_dict)
            energy = product_dict["transportation energy (J/km/tonne)"]
        end

        if "transportation emissions (tonne/km/tonne)" in keys(product_dict)
            emissions = product_dict["transportation emissions (tonne/km/tonne)"]
        end

        product = Product(product_name, cost, energy, emissions)
        push!(products, product)
        prod_name_to_product[product_name] = product

        # Create collection centers
        if "initial amounts" in keys(product_dict)
            for (center_name, center_dict) in product_dict["initial amounts"]
                if "location" in keys(center_dict)
                    region = geodb_query(center_dict["location"])
                    center_dict["latitude (deg)"] = region.centroid.lat
                    center_dict["longitude (deg)"] = region.centroid.lon
                end
                center = CollectionCenter(
                    length(collection_centers) + 1,
                    center_name,
                    center_dict["latitude (deg)"],
                    center_dict["longitude (deg)"],
                    product,
                    center_dict["amount (tonne)"],
                )
                push!(collection_centers, center)
            end
        end
    end

    # Create plants
    for (plant_name, plant_dict) in json["plants"]
        input = prod_name_to_product[plant_dict["input"]]
        output = Dict()

        # Plant outputs
        if "outputs (tonne/tonne)" in keys(plant_dict)
            output = Dict(
                prod_name_to_product[key] => value for
                (key, value) in plant_dict["outputs (tonne/tonne)"] if value > 0
            )
        end

        energy = zeros(T)
        emissions = Dict()

        if "energy (GJ/tonne)" in keys(plant_dict)
            energy = plant_dict["energy (GJ/tonne)"]
        end

        if "emissions (tonne/tonne)" in keys(plant_dict)
            emissions = plant_dict["emissions (tonne/tonne)"]
        end

        for (location_name, location_dict) in plant_dict["locations"]
            sizes = PlantSize[]
            disposal_limit = Dict(p => [0.0 for t = 1:T] for p in keys(output))
            disposal_cost = Dict(p => [0.0 for t = 1:T] for p in keys(output))

            # GeoDB
            if "location" in keys(location_dict)
                region = geodb_query(location_dict["location"])
                location_dict["latitude (deg)"] = region.centroid.lat
                location_dict["longitude (deg)"] = region.centroid.lon
            end

            # Disposal
            if "disposal" in keys(location_dict)
                for (product_name, disposal_dict) in location_dict["disposal"]
                    limit = [1e8 for t = 1:T]
                    if "limit (tonne)" in keys(disposal_dict)
                        limit = disposal_dict["limit (tonne)"]
                    end
                    disposal_limit[prod_name_to_product[product_name]] = limit
                    disposal_cost[prod_name_to_product[product_name]] =
                        disposal_dict["cost (\$/tonne)"]
                end
            end

            # Capacities
            for (capacity_name, capacity_dict) in location_dict["capacities (tonne)"]
                push!(
                    sizes,
                    PlantSize(
                        Base.parse(Float64, capacity_name),
                        capacity_dict["variable operating cost (\$/tonne)"],
                        capacity_dict["fixed operating cost (\$)"],
                        capacity_dict["opening cost (\$)"],
                    ),
                )
            end
            length(sizes) > 1 || push!(sizes, sizes[1])
            sort!(sizes, by = x -> x.capacity)

            # Storage
            storage_limit = 0
            storage_cost = zeros(T)
            if "storage" in keys(location_dict)
                storage_dict = location_dict["storage"]
                storage_limit = storage_dict["limit (tonne)"]
                storage_cost = storage_dict["cost (\$/tonne)"]
            end

            # Validation: Capacities
            if length(sizes) != 2
                throw("At most two capacities are supported")
            end
            if sizes[1].variable_operating_cost != sizes[2].variable_operating_cost
                throw("Variable operating costs must be the same for all capacities")
            end

            plant = Plant(
                length(plants) + 1,
                plant_name,
                location_name,
                input,
                output,
                location_dict["latitude (deg)"],
                location_dict["longitude (deg)"],
                disposal_limit,
                disposal_cost,
                sizes,
                energy,
                emissions,
                storage_limit,
                storage_cost,
            )

            push!(plants, plant)
        end
    end

    @info @sprintf("%12d collection centers", length(collection_centers))
    @info @sprintf("%12d candidate plant locations", length(plants))

    return Instance(T, products, collection_centers, plants, building_period)
end


export buildnet

function buildnet(;name="temp",model=NetworkModel(),method="Euler1",duration=0.1)
  iteration=round(Int64,duration/model.dt)
  syncIter=model.synctime>0?round(Int64,model.synctime÷model.dt):1
  iteration=iteration÷syncIter
  model.nCores=length(model.groups[:index])+1
  ##############################################
  # I.parse

  # equations
  eqs=map(strip,split(replace(model.equations,";","\n"),"\n"));
  vars=[]
  diffEqs=[]

  for k=1:length(eqs)
    if eqs[k][1]=='d'
      vars=[vars;match(r"(?<=d).*(?=/dt)",eqs[k]).match]
      diffEq=vectorize(match(r"(?<==).*$",eqs[k]).match)
      diffEqs=[diffEqs;diffEq]
    else
      funcName=match(r"^.*(?==)",eqs[k]).match
      funcExpr=vectorize(match(r"(?<==).*$",eqs[k]).match)
      diffEqs=map(diffEq->replace(diffEq,funcName,funcExpr),diffEqs)
    end
  end

  # report
  report="""
  k%(1000÷syncIter)==0?@printf("%s%4.2f finished\\n","%",100*k/iteration):nothing
  """

  # integrate
  argument=join([vars;"t"],",")
  splitArg=split(argument,",")
  argumentTemp=join([splitArg[1]*"+Δ$(vars[1])";splitArg[2:end]],",")

  func=map(var->"F$var($argument)",vars)
  tempFunc="F$(vars[1])($argumentTemp)"

  if model.solver=="Euler1"
    integrate="Δ$(vars[1])=dt*$(func[1])\n"
  elseif model.solver=="Euler2"
    integrate="Δ$(vars[1])=dt/2.0*$(func[1])\n"
    integrate*="Δ$(vars[1])+=dt/2.0*$tempFunc\n"
  elseif model.solver=="Euler4"
    integrate="Δ$(vars[1])=dt/4.0*$(func[1])\n"
    integrate*="Δ$(vars[1])+=dt/4.0*$tempFunc\n"
    integrate*="Δ$(vars[1])+=dt/4.0*$tempFunc\n"
    integrate*="Δ$(vars[1])+=dt/4.0*$tempFunc\n"
  end

  for k=2:length(vars)
    integrate*="Δ$(vars[k])=dt*$(func[k])\n"
  end
  for k=1:length(vars)
    integrate*="$(vars[k])+=Δ$(vars[k])\n"
  end

  # spike
  if !isempty(model.spike)
    spike=split(replace(model.spike,";","\n"),"\n")
    threshold=spike[1]
    length(spike)==2? refractory=spike[2]:refractory=""

    spike=""
    spike*="isExceedThreshold=$threshold\n"
    if !isempty(refractory)
      spike*="isNotRefractory=$refractory\n"
      spike*="isSpike=isExceedThreshold & isNotRefractory\n"
    else
      spike*="isSpike=isExceedThreshold\n"
    end
    spike*="spikeIndex=find(isSpike)\n"
    spike*="nSpike=length(spikeIndex)\n"
    spike*="tSpike[state:state+nSpike-1]=t\n"
    spike*="ts[spikeIndex]=t\n"
    spike*="spikeNeuron[state:state+nSpike-1]=spikeIndex\n"
    spike*="state+=nSpike\n"
  else spike=""
  end

  # reset
  if !isempty(model.reset)
    reset=""
    for line in split(replace(model.reset,";","\n"),"\n")
      var=match(r"^.*(?==)",line).match
      expr=match(r"(?<==).*",line).match
      reset*="$(var)[spikeIndex]=$(expr)\n"
    end
  else reset=""
  end

  # synapse
  if !isempty(model.synapse)
    synapse=""
    for line in split(replace(model.synapse,";","\n"),"\n")
      var1=match(r"^.*(?=\+)",line).match
      var2=match(r"(?<==).*",line).match
      synapse*="$(var1)+=sum($(var2)[:,totalSpikeIndex],2)\n"
   end
  else synapse=""
  end

  # plasticity
  if !isempty(model.plasticity)
    str=split(replace(model.plasticity,";","\n"),"\n");
    plasticity1=join(str[1:end÷2],"\n");plasticity2=join(str[end÷2+1:end],"\n")
    plasticity1=replace(replace(plasticity1,"j",":"),"i","totalSpikeIndex")*"\n"
    plasticity2=replace(replace(plasticity2,"j",":"),"i","spikeIndex")*"\n"
    plasticity=plasticity1*plasticity2
  else plasticity=""
  end

  # record
  if !isempty(model.record)
    record=""
    for i=1:length(vars)
      record*="$(vars[i])Record[:,\$k]=$(vars[i])[localRecord]\n"
    end
  else
    record=""
  end
  # end of parse
  ####################################################

  ####################################################
  # II.print

  # open a file
  fid=open(name*".jl","w")

  # set up parameters
  paramsStr="""
  println("\\nExecuting $(name).jl")

  function sendto(p::Int; args...)
    for (nm, val) in args
        @spawnat(p, eval(Main, Expr(:(=), nm, val)))
    end
  end
  function sendto(ps::Vector{Int}; args...)
      for p in ps
          sendto(p; args...)
      end
  end

  using JLD

  if $(model.nCores)>0 && (addCores=$(model.nCores)-nworkers())>0
    nworkers()==1?addprocs(addCores):addprocs(addCores-1)
  end
  println("Using \$(nworkers()) workers")

  @everywhere begin
  index=$(model.groups[:index])
  # units
  metre=1e2;meter=1e2;cm=metre/1e2;mm=metre/1e3;um=metre/1e6;nm=metre/1e9;
  second=1e3;ms=second/1e3;
  Hz=1/second;
  voltage=1e3;mV=voltage/1e3;
  ampere=1e6;mA=ampere/1e3;uA=ampere/1e6;nA=ampere/1e9;
  farad=1e6;uF=ufarad=farad/1e6;
  siemens=1e3;mS=msiemens=siemens/1e3;nS=nsiemens=siemens/1e9
  # parameters
  $(model.parameters)"""
  println(fid,paramsStr)

  varsType=map(var->var*"::Array{Float64}",vars)
  argumentType=join([varsType;"t::Float64"],",");
  println(fid,"# function")
  for k=1:length(vars)
    println(fid,"F",vars[k],"(",argumentType,")=",diffEqs[k])
  end
  println(fid,"end # end of @everywhere\n")

  # function
  initVars=""
  for var in vars
    initVars*="@parallel for p in workers()  @eval $var=zeros(N) end\n"
  end

  initRecord=""
  if !isempty(model.record)
    for k=1:length(vars)
      initRecord*="$(vars[k])Record=zeros(length(localRecord),iteration)\n"
    end
  end

  collectVars="""
  spikeNeuron=@parallel (vcat) for p in workers() @eval spikeNeuron[1:state-1]+localIndex[1]-1  end
  tSpike=@parallel (vcat) for p in workers() @eval tSpike[1:state-1]  end
  """
  if !isempty(model.record)
    collectVars*="$(vars[1])Record=@parallel (vcat) for p in workers() @eval $(vars[1])Record end\n"
  end
  for var in model.saveVars
    collectVars*="$var=@parallel (vcat) for p in workers() @eval $var end\n"
  end

  saveVars="""
  try jldopen("$name.jld","w") do fid
    try  write(fid,"spikeNeuron",spikeNeuron)  end
    try  write(fid,"tSpike",tSpike)            end
    try  write(fid,"$(vars[1])Record",$(vars[1])Record)   end
    """
  for var in model.saveVars
  saveVars*="try write(fid,\"$(var)\",$var) end\n"
  end
  saveVars*="end #end of do\n"
  saveVars*="end #end of try jldopen\n"

  funcStr="""
  function $name(iteration::Int64=10)
  syncIter=$syncIter
  @sync @parallel for n=1:nworkers()
  @eval begin
  t=0.0;dt=$(model.dt); N_tot=$(model.groups[:N]); syncIter=$syncIter;
  localIndex=index[\$n];N=length(localIndex); iteration=\$iteration
  localRecord=intersect($(model.record),localIndex)-localIndex[1]+1
  # initialize record
  $initRecord
  # initialize tSpike
  max_tSpike=round(Int64,N*iteration*dt)
  tSpike=zeros(max_tSpike);spikeNeuron=zeros(Int64,max_tSpike)
  ts=zeros(N);state=1;spikeIndex=[]
  end # end of @eval begin
  end # end of @parallel for

  @sync begin
  # initialize variables
  $initVars
  $(model.initialization)
  # initialize connections
  $(model.connection)
  end # end of @sync begin

  # print status
  iteration>=50?println("simulation starts..."):println("compiling...")

  # for-loop
  tic()
  @fastmath @inbounds for k=1:iteration
  # report
  $report
  # run until sync
  @sync @parallel for p in workers()
  @eval begin
  for idx=1:syncIter
  # integrate
  $integrate
  t+=dt
  end # end of for idx=1:syncIter
  # spike
  $spike
  # reset
  $reset
  # record
  $record
  end # end of @eval begin
  end # @sync for p in workers()

  totalSpikeIndex=@parallel (vcat) for p in workers()
  @eval spikeIndex+localIndex[1]-1
  end
  sendto(workers(),totalSpikeIndex=totalSpikeIndex)
  @sync @parallel for p in workers()
  @eval begin
  # synapse
  $synapse
  # plasticity
  $plasticity
  end
  end

  end # end of @fastmath @inbounds for k=1:iteration
  toc()

  # collect variables
  $collectVars
  # save variables
  $saveVars

  end # end of main function

  # precompile and run
  precompile($name,($iteration,))
  $name(20)
  $name(20)
  $name($iteration)

  # post-processing
  $(model.postProcs)
  """
  println(fid,funcStr)

  # close file id, try to open file in atom
  close(fid)
  try run(`atom $(abspath("$name.jl"))`) end

end # end of function buildNetwork


################################################################################
#                             Lines, LineSegments                              #
################################################################################

function draw_atomic(scene::Scene, screen::CairoScreen, primitive::Union{Lines, LineSegments})
    fields = @get_attribute(primitive, (color, linewidth, linestyle))
    linestyle = Makie.convert_attribute(linestyle, Makie.key"linestyle"())
    ctx = screen.context
    model = primitive[:model][]
    positions = primitive[1][]

    isempty(positions) && return

    # workaround for a LineSegments object created from a GLNormalMesh
    # the input argument is a view of points using faces, which results in
    # a vector of tuples of two points. we convert those to a list of points
    # so they don't trip up the rest of the pipeline
    # TODO this shouldn't be necessary anymore!
    if positions isa SubArray{<:Point3, 1, P, <:Tuple{Array{<:AbstractFace}}} where P
        positions = let
            pos = Point3f0[]
            for tup in positions
                push!(pos, tup[1])
                push!(pos, tup[2])
            end
            pos
        end
    end

    projected_positions = project_position.(Ref(scene), positions, Ref(model))

    if color isa AbstractArray{<: Number}
        color = numbers_to_colors(color, primitive)
    end

    # color is now a color or an array of colors
    # if it's an array of colors, each segment must be stroked separately

    # The linestyle can be set globally, as we do here.
    # However, there is a discrepancy between Makie
    # and Cairo when it comes to linestyles.
    # For Makie, the linestyle array is cumulative,
    # and defines the "absolute" endpoints of segments.
    # However, for Cairo, each value provides the length of
    # alternate "on" and "off" portions of the stroke.
    # Therefore, we take the diff of the given linestyle,
    # to convert the "absolute" coordinates into "relative" ones.
    if !isnothing(linestyle) && !(linewidth isa AbstractArray)
        Cairo.set_dash(ctx, diff(Float64.(linestyle)) .* linewidth)
    end
    if color isa AbstractArray || linewidth isa AbstractArray
        # stroke each segment separately, this means disjointed segments with probably
        # wonky dash patterns if segments are short

        # we can hide the gaps by setting the line cap to round
        Cairo.set_line_cap(ctx, Cairo.CAIRO_LINE_CAP_ROUND)
        draw_multi(
            primitive, ctx,
            projected_positions,
            color, linewidth,
            isnothing(linestyle) ? nothing : diff(Float64.(linestyle))
        )
    else
        # stroke the whole line at once if it has only one color
        # this allows correct linestyles and line joins as well and will be the
        # most common case
        Cairo.set_line_width(ctx, linewidth)
        Cairo.set_source_rgba(ctx, red(color), green(color), blue(color), alpha(color))
        draw_single(primitive, ctx, projected_positions)
    end
    nothing
end

function draw_single(primitive::Lines, ctx, positions)
    n = length(positions)
    @inbounds for i in 1:n
        p = positions[i]
        # only take action for non-NaNs
        if !isnan(p)
            # new line segment at beginning or if previously NaN
            if i == 1 || isnan(positions[i-1])
                Cairo.move_to(ctx, p...)
            else
                Cairo.line_to(ctx, p...)
                # complete line segment at end or if next point is NaN
                if i == n || isnan(positions[i+1])
                    Cairo.stroke(ctx)
                end
            end
        end
    end
end

function draw_single(primitive::LineSegments, ctx, positions)

    @assert iseven(length(positions))

    @inbounds for i in 1:2:length(positions)-1
        p1 = positions[i]
        p2 = positions[i+1]

        if isnan(p1) || isnan(p2)
            continue
        else
            Cairo.move_to(ctx, p1...)
            Cairo.line_to(ctx, p2...)
            Cairo.stroke(ctx)
        end
    end
end

# if linewidth is not an array
function draw_multi(primitive, ctx, positions, colors::AbstractArray, linewidth, dash)
    draw_multi(primitive, ctx, positions, colors, [linewidth for c in colors], dash)
end

# if color is not an array
function draw_multi(primitive, ctx, positions, color, linewidths::AbstractArray, dash)
    draw_multi(primitive, ctx, positions, [color for l in linewidths], linewidths, dash)
end

function draw_multi(primitive::Union{Lines, LineSegments}, ctx, positions, colors::AbstractArray, linewidths::AbstractArray, dash)
    if primitive isa LineSegments
        @assert iseven(length(positions))
    end
    @assert length(positions) == length(colors)
    @assert length(linewidths) == length(colors)

    iterator = if primitive isa Lines
        1:length(positions)-1
    elseif primitive isa LineSegments
        1:2:length(positions)
    end

    for i in iterator
        if isnan(positions[i+1]) || isnan(positions[i])
            continue
        end
        Cairo.move_to(ctx, positions[i]...)

        Cairo.line_to(ctx, positions[i+1]...)
        if linewidths[i] != linewidths[i+1]
            error("Cairo doesn't support two different line widths ($(linewidths[i]) and $(linewidths[i+1])) at the endpoints of a line.")
        end
        Cairo.set_line_width(ctx, linewidths[i])
        !isnothing(dash) && Cairo.set_dash(ctx, dash .* linewidths[i])
        c1 = colors[i]
        c2 = colors[i+1]
        # we can avoid the more expensive gradient if the colors are the same
        # this happens if one color was given for each segment
        if c1 == c2
            Cairo.set_source_rgba(ctx, red(c1), green(c1), blue(c1), alpha(c1))
            Cairo.stroke(ctx)
        else
            pat = Cairo.pattern_create_linear(positions[i]..., positions[i+1]...)
            Cairo.pattern_add_color_stop_rgba(pat, 0, red(c1), green(c1), blue(c1), alpha(c1))
            Cairo.pattern_add_color_stop_rgba(pat, 1, red(c2), green(c2), blue(c2), alpha(c2))
            Cairo.set_source(ctx, pat)
            Cairo.stroke(ctx)
            Cairo.destroy(pat)
        end
    end
end

################################################################################
#                                   Scatter                                    #
################################################################################

function draw_atomic(scene::Scene, screen::CairoScreen, primitive::Scatter)
    fields = @get_attribute(primitive, (color, markersize, strokecolor, strokewidth, marker, marker_offset, rotations))
    @get_attribute(primitive, (transform_marker,))

    ctx = screen.context
    model = primitive[:model][]
    positions = primitive[1][]
    isempty(positions) && return
    size_model = transform_marker ? model : Mat4f0(I)

    font = to_font(to_value(get(primitive, :font, Makie.defaultfont())))

    colors = if color isa AbstractArray{<: Number}
        numbers_to_colors(color, primitive)
    else
        color
    end

    broadcast_foreach(primitive[1][], colors, markersize, strokecolor,
                      strokewidth, marker, marker_offset, remove_billboard(rotations)) do point, col,
                          markersize, strokecolor, strokewidth, marker, mo, rotation

        # if we give size in pixels, the size is always equal to that value
        is_pixelspace = haskey(primitive, :markerspace) && primitive.markerspace[] == Makie.Pixel
        scale = if is_pixelspace
            Makie.to_2d_scale(markersize)
        else
            # otherwise calculate a scaled size
            project_scale(scene, markersize, size_model)
        end
        offset = if is_pixelspace
            Makie.to_2d_scale(mo)
        else
            project_scale(scene, mo, size_model)
        end

        pos = project_position(scene, point, model)

        isnan(pos) && return

        Cairo.set_source_rgba(ctx, rgbatuple(col)...)
        m = convert_attribute(marker, key"marker"(), key"scatter"())
        Cairo.save(ctx)
        if m isa Char
            draw_marker(ctx, m, best_font(m, font), pos, scale, strokecolor, strokewidth, offset, rotation)
        else
            draw_marker(ctx, m, pos, scale, strokecolor, strokewidth, offset, rotation)
        end
        Cairo.restore(ctx)
    end
    nothing
end

function draw_marker(ctx, marker::Char, font, pos, scale, strokecolor, strokewidth, marker_offset, rotation)
    # Marker offset is meant to be relative to the
    # bottom left corner of the box centered at
    # `pos` with sides defined by `scale`, but
    # this does not take the character's dimensions
    # into account.
    # Here, we reposition the marker offset to be
    # relative to the center of the char.
    marker_offset = marker_offset .+ scale ./ 2

    cairoface = set_ft_font(ctx, font)

    charextent = Makie.FreeTypeAbstraction.internal_get_extent(font, marker)
    inkbb = Makie.FreeTypeAbstraction.inkboundingbox(charextent)

    # scale normalized bbox by font size
    inkbb_scaled = FRect2D(origin(inkbb) .* scale, widths(inkbb) .* scale)

    # flip y for the centering shift of the character because in Cairo y goes down
    centering_offset = [1, -1] .* (-origin(inkbb_scaled) .- 0.5 .* widths(inkbb_scaled))
    # this is the origin where we actually have to place the glyph so it can be centered
    charorigin = pos .+ Vec2f0(marker_offset[1], -marker_offset[2])
    old_matrix = get_font_matrix(ctx)
    set_font_matrix(ctx, scale_matrix(scale...))

    # First, we translate to the point where the
    # marker is supposed to go.
    Cairo.translate(ctx, charorigin...)
    # Then, we rotate the context by the
    # appropriate amount,
    Cairo.rotate(ctx, to_2d_rotation(rotation))
    # and apply a centering offset to account for
    # the fact that text is shown from the (relative)
    # bottom left corner.
    Cairo.translate(ctx, centering_offset...)

    Cairo.move_to(ctx, 0, 0)
    Cairo.text_path(ctx, string(marker))
    Cairo.fill_preserve(ctx)
    # stroke
    Cairo.set_line_width(ctx, strokewidth)
    Cairo.set_source_rgba(ctx, rgbatuple(strokecolor)...)
    Cairo.stroke(ctx)

    # if we use set_ft_font we should destroy the pointer it returns
    cairo_font_face_destroy(cairoface)

    set_font_matrix(ctx, old_matrix)
end

function draw_marker(ctx, marker::Circle, pos, scale, strokecolor, strokewidth, marker_offset, rotation)

    marker_offset = marker_offset + scale ./ 2
    pos += Point2f0(marker_offset[1], -marker_offset[2])
    Cairo.arc(ctx, pos[1], pos[2], scale[1]/2, 0, 2*pi)
    Cairo.fill_preserve(ctx)

    Cairo.set_line_width(ctx, Float64(strokewidth))

    sc = to_color(strokecolor)
    Cairo.set_source_rgba(ctx, rgbatuple(sc)...)
    Cairo.stroke(ctx)
end

function draw_marker(ctx, marker::Rect, pos, scale, strokecolor, strokewidth, marker_offset, rotation)
    s2 = Point2((widths(marker) .* scale .* (1, -1))...)
    pos = pos .+ Point2f0(marker_offset[1], -marker_offset[2])
    Cairo.rotate(ctx, to_2d_rotation(rotation))
    Cairo.rectangle(ctx, pos[1], pos[2], s2...)
    Cairo.fill_preserve(ctx)
    Cairo.set_line_width(ctx, Float64(strokewidth))
    sc = to_color(strokecolor)
    Cairo.set_source_rgba(ctx, rgbatuple(sc)...)
    Cairo.stroke(ctx)
end

################################################################################
#                                     Text                                     #
################################################################################

function p3_to_p2(p::Point3{T}) where T
    if p[3] == 0 || isnan(p[3])
        Point2{T}(p[1:2]...)
    else
        error("Can't reduce Point3 to Point2 with nonzero third component $(p[3]).")
    end
end

function draw_atomic(scene::Scene, screen::CairoScreen, primitive::Text{<:Tuple{<:G}}) where G <: Union{AbstractArray{<:Makie.GlyphCollection}, Makie.GlyphCollection}
    ctx = screen.context
    @get_attribute(primitive, (rotation, model, space, offset))
    position = primitive.attributes[:position][]
    # use cached glyph info
    glyph_collection = to_value(primitive[1])

    draw_glyph_collection(scene, ctx, position, glyph_collection, remove_billboard(rotation), model, space, offset)

    nothing
end

function draw_glyph_collection(scene, ctx, positions, glyph_collections::AbstractArray, rotation, model::SMatrix, space, offset)

    # TODO: why is the Ref around model necessary? doesn't broadcast_foreach handle staticarrays matrices?
    broadcast_foreach(positions, glyph_collections, rotation,
        Ref(model), space, offset) do pos, glayout, ro, mo, sp, off

        draw_glyph_collection(scene, ctx, pos, glayout, ro, mo, sp, off)
    end
end

function draw_glyph_collection(scene, ctx, position, glyph_collection, rotation, model, space, offsets)

    glyphs = glyph_collection.glyphs
    glyphoffsets = glyph_collection.origins
    fonts = glyph_collection.fonts
    rotations = glyph_collection.rotations
    scales = glyph_collection.scales
    colors = glyph_collection.colors
    strokewidths = glyph_collection.strokewidths
    strokecolors = glyph_collection.strokecolors

    Cairo.save(ctx)

    broadcast_foreach(glyphs, glyphoffsets, fonts, rotations, scales, colors, strokewidths, strokecolors, offsets) do glyph,
        glyphoffset, font, rotation, scale, color, strokewidth, strokecolor, offset

        cairoface = set_ft_font(ctx, font)
        old_matrix = get_font_matrix(ctx)

        p3_offset = to_ndim(Point3f0, offset, 0)

        glyph in ('\r', '\n') && return

        Cairo.save(ctx)
        Cairo.set_source_rgba(ctx, rgbatuple(color)...)

        if space == :data
            # in data space, the glyph offsets are just added to the string positions
            # and then projected

            # glyph position in data coordinates (offset has rotation applied already)
            gpos_data = to_ndim(Point3f0, position, 0) .+ glyphoffset .+ p3_offset

            scale3 = scale isa Number ? Point3f0(scale, scale, 0) : to_ndim(Point3f0, scale, 0)

            # this could be done better but it works at least

            # the CairoMatrix is found by transforming the right and up vector
            # of the character into screen space and then subtracting the projected
            # origin. The resulting vectors give the directions in which the character
            # needs to be stretched in order to match the 3D projection

            xvec = rotation * (scale3[1] * Point3f0(1, 0, 0))
            yvec = rotation * (scale3[2] * Point3f0(0, -1, 0))

            glyphpos = project_position(scene, gpos_data, Mat4f0(I))
            xproj = project_position(scene, gpos_data + xvec, Mat4f0(I))
            yproj = project_position(scene, gpos_data + yvec, Mat4f0(I))

            xdiff = xproj - glyphpos
            ydiff = yproj - glyphpos

            mat = Cairo.CairoMatrix(
                xdiff[1], xdiff[2],
                ydiff[1], ydiff[2],
                0, 0,
            )

        elseif space == :screen
            # in screen space, the glyph offsets are added after projecting
            # the string position into screen space
            glyphpos = project_position(
                scene,
                position,
                Mat4f0(I)) .+ (p3_to_p2(glyphoffset .+ p3_offset)) .* (1, -1) # flip for Cairo
            # and the scale is just taken as is
            scale = length(scale) == 2 ? scale : SVector(scale, scale)

            mat = scale_matrix(scale...)
        else
            error()
        end

        Cairo.save(ctx)
        Cairo.move_to(ctx, glyphpos...)
        set_font_matrix(ctx, mat)
        if space == :screen
            Cairo.rotate(ctx, to_2d_rotation(rotation))
        end
        Cairo.show_text(ctx, string(glyph))
        Cairo.restore(ctx)

        if strokewidth > 0 && strokecolor != RGBAf0(0, 0, 0, 0)
            Cairo.save(ctx)
            Cairo.move_to(ctx, glyphpos...)
            set_font_matrix(ctx, mat)
            if space == :screen
                Cairo.rotate(ctx, to_2d_rotation(rotation))
            end
            Cairo.text_path(ctx, string(glyph))
            Cairo.set_source_rgba(ctx, rgbatuple(strokecolor)...)
            Cairo.set_line_width(ctx, strokewidth)
            Cairo.stroke(ctx)
            Cairo.restore(ctx)
        end
        Cairo.restore(ctx)

        cairo_font_face_destroy(cairoface)
        set_font_matrix(ctx, old_matrix)
    end

    Cairo.restore(ctx)

    nothing
end

################################################################################
#                                Heatmap, Image                                #
################################################################################

"""
    regularly_spaced_array_to_range(arr)
If possible, converts `arr` to a range.
If not, returns array unchanged.
"""
function regularly_spaced_array_to_range(arr)
    diffs = unique!(sort!(diff(arr)))
    step = sum(diffs) ./ length(diffs)
    if all(x-> x ≈ step, diffs)
        m, M = extrema(arr)
        if step < zero(step)
            m, M = M, m
        end
        # don't use stop=M, since that may not include M
        return range(m; step=step, length=length(arr))
    else
        return arr
    end
end

regularly_spaced_array_to_range(arr::AbstractRange) = arr

"""
    interpolation_flag(is_vector, interp, wpx, hpx, w, h)

* is_vector: if we're using vector backend
* interp: does the user want to interpolate?
* wpx, hpx: projected size of the image in pixels, so the actual width in pixels on screen
* w, h: size of image in pixels
"""
function interpolation_flag(is_vector, interp, wpx, hpx, w, h)
    if interp
        if is_vector
            return Cairo.FILTER_BILINEAR
        else
            return Cairo.FILTER_BEST
        end
    else
        if wpx < w || hpx < h
            # if size of image size in pixels is larger then the rectangle it gets drawn into,
            # the pixels will be smaller than what ends up on screen, so one won't be able to see rectangles.
            # In that case, we need to apply filtering, or we get artifacts from incorrectly downsampling!
            return interpolation_flag(is_vector, true, wpx, hpx, w, h)
        else
            return Cairo.FILTER_NEAREST
        end
    end
end

function draw_atomic(scene::Scene, screen::CairoScreen, primitive::Union{Heatmap, Image})
    ctx = screen.context
    image = primitive[3][]
    xs, ys = primitive[1][], primitive[2][]
    if !(xs isa AbstractVector)
        l, r = extrema(xs)
        N = size(image, 1)
        xs = range(l, r, length = N+1)
    else
        xs = regularly_spaced_array_to_range(xs)
    end
    if !(ys isa AbstractVector)
        l, r = extrema(ys)
        N = size(image, 2)
        ys = range(l, r, length = N+1)
    else
        ys = regularly_spaced_array_to_range(ys)
    end
    model = primitive[:model][]
    interp = to_value(get(primitive, :interpolate, true))
    weird_cairo_limit = (2^15) - 23

    # Debug attribute we can set to disable fastpath
    # probably shouldn't really be part of the interface
    fast_path = to_value(get(primitive, :fast_path, true))
    # Vector backends don't support FILTER_NEAREST for interp == false, so in that case we also need to draw rects
    is_vector = is_vector_backend(ctx)
    t = Makie.transform_func_obs(primitive)[]
    identity_transform = t === identity || t isa Tuple && all(x-> x === identity, t)
    if fast_path && xs isa AbstractRange && ys isa AbstractRange && !(is_vector && !interp) && identity_transform
        imsize = ((first(xs), last(xs)), (first(ys), last(ys)))

        # find projected image corners
        # this already takes care of flipping the image to correct cairo orientation
        xy = project_position(scene, Point2f0(first.(imsize)), model)
        xymax = project_position(scene, Point2f0(last.(imsize)), model)
        w, h = xymax .- xy

        s = to_cairo_image(image, primitive)

        interp_flag = interpolation_flag(is_vector, interp, abs(w), abs(h), s.width, s.height)

        if s.width > weird_cairo_limit || s.height > weird_cairo_limit
            error("Cairo stops rendering images bigger than $(weird_cairo_limit), which is likely a bug in Cairo. Please resample your image/heatmap with e.g. `ImageTransformations.imresize`")
        end

        Cairo.rectangle(ctx, xy..., w, h)
        Cairo.save(ctx)
        Cairo.translate(ctx, xy...)
        Cairo.scale(ctx, w / s.width, h / s.height)
        Cairo.set_source_surface(ctx, s, 0, 0)
        p = Cairo.get_source(ctx)
        # this is needed to avoid blurry edges
        Cairo.pattern_set_extend(p, Cairo.EXTEND_PAD)
        # Set filter doesn't work!?
        Cairo.pattern_set_filter(p, interp_flag)
        Cairo.fill(ctx)
        Cairo.restore(ctx)

    else
        # We need to draw rectangles for vector backends, or irregular grids
        if interp
            error("Interpolation for non gridded heatmaps/images isn't supported right now. Please use interpolate=false for this plot")
        end
        # find projected image corners
        # this already takes care of flipping the image to correct cairo orientation
        xys = [project_position(scene, Point2f0(x, y), model) for x in xs, y in ys]
        colors = to_rgba_image(image, primitive)

        # Note: xs and ys should have size ni+1, nj+1
        ni, nj = size(image)
        if ni + 1 != length(xs) || nj + 1 != length(ys)
            error("Error in conversion pipeline. xs and ys should have size ni+1, nj+1. Found: xs: $(length(xs)), ys: $(length(ys)), ni: $(ni), nj: $(nj)")
        end
        @inbounds for i in 1:ni, j in 1:nj
            x0, y0 = xys[i, j]
            x1, y1 = xys[i+1, j+1]
            w = x1 - x0; h = y1 - y0

            # there are usually white lines between directly adjacent rectangles
            # in vector graphics because of anti-aliasing

            # if we let each cell stick out (bulge) a little bit (half a point) under its neighbors
            # those lines disappear

            # we heuristically only do this if the adjacent cells are fully opaque
            # and if we're not in the last row / column so the overall heatmap doesn't get bigger

            # this should be the most common case by far, though

            xbulge = if i < ni && alpha(colors[i+1, j]) == 1
                0.5
            else
                0.0
            end
            ybulge = if j < nj && alpha(colors[i, j+1]) == 1
                0.5
            else
                0.0
            end

            # we add the bulge in the direction of cell width / height in case the axes are reversed
            Cairo.rectangle(ctx, x0, y0, w + sign(w) * xbulge, h + sign(h) * ybulge)
            Cairo.set_source_rgba(ctx, rgbatuple(colors[i, j])...)
            Cairo.fill(ctx)
        end
    end
end

################################################################################
#                                     Mesh                                     #
################################################################################

function draw_atomic(scene::Scene, screen::CairoScreen, primitive::Makie.Mesh)
    if scene.camera_controls[] isa Union{Camera2D, Makie.PixelCamera}
        draw_mesh2D(scene, screen, primitive)
    else
        if !haskey(primitive, :faceculling)
            primitive[:faceculling] = Node(-10)
        end
        draw_mesh3D(scene, screen, primitive)
    end
    return nothing
end

function draw_mesh2D(scene, screen, primitive)
    @get_attribute(primitive, (color,))

    colormap = get(primitive, :colormap, nothing) |> to_value |> to_colormap
    colorrange = get(primitive, :colorrange, nothing) |> to_value
    ctx = screen.context
    model = primitive.model[]
    mesh = GeometryBasics.mesh(primitive[1][])
    # Priorize colors of the mesh if present
    # This is a hack, which needs cleaning up in the Mesh plot type!
    color = hasproperty(mesh, :color) ? mesh.color : color
    vs =  decompose(Point, mesh); fs = decompose(TriangleFace, mesh)
    uv = hasproperty(mesh, :uv) ? mesh.uv : nothing
    pattern = Cairo.CairoPatternMesh()

    cols = per_face_colors(color, colormap, colorrange, nothing, vs, fs, nothing, uv)
    for (f, (c1, c2, c3)) in zip(fs, cols)
        t1, t2, t3 =  project_position.(scene, vs[f], (model,)) #triangle points
        Cairo.mesh_pattern_begin_patch(pattern)

        Cairo.mesh_pattern_move_to(pattern, t1...)
        Cairo.mesh_pattern_line_to(pattern, t2...)
        Cairo.mesh_pattern_line_to(pattern, t3...)

        mesh_pattern_set_corner_color(pattern, 0, c1)
        mesh_pattern_set_corner_color(pattern, 1, c2)
        mesh_pattern_set_corner_color(pattern, 2, c3)

        Cairo.mesh_pattern_end_patch(pattern)
    end
    Cairo.set_source(ctx, pattern)
    Cairo.close_path(ctx)
    Cairo.paint(ctx)
    return nothing
end

function average_z(positions, face)
    vs = positions[face]
    sum(v -> v[3], vs) / length(vs)
end

nan2zero(x) = !isnan(x) * x

function draw_mesh3D(
        scene, screen, primitive;
        mesh = primitive[1][], pos = Vec4f0(0), scale = 1f0
    )
    @get_attribute(primitive, (color, shading, lightposition, ambient, diffuse,
        specular, shininess, faceculling))

    colormap = get(primitive, :colormap, nothing) |> to_value |> to_colormap
    colorrange = get(primitive, :colorrange, nothing) |> to_value
    matcap = get(primitive, :matcap, nothing) |> to_value
    # Priorize colors of the mesh if present
    color = hasproperty(mesh, :color) ? mesh.color : color

    ctx = screen.context

    model = primitive.model[]
    view = scene.camera.view[]
    projection = scene.camera.projection[]
    normalmatrix = get(
        scene.attributes, :normalmatrix, let
            i = SOneTo(3)
            transpose(inv(view[i, i] * model[i, i]))
        end
    )

    # Mesh data
    # transform to view/camera space
    func = Makie.transform_func_obs(scene)[]
    # pass func as argument to function, so that we get a function barrier
    # and have `func` be fully typed inside closure
    vs = broadcast(decompose(Point, mesh), (func,)) do v, f
        # Should v get a nan2zero?
        v = Makie.apply_transform(f, v)
        p4d = to_ndim(Vec4f0, scale .* to_ndim(Vec3f0, v, 0f0), 1f0)
        view * (model * p4d .+ to_ndim(Vec4f0, pos, 0f0))
    end
    fs = decompose(GLTriangleFace, mesh)
    uv = hasproperty(mesh, :uv) ? mesh.uv : nothing
    ns = map(n -> normalize(normalmatrix * n), normals(mesh))
    cols = per_face_colors(
        color, colormap, colorrange, matcap, vs, fs, ns, uv,
        get(primitive, :lowclip, nothing) |> to_value |> color_or_nothing,
        get(primitive, :highclip, nothing) |> to_value |> color_or_nothing,
        get(primitive, :nan_color, nothing) |> to_value |> color_or_nothing
    )

    # Liight math happens in view/camera space
    if lightposition == :eyeposition
        lightposition = scene.camera.eyeposition[]
    end
    lightpos = (view * to_ndim(Vec4f0, lightposition, 1.0))[Vec(1, 2, 3)]

    # Camera to screen space
    ts = map(vs) do v
        clip = projection * v
        @inbounds begin
            p = (clip ./ clip[4])[Vec(1, 2)]
            p_yflip = Vec2f0(p[1], -p[2])
            p_0_to_1 = (p_yflip .+ 1f0) / 2f0
        end
        p = p_0_to_1 .* scene.camera.resolution[]
        Vec3f0(p[1], p[2], clip[3])
    end

    # Approximate zorder
    zorder = sortperm(fs, by = f -> average_z(ts, f))

    # Face culling
    zorder = filter(i -> any(last.(ns[fs[i]]) .> faceculling), zorder)

    pattern = Cairo.CairoPatternMesh()
    for k in reverse(zorder)
        f = fs[k]
        t1, t2, t3 = ts[f]

        # light calculation
        c1, c2, c3 = if shading
            map(ns[f], vs[f], cols[k]) do N, v, c
                L = normalize(lightpos .- v[Vec(1,2,3)])
                diff_coeff = max(dot(L, N), 0.0)
                H = normalize(L + normalize(-v[SOneTo(3)]))
                spec_coeff = max(dot(H, N), 0.0)^shininess
                c = RGBA(c)
                new_c = (ambient .+ diff_coeff .* diffuse) .* Vec3f0(c.r, c.g, c.b) .+
                        specular * spec_coeff
                RGBA(new_c..., c.alpha)
            end
        else
            cols[k]
        end
        # debug normal coloring
        # n1, n2, n3 = Vec3f0(0.5) .+ 0.5ns[f]
        # c1 = RGB(n1...)
        # c2 = RGB(n2...)
        # c3 = RGB(n3...)

        Cairo.mesh_pattern_begin_patch(pattern)

        Cairo.mesh_pattern_move_to(pattern, t1[1], t1[2])
        Cairo.mesh_pattern_line_to(pattern, t2[1], t2[2])
        Cairo.mesh_pattern_line_to(pattern, t3[1], t3[2])

        mesh_pattern_set_corner_color(pattern, 0, c1)
        mesh_pattern_set_corner_color(pattern, 1, c2)
        mesh_pattern_set_corner_color(pattern, 2, c3)

        Cairo.mesh_pattern_end_patch(pattern)
    end
    Cairo.set_source(ctx, pattern)
    Cairo.close_path(ctx)
    Cairo.paint(ctx)
    return nothing
end

################################################################################
#                                   Surface                                    #
################################################################################

function draw_atomic(scene::Scene, screen::CairoScreen, primitive::Makie.Surface)
    # Pretend the surface plot is a mesh plot and plot that instead
    mesh = surface2mesh(primitive[1][], primitive[2][], primitive[3][])
    old = primitive[:color]
    if old[] === nothing
        primitive[:color] = primitive[3]
    end
    if !haskey(primitive, :faceculling)
        primitive[:faceculling] = Node(-10)
    end
    draw_mesh3D(scene, screen, primitive, mesh=mesh)
    primitive[:color] = old
    return nothing
end

function surface2mesh(xs::Makie.ClosedInterval, ys, zs::AbstractMatrix)
    surface2mesh(range(xs.left, xs.right, length = size(zs, 1)), ys, zs)
end

function surface2mesh(xs, ys::Makie.ClosedInterval, zs::AbstractMatrix)
    surface2mesh(xs, range(ys.left, ys.right, length = size(zs, 2)), zs)
end

function surface2mesh(xs::Makie.ClosedInterval, ys::Makie.ClosedInterval, zs::AbstractMatrix)
    surface2mesh(
        range(xs.left, xs.right, length = size(zs, 1)),
        range(ys.left, ys.right, length = size(zs, 2)),
        zs)
end

function surface2mesh(xs::AbstractVector, ys::AbstractVector, zs::AbstractMatrix)
    ps = [nan2zero.(Point3f0(xs[i], ys[j], zs[i, j])) for j in eachindex(ys) for i in eachindex(xs)]
    idxs = LinearIndices(size(zs))
    faces = [
        QuadFace(idxs[i, j], idxs[i+1, j], idxs[i+1, j+1], idxs[i, j+1])
        for j in 1:size(zs, 2)-1 for i in 1:size(zs, 1)-1
    ]
    normal_mesh(ps, faces)
end

function surface2mesh(xs::AbstractMatrix, ys::AbstractMatrix, zs::AbstractMatrix)
    ps = [nan2zero.(Point3f0(xs[i, j], ys[i, j], zs[i, j])) for j in 1:size(zs, 2) for i in 1:size(zs, 1)]
    idxs = LinearIndices(size(zs))
    faces = [
        QuadFace(idxs[i, j], idxs[i+1, j], idxs[i+1, j+1], idxs[i, j+1])
        for j in 1:size(zs, 2)-1 for i in 1:size(zs, 1)-1
    ]
    normal_mesh(ps, faces)
end

################################################################################
#                                 MeshScatter                                  #
################################################################################

function draw_atomic(scene::Scene, screen::CairoScreen, primitive::Makie.MeshScatter)
    @get_attribute(primitive, (color, model, marker, markersize, rotations))

    if color isa AbstractArray{<: Number}
        color = numbers_to_colors(color, primitive)
    end

    m = convert_attribute(marker, key"marker"(), key"meshscatter"())
    pos = primitive[1][]
    # For correct z-ordering we need to be in view/camera or screen space
    model = copy(model)
    view = scene.camera.view[]

    zorder = sortperm(pos, by = p -> begin
        p4d = to_ndim(Vec4f0, to_ndim(Vec3f0, p, 0f0), 1f0)
        cam_pos = view * model * p4d
        cam_pos[3] / cam_pos[4]
    end, rev=false)

    submesh = Attributes(
        model=model,
        color=color,
        shading=primitive.shading, lightposition=primitive.lightposition,
        ambient=primitive.ambient, diffuse=primitive.diffuse,
        specular=primitive.specular, shininess=primitive.shininess,
        faceculling=get(primitive, :faceculling, -10)
    )

    if !(rotations isa Vector)
        R = Makie.rotationmatrix4(to_rotation(rotations))
        submesh[:model] = model * R
    end
    scales = primitive[:markersize][]

    for i in zorder
        p = pos[i]
        if color isa AbstractVector
            submesh[:color] = color[i]
        end
        if rotations isa Vector
            R = Makie.rotationmatrix4(to_rotation(rotations[i]))
            submesh[:model] = model * R
        end
        scale = markersize isa Vector ? markersize[i] : markersize

        draw_mesh3D(
            scene, screen, submesh, mesh = m, pos = p,
            scale = scale isa Real ? Vec3f0(scale) : to_ndim(Vec3f0, scale, 1f0)
        )
    end

    return nothing
end


using Oceananigans: znodes

"""
    animate_training_data(ds, ds_coarse; filepath, frameskip=1, fps=15)

Create an animation of the variable `v` and coarse-grained variable `v̄`.
`xlabel` should be specified for the plot. A `filepath should also be specified`.
A `.gif` and `.mp4` will be saved. `frameskip` > 1 can be used to skip some frames
to produce the animation more quickly. The output animation will use `fps` frames
per second.
"""
function animate_training_data(ds, ds_coarse; filepath, frameskip=1, fps=15)
    T = ds["T"]
    wT = ds["wT"]
    zc = znodes(T)
    zf = znodes(wT)
    times = T.times

    T_coarse = ds_coarse["T"]
    wT_coarse = ds_coarse["wT"]
    z̄c = znodes(T_coarse)
    z̄f = znodes(wT_coarse)

    kwargs = (linewidth=3, linealpha=0.8, ylims=extrema(zf),
              grid=false, legend=:bottomright, framestyle=:box,
              foreground_color_legend=nothing, background_color_legend=nothing)

    T_lims = interior(T) |> extrema
    wT_lims = interior(wT) |> extrema

    anim = @animate for n=1:frameskip:length(times)
        @info "Plotting free convection data [frame $n/$(length(times))]: $filepath..."

        time_str = @sprintf("%.2f days", times[n] / days)

        wT_profile = interior(wT)[1, 1, :, n]
        wT_coarse_profile = interior(wT_coarse)[1, 1, :, n]

        wT_plot = plot()
        plot!(wT_plot, wT_profile, zf, label="LES (Nz=$(length(zf)))", xlabel="Heat flux wT (m/s °C)",
              xlims=wT_lims, ylabel="Depth z (meters)", title="Free convection: $time_str"; kwargs...)
        plot!(wT_plot, wT_coarse_profile, z̄f, label="coarse (Nz=$(length(z̄f)))"; kwargs...)

        T_profile = interior(T)[1, 1, :, n]
        T_coarse_profile = interior(T_coarse)[1, 1, :, n]

        T_plot = plot()
        plot!(T_plot, T_profile, zc, label="LES (Nz=$(length(zc)))", xlabel="Temperature T (°C)",
              xlims=T_lims; kwargs...)
        plot!(T_plot, T_coarse_profile, z̄c, label="coarse (Nz=$(length(z̄c)))"; kwargs...)

        plot(wT_plot, T_plot, dpi=200)
    end

    @info "Saving $filepath..."
    mp4(anim, filepath * ".mp4", fps=fps)
    gif(anim, filepath * ".gif", fps=fps)

    return nothing
end

function animate_learned_free_convection(ds, NN, NN_function, NDEType, algorithm, T_scaling, wT_scaling; filepath, frameskip=1, fps=15)
    T = ds["T"]
    wT = ds["wT"]
    Nz = size(T, 3)
    Nt = size(T, 4)
    zc = znodes(T)
    zf = znodes(wT)
    times = T.times

    nde_params = FreeConvectionNDEParameters(ds, T_scaling, wT_scaling)
    T₀ = T_scaling.(interior(T)[1, 1, :, 1])
    nde = NDEType(NN, ds)
    nde_sol = solve_nde(nde, NN, T₀, algorithm, nde_params)

    kwargs = (linewidth=3, linealpha=0.8, ylims=extrema(zf),
              grid=false, legend=:bottomright, framestyle=:box,
              foreground_color_legend=nothing, background_color_legend=nothing)

    T_lims = interior(T) |> extrema
    wT_lims = interior(wT) |> extrema

    anim = @animate for n=1:frameskip:Nt
        @info "Plotting learned free convection [frame $n/$Nt]: $filepath ..."

        time_str = @sprintf("%.2f days", times[n] / days)

        wT_plot = plot()
        T_plot = plot()

        T_profile = interior(T)[1, 1, :, n]
        wT_profile = interior(wT)[1, 1, :, n]

        plot!(wT_plot, wT_profile, zf, label="Oceananigans", xlabel="Heat flux w'T' (m/s K)",
              xlims=wT_lims, ylabel="Depth z (meters)", title="Free convection: $time_str"; kwargs...)

        plot!(T_plot, T_profile, zc[:], label="Oceananigans", xlabel="Temperature T (°C)",
              xlims=T_lims; kwargs...)

        T_NN = nde_sol[n]
        bottom_flux = wT_scaling(interior(wT)[1, 1, 1, n])
        top_flux = wT_scaling(interior(wT)[1, 1, Nz+1, n])
        input = FreeConvectionTrainingDataInput(T_NN, bottom_flux, top_flux)
        wT_NN = unscale.(NN_function(input), Ref(wT_scaling))
        T_NN = unscale.(T_NN, Ref(T_scaling))

        plot!(wT_plot, wT_NN, zf, label="neural network"; kwargs...)
        plot!(T_plot, T_NN, zc, label="NDE"; kwargs...)

        plot(wT_plot, T_plot, dpi=200)
    end

    @info "Saving $filepath..."
    mp4(anim, filepath * ".mp4", fps=fps)
    gif(anim, filepath * ".gif", fps=fps)

    return nothing
end