text:
using LinearAlgebra, GeometricMultigrid, BenchmarkTools, Plots

function setup_2D(n=128,T::Type=Float64)
    L = zeros(T,n+2,n+2,2); L[3:n+1,2:n+1,1] .= 1; L[2:n+1,3:n+1,2] .= 1;
    x = T[i-1 for i ∈ 1:n+2, j ∈ 1:n+2]
    Poisson(L),FieldVector(x)
end

begin
    suite = BenchmarkGroup()
    nlist = 2 .^ (3:10)
    for n ∈ nlist
        A,x = setup_2D(n,Float32)
        suite["mg-gs",n] = @benchmarkable mg!(st) setup=(st=mg_state($A,zero($x),$A*$x))
#         suite["pseudo",n] = @benchmarkable mg!(st) setup=(st=mg_state($A,zero($x),$A*$x,pseudo=true))
    end
end

results = run(suite)

begin
    nlogn(n) = n*log(n)
    N = nlist.^2
    time = [minimum(results["mg-gs",n]).time for n ∈ nlist]
    ptime = [minimum(results["pseudo",n]).time for n ∈ nlist]
    plot(N,time[end]/N[end] .* N, color=:grey, label="N", legend=:bottomright)
    plot!(N,time[end]/nlogn(N[end]) .* nlogn.(N) , color=:black, label="NlogN")
    scatter!(N,time, xaxis=("length x",:log10), yaxis=("time (ns)",:log10), label="mg!")
    # scatter!(N,ptime, xaxis=("length x",:log10), yaxis=("time (ns)",:log10), label="pseudo=true")
end

savefig("MGscaling.png")

begin
    A,x = setup_2D()
    b = zero(x)
    st = SolveState(A,zero(x),A*x)
#     st.P = PseudoInv(A)
end
@btime norm($x)  # 5.3 μs
@btime dot($b,$x) # 7.8 μs
@btime dot($b,$A,$x) # 29.5 μs
@btime mul!($b,$A,$x) # 30.6 μs
@btime GeometricMultigrid.increment!($st)  # 45.4 μs
@btime GeometricMultigrid.GS!($st;inner=1) # 183 μs
@btime GeometricMultigrid.GS!($st;inner=2) # 306 μs
# @btime GeometricMultigrid.pseudo!($st)     # 150 μs


text:
using Luxor, Test, Random
Random.seed!(42)

function polygon_intersection_test(fname)
    Drawing(2000, 2000, fname)
    origin()
    background("white")
    sethue("magenta")
    setline(.2)
    fontsize(4)
    tiles = Tiler(2000, 2000, 5, 5)
    for (pos, n) in tiles
        @layer begin
            randomhue()
            translate(pos)
            setline(1)
            p1 = randompointarray(BoundingBox(box(O, tiles.tilewidth, tiles.tileheight, vertices=true))..., 3)
            p2 = randompointarray(BoundingBox(box(O, tiles.tilewidth, tiles.tileheight, vertices=true))..., 3)

            sethue("grey80")
            poly(p1, close=true, :stroke)
            poly(p2, close=true, :stroke)

            sethue("red")
            prettypoly(p1, close=false, :stroke, vertexlabels = (n, l) -> (text(string(n))))

            sethue("blue")
            prettypoly(p2, close=false, :stroke, vertexlabels = (n, l) -> (text(string(n))))

            pint = polyintersect(p1, p2)
            sethue("purple")
            setopacity(0.5)
            circle.(pint, 6, :fill)

            pint = polyintersect(p1, p2, closed=false)
            sethue("blue")
            setopacity(0.8)
            circle.(pint, 10, :stroke)
        end
    end
    @test finish() == true
end

fname = "polygon_intersection_test.png"
polygon_intersection_test(fname)
println("...finished test: output in $(fname)")

ext: jl


text:
export BSONObject, append, append_null, append_minkey, append_maxkey, dict

mutable struct BSONObject
    _wrap_::VoidPtr
    _owner_::Any

    function BSONObject()
        _wrap_ = ccall((:bson_new, libbson), VoidPtr, ())
        bsonObject = new(_wrap_, Union{})
        finalizer(destroy, bsonObject)
        bsonObject
    end

    function BSONObject(other::BSONObject)
        _owner_ = alloc_buf(128)
        ccall((:bson_copy_to, libbson), Cvoid, (VoidPtr, BytePtr), other._wrap_,
                pointer(_owner_))
        bsonObject = new(VoidPtr(pointer(_owner_)), _owner_)
        finalizer(destroy, bsonObject)
        bsonObject
    end

    function BSONObject(dict::AbstractDict)
        bsonObject = BSONObject()
        for (k, v) in dict
            append(bsonObject, k, v)
        end
        bsonObject
    end

    function BSONObject(jsonString::AbstractString)
        jsonCStr = string(jsonString)
        bsonError = BSONError()
        _wrap_ = ccall((:bson_new_from_json, libbson), VoidPtr, (BytePtr, Csize_t, BytePtr),
                       jsonCStr, length(jsonCStr), bsonError._wrap_)
        _wrap_ != C_NULL || error(bsonError)
        bsonObject = new(_wrap_, Union{})
        finalizer(destroy, bsonObject)
        bsonObject
    end

    function BSONObject(data::BytePtr, length::Integer, _ref_::Any)
        buffer = alloc_buf(128)
        ccall((:bson_init_static, libbson), Bool, (VoidPtr, BytePtr, UInt32),
            buffer, data, length) || error("bson_init_static: failure")
        b = Base.unsafe_convert(VoidPtr, buffer)
        new(b, (_ref_, buffer))
    end

    BSONObject(_wrap_::VoidPtr, _owner_::Any) = new(_wrap_, _owner_)
end

function Base.convert(::Type{AbstractString}, bsonObject::BSONObject)
    cstr = ccall((:bson_as_json, libbson), BytePtr, (VoidPtr, BytePtr),
                 bsonObject._wrap_, C_NULL)
    result = String(unsafe_string(cstr))
    ccall((:bson_free, libbson), Cvoid, (VoidPtr,), cstr)
    result
end

Base.string(bsonObject::BSONObject) = convert(AbstractString, bsonObject)

Base.show(io::IO, bsonObject::BSONObject) =
    print(io, "BSONObject($(convert(AbstractString, bsonObject)))")

Base.length(bsonObject::BSONObject) =
    ccall((:bson_count_keys, libbson), UInt32, (VoidPtr,), bsonObject._wrap_)

function append(bsonObject::BSONObject, key::AbstractString, val::Bool)
    keyCStr = string(key)
    ccall((:bson_append_bool, libbson), Bool, (VoidPtr, BytePtr, Cint, Bool),
          bsonObject._wrap_, keyCStr, length(keyCStr), val) || bson_ovf()
end
function append(bsonObject::BSONObject, key::AbstractString, val::Real)
    keyCStr = string(key)
    ccall((:bson_append_double, libbson), Bool, (VoidPtr, BytePtr, Cint, Cdouble),
          bsonObject._wrap_, keyCStr, length(keyCStr), val) || bson_ovf()
end

function append(bsonObject::BSONObject, key::AbstractString, val::DateTime)
    keyCStr = string(key)
    ts = round(Int64, datetime2unix(val)*1000)
    ccall((:bson_append_date_time, libbson), Bool, (VoidPtr, BytePtr, Cint, Clonglong),
          bsonObject._wrap_, keyCStr, length(keyCStr), ts) || bson_ovf()
end

append(bsonObject::BSONObject, key::AbstractString, val::Date) =
    append(bsonObject, key, DateTime(val))

function append(bsonObject::BSONObject, key::AbstractString, val::BSONObject)
    keyCStr = string(key)
    ccall((:bson_append_document, libbson), Bool, (VoidPtr, BytePtr, Cint, BytePtr),
          bsonObject._wrap_, keyCStr, length(keyCStr), val._wrap_) || bson_ovf()
end

function append(bsonObject::BSONObject, key::AbstractString,
                val::Union{Int8, UInt8, Int16, UInt16, Int32, UInt32})
    keyCStr = string(key)
    ccall((:bson_append_int32, libbson), Bool, (VoidPtr, BytePtr, Cint, Cint),
          bsonObject._wrap_, keyCStr, length(keyCStr), val) || bson_ovf()
end

function append(bsonObject::BSONObject, key::AbstractString, val::Union{Int64, UInt64})
    keyCStr = string(key)
    ccall((:bson_append_int64, libbson), Bool, (VoidPtr, BytePtr, Cint, Clong),
          bsonObject._wrap_, keyCStr, length(keyCStr), val) || bson_ovf()
end

function append(bsonObject::BSONObject, key::AbstractString, val::BSONOID)
    keyCStr = string(key)
    ccall((:bson_append_oid, libbson), Bool, (VoidPtr, BytePtr, Cint, BytePtr),
          bsonObject._wrap_, keyCStr, length(keyCStr), val._wrap_) || bson_ovf()
end

append(bsonObject::BSONObject, key::AbstractString, val::Char) =
    append(bsonObject, key, string(val))

function append(bsonObject::BSONObject, key::AbstractString, val::AbstractString)
    keyCStr = string(key)
    valUTF8 = String(val)
    ccall((:bson_append_utf8, libbson), Bool, (VoidPtr, BytePtr, Cint, BytePtr, Cint),
          bsonObject._wrap_, keyCStr, length(keyCStr), valUTF8, sizeof(valUTF8)) || bson_ovf()
end

append(bsonObject::BSONObject, key::AbstractString, val::Cvoid) =
    append_null(bsonObject, key)

function append(bsonObject::BSONObject, key::AbstractString, val::Symbol)
    if val == :null
        append_null(bsonObject, key)
    elseif val == :minkey
        append_minkey(bsonObject, key)
    elseif val == :maxkey
        append_maxkey(bsonObject, key)
    else
        append(bsonObject, key, string(val))
    end
end

function append(bsonObject::BSONObject, key::AbstractString, val::AbstractDict)
    keyCStr = string(key)
    childBuffer = alloc_buf(128)
    ccall((:bson_append_document_begin, libbson), Bool, (VoidPtr, BytePtr, Cint, VoidPtr),
          bsonObject._wrap_, keyCStr, length(keyCStr), childBuffer) ||
              error("bson_append_document_begin: failure")
    childBSON = BSONObject(Base.unsafe_convert(VoidPtr, childBuffer), childBuffer)
    for (k, v) in val
        append(childBSON, k, v)
    end
    ccall((:bson_append_document_end, libbson), Bool, (VoidPtr, VoidPtr),
          bsonObject._wrap_, childBuffer) || error("bson_append_document_end: failure")
end

function append(bsonObject::BSONObject, key::AbstractString, valBinary::Vector{UInt8})
    keyCStr = string(key)
    ccall((:bson_append_binary, libbson), Bool,
          (VoidPtr, BytePtr, Cint, Cint, BytePtr, Cint),
          bsonObject._wrap_, keyCStr, length(keyCStr), 0, #BSON_SUBTYPE_BINARY
          valBinary, length(valBinary)) || bson_ovf()
end

function append(bsonObject::BSONObject, key::AbstractString, val::Vector)
    keyCStr = string(key)
    childBuffer = alloc_buf(128)
    ccall((:bson_append_array_begin, libbson), Bool, (VoidPtr, BytePtr, Cint, VoidPtr),
          bsonObject._wrap_, keyCStr, length(keyCStr), childBuffer) ||
              error("bson_append_array_begin: failure")
    childBSONArray = BSONArray(Base.unsafe_convert(VoidPtr, childBuffer), childBuffer)
    for element in val
        append(childBSONArray, element)
    end
    ccall((:bson_append_array_end, libbson), Bool, (VoidPtr, VoidPtr),
          bsonObject._wrap_, childBuffer) || error("bson_append_array_end: failure")
end

function append_null(bsonObject::BSONObject, key::AbstractString)
    keyCStr = string(key)
    ccall((:bson_append_null, libbson), Bool, (VoidPtr, BytePtr, Cint),
          bsonObject._wrap_, keyCStr, length(keyCStr)) || bson_ofv()
end

function append_minkey(bsonObject::BSONObject, key::AbstractString)
    keyCStr = string(key)
    ccall((:bson_append_minkey, libbson), Bool, (VoidPtr, BytePtr, Cint),
          bsonObject._wrap_, keyCStr, length(keyCStr)) || bson_ovf()
end

function append_maxkey(bsonObject::BSONObject, key::AbstractString)
    keyCStr = string(key)
    ccall((:bson_append_maxkey, libbson), Bool, (VoidPtr, BytePtr, Cint),
          bsonObject._wrap_, keyCStr, length(keyCStr)) || bson_ovf()
end

function dict(bsonObject::BSONObject)
    d = Dict{Any, Any}()
    @static VERSION < v"0.7-" || dump(bsonObject)
    for (k, v) in bsonObject
        if isa(v, BSONObject)
            d[k] = dict(v)
        elseif isa(v, BSONArray)
            d[k] = vector(v)
        else
            @static VERSION < v"0.7-" || println("k = $k, v = $v")
            d[k] = v
        end
    end
    d
end

# Private

destroy(bsonObject::BSONObject) =
    ccall((:bson_destroy, libbson), Cvoid, (VoidPtr,), bsonObject._wrap_)


text:
"""
Generate an output column by evaluating an `Expression` against `source`.
"""
function generate_column(source::Any, e::Expression)
    if is_constant(e.body)
        if isa(e.body, Expr)
            new_column(materialize(source), e.tuple_form(()))
        else
            new_column(materialize(source), e.body)
        end
    elseif is_column(e.body)
        # TODO: Make it possible to avoid copying here.
        copy_column(materialize(source), e.body)
    else
        gen_column(
            materialize(source),
            e.tuple_form,
            e.input_columns,
        )
    end
end


text:
@testset "GammaTrickHomotopy" begin

    PolyImpl.@polyvar x y

    H = GammaTrickHomotopy([y^2-x], [x^2+3y], 1.0+0im)
    @test typeof(H)<:GammaTrickHomotopy{Int64, PolySystem{Int64}, PolySystem{Int64}}

    @test evaluate(H, [2.0+0im, 1.0], 1.0) == [-1.0+0im]
    @test H([2.0+0im, 1.0], 1.0) == [-1.0+0im]
    @test evaluate(H, [2.0+0im, 1.0], 0.0) == [7.0+0im]

    K = substitute(H, :y=>1.0)
    @test K([1.0], 0.0) == [4.0]
    @test K([1.0], 1.0) == [0.0]

    J_H = differentiate(H)

    @test J_H([1.0+0im, 1.0], 0.0) == [2 3+0im]
    @test J_H([1.0+0im, 1.0], 1.0) == [-1 2+0im]

    ∂H∂t = dt(H)

    ## time derivate is independent of t
    @test ∂H∂t([1.0+0im,1.0], 0.23) == [-4+0im]

    @test degrees(H) == [2]
    @test startsystem(H) == PolySystem([y^2-x])
    @test targetsystem(H) == PolySystem([x^2+3.0y])

    f = [x^2+3y, 2y-3x]
    g = [y^2-x, 3x + y]
    H = GammaTrickHomotopy(g, f)
    @test homogenized(H) == false
    K = homogenize(H)
    @test ishomogenous(H) == false
    @test homogenized(K)
    @test ishomogenous(K)
    @test nvariables(K) == 3
    @test H.γ == K.γ
    @test length(removepoly(H, 2)) == 1
    @test removepoly(H, 1).γ == H.γ

    H = GammaTrickHomotopy([y^2-x], [x^2+3y])
    @test norm(H.γ) ≈ 1.0 atol=1e-8

    H = GammaTrickHomotopy(g, f, 212)
    K = GammaTrickHomotopy(g, f, 212)
    @test H.γ ≈ K.γ

    H = GammaTrickHomotopy(g, f, 1.3+4im)
    @test H.γ ≈ 1.3+4im

    PolyImpl.@polyvar a

    @test_throws ErrorException GammaTrickHomotopy([x^2+3y], [a], 1.0+0im)
    @test_throws ErrorException GammaTrickHomotopy([x^2+3y, y^2-x], [x^2+3y], 1.0+0im)
end


text:
profile_and_watch_file(args...; kwargs...) = profile_and_watch_file(stdout, args...; kwargs...)
function profile_and_watch_file(io::IO,
                                filename::AbstractString,
                                args...;
                                profiling_logger::Union{Nothing,IO} = io, # enable logger by default for watch mode
                                kwargs...)
    while true
        do_prehooks(io,
                    filename,
                    args...;
                    profiling_logger,
                    kwargs...)

        included_files, _ = try
            println(io)
            profile_file(io, filename, args...;
                         profiling_logger,
                         kwargs...)
        catch err
            println(io)
            @error "internal error occured:"
            showerror(stderr, err, catch_backtrace())
            [filename], true
        end

        # TODO: is there a way to abort tasks that are defeated in the race ?
        watch_chn = Channel{FileWatching.FileEvent}()
        for file in included_files
            @async put!(watch_chn, watch_file(file))
        end
        take!(watch_chn) # wait until any of the scheduled tasks fulfilled

        if !isfile(filename)
            @info "$(filename) has been removed"
            return
        end

        do_posthooks(io,
                     filename,
                     args...;
                     profiling_logger,
                     kwargs...)
    end
end

# hooks
# -----

const PRE_HOOKS = Function[]
push_prehook!(f) = push!(PRE_HOOKS, f)
do_prehooks(args...; kwargs...) = foreach(f->f(args...; kwargs...), PRE_HOOKS)

const POST_HOOKS = Function[]
push_posthook!(f) = push!(POST_HOOKS, f)
do_posthooks(args...; kwargs...) = foreach(f->f(args...; kwargs...), POST_HOOKS)

function __init_revise__()
    @require Revise = "295af30f-e4ad-537b-8983-00126c2a3abe" begin
        using .Revise

        function __hook_revise__(args...; kwargs...)
            revise()
        end |> push_prehook!
    end
end |> push_inithook!


text:
@testset "Monomial selection for certificate" begin
    @polyvar x y z
    @ncpolyvar a b
    @testset "Multipartite error not commutative" for parts in [([a],), ([a], [b])]
        err = ArgumentError("Multipartite Newton polytope not supported with noncommutative variables.")
        @test_throws err SumOfSquares.Certificate.monomials_half_newton_polytope([a*b, b^2], parts)
    end
    @testset "Multipartite error not disjoint: $parts" for parts in [([x], [x]),
                                                                     ([x], [x, y]),
                                                                     ([x], [y], [x, y])]
        err = ArgumentError("Parts are not disjoint in multipartite Newton polytope estimation: $parts.")
        @test_throws err SumOfSquares.Certificate.monomials_half_newton_polytope([x*y, y^2], parts)
    end
    @testset "Unipartite" begin
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x*y, y^2], tuple()) == [y]
        @test isempty(SumOfSquares.Certificate.monomials_half_newton_polytope([x, y], tuple()))
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^2, y^2], tuple()) == [x, y]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^2, y^2], ([x, y],)) == [x, y]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^2, y^2], ([y, x],)) == [x, y]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^2, x^3*y^2, x^4*y^4], tuple()) == [x^2*y^2, x]
    end
    @testset "Non-commutative" begin
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([a^4, a^3*b, a*b*a^2, a*b*a*b], tuple()) == [a^2, a*b]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([a^2, a^10*b^20*a^11, a^11*b^20*a^10, a^10*b^20*a^20*b^20*a^10], tuple()) == [a^10*b^20*a^10, a]
    end
    @testset "Multipartite" begin
        # In the part [y, z], the degree is between 0 and 2
        X = [x^4, x^2*y^2, x^2*z^2, x^2*y*z, y*z]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, tuple(), apply_post_filter = false) == [x^2, x*y, x*z, y*z, x, y, z]
        function full_test(X, Y, part1, part2)
            @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part1,), apply_post_filter = false) == Y
            @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part2,), apply_post_filter = false) == Y
            a = SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part2,), apply_post_filter = false)
            @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part1, part2), apply_post_filter = false) == Y
            @test SumOfSquares.Certificate.monomials_half_newton_polytope(X, (part2, part1), apply_post_filter = false) == Y
        end
        full_test(X, monovec([x^2, x*y, x*z,      x, y, z]), [x], [y, z])
        full_test(X, monovec([x^2, x*y, x*z, y*z, x,    z]), [y], [x, z])
        full_test(X, monovec([x^2, x*y, x*z, y*z, x, y   ]), [z], [x, y])
        # FIXME: With recursive merging, it should give [x^2, x*y, x*z, x]
        @test SumOfSquares.Certificate.monomials_half_newton_polytope([x^4, x^2*y^2, x^2*z^2, x^2*y*z, y*z], ([x], [y], [z]), apply_post_filter = false) == [x^2, x*y, x*z, y*z, x, y, z]
    end
end

@testset "Random SOS should be SOS with $(factory.optimizer_constructor)" for factory in sdp_factories
    @polyvar x y
    x = [1, x, y, x^2, y^2, x*y]
    @test_throws ArgumentError randsos(x, monotype=:Unknown)
    for i in 1:10
        for monotype in [:Classic, :Gram]
            p = randsos(x, monotype=monotype)

            m = SOSModel(factory)

            @constraint m p >= 0

            JuMP.optimize!(m)

            @test JuMP.primal_status(m) == MOI.FEASIBLE_POINT
        end
    end
end

include("ceg_test.jl")
include("csp_test.jl")


text:
module PolrfitTest

using Test, OrdinalMultinomialModels, RDatasets, MathOptInterface
const MOI = MathOptInterface

housing = dataset("MASS", "housing")
@info "Housing example for `polr` function in R package MASS"

@testset "logit link" begin
    Ipopt_solver = Ipopt.Optimizer()
    MOI.set(Ipopt_solver, MOI.RawOptimizerAttribute("print_level"), 0)
    NLopt_solver = NLopt.Optimizer()
    MOI.set(NLopt_solver, MOI.RawOptimizerAttribute("algorithm"), :LD_SLSQP)
    for solver in [Ipopt_solver, NLopt_solver]
        houseplr = polr(@formula(Sat ~ Infl + Type + Cont), housing,
            LogitLink(), solver; wts = housing[!, :Freq])
        @test nobs(houseplr) == 72
        @test dof(houseplr) == 8
        @test isapprox(deviance(houseplr), 3479.149; rtol=1e-4)
        @test isapprox(coef(houseplr), [-0.4961353, 0.6907083, 0.5663937, 1.2888191, -0.5723501, -0.3661866, -1.0910149, 0.3602841]; rtol=1e-4)
        @test isapprox(stderror(houseplr), [0.124541, 0.125212, 0.104963, 0.126705, 0.118747, 0.156766, 0.151514, 0.0953575]; rtol=1e-4)
        @test loglikelihood(houseplr.model) ≈ -1739.5746495294738
        @test fitted(houseplr.model) == nothing # TODO
        @test predict(houseplr.model) == nothing # TODO
        @test response(houseplr.model) === houseplr.model.Y
        @test modelmatrix(houseplr.model) === houseplr.model.X
        @test score(houseplr.model) === houseplr.model.∇
        @test vcov(houseplr.model) === houseplr.model.vcov
        @test all(abs.(cor(houseplr.model)) .≤ 1 + eps(Float64))
        @test all(weights(houseplr.model) .== housing[!, :Freq])
        @show coeftable(houseplr)
        @show confint(houseplr.model)
    end
end

@testset "probit link" begin
    Ipopt_solver = Ipopt.Optimizer()
    MOI.set(Ipopt_solver, MOI.RawOptimizerAttribute("print_level"), 0)
    NLopt_solver = NLopt.Optimizer()
    MOI.set(NLopt_solver, MOI.RawOptimizerAttribute("algorithm"), :LD_SLSQP)
    for solver in [Ipopt_solver, NLopt_solver]
        houseplr = polr(@formula(Sat ~ Infl + Type + Cont), housing,
            ProbitLink(), solver; wts = housing[!, :Freq])
        @test nobs(houseplr) == 72
        @test dof(houseplr) == 8
        @test isapprox(deviance(houseplr), 3479.689; rtol=1e-4)
        @test isapprox(coef(houseplr), [-0.29983, 0.42672, 0.34642, 0.78291, -0.34754, -0.21789, -0.66417, 0.22239]; rtol=1e-4)
        @test isapprox(stderror(houseplr), [0.0761614, 0.0763991, 0.0641796, 0.0762645, 0.0722116, 0.0955741, 0.0919294, 0.0581214]; rtol=1e-4)
    end
end

@testset "cloglog link" begin
    Ipopt_solver = Ipopt.Optimizer()
    MOI.set(Ipopt_solver, MOI.RawOptimizerAttribute("print_level"), 0)
    NLopt_solver = NLopt.Optimizer()
    MOI.set(NLopt_solver, MOI.RawOptimizerAttribute("algorithm"), :LD_SLSQP)
    for solver in [Ipopt_solver, NLopt_solver]
        houseplr = polr(@formula(Sat ~ Infl + Type + Cont), housing,
            CloglogLink(), solver; wts = housing[!, :Freq])
        @test nobs(houseplr) == 72
        @test dof(houseplr) == 8
        @test isapprox(deviance(houseplr), 3484.053; rtol=1e-4)
        @test isapprox(coef(houseplr), [-0.79622, 0.05537, 0.38204, 0.91536, -0.40720, -0.28053, -0.74245, 0.20922]; rtol=1e-3)
        @test isapprox(stderror(houseplr), [0.0904734, 0.0866657, 0.0701218, 0.0924964, 0.0861325, 0.112948, 0.102084, 0.0653755]; rtol=1e-4)
    end
end

end


text:
function resize_center_crop(image, len=224)
    a0 = image
    new_size = ntuple(i->div(size(a0,i)*len,minimum(size(a0))),2)
    a1 = Images.imresize(a0, new_size)
    i1 = div(size(a1,1)-len,2)
    j1 = div(size(a1,2)-len,2)
    b1 = a1[i1+1:i1+len,j1+1:j1+len]
    RGB{Float64}.(b1)
end

function vgg_image(cropped, average_image=F(0.0))
    b1 = convert(Array{FixedPointNumbers.Normed{UInt8,8},3}, channelview(cropped))
    c1 = permutedims(b1, (3,2,1))
    d1 = convert(Array{Float32}, c1)
    e1 = reshape(d1[:,:,1:3], (224,224,3,1))
    f1 = (255 * e1 .- average_image)
    g1 = permutedims(f1, [2,1,3,4])
end

function preprocess_image(img, average_image=F(0.0))
    vgg_image(resize_center_crop(img), average_image)
end


text:
eye(N) = Matrix(1.0I, N, N)

function update!(old,new,ρ=0.0)
    r=norm(new-old,Inf);
	old .*=ρ ;
	old .+= (1-ρ)*new;
    return r
end

"""
Density consistency algorithm

Compute approximate marginals of
``p( \\bf{x} )=\\frac1Z \\Phi(\\bf{x} ) \\prod_a \\psi_{a}(x_{\\partial a})``

where `\\Phi(\\bf{x} )` is an optional multivariate Gaussian distribution

Input:

* `Ψ::Vector{<:Factor}`: Array of factor nodes ``\\psi_{a}``
* `N::Int`: number of spins

Optional named arguments:

* `closure::Symbol = :DC`: closure protocol
* `maxiter::Int64 = 2000`: max number of iterations
* `η::Float64 = 1.0`: interpolation parameter (valid on :DC closure) -> ``η=0`` gives BP fixed points
* `γ0::Float64 = 0.0`: reinforcement -> at each iteration t adds a field vector ``h^{t}_{i} = \\gamma_0 t \\mu^{t}_{i}``
* `epsconv::Float64 = 1e-6`: precision convergence
* `Λ::Float64 = 1e-15`: adds a diagonal matrix ``\\Lambda \\mathbb{I}`` in the inversion of the full correlation matrix
* `λ::Float64 = 1e-15`: adds a diagonal matrix ``\\lambda \\mathbb{I}`` in the inversion of the cavity and tilted correlation matrices
* `update::Symbol = :par`: type of update -> parallel (:par) or random sequential (:seq)
* `rndamp::Bool = false`: a random damping (might be useful in parallel update)
* `epsclamp::Float64 = 1e-15:` clamp tilted moments in [-1+epsclamp, 1-epsclamp]
* `ρ::Float64 = 0.9`: damping for parameters' update (ρ = 0 means no damping)
* `verbose::Bool = true` print at (un)convergence
* `Hg::Vector{Float64} = zeros(N)`: field vector of `\\Phi(\\bf{x} )`
* `Ag::Matrix{Float64} = zeros(N,N)`: inverse covariance matrix of `\\Phi(\\bf{x} )`
* `convtype::Symbol = :params`: convergence criterion on :moments or  gaussian factor :params
* `seed :: Int64 = -1`: seed for RNG
* `callback::Function  = (x...)->nothing`: callback function to print progress report
* `h0scale::Float64 = 1e-3`: field scale for gaussian parameter initialization
* `state::DCState = DCState( N, [length(Ψ[a].idx) for a=1:length(Ψ)]; h0 = h0scale)`: initial state of DC run

Output:

* `state::DCState`: final DCstate -> see [`DCState`](@ref)
* `convergence_flag::Symbol`: :converged or :unconverged with criterion `convtype`
* `iter::Int`: number of iterations needed to converge
* `errors::Vector{Float64}`: errors on params and moments at convergence `[error_params , error_moments]`

"""
function density_consistency(Ψ::Vector{<:Factor},
			N::Int;
			closure::Symbol = :DC,
			maxiter::Int64 = 2000,
			η::Float64 = 1.0,
			γ0::Float64 = 0.0,
			epsconv::Float64 = 1e-6,
			Λ::Float64 = 1e-15,
			λ::Float64 = 1e-15,
			update::Symbol = :par,
			rndamp::Bool = false,
			epsclamp::Float64 = 1e-15,
			ρ::Float64 = 0.9,
			verbose::Bool = true,
			Hg::Vector{Float64} = zeros(N),
			Ag::Matrix{Float64} = zeros(N,N),
			convtype::Symbol = :params,
			seed :: Int64 = -1,
			callback::Function  = (x...)->nothing,
			h0scale::Float64 = 1e-3,
			state::DCState = DCState( N, [length(Ψ[a].idx) for a=1:length(Ψ)]; h0 = h0scale),
			)

    @extract state : μtl Σtl μt Σt µ Σ h S yc Sc

	seed != -1 && Random.seed!(seed) #set the seed

	A,y = zeros(N,N), zeros(N)
    Id = Matrix(1.0I, N, N); M = length(Ψ);
    ε,εₘ = 0.0 , 0.0
    iter = 1
    γ = 0.0
    while iter<=maxiter
    	γ += γ0
        A[:] .= 0.0; y .= μt * γ
        A .+= Ag; y .+= Hg
        μt[:] .= 0.0; Σt[:] .= 0.0;
        ε,εₘ = 0.0 ,  0.0
        for a in 1:M
            ∂a = Ψ[a].idx
            A[∂a,∂a] .+= S[a]; y[∂a] .+= h[a]
        end
        perm = collect(1:M)
        randperm!(perm)
      	Σ .= inv(A + Λ * I); μ .= Σ * y; # also defined in rndseq update
        for a in perm
            ψₐ=Ψ[a]; ∂a = ψₐ.idx
	    	if update == :seq
	    		A[∂a, ∂a] .-= S[a]; y[∂a] .-= h[a]
            	Σa = ((A + Λ * I) \ Id[:, ∂a])[∂a,:]; Sc[a] .= (Σa + λ * I) \ eye(length(∂a))
            	μa = ((A + Λ * I) \ y)[∂a]; yc[a] .= Sc[a] * μa
	    	elseif update == :par
                Sc[a] .= (Σ[∂a, ∂a] + λ * I)\eye(length(∂a)) - S[a]
                yc[a] .= (Σ[∂a, ∂a] + λ * I)\μ[∂a] - h[a]
	    	end
            ynew, Snew, εₘ = setclosure!(ψₐ,yc[a],Sc[a],µtl[a],Σtl[a],closure,η,λ,epsclamp,εₘ)
            if update == :seq
                ε = max(ε, update!(S[a], Snew, ρ), update!(h[a], ynew, ρ))
				A[∂a,∂a] .+= S[a]; y[∂a] .+= h[a]
            elseif update == :par
                ρr = rndamp*(1-ρ)*(2*rand()-1) + ρ    # ρr ~ U(2ρ-1,1) (with <ρ̃>=ρ) if rndamp true
                ε = max(ε, update!(S[a], Snew,ρr), update!(h[a], ynew,ρr))
            end
            μt[∂a] .= µtl[a] ; Σt[∂a,∂a] .= Σtl[a]
        end
        callback(state,iter,[ε, εₘ])
        if check_convergence(ε,εₘ,epsconv,convtype)
			verbose && println("Closure $closure (η = $η) converged in $iter iterations w.r.t $convtype - err params = ", ε, " - err moments = ", εₘ)
			return state, :converged, iter, [ε, εₘ]
		end
		iter+=1
    end
	verbose && println("Closure $closure (η = $η) NOT converged w.r.t $convtype - err params = ", ε, " - err moments = ", εₘ)
    return state, :unconverged, iter, [ε, εₘ]
end

function check_convergence(err_params::Float64,err_moments::Float64,epsconv::Float64,convtype::Symbol)
	convtype == :params && return err_params < epsconv
	convtype == :moments && return err_moments < epsconv
end

max_stars_repo_name: giovact/DensityConsistency

max_issues_repo_name: giovact/DensityConsistency

max_forks_repo_name: giovact/DensityConsistency


text:
using Abaco
using Test

x = 2.0
y = 8.0
z = 2.0
t = 10.0

r = (x + y) / exp(z)
w = (x + y) / z
sum = x + y + z + t

expected_triggers = 3
actual_triggers = 0

function onresult(ts, target_sn, target_name, value, inputs)
    global actual_triggers
    @debug "age [$ts]: scope: [$target_sn] $target_name = $value"
    @debug "inputs: $inputs"

    @test target_sn == ne
    if target_name === "r"
        @test value == r
    end
    if target_name === "sum"
        @test value == sum
    end
    if target_name === "w"
        @test value == w
    end
    actual_triggers += 1
end

# My time zone: Saturday, October 30, 2021 3:43:35 PM GMT+02:00 DST
ts = 1635601415

interval = 5
ages = 4
abaco = abaco_init(onresult, interval=interval, ages=ages)

ne = "Mulaz"

# add a formula
formula(abaco, "r = (x + y) / exp(z)")
formula(abaco, "w = (x + y) / z")
formula(abaco, "sum = x + y + z + t")

ingest!(abaco, Dict(
                "ne" => ne,
                "ts" => ts,
                "x" => x,
            ))

ingest!(abaco, Dict(
                "ne" => ne,
                "ts" => ts+2,
                "y" => y,
                "z" => z
            ))

ingest!(abaco, Dict(
                "ne" => ne,
                "ts" => ts+2,
                "t" => t,
            ))

ingest!(abaco, Dict(
                "ne" => ne,
                "ts" => ts+2,
                "w" => nothing,
            ))

@debug "[$ne]: $(abaco.node[ne])"

result = get_values(abaco, ne, "r")
@debug "get_values(r) = $result"

@test actual_triggers === expected_triggers


text:
function run_2D_benchmark(;algorithm = MetropolisHastings(),nchains = 8,nsteps = 10^5)
    for i in 1:length(testfunctions_2D)
        sample_stats_all = run2D(
            collect(keys(testfunctions_2D))[i], #There might be a nicer way but I need the name to save the plots
            testfunctions_2D,
            sample_stats2D[i],
            run_stats2D[i],
            algorithm,
            nsteps,
            nchains
        )
    end
    make_2D_results(testfunctions_2D,sample_stats2D)
    save_stats_2D(collect(keys(testfunctions_2D)),run_stats2D,run_stats_names2D)
end


text:
# [1](https://projecteuler.net/problem=1)

# \sum_i ia + \sum_j jb  \text{ ,where } i \neq nb , j \neq ma

@time s = sum([ i for i in filter(x -> (x%3==0) || (x%5==0), 1:999) ])
println(s)
# slow
# 233168

# equivalent to \sum_i ia + \sum_j jb - \sum_k kab
# where kab is the repeated numbers among ia and jb.

@time s = sum([ 3:3:999; 5:5:999; -(15:15:999) ])
println(s)
# fast
# 233168

max_stars_repo_name: cdluminate/MyNotes

max_issues_repo_name: cdluminate/MyNotes

max_forks_repo_name: cdluminate/MyNotes


text:
using BenchmarkTools

function result(bench, tool, time)
  open("results/bench.csv", "a") do io
    println(io, bench, ", ", tool, ", ", time)
  end
end

macro result(bench, tool, ex)
  quote
    t = @belapsed $ex
    result($(string(bench)), $(string(tool)), t)
  end
end


text:
# Latent ODE model
#
# Based on
# https://arxiv.org/abs/1806.07366
# https://arxiv.org/abs/2003.10775

struct LatentODE <: LatentDE end

apply_feature_extractor(encoder::Encoder{LatentODE}, x) = encoder.feature_extractor.(x)

function apply_pattern_extractor(encoder::Encoder{LatentODE}, fe_out)
    pe_z₀ = encoder.pattern_extractor

    # reverse sequence
    fe_out_rev = reverse(fe_out)

    # pass it through the recurrent layers
    pe_out = map(pe_z₀, fe_out_rev)[end]

    # reset hidden states
    Flux.reset!(pe_z₀)

    return pe_out
end

function apply_latent_in(encoder::Encoder{LatentODE}, pe_out)
    li_μ_z₀, li_logσ²_z₀ = encoder.latent_in

    z₀_μ = li_μ_z₀(pe_out)
    z₀_logσ² = li_logσ²_z₀(pe_out)

    return z₀_μ, z₀_logσ²
end

apply_latent_out(decoder::Decoder{LatentODE}, z̃₀) = decoder.latent_out(z̃₀)

function diffeq_layer(decoder::Decoder{LatentODE}, ẑ₀, t)
    dudt = decoder.diffeq.dudt
    solver = decoder.diffeq.solver
    neural_model = decoder.diffeq.neural_model
    augment_dim = decoder.diffeq.augment_dim
    kwargs = decoder.diffeq.kwargs
    # sensealg = decoder.diffeq.sensealg

    # nODE = neural_model(dudt, (t[1], t[end]), solver, sensealg = sensealg, saveat = t)
    nODE = neural_model(dudt, (t[1], t[end]), solver; saveat = t, kwargs...)
    nODE = augment_dim == 0 ? nODE : AugmentedNDELayer(nODE, augment_dim)
    ẑ = Array(nODE(ẑ₀))

    # Transform the resulting output (mainly used for Kuramoto-like systems)
    ẑ = transform_after_diffeq(ẑ, decoder.diffeq)
    ẑ = Flux.unstack(ẑ, 3)

    return ẑ
end

apply_reconstructor(decoder::Decoder{LatentODE}, ẑ) = decoder.reconstructor.(ẑ)

function sample(μ::T, logσ²::T, model::LatentDiffEqModel{LatentODE}) where T <: Array
    z₀_μ = μ
    z₀_logσ² = logσ²

    ẑ₀ = z₀_μ + randn(Float32, size(z₀_logσ²)) .* exp.(z₀_logσ²/2f0)

    return ẑ₀
end

function sample(μ::T, logσ²::T, model::LatentDiffEqModel{LatentODE}) where T <: Flux.CUDA.CuArray
    z₀_μ = μ
    z₀_logσ² = logσ²

    ẑ₀ = z₀_μ + gpu(randn(Float32, size(z₀_logσ²))) .* exp.(z₀_logσ²/2f0)

    return ẑ₀
end

function default_layers(model_type::LatentODE, input_dim, diffeq; device = cpu,
                            hidden_dim_resnet = 200, rnn_input_dim = 32,
                            rnn_output_dim = 32, latent_to_diffeq_dim = 200,
                            output_activation = σ)

    latent_dim_in = diffeq.latent_dim_in
    latent_dim_out = diffeq.latent_dim_out

    ######################
    ### Encoder layers ###
    ######################

    # Resnet
    l1 = Dense(input_dim, hidden_dim_resnet, relu)
    l2 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu)
    l3 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu)
    l4 = Dense(hidden_dim_resnet, rnn_input_dim, relu)
    feature_extractor = Chain(l1,
                                SkipConnection(l2, +),
                                SkipConnection(l3, +),
                                l4) |> device

    # RNN
    pattern_extractor = Chain(RNN(rnn_input_dim, rnn_output_dim, relu),
                                RNN(rnn_output_dim, rnn_output_dim, relu)) |> device

    # final fully connected layers before sampling
    li_μ_z₀ = Dense(rnn_output_dim, latent_dim_in) |> device
    li_logσ²_z₀ = Dense(rnn_output_dim, latent_dim_in) |> device

    latent_in = (li_μ_z₀, li_logσ²_z₀)

    encoder_layers = (feature_extractor, pattern_extractor, latent_in)

    ######################
    ### Decoder layers ###
    ######################

    # going back to the input dimensions
    # Resnet
    l1 = Dense(latent_dim_out, hidden_dim_resnet, relu)
    l2 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu)
    l3 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu)
    l4 = Dense(hidden_dim_resnet, input_dim, output_activation)
    reconstructor = Chain(l1,
                            SkipConnection(l2, +),
                            SkipConnection(l3, +),
                            l4)  |> device

    decoder_layers = (x -> x, diffeq, reconstructor)

    return encoder_layers, decoder_layers
end


text:
export hclt

using DataStructures: PriorityQueue, enqueue!, dequeue!, dequeue_pair!

####################
# Helper functions #
####################

num_bits_for_cats(num_cats) = ceil(Int,log2(num_cats))

num_categories(voc) = count(x -> !iszero(x), voc)

function as_bits(category, num_bits, bits = BitVector(undef, num_bits))
    for bit_idx = 1:num_bits
        @inbounds bits[bit_idx] = (category & 1)
        category = category >> 1
    end
    bits
end

function as_cats(bits)
    category = 0
    for bit_idx = length(bits):-1:1
        category = category << 1
        @inbounds category |= bits[bit_idx]
    end
    category
end

function categorical_leafs(num_vars, num_cats, ::Type{T} = ProbCircuit; mode = "pos", var_idx_offset = 0) where T
    num_bits = num_bits_for_cats(num_cats)

    plits = pos_literals(ProbCircuit, num_vars * num_bits)
    nlits = neg_literals(ProbCircuit, num_vars * num_bits)

    offset_var_idx(node) = PlainProbLiteralNode(node.literal + sign(node.literal) * var_idx_offset)
    plits = [offset_var_idx(n) for n in plits]
    nlits = [offset_var_idx(n) for n in nlits]

    cat_leaf(var,cat) = begin
        bits = as_bits(cat, num_bits)
        binary_leafs = Vector{T}(undef, num_bits)
        for i = 1 : num_bits
            bit_index = (var-1) * num_bits + i
            lit = bits[i] ? plits[bit_index] : nlits[bit_index]
            binary_leafs[i] = lit
        end
        if num_bits >= 2
            multiply(binary_leafs...)
        else
            binary_leafs[1]
        end
    end

    cat_leaf.(1:num_vars, (1:num_cats)')
end

function as_categorical_data(dataset::DataFrame, num_vars, num_cats)

    # Convert binary data to categorical data
    if isweighted(dataset)
        dataset, weights = split_sample_weights(dataset)
        weights = convert(Vector{Float64}, weights)
    else
        weights = nothing
    end

    num_samples = num_examples(dataset)
    num_bits = ceil(Int,log2(num_cats))

    # Get categorical dataset from the binarized dataset
    categorical_dataset = Matrix{UInt32}(undef, num_samples, num_vars)
    if iscomplete(dataset)
        for sample_idx = 1 : num_samples
            for variable_idx = 1 : num_vars
                @inbounds categorical_dataset[sample_idx, variable_idx] = as_cat(dataset[sample_idx, (variable_idx - 1) * num_bits + 1 : variable_idx * num_bits]; complete = true)
            end
        end
    else # If the dataset contains missing values, we impute the missing values with the mode of each column
        for variable_idx = 1 : num_vars
            cat_counts::Array{UInt32} = zeros(UInt32, num_cats)
            for sample_idx = 1 : num_samples
                category = as_cat(dataset[sample_idx, (variable_idx - 1) * num_bits + 1 : variable_idx * num_bits]; complete = false)
                if category != typemax(UInt32)
                    cat_counts[category] += 1
                end
                @inbounds categorical_dataset[sample_idx, variable_idx] = category
            end
            cat_mode = argmax(cat_counts)
            for sample_idx = 1 : num_samples
                if categorical_dataset[sample_idx, variable_idx] == typemax(UInt32)
                    @inbounds categorical_dataset[sample_idx, variable_idx] = cat_mode
                end
            end
        end
    end
    return categorical_dataset, weights
end

function as_cat(bits; complete)
    if !complete && !iscomplete(bits)
        return typemax(UInt32)
    end

    category::UInt32 = UInt32(0)
    for bit_idx = length(bits) : -1 : 1
        category = (category << 1) + bits[bit_idx]
    end

    (category == 0) ? 2^(length(bits)) : category
end

"Compute the Chow-Liu Tree given a binary dataset.
 Automatically convert to categorical dataset if specified by `num_vars' and `num_cats'.
 If `num_trees` is greater than 1, the algorithm returns the top-K maximum spanning trees
 with respect to the pairwise_MI weights.
 Reference: Listing all the minimum spanning trees in an undirected graph
 http://www.nda.ac.jp/~yamada/paper/enum-mst.pdf"
function chow_liu_tree(data, num_vars, num_cats; pseudocount = 0.1, num_trees::Integer = 1, dropout_prob::Float64 = 0.0)
    # Compute pairwise mutual information between variables
    MI = pairwise_MI(data, num_vars, num_cats; pseudocount)

    # Priority queue that maintain candidate MSTs
    candidates = PriorityQueue{Tuple{Vector{SimpleWeightedEdge}, Vector{SimpleWeightedEdge}, Vector{SimpleWeightedEdge}}, Float32}()

    # The fully connect graph and its weight
    g = SimpleWeightedGraph(complete_graph(num_vars))
    weights = -MI

    included_edges::Vector{SimpleWeightedEdge} = Vector{SimpleWeightedEdge}()
    excluded_edges::Vector{SimpleWeightedEdge} = Vector{SimpleWeightedEdge}()
    reuse = Matrix{Float64}(undef, num_vars, num_vars)
    topk_msts::Vector{Vector{SimpleWeightedEdge}} = Vector{Vector{SimpleWeightedEdge}}()

    # Initialize `candidate` with the global MST
    mst_edges, total_weight = MST(g, weights, included_edges, excluded_edges; reuse, dropout_prob = 0.0)
    enqueue!(candidates, (mst_edges, included_edges, excluded_edges), total_weight)

    if Threads.nthreads() == 1

        # Sequential code
        for idx = 1 : num_trees
            if isempty(candidates)
                break
            end

            (mst_edges, included_edges, excluded_edges), total_weight = dequeue_pair!(candidates)

            # Record the current ST into `topk_msts`
            push!(topk_msts, mst_edges)

            if idx == num_trees
                break
            end

            edge_added = false
            for edge_idx = 1 : length(mst_edges)
                if mst_edges[edge_idx] in included_edges
                    continue
                end

                if edge_added
                    push!(included_edges, pop!(excluded_edges))
                end
                push!(excluded_edges, mst_edges[edge_idx])
                edge_added = true

                candidate_mst, total_weight = MST(g, weights, included_edges, excluded_edges; reuse, dropout_prob)
                if candidate_mst !== nothing
                    # A shallow copy of the vectors `included_edges` and `excluded_edges` is sufficient
                    enqueue!(candidates, (candidate_mst, copy(included_edges), copy(excluded_edges)), total_weight)
                end
            end
        end

    else

        # Parallel code
        reuse = map(1:Threads.nthreads()) do idx
            Matrix{Float64}(undef, num_vars, num_vars)
        end
        g = map(1:Threads.nthreads()) do idx
            deepcopy(g)
        end
        weights = map(1:Threads.nthreads()) do idx
            deepcopy(weights)
        end

        l = ReentrantLock()

        for idx = 1 : num_trees
            if isempty(candidates)
                break
            end

            (mst_edges, included_edges, excluded_edges), total_weight = dequeue_pair!(candidates)

            # Record the current ST into `topk_msts`
            push!(topk_msts, mst_edges)

            if idx == num_trees
                break
            end

            Threads.@threads for edge_idx = 1 : length(mst_edges)
                curr_included_edges = copy(included_edges)
                curr_excluded_edges = copy(excluded_edges)
                for edge in mst_edges[1:edge_idx-1]
                    if !(edge in included_edges)
                        push!(curr_included_edges, edge)
                    end
                end
                if !(mst_edges[edge_idx] in excluded_edges)
                    push!(curr_excluded_edges, mst_edges[edge_idx])
                end

                id = Threads.threadid()
                candidate_mst, total_weight = MST(g[id], weights[id], curr_included_edges, curr_excluded_edges; reuse = reuse[id], dropout_prob)

                lock(l)
                if candidate_mst !== nothing
                    # A shallow copy of the vectors `included_edges` and `excluded_edges` is sufficient
                    enqueue!(candidates, (candidate_mst, copy(included_edges), copy(excluded_edges)), total_weight)
                end
                unlock(l)
            end
        end

    end

    # Post-process the top-K Spanning Trees
    map(topk_msts) do mst_edges
        MStree = SimpleGraph(num_vars)
        map(mst_edges) do edge
            add_edge!(MStree, src(edge), dst(edge))
        end

        # Use the graph center of `MStree' as the root node of the CLT
        clt = SimpleDiGraph(num_vars)
        for c in filter(c -> (length(c) > 1), connected_components(MStree))
            sg, vmap = induced_subgraph(MStree, c)
            sub_root = vmap[Graphs.center(sg)[1]]
            clt = union(clt, bfs_tree(MStree, sub_root))
        end

        MetaDiGraph(clt)
    end
end

"Return a vertex sequence to traverse the tree `g', where children are accessed before parent."
function bottom_up_traverse_node_seq(g::MetaDiGraph)
    num_nodes = length(vertices(g))
    node_seq = Array{UInt32, 1}(undef, num_nodes)
    idx = 1

    function dfs(node_idx)
        out_neighbors = outneighbors(g, node_idx)

        for out_neighbor in out_neighbors
            dfs(out_neighbor)
        end

        node_seq[idx] = node_idx
        idx += 1
    end

    root_node_idx = findall(x->x==0, indegree(g))[1]
    dfs(root_node_idx)

    @assert idx == num_nodes + 1

    node_seq
end

"Compute the Minimum Spanning Tree (MST) of graph g with weights `weights`, with
 constraints such that `included_edges` should be included while `excluded_edges`
 should be excluded."
function MST(g::SimpleWeightedGraph, weights::Matrix{T},
             included_edges::Vector{SimpleWeightedEdge},
             excluded_edges::Vector{SimpleWeightedEdge}; reuse::Matrix{T}, dropout_prob = 0.0) where T <: AbstractFloat
    @inbounds @views reuse[:, :] .= weights[:, :]

    # Dropout
    if dropout_prob > 1e-8
        dropped_mask = rand(Bernoulli(dropout_prob), size(reuse, 1), size(reuse, 2))
        @inbounds @views reuse[dropped_mask] .= 10000.0
    end

    # Add constraints
    map(included_edges) do edge
        reuse[src(edge), dst(edge)] = -10000.0
        reuse[dst(edge), src(edge)] = -10000.0
        nothing # Return nothing to save some effort
    end
    map(excluded_edges) do edge
        reuse[src(edge), dst(edge)] = 10000.0
        reuse[dst(edge), src(edge)] = 10000.0
        nothing # Return nothing to save some effort
    end

    mst_edges = kruskal_mst(g, reuse)

    # Senity check
    valid_tree::Bool = true

    edges = Set(mst_edges)
    map(included_edges) do edge
        if !(edge in edges)
            valid_tree = false
        end
        nothing
    end
    map(excluded_edges) do edge
        if (edge in edges)
            valid_tree = false
        end
        nothing
    end

    if valid_tree
        # Compute the tree weight
        total_weight::T = 0.0
        map(mst_edges) do edge
            total_weight += weights[src(edge), dst(edge)]
            nothing
        end
        mst_edges, total_weight
    else
        nothing, nothing
    end
end

"Compute pairwise Mutual Information given binary/categorical data."
function pairwise_MI(dataset::DataFrame, num_vars, num_cats; pseudocount = 1.0)
    categorical_dataset, weights = as_categorical_data(dataset::DataFrame, num_vars, num_cats)
    pairwise_MI(categorical_dataset, num_vars, num_cats, weights; pseudocount = pseudocount)
end
function pairwise_MI(dataset::Matrix, num_vars, num_cats, weights = nothing; pseudocount = 1.0)
    num_samples = size(dataset, 1)
    num_vars = size(dataset, 2)

    if weights === nothing
        weights = ones(Int32, num_samples)
    else
        pseudocount = pseudocount * sum(weights) / num_samples
    end

    # Sum of the weights
    sum_weights::Float64 = Float64(sum(weights) + num_cats^2 * pseudocount)

    # `joint_cont[i, j, k, w]' is the total weight of samples whose i- and j-th variable are k and w, respectively
    joint_cont = pairwise_marginals(dataset, weights, num_cats; pseudocount)

    # `marginal_cont[i, j]' is the total weight of sample whose i-th variable is j
    marginal_cont = zeros(Float64, num_vars, num_cats)
    for i = 1:num_vars
        for j = 1:num_cats
            @inbounds marginal_cont[i,j] = joint_cont[i,i,j,j]
        end
    end

    # Compute mutual information
    MI = zeros(Float64, num_vars, num_vars)
    for var1_idx = 1 : num_vars
        for var2_idx = var1_idx : num_vars
            @inbounds MI[var1_idx, var2_idx] = sum(joint_cont[var1_idx, var2_idx, :, :] .* (@. log(sum_weights .* joint_cont[var1_idx, var2_idx, :, :] / (marginal_cont[var1_idx, :] .* marginal_cont[var2_idx, :]')))) / sum_weights
        end
    end

    for var1_idx = 2 : num_vars
        for var2_idx = 1 : var1_idx - 1
            @inbounds MI[var1_idx, var2_idx] = MI[var2_idx, var1_idx]
        end
    end

    MI
end

function hclt(num_vars, num_cats = 2, ::Type{T} = ProbCircuit; data::DataFrame,
              num_hidden_cats::Integer = 16, num_trees::Integer = 1,
              num_tree_candidates::Integer = 1, tree_sample_type::String = "fixed_interval",
              dropout_prob::Float64 = 0.0) where T
    # Chow-Liu Tree (CLT) given data
    clts = chow_liu_tree(data, num_vars, num_cats; num_trees = num_tree_candidates, dropout_prob)

    # Sample `num_trees` trees from the `num_tree_candidates` candidates
    if tree_sample_type == "random"
        clts = clts[randperm(num_tree_candidates)[1:num_trees]]
    elseif tree_sample_type == "fixed_interval"
        clts = clts[Int.(round.(LinRange(1, num_tree_candidates, num_trees)))]
    end

    observed_leafs = categorical_leafs(num_vars, num_cats, T)
    var_idx_offset = num_vars * num_bits_for_cats(num_cats)
    circuits = map(clts) do clt
        pc = hclt(clt, num_vars, num_cats, observed_leafs; data = data, num_hidden_cats = num_hidden_cats,
                  var_idx_offset = var_idx_offset)
        var_idx_offset += num_vars * num_bits_for_cats(num_hidden_cats)

        pc
    end

    children::Array{T} = Array{T}(undef, 0)
    for circuit in circuits
        append!(children, circuit.children)
    end
    summate(children...)
end
function hclt(clt::CLT, num_vars, num_cats, observed_leafs, ::Type{T} = ProbCircuit;
              data::DataFrame, var_idx_offset::Integer = 0, num_hidden_cats::Integer = 4) where T
    # Circuits representing the leaves
    hidden_leafs = categorical_leafs(num_vars, num_hidden_cats, T; var_idx_offset)

    # meaning: `joined_leafs[i,j]` is a distribution of the hidden variable `i` having value `j`
    # conditioned on the observed variable `i`
    gen_joined_leaf(var_idx, hidden_cat_idx) = begin
        # This line encodes the hidden leafs explicitly
        # summate([multiply(hidden_leafs[var_idx, hidden_cat_idx], observed_leafs[var_idx, i]) for i = 1 : num_cats])

        # This line does not encode the hidden leafs
        summate(observed_leafs[var_idx, :])
    end
    joined_leafs = gen_joined_leaf.(1:num_vars, (1:num_hidden_cats)')

    # Construct the CLT circuit bottom-up
    node_seq = bottom_up_traverse_node_seq(clt)
    for curr_node in node_seq
        out_neighbors = outneighbors(clt, curr_node)

        # meaning: `circuits' of leaf CLT nodes refer to a collection of marginal distribution Pr(X);
        #          `circuits' of an inner CLT node (corr. var Y) is a collection of joint distributions
        #              over itself and its child vars (corr. var X_1, ..., X_k): Pr(Y)Pr(X_1|Y)...Pr(X_k|Y)

        if length(out_neighbors) == 0
            # Leaf node
            # We do not add hidden variables for leaf nodes
            circuits = [summate(observed_leafs[curr_node, :]) for idx = 1 : num_hidden_cats]
            set_prop!(clt, curr_node, :circuits, circuits)
        else
            # Inner node

            # Each element in `child_circuits' represents the joint distribution of the child nodes,
            # i.e., Pr(X_1)...Pr(X_k)
            child_circuits = [get_prop(clt, child_node, :circuits) for child_node in out_neighbors]
            if length(out_neighbors) > 1
                child_circuits = [summate(multiply([child_circuit[cat_idx] for child_circuit in child_circuits])) for cat_idx = 1 : num_hidden_cats]
            else
                child_circuits = child_circuits[1]
            end
            # Pr(X_1)...Pr(X_k) -> Pr(Y)Pr(X_1|Y)...Pr(X_k|Y)
            circuits = [summate(multiply.(child_circuits, joined_leafs[curr_node, :])) for cat_idx = 1 : num_hidden_cats]
            set_prop!(clt, curr_node, :circuits, circuits)
        end
    end

    get_prop(clt, node_seq[end], :circuits)[1] # A ProbCircuit node
end


text: const Graph{T,U} = Dict{T, Vector{Tuple{T, U}}} where {T, U}


text:
"""
    whittaker1(y, w, lambda) -> z

Perform a first-order differences Whittaker-Henderson smoothing and interpolation.
# Citation
'Smoothing and interpolation with finite differences' [Eilers P. H. C, 1994]
(URL: http://dl.acm.org/citation.cfm?id=180916)
"""
function whittaker1{T <: Number}(y::Vector{Float64}, w::Vector{T}, lambda::Float64)
  z = similar(y)
  return smooth1!(y, w, lambda, z)
end

"""
    whittaker1!(y, w, lambda) -> y

Perform an in-place first-order differences Whittaker-Henderson smoothing and interpolation on `y`.
# Citation
'Smoothing and interpolation with finite differences' [Eilers P. H. C, 1994]
(URL: http://dl.acm.org/citation.cfm?id=180916)
"""
function whittaker1!{T <: Number}(y::Vector{Float64}, w::Vector{T}, lambda::Float64)
  return smooth1!(y, w, lambda, y)
end

"""
    smooth1!(y, w, lambda, z) -> z

Smooth and interpolate with first-order differences.
"""
function smooth1!{T <: Number}(y::Vector{Float64}, w::Vector{T}, lambda::Float64, z::Vector{Float64})
  # init
  m = length(y)
  c = Array{Float32}(m)
  d = Array{Float32}(m)
  @inbounds @fastmath begin
    w0 = w[1]
    d1 = d[1] = w0 + lambda
    c1 = c[1] = -lambda/d1
    z1 = z[1] = w0*y[1]
  end

  # compute intermediate values
  @inbounds @fastmath for i = 2:(m-1)
    w0 = w[i]
    d1 = d[i] = w0 + 2*lambda - c1^2*d1
    z1 = z[i] = w0*y[i] - c1*z1
    c1 = c[i] = -lambda/d1
  end

  # compute last values
  @inbounds @fastmath begin
    w0 = w[m]
    d1 = d[m] = w0 + lambda - c1^2*d1
    z1 = z[m] = (w0*y[m] - c1*z1)/d1
  end

  # compute result
  @inbounds @fastmath for i = (m-1):-1:1
    z1 = z[i] = z[i]/d[i] - c[i]*z1
  end

  return z
end


text:
module TestUtil
using Test
using SARA
using SARA: interpolate_xrd_stripe

@testset "util" begin
    n = 128
    x = range(-1, 1, length = n)
    m = 5
    Y = randn(n, m)
    Yx = interpolate_xrd_stripe(x, Y, x)
    @test Yx ≈ Y
end

end


text:
module Behavior

# Module Description:
#   A behavior is a set of conditional probabilities which describe the relative
#   outcome frequencies for each of unique experimental bell test setting. This
#   module provides utilities for manipulating behaviors for bipartite bell scenarios.
#
#   A behavior can be represented in different subspaces depending on system
#   constraints.

using LinearAlgebra
using QBase: QMath

# """
# Description:
#     Determines the index at which information about the specified outcome is
#     stored.
#
# Inputs:
#     outcome: Tuple, (a,b,x,y) values
#     (α/β)_expt: Tuple, experiment parameters
#     rep: String, behavior representation (generalized,normalized,fixed-direction,no-signaling)
#
# Output:
#     index: Integer, id that indexes into the specified behavior vector
# """
function index(outcome, α_expt, β_expt, rep="generalized")
    (a,b,x,y) = outcome
    (x_in, a_out) = α_expt
    (y_in, b_out) = β_expt

    index = 1
    if rep == "generalized"
        index += (a-1)*b_out*x_in*y_in + (b-1)*x_in*y_in + (x-1)*y_in + y
    elseif rep == "normalized"
        if a == a_out & b == b_out
            throw(ArgumentError("outcome coordinate 'a' or 'b' is not in normalized representation"))
        else
            index += (a-1)*(b_out)*x_in*y_in + (b-1)*x_in*y_in + (x-1)*y_in + y
        end
    elseif rep == "no-signaling"
        if a == a_out | b == b_out
            throw(ArgumentError("outcome coordinate 'a' or 'b' is not in no-signaling representation"))
        elseif (b == 0 & y == 0) & (in(a,1:(a_out-1)) & in(x,1:(x_in)))
            index += (a-1)*x_in + x
        elseif (a == 0 & x == 0) & (in(b,1:(b_out-1)) & in(y,1:(y_in)))
            index += x_in*(a_out - 1) + (b-1)*y_in + y
        elseif (in(a,1:(a_out-1)) & in(x,1:(x_in))) & (in(b,1:(b_out-1)) & in(y,1:(y_in)))
            index += x_in*(a_out - 1) + y_in*(b_out - 1) + (a-1)*(b_out-1)*x_in*y_in + (b-1)*x_in*y_in + (x-1)*y_in + y
        else
            throw(ArgumentError("outcome coordinate ($a,$b,$x,$y) is not valid in no-signaling representation"))
        end
    elseif rep == "fixed-direction"
        if b == b_out
            throw(ArgumentError("outcome coordinate is not in fixed-direction representation"))
        elseif (b == 0 & y == 0) & (in(a,1:(a_out-1)) & in(x,1:(x_in)))
            index += (a-1)*x_in + x
        elseif (in(a,1:a_out) & in(x,1:(x_in))) & (in(b,1:(b_out-1)) & in(y,1:(y_in)))
            index += x_in*(a_out - 1) + (a-1)*(b_out-1)*x_in*y_in + (b-1)*x_in*y_in + (x-1)*y_in + y
        else
            throw(ArgumentError("outcome ($a,$b,$x,$y) is not valid"))
        end
    end

    index
end

# """
# dimension(α_expt,β_expt,rep)
#     Returns the dimension of a behavior vector in the specified representation
#     ("generalized", "normalized", "fixed-direction", "no-signaling").
#
# Inputs:
#     (α/β)_expt: Tuple, bell-test parameters
#     rep: String, behavior representation
#
# Outputs:
#     dim: Integer, behavior dimension
# """
function dimension(α_expt,β_expt,rep)
    (x_num, a_num) = α_expt
    (y_num, b_num) = β_expt
    (num_in, num_out) = α_expt .* β_expt

    dim = 1
    if rep == "generalized"
        dim += num_in * num_out
    elseif rep == "normalized"
        dim += num_in * (num_out - 1)
    elseif rep == "fixed-direction"
        dim += x_num * (a_num - 1) + num_in * a_num * (b_num - 1)
    elseif rep == "no-signaling"
        dim += x_num * (a_num - 1) + y_num * (b_num - 1) + num_in * (a_num - 1) * (b_num - 1)
    end

    dim
end

# """
# Description:
#     Applies tests checking the validity of a behavior in the specified representation.
#
# Input:
#     behavior: Vector
#     (α/β)_expt: experiment parameters
#     rep: String, behavior representation
#
# Output:
#     is_valid: Boolean
# """
function is_valid( behavior, α_expt, β_expt, rep="generalized")
    (x_in, a_out) = α_expt
    (y_in, b_out) = β_expt
    (num_inputs, num_outputs) = α_expt .* β_expt

    is_valid = false

    false_count = 0

    for i in 1:size(behavior)[1]
        if !((0 <= behavior[i]) & (behavior[i] <= 1))
            false_count += 1
        end
    end
    if rep == "generalized"
        # verify normalization
        for x in 1:x_in
            for y in 1:y_in
                norm = 0
                for a in 1:a_out
                    for b in 1:b_out
                        id = index((a,b,x,y),α_expt,β_expt)
                        norm += behavior[id]
                    end
                end

                if !(norm ≈ 1)
                    false_count += 1
                end
            end
        end
    elseif rep == "normalized"
        max_length = Behavior.dimension(α_expt,β_expt,rep) - 1

        valid_checks = map(
            (i) -> begin
                indices = 1 .+ collect(i:num_inputs:max_length)
                norm_sum = sum(behavior[indices])
                (0 <= norm_sum <= 1)
            end,
            1:num_inputs
        )
        if !(all(valid_checks))
            false_count += 1
        end
    elseif rep == "no-signaling"
        #verify normalizability
        for x in 1:x_in
            a_norm = 0
            for a in 1:(a_out-1)
                id = index((a,0,x,0),α_expt,β_expt,"no-signaling")
                a_norm += behavior[id]
            end
            if (1 < a_norm) & !(1 ≈ a_norm)
                false_count += 1
            end
        end

        for y in 1:y_in
            b_norm = 0
            for b in 1:(b_out-1)
                id = index((0,b,0,y),α_expt,β_expt,"no-signaling")
                b_norm += behavior[id]
            end
            if (1 < b_norm) & !(1 ≈ b_norm)
                false_count += 1
            end
        end

        # verify joint construction
        for x in 1:x_in
            for y in 1:y_in
                for a in 1:(a_out-1)
                    for b in 1:(b_out-1)
                        id = index((a,b,x,y),α_expt,β_expt,rep)
                        α_id = index((a,0,x,0),α_expt,β_expt,rep)
                        β_id = index((0,b,0,y),α_expt,β_expt,rep)

                        if !(
                            round(behavior[α_id]*behavior[β_id],sigdigits=3) == round(behavior[id],sigdigits = 3)
                        )
                            false_count += 1
                        end
                    end
                end
            end
        end
    elseif rep == "fixed-direction"
        #verify normalizability
        a_comps = zeros(x_in)
        for x in 1:x_in
            a_norm = 0
            for a in 1:(a_out-1)
                id = index((a,0,x,0),α_expt,β_expt,"fixed-direction")
                a_norm += behavior[id]

            end
            a_comps[x] = 1-a_norm

            if (1 < a_norm) & !(1 ≈ a_norm)
                false_count += 1
            end
        end

        if !all(map((a_comp) -> a_comp >= 0 ,a_comps))
            false_count += 1
        end

        # storing b marginals "separated" from a
        b_margs = fill(-1.0, (x_in,y_in,(b_out-1)))
        for x in 1:x_in
            for y in 1:y_in

                ab_norm = 0
                for a in 1:a_out
                    ax_id = 0
                    ax_val = 0

                    # only test the b_marg if we aren't going to divide by zero
                    test_b_marg = false
                    if a < a_out
                        ax_id = index((a,0,x,0),α_expt, β_expt, "fixed-direction")
                        ax_val = behavior[ax_id]
                        if behavior[ax_id] != 0
                            test_b_marg = true
                        end
                    else
                        if a_comps[x] != 0
                            test_b_marg = true
                        end
                    end

                    for b in 1:(b_out-1)
                        # verifying normalization of joint outputs
                        id = index((a,b,x,y), α_expt, β_expt, "fixed-direction")
                        ab_norm += behavior[id]

                        # verifying that the b_marginal matches the rest
                        if (a < a_out) & test_b_marg
                            if b_margs[x,y,b] == -1
                                b_margs[x,y,b] = (behavior[id]/ax_val)
                            else
                                if !(behavior[id]/ax_val ≈ b_margs[x,y,b])
                                    false_count += 1
                                end
                            end
                        elseif (a == a_out) & test_b_marg
                            if b_margs[x,y,b] == -1
                                b_margs[x,y,b] = behavior[id]/a_comps[x]
                            else
                                if !(behavior[id]/a_comps[x] ≈ b_margs[x,y,b])
                                    false_count += 1
                                end
                            end
                        end

                    end
                end

                if (1 < ab_norm) & !(1 ≈ ab_norm)
                    false_count += 1
                end
            end
        end

    end

    if false_count == 0
        is_valid = true
    end

    is_valid
end

# """
# gen_to_norm_proj(α_expt, β_expt)
#     Constructs the projection matrix from generalized to normalized
#     representation.
#
# Inputs:
#     (α/β)_expt: Tuple, bell-test parameters
#
# Outputs:
#     proj: Matrix that projects gneralized coordinate behaviors to normalized
#        behavior vectors.
# """
function gen_to_norm_proj(α_expt, β_expt)
    (x_num, a_num) = α_expt
    (y_num, b_num) = β_expt
    (num_in, num_out) = α_expt .* β_expt

    gen_dim = num_in * num_out

    proj = zeros(Int64, (1 + num_in*(num_out - 1), gen_dim + 1))

    proj[1,1] = 1

    # fill in 1:1 mappings
    for x in 1:x_num
        for y in 1:y_num
            for a in 1:a_num
                for b in 1:b_num
                    if (a < a_num) | (b < b_num)
                        row = index((a,b,x,y),α_expt,β_expt,"normalized")
                        col = index((a,b,x,y),α_expt,β_expt)
                        proj[row,col] = 1
                    end
                end
            end
        end
    end

    proj
end

# """
# gen_to_fd_proj(α_expt, β_expt)
#     Constructs the projection matrix from generalized to fixed-direction
#     representation.
#
# Inputs:
#     (α/β)_expt: Tuple, bell-test parameters
#
# Outputs:
#     proj: Matrix that projects gneralized coordinate behaviors to fixed-direction
#        behavior vectors.
# """
function gen_to_fd_proj(α_expt, β_expt)
    (x_num, a_num) = α_expt
    (y_num, b_num) = β_expt
    (num_in, num_out) = α_expt .* β_expt

    gen_dim = num_in * num_out
    α_dim = x_num*(a_num - 1)
    β_dim = y_num*(b_num - 1)

    proj = zeros(Int64, (1 + α_dim + x_num*a_num*β_dim, gen_dim + 1))

    proj[1,1] = 1

    # filling in p(a|x) using no-signaling constraint
    for x in 1:x_num
        for a in 1:(a_num - 1)

            row = index((a,0,x,0),α_expt,β_expt,"fixed-direction")

            for b in 1:b_num
                col = index((a,b,x,1),α_expt,β_expt)
                proj[row, col] = 1
            end
        end
    end

    # fill in 1:1 mappings
    for x in 1:x_num
        for y in 1:y_num
            for a in 1:a_num
                for b in 1:(b_num - 1)
                    row = index((a,b,x,y),α_expt,β_expt,"fixed-direction")
                    col = index((a,b,x,y),α_expt,β_expt)
                    proj[row,col] = 1
                end
            end
        end
    end

    proj
end

# """
# gen_to_ns_proj:
#     Returns a projection matrix which maps generalized behaviors to no-signaling
#     behaviors.
#
# Inputs
#     (α/β)_expt: experiment parameters
#
# Output:
#     proj: Matrix which projects generalized to no-signaling representations
# """
function gen_to_ns_proj(α_expt, β_expt)
    (α_num_inputs, α_num_outputs) = α_expt
    (β_num_inputs, β_num_outputs) = β_expt
    (num_inputs, num_outputs) = α_expt .* β_expt

    gen_dim = num_inputs * num_outputs
    α_dim = α_num_inputs*(α_num_outputs - 1)
    β_dim = β_num_inputs*(β_num_outputs - 1)
    joint_dim = α_dim*β_dim

    proj = zeros(Int64, (1 + α_dim + β_dim + joint_dim, 1 + gen_dim))

    proj[1,1] = 1

    # fill in p(a|x)
    for x in 1:α_num_inputs
        for a in 1:(α_num_outputs - 1)

            row = index((a,0,x,0),α_expt,β_expt,"no-signaling")

            for b in 1:β_num_outputs
                col = index((a,b,x,1),α_expt,β_expt)
                proj[row, col] = 1
            end
        end
    end

    # fill in p(b|y)
    for y in 1:β_num_inputs
        for b in 1:(β_num_outputs - 1)

            row = index((0,b,0,y),α_expt,β_expt,"no-signaling")

            for a in 1:(α_num_outputs)
                col = index((a,b,1,y),α_expt,β_expt)
                proj[row, col] = 1
            end
        end
    end

    # fill in 1:1 mappings
    for x in 1:α_num_inputs
        for y in 1:β_num_inputs
            for a in 1:(α_num_outputs - 1)
                for b in 1:(β_num_outputs - 1)
                    row = index((a,b,x,y),α_expt,β_expt,"no-signaling")
                    col = index((a,b,x,y),α_expt,β_expt)
                    proj[row,col] = 1
                end
            end
        end
    end

    proj
end

function norm_to_gen_proj(α_expt, β_expt)
    (x_num, a_num) = α_expt
    (y_num, b_num) = β_expt
    (num_in, num_out) = α_expt .* β_expt

    gen_dim = num_in * num_out

    proj = zeros(Int64, (gen_dim + 1, 1 + num_in*(num_out - 1)))
    proj[1,1] = 1

    for x in 1:x_num
        for y in 1:y_num
            for a in 1:a_num
                for b in 1:b_num
                    # 1:1 mappings
                    if (a < a_num) | (b < b_num)
                        row = index((a,b,x,y),α_expt,β_expt)
                        col = index((a,b,x,y),α_expt,β_expt,"normalized")
                        proj[row,col] = 1
                    else
                        row = index((a_num,b_num,x,y),α_expt,β_expt)
                        proj[row,1] = 1
                        for α in 1:a_num
                            for β in 1:b_num
                                if (α < a_num) | (β < b_num)
                                    col = index((α,β,x,y),α_expt,β_expt,"normalized")
                                    proj[row,col] = -1
                                end
                            end
                        end
                    end
                end
            end
        end
    end

    proj
end

# """
# fd_to_gen_proj(α_expt,β_expt)
#     Constructs the map from fixed-direction to generalized coordinates.
#
# Inputs:
#     (α/β)_expt: Tuple, bell-test parameters
#
# Outputs:
#     proj: Matrix, transforms fixed-directino into generalized coordinates.
# """
function fd_to_gen_proj(α_expt, β_expt)
    (x_num, a_num) = α_expt
    (y_num, b_num) = β_expt
    (num_in, num_out) = α_expt .* β_expt

    gen_dim = num_in * num_out
    α_dim = x_num*(a_num - 1)
    β_dim = y_num*(b_num - 1)

    proj = zeros(Int64, (gen_dim + 1,1 + α_dim + x_num*a_num*β_dim))
    proj[1,1] = 1

    for x in 1:x_num
        for y in 1:y_num

            # 1:1 maps
            for a in 1:a_num
                for b in 1:(b_num - 1)
                    row = index((a,b,x,y),α_expt,β_expt)
                    col = index((a,b,x,y),α_expt,β_expt,"fixed-direction")
                    proj[row,col] = 1
                end
            end

            # p(ak|xy)
            for a in 1:(a_num -1)

                row = index((a,b_num,x,y),α_expt,β_expt)

                α_col = index((a,0,x,0),α_expt,β_expt,"fixed-direction")
                proj[row,α_col] = 1

                for b in 1:(b_num - 1)
                    col = index((a,b,x,y),α_expt,β_expt,"fixed-direction")
                    proj[row, col] = -1
                end
            end

            # populating P(KK|xy)
            row = index((a_num,b_num,x,y),α_expt,β_expt)
            proj[row,1] = 1

            for α in 1:(a_num-1)
                α_col = index((α,0,x,0),α_expt,β_expt,"fixed-direction")
                proj[row,α_col] = -1
            end

            for β in 1:(b_num-1)
                β_col = index((a_num,β,x,y),α_expt,β_expt,"fixed-direction")
                proj[row,β_col] = -1
            end
        end
    end

    proj
end

# """
# ns_to_gen_proj:
#     Returns a projection matrix which maps vectors from the no-signaling representation
#     to the generalized representation.
#
# Inputs
#     (α/β)_expt: experiment parameters
#
# Output:
#     proj: matrix which maps no-signaling to generalized representations
# """
function ns_to_gen_proj(α_expt, β_expt)
    (α_num_inputs, α_num_outputs) = α_expt
    (β_num_inputs, β_num_outputs) = β_expt
    (num_inputs, num_outputs) = α_expt .* β_expt

    gen_dim = num_inputs * num_outputs
    α_dim = α_num_inputs*(α_num_outputs - 1)
    β_dim = β_num_inputs*(β_num_outputs - 1)

    proj = zeros(Int64, (gen_dim + 1,1 + α_dim + β_dim + α_dim*β_dim))
    proj[1,1] = 1

    for x in 1:α_num_inputs
        for y in 1:β_num_inputs
            for a in 1:(α_num_outputs - 1)
                for b in 1:(β_num_outputs - 1)
                    row = index((a,b,x,y),α_expt,β_expt)
                    col = index((a,b,x,y),α_expt,β_expt,"no-signaling")
                    proj[row,col] = 1
                end
            end
        end
    end

    # p(an|xy)
    for x in 1:α_num_inputs
        for y in 1:β_num_inputs
            # p(ak|xy)
            for a in 1:(α_num_outputs -1)

                row = index((a,β_num_outputs,x,y),α_expt,β_expt)

                α_col = index((a,0,x,0),α_expt,β_expt,"no-signaling")
                proj[row,α_col] = 1

                for b in 1:(β_num_outputs - 1)
                    col = index((a,b,x,y),α_expt,β_expt,"no-signaling")
                    proj[row, col] = -1
                end
            end

            # p(kb|xy)
            for b in 1:(β_num_outputs -1)

                row = index((α_num_outputs,b,x,y),α_expt,β_expt)

                β_col = index((0,b,0,y),α_expt,β_expt,"no-signaling")
                proj[row,β_col] = 1

                for a in 1:(α_num_outputs - 1)
                    col = index((a,b,x,y),α_expt,β_expt,"no-signaling")
                    proj[row, col] = -1
                end
            end

            # populating P(KK|xy)
            a = α_num_outputs
            b = β_num_outputs

            row = index((α_num_outputs,b,x,y),α_expt,β_expt)
            proj[row,1] = 1

            for α in 1:(α_num_outputs-1)
                α_col = index((α,0,x,0),α_expt,β_expt,"no-signaling")
                proj[row,α_col] = -1
            end

            for β in 1:(β_num_outputs-1)
                β_col = index((0,β,0,y),α_expt,β_expt,"no-signaling")
                proj[row,β_col] = -1
            end

            for α in 1:(α_num_outputs-1)
                for β in 1:(β_num_outputs-1)
                    col = index((α,β,x,y),α_expt,β_expt,"no-signaling")
                    proj[row,col] = 1
                end
            end
        end
    end

    proj
end

# """
# conditionals(α_expt, β_expt, behavior; rep="generalized")
#
#     Transforms a behavior vector into a conditionals matrix.
#
# Inputs:
#     α/β_expt: Tuple, (num_in, num_out), integer values
#     behavior: Col Vector, a bell-test behavior.
#     rep: String, current behavior representation
#
# Output:
#     conditionals: Matrix, the equivalent conditional probabilities in matrix form.
# """
function conditionals(α_expt, β_expt, behavior; rep="generalized")

    expected_dim = Behavior.dimension(α_expt,β_expt,rep)
    gen_dim = Behavior.dimension(α_expt,β_expt,"generalized")

    (num_in, num_out) = α_expt .* β_expt

    if length(behavior) + 1 == expected_dim
        behavior = reshape(cat([1],behavior,dims=1),(expected_dim,1))
    end

    if !(Behavior.is_valid(behavior, α_expt, β_expt, rep))
        throw(ArgumentError("cannot make conditional probabilities from an invalid behvior"))
    end

    gen_behavior = zeros(gen_dim,1)
    if rep == "generalized"
        gen_behavior = behavior
    elseif rep == "normalized"
        gen_behavior = Behavior.norm_to_gen_proj(α_expt,β_expt)*behavior
    elseif rep == "fixed-direction"
        gen_behavior = Behavior.fd_to_gen_proj(α_expt,β_expt)*behavior
    elseif rep == "no-signaling"
        gen_behavior = Behavior.ns_to_gen_proj(α_expt,β_expt)*behavior
    else
        throw(ArgumentError("invalid behavior representation specified"))
    end

    QMath.Conditionals(Matrix(reshape(gen_behavior[2:end], (num_out,num_in))'))
end

# """
# has_constant(α, β, behavior; rep="normalized")
#
#     Checks if the specified behavior has constant required for performing
#     affine transformations.
#
# Input:
#     α/β: Tuple, specifies inputs and outputs of bipartite bell scenario.
#     behavior: Col vector, may contain the affine constant or not.
#     rep: String, optional, the behavior's representation.
#
# Output:
#     has_constant: Bool, true if the behavior has the correct dimension for
#                   including the constant
# """
function has_constant(α, β, behavior; rep="normalized")
    dim = Behavior.dimension(α, β, rep)
    behavior_dim = length(behavior)

    has_constant = NaN
    if dim == behavior_dim
        has_constant = true
    elseif (dim - 1) == behavior_dim
        has_constant = false
    else
        throw(ArgumentError("invalid behavior dimension"))
    end

    has_constant
end

# """
# add_constants(α, β, behaviors; rep="normalized")
#
#     Appends the constant to each behavior if it does not already exist.
#
# Input:
#     α/β: Tuple, number of inputs and outputs for each party.
#     behaviors: Array, contains behavior column vector
#     rep: String, default "normalized", specifies the representation of the behavior
#
# Output:
#     corrected_behaviors: Array, containsCol Vector, the behavior with the constant
# """
function add_constants(α, β, behaviors; rep="normalized")
    map(
        (behavior) -> begin
            has_constant = Behavior.has_constant(α, β, behavior, rep=rep)

            corrected_behavior = behavior
            if !has_constant
                corrected_behavior = cat([1], behavior, dims=1)
            end

            corrected_behavior
        end,
        behaviors
    )
end

# """
# remove_constants(α, β, behavior; rep="normalized")
#
#     Removes the constant to the behavior if it does not already exist.
#
# Input:
#     α/β: Tuple, number of inputs and outputs for each party.
#     behaviors: Array, contains behavior column vectors.
#     rep: String, default "normalized", specifies the representation of the behavior
#
# Output:
#     corrected_behaviors: Array, contains corrected behavior column vectors.
# """
function remove_constants(α, β, behaviors; rep="normalized")
    map(
        (behavior) -> begin
            has_constant = Behavior.has_constant(α, β, behavior, rep=rep)

            corrected_behavior = behavior
            if has_constant
                corrected_behavior = behavior[2:end]
            end

            corrected_behavior
        end,
        behaviors
    )
end

end


text:
# this object represents a slice of the hamiltonian at a single position.
# you can then do things like transfer_left(vector,ham[5],st.AL[5],st.AR[5])

struct MPOHamSlice{O<:MPOHamiltonian}
    ham::O
    i::Int
end

function Base.getproperty(x::MPOHamSlice,s::Symbol)
    if s == :imspaces
        return x.ham.imspaces[x.i,:]
    elseif s == :domspaces
        return x.ham.domspaces[x.i,:]
    elseif s == :odim
        return x.ham.odim
    elseif s == :Os
        return x.ham.Os[x.i,:,:]
    else
        return getfield(x,s)
    end
end

Base.getindex(x::MPOHamiltonian,a::AbstractVector{Int}) = [x[i] for i in a];
Base.getindex(x::MPOHamiltonian,a::Int) = MPOHamSlice(x,a);

Base.getindex(x::MPOHamSlice,args...) = Base.getindex(x.ham,x.i,args...);
isscal(x::MPOHamSlice,j,k) = isscal(x.ham,x.i,j,k);
Base.keys(x::MPOHamSlice) = keys(x.ham,x.i);
scalkeys(x::MPOHamSlice) = scalkeys(x.ham,x.i);
opkeys(x::MPOHamSlice) = opkeys(x.ham,x.i);
Base.contains(x::MPOHamSlice,j,k) = contains(x.ham,x.i,j,k);

#the usual mpoham transfer
function transfer_left(vec::Vector{V},ham::MPOHamSlice,A::V,Ab::V=A) where V<:MPSTensor
    toret = [TensorMap(zeros,eltype(A),_lastspace(Ab)'*ham.imspaces[i],_lastspace(A)') for i in 1:ham.odim]::Vector{V}

    @sync for k in 1:ham.odim
        @Threads.spawn toret[k] = foldl(+, 1:ham.odim |>
            Filter(j->contains(ham,j,k)) |>
            Map() do j
                if isscal(ham,j,k)
                    ham.Os[j,k]*transfer_left(vec[j],A,Ab)
                else
                    transfer_left(vec[j],ham[j,k],A,Ab)
                end
            end,init=toret[k]);
    end

    return toret
end
function transfer_right(vec::Vector{V},ham::MPOHamSlice,A::V,Ab::V=A) where V<:MPSTensor
    toret = [TensorMap(zeros,eltype(A),_firstspace(A)*ham.domspaces[i],_firstspace(Ab)) for i in 1:ham.odim]::Vector{V}

    @sync for j in 1:ham.odim
        @Threads.spawn toret[j] = foldl(+, 1:ham.odim |>
            Filter(k->contains(ham,j,k)) |>
            Map() do k
                if isscal(ham,j,k)
                    ham.Os[j,k]*transfer_right(vec[k],A,Ab)
                else
                    transfer_right(vec[k],ham[j,k],A,Ab)
                end
            end,init=toret[j]);
    end
    return toret
end


text:
@testset "test base" begin
    @testset "huber_loss" begin
        @testset "all linear delta" begin
            δ = 1.
            predictions = [1.5, -1.4, -1.0, 0.0]
            labels = [0.0, 1.0, 0.0, 1.5]
            expected = mean(δ .* [1.5, 2.4, 1.0, 1.5]) .- 0.5 .* δ ^2
            @test expected ≈ mean(huber_loss(labels, predictions;δ=δ))
        end
        @testset "all quadratic delta" begin
            δ = 1.
            predictions = [1.5, -1.4, -0.5, 0.0]
            labels = [1.0, -1.0, 0.0, 0.5]
            expected = mean(0.5 .* [0.5, 0.4, 0.5, 0.5] .^2)
            @test expected ≈ mean(huber_loss(labels, predictions;δ=δ))
        end
    end

    @testset "findallmax" begin
        @test findallmax([-Inf, -Inf, -Inf]) == (-Inf, [1, 2, 3])
        @test findallmax([Inf, Inf, Inf]) == (Inf, [1, 2, 3])
        @test findallmax([0, 1, 2, 1, 2, 1, 0]) == (2, [3, 5])
        @test begin
            max_val, inds = findallmax([NaN, NaN, NaN])
            isnan(max_val) && inds == [1, 2, 3]
        end
    end

    @testset "discount_rewards" begin
        reward, γ = Float64[1, 2, 3], 0.5
        @test discount_rewards(reward, γ) ≈ [2.75, 3.5, 3.0]

        reward, γ = [1.0], 0.5
        @test discount_rewards(reward, γ) ≈ reward
    end

    @testset "SampleAvg" begin
        f = SampleAvg()
        @test f(2) ≈ 2.0
        @test f(3) ≈ (2 + 3) / 2
        @test f(5) ≈ (2 + 3 + 5) / 3
    end

    @testset "CachedSampleAvg" begin
        f = CachedSampleAvg()
        @test f(:a, 3) ≈ 3
        @test f(:a, 5) ≈ (3 + 5) / 2
        @test f(:a, 8) ≈ (3 + 5 + 8) / 3
        @test f(:b, 0) ≈ 0
    end

    @testset "CachedSum" begin
        f = CachedSum()
        @test f(:a, 1) ≈ 1
        @test f(:a, 2) ≈ 3
        @test f(:a, 3) ≈ 6
        @test f(:b, 0.3) ≈ 0.3
    end
end


text:
using OptimizationDynamics
const iLQR = OptimizationDynamics.IterativeLQR
const RoboDojo = OptimizationDynamics.RoboDojo
using LinearAlgebra
using Random

# ## visualize
vis = Visualizer()
render(vis);

# ## state-space model
T = 21
h = 0.05
hopper = RoboDojo.hopper

struct ParameterOptInfo{T}
	idx_q1::Vector{Int}
	idx_q2::Vector{Int}
	idx_u1::Vector{Int}
	idx_uθ::Vector{Int}
	idx_uθ1::Vector{Int}
	idx_uθ2::Vector{Int}
	idx_xθ::Vector{Int}
	v1::Vector{T}
end

info = ParameterOptInfo(
	collect(1:hopper.nq),
	collect(hopper.nq .+ (1:hopper.nq)),
	collect(1:hopper.nu),
	collect(hopper.nu .+ (1:2 * hopper.nq)),
	collect(hopper.nu .+ (1:hopper.nq)),
	collect(hopper.nu + hopper.nq .+ (1:hopper.nq)),
	collect(2 * hopper.nq .+ (1:2 * hopper.nq)),
	zeros(hopper.nq)
)

im_dyn1 = ImplicitDynamics(hopper, h,
	eval(RoboDojo.residual_expr(hopper)),
	eval(RoboDojo.jacobian_var_expr(hopper)),
	eval(RoboDojo.jacobian_data_expr(hopper));
    r_tol=1.0e-8, κ_eval_tol=1.0e-4, κ_grad_tol=1.0e-3,
	n=(2 * hopper.nq), m=(hopper.nu + 2 * hopper.nq), nc=4, nb=2, info=info)

im_dynt = ImplicitDynamics(hopper, h,
	eval(RoboDojo.residual_expr(hopper)),
	eval(RoboDojo.jacobian_var_expr(hopper)),
	eval(RoboDojo.jacobian_data_expr(hopper));
    r_tol=1.0e-8, κ_eval_tol=1.0e-4, κ_grad_tol=1.0e-3,
	n=4 * hopper.nq, m=hopper.nu, nc=4, nb=2, info=info)

function f1(d, model::ImplicitDynamics, x, u, w)

	θ = @views u[model.info.idx_uθ]
	q1 = @views u[model.info.idx_uθ1]
	q2 = @views u[model.info.idx_uθ2]
	u1 = @views u[model.info.idx_u1]

	model.info.v1 .= q2
	model.info.v1 .-= q1
	model.info.v1 ./= model.eval_sim.h

	q3 = RoboDojo.step!(model.eval_sim, q2, model.info.v1, u1, 1)

	d[model.info.idx_q1] = q2
	d[model.info.idx_q2] = q3
	d[model.info.idx_xθ] = θ

	return d
end

function f1x(dx, model::ImplicitDynamics, x, u, w)
	dx .= 0.0
	return dx
end

function f1u(du, model::ImplicitDynamics, x, u, w)
	nq = model.grad_sim.model.nq

	θ = @views u[model.info.idx_uθ]
	q1 = @views u[model.info.idx_uθ1]
	q2 = @views u[model.info.idx_uθ2]
	u1 = @views u[model.info.idx_u1]

	model.info.v1 .= q2
	model.info.v1 .-= q1
	model.info.v1 ./= model.grad_sim.h

	RoboDojo.step!(model.grad_sim, q2, model.info.v1, u1, 1)

	for i = 1:nq
		du[model.info.idx_q1[i], model.info.idx_uθ[i]] = 1.0
	end
	du[model.info.idx_q2, model.info.idx_u1] = model.grad_sim.grad.∂q3∂u1[1]
	du[model.info.idx_q2, model.info.idx_uθ1] = model.grad_sim.grad.∂q3∂q1[1]
	du[model.info.idx_q2, model.info.idx_uθ2] = model.grad_sim.grad.∂q3∂q2[1]

	return du
end

function ft(d, model::ImplicitDynamics, x, u, w)

	θ = @views x[model.info.idx_xθ]
	q1 = @views x[model.info.idx_q1]
	q2 = @views x[model.info.idx_q2]
	u1 = u

	model.info.v1 .= q2
	model.info.v1 .-= q1
	model.info.v1 ./= model.eval_sim.h

	q3 = RoboDojo.step!(model.eval_sim, q2, model.info.v1, u1, 1)

	d[model.info.idx_q1] = q2
	d[model.info.idx_q2] = q3
	d[model.info.idx_xθ] = θ

	return d
end

function ftx(dx, model::ImplicitDynamics, x, u, w)
	nq = model.grad_sim.model.nq

	θ = @views x[model.info.idx_xθ]
	q1 = @views x[model.info.idx_q1]
	q2 = @views x[model.info.idx_q2]
	u1 = u

	model.info.v1 .= q2
	model.info.v1 .-= q1
	model.info.v1 ./= model.grad_sim.h

	q3 = RoboDojo.step!(model.grad_sim, q2, model.info.v1, u1, 1)

	for i = 1:nq
		dx[model.info.idx_q1[i], model.info.idx_q2[i]] = 1.0
	end
	dx[model.info.idx_q2, model.info.idx_q1] = model.grad_sim.grad.∂q3∂q1[1]
	dx[model.info.idx_q2, model.info.idx_q2] = model.grad_sim.grad.∂q3∂q2[1]
	for i in model.info.idx_xθ
		dx[i, i] = 1.0
	end

	return dx
end

function ftu(du, model::ImplicitDynamics, x, u, w)
	θ = @views x[model.info.idx_xθ]
	q1 = @views x[model.info.idx_q1]
	q2 = @views x[model.info.idx_q2]
	u1 = u

	model.info.v1 .= q2
	model.info.v1 .-= q1
	model.info.v1 ./= model.grad_sim.h

	q3 = RoboDojo.step!(model.grad_sim, q2, model.info.v1, u1, 1)

	du[model.info.idx_q2, model.info.idx_u1] = model.grad_sim.grad.∂q3∂u1[1]

	return du
end

# ## iLQR model
ilqr_dyn1 = iLQR.Dynamics((d, x, u, w) -> f1(d, im_dyn1, x, u, w),
					(dx, x, u, w) -> f1x(dx, im_dyn1, x, u, w),
					(du, x, u, w) -> f1u(du, im_dyn1, x, u, w),
					4 * hopper.nq, 2 * hopper.nq, hopper.nu + 2 * hopper.nq)

ilqr_dynt = iLQR.Dynamics((d, x, u, w) -> ft(d, im_dynt, x, u, w),
	(dx, x, u, w) -> ftx(dx, im_dynt, x, u, w),
	(du, x, u, w) -> ftu(du, im_dynt, x, u, w),
	4 * hopper.nq, 4 * hopper.nq, hopper.nu)

model = [ilqr_dyn1, [ilqr_dynt for t = 2:T-1]...];

# ## initial conditions
q1 = [0.0; 0.5 + hopper.foot_radius; 0.0; 0.5]
qM = [0.5; 0.5 + hopper.foot_radius; 0.0; 0.5]
qT = [1.0; 0.5 + hopper.foot_radius; 0.0; 0.5]
q_ref = [0.5; 0.75 + hopper.foot_radius; 0.0; 0.25]

x1 = [q1; q1]
xM = [qM; qM]
xT = [qT; qT]
x_ref = [q_ref; q_ref]

# ## objective

GAIT = 1
## GAIT = 2
## GAIT = 3

if GAIT == 1
	r_cost = 1.0e-1
	q_cost = 1.0e-1
elseif GAIT == 2
	r_cost = 1.0
	q_cost = 1.0
elseif GAIT == 3
	r_cost = 1.0e-3
	q_cost = 1.0e-1
end

function obj1(x, u, w)
	J = 0.0
	J += 0.5 * transpose(x - x_ref) * Diagonal([1.0; 10.0; 1.0; 10.0; 1.0; 10.0; 1.0; 10.0]) * (x - x_ref)
	J += 0.5 * transpose(u) * Diagonal([r_cost * ones(hopper.nu); 1.0e-1 * ones(hopper.nq); 1.0e-5 * ones(hopper.nq)]) * u
	return J
end

function objt(x, u, w)
	J = 0.0
	J += 0.5 * transpose(x - [x_ref; zeros(2 * hopper.nq)]) * q_cost * Diagonal([1.0; 10.0; 1.0; 10.0; 1.0; 10.0; 1.0; 10.0; zeros(2 * hopper.nq)]) * (x - [x_ref; zeros(2 * hopper.nq)])
	J += 0.5 * transpose(u) * Diagonal(r_cost * ones(hopper.nu)) * u
	return J
end

function objT(x, u, w)
	J = 0.0
	J += 0.5 * transpose(x - [x_ref; zeros(2 * hopper.nq)]) * Diagonal([1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; 1.0; zeros(2 * hopper.nq)]) * (x - [x_ref; zeros(2 * hopper.nq)])
	return J
end

c1 = iLQR.Cost(obj1, 2 * hopper.nq, hopper.nu + 2 * hopper.nq)
ct = iLQR.Cost(objt, 4 * hopper.nq, hopper.nu)
cT = iLQR.Cost(objT, 4 * hopper.nq, 0)
obj = [c1, [ct for t = 2:T-1]..., cT];

# ## constraints
ul = [-10.0; -10.0]
uu = [10.0; 10.0]

function stage1_con(x, u, w)
    [
    ul - u[1:hopper.nu]; # control limit (lower)
    u[1:hopper.nu] - uu; # control limit (upper)

	u[hopper.nu .+ (1:hopper.nq)] - x1[1:hopper.nq];

	RoboDojo.kinematics_foot(hopper, u[hopper.nu .+ (1:hopper.nq)]) - RoboDojo.kinematics_foot(hopper, x1[1:hopper.nq]);
	RoboDojo.kinematics_foot(hopper, u[hopper.nu + hopper.nq .+ (1:hopper.nq)]) - RoboDojo.kinematics_foot(hopper, x1[hopper.nq .+ (1:hopper.nq)])
    ]
end

function staget_con(x, u, w)
    [
    ul - u[collect(1:hopper.nu)]; # control limit (lower)
    u[collect(1:hopper.nu)] - uu; # control limit (upper)
    ]
end

function terminal_con(x, u, w)
	x_travel = 0.5
	θ = x[2 * hopper.nq .+ (1:(2 * hopper.nq))]
    [
	x_travel - (x[1] - θ[1])
	x_travel - (x[hopper.nq + 1] - θ[hopper.nq + 1])
	x[1:hopper.nq][collect([2, 3, 4])] - θ[1:hopper.nq][collect([2, 3, 4])]
	x[hopper.nq .+ (1:hopper.nq)][collect([2, 3, 4])] - θ[hopper.nq .+ (1:hopper.nq)][collect([2, 3, 4])]
    ]
end

con1 = iLQR.Constraint(stage1_con, 2 * hopper.nq, hopper.nu + 2 * hopper.nq, idx_ineq=collect(1:4))
cont = iLQR.Constraint(staget_con, 4 * hopper.nq, hopper.nu, idx_ineq=collect(1:4))
conT = iLQR.Constraint(terminal_con, 4 * hopper.nq, 0, idx_ineq=collect(1:2))
cons = [con1, [cont for t = 2:T-1]..., conT];

# ## rollout
ū_stand = [t == 1 ? [0.0; hopper.gravity * hopper.mass_body * 0.5 * h; x1] : [0.0; hopper.gravity * hopper.mass_body * 0.5 * h] for t = 1:T-1]
x̄ = iLQR.rollout(model, x1, ū_stand)
q̄ = state_to_configuration(x̄)
RoboDojo.visualize!(vis, hopper, x̄, Δt=h);

# ## solver
solver = iLQR.solver(model, obj, cons,
	opts=iLQR.Options(linesearch = :armijo,
		α_min=1.0e-5,
		obj_tol=1.0e-3,
		grad_tol=1.0e-3,
		max_iter=10,
		max_al_iter=15,
		con_tol=0.001,
		ρ_init=1.0,
		ρ_scale=10.0,
		verbose=false))
iLQR.initialize_controls!(solver, ū_stand)
iLQR.initialize_states!(solver, x̄);

# ## solve
iLQR.reset!(solver.s_data)
@time iLQR.solve!(solver);

@show iLQR.eval_obj(solver.m_data.obj.costs, solver.m_data.x, solver.m_data.u, solver.m_data.w)
@show solver.s_data.iter[1]
@show norm(terminal_con(solver.m_data.x[T], zeros(0), zeros(0))[3:4], Inf)
@show solver.s_data.obj[1] # augmented Lagrangian cost

# ## solution
x_sol, u_sol = iLQR.get_trajectory(solver)
q_sol = state_to_configuration(x_sol)
RoboDojo.visualize!(vis, hopper, q_sol, Δt=h);

# ## benchmark (NOTE: gate 3 seems to break @benchmark, just run @time instead...)
solver.options.verbose = false
@benchmark iLQR.solve!($solver, $x̄, $ū_stand);


text:
Base.@kwdef struct MOEAD <: AbstractOptimizer
    # number of the subproblems
    N::Int = 100
    # number of the weight vectors in the neighborhoor of each weight vector
    T::Int = ceil(N/10)
    crossover = BINX(0.5)
    mutation = PLM()
    metrics::ConvergenceMetrics = ConvergenceMetric[GD(), GD(true)]
end

population_size(method::MOEAD) = method.N
default_options(method::MOEAD) = (iterations = 1000,)
summary(m::MOEAD) =
    "MOEA/D[P=$(population_size(m)),]"
show(io::IO, m::MOEAD) = print(io, summary(m))

mutable struct MOEADState <: AbstractOptimizerState
    iteration
    start_time
    stop_time
    metrics::ConvergenceMetrics # collection of convergence metrics

    population                  # population
    m                           # number of objectives
    W                           # weight vectors
    B                           # neighbours of each solution
    z                           # the best value found so far
    EP                          # nondominated solutions
end
pfront(s::MOEADState) = s.EP
value(s::MOEADState) = objectives(pfront(s))
minimizer(s::MOEADState) = objectives(pfront(s))
copy(s::MOEADState) = MOEADState(s.iteration, s.start_time, s.stop_time, copy(s.metrics),
                                 copy(s.population),
                                 copy(s.m),
                                 copy(s.W),
                                 copy(s.B),
                                 copy(s.z),
                                 copy(s.EP))

function initial_state(method::MOEAD, objective, population, options)
    value!(objective, population)

    m = length(objectives(first(population)))

    # Generate weight vectors
    W = uniform_point(population_size(method), m)

    # Set the neighbours of each weight vector
    B = pairwise(Euclidean(), W, W, dims=1)
    B = sortperm(B)
    B = B[:,2:method.T]

    # Set reference point
    z = Matrix(minimum(objectives(population), dims=1)')

    EP = Individual[]
    return MOEADState(0, 0, 0, copy(method.metrics), population, m, W, B, z, EP)
end

function update_state!(
    method::MOEAD,
    state,
    objective,
    constraints,
    options
)
    for i = 1:method.N
        x = state.population[i]

        # Reproduction
        xi = rand(options.rng, state.B[i,:], 2)
        xs = state.population[xi]
        y,_ = method.crossover(xs[1], xs[2], rng=options.rng)

        # Improvement
        y = method.mutation(y)

        apply!(constraints, variables(y))
        value!(objective, y)

        # Update reference point
        for j = 1:state.m
            if state.z[j] > objectives(y)[j]
                state.z[j] = objectives(y)[j]
            end
        end

        # Update of neighboring solutions
        for j = state.B[i,:]
            x = state.population[j]
            g_x, _ = findmax(state.W[j,:] .* abs.(objectives(x) - state.z))
            g_y, _ = findmax(state.W[j,:] .* abs.(objectives(y) - state.z))
            if g_y <= g_x
                state.population[j] = y
            end
        end

        # Update of EP
        deleteat!(state.EP, findall((p) -> dominate(y, p) == 1, state.EP))
        if isempty(findall((p) -> dominate(p, y) == 1, state.EP))
            pushfirst!(state.EP, y)
        end
    end

    return false
end

# 参考 https://github.com/stevengj/Sobol.jl ？
# 这里使用 Das and Dennis's method
function uniform_point(N, M)
    # gen_ref_dirs 会生成 N+1 个...，需修复
    reduce(vcat, gen_ref_dirs(M, N-1)')
    # gen_ref_dirs(M, N-1)
end

function gen_ref_dirs(dimension, n_paritions)
    gen_weights(dimension, n_paritions)
end

function gen_weights(a, b)
    nobj = a;
    H    = b;
    a    = zeros(nobj);
    d    = H;
    w    = [];
    produce_weight!(a, 1, d, H, nobj, w)
    return Array.(w)
end

function  produce_weight!(a, i, d, H, nobj, w)
    for k=0:d
        if i<nobj
            a[i] = k;
            d2   = d - k;
            produce_weight!(a, i+1, d2, H, nobj, w);
        else
            a[i] = d;
            push!(w, a/H)
            break;
        end
    end
end

function Base.sortperm(M::Matrix)
    reduce(vcat, map(sortperm, eachrow(M))')
end

function Base.sortperm!(M::Matrix)
    error("Not implemented")
end


text:
# Piracy that used to be requires, allowing ReverseDiff.jl to be specialized for SciML

DiffEqBase.value(x::ReverseDiff.TrackedReal) = x.value
DiffEqBase.value(x::ReverseDiff.TrackedArray) = x.value

DiffEqBase.promote_u0(u0::ReverseDiff.TrackedArray, p::ReverseDiff.TrackedArray, t0) = u0
DiffEqBase.promote_u0(u0::AbstractArray{<:ReverseDiff.TrackedReal}, p::ReverseDiff.TrackedArray, t0) = u0
DiffEqBase.promote_u0(u0::ReverseDiff.TrackedArray, p::AbstractArray{<:ReverseDiff.TrackedReal}, t0) = u0
DiffEqBase.promote_u0(u0::AbstractArray{<:ReverseDiff.TrackedReal}, p::AbstractArray{<:ReverseDiff.TrackedReal}, t0) = u0
DiffEqBase.promote_u0(u0, p::ReverseDiff.TrackedArray, t0) = ReverseDiff.track(u0)
DiffEqBase.promote_u0(u0, p::AbstractArray{<:ReverseDiff.TrackedReal}, t0) = eltype(p).(u0)

# Support adaptive with non-tracked time
@inline function DiffEqBase.ODE_DEFAULT_NORM(u::ReverseDiff.TrackedArray, t) where {N}
    sqrt(sum(abs2, DiffEqBase.value(u)) / length(u))
end
@inline function DiffEqBase.ODE_DEFAULT_NORM(u::AbstractArray{<:ReverseDiff.TrackedReal,N}, t) where {N}
    sqrt(sum(x -> DiffEqBase.ODE_DEFAULT_NORM(x[1], x[2]), zip((DiffEqBase.value(x) for x in u), Iterators.repeated(t))) / length(u))
end
@inline function DiffEqBase.ODE_DEFAULT_NORM(u::Array{<:ReverseDiff.TrackedReal,N}, t) where {N}
    sqrt(sum(x -> DiffEqBase.ODE_DEFAULT_NORM(x[1], x[2]), zip((DiffEqBase.value(x) for x in u), Iterators.repeated(t))) / length(u))
end
@inline DiffEqBase.ODE_DEFAULT_NORM(u::ReverseDiff.TrackedReal, t) = abs(DiffEqBase.value(u))

# Support TrackedReal time, don't drop tracking on the adaptivity there
@inline function DiffEqBase.ODE_DEFAULT_NORM(u::ReverseDiff.TrackedArray, t::ReverseDiff.TrackedReal) where {N}
    sqrt(sum(abs2, u) / length(u))
end
@inline function DiffEqBase.ODE_DEFAULT_NORM(u::AbstractArray{<:ReverseDiff.TrackedReal,N}, t::ReverseDiff.TrackedReal) where {N}
    sqrt(sum(x -> DiffEqBase.ODE_DEFAULT_NORM(x[1], x[2]), zip(u, Iterators.repeated(t))) / length(u))
end
@inline function DiffEqBase.ODE_DEFAULT_NORM(u::Array{<:ReverseDiff.TrackedReal,N}, t::ReverseDiff.TrackedReal) where {N}
    sqrt(sum(x -> DiffEqBase.ODE_DEFAULT_NORM(x[1], x[2]), zip(u, Iterators.repeated(t))) / length(u))
end
@inline DiffEqBase.ODE_DEFAULT_NORM(u::ReverseDiff.TrackedReal, t::ReverseDiff.TrackedReal) = abs(u)

function DiffEqBase.solve_up(prob::DiffEqBase.DEProblem, sensealg::Union{DiffEqBase.AbstractSensitivityAlgorithm,Nothing}, u0::ReverseDiff.TrackedArray, p::ReverseDiff.TrackedArray, args...; kwargs...)
    ReverseDiff.track(solve_up, prob, sensealg, u0, p, args...; kwargs...)
end

function DiffEqBase.solve_up(prob::DiffEqBase.DEProblem, sensealg::Union{DiffEqBase.AbstractSensitivityAlgorithm,Nothing}, u0, p::ReverseDiff.TrackedArray, args...; kwargs...)
    ReverseDiff.track(solve_up, prob, sensealg, u0, p, args...; kwargs...)
end

function DiffEqBase.solve_up(prob::DiffEqBase.DEProblem, sensealg::Union{DiffEqBase.AbstractSensitivityAlgorithm,Nothing}, u0::ReverseDiff.TrackedArray, p, args...; kwargs...)
    ReverseDiff.track(solve_up, prob, sensealg, u0, p, args...; kwargs...)
end

@inline function DiffEqNoiseProcess.wiener_randn(rng::Random.AbstractRNG, proto::ReverseDiff.TrackedArray)
    ReverseDiff.track(convert.(eltype(proto.value), randn(rng, size(proto))))
end
@inline function DiffEqNoiseProcess.wiener_randn!(rng::AbstractRNG, rand_vec::Array{<:ReverseDiff.TrackedReal})
    rand_vec .= ReverseDiff.track.(randn.((rng,), typeof.(DiffEqBase.value.(rand_vec))))
end
@inline function DiffEqNoiseProcess.wiener_randn!(rng::AbstractRNG, rand_vec::AbstractArray{<:ReverseDiff.TrackedReal})
    rand_vec .= ReverseDiff.track.(randn.((rng,), typeof.(DiffEqBase.value.(rand_vec))))
end