text:
# implementation of the GPUCompiler interfaces for generating eBPF code

## target

export BPFCompilerTarget

Base.@kwdef struct BPFCompilerTarget <: AbstractCompilerTarget
    function_pointers::UnitRange{Int}=1:1000 # set of valid function "pointers"
end

llvm_triple(::BPFCompilerTarget) = "bpf-bpf-bpf"
llvm_datalayout(::BPFCompilerTarget) = "e-m:e-p:64:64-i64:64-n32:64-S128"

function llvm_machine(target::BPFCompilerTarget)
    triple = llvm_triple(target)
    t = Target(;triple=triple)

    cpu = ""
    feat = ""
    tm = TargetMachine(t, triple, cpu, feat)
    asm_verbosity!(tm, true)

    return tm
end

## job

runtime_slug(job::CompilerJob{BPFCompilerTarget}) = "bpf"

const bpf_intrinsics = () # TODO
isintrinsic(::CompilerJob{BPFCompilerTarget}, fn::String) = in(fn, bpf_intrinsics)

valid_function_pointer(job::CompilerJob{BPFCompilerTarget}, ptr::Ptr{Cvoid}) =
    reinterpret(UInt, ptr) in job.target.function_pointers


text:
function lyapunovs_convergence(integ, N, dt::Real, Ttr::Real = 0.0)

    T = stateeltype(integ)
    t0 = integ.t
    if Ttr > 0
        while integ.t < t0 + Ttr
            step!(integ, dt)
            qrdec = LinearAlgebra.qr(get_deviations(integ))
            set_deviations!(integ, qrdec.Q)
        end
    end
    k = size(get_deviations(integ))[2]
    t0 = integ.t; t = zeros(T, N); t[1] = t0
    λs = [zeros(T, k) for i in 1:N];

    for i in 2:N
        step!(integ, dt)
        qrdec = LinearAlgebra.qr(get_deviations(integ))
        for j in 1:k
            @inbounds λs[i][j] = λs[i-1][j] + log(abs(qrdec.R[j,j]))
        end
        t[i] = integ.t
        set_deviations!(integ, qrdec.Q)
    end
    popfirst!(λs); popfirst!(t)
    for j in eachindex(t)
        λs[j] ./= (t[j] - t0)
    end
    return λs, t
end

ext: jl


text:
module Leiden

using Random:
    shuffle,
    shuffle!
using SparseArrays:
    sparse,
    spzeros
using StatsFuns:
    logsumexp

include("PartitionedGraphs.jl")
using .PartitionedGraphs:
    PartitionedGraph,
    Partition,
    WeightMatrix,
    nv, nc,
    neighbors,
    move_node!,
    drop_empty_communities!,
    reset_partition!,
    create_singleton_partition

# The Leiden algorithm
# --------------------

function leiden(adjmat::AbstractMatrix{<:Real};
                resolution::Real = 1.0,
                randomness::Real = 0.01,
                partition::Partition = create_singleton_partition(size(adjmat, 1)))
    graph = PartitionedGraph(adjmat, partition = partition)
    return _leiden!(graph, Float64(resolution), Float64(randomness))
end

function _leiden!(graph::PartitionedGraph, γ::Float64, θ::Float64)
    stack = Partition[]
    @label loop
    move_nodes_fast!(graph, γ)
    @debug "nv = $(nv(graph)); nc = $(nc(graph)); H = $(H(graph, γ))"
    if nc(graph) != nv(graph)
        refined = refine_partition(graph, γ, θ)
        if nc(refined) == nv(refined)
            #@warn "not refined" nc(refined)
            @goto finish
        end
        push!(stack, refined.partition)
        graph′ = aggregate_graph(refined)
        partition = [Int[] for _ in 1:nc(graph)]
        for (i, community) in enumerate(refined.partition)
            u = first(community)
            j = graph.membership[u]
            push!(partition[j], i)
        end
        reset_partition!(graph′, partition)
        graph = graph′
        @goto loop
    end
    @label finish
    push!(stack, graph.partition)
    return (quality = H(graph, γ), partition = flatten(stack))
end

# The Louvain algorithm
# ---------------------

function louvain(adjmat::AbstractMatrix{<:Real};
                 resolution::Real = 1.0,
                 partition::Partition = create_singleton_partition(size(adjmat, 1)))
    graph = PartitionedGraph(adjmat, partition = partition)
    return _louvain!(graph, Float64(resolution))
end

function _louvain!(graph::PartitionedGraph, γ::Float64)
    stack = Partition[]
    @label loop
    move_nodes!(graph, γ)
    @debug "nv = $(nv(graph)); nc = $(nc(graph)); H = $(H(graph, γ))"
    push!(stack, graph.partition)
    if nc(graph) != nv(graph)
        graph = aggregate_graph(graph)
        @goto loop
    end
    return (quality = H(graph, γ), partition = flatten(stack))
end

function move_nodes!(graph::PartitionedGraph, γ::Float64)
    H_old = H(graph, γ)
    nodes = collect(1:nv(graph))
    connected = Int[]
    total_weights = zeros(Float64, nc(graph))
    @label loop
    shuffle!(nodes)
    for u in nodes
        # compute total edge weights for each community
        empty!(connected)
        for v in neighbors(graph, u)
            i = graph.membership[v]
            if total_weights[i] == 0
                push!(connected, i)
            end
            total_weights[i] += graph.edge_weight[u,v]
        end

        # find the best community to which `u` belongs
        c_u = graph.cardinality[u]
        weight_u = graph.edge_weight[u,u]
        src = dst = graph.membership[u]
        weight_src = total_weights[src]
        size_src = graph.size[src]
        maxgain = 0.0
        for i in connected
            i == src && continue
            gain = total_weights[i] + weight_u - weight_src - γ * (graph.size[i] - size_src + c_u) * c_u
            if gain > maxgain
                dst = i
                maxgain = gain
            end
            total_weights[i] = 0
        end
        total_weights[src] = 0
        #@assert all(x == 0 for x in total_weights)

        if src != dst
            # move `u` to the best community and add its neighbors to the queue if needed
            move_node!(graph, u => dst)
        end
    end
    H_new = H(graph, γ)
    if H_new > H_old
        H_old = H_new
        @goto loop
    end
    return drop_empty_communities!(graph)
end

function move_nodes_fast!(graph::PartitionedGraph, γ::Float64)
    n = nv(graph)
    queue = shuffle(1:n)
    queued = BitSet(1:n)
    connected = Int[]
    total_weights = zeros(Float64, nc(graph))
    while !isempty(queue)
        u = popfirst!(queue)
        delete!(queued, u)

        # compute total edge weights for each community
        empty!(connected)
        for v in neighbors(graph, u)
            i = graph.membership[v]
            if total_weights[i] == 0
                push!(connected, i)
            end
            total_weights[i] += graph.edge_weight[u,v]
        end

        # find the best community to which `u` belongs
        c_u = graph.cardinality[u]
        weight_u = graph.edge_weight[u,u]
        src = dst = graph.membership[u]
        weight_src = total_weights[src]
        size_src = graph.size[src]
        maxgain = 0.0
        for i in connected
            i == src && continue
            gain = total_weights[i] + weight_u - weight_src - γ * (graph.size[i] - size_src + c_u) * c_u
            if gain > maxgain
                dst = i
                maxgain = gain
            end
            total_weights[i] = 0
        end
        total_weights[src] = 0
        #@assert all(x == 0 for x in total_weights)

        if src != dst
            # move `u` to the best community and add its neighbors to the queue if needed
            move_node!(graph, u => dst)
            for v in neighbors(graph, u)
                if graph.membership[v] != graph.membership[u] && v ∉ queued
                    push!(queue, v)
                    push!(queued, v)
                end
            end
        end
    end
    @assert isempty(queued)
    return drop_empty_communities!(graph)
end

function refine_partition(graph::PartitionedGraph, γ::Float64, θ::Float64)
    @assert γ > 0
    @assert θ > 0
    refined = PartitionedGraph(graph.edge_weight, cardinality = graph.cardinality)
    function is_well_connected(u::Int)
        i = graph.membership[u]
        c = graph.cardinality[u]
        threshold = γ * c * (graph.size[i] - c)
        x = 0.0
        for v in graph.partition[i]
            v == u && continue
            x += graph.edge_weight[u,v]
            if x ≥ threshold  # return as early as possible
                return true
            end
        end
        return false
    end
    function is_well_connected(u::Int, i::Int, between_weights::Vector{Float64})
        sz = refined.size[i]
        return between_weights[i] ≥ γ * sz * (graph.size[graph.membership[u]] - sz)
    end
    function is_singleton(u::Int)
        return length(refined.partition[refined.membership[u]]) == 1
    end
    total_weights = zeros(Float64, nc(refined))
    between_weights = zeros(Float64, nc(refined))
    for subset in graph.partition
        communities = Int[]
        logprobs = Float64[]
        indexes = Int[]
        for u in subset
            weight = 0.0
            for v in subset
                if v != u
                    weight += refined.edge_weight[u,v]
                end
            end
            between_weights[refined.membership[u]] = weight
        end
        for u in shuffle(subset)
            if !is_well_connected(u) || !is_singleton(u)
                continue
            end

            empty!(communities)
            for v in neighbors(refined, u)
                if v ∉ subset
                    continue
                end
                i = refined.membership[v]
                if total_weights[i] == 0
                    push!(communities, i)
                end
                total_weights[i] += refined.edge_weight[u,v]
            end

            c_u = refined.cardinality[u]
            weight_u = refined.edge_weight[u,u]
            src = refined.membership[u]
            weight_src = total_weights[src]
            size_src = refined.size[src]
            empty!(logprobs)
            empty!(indexes)
            for i in communities
                (i == src || !is_well_connected(u, i, between_weights)) && continue
                gain = total_weights[i] + weight_u - weight_src - γ * (refined.size[i] - size_src + c_u) * c_u
                if gain ≥ 0
                    push!(logprobs, 1/θ * gain)
                    push!(indexes, i)
                end
            end
            total_weights[communities] .= 0
            if isempty(indexes)
                continue
            end

            probs = exp.(logprobs .- logsumexp(logprobs))
            dst = indexes[sample(probs)]
            move_node!(refined, u => dst)

            for v in neighbors(refined, u)
                (v == u || v ∉ subset) && continue
                i = refined.membership[v]
                weight = refined.edge_weight[u,v]
                if i == dst
                    between_weights[src] -= weight
                    between_weights[dst] -= weight
                elseif i == src
                    between_weights[src] += weight
                    between_weights[dst] += weight
                else
                    between_weights[src] -= weight
                    between_weights[dst] += weight
                end
            end
        end
        for u in subset
            between_weights[refined.membership[u]] = 0
        end
        #@assert all(x == 0 for x in between_weights)
    end
    return drop_empty_communities!(refined)
end

function sample(probs::Vector{Float64})
    r = rand()
    p = 0.0
    i = 1
    while i < lastindex(probs)
        p += probs[i]
        if p > r
            return i
        end
        i += 1
    end
    return lastindex(probs)
end

function aggregate_graph(graph::PartitionedGraph)
    n = nc(graph)
    I = Int[]; J = Int[]; V = Float64[]
    cardinality = zeros(Int, n)
    connected = Int[]
    total_weights = zeros(Float64, n)
    for (i, community) in enumerate(graph.partition)
        cardinality[i] = graph.size[i]
        empty!(connected)
        for u in community, v in neighbors(graph, u)
            j = graph.membership[v]
            if i == j && u > v
                continue
            end
            if total_weights[j] == 0
                push!(connected, j)
            end
            total_weights[j] += graph.edge_weight[u,v]
        end
        for j in connected
            v = total_weights[j]
            push!(I, i)
            push!(J, j)
            push!(V, v)
            total_weights[j] = 0
        end
        #@assert all(x == 0 for x in total_weights)
    end
    return PartitionedGraph(sparse(I, J, V, n, n)::WeightMatrix, cardinality = cardinality)
end

function H(graph::PartitionedGraph, γ::Float64)
    quality = 0.0
    for (i, community) in enumerate(graph.partition)
        for u in community, v in community
            if u ≤ v
                quality += graph.edge_weight[u,v]
            end
        end
        n = graph.size[i]
        quality -= γ * div(n * (n - 1), 2)
    end
    return quality
end

function flatten(stack::Vector{Partition})
    k = lastindex(stack)
    result = copy(stack[k])
    k -= 1
    while k ≥ 1
        partition = stack[k]
        for (i, indices) in enumerate(result)
            result[i] = mapfoldl(i -> partition[i], vcat, indices, init = Int[])
        end
        k -= 1
    end
    sort!(result, by = length)
    foreach(sort!, result)
    return result
end

end # module


text:
#Demo 5: Logy plots
#-------------------------------------------------------------------------------
using InspectDR
using Colors

#==Input
===============================================================================#

#Constants
#-------------------------------------------------------------------------------
black = RGB24(0, 0, 0)
white = RGB24(1, 1, 1)
red = RGB24(1, 0, 0)
green = RGB24(0, 1, 0)
blue = RGB24(0, 0, 1)

#Input data
#-------------------------------------------------------------------------------
x = collect(1.0:1.0:100.0)
b = 2.0 #Base
k = 0.2 #Exponential coefficient

#==Generate plot
===============================================================================#
mplot = InspectDR.Multiplot(title="Log-Y Tests")
mplot.ncolumns = 1

plot_liny = InspectDR.Plot2D()
	plot_liny.axes = InspectDR.axes(:lin, :lin)
plot_logy = InspectDR.Plot2D()
	plot_logy.axes = InspectDR.axes(:lin, :log10)
	plot_logy.layout.grid = grid(vmajor=true, hmajor=true, hminor=true)

plotlist = [plot_liny, plot_logy]

for plot in plotlist
	wfrm = add(plot, x, b.^(k*x))
	wfrm.line = line(color=blue, width=3)

	add(mplot, plot)
	a = plot.annotation
	a.title = "y = $b^($(k)x)"
	a.xlabel = "x"
	a.ylabel = "y"
end

gplot = display(InspectDR.GtkDisplay(), mplot)

:DONE


text:
using SpectralMethodTools
using Test

@testset "SpectralMethodTools.jl" begin
    # Write your tests here.
end


text:
using FaultTolerantControl
const FTC = FaultTolerantControl
using UnPack

function test()
    multicopter = LeeHexacopterEnv()
    @unpack m, g, B = multicopter
    m = copy(m)
    g = copy(g)
    pos0 = zeros(3)
    controller = BacksteppingPositionControllerEnv(copy(m))
    allocator = PseudoInverseAllocator(copy(B))
    cs = FTC.BacksteppingConrtol_PseudoInverseCA_ControlSystem(controller, allocator)
    x0 = State(cs)(pos0, m, g)
    # simulation
    tf = 10.0
    prob, df = sim(x0,
                   apply_inputs(Dynamics!(cs); pos_cmd=[2, 1, 3], Ṫd=0.0);
                   tf=tf,
                  )
    df
end


text:
###
# Test using simulated data.  Only tests the Shrunken SAA with CV Methods
#
# For increasing K, simulate datasets and performance
# of various methods SAA, Alpha LOO, and Alpha OR
###
using Distributions, Random, DelimitedFiles
include("../src/JS_SAA_main.jl")

#supp_full, ps_full are d x K matrices with true info per problem
#adds new subproblems in order appear in files
#s is the service level, N is the average amount of data per problem
function cvtests(numRuns, K_grid, supp_full, ps_full, outPath, N_grid, s;
						usePoisson=false, seed=8675309, alpha_max = 180, alpha_len=120)
	Random.seed!(seed)
	Kmax = maximum(K_grid)
	@assert Kmax <= size(supp_full, 2) "K_grid exceeds available subproblems"
	@assert size(supp_full) == size(ps_full) "supp_full and ps_full have incompatible dimensions"

	d = size(supp_full, 1)

	#For safety, trim inputs to size Kmax
	supp_full = view(supp_full, 1:d, 1:Kmax)
	ps_full = view(ps_full, 1:d, 1:Kmax)

	p0 = ones(d)/d
	alpha_grid = range(0., stop=alpha_max, length=alpha_len)
	Gamma_grid = range(0, stop = (1-s), length=51)
	theta2_grid = range(1e-6, stop=3, length=20)
	mu_grid = range(1e-6, stop=1, length=20)

	#set up output file
	f = open("$(outPath).csv", "w")
	writedlm(f, ["Run" "K" "d" "N" "Method" "TruePerf" "time" "alpha"], ',')

	#generate all Kmax subproblems upfront and store in memory
	cs_full = JS.getNewsVendorCosts(supp_full, s, Kmax)
	xs_full = JS.genSSAAtrainers(supp_full, s)
	xsKS_full = JS.genKSTrainers(supp_full, s, Kmax, :crossVal)

	lam_full = ones(Kmax)

	for (iRun, N) = Iterators.product(1:numRuns, N_grid)
		#simulate data for the run
		Nhats_full = usePoisson ? rand(Poisson(N), Kmax) : N * ones(Int, Kmax)

		#for now, cludge the Nhatk == 0 issue
		Nhats_full[Nhats_full .== 0] .= 1
		mhats_full = JS.sim_path(ps_full, Nhats_full)

		for K in K_grid
			#Take views on evrything for simplicity
			lams = view(lam_full, 1:K)
			supp = view(supp_full, 1:d, 1:K)
			ps = view(ps_full, 1:d, 1:K)
			Nhats = view(Nhats_full, 1:K)
			mhats = view(mhats_full, 1:d, 1:K)
			cs = view(cs_full, 1:d, 1:K)
			xs = view(xs_full, 1:K)
			xsKS = view(xsKS_full, 1:K)

			#for data-driven shrinkage anchor
			#@VG Consider changing below to get_GM_anchor
			phat_avg = vec(mean(mhats ./ Nhats', dims=2))

			#Compute the full-info value once for reference
			if iRun == 1
				t =
				  @elapsed full_info = JS.zstar(xs, cs, ps, lams)
				writedlm(f, [1 K d N "FullInfo" full_info t 0.], ',')
			end

			#SAA
			t =
			  @elapsed perf_SAA = JS.zbar(xs, cs, mhats, ps, lams, (p0, 0.))
			writedlm(f, [iRun K d N "SAA" perf_SAA t 0.0], ',')

			#Gen the Oracle cost with 1/d anchor
			t =
			  @elapsed alphaOR, min_indx, or_alpha_curve = JS.oracle_alpha(xs, cs, mhats, ps, lams, p0, alpha_grid)
			writedlm(f, [iRun K d N "Oracle" or_alpha_curve[min_indx] t alphaOR], ',')

			#Gen the LOO cost with 1/d anchor
			t =
			  @elapsed alphaLOO, min_indx, looUnsc_curve = JS.loo_alpha(xs, cs, mhats, p0, alpha_grid)
			writedlm(f, [iRun K d N "LOO_unif" or_alpha_curve[min_indx] t alphaLOO], ',')

			#Gen the 5-Fold Cost with 1/d anchor
			t =
			  @elapsed alphaCV, min_indx, CVUnsc_curve = JS.cv_alpha(xs, cs, mhats, p0, alpha_grid, 5)
			writedlm(f, [iRun K d N "CV5_unif" or_alpha_curve[min_indx] t alphaCV], ',')

			#Gen the 2-Fold Cost with 1/d anchor
			t =
			  @elapsed alphaCV, min_indx, CVUnsc_curve = JS.cv_alpha(xs, cs, mhats, p0, alpha_grid, 2)
			writedlm(f, [iRun K d N "CV2_unif" or_alpha_curve[min_indx] t alphaCV], ',')

			#Gen the 10-Fold Cost with 1/d anchor
			t =
			  @elapsed alphaCV, min_indx, CVUnsc_curve = JS.cv_alpha(xs, cs, mhats, p0, alpha_grid, 10)
			writedlm(f, [iRun K d N "CV10_unif" or_alpha_curve[min_indx] t alphaCV], ',')

			#Gen the Oracle cost with GM Anchor
			t =
			  @elapsed alphaOR_GM, min_indx, or_alpha_curve_GM = JS.oracle_alpha(xs, cs, mhats, ps, lams, phat_avg, alpha_grid)
			writedlm(f, [iRun K d N "OraclePhat" or_alpha_curve_GM[min_indx] t alphaOR_GM], ',')

			#Gen the LOO cost with the GM Anchor
			t =
			  @elapsed alphaLOO, min_indx, looUnsc_curve = JS.loo_alpha(xs, cs, mhats, phat_avg, alpha_grid)
			writedlm(f, [iRun K d N "LOO_avg" or_alpha_curve_GM[min_indx] t alphaLOO], ',')
			temp_GM = looUnsc_curve[min_indx]

			#Gen the CV5 cost with the GM Anchor
			t =
			  @elapsed alphaCV, min_indx, CVUnsc_curve = JS.cv_alpha(xs, cs, mhats, phat_avg, alpha_grid, 5)
			writedlm(f, [iRun K d N "CV5_avg" or_alpha_curve_GM[min_indx] t alphaCV], ',')

			#Gen the CV2 cost with the GM Anchor
			t =
			  @elapsed alphaCV, min_indx, CVUnsc_curve = JS.cv_alpha(xs, cs, mhats, phat_avg, alpha_grid, 2)
			writedlm(f, [iRun K d N "CV2_avg" or_alpha_curve_GM[min_indx] t alphaCV], ',')

			#Gen the CV10 cost with the GM Anchor
			t =
			  @elapsed alphaCV, min_indx, CVUnsc_curve = JS.cv_alpha(xs, cs, mhats, phat_avg, alpha_grid, 10)
			writedlm(f, [iRun K d N "CV10_avg" or_alpha_curve_GM[min_indx] t alphaCV], ',')

		end  #end K Loop
		flush(f)
	end #end run loop
	close(f)
	"$(outPath).csv"
end #end function


text:
using Flux, GeometricFlux, LightGraphs, BenchmarkTools, CUDA
using DataFrames, Statistics, JLD2, SparseArrays

BenchmarkTools.ratio(::Missing, x) = Inf
BenchmarkTools.ratio(x, ::Missing) = 0.0
BenchmarkTools.ratio(::Missing, ::Missing) = missing

adjlist(g) = [neighbors(g, i) for i in 1:nv(g)]

function run_single_benchmark(N, c, D, CONV; gtype=:lg)
    @assert gtype ∈ [:lightgraph, :adjlist, :dense, :sparse]
    g = erdos_renyi(N, c / (N-1), seed=17)
    if gtype == :adjlist
        g = adjlist(g)
    elseif gtype == :dense
        g = Array(adjacency_matrix(g))
    elseif gtype == :sparse
        g = adjacency_matrix(g) # lightgraphs returns sparse adj mats
    end

    res = Dict()
    X = randn(Float32, D, N)
    res["FG"] = @benchmark FeaturedGraph($g, nf=$X)

    fg = FeaturedGraph(g, nf=X)
    fg_gpu = fg |> gpu

    m = CONV(D => D)
    m_gpu = m |> gpu
    try
        res["CPU"] = @benchmark $m($fg)
    catch
        res["CPU"] = missing
    end

    try
        res["GPU"] = @benchmark CUDA.@sync($m_gpu($fg_gpu)) teardown=(GC.gc(); CUDA.reclaim())
    catch
        res["GPU"] = missing
    end

    return res
end
"""
    run_benchmarks(;
        Ns = [10, 100, 1000, 10000],
        c = 6,
        D = 100)

Benchmark GNN layers on random regular graphs
of mean connectivity `c` and number of nodes in the list `Ns`.
`D` is the number of node features.
"""
function run_benchmarks(;
        Ns = [10, 100, 1000, 10000],
        c = 6.0,
        D = 100)

    df = DataFrame(N=Int[], c=Float64[], layer=String[], gtype=Symbol[],
                   time_fg=Any[], time_cpu=Any[], time_gpu=Any[]) |> allowmissing

    for gtype in [:lightgraph, :adjlist, :dense, :sparse]
        for N in Ns
            println("## GRAPH_TYPE = $gtype  N = $N")
            for CONV in [GCNConv, GraphConv, GATConv]
                res = run_single_benchmark(N, c, D, CONV; gtype)
                row = (; N = N,
                        c = c,
                        layer = "$CONV",
                        gtype = gtype,
                        time_fg = median(res["FG"]),
                        time_cpu = ismissing(res["CPU"]) ? missing : median(res["CPU"]),
                        time_gpu = ismissing(res["GPU"]) ? missing : median(res["GPU"]),
                    )
                push!(df, row)
            end
        end
    end
    df.gpu_to_cpu = ratio.(df.time_gpu, df.time_cpu)
    sort!(df, [:layer, :N, :c, :gtype])
    return df
end

# df = run_benchmarks()
# for g in groupby(df, :layer); println(g, "\n"); end

# @save "perf/perf_master_20210803_carlo.jld2" dfmaster=df
## or
# @save "perf/perf_pr.jld2" dfpr=df

function compare(dfpr, dfmaster; on=[:N, :c, :layer])
    df = outerjoin(dfpr, dfmaster; on=on, makeunique=true, renamecols = :_pr => :_master)
    df.pr_to_master_cpu = ratio.(df.time_cpu_pr, df.time_cpu_master)
    df.pr_to_master_gpu = ratio.(df.time_cpu_pr, df.time_gpu_master)
    return df[:,[:N, :c, :gtype_pr, :gtype_master, :layer, :pr_to_master_cpu, :pr_to_master_gpu]]
end

# @load "perf/perf_pr.jld2" dfpr
# @load "perf/perf_master.jld2" dfmaster
# compare(dfpr, dfmaster)


text:
# Routines for 2D grids

function Vandermonde(basis, quadrature)
    V = zeros(length(quadrature.locations), basis.order)
    for j=1:basis.order
        V[:,j] = basisFunction(basis, quadrature.locations, j)
    end
end

"Compute scaled warp function at order N based on rout
interpolation nodes"
function Warpfactor(basis, quadrature, rout)
    N = len(quadrature.locations)

    # Equidistant distribution
    eq = equidistant(N)

    # Compute V based on equidistant nodes
    Veq = Vandermonde(basis, eq)

    # Evaluate basis function at rout
    Pmat = zeros(basis.order, length(rout))
    for i=1:basis.order
        Pmat[i,:] = basisFunction(basis, rout, i)
    end
    Lmat = Veq' / Pmat

    # Compute warp factor
    warp = Lmat' * (quadrature.locations - req)

    # Scale factor

end

"Compute (x,y) nodes in equilateral triangle for polynomial
of order N"
function Nodes2D(N::Int)
    # Optimum correction to warp factor, to minimise
    # Lebesque constant for interpolant
    alpopt = [0.0000, 0.0000, 1.4152, 0.1001, 0.2751,
              0.9800, 1.0999, 1.2832, 1.3648, 1.4773,
              1.4959, 1.5743, 1.5770, 1.6223, 0.6258]

    # Set optimised parameter, alpha, depending on order
    alpha = N < 16 ? alpopt[N] : 5.0 / 3

    # Total number of nodes
    Np = (N + 1) * (N + 2) / 2

    # Create equidistributed nodes on equilateral triangle
    L1 = zeros(Np, 1)
    L3 = zeros(Np, 1)
    sk = 1
    for n = 1:N+1
        for m = 1:N+2-n
            L1[sk] = (n - 1) / N
            L3[sk] = (m - 1) / N
            sk += 1
        end
    end
    L2 = 1.0 - L1 - L3
    x = -L2 + L3
    y = (-L2 - L3 + 2 * L1)/sqrt(3.0)

    # Compute blending function at each node for each edge
    blend1 = 4 * L2 * L3
    blend2 = 4 * L1 * L3
    blend3 = 4 * L1 * L2

    # Amount of warp for each node, for each edge
    warpf1 = Warpfactor(N, L3 - L2)
    warpf2 = Warpfactor(N, L1 - L3)
    warpf3 = Warpfactor(N, L2 - L1)

end


text:
"""
Atomic solution for an ODE.

### Parameters

* `DT`: data type
* `TT`: time step type
* `AT`: array type
* `IT`: internal variable types

### Fields

* `t`: time of current time step
* `t̄`: time of previous time step
* `q`: current solution of q
* `q̄`: previous solution of q
* `q̃`: compensated summation error of q
* `v`: vector field of q
* `v̄`: vector field of q̄
* `internal`: internal variables of the integrator (e.g., internal stages of a Runge-Kutta methods or solver output)

### Constructors

```julia
AtomicSolutionODE{DT,TT,AT,IT}(nd, internal::IT)
AtomicSolutionODE{DT,TT,AT,IT}(t::TT, q::AT, internal::IT)
AtomicSolutionODE(DT, TT, AT, nd, internal::IT=NamedTuple())
AtomicSolutionODE(t::TT, q::AT, internal::IT=NamedTuple())
```

* `nd`: dimension of the state vector

"""
mutable struct AtomicSolutionODE{DT <: Number, TT <: Real, AT <: AbstractArray{DT}, IT <: NamedTuple} <: AtomicSolution{DT,TT,AT}
    t::TT
    t̄::TT

    q::AT
    q̄::AT
    q̃::AT

    v::AT
    v̄::AT

    internal::IT

    function AtomicSolutionODE{DT,TT,AT,IT}(nd, internal::IT) where {DT,TT,AT,IT}
        new(zero(TT), zero(TT),
            AT(zeros(DT, nd)), AT(zeros(DT, nd)), AT(zeros(DT, nd)),
            AT(zeros(DT, nd)), AT(zeros(DT, nd)),
            internal)
    end

    function AtomicSolutionODE{DT,TT,AT,IT}(t::TT, q::AT, internal::IT) where {DT,TT,AT,IT}
        new(zero(t), zero(t),
            zero(q), zero(q), zero(q),
            zero(q), zero(q),
            internal)
    end
end

AtomicSolutionODE(DT, TT, AT, nd, internal::IT=NamedTuple()) where {IT} = AtomicSolutionODE{DT,TT,AT,IT}(nd, internal)
AtomicSolutionODE(t::TT, q::AT, internal::IT=NamedTuple()) where {DT, TT, AT <: AbstractArray{DT}, IT} = AtomicSolutionODE{DT,TT,AT,IT}(t, q, internal)

function set_initial_conditions!(asol::AtomicSolutionODE, equ::AbstractEquationODE, i::Int=1)
    @assert i ≥ nsamples(equ)
    t, q = initial_conditions(equ)
    asol.t  = t
    asol.q .= q[i]
    asol.v .= 0
end

function set_solution!(asol::AtomicSolutionODE, sol)
    t, q = sol
    asol.t  = t
    asol.q .= q
    asol.v .= 0
end

function get_solution(asol::AtomicSolutionODE)
    (asol.t, asol.q)
end

function GeometricBase.reset!(asol::AtomicSolutionODE, Δt)
    asol.t̄  = asol.t
    asol.q̄ .= asol.q
    asol.v̄ .= asol.v
    asol.t += Δt
end

function update!(asol::AtomicSolutionODE{DT}, y::Vector{DT}) where {DT}
    for k in eachindex(y)
        update!(asol, y[k], k)
    end
end

function update!(asol::AtomicSolutionODE{DT}, y::DT, k::Union{Int,CartesianIndex}) where {DT}
    asol.q[k], asol.q̃[k] = compensated_summation(y, asol.q[k], asol.q̃[k])
end


text:
#!/usr/bin/env julia

function fact(n::Int)::Int64
    n >= 0 || throw(DomainError(n, "argument must be positive"))
    result = 1
    for i in 2:n
        result *= i
    end
    return result
end

for n in 5:-1:-2
    try
        println("fact($n) = $(fact(n))")
    catch err
        if isa(err, DomainError)
            println("sorry, problem with $(err.val): $(err.msg)")
        else
            rethrow(err)
        end
    end
end


text:
function get_absolute_accuracy(obj::UnivariateIntegrator)
    return jcall(obj, "getAbsoluteAccuracy", jdouble, ())
end

function get_evaluations(obj::UnivariateIntegrator)
    return jcall(obj, "getEvaluations", jint, ())
end

function get_iterations(obj::UnivariateIntegrator)
    return jcall(obj, "getIterations", jint, ())
end

function get_maximal_iteration_count(obj::UnivariateIntegrator)
    return jcall(obj, "getMaximalIterationCount", jint, ())
end

function get_minimal_iteration_count(obj::UnivariateIntegrator)
    return jcall(obj, "getMinimalIterationCount", jint, ())
end

function get_relative_accuracy(obj::UnivariateIntegrator)
    return jcall(obj, "getRelativeAccuracy", jdouble, ())
end

function integrate(obj::UnivariateIntegrator, arg0::jint, arg1::UnivariateFunction, arg2::jdouble, arg3::jdouble)
    return jcall(obj, "integrate", jdouble, (jint, UnivariateFunction, jdouble, jdouble), arg0, arg1, arg2, arg3)
end


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

# OpticalInterface.jl is included in Geometry.jl so is omitted here

include("OpticalRay.jl")
include("RayGenerator.jl")
include("Emitters/Emitters.jl")
include("Fresnel.jl")
include("Paraxial.jl")
include("ParaxialAnalysis/ParaxialAnalysis.jl")
include("Grating.jl")
include("LensAssembly.jl")
include("HierarchicalImage.jl")
include("OpticalSystem.jl")
include("Lenses.jl")
include("Eye.jl")
include("Fields.jl")


text:
using Documenter
using TensorPolynomialBases

makedocs(
    sitename = "TensorPolynomialBases.jl",
    format = Documenter.HTML(),
    modules = [TensorPolynomialBases],
    pages = ["Home" => "index.md","API"=>"pages/api.md"]
)

# Documenter can also automatically deploy documentation to gh-pages.
# See "Hosting Documentation" and deploydocs() in the Documenter manual
# for more information.
#=deploydocs(
    repo = "<repository url>"
)=#

deploydocs(repo="github.com/gridap/TensorPolynomialBases.jl.git")


text:
using Images

function erosion_(image::AbstractMatrix, mask_order::Number=3)

  tipo = typeof(image)

  try
    image = convert(Array{Float32}, image)
  catch er
    error(er," --- Problem converting the single band image into an array.")
  end

  d = size(image)

  img1 = zeros(d[1]-mask_order+1,d[2]-mask_order+1)

  for i = 1:(d[1]-mask_order+1)
    for j = 1:(d[2]-mask_order+1)
      ai=i
      ia=i+mask_order-1
      aj=j
      ja=j+mask_order-1

      img1[i,j] = minimum(image[ai:ia,aj:ja])
    end
  end

  img1 = convert(tipo, img1)

  return(img1)

end


text:
import JSON
export BrainFlowInputParams

@enum BoardIds begin

    PLAYBACK_FILE_BOARD = -3
    STREAMING_BOARD = -2
    SYNTHETIC_BOARD = -1
    CYTON_BOARD = 0
    GANGLION_BOARD = 1
    CYTON_DAISY_BOARD = 2
    GALEA_BOARD = 3
    GANGLION_WIFI_BOARD = 4
    CYTON_WIFI_BOARD = 5
    CYTON_DAISY_WIFI_BOARD = 6
    BRAINBIT_BOARD = 7
    UNICORN_BOARD = 8
    CALLIBRI_EEG_BOARD = 9
    CALLIBRI_EMG_BOARD = 10
    CALLIBRI_ECG_BOARD = 11
    FASCIA_BOARD = 12
    NOTION_1_BOARD = 13
    NOTION_2_BOARD = 14
    IRONBCI_BOARD = 15
    GFORCE_PRO_BOARD = 16
    FREEEEG32_BOARD = 17

end

BoardIdType = Union{BoardIds, Integer}

@enum IpProtocolType begin

    NONE = 0
    UDP = 1
    TCP = 2

end

@Base.kwdef mutable struct BrainFlowInputParams
    serial_port::String = ""
    mac_address::String = ""
    ip_address::String = ""
    ip_port::Int32 = 0
    ip_protocol::Int32 = Integer(NONE)
    other_info::String = ""
    timeout::Int32 = 0
    serial_number::String = ""
    file::String = ""
end

@brainflow_rethrow function get_sampling_rate(board_id::BoardIdType)
    val = Vector{Cint}(undef, 1)
    ccall((:get_sampling_rate, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{Cint}), Int32(board_id), val)
    value = val[1]
    return value
end

@brainflow_rethrow function get_num_rows(board_id::BoardIdType)
    val = Vector{Cint}(undef, 1)
    ccall((:get_num_rows, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{Cint}), Int32(board_id), val)
    # here dont need to add 1, last element included
    value = val[1]
    return value
end

@brainflow_rethrow function get_eeg_names(board_id::BoardIdType)
    names_string = Vector{Cuchar}(undef, 4096)
    len = Vector{Cint}(undef, 1)
    ccall((:get_eeg_names, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{UInt8}, Ptr{Cint}), Int32(board_id), names_string, len)
    sub_string = String(names_string)[1:len[1]]
    value = split(sub_string, ',')
    return value
end

@brainflow_rethrow function get_device_name(board_id::BoardIdType)
    names_string = Vector{Cuchar}(undef, 4096)
    len = Vector{Cint}(undef, 1)
    ccall((:get_device_name, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{UInt8}, Ptr{Cint}), Int32(board_id), names_string, len)
    sub_string = String(names_string)[1:len[1]]
    return sub_string
end

single_channel_function_names = (
    :get_timestamp_channel,
    :get_package_num_channel,
    :get_battery_channel,
    :get_marker_channel,
)

# generating the channels functions
for func_name = single_channel_function_names
    cglobal_expr = Meta.parse("cglobal((:$func_name, BOARD_CONTROLLER_INTERFACE))")
    @eval @brainflow_rethrow function $func_name(board_id::BoardIdType)
        channel = Vector{Cint}(undef, 1)
        ccall($cglobal_expr, Cint, (Cint, Ptr{Cint}), Int32(board_id), channel)
        # julia counts from 1
        @inbounds value = channel[1] + 1
        return value
    end
end

channel_function_names = (
    :get_eeg_channels,
    :get_exg_channels,
    :get_emg_channels,
    :get_ecg_channels,
    :get_eog_channels,
    :get_eda_channels,
    :get_ppg_channels,
    :get_accel_channels,
    :get_analog_channels,
    :get_gyro_channels,
    :get_other_channels,
    :get_temperature_channels,
    :get_resistance_channels,
)

# generating the channels functions
for func_name = channel_function_names
    cglobal_expr = Meta.parse("cglobal((:$func_name, BOARD_CONTROLLER_INTERFACE))")
    @eval @brainflow_rethrow function $func_name(board_id::BoardIdType)
        channels = Vector{Cint}(undef, 512)
        len = Vector{Cint}(undef, 1)
        ccall($cglobal_expr, Cint, (Cint, Ptr{Cint}, Ptr{Cint}), Int32(board_id), channels, len)
        # julia counts from 1
        @inbounds value = channels[1:len[1]] .+ 1
        return value
    end
end

struct BoardShim

    master_board_id::Int32
    board_id::Int32
    input_json::String

    function BoardShim(id::Integer, params::BrainFlowInputParams)
        master_id = id
        if id == Integer(STREAMING_BOARD) || id == Integer(PLAYBACK_FILE_BOARD)
            try
                master_id = parse(Int, params.other_info)
            catch
                throw(BrainFlowError("you need to provide master board id to other_info field of BrainFlowInputParams", Integer(INVALID_ARGUMENTS_ERROR)))
            end
        end
        new(master_id, id, JSON.json(params))
    end

end
BoardShim(id::BoardIds, params::BrainFlowInputParams) = BoardShim(Integer(id), params)
BoardShim(id) = BoardShim(id, BrainFlowInputParams())

@brainflow_rethrow function prepare_session(board_shim::BoardShim)
    ccall((:prepare_session, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{UInt8}), board_shim.board_id, board_shim.input_json)
end

@brainflow_rethrow function start_stream(board_shim::BoardShim, num_samples::Int, streamer_params::String)
    ccall((:start_stream, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{UInt8}, Cint, Ptr{UInt8}), num_samples, streamer_params,
            board_shim.board_id, board_shim.input_json)
end

function start_stream(board_shim::BoardShim)
    start_stream(board_shim, 450000, "")
end

@brainflow_rethrow function is_prepared(board_shim::BoardShim)
    val = Vector{Cint}(undef, 1)
    ccall((:is_prepared, BOARD_CONTROLLER_INTERFACE), Cint, (Ptr{Cint}, Cint, Ptr{UInt8}), val, board_shim.board_id, board_shim.input_json)
    value = Bool(val[1])
    return value
end

@brainflow_rethrow function get_board_data_count(board_shim::BoardShim)
    val = Vector{Cint}(undef, 1)
    ccall((:get_board_data_count, BOARD_CONTROLLER_INTERFACE), Cint, (Ptr{Cint}, Cint, Ptr{UInt8}), val, board_shim.board_id, board_shim.input_json)
    value = val[1]
    return value
end

@brainflow_rethrow function insert_marker(value::Float64, board_shim::BoardShim)
    ccall((:insert_marker, BOARD_CONTROLLER_INTERFACE), Cint, (Float64, Cint, Ptr{UInt8}), value, board_shim.board_id, board_shim.input_json)
end

@brainflow_rethrow function stop_stream(board_shim::BoardShim)
    ccall((:stop_stream, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{UInt8}), board_shim.board_id, board_shim.input_json)
end

@brainflow_rethrow function release_session(board_shim::BoardShim)
    ccall((:release_session, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{UInt8}), board_shim.board_id, board_shim.input_json)
end

@brainflow_rethrow function config_board(config::String, board_shim::BoardShim)
    resp_string = Vector{Cuchar}(undef, 4096)
    len = Vector{Cint}(undef, 1)
    ccall((:config_board, BOARD_CONTROLLER_INTERFACE), Cint, (Ptr{UInt8}, Ptr{UInt8}, Ptr{Cint}, Cint, Ptr{UInt8}),
            config, resp_string, len, board_shim.board_id, board_shim.input_json)
    sub_string = String(resp_string)[1:len[1]]
    return sub_string
end

@brainflow_rethrow function get_board_data(board_shim::BoardShim)
    data_size = get_board_data_count(board_shim)
    num_rows = get_num_rows(board_shim.master_board_id)
    val = Vector{Float64}(undef, num_rows * data_size)
    ccall((:get_board_data, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{Float64}, Cint, Ptr{UInt8}),
            data_size, val, board_shim.board_id, board_shim.input_json)
    value = transpose(reshape(val, (data_size, num_rows)))
    return value
end

@brainflow_rethrow function get_current_board_data(num_samples::Integer, board_shim::BoardShim)
    data_size = Vector{Cint}(undef, 1)
    num_rows = get_num_rows(board_shim.master_board_id)
    val = Vector{Float64}(undef, num_rows * num_samples)
    ccall((:get_current_board_data, BOARD_CONTROLLER_INTERFACE), Cint, (Cint, Ptr{Float64}, Ptr{Cint}, Cint, Ptr{UInt8}),
            num_samples, val, data_size, board_shim.board_id, board_shim.input_json)
    value = transpose(reshape(val[1:data_size[1] * num_rows], (data_size[1], num_rows)))
    return value
end

max_stars_repo_name: evaesteban/brainflow

max_issues_repo_name: evaesteban/brainflow

max_forks_repo_name: xloem/brainflow


text:
"""
    `x, u, traj_new, Vx, Vxx, cost, trace = iLQGkl(dynamics,costfun,derivs, x0, traj_prev, model;
        constrain_per_step = false,
        kl_step            = 0,
        lims               = [],                    # Control signal limits ::Matrix ∈ R(m,2)
        tol_fun            = 1e-7,
        tol_grad           = 1e-4,
        max_iter           = 50,
        print_head         = 10,                    # Print headers this often
        print_period       = 1,                     # Print this often
        reduce_ratio_min   = 0,                     # Not used ATM
        diff_fun           = -,
        verbosity          = 2,                     # ∈ (0,3)
        plot_fun           = x->0,                  # Not used
        cost               = [],                    # Supply if pre-rolled trajectory supplied
        ηbracket           = [1e-8,1,1e16],         # dual variable bracket [min_η, η, max_η]
        del0               = 0.0001,                # Start of dual variable increase
        gd_alpha           = 0.01                   # Step size in GD (ADAMOptimizer) when constrain_per_step is true
        )`

Solves the iLQG problem with constraints on control signals `lims` and bound on the KL-divergence `kl_step` from the old trajectory distribution `traj_prev::GaussianPolicy`.

To solve the maximum entropy problem, use controller `controller(xi,i)  = u[:,i] + K[:,:,i]*(xi-x[:,i]) + chol(Σ)*randn(m)` where `K` comes from `traj_new`. Note that multiplying the cost by a constant changes the relative weight between the cost term and the entropy term, i.e., higher cost produces less noise through chol(Σ) since (Σ = Qᵤᵤ⁻¹).
"""
function iLQGkl(dynamics,costfun,derivs, x0, traj_prev, model;
    constrain_per_step = false,
    kl_step            = 1,
    lims               = [],
    tol_fun            = 1e-7,
    tol_grad           = 1e-4,
    max_iter           = 50,
    print_head         = 10,
    print_period       = 1,
    reduce_ratio_min   = 0,
    diff_fun           = -,
    verbosity          = 2,
    plot_fun           = x->0,
    cost               = [],
    ηbracket           = [1e-8,1,1e16], # min_η, η, max_η
    del0               = 0.0001,
    gd_alpha           = 0.01
    )
    debug("Entering iLQG")
    local fx,fu,fxx,fxu,fuu,cx,cu,cxx,cxu,cuu,xnew,unew,costnew,sigmanew,g_norm,Vx,Vxx,dV

    # --- initial sizes and controls
    u            = copy(traj_prev.k) # initial control sequence
    n            = size(x0, 1) # dimension of state vector
    m,N          = size(u) # dimension of control vector and number of state transitions
    traj_new     = GaussianPolicy(Float64)
    k_old = copy(traj_prev.k)
    traj_prev.k *= 0 # We are adding new k to u, so must set this to zero for correct kl calculations
    ηbracket     = copy(ηbracket) # Because we do changes in this Array
    if constrain_per_step
        ηbracket = ηbracket.*ones(1,N)
        kl_step = kl_step*ones(N)
    end
    η = view(ηbracket,2,:)

    # --- initialize trace data structure
    trace = MVHistory()

    # --- initial trajectory
    debug("Checking initial trajectory")
    if size(x0,2) == N
        debug("# pre-rolled initial forward pass, initial traj provided")
        x        = x0
        diverge  = false
        isempty(cost) && error("Initial trajectory supplied, initial cost must also be supplied")
    else
        error("pre-rolled initial trajectory must be of correct length (size(x0,2) == N)")
    end

    trace(:cost, 0, sum(cost))

    # constants, timers, counters
    Δcost              = 0.
    expected_reduction = 0.
    divergence         = 0.
    step_mult          = 1.
    iter               = 0
    last_head          = print_head
    t_start            = time()
    verbosity > 0 && @printf("\n---------- begin iLQG ----------\n")
    satisfied          = false # Indicating KL-constraint satisfied

    # ====== STEP 1: differentiate dynamics and cost along new trajectory
    _t = @elapsed fx,fu,fxx,fxu,fuu,cx,cu,cxx,cxu,cuu = derivs(x, u)
    trace(:time_derivs, 0, _t)

    reduce_ratio     = 0.
    kl_cost_terms    = (∇kl(traj_prev), ηbracket) # This tuple is sent into back_pass, elements in ηbracket are mutated.
    for outer iter = 1:(constrain_per_step ? 0 : max_iter) # Single KL constraint
        diverge = 1
        # ====== STEP 2: backward pass, compute optimal control law and cost-to-go
        back_pass_done = false
        while diverge > 0 # Done when regularization (through 1/η) for Quu is high enough
            # debug("Entering back_pass with η=$ηbracket")
            # η is the only regularization when optimizing KL, hence λ = 0 and regType arbitrary
            _t = @elapsed diverge, traj_new,Vx, Vxx,dV =  back_pass_gps(cx,cu,cxx,cxu,cuu,fx,fu, lims,x,u,kl_cost_terms) # Set λ=0 since we use η
            trace(:time_backward, iter, _t)

            if diverge > 0
                ηbracket[2] += del0 # η increased, used in back_pass through kl_cost_terms
                # Higher η downweights the original Q function and upweights KL-cost terms
                del0 *= 2
                if verbosity > 2
                    println("Inversion failed at timestep $diverge. η-bracket: ", ηbracket)

                end
                # if ηbracket[2] >  0.999ηbracket[3] #  terminate ?
                #     if verbosity > 0
                #         @printf("\nEXIT: η > ηmax (back_pass failed)\n")
                #         @show kl_cost_terms[1][5]
                #         @show diverge
                #         @show traj_new.Σ[:,:,diverge]
                #         @show traj_new.Σi[:,:,diverge]
                #         @show ηbracket[2], 0.999ηbracket[3]
                #     end
                #     break
                # end

            end
        end

        #  check for termination due to small gradient
        g_norm = mean(maximum(abs.(traj_new.k) ./ (abs.(u) .+1),dims=1))
        trace(:grad_norm, iter, g_norm)

        # ====== STEP 3: Forward pass

        _t = @elapsed begin
            # debug("#  entering forward_pass")
            xnew,unew,costnew = forward_pass(traj_new, x0[:,1] ,u, x,1,dynamics,costfun, lims, diff_fun)
            sigmanew = forward_covariance(model, x, u, traj_new)
            traj_new.k .+= traj_prev.k # unew = k_new + k_old + Knew*Δx, this doesn't matter since traj_prev.k set to 0 above
            Δcost    = sum(cost) - sum(costnew)
            expected_reduction = -(dV[1] + dV[2]) # According to second order approximation

            reduce_ratio = Δcost/expected_reduction

            # calc_η modifies the dual variables η according to current constraint_violation
            ηbracket, satisfied, divergence = calc_η(xnew,x,sigmanew,ηbracket, traj_new, traj_prev, kl_step)
        end
        trace(:time_forward, iter, _t)
        debug("Forward pass done: η: $ηbracket")

        # ====== STEP 4: accept step (or not), print status

        #  print headings
        if verbosity > 1 && iter % print_period == 0
            if last_head == print_head
                last_head = 0
                @printf("%-12s", "iteration     est. cost    reduction     expected    gradient    log10(η)    divergence      entropy\n")
            end
            @printf("%-14d%-14.6g%-14.3g%-14.3g%-12.3g%-12.2f%-14.3g%-12.3g\n",
            iter, sum(costnew), Δcost, expected_reduction, g_norm, log10(mean(η)), mean(divergence), entropy(traj_new))
            last_head += 1
        end
        #  update trace
        trace(:alpha, iter, 1)
        trace(:improvement, iter, Δcost)
        trace(:cost, iter, sum(costnew))
        trace(:reduce_ratio, iter, reduce_ratio)
        trace(:divergence, iter, mean(divergence))
        trace(:η, iter, ηbracket[2])

        # Termination checks
        # if g_norm <  tol_grad && divergence-kl_step > 0 # In this case we're only going to get even smaller gradients and might as well quit
        #     verbosity > 0 && @printf("\nEXIT: gradient norm < tol_grad while constraint violation too large\n")
        #     break
        # end
        if satisfied # KL-constraint is satisfied and we're happy (at least if Δcost is positive)
            plot_fun(x)
            verbosity > 0 && @printf("\nSUCCESS: abs(KL-divergence) < kl_step\n")
            break
        end
        if ηbracket[2] >  0.999ηbracket[3]
            verbosity > 0 && @printf("\nEXIT: η > ηmax\n")
            break
        end
        # graphics(xnew,unew,cost,traj_new.K,Vx,Vxx,fx,fxx,fu,fuu,trace[1:iter],0)
    end # !constrain_per_step

    if constrain_per_step # This implements the gradient descent procedure for η
        optimizer = ADAMOptimizer(kl_step, α=gd_alpha)
        for outer iter = 1:max_iter
            diverge = 1
            del = del0*ones(N)
            while diverge > 0
                diverge, traj_new,Vx, Vxx,dV = back_pass_gps(cx,cu,cxx,cxu,cuu,fx,fu, lims,x,u,kl_cost_terms)
                if diverge > 0
                    delind = diverge # This is very inefficient since back_pass only returs a single diverge per call.
                    ηbracket[2,delind] .+= del[delind]
                    del[delind] *= 2
                    if verbosity > 2; println("Inversion failed at timestep $diverge. η-bracket: ", mean(η)); end
                    if all(ηbracket[2,:] .>  0.999ηbracket[3,:])
                        # TODO: This termination criteria could be improved
                        verbosity > 0 && @printf("\nEXIT: η > ηmax\n")
                        break
                    end
                end
            end

            xnew,unew,costnew = forward_pass(traj_new, x0[:,1] ,u, x,1,dynamics,costfun, lims, diff_fun)
            sigmanew = forward_covariance(model, x, u, traj_new)
            traj_new.k .+= traj_prev.k # unew = k_new + k_old + Knew*Δx
            Δcost                 = sum(cost) - sum(costnew)
            expected_reduction    = -(dV[1] + dV[2])
            reduce_ratio          = Δcost/expected_reduction
            divergence            = kl_div_wiki(xnew,x,sigmanew, traj_new, traj_prev)
            constraint_violation  = divergence - kl_step
            lη                    = log.(η) # Run GD in log-space (much faster)
            η                    .= exp.(optimizer(lη, -constraint_violation, iter))
            # η                    .= optimizer(η, -constraint_violation, iter)
            # println(maximum(constraint_violation), " ", extrema(η), " ", indmax(constraint_violation))
            # println(round.(constraint_violation,4))
            η                    .= clamp.(η, ηbracket[1,:], ηbracket[3,:])
            g_norm                = mean(maximum(abs.(traj_new.k) ./ (abs.(u) .+1),dims=1))
            trace(:grad_norm, iter, g_norm)
            # @show maximum(constraint_violation)
            if all(divergence .< 2*kl_step) && mean(constraint_violation) < 0.1*kl_step[1]
                satisfied = true
                break
            end
            if verbosity > 1 && iter % print_period == 0
                if last_head == print_head
                    last_head = 0
                    @printf("%-12s", "iteration     est. cost    reduction     expected    log10(η)    divergence      entropy\n")
                end
                @printf("%-14d%-14.6g%-14.3g%-14.3g%-12.3f%-12.3g%-14.3g\n",
                iter, sum(costnew), Δcost, expected_reduction, mean(log10.(η)), mean(divergence), entropy(traj_new))
                last_head += 1
            end
        end
    end

    iter ==  max_iter &&  verbosity > 0 && @printf("\nEXIT: Maximum iterations reached.\n")
    # if costnew < 1.1cost # In this case we made an (approximate) improvement under the model and accept the changes
        x,u,cost  = xnew,unew,costnew
        traj_new.k = copy(u)

    # else
    #     traj_new = traj_prev
    #     verbosity > 0 && println("Cost (under model) increased, did not accept changes to u")
    # end
    traj_prev.k = k_old
    any((divergence .> kl_step) .& (abs.(divergence - kl_step) .> 0.1*kl_step)) && @warn("KL divergence too high for some time steps when done")
    verbosity > 0 && print_timing(trace,iter,t_start,cost,g_norm,mean(ηbracket[2,:]))

    return x, u, traj_new, Vx, Vxx, cost, trace
end


text:
using Documenter, Optim

# use include("Rosenbrock.jl") etc

# assuming linux.
#run('mv ../LICENSE.md ./LICENSE.md')
#run('mv ../CONTRIBUTING.md ./dev/CONTRIBUTING.md')
makedocs(
    doctest = false
)

deploydocs(
    repo = "github.com/JuliaOpt/Optim.jl.git",
    julia = "0.5"
)


text:
import JSON
using ArgCheck: @argcheck
using BangBang
using BenchmarkTools: @benchmarkable, BenchmarkTools
using ProgressLogging: @progress, @withprogress
using Random: MersenneTwister
using Referenceables: referenceable
using ThreadsX

ntproduct(; kwargs...) =
    Base.Generator(NamedTuple{keys(kwargs)}, Iterators.product(values(kwargs)...))

@inline function setavg!(a, b, c)
    a[] = (b + c) / 2
end

const BENCHMARKS = Dict{Symbol,NamedTuple{(:prepare, :run, :paramasdata, :paramaxes)}}(
    :sort => (
        prepare = (p, rng) -> rand(rng, p.dist.value, p.datasize),
        run =
            (p, r) -> ThreadsX.sort!(
                r;
                alg = p.alg.value,
                basesize = p.basesize,
                smallsize = p.smallsize,
            ),
        paramasdata =
            p -> (
                datasize = p.datasize,
                distribution = p.dist.label,
                alg = p.alg.label,
                basesize = p.basesize,
                smallsize = p.smallsize,
            ),
        paramaxes = (
            datasize = [1_000_000, 10_000_000],
            dist = [
                (label = "wide", value = Float64),
                (label = "narrow", value = 0:0.01:1),
            ],
            alg = [
                (label = "ThreadsX.MergeSort", value = ThreadsX.MergeSort),
                (label = "ThreadsX.QuickSort", value = ThreadsX.QuickSort),
            ],
            basesize = [nothing],
            # basesize = [nothing, 20_000, 10_000, 5_000],
            smallsize = [nothing],
        ),
    ),
    :foreach_symmetrize => (
        prepare = function (p, rng)
            B = randn(rng, (p.datasize, p.datasize))
            A = similar(B)
            return (A = A, B = B)
        end,
        run = (p, r) -> ThreadsX.foreach(setavg!, referenceable(r.A), r.B, r.B'),
        paramasdata = identity,
        paramaxes = (datasize = [2000, 6000],),
    ),
    :sum_sin => (
        prepare = (p, rng) -> 1:p.datasize,
        run = (p, r) -> ThreadsX.sum(sin, r),
        paramasdata = identity,
        paramaxes = (datasize = [1_000_000, 10_000_000],),
    ),
    :findfirst => (
        prepare = function (p, rng)
            xs = rand(rng, p.datasize)
            xs[2*end÷3] = -1
            xs
        end,
        run = (p, r) -> ThreadsX.findfirst(==(-1), r),
        paramasdata = identity,
        paramaxes = (datasize = [2^22, 2^26],),
    ),
    :unique => (
        prepare = (p, rng) -> rand(rng, 1:10, p.datasize),
        run = (p, r) -> ThreadsX.unique(r),
        paramasdata = identity,
        paramaxes = (datasize = [1_000_000, 10_000_000],),
    ),
)

# Workaround custom JSON lowering
asdict(x) = x
asdict(x::BenchmarkTools.Trial) = _asdict(x)
asdict(x::BenchmarkTools.Parameters) = _asdict(x)
_asdict(x) = Dict(n => asdict(getfield(x, n)) for n in fieldnames(typeof(x)))

function main(
    ARGS = ARGS;
    benchmark_definitions = BENCHMARKS,
    scriptname::AbstractString = basename(@__FILE__),
    tags::AbstractVector{<:AbstractString} = ["nthreads"],
)
    @argcheck length(ARGS) == 1
    outputstem, = ARGS

    benchmarks = []
    @progress "Benchmark" for (benchname, bench) in benchmark_definitions
        @info "Benchmark: $benchname"

        results = []
        @progress "Sweep" for p in ntproduct(; bench.paramaxes...)
            @info "Parameter: $(bench.paramasdata(p))"
            b = @benchmarkable(
                $(bench.run)(p, r),
                setup = begin
                    p = $p
                    rng = $(MersenneTwister(1234))
                    r = $(bench.prepare)(p, rng)
                end,
                seconds = 20.0,
            )
            BenchmarkTools.warmup(b)
            BenchmarkTools.tune!(b; verbose = true)
            trial = run(b; verbose = true)
            push!(
                results,
                Dict(:trial => asdict(trial), :parameter => bench.paramasdata(p)),
            )
        end

        push!(benchmarks, Dict(:benchname => benchname, :results => results))
    end

    outputpath = "$outputstem.json"
    @info "Writing results to `$outputpath`"
    open(outputpath, write = true) do io
        JSON.print(
            io,
            Dict(
                :script => scriptname,
                :tags => append!([scriptname], tags),
                :benchmarks => benchmarks,
                :VERSION => string(VERSION),
                :nthreads => Threads.nthreads(),
            ),
        )
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


text:
#get

function get_int_param(m::Model,param::Int)

	ipar = convert(Cint,param)
	igval = Array(Cint,1)

	ret = @xprs_ccall(getintcontrol, Cint, (Ptr{Void},Cint,Ptr{Cint}),
		m.ptr_model, ipar, igval)
	if ret != 0
		throw(XpressError(m))
	end
	convert(Int, igval[1])
end

function get_dbl_param(m::Model,param::Int)

	ipar = convert(Cint,param)
	dgval = Array(Cdouble,1)

	ret = @xprs_ccall(getdblcontrol, Cint, (Ptr{Void},Cint,Ptr{Cdouble}),
		m.ptr_model, ipar, dgval)
	if ret != 0
		throw(XpressError(m))
	end
	convert(Float64, dgval[1])
end

function get_str_param(m::Model,param::Int)

	ipar = convert(Cint,param)
	cgval = Array(Cchar, 256)

	ret = @xprs_ccall(getstrcontrol, Cint, (Ptr{Void},Cint,Ptr{Cchar}),
		m.ptr_model, ipar, cgval)
	if ret != 0
		throw(XpressError(m))
	end

	unsafe_string(pointer(out))
end

#set

function set_int_param(m::Model,ipar::Int,isval::Int)

	ipar = convert(Cint,ipar)
	isval = convert(Cint,isval)

	ret = @xprs_ccall(setintcontrol, Cint, (Ptr{Void},Cint,Cint),
		m.ptr_model, ipar, isval)
	if ret != 0
		throw(XpressError(m))
	end

	nothing
end

function set_dbl_param(m::Model,ipar::Int,dsval::Float64)

	ipar = convert(Cint,ipar)
	dsval = convert(Cdouble,dsval)

	ret = @xprs_ccall(setdblcontrol, Cint, (Ptr{Void},Cint,Cdouble),
		m.ptr_model, ipar, dsval)
	if ret != 0
		throw(XpressError(m))
	end

	nothing
end

function set_str_param(m::Model,ipar::Int,csval::AbstractString)

	ipar = convert(Cint,ipar)
	csval = convert(ASCIIString,csval)

	ret = @xprs_ccall(setstrcontrol, Cint, (Ptr{Void},Cint,Ptr{Cchar}),
		m.ptr_model, ipar, csval)
	if ret != 0
		throw(XpressError(m))
	end

	nothing
end

function getparam(m::Model,param::Int)
    if convert(Int,param) in XPRS_INT_CONTROLS
        return get_int_param(m, param)
    elseif convert(Int,param) in XPRS_DBL_CONTROLS
        return get_dbl_param(m, param)
    elseif convert(Int,param) in XPRS_STR_CONTROLS
        return get_str_param(m, param)
    else
        error("Unrecognized parameter number: $(param).")
    end
end

function setparam!(m::Model,param::Int,val::Any)
    if convert(Int,param) in XPRS_INT_CONTROLS
        set_int_param(m, param, val)
    elseif convert(Int,param) in XPRS_DBL_CONTROLS
        set_dbl_param(m, param, val)
    elseif convert(Int,param) in XPRS_STR_CONTROLS
        set_str_param(m, param, val)
    else
        error("Unrecognized parameter number: $(param).")
    end
end

function setparams!(m::Model;args...)
	for (param,val) in args
		setparam!(m,XPRS_CONTROLS_DICT[param],val)
	end
end


text:
using Fast12thRoot

function accuracystats(N::Int, interval::Tuple; seed=123456789)
    srand(seed)
    a = interval[1]
    s = (interval[2]-interval[1])
    epsrng = [0.5,0.51,0.52,0.55,0.6,0.65,0.7,0.75,1.0]
    abserr = Array{Float64}(3)
    relerr = Array{Float64}(3)

    epshist_abs = zeros(Float64,(3,51))
    epshist_rel = zeros(Float64,(3,51))

    fname=[:base, :fast, :pow]
    println("Accuracy analysis for interval: [$(interval[1]), $(interval[2])]")
    println("using $N random numbers with initial seed $seed")

    for i=1:N
        # if mod(i,500000) == 0
        #     println("processed: $i")
        # end
        rv = a + s*rand()
        # "true" value
        bigval = big(rv)^(one(BigFloat)/big"12.0")
        # cbrt(sqrt(sqrt(x)))
        baseval = cbrt(sqrt(sqrt(rv)))
        # fast: rational polynomial
        fastval = fast12throot(rv)
        # call to power function [x^a = exp(aln(x))]
        powval = rv^(1/12)

        # absolute errors versus "true" bigval
        abserr[1] = abs(bigval-baseval)
        abserr[2] = abs(bigval-fastval)
        abserr[3] = abs(bigval-powval)

        # relative errors versus "true" bigval
        relerr[1] = abs((bigval - baseval)/bigval)
        relerr[2] = abs((bigval - fastval)/bigval)
        relerr[3] = abs((bigval - powval)/bigval)

        for i=1:3
            idx = 0
            epsfac0 = 0.5
            for epsfac in epsrng
                idx += 1
                if abserr[i]<= epsfac*eps(Float64)
                    epshist_abs[i,idx] += 1
                end
                if relerr[i]<= epsfac*eps(Float64)
                    epshist_rel[i,idx] += 1
                end
            end
        end
    end
    #println("123456789012345678901234567890123456789012345678901234567890123456789012345678901234567890")
    println("======================================================================================================================")
    println("|| error       ||                relative error                  ||                 absolute error                  ||")
    println("|| range       ||    pow        |    fast        |    simple     ||    pow        |    fast        |    simple      ||")
    println("||=============||===============|================|===============||===============|================|================||")
    for idx=1:length(epsrng)
        @printf("|| <= %.2feps  ||  %010.6f%%  |  %010.6f%%   |  %010.6f%%  ||  %010.6f%%  |  %010.6f%%   |  %010.6f%%   ||\n", epsrng[idx], epshist_rel[3,idx]*100/N,epshist_rel[2,idx]*100/N,epshist_rel[1,idx]*100/N, epshist_abs[3,idx]*100/N,epshist_abs[2,idx]*100/N,epshist_abs[1,idx]*100/N)
    end
    println("======================================================================================================================")
    println()

end

#accuracystats(5_000_000, (0.0,0.5))
#accuracystats(5_000_000, (0.5,1.0))
accuracystats(1_000_000, (0.0,1.0))
accuracystats(1_000_000, (1.0,2.0))
#accuracystats(1_000_000, (2.0,5.0))
#accuracystats(1_000_000, (5.0,10.0))
#accuracystats(1_000_000, (10.0,100.0))
accuracystats(1_000_000, (100.0,1000.0))
accuracystats(1_000_000, (1000.0,1_000_000.0))


text:
# TODO [tuples]
# - document eval_basis and move it into JuLIP

import JuLIP: energy!, forces!, virial!, alloc_temp, alloc_temp_d
import JuLIP.Potentials: evaluate!, evaluate_d!

function eval_basis! end
function eval_basis_d! end

eval_basis(B, x, args...) =
   eval_basis!(alloc_B(B, x), alloc_temp(B, x), B, x, args...)

function eval_basis_d(B, x, args...)
   b = alloc_B(B, x)
   db = alloc_dB(B, x)
   tmp = alloc_temp_d(B, x)
   eval_basis_d!(b, db, tmp, B, x, args...)
   return b, db
end

function alloc_B end
function alloc_dB end

function transform end
function transform_d end
function fcut end
function fcut_d end

max_stars_repo_name: cortner/SHIPs

max_issues_repo_name: cortner/SHIPs

max_forks_repo_name: cortner/SHIPs


text:
struct Hold{Any} <: State{Any}
    name::Symbol
    alias::Union{Symbol,Nothing}
end

Hold(; _name, _alias, _...) = begin
    Hold{Any}(_name, _alias)
end

constructortags(::Val{:Hold}) = ()

value(s::Hold) = error("cannot read variable on hold: $(s.name) $(isnothing(s.alias) ? "" : "($(s.alias))")")
store!(s::Hold, _) = error("cannot store variable on hold: $(s.name) $(isnothing(s.alias) ? "" : "($(s.alias))")")

genvartype(v::VarInfo, ::Val{:Hold}; _...) = @q Hold{Any}

gendefault(v::VarInfo, ::Val{:Hold}) = nothing


text:
@testset "with_lengthscale" begin
    @testset "ScaleTransform" begin
        l = exp(rand())
        kernel = @inferred(with_lengthscale(SqExponentialKernel(), l))

        @test kernel isa TransformedKernel{<:SqExponentialKernel,<:ScaleTransform}
        @test kernel.transform.s[1] ≈ inv(l)
    end

    @testset "ARDTransform" begin
        l = map(exp, rand(5))
        kernel = @inferred(with_lengthscale(SqExponentialKernel(), l))

        @test kernel isa TransformedKernel{<:SqExponentialKernel,<:ARDTransform}
        @test kernel.transform.v ≈ map(inv, l)
    end
end


text:
module Visualisations

# nothing to see here; just some runner scripts.

end


text:
module TransitionTest

using Test
using ForneyLab
using ForneyLab: outboundType, isApplicable
using ForneyLab: SPTransitionOutNPP, SPTransitionIn1PNP, SPTransitionOutNCP, SPTransitionIn1CNP, VBTransitionOut, VBTransitionIn1, VBTransitionA, SVBTransitionOutVCD, SVBTransitionIn1CVD, SVBTransitionADV, MTransitionCCD, MTransitionCCN

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

@testset "SPTransitionOutNPP" begin
    @test SPTransitionOutNPP <: SumProductRule{Transition}
    @test outboundType(SPTransitionOutNPP) == Message{Categorical}
    @test isApplicable(SPTransitionOutNPP, [Nothing, Message{PointMass}, Message{PointMass}])
    @test !isApplicable(SPTransitionOutNPP, [Message{PointMass}, Nothing, Message{PointMass}])

    @test ruleSPTransitionOutNPP(nothing, Message(Multivariate, PointMass, m=[0.1, 0.4, 0.5]), Message(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == Message(Univariate, Categorical, p=[0.3660714285714285, 0.27678571428571425, 0.35714285714285715])
end

@testset "SPTransitionIn1PNP" begin
    @test SPTransitionIn1PNP <: SumProductRule{Transition}
    @test outboundType(SPTransitionIn1PNP) == Message{Categorical}
    @test isApplicable(SPTransitionIn1PNP, [Message{PointMass}, Nothing, Message{PointMass}])

    @test ruleSPTransitionIn1PNP(Message(Multivariate, PointMass, m=[0.1, 0.4, 0.5]), nothing, Message(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == Message(Univariate, Categorical, p=[0.23000000000000004, 0.43, 0.33999999999999997])
end

@testset "SPTransitionOutNCP" begin
    @test SPTransitionOutNCP <: SumProductRule{Transition}
    @test outboundType(SPTransitionOutNCP) == Message{Categorical}
    @test isApplicable(SPTransitionOutNCP, [Nothing, Message{Categorical}, Message{PointMass}])
    @test !isApplicable(SPTransitionOutNCP, [Message{Categorical}, Nothing, Message{PointMass}])

    @test ruleSPTransitionOutNCP(nothing, Message(Univariate, Categorical, p=[0.1, 0.4, 0.5]), Message(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == Message(Univariate, Categorical, p=[0.3660714285714285, 0.27678571428571425, 0.35714285714285715])
end

@testset "SPTransitionIn1CNP" begin
    @test SPTransitionIn1CNP <: SumProductRule{Transition}
    @test outboundType(SPTransitionIn1CNP) == Message{Categorical}
    @test isApplicable(SPTransitionIn1CNP, [Message{Categorical}, Nothing, Message{PointMass}])

    @test ruleSPTransitionIn1CNP(Message(Univariate, Categorical, p=[0.1, 0.4, 0.5]), nothing, Message(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == Message(Univariate, Categorical, p=[0.23000000000000004, 0.43, 0.33999999999999997])
end

@testset "VBTransitionOut" begin
    @test VBTransitionOut <: NaiveVariationalRule{Transition}
    @test outboundType(VBTransitionOut) == Message{Categorical}
    @test isApplicable(VBTransitionOut, [Nothing, ProbabilityDistribution, ProbabilityDistribution])

    @test ruleVBTransitionOut(nothing, ProbabilityDistribution(Univariate, Categorical, p=[0.1, 0.4, 0.5]), ProbabilityDistribution(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == Message(Univariate, Categorical, p=[0.2994385327821292, 0.3260399327754116, 0.37452153444245917])
end

@testset "VBTransitionIn1" begin
    @test VBTransitionIn1 <: NaiveVariationalRule{Transition}
    @test outboundType(VBTransitionIn1) == Message{Categorical}
    @test isApplicable(VBTransitionIn1, [ProbabilityDistribution, Nothing, ProbabilityDistribution])

    @test ruleVBTransitionIn1(ProbabilityDistribution(Univariate, Categorical, p=[0.1, 0.4, 0.5]), nothing, ProbabilityDistribution(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == Message(Univariate, Categorical, p=[0.208905960977562, 0.4255474215521219, 0.365546617470316])
end

@testset "VBTransitionA" begin
    @test VBTransitionA <: NaiveVariationalRule{Transition}
    @test outboundType(VBTransitionA) == Message{Dirichlet}
    @test isApplicable(VBTransitionA, [ProbabilityDistribution, ProbabilityDistribution, Nothing])

    @test ruleVBTransitionA(ProbabilityDistribution(Univariate, Categorical, p=[0.2, 0.7, 0.1]), ProbabilityDistribution(Univariate, Categorical, p=[0.1, 0.4, 0.5]), nothing) == Message(MatrixVariate, Dirichlet, a=[1.02 1.08 1.1; 1.07 1.28 1.35; 1.01 1.04 1.05])
end

@testset "SVBTransitionOutVCD" begin
    @test SVBTransitionOutVCD <: StructuredVariationalRule{Transition}
    @test outboundType(SVBTransitionOutVCD) == Message{Categorical}
    @test isApplicable(SVBTransitionOutVCD, [Nothing, Message{Categorical}, ProbabilityDistribution])

    @test ruleSVBTransitionOutVCD(nothing, Message(Univariate, Categorical, p=[0.1, 0.4, 0.5]), ProbabilityDistribution(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == Message(Univariate, Categorical, p=[0.3660714285714286, 0.27678571428571425, 0.35714285714285715])
end

@testset "SVBTransitionIn1CVD" begin
    @test SVBTransitionIn1CVD <: StructuredVariationalRule{Transition}
    @test outboundType(SVBTransitionIn1CVD) == Message{Categorical}
    @test isApplicable(SVBTransitionIn1CVD, [Message{Categorical}, Nothing, ProbabilityDistribution])

    @test ruleSVBTransitionIn1CVD(Message(Univariate, Categorical, p=[0.1, 0.4, 0.5]), nothing, ProbabilityDistribution(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == Message(Univariate, Categorical, p=[0.23000000000000007, 0.42999999999999994, 0.33999999999999997])
end

@testset "SVBTransitionADV" begin
    @test SVBTransitionADV <: StructuredVariationalRule{Transition}
    @test outboundType(SVBTransitionADV) == Message{Dirichlet}
    @test isApplicable(SVBTransitionADV, [ProbabilityDistribution, Nothing])

    @test ruleSVBTransitionADV(ProbabilityDistribution(Multivariate, Contingency, p=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3]), nothing) == Message(MatrixVariate, Dirichlet, a=[1.2 1.1 1.7; 1.4 1.3 1.3; 1.1 1.6 1.3])
end

@testset "MTransitionCCD" begin
    @test MTransitionCCD <: MarginalRule{Transition}
    @test isApplicable(MTransitionCCD, [Message{Categorical}, Message{Categorical}, ProbabilityDistribution])
    @test !isApplicable(MTransitionCCD, [Message{Categorical}, Message{Categorical}, Nothing])

    @test ruleMTransitionCCD(Message(Univariate, Categorical, p=[0.2, 0.7, 0.1]), Message(Univariate, Categorical, p=[0.1, 0.4, 0.5]), ProbabilityDistribution(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == ProbabilityDistribution(Multivariate, Contingency, p=[0.011799410029498528 0.023598820058997057 0.20648967551622419; 0.08259587020648967 0.247787610619469 0.3097345132743362; 0.002949852507374632 0.07079646017699116 0.04424778761061947])
end

@testset "MTransitionCCN" begin
    @test MTransitionCCN <: MarginalRule{Transition}
    @test isApplicable(MTransitionCCN, [Message{Categorical}, Message{Categorical}, Nothing])
    @test !isApplicable(MTransitionCCN, [Message{Categorical}, Message{Categorical}, ProbabilityDistribution])

    @test ruleMTransitionCCN(Message(Univariate, Categorical, p=[0.2, 0.7, 0.1]), Message(Univariate, Categorical, p=[0.1, 0.4, 0.5]), Message(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == ProbabilityDistribution(Multivariate, Contingency, p=[0.011799410029498528 0.023598820058997057 0.20648967551622419; 0.08259587020648967 0.247787610619469 0.3097345132743362; 0.002949852507374632 0.07079646017699116 0.04424778761061947])
end

@testset "averageEnergy" begin
    @test averageEnergy(Transition, ProbabilityDistribution(Univariate, Categorical, p=[0.2, 0.7, 0.1]), ProbabilityDistribution(Univariate, Categorical, p=[0.1, 0.4, 0.5]), ProbabilityDistribution(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == 1.1783637941354863
    @test averageEnergy(Transition, ProbabilityDistribution(Multivariate, Contingency, p=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3]), ProbabilityDistribution(MatrixVariate, PointMass, m=[0.2 0.1 0.7; 0.4 0.3 0.3; 0.1 0.6 0.3])) == 2.7886642527453405
end

end


text:
include("test_params.jl")
using SymbolicRegression, Test
import SymbolicUtils: simplify, Symbolic
import Random: MersenneTwister

binary_operators = (+, -, /, *)

index_of_mult = [i for (i, op) in enumerate(binary_operators) if op == *][1]

options = Options(; binary_operators=binary_operators)

tree = Node("x1") + Node("x1")

# Should simplify to 2*x1:
eqn = convert(Symbolic, tree, options)
eqn2 = simplify(eqn)
# Should correctly simplify to 2 x1:
# (although it might use 2(x1^1))
@test occursin("2", "$(repr(eqn2)[1])")

# Let's convert back the simplified version.
# This should remove the ^ operator:
tree = convert(Node, eqn2, options)
# Make sure one of the nodes is now 2.0:
@test (tree.l.constant ? tree.l : tree.r).val == 2
# Make sure the other node is x1:
@test (!tree.l.constant ? tree.l : tree.r).feature == 1

# Finally, let's try converting a product, and ensure
# that SymbolicUtils does not convert it to a power:
tree = Node("x1") * Node("x1")
eqn = convert(Symbolic, tree, options)
@test repr(eqn) == "x1*x1"
# Test converting back:
tree_copy = convert(Node, eqn, options)
@test repr(tree_copy) == "(x1 * x1)"

# Let's test a much more complex function,
# with custom operators, and unary operators:
x1, x2, x3 = Node("x1"), Node("x2"), Node("x3")
pow_abs(x, y) = abs(x)^y
custom_cos(x) = cos(x)^2

# Define for Node (usually these are done internally to Options)
pow_abs(l::Node, r::Node)::Node =
    (l.constant && r.constant) ? Node(pow_abs(l.val, r.val)::AbstractFloat) : Node(5, l, r)
pow_abs(l::Node, r::AbstractFloat)::Node =
    l.constant ? Node(pow_abs(l.val, r)::AbstractFloat) : Node(5, l, r)
pow_abs(l::AbstractFloat, r::Node)::Node =
    r.constant ? Node(pow_abs(l, r.val)::AbstractFloat) : Node(5, l, r)
custom_cos(x::Node)::Node = x.constant ? Node(custom_cos(x.val)::AbstractFloat) : Node(1, x)

options = Options(;
    binary_operators=(+, *, -, /, pow_abs), unary_operators=(custom_cos, exp, sin)
)
tree = (
    ((x2 + x2) * ((-0.5982493 / pow_abs(x1, x2)) / -0.54734415)) + (
        sin(
            custom_cos(
                sin(1.2926733 - 1.6606787) /
                sin(((0.14577048 * x1) + ((0.111149654 + x1) - -0.8298334)) - -1.2071426),
            ) * (custom_cos(x3 - 2.3201916) + ((x1 - (x1 * x2)) / x2)),
        ) / (0.14854191 - ((custom_cos(x2) * -1.6047639) - 0.023943262))
    )
)
# We use `index_functions` to avoid converting the custom operators into the primitives.
eqn = convert(Symbolic, tree, options; index_functions=true)

tree_copy = convert(Node, eqn, options)
tree_copy2 = convert(Node, simplify(eqn), options)
# Too difficult to check the representation, so we check by evaluation:
N = 100
X = rand(MersenneTwister(0), 3, N) .+ 0.1
output1, flag1 = eval_tree_array(tree, X, options)
output2, flag2 = eval_tree_array(tree_copy, X, options)
output3, flag3 = eval_tree_array(tree_copy2, X, options)

@test isapprox(output1, output2, atol=1e-4 * sqrt(N))
# Simplified equation may give a different answer due to rounding errors,
# so we weaken the requirement:
@test isapprox(output1, output3, atol=1e-2 * sqrt(N))