text:
# make your scripts automatically re-activate your project
cd(@__DIR__)
using Pkg; Pkg.activate("../..")

# load packages
using ml_for_pt
using DelimitedFiles
using LaTeXStrings
using Plots
using Random
using Flux
ENV["GKSwstype"]="nul"

# set linear system size and corresponding data folder
L = 6
# L = 8
data_folder = "../../data/mbl_bose_hubbard/L="*string(L)*"/"

# set path to save folder
save_folder = "./results/L="*string(L)*"/"

# critical parameter values for many-body localization transition from literature
p_crit1 = 4.0
p_crit2 = 7.0

# define parameter ranges
p_min = 0.1
p_max = 20.0
n_points = 200
p_range = collect(LinRange(p_min, p_max, n_points))
dp = p_range[2]-p_range[1]
p_range_LBC = collect(p_min-dp/2:dp:p_max+dp/2)

# import data

# probabilities associated with each basis state (Fock basis) at p=p_1
prob = vcat(readdlm(data_folder*"probs_W="*string(1)*".txt")...)

# iterate over entire parameter range
probss = [prob]
for index in 2:201-1
  probb = vcat(readdlm(data_folder*"probs_W="*string(index)*".txt")...)
  push!(probss, probb)
end
samples = collect(1:length(probss[1]))
n_samples = length(samples)

# construct probability distribution from data
# distr(x,p) gives probability to draw the sample x at parameter value p
# samples correspond to enumeration of unique basis states with integers
function distribution(sample, p, p_range, probss)
  p_indx = Int(round((p-p_range[1])/(p_range[2]-p_range[1])))+1
  return probss[p_indx][Int(sample)]
end
distr = (x,p)->distribution(x, p, p_range, probss)

# optionally run methods using the analytical expressions to compare with neural network results

# supervised learning using analytical expression
# returns optimal predictions, indicator, and loss
pred_opt_SL, indicator_opt_SL, loss_opt_SL = MLP.get_indicators_SL_analytical(samples, distr, p_range, dp, p_min, p_max)

# prediction-based method using analytical expression
# returns optimal predictions, indicator, and loss
pred_opt_PBM, indicator_opt_PBM, loss_opt_PBM = MLP.get_indicators_PBM_analytical(samples, distr, p_range, dp)

# learning by confusion using analytical expression
# returns optimal indicator and loss
indicator_opt_LBC, loss_opt_LBC = MLP.get_indicators_LBC_analytical(samples, distr, p_range, p_range_LBC)

# construct dataset for training neural networks
dataset = MLP.get_dataset(p_range, distr, samples)
dataset_train_SL = MLP.get_training_data_SL(dataset, p_range, p_max, p_min)

# standardize inputs
inputs_one_hot = transpose(readdlm(data_folder*"strings.txt", eltype(dp)))
mean_train, std_train = MLP.get_dataset_stats(dataset, inputs_one_hot, length(p_range))
inputs_one_hot_stand = (inputs_one_hot.-mean_train)./std_train
mean_train_SL, std_train_SL = MLP.get_dataset_stats(dataset_train_SL, inputs_one_hot, length(p_range))
inputs_one_hot_stand_SL = (inputs_one_hot.-mean_train_SL)./std_train_SL

# supervised learning using neural networks

# initialize neural network
n_nodes_SL = [size(inputs_one_hot)[1], 128, 128, 64, 64, 64, 2]
NN_SL = Chain(
  Dense(n_nodes_SL[1], n_nodes_SL[2], relu),
  Dense(n_nodes_SL[2], n_nodes_SL[3], relu),
  Dense(n_nodes_SL[3], n_nodes_SL[4], relu),
  Dense(n_nodes_SL[4], n_nodes_SL[5], relu),
  Dense(n_nodes_SL[5], n_nodes_SL[6], relu),
  Dense(n_nodes_SL[6], n_nodes_SL[end], identity))
pnn_SL, re_SL = Flux.destructure(NN_SL)

# set hyperparameters
lr_SL = 0.001f0
epochs_SL = 1000
saveat_SL = 100
opt_SL = ADAM(lr_SL)
verbose = true

# train neural network
# returns predictions, indicators, and loss saved at epochs specified by saveat_SL variable
pred_logger_SL, losses_SL, NN_logger_SL = MLP.get_indicators_SL_numerical(pnn_SL,  re_SL, dataset_train_SL, dataset, epochs_SL, p_range, dp, p_min, p_max, opt_SL, inputs_one_hot_stand, verbose=verbose, saveat=saveat_SL)

# extract results at last save point
pred_NN_SL = pred_logger_SL[end][1]
indicator_NN_SL = pred_logger_SL[end][2]
loss_NN_SL = pred_logger_SL[end][3]

# save results of supervised learning in save_folder
open(save_folder*"hyperparams_SL.txt", "w") do io
  writedlm(io, [saveat_SL, lr_SL, epochs_SL])
end

open(save_folder*"NN_logger_SL.txt", "w") do io
  writedlm(io, NN_logger_SL)
end

for i in 1:length(pred_logger_SL)
  open(save_folder*"pred_logger_SL_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_SL[i][1])
  end

  open(save_folder*"indicator_logger_SL_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_SL[i][2])
  end

  open(save_folder*"loss_logger_SL_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_SL[i][3])
  end

end

open(save_folder*"pred_NN_SL.txt", "w") do io
  writedlm(io, pred_NN_SL)
end

open(save_folder*"indicator_NN_SL.txt", "w") do io
  writedlm(io, indicator_NN_SL)
end

open(save_folder*"loss_NN_SL.txt", "w") do io
  writedlm(io, loss_NN_SL)
end

open(save_folder*"NN_structure_SL.txt", "w") do io
  writedlm(io, n_nodes_SL)
end

open(save_folder*"losses_SL.txt", "w") do io
  writedlm(io, losses_SL)
end

# plot and save results of supervised learning in save_folder
plt = plot(p_range, pred_NN_SL, dpi=300, c="black", label="NN", ylims=(0.0,1.0))
plot!(p_range ,pred_opt_SL, dpi=300, c="blue", label="analytical", ylims=(0.0,1.0))
vline!([p_crit1, p_crit2], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\hat{y}_{\mathrm{SL}}$")
savefig(save_folder*"pred_NN_SL.png")

plt = plot(p_range[2:end-1], indicator_NN_SL, dpi=300, c="black", label="NN")
plot!(p_range[2:end-1], indicator_opt_SL, dpi=300, c="blue", label="analytical")
vline!([p_crit1, p_crit2], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I_{\mathrm{SL}}$")
savefig(save_folder*"indicator_NN_SL.png")

plt = plot(collect(1:length(losses_SL)), losses_SL, dpi=300, c="black", label="NN")
hline!([loss_opt_SL], label="analytical", c="blue")
xlabel!("Epochs")
ylabel!(L"$\mathcal{L}_{\mathrm{SL}}$")
savefig(save_folder*"loss_NN_SL.png")

# prediction-based method using neural networks

# initialize neural network
n_nodes_PBM = [size(inputs_one_hot)[1], 128, 128, 64, 64, 64, 1]
NN_PBM = Chain(
  Dense(n_nodes_PBM[1], n_nodes_PBM[2], relu),
  Dense(n_nodes_PBM[2], n_nodes_PBM[3], relu),
  Dense(n_nodes_PBM[3], n_nodes_PBM[4], relu),
  Dense(n_nodes_PBM[4], n_nodes_PBM[5], relu),
  Dense(n_nodes_PBM[5], n_nodes_PBM[6], relu),
  Dense(n_nodes_PBM[6], n_nodes_PBM[end], identity))
pnn_PBM, re_PBM = Flux.destructure(NN_PBM)

# set hyperparameters
lr_PBM = 0.001f0
epochs_PBM = 100
saveat_PBM = 100
opt_PBM = ADAM(lr_PBM)
verbose = true

# train neural network
# returns predictions, indicators, and loss saved at epochs specified by saveat_PBM variable
pred_logger_PBM, losses_PBM, NN_logger_PBM = MLP.get_indicators_PBM_numerical(pnn_PBM, re_PBM, dataset, epochs_PBM, p_range, dp, opt_PBM, inputs_one_hot_stand, verbose=verbose, saveat=saveat_PBM)

# extract results at last save point
pred_NN_PBM = pred_logger_PBM[end][1]
indicator_NN_PBM = pred_logger_PBM[end][2]
loss_NN_PBM = pred_logger_PBM[end][3]

# save results of prediction-based method in save_folder
open(save_folder*"hyperparams_PBM.txt", "w") do io
  writedlm(io, [saveat_PBM, lr_PBM, epochs_PBM])
end

open(save_folder*"NN_logger_PBM.txt", "w") do io
  writedlm(io, NN_logger_PBM)
end

for i in 1:length(pred_logger_PBM)
  open(save_folder*"pred_logger_PBM_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_PBM[i][1])
  end

  open(save_folder*"indicator_logger_PBM_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_PBM[i][2])
  end

  open(save_folder*"loss_logger_PBM_"*string(i)*".txt", "w") do io
    writedlm(io, pred_logger_PBM[i][3])
  end

end

open(save_folder*"pred_NN_PBM.txt", "w") do io
  writedlm(io, pred_NN_PBM)
end

open(save_folder*"indicator_NN_PBM.txt", "w") do io
  writedlm(io, indicator_NN_PBM)
end

open(save_folder*"loss_NN_PBM.txt", "w") do io
  writedlm(io, loss_NN_PBM)
end

open(save_folder*"NN_structure_PBM.txt", "w") do io
  writedlm(io, n_nodes_PBM)
end

open(save_folder*"losses_PBM.txt", "w") do io
  writedlm(io, losses_PBM)
end

# plot and save results of prediction-based method in save_folder
plt = plot(p_range, pred_NN_PBM, dpi=300, c="black", label="NN", ylims=(p_range[1], p_range[end]))
plot!(p_range, pred_opt_PBM, dpi=300, c="blue", label="analytical", ylims=(p_range[1], p_range[end]))
vline!([p_crit1, p_crit2], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\hat{y}_{\mathrm{PBM}}$")
savefig(save_folder*"pred_NN_PBM.png")

plt = plot(p_range[2:end-1], indicator_NN_PBM, dpi=300, c="black", label="NN")
plot!(p_range[2:end-1], indicator_opt_PBM, dpi=300, c="blue", label="analytical")
vline!([p_crit1, p_crit2], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I^{\mathrm{NN}}_{\mathrm{PBM}}$")
savefig(save_folder*"indicator_NN_PBM.png")

plt = plot(collect(1:length(losses_PBM)), losses_PBM, dpi=300, c="black", label="NN")
hline!([loss_opt_PBM], label="analytical", c="blue")
xlabel!("Epochs")
ylabel!(L"$\mathcal{L}_{\mathrm{PBM}}$")
savefig(save_folder*"loss_NN_PBM.png")

# learning by confusion using neural networks

# initialize neural network
n_nodes_LBC = [size(inputs_one_hot)[1], 128, 128, 64, 64, 64, 2]
NN_LBC = Chain(
  Dense(n_nodes_LBC[1], n_nodes_LBC[2], relu),
  Dense(n_nodes_LBC[2], n_nodes_LBC[3], relu),
  Dense(n_nodes_LBC[3], n_nodes_LBC[4], relu),
  Dense(n_nodes_LBC[4], n_nodes_LBC[5], relu),
  Dense(n_nodes_LBC[5], n_nodes_LBC[6], relu),
  Dense(n_nodes_LBC[6], n_nodes_LBC[end], identity))
pnn_LBC, re_LBC = Flux.destructure(NN_LBC)

# set hyperparameters
lr_LBC = 0.001f0
epochs_LBC = 100
saveat_LBC = epochs_LBC
opt_LBC = ADAM(lr_LBC)
verbose = true

indicator_NN_LBC = []
loss_NN_LBC = []
for indx in collect(1:length(p_range_LBC))
  pred_logger_LBC, losses_LBC, NN_logger_LBC = MLP.get_indicators_LBC_numerical_fixed_p(deepcopy(pnn_LBC), re_LBC, dataset, epochs_LBC, p_range, dp, opt_LBC, p_range_LBC, indx, inputs_one_hot_stand, saveat=saveat_LBC, verbose=verbose)

  indicator_NN_LBC_p = pred_logger_LBC[end][1][1]
  push!(indicator_NN_LBC, indicator_NN_LBC_p)
  loss_NN_LBC_p = pred_logger_LBC[end][2][1]
  push!(loss_NN_LBC, loss_NN_LBC_p)
end

# save results of learning by confusion in save_folder
open(save_folder*"hyperparams_LBC.txt", "w") do io
  writedlm(io, [saveat_LBC, lr_LBC, epochs_LBC])
end

open(save_folder*"indicator_NN_LBC.txt", "w") do io
  writedlm(io, indicator_NN_LBC)
end

open(save_folder*"loss_NN_LBC.txt", "w") do io
  writedlm(io, loss_NN_LBC)
end

open(save_folder*"NN_structure_LBC.txt", "w") do io
  writedlm(io, n_nodes_LBC)
end

# plot and save results of learning by confusion in save_folder
plt = plot(p_range_LBC, indicator_NN_LBC, dpi=300, c="black", label="NN")
plot!(p_range_LBC, indicator_opt_LBC, dpi=300, c="blue", label="analytical")
vline!([p_crit1, p_crit2], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$I_{\mathrm{LBC}}$")
savefig(save_folder*"indicator_NN_LBC.png")

plt = plot(p_range_LBC, loss_NN_LBC, dpi=300, c="black", label="NN")
plot!(p_range_LBC, loss_opt_LBC, dpi=300, c="blue", label="analytical")
vline!([p_crit1, p_crit2], label=L"p_{c}", c="red")
xlabel!(L"$p$")
ylabel!(L"$\mathcal{L}_{\mathrm{LBC}}$")
savefig(save_folder*"loss_NN_LBC.png")


text:
# Most code below is taken from https://github.com/richardkwo/InvariantCausal.jl,
# specifically SEM.jl with minor modifications to update for Julia 1.0
using LinearAlgebra # For identity matrix, I
using Distributions

abstract type SCM end

struct LinAddSCM <: SCM
   p ::Int64 # Number of variables
   B ::Matrix{Float64} # Weighted Adjacency Matrix of coefficients
   variances ::Vector{Float64} # Variances for each variables error term
   function LinAddSCM(B, variances)
       @assert length(variances) == size(B,1) == size(B,2)
       @assert all(variances .> 0)
       p = length(variances)
       new(p, B, variances)
   end
end

struct NonLinAddSCM <: SCM
   p ::Int64 # Number of variables
   A ::Matrix{Float64} # Adjacency Matrix (signed so weights can only be {-1,0,1})
   fs::Vector{Function} # Vector of Non-linear Functions for each variable
   # variances ::Vector{Float64} # Variances for each variables error term; NOTE: will reintroduce after the first sim
   function NonLinAddSCM(A, fs)#, variances)
       @assert size(A,1) == size(A,2) == length(fs)# == length(variances)
       # @assert all(variances .> 0)
       # p = length(variances)
       p = size(A,1)
       new(p, A, fs)#, variances)
   end
end

function Base.show(io::IO, scm::LinAddSCM)
    println(io, "LinAdd SCM with $(scm.p) variables.")
    println(io, "σ² = $(scm.variances).")
end

"""
    simulate(scm; d::UnivariateDistribution)
    simulate(scm, [n]; d::UnivariateDistribution)
    simulate(scm, [do_variables, do_values], [shift, additive], [n]; d::UnivariateDistribution)

Simulate from a Linear Additive (LinAdd) scm `scm`. `n` is the sample size.
`d` is a keyword argument that specifies which distribution to sample from for the error terms.
do-interventions can be performed by specifying vectors of `do_variables` and `do_values`.
shift-interventions can be performed by setting `shift=true` with `shift_values = do_variables`
and `shift_values = do_values`.
shift-interventions can be either additive by letting `additive=true` or multiplicative `additive=false`.
"""
function simulate(scm::LinAddSCM, d::UnivariateDistribution)
    ϵ = rand(d, scm.p) .* sqrt.(scm.variances) # Scale variances to our chosen variances
    return (I - scm.B) \ ϵ
end

function simulate(scm::LinAddSCM, n::Int64; d::UnivariateDistribution=Normal(0,1))
    return vcat(map(i -> simulate(scm, d), 1:n)'...)
end

function simulate(scm::LinAddSCM, d::UnivariateDistribution,
                  interventions::Tuple{Vector{Int64}, Vector{Float64}},
                  shift, additive)
    do_variables, do_values = interventions
    @assert length(do_variables) == length(do_values)
    p = scm.p
    ϵ = rand(d, p) .* sqrt.(scm.variances)
    if !shift # If shift intervention is false, perform an atomic intervention
        ϵ[do_variables] .= do_values
        B = copy(scm.B)
        B[do_variables, :] .= 0
    else
        shift_variables, shift_values = do_variables, do_values
        if additive
            ϵ[shift_variables] .+= shift_values
        else
            ϵ[shift_variables] .*= shift_values
        end
        B = copy(scm.B)
    end
    return (I - B) \ ϵ
end

function simulate(scm::LinAddSCM, interventions::Tuple{Vector{Int64}, Vector{Float64}}, n::Int64;
                  d::UnivariateDistribution=Normal(0,1),
                  shift=false, additive=true)
    return vcat(map(i -> simulate(scm, d, interventions, shift, additive), 1:n)'...)
end

# simulate observationa data, but allowing to give the full vector of distributions for every noises separately in the SCM
function simulate(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, n::Int64)
    # check
    @assert length(noise_dist) == scm.p

    # interventions
    jointdist = product_distribution(noise_dist)
    B = copy(scm.B)

    # simulate
    ϵ = rand(jointdist, n) .* sqrt.(scm.variances)

    return ((I - B) \ ϵ)'
end

function simulate(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, n,
                  noise=nothing, sum_or_prod=:sum, top_order=collect(1:6))
    # check
    @assert length(noise_dist) == scm.p
    if sum_or_prod == :sum
        g = xs -> dropdims(sum(xs, dims=2), dims=2)
    else
        g = xs -> dropdims(prod(xs, dims=2), dims=2)
    end

    # simulate
    if isnothing(noise)
        jointdist = product_distribution(noise_dist)
        ϵ = rand(jointdist, n) #.* sqrt.(scm.variances) # p × n; will reintroduce this term after the first sim we do
    else
        ϵ = noise
    end
    A = scm.A
    fs = scm.fs
    x = zeros(n, p)
    for j in top_order
        inds = findall(!iszero, A[:, j])
        if length(inds) == 0
            x[:,j] .= ϵ[j, :]
        else
            signs = A[inds,j]
            x[:, j] .= g(fs[j].(x[:, inds] .* signs')) .+ ϵ[j, :]
        end
    end
    return x
end

"""
        stochastic_inter_sim(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                             Vector{<:UnivariateDistribution}}, n::Int64)

Simulate from an SCM with stochastic interventions, that is do(X:=N), with N indep. of the rest, for a set list of variables X.

* `scm`         : a linear Structural Causal Model
* `noise_dist`  : list of the distributions of the noises in observational setting (assumed to be indep., with mean = 0 and var = 1)
* `inter_vect`  : tuple of indexes of variables to intervene on, and their corresponding interventional distr. (~N)
* `n`           : nbr of samples
"""
function stochastic_inter_sim(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                                Vector{<:UnivariateDistribution}}, n::Int64)
    # check
    inter_variables, inter_distr = inter_vect
    @assert length(noise_dist) == scm.p >= length(inter_variables) == length(inter_distr)

    # interventions
    dist_vect = copy(noise_dist)
    dist_vect[inter_variables] = inter_distr
    jointdist = product_distribution(dist_vect)
    B = copy(scm.B)
    B[inter_variables,:] .= 0

    # simulate
    ϵ = rand(jointdist, n) .* sqrt.(scm.variances)

    return ((I - B) \ ϵ)'
end

"""
        stochastic_inter_sim(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                             Vector{<:UnivariateDistribution}}, n::Int64, fs::Vector{Function})

Simulate from a non-linear SCM with stochastic interventions, that is do(X:=N), with N indep. of the rest, for a set list of variables X.

* `scm`         : a non-linear Structural Causal Model
* `noise_dist`  : list of the distributions of the noises in observational setting (assumed to be indep., with mean = 0 and var = 1)
* `inter_vect`  : tuple of indexes of variables to intervene on, and their corresponding interventional distr. (~N)
* `n`           : nbr of samples
* `fs`          : vector of non-linear functions used for the graph
"""
function stochastic_inter_sim(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect, n::Int64, fs::Vector{Function})
    # interventions
    dist_vect = copy(noise_dist)
    inter_variables, inter_distr = inter_vect
    dist_vect[inter_variables] = inter_distr
    A = copy(scm.A)
    A[:,inter_variables] .= 0
    new_scm = NonLinAddSCM(A, fs)
    return simulate(new_scm, dist_vect, n)
end

"""
        mech_changes_sim(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                             Vector{<:UnivariateDistribution}}, n::Int64)

Simulate from an SCM with mechanism changes, by either rescaling (i.e. X = β*X' + A*ϵ, with A random) or shifting the noise (i.e. X = β*X' + C+ϵ, with C random) or changing the structural equation's coefficients (i.e. X = β'*X' + ϵ) - or both coeff change + shift/rescale noise.

* `scm`         : a linear Structural Causal Model
* `noise_dist`  : list of the distributions of the noises in observational setting (assumed to be indep., with mean = 0 and var = 1)
* `noise_change`: tuple of indexes of variables for scale/shift noise change, and corresponding distr. (~A/C)
* `coeff_change`: tuple of indexes of variables for coeff. changes, and corresponding new coeffs (β')
* `n`           : nbr of samples
* `additive`    : whether we shift (=true) or rescale (=false) the intervened noise
"""
function mech_changes_sim(scm::LinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, noise_change::Tuple{Vector{Int64},
                                Vector{<:UnivariateDistribution}}, coeff_change::Tuple{Vector{Int64}, Array{Float64,2}}, n::Int64;
                                additive = true)
    # check
    noise_change_variables, noise_change_distr::Array{UnivariateDistribution{Continuous},1} = noise_change
    coeff_variables, coeff_mat = coeff_change
    @assert size(coeff_mat)[1] == length(coeff_variables) <= size(coeff_mat)[2] == length(noise_dist) == scm.p >=                      length(noise_change_variables) == length(noise_change_distr)

    # interventions
    joint_change_distr = product_distribution(noise_change_distr)
    B = copy(scm.B)
    B[coeff_variables,:] = coeff_mat

    # simulate
    jointdist = product_distribution(noise_dist)
    ϵ = rand(jointdist, n) .* sqrt.(scm.variances)
    noise_mat = (additive ? zeros(scm.p, n) : ones(scm.p, n))
    noise_mat[noise_change_variables,:] = rand(joint_change_distr, n)

    return ((I - B) \ (additive ? ϵ .+ noise_mat : ϵ .* noise_mat))'

end

"""
        mech_changes_sim(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, inter_vect::Tuple{Vector{Int64},
                             Vector{<:UnivariateDistribution}}, n::Int64)

Simulate from an SCM with mechanism changes, by either rescaling (i.e. X = f(X) + A*ϵ, with A random) or shifting the noise
(i.e. X = f(X) + C+ϵ, with C random) or changing the structural equation's coefficients (i.e. X = β'*X' + ϵ) - or both coeff change + shift/rescale noise.

* `scm`         : a non-linear Structural Causal Model
* `noise_dist`  : list of the distributions of the noises in observational setting (assumed to be indep., with mean = 0 and var = 1)
* `noise_change`: tuple of indexes of variables for scale/shift noise change, and corresponding distr. (~A/C)
* `n`           : nbr of samples
* `additive`    : whether we shift (=true) or rescale (=false) the intervened noise
"""
function mech_changes_sim(scm::NonLinAddSCM, noise_dist::Vector{<:UnivariateDistribution}, noise_change::Tuple{Vector{Int64},
                                Vector{<:UnivariateDistribution}}, n::Int64;
                                additive = true)
    # check
    noise_change_variables, noise_change_distr::Array{UnivariateDistribution{Continuous},1} = noise_change
    @assert length(noise_dist) == scm.p >= length(noise_change_variables) == length(noise_change_distr)

    # interventions
    joint_change_distr = product_distribution(noise_change_distr)

    # simulate
    jointdist = product_distribution(noise_dist)
    ϵ = rand(jointdist, n) #.* sqrt.(scm.variances) # NOTE: pretend it doesnt have variances for now
    noise_mat = (additive ? zeros(scm.p, n) : ones(scm.p, n))
    noise_mat[noise_change_variables,:] = rand(joint_change_distr, n)
    # although it seems like noise_dist is getting used here, it is isn't and the noise distribution is being passed in.
    return simulate(scm, noise_dist, n, (additive ? ϵ .+ noise_mat : ϵ .* noise_mat))
end

# get direct causes for i
function causes(scm::LinAddSCM, i::Int64)
    @assert 1 <= i <= scm.p
    return (1:scm.p)[scm.B[i, :].!=0]
end

# get markov blannket for i
function markov_blanket(scm::LinAddSCM, i::Int64)
    blanket = causes(scm, i)
    children = (1:scm.p)[scm.B[:,i].!=0]
    blanket = vcat(blanket, children)
    for j in children
        blanket = vcat(blanket,(1:scm.p)[scm.B[j, :].!=0])
    end
    return unique(blanket[blanket .!= i])
end

# function returning the structural equation coeff. for variable i
function causalcoeff(scm::LinAddSCM, i::Int64)
    @assert 1 <= i <= scm.p
    return (scm.B[i, 1:scm.p .!= i])
end

"""
    fixed_linear_SCM(p, k; [lb=-2, ub=2, var_min=0.5, var_max=2])

Generate a fixed-graph acyclic SCM with `p` variables with `edges` connecting them, and random coefficients.
* Edges should be of form [(child, parent), ...]
* `lb`, `ub`: coeff  ~ unif[`lb`, `ub`] with random sign
* `var_min`, `var_max`: var of error ~ unif[`var.min`, `var.max`]
"""
function fixed_linear_SCM(edges::Array{Tuple{Int64, Int64},1}, p::Int64; lb=-2, ub=2, var_min=0.5^2, var_max=2^2)
    B = zeros(p, p)
    for e in edges
        B[e...] = (rand() * (ub - lb) + lb) * sign(randn())
    end
    variances = rand(p) * (var_max - var_min) .+ var_min
    return LinAddSCM(B, variances)
end

"""
    random_linear_SCM(p, k; [lb=-2, ub=2, var_min=0.5, var_max=2])

Generate a random-graph acyclic SCM with `p` variables and `k` average degree, and random coefficients.
* `lb`, `ub`: coeff  ~ unif[`lb`, `ub`] with random sign
* `var_min`, `var_max`: var of error ~ unif[`var.min`, `var.max`]
"""
function random_linear_SCM(p::Int64, k::Int64; lb=-2, ub=2, var_min=0.5^2, var_max=2^2)
    B = zeros(p, p)
    B[rand(p, p) .< k / (p-1)] .= 1
    B[UpperTriangular(B).!=0] .= 0
    m = sum(B.==1)
    B[B.==1] .= (rand(m) * (ub - lb) .+ lb) .* sign.(randn(m))
    err_var = rand(p) * (var_max - var_min) .+ var_min
    _order = sample(1:p, p, replace=false)
    B = B[_order, _order]
    return LinAddSCM(B, err_var)
end

"""
    random_noise_intervened_SCM(scm::LinAddSCM, [p_intervened=2, noise_multiplier_min=0.5, noise_multiplier_max=2., avoid=[],
                                                   prob_coeff_unchanged=2/3, lb=-2, ub=2])

Produce a new SCM based on original SCM by changing coefficients and noise variances.
* `p_intervened`: randomly choose `p_intervened` variables to intervene; will avoid those specified in `avoid`
* [`noise_multiplier_min`, `noise_multiplier_max`]: interval that noise multiplier is uniformly sampled from
* `prob.coeff.unchanged`: probability that coefficient is not changed
* `[lb, ub]`: if to change, coefficient is drawn uniformly from this interval with random sign

Return: `scm_new`, `intervened_variables`
"""
function random_noise_intervened_SCM(scm::LinAddSCM;
                                     p_intervened=2, noise_multiplier_min=0.5, noise_multiplier_max=2., avoid=[],
                                     prob_coeff_unchanged=2/3, lb=-2, ub=2)
    B = copy(scm.B)
    p = scm.p
    err_var = copy(scm.variances)
    vars = sample(setdiff(collect(1:p), avoid), p_intervened, replace=false)
    for i in vars
        noise_multiplier = rand() * (noise_multiplier_max - noise_multiplier_min) + noise_multiplier_min
        err_var[i] = err_var[i] * noise_multiplier
        if rand() > prob_coeff_unchanged
            _J = (1:p)[B[i, :].!=0]
            B[i, _J] .= rand(length(_J)) * (ub - lb) .+ lb
        end
    end
    return LinAddSCM(B, err_var), vars
end

"""
        high_dim_SCM(p_anterior::Int64, p_posterior::Int64, k::Int64, parents_nbr::Int64 ;
                            lb=-2, ub=2, var_min=0.5, var_max=2, shuffle = false)

Generate a linear SCM, by allowing to fix the number of the direct causes (i.e. parents) of the target.
Mainly used for high-dimensional simulations. The variables are returned in their causal order.

* `p_anterior`          : nbr of non-desc variables to the target (which is the (`p_anterior`+1)th variable)
* `p_posterior`         : nbr of variables non-asc variables to the target
* `k`                   : averaged degree of every other nodes
* `parents_nbr`         : nbr of direct causes of the target
* `lb`, `ub`            : coeff  ~ unif[`lb`, `ub`] with random sign
* `var_min`, `var_max`  : var of error ~ unif[`var.min`, `var.max`]
"""
function high_dim_SCM(p_anterior::Int64, p_posterior::Int64, k::Int64, parents_nbr::Int64 ;
                      lb=-2, ub=2, var_min=0.5^2, var_max=2^2)
    # check
    @assert parents_nbr <= p_anterior

    # vars
    p = p_anterior + p_posterior + 1
    target_idx = p_anterior + 1

    # adjacency matrix
    B = zeros(p, p)
    draws = rand(p, p)
    draws[UpperTriangular(draws).> 0] .= 0
    B[draws .> 1 - min(k,p-1) / (p-1)] .= 1

    # target parents
    B[target_idx,sortperm(draws[target_idx,:])] = vcat(zeros(p - parents_nbr), ones(parents_nbr))
    m = sum(B.==1)

    # generate the coeff
    B[B.==1] .= (rand(m) * (ub - lb) .+ lb) .* rand([-1,1],m)

    # generate the variances
    err_var = rand(p) * (var_max - var_min) .+ var_min

    # Normalize X*β so that var(X*β) ∝ var(ϵ)
    B_inv = inv(I - B)
    for i in 2:p
        b²ᵢ = B_inv[i, :].^2
        vᵢ = sqrt(dot(b²ᵢ, err_var))
        B[i,:] ./= vᵢ
    end

    return LinAddSCM(B, err_var)
end

ext: jl


text:
# Estimates model parameters for proof-of-concept biochemical model -
using POETs
include("complement_lib.jl")

function estimate_complement_ensemble()

  number_of_subdivisions = 20
  number_of_parameters = 28
  number_of_objectives = 2

  ec_array_full = readdlm("./data/ec_array_O1_O2_O3.dat")
  total_error = sum(ec_array_full[:,1:end],1)
  min_index = indmin(total_error)
  pc_array_full = readdlm("./data/pc_array_O1_O2_O3.dat")
  initial_parameter_array = pc_array_full[:,min_index]
  #initial_parameter_array = 1.0*ones(28)
  #initial_parameter_array = vec(readdlm("best_solution.txt"))

  ec_array = zeros(number_of_objectives)
  pc_array = zeros(number_of_parameters)
  for index in collect(1:number_of_subdivisions)

    # Run JuPOETs -
    (EC,PC,RA) = estimate_ensemble(objective_function,neighbor_function,acceptance_probability_function,cooling_function,initial_parameter_array;rank_cutoff=4,maximum_number_of_iterations=20,show_trace=true)

    # Package -
    ec_array = [ec_array EC]
    pc_array = [pc_array PC]

    # Take the *best* value from the current ensemble, refine it and go around again -
    total_error = sum(ec_array[:,2:end],1)

    # Which col is the min error?
    min_index = indmin(total_error)
    @show index,total_error[min_index]

    # Refine the best solution -
    initial_parameter_array = pc_array[:,min_index].*(1+0.15*randn(number_of_parameters))
    #initial_parameter_array = local_refienment_step(initial_parameter_array)
  end

  return (ec_array,pc_array)
end


text:
# Autogenerated wrapper script for ZPares_jll for armv7l-linux-gnueabihf-libgfortran4
export libzpares

using OpenBLAS32_jll
using CompilerSupportLibraries_jll
JLLWrappers.@generate_wrapper_header("ZPares")
JLLWrappers.@declare_library_product(libzpares, "libzpares.so")
function __init__()
    JLLWrappers.@generate_init_header(OpenBLAS32_jll, CompilerSupportLibraries_jll)
    JLLWrappers.@init_library_product(
        libzpares,
        "lib/libzpares.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
using Documenter, MNPDynamics

makedocs(
    modules = [MNPDynamics],
    format = :html,
    checkdocs = :exports,
    sitename = "MNPDynamics.jl",
    pages = Any["index.md"]
)

deploydocs(
    repo = "github.com/tknopp/MNPDynamics.jl.git",
)


text:
using Plots
"""
    newton(f, df)
    newton(f, df, x0)
    newton(f, df, x0, tol)

Find roots of the noninear function `f` using Netwon's method.

# Argument
- `f::Function`: the nonlinear function.
- `df::function`: the derivative of the nonlinear function.
- `x0::Float64`: initial guess.
- 'tol::Float64': error tolerance.

# Example
```julia-repl
julia> newton(x->cos(x)-x, x->-sin(x)-1)
```
"""
function newton(f::Function, df::Function, x0::Float64=1.0, tol::Float64=1e-7)
    err = 1.0
    it = 0
    xp = x0
    while err > tol
        xc = x0 - f(xp)/df(xp)
        err = abs(xc-xp)
        xp = xc
        it += 1
    end
    return xp, it
end

"""
    secant(f)
    secant(f, x0, x1)
    secant(f, x0, x1, tol)

Find roots of the noninear function `f` using secant method.

# Argument
- `f::Function`: the nonlinear function.
- `x0::Float64`: the first initial guess.
- `x1::Float64`: the second initial guess.
- `tol::Float64`: error tolerance.

# Example
```julia-repl
julia> secant(x->cos(x)-x)
```
"""
function secant(f::Function, x0=1, x1 = 2, tol=1e-14)
    err = 1.0
    it = 0
    xp1 = x1
    xp2 = x0
    while err > tol
        xc = xp1 - f(xp1)*(xp1-xp2)/(f(xp1)-f(xp2))
        err = abs(xp1-xc)
        xp2 = xp1
        xp1 = xc
        it += 1
    end
    return xp1, it
end

"""
    bisect(f, a, b)
    bisect(f, a, b, tol)
    bisect(f, a, b, tol, maxit)

Find roots of the noninear function `f` using bisection.

# Argument
- `f::Function`: the nonlinear function.
- `a::Float64`: the lower bound of the interval.
- `b::Float64`: the upper bound of the interval.
- `tol::Float64`: error tolerance.
- `maxit::Int64`: maximum number of iteration.

# Example
```julia-repl
julia> bisection(x->cos(x)-x, 0, pi)
```
"""
function bisect(f::Function, a::Number, b::Number, tol=1e-7, maxit=100)
    if b < a
        error("the input of interval is not correct")
    end
    if f(a)*f(b) > 0
        error("You must give an input with difference sign")
    end
    it = 0
    while (b-a)>tol && it < maxit
        c = (a+b)/2
        if f(c) == 0
            return c
        elseif f(a)*f(c) < 0
            b = c
        else
            a = c
        end
        it += 1
    end
    return (a+b)/2, it
end

"""
    fixedpoint(g, x0)
    fixedpoint(g, x0, tol)
    fixedpoint(g, x0, tol, maxit)

Find roots of the noninear function `x=g(x)` using fixed point iteration.

# Argument
- `g::Function`: the fixed point function.
- `x0::Float64`: the initial guess.
- `tol::Float64`: error tolerance.
- `maxit::Int64`: maximum number of iteration.

# Example
```julia-repl
julia> fixedpoint(x->cos(x)-x, 0, pi)
```
"""
function fixedpoint(g::Function, x0, tol=1e-15, maxIt=50)
    it = 0
    x1 = x0 + 1
    while abs(x1-x0)>tol && it < maxIt
        x1 = x0
        x0 = g(x0)
        it += 1
    end
    return x0, it
end

"""
    cobweb(g, a, b, x0, N)

Generate the cobweb plot associated with the orbits x_n+1=g(x_n).

# Argument
- `g::Function`: the fixed point function.
- `a::Float64`: the lower bound of the interval.
- `b::Float64`: the upper bound of the interval.
- `x0::Float64`: the initial guess.
- `N::Int64`: the number of iteration.

# Example
```julia-repl
julia> cobweb(x->cos(x), 0, pi, 1, 30)
```
"""
function cobweb(g::Function,a,b,x0,N)
    #generate N linearly space values on (a,b)
    x = collect(range(a, b, length=N+1))
    y = g.(x);
    plot(x,y,linecolor=:black, legend = false)
    plot!(x,x,linecolor=:black, legend = false)
    x[1] = x0
    #@gif for i in 1:N
    anim = @animate  for i in 1:N
        x[i+1] = g(x[i])
        plot!([x[i]; x[i]], [x[i], x[i+1]], linecolor=:red,legend = false)
        plot!([x[i]; x[i+1]], [x[i+1], x[i+1]], linecolor=:red,legend = false)
    end
    return anim
end

"""
    viznewton(f, df, a, b, x0, N)

Visualize the Newton's method for the nonlinear equation f(x)=0.

# Argument
- `f::Function`: the nonlinear function.
- `df::function`: the derivative of the nonlinear function.
- `a::Float64`: the lower bound of the interval.
- `b::Float64`: the upper bound of the interval.
- `x0::Float64`: the initial guess.
- `N::Int64`: the number of iteration.

# Example
```julia-repl
julia> cobweb(x->cos(x), 0, pi, 1, 30)
```
"""
function viznewton(f::Function, df::Function, a,b,x0,N)
    #generate N linearly space values on (a,b)
    x = collect(range(a, b, length=N+1))
    y = f.(x);
    plot(x,y,linecolor=:black, legend = false)
    plot!(x,zeros(N+1),linecolor=:black, legend = false)
    x[1] = x0
    #@gif for i in 1:N
    anim = @animate  for i in 1:N
        x[i+1] = x[i] - f(x[i])/df(x[i])
        plot!([x[i]; x[i]], [0, f(x[i])], linecolor=:red,legend = false)
        plot!([x[i]; x[i+1]], [f(x[i]), 0], linecolor=:red,legend = false)
    end
    return anim
end

"""
    brent(f, a, b, tol, maxit)

Find roots of the noninear function `f` using Brent algorithm

# Argument
- `f::Function`: the nonlinear function.
- `a::Float64`: the lower bound of the interval.
- `b::Float64`: the upper bound of the interval.
- `tol::Float64`: error tolerance.
- `maxit::Int64`: maximum number of iteration.

# Example
```julia-repl
julia> brent(x->cos(x)-x, 0, pi)
```
"""
function brent(f::Function, a::Float64, b::Float64, tol::Float64=1e-13)
    if f(a)*f(b) >= 0.0
        error("Error: the root is not bracket")
    end
    if abs(f(a)) < abs(f(b))
        a, b = b, a
    end
    c = a
    d = c
    mflag = 1
    while abs(f(b)) > tol || abs(b-a) > tol
        if ~isapprox(f(a), f(c)) && ~isapprox(f(b), f(c))
            # inverse quadratic interpolation
            s = a*f(b)*f(c)/(f(a)-f(b))/(f(a)-f(c)) -
                b*f(a)*f(c)/(f(b)-f(a))/(f(b)-f(c)) -
                c*f(a)*f(b)/(f(c)-f(a))/(f(c)-f(b))
        else
            # secan method
            s = b - f(b)*(b-a)/(f(b)-f(a))
        end
        if (s > max((3a+b)/4, b) || s < min((3a+b)/4, b)) ||
            ( mflag == 1 && abs(s-b) >= abs(b-c)/2.0) ||
            ( mflag == 0 && abs(s-b) >= abs(c-d)/2.0) ||
            ( mflag == 1 && abs(b-c) < tol) ||
            ( mflag == 0 && abs(c-d) < tol)
            s = (a+b)/2
            mflag = 1
        else
            mflag = 0
        end
        d = c
        c = b
        if f(a)*f(s) < 0
            b = s
        else
            a = s
        end
        if abs(f(a)) < abs(f(b))
            a, b = b, a
        end
    end
    return b
end


text:
struct Loop
    master::Master
    learner::Learner
    iterhook::Function
end
Loop(master::Master,learner::Learner) = Loop(master,learner,x->nothing)

function evaluate!(loop::Loop,inch,outch)
    np = nparalel(loop.master)
    unresolved = 0
    x = 0 ### Needs to be some value of ask! output type
    while true
        if x!=nothing && unresolved<np && isopen(inch)
            i = take!(inch)
            x = ask!(loop.learner,i)

            if x!=nothing
                put!(loop.master.tasks,x)
                unresolved+=1
            else
                close(inch)
            end
        else
            res = take!(loop.master.results)
            tell!(loop.learner,res)
            put!(outch,res)
            unresolved -= 1

            if unresolved==0
                break
            end
        end
        loop.iterhook(loop)
    end
    close(outch)
end

Looptr(master,learner) = inch -> Channel() do outch
    loop = Loop(master,learner)
    evaluate!(loop, inch, outch)
end

"""
Evaluates until evaluate closes inchannel.
"""
function evaluate!(loop::Loop)
    inch = Channel(1)
    @async while true
        put!(inch,true)
    end
    inch |> Looptr(loop.master,loop.learner) |> collect
end

"""
Easy way to evaluate master with learner and a stopping condition.
"""
function evaluate!(loop::Loop,stop::Function)
    wl = WrappedLearner(loop.learner, stop, (x,y)->nothing, (x,y)->nothing)
    loop = Loop(loop.master,wl,loop.iterhook)
    evaluate!(loop)
end

### Channel(Map(identity)) is the only place where I use transducers.
"""
Evaluates until evaluate closes inchannel or iterator ends. Could be also used to pass random numbers.
"""
function evaluate!(loop::Loop,stop::Function,iter)

    wl = WrappedLearner(loop.learner, stop, (x,y)->nothing, (x,y)->nothing)
    loop = Loop(loop.master,wl,loop.iterhook)

    inch = Channel(1)
    @async begin
        for i in iter
            put!(inch,i)
        end
        close(inch)
    end
    inch |> Looptr(loop.master,loop.learner) |> collect
end

evaluate!(loop::Loop,iter) = evaluate!(loop,learner->false,iter)


text:
module BayesianLinearRegression

include("../src/Jaynes.jl")
using .Jaynes

function bayesian_linear_regression(x::Vector{Float64})
    σ = rand(:σ, InverseGamma(2, 3))
    β = rand(:β, Normal(0.0, 1.0))
    y = Vector{Float64}(undef, length(x))
    for i in 1 : length(x)
        push!(y, rand(:y => i, Normal(β * x[i], σ)))
    end
    return y
end

x = [Float64(i) for i in 1 : 100]
obs = selection(map(1 : 100) do i
                    (:y => i, ) => 3.0 * x[i] + rand()
                end)

n_samples = 5000
@time ps, lnw = importance_sampling(obs, n_samples, bayesian_linear_regression, (x, ))

mean_σ = sum(map(ps.calls) do cl
                 get_ret(cl[:σ])
             end) / n_samples
println("Mean σ: $mean_σ")

mean_β = sum(map(ps.calls) do cl
                 get_ret(cl[:β])
             end) / n_samples
println("Mean β: $mean_β")

end # module


text:
module ProbitTest

using Test
using ForneyLab
using ForneyLab: outboundType, isApplicable, requiresBreaker, breakerParameters
using ForneyLab: SPProbitOutNG, SPProbitIn1PN, EPProbitIn1BG, EPProbitIn1CG, EPProbitIn1PG
using StatsFuns: normcdf

@testset "requiresBreaker and breakerParameters" begin
    fg = FactorGraph()
    x = Variable()
    y = Variable()
    ng = GaussianMeanVariance(x, 0.0, 1.0)
    np = Probit(y, x)

    @test !requiresBreaker(np.i[:out]) # Dangling
    @test !requiresBreaker(np.i[:in1])
    @test !requiresBreaker(ng.i[:m])
    @test !requiresBreaker(ng.i[:out]) # requiresBreaker not applicable for EP

    @test_throws Exception breakerParameters(np.i[:out])
    @test breakerParameters(ng.i[:out]) == (Message{GaussianMeanVariance, Univariate}, ())
end

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

@testset "SPProbitOutNG" begin
    @test SPProbitOutNG <: SumProductRule{Probit}
    @test outboundType(SPProbitOutNG) == Message{Bernoulli}
    @test isApplicable(SPProbitOutNG, [Nothing, Message{Gaussian}])
    @test !isApplicable(SPProbitOutNG, [Message{Bernoulli}, Nothing])

    @test ruleSPProbitOutNG(nothing, Message(Univariate, GaussianMeanVariance, m=1.0, v=0.5)) == Message(Univariate, Bernoulli, p=normcdf(1/sqrt(1+0.5)))
end

@testset "SPProbitIn1PN" begin
    @test SPProbitIn1PN <: SumProductRule{Probit}
    @test outboundType(SPProbitIn1PN) == Message{Function}
    @test isApplicable(SPProbitIn1PN, [Message{PointMass}, Nothing])

    @test isa(ruleSPProbitIn1PN(Message(Univariate, PointMass, m=1.0), nothing), Message{Function, Univariate})
end

@testset "EPProbitIn1BG" begin
    @test EPProbitIn1BG <: ExpectationPropagationRule{Probit}
    @test outboundType(EPProbitIn1BG) == Message{GaussianWeightedMeanPrecision}
    @test isApplicable(EPProbitIn1BG, [Message{Bernoulli}, Message{Gaussian}], 2)
    @test !isApplicable(EPProbitIn1BG, [Message{PointMass}, Message{Gaussian}], 2)

    @test ruleEPProbitIn1BG(Message(Univariate, Bernoulli, p=1.0), Message(Univariate, GaussianMeanVariance, m=1.0, v=0.5)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=0.6723616582693994, w=0.3295003993960708)
    @test ruleEPProbitIn1BG(Message(Univariate, Bernoulli, p=0.8), Message(Univariate, GaussianMeanVariance, m=1.0, v=0.5)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=0.4270174959448596, w=0.19914199922339604)
    @test ruleEPProbitIn1BG(Message(Univariate, Bernoulli, p=0.5), Message(Univariate, GaussianMeanVariance, m=1.0, v=0.5)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=4e-12, w=4e-12)
end

@testset "EPProbitIn1CG" begin
    @test EPProbitIn1CG <: ExpectationPropagationRule{Probit}
    @test outboundType(EPProbitIn1CG) == Message{GaussianWeightedMeanPrecision}
    @test isApplicable(EPProbitIn1CG, [Message{Categorical}, Message{Gaussian}], 2)
    @test !isApplicable(EPProbitIn1CG, [Message{PointMass}, Message{Gaussian}], 2)

    @test ruleEPProbitIn1CG(Message(Univariate, Categorical, p=[1.0, 0.0]), Message(Univariate, GaussianMeanVariance, m=1.0, v=0.5)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=0.6723616582693994, w=0.3295003993960708)
    @test ruleEPProbitIn1CG(Message(Univariate, Categorical, p=[0.8, 0.2]), Message(Univariate, GaussianMeanVariance, m=1.0, v=0.5)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=0.4270174959448596, w=0.19914199922339604)
    @test ruleEPProbitIn1CG(Message(Univariate, Categorical, p=[0.5, 0.5]), Message(Univariate, GaussianMeanVariance, m=1.0, v=0.5)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=4e-12, w=4e-12)
end

@testset "EPProbitIn1PG" begin
    @test EPProbitIn1PG <: ExpectationPropagationRule{Probit}
    @test outboundType(EPProbitIn1PG) == Message{GaussianWeightedMeanPrecision}
    @test isApplicable(EPProbitIn1PG, [Message{PointMass}, Message{Gaussian}], 2)
    @test !isApplicable(EPProbitIn1PG, [Message{Bernoulli}, Message{Gaussian}], 2)

    @test ruleEPProbitIn1PG(Message(Univariate, PointMass, m=true), Message(Univariate, GaussianMeanVariance, m=1.0, v=0.5)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=0.6723616582693994, w=0.3295003993960708)
    @test ruleEPProbitIn1PG(Message(Univariate, PointMass, m=NaN), Message(Univariate, GaussianMeanVariance, m=1.0, v=0.5)) == Message(Univariate, GaussianWeightedMeanPrecision, xi=4e-12, w=4e-12)
end

@testset "averageEnergy" begin
    @test averageEnergy(Probit, ProbabilityDistribution(Univariate, Bernoulli, p=1.0), ProbabilityDistribution(Univariate, GaussianMeanVariance, m=0.0, v=1.0)) == 1.0000000289419368
    @test averageEnergy(Probit, ProbabilityDistribution(Univariate, PointMass, m=1.0), ProbabilityDistribution(Univariate, GaussianMeanVariance, m=0.0, v=1.0)) == 1.0000000289419368
end

end # module


text:
"""
Random Policy factored for each agent. Avoids exploding action space.
"""
struct FactoredRandomPolicy{RNG<:AbstractRNG,P<:JointMDP, U<:Updater} <: Policy
    rng::RNG
    problem::P
    updater::U
end

FactoredRandomPolicy(problem::JointMDP; rng=Random.GLOBAL_RNG, updater=NothingUpdater()) = FactoredRandomPolicy(rng, problem, updater)

function POMDPs.action(policy::FactoredRandomPolicy, s)
    return [rand(policy.rng, agent_actions(policy.problem, i, si)) for (i, si) in enumerate(s)]
end

POMDPs.solve(solver::RandomSolver, problem::JointMDP) = FactoredRandomPolicy(solver.rng, problem, NothingUpdater())


text:
import Base: eof, close
import Statistics: mean
export MatrixWriter, close, writeMatrix, MatrixReader, value, eof, advance!, mean, cor, cov

mutable struct MatrixWriter
    fileStream
    zerocount::Int
    maxnum::Int ##Somewhat of unused number
    expsize::Int
    datatype::DataType
end

function MatrixWriter(fileName::String, expsize::Int64, datatype::DataType)
    f = open(fileName, "w")
    assert(datatype in [UInt8, UInt16])
    if datatype == UInt8 #UInt8 mode stores up to 100
        mw = MatrixWriter(f, 1, 100, expsize, datatype)
    elseif datatype == UInt16 #UInt16 mode stores up to 60000
        mw = MatrixWriter(f, 1, 60000, expsize, datatype)
    end
    writeHeader(mw)
    mw
end

function writeHeader(mw::MatrixWriter)
    #Write other information.
    write(mw.fileStream, UInt16(mw.expsize))
    write(mw.fileStream, UInt16(mw.maxnum))

    #Indicate datatype.
    if mw.datatype == UInt8
        write(mw.fileStream, UInt16(0))
    elseif mw.datatype == UInt16
        write(mw.fileStream, UInt16(1))
    end
end

close(mw::MatrixWriter) = close(mw.fileStream)

function writeMatrix(mw::MatrixWriter, matrix::Array{Int64,2})

    skipnum = typemax(mw.datatype)-mw.maxnum
    for j in 1:size(matrix)[1]
        for i in 1:size(matrix)[2]
            if matrix[j, i] == 0
                mw.zerocount += 1
            elseif matrix[j, i] != 0

                # Figure out how much entry you have skipped and write it out.
                if mw.zerocount != 1
                    temp = mw.zerocount-1
                    _multi = floor(Int, temp/skipnum)
                    _const = temp%skipnum

                    for _ in 1:_multi
                        write(mw.fileStream, mw.datatype(typemax(mw.datatype)))
                    end

                    if _const != 0
                        write(mw.fileStream, mw.datatype(mw.maxnum+_const))
                    end

                end

                # Write non-zero number.
                if matrix[j, i] < mw.maxnum+1
                    # Write actual number
                    write(mw.fileStream, mw.datatype(matrix[j, i]))
                else
                    println("$(matrix[j, i]) is too big. Writing $(mw.maxnum) instead")
                    write(mw.fileStream, mw.datatype(mw.maxnum))
                end

                # Start new zero count
                mw.zerocount = 1
            end
        end
    end
end

mutable struct MatrixReader
    # For basic matrix reading
    fileStream
    expsize::Int
    maxnum::Int
    binsize::Int
    datatype::DataType
    blocksize::Int
    data::Array{Int64, 2}
    offset::Int

    # For buffering bites
    bufferpointer::Int
    buffersize::Int
    buffer
end

function MatrixReader(fileName::String, blocksize; buffsize=10000000)
    f = open(fileName)

    # These are read as header of file.
    _expsize = Int(read!(f, Array{UInt16}(undef, 1))[1])
    _maxnum = Int(read!(f, Array{UInt16}(undef, 1))[1])
    _datatype = Int(read!(f, Array{UInt16}(undef, 1))[1])

    if _datatype == 0
        dt = UInt8
    elseif _datatype == 1
        dt = UInt16
    end

    _buffsize = buffsize # This is 10Mb of buffer size for default
    mr = MatrixReader(f, _expsize, _maxnum, _maxnum, dt, blocksize, zeros(Int64, (blocksize, _expsize)), 0, _buffsize+1, _buffsize, zeros(dt, _buffsize))
end

close(mr::MatrixReader) = close(mr.fileStream)
value(mr::MatrixReader) = mr.data
eof(mr::MatrixReader) = eof(mr.fileStream) && mr.bufferpointer > length(mr.buffer)

function advance!(mr::MatrixReader)
    temp = zeros(Int64, (mr.blocksize, mr.expsize))
    pointer = mr.offset
    while !eof(mr) && pointer < mr.blocksize*mr.expsize
        #Reload buffer
        if mr.bufferpointer > length(mr.buffer)
            if mr.datatype == UInt8
                mr.buffer = read(mr.fileStream, mr.buffersize)
            elseif mr.datatype == UInt16
                tempbuffer = read(mr.fileStream, mr.buffersize*2)
                assert(length(tempbuffer)%2==0)
                mr.buffer = reinterpret(UInt16, tempbuffer)
            end
            mr.bufferpointer = 1
        end

        # Read in from buffer
        value = mr.buffer[mr.bufferpointer]
        mr.bufferpointer += 1

        # Fill matrix
        if value < mr.maxnum+1
            temp[ceil(Int, (pointer+1)/mr.expsize), pointer%mr.expsize+1] = value
            pointer += 1
        else
            pointer += value-mr.maxnum
        end
    end
    mr.offset = pointer - mr.blocksize*mr.expsize
    mr.data = temp
end

function mean(mr::MatrixReader; verbose=0)
    total = zeros(1, mr.expsize)
    advance!(mr)

    count = 0
    while !eof(mr)
        total += sum(mr.data, 1)
        count += mr.blocksize
        advance!(mr)
        if verbose>0 && count % (verbose*mr.blocksize) == 0
            println(count)
        end
    end
    (total/count)[1, 1:end]
end

function cor(mr::MatrixReader, means; verbose=0)
    # loop through all the chunks
    XtX = zeros(Float64, mr.expsize, mr.expsize)
    advance!(mr)

    count = 0
    while !eof(mr)
        temp = mr.data
        _x = temp'.-means
        BLAS.syrk!('U', 'N', 1.0, _x, 1.0, XtX)
        advance!(mr)
        if verbose>0 && count % (verbose) == 0
            println(count, ":", size(XtX))
        end
        count += 1
    end

    # Converting top half covariance to correlation.
    uppercov = XtX
    uppercor = cov2cor!(uppercov)
    cor = uppercor+uppercor'-eye(size(uppercor)[1])
    cor
end

function cov(mr::MatrixReader, means; verbose=0)
    # loop through all the chunks
    XtX = zeros(Float64, mr.expsize, mr.expsize)
    advance!(mr)

    count = 0
    while !eof(mr)
        temp = mr.data
        _x = temp'.-means
        BLAS.syrk!('U', 'N', 1.0, _x, 1.0, XtX)
        advance!(mr)
        if verbose>0 && count % (verbose) == 0
            println(count, ":", size(XtX))
        end
        count += 1
    end

    cov = XtX+XtX'
    for i in 1:size(XtX)[1]
       cov[i, i] ./= 2
    end
    cov
end


text:
using BenchmarkTools, Random, DataStructures, Printf, Plots, Statistics

include("../src/robin_dict.jl")

function add_entries(h::AbstractDict, entries::Vector{Pair{K, V}}) where {K, V}
	for (k, v) in entries
		h[k] = v
	end
end

@printf(".\nSample #1 Key => Integer , Size => 10^6 entries\n.\n")

sample1 = rand(Int, 10^6, 2)
entries1 = Vector{Pair{Int, Int}}()
sizehint!(entries1, 10^6)
for i = 1 : 10^6
	push!(entries1, Pair{Int, Int}(sample1[i, 1], sample1[i, 2]))
end

@printf("	add_entries for RobinDict()\n")
@btime add_entries(RobinDict(), entries1)

@printf("	add_entries for Dict()\n")
@btime add_entries(Dict(), entries1)

@printf("	add_entries for RobinDict{Int, Int}()\n")
@btime add_entries(RobinDict{Int, Int}(), entries1)

@printf("	add_entries for Dict{Int, Int}()\n")
@btime add_entries(Dict{Int, Int}(), entries1)

@printf(".\nSample #2 Key => Float32 , Size => 10^6 entries\n.\n")

sample2 = rand(Float32, 10^6, 2)
entries2 = Vector{Pair{Float32, Float32}}()
sizehint!(entries2, 10^6)
for i = 1 : 10^6
	push!(entries2, Pair{Float32, Float32}(sample1[i, 1], sample1[i, 2]))
end

@printf("	add_entries for RobinDict()\n")
@btime add_entries(RobinDict(), entries2)

@printf("	add_entries for Dict()\n")
@btime add_entries(Dict(), entries2)

@printf("	add_entries for RobinDict{Float32, Float32}()\n")
@btime add_entries(RobinDict{Float32, Float32}(), entries2)

@printf("	add_entries for Dict{Float32, Float32}()\n")
@btime add_entries(Dict{Float32, Float32}(), entries2)

@printf(".\nSample #3 Key => String , Size => 10^6 entries\n.\n")

entries3 = Vector{Pair{String, String}}()
sizehint!(entries3, 10^6)
for i = 1 : 10^6
	push!(entries3, Pair{String, String}(randstring(), randstring()))
end

@printf("	add_entries for RobinDict()\n")
@btime add_entries(RobinDict(), entries3)

@printf("	add_entries for Dict()\n")
@btime add_entries(Dict(), entries3)

@printf("	add_entries for RobinDict{String, String}()\n")
@btime add_entries(RobinDict{String, String}(), entries3)

@printf("	add_entries for Dict{String, String}()\n")
@btime add_entries(Dict{String, String}(), entries3)

@printf(".\nSample #4 Key => Integer , Size => 10^7 entries\n.\n")

sample1 = rand(Int, 10^7, 2)
entries1 = Vector{Pair{Int, Int}}()
sizehint!(entries1, 10^7)
for i = 1 : 10^7
	push!(entries1, Pair{Int, Int}(sample1[i, 1], sample1[i, 2]))
end

@printf("	add_entries for RobinDict()\n")
@btime add_entries(RobinDict(), entries1)

@printf("	add_entries for Dict()\n")
@btime add_entries(Dict(), entries1)

@printf("	add_entries for RobinDict{Int, Int}()\n")
@btime add_entries(RobinDict{Int, Int}(), entries1)

@printf("	add_entries for Dict{Int, Int}()\n")
@btime add_entries(Dict{Int, Int}(), entries1)

@printf(".\nSample #5 Key => Integer , Size => 10^5 entries\n.\n")

sample1 = rand(Int, 10^5, 2)
entries1 = Vector{Pair{Int, Int}}()
sizehint!(entries1, 10^5)
for i = 1 : 10^5
	push!(entries1, Pair{Int, Int}(sample1[i, 1], sample1[i, 2]))
end

@printf("	add_entries for RobinDict()\n")
@btime add_entries(RobinDict(), entries1)

@printf("	add_entries for Dict()\n")
@btime add_entries(Dict(), entries1)

@printf("	add_entries for RobinDict{Int, Int}()\n")
@btime add_entries(RobinDict{Int, Int}(), entries1)

@printf("	add_entries for Dict{Int, Int}()\n")
@btime add_entries(Dict{Int, Int}(), entries1)

@printf(".\nSample #6 Key => Float32 , Size => 10^5 entries\n.\n")

sample2 = rand(Float32, 10^5, 2)
entries2 = Vector{Pair{Float32, Float32}}()
sizehint!(entries2, 10^5)
for i = 1 : 10^5
	push!(entries2, Pair{Float32, Float32}(sample1[i, 1], sample1[i, 2]))
end

@printf("	add_entries for RobinDict()\n")
@btime add_entries(RobinDict(), entries2)

@printf("	add_entries for Dict()\n")
@btime add_entries(Dict(), entries2)

@printf("	add_entries for RobinDict{Float32, Float32}()\n")
@btime add_entries(RobinDict{Float32, Float32}(), entries2)

@printf("	add_entries for Dict{Float32, Float32}()\n")
@btime add_entries(Dict{Float32, Float32}(), entries2)

@printf(".\nSample #7 Key => String , Size => 10^5 entries\n.\n")

entries3 = Vector{Pair{String, String}}()
sizehint!(entries3, 10^5)
for i = 1 : 10^5
	push!(entries3, Pair{String, String}(randstring(), randstring()))
end

@printf("	add_entries for RobinDict()\n")
@btime add_entries(RobinDict(), entries3)

@printf("	add_entries for Dict()\n")
@btime add_entries(Dict(), entries3)

@printf("	add_entries for RobinDict{String, String}()\n")
@btime add_entries(RobinDict{String, String}(), entries3)

@printf("	add_entries for Dict{String, String}()\n")
@btime add_entries(Dict{String, String}(), entries3)

## Plots

get_load_factor(h::AbstractDict) = (h.count / length(h.keys))

function get_mean_dibs(h::RobinDict)
	sz = length(h.keys)
	dibs = zeros(Int8, sz)
	for i = 1:sz
		if isslotfilled(h, i)
			dibs[i] = calculate_distance(h, i)
		end
	end
	mean(dibs)
end

function get_variance_dibs(h::RobinDict)
	sz = length(h.keys)
	dibs = zeros(Int8, sz)
	for i = 1:sz
		if isslotfilled(h, i)
			dibs[i] = calculate_distance(h, i)
		end
	end
	var(dibs)
end

function plot_helper_add_entries(h::RobinDict, entries::Vector{Pair{K, V}}) where {K, V}
	num = 0
	sz = length(entries)
	sq = floor(sqrt(sz))
	x = Int[]
	lf = Float32[]
	xx = Int[]
	mean_dibs = Float32[]
	var_dibs = Float32[]
	for (k, v) in entries
		push!(x, num)
		push!(lf, get_load_factor(h))
		if (num % sq == 0)
			push!(mean_dibs, get_mean_dibs(h))
			push!(var_dibs, get_variance_dibs(h))
			push!(xx, num)
		end
		h[k] = v
		num += 1
	end
	push!(x, num)
	push!(lf, get_load_factor(h))
	if (num % sq == 0)
		push!(mean_dibs, get_mean_dibs(h))
		push!(var_dibs, get_variance_dibs(h))
		push!(xx, num)
	end
	y = [lf]
	yy = [mean_dibs var_dibs]
	png(plot(x, y, label = ["Load Factor"]), "lf_$(Int(sz/1000))K @$(ROBIN_DICT_LOAD_FACTOR) L.F")
	png(plot(xx, yy, label = ["mean(DIB)", "var(DIB)"]), "dibs_$(Int(sz/1000))K  @$(ROBIN_DICT_LOAD_FACTOR) L.F")
end

sample1 = rand(Int, 10^6, 2)
entries1 = Vector{Pair{Int, Int}}()
sizehint!(entries1, 10^6)
for i = 1 : 10^6
	push!(entries1, Pair{Int, Int}(sample1[i, 1], sample1[i, 2]))
end

plot_helper_add_entries(RobinDict{Int, Int}(), entries1)


text:
using PkgBenchmark
using IndexedTables
using IntervalSets: (..)

srand(666)

@benchgroup "construction" begin
    N = 10000
    a = rand(0:0.001:2, N)
    b = rand(0:0.001:2, N)
    c = rand(N)

    sorted = sort!(Columns(a, b))

    @bench "vectors" IndexedTable(a, b, c)
    @bench "vectors-agg" IndexedTable(a, b, c, agg=+)
    #@bench "sorted" IndexedTable(sorted, agg=+, presorted=true)
    @bench "vectors-in-place" IndexedTable(a, b, c, copy=false)
end

@benchgroup "indexing" begin
    N = 10000
    a = rand(0:0.001:2, N)
    b = rand(0:0.001:2, N)
    c = rand(N)

    t1 = IndexedTable(a, b, c, agg=+)
    i, j = t1.index[rand(1:length(t1))]

    # fast cases:
    @bench "const-const" $t1[$i, $j]
    @bench "colon-const" $t1[:, $j]
    @bench "colon-colon" $t1[:, :]

    @bench "interval-colon" $t1[0.2..0.8, :]
    @bench "interval-const" $t1[0.2..0.8, $j]
    @bench "const-interval" $t1[$i, 0.2..0.8]

    # slow cases
    @bench "range-colon" $t1[0.2:eps():0.8, :]
    @bench "range-const" $t1[0.2:eps():0.8, $j]
    @bench "interval-interval" $t1[0.2..0.8, 0.2..0.8]
    @bench "range-interval" $t1[0.2:eps():0.8, 0.2..0.8]

    x = sort!(a[200:800])
    @bench "vector-interval" $t1[$x, 0.2..0.8]
    y = sort!(b[400:600])
    @bench "vector-vector" $t1[$x, $y]
end

let
    N = 10000
    a = rand(0:0.001:2, N)
    b = rand(0:0.001:2, N)
    c = rand(N)
    t1 = IndexedTable(a, b, c, agg=+)
    @benchgroup "select" begin
        @bench "select-filter" select($t1, 1 => x -> x > 0.5, 2 => x -> x < 0.5)
        @bench "select-pick-1" select($t1, 1, agg=+)
        @bench "select-pick-2" select($t1, 2, agg=+)
    end

    t2 = IndexedTable(a, b, Columns(c, rand(N)))
    t3 = IndexedTable(a, b, Columns(c, [randstring(2) for i=1:N]))
    @benchgroup "filter" begin
        @bench "filter-float" filter(x->x>0.5, $t1)
        @bench "filter-float-float" filter(x->x[1]>0.5, $t2)
        # the below is much slower than the above
        # since each x is stack allocated
        @bench "filter-float-string" filter(x->x[1]>0.5, $t3)
    end

    @benchgroup "aggregate" begin
        N = 10000
        a = rand(1:10, N)
        b = rand(1:10, N)
        c = rand(N)
        t4 = IndexedTable(a, b, c)
        @bench "in-place" aggregate!(+, $(copy(t4)))
        @bench "out-of-place" aggregate(+, $(copy(t4)))
        @bench "vec" aggregate_vec(sum, $t4)
    end

    @benchgroup "reducedim" begin
        N = 10000
        a = rand(0:0.001:2, N)
        b = rand(0:0.001:2, N)
        c = rand(N)
        t4 = IndexedTable(a, b, c)
        @bench "dim-1" reducedim(+, $t4, 1)
        @bench "dim-2" reducedim(+, $t4, 2)
        @bench "vec-dim-1" reducedim_vec(+, $t4, 1)
        @bench "vec-dim-2" reducedim_vec(+, $t4, 2)
    end

    let

        function maketable(N)
            a = rand(0:0.001:2, N)
            b = rand(0:0.001:2, N)
            c = rand(N)
            IndexedTable(a, b, c)
        end
        big1 = maketable(10000)
        big2 = maketable(10000)
        small1 = maketable(1000)
        small2 = maketable(1000)
        @benchgroup "innerjoin" begin
            @bench "small-small" innerjoin($small1, $small2)
            @bench "small-big"   innerjoin($small1, $big1)
            @bench "big-small"   innerjoin($big1, $small1)
            @bench "big-big"     innerjoin($big1, $big2)
        end

        @benchgroup "leftjoin" begin
            @bench "small-small" leftjoin($small1, $small2)
            @bench "small-big"   leftjoin($small1, $big1)
            @bench "big-small"   leftjoin($big1, $small1)
            @bench "big-big"     leftjoin($big1, $big2)

            @bench "in-place-small-small" leftjoin!($(copy(small1)), $small2)
            @bench "in-place-small-big"   leftjoin!($(copy(small1)), $big1)
            @bench "in-place-big-small"   leftjoin!($(copy(big1)), $small1)
            @bench "in-place-big-big"     leftjoin!($(copy(big1)), $big2)
        end
    end

    @benchgroup "merge" begin
        @bench "small-small" merge($small1, $small2)
        @bench "small-big"   merge($small1, $big1)
        @bench "big-small"   merge($big1, $small1)
        @bench "big-big"     merge($big1, $big2)

        @bench "small-small-agg" merge($(copy(small1)),$(copy(small2)), agg=+)
        @bench "small-big-agg"   merge($(copy(small1)),$(copy(big1)), agg=+)
        @bench "big-small-agg"   merge($(copy(big1)),  $(copy(small1)), agg=+)
        @bench "big-big-agg"     merge($(copy(big1)),  $(copy(big2)), agg=+)
    end

    @benchgroup "asofjoin" begin
        @bench "small-small" asofjoin($small1, $small2)
        @bench "small-big"   asofjoin($small1, $big1)
        @bench "big-small"   asofjoin($big1, $small1)
        @bench "big-big"     asofjoin($big1, $big2)
    end
end


text:
#############################
##   triQgen: Generate Q operator which represents integration with regards to y
##               on JacobiTriangle().
####
function triQgen(N::Integer)
    # Generates Q operator as a block banded matrix.
    # Q*f(x,y) can be used to integrate f from 0 to 1-u via Jacobi polynomials on triangle.
    Q = BlockBandedMatrix(Zeros(N, sum(1:N)), Ones{Int}(N), 1:N, (0,0))
        for n=1:N
            view(Q, Block(n,n))[1,1] = 1.
        end
    return Q
end
#############################
##   triEygenP01: Extension operator for y from Jacobi P^(0,1) basis to canonical Jacobi triangle
####
function triEygenP01(N::Integer)
    Ey = BlockBandedMatrix(Zeros(sum(1:N),N), 1:N,Ones{Int}(N), (0,0))
        for n=1:N
            for j=1:n
                view(Ey, Block(n,n))[j,1] = (-1)^n*(-1)^j*(2*j-1)/n
            end
        end
    return Ey
end
#############################
##   triQEygenP01: This directly generates Q*EyP01 as a diagonal matrix
##                 without first computing the operators Q and RyP01
####
function triQEygenP01(N::Integer)
    QEy=BandedMatrix{Float64}(undef, (N, N), (0,0))
    for n=1:N
        QEy[n,n]=(-1)^(n+1)/n
    end
    return QEy
end
#############################
##   WLoweringP01P00: Weighted lowering operator which also multiplies by (1-x) from
##                      Jacobi(0,1) to Jacobi(0,0)
function WLoweringP01P00(N::Integer)
    WLow=I+BandedMatrix{Float64}(-0.5*Ones(N,N), (1,0))
    return WLow
end
#############################
##   reflectPabtoPba: Reflects coefficient vector from P^(a,b)(1-x) basis into P^(b,a)(x) basis
####
function reflectPabtoPba(N::Integer)
    Refl=BandedMatrix{Float64}(undef, (N, N), (0,0))
    for n=1:N
        Refl[n,n]=(-1)^(n+1)
    end
    return Refl
end
#############################################################
## The functions below are not in use outside of testing purposes
#############################
##   WLoweringP01Pm11: Weighted lowering operator which also multiplies by (1-x) from
##                      Jacobi(0,1) to Jacobi(1,-1)
function WLoweringP01Pm11(N::Integer)
    WLow=BandedMatrix{Float64}(undef, (N, N), (1,0))
    for c=1:N-1
        WLow[c,c]=(1-c)/(2*c)
        WLow[c+1,c]=-0.5
    end
    WLow[N,N]=(1-N)/(2*N)
    return WLow
end
##   ReflectorP01: Reflects coefficient vector from P^(1,0)(x) basis into P^(1,0)(1-t) basis
##      avoid use because it is upper triangular
####
function ReflectorP01(N::Integer)
    Reflector=BandedMatrix{Float64}(undef, (N, N), (0,N))
    for r=1:N
        for col=r+1:N
            Reflector[r,col]=(-1)^(r+1)*2*r/col
        end
            Reflector[r,r]=(-1)^(r+1)
    end
    return Reflector
end
#############################
##   triEygenP00: Extension operator for y from Legendre P^(0,0) basis to canonical Jacobi triangle
####
function triEygenP00(N::Integer)
    Ey = BlockBandedMatrix(Zeros(sum(1:N),N), 1:N,Ones{Int}(N), (0,1))
        for n=1:N
            for j=1:n-1
                view(Ey, Block(n-1,n))[j,1] = (-1)^n*(-1)^j*(2*j-1)/(2*n-1)
                view(Ey, Block(n,n))[j,1] = (-1)^n*(-1)^j*(2*j-1)/(2*n-1)
            end
            view(Ey, Block(n,n))[n,1] = 1.0
        end
    return Ey
end
#############################
##   triEygenP11: Extension operator for y from Jacobi P^(1,1) basis into T = JacobiTriangle(1,0,0,Triangle([0.0,1.0],[0.0,0.0],[1.0,1.0]))
####
function triEygenP11(N::Integer)
    Ey = BlockBandedMatrix(Zeros(sum(1:N),N), 1:N,Ones{Int}(N), (0,0))
        for n=1:N
                view(Ey, Block(n,n))[1,1] = (-1)^(n+1)
        end
    return Ey
end
#############################
##   triExgenP10: Extension operator for x from Jacobi P^(1,0) basis to canonical Jacobi triangle
####
function triExgenP01(N::Integer)
    Ex = BlockBandedMatrix(Zeros(sum(1:N), N), 1:N, Ones{Int}(N), (0,0))
        for n=1:N
            view(Ex, Block(n,n))[1,1] = 1
        end
    return Ex
end


text:
# Solve a trivial problem
A = reshape([1.0],(1,1))
solution = SCS_solve(SCS.Indirect, 1, 1, A, [1.0], [1.0], 1, 0, Int[], Int[], 0, 0, Float64[]);
@test solution.ret_val == 1

feasible_basic_conic(SCS.Indirect)

feasible_exponential_conic(SCS.Indirect);

feasible_sdp_conic(SCS.Indirect)

feasible_pow_conic(SCS.Indirect)


text:
# Isomap
# ------
# A Global Geometric Framework for Nonlinear Dimensionality Reduction,
# J. B. Tenenbaum, V. de Silva and J. C. Langford, Science 290 (5500): 2319-2323, 22 December 2000

"""

    Isomap{NN <: AbstractNearestNeighbors} <: AbstractDimensionalityReduction

The `Isomap` type represents an isometric mapping model constructed with a help of the `NN` nearest neighbor algorithm.
"""
struct Isomap{NN <: AbstractNearestNeighbors} <: AbstractDimensionalityReduction
    nearestneighbors::NN
    component::AbstractVector{Int}
    model::KernelPCA
end
Isomap(nn::NN, model::KernelPCA) where {NN <: AbstractNearestNeighbors} = Isomap(nn, model, Int[])

## properties
outdim(R::Isomap) = outdim(R.model)
eigvals(R::Isomap) = principalvars(R.model)
neighbors(R::Isomap) = R.nearestneighbors.k
vertices(R::Isomap) = R.component

## show
summary(io::IO, R::Isomap{T}) where T =
    print(io, "Isomap{$T}(outdim = $(outdim(R)), neighbors = $(neighbors(R)))")

## interface functions
"""
    fit(Isomap, data; k=12, maxoutdim=2, nntype=BruteForce)

Fit an isometric mapping model to `data`.

# Arguments
* `data`: a matrix of observations. Each column of `data` is an observation.

# Keyword arguments
* `k`: a number of nearest neighbors for construction of local subspace representation
* `maxoutdim`: a dimension of the reduced space.
* `nntype`: a nearest neighbor construction class (derived from `AbstractNearestNeighbors`)

# Examples
```julia
M = fit(Isomap, rand(3,100)) # construct Isomap model
R = transform(M)             # perform dimensionality reduction
```
"""
function fit(::Type{Isomap}, X::AbstractMatrix{T};
             k::Int=12, maxoutdim::Int=2, nntype=BruteForce) where {T<:Real}
    # Construct NN graph
    NN = fit(nntype, X, k)
    D, E = knn(NN, X)
    G, C = largest_component(SimpleWeightedGraph(adjmat(D,E)))

    # Compute shortest path for every point
    n = length(C)
    DD = zeros(T, n, n)
    for i in 1:n
        dj = dijkstra_shortest_paths(G, i)
        DD[i,:] = dj.dists
    end

    M = fit(KernelPCA, dmat2gram(DD), kernel=nothing, maxoutdim=maxoutdim)

    return Isomap{nntype}(NN, C, M)
end

"""
    transform(R::Isomap)

Transforms the data fitted to the Isomap model `R` into a reduced space representation.
"""
transform(R::Isomap) = transform(R.model)

"""
    transform(R::Isomap, X::AbstractVecOrMat)

Returns a transformed out-of-sample data `X` given the Isomap model `R` into a reduced space representation.
"""
function transform(R::Isomap, X::AbstractVecOrMat{T}) where {T<:Real}
    n = size(X,2)
    D, E = knn(R.nearestneighbors, X, self = true)
    DD = gram2dmat(R.model.X)
    G = zeros(size(R.model.X,2), n)
    for i in 1:n
        G[:,i] = minimum(DD[:,E[:,i]] .+ D[:,i]', dims=2)
    end

    broadcast!(x->-x*x/2, G, G)
    transform!(R.model.center, G)
    return projection(R.model)'*G'
end


text:
using Test
using Plots
using StatPlots; gr()
using ForestBiometrics
using DataFrames

#@test amd_convert(7) == 6.552 #how to test this since its in function?
@test typeof(gingrich_chart(400,80)) == Plots.Plot{Plots.GRBackend}


text:
function handle_color!(uniform_dict, instance_dict)
    color, udict = if haskey(uniform_dict, :color)
        to_value(uniform_dict[:color]), uniform_dict
    elseif haskey(instance_dict, :color)
        to_value(instance_dict[:color]), instance_dict
    else
        nothing, uniform_dict
    end
    if color isa Colorant ||
              color isa AbstractVector{<:Colorant} ||
              color === nothing
        delete!(uniform_dict, :colormap)
    elseif color isa AbstractArray{<:Real}
        udict[:color] = lift(x -> convert(Vector{Float32}, x), udict[:color])
        # udict[:color] = lift(x -> convert(Vector{Float32}, x), udict[:color])
        uniform_dict[:color_getter] = """
            vec4 get_color(){
                vec2 norm = get_colorrange();
                float normed = (color - norm.x) / (norm.y - norm.x);
                return texture(colormap, vec2(normed, 0));
            }
        """
    end
end

const IGNORE_KEYS = Set([:shading, :overdraw, :rotation, :distancefield, :markerspace,
                         :fxaa, :visible, :transformation, :alpha, :linewidth,
                         :transparency, :marker, :lightposition, :cycle, :label])

function create_shader(scene::Scene, plot::MeshScatter)
    # Potentially per instance attributes
    per_instance_keys = (:rotations, :markersize, :color, :intensity)
    per_instance = filter(plot.attributes.attributes) do (k, v)
        return k in per_instance_keys && !(isscalar(v[]))
    end
    per_instance[:offset] = apply_transform(transform_func_obs(plot),  plot[1])

    for (k, v) in per_instance
        per_instance[k] = Buffer(lift_convert(k, v, plot))
    end

    uniforms = filter(plot.attributes.attributes) do (k, v)
        return (!haskey(per_instance, k)) && isscalar(v[])
    end

    uniform_dict = Dict{Symbol,Any}()
    for (k, v) in uniforms
        k in IGNORE_KEYS && continue
        uniform_dict[k] = lift_convert(k, v, plot)
    end

    handle_color!(uniform_dict, per_instance)
    instance = convert_attribute(plot.marker[], key"marker"(), key"meshscatter"())

    if !hasproperty(instance, :uv)
        uniform_dict[:uv] = Vec2f(0)
    end

    uniform_dict[:depth_shift] = get(plot, :depth_shift, Observable(0f0))
    get!(uniform_dict, :ambient, Vec3f(1))

    return InstancedProgram(WebGL(), lasset("particles.vert"), lasset("particles.frag"),
                            instance, VertexArray(; per_instance...); uniform_dict...)
end

@enum Shape CIRCLE RECTANGLE ROUNDED_RECTANGLE DISTANCEFIELD TRIANGLE

primitive_shape(::Union{String,Char,Vector{Char}}) = Cint(DISTANCEFIELD)
primitive_shape(x::X) where {X} = Cint(primitive_shape(X))
primitive_shape(::Type{<:Circle}) = Cint(CIRCLE)
primitive_shape(::Type{<:Rect2}) = Cint(RECTANGLE)
primitive_shape(::Type{T}) where {T} = error("Type $(T) not supported")
primitive_shape(x::Shape) = Cint(x)

function char_scale_factor(char, font)
    # uv * size(ta.data) / Makie.PIXELSIZE_IN_ATLAS[] is the padded glyph size
    # normalized to the size the glyph was generated as.
    ta = Makie.get_texture_atlas()
    lbrt = glyph_uv_width!(ta, char, font)
    width = Vec(lbrt[3] - lbrt[1], lbrt[4] - lbrt[2])
    width * Vec2f(size(ta.data)) / Makie.PIXELSIZE_IN_ATLAS[]
end

# This works the same for x being widths and offsets
rescale_glyph(char::Char, font, x) = x * char_scale_factor(char, font)
function rescale_glyph(char::Char, font, xs::Vector)
    f = char_scale_factor(char, font)
    map(xs -> f * x, xs)
end
function rescale_glyph(str::String, font, x)
    [x * char_scale_factor(char, font) for char in collect(str)]
end
function rescale_glyph(str::String, font, xs::Vector)
    map((char, x) -> x * char_scale_factor(char, font), collect(str), xs)
end

using Makie: to_spritemarker

function scatter_shader(scene::Scene, attributes)
    # Potentially per instance attributes
    per_instance_keys = (:offset, :rotations, :markersize, :color, :intensity,
                         :uv_offset_width, :marker_offset)
    uniform_dict = Dict{Symbol,Any}()

    if haskey(attributes, :marker) && attributes[:marker][] isa Union{Char, Vector{Char},String}
        font = get(attributes, :font, Observable(Makie.defaultfont()))
        attributes[:markersize] = map(rescale_glyph, attributes[:marker], font, attributes[:markersize])
        attributes[:marker_offset] = map(rescale_glyph, attributes[:marker], font, attributes[:marker_offset])
    end

    if haskey(attributes, :marker) && attributes[:marker][] isa Union{Vector{Char},String}
        x = pop!(attributes, :marker)
        attributes[:uv_offset_width] = lift(x -> Makie.glyph_uv_width!.(collect(x)), x)
        uniform_dict[:shape_type] = Cint(3)
    end

    per_instance = filter(attributes) do (k, v)
        return k in per_instance_keys && !(isscalar(v[]))
    end

    for (k, v) in per_instance
        per_instance[k] = Buffer(lift_convert(k, v, nothing))
    end

    uniforms = filter(attributes) do (k, v)
        return !haskey(per_instance, k)
    end

    for (k, v) in uniforms
        k in IGNORE_KEYS && continue
        uniform_dict[k] = lift_convert(k, v, nothing)
    end

    get!(uniform_dict, :shape_type) do
        return lift(x -> primitive_shape(to_spritemarker(x)), attributes[:marker])
    end
    if uniform_dict[:shape_type][] == 3
        atlas = Makie.get_texture_atlas()
        uniform_dict[:distancefield] = Sampler(atlas.data, minfilter=:linear,
                                               magfilter=:linear, anisotropic=16f0)
        uniform_dict[:atlas_texture_size] = Float32(size(atlas.data, 1)) # Texture must be quadratic
    else
        uniform_dict[:atlas_texture_size] = 0f0
        uniform_dict[:distancefield] = Observable(false)
    end

    if !haskey(per_instance, :uv_offset_width)
        get!(uniform_dict, :uv_offset_width) do
            return if haskey(attributes, :marker) &&
                      to_spritemarker(attributes[:marker][]) isa Char
                lift(x -> Makie.glyph_uv_width!(to_spritemarker(x)),
                     attributes[:marker])
            else
                Vec4f(0)
            end
        end
    end

    space = get(uniforms, :markerspace, Observable(SceneSpace))
    uniform_dict[:use_pixel_marker] = map(space) do space
        return space == Pixel
    end
    handle_color!(uniform_dict, per_instance)

    instance = uv_mesh(Rect2(-0.5f0, -0.5f0, 1f0, 1f0))
    uniform_dict[:resolution] = scene.camera.resolution

    return InstancedProgram(WebGL(), lasset("simple.vert"), lasset("sprites.frag"),
                            instance, VertexArray(; per_instance...); uniform_dict...)
end

function create_shader(scene::Scene, plot::Scatter)
    # Potentially per instance attributes
    per_instance_keys = (:offset, :rotations, :markersize, :color, :intensity,
                         :marker_offset)
    per_instance = filter(plot.attributes.attributes) do (k, v)
        return k in per_instance_keys && !(isscalar(v[]))
    end
    attributes = copy(plot.attributes.attributes)
    attributes[:offset] = apply_transform(transform_func_obs(plot),  plot[1])
    attributes[:billboard] = map(rot -> isa(rot, Billboard), plot.rotations)
    attributes[:pixelspace] = getfield(scene.camera, :pixel_space)
    attributes[:model] = plot.model
    attributes[:markerspace] = plot.markerspace
    attributes[:depth_shift] = get(plot, :depth_shift, Observable(0f0))

    delete!(attributes, :uv_offset_width)
    return scatter_shader(scene, attributes)
end

value_or_first(x::AbstractArray) = first(x)
value_or_first(x::StaticArray) = x
value_or_first(x) = x

function create_shader(scene::Scene, plot::Makie.Text{<:Tuple{<:Union{<:Makie.GlyphCollection, <:AbstractVector{<:Makie.GlyphCollection}}}})
    glyphcollection = plot[1]
    res = map(x->Vec2f(widths(x)), pixelarea(scene))
    projview = scene.camera.projectionview
    transfunc =  Makie.transform_func_obs(scene)
    pos = plot.position
    space = plot.space
    offset = plot.offset

    # TODO: This is a hack before we get better updating of plot objects and attributes going.
    # Here we only update the glyphs when the glyphcollection changes, if it's a singular glyphcollection.
    # The if statement will be compiled away depending on the parameter of Text.
    # This means that updates of a text vector and a separate position vector will still not work if only the text
    # vector is triggered, but basically all internal objects use the vector of tuples version, and that triggers
    # both glyphcollection and position, so it still works
    if glyphcollection[] isa Makie.GlyphCollection
        # here we use the glyph collection observable directly
        gcollection = glyphcollection
    else
        # and here we wrap it into another observable
        # so it doesn't trigger dimension mismatches
        # the actual, new value gets then taken in the below lift with to_value
        gcollection = Observable(glyphcollection)
    end

    glyph_data = lift(pos, gcollection, space, projview, res, offset, transfunc) do pos, gc, args...
        Makie.preprojected_glyph_arrays(pos, to_value(gc), args...)
    end
    # unpack values from the one signal:
    positions, offset, uv_offset_width, scale = map((1, 2, 3, 4)) do i
        lift(getindex, glyph_data, i)
    end

    uniform_color = lift(glyphcollection) do gc
        if gc isa AbstractArray
            reduce(vcat, (Makie.collect_vector(g.colors, length(g.glyphs)) for g in gc),
                init = RGBAf[])
        else
            Makie.collect_vector(gc.colors, length(gc.glyphs))
        end
    end

    uniform_rotation = lift(glyphcollection) do gc
        if gc isa AbstractArray
            reduce(vcat, (Makie.collect_vector(g.rotations, length(g.glyphs)) for g in gc),
                init = Quaternionf[])
        else
            Makie.collect_vector(gc.rotations, length(gc.glyphs))
        end
    end

    uniforms = Dict(
        :model => plot.model,
        :shape_type => Observable(Cint(3)),
        :color => uniform_color,
        :rotations => uniform_rotation,
        :markersize => scale,
        :markerspace => Observable(Pixel),
        :marker_offset => offset,
        :offset => positions,
        :uv_offset_width => uv_offset_width,
        :transform_marker => Observable(false),
        :billboard => Observable(false),
        :pixelspace => getfield(scene.camera, :pixel_space),
        :depth_shift => get(plot, :depth_shift, Observable(0f0))
    )

    return scatter_shader(scene, uniforms)
end


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

# temp2 move-up

# initialise data relevant to move up
function initTemp2!(sim::Simulation;
	busyFraction::Float = 0.5, travelTimeCost::Float = 10.0, maxIdleAmbTravelTime::Float = 1.0, maxNumNearestStations::Int = 99)

	@warn("Move up module 'temp2' will not work with JuMP v0.19 and above.")

	# shorthand names:
	tmp = sim.moveUpData.temp2Data
	@unpack stations, numStations = sim

	# parameters:
	tmp.busyFraction = busyFraction
	tmp.travelTimeCost = travelTimeCost
	tmp.maxIdleAmbTravelTime = maxIdleAmbTravelTime # (days)
	tmp.maxNumNearestStations = maxNumNearestStations

	# create station pairs
	@assert(sim.travel.numSets == 1) # otherwise, would need to create station pairs for each travel set
	travelMode = getTravelMode!(sim.travel, lowPriority, sim.time)
	stationPairs = createStationPairs(sim, travelMode;
		maxPairsPerStation = 10, maxPairSeparation = 60/(24*60))

	# for easy indexing:
	stationSingles = Vector{Vector{Int}}()
	for i = 1:numStations
		push!(stationSingles, [i,i])
	end

	# testing / temporary:
	marginalBenefit = Array{Array{Float,2},2}(undef,numStations,numStations)
	for (i,j) in [stationSingles; stationPairs]
		m = stations[i].capacity
		n = stations[j].capacity
		mb = Array{Float,2}(undef,m,n)
		mb[1:m,1] = (1-busyFraction)*busyFraction.^[0:m-1;]
		mb[1,1:n] = (1-busyFraction)*busyFraction.^[0:n-1;]
		if i != j
			mb *= 0.2
		end
		for k1 = 2:m, k2 = 2:n
			mb[k1,k2] = min(mb[k1,k2-1], mb[k1-1,k2]) * 0.5
		end
		marginalBenefit[i,j] = mb
	end
	benefit = deepcopy(marginalBenefit)
	for (i,j) in [stationSingles; stationPairs]
		m = stations[i].capacity
		n = stations[j].capacity
		for k1 = 1:m, k2 = 1:n
			benefit[i,j][k1,k2] = sum(marginalBenefit[i,j][1:k1,1:k2])
		end
	end

	# note that for 'benefit' and 'marginalBenefit', if the station indices are the same, only the first row/column should be used
	tmp.benefit = benefit
	tmp.marginalBenefit = marginalBenefit
	tmp.stationPairs = stationPairs
end

# determine move ups to make at current time
# returns list of ambulances to be moved, and list of their destinations (stations)
function temp2MoveUp(sim::Simulation)

	# shorthand names:
	tmp = sim.moveUpData.temp2Data
	@unpack stationPairs, travelTimeCost, maxIdleAmbTravelTime, maxNumNearestStations, marginalBenefit = tmp
	@unpack ambulances, stations, numAmbs, numStations = sim

	# get movable ambulances (movableAmbs)
	ambMovable = Vector{Bool}(undef, numAmbs) # ambMovable[i] = true if ambulances[i] can move-up
	for i = 1:numAmbs
		ambMovable[i] = isAmbMovable(ambulances[i])
	end
	movableAmbs = ambulances[ambMovable]
	numMovableAmbs = length(movableAmbs)

	# calculate travel time for each movable ambulance to reach every station
	ambToStationTimes = Array{Float,2}(undef, numMovableAmbs, numStations)
	for i = 1:numMovableAmbs
		ambToStationTimes[i,:] = ambMoveUpTravelTimes!(sim, movableAmbs[i])
	end
	# edit travel times so that an ambulance travelling to its own station has no travel time
	for i = 1:numMovableAmbs
		j = movableAmbs[i].stationIndex
		ambToStationTimes[i,j] = 0.0
	end

	# restrict which stations each ambulance can be moved to
	# ambMovableToStation[i,j] = true if movableAmbs[i] can be moved to stations[j]; false otherwise
	ambMovableToStation = Array{Bool,2}(undef, numMovableAmbs, numStations)
	ambMovableToStation[:,:] .= true

	# limit ambulance move-up to nearest stations
	numNearestStations = min(maxNumNearestStations, numStations)
	sortedTimes = sort(ambToStationTimes, dims = 2) # sortedTimes[i,:] gives travel times for ith movable ambulance to all other stations, sorted
	ambMovableToStation = (ambMovableToStation .* (ambToStationTimes .<= sortedTimes[:, numNearestStations]))

	# for ambulances idle at station, limit travel time
	for i = 1:numMovableAmbs
		ambulance = movableAmbs[i]
		if ambulance.status == ambIdleAtStation
			ambMovableToStation[i,:] = (ambMovableToStation[i,:] .* (ambToStationTimes[i,:] .<= maxIdleAmbTravelTime))
		end
	end

	# useful lists for IP
	I = findall(ambMovableToStation)
	(ambList, stationList) = (getindex.(I, 1), getindex.(I, 2))
	# ambList and stationList together have all the information of ambMovableToStation:
	# - movableAmbs[i] can move to stations stationList[ambList .== i]
	# - stations[j] can have any of the ambulances in ambList[stationList .== j]
	m = length(ambList) # number of variables needed for assignment of ambulances to stations
	travelCostList = Vector{Float}(undef, m)
	for k = 1:m
		travelCostList[k] = ambToStationTimes[ambList[k], stationList[k]] * travelTimeCost
	end

	# counting number of ambulances at each station
	stationSlots = Vector{Int}()
	benefitSlots = Vector{Float}()
	for j = 1:numStations
		numSlots = min(stations[j].capacity, sum(stationList .== j))
		for k = 1:numSlots
			push!(stationSlots, j)
			push!(benefitSlots, marginalBenefit[j,j][1,k])
		end
	end

	# counting product of number of ambulances at station pairs
	stationPairSlots = Vector{Int}()
	benefitPairSlots = Vector{Float}()
	mapzy = Vector{Vector{Int}}() # for mapping stationPairSlots to pair of stationSlots
	mapmbz = Array{Array{Int,2},2}(undef, size(marginalBenefit)) # for mapping marginalBenefit to stationPairSlots
	for (p, (i,j)) in enumerate(stationPairs)
		ssi = findall(stationSlots .== i)
		ssj = findall(stationSlots .== j)
		mapmbz[i,j] = nullIndex * ones(Int, size(marginalBenefit[i,j]))
		for k1 = 1:length(ssi), k2 = 1:length(ssj)
			push!(stationPairSlots, p)
			push!(benefitPairSlots, marginalBenefit[i,j][k1,k2])
			push!(mapzy, [ssi[k1], ssj[k2]])
			mapmbz[i,j][k1,k2] = length(stationPairSlots)
		end
	end

	# testing: calculate lower bound on total benefit
	numAmbsAtStation = zeros(Int,numStations)
	for i = 1:numMovableAmbs
		numAmbsAtStation[movableAmbs[i].stationIndex] += 1
	end
	totalBenefitLowerBound = 0.0
	for i = 1:numStations
		totalBenefitLowerBound += sum(marginalBenefit[i,i][1,1:numAmbsAtStation[i]])
	end
	for (i,j) in stationPairs
		totalBenefitLowerBound += sum(sum(marginalBenefit[i,j][1:numAmbsAtStation[i],1:numAmbsAtStation[j]]))
	end

	######################
	# IP

	# shorthand variable names:
	a = numMovableAmbs
	s = numStations
	m = length(stationList)
	n = length(stationSlots) # <= m
	q = length(stationPairSlots) # = length(benefitPairSlots)

	# using JuMP
	model = Model()

	setsolver(model, GLPKSolverMIP(presolve=true))

	@variables(model, begin
		(x[k=1:m], Bin) # x[k] = 1 if ambulance: ambList[k] should be moved to station: stationList[k]
		(0 <= y[k=1:n] <= 1) # y should be naturally binary; sum(y[stationSlots .== k]) = number of ambulances assigned to stations[k]
		(0 <= z[k=1:q] <= 1) # z should be naturally binary
	end)

	@expressions(model, begin
		totalAmbTravelCosts, sum(x[k] * travelCostList[k] for k=1:m)
		totalBenefitAtStations, sum(y[k] * benefitSlots[k] for k=1:n)
		totalBenefitAtStationPairs, sum(z[k] * benefitPairSlots[k] for k=1:q)
		# totalBenefit, totalBenefitAtStations + totalBenefitAtStationPairs
	end)

	@constraints(model, begin
		(ambAtOneLocation[i=1:a], sum(x[k] for k=findall(ambList .== i)) == 1) # each ambulance must be assigned to one station
		(stationAmbCounts[j=1:s], sum(x[k] for k=findall(stationList .== j)) == sum(y[k] for k=findall(stationSlots .== j)))
		(stationPairAmbCounts[k=1:q, l=1:2], z[k] <= y[mapzy[k][l]])
		# (totalBenefit >= totalBenefitLowerBound)
	end)

	@objective(model, :Max, totalBenefitAtStations + totalBenefitAtStationPairs - totalAmbTravelCosts)

	# # testing: giving back fake results, for testing runtime without solving IP model
	# if true
		# return moveUpNull()
	# end

	solve(model)

	# extract solution
	sol = convert(Vector{Bool}, round.(getvalue(x)))
	ambStations = Vector{Station}(undef, numMovableAmbs)
	for k = 1:m
		if sol[k] == 1
			ambStations[ambList[k]] = stations[stationList[k]]
		end
	end

	if checkMode
		# check that y values are ordered correctly
		stationSlotsFilled = convert(Vector{Bool}, round.(getvalue(y)))
		for i = 1:numStations
			@assert(issorted(stationSlotsFilled[stationSlots .== i], rev=true))
		end

		# check that z values are ordered correctly
		stationPairSlotsFilled = convert(Vector{Bool}, round.(getvalue(z)))
		for (i,j) in stationPairs
			f(k1,k2) = mapmbz[i,j][k1,k2] == nullIndex ? 0 : stationPairSlotsFilled[mapmbz[i,j][k1,k2]]
			(m,n) = size(marginalBenefit[i,j])
			for k1 = 1:m, k2 = 1:n
				if k1 > 1
					@assert(f(k1,k2) <= f(k1-1,k2))
				end
				if k2 > 1
					@assert(f(k1,k2) <= f(k1,k2-1))
				end
			end
		end
	end

	return movableAmbs, ambStations
end


text:
function quadgk_PV(f, a, c, b; kwargs...)
    if b > c > a
        fc = f(c)
        g(x) = 2 * x * (-c^2 + x^2)^-1 * (f(x) - fc)
        return quadgk(g, a, prevfloat(c); kwargs...)[1] +
               quadgk(g, nextfloat(c), b; kwargs...)[1] +
               fc * log(abs((b^2 - c^2) / (a^2 - c^2)))
    else
        return quadgk(x -> 2 * x * (-c^2 + x^2)^-1 * f(x), a, b; kwargs...)[1]
    end
end

function quadgk_PV2(f, a, c, b; initdiv=1,kwargs...)
    fc = f(c)
    function g(x)
        if x == c
            2 * x * (-nextfloat(c)^2 + x^2)^-1 * (f(x) - fc)
        else
            return 2 * x * (-c^2 + x^2)^-1 * (f(x) - fc)
        end
    end
    return quadgk(g, a, b; order=ceil(Int64,initdiv/5),kwargs...)[1] + fc * log(abs((b^2 - c^2) / (a^2 - c^2)))
end

function quadgk_PV3(f, a, c, b;initdiv=1,kwargs...)
    if a < c < b
        fc = f(c)
        function g(x)
            if x == c
                2 * x * (-nextfloat(c)^2 + x^2)^-1 * (f(x) - fc)
            else
                return 2 * x * (-c^2 + x^2)^-1 * (f(x) - fc)
            end
        end
        return quadgk(g, a, b; order=ceil(Int64,initdiv/5),kwargs...)[1]
        # + fc * log(abs((b^2 - c^2) / (a^2 - c^2)))
    else
        fc2 = f(nextfloat(a))
        function g2(x)
            if x == c
                2 * x * (-nextfloat(c)^2 + x^2)^-1 * (f(x) - fc2)
            else
                return 2 * x * (-c^2 + x^2)^-1 * (f(x) - fc2)
            end
        end
        return quadgk(g2, a, b; order=ceil(Int64,initdiv/5),kwargs...)[1]
        # + fc * log(abs((b^2 - c^2) / (a^2 - c^2)))
    end
end

function hquadrature_PV(f, a, c, b; kwargs...)
    fc = f(c)
    g(x) = 2 * x * (-c^2 + x^2)^-1 * (f(x) - fc)
    return hquadrature(g, a, prevfloat(c); kwargs...)[1] +
           hquadrature(g, nextfloat(c), b; kwargs...)[1] +
           fc * log(abs((b^2 - c^2) / (a^2 - c^2)))
end

# function quadgk_PV2(f, c; kwargs...)
#     g(x) = (2*f(x))/(x)
#     return quadgk(g, 0.0, Inf; kwargs...)[1]
# end

# quadgk_PV(x->sin(x),10.0,15.0,20.0,atol=0.0,rtol=1e-4)
# PvdkF1Tildeps(p0, k, kprim, m, T)


text:
include("utils.jl")

function test_arraymanifold()
    M = Sphere(2)
    A = ArrayManifold(M)
    x = [1., 0., 0.]
    y = 1/sqrt(2)*[1., 1., 0.]
    z = [0., 1., 0.]
    v = log(M,x,y)
    v2 = log(A,x,y)
    y2 = exp(A,x,v2)
    w = log(M,x,z)
    w2 = log(A,x,z; atol=10^(-15))
    @test isapprox(y2.value,y)
    @test distance(A,x,y) == distance(M,x,y)
    @test norm(A,x,v) == norm(M,x,v)
    @test inner(A,x,v2,w2; atol=10^(-15)) == inner(M,x,v,w)
    @test_throws DomainError is_manifold_point(M,2*y)
    @test_throws DomainError is_tangent_vector(M,y,v; atol=10^(-15))

    test_manifold(A, [x, y, z],
        test_tangent_vector_broadcasting = false)
end

@testset "Sphere" begin
    M = Sphere(2)
    types = [Vector{Float64},
             SizedVector{3, Float64},
             MVector{3, Float64},
             Vector{Float32},
             SizedVector{3, Float32},
             MVector{3, Float32},
             Vector{Double64},
             MVector{3, Double64},
             SizedVector{3, Double64}]
    for T in types
        @testset "Type $T" begin
            pts = [convert(T, [1.0, 0.0, 0.0]),
                   convert(T, [0.0, 1.0, 0.0]),
                   convert(T, [0.0, 0.0, 1.0])]
            test_manifold(M,
                          pts,
                          test_reverse_diff = isa(T, Vector),
                          test_project_tangent = true,
                          point_distributions = [Manifolds.uniform_distribution(M, pts[1])],
                          tvector_distributions = [Manifolds.normal_tvector_distribution(M, pts[1], 1.0)])

            @test isapprox(-pts[1], exp(M, pts[1], log(M, pts[1], -pts[1])))
        end
    end

    @testset "Distribution tests" begin
        usd_mvector = Manifolds.uniform_distribution(M, @MVector [1.0, 0.0, 0.0])
        @test isa(rand(usd_mvector), MVector)

        gtsd_mvector = Manifolds.normal_tvector_distribution(M, (@MVector [1.0, 0.0, 0.0]), 1.0)
        @test isa(rand(gtsd_mvector), MVector)
    end

    @testset "log edge case" begin
        n = manifold_dimension(M)
        x = normalize(randn(n + 1))
        v = log(M, x, -x)
        @test norm(v) ≈ π
        @test isapprox(dot(x, v), 0; atol=1e-12)
        vexp = normalize(project_tangent(M, x, [1, zeros(n)...]))
        @test v ≈ π * vexp

        x = [1, zeros(n)...]
        v = log(M, x, -x)
        @test norm(v) ≈ π
        @test isapprox(dot(x, v), 0; atol=1e-12)
        vexp = normalize(project_tangent(M, x, [0, 1, zeros(n - 1)...]))
        @test v ≈ π * vexp
    end

    test_arraymanifold()
end


text:
function bar()
  return 404
end


text:
function Statistics.mean(states::Vector{SingleCovidState}, N::Int)
    n_states = length(states)
    sumS = 0
    sumI = zeros(Int,length(first(states).I))
    sumT = zeros(Int,size(first(states).T))
    for s in states
        sumS += s.S
        sumI .+= s.I
        sumT .+= s.T
    end
    avgS = floor(Int,sumS/n_states)
    avgI = floor.(Int, sumI./n_states)
    avgR = N - (avgS + sum(avgI))
    @assert avgR ≥ 0
    avgT = floor.(Int, sumT./n_states)

    return SingleCovidState(
        avgS,
        avgI,
        avgR,
        avgT,
        CovidPOMDPs.mean_params(states),
        first(states).prev_action,
        )
end


text:
function add_kepler_contribution(obj::Orbit, arg0::PositionAngle, arg1::jdouble, arg2::Vector{jdouble})
    return jcall(obj, "addKeplerContribution", void, (PositionAngle, jdouble, Vector{jdouble}), arg0, arg1, arg2)
end

function get_a(obj::Orbit)
    return jcall(obj, "getA", jdouble, ())
end

function get_a_dot(obj::Orbit)
    return jcall(obj, "getADot", jdouble, ())
end

function get_e(obj::Orbit)
    return jcall(obj, "getE", jdouble, ())
end

function get_e_dot(obj::Orbit)
    return jcall(obj, "getEDot", jdouble, ())
end

function get_equinoctial_ex(obj::Orbit)
    return jcall(obj, "getEquinoctialEx", jdouble, ())
end

function get_equinoctial_ex_dot(obj::Orbit)
    return jcall(obj, "getEquinoctialExDot", jdouble, ())
end

function get_equinoctial_ey(obj::Orbit)
    return jcall(obj, "getEquinoctialEy", jdouble, ())
end

function get_equinoctial_ey_dot(obj::Orbit)
    return jcall(obj, "getEquinoctialEyDot", jdouble, ())
end

function get_hx(obj::Orbit)
    return jcall(obj, "getHx", jdouble, ())
end

function get_hx_dot(obj::Orbit)
    return jcall(obj, "getHxDot", jdouble, ())
end

function get_hy(obj::Orbit)
    return jcall(obj, "getHy", jdouble, ())
end

function get_hy_dot(obj::Orbit)
    return jcall(obj, "getHyDot", jdouble, ())
end

function get_i(obj::Orbit)
    return jcall(obj, "getI", jdouble, ())
end

function get_i_dot(obj::Orbit)
    return jcall(obj, "getIDot", jdouble, ())
end

function get_le(obj::Orbit)
    return jcall(obj, "getLE", jdouble, ())
end

function get_le_dot(obj::Orbit)
    return jcall(obj, "getLEDot", jdouble, ())
end

function get_lm(obj::Orbit)
    return jcall(obj, "getLM", jdouble, ())
end

function get_lm_dot(obj::Orbit)
    return jcall(obj, "getLMDot", jdouble, ())
end

function get_lv(obj::Orbit)
    return jcall(obj, "getLv", jdouble, ())
end

function get_lv_dot(obj::Orbit)
    return jcall(obj, "getLvDot", jdouble, ())
end

function get_type(obj::Orbit)
    return jcall(obj, "getType", OrbitType, ())
end

function interpolate(obj::TimeInterpolable, arg0::AbsoluteDate, arg1::Stream)
    return jcall(obj, "interpolate", TimeInterpolable, (AbsoluteDate, Stream), arg0, arg1)
end

function shifted_by(obj::Orbit, arg0::jdouble)
    return jcall(obj, "shiftedBy", Orbit, (jdouble,), arg0)
end


text:
if true # create local scope
	# increment_alphanumeric_string
	pairs = [("a","b"),("z","aa"),("aa","ba"),("bb","cb"),("za","ab"),("zz","aaa")]
	for pair in pairs
		@test Smile.increment_alphanumeric_string(pair[1]) == pair[2]
	end

	# test alphanumeric_names
	@test alphanumeric_names(5)==["a","b","c","d","e"]
	@test_throws AssertionError alphanumeric_names(0)

	# test learn_bayesian_search + !
	mat = rand(1:4, 5000, 4)
	@test learn( mat, DSL_BayesianSearch )[2]
	net = Network()
	@test learn!( net, mat, DSL_BayesianSearch )
	params = LearnParams_BayesianSearch()
	@test learn!( net, mat, params )

	# test setting params
	params.maxparents           = 1
	params.maxsearchtime        = 10
	params.niterations          = 5
	params.linkprobability      = 0.5
	params.priorlinkprobability = 0.5
	params.priorsamplesize      = 5000
	params.seed                 = 1
	push!(params.forced_arcs, (1,2))
	push!(params.forbidden_arcs, (1,3))
	@test learn!( net, mat, params )
end