text:
"""
   Solve linear least square problem by LSQR method, the method solve
   min ||Ax-b||₂² if damp = 0 or
   min ||Ax-b||₂² + damp² ||x||₂² otherwise;
   A can be an explicit matrix or a function(linear operator)
   A(x, 1; params...) forward operator
   A(x, 2; params...) adjoint operator; the input and output are standard vector
   atol: the stopping tolerance for the residue norm
   conlim: the limite of the estimated condition number
   params is a dictionary contain the parameters for forward modeling operator

   x: final solution
   rnorm: sqrt(||Ax-b||₂² + damp² ||x||₂²)
   Anorm: the Frobenous norm of A
   Acond: the estimated condition number of A
   diagInv: diag(Dk * Dk'), where Dk is n*k matrix of search directions after k iterations,
          it satisfy Dk' (A'A + damp²* I) Dk = I, so the var is the estimated diagonals
          of the inverse matrix.
"""
function LSQR(A, b::Vector{Tv}; params=[], damp=0.0,
              atol=1e-9, btol=1e-9, conlim=1e8, maxIter=100,
              print_flag=false, diagInv_flag=false) where {Tv<:Number}

    # determine whether A is a matrix or operator
    if isa(A, Matrix)
       (m, n) = size(A)
       explicitA = true
    elseif isa(A, Function)
       m = length(b)
       explicitA = false
    else
       error("A either to be a matrix or operator")
    end

    msg = ["the exact solution is x = 0"
           "Ax-b is small enough given atol and btol"
           "the estimated cond(Abar) has exceeded conlim"
           "the max iterations has been reached"]

    # initialize some intermediate variables
    itn   = 0     # count number of iterations
    istop = 0     # the indicater for stop reason
    ctol  = conlim > 0 ? 1.0/conlim : 0.0
    Anorm = 0.0   # the estimated norm of vcat(A, damp I)
    Acond = 0.0   # the estimated condition number of vcat(A, damp I)
    xnorm = 0.0   # the norm of computed x
    dampsq = damp * damp;
    ddnorm = 0.0; res2 = 0.0; xxnorm = 0.0;
    z      = 0.0; cs2  =-1.0; sn2    = 0.0;

    # set up the first vectors for lanczos bidiagonalization
    # beta * u = b; alpha * v = A' * u
    u    = copy(b)
    beta = norm(u)
    alpha= 0.0

    if beta > 0.0
       u = u / beta
       if explicitA
          v = A' * u
       else
          v = A(u, 2; params...)
       end
       n = length(v)
       x = zeros(eltype(b), n)
       alpha = norm(v)
    end

    if alpha > 0.0
       v = v / alpha
       w = copy(v)
    end

    # estimate the diagonal element of inverse (A'A + dmap^2 I)^(-1)
    if diagInv_flag
       diagInv = zeros(eltype(b), n)
    end

    Arnorm = alpha * beta
    if Arnorm == 0
       println(msg[1])
       return
    end

    rhobar = alpha; phibar = beta ; bnorm  = beta ;
    rnorm  = beta ; r1norm = rnorm; r2norm = rnorm;

    # some header for print info
    test1 = 1.0; test2 = alpha / beta;
    header = "Itn               r1norm               r2norm               norm A               cond A"
    if print_flag
       println("")
       println(header)
       @printf("%3g %20.10e %20.10e %20.10e %20.10e\n", itn, r1norm, r2norm, test1, test2)
    end

     while itn < maxIter

           itn = itn + 1

           # do next step lanczos bidiagonalization
           if explicitA
              u = A * v              - alpha * u
           else
              u = A(v, 1; params...) - alpha * u
           end
           beta = norm(u)

           if beta > 0.0
              u = u / beta
              Anorm = norm([Anorm; alpha; beta; damp])
              if explicitA
                 v = A' * u             - beta * v
              else
                 v = A(u, 2; params...) - beta * v
              end

              alpha = norm(v)
              if alpha > 0.0
                 v = v / alpha
              end
           end

           # plane rotation to eliminate damping parameter
           rhobar1 =   norm([rhobar; damp])
           cs1     =   rhobar  / rhobar1
           sn1     =   damp    / rhobar1
           psi     =   sn1     * phibar
           phibar  =   cs1     * phibar

           # plane rotation
           rho     =   norm([rhobar1; beta])
           cs      =   rhobar1 / rho
           sn      =   beta    / rho
           theta   =   sn      * alpha
           rhobar  = - cs      * alpha
           phi     =   cs      * phibar
           phibar  =   sn      * phibar
           tau     =   sn      * phi

           # update x and w
           t1      =   phi     / rho
           t2      = - theta   / rho
           dk      =   w       / rho
           x       =   x +  t1 * w
           w       =   v +  t2 * w
           ddnorm  =   ddnorm  + norm(dk)^2

           # estimate the diagonal element of Inverse matrix
           if diagInv_flag
              diagInv = diagInv + dk .* dk
           end

           # plane rotation
           delta  =   sn2    * rho
           gambar = - cs2    * rho
           rhs    =   phi    - delta * z
           zbar   =   rhs    / gambar
           xnorm  =   sqrt(xxnorm + zbar^2)
           gamma  =   norm([gambar; theta])
           cs2    =   gambar / gamma
           sn2    =   theta  / gamma
           z      =   rhs    / gamma
           xxnorm =   xxnorm + z^2

           # test for convergence
           Acond  = Anorm * sqrt(ddnorm)
           res1   = phibar^2
           res2   = res2 + psi^2
           rnorm  = sqrt(res1 + res2)
           Arnorm = alpha * abs(tau)

           r1sq   = rnorm^2 - dampsq * xxnorm
           r1norm = sqrt(abs(r1sq)); r1sq = r1sq > 0.0 ? r1sq : -r1sq;
           r2norm = rnorm

           test1  = rnorm  / bnorm
           test2  = Arnorm / (Anorm * rnorm)
           test3  = 1.     / Acond
           t1     = test1  / (1   + Anorm * xnorm / bnorm)
           rtol   = btol   + atol * Anorm * xnorm / bnorm

           istop  = itn        >= maxIter ? 7 : 0
           istop  = 1. + test3 <= 1.      ? 6 : 0
           istop  = 1. + test2 <= 1.      ? 5 : 0
           istop  = 1. + t1    <= 1.      ? 4 : 0
           istop  = test3      <= ctol    ? 3 : 0
           istop  = test2      <= atol    ? 2 : 0
           istop  = test1      <= rtol    ? 1 : 0

           if istop > 0
              break
           end

           if print_flag
              @printf("%3g %20.10e %20.10e %20.10e %20.10e\n", itn, r1norm, r2norm, test1, test2)
           end

     end

     if diagInv_flag
        optout = Dict(:r1norm=>r1norm, :r2norm=>r2norm,
                      :Anorm=>Anorm  , :Acond=>Acond,
                      :Arnorm=>Arnorm, :diagInv=>diagInv)
     else
        optout = Dict(:r1norm=>r1norm, :r2norm=>r2norm,
                      :Anorm=>Anorm  , :Acond=>Acond,
                      :Arnorm=>Arnorm                 )
     end

     return x, optout
end

# the LSQR method for LSRTM
function LSQR(A::Function, b::Ts; params=[], path="NULL", damp=0.0,
              atol=1e-9, btol=1e-9, conlim=1e8, maxIter=100,
              print_flag=false, save_process_flag=false, diagInv_flag=false) where {Ts<:Union{String, Vector{String}}}

    # set the work space
    if path == "NULL"
       path = pwd()
    end

    # number of observation files
    if typeof(b) == String
       num_file = 1
       u = join([path "/u.bin"])
       q = join([path "/q.bin"])

    elseif typeof(b) == Vector{String}
       num_file = length(b)
       path_tmp = joinpath(path, "mu")
       if !isdir(path_tmp)
          mkdir(path_tmp)
          if !isdir(path_tmp)
             error("could not create directory")
          end
       end
       u = Vector{String}(num_file)
       q = Vector{String}(num_file)
       for i = 1 : num_file
           u[i] = join([path_tmp "/u" "_" "$i" ".bin"])
           q[i] = join([path_tmp "/q" "_" "$i" ".bin"])
       end
    end

    v = join([path "/v.bin"])
    x = join([path "/x.bin"])
    w = join([path "/w.bin"])
    p = join([path "/p.bin"])
    dk= join([path "/dk.bin"])

    msg = ["the exact solution is x = 0"
           "Ax-b is small enough given atol and btol"
           "the estimated cond(Abar) has exceeded conlim"
           "the max iterations has been reached"]

    # initialize some intermediate variables
    itn   = 0     # count number of iterations
    istop = 0     # the indicater for stop reason
    ctol  = conlim > 0 ? 1.0/conlim : 0.0
    Anorm = 0.0   # the estimated norm of vcat(A, damp I)
    Acond = 0.0   # the estimated condition number of vcat(A, damp I)
    xnorm = 0.0   # the norm of computed x
    dampsq = damp * damp
    ddnorm = 0.0; res2 = 0.0; xxnorm = 0.0;
    z      = 0.0; cs2  =-1.0; sn2    = 0.0;

    # set up the first vectors for lanczos bidiagonalization
    # beta * u = b; alpha * v = A' * u
    cp(b, u, remove_destination=true)
    beta = norm(u)
    alpha= 0.0

    if beta > 0.0
       scaling!(u, beta)

       if typeof(params) <: Dict{Symbol, Any}
          params[:path_m] = v
       elseif typeof(params) == Vector{Dict}
         for i = 1 : length(params)
             params[i][:path_m] = v
         end
       end
       v = A(u, 2; params=params)

       # x = zeros(v)
       (h, d) = read_USdata(v)
       write_USdata(x, h, zeros(d))
       alpha = norm(v)
    end

    if alpha > 0.0
       scaling!(v, alpha)
       cp(v, w, remove_destination=true)
    end

    # estimate the diagonal element of inverse (A'A + dmap^2 I)^(-1)
    if diagInv_flag
       diagInv = join([path "/diagInv.bin"])
       (h, d) = read_USdata(v)
       write_USdata(diagInv, h, zeros(d))
    end

    Arnorm = alpha * beta
    if Arnorm == 0
       println(msg[1])
       return
    end

    rhobar = alpha; phibar = beta ; bnorm  = beta ;
    rnorm  = beta ; r1norm = rnorm; r2norm = rnorm;

    # some header for print info
    test1 = 1.0; test2 = alpha / beta;
    header = "Itn               r1norm               r2norm               norm A               cond A"
    if print_flag
       println("")
       println(header)
       @printf("%3g %20.10e %20.10e %20.10e %20.10e\n", itn, r1norm, r2norm, test1, test2)
    end

     while itn < maxIter

           itn = itn + 1

           # do next step lanczos bidiagonalization; forward
           if typeof(params) <: Dict{Symbol, Any}
              params[:path_fwd] = q
           elseif typeof(params) == Vector{Dict}
              for i = 1 : length(params)
                  params[i][:path_fwd] = q[i]
              end
           end
           q = A(v, 1; params=params)
           x_plus_alpha_q!(q, -alpha, u; iflag=2)
           beta = norm(u)

           if beta > 0.0
              scaling!(u, beta)
              Anorm = norm([Anorm; alpha; beta; damp])

              # adjoint
              if typeof(params) <: Dict{Symbol, Any}
                 params[:path_m] = p
              elseif typeof(params) == Vector{Dict}
                 for i = 1 : length(params)
                     params[i][:path_m] = p
                 end
              end
              p = A(u, 2; params=params)
              x_plus_alpha_q!(p, -beta, v; iflag=2)

              alpha = norm(v)
              if alpha > 0.0
                 scaling!(v, alpha)
              end
           end

           # plane rotation to eliminate damping parameter
           rhobar1 =   norm([rhobar; damp])
           cs1     =   rhobar  / rhobar1
           sn1     =   damp    / rhobar1
           psi     =   sn1     * phibar
           phibar  =   cs1     * phibar

           # plane rotation
           rho     =   norm([rhobar1; beta])
           cs      =   rhobar1 / rho
           sn      =   beta    / rho
           theta   =   sn      * alpha
           rhobar  = - cs      * alpha
           phi     =   cs      * phibar
           phibar  =   sn      * phibar
           tau     =   sn      * phi

           # update x and w
           t1      =   phi     / rho
           t2      = - theta   / rho
           scaling!(dk, w, rho)
           x_plus_alpha_q!(x, t1, w; iflag=1)
           x_plus_alpha_q!(v, t2, w; iflag=2)
           ddnorm  =   ddnorm  + norm(dk)^2

           # save the result of each iteration
           if save_process_flag
              path_process = join([path "/iter" "_" "$itn" ".bin"])
              cp(x, path_process, remove_destination=true)
           end

           # estimate the diagonal element of Inverse matrix
           if diagInv_flag
              diagInv = diagInv + dk .* dk
           end

           # plane rotation
           delta  =   sn2    * rho
           gambar = - cs2    * rho
           rhs    =   phi    - delta * z
           zbar   =   rhs    / gambar
           xnorm  =   sqrt(xxnorm + zbar^2)
           gamma  =   norm([gambar; theta])
           cs2    =   gambar / gamma
           sn2    =   theta  / gamma
           z      =   rhs    / gamma
           xxnorm =   xxnorm + z^2

           # test for convergence
           Acond  = Anorm * sqrt(ddnorm)
           res1   = phibar^2
           res2   = res2 + psi^2
           rnorm  = sqrt(res1 + res2)
           Arnorm = alpha * abs(tau)

           r1sq   = rnorm^2 - dampsq * xxnorm
           r1norm = sqrt(abs(r1sq)); r1sq = r1sq > 0.0 ? r1sq : -r1sq;
           r2norm = rnorm

           test1  = rnorm  / bnorm
           test2  = Arnorm / (Anorm * rnorm)
           test3  = 1.     / Acond
           t1     = test1  / (1   + Anorm * xnorm / bnorm)
           rtol   = btol   + atol * Anorm * xnorm / bnorm

           istop  = itn        >= maxIter ? 7 : 0
           istop  = 1. + test3 <= 1.      ? 6 : 0
           istop  = 1. + test2 <= 1.      ? 5 : 0
           istop  = 1. + t1    <= 1.      ? 4 : 0
           istop  = test3      <= ctol    ? 3 : 0
           istop  = test2      <= atol    ? 2 : 0
           istop  = test1      <= rtol    ? 1 : 0

           if istop > 0
              break
           end

           if print_flag
              @printf("%3g %20.10e %20.10e %20.10e %20.10e\n", itn, r1norm, r2norm, test1, test2)
           end

     end

     if diagInv_flag
        optout = Dict(:r1norm=>r1norm, :r2norm=>r2norm,
                      :Anorm=>Anorm  , :Acond=>Acond,
                      :Arnorm=>Arnorm, :diagInv=>diagInv)
     else
        optout = Dict(:r1norm=>r1norm, :r2norm=>r2norm,
                      :Anorm=>Anorm  , :Acond=>Acond,
                      :Arnorm=>Arnorm                 )
     end

     return x, optout
end

# ==========================support functions for CGLS==========================
function cp(ps::Ts, pd::Ts; remove_destination=false) where {Ts<:Vector{String}}
    if length(ps) != length(pd)
       error("number of file does not match")
    end
    for i = 1 : length(ps)
        cp(ps[i], pd[i], remove_destination=remove_destination)
    end
    return nothing
end

function norm(s::Ts) where {Ts<:Union{String, Vector{String}}}

    if typeof(s) == String
       (hdr, d) = read_USdata(s)
       return vecnorm(d)
    elseif typeof(s) == Vector{String}
       r = 0.0
       for i = 1 : length(s)
           (hdr, d) = read_USdata(s[i])
           tmp = vecnorm(d)
           r = r + tmp * tmp
       end
       return sqrt(r)
    end
end

function x_plus_alpha_q!(x::Ts, alpha::Tv, q::Ts; iflag=1) where {Ts<:Union{String, Vector{String}}, Tv<:AbstractFloat}

    if typeof(x) == String
       (hdr_x, dx) = read_USdata(x)
       (hdr_q, dq) = read_USdata(q)
       for i = 1 : length(dx)
           dx[i] = dx[i] + alpha * dq[i]
       end

       if iflag == 1
          write_USdata(x, hdr_x, dx)
       elseif iflag == 2
          write_USdata(q, hdr_x, dx)
       end

    elseif typeof(x) == Vector{String}
       if length(x) != length(q)
          error("number of file does not match")
       end
       for i = 1 : length(x)
           (hdr_x, dx) = read_USdata(x[i])
           (hdr_q, dq) = read_USdata(q[i])
           for j = 1 : length(dx)
               dx[j] = dx[j] + alpha * dq[j]
           end

           if iflag == 1
              write_USdata(x[i], hdr_x, dx)
           elseif iflag == 2
              write_USdata(q[i], hdr_x, dx)
           end
       end
    end

    return nothing
end

function scaling!(u::Ts, beta::Tv) where {Ts<:Union{String, Vector{String}}, Tv<:AbstractFloat}

    beta_inv = 1.0 / beta

    if typeof(u) == String
       (hdr, d) = read_USdata(u)
       for i = 1 : length(d)
           d[i] = d[i] * beta_inv
       end
       write_USdata(u, hdr, d)

    elseif typeof(u) == Vector{String}

       for i = 1 : length(u)
           (hdr, d) = read_USdata(u[i])

           for j = 1 : length(d)
               d[j] = d[j] * beta_inv
           end

           write_USdata(u[i], hdr, d)
       end
    end
    return nothing
end

function scaling!(a::Ts, u::Ts, beta::Tv) where {Ts<:Union{String, Vector{String}}, Tv<:AbstractFloat}

    beta_inv = 1.0 / beta

    if typeof(u) == String
       (hdr, d) = read_USdata(u)
       for i = 1 : length(d)
           d[i] = d[i] * beta_inv
       end
       write_USdata(a, hdr, d)

    elseif typeof(u) == Vector{String}

       for i = 1 : length(u)
           (hdr, d) = read_USdata(u[i])

           for j = 1 : length(d)
               d[j] = d[j] * beta_inv
           end

           write_USdata(a[i], hdr, d)
       end
    end
    return nothing
end


text:
"""
Squared distance of two vertices.
"""
@i @inline function sqdistance(dist!, x1::AbstractVector{T}, x2::AbstractVector) where T
    @invcheckoff @inbounds for i=1:length(x1)
        @routine begin
            d ← zero(T)
            d += x1[i] - x2[i]
        end
        dist! += d^2
        ~@routine
    end
end

"""
    var_and_mean_sq(var!, mean!, sqv)

the variance and mean value from squared values.
"""
@i function var_and_mean_sq(var!, var_sum, mean!, mean_sum, sqv::AbstractVector{T}) where T
    @inbounds for i=1:length(sqv)
        mean_sum += sqv[i] ^ 0.5
        var_sum += sqv[i]
    end
    mean! += mean_sum / (@const length(sqv))
    @routine begin
        @zeros T var_anc1 var_anc2 sqmean
        var_anc1 += var_sum / (@const length(sqv))
        sqmean += mean! ^ 2
        var_anc1 -= sqmean
        var_anc2 += var_anc1 * (@const length(sqv))
    end
    var! += var_anc2 / (@const length(sqv)-1)
    ~@routine
end

"""The loss of graph embedding problem."""
@i function embedding_loss(out!::T, x) where T
    @routine @invcheckoff begin
        d1 ← zeros(T, length(L1))
        d2 ← zeros(T, length(L2))
        for i=1:length(L1)
            @inbounds sqdistance(d1[i], x[:,L1[i][1]],x[:,L1[i][2]])
        end
        for i=1:length(L2)
            @inbounds sqdistance(d2[i], x[:,L2[i][1]],x[:,L2[i][2]])
        end
        @zeros T v1 m1 v2 m2 vacc1 vacc2 acc1 acc2
        var_and_mean_sq(v1, vacc1, m1, acc1, d1)
        var_and_mean_sq(v2, vacc2, m2, acc2, d2)
        m1 -= m2 - 0.1
    end
    out! += v1 + v2
    if m1 > 0
        # to ensure mean(v2) > mean(v1)
        # if mean(v1)+0.1 - mean(v2) > 0, punish it.
        out! += exp(m1)
        out! -= 1
    end
    ~@routine
end

ext: jl


text:
# Translates SPN to Bayesian Networks and runs Hybrid-Product Belief Propagation to obtain MAP configuration
import GraphicalModels: FactorGraph, FGNode, VariableNode, FactorNode
import GraphicalModels.MessagePassing: HybridBeliefPropagation, update!, decode, setevidence!, setmapvar!
import GraphicalModels.MessagePassing: rmevidence!
# TODO: dichotomize nodes (make max no. of children for inner nodes = 2)

"""
    spn2bn(spn::SumProductNetwork)

Returns factor graph representing the computation graph of `spn`.
"""
function spn2bn(spn::SumProductNetwork)
    # recover dimensionality for each variable
    vdims = vardims(spn)
    variables = Dict{String,VariableNode}()
    # sizehint!(variables,length(vdims)+length(spn))
    factors = Dict{String,FactorNode}()
    # sizehint!(factors,length(spn))
    # Add root nodes for manifest variables
    for (v,d) in vdims
        variables["X" * string(v)] = VariableNode(d)
    end
    for i = length(spn):-1:1
        node = spn[i]
        # println(i, " ", node)
        if issum(node)
            # process sum node
            if length(node.children) > 4
                @warn "node $i indegree is too large: $(length(node.children)). It is highly recommend to split nodes before running this. Try running binarize! first."
            end
            var = VariableNode(2)
            variables["Y"*string(i)] = var
            factor = FactorNode(
                Array{Float64}(undef, Tuple(2 for _ = 1:(length(node.children)+1)) ), # factor
                VariableNode[var ; map(j -> variables["Y"*string(j)], node.children)] # neighbors
                )
            # P(Y=2|parents=z) = sum( node.weights[i] | z[i] = 2 )
            for z in CartesianIndices(Tuple(2 for _ = 1:length(node.children)))
                prob = min(1.0,mapreduce(j -> node.weights[j]*(z[j]-1), +, 1:length(z)))
                factor.factor[2,z] = log(prob)
                factor.factor[1,z] = log(1 - prob)
            end
            factors[string(i)] = factor
            # display(factor.neighbors)
        elseif isprod(node)
            # process product node
            if length(node.children) > 4
                @warn "node $i indegree is too large: $(length(node.children)). It is highly recommend to split nodes before running this. Try running binarize! first."
            end
            var = VariableNode(2)
            variables["Y"*string(i)] = var
            factor = FactorNode(
                zeros(Float64, Tuple(2 for _ = 1:(length(node.children)+1)) ), # factor
                VariableNode[var ; map(j -> variables["Y"*string(j)], node.children)] # neighbors
                )
            # P(Y=2|parents) = { 1 if all parents = 1, 0 otherwise }
            factor.factor[2:2:(end-1)] .= -Inf
            factor.factor[end-1] = -Inf
            factors[string(i)] = factor
        elseif isa(node, IndicatorFunction)
            # process indicator leaf node
            var = VariableNode(2)
            variables["Y"*string(i)] = var
            parent = variables["X"*string(node.scope)]
            factor = FactorNode(zeros(Float64, 2, vdims[node.scope]), [var, parent])
            # P(Y=2|parent) = { 1 if parent = node.value, 0 otherwise }
            factor.factor[2:2:end] .= -Inf
            factor.factor[1,convert(Int,node.value)] = -Inf
            factor.factor[2,convert(Int,node.value)] = 0.0
            factors[string(i)] = factor
        elseif isa(node, CategoricalDistribution)
            # process categorical leaf node
            var = VariableNode(2)
            variables["Y"*string(i)] = var
            parent = variables["X"*string(node.scope)]
            factor = FactorNode(Array{Float64}(undef, 2, vdims[node.scope]), [var, parent])
            # P(Y=2|parent=j) = node.values[j]
            for j = 1:vdims[node.scope]
                factor.factor[2,j] = log(node.values[j])
                factor.factor[1,j] = log(1.0 - node.values[j])
            end
            factors[string(i)] = factor
        else
            @error "Unsupported node type: $(typeof(node))"
        end
    end
    # Create and return factor graph
    FactorGraph(variables, factors)
end

"""
    beliefpropagation!(x, spn, query; kwargs)

Runs mixed belief propagation algorithm with given `spn`, `query` variables and evidence `x`.

### Keyword Arguments
- `maxiterations`: positive integer value specifying maximum number of iterations [default: `10`]
- `verbose`: whether to output information during inference run [default: `true`]
- `normalize`: whether to normalize messages [defulat: `true`]
- `lowerbound`: whether to use the probability of `x` as initial lower bound [default: `false`]
- `warmstart`: whether to use the configuration of query variables in `x` to initilize messages [default: `false`]
- `earlystop`: early stop algorithm if residual is below that threshold [default: `0.001`]
- `rndminit`: whether to use random initialization of messages (rather than constant initialization) [default: `false`]
"""
function beliefpropagation!(x::AbstractArray{<:Real}, spn::SumProductNetwork, query;
        maxiterations = 10,
        verbose = true,
        normalize = true,
        lowerbound = false,
        earlystop = 0.001,
        rndminit = false,
        warmstart = false )
    # Translate SPN into distribution-equivalent Factor Graph
    if verbose
        fg, t, bytes, gctime, mallocs = @timed spn2bn(spn)
        @info "Generated factor graph ($(round(bytes/1048576,digits=1))MiB) in $t secs."
    else
        fg = spn2bn(spn)
    end
    # consistency checks
    @assert length(fg.variables) == (length(spn) + length(scope(spn)))
    @assert length(fg.factors) == length(spn)

    # Initialize belief propagation
    bp = HybridBeliefPropagation(fg; rndinit = rndminit) # rndinit = true generates random initial message; = true sets all to 1.
    bp.normalize = normalize # normalize messages (sum = 1)?
    # Set evidence and query
    setevidence!(bp, "Y1", 2)
    for v in findall(isfinite,x) # filter non-marginalized variables
        if v ∉ query # filter out query variables
            setevidence!(bp, "X$v", x[v]) # remaining are evidence
        end
    end
    for v in query
        setmapvar!(bp, "X$v")
    end
    # run Inference
    if verbose
        println("Running belief propagation for $maxiterations iterations")
    end
    cpad(s, n::Integer, p=" ") = rpad(lpad(s,div(n+length(s),2),p),n,p) # for printing centered
    columns = ["it", "time (s)", "residual", "value", "best?"]
    widths = [4, 10, 10, 24, 7]
    bordercolor = :light_cyan
    headercolor = :cyan
    fieldcolor = :normal
    if verbose
        printstyled('╭', join(map(w -> repeat('─',w), widths), '┬'), '╮', '\n' ;color = bordercolor)
        for (col,w) in zip(columns,widths)
            printstyled("│"; color = bordercolor)
            printstyled(cpad(col, w); color = headercolor)
        end
        printstyled("│\n"; color = bordercolor)
        printstyled('├', join(map(w -> repeat('─',w), widths), '┼'), '┤', '\n' ;color = bordercolor)
    end
    # for computing logpdf of solution
    values = Vector{Float64}(undef,length(spn))
    # value of best incumbent solution
    best = -Inf
    if lowerbound
        # best = spn(x)
        best = logpdf!(values,spn,x)
    end
    if warmstart
        # bias messages towards initial solution
        for v in query
            var = fg.variables["X$v"]
            for ne in var.neighbors
                bp.messages[(var,ne)] .+= 1
            end
        end
    end
    y = copy(x) # incumbent solution
    start = time_ns()
    for it=1:maxiterations
        # if warmstart
        #     if it == 1
        #         for v in query
        #             setevidence!(bp, "X$v", x[v])
        #         end
        #     elseif it == 2
        #         for v in query
        #             rmevidence!(bp, "X$v")
        #         end
        #     end
        # end
        # downward propagation
        res = 0.0
        for i = length(spn):-1:1
            node = spn[i]
            # compute incoming messages from children
            if isleaf(node)
                factor = fg.factors[string(i)]
                update!(bp, fg.variables["X$(node.scope)"], factor)
                res = max(res, update!(bp, factor, fg.variables["Y$i"]))
            else
                factor = fg.factors[string(i)]
                for j in node.children
                    update!(bp, fg.variables["Y$j"], factor)
                end
                res = max(res, update!(bp, factor, fg.variables["Y$i"]))
            end
        end
        # upward propagation
        for i = 1:length(spn)
            node = spn[i]
            # compute outgoing messages to children
            if isleaf(node)
                factor = fg.factors[string(i)]
                update!(bp, fg.variables["Y$i"], factor)
                res = max(res, update!(bp, factor, fg.variables["X$(node.scope)"]))
            else
                factor = fg.factors[string(i)]
                update!(bp, fg.variables["Y$i"], factor)
                for j in node.children
                    res = max(res, update!(bp, factor, fg.variables["Y$j"]))
                end
            end
        end
        for i in query
            y[i] = decode(bp, "X$i")
        end
        # prob = spn(y)
        prob = logpdf!(values,spn,y)
        if verbose
            etime = (time_ns()-start)/1e9;
            for (col,w) in zip([it, round(etime, digits=2), round(res,digits=2), prob, prob >= best ? "*" : " "],widths)
                printstyled("│"; color = bordercolor)
                printstyled(lpad(col, w-1), ' '; color = fieldcolor)
            end
            printstyled("│\n"; color = bordercolor)
        end
        if prob > best
            best = prob
            x .= y
        end
        if res < earlystop break end # early stop
    end
    if verbose
        printstyled('╰', join(map(w -> repeat('─',w), widths), '┴'), '╯', '\n' ;color = bordercolor)
    end
    exp(best)
end

"""
    treebeliefpropagation!(x::AbstractArray{<:Real}, spn::SumProductNetwork, query; kwargs)

Speciliazed hybrid belief propagation for tree-shaped SPNs where innner nodes have two children, leaves are indicators.
Warning: this implementation is not numerically stable. It require some care when running on large networks.
"""
function treebeliefpropagation!(x::AbstractArray{<:Real}, spn::SumProductNetwork, query;
    maxiterations = 10,
    verbose = true,
    normalize = true,
    lowerbound = false,
    earlystop = 0.001,
    rndminit = false,
    warmstart = false  )
    # For conveniency
    leaves = filter(i -> isleaf(spn[i]), 1:length(spn))
    gates = filter(i -> !isleaf(spn[i]), 1:length(spn))
    nleaves, ngates = length(leaves), length(gates)
    evidence = setdiff(Set(Base.filter(i -> isfinite(x[i]), 1:length(x))), query)
    # Compute log( exp(x1)/(exp(x1)+exp(x2)) )
    lognorm(x1, x2) = x1 >= x2 ? -log1p(exp(x2 - x1)) : x1 - x2 - log1p(exp(x1 - x2))
    # Messages
    π = Dict{Int,Float64}()  # upstream messages from node n
    π_v = Dict{Tuple{Int, Int}, Array{Float64}}()  # upstream messages from manifest variables
    π_i = Dict{Tuple{Int, Int}, Float64}()  # upstream messages from indicator nodes
    λ = Dict{Int, Float64}()  # downstream messages
    λ_i = Dict{Tuple{Int, Int}, Float64}()  # downstream messages to indicator nodes
    λ_v = Dict{Tuple{Int, Int}, Array{Float64}}()  # downstream messages to manifest variables
    out = Dict{Int,Int}()  # parent of nodes
    out_i = Dict{Int,Array{Int}}() # parents of indicator nodes
    out_v = Dict{Int,Array{Int}}() # parents of manifest variables
    bel = Dict{Int, Array{Float64}}()  # beliefs for manifest variables
    ### Initialization ################################################################
    vdims = vardims(spn)
    for (var, ncat) in vdims
        # Initialize beliefs
        # evidence or warm start
        if var in evidence || (warmstart && (var in query) && !isnan(x[var]))
            bel[var] = zeros(Float64, ncat)
            bel[var][Int(x[var])] = 1.0
        else
            bel[var] = ones(Float64, ncat)/ncat
        end
        out_v[var] = []
    end
    for n in leaves
        node = spn[n]
        @assert isa(node, IndicatorFunction)
        ncat = vdims[node.scope]
        λ_v[(n, node.scope)] = ones(Float64, ncat)/ncat
        π_v[(node.scope, n)] = ones(Float64, ncat)/ncat
        push!(out_v[node.scope], n)
        out_i[n] = []
    end
    out[1] = -1 # root/output node has no parents
    for n in gates
        node = spn[n]
        @assert length(node.children) == 2
        for ch in node.children
            if isa(spn[ch],IndicatorFunction)
                push!(out_i[ch], n)
                π_i[(ch,n)] = 0.5
                λ_i[(n,ch)] = 0.5
            else
                λ[ch] = 0.5
                π[ch] = 0.5
                out[ch] = n
            end
        end
    end
    λ[1] = 1.0 # Root/output node receives "evidence"
    ### Run message passing ###########################################################
    best = -Inf
    y = copy(x) # candidate solution
    for iteration = 1:maxiterations
        ### upstream propagation ##########################################################
        ## compute messages π[N -> P] -- assumes incoming messages are normalized
        # compute messages from variables to indicator nodes
        for (edge, message) in π_v
            var, n = edge
            if var in evidence
                # variable is evidence
                fill!(message, 0.0)
                message[Int(x[var])] = 1.0
            else
                Z = 0.0
                for j = 1:vdims[var]
                    message[j] = prod(λ_v[(n, var)][j] for k in out_v[var] if k ≠ n)
                    Z += message[j]
                end
                @assert Z > 0
                # normalize messages (seems to be important only for marginalized variables)
                message ./= Z
            end
            # print("root",v,node.assignment,values,Z)
        end
        # compute messages from indicator nodes
        for n in leaves
            node = spn[n]
            vid = node.scope
            for pa in out_i[n]
                @assert length(out_i[n]) > 0
                if x[vid] == NaN && (var ∉ query) # marginalized
                    π1 = π_v[(vid, n)][Int(node.value)]
                    for opa in out_i[n]
                        if opa ≠ pa
                            π1 *= λ_i[(opa, n)]
                        end
                    end
                    π0 = (1 - π_v[(vid, n)][Int(node.value)])
                    for opa in out_i[n]
                        if opa ≠ pa
                            π0 *= 1.0 - λ_i[(opa, n)]
                        end
                    end
                else
                    maxbel = maximum(bel[vid])
                    # TODO: handle evidence more efficiently
                    if bel[vid][Int(node.value)] == maxbel
                        π1 = π_v[(vid, n)][Int(node.value)]
                        for opa in out_i[n]
                            if opa ≠ pa
                                π1 *= λ_i[(opa, n)]
                            end
                        end
                    else
                        π1 = 0.0
                    end
                    π0 = 0.0
                    for j = 1:vdims[vid]
                        if j ≠ node.value && bel[vid][j] == maxbel
                            π0 += π_v[(vid, n)][j]
                        end
                    end
                    for opa in out_i[n]
                        if opa ≠ pa
                            π0 *= 1.0 - λ_i[(opa, n)]
                        end
                    end
                end
                @assert π0 + π1 > 0
                π_i[(n,pa)] = π1 / (π1 + π0)  # normalization
            end
        end
        # compute remaining messages
        for n in reverse(gates)
            node = spn[n]
            if issum(node)
                π[n] = 0.0
                for j = 1:2
                    if isleaf(spn[node.children[j]])
                        π[n] += node.weights[j] * π_i[(node.children[j],n)]
                    else
                        π[n] += node.weights[j] * π[node.children[j]]
                    end
                end
            else  # Product
                π[n] = 1.0
                for j = 1:2
                    if isleaf(spn[node.children[j]])
                        π[n] *= π_i[(node.children[j],n)]
                    else
                        π[n] *= π[node.children[j]]
                    end
                end
            end
        end
        ### Upward propagation ##########################################################
        for n in gates
            # compute messages λ[N -> C] -- assumes incoming messages are normalized
            node = spn[n]
            ch1, ch2 = node.children
            λp = λ[n]
            if isa(spn[ch2],IndicatorFunction)
                π2 =  π_i[(ch2,n)]
            else
                π2 =  π[ch2]
            end
            # left child
            if issum(node)
                λ1 = (1 - λp) * node.weights[2] * (1 - π2) + λp * (node.weights[1] * (1 - π2) + π2)
                λ0 = (1 - λp) * (1 - π2 + node.weights[1] * π2) + λp * node.weights[2] * π2
            else # product node
                λ1 = (1.0 - λp) * (1.0 - π2) + λp * π2
                λ0 = 1 - λp
            end
            @assert λ0 + λ1 > 0
            if isa(spn[ch1],IndicatorFunction)
                λ_i[(n,ch1)] = λ1 / (λ0 + λ1)
            else
                λ[ch1] = λ1 / (λ0 + λ1)
            end
            if isa(spn[ch1],IndicatorFunction)
                π1 =  π_i[(ch1,n)]
            else
                π1 =  π[ch1]
            end
            # right child
            if issum(node)
                λ1 = (1 - λp) * node.weights[1] * (1 - π1) + λp * (node.weights[2] * (1 - π1) + π1)
                λ0 = (1 - λp) * (1 - π1 + node.weights[2] * π1) + λp * node.weights[1] * π1
            else # product
                λ1 = (1.0 - λp) * (1.0 - π1) + λp * π1
                λ0 = 1 - λp
            end
            @assert λ0 + λ1 > 0
            if isa(spn[ch1],IndicatorFunction)
                λ_i[(n,ch2)] = λ1 / (λ0 + λ1)
            else
                λ[ch2] = λ1 / (λ0 + λ1)
            end
        end
        # compute messages from indicator nodes to root nodes
        for (p, message) in λ_v
            n, v = p
            λ1 = prod(λ_i[(pa, n)] for pa in out_i[n])
            λ0 = prod(1.0 - λ_i[(pa, n)] for pa in out_i[n])
            Z = λ1 + λ0 * (vdims[v] - 1) # normalization constant
            @assert Z > 0
            message .= (λ0 / Z)
            message[Int(spn[n].value)] = λ1 / Z
        end
        ### Compute beliefs ##########################################################
        for var in keys(vdims)
            if var in evidence
                fill!(bel[var], 0.0)
                bel[var][Int(x[var])] = 1.0
            else
                Z = 0.0
                for j = 1:vdims[var]
                    b = prod(λ_v[(n, var)][j] for n in out_v[var])
                    Z += b
                    bel[var][j] = b
                end
                @assert Z > 0
                bel[var] ./= Z
                if var in query
                    # decode solution
                    m, y[var] = findmax(bel[var])
                    @assert isfinite(m) && m > 0
                end
            end
        end
        # compute value at root/output
        prob = π[1]
        value = spn(y)
        if verbose
            printstyled("[$iteration/$maxiterations] "; color = :blue)
            printstyled("Root Prob: "; color = :light_cyan)
            print(prob)
            printstyled("\t Value: "; color = :light_cyan)
            println(value)
        end
        if value > best
            # found improving solution
            best = value
            x .= y
        end
    end
    best
end


text:
## ODEProblem
"""
    ODEProblem(env::AbstractEnv, f, x0, tspan; kwargs...)

An API for DifferentialEquations.ODEProblem(f, x0, tspan; kwargs...).
"""
function DifferentialEquations.ODEProblem(env::AbstractEnv, f, x0, tspan; kwargs...)
    warn_is_registered(env)
    _x0 = raw(env, x0)
    _f(_x, p, t) = raw(env, f(readable(env, _x), p, t))
    ODEProblem(_f, _x0, tspan; kwargs...)
end
function DifferentialEquations.ODEProblem(env::AbstractEnv, f, x0, tspan, p; kwargs...)
    warn_is_registered(env)
    _x0 = raw(env, x0)
    _f(_x, p, t) = raw(env, f(readable(env, _x), p, t))
    ODEProblem(_f, _x0, tspan, p; kwargs...)
end

## initial condition
# automatic completion of initial condition
"""
    initial_condition(env::AbstractEnv)

Automatic completion of readable initial_condition if methods of its sub-environments are given.
"""
function initial_condition(env::AbstractEnv)
    env_names = names(env)
    values = env_names |> Map(name -> initial_condition(getfield(env, name)))
    return (; zip(env_names, values)...)  # NamedTuple
end

## register env
"""
    @reg_env(env, x0)

Register given environment.
`x0` is an example of readable initial_condition for `env` so that `NestedEnvironments` infers the information of `env`.

# Notes
- Must be used in the global scope.
- `x0` can be any dummy NamedTuple-collection.
"""
macro reg_env(env, x0)
    ex = quote
        if typeof($(env)) <: AbstractEnv
            local env_index_nt, env_size_nt = NestedEnvironments._preprocess($(env), $(x0))
            # raw & readable
            NestedEnvironments.readable(env::typeof($(env)), _x) = NestedEnvironments._readable(_x, env_index_nt, env_size_nt)
            NestedEnvironments.raw(env::typeof($(env)), x) = NestedEnvironments._raw($(env), x)
            # register env
            push!(NestedEnvironments.__REGISTERED_ENVS, $(env), $(x0))
        else
            error("Invalid environment")
        end
    end
    esc(ex)
end
# aux
function Base.push!(__REGISTERED_ENVS::RegisteredEnvs, env::AbstractEnv, x0)
    num_of_already_reg_envs = __REGISTERED_ENVS.__envs |> Filter(__env -> typeof(__env) == typeof(env)) |> collect |> length
    if num_of_already_reg_envs == 0
        push!(__REGISTERED_ENVS.__envs, env)
        push!(__REGISTERED_ENVS.__xs, x0)
        println("$(typeof(env)): registered")
    else
        println("$(typeof(env)): overwrite existing registered env")
    end
end

## macros
# macro for transformation from readable to raw
"""
    @raw(env, x)

A macro that transforms a readable (NamedTuple) value to raw (flattened array) value.
"""
macro raw(env, x)
    ex = quote
        if typeof($(env)) <: AbstractEnv
            _x = NestedEnvironments.raw($(env), $(x))
        else
            error("Invalid environment")
        end
    end
    esc(ex)
end
# auto-completion
"""
    @raw(x)

A macro that transforms a readable (NamedTuple) value to raw (flattened array) value by corresponding an appropriate environment from the registered environments.
"""
macro raw(x)
    ex = quote
        local env_x_cands = zip(NestedEnvironments.__REGISTERED_ENVS.__envs, NestedEnvironments.__REGISTERED_ENVS.__xs) |> Transducers.Filter(env_x -> size(env_x[2]) == size($(x))) |> collect
        if length(env_x_cands) == 0
            error("There is no matched registered envrionment")
        elseif length(env_x_cands) > 1
            error("It is ambiguous; too many matched registered environments")
        else
            @raw(env_x_cands[1][1], $(x))
        end
    end
    esc(ex)
end

# macro for transformation from raw to readable
"""
    @readable(env, _x)

A macro that transforms a raw (flattened array) value to readable (NamedTuple) value.
"""
macro readable(env, _x)
    ex = quote
        if typeof($(env)) <: AbstractEnv
            x = NestedEnvironments.readable($(env), $(_x))
        else
            error("Invalid environment")
        end
    end
    esc(ex)
end
# auto-completion
"""
    @readable(x)

A macro that transforms a raw (flattened array) value to readable (NamedTuple) value by corresponding an appropriate environment from the registered environments.
"""
macro readable(_x)
    ex = quote
        local env_x_cands = zip(NestedEnvironments.__REGISTERED_ENVS.__envs, NestedEnvironments.__REGISTERED_ENVS.__xs) |> Transducers.Filter(env_x -> NestedEnvironments.flatten_length(env_x[2]) == NestedEnvironments.flatten_length($(_x))) |> collect
        if length(env_x_cands) == 0
            error("There is no matched registered envrionment")
        elseif length(env_x_cands) > 1
            error("It is ambiguous; too many matched registered environments")
        else
            @readable(env_x_cands[1][1], $(_x))
        end
    end
    esc(ex)
end


text:
export @dateformat_str, DateFormat

import Base.Dates: year, month, day, hour, minute, second, millisecond,
                   monthabbr, monthname, dayabbr, dayname, Year, Month,
                   Day, Hour, Minute, Second, Millisecond, TimeType, Date,
                   DateTime

include("date-locale.jl")
include("date-io.jl")
include("date-parse.jl")


text:
# Based on this guide: https://ase.tufts.edu/chemistry/lin/images/FortranMD_TeachersGuide.pdf
# Uses DFTK in place of LJ for the production stage as a proof-of-concept ab initio MD simulation
# Note that the choice of parameters is for demonstration purposes, and the results are non-physical

using Atomistic
using DFTK
using GLMakie
using InteratomicPotentials
using Molly

setup_threading(n_blas=4)

N = 8
box_size = 1.5u"nm" # this number was chosen arbitrarily
reference_temp = 94.4u"K"
coupling_factor = 10 # this number was chosen arbitrarily
Δt = 1e-2u"ps"

initial_system = attach_psp(generate_atoms_in_cubic_cell(N, :Ar, box_size, reference_temp); functional="lda")

eq_steps = 20000
eq_thermostat = Molly.AndersenThermostat(reference_temp, Δt * coupling_factor)
eq_simulator = MollySimulator(Δt, eq_steps, coupling=eq_thermostat)
potential = InteratomicPotentials.LennardJones(1.657e-21u"J", 0.34u"nm", 0.765u"nm", [:Ar])

eq_result = @time simulate(initial_system, eq_simulator, potential)

display(plot_temperature(eq_result, eq_simulator.steps ÷ 200))
display(plot_energy(eq_result, eq_simulator.steps ÷ 200))
display(plot_rdf(eq_result, potential.σ))

ab_initio_steps = 200
ab_initio_simulator = MollySimulator(Δt, ab_initio_steps, t₀=get_time(eq_result))
# very non-physical but fast for demonstration purposes
dftk_potential = DFTKPotential(5u"hartree", [1, 1, 1]; functionals=[:lda_x, :lda_c_pw], scf_kwargs=Dict(:damping => 0.7))

ab_initio_result = @time simulate(get_system(eq_result), ab_initio_simulator, dftk_potential)

# Plotting on separate plots because the timespan is so much smaller than in the first phase
display(plot_temperature(ab_initio_result, 1))
display(plot_energy(ab_initio_result, 1))
display(plot_rdf(ab_initio_result, potential.σ))

visualize(ab_initio_result, "artifacts/argon_ab_initio_molly.mp4")

;


text:
@testset "ModelLike" begin
    model = DummyModelWithAdd()
    MOIT.failcopytestc(model)
    MOIT.failcopytestia(model)
    MOIT.failcopytestva(model)
    MOIT.failcopytestca(model)
end


text:
# This file defines the EmpiricalOrthogonalFunction types

# abstract EOF type
"""
Abstract type to build subtypes from
"""
abstract type AbstractEmpiricalOrthogonalFunction end

# eof struct for 2d data (space, time)
"""
    EmpiricalOrthogonalFunction(dataset; center=true, ddof=1)

Create an Empirical Orthogonal Function object.
The EOF solution is computed at initialization time. Method
calls are used to retrieve or update computed quantities.
"""
mutable struct EmpiricalOrthogonalFunction <: AbstractEmpiricalOrthogonalFunction
    ""
    dataset::AbstractArray{<:Any,2}

    ""
    center::AbstractArray{<:Any,2}

    valididx::AbstractArray{Int64,1}

    eofs::AbstractArray{<:Any,2}

    eigenvals::AbstractArray{<:Any,1}

    pcs::AbstractArray{<:Any,2}

    EmpiricalOrthogonalFunction(
        dataset::AbstractArray{<:Any,2},
        center::AbstractArray{<:Any,2},
        valididx::AbstractArray{Int64,1},
        eofs::AbstractArray{<:Any,2},
        eigenvals::AbstractArray{<:Any,1},
        pcs::AbstractArray{<:Any,2},
    ) = new(dataset, center, valididx, eofs, eigenvals, pcs)

    """
    """
    function EmpiricalOrthogonalFunction(
        dataset::AbstractArray{<:Any,2};
        center::Bool = true,
        ddof::Int = 1,
    )

        # Store information about the shape/size of the input data.
        shape = size(dataset)
        timelen = shape[1]

        if shape[1] > shape[2]
            @error(
                "time dimension has more records than the space dimensions and an EOF cannot be calculated, time = $(shape[1]), space= $(shape[2])"
            )
            return
        end

        # Remove the time mean of the input data unless explicitly told
        # not to by the "center" argument.
        if center
            spatial_mean = Statistics.mean(dataset, dims = 1)
            centered = @. dataset - spatial_mean
        else
            spatial_mean = zeros(shape[1])
            centered = dataset
        end

        # Find the indices of values that are not missing in one row. All the
        # rows will have missing values in the same places provided the
        # array was centered. If it wasn't then it is possible that some
        # missing values will be missed and the singular value decomposition
        # will produce not a number for everything.
        isnotmissing = findall(x -> x !== missing, spatial_mean[1, :])

        # Compute the singular value decomposition of the design matrix.
        U, S, V = svd(centered[:, isnotmissing])

        # Singular values are the square-root of the eigenvalues of the
        # covariance matrix. Construct the eigenvalues appropriately and
        # normalize by N-ddof where N is the number of observations. This
        # corresponds to the eigenvalues of the normalized covariance matrix.
        norm = timelen - 1
        eigenvals = @. S * S / norm

        eofs = Array{Union{Missing,Float64}}(missing, shape[2], shape[1])
        eofs[isnotmissing, :] = V[:, :]

        # Remove the scaling on the principal component time-series that is
        # implicitily introduced by using SVD instead of eigen-decomposition.
        # The PCs may be re-scaled later if required.
        pcs = similar(U)
        for i = 1:length(S)
            pcs[:, i] = U[:, i] * S[i]
        end

        EmpiricalOrthogonalFunction(
            dataset,
            spatial_mean,
            isnotmissing,
            eofs,
            eigenvals,
            pcs,
        )
    end

    """

    """
    function EmpiricalOrthogonalFunction(
        dataset::AbstractArray{<:Any,3};
        timedim::Int = 3,
        kwargs...,
    )
        geodims = [x for x in 1:3 if x != timedim]
        shape = size(dataset)
        spatial_shape = shape[geodims]
        time_shape = shape[timedim]

        if timedim != 1
            dataset = permutedims(dataset, (timedim, geodims...))
        end

        flattened = reshape(dataset, (time_shape, prod(spatial_shape)))

        EmpiricalOrthogonalFunction(flattened; kwargs...)
    end

end


text:
import Base.GC: gc

# dynamic memory pool allocator
#
# this allocator sits between CuArray constructors
# and the actual memory allocation in CUDAdrv.Mem
#
# the core design is a pretty simple:
# - bin allocations into multiple pools according to their size (see `poolidx`)
# - when requested memory, check the pool for unused memory, or allocate dynamically
# - conversely, when released memory, put it in the appropriate pool for future use
#
# to avoid memory hogging and/or trashing the Julia GC:
# - keep track of used and available memory, in order to determine the usage of each pool
# - keep track of each pool's usage, as well as a window of previous usages
# - regularly release memory from underused pools (see `reclaim(false)`)
#
# possible improvements:
# - context management: either switch contexts when performing memory operations,
#                       or just use unified memory for all allocations.
# - per-device pools

const pool_lock = ReentrantLock()

using TimerOutputs
const to = TimerOutput()

## infrastructure

const pools_used = Vector{Set{Mem.Buffer}}()
const pools_avail = Vector{Vector{Mem.Buffer}}()

poolidx(n) = ceil(Int, log2(n))+1
poolsize(idx) = 2^(idx-1)

function create_pools(idx)
  if length(pool_usage) >= idx
    # fast-path without taking a lock
    return
  end

  lock(pool_lock) do
    while length(pool_usage) < idx
      push!(pool_usage, 1)
      push!(pool_history, initial_usage)
      push!(pools_used, Set{Mem.Buffer}())
      push!(pools_avail, Vector{Mem.Buffer}())
    end
  end
end

## management

const USAGE_WINDOW = 5
const initial_usage = Tuple(1 for _ in 1:USAGE_WINDOW)

const pool_usage = Vector{Float64}()
const pool_history = Vector{NTuple{USAGE_WINDOW,Float64}}()

# min and max time between successive background task iterations.
# when the pool usages don't change, scan less regularly.
#
# together with USAGE_WINDOW, this determines how long it takes for objects to get reclaimed
const MIN_DELAY = 1.0
const MAX_DELAY = 5.0

# debug stats
mutable struct PoolStats
  # allocation requests
  req_nalloc::Int
  req_nfree::Int
  ## in bytes
  req_alloc::Int
  user_free::Int

  # actual allocations
  actual_nalloc::Int
  actual_nfree::Int
  ## in bytes
  actual_alloc::Int
  actual_free::Int

  cuda_time::Float64
  total_time::Float64
end
const stats = PoolStats(0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
Base.copy(stats::PoolStats) =
  PoolStats((getfield(stats, field) for field in fieldnames(PoolStats))...)

# allocation traces
const alloc_sites = Dict{Mem.Buffer, Tuple{Int, Vector{Union{Base.InterpreterIP,Ptr{Cvoid}}}}}()
const tracing = parse(Bool, get(ENV, "CUARRAYS_TRACE_POOL", "false"))

function __init_memory__()
  create_pools(30) # up to 512 MiB

  managed = parse(Bool, get(ENV, "CUARRAYS_MANAGED_POOL", "true"))
  if managed
    delay = MIN_DELAY
    @async begin
      while true
        @timeit to "background task" lock(pool_lock) do
          if scan()
            delay = MIN_DELAY
          else
            delay = min(delay*2, MAX_DELAY)
          end

          reclaim()
        end

        sleep(delay)
      end
    end
  end

  verbose = haskey(ENV, "CUARRAYS_MANAGED_POOL")
  if verbose
    atexit(()->begin
      Core.println("""
        Pool statistics (managed: $(managed ? "yes" : "no")):
         - requested alloc/free: $(stats.req_nalloc)/$(stats.req_nfree) ($(Base.format_bytes(stats.req_nalloc))/$(Base.format_bytes(stats.req_free)))
         - actual alloc/free: $(stats.actual_nalloc)/$(stats.actual_nfree) ($(Base.format_bytes(stats.actual_alloc))/$(Base.format_bytes(stats.actual_free)))""")
    end)
  end
end

# scan every pool and manage the usage history
#
# returns a boolean indicating whether any pool is active (this can be a false negative)
function scan()
  gc(false) # quick, incremental collection

  active = false

  @inbounds for pid in 1:length(pool_history)
    nused = length(pools_used[pid])
    navail = length(pools_avail[pid])
    history = pool_history[pid]

    if nused+navail > 0
      usage = pool_usage[pid]
      current_usage = nused / (nused + navail)

      # shift the history window with the recorded usage
      history = pool_history[pid]
      pool_history[pid] = (Base.tail(pool_history[pid])..., usage)

      # reset the usage with the current one
      pool_usage[pid] = current_usage

      if usage != current_usage
        active = true
      end
    else
      pool_usage[pid] = 1
      pool_history[pid] = initial_usage
    end
  end

  active
end

# reclaim unused buffers
function reclaim(full::Bool=false, target_bytes::Int=typemax(Int))
  stats.total_time += Base.@elapsed begin
    # find inactive buffers
    @timeit to "scan" begin
      pools_inactive = Vector{Int}(undef, length(pools_avail)) # pid => buffers that can be freed
      if full
        # consider all currently unused buffers
        for (pid, avail) in enumerate(pools_avail)
          pools_inactive[pid] = length(avail)
        end
      else
        # only consider inactive buffers
        @inbounds for pid in 1:length(pool_usage)
          nused = length(pools_used[pid])
          navail = length(pools_avail[pid])
          recent_usage = (pool_history[pid]..., pool_usage[pid])

          if navail > 0
            # reclaim as much as the usage allows
            reclaimable = floor(Int, (1-maximum(recent_usage))*(nused+navail))
            pools_inactive[pid] = reclaimable
          else
            pools_inactive[pid] = 0
          end
        end
      end
    end

    # reclaim buffers (in reverse, to discard largest buffers first)
    @timeit to "reclaim" begin
      for pid in reverse(eachindex(pools_inactive))
        bytes = poolsize(pid)
        avail = pools_avail[pid]

        bufcount = pools_inactive[pid]
        @assert bufcount <= length(avail)
        for i in 1:bufcount
          buf = pop!(avail)

          stats.actual_nfree += 1
          stats.cuda_time += Base.@elapsed Mem.free(buf)
          stats.actual_free += bytes

          target_bytes -= bytes
          target_bytes <= 0 && return true
        end
      end
    end
  end

  return false
end

## allocator state machine

function try_cuda_alloc(bytes)
  buf = nothing
  try
    stats.cuda_time += Base.@elapsed begin
      buf = Mem.alloc(bytes)
    end
    stats.actual_nalloc += 1
    stats.actual_alloc += bytes
  catch ex
    ex == CUDAdrv.ERROR_OUT_OF_MEMORY || rethrow()
  end

  return buf
end

function try_alloc(bytes, pid=-1)
  # NOTE: checking the pool is really fast, and not included in the timings
  if pid != -1 && !isempty(pools_avail[pid])
    return pop!(pools_avail[pid])
  end

  @timeit to "1 try alloc" begin
    let buf = try_cuda_alloc(bytes)
      buf !== nothing && return buf
    end
  end

  @timeit to "2 gc(false)" begin
    gc(false) # incremental collection
  end

  if pid != -1 && !isempty(pools_avail[pid])
    return pop!(pools_avail[pid])
  end

  # TODO: we could return a larger allocation here, but that increases memory pressure and
  #       would require proper block splitting + compaction to be any efficient.

  @timeit to "3 reclaim unused" begin
    reclaim(true, bytes)
  end

  @timeit to "4 try alloc" begin
    let buf = try_cuda_alloc(bytes)
      buf !== nothing && return buf
    end
  end

  if tracing
    alloc_sites_old = copy(alloc_sites)
  end
  @timeit to "5 gc(true)" begin
    gc(true) # full collection
  end
  if tracing
    # report on buffers that we collected -- these could benefit from early freeing
    for (buf, info) in sort(collect(alloc_sites_old), by=x->x[2][1])
      bytes, bt = info
      if !haskey(alloc_sites, buf)
        st = stacktrace(bt, false)
        Core.print(Core.stderr, "WARNING: force-collected a GPU allocation of $(Base.format_bytes(bytes))")
        Base.show_backtrace(Core.stderr, st)
        Core.println(Core.stderr)
      end
    end
  end

  if pid != -1 && !isempty(pools_avail[pid])
    return pop!(pools_avail[pid])
  end

  @timeit to "6 reclaim unused" begin
    reclaim(true, bytes)
  end

  @timeit to "7 try alloc" begin
    let buf = try_cuda_alloc(bytes)
      buf !== nothing && return buf
    end
  end

  @timeit to "8 reclaim everything" begin
    reclaim(true)
  end

  @timeit to "9 try alloc" begin
    let buf = try_cuda_alloc(bytes)
      buf !== nothing && return buf
    end
  end

  if tracing
    for buf in keys(alloc_sites)
      bytes, bt = alloc_sites[buf]
      st = stacktrace(bt, false)
      Core.print(Core.stderr, "WARNING: outstanding a GPU allocation of $(Base.format_bytes(bytes))")
      Base.show_backtrace(Core.stderr, st)
      Core.println(Core.stderr)
    end
  end

  throw(OutOfMemoryError())
end

## interface

const MAX_POOL = 100*1024^2 # 100 MiB

function alloc(bytes)
  # 0-byte allocations shouldn't hit the pool
  bytes == 0 && return Mem.alloc(0)

  stats.req_nalloc += 1
  stats.req_alloc += bytes
  stats.total_time += Base.@elapsed begin
    # only manage small allocations in the pool
    if bytes <= MAX_POOL
      pid = poolidx(bytes)
      create_pools(pid)
      alloc_bytes = poolsize(pid)

      @inbounds used = pools_used[pid]
      @inbounds avail = pools_avail[pid]

      lock(pool_lock) do
        buf = @timeit to "pooled alloc" try_alloc(alloc_bytes, pid)

        # mark the buffer as used
        push!(used, buf)

        # update pool usage
        current_usage = length(used) / (length(avail) + length(used))
        pool_usage[pid] = max(pool_usage[pid], current_usage)
      end
    else
      buf = @timeit to "large alloc" try_alloc(bytes)
    end
  end

  if tracing
    alloc_sites[buf] = (bytes, backtrace())
  end

  buf
end

function dealloc(buf, bytes)
  # 0-byte allocations shouldn't hit the pool
  bytes == 0 && return

  stats.req_nfree += 1
  stats.user_free += bytes
  stats.total_time += Base.@elapsed begin
    # was this a pooled buffer?
    if bytes <= MAX_POOL
      pid = poolidx(bytes)
      @assert pid <= length(pools_used)

      @inbounds used = pools_used[pid]
      @inbounds avail = pools_avail[pid]

      lock(pool_lock) do
        # mark the buffer as available
        delete!(used, buf)
        push!(avail, buf)

        # update pool usage
        current_usage = length(used) / (length(used) + length(avail))
        pool_usage[pid] = max(pool_usage[pid], current_usage)
      end
    else
      @timeit to "large dealloc" Mem.free(buf)
    end
  end

  if tracing
    delete!(alloc_sites, buf)
  end

  return
end

## utilities

using Printf

macro allocated(ex)
    quote
        let
            local f
            function f()
                b0 = stats.req_alloc
                $(esc(ex))
                stats.req_alloc - b0
            end
            f()
        end
    end
end

macro time(ex)
    quote
        local gpu_mem_stats0 = copy(stats)
        local cpu_mem_stats0 = Base.gc_num()
        local cpu_time0 = time_ns()

        local val = $(esc(ex))

        local cpu_time1 = time_ns()
        local cpu_mem_stats1 = Base.gc_num()
        local gpu_mem_stats1 = copy(stats)

        local cpu_time = (cpu_time1 - cpu_time0) / 1e9
        local gpu_gc_time = gpu_mem_stats1.total_time - gpu_mem_stats0.total_time
        local gpu_lib_time = gpu_mem_stats1.cuda_time - gpu_mem_stats0.cuda_time
        local gpu_alloc_count = gpu_mem_stats1.req_nalloc - gpu_mem_stats0.req_nalloc
        local gpu_alloc_size = gpu_mem_stats1.req_alloc - gpu_mem_stats0.req_alloc
        local cpu_mem_stats = Base.GC_Diff(cpu_mem_stats1, cpu_mem_stats0)
        local cpu_gc_time = cpu_mem_stats.total_time / 1e9
        local cpu_alloc_count = Base.gc_alloc_count(cpu_mem_stats)
        local cpu_alloc_size = cpu_mem_stats.allocd

        Printf.@printf("%10.6f seconds", cpu_time)
        for (typ, gctime, libtime, bytes, allocs) in
            (("CPU", cpu_gc_time, 0, cpu_alloc_size, cpu_alloc_count),
             ("GPU", gpu_gc_time, gpu_lib_time, gpu_alloc_size, gpu_alloc_count))
          if bytes != 0 || allocs != 0
              allocs, ma = Base.prettyprint_getunits(allocs, length(Base._cnt_units), Int64(1000))
              if ma == 1
                  Printf.@printf(" (%d%s %s allocation%s: ", allocs, Base._cnt_units[ma], typ, allocs==1 ? "" : "s")
              else
                  Printf.@printf(" (%.2f%s %s allocations: ", allocs, Base._cnt_units[ma], typ)
              end
              print(Base.format_bytes(bytes))
              if gctime > 0
                  Printf.@printf(", %.2f%% gc time", 100*gctime/cpu_time)
                if libtime > 0
                    Printf.@printf(" of which %.2f%% spent allocating", 100*libtime/gctime)
                end
              end
              print(")")
          elseif gctime > 0
              Printf.@printf(", %.2f%% %s gc time", 100*gctime/cpu_time, typ)
          end
        end
        println()

        val
    end
end

function pool_status()
  used_pool_buffers = 0
  used_pool_bytes = 0
  for (pid, pl) in enumerate(pools_used)
    bytes = poolsize(pid)
    used_pool_buffers += length(pl)
    used_pool_bytes += bytes * length(pl)
  end

  avail_pool_buffers = 0
  avail_pool_bytes = 0
  for (pid, pl) in enumerate(pools_avail)
    bytes = poolsize(pid)
    avail_pool_buffers += length(pl)
    avail_pool_bytes += bytes * length(pl)
  end

  free_bytes, total_bytes = CUDAdrv.Mem.info()
  used_bytes = total_bytes - free_bytes
  used_ratio = used_bytes / total_bytes

  pool_ratio = (used_pool_bytes + avail_pool_bytes) / used_bytes

  println("Total GPU memory usage: $(100*round(used_ratio; digits=2))% ($(Base.format_bytes(used_bytes))/$(Base.format_bytes(total_bytes)))")
  println("CuArrays.jl pool usage: $(100*round(pool_ratio; digits=2))% ($(Base.format_bytes(used_pool_bytes)) in use by $used_pool_buffers buffer(s), $(Base.format_bytes(avail_pool_bytes)) idle)")

  return
end

pool_timings() = println(to)


text:
mutable struct Feature
  id::Int
  age::Int
  lastzage::Int
  lastz::Array{Float64,1}
  bel::BallTreeDensity
end

locpid=1
function upplocal()
    global locpid
    N = nprocs()
    locpid = (N > 1 ? ((locpid < N && locpid != 0 && N > 2) ? (locpid+1)%(N+1) : 2) : (N == 1 ? 1 : 2))
    return locpid
end

# trackers = Dict{Int,Feature}()

function propagate(b1Xl::BallTreeDensity, bDxb1::Array{Float64,1}, s::Array{Float64,1})
  pts1 = getPoints(b1Xl)
  pts = zeros(size(pts1))
  for i in 1:size(pts,2)
    ent = [s[1]*randn();s[2]*randn();s[3]*randn()]
    b1Tb = inv(SE2(bDxb1+ent))
    pt = vec(pts1[:,i])
    bTl = SE2([pt[1];pt[2];0.0])
    pts[:,i] = vec(se2vee(b1Tb*bTl))[1:2]
  end
  return kde!( pts, "lcv");
end

function discardOldFeatures!(trkrs::Dict{Int,Feature})
  for ft in trkrs
    if ft[2].lastzage > 30 # TODO, still hand tuned for odo update rate
      println("deleting $(ft[1])")
      delete!(trkrs, ft[1])
    end
  end
  nothing
end

function propAllTrackers!(trkrs::Dict{Int,Feature}, bDxb1::Array{Float64,1}, s::Array{Float64,1})
  discardOldFeatures!(trkrs)
  rr = RemoteRef[]
  trkID = Int[]
  for t in trkrs
    # trkrs[t[1]].bel = propagate(t[2].bel, bDxb1, s)
    push!(rr, remotecall(upplocal(), propagate, t[2].bel, bDxb1, s))
    push!(trkID, t[1])

    trkrs[t[1]].age += 1
    trkrs[t[1]].lastzage += 1
  end

  @sync begin
    for i in 1:length(trkID)
      @async trkrs[trkID[i]].bel = fetch(rr[i])
    end
  end

  nothing
end

# Input [range; bearing]
# Output [x,y], R(bearing)
function p2c(z::Array{Float64,1})
  Rt = R(z[2])
  return Rt*[z[1];0.0], Rt
end

# Input [x;y]
# Ouput range, bearing
function c2p(x::Array{Float64,1})
  b = atan(x[2],x[1])
  r = norm(x)
  return r, b
end

# Input z=[range; bearing], s=diag(Cov_pol)
# Output [x;y], Cov_cart
function pol2cart(z::Array{Float64,1}, s::Array{Float64,1})
  u, Rt = p2c(z)
  Pp2 = Matrix(Diagonal(s.^2))
  P = abs(sqrtm(Rt*Pp2*(Rt')))
  return u, P
end

# Input z=[x;y], s=diag(Cov_cart)
# Output [bearing; range], Cov_pol
function cart2pol(z::Array{Float64,1}, s::Array{Float64,1})
  r, b = c2p(z)
  Rt = R(b)
  Pp2 = Matrix(Diagonal(s.^2))
  P = abs(sqrtm(Rt'*Pp2*(Rt)))
  return [b;r], P
end

# Input z=[range; bearing], s=diag(Cov_pol), N=#points to generate for KDE
# Output KDE_cart
function p2cPtsKDE(z::Array{Float64,1}, s::Array{Float64,1}; N::Int=50)
  u, P = pol2cart(z, s)
  zPts = rand(MvNormal(u,P),N)
  return kde!(zPts, "lcv")
end

featid = map(Int,0)

function addNewFeatTrk!(trkrs::Dict{Int,Feature}, z::Array{Float64,1}, s::Array{Float64,1})
  global featid
  # len = length(trkrs)
  featid+=1
  # if length(z) < 3
  #   error("addNewFeat -- z not big enough")
  # end
  trkrs[featid] = Feature(featid, 0, 0, z, p2cPtsKDE(z, s)) #len+1
  return featid # len+1
end

function initTrackersFrom(bearan::Array{Float64,2})
  global featid
  tr = Dict{Int, Feature}()
  featid = 0
  for i in 1:size(bearan,2)
    addNewFeatTrk!(tr, bearan[:,i], [0.5;0.03])
  end
  return tr
end

function insertLkhdEvals!(lkhds::Array{Float64,2}, currFeats::Dict{Int,Feature}, sightCart::Array{Float64,2})
  #@show size(lkhds), length(currFeats)
  lkpidx = Array{Int,1}(length(currFeats))
  idx = 0
  for f in currFeats
      # @show "insertLkhdEvals", idx, size(sightCart)
      yV = evaluateDualTree(f[2].bel, sightCart[1:2,:])
      idx+=1
      lkpidx[idx]=f[1]
      lkhds[:,idx] = yV
  end
  return lkpidx
end

function evalAllLikelihoods(currFeats::Dict{Int, Feature}, sightFeats::Array{Float64,2})
  numfe = length(currFeats)
  numz = size(sightFeats,2)
  lkhds = zeros(numz, numfe)

  sightCart = zeros(2,numz)
  for i in 1:numz
    u,R = p2c(sightFeats[1:2,i])
    sightCart[:,i] = u
  end

  lkpidx = insertLkhdEvals!(lkhds, currFeats, sightCart)

  return lkhds, lkpidx
end

function addNewFeats!(trckr::Dict{Int,Feature}, lkhds::Array{Float64,2}, lkpidx::Array{Int,1},
                      allmeas::Array{Float64,2}, nidx::Array{Int,1})

  newfeIds = Int[]
  # newfeats = Dict{Int, Feature}()
  newmeas = Float64[]
  if length(nidx) != 0
    if nidx[1] != -1
      newmeas = allmeas[:,nidx]
    else
      newmeas = allmeas
    end
  else
    newmeas = allmeas[:,nidx]
  end
  for i in 1:size(newmeas,2)
      id = addNewFeatTrk!(trckr, vec(newmeas[:,i]), [0.4;0.02])
      # newfeats[id] = trckr[id]
      push!(newfeIds, id)
  end
  # increase size of likelihoods
  # r,c = size(lkhds)
  # nlkhds = zeros(r,c+length(newfeIds))
  # nlkhds[:,1:c] = lkhds
  # insertLkhdEvals!(nlkhds, newfeats, allmeas)

  # TODO -- lazy re-evaluation of all likelihoods -- big speedup possible here
  # return evalAllLikelihoods(trckr, allmeas)
  nothing
end

# agnostic to feature idex permutations
function divMaxAcross(lk::Array{Float64,2})
  rlk = round.(lk, digits=5)
  m1 = maximum(rlk, dims=1)
  m1[m1.==0.0] = 1.0
  return rlk./m1
end

# agnostic to featidxpermutations
function divMaxAlong(lk::Array{Float64,2})
  rlk = round.(lk, digits=5)
  m2 = maximum(rlk, dims=2)
  m2[m2.==0.0] = 1.0
  return rlk./m2
end

function hardMatches!(ha::Dict{Int,Array{Float64,1}}, dmdm::Array{Float64,2}, lkpidx::Array{Int,1}, allmeas::Array{Float64,2})
  # check second maximum is much smaller
  dmdm[dmdm.==2.0]=-1.0
  m1b = (maximum(dmdm,1) .< 0.1)

  for i in 1:size(dmdm,2)
    fid = lkpidx[i]
    if m1b[i] #fid
      zid = findfirst(vec(dmdm[:,i]),-1) # fid
      if zid > 0
        ha[fid] = vec(allmeas[:,zid])
      end
    end
  end
  nothing
end

function findMatches(lk::Array{Float64,2}, lkpidx::Array{Int,1}, allmeas::Array{Float64,2})
  dmac = divMaxAlong(lk)
  dmal = divMaxAcross(lk)
  dmdm = dmac+dmal
  hardassoc = Dict{Int,Array{Float64,1}}()
  hardMatches!(hardassoc, deepcopy(dmdm), lkpidx, allmeas)
  return hardassoc
end

# agnostic to lkindex permutations
function findNewFeats(lkhds::Array{Float64,2}; thr=0.00001)
  numz = size(lkhds, 1)
  if size(lkhds, 2) == 0
    return [-1]
  end
  low = vec(maximum(lkhds,2) .< thr)
  return (1:numz)[low]
end

# will also add new features at the end of the tracking pool
# currently only does hard associations, soft multihypothesis tracking work to follow
function assocMeasWFeats!(trkrs::Dict{Int, Feature}, fez::Array{Float64,2})
  if size(fez,2) == 0
    return Dict{Int,Array{Float64,1}}() # no matches possible
  end
  lk, lkpidx = evalAllLikelihoods(trkrs, fez)
  hardassoc = findMatches(lk, lkpidx, fez)
  # Also add new features into the mix here
  nidx = findNewFeats(lk)
  addNewFeats!(trkrs, lk, lkpidx, fez, nidx)
  return hardassoc
end

## Moved to RecursiveFiltering
function update(bhatXl::BallTreeDensity, z::Array{Float64,1}, s::Array{Float64,1}; N::Int=75)

  bXl = p2cPtsKDE(z,s, N=N)
  # take the product between predicted and measured position
  dummy = kde!(rand(2,N),[1.0])
  pGM, = prodAppxMSGibbsS(dummy, [bhatXl, bXl], Union{}, Union{}, 5)

  # error("update -- Not ready")
  return kde!(pGM, "lcv")
end
function updatelin(bhatXl::BallTreeDensity, z::Array{Float64,1}, s::Array{Float64,1}; N::Int=75)

  bXl = resample(kde!(vectoarr2(z),s),N)
  # take the product between predicted and measured position
  dummy = kde!(rand(length(z),N),[1.0])
  pGM, = prodAppxMSGibbsS(dummy, [bhatXl, bXl], Union{}, Union{}, 5)

  # error("update -- Not ready")
  return kde!(pGM, "lcv")
end
function updatelin(bhatXl::BallTreeDensity, bXl::BallTreeDensity; N=75)
  dummy = kde!(rand(bhatXl.bt.dims,N),[1.0])
  pGM, = prodAppxMSGibbsS(dummy, [bhatXl, bXl], Union{}, Union{}, 5)
  # error("update -- Not ready")
  return kde!(pGM, "lcv")
end
## Moved to RecursiveFiltering

function doAsyncUpdate(oldFeat::Feature, meas::Array{Float64,1}, s::Array{Float64,1})
  bel = update(oldFeat.bel, meas, s)
  return Feature(oldFeat.id, oldFeat.age, 0, meas[1:3], bel)
end

function measUpdateTrackers!(trkrs::Dict{Int,Feature}, assoc::Dict{Int,Array{Float64,1}}, s::Array{Float64,1})
  # loop through all associated measurements
  # TODO --this can all be done in parallel to reduce computation time
  # dbg = Dict{Int,Array{BallTreeDensity,1}}()
  rr = RemoteRef[]
  trkID = Int[]
  for meas in assoc
    # bel = deepcopy(trkrs[meas[1]].bel)
    # trkrs[meas[1]].bel = update(bel, meas[2], s)
    # trkrs[meas[1]].lastzage = 0
    # trkrs[meas[1]].lastz = meas[2]

    # trkrs[meas[1]] = doAsyncUpdate(trkrs[meas[1]], meas[2], s)
    push!(rr, remotecall(upplocal(), doAsyncUpdate, trkrs[meas[1]], meas[2], s))
    push!( trkID, meas[1])

    # arr = Array{BallTreeDensity,1}()
    # push!(arr,bel)
    # push!(arr,p2cPtsKDE(meas[2],[1.0;0.02]))
    # push!(arr,deepcopy(trkrs[meas[1]].bel))
    # dbg[meas[1]] = arr
  end

  @sync begin
    for i in 1:length(trkID)
      @async trkrs[trkID[i]] = fetch(rr[i])
    end
  end

  # return dbg
  nothing
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct IoK8sApiCoreV1PortworxVolumeSource <: SwaggerModel
    fsType::Any # spec type: Union{ Nothing, String } # spec name: fsType
    readOnly::Any # spec type: Union{ Nothing, Bool } # spec name: readOnly
    volumeID::Any # spec type: Union{ Nothing, String } # spec name: volumeID

    function IoK8sApiCoreV1PortworxVolumeSource(;fsType=nothing, readOnly=nothing, volumeID=nothing)
        o = new()
        validate_property(IoK8sApiCoreV1PortworxVolumeSource, Symbol("fsType"), fsType)
        setfield!(o, Symbol("fsType"), fsType)
        validate_property(IoK8sApiCoreV1PortworxVolumeSource, Symbol("readOnly"), readOnly)
        setfield!(o, Symbol("readOnly"), readOnly)
        validate_property(IoK8sApiCoreV1PortworxVolumeSource, Symbol("volumeID"), volumeID)
        setfield!(o, Symbol("volumeID"), volumeID)
        o
    end
end # type IoK8sApiCoreV1PortworxVolumeSource

const _property_map_IoK8sApiCoreV1PortworxVolumeSource = Dict{Symbol,Symbol}(Symbol("fsType")=>Symbol("fsType"), Symbol("readOnly")=>Symbol("readOnly"), Symbol("volumeID")=>Symbol("volumeID"))
const _property_types_IoK8sApiCoreV1PortworxVolumeSource = Dict{Symbol,String}(Symbol("fsType")=>"String", Symbol("readOnly")=>"Bool", Symbol("volumeID")=>"String")
Base.propertynames(::Type{ IoK8sApiCoreV1PortworxVolumeSource }) = collect(keys(_property_map_IoK8sApiCoreV1PortworxVolumeSource))
Swagger.property_type(::Type{ IoK8sApiCoreV1PortworxVolumeSource }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiCoreV1PortworxVolumeSource[name]))}
Swagger.field_name(::Type{ IoK8sApiCoreV1PortworxVolumeSource }, property_name::Symbol) =  _property_map_IoK8sApiCoreV1PortworxVolumeSource[property_name]

function check_required(o::IoK8sApiCoreV1PortworxVolumeSource)
    (getproperty(o, Symbol("volumeID")) === nothing) && (return false)
    true
end

function validate_property(::Type{ IoK8sApiCoreV1PortworxVolumeSource }, name::Symbol, val)
end


text:
"""
A `ConstraintFunction` contains the created forward and backward
code
"""
mutable struct ConstraintFunction{F <: Function, G <: Function}
    input::Vector{Symbol}  # input arguments for forward function
    output::Vector{Symbol} # output arguments for forward function
    forward::F
    backward::G
    forward_code::Expr
    backward_code::Expr
    expression::Expr
end

function (C::ConstraintFunction)(args...)
    return C.forward(args)[1]
end

function Base.show(io::IO, f::ConstraintFunction{F,G}) where {F,G}
    println(io, "ConstraintFunction:")
    println(io, "  - input arguments: $(f.input)")
    println(io, "  - output arguments: $(f.output)")
    print(io, "  - expression: $(MacroTools.striplines(f.expression))")
end

struct FunctionArguments
    input
    return_arguments
    intermediate
end

const registered_functions = Dict{Symbol, FunctionArguments}()

# """
# `@function` registers a function to be used in forwards and backwards mode.
#
# Example: `@function f(x, y) = x^2 + y^2`
# """  # this docstring does not work!
@eval macro ($(:function))(ex)   # workaround to define macro @function

    (f, args, code) = match_function(ex)

    return_arguments, flatAST = flatten(code)

    # make into an array:
    if !(isa(return_arguments, Array))
        return_arguments = [return_arguments]
    end

    # rearrange so actual return arguments come first:

    flatAST.intermediate = setdiff(flatAST.intermediate, return_arguments)

    # println("HERE")
    # flatAST.intermediate = [return_arguments; flatAST.intermediate]

    forward, backward = forward_backward(flatAST)

    registered_functions[f] = FunctionArguments(
            flatAST.variables, return_arguments, flatAST.intermediate)

    return quote

        $(esc(f)) =
            ConstraintFunction($(flatAST.variables),
                                $(flatAST.intermediate),
                                $(forward),
                                $(backward),
                                $(Meta.quot(forward)),
                                $(Meta.quot(backward)),
                                $(Meta.quot(ex))
                                )

    end

end

function match_function(ex)

    try
        @capture ex begin
            (   (f_(args__) = body_)
                  | (function f_(args__) body_ end) )
           end

         return (f, args, rmlines(body))  # rmlines is from MacroTools package

    catch
        throw(ArgumentError("$ex does not have the form of a function"))
    end
end


text:
"""
    PoisLik <: Likelihood

Poisson likelihood
```math
p(yᵢ = k | fᵢ) = θᵏ\\exp(-θ)/k!
```
for ``k ∈ N₀``, where ``θ = \\exp(f)`` and ``f`` is the latent Gaussian process.
"""
struct PoisLik <: Likelihood end

#log of probability density
function log_dens(poisson::PoisLik, f::AbstractVector, y::Vector{Int})
    #where we exponentiate for positivity f = exp(fi)
    return y.*f - exp.(f) - lgamma.(1.0 .+ y)
end

#derivative of pdf wrt latent function
function dlog_dens_df(poisson::PoisLik, f::AbstractVector, y::Vector{Int})
    return y - exp.(f)
end

#mean and variance under likelihood
mean_lik(poisson::PoisLik, f::AbstractVector) = exp.(f)
var_lik(poisson::PoisLik, f::AbstractVector) = exp.(f)

get_params(poisson::PoisLik) = []
num_params(poisson::PoisLik) = 0

function var_exp(ll::PoisLik, y::AbstractArray, m::AbstractArray, V::AbstractMatrix)
    tot = 0
    V_diag = diag(V)
    for (a, b, c) in zip(y, m, V_diag)
        tot +=  a*b - exp(b + c/2) - log(factorial(convert(Int64, a))) # convert to lgamma(y+1)
    end
    return tot
end

function var_exp(ll::PoisLik, y::AbstractArray, m::AbstractArray, V::AbstractArray)
    tot = 0
    for (a, b, c) in zip(y, m, V)
        tot +=  a*b - exp(b + c/2) - log(factorial(convert(Int64, a))) # convert to lgamma(y+1)
    end
    return tot
end

function var_exp(ll::PoisLik, y::Number, m::Number, V::Number)
    return y*m - exp(m + V/2) - log(factorial(convert(Int64, y))) # convert to lgamma(y+1)
end

function dv_var_exp(ll::PoisLik, y::Number, m::Number, V::Number)
    return gradient(x -> var_exp(ll, y, m, x), V)[1]
end


text:
using ConstrainedOptim
COpt = ConstrainedOptim
using Optim
using Base.Test

@testset "Finite Difference Testing" begin
    f(x) = x[1]^2 + x[2]^2 + x[2]^4
    df(x) = [2*x[1], 2*x[2]+4*x[2]^3]
    c(x) = x[1]^2 + (x[2]-0.5)^2 - 2.0
    dc(x) = [ 2*x[1] 2*(x[2]-0.5) ]

    # the minimiser should be [0.0,0.0]
    x0 = [1.0, 0.0]

    F = DifferentiableFunction(f, (x,g) -> copy!(g, df(x)) )
    C = EqualityConstraint(c, (x,g) -> copy!(g, dc(x)) )

    println("------------------------------------------------------------")
    println("       Finite Difference Testing the AL ")
    @printf("   h   |   err \n")
    @printf("-------|-------------\n")
    # finite-difference test of the augmented Lagrangian implementation
    al = COpt.AugmentedLagrangian(F, C, x0)
    al.lambda = 0.0
    al.mu = 1.0
    A = COpt.evaluate(x0, al)
    dA = COpt.gradient(x0, al)
    err = Float64[]
    for p = 2:12
       h = 0.1^p
       dAh = zeros(dA)
       for n = 1:length(x0)
          x0[n] += h
          dAh[n] = (COpt.evaluate(x0, al) - A)  / h
          x0[n] -= h
       end
       push!(err, vecnorm(dA - dAh, Inf))
       @printf(" 1e-%2d |  %1.4e \n", p, err[end])
    end
    @test minimum(err) < 1e-4 * err[1]
    if minimum(err) < 1e-4 * err[1]
       println("looks like the FD test has passed...")
    else
       warn("""the finite difference test for the augmented Lagrangien didn't
             pass; please check visually what happened and debug""")
    end

    println("------------------------------------------------------------")
    println("      Try to optimise something simple")
    x, al = ConstrainedOptim.optimize(F, C, x0)
    println("Converged to ", x, "; λ = ", al.lambda)
    println("First-order optimality: ")
    al.mu = 0.0; g = COpt.gradient(x, al); C = al.C.c(x)
    println("   ∇ₓL(x, λ) = ", g)
    println("        c(x) = ", C)
    println("  |∇L(x, λ)| = ", max(vecnorm(g), vecnorm(C)))
    @test vecnorm(g) < 1e-6
    @test vecnorm(C) < 1e-6
    end

# Another test case
@testset "Simple problems" begin
    let
        f(x) = 2x[1]^2+x[2]^2
        ∇f(x) = [4x[1], 2x[2]]
        c(x) = sum(x)-1
        ∇c(x) = ones(1,2)
        initial_x = [0.3,0.75]
        solution_x = [1/3, 2/3]
        F = DifferentiableFunction(f, (x,g) -> copy!(g, ∇f(x)))
        C = EqualityConstraint(c, (x,g) -> copy!(g, ∇c(x)) )

        x, al = ConstrainedOptim.optimize(F, C, initial_x)
        @test norm(x - solution_x, Inf) < 1e-6
    end

    let
        f(x) = -prod(x)
        ∇f(x) = [-x[2]; -x[1]]
        c(x) = sum(x)-6
        ∇c(x) = ones(1,2)
        initial_x = [2.0, 2.0]
        solution_x = [3.0, 3.0]
        F = DifferentiableFunction(f, (x,g) -> copy!(g, ∇f(x)))
        C = EqualityConstraint(c, (x,g) -> copy!(g, ∇c(x)) )

        x, al = ConstrainedOptim.optimize(F, C, initial_x)
        @test norm(x - solution_x, Inf) < 1e-6
    end

    @testset "Cobb-Douglas Cost Minimization" begin
        # Ref: any microeconomics text book

        # Factor prices
        w = [1., 1.]

        # Output elasticities of the factors
        γ = [0.25, 0.75]

        # Quantity to be produced
        q = 2.

        # Cost
        f(x) = dot(w,x)
        ∇f(x) = [w[1], w[2]]

        # Cobb-Douglas Production function
        c(x) = (x[1]^γ[1])*(x[2]^γ[2])-q
        ∇c(x) = [γ[1]*(x[1]^(γ[1]-1))*(x[2]^γ[2]) (x[1]^γ[1])*γ[2]*(x[2]^(γ[2]-1))]

        # Some guess (it doesn't even produce the correct amount...)
        initial_x = [1., 2.]

        # Analytical solution
        solution_x = [(w[2]/w[1])*(γ[1]/γ[2])^(γ[2]/sum(γ))*q^(1/sum(γ));(w[1]/w[2])*(γ[2]/γ[1])^(γ[1]/sum(γ))*q^(1/sum(γ))]

        F = DifferentiableFunction(f, (x,g) -> copy!(g, ∇f(x)))
        C = EqualityConstraint(c, (x,g) -> copy!(g, ∇c(x)) )

        x, al = ConstrainedOptim.optimize(F, C, initial_x)
        @test norm(x - solution_x, Inf) < 1e-6
    end

    @testset "Nocedal-Wright examples" begin
        let # 17.1
            f(x) = x[1]+x[2]
            ∇f(x) = [1., 1.]
            c(x) = x[1]^2+x[2]^2-2
            ∇c(x) = [2*x[1] 2*x[2]]

            # Initial value is arbitrary...
            initial_x = [-0.3, -0.5]
            solution_x = [-1.0, -1.0]
            F = DifferentiableFunction(f, (x,g) -> copy!(g, ∇f(x)))
            C = EqualityConstraint(c, (x,g) -> copy!(g, ∇c(x)) )

            x, al = ConstrainedOptim.optimize(F, C, initial_x)
            @test norm(x - solution_x, Inf) < 1e-6
        end

        let # Page 500
            f(x) = -5x[1]^2+x[2]^2
            ∇f(x) = [-10.0x[1], 2.0x[2]]
            c(x) = x[1]-1
            ∇c(x) = [1  0]

            # Initial value is arbitrary...
            initial_x = [3., 0.5]
            solution_x = [1.0, 0.0]
            F = DifferentiableFunction(f, (x,g) -> copy!(g, ∇f(x)))
            C = EqualityConstraint(c, (x,g) -> copy!(g, ∇c(x)) )

            x, al = ConstrainedOptim.optimize(F, C, initial_x)
            @test norm(x - solution_x, Inf) < 1e-6
        end

        @testset "Simple Convex Constraints on x" begin
            eta = 0.9

            function f_gd_2(x)
              (1.0 / 2.0) * (x[1]^2 + eta * x[2]^2)
            end

            function g_gd_2(x, storage)
              storage[1] = x[1]
              storage[2] = eta * x[2]
            end

            d = DifferentiableFunction(f_gd_2, g_gd_2)
            box_c = BoxConstraint(fill(-1.0, 2), fill(1.0, 2))
            ball_c = BallConstraint(fill(0.0, 2), 1.0, 2)
            res_unc = Optim.optimize(d, [1.0, 1.0], GradientDescent(linesearch! = Optim.backtracking_linesearch!))
            res_con_box = optimize(d, [1.0, 1.0], box_c, ProjectedGradientDescent(), OptimizationOptions())
            res_con_ball = optimize(d, [1.0, 1.0], ball_c, ProjectedGradientDescent(), OptimizationOptions())

            @test Optim.minimum(res_unc) == Optim.minimum(res_con_box) == Optim.minimum(res_con_ball)
        end
    end
end

max_stars_repo_name: cortner/AugmentedLagrangian

max_issues_repo_name: cortner/AugmentedLagrangian

max_forks_repo_name: cortner/AugmentedLagrangian


text:
"""
    struct NodeToDofGlue{T}
      free_dof_to_node::Vector{Int32}
      free_dof_to_comp::Vector{Int16}
      dirichlet_dof_to_node::Vector{Int32}
      dirichlet_dof_to_comp::Vector{Int16}
      node_and_comp_to_dof::Vector{T}
    end
"""
struct NodeToDofGlue{T}
  free_dof_to_node::Vector{Int32}
  free_dof_to_comp::Vector{Int16}
  dirichlet_dof_to_node::Vector{Int32}
  dirichlet_dof_to_comp::Vector{Int16}
  node_and_comp_to_dof::Vector{T}
end

"""
    CLagrangianFESpace(::Type{T},grid::Grid) where T
"""
function CLagrangianFESpace(::Type{T},grid::Grid) where T
  vector_type = Vector{_dof_type(T)}
  node_to_tag = fill(Int8(UNSET),num_nodes(grid))
  tag_to_mask = fill(_default_mask(T),0)
  CLagrangianFESpace(T,grid,vector_type,node_to_tag,tag_to_mask)
end

"""
    CLagrangianFESpace(
    ::Type{T},
    grid::Grid,
    vector_type::Type,
    node_to_tag::AbstractVector,
    tag_to_mask::AbstractVector) where T
"""
function CLagrangianFESpace(
  ::Type{T},
  grid::Grid,
  vector_type::Type,
  node_to_tag::AbstractVector{<:Integer},
  tag_to_masks::AbstractVector) where T

  z = zero(T)
  glue, dirichlet_dof_tag = _generate_node_to_dof_glue_component_major(
    z,node_to_tag,tag_to_masks)
  cell_reffe = _generate_cell_reffe_clagrangian(z,grid)
  cell_dofs_ids = _generate_cell_dofs_clagrangian(z,grid,glue,cell_reffe)
  cell_shapefuns = lazy_map(get_shapefuns,cell_reffe)
  cell_dof_basis = lazy_map(get_dof_basis,cell_reffe)
  rd = ReferenceDomain()
  fe_basis, fe_dof_basis = compute_cell_space(
    cell_shapefuns,cell_dof_basis,rd,rd,grid)
  cell_is_dirichlet = _generate_cell_is_dirichlet(cell_dofs_ids)

  nfree = length(glue.free_dof_to_node)
  ndirichlet = length(glue.dirichlet_dof_to_node)
  ntags = length(tag_to_masks)
  dirichlet_cells = collect(Int32,findall(cell_is_dirichlet))

  space = UnconstrainedFESpace(
    vector_type,
    nfree,
    ndirichlet,
    cell_dofs_ids,
    fe_basis,
    fe_dof_basis,
    cell_is_dirichlet,
    dirichlet_dof_tag,
    dirichlet_cells,
    ntags,
    glue)

  space
end

function _use_clagrangian(trian,cell_reffe,conf)
  false
end

function _use_clagrangian(trian::Grid,cell_reffe,conf::H1Conformity)
  ctype_reffe1, cell_ctype1 = compress_cell_data(cell_reffe)
  ctype_reffe2 = get_reffes(trian)
  if length(ctype_reffe1) != 1 || length(ctype_reffe2) != 1
    return false
  end
  reffe1 = first(ctype_reffe1)
  reffe2 = first(ctype_reffe2)
  if get_orders(reffe1) != get_orders(reffe2)
    return false
  end
  if get_order(reffe1) != 1 # This can be relaxed in the future
    return false
  end
  return true
end

function _unsafe_clagrangian(
  cell_reffe,
  grid,
  labels,
  vector_type,
  dirichlet_tags,
  dirichlet_masks)

  ctype_reffe, cell_ctype = compress_cell_data(cell_reffe)
  prebasis = get_prebasis(first(ctype_reffe))
  T = return_type(prebasis)
  # Next line assumes linear grids
  node_to_tag = get_face_tag_index(labels,dirichlet_tags,0)
  _vector_type = vector_type === nothing ? Vector{Float64} : vector_type
  tag_to_mask = dirichlet_masks === nothing ? fill(_default_mask(T),length(dirichlet_tags)) : dirichlet_masks
  CLagrangianFESpace(T,grid,_vector_type,node_to_tag,tag_to_mask)
end

# Helpers

_default_mask(::Type) = true
_default_mask(::Type{T}) where T <: MultiValue = ntuple(i->true,Val{length(T)}())

_dof_type(::Type{T}) where T = T
_dof_type(::Type{T}) where T<:MultiValue = eltype(T)

function _generate_cell_is_dirichlet(cell_dofs)
  cell_is_dirichlet = Vector{Bool}(undef,length(cell_dofs))
  cache = array_cache(cell_dofs)
  isnegative(dof) = dof<0
  for cell in 1:length(cell_dofs)
    dofs = getindex!(cache,cell_dofs,cell)
    cell_is_dirichlet[cell] = any(isnegative,dofs)
  end
  cell_is_dirichlet
end

function _generate_node_to_dof_glue_component_major(z,node_to_tag,tag_to_mask)
  @notimplementedif length(z) != 1
  nfree_dofs = 0
  ndiri_dofs = 0
  for tag in node_to_tag
    if tag == UNSET
      nfree_dofs += 1
    else
      mask = tag_to_mask[tag]
      if mask
        ndiri_dofs += 1
      else
        nfree_dofs += 1
      end
    end
  end
  free_dof_to_node = Vector{Int32}(undef,nfree_dofs)
  free_dof_to_comp = ones(Int16,nfree_dofs)
  diri_dof_to_node = Vector{Int32}(undef,ndiri_dofs)
  diri_dof_to_comp = ones(Int16,ndiri_dofs)
  diri_dof_to_tag = ones(Int8,ndiri_dofs)
  nnodes = length(node_to_tag)
  node_and_comp_to_dof = Vector{Int32}(undef,nnodes)
  nfree_dofs = 0
  ndiri_dofs = 0
  for (node,tag) in enumerate(node_to_tag)
    if tag == UNSET
      nfree_dofs += 1
      free_dof_to_node[nfree_dofs] = node
      node_and_comp_to_dof[node] = nfree_dofs
    else
      mask = tag_to_mask[tag]
      if mask
        ndiri_dofs += 1
        diri_dof_to_node[ndiri_dofs] = node
        diri_dof_to_tag[ndiri_dofs] = tag
        node_and_comp_to_dof[node] = -ndiri_dofs
      else
        nfree_dofs += 1
        free_dof_to_node[nfree_dofs] = node
        node_and_comp_to_dof[node] = nfree_dofs
      end
    end
  end
  glue = NodeToDofGlue(
   free_dof_to_node,
   free_dof_to_comp,
   diri_dof_to_node,
   diri_dof_to_comp,
   node_and_comp_to_dof)

  glue, diri_dof_to_tag
end

function _generate_node_to_dof_glue_component_major(
  z::MultiValue,node_to_tag,tag_to_masks)
  nfree_dofs = 0
  ndiri_dofs = 0
  ncomps = length(z)
  @assert length(testitem(tag_to_masks)) == ncomps
  for (node,tag) in enumerate(node_to_tag)
    if tag == UNSET
      nfree_dofs += ncomps
    else
      for mask in tag_to_masks[tag]
        if mask
          ndiri_dofs += 1
        else
          nfree_dofs += 1
        end
      end
    end
  end
  free_dof_to_node = Vector{Int32}(undef,nfree_dofs)
  free_dof_to_comp = ones(Int16,nfree_dofs)
  diri_dof_to_node = Vector{Int32}(undef,ndiri_dofs)
  diri_dof_to_comp = ones(Int16,ndiri_dofs)
  diri_dof_to_tag = ones(Int8,ndiri_dofs)
  T = change_eltype(z,Int32)
  nnodes = length(node_to_tag)
  node_and_comp_to_dof = zeros(T,nnodes)
  nfree_dofs = 0
  ndiri_dofs = 0
  m = zero(Mutable(T))
  for (node,tag) in enumerate(node_to_tag)
    if tag == UNSET
      for comp in 1:ncomps
        nfree_dofs += 1
        free_dof_to_node[nfree_dofs] = node
        free_dof_to_comp[nfree_dofs] = comp
        m[comp] = nfree_dofs
      end
    else
      masks = tag_to_masks[tag]
      for comp in 1:ncomps
        mask = masks[comp]
        if mask
          ndiri_dofs += 1
          diri_dof_to_node[ndiri_dofs] = node
          diri_dof_to_comp[ndiri_dofs] = comp
          diri_dof_to_tag[ndiri_dofs] = tag
          m[comp] = -ndiri_dofs
        else
          nfree_dofs += 1
          free_dof_to_node[nfree_dofs] = node
          free_dof_to_comp[nfree_dofs] = comp
          m[comp] = nfree_dofs
        end
      end
    end
    node_and_comp_to_dof[node] = m
  end
  glue = NodeToDofGlue(
   free_dof_to_node,
   free_dof_to_comp,
   diri_dof_to_node,
   diri_dof_to_comp,
   node_and_comp_to_dof)

  glue, diri_dof_to_tag
end

function _generate_cell_reffe_clagrangian(z,grid)
  ctype_to_reffe = get_reffes(grid)
  cell_to_ctype = get_cell_type(grid)
  function fun(reffe)
    p = get_polytope(reffe)
    orders = get_orders(reffe)
    LagrangianRefFE(typeof(z),p,orders)
  end
  # using map instead of lazy_map seems to kill the compiler.
  ctype_to_reffe = collect(lazy_map(fun,ctype_to_reffe))
  cell_to_reffe = expand_cell_data(ctype_to_reffe,cell_to_ctype)
end

function _generate_cell_dofs_clagrangian(z,grid,glue,cell_to_reffe)
  ctype_to_reffe, cell_to_ctype = compress_cell_data(cell_to_reffe)
  cell_to_nodes = get_cell_node_ids(grid)
  cell_to_dofs = _generate_cell_dofs_clagrangian(
    z,
    cell_to_nodes,
    ctype_to_reffe,
    cell_to_ctype,
    glue.node_and_comp_to_dof)
  cell_to_dofs
end

function _generate_cell_dofs_clagrangian(
  z,
  cell_to_nodes,
  ctype_to_reffe,
  cell_to_ctype,
  node_and_comp_to_dof)

  cell_to_dofs = lazy_map(
    Broadcasting(Reindex(node_and_comp_to_dof)),cell_to_nodes)
  Table(cell_to_dofs)
end

function _generate_cell_dofs_clagrangian(
  z::MultiValue,
  cell_to_nodes,
  ctype_to_reffe,
  cell_to_ctype,
  node_and_comp_to_dof)

  ncomps = num_components(z)

  ctype_to_lnode_to_comp_to_ldof = map(get_node_and_comp_to_dof,ctype_to_reffe)
  ctype_to_num_ldofs = map(num_dofs,ctype_to_reffe)

  cell_to_dofs = _allocate_cell_dofs_clagrangian(
    ctype_to_num_ldofs,cell_to_ctype)

  cache = array_cache(cell_to_nodes)
  _fill_cell_dofs_clagrangian!(
    cell_to_dofs,
    cache,
    cell_to_nodes,
    ctype_to_lnode_to_comp_to_ldof,
    cell_to_ctype,
    node_and_comp_to_dof,
    ncomps)

  cell_to_dofs
end

function _allocate_cell_dofs_clagrangian(ctype_to_num_ldofs,cell_to_ctype)
  ncells = length(cell_to_ctype)
  ptrs = zeros(Int32,ncells+1)
  for cell in 1:ncells
    ctype = cell_to_ctype[cell]
    nldofs = ctype_to_num_ldofs[ctype]
    ptrs[cell+1] = nldofs
  end
  length_to_ptrs!(ptrs)
  ndata = ptrs[end]-1
  data = zeros(Int32,ndata)
  Table(data,ptrs)
end

function _fill_cell_dofs_clagrangian!(
  cell_to_dofs,
  cache,
  cell_to_nodes,
  ctype_to_lnode_and_comp_to_ldof,
  cell_to_ctype,
  node_and_comp_to_dof,
  ncomps)

  ncells = length(cell_to_ctype)
  for cell in 1:ncells
    nodes = getindex!(cache,cell_to_nodes,cell)
    ctype = cell_to_ctype[cell]
    lnode_and_comp_to_ldof = ctype_to_lnode_and_comp_to_ldof[ctype]
    p = cell_to_dofs.ptrs[cell]-1
    for (lnode, node) in enumerate(nodes)
      for comp in 1:ncomps
        ldof = lnode_and_comp_to_ldof[lnode][comp]
        dof = node_and_comp_to_dof[node][comp]
        cell_to_dofs.data[p+ldof] = dof
      end
    end
  end
end


text:
export convert, conj, real, imag, abs2, norm

## convert functions
function convert(::Type{Polynomial}, expo::Exponent)
    return Polynomial(SortedDict{Exponent, Number}(expo=>1.0))
end

function convert(::Type{Polynomial}, x::Variable)
    return Polynomial(SortedDict{Exponent, Number}(Exponent(SortedDict(x=>Degree(1,0)))=>1.0))
end

function convert(::Type{Polynomial}, λ::Number)
    return Polynomial(SortedDict{Exponent, Number}(Exponent()=>λ))
end

function convert(::Type{Point}, pt::SortedDict{Variable, T}) where T
    return Point(convert(SortedDict{Variable, Number}, pt))
end

function convert(::Type{AbstractPolynomial}, λ::Number)
    return Polynomial(SortedDict{Exponent, Number}(Exponent()=>λ))
end

## Conjugate
function conj(d::Degree)
    return Degree(d.conjvar, d.explvar)
end

function conj(x::Variable)
    if x.kind<:Complex
        return Exponent(SortedDict(x=>Degree(0,1)))
    else
        return Exponent(SortedDict(x=>Degree(1,0)))
    end
end

function conj(expo::Exponent)
    expodict = SortedDict{Variable, Degree}()
    for (var, deg) in expo.expo
        if iscomplex(var)
            expodict[var] = conj(deg)
        else
            expodict[var] = deg
        end
    end
    return Exponent(expodict)
end

function conj(p::Polynomial)
    pdict = SortedDict{Exponent, Number}()
    for (expo, λ) in p
        pdict[conj(expo)] = conj(λ)
    end
    return Polynomial(pdict)
end

## Real, imag
real(p::Polynomial) = (p + conj(p))/2
imag(p::Polynomial) = (p - conj(p))/(2im)

function real(p::T) where T<:AbstractPolynomial
    return (p + conj(p))/2
end

function imag(p::T) where T<:AbstractPolynomial
    return (p + -conj(p))/(2im)
end

## Norm, abs2
function abs2(p::T) where T<:AbstractPolynomial
    return p*conj(p)
end

function norm(pt::Point, p::Real = 2)
    p > 0 || error(LOGGER, "norm(): p should be positive ($p ≤ 0).")
    if p == Inf
        maxi = -Inf
        for (var, val) in pt
            # iscomplex(var) && warn("norm(): var $(var) is not real, splitting real and imag parts.")
            if !iscomplex(var) && (val > maxi)
                maxi = val
            elseif iscomplex(var) && max(real(val), imag(val)) > maxi
                maxi = max(real(val), imag(val))
            end
        end
        return maxi
    elseif p == 1
      s = 0
      for (var, val) in pt
          if !iscomplex(var)
            s += abs(val)
          else
            # warn("norm(): var $(var) is not real, splitting real and imag parts.")
            s += abs(real(val)) + abs(imag(val))
          end
      end
      return s
    else
      s = 0
      for (var, val) in pt
          if !iscomplex(var)
            s += val^p
          else
            # warn("norm(): var $(var) is not real, splitting real and imag parts.")
            s += real(val)^p + imag(val)^p
          end
      end
      return s^(1/p)
  end
end