text:
"""
    simplify!(com)

Simplify constraints i.e by adding new constraints which uses an implicit connection between two constraints.
i.e an `all_different` does sometimes include information about the sum.
Return a list of newly added constraint ids
"""
function simplify!(com)
    added_constraint_idxs = Int[]
    # check if we have all_different and sum constraints
    b_all_different = false
    # (all different where every value is used)
    b_all_different_sum = false
    b_equal_to = false
    # != constraint
    b_not_equal_to = false
    b_svc_less_than = false
    for constraint in com.constraints
        if isa(constraint.set, CPE.AllDifferent)
            b_all_different = true
            if length(constraint.indices) == length(constraint.pvals)
                b_all_different_sum = true
            end
        elseif isa(constraint.fct, SAF) && isa(constraint.set, MOI.EqualTo)
            b_equal_to = true
        elseif isa(constraint.fct, SAF) && isa(constraint.set, CPE.DifferentFrom)
            b_not_equal_to = true
        elseif isa(constraint, SingleVariableConstraint) &&
               isa(constraint.set, MOI.LessThan)
            b_svc_less_than = true
        end
    end
    if b_all_different_sum && b_equal_to
        append!(added_constraint_idxs, simplify_all_different_and_equal_to(com))
    end
    if b_not_equal_to
        append!(added_constraint_idxs, simplify_not_equal_to_cliques(com))
    end
    if b_svc_less_than
        append!(added_constraint_idxs, simplify_svc_less_than(com))
    end

    if length(added_constraint_idxs) > 0 && :Info in com.options.logging
        println("Added $(length(added_constraint_idxs)) new constraints")
    end
    return added_constraint_idxs
end

"""
    simplify_svc_less_than(com)

Simplify several a >= X constraints (with the same a) to a GeqSetConstraint.
TODO: A LeqSetConstraint is missing to do the same for a <= X
"""
function simplify_svc_less_than(com)
    added_constraint_idxs = Int[]
    # save all variables, constraints of the form a <= b
    lhs_vars = Dict{Int,Vector{Int}}()
    lhs_cons = Dict{Int,Vector{Int}}()
    # save all variables, constraints of the form a >= b
    rhs_vars = Dict{Int,Vector{Int}}()
    rhs_cons = Dict{Int,Vector{Int}}()
    for constraint_idx in 1:length(com.constraints)
        constraint = com.constraints[constraint_idx]
        if isa(constraint, SingleVariableConstraint) &&
           isa(constraint.set, MOI.LessThan)
            if haskey(lhs_vars, constraint.lhs)
                push!(lhs_vars[constraint.lhs], constraint.rhs)
                push!(lhs_cons[constraint.lhs], constraint.idx)
            else
                lhs_vars[constraint.lhs] = [constraint.rhs]
                lhs_cons[constraint.lhs] = [constraint.idx]
            end

            if haskey(rhs_vars, constraint.rhs)
                push!(rhs_vars[constraint.rhs], constraint.lhs)
                push!(rhs_cons[constraint.rhs], constraint.idx)
            else
                rhs_vars[constraint.rhs] = [constraint.lhs]
                rhs_cons[constraint.rhs] = [constraint.idx]
            end
        end
    end

    # check for >= which constraint appears >= 5 times on the left side
    # Todo: Do the same for lhs_vars with LeqSet
    for (key, value) in rhs_vars
        if length(value) >= 5
            # add new all different constraint
            set = GeqSetInternal(1 + length(value))
            vars = MOI.VectorOfVariables([
                MOI.VariableIndex(key),
                [MOI.VariableIndex(vidx) for vidx in value]...,
            ])
            internals = ConstraintInternals(
                length(com.constraints) + 1,
                vars,
                set,
                Int[v.value for v in vars.variables],
            )

            constraint = init_constraint_struct(com, set, internals)
            add_constraint!(com, constraint)
            push!(added_constraint_idxs, length(com.constraints))

            # deactivate the single constraints
            for cidx in rhs_cons[key]
                com.constraints[cidx].is_deactivated = true
            end
        end
    end

    return added_constraint_idxs
end

"""
    simplify_not_equal_to_cliques(com)

Combine simple `a != b` constraints to `alldifferent` constraints using maximal cliques
"""
function simplify_not_equal_to_cliques(com)
    added_constraint_idxs = Int[]
    simple_not_equal_constraints = Int[]
    # one node per variable which is part of != constraint
    nodes = Set()
    # check for simple != constraints like `a != b`
    for constraint_idx in 1:length(com.constraints)
        constraint = com.constraints[constraint_idx]

        if isa(constraint.set, CPE.DifferentFrom) && length(constraint.indices) == 2
            if constraint.fct.terms[1].coefficient ==
               -constraint.fct.terms[2].coefficient &&
               abs(constraint.fct.terms[1].coefficient) == 1

                if constraint.fct.constant == 0 && constraint.set.value == 0
                    push!(simple_not_equal_constraints, constraint_idx)
                    for vidx in constraint.indices
                        push!(nodes, vidx)
                    end
                end
            end
        end
    end

    g = SimpleGraph(length(nodes))
    variables_to_constraint = Dict{Tuple{Int,Int},Int}()
    for constraint_idx in simple_not_equal_constraints
        constraint = com.constraints[constraint_idx]
        add_edge!(g, constraint.indices[1], constraint.indices[2])
        variables_to_constraint[(constraint.indices[1], constraint.indices[2])] =
            constraint_idx
    end
    cliques = maximal_cliques(g)
    for clique in cliques
        clique_size = length(clique)
        # use cliques if they contain at least 4 nodes
        if clique_size >= 4
            # add new all different constraint
            set = CPE.AllDifferent(clique_size)
            vars = MOI.VectorOfVariables([MOI.VariableIndex(vidx) for vidx in clique])
            internals = ConstraintInternals(
                length(com.constraints) + 1,
                vars,
                set,
                Int[v.value for v in vars.variables],
            )

            constraint = init_constraint_struct(com, set, internals)
            add_constraint!(com, constraint)
            push!(added_constraint_idxs, length(com.constraints))

            # deactivate old constraints
            for first_idx in 1:clique_size
                for second_idx in (first_idx + 1):clique_size
                    constraint_idx = get(
                        variables_to_constraint,
                        (clique[first_idx], clique[second_idx]),
                        0,
                    )
                    if constraint_idx != 0
                        com.constraints[constraint_idx].is_deactivated = true
                    end
                end
            end
        end
    end

    return added_constraint_idxs
end

"""
    simplify_all_different_and_equal_to(com)

Several functions to combine all_different constraints and equal to constraints
to add new equal to constraints
"""
function simplify_all_different_and_equal_to(com)
    added_constraint_idxs = Int[]
    # for each all_different constraint
    # which has an implicit sum constraint
    # check which sum constraints are completely inside all different
    # which are partially inside
    # compute inside sum and total sum
    n_constraints_before = length(com.constraints)
    for constraint_idx in 1:length(com.constraints)
        constraint = com.constraints[constraint_idx]

        if isa(constraint.set, CPE.AllDifferent)
            # check that the all different constraint uses every value
            if length(constraint.indices) == length(constraint.pvals)
                append!(
                    added_constraint_idxs,
                    simplify_all_different_inner_equal_to(com, constraint),
                )
                append!(
                    added_constraint_idxs,
                    simplify_all_different_outer_equal_to(
                        com,
                        constraint,
                        n_constraints_before,
                    ),
                )
            end
        end
    end
    return added_constraint_idxs
end

"""
    simplify_all_different_inner_equal_to(com, constraint::AllDifferentConstraint)

Add new sum constraint inside an alldifferent constraint to get restrict some variables more.
This function checks all sum constraints which are completely inside of the alldifferent constraint.

# Example
- [a,b,c,d,e] is the `AllDifferentConstraint`
- [a,b,c] is a sum constraint
- Add a new sum constraint for [d,e]
"""
function simplify_all_different_inner_equal_to(com, constraint::AllDifferentConstraint)
    @assert length(constraint.indices) == length(constraint.pvals)
    added_constraint_idxs = Int[]
    # make sure that we don't use an index more than once
    used_indices = falses(length(com.search_space))
    # we don't try to maximize the sum here so one might be able to use
    # a different constraint to start with to get a better sum.
    all_diff_sum = sum(constraint.pvals)
    in_sum = 0
    found_possible_constraint = false
    outside_indices = constraint.indices
    # go over all constraints that are completely inside alldifferent
    for sc_idx in constraint.sub_constraint_idxs
        sub_constraint = com.constraints[sc_idx]
        # if it's an equal_to constraint
        if isa(sub_constraint.fct, SAF) && isa(sub_constraint.set, MOI.EqualTo)
            # the coefficients must be all 1
            if all(t.coefficient == 1 for t in sub_constraint.fct.terms)
                # only use constraint if none of the indices was already used for
                # computing the inner sum
                if !any(used_indices[sub_constraint.indices])
                    used_indices[sub_constraint.indices] .= true
                    # compute sum inside all sum constraints
                    found_possible_constraint = true
                    in_sum += sub_constraint.set.value - sub_constraint.fct.constant
                    # for sum which are in alldifferent but not in sum constraints
                    outside_indices = setdiff(outside_indices, sub_constraint.indices)
                end
            end
        end
    end
    if found_possible_constraint && length(outside_indices) <= 4
        constraint_idx = length(com.constraints) + 1
        # all_diff_sum is Int and in_sum must be as well as all variables are Int and coefficients 1
        lc = LinearConstraint(
            constraint_idx,
            outside_indices,
            ones(Int, length(outside_indices)),
            0,
            MOI.EqualTo{Int}(all_diff_sum - in_sum),
        )
        add_constraint!(com, lc)
        push!(added_constraint_idxs, constraint_idx)
    end
    return added_constraint_idxs
end

"""
    simplify_all_different_outer_equal_to(com, constraint::AllDifferentConstraint)

check if several sum constraints completely fill the sum constraint
if this is the case we can use the sum of the all different constraint
to get useful information about the sum of variables outside the
all different constraint but inside one of the sum constraints

# Example
[a,b,c,d] where [a,b,c] are in alldifferent and sum is 6
[a,b] in sum constraint == 3
[c,d] in sum constraint == 9
=> 6+d == 12 => d=6
"""
function simplify_all_different_outer_equal_to(
    com,
    constraint::AllDifferentConstraint,
    n_constraints_before,
)
    @assert length(constraint.indices) == length(constraint.pvals)
    add_sum_constraint = true
    all_diff_sum = sum(constraint.pvals)
    added_constraint_idxs = []

    total_sum = 0
    outside_indices = Int[]
    cons_indices_dict = arr2dict(constraint.indices)
    for variable_idx in keys(cons_indices_dict)
        found_sum_constraint = false
        for sub_constraint_idx in com.subscription[variable_idx]
            # don't mess with constraints added later on
            if sub_constraint_idx > n_constraints_before
                continue
            end
            sub_constraint = com.constraints[sub_constraint_idx]
            # it must be an equal constraint and all coefficients must be 1 otherwise we can't add a constraint
            if isa(sub_constraint.fct, SAF) && isa(sub_constraint.set, MOI.EqualTo)
                if all(t.coefficient == 1 for t in sub_constraint.fct.terms)
                    found_sum_constraint = true
                    total_sum += sub_constraint.set.value - sub_constraint.fct.constant
                    all_inside = true
                    for sub_variable_idx in sub_constraint.indices
                        if !haskey(cons_indices_dict, sub_variable_idx)
                            all_inside = false
                            push!(outside_indices, sub_variable_idx)
                        else
                            delete!(cons_indices_dict, sub_variable_idx)
                        end
                    end
                    break
                end
            end
        end
        if !found_sum_constraint
            add_sum_constraint = false
            break
        end
    end

    # make sure that there are not too many outside indices
    if add_sum_constraint && length(outside_indices) <= 4
        constraint_idx = length(com.constraints) + 1
        lc = LinearConstraint(
            constraint_idx,
            outside_indices,
            ones(Int, length(outside_indices)),
            0,
            MOI.EqualTo{Int}(total_sum - all_diff_sum),
        )
        add_constraint!(com, lc)
        push!(added_constraint_idxs, constraint_idx)
    end
    return added_constraint_idxs
end

"""
    recompute_subscriptions(com)

Reset all subscriptions and recomputes them to only point to constraints that aren't deactivated
"""
function recompute_subscriptions(com)
    for variable in com.search_space
        empty!(com.subscription[variable.idx])
    end
    for constraint in com.constraints
        constraint.is_deactivated && continue
        for vidx in constraint.indices
            push!(com.subscription[vidx], constraint.idx)
        end
    end
end


text:
using ForwardDiff

"""
Evaluates the approximate contouring error ϵc = sin(Φ)(X - X_ref(s)) - cos(Φ)(Y - Y_ref(s)). Argument
P should represent the positional states of the systems (X, Y, s) corresponding to the position of the vehicle
(X, Y) and the arc length of the (approximate) projection point of the vehicle onto the reference path.
"""
function approx_contouring_error(P, path::SplinePath, spline_idx::Integer)
    X = P[1]; Y = P[2]; s = P[3]

    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    x_ref = spline_x(s, path, spline_idx)
    y_ref = spline_y(s, path, spline_idx)
    Φ = heading(s, path, spline_idx)
    sin(Φ) * (X - x_ref) - cos(Φ) * (Y - y_ref)
end

"""
Evaluates the approximate lag error ϵl = -cos(Φ)(X - X_ref(s)) - sin(Φ)(Y - Y_ref(s)). Argument
P should represent the positional states of the systems (X, Y, s) corresponding to the position of the vehicle
(X, Y) and the arc length of the (approximate) projection point of the vehicle onto the reference path.
"""
function approx_lag_error(P, path::SplinePath, spline_idx::Integer)
    X = P[1]; Y = P[2]; s = P[3]

    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    x_ref = spline_x(s, path, spline_idx)
    y_ref = spline_y(s, path, spline_idx)
    Φ = heading(s, path, spline_idx)
    -cos(Φ) * (X - x_ref) - sin(Φ) * (Y - y_ref)
end

"""
Returns [dϵc/dX, dϵc/dY, dϵc/ds].
"""
function contouring_error_gradient(P0::Array{Float64,1}, path::SplinePath, spline_idx::Integer)
    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    ForwardDiff.gradient(P -> approx_contouring_error(P, path, spline_idx), P0)
end

"""
Returns [dϵl/dX, dϵl/dY, dϵl/ds].
"""
function lag_error_gradient(P0::Array{Float64,1}, path::SplinePath, spline_idx::Integer)
    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    ForwardDiff.gradient(P -> approx_lag_error(P, path, spline_idx), P0)
end

#=
NOTE regarding the following two functions:
Define the tracking error cost at some point (X, Y, s) as (qc * ϵc(X, Y, s)^2 + ql * ϵl(X, Y, s)^2). As this is
nonlinear, approximate this cost by linearizing ϵc and ϵl individually about some point P0 = (X0, Y0, s0) and
plugging into the tracking error cost expression. The approximate tracking cost may then be arranged as the
sum of a linear term, cᵀP, and a quadratic term, PᵀQP.
=#

"""
Returns the coefficient vector, c, of the linear term described above.
"""
function tracking_linear_term(P0::Array{Float64,1}, qc::Float64, ql::Float64, path::SplinePath, spline_idx::Integer)
    X0 = P0[1]; Y0 = P0[2]; s0 = P0[3]

    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    ϵc0 = approx_contouring_error(P0, path, spline_idx)
    ϵl0 = approx_lag_error(P0, path, spline_idx)
    dϵc_dx, dϵc_dy, dϵc_ds = contouring_error_gradient(P0, path, spline_idx)
    dϵl_dx, dϵl_dy, dϵl_ds = lag_error_gradient(P0, path, spline_idx)

    cc = zeros(3)
    cc[1] = 2*qc*ϵc0*dϵc_dx - 2*qc*X0*dϵc_dx^2 - 2*qc*Y0*dϵc_dx*dϵc_dy - 2*qc*s0*dϵc_dx*dϵc_ds
    cc[2] = 2*qc*ϵc0*dϵc_dy - 2*qc*Y0*dϵc_dy^2 - 2*qc*X0*dϵc_dx*dϵc_dy - 2*qc*s0*dϵc_dy*dϵc_ds
    cc[3] = 2*qc*ϵc0*dϵc_ds - 2*qc*s0*dϵc_ds^2 - 2*qc*X0*dϵc_dx*dϵc_ds - 2*qc*Y0*dϵc_dy*dϵc_ds

    cl = zeros(3)
    cl[1] = 2*ql*ϵl0*dϵl_dx - 2*ql*X0*dϵl_dx^2 - 2*ql*Y0*dϵl_dx*dϵl_dy - 2*ql*s0*dϵl_dx*dϵl_ds
    cl[2] = 2*ql*ϵl0*dϵl_dy - 2*ql*Y0*dϵl_dy^2 - 2*ql*X0*dϵl_dx*dϵl_dy - 2*ql*s0*dϵl_dy*dϵl_ds
    cl[3] = 2*ql*ϵl0*dϵl_ds - 2*ql*s0*dϵl_ds^2 - 2*ql*X0*dϵl_dx*dϵl_ds - 2*ql*Y0*dϵl_dy*dϵl_ds

    c = cc + cl
end

"""
Returns the coefficient matrix, Q, of the quadratic term described above.
"""
function tracking_quadratic_term(P0::Array{Float64,1}, qc::Float64, ql::Float64, path::SplinePath, spline_idx::Integer)
    X0 = P0[1]; Y0 = P0[2]; s0 = P0[3]
    ϵc0 = approx_contouring_error(P0, path, spline_idx)
    ϵl0 = approx_lag_error(P0, path, spline_idx)
    dϵc_dx, dϵc_dy, dϵc_ds = contouring_error_gradient(P0, path, spline_idx)
    dϵl_dx, dϵl_dy, dϵl_ds = lag_error_gradient(P0, path, spline_idx)

    Qc = zeros(3, 3)
    a = qc*dϵc_dx^2
    b = d = (2*qc*dϵc_dx*dϵc_dy)/2
    c = g = (2*qc*dϵc_dx*dϵc_ds)/2
    e = qc*dϵc_dy^2
    f = h = (2*qc*dϵc_dy*dϵc_ds)/2
    i = qc*dϵc_ds^2
    Qc = [a b c; d e f; g h i]

    Ql = zeros(3, 3)
    a = ql*dϵl_dx^2
    b = d = (2*ql*dϵl_dx*dϵl_dy)/2
    c = g = (2*ql*dϵl_dx*dϵl_ds)/2
    e = ql*dϵl_dy^2
    f = h = (2*ql*dϵl_dy*dϵl_ds)/2
    i = ql*dϵl_ds^2
    Ql = [a b c; d e f; g h i]

    Q = Qc + Ql
end

function tracking_cost_matrices!(qs, vals)
    S_q = vals.k_c + vals.n_modes * (vals.N - vals.k_c)

    for k in 1:S_q
        P0 = [qs[1, k], qs[2, k], qs[5, k]] #X, Y, s of the initial guess
        spline_idx = find_spline_interval(P0[3], vals.path)
        if k in end_horizon_idces(vals)
            ck = tracking_linear_term(P0, vals.qcterm, vals.ql, vals.path, spline_idx)
            Γk = tracking_quadratic_term(P0, vals.qcterm, vals.ql, vals.path, spline_idx) # TODO: fix function signature
        else
            ck = tracking_linear_term(P0, vals.qc, vals.ql, vals.path, spline_idx)
            Γk = tracking_quadratic_term(P0, vals.qc, vals.ql, vals.path, spline_idx) # TODO: fix function signature
        end
        vals.cs[k] = ck
        vals.Γs[k] = Γk
    end

    vals.c = vcat(vals.cs...)
    vals.Γ = Array(blockdiag(sparse.(vals.Γs)...))
end

ext: jl

max_stars_repo_name: StanfordASL/MATS

max_issues_repo_name: StanfordASL/MATS

max_forks_repo_name: StanfordASL/MATS


text:
#!/usr/bin/julia
#
# Implementation of split step propagation method.
#
using Plots: Animation, plot, frame, gif;

export propagate, find_groundstate;

"""
    gen_opr(H::SeparableHamiltonian, par::SimParams)

Generate kinetic and potential operators in momentum and configuration space
(respectively) assuming that Strang splitting will be used (with the half
exponential factor on the potential term).

Assumes ħ = 1 if values of type `Unitful.Quantity` are not passed.
"""
function gen_opr(H::SeparableHamiltonian, par::SimParams)
  # Compensate for factor of -1 introduced by imaginary time propagation.
  if (imag(par.Δt) != 0) fac=-1
  else fac = 1;
  end

  # This version assumes ħ = 1 and that you'll resolve the unit conversions
  # yourself.
  return (exp.(ustrip.(-im*H.T*par.Δt*fac)), exp.(ustrip.(-0.5im*H.V*par.Δt*fac)));
end
function gen_opr(H::SeparableHamiltonian{T1}, par::SimParams{T2}) where
                 {N1<:Number,T1<:Unitful.Quantity{N1,Unitful.𝐌*Unitful.𝐋^2/Unitful.𝐓^2},
                  N2<:Number, T2<:Unitful.Quantity{N2,Unitful.𝐓}}
  # Compensate for factor of -1 introduced by imaginary time propagation.
  if (imag(par.Δt) != 0) fac=-1
  else fac = 1;
  end

  return (exp.(-im*H.T*par.Δt*fac/UnitfulAtomicHarmonic.ħ_u),
          exp.(-0.5im*H.V*par.Δt*fac/UnitfulAtomicHarmonic.ħ_u));
end

"""
    propagate(sim::SplitStep,
              xgrid::G=Grid(10, 0.01),
              kgrid::UniformGrid{M,1}=Grid(xgrid),
              par::SimParams=SimParams(0.05, 100, 1e-8),
              H::SeparableHamiltonian=SeparableHamiltonian(xgrid, kgrid, 0.5*xgrid.^2),
              ψ::WaveFunction{T,1,C,G}=WaveFunction(xgrid, exp.(-xgrid.^2/2));
              debug::Bool = true) where
        {T,C<:ConfigurationSpace,M<:MomentumSpace,G<:UniformGrid{C,1}}

Propagate the given wave function one time step using the split-step method.

# Arguments
- `sim::SplitStep`: the propagation method. Primarily used for dispatch, but can
  contain method specific simulation parameters.
- `xgrid::UniformGrid{ConfigurationSpace,1}`: grid of configuration space to evolve over.
- `kgrid::UniformGrid{M,1}`: 1D discretised momentum space.
- `par::SimParams`: Simulation parameters common to all simulation methods.
- `H::SeparableHamiltonian`: Hamiltonian to evolve system with respect to.
- `ψ::WaveFunction{T,1,C,G}`: 1D configuration space wave function to evolve.
- `debug::Bool`: If additional debug messages should be shown.
"""
function propagate(sim::SplitStep,
                   xgrid::G=Grid(10, 0.01),
                   kgrid::UniformGrid{M,1}=Grid(xgrid),
                   par::SimParams=SimParams(0.05, 100, 1e-8),
                   H::SeparableHamiltonian=SeparableHamiltonian(xgrid, kgrid, 0.5*xgrid.^2),
                   ψ::WaveFunction{T,1,C,G}=WaveFunction(xgrid, exp.(-xgrid.^2/2));
                   debug::Bool = false) where
        {T,C<:ConfigurationSpace,M<:MomentumSpace,G<:UniformGrid{C,1}}
  # Construct evolution operators from the passed Hamiltonian.
  (opT, opV) = gen_opr(H, par);

  # Work on a copy of the wave function passed with all units stripped.
  #ψ = WaveFunction{Complex{Float64},1,C,G}(ustrip.(ψi), ustrip(ψi.Δ));

  # Half-step in real space.
  ψ .*= opV;

  # FFT to momentum space.
  ϕ = transform(ψ, kgrid);

  # Full step in momentum space.
  ϕ .*= opT;

  # FFT back to position space.
  ψ = transform(ϕ, xgrid);

  # Half-step in real space.
  ψ .*= opV;

  # Renormalise wave function if running in imaginary time.
  if (ustrip(imag(par.Δt)) != 0)
    if debug println("Pre-normalisation norm was ", norm(ψ)); end
    normalise!(ψ);
  end

  # Return wave function of the same type as that passed.
  #return WaveFunction{T,1,C,G}(ψ, ψ.Δ);
  return ψ;
end

"""
    find_groundstate(sim::SplitStep,
                     xgrid::G,
                     kgrid::UniformGrid{M,1},
                     H::SeparableHamiltonian{C},
                     ψi::WaveFunction{T,1,C,G}=WaveFunction(xgrid, exp.(-xgrid.^2/2)),
                     par::SimParams=SimParams(0.1im, 250, 1e-8);
                     debug::Bool = false,
                     debugpath::String = "simrevamp/dev",
                     printstats::Bool = true) where
        {T,C<:ConfigurationSpace,M<:MomentumSpace,G<:UniformGrid{C,1}}

Determine the lowest energy eigenstate of the given potential.
"""
function find_groundstate(sim::SplitStep,
                         xgrid::G,
                         kgrid::UniformGrid{M,1},
                         H::SeparableHamiltonian{K,BT,B,C},
                         ψi::WaveFunction{T,1,C,G}=WaveFunction(xgrid, exp.(-xgrid.^2/2)),
                         par::SimParams=SimParams(0.1im, 250, 1e-8);
                         debug::Bool = false,
                         debugpath::String = "simrevamp/dev",
                         debugpltopts::SplitStepPlotOptions=SplitStepPlotOptions(true, 0.1),
                         printstats::Bool = true) where
        {T,K,BT,B,C<:ConfigurationSpace,M<:MomentumSpace,G<:UniformGrid{C,1}}

  if real(par.Δt) > 0*unit(par.Δt)
    throw(DomainError(par.Δt, "time step must be imaginary."));
  end

  # Create animation for convergence process if debugging.
  if debug
    anim = Animation();
    plt = SimPlot();
  end

  # Allow future expansion for determining multiple eigenstates.
  states = 1;

  # Define WaveBasis to store result in and make copy of passed WaveFunction for
  # operating on.
  Φ = WaveBasis{T}(xgrid, states);
  ψ = deepcopy(ψi);

  for i ∈ 1:states
    # Define vector to store wave function energy at each propagation step.
    E = zeros(typeof(1.0*unit(K)),par.steps+1);
    E[1] = energy(ψ, H);

    for j ∈ 1:par.steps
      if (debug && j%25 == 0) println("Iteration ", j); end

      ψ[:] = propagate(sim, xgrid, kgrid, par, H, ψ);
      E[j+1] = energy(ψ, H);

      if debug
        plot(plt, debugpltopts, xgrid, ψ, H,
            title=string("Step ", j, ", E = ", round(typeof(E[j+1]), E[j+1], sigdigits=6)));
        frame(anim);
      end

      # Stop propagation if energy has converged.
      if abs(E[j+1] - E[j]) < par.ϵ
        if printstats
          println("Converged to eigenstate ", i, " in ", j, " steps.");
          println("Energy found was ", round(typeof(E[j+1]), E[j+1], sigdigits=6));
        end
        break;
      end
    end
    normalise!(ψ);
    Φ[:,i] = ψ;
    #ψ[:] = orthog(ψi, Φ, i);
  end

  # Plot debug animation if requested.
  if debug
    check_path(debugpath);
    gif(anim, string(debugpath, gen_filename(par, xgrid, SplitStep(), debugpltopts), ".gif"), fps=16);
  end

  return Φ;
end


text:
using Test
using TightBindingLattice

@testset "OrthoCube" begin

    @testset "failures" begin
        @test_throws DimensionMismatch OrthoCube([1 2 3; 4 5 6])
        @test_throws ArgumentError OrthoCube([1 1; 1 1])
    end

    @testset "orthogonal" begin
        orthocube = OrthoCube([3 0; 0 3])
        @test dimension(orthocube) == 2
        @test isequiv(orthocube, OrthoCube([3 0; 3 3]))
        @test !isequiv(orthocube, OrthoCube([3 0; 1 3]))
        @test !isequiv(orthocube, OrthoCube([6 0; 0 6]))
        @test !isequiv(OrthoCube([6 0; 0 6]), orthocube)

        for i in 0:6, j in 0:6
            @test orthocube.wrap([i, j]) == ([i÷3, j÷3], [i%3, j%3])
        end
        let r = hcat([[i,j] for i in 0:6 for j in 0:6]...),
            R = hcat([[i÷3,j÷3] for i in 0:6 for j in 0:6]...),
            ρ = hcat([[i%3,j%3] for i in 0:6 for j in 0:6]...)
            @test orthocube.wrap(r) == (R, ρ)
        end
    end

    @testset "nonorthogonal" begin
        orthocube = OrthoCube([2 1; 1 2])
        @test orthocube.inverse_shape_matrix == [2//3 -1//3; -1//3 2//3]
        @test dimension(orthocube) == 2
        @test orthocube.wrap([0,0]) == ([0,  0], [0, 0])
        @test orthocube.wrap([1,0]) == ([0, -1], [2, 2])
        @test orthocube.wrap([2,0]) == ([1, -1], [1, 1])
        @test orthocube.wrap([3,0]) == ([2, -1], [0, 0])
    end

    @testset "generator" begin
        ortho = OrthoCube([4 0; 0 3])
        generate_coordinates(ortho, [ 1 4; -1 -3])
        generate_coordinates(ortho, [ 1 0; 0 1])
        @test_throws DimensionMismatch generate_coordinates(ortho, [ 1 0 0; 0 1 0]) # not square
        @test_throws ArgumentError generate_coordinates(ortho, [ 2 0; 0 1]) # unimodular
        @test_throws ArgumentError generate_coordinates(ortho, [ 1 0; 1 1])
        # c2 = generate_coordinates(ortho, [1 0; 0 1])
    end

end


text:
# when Tmatrix size grows, I can see error in inverting Q_matrix
using Tmatrix
using DataStructures

# inputs
rx, rz = 1e-6, 1.3e-6
n_max = 15
k1_r = 1e7; k1_i = 0.0
k2_r = 1.5e7; k2_i = 1e3

k1 = Complex(k1_r, k1_i)
k2 = Complex(k2_r, k2_i)

T = SortedDict()

kk = 10
for nn = 2 : n_max
    # calculating Tmatrix for different sizes
    println()
    @time T[nn] = calculate_Tmatrix_for_spheroid(rx, rz, nn, k1, k2; n_θ_points=100, rotationally_symmetric=false, symmetric_about_plane_perpendicular_z=false, BigFloat_precision=nothing);
    println("n_max = $nn, matrix to be inverted has size: $(size(T[nn]))")
    println("max error compared with n_max = 2: $(max(abs.(T[nn][1:kk ,1:kk] - T[2][1:kk ,1:kk])...))")
    display(T[nn][1:kk, 1:kk])
end


text:
struct StratumProxyServer
    left::StratumEndpoint
    right::StratumEndpoint

    err_handler::Union{Nothing,Function}

    left_msg_queue::Channel{Any}
    right_msg_queue::Channel{Any}

    status::Symbol

    left_task::Union{Nothing,Task}
    right_task::Union{Nothing,Task}
end
StratumProxyServer(left, right, err_handler = nothing) = StratumProxyServer(left, right, err_handler, Channel{Any}(Inf), Channel{Any}(Inf), :idle, nothing, nothing)

function Base.run(x::StratumProxyServer)
    x.status == :idle || error("Proxy is not idle.")
    x.task = @async try
        try
            while true
                msg = get_next_message(x.from)
                put!(x.msg_queue, msg)
                put_message!(x.to, msg)
            end
        finally
            close(x.msg_queue)
        end
    catch err
        bt = catch_backtrace()
        if x.err_handler !== nothing
            x.err_handler(err, bt)
        else
            Base.display_error(stderr, err, bt)
        end
    end
end

function proxy(proxy::StratumProxyServer, from::StratumEndpoint, from_msg_queue::Channel, to::StratumEndpoint)
    try
        while true
            msg = get_next_message(from)
            # put!(from_msg_queue,
            # put_message!(to, msg)
        end
    catch err
        # TODO handle InvalidStateException?
        bt = catch_backtrace()
        if proxy.err_handler !== nothing
            proxy.err_handler(err, bt)
        else
            Base.display_error(stderr, err, bt)
        end
    end
end


text:
module Utils
export report
using Gadfly
using DataFrames
using Printf
using StatsBase
using Statistics

set_default_plot_size(16cm, 6cm)

function report(result::DataFrame, layouts...; benchfield::Symbol)
    result = copy(result)
    result[!, :label] = (x -> @sprintf "%.2f" x).(result[!, benchfield])
    means = []

    case_names = unique(result.case)
    case_names_with_unit = Dict()
    for case_name in case_names
        idx = result[!, :case] .== case_name
        tmp = result[idx, benchfield]
        minval = minimum(tmp)
        gmean = geomean(tmp)
        result[idx, benchfield] = tmp ./ (minval / 10.0)
        push!(means, gmean)
        minval_repr = @sprintf "%.2f" minval
        case_names_with_unit[case_name] = Symbol(case_name, "(1 unit=$minval_repr)")
    end

    casemean = DataFrame(case = case_names, geomean = means)
    benchmarks = join(result, casemean, on = :case)

    benchmarks[:, :case] = map(x -> case_names_with_unit[x], benchmarks[:, :case])
    casemean[:, :case] = map(x -> case_names_with_unit[x], casemean[:, :case])

    rename!(benchmarks, :implementation => :Implementation)

    plot(
        benchmarks,
        x = :Implementation,
        y = benchfield,
        label = :label,
        color = :Implementation,
        xgroup = :case,
        # Theme(bar_spacing=2mm),
        Guide.ylabel(nothing),
        Guide.xlabel(nothing),
        # Coord.cartesian(ymin=0, ymax=3),
        Scale.y_log10,
        # Geom.label(position=:above),
        # Stat.dodge(axis=:y),
        Geom.subplot_grid(Geom.label(position=:above), Geom.bar(position=:dodge)),
        # Scale.x_discrete,
        Scale.color_discrete(levels=unique(benchmarks[:Implementation])),
        layouts...,
    ), casemean[:, 1:2]

end
end


text:
# code for handling conjugacy, natural parameters, and conjugate updates

function check_conjugate{D}(n::RandomNode{D}, m::VBModel)
    is_conj = true
    for (f, s) in m.graph[n]
        ttype = Tuple{Type{typeof(f)}, Type{Val{s}}, Type{D.name.primary}}
        isentropy = isa(f, EntropyFactor)
        if !isentropy && !method_exists(natural_formula, ttype)
            is_conj = false
            break
        end
    end
    is_conj
end

function update!{D}(n::RandomNode{D}, m::VBModel, ::Type{Val{:conjugate}})
    # get natural parameter vectors for each factor
    hasnaturals = filter(x -> !isa(x[1], EntropyFactor), m.graph[n])
    messages = [naturals(f, n) for (f, s) in hasnaturals]

    # update each distribution in the array
    for idx in eachindex(n.data)
        # get all messages corresponding to this element
        this_messages = Any[msg[idx] for msg in messages]

        # sum respective tuple elements
        natpars = map(+, this_messages...)

        # convert natural parameters to Distributions.jl parameters
        n[idx] = D(naturals_to_params(natpars, D)...)
    end
end

function naturals(f::Factor, n::RandomNode)
    fsym = f.namemap[n.name]
    _naturals(f, Val{fsym}, n)
end

@generated function _naturals{N, S, D}(f::Factor{N}, fsym::Type{Val{S}}, n::RandomNode{D})
    vars = fieldnames(f)

    # get expression corresponding to the natural parameters
    nat_expr = natural_formula(f, Val{S}, D)

    quote
        # init array of natural parameter tuples
        # should have the same dimension as fsym
        # get the type of the naturals by calling the
        # function on the first element of the array
        nats = naturals(n[1])
        η_type = typeof(nats)
        η = Array{η_type}(size(n)...)
        for i in eachindex(η)
            η[i] = map(zero_like, nats)
        end

        @nloops $N i d -> 1:f.inds.maxvals[d] begin
            nats = @wrapvars $vars $nat_expr project f (@ntuple $N i)
            nat_tup = project_inds(S, f, (@ntuple $N i))
            η[nat_tup...] = add_nats(η[nat_tup...], nats)
        end
        η
    end
end

#=
The strategy for combining natural parameters is:
- If η and nats are both tuples, add elementwise
- If η and nats are arrays of the same size, add the tuples at each position elementwise
=#
@inline function add_nats{N}(η::NTuple{N}, nats::NTuple{N})
    map(+, η, nats)
end

@inline function add_nats{T <: NTuple, N}(η::Array{T, N}, nats::Array{T, N})
    map(add_nats, η, nats)
end

"""
Define natural parameters for a given factor type, variable within that
factor, and form of the conjugate distribution for that variable.
"""
macro defnaturals(factortype, varname, disttype, expr)
    nat_expr = Expr(:quote, expr)
    varsym = Expr(:quote, varname)
    ex = quote
        natural_formula{N}(::Type{$(factortype){N}}, ::Type{Val{$varsym}}, ::Type{$disttype}) = $nat_expr
    end
    esc(ex)
end


text:
using Yao, Yao.Blocks
using NLopt

include("maxcut_gw.jl")

HB(nbit::Int) = sum([put(nbit, i=>X) for i=1:nbit])
tb = TimeEvolution(HB(3), 0.1)
function HC(W::AbstractMatrix)
    nbit = size(W, 1)
    ab = add(nbit)
    for i=1:nbit,j=i+1:nbit
        if W[i,j] != 0
            push!(ab, 0.5*W[i,j]*repeat(nbit, Z, [i,j]))
        end
    end
    ab
end

function qaoa_circuit(W::AbstractMatrix, depth::Int; use_cache::Bool=false)
    nbit = size(W, 1)
    hb = HB(nbit)
    hc = HC(W)
	use_cache && (hb = hb |> cache; hc = hc |> cache)
    c = chain(nbit, [repeat(nbit, H, 1:nbit)])
    append!(c, [chain(nbit, [timeevolve(hc, 0.0, tol=1e-5), timeevolve(hb, 0.0, tol=1e-5)]) for i=1:depth])
end

function cobyla_optimize(circuit::MatrixBlock{N}, hc::MatrixBlock; niter::Int) where N
    function f(params, grad)
        reg = zero_state(N) |> dispatch!(circuit, params)
        loss = expect(hc, reg) |> real
        #println(loss)
        loss
    end
    opt = Opt(:LN_COBYLA, nparameters(circuit))
    min_objective!(opt, f)
    maxeval!(opt, niter)
    cost, params, info = optimize(opt, parameters(circuit))
    pl = zero_state(N) |> circuit |> probs
    cost, params, pl
end

using Random, Test
@testset "qaoa circuit" begin
	Random.seed!(2)
	nbit = 5
	W = [0 5 2 1 0;
		 5 0 3 2 0;
		 2 3 0 0 0;
		 1 2 0 0 4;
		 0 0 0 4 0]

	# the exact solution
	exact_cost, sets = goemansWilliamson(W)

	hc = HC(W)
	@test ishermitian(hc)
	# the actual loss is -<Hc> + sum(wij)/2
	@test -expect(hc, product_state(5, bmask(sets[1]...)))+sum(W)/4 ≈ exact_cost

	# build a QAOA circuit and start training.
	qc = qaoa_circuit(W, 20; use_cache=false)
	opt_pl = nothing
	opt_cost = Inf
	for i = 1:10
		dispatch!(qc, :random)
		cost, params, pl = cobyla_optimize(qc, hc; niter=2000)
		@show cost
		cost < opt_cost && (opt_pl = pl)
	end

	# check the correctness of result
	config = argmax(opt_pl)-1
	@test config in [bmask(set...) for set in sets]
end


text:
include("ranking_criteria.jl")

abstract type AbstractRanks end

"""
    struct MeasurementRanks

Type for the results of the measurement ranking.

Fields:
* `names::Vector{Symbol}`: Names of the measurements.
* `values::Vector{Float64}`: Values of the ranking, calculated according to the criterion.
* `criterion<:AbstractRankingCriterion`: Criterion that was used for ranking.

Constructors:

```julia
function MeasurementRanks(
    names::Vector{Symbol},
    values::Vector{Float64},
    criterion::AbstractRankingCriterion;
    order::Symbol = :values,
    rev::Bool = false
)
```
Keyword arguments:
* `order::Symbol = :values`: Specifies how the `names` and `values` are sorted. Further options are
`:names` for alphabetical sorting based on the measurement names or `:none` for keeping the initial order of measurements.
* `rev::Bool=false`: Switch to invert the order.
"""
struct MeasurementRanks <: AbstractRanks
    names::Vector{Symbol}
    values::Vector{Float64}
    criterion::AbstractRankingCriterion

    function MeasurementRanks(
        names::Vector{Symbol},
        values::Vector{<:Real},
        criterion::AbstractRankingCriterion;
        order::Symbol = :values,
        rev::Bool = false
    )
        sorted_names, sorted_values = sort_ranking_results(order, names, values, rev=rev)
        return new(sorted_names, sorted_values, criterion)
    end
end

"""
    struct UncertaintyRanks

Type for the results of the unceertainty ranking.

Fields:
* `names::Vector{Symbol}`: Names of the uncertainty categories.
* `values::Vector{Float64}`: Values of the ranking, calculated according to the criterion.
* `criterion<:AbstractRankingCriterion`: Criterion that was used for ranking.

Constructors:

```julia
function UncertaintyRanks(
    names::Vector{Symbol},
    values::Vector{Float64},
    criterion::AbstractRankingCriterion;
    order::Symbol = :values,
    rev::Bool = false
)
```
Keyword arguments:
* `order::Symbol = :values`: Specifies how the `names` and `values` are sorted. Further options are
`:names` for alphabetical sorting based on the uncertainty names or `:none` for keeping the initial order of uncertainty categories.
* `rev::Bool=false`: Switch to invert the order.
"""
struct UncertaintyRanks <: AbstractRanks
    names::Vector{Symbol}
    values::Vector{Float64}
    criterion::AbstractRankingCriterion

    function UncertaintyRanks(
        names::Vector{Symbol},
        values::Vector{<:Real},
        criterion::AbstractRankingCriterion;
        order::Symbol = :values,
        rev::Bool = false
    )
        sorted_names, sorted_values = sort_ranking_results(order, names, values, rev=rev)
        return new(sorted_names, sorted_values, criterion)
    end
end

"""
    rank_measurements(
        model::EFTfitterModel;
        sampling_algorithm::BAT.AbstractSamplingAlgorithm = MCMCSampling(),
        criterion = SmallestIntervalsSum(p=0.9),
        order = :values,
        rev = true
    )

Computes a ranking of the individual measurement in the `EFTfitterModel` based
on a certain ranking criterion by performing the sampling according to the specified `sampling_algorithm`.
Returns a `MeasurementRank` object.

By default, the summed width of all one-dimensional marginalized smallest intervals containing 90% of the posterior
probability is used as the ranking criterion.
The measurements are ranked by the relative increase of this value.

Available ranking criteria:

* `SmallestIntervalsSum(p=0.9)`: summed width of all one-dimensional marginalized smallest intervals containing p=90% of the posterior probability
* `SmallestInterval(key=:C1, p=0.9)`: width of the one-dimensional marginalized smallest interval of parameter `:C1` containing p=90% of the posterior probability.
"""
function rank_measurements(
    model::EFTfitterModel;
    sampling_algorithm::BAT.AbstractSamplingAlgorithm = MCMCSampling(),
    criterion = SumOfSmallestIntervals(),
    order = :values,
    rev = true
)
    @info "Begin ranking of measurements"
    original = criterion_value(criterion, model, sampling_algorithm=sampling_algorithm)

    deactivated_models, measurement_keys = measurement_models(model)
    values = [criterion_value(criterion, m, sampling_algorithm=sampling_algorithm) for m in deactivated_models]

    relative_increase = calculate_ranks(original, values)

    return MeasurementRanks(collect(measurement_keys), relative_increase, criterion, order=order, rev=rev)
end

"""
    rank_uncertainties(
        model::EFTfitterModel;
        sampling_algorithm::BAT.AbstractSamplingAlgorithm = MCMCSampling(),
        criterion = SmallestIntervalsSum(p=0.9),
        order = :values,
        rev = true
    )

Computes a ranking of the types of uncertainty in the `EFTfitterModel` based
on a certain ranking criterion by performing the sampling according to the specified `sampling_algorithm`.
Returns a `MeasurementRank` object.

By default, the summed width of all one-dimensional marginalized smallest intervals containing 90% of the posterior
probability is used as the ranking criterion.
The uncertainty types are ranked by the relative decrease of this value.

Available ranking criteria:

* `SmallestIntervalsSum(p=0.9)`: summed width of all one-dimensional marginalized smallest intervals containing p=90% of the posterior probability
* `SmallestInterval(key=:C1, p=0.9)`: width of the one-dimensional marginalized smallest interval of parameter `key=:C1` containing p=90% of the posterior probability
"""
function rank_uncertainties(
    model::EFTfitterModel;
    sampling_algorithm::BAT.AbstractSamplingAlgorithm = MCMCSampling(),
    criterion = SumOfSmallestIntervals(),
    order = :values,
    rev = true
)
    @info "Begin ranking of uncertainties"
    original = criterion_value(criterion, model, sampling_algorithm=sampling_algorithm)

    deactivated_models, uncertainty_keys = uncertainty_models(model)
    values = [criterion_value(criterion, m, sampling_algorithm=sampling_algorithm) for m in deactivated_models]
    relative_decrease = -1*calculate_ranks(original, values)

    return UncertaintyRanks(collect(uncertainty_keys), relative_decrease, criterion, order=order, rev=rev)
end

"""
    function measurement_models(model::EFTfitterModel)

Creates a `Vector` of `EFTfitterModel` where always one of the initially active measurements is deactivated at a time.
Returns a `Vector` of `EFTfitterModel` and a `Vector` of `Symbol` with the names of the currently deactivated measurement.
"""
function measurement_models(model::EFTfitterModel)
    m = model.measurements
    measurement_keys = collect(keys(m))

    if length(measurement_keys) < 2
         throw(ArgumentError("Only one measurement active. Ranking of measurements not possible."))
    end

    models = Vector{EFTfitterModel}(undef, length(m))
    for i in 1:length(m)
        meas = collect(values(model.measurements))
        meas[i] = deactivate_measurement(meas[i])
        meas_nt = namedtuple(measurement_keys, meas)
        models[i] = EFTfitterModel(model.parameters, meas_nt, model.correlations)
    end

    return models, measurement_keys
end

"""
    function uncertainty_models(model::EFTfitterModel)

Creates a `Vector` of `EFTfitterModel` where always one of the initially active uncertainty types is deactivated at a time.
Returns a `Vector` of `EFTfitterModel` and a `Vector` of `Symbol` with the names of the currently deactivated uncertainty types.
"""
function uncertainty_models(model::EFTfitterModel)
    cors = model.correlations
    uncertainty_keys = collect(keys(cors))

    if length(uncertainty_keys) < 2
         throw(ArgumentError("Only one type of uncertainties active. Ranking of uncertainties not possible."))
    end

    models = Vector{EFTfitterModel}(undef, length(cors))
    for i in 1:length(cors)
        cors = collect(values(model.correlations))
        cors[i] = deactivate_uncertainty(cors[i])
        cors_nt = namedtuple(uncertainty_keys, cors)
        models[i] = EFTfitterModel(model.parameters, model.measurements, cors_nt)
    end

    return models, uncertainty_keys
end

function deactivate_measurement(meas::Measurement)
    return Measurement(meas.observable, meas.value, meas.uncertainties, false)
end

function deactivate_uncertainty(cor::Correlation)
    return Correlation(cor.matrix, false)
end

"""
    function criterion_value(
        criterion::AbstractRankingCriterion,
        model::EFTfitterModel;
        sampling_algorithm::BAT.AbstractSamplingAlgorithm=MCMCSampling()
    )

Computes the value of the `criterion` for the `model` using the `sampling_algorithm`.

Available ranking criteria:

* `SmallestIntervalsSum(p=0.9)`: summed width of all one-dimensional marginalized smallest intervals containing p=90% of the posterior probability
* `SmallestInterval(key=:C1, p=0.9)`: width of the one-dimensional marginalized smallest interval of parameter `key=:C1` containing p=90% of the posterior probability
"""
function criterion_value(
    criterion::AbstractRankingCriterion,
    model::EFTfitterModel;
    sampling_algorithm::BAT.AbstractSamplingAlgorithm=MCMCSampling()
)
    posterior = PosteriorDensity(EFTfitterDensity(model), model.parameters)
    samples = BAT.bat_sample(posterior, sampling_algorithm).result
    criterion_value = apply_criterion(criterion, samples)
end

function calculate_ranks(original_value::Real, values::AbstractArray{<:Real})
    return relative_increase = (values.-original_value)./original_value
end

function calculate_ranks(original_values::AbstractArray{<:Real}, values::AbstractArray{<:AbstractArray{<:Real, 1}, 1})
    relative_increase = [sum((val.-original_values)./original_values) for val in values]
end

function sort_ranking_results(
    order::Symbol,
    names::Vector{Symbol},
    values::Vector{<:Real};
    rev=false
)
    if order == :names
        sorted = sortperm(names, rev=rev)
        return names[sorted], values[sorted]

    elseif order == :values
        sorted = sortperm(values, rev=rev)
        return names[sorted], values[sorted]

    elseif order == :none
        return names, values

    else
        return names, values
    end
end


text:
#=

  Linear regression

  This is actually a WebPPL example (example/linear_regression.wppl) but translated via my BLOG model.

  xs = [0,1,2,3]
  ys = [0,1,4,6]

  mu should be ~ 2 and b ~ 0.
  y4: y for x=4 should be 2*4+0 = 8
  y5: y for x=5 should be 2*5+0 = 10

  Summary Statistics
  parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

           m    2.0134    0.3960     0.0040    0.0088   1980.8480    1.0001      445.5349
           b   -0.2678    0.7118     0.0071    0.0135   2785.0095    0.9999      626.4079
       sigma    0.8463    0.4970     0.0050    0.0155    929.0047    0.9999      208.9529
          y4    7.7715    1.5278     0.0153    0.0314   2240.9116    1.0001      504.0287
          y5    9.7916    1.8747     0.0187    0.0401   2060.8303    1.0000      463.5246

  With 5 data points (x=[0,1,2,3,4] and y=[0,1,4,6,8]) with got a much better result:
  Summary Statistics
  parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

           m    2.0000    0.0001     0.0000    0.0000    684.0973    0.9999      170.6830
           b    0.0000    0.0004     0.0000    0.0000   1140.3205    1.0005      284.5111
       sigma    0.0002    0.0003     0.0000    0.0000    111.5388    1.0486       27.8290
          y4    8.0000    0.0005     0.0000    0.0000    713.2871    1.0046      177.9658
          y5   10.0000    0.0006     0.0000    0.0000    553.0126    1.0016      137.9772

  cf ~/blog/linear_regression2.blog
     ~/webppl/linear_regression2.wppl
=#

using Turing
include("jl_utils.jl")

@model function linear_regression2(xs,ys)
    n = length(xs)

    m ~ Normal(0,2)
    b ~ Normal(0,2)
    sigma ~ Gamma(1,1)

    xs ~ filldist(Uniform(0,14),n)

    for i in 1:n
        ys[i] ~ Normal(m*xs[i] + b, sigma)
    end

    # Prediction for x=4 and x=5
    y4 ~ Normal(m*4 + b, sigma)
    y5 ~ Normal(m*5 + b, sigma)

end

xs = [0,1,2,3]
ys = [0,1,4,6]
n = 4
xs = [i for i in 0:n]
ys = [2*x for x in xs]
println("xs: ", xs)
println("ys: ", ys)
model = linear_regression2(xs,ys)

# chns = sample(model, Prior(), 100_000)
# chns = sample(model, MH(), 100_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 100_000)
chns = sample(model, NUTS(), 10_000)

display(chns)

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
using DataFrames
using VegaLite
using CSV

const inputFile = "test.csv";

# Loading csv-File with Portfolio-Data
rawDataDf = CSV.File(
    inputFile,
    delim = ';',
    decimal = '.',
    missingstrings = ["", "0.0"],
) |> DataFrame!

show(rawDataDf);

plt = rawDataDf |> @vlplot(
        mark=:point,
        x=:profit,
        y=:max_dd,
        title="Title Of chart",
        width = 800,
        height = 400,
        color=:portfolio_Name
        )

plt |> display

save("C:\\myplot.png", plt)

max_stars_repo_name: RainerNeu/juliaExamplesWithJuliaPro

max_issues_repo_name: RainerNeu/juliaExamplesWithJuliaPro

max_forks_repo_name: RainerNeu/juliaExamplesWithJuliaPro


text:
module XMLTools

using LightXML, Statistics, LinearAlgebra

export collect_data,
    data_stats,
    remove_nans,
    replace_data!,
    introduce_sparsity!,
    get_corrs,
    complete_cases!,
    replace_newick!,
    replace_filelog!,
    replace_chainLength!,
    replace_logEvery!,
    remove_all_data!,
    remove_dates!,
    get_element_by_id,
    standardize_data!,
    replace_all_data!,
    get_trait_dimmension,
    add_note!,
    transfer_data!,
    get_newick,
    set_diffusion_precision!,
    set_rm_precision!,
    set_operator_weight!,
    setup_trait_validation!,
    multivariateDistributionLikelihood,
    leafTraitParameter,
    insert_element,
    get_likelihood,
    get_diffusion_precision,
    get_rm_precision,
    get_extension_precision,
    get_loadings_mat,
    get_if_precision,
    collect_dates,
    make_matrix_parameter

const TAXA = "taxa"
const TAXON = "taxon"
const ID = "id"
const NAME = "name"
const MISSING_VAL = "NA"
const NEWICK = "newick"
const MCMC = "mcmc"
const LOG = "log"
const FILENAME = "fileName"
const CHAINLENGTH = "chainLength"
const LOGEVERY = "logEvery"
const DATE = "date"
const OPERATORS = "operators"
const ATTR = "attr"
const MULTIVARIATE_DIFFUSION_MODEL = "multivariateDiffusionModel"
const PRECISION_MATRIX = "precisionMatrix"
const MATRIX_PARAMETER = "matrixParameter"
const PARAMETER = "parameter"
const VALUE = "value"
const SAMPLING_PRECISION = "samplingPrecision"
const REPEATED_MEASURES = "repeatedMeasuresModel"
const WEIGHT = "weight"
const BEAST = "beast"
const TRAIT_VALIDATION = "traitValidation"
const TRAIT_NAME = "traitName"
const INFERRED_TRAIT = "inferredTrait"
const USE_TREE_TRAITS = "useTreeTraits"
const TRAIT_DATA_LIKELIHOOD = "traitDataLikelihood"
const IDREF = "idref"
const TRAIT_PARAMETER = "traitParameter"
const FILELOG = "fileLog"
const MULT_DIST_LIKELIHOOD = "multivariateDistributionLikelihood"
const DISTRIBUTION = "distribution"
const MVM_DISTRIBUTION = "multivariateNormalDistributionModel"
const MEAN_PARAMETER = "meanParameter"
const PRECISION_PARAMETER = "precisionParameter"
const DATA = "data"
const LEAF_TRAIT_PARAMETER = "leafTraitParameter"
const TREE_MODEL = "treeModel"
const LIKELIHOOD = "likelihood"
const POSTERIOR = "posterior"
const INTEGRATED_FACTOR_MODEL = "integratedFactorModel"
const PRECISION = "precision"
const LOADINGS = "loadings"
const DIAGONAL_MATRIX = "DiagonalMatrix"
const DIMENSION = "dimension"
const DATE = "date"

function get_child_by_attribute(xel::XMLElement, attr::String, attr_value::String)
    children = child_elements(xel)
    for child in children
        if attribute(child, attr) == attr_value
            return child
        end
    end
    throw("ERROR: No children have attribute $attr with value \"$attr_value\".")
end

function get_element_by_id(el::XMLElement, id::String)
    children = child_elements(el)
    final = nothing
    for child in children
        if attribute(child, ID) == id
            final = child
            break
        else
            potential = get_element_by_id(child, id)
            if potential != nothing
                final = potential
                break
            end
        end
    end
    return final
end

function get_element_by_id(xdoc::XMLDocument, id::String)
    el = root(xdoc)
    return get_element_by_id(el, id)
end

function collect_data(xdoc::XMLDocument, trait_name::String)
    taxa_el = find_element(root(xdoc), TAXA)
    taxa_children = collect(child_elements(taxa_el))
    n = length(taxa_children)
    taxa = Vector{String}(undef, n)
    p = length(split(content(get_child_by_attribute(taxa_children[1], NAME,
            trait_name))))
    data = Matrix{Float64}(undef, n, p)
    for i = 1:n
        taxa[i] = attribute(taxa_children[i], ID)
        taxa_vals = split(content(get_child_by_attribute(taxa_children[i],
                NAME, trait_name)))
        for j = 1:p
            if taxa_vals[j] == MISSING_VAL
                data[i, j] = NaN
            else
                data[i, j] = parse(Float64, taxa_vals[j])
            end
        end
    end

    return taxa, data
end

function collect_data(path::String, trait_name::String)
    xdoc = parse_file(path)
    to_return = collect_data(xdoc, trait_name)
    free(xdoc)
    return to_return
end

function collect_dates(xdoc::XMLDocument)
    xroot = root(xdoc)
    taxa_el = find_element(xroot, TAXA)
    taxa_children = collect(child_elements(taxa_el))
    n = length(taxa_children)
    dates = zeros(n)
    taxa = Vector{String}(undef, n)
    for i = 1:n
        taxa[i] = attribute(taxa_children[i], ID)

        date_el = find_element(taxa_children[i], DATE)
        dates[i] = parse(Float64, attribute(date_el, VALUE))
    end
    return taxa, dates
end

function collect_dates(path::String)
    xdoc = parse_file(path)
    to_return = collect_dates(xdoc)
    free(xdoc)
    return to_return
end

function replace_data!(xdoc::XMLDocument, trait_name::String,
        replacement_data::Dict{String, Vector{Float64}})
    taxa_el = find_element(root(xdoc), TAXA)
    for taxon_el in child_elements(taxa_el)
        taxon = attribute(taxon_el, ID, required=true)
        trait_el = get_child_by_attribute(taxon_el, NAME, trait_name)
        set_content(trait_el, make_data_string(replacement_data[taxon]))
    end
end

function remove_all_data!(xdoc::XMLDocument)
    taxa_el = find_element(root(xdoc), TAXA)
    unlink(taxa_el)
    free(taxa_el)
end

function transfer_data!(src_doc::XMLDocument, dest_doc::XMLDocument)
    src_taxa = find_element(root(src_doc), TAXA)
    dest_taxa = find_element(root(dest_doc), TAXA)
    for node in child_nodes(dest_taxa)
        unlink(node)
        free(node)
    end
    for node in child_nodes(src_taxa)
        add_child(dest_taxa, node)
    end
end

function remove_dates!(xdoc::XMLDocument)
    taxa_el = find_element(root(xdoc), TAXA)
    for taxon in child_elements(taxa_el)
        date = find_element(taxon, DATE)
        unlink(date)
        free(date)
    end
end

function introduce_sparsity!(xdoc::XMLDocument, trait_name::String,
        sparsity::Vector{Float64})

    child_els = child_elements(find_element(root(xdoc), TAXA))
    n  = length(collect(child_els))
    int_sparsity = [Int(floor(x * n)) for x in sparsity]
    return introduce_sparsity!(xdoc, trait_name, int_sparsity)
end

function introduce_sparsity!(xdoc::XMLDocument, trait_name::String,
        sparsity::Vector{Int})

    taxa_el = find_element(root(xdoc), TAXA)
    taxa = collect(child_elements(taxa_el))
    n = length(taxa)
    removed_inds = get_removed_inds(sparsity, n)

    p = length(sparsity)
    removed_taxa = [Vector{String}(undef, sparsity[i]) for i = 1:p]
    for j = 1:p
        ind = 0
        for i in removed_inds[j]
            ind += 1
            taxon_el = taxa[i]
            trait_el = get_child_by_attribute(taxon_el, NAME, trait_name)
            data = split(content(trait_el))
            data[j] = MISSING_VAL
            set_content(trait_el, join(data, ' '))
            removed_taxa[j][ind] = attribute(taxon_el, ID, required = true)
        end
    end
    return removed_taxa
end

function get_removed_inds(num_removed::Vector{Int}, n::Int)
    p = length(num_removed)
    for i = 1:p
        @assert num_removed[i] <= n
    end
    removed_inds = [Vector{Int}(undef, num_removed[i]) for i = 1:p]
    for i = 1:p
        for j = 1:num_removed[i]
            ok = false
            while !ok
                draw = rand(1:n)
                if !(draw in removed_inds[i])
                    removed_inds[i][j] = draw
                    ok = true
                end
            end
        end
        sort!(removed_inds[i])
    end
    return removed_inds
end

function make_data_string(data::Vector{Float64}; missing_val = MISSING_VAL)
    p = length(data)
    s_data = Vector{String}(undef, p)
    for i = 1:p
        if isnan(data[i])
            s_data[i] = missing_val
        else
            s_data[i] = string(data[i])
        end
    end
    return join(s_data, " ")
end

function data_stats(data::Matrix{Float64})
    n, p = size(data)
    n_observed = zeros(Int, p)
    mean_sums = zeros(p)
    var_sums = zeros(p)
    for i = 1:n
        for j = 1:p
            if !isnan(data[i, j])
                n_observed[j] += 1
                mean_sums[j] += data[i, j]
                var_sums[j] += data[i, j]^2
            end
        end
    end

    for j = 1:p
        mean_sums[j] = mean_sums[j] / n_observed[j]
        var_sums[j] = var_sums[j] / n_observed[j] - mean_sums[j]^2
    end
    cormat = get_corrs(data)
    perc_missing = [1 - n_observed[i] / n for i = 1:p]
    return mean_sums, var_sums, cormat, perc_missing
end

function remove_nans(x::Vector{Float64})
    removed = Vector{Float64}()
    for n in x
        if !isnan(n)
            push!(removed, n)
        end
    end
    return removed
end

function remove_nans(x1::Vector{Float64}, x2::Vector{Float64})
    removed = Matrix{Float64}(undef, 0, 2)
    n = length(x1)
    for i = 1:n
        if !isnan(x1[i]) && !isnan(x2[i])
            removed = [removed; [x1[i] x2[i]]]
        end
    end
    return removed
end

function get_corrs(data::Matrix{Float64})
    n, p = size(data)
    cormat = zeros(p, p)
    for i = 1:n
        for j = i:p
            pairwise = remove_nans(data[:, i], data[:, j])
            cormat[i, j] = cor(pairwise[:, 1], pairwise[:, 2])
        end
    end
    return cormat
end

function complete_cases!(xdoc::XMLDocument, trait_name::String)
    taxa_el = find_element(root(xdoc), TAXA)
    removed_taxa = Vector{String}(undef, 0)
    for taxon_el in child_elements(taxa_el)
        trait_el = get_child_by_attribute(taxon_el, NAME, trait_name)
        data = split(content(trait_el))
        if MISSING_VAL in data
            unlink(taxon_el)
            push!(removed_taxa, attribute(taxon_el, ID, required = true))
        end
    end
    return removed_taxa
end

function replace_newick!(xdoc::XMLDocument, new_newick::AbstractString)
    newick_el = find_element(root(xdoc), NEWICK)
    strip(new_newick)
    strip(new_newick, ';')
    set_content(newick_el, "$new_newick;")
end

function get_newick(xdoc::XMLDocument)
    newick_el = find_element(root(xdoc), NEWICK)
    newick = content(newick_el)
    newick = strip(newick)
    # newick = strip(newick, ';')
    return String(newick)
end

function get_newick(path::String)
    xdoc = parse_file(path)
    newick = get_newick(xdoc)
    free(xdoc)
    return newick
end

function replace_filelog!(xdoc::XMLDocument, new_filename::String)
    mcmc_el = find_element(root(xdoc), MCMC)
    changed = false
    for element in child_elements(mcmc_el)
        if name(element) == LOG && has_attribute(element, FILENAME)
            set_attribute(element, FILENAME, new_filename)
            changed = true
            break
        end
    end
    @assert changed
end

function replace_chainLength!(xdoc::XMLDocument, new_chainLength::Int)
    mcmc_el = find_element(root(xdoc), MCMC)
    set_attribute(mcmc_el, CHAINLENGTH, string(new_chainLength))
end

function replace_logEvery!(xdoc::XMLDocument, new_logEvery::Int)
    mcmc_el = find_element(root(xdoc), MCMC)
    changed = false
    for element in child_elements(mcmc_el)
        if name(element) == LOG && has_attribute(element, LOGEVERY) &&
                has_attribute(element, FILENAME)
            set_attribute(element, LOGEVERY, string(new_logEvery))
            changed = true
            break
        end
    end
    @assert changed
end

function standardize_data!(xdoc::XMLDocument, trait::String;
        keep_zeros::Bool = false)
    taxa, data = collect_data(xdoc, trait)
    n, p = size(data)
    standardize_data!(data, keep_zeros = keep_zeros)
    d = Dict{String, Vector{Float64}}()
    for i = 1:n
        d[taxa[i]] = data[i, :]
    end
    replace_data!(xdoc, trait, d)
end

function standardize_data!(data::Matrix{Float64}; keep_zeros::Bool = false)
    n, p = size(data)
    means, vars = missing_MVs(data, keep_zeros = keep_zeros)
    for j = 1:p
        for i = 1:n
            if !(keep_zeros && data[i, j] == 0.0)
                data[i, j] = (data[i, j] - means[j]) / sqrt(vars[j])
            end
        end
    end
end

function missing_means(X::Array{Float64, 2}; keep_zeros::Bool = false)
    n, p = size(X)
    means = Vector{Float64}(undef, p)
    for j = 1:p
        t = 0.0
        N = 0
        for i = 1:n
            if !isnan(X[i, j])
                if !(keep_zeros && X[i, j] == 0.0)
                    t += X[i, j]
                    N += 1
                end
            end
        end
        means[j] = t / N
    end
    return means
end

function missing_vars(X::Matrix{Float64};
        means::Vector{Float64} = Vector{Float64}(undef, 0),
        keep_zeros::Bool = false)

    if length(means) == 0
        means = missing_means(X)
    end
    n, p = size(X)
    vars = Vector{Float64}(undef, p)
    for j = 1:p
        t = 0.0
        N = 0
        for i = 1:n
            if !isnan(X[i, j])
                if !(keep_zeros && X[i, j] == 0.0)
                    t += (X[i, j] - means[j])^2
                    N += 1
                end
            end
        end
        vars[j] = t / (N - 1)
    end
    return vars
end

function missing_MVs(X::Matrix{Float64}; keep_zeros::Bool = false)
    means = missing_means(X, keep_zeros = keep_zeros)
    vars = missing_vars(X, means = means, keep_zeros = keep_zeros)
    return means, vars
end

function replace_all_data!(xdoc::XMLDocument, taxa::Vector{String},
        data::Matrix{Float64};
        traits::Vector{String} = Vector{String}(undef, 0),
        all_name::String = "traits",
        dates::Vector{Float64} = Float64[])

    if length(dates) > 0
        error("not implemented")
    end
    n, p = size(data)
    use_traits = false
    if length(traits) != 0
        use_traits = true
    end
    @assert length(taxa) == n
    if use_traits
        @assert length(traits) == p
    end
    taxa_el = find_element(root(xdoc), TAXA)
    children = child_nodes(taxa_el)
    for child in children
        unlink(child)
        free(child)
    end
    all_traits = zeros(p)
    for i = 1:n

        new_taxon = new_element(TAXON)
        set_attribute(new_taxon, ID, taxa[i])
        if use_traits
            for j = 1:p
                new_attr = new_element(ATTR)
                set_attribute(new_attr, NAME, traits[j])
                set_content(new_attr, string(data[i, j]))
                add_child_level(new_taxon, new_attr, 3)
            end
        end
        for j = 1:p
            all_traits[j] = data[i, j]
        end
        all_attr = new_element(ATTR)
        set_attribute(all_attr, NAME, all_name)
        set_content(all_attr, make_data_string(all_traits))
        add_child_level(new_taxon, all_attr, 3)
        # add_text(taxa_el, "\t\t")
        add_child_level(taxa_el, new_taxon, 2)
        # @show has_children(new_taxon)
        # @show length(collect(child_elements(new_taxon)))
        # add_text(taxa_el, "\n")
    end
end

function replace_all_data!(xdoc::XMLDocument, taxa::Vector{String},
        data::Vector{Array{Float64, 2}}, traits::Vector{String};
        dates::Vector{Float64} = Float64[])

    use_dates = false
    if length(dates) > 0
        @assert length(dates) == length(taxa)
        use_dates = true
    end

    p = length(data)
    @assert length(traits) == p
    n = length(taxa)

    taxa_el = find_element(root(xdoc), TAXA)
    children = child_nodes(taxa_el)
    for child in children
        unlink(child)
        free(child)
    end
    all_traits = zeros(p)
    for i = 1:n
        new_taxon = new_element(TAXON)
        set_attribute(new_taxon, ID, taxa[i])
        if use_dates
            date = new_element(DATE)
            set_attribute(date, VALUE, dates[i])
            set_attribute(date, "direction", "forwards")
            set_attribute(date, "units", "years")
            add_child_level(new_taxon, date, 3)
        end
        for j = 1:p
            new_attr = new_element(ATTR)
            set_attribute(new_attr, NAME, traits[j])
            set_content(new_attr, make_data_string(data[j][i, :]))
            add_child_level(new_taxon, new_attr, 3)
        end
        add_child_level(taxa_el, new_taxon, 2)
    end
end

function get_trait_dimmension(xdoc::XMLDocument, trait::String)
    taxa_el = find_element(root(xdoc), TAXA)
    child = collect(child_elements(taxa_el))[1]
    trait_el = get_child_by_attribute(child, NAME, trait)
    p = length(split(content(trait_el)))
    return p
end

function add_child_level(parent_el::XMLElement, child_el::XMLElement,
        level::Int; line_below::Bool = false)
    add_text(parent_el, "\n")
    tabs = join(["\t" for i = 1:level])
    add_text(parent_el, tabs)
    # @show has_children(child_el)
    if length(collect(child_elements(child_el))) > 0
        add_text(child_el, "\n$tabs")
    end
    add_child(parent_el, child_el)
    if line_below
        add_text(parent_el, "\n")
    end
end

function add_note!(xdoc::XMLDocument, note::String)
    xroot = root(xdoc)
    comment_note = "<!--\n$note\n-->\n"
    t = new_textnode(comment_note)
    add_child(xroot, t)
end

function set_matrix_parameter!(xel::XMLElement, P::Matrix{Float64})
    p = size(P, 1)
    i = 0
    for child in child_elements(xel)
        i += 1
        set_attribute(child, VALUE, join(P[i, :], " "))
    end
    @assert i == p
end

function set_diffusion_precision!(xdoc::XMLDocument,
        P::Matrix{Float64})

    @assert isposdef(P)
    mp_el = get_diffusion_precision_el(xdoc)
    set_matrix_parameter!(mp_el, P)
end

function get_diffusion_precision(xdoc::XMLDocument)
    mp_el = get_diffusion_precision_el(xdoc)
    P = get_matrix_parameter(mp_el)
    return P
end

function get_diffusion_precision_el(xdoc::XMLDocument)
    diffusion_el = find_element(root(xdoc), MULTIVARIATE_DIFFUSION_MODEL)
    pmat_el = find_element(diffusion_el, PRECISION_MATRIX)
    mp_el = find_element(pmat_el, MATRIX_PARAMETER)
    if mp_el == nothing
        mp_el = find_element(pmat_el, DIAGONAL_MATRIX)
    end
    return mp_el
end

function set_rm_precision!(xdoc::XMLDocument, P::Matrix{Float64})

    @assert isposdef(P)
    mp_el = get_rm_precision_el(xdoc)
    set_matrix_parameter!(mp_el, P)
end

function get_rm_precision(xdoc::XMLDocument)
    mp_el = get_rm_precision_el(xdoc)
    return get_matrix_parameter(mp_el)
end

function get_rm_precision_el(xdoc::XMLDocument)
    rm_el = find_element(root(xdoc), REPEATED_MEASURES)
    pmat_el = find_element(rm_el, SAMPLING_PRECISION)
    mp_el = find_element(pmat_el, MATRIX_PARAMETER)
    return mp_el
end

function get_extension_precision(xdoc::XMLDocument)
    xroot = root(xdoc)
    rm_el = find_element(xroot, REPEATED_MEASURES)
    if_el = find_element(xroot, INTEGRATED_FACTOR_MODEL)
    if rm_el == nothing
        if if_el == nothing
            error("Could not find $(REPEATED_MEASURES) or $(INTEGRATED_FACTOR_MODEL) elements under $(name(xroot))")
        else
            return get_if_precision(xdoc) #TODO
        end
    else
        if if_el == nothing
            return get_rm_precision(xdoc)
        else
            error("XML contained both $(REPEATED_MEASURES) and $(INTEGRATED_FACTOR_MODEL) elements. Did not know which one to get precision from.")
        end
    end
end

function get_if_precision(xdoc::XMLDocument)
    if_el = find_element(root(xdoc), INTEGRATED_FACTOR_MODEL)
    prec_el = find_element(if_el, PRECISION)
    param_el = find_element(prec_el, PARAMETER)
    prec_vec = get_vector_parameter(param_el)
    return diagm(0 => prec_vec)
end

function get_loadings_mat(xdoc::XMLDocument)
    xroot = root(xdoc)
    if_el = find_element(xroot, INTEGRATED_FACTOR_MODEL)
    loadings_el = find_element(if_el, LOADINGS)
    mp_el = find_element(loadings_el, MATRIX_PARAMETER)
    if has_children(mp_el)
        return get_matrix_parameter(mp_el)
    else
        id = attribute(mp_el, IDREF)
    end
    mp_el = get_element_by_id(xdoc, id)
    return get_matrix_parameter(mp_el)
end

function set_operator_weight!(xdoc::XMLDocument, operator_name::String,
        weight::Int)
    operator_el = find_element(root(xdoc), OPERATORS)
    xel = find_element(operator_el, operator_name)
    set_attribute(xel, WEIGHT, weight)
end

function insert_element(xdoc::XMLDocument,
        element::XMLElement,
        reference_name::String;
        above::Bool = true)

    # xdoc = parse_file(src_path)

    xdoc2 = XMLDocument()
    create_root(xdoc2, BEAST)
    xroot2 = root(xdoc2)
    original_els = child_elements(root(xdoc))
    n_element = deepcopy(element)
    for el in original_els
        # free(xdoc)
        # free(xdoc2)
        # return 0
        # @show name(el)
        new_el = xml_deepcopy(el)
        if name(el) != reference_name
            add_child(xroot2, new_el)
        else
            if above
                add_child(xroot2, n_element)
                add_child(xroot2, new_el)
            else
                add_child(xroot2, new_el)
                add_child(xroot2, n_element)
            end
        end

    end
#
    for el in child_nodes(root(xdoc))
        unlink(el)
        free(el)
    end
    for el in child_elements(root(xdoc2))
        add_child(root(xdoc), xml_deepcopy(el))

    end
    free(xdoc2)
    return 0
end

function xml_deepcopy(el::XMLElement)
    xdoc = parse_string(string(el))
    new_el = root(xdoc)
    unlink(new_el)
    free(xdoc)
    return new_el
end

# function setup_trait_validation!(xdoc::XMLDocument;
#         id::String = "validation",
#         true_name::String = "traitsTrue",
#         true_param::String = "leafTraitsTrue",
#         use_mask::Bool = false,
#         use_tree_traits::Bool = false,
#         inferred_trait::String = "")
#
#     xroot = root(xdoc)
#     @show 1
#
#     validation = new_element(TRAIT_VALIDATION)
#     set_attribute(validation, TRAIT_NAME, true_name)
#     set_attribute(validation, "logSum", "false")
#     if use_tree_traits
#         set_attribute(validation, USE_TREE_TRAITS, "true")
#         set_attribute(validation, INFERRED_TRAIT, inferred_trait)
#     end
#     @show 2
#     trait_likelihood_el = pull_element(xdoc, TRAIT_DATA_LIKELIHOOD)
#     add_child(validation, trait_likelihood_el)
#     trait_param_el = new_element(TRAIT_PARAMETER)
#     param_el = new_element(PARAMETER)
#     set_attribute(param_el, IDREF, true_param)
#     add_child_level(trait_param_el, param_el, 3)
#     add_child_level(validation, trait_param_el, 2)
#
#     xdoc = insert_element!(xdoc, validation, MCMC)
#     save_file(xdoc, joinpath(Directories.desktop, "test.xml"))
#
#
#     log_el = get_element_by_id(xdoc, FILELOG)
#     new_el = new_element(TRAIT_VALIDATION)
#     set_attribute(new_el, IDREF, id)
#     add_child(log_el, new_el)
# end

function pull_element(xdoc::XMLDocument, el_name::String;
        parent::XMLElement = root(xdoc))

    el = find_element(parent, el_name)
    id = attribute(el, ID)
    new_el = new_element(el_name)
    set_attribute(new_el, IDREF, id)
    return new_el
end

function multivariateDistributionLikelihood(id::String, mean_param_id::String,
        prec_param_id::String,
        data::Vector{T}) where T <: Real

    el = new_element(MULT_DIST_LIKELIHOOD)
    set_attribute(el, ID, id)
    dist_el = new_element(DISTRIBUTION)
    mvn_el = new_element(MVM_DISTRIBUTION)
    mean_param = new_element(MEAN_PARAMETER)
    param = new_element(PARAMETER)
    set_attribute(param, IDREF, mean_param_id)
    add_child_level(mean_param, param, 4)
    add_child_level(mvn_el, mean_param, 3)
    prec_param = new_element(PRECISION_PARAMETER)
    mat_param = new_element(MATRIX_PARAMETER)
    set_attribute(mat_param, IDREF, prec_param_id)
    add_child_level(prec_param, mat_param, 4)
    add_child_level(mvn_el, prec_param, 3)
    add_child_level(dist_el, mvn_el, 2)
    add_child_level(el, dist_el, 1)

    data_el = new_element(DATA)
    param_el = new_element(PARAMETER)
    set_attribute(param_el, VALUE, join(data, " "))
    add_child_level(data_el, param_el, 2)
    add_child_level(el, data_el, 1)
    return el
end

function leafTraitParameter(id::String, taxon::String, tree::String,
        param::String)

    el = new_element(LEAF_TRAIT_PARAMETER)
    set_attribute(el, ID, id)
    set_attribute(el, TAXON, taxon)
    tree_el = new_element(TREE_MODEL)
    set_attribute(tree_el, IDREF, tree)
    add_child_level(el, tree_el, 2)
    param_el = new_element(PARAMETER)
    set_attribute(param_el, IDREF, param)
    add_child_level(el, param_el, 2)
    return el
end

function get_likelihood(xdoc::XMLDocument)
    mcmc_el = find_element(root(xdoc), MCMC)
    posterior_el = find_element(mcmc_el, POSTERIOR)
    likelihood_el = find_element(posterior_el, LIKELIHOOD)
    return likelihood_el
end

function get_matrix_parameter(mat_el::XMLElement; T::DataType = Float64)
    if name(mat_el) == MATRIX_PARAMETER
        children = collect(child_elements(mat_el))
        n = length(children)
        v1 = get_vector_parameter(children[1], T = T)
        p = length(v1)

        M = Matrix{T}(undef, n, p)
        M[1, :] .= v1
        for i = 2:n
            M[i, :] = get_vector_parameter(children[i], T = T)
        end
        return M
    elseif name(mat_el) == DIAGONAL_MATRIX
        param_el = find_element(mat_el, PARAMETER)
        dim = parse(Int, attribute(param_el, DIMENSION))
        val = parse(Float64, attribute(param_el, VALUE))
        return diagm(0 => fill(val, dim))
    else
        error("Not yet implemented")
    end
end

function get_vector_parameter(vec_el::XMLElement; T::DataType = Float64)
    buffer = Vector{T}(undef, 0)
    return get_vector_parameter(vec_el, buffer)
end

function get_vector_parameter(vec_el::XMLElement,
        buffer::Vector{T}) where T <: Any
    vals_attr = attribute(vec_el, VALUE, required = true)
    vals = split(vals_attr, ' ')
    p = length(vals)
    if length(buffer) == 0
        buffer = Vector{T}(undef, p)
    else
        @assert length(buffer) == p
    end
    for i = 1:p
        buffer[i] = parse(T, vals[i])
    end
    return buffer
end

function make_sub_ids(mat::Matrix{Float64}, id::String)
    n = size(mat, 1)
    return ["$(id).$i" for i = 1:n]
end

function make_matrix_parameter(mat::Matrix{Float64}, id::String;
                            sub_ids::Vector{String} = make_sub_ids(mat, id))
    el = new_element(MATRIX_PARAMETER)
    set_attribute(el, ID, id)
    n = size(mat, 1)
    for i = 1:n
        p_el = new_child(el, PARAMETER)
        set_attribute(p_el, ID, sub_ids[i])
        set_attribute(p_el, VALUE, join(mat[i, :], ' '))
    end
    return string(el)
end

end