text:
module TestModel

using Test

import VecMathOptInterface

const MOI = VecMathOptInterface

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

###
### Test helpers
###

"""
    TestExternalModel

We need to define this in a module at the top level because it can't be
defined in a testset. If it runs without error, then we're okay.
"""
module TestExternalModel

using VecMathOptInterface

struct NewSet <: VecMathOptInterface.AbstractScalarSet end

struct NewFunction <: VecMathOptInterface.AbstractScalarFunction end

VecMathOptInterface.Utilities.canonicalize!(f::NewFunction) = f

Base.copy(::NewFunction) = NewFunction()

Base.copy(::NewSet) = NewSet()

VecMathOptInterface.Utilities.@model(
    ExternalModel,
    (VecMathOptInterface.ZeroOne, NewSet),
    (),
    (),
    (),
    (NewFunction,),
    (VecMathOptInterface.ScalarAffineFunction,),
    (),
    (),
)

VecMathOptInterface.Utilities.@model(
    ExternalOptimizer,
    (VecMathOptInterface.ZeroOne, NewSet),
    (),
    (),
    (),
    (NewFunction,),
    (),
    (),
    (),
    true,
)

end

struct DummyFunction <: MOI.AbstractScalarFunction end

struct DummySet <: MOI.AbstractScalarSet end

# We create a new function and set to test catching errors if users create their
# own sets and functions
struct FunctionNotSupportedBySolvers <: MOI.AbstractScalarFunction
    x::MOI.VariableIndex
end

struct SetNotSupportedBySolvers <: MOI.AbstractSet end

###
### The tests start here.
###

function test_MOI_Test()
    MOI.Test.runtests(
        MOI.Utilities.Model{Float64}(),
        MOI.Test.Config(exclude = Any[MOI.optimize!,]),
    )
    return
end

function test_TestExternalModel_fields()
    @test isdefined(TestExternalModel, :ExternalModelScalarConstraints)
    @test !isdefined(TestExternalModel, :ExternalModelVectorConstraints)
    @test isdefined(TestExternalModel, :ExternalModelFunctionConstraints)
    @test isdefined(TestExternalModel, :ExternalOptimizerScalarConstraints)
    @test !isdefined(TestExternalModel, :ExternalOptimizerVectorConstraints)
    @test !isdefined(TestExternalModel, :ExternalOptimizerFunctionConstraints)
    @test TestExternalModel.ExternalModel{Int} == MOI.Utilities.GenericModel{
        Int,
        MOI.Utilities.ObjectiveContainer{Int},
        MOI.Utilities.VariablesContainer{Int},
        TestExternalModel.ExternalModelFunctionConstraints{Int},
    }
    model = MOI.Utilities.GenericOptimizer{
        Int,
        MOI.Utilities.ObjectiveContainer{Int},
        MOI.Utilities.VariablesContainer{Int},
        TestExternalModel.ExternalOptimizerScalarConstraints{
            Int,
            MOI.Utilities.VectorOfConstraints{
                TestExternalModel.NewFunction,
                MOI.ZeroOne,
            },
            MOI.Utilities.VectorOfConstraints{
                TestExternalModel.NewFunction,
                TestExternalModel.NewSet,
            },
        },
    }
    @test TestExternalModel.ExternalOptimizer{Int} == model

    return
end

function test_TestExternalModel_supertypes()
    model = TestExternalModel.ExternalModel{Float64}()
    optimizer = TestExternalModel.ExternalOptimizer{Float64}()
    @test isa(model, MOI.Utilities.AbstractModelLike{Float64})
    @test !isa(model, MOI.Utilities.AbstractOptimizer{Float64})
    @test !isa(model, MOI.AbstractOptimizer)
    @test !isa(optimizer, MOI.Utilities.AbstractModelLike{Float64})
    @test isa(optimizer, MOI.Utilities.AbstractOptimizer{Float64})
    @test isa(optimizer, MOI.AbstractOptimizer)
    return
end

function test_TestExternalModel()
    model = TestExternalModel.ExternalModel{Float64}()
    c = MOI.add_constraint(
        model,
        TestExternalModel.NewFunction(),
        TestExternalModel.NewSet(),
    )
    @test typeof(c) == MOI.ConstraintIndex{
        TestExternalModel.NewFunction,
        TestExternalModel.NewSet,
    }
    c2 = MOI.add_constraint(
        model,
        TestExternalModel.NewFunction(),
        MOI.ZeroOne(),
    )
    @test typeof(c2) ==
          MOI.ConstraintIndex{TestExternalModel.NewFunction,MOI.ZeroOne}
    return
end

function _test_ObjectiveFunction(T)
    model = MOI.Utilities.Model{T}()
    @test !MOI.supports(model, MOI.ObjectiveFunction{DummyFunction}())
    for F in [
        MOI.VariableIndex,
        MOI.ScalarAffineFunction{T},
        MOI.ScalarQuadraticFunction{T},
    ]
        @test MOI.supports(model, MOI.ObjectiveFunction{F}())
    end
    U = Float32
    for F in [MOI.ScalarAffineFunction{U}, MOI.ScalarQuadraticFunction{U}]
        @test !MOI.supports(model, MOI.ObjectiveFunction{F}())
    end
    return
end

function test_ObjectiveFunction()
    _test_ObjectiveFunction(Float64)
    _test_ObjectiveFunction(Int)
    return
end

function _test_VariableIndex(T)
    model = MOI.Utilities.Model{T}()
    @test !MOI.supports_constraint(model, DummyFunction, DummySet)
    @test !MOI.supports_constraint(model, MOI.VariableIndex, DummySet)
    @test !MOI.supports_add_constrained_variable(model, DummySet)
    for S in [
        MOI.EqualTo{T},
        MOI.GreaterThan{T},
        MOI.LessThan{T},
        MOI.Interval{T},
        MOI.Integer,
        MOI.ZeroOne,
        MOI.Semicontinuous{T},
        MOI.Semiinteger{T},
    ]
        @test MOI.supports_constraint(model, MOI.VariableIndex, S)
        @test MOI.supports_add_constrained_variable(model, S)
    end
    U = Float32
    for S in [
        MOI.EqualTo{U},
        MOI.GreaterThan{U},
        MOI.LessThan{U},
        MOI.Interval{U},
        MOI.Semicontinuous{U},
        MOI.Semiinteger{U},
    ]
        @test !MOI.supports_constraint(model, MOI.VariableIndex, S)
        @test !MOI.supports_add_constrained_variable(model, S)
    end
    return
end

function test_VariableIndex()
    _test_VariableIndex(Float64)
    _test_VariableIndex(Int)
    return
end

function test_quadratic_functions()
    model = MOI.Utilities.Model{Int}()
    x, y = MOI.add_variables(model, 2)
    @test 2 == @inferred MOI.get(model, MOI.NumberOfVariables())
    f1 = MOI.ScalarQuadraticFunction(
        MOI.ScalarQuadraticTerm.([1, 2, 3], [x, y, x], [x, y, y]),
        MOI.ScalarAffineTerm.([3], [x]),
        7,
    )
    c1 = MOI.add_constraint(model, f1, MOI.Interval(-1, 1))
    @test 1 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{
            MOI.ScalarQuadraticFunction{Int},
            MOI.Interval{Int},
        }(),
    )
    @test (@inferred MOI.get(
        model,
        MOI.ListOfConstraintIndices{
            MOI.ScalarQuadraticFunction{Int},
            MOI.Interval{Int},
        }(),
    )) == [c1]
    f2 = MOI.VectorQuadraticFunction(
        MOI.VectorQuadraticTerm.(
            [1, 1, 2],
            MOI.ScalarQuadraticTerm.([1, 2, 3], [x, y, x], [x, y, y]),
        ),
        MOI.VectorAffineTerm.(
            [1, 2, 2],
            MOI.ScalarAffineTerm.([3, 1, 2], [x, x, y]),
        ),
        [7, 3, 4],
    )
    c2 = MOI.add_constraint(model, f2, MOI.PositiveSemidefiniteConeTriangle(3))
    @test 1 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{
            MOI.VectorQuadraticFunction{Int},
            MOI.PositiveSemidefiniteConeTriangle,
        }(),
    )
    @test (@inferred MOI.get(
        model,
        MOI.ListOfConstraintIndices{
            MOI.VectorQuadraticFunction{Int},
            MOI.PositiveSemidefiniteConeTriangle,
        }(),
    )) == [c2]
    loc = MOI.get(model, MOI.ListOfConstraintTypesPresent())
    @test length(loc) == 2
    @test (
        MOI.VectorQuadraticFunction{Int},
        MOI.PositiveSemidefiniteConeTriangle,
    ) in loc
    @test (
        MOI.VectorQuadraticFunction{Int},
        MOI.PositiveSemidefiniteConeTriangle,
    ) in loc
    c3 = MOI.add_constraint(model, f1, MOI.Interval(-1, 1))
    change_1 = MOI.ScalarConstantChange(9)
    f3 = MOI.Utilities.modify_function(f1, change_1)
    change_2 = MOI.ScalarCoefficientChange(y, 2)
    f3 = MOI.Utilities.modify_function(f3, change_2)
    @test !(MOI.get(model, MOI.ConstraintFunction(), c1) ≈ f3)
    @test !(MOI.get(model, MOI.ConstraintFunction(), c3) ≈ f3)
    MOI.set(model, MOI.ConstraintFunction(), c1, f3)
    MOI.modify(model, c3, change_1)
    MOI.modify(model, c3, change_2)
    F1 = MOI.get(model, MOI.CanonicalConstraintFunction(), c1)
    @test F1 ≈ f3
    @test F1 !== MOI.get(model, MOI.ConstraintFunction(), c1)
    @test MOI.Utilities.is_canonical(F1)
    F3 = MOI.get(model, MOI.CanonicalConstraintFunction(), c3)
    @test F3 ≈ f3
    @test F3 === MOI.get(model, MOI.ConstraintFunction(), c3)
    @test MOI.Utilities.is_canonical(F3)
    @test MOI.get(model, MOI.CanonicalConstraintFunction(), c1) ≈ f3
    @test MOI.Utilities.is_canonical(
        MOI.get(model, MOI.CanonicalConstraintFunction(), c1),
    )
    f4 = MOI.VectorAffineFunction(
        MOI.VectorAffineTerm.(
            [1, 1, 2],
            MOI.ScalarAffineTerm.([2, 4, 3], [x, y, y]),
        ),
        [5, 7],
    )
    c4 = MOI.add_constraint(model, f4, MOI.SecondOrderCone(2))
    @test 1 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{
            MOI.VectorAffineFunction{Int},
            MOI.SecondOrderCone,
        }(),
    )
    f5 = MOI.VectorOfVariables([x, x])
    c5 = MOI.add_constraint(model, f5, MOI.RotatedSecondOrderCone(2))
    @test 1 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{
            MOI.VectorOfVariables,
            MOI.RotatedSecondOrderCone,
        }(),
    )
    f6 = MOI.VectorAffineFunction(
        MOI.VectorAffineTerm.([1, 2], MOI.ScalarAffineTerm.([2, 9], [x, y])),
        [6, 8],
    )
    c6 = MOI.add_constraint(model, f6, MOI.SecondOrderCone(2))
    @test 2 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{
            MOI.VectorAffineFunction{Int},
            MOI.SecondOrderCone,
        }(),
    )
    f7 = MOI.VectorOfVariables([x, y])
    c7 = MOI.add_constraint(model, f7, MOI.Nonpositives(2))
    @test 1 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{MOI.VectorOfVariables,MOI.Nonpositives}(),
    )

    f8 = MOI.VectorOfVariables([x, y])
    c8 = MOI.add_constraint(model, f8, MOI.SecondOrderCone(2))
    @test 1 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{MOI.VectorOfVariables,MOI.SecondOrderCone}(),
    )
    loc1 = MOI.get(model, MOI.ListOfConstraintTypesPresent())
    loc2 = Vector{Tuple{Type,Type}}()
    function _pushloc(v::MOI.Utilities.VectorOfConstraints{F,S}) where {F,S}
        if !MOI.is_empty(v)
            push!(loc2, (F, S))
        end
    end
    _pushloc(::Nothing) = nothing
    function _pushloc(model::MOI.Utilities.StructOfConstraints)
        return MOI.Utilities.broadcastcall(_pushloc, model)
    end
    MOI.Utilities.broadcastcall(_pushloc, model.constraints)
    for loc in (loc1, loc2)
        @test length(loc) == 6
        @test (
            MOI.VectorQuadraticFunction{Int},
            MOI.PositiveSemidefiniteConeTriangle,
        ) in loc
        @test (
            MOI.VectorQuadraticFunction{Int},
            MOI.PositiveSemidefiniteConeTriangle,
        ) in loc
        @test (MOI.VectorOfVariables, MOI.RotatedSecondOrderCone) in loc
        @test (MOI.VectorAffineFunction{Int}, MOI.SecondOrderCone) in loc
        @test (MOI.VectorOfVariables, MOI.Nonpositives) in loc
    end
    @test MOI.is_valid(model, c4)
    MOI.delete(model, c4)
    @test !MOI.is_valid(model, c4)

    @test 1 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{
            MOI.VectorAffineFunction{Int},
            MOI.SecondOrderCone,
        }(),
    )
    @test MOI.get(model, MOI.ConstraintFunction(), c6).constants == f6.constants
    obj = 1x + 2y
    MOI.set(model, MOI.ObjectiveFunction{typeof(obj)}(), obj)
    message = string(
        "Cannot delete variable as it is constrained with other",
        " variables in a `MOI.VectorOfVariables`.",
    )
    err = MOI.DeleteNotAllowed(y, message)
    @test_throws err MOI.delete(model, y)
    @test MOI.get(model, MOI.ObjectiveFunction{typeof(obj)}()) ≈ 1x + 2y
    @test MOI.is_valid(model, c8)
    MOI.delete(model, c8)
    @test !MOI.is_valid(model, c8)
    @test 0 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{MOI.VectorOfVariables,MOI.SecondOrderCone}(),
    )
    MOI.delete(model, y)
    @test MOI.get(model, MOI.ObjectiveFunction{typeof(obj)}()) ≈ 1x
    f = MOI.get(model, MOI.ConstraintFunction(), c2)
    @test f.affine_terms ==
          MOI.VectorAffineTerm.([1, 2], MOI.ScalarAffineTerm.([3, 1], [x, x]))
    @test f.quadratic_terms ==
          MOI.VectorQuadraticTerm.([1], MOI.ScalarQuadraticTerm.([1], [x], [x]))
    @test f.constants == [7, 3, 4]
    f = MOI.get(model, MOI.ConstraintFunction(), c6)
    @test f.terms == MOI.VectorAffineTerm.([1], MOI.ScalarAffineTerm.([2], [x]))
    @test f.constants == [6, 8]
    @test 1 == @inferred MOI.get(
        model,
        MOI.NumberOfConstraints{MOI.VectorOfVariables,MOI.Nonpositives}(),
    )
    @test [c7] == @inferred MOI.get(
        model,
        MOI.ListOfConstraintIndices{MOI.VectorOfVariables,MOI.Nonpositives}(),
    )
    f = MOI.get(model, MOI.ConstraintFunction(), c7)
    @test f.variables == [x]
    s = MOI.get(model, MOI.ConstraintSet(), c7)
    @test MOI.dimension(s) == 1
end

function test_default_fallbacks()
    model = MOI.Utilities.Model{Float64}()
    x = MOI.add_variable(model)
    func = convert(MOI.ScalarAffineFunction{Float64}, x)
    c = MOI.add_constraint(model, func, MOI.LessThan(0.0))
    @test !MOI.supports_constraint(
        model,
        FunctionNotSupportedBySolvers,
        SetNotSupportedBySolvers,
    )

    # set of different type
    @test_throws Exception MOI.set(
        model,
        MOI.ConstraintSet(),
        c,
        MOI.GreaterThan(0.0),
    )
    # set not implemented
    @test_throws Exception MOI.set(
        model,
        MOI.ConstraintSet(),
        c,
        SetNotSupportedBySolvers(),
    )

    # function of different type
    @test_throws Exception MOI.set(
        model,
        MOI.ConstraintFunction(),
        c,
        MOI.VectorOfVariables([x]),
    )
    # function not implemented
    @test_throws Exception MOI.set(
        model,
        MOI.ConstraintFunction(),
        c,
        FunctionNotSupportedBySolvers(x),
    )
end

function test_extension_dictionary()
    model = MOI.Utilities.Model{Float64}()
    model.ext[:my_store] = 1
    @test model.ext[:my_store] == 1
    MOI.empty!(model)
    @test model.ext[:my_store] == 1
    model.ext[:my_store] = 2
    dest = MOI.Utilities.Model{Float64}()
    MOI.copy_to(dest, model)
    @test !haskey(dest.ext, :my_store)
    @test model.ext[:my_store] == 2
    return
end

end  # module

TestModel.runtests()


text:
# A Classified Collection is a map which associates a collection to each key
#
# The collection can be either an array or a set, a counter, or other data structures
# that support the push! method
#

mutable struct ClassifiedCollections{K, Collection}
    map::Dict{K, Collection}
end

## constructors

ClassifiedCollections(K::Type, C::Type) = ClassifiedCollections{K, C}(Dict{K,C}())

classified_lists(K::Type, V::Type) = ClassifiedCollections(K, Vector{V})
classified_sets(K::Type, V::Type) = ClassifiedCollections(K, Set{V})
classified_counters(K::Type, T::Type) = ClassifiedCollections(K, Accumulator{T, Int})

_create_empty(::Type{Vector{T}}) where {T} = Vector{T}()
_create_empty(::Type{Set{T}}) where {T} = Set{T}()
_create_empty(::Type{Accumulator{T,V}}) where {T,V} = Accumulator(T, V)

copy(cc::ClassifiedCollections{K, C}) where {K, C} = ClassifiedCollections{K, C}(copy(cc.map))

length(cc::ClassifiedCollections) = length(cc.map)

## retrieval

getindex(cc::ClassifiedCollections{T,C}, x::T) where {T,C} = cc.map[x]

haskey(cc::ClassifiedCollections{T,C}, x::T) where {T,C} = haskey(cc.map, x)

keys(cc::ClassifiedCollections) = keys(cc.map)

## iteration

start(cc::ClassifiedCollections) = start(cc.map)
next(cc::ClassifiedCollections, state) = next(cc.map, state)
done(cc::ClassifiedCollections, state) = done(cc.map, state)

# manipulation

function push!(cc::ClassifiedCollections{K, C}, key::K, e) where {K, C}
    c = get(cc.map, key, nothing)
    if c === nothing
        c = _create_empty(C)
        cc.map[key] = c
    end
    push!(c, e)
end

pop!(cc::ClassifiedCollections{K}, key::K) where {K} = pop!(cc.map, key)

ext: jl


text:
#!/usr/bin/env julia

# Generate a set of model and residual visibilities and then write them to UVHDF5 file format
# More about this format, including scripts to convert to UVFITS and CASA Measurement set,
# can be found at https://github.com/Astrochem/UVHDF5

using ArgParse

s = ArgParseSettings()
@add_arg_table s begin
    "--config"
    help = "a YAML configuration file"
    default = "config.yaml"
    "--out-model"
    help = "The file name to write out the model visibilities."
    default = "model.hdf5"
    "--out-resid"
    help = "The file name to write out the residual visibilities."
    default = "resid.hdf5"
end

parsed_args = parse_args(ARGS, s)

using DiskJockey.constants
using DiskJockey.image
using DiskJockey.model
using DiskJockey.visibilities
using DiskJockey.gridding

using HDF5

import YAML
config = YAML.load(open(parsed_args["config"]))

# The model writing is taken care of by `DJ_initialize.jl`, the RADMC-3D synthesis is taken care of by `DJ_plot_model.jl` and then this simply reads in the image, does FFT, downsamples, etc.

# read the wavelengths for all channels
fid = h5open(config["data_file"], "r")
nchan = length(read(fid["freqs"]))
close(fid)

# Read the parameters from the config file
pars = convert_dict(config["parameters"], config["model"])
grid = Grid(config["grid"])

# Mention the contribution of the prior to the lnprob
ln_prior = lnprior(pars, grid)

im = imread()
skim = imToSky(im, pars.dpc)
corrfun!(skim) # alpha = 1.0

# Determine dRA and dDEC from the image and distance
dRA = abs(skim.ra[2] - skim.ra[1]) / 2. # [arcsec] the half-size of a pixel
println("dRA is ", dRA, " arcsec")

# For *this purpose only*, read in the flagged data in addition to the unflagged data
# so that we can export a model for these otherwise flagged visibilities
dvarr = DataVis(config["data_file"], true)

for dset in dvarr
    # Conjugation is necessary for the SMA and ALMA
    visibilities.conj!(dset) # Swap UV convention
end

mvarr = Array{DataVis}(undef, nchan)
chi2s = Array{Float64}(undef, nchan)
lnprobs = Array{Float64}(undef, nchan)

for i = 1:nchan
    dv = dvarr[i]

    # FFT the appropriate image channel
    vis_fft = transform(skim, i)

    # Interpolate the `vis_fft` to the same locations as the DataSet
    mvis = ModelVis(dv, vis_fft)

    # Apply the phase correction here, since there are fewer data points
    phase_shift!(mvis, pars.mu_RA + dRA, pars.mu_DEC - dRA)

    dvis = visibilities.ModelVis2DataVis(mvis)

    mvarr[i] = dvis

    chi2s[i] = visibilities.chi2(dv, mvis)
    lnprobs[i] = visibilities.lnprob(dv, mvis)
end

# Now generate the residual visibilities
rvarr = ResidVis(dvarr, mvarr)

# Now swap the model and residual visibilities back to ALMA/SMA convetion
for i = 1:nchan
    visibilities.conj!(mvarr[i])
    visibilities.conj!(rvarr[i])
end

N = nchan * 2 * length(dvarr[1].VV)

# Only use the unmasked channels in the chi2 calculation
# if haskey(config, "exclude")
#     exclude = config["exclude"]
#     # which channels of the dset to fit
#     keylist = filter(x->(!in(x, exclude)), Int[i for i=1:nchan])
# else
#     keylist = Int[i for i=1:nchan]
# end

# The data are stored in increasing frequency, so
# exclude: [1] means exclude the most redshifted channel
# whereas
# exclude : [nchan] excludes the most blueshifted channel
if haskey(config, "exclude")
    exclude = config["exclude"]
    # which channels of the dset to fit
    # keylist = filter(x->(!in(x, exclude)), Int[i for i=1:nchan])

    lam0 = lam0s[config["species"] * config["transition"]]
    # calculate the velocities corresponding to dvarr
    lams = Float64[dv.lam for dv in dvarr]
    vels = c_kms * (lams .- lam0) / lam0
    # get the mask
    vel_mask = generate_vel_mask(exclude, vels)
else
    # keylist = Int[i for i=1:nchan]
    vel_mask = trues(nchan)
end

println("Note: may include flagged visibilities! Use DJ_calc_lnprob_resid.jl to calculate lnprob with correct flags.")
chi2s = chi2s[vel_mask]
# println("Chi^2 :", sum(chi2s))
println("Reduced Chi^2 ", sum(chi2s) / N)

# Calculate the lnprob between these two
lnprobs = lnprobs[vel_mask]

println("lnprior ", ln_prior)
println("lnlikelihood ", sum(lnprobs))
println("lnprob ", ln_prior + sum(lnprobs))

# actually write the datasets (including flagged points) back to the HDF5 files
visibilities.write(mvarr, parsed_args["out-model"])
visibilities.write(rvarr, parsed_args["out-resid"])

# copy visibility flags from the original dataset to this new one
visibilities.copy_flags(config["data_file"], parsed_args["out-model"])
visibilities.copy_flags(config["data_file"], parsed_args["out-resid"])


text:
function print_header(verbose::Int)
    if verbose > 0
        println("==============================================================")
        println("                  StateSpaceModels.jl v0.3.0                  ")
        println(" (c) Raphael Saavedra, Guilherme Bodin, and Mario Souto, 2019 ")
        println("--------------------------------------------------------------")
        println("            Starting state-space model estimation.            ")

    end
end

function print_estimation_start(verbose::Int, nseeds::Int)
    if verbose > 0
        println("    Initiating maximum likelihood estimation with $(nseeds) seeds.    ")
        println("--------------------------------------------------------------")
        println("||    seed    |     log-likelihood      |      time (s)     ||")
    end
end

function print_loglikelihood(verbose::Int, iseed::Int, loglikelihood::Vector{T}, t0::DateTime) where T
    if verbose > 0
        t1 = now() - t0

        str_seed   = @sprintf("%d", iseed) * "    |"
        str_loglik = @sprintf("%.4f", loglikelihood[iseed]) * "      |"
        str_time   = @sprintf("%.2f", t1.value/1000) * "     ||"

        str = "||"
        str *= " "^max(0, 13 - length(str_seed)) * str_seed
        str *= " "^max(0, 26 - length(str_loglik)) * str_loglik
        str *= " "^max(0, 21 - length(str_time)) * str_time

        println(str)
    end
end

function print_estimation_end(verbose::Int, log_lik::T, aic::T, bic::T) where T
    if verbose > 0
        str_1 = "                  Log-likelihood: "
        str_log_lik = @sprintf("%.4f", log_lik) * "                  "
        str_1 *= " "^max(0, 62 - length(str_log_lik) - length(str_1)) * str_log_lik

        str_2 = "                             AIC: "
        str_aic = @sprintf("%.4f", aic) * "                  "
        str_2 *= " "^max(0, 62 - length(str_aic) - length(str_2)) * str_aic

        str_3 = "                             BIC: "
        str_bic = @sprintf("%.4f", bic) * "                  "
        str_3 *= " "^max(0, 62 - length(str_bic) - length(str_3)) * str_bic

        println("--------------------------------------------------------------")
        println("           Maximum likelihood estimation complete.            ")
        println(str_1)
        println(str_2)
        println(str_3)
    end
end

function print_bottom(verbose::Int)
    if verbose > 0
        println("             End of state-space model estimation.             ")
        println("==============================================================")
    end
end


text:
# Sampler and model implementations

struct Nested{B, P <: AbstractProposal} <: AbstractSampler
    ndims::Int
    nactive::Int
    bounds::B
    enlarge::Float64
    update_interval::Int
    min_ncall::Int
    min_eff::Float64
    proposal::P
    dlv::Float64
end

"""
    Nested(ndims, nactive;
        bounds=Bounds.MultiEllipsoid,
        proposal=:auto,
        enlarge=1.25,
        update_interval=default_update_interval(proposal, ndims),
        min_ncall=2nactive,
        min_eff=0.10)

Static nested sampler with `nactive` active points and `ndims` parameters.

`ndims` is equivalent to the number of parameters to fit, which defines the dimensionality of the prior volume used in evidence sampling. `nactive` is the number of live or active points in the prior volume. This is a static sampler, so the number of live points will be constant for all of the sampling.

## Bounds and Proposals

`bounds` declares the Type of [`Bounds.AbstractBoundingSpace`](@ref) to use in the prior volume. The available bounds are described by [`Bounds`](@ref). `proposal` declares the algorithm used for proposing new points. The available proposals are described in [`Proposals`](@ref). If `proposal` is `:auto`, will choose the proposal based on `ndims`
* `ndims < 10` - [`Proposals.Rejection`](@ref)
* `10 ≤ ndims ≤ 20` - [`Proposals.RWalk`](@ref)
* `ndims > 20` - [`Proposals.Slice`](@ref)

The original nested sampling algorithm is roughly equivalent to using `Bounds.Ellipsoid` with `Proposals.Rejection`. The MultiNest algorithm is roughly equivalent to `Bounds.MultiEllipsoid` with `Proposals.Rejection`. The PolyChord algorithm is roughly equivalent to using `Proposals.RSlice`.

## Other Parameters
* `enlarge` - When fitting the bounds to live points, they will be enlarged (in terms of volume) by this linear factor.
* `update_interval` - How often to refit the live points with the bounds as a fraction of `nactive`. By default this will be determined using `default_update_interval` for the given proposal
    * `Proposals.Rejection` - `1.5`
    * `Proposals.RWalk` and `Proposals.RStagger` - `0.15 * walks`
    * `Proposals.Slice` - `0.9 * ndims * slices`
    * `Proposals.RSlice` - `2 * slices`
* `min_ncall`: The minimum number of iterations before fitting the first bound; used to
avoid shrinking the bounds before burn-in is completed. By default 2*`nactive`.
* `min_eff`: Minimum efficiency `(samples accepted / samples generated)` before fitting the
first bound; used to avoid shrinking the bounds before burn-in is completed By default 0.1.
"""
function Nested(ndims,
    nactive;
    bounds = Bounds.MultiEllipsoid,
    proposal = :auto,
    enlarge = 1.25,
    min_ncall=2nactive,
    min_eff=0.10,
    kwargs...)

    nactive < 2ndims && @warn "Using fewer than 2ndim ($(2ndims)) active points is discouraged"

    # get proposal
    if proposal === :auto
        proposal = if ndims < 10
            Proposals.Rejection()
        elseif 10 ≤ ndims ≤ 20
            Proposals.RWalk()
        else
            Proposals.Slice()
        end
    end

    dlv = log(nactive + 1) - log(nactive)

    update_interval_frac = get(kwargs, :update_interval, default_update_interval(proposal, ndims))
    update_interval = round(Int, update_interval_frac * nactive)
    return Nested(ndims,
        nactive,
        bounds,
        enlarge,
        update_interval,
        min_ncall,
        min_eff,
        proposal,
        dlv)
end

default_update_interval(p::Proposals.Rejection, ndims) = 1.5
default_update_interval(p::Proposals.RWalk, ndims) = 0.15 * p.walks
default_update_interval(p::Proposals.RStagger, ndims) = 0.15 * p.walks
default_update_interval(p::Proposals.Slice, ndims) = 0.9 * ndims * p.slices
default_update_interval(p::Proposals.RSlice, ndims) = 2.0 * p.slices


text:
### CACHE

cacheptr!(c, x::Ptr) = x
cacheptr!(c, x::String) = (push!(c, x); pointer(x))
cacheptr!(c, x::AbstractString) = cacheptr!(c, String(x))
cacheptr!(c, x::Array) = (push!(c, x); pointer(x))
cacheptr!(c, x::AbstractArray) = cacheptr!(c, Array(x))
cacheptr!(c, x::PyObject) = (push!(c, x); pyptr(x))
cacheptr!(c, x::Base.CFunction) = (push!(c, x); Base.unsafe_convert(Ptr{Cvoid}, x))

cachestrptr!(c, x::Ptr) = x
cachestrptr!(c, x) = cacheptr!(c, string(x))

macro cachefuncptr!(c, f, R, Ts)
    quote
        let c = $(esc(c)), f = $(esc(f))
            if f isa Ptr
                f
            elseif f isa Base.CFunction
                cacheptr!(c, f)
            else
                cacheptr!(
                    c,
                    @cfunction($(Expr(:$, :f)), $(esc(R)), ($(map(esc, Ts.args)...),))
                )
            end
        end
    end
end

### PROTOCOLS

PyNumberMethods_Create(c, x::PyNumberMethods) = x
PyNumberMethods_Create(
    c;
    add = C_NULL,
    subtract = C_NULL,
    multiply = C_NULL,
    remainder = C_NULL,
    divmod = C_NULL,
    power = C_NULL,
    negative = C_NULL,
    positive = C_NULL,
    absolute = C_NULL,
    bool = C_NULL,
    invert = C_NULL,
    lshift = C_NULL,
    rshift = C_NULL,
    and = C_NULL,
    or = C_NULL,
    xor = C_NULL,
    int = C_NULL,
    float = C_NULL,
    inplace_add = C_NULL,
    inplace_subtract = C_NULL,
    inplace_multiply = C_NULL,
    inplace_remainder = C_NULL,
    inplace_power = C_NULL,
    inplace_lshift = C_NULL,
    inplace_rshift = C_NULL,
    inplace_and = C_NULL,
    inplace_xor = C_NULL,
    inplace_or = C_NULL,
    floordivide = C_NULL,
    truedivide = C_NULL,
    inplace_floordivide = C_NULL,
    inplace_truedivide = C_NULL,
    index = C_NULL,
    matrixmultiply = C_NULL,
    inplace_matrixmultiply = C_NULL,
) = PyNumberMethods(
    add = @cachefuncptr!(c, add, PyPtr, (PyPtr, PyPtr)),
    subtract = @cachefuncptr!(c, subtract, PyPtr, (PyPtr, PyPtr)),
    multiply = @cachefuncptr!(c, multiply, PyPtr, (PyPtr, PyPtr)),
    remainder = @cachefuncptr!(c, remainder, PyPtr, (PyPtr, PyPtr)),
    divmod = @cachefuncptr!(c, divmod, PyPtr, (PyPtr, PyPtr)),
    power = @cachefuncptr!(c, power, PyPtr, (PyPtr, PyPtr, PyPtr)),
    negative = @cachefuncptr!(c, negative, PyPtr, (PyPtr,)),
    positive = @cachefuncptr!(c, positive, PyPtr, (PyPtr,)),
    absolute = @cachefuncptr!(c, absolute, PyPtr, (PyPtr,)),
    bool = @cachefuncptr!(c, bool, Cint, (PyPtr,)),
    invert = @cachefuncptr!(c, invert, PyPtr, (PyPtr,)),
    lshift = @cachefuncptr!(c, lshift, PyPtr, (PyPtr, PyPtr)),
    rshift = @cachefuncptr!(c, rshift, PyPtr, (PyPtr, PyPtr)),
    and = @cachefuncptr!(c, and, PyPtr, (PyPtr, PyPtr)),
    xor = @cachefuncptr!(c, xor, PyPtr, (PyPtr, PyPtr)),
    or = @cachefuncptr!(c, or, PyPtr, (PyPtr, PyPtr)),
    int = @cachefuncptr!(c, int, PyPtr, (PyPtr,)),
    float = @cachefuncptr!(c, float, PyPtr, (PyPtr,)),
    inplace_add = @cachefuncptr!(c, inplace_add, PyPtr, (PyPtr, PyPtr)),
    inplace_subtract = @cachefuncptr!(c, inplace_subtract, PyPtr, (PyPtr, PyPtr)),
    inplace_multiply = @cachefuncptr!(c, inplace_multiply, PyPtr, (PyPtr, PyPtr)),
    inplace_remainder = @cachefuncptr!(c, inplace_remainder, PyPtr, (PyPtr, PyPtr)),
    inplace_power = @cachefuncptr!(c, inplace_power, PyPtr, (PyPtr, PyPtr, PyPtr)),
    inplace_lshift = @cachefuncptr!(c, inplace_lshift, PyPtr, (PyPtr, PyPtr)),
    inplace_rshift = @cachefuncptr!(c, inplace_rshift, PyPtr, (PyPtr, PyPtr)),
    inplace_and = @cachefuncptr!(c, inplace_and, PyPtr, (PyPtr, PyPtr)),
    inplace_xor = @cachefuncptr!(c, inplace_xor, PyPtr, (PyPtr, PyPtr)),
    inplace_or = @cachefuncptr!(c, inplace_or, PyPtr, (PyPtr, PyPtr)),
    floordivide = @cachefuncptr!(c, floordivide, PyPtr, (PyPtr, PyPtr)),
    truedivide = @cachefuncptr!(c, truedivide, PyPtr, (PyPtr, PyPtr)),
    inplace_floordivide = @cachefuncptr!(c, inplace_floordivide, PyPtr, (PyPtr, PyPtr)),
    inplace_truedivide = @cachefuncptr!(c, inplace_truedivide, PyPtr, (PyPtr, PyPtr)),
    index = @cachefuncptr!(c, index, PyPtr, (PyPtr,)),
    matrixmultiply = @cachefuncptr!(c, matrixmultiply, PyPtr, (PyPtr, PyPtr)),
    inplace_matrixmultiply = @cachefuncptr!(
        c,
        inplace_matrixmultiply,
        PyPtr,
        (PyPtr, PyPtr)
    ),
)
PyNumberMethods_Create(c, x::Dict) = PyNumberMethods_Create(c; x...)
PyNumberMethods_Create(c, x::NamedTuple) = PyNumberMethods_Create(c; x...)

PyMappingMethods_Create(c, x::PyMappingMethods) = x
PyMappingMethods_Create(c; length = C_NULL, subscript = C_NULL, ass_subscript = C_NULL) =
    PyMappingMethods(
        length = @cachefuncptr!(c, length, Py_ssize_t, (PyPtr,)),
        subscript = @cachefuncptr!(c, subscript, PyPtr, (PyPtr, PyPtr)),
        ass_subscript = @cachefuncptr!(c, ass_subscript, Cint, (PyPtr, PyPtr, PyPtr)),
    )
PyMappingMethods_Create(c, x::Dict) = PyMappingMethods_Create(c; x...)
PyMappingMethods_Create(c, x::NamedTuple) = PyMappingMethods_Create(c; x...)

PySequenceMethods_Create(c, x::PySequenceMethods) = x
PySequenceMethods_Create(
    c;
    length = C_NULL,
    concat = C_NULL,
    repeat = C_NULL,
    item = C_NULL,
    ass_item = C_NULL,
    contains = C_NULL,
    inplace_concat = C_NULL,
    inplace_repeat = C_NULL,
) = PySequenceMethods(
    length = @cachefuncptr!(c, length, Py_ssize_t, (PyPtr,)),
    concat = @cachefuncptr!(c, concat, PyPtr, (PyPtr, PyPtr)),
    repeat = @cachefuncptr!(c, repeat, PyPtr, (PyPtr, Py_ssize_t)),
    item = @cachefuncptr!(c, item, PyPtr, (PyPtr, Py_ssize_t)),
    ass_item = @cachefuncptr!(c, ass_item, Cint, (PyPtr, Py_ssize_t, PyPtr)),
    contains = @cachefuncptr!(c, contains, Cint, (PyPtr, PyPtr)),
    inplace_concat = @cachefuncptr!(c, inplace_concat, PyPtr, (PyPtr, PyPtr)),
    inplace_repeat = @cachefuncptr!(c, inplace_repeat, PyPtr, (PyPtr, Py_ssize_t)),
)
PySequenceMethods_Create(c, x::Dict) = PySequenceMethods_Create(c; x...)
PySequenceMethods_Create(c, x::NamedTuple) = PySequenceMethods_Create(c; x...)

PyBufferProcs_Create(c, x::PyBufferProcs) = x
PyBufferProcs_Create(c; get = C_NULL, release = C_NULL) = PyBufferProcs(
    get = @cachefuncptr!(c, get, Cint, (PyPtr, Ptr{Py_buffer}, Cint)),
    release = @cachefuncptr!(c, release, Cvoid, (PyPtr, Ptr{Py_buffer})),
)
PyBufferProcs_Create(c, x::Dict) = PyBufferProcs_Create(c; x...)
PyBufferProcs_Create(c, x::NamedTuple) = PyBufferProcs_Create(c; x...)

PyMethodDef_Create(c, x::PyMethodDef) = x
PyMethodDef_Create(c; name = C_NULL, meth = C_NULL, flags = 0, doc = C_NULL) = PyMethodDef(
    name = cachestrptr!(c, name),
    meth = iszero(flags & Py_METH_KEYWORDS) ?
           @cachefuncptr!(c, meth, PyPtr, (PyPtr, PyPtr)) :
           @cachefuncptr!(c, meth, PyPtr, (PyPtr, PyPtr, PyPtr)),
    flags = flags,
    doc = cachestrptr!(c, doc),
)
PyMethodDef_Create(c, x::Dict) = PyMethodDef_Create(c; x...)
PyMethodDef_Create(c, x::NamedTuple) = PyMethodDef_Create(c; x...)

PyGetSetDef_Create(c, x::PyGetSetDef) = x
PyGetSetDef_Create(
    c;
    name = C_NULL,
    get = C_NULL,
    set = C_NULL,
    doc = C_NULL,
    closure = C_NULL,
) = PyGetSetDef(
    name = cachestrptr!(c, name),
    get = @cachefuncptr!(c, get, PyPtr, (PyPtr, Ptr{Cvoid})),
    set = @cachefuncptr!(c, set, Cint, (PyPtr, PyPtr, Ptr{Cvoid})),
    doc = cachestrptr!(c, doc),
    closure = closure,
)
PyGetSetDef_Create(c, x::Dict) = PyGetSetDef_Create(c; x...)
PyGetSetDef_Create(c, x::NamedTuple) = PyGetSetDef_Create(c; x...)

PyType_Create(
    c;
    type = C_NULL,
    name,
    as_number = C_NULL,
    as_mapping = C_NULL,
    as_sequence = C_NULL,
    as_buffer = C_NULL,
    methods = C_NULL,
    getset = C_NULL,
    dealloc = C_NULL,
    getattr = C_NULL,
    setattr = C_NULL,
    repr = C_NULL,
    hash = C_NULL,
    call = C_NULL,
    str = C_NULL,
    getattro = C_NULL,
    setattro = C_NULL,
    doc = C_NULL,
    iter = C_NULL,
    iternext = C_NULL,
    richcompare = C_NULL,
    opts...,
) = begin
    type = cacheptr!(c, type)
    name = cachestrptr!(c, name)
    as_number =
        as_number isa Ptr ? as_number :
        cacheptr!(c, fill(PyNumberMethods_Create(c, as_number)))
    as_mapping =
        as_mapping isa Ptr ? as_mapping :
        cacheptr!(c, fill(PyMappingMethods_Create(c, as_mapping)))
    as_sequence =
        as_sequence isa Ptr ? as_sequence :
        cacheptr!(c, fill(PySequenceMethods_Create(c, as_sequence)))
    as_buffer =
        as_buffer isa Ptr ? as_buffer :
        cacheptr!(c, fill(PyBufferProcs_Create(c, as_buffer)))
    methods = if methods isa Ptr
        methods
    else
        cacheptr!(c, [[PyMethodDef_Create(c, m) for m in methods]; PyMethodDef()])
    end
    getset = if getset isa Ptr
        getset
    else
        cacheptr!(c, [[PyGetSetDef_Create(c, m) for m in getset]; PyGetSetDef()])
    end
    dealloc = @cachefuncptr!(c, dealloc, Cvoid, (PyPtr,))
    getattr = @cachefuncptr!(c, getattr, PyPtr, (PyPtr, Cstring))
    setattr = @cachefuncptr!(c, setattr, Cint, (PyPtr, Cstring, PyPtr))
    repr = @cachefuncptr!(c, repr, PyPtr, (PyPtr,))
    hash = @cachefuncptr!(c, hash, Py_hash_t, (PyPtr,))
    call = @cachefuncptr!(c, call, PyPtr, (PyPtr, PyPtr, PyPtr))
    str = @cachefuncptr!(c, str, PyPtr, (PyPtr,))
    getattro = @cachefuncptr!(c, getattro, PyPtr, (PyPtr, PyPtr))
    setattro = @cachefuncptr!(c, setattro, Cint, (PyPtr, PyPtr, PyPtr))
    doc = cachestrptr!(c, doc)
    iter = @cachefuncptr!(c, iter, PyPtr, (PyPtr,))
    iternext = @cachefuncptr!(c, iternext, PyPtr, (PyPtr,))
    richcompare = @cachefuncptr!(c, richcompare, PyPtr, (PyPtr, PyPtr, Cint))
    PyTypeObject(;
        ob_base = PyVarObject(ob_base = PyObject(type = type)),
        name = name,
        as_number = as_number,
        as_mapping = as_mapping,
        as_sequence = as_sequence,
        as_buffer = as_buffer,
        methods = methods,
        getset = getset,
        dealloc = dealloc,
        getattr = getattr,
        setattr = setattr,
        repr = repr,
        hash = hash,
        call = call,
        str = str,
        getattro = getattro,
        setattro = setattro,
        iter = iter,
        iternext = iternext,
        richcompare = richcompare,
        opts...,
    )
end


text:
@doc raw"""
    LaplaceLikelihood(β::T=1.0)  #  Laplace likelihood with scale β

Laplace likelihood for regression:
```math
\frac{1}{2\beta} \exp(-\frac{|y-f|}{β})
```
see [wiki page](https://en.wikipedia.org/wiki/Laplace_distribution)
---
For the analytical solution, it is augmented via:
```math
p(y|f,ω) = N(y|f,ω⁻¹)
```
where ``ω \sim \text{Exp}(ω | 1/(2 β^2))``, and ``\text{Exp}`` is the [Exponential distribution](https://en.wikipedia.org/wiki/Exponential_distribution)
We use the variational distribution ``q(ω) = GIG(ω|a,b,p)``
"""
struct LaplaceLikelihood{T<:Real} <: RegressionLikelihood
    β::T
    a::T
    p::T
    function LaplaceLikelihood{T}(β::T) where {T<:Real}
        return new{T}(β, β^-2, 0.5)
    end
end

function LaplaceLikelihood(β::T=1.0) where {T<:Real}
    return LaplaceLikelihood{T}(β)
end

function implemented(
    ::LaplaceLikelihood, ::Union{<:AnalyticVI,<:QuadratureVI,<:GibbsSampling}
)
    return true
end

function (l::LaplaceLikelihood)(y::Real, f::Real)
    return pdf(Laplace(f, l.β), y)
end

function Distributions.loglikelihood(l::LaplaceLikelihood, y::Real, f::Real)
    return logpdf(Laplace(f, l.β), y)
end

function Base.show(io::IO, l::LaplaceLikelihood)
    return print(io, "Laplace likelihood (β=$(l.β))")
end

function compute_proba(
    l::LaplaceLikelihood, μ::AbstractVector{<:Real}, σ²::AbstractVector{<:Real}
)
    return μ, max.(σ², zero(eltype(σ²))) .+ 2 * l.β^2
end

## Local Updates ##
# b : Variational parameter b of GIG
# θ : Expected value of ω
function init_local_vars(::LaplaceLikelihood, batchsize::Int, T::DataType=Float64)
    return (; b=rand(T, batchsize), θ=zeros(T, batchsize))
end

function local_updates!(
    local_vars,
    l::LaplaceLikelihood,
    y::AbstractVector,
    μ::AbstractVector,
    diagΣ::AbstractVector,
)
    @. local_vars.b = diagΣ + abs2(μ - y)
    @. local_vars.θ = sqrt(l.a) / sqrt.(local_vars.b)
    return local_vars
end

function sample_local!(
    local_vars, l::LaplaceLikelihood, y::AbstractVector, f::AbstractVector
)
    @. local_vars.b = rand(GeneralizedInverseGaussian(1 / l.β^2, abs2(f - y), 0.5))
    @. local_vars.θ = inv(local_vars.b)
    return local_vars
end

@inline function ∇E_μ(::LaplaceLikelihood, ::AOptimizer, y::AbstractVector, state)
    return (state.θ .* y,)
end
@inline function ∇E_Σ(::LaplaceLikelihood, ::AOptimizer, ::AbstractVector, state)
    return (0.5 * state.θ,)
end

## ELBO ##
function expec_loglikelihood(
    ::LaplaceLikelihood,
    ::AnalyticVI,
    y::AbstractVector,
    μ::AbstractVector,
    diag_cov::AbstractVector,
    state,
)
    tot = -0.5 * length(y) * log(twoπ)
    tot += 0.5 * Zygote.@ignore(sum(log, state.θ))
    tot +=
        -0.5 * (
            dot(state.θ, diag_cov) + dot(state.θ, abs2.(μ)) - 2.0 * dot(state.θ, μ .* y) +
            dot(state.θ, abs2.(y))
        )
    return tot
end

function AugmentedKL(l::LaplaceLikelihood, state, ::Any)
    return GIGEntropy(l, state) - expecExponentialGIG(l, state)
end

GIGEntropy(l::LaplaceLikelihood, state) = GIGEntropy(l.a, state.b, l.p)

function expecExponentialGIG(l::LaplaceLikelihood, state)
    return sum(
        -log(2 * l.β^2) .-
        0.5 * (l.a .* sqrt.(state.b) + state.b .* sqrt(l.a)) ./ (l.a .* state.b * l.β^2),
    )
end

## PDF and Log PDF Gradients ##

function grad_quad(
    likelihood::LaplaceLikelihood, y::Real, μ::Real, σ²::Real, inference::AbstractInference
)
    nodes = inference.nodes * sqrt(σ²) .+ μ
    Edloglike = dot(inference.weights, ∇loglikehood.(likelihood, y, nodes))
    Ed²loglike = (1 / sqrt(twoπ * σ²)) / (likelihood.β^2)
    return -Edloglike, Ed²loglike
end

@inline function ∇loglikehood(l::LaplaceLikelihood, y::Real, f::Real)
    return sign(y - f) ./ l.β
end

@inline hessloglikelihood(::LaplaceLikelihood, ::Real, f::T) where {T<:Real} = zero(T)


text:
"""
Computes the maximum flow between the source and target vertexes in a flow
graph using the [Edmonds-Karp algorithm](https://en.wikipedia.org/wiki/Edmondss%E2%80%93Karp_algorithm).
Returns the value of the maximum flow as well as the final flow matrix.
Use a default capacity of 1 when the capacity matrix isn\'t specified.
Requires arguments:
- residual_graph::ADiGraph                # the input graph
- source                            # the source vertex
- target                            # the target vertex
- capacity_matrix::AbstractMatrix{T}    # edge flow capacities
"""
function edmonds_karp_impl(
        residual_graph::ADiGraph,               # the input graph
        source,                           # the source vertex
        target,                           # the target vertex
        capacity_matrix::AbstractMatrix{T}    # edge flow capacities
    ) where T<:Number
    n = nv(residual_graph)                     # number of vertexes
    flow = 0
    flow_matrix = zeros(T, n, n)           # initialize flow matrix

    P = zeros(Int, n)
    S = zeros(Int, n)
    while true
        fill!(P, -1)
        fill!(S, -1)
        v, P, S, flag = fetch_path!(residual_graph, source, target, flow_matrix, capacity_matrix, P,S)

        if flag != 0                       # no more valid paths
            break
        else
            path = [v]                     # initialize path
            sizehint!(path, n)

            u = v
            while u!=source                # trace path from v to source
                u = P[u]
                push!(path, u)
            end
            reverse!(path)

            u = v                          # trace path from v to target
            while u!=target
                u = S[u]
                push!(path, u)
            end
                                           # augment flow along path
            flow += augment_path!(path, flow_matrix, capacity_matrix)
        end
    end

    return flow, flow_matrix
end

"""
Calculates the amount by which flow can be augmented in the given path.
Augments the flow and returns the augment value.
Requires arguments:
- path::Vector{Int}                      # input path
- flow_matrix::AbstractMatrix{T}        # the current flow matrix
- capacity_matrix::AbstractMatrix{T}    # edge flow capacities
"""
function augment_path!(
    path::Vector,                     # input path
    flow_matrix::AbstractMatrix{T},       # the current flow matrix
    capacity_matrix::AbstractMatrix{T}    # edge flow capacities
    ) where T<:Number
    augment = typemax(T)                   # initialize augment
    for i in 1:length(path)-1              # calculate min capacity along path
        u = path[i]
        v = path[i+1]
        augment = min(augment,capacity_matrix[u,v] - flow_matrix[u,v])
    end

    for i in 1:length(path)-1              # augment flow along path
        u = path[i]
        v = path[i+1]
        flow_matrix[u,v] += augment
        flow_matrix[v,u] -= augment
    end

    return augment
end

"""
Uses Bidirectional BFS to look for augmentable-paths. Returns the vertex where
the two BFS searches intersect, the Parent table of the path, the
Successor table of the path found, and a flag indicating success
Flag Values:
0 => success
1 => No Path to target
2 => No Path to source
"""
function fetch_path(
    residual_graph::ADiGraph,               # the input graph
    source,                           # the source vertex
    target,                           # the target vertex
    flow_matrix::AbstractMatrix{T},       # the current flow matrix
    capacity_matrix::AbstractMatrix{T}    # edge flow capacities
    ) where T<:Number
    n = nv(residual_graph)
    P = -1 * ones(Int, n)
    S = -1 * ones(Int, n)
    return fetch_path!(residual_graph,
                       source,
                       target,
                       flow_matrix,
                       capacity_matrix,
                       P,
                       S)
end

"""
A preallocated version of fetch_paths. The parent and successor tables are pre-allocated.
Uses Bidirectional BFS to look for augmentable-paths. Returns the vertex where
the two BFS searches intersect, the Parent table of the path, the
Successor table of the path found, and a flag indicating success
Flag Values:
0 => success
1 => No Path to target
2 => No Path to source
Requires arguments:
    residual_graph::ADiGraph                # the input graph
    source                            # the source vertex
    target                            # the target vertex
    flow_matrix::AbstractMatrix{T}        # the current flow matrix
    capacity_matrix::AbstractMatrix{T}    # edge flow capacities
    P::Vector{Int}                         # parent table of path init to -1s
    S::Vector{Int}                         # successor table of path init to -1s
"""
function fetch_path!(
    residual_graph::ADiGraph,               # the input graph
    source,                           # the source vertex
    target,                           # the target vertex
    flow_matrix::AbstractMatrix{T},       # the current flow matrix
    capacity_matrix::AbstractMatrix{T},   # edge flow capacities
    P::Vector{Int},                        # parent table of path init to -1s
    S::Vector{Int}                         # successor table of path init to -1s
    ) where T<:Number
    n = nv(residual_graph)

    P[source] = -2
    S[target] = -2

    Q_f = [source]                         # forward queue
    sizehint!(Q_f, n)

    Q_r = [target]                         # reverse queue
    sizehint!(Q_r, n)

    while true

        if length(Q_f) <= length(Q_r)
            u = pop!(Q_f)
            for v in out_neighbors(residual_graph, u)
                if capacity_matrix[u,v] - flow_matrix[u,v] > 0 && P[v] == -1
                    P[v] = u
                    if S[v] == -1
                        pushfirst!(Q_f, v)
                    else
                        return v, P, S, 0 # 0 indicates success
                    end
                end
            end

            length(Q_f) == 0 && return 0, P, S, 1 # No paths to target
        else
            v = pop!(Q_r)
            for u in in_neighbors(residual_graph, v)
                if capacity_matrix[u,v] - flow_matrix[u,v] > 0 && S[u] == -1
                    S[u] = v
                    P[u] != -1 && return  u, P, S, 0 # 0 indicates success

                    pushfirst!(Q_r, u)
                end

            end

            length(Q_r) == 0 && return 0, P, S, 2 # No paths to source
        end
    end
end


text:
using IonSim
using QuantumOptics
using DifferentialEquations
using Distributions
using DataStructures

num_ions = 1
b_field = 3.799953549
carrier_frequency = 2127405.1614560783

for duration in 0:5:50  #0:2e-6:50e-6

    #############################################
    # Create and load the ions
    S = ["S-1/2", "S+1/2"]
    D = ["D-5/2", "D-3/2", "D-1/2" ,"D+1/2", "D+3/2", "D+5/2"]
    ions = Array{Ca40}(undef, num_ions)
    for i = 1:num_ions
        ions[i] = Ca40([S; D])
    end

    axial_frequency = 900e3
    radial_frequency_1 = 1.3e6
    radial_frequency_2 = 1.6e6
    chain = LinearChain(
        ions=ions,
        com_frequencies=(x=radial_frequency_1, y=radial_frequency_2, z=axial_frequency),
        vibrational_modes=(;z=[1]))

    #############################################
    # Set up the lasers and the trap
    laser = Laser(
        k=(x̂ + ẑ)/√2,
        ϵ=(x̂ - ẑ)/√2,
        Δ=carrier_frequency + axial_frequency,
        pointing=[(1, 1.)],
    )

    trap = Trap(configuration=chain, B=b_field*1e-4, Bhat=ẑ, δB=0, lasers=[laser])
    mode = trap.configuration.vibrational_modes.z[1]

    #laser.Δ = transition_frequency(trap, 1, ("S-1/2", "D-1/2")) + axial_frequency
    global_beam!(trap, laser)

    simulation_start_time = 0
    simulation_stop_time = duration
    simulation_tstops = [simulation_start_time, simulation_stop_time]
    simulation_tspan = range(simulation_start_time - 1e-3, simulation_stop_time + 1e-3, length=2)

    #############################################
    # Set up the time-dependent E-field
    intensity_factor = 10^(-0.5) # this corresponds to pi time 3 μs
    pi_min = 3e-6
    amp = 1.0
    att = 6.0

    function step_interval(t, t_begin, t_end)
        t >= t_begin && t < t_end ? 1 : 0
    end

    t_pi = pi_min * sqrt(intensity_factor / (amp * 10^(-1.5 * att / 10)))
    E = Efield_from_pi_time(t_pi, trap.Bhat, laser, ions[1], ("S-1/2", "D-1/2"))
    laser.E = t -> E * step_interval(t, simulation_start_time, simulation_stop_time)

    #############################################
    # Run the simulation
    initial_state = ionstate(trap, "S-1/2")
    h = hamiltonian(trap, lamb_dicke_order=1, timescale=1e-6, rwa_cutoff=1e5)
    @time tout, solution = timeevolution.schroedinger_dynamic(
        simulation_tspan,
        tensor(initial_state, fockstate(mode, 0)),
        h,
        callback=PresetTimeCallback(simulation_tstops, integrator -> return)
        );
    solution = solution[end]

    #############################################
    # Print the result
    function project(solution, states...)
        real.(expect(ionprojector(trap, states...), solution))
    end

    result = Dict(
        "S" => sum([project(solution, state) for state=S]),
        "D" => sum([project(solution, state) for state=D]),
        )

    println("Results: $result")

end

max_stars_repo_name: rmshaffer/pulse-sequence-simulator

max_issues_repo_name: rmshaffer/pulse-sequence-simulator

max_forks_repo_name: rmshaffer/pulse-sequence-simulator


text:
#
# This small part introduces both debug and record for C(k)
# including this file, you can use
# DebugCk() within debug = and RecordCk() within record =
#
# you have to have x_hat defined.
function tilde_x_old(p, o, x_old, ξbar_old)
    return exp(
        p.M,
        x_old,
        vector_transport_to(
            p.M,
            o.m,
            -o.primal_stepsize * adjoint_linearized_operator(p, o.m, o.n, ξbar_old),
            x_old,
        ),
    )
end
function ζk(p, o, x_old, ξbar_old)
    return vector_transport_to(
        p.M,
        x_old,
        log(p.M, x_old, o.x) - vector_transport_to(
            p.M,
            tilde_x_old(p, o, x_old, ξbar_old),
            log(p.M, tilde_x_old(p, o, x_old, ξbar_old), x_hat),
            x_old,
        ),
        o.m,
    ) - log(p.M, o.m, o.x) + log(p.M, o.m, x_hat)
end
function Ck(p, o, x_old, ξ_bar_old)
    return 1 / o.primal_stepsize *
           distance(p.M, x_old, tilde_x_old(p, o, x_old, ξ_bar_old))^2 + inner(
        p.N,
        o.n,
        ξ_bar_old,
        linearized_forward_operator(p, o.m, ζk(p, o, x_old, ξ_bar_old), o.n),
    )
end

struct DebugCk <: DebugAction
    io::IO
    prefix::String
    storage::StoreOptionsAction
    function DebugCk(a::StoreOptionsAction=StoreOptionsAction((:x, :ξbar)), io::IO=stdout)
        return new(io, "C(k): ", a)
    end
end
function (d::DebugCk)(p::P, o::ChambollePockOptions, i::Int) where {P<:PrimalDualProblem}
    if all(has_storage.(Ref(d.storage), [:x, :ξbar])) && i > 0 # all values stored
        x_old, ξ_bar_old = get_storage.(Ref(d.storage), [:x, :ξbar]) #fetch
        print(d.io, d.prefix * "$(Ck(p, o, x_old,ξ_bar_old))")
    end
    return d.storage(p, o, i)
end

struct RecordCk <: RecordAction
    recorded_values::Array{Float64,1}
    storage::StoreOptionsAction
    function RecordCk(a::StoreOptionsAction=StoreOptionsAction((:x, :ξbar)))
        return new(Array{Float64,1}(), a)
    end
end
function (r::RecordCk)(p::P, o::ChambollePockOptions, i::Int) where {P<:PrimalDualProblem}
    if all(has_storage.(Ref(r.storage), [:x, :ξbar])) && i > 0 # all values stored
        x_old = get_storage(r.storage, :x)
        ξ_bar_old = get_storage(r.storage, :ξbar)
        Manopt.record_or_reset!(r, Ck(p, o, x_old, ξ_bar_old), i)
    end
    return r.storage(p, o, i)
end


text:
#GenSPGL

export NormL1_project

#DEVNOTE# Backtrace weights so only oneprojectormex is type tight
"""
Use: x,itn = NormL1_project(x,weights,tau)

"""
function NormL1_project(x::AbstractVector{<:Real}, tau::Number, weights, params)

    x_out,itn = oneprojector(x, weights, tau)

    return x_out,itn
end

"""
For Complex
"""
function NormL1_project(x::AbstractVector{<:Complex}, tau::Number, weights, params)

    xa = abs.(x)
    idx = findall(xa .< eps())
    xc,itn = oneprojector(xa, weights, tau)
    xc = xc ./ xa
    xc[idx] .= 0
    x_out = x.*xc

    return x_out,itn
end


text:
"""
    imshow([stream], img, [depth::TermColorDepth], [maxsize])

Displays the given image `img` using unicode characters and
terminal colors (defaults to 256 colors).
`img` has to be an array of `Colorant`.

If working in the REPL, the function tries to choose the encoding
based on the current display size. The image will also be
downsampled to fit into the display (using `restrict`).
"""
function imshow(
        io::IO,
        img::AbstractArray{<:Colorant},
        colordepth::TermColorDepth,
        maxsize::Tuple = displaysize(io))
    print_matrix(io, x) = imshow(io, x, colordepth, maxsize)
    Base.show_nd(io, img, print_matrix, true)
end

function imshow(
        io::IO,
        img::AbstractMatrix{<:Colorant},
        colordepth::TermColorDepth,
        maxsize::Tuple = displaysize(io))
    io_h, io_w = maxsize
    img_h, img_w = map(length, axes(img))
    str = if img_h <= io_h-4 && 2img_w <= io_w
        first(encodeimg(BigBlocks(),   colordepth, img, io_h-4, io_w))
    else
        first(encodeimg(SmallBlocks(), colordepth, img, io_h-4, io_w))
    end
    for (idx, line) in enumerate(str)
        print(io, line)
        idx < length(str) && println(io)
    end
end

# colorant vector
function imshow(
        io::IO,
        img::AbstractVector{<:Colorant},
        colordepth::TermColorDepth,
        maxsize::Tuple = displaysize(io))
    io_h, io_w = maxsize
    img_w = length(img)
    str = if 3img_w <= io_w
        first(encodeimg(BigBlocks(),   colordepth, img, io_w))
    else
        first(encodeimg(SmallBlocks(), colordepth, img, io_w))
    end
    for (idx, line) in enumerate(str)
        print(io, line)
        idx < length(str) && println(io)
    end
end

imshow(io::IO, img, args...) = imshow(io, img, colormode[1], args...)
imshow(img, args...) = imshow(stdout, img, colormode[1], args...)
imshow(io::IO, img, colordepth::TermColorDepth, args...) = throw(ArgumentError("imshow only supports colorant arrays with 1 or 2 dimensions"))

"""
    imshow256([stream], img, [maxsize])

Displays the given image `img` using unicode characters and
the widely supported 256 terminal colors.
`img` has to be an array of `Colorant`.

If working in the REPL, the function tries to choose the encoding
based on the current display size. The image will also be
downsampled to fit into the display (using `restrict`).
"""
imshow256(io::IO, img, args...) = imshow(io, img, TermColor256(), args...)
imshow256(img, args...) = imshow256(stdout, img, args...)

"""
    imshow24bit([stream], img, [maxsize])

Displays the given image `img` using unicode characters and
the 24 terminal colors that some modern terminals support.
`img` has to be an array of `Colorant`.

If working in the REPL, the function tries to choose the encoding
based on the current display size. The image will also be
downsampled to fit into the display (using `restrict`).
"""
imshow24bit(io::IO, img, args...) = imshow(io, img, TermColor24bit(), args...)
imshow24bit(img, args...) = imshow24bit(stdout, img, args...)


text:
using JuliaWeBWorK

using SymPy # useful for formula
using Plots # for plotq

meta=(AuthorText="JuliaWeBWorK",
      Keywords="Kitchen sink, everything but",
      Question="1"
      )

qs = JuliaWeBWorK.QUESTIONS()

# text formatting example
intro = raw"""

[WeBWorK](https://webwork.maa.org/)

----

A *silly* example of the different questions that can readily be asked.

"""

# Randomized questions

## numeric

###  No randomization
numericq("Enter 1", ()->1, ()) |> qs
numericq("Enter 1, 2", () -> List((1,2)), ()) |> qs
numericq(raw"Enter \([\sin(x)]'\)", () ->  (@vars  x; Formula(diff(sin(x),x))), ()) |> qs
numericq(raw"Enter the interval \(  (1,2) \)", () -> Interval(1,2),  ()) |> qs
numericq(raw"Enter the intervals \(  (1,2) \) and  \( (3,4)  \)", () -> List(Interval(1,2), Interval(3,4)),  ()) |> qs

p = plot(x->x^5 -x -1, -1, 1.5, legend=false)
plot!(p, zero)
numericq("How many  zeros in the plot: ![plot]($(Plot(p)))", () -> 1,()) |> qs

## randomization per problem
numericq("Enter `2*{{:a1}}`", (a) -> 2a, (1:3, )) |> qs
numericq("Enter  `{{:a1}}`, `{{:a2}}`", (a,b) ->  List(a,b),  (1:3, 1:3)) |> qs
numericq(raw"Enter \([\sin({{:a1}}x)]'\)", (a) ->  (@vars  x; Formula(diff(sin(a*x),x))), (2:3,)) |> qs
numericq(raw"Enter  the interval \( ({{:a1}}, {{:a2}}) \)", (a,b) ->  Interval(a,b),  (1:3, 4:6)) |> qs
numericq(raw"Enter  the intervals \( ({{:a1}}, {{:a2}}) \) and \( ({{:a3}}, {{:a4}}) \)",
               (a,b,c,d) ->  List(Interval(a,b),  Interval(c,d)),  (1:3, 4:6, 1:2, 3:4)) |> qs
stringq("Spell out {{:a1}}", (a) -> ("one", "two", "three")[a], (1:3, )) |> qs

## Shared randomization
r = randomizer(2:3)
numericq("Enter `2*{{:a1}}`", (a) -> 2a,r) |> qs
numericq(raw"Enter \([\sin({{:a1}}x)]'\)", (a) ->  (@vars  x; Formula(diff(sin(a*x),x))), r) |> qs
function n36_fn(a)
    p = plot(x->x^5 - x - (1+1/a), -1, 1.5)
    plot!(p,zero)
    Plot(p)
end
plotq(raw"Plot of \(f \)", n36_fn, r) |> qs
numericq(raw"How many zeros does the graph ofa \( f\) show?", (a) -> 1, r) |> qs
stringq("Spell out {{:a1}}", (a) -> ("one", "two", "three")[a], r) |> qs

# choice questions
radioq("Select 1", (("one",  jmt"two", "three"), "not available"), 1) |> qs
multiplechoiceq("Select all  bigger than 2", ((1,2,5,6),(7,8, jmt"0")), (3,4,5,6)) |> qs

# output only
stringq("Who is your father?", (a) -> "Darth Vader", (1:1,)) |> qs  # string q needs some randomization
hint("Use  the force  Luke") |> qs

essay = essayq("""
How  was this exercise?
""") |> qs

# create  Page object
page = Page(intro, qs; meta...)

# create pg file
# open("filename",  "w")  do  io
#     print(io, page)
# end


text:
function generateInitialPopPlot( report::DataFrame, showPlot::Bool,
    savePlot::Bool, filename::String )

    plotTitle = "Initial population composition plot"
    initSize = sum( report[:, 2] )
    hoverValues = string.( report[:, 1], ": ", report[:, 2], " (",
        round.( 100.0 * report[:, 2] / initSize, digits=2 ), "%)" )
    plt = pie( report[:, 1], report[:, 2], show=false, size=(960,540),
        title=plotTitle, hover=hoverValues )

    # Show plot if needed.
    if showPlot
        gui( plt )
    end  # if showPlot

    # Save plot if needed.
    if savePlot
        # If the user doesn't specify a filename, a random one is generated.
        if filename == ""
            dateStr = Dates.format( now(), "yyyymmdd HHMMSS" )
            filename = string( "Initial population composition plot  ",
                dateStr )
        end  # if filename == ""

        if !ispath( dirname( filename ) )
            mkpath( dirname( filename ) )
        end  # if !ispath( dirname( filename ) )

        extension = split( filename, "." )
        extension = length( extension ) == 1 ? "" : extension[end]
        filename = string( filename, extension ∈ extensions ? "" : ".svg" )
        savefig( plt, filename )
    end  # if savePlot

end  # generateInitialPopPlot( report, showPlot, savePlot, filename )

function generateInitialPopAgePlot( report::DataFrame, timeFactor::Float64, showPlot::Bool, savePlot::Bool, filename::String )

    # Generate the plot.
    plotTitle = string( "Initial population ", names( report )[1], " plot" )
    ages = report[1:(end-5), 1] / timeFactor

    if length( ages ) > 1
        ages .+= ( ages[2] - ages[1] ) / 2
    end  # if length( ages ) > 1

    plotData = report[1:(end-5), 2]
    ymax = max( maximum( plotData ), 20 )

    plt = bar( ages, plotData, show=false, size=(960, 540), title=plotTitle,
        bar_width=length( ages ) == 1 ? .5 : ages[2] - ages[1], labels=nothing,
        ylim=[0, ymax], fa=0.5 )
    plt = vline!( [report[(end-4), 2] / timeFactor], labels="mean", lw=2,
        ls=:dash )
    plt = vline!( [report[(end-3), 2] / timeFactor], labels="median", lw=2,
        ls=:dash )

    # Show plot if needed.
    if showPlot
        gui( plt )
    end  # if showPlot

    # Save plot if needed.
    if savePlot
        # If the user doesn't specify a filename, a random one is generated.
        if filename == ""
            dateStr = Dates.format( now(), "yyyymmdd HHMMSS" )
            filename = string( "Initial population ", names( report )[1],
                " plot  ", dateStr )
        end  # if filename == ""

        if !ispath( dirname( filename ) )
            mkpath( dirname( filename ) )
        end  # if !ispath( dirname( filename ) )

        extension = split( filename, "." )
        extension = length( extension ) == 1 ? "" : extension[end]
        filename = string( filename, extension ∈ extensions ? "" : ".svg" )
        savefig( plt, filename )
    end  # if savePlot

end  # generateInitialPopAgePlot( report, timeFactor, showPlot, savePlot,
     #   filename )


text:
# Hutchinson Trace Estimator for Inverse of Matrix
# Based on Hutchinson's Stochastic Estimation Proof
# Values from Extrapolation Methods for Estimating Trace of the Matrix Inverse by P. Fika
# (This method works for SPD with low condition number)

export hutch, hutch!, HutchWorkspace

using LinearAlgebra

struct HutchWorkspace{T, TM <: AbstractMatrix{T}, F <: Function, TV <: AbstractVector{T}}
    A::TM
    randfunc::F
    x::TV
    y::TV
    N::Int64
    skipverify::Bool
end

"""
    HutchWorkspace(A; N = 30, skipverify = false)
    HutchWorkspace(A, randfunc::Function; N = 30, skipverify = false)

# Arguments
- `A` : Symmetric Positive Definite Matrix with Low Condtion Number (k < 500)
- `randfunc` : Function to generate random values for x distributed uniformly
                (Base: rand(-1.0:2.0:1.0, size(A,1))
                 Example: f(n) = rand(-1.0:2.0:1.0, n)
- `N` : Number of iterations (Default: 30)
- `skipverify` : If false, it will check isposdef(A) (Default: false)
"""
function HutchWorkspace(A, randfunc::Function; N = 30, skipverify = false)
    x = randfunc(size(A,1))
    y = similar(x)
    return HutchWorkspace(A, randfunc, x, y, N, skipverify)
end

function HutchWorkspace(A; N = 30, skipverify = false)
    randfunc(n) = rand(-1:2:1, n)
    x = rand(eltype(A) <: Integer ? Float64 : eltype(A), size(A,1))
    y = similar(x)
    return HutchWorkspace(A, randfunc, x, y, N, skipverify)
end

# Calculating the moments and extrapolating c-1
function ev(w::HutchWorkspace)
    # Create new random values for x
    copyto!(w.x, w.randfunc(size(w.A, 1)))

    # Find c(r) = x' * A^r * x = dot(x', A^r * x)
    c0 = dot(w.x, w.x)
    mul!(w.y, w.A, w.x)
    c1 = dot(w.y, w.x)
    c2 = dot(w.y, w.y)
    mul!(w.x, w.A, w.y)
    c3 = dot(w.y, w.x)

    # intermediate arguments for v0
    nVal = (c1^2)/(c0 * c2)
    dVal = (c1 * c3)/(c2^2)

    if (nVal > 0) && (dVal > 0)
        # Find p (Page 175)
        p = (c0 * c2)/(c1^2)
        # Find v0 (Page 179, Numerical Example) -> Corollary 4 in [16]
        v0 = log10(nVal) / log10(dVal)
        # Finally find ev (Page 175)
        ev0 = (c0^2)/(p^v0 * c1)
        return ev0
    else
        @warn "v0 cannot be calculated. Aitken's process may give inaccurate results!"
        w1 = HutchWorkspace(w.A, w.randfunc, w.x, w.y, w.N, w.skipverify)
        return hutch!(w1, aitken=true)
    end
end

# Aitken's Process to Predict the negative moment
# (Page 176, eq 4)
function gfun(w::HutchWorkspace)
    copyto!(w.x, w.randfunc(size(w.A, 1)))

    c0 = dot(w.x, w.x)
    mul!(w.y, w.A, w.x)
    c1 = dot(w.y, w.x)
    c2 = dot(w.y, w.y)
    mul!(w.x, w.A, w.y)
    c3 = dot(w.y, w.x)

    #g1 = c0 - (((c1 - c0)^2) / (c2 - c1))
    g2 = c1 - (((c2 - c0) * (c2 - c1)) / (c3 - c2))
    #g = (g1 + g2)/2
end

"""
    hutch(A; N = 30, skipverify = false)

Take a HutchWorkspace object as input, apply hutchinson estimation algorithm and solve for trace
of inverse of the matrix. (in-place version also available)
"""
# Hutchinson trace estimation using extrapolation technique (Page 177)
function hutch(A; N=30, skipverify = false, aitken = false)
    w = HutchWorkspace(A, N = N, skipverify = skipverify)
    return hutch!(w, aitken)
end

"""
    hutch!(w::HutchWorkspace; aitken = false)
    hutch!(A::AbstractArray; N = 30, skipverify = false, aitken = false)

Take a HutchWorkspace object as input, apply hutchinson estimation algorithm and solve for trace
of inverse of the matrix. (in-place version of hutch)
"""
function hutch!(w::HutchWorkspace; aitken = false)
    if w.skipverify == true || isposdef(w.A)
        tr = zero(eltype(w.A))
        sum = zero(eltype(w.A))
        for i in 1:w.N
            # if aitken == true => use aitken process to predict the terms via gfun
            sum = sum + (aitken ? gfun(w) : ev(w))
        end
        tr = sum/w.N
    end
end

function hutch!(A::AbstractArray{<:Any, 2}; N = 30, skipverify = false, aitken = false)
    w = HutchWorkspace(A, N, skipverify)
    return hutch!(w, aitken)
end


text:
# test case to enforce consistency in joint gibbs

using Test
using IncrementalInference

##

@testset "test disjoint up clique message joint subgraph" begin

## test case with disjoint clique joint subgraph

fg = initfg()

addVariable!(fg, :x0, ContinuousEuclid{2})
addVariable!(fg, :x1, ContinuousEuclid{2})
addVariable!(fg, :x2, ContinuousEuclid{2})

initManual!(fg, :x0, randn(2,100))
initManual!(fg, :x1, randn(2,100) .+ 10)
initManual!(fg, :x2, randn(2,100) .+ 20)

addFactor!(fg , [:x0; :x1], LinearRelative(MvNormal([10.0;10], diagm([1.0;1]))))
addFactor!(fg , [:x1; :x2], LinearRelative(MvNormal([10.0;10], diagm([1.0;1]))))

addVariable!(fg, :x3, ContinuousEuclid{2})
addFactor!( fg, [:x2; :x3], EuclidDistance(Normal(10, 1)) )

addFactor!( fg, [:x0; :x3], EuclidDistance(Normal(30, 1)), graphinit=false )

##

ensureAllInitialized!(fg)
# drawGraph(fg, show=true)

## test shortest path

# first path
tp1_ = [ :x0;:x0x1f1;:x1;:x1x2f1;:x2]
tp2_ = [ :x0;:x0x3f1;:x3;:x2x3f1;:x2]

pth = findShortestPathDijkstra(fg, :x0, :x2)
@test pth == tp1_ || pth == tp2_

pth = findShortestPathDijkstra(fg, :x0, :x2, typeFactors=[LinearRelative;])
@test pth == tp1_

# different path
pth = findShortestPathDijkstra(fg, :x0, :x2, typeFactors=[EuclidDistance;])
@test pth == tp2_

##

isHom, typeName = isPathFactorsHomogeneous(fg, :x0, :x2)

@test isHom
@test length(typeName) == 1
@test typeName[1].name == :LinearRelative

## use a specific solve order

vo = [:x3; :x1; :x2; :x0]  # getEliminationOrder(fg)
tree = buildTreeReset!(fg, vo)

##

# drawTree(tree, show=true)

## Child clique subgraph

cliq2 = getClique(tree,:x3)
# drawGraph(cfg2, show=true)

##

cfg2 = buildCliqSubgraph(fg, cliq2)

jointmsg = IIF._generateMsgJointRelativesPriors(cfg2, cliq2)

@info "update these jointmsg test after #1010"
@test intersect( keys(jointmsg.priors) , [:x0; :x2] ) |> length == 2
@test length(jointmsg.relatives) == 0

##

# retlist = addLikelihoodsDifferentialCHILD!([:x2; :x0], cfg2)

##

tree, _, = solveTree!(fg, eliminationOrder=vo);

## get up message from child clique

msgBuf = IIF.getMessageBuffer(getClique(tree, :x3))
msg = msgBuf.upTx

# only priors for these mixed factors, before TAF
@test msg.jointmsg.priors |> length == 2
@test msg.jointmsg.relatives |> length == 0

##

end

@testset "test homogeneous up clique message joint subgraph" begin

## test case with disjoint clique joint subgraph

fg = initfg()

addVariable!(fg, :x0, ContinuousEuclid{2})
addVariable!(fg, :x1, ContinuousEuclid{2})
addVariable!(fg, :x2, ContinuousEuclid{2})

initManual!(fg, :x0, randn(2,100))
initManual!(fg, :x1, randn(2,100) .+ 10)
initManual!(fg, :x2, randn(2,100) .+ 20)

addFactor!(fg , [:x0; :x1], LinearRelative(MvNormal([10.0;10], diagm([1.0;1]))))
addFactor!(fg , [:x1; :x2], LinearRelative(MvNormal([10.0;10], diagm([1.0;1]))))

addVariable!(fg, :x3, ContinuousEuclid{2})
addFactor!( fg, [:x2; :x3], LinearRelative(MvNormal([10.0;10], diagm([1.0;1]))) )

addFactor!( fg, [:x0; :x3], LinearRelative(MvNormal([10.0;10], diagm([1.0;1]))) )

##

ensureAllInitialized!(fg)
# drawGraph(fg, show=true)

## test shortest path

# first path
tp1_ = [ :x0;:x0x1f1;:x1;:x1x2f1;:x2]
tp2_ = [ :x0;:x0x3f1;:x3;:x2x3f1;:x2]

pth = findShortestPathDijkstra(fg, :x0, :x2)
@test pth == tp1_ || pth == tp2_

pth = findShortestPathDijkstra(fg, :x0, :x2, typeFactors=[LinearRelative;])
@test pth == tp1_ || pth == tp2_

# different path
pth = findShortestPathDijkstra(fg, :x0, :x2, typeFactors=[EuclidDistance;])
@test pth == Symbol[]

##

vo = [:x3; :x1; :x2; :x0]  # getEliminationOrder(fg)
tree = buildTreeReset!(fg, vo)
cliq2 = getClique(tree,:x3)

cfg2 = buildCliqSubgraph(fg, cliq2)

# drawGraph(cfg2, show=true)

##

separators = getCliqSeparatorVarIds(cliq2)
jointrelatives = addLikelihoodsDifferentialCHILD!( cfg2, separators )
allClasses = IIF._findSubgraphsFactorType( cfg2, jointrelatives, separators )

@test length(allClasses) == 1
@test intersect(allClasses[1], separators) |> length == 2

##

jointmsg = IIF._generateMsgJointRelativesPriors(cfg2, cliq2)

##

@info "update these jointmsg test after #1010"
@test intersect( jointmsg.relatives[1].variables , [:x0; :x2] ) |> length == 2
@test length(keys(jointmsg.priors)) == 0

## similar during solving

tree, _, = solveTree!(fg, eliminationOrder=vo);

## get up message from child clique

msgBuf = IIF.getMessageBuffer(getClique(tree, :x3))
msg = msgBuf.upTx

## only priors for these mixed factors, before TAF

@test intersect( msg.jointmsg.relatives[1].variables , [:x0; :x2] ) |> length == 2
@test length(keys(msg.jointmsg.priors)) == 0

##

end

#


text:
using GeometricIntegrators.Equations
using GeometricIntegrators.Equations: function_v_dummy, get_λ₀
using Test

t₀ = 0.
q₀ = [1.]
p₀ = [1.]
x₀ = [1., 1.]
λ₀ = [0.]

@testset "$(rpad("General equation functionality",80))" begin

    @test function_v_dummy(t₀, q₀, p₀, λ₀) == nothing

    @test get_λ₀(zeros(3), zeros(3))   == zeros(3)
    @test get_λ₀(zeros(3), zeros(3,3)) == zeros(3)
    @test get_λ₀(zeros(3,3), zeros(3)) == zeros(3,3)

    struct TestEquation{DT,TT} <: Equation{DT,TT} end

    @test_throws ErrorException ndims(TestEquation{Float64,Float64}())

end

include("deterministic_equations_tests.jl")
include("stochastic_equations_tests.jl")


text:
"""
Describes a dam purpose, such as providing drinking water (i.e. water withdrawn from the reservoir), or ensuring inland
navigation (i.e. constraints on the reservoir level behind the dam).

Each purpose must have a unique identifier, whose type is `Symbol`. It is always available with the `getId()` function.
For implementers, a default implementation uses the field `id`.
"""
abstract Purpose

getId(p::Purpose) = p.id
isDeterministic(p::Purpose) = false

id(p::Purpose) = getId(p)

"""
Described a dam purpose that is fulfilled by taking water from the reservoir, such as providing drinking water.

The given discharges are expressed as 10^6 m^3/s.
"""
abstract WaterWithdrawalPurpose <: Purpose

isWaterWithdrawal(p::Purpose) = false
isWaterWithdrawal(p::WaterWithdrawalPurpose) = true

"""
Describes a deterministic dam purpose, without any uncertainty, which only withdraws water.

Each such purpose corresponds to a given needed discharge. It is always available with the `getNeed()` function.
For implementers, a default implementation uses the field `need`. The given discharges are expressed as 10^6 m^3/s.

The rationale behind this layer of abstraction (instead of making it concrete, with the various purposes having
a different id) is that some purposes may have multiple instances for the same dam. For example, a conveyance network
may extract water from multiple dams for a given quantity of water to reach; and the same dam may participate in multiple
such summations.
"""
abstract DeterministicWaterWithdrawalPurpose <: WaterWithdrawalPurpose

getNeed(d::DeterministicWaterWithdrawalPurpose) = d.need
getNeed(d::DeterministicWaterWithdrawalPurpose, p::Period) = getNeed(d) * Dates.days(p) * 86400.
isDeterministic(d::DeterministicWaterWithdrawalPurpose) = true

need(d::DeterministicWaterWithdrawalPurpose) = getNeed(d)
need(d::DeterministicWaterWithdrawalPurpose, p::Period) = getNeed(d, p)

"""
Represents a deterministic drinking water consumption from the reservoir. It can only be fulfilled by one reservoir.

By default, the id is `:drinkingWater`.
"""
immutable DeterministicDrinkingWater <: DeterministicWaterWithdrawalPurpose
  id::Symbol
  need::Float64 # 10^6 m^3/s
  DeterministicDrinkingWater(v::Discharge) = new(:drinkingWater, _from_unitful(v))
  DeterministicDrinkingWater(id::Symbol, v::Discharge) = new(id, _from_unitful(v))
end

"""
Represents a deterministic environmental flow that must be released from the reservoir into its river.

By default, the id is `:environmentalFlow`.
"""
immutable DeterministicEnvironmentalFlow <: DeterministicWaterWithdrawalPurpose
  id::Symbol
  need::Float64 # 10^6 m^3/s
  DeterministicEnvironmentalFlow(v::Discharge) = new(:environmentalFlow, _from_unitful(v))
  DeterministicEnvironmentalFlow(id::Symbol, v::Discharge) = new(id, _from_unitful(v))
end

"""
Represents a hydropower unit, which has no need for water as the exact discharge is a degree of freedom.

The id is computed with that of the hydropower unit as `getPurposeId(hp)`.

This object is automatically created when creating a `Reservoir` data structure, and thus does not need to be used
outside internal code.
"""
immutable DeterministicHydropower <: DeterministicWaterWithdrawalPurpose
  hp::HydropowerUnit
end

getId(hp::DeterministicHydropower) = getPurposeId(hp.hp)
getNeed(hp::DeterministicHydropower) = 0.
getHydropowerUnit(hp::DeterministicHydropower) = hp.hp
isHydropower(p::Purpose) = false
isHydropower(p::DeterministicHydropower) = true

hydropower(hp::DeterministicHydropower) = hp.hp

"""
Builds an array of deterministic purposes for the most common cases, i.e. one `DeterministicDrinkingWater`
and one `DeterministicEnvironmentalFlow` at most.
"""
function DeterministicPurposes(;drinkingWater::Discharge=0.m^3/s, environmentalFlow::Discharge=0.m^3/s)
  list = Purpose[] # All data structures expect a generic array of purposes, be they deterministic or not.
  if drinkingWater != 0.m^3/s
    push!(list, DeterministicDrinkingWater(drinkingWater))
  end
  if environmentalFlow != 0.m^3/s
    push!(list, DeterministicEnvironmentalFlow(environmentalFlow))
  end
  return list
end


text:
function display_arrow(prob::Algames.GameProblem, sce::Scenario, ref_pdtraj::PrimalDualTraj, val::T,
	eig_pdtraj::PrimalDualTraj; vis_opts::PlayerVisualizationOptions=PlayerVisualizationOptions{T}(),
	color=:orange, key::Int=0, β::T=1.2*10^5, height::T=0.01) where {T}

	probsize = prob.probsize
	N = probsize.N
	α = vis_opts.α
	for k = 2:N
		vs_ref = standardize(prob.model, Algames.state(ref_pdtraj.pr[k]))
		vs_eig = standardize(prob.model, Algames.state(eig_pdtraj.pr[k]))
		for i = 1:p
			player = sce.player[i]
			arr_vis = ArrowVisualizer(vis["env/arrow/arrow_traj$key/player$(player.id)/$k"])
			col = vis_opts.colors[color]
			mat = MeshPhongMaterial(color=RGBA(col, α))

			setobject!(arr_vis, mat)
			settransform!(arr_vis,
				Point(0.0, 0.0, 0.0),
				Vec(β*val*norm([vs_eig[i].x, vs_eig[i].y]), 0.0, 0.0),
				shaft_radius=0.008,
				max_head_radius=0.020)

			t = Translation(vs_ref[i].x, vs_ref[i].y, height)
			r1 = LinearMap(AngleAxis(pi/2, 0.0, 1.0, 0.0))
			r2 = LinearMap(AngleAxis(-atan(vs_eig[i].y,vs_eig[i].x), 1.0, 0.0, 0.0))
			tr = compose(t, compose(r1, r2))
			settransform!(vis["env/arrow/arrow_traj$key/player$(player.id)/$k"], tr)
		end
	end
	return nothing
end

function display_ghost(vis::Visualizer, prob::Algames.GameProblem, sce::Scenario)
    N = prob.probsize.N
    for k ∈ vcat(Vector((1:10:N)),N)
        @show k
        vis_opts = PlayerVisualizationOptions{T}(α=0.3+0.2*k/10)
        set_player!(vis, sce.player, vis_opts=vis_opts, key=k)
        set_state!(vis, prob.model, sce, Algames.state(prob.pdtraj.pr[k]), key=k)
    end
    return nothing
end

function get_num_active_constraint(prob::Algames.GameProblem)
    probsize = prob.probsize
    p = probsize.p
	update_active_set!(prob.game_con, prob.pdtraj.pr)
    na = 0
    for i = 1:p
        for j = i+1:p
            conval, v = get_collision_conval(prob.game_con, i, j)
            if v
				for (l,k) in enumerate(conval.inds)
					na += conval.active[l][1]
				end
			end
        end
    end
    return na
end

function update_λ!(ascore::ActiveSetCore, prob::Algames.GameProblem, Δλcol::AbstractVector)
    probsize = prob.probsize
    N = probsize.N
    p = probsize.p
    px = probsize.px

    stamp = CStamp()
    for i = 1:p
        for j ∈ setdiff(1:p,i)
            conval, v = get_collision_conval(prob.game_con, i, j)
            if v
                for (l,k) in enumerate(conval.inds)
                    stampify!(stamp, :h, :col, i, j, k)
                    if conval.active[l][1] == 1
                        conval.λ[l] += Δλcol[horizontal_idx(ascore, stamp) .- probsize.S]
                        conval.λ[l] = max.(0, conval.λ[l])
                    end
                end
            end
        end
    end
    return nothing
end

function subspace_dimension(prob::Algames.GameProblem; α::T=2e-3) where {T}
    prob_ = deepcopy(prob)
    probsize = prob.probsize
    ascore = ActiveSetCore(probsize)
    update_nullspace!(ascore, prob, prob.pdtraj)
    ns = length(ascore.null.vec)
    val = 0
    subspace = Vector{PrimalDualTraj}([deepcopy(prob.pdtraj)])
    plt = plot(legend=false)
    for l = 1:ns
        @show l
        prob_ = deepcopy(prob)
        prob_.opts.shift = 0
        prob_.opts.α_dual = 0.0
        prob_.opts.αx_dual = 1e-3*ones(p)
        prob_.opts.dual_reset = false
        prob_.opts.ρ_0 = prob_.pen.ρ[1]
        prob_.opts.ρ_max = prob_.pen.ρ[1]
        prob_.opts.inner_print = false
        prob_.opts.outer_print = false
        prob_.opts.inner_iter = 20
        prob_.opts.outer_iter = 40
        prob_.opts.reg_0 = 1e-6
        ϵ = 1e-6
        prob_.opts.ϵ_opt = ϵ
        prob_.opts.ϵ_dyn = ϵ
        prob_.opts.ϵ_con = ϵ
        prob_.opts.ϵ_sta = ϵ
        set_traj!(prob_.core, prob_.Δpdtraj, α*ascore.null.Δtraj[l])
        update_traj!(prob.pdtraj, prob_.pdtraj, 1.0, prob_.Δpdtraj)
        update_λ!(ascore, prob_, α*ascore.null.Δλ[l])
        newton_solve!(prob_)

        opt = mean(abs.(prob_.stats.opt_vio[end].vio))
        sta = mean(abs.(prob_.stats.sta_vio[end].vio))
        dyn = mean(abs.(prob_.stats.dyn_vio[end].vio))
        con = mean(abs.(prob_.stats.con_vio[end].vio))
        opts = prob_.opts
        if opt < opts.ϵ_opt && sta < opts.ϵ_sta && dyn < opts.ϵ_dyn && con < opts.ϵ_con
            val += 1
			@show val
            push!(subspace, deepcopy(prob_.pdtraj))
        end
    end
    @show ns
    @show val
    return ascore, subspace
end

function pca(prob::Algames.GameProblem, subspace::Vector{PrimalDualTraj}; β::T=2.0) where {T}
    probsize = prob.probsize
    N = probsize.N
    n = probsize.n
    m = probsize.m
    p = probsize.p
    px = probsize.px
    S = probsize.S
    M = length(subspace)-1

    V = []
    v = zeros(S)
    v_ref = zeros(S)
    get_traj!(prob.core, subspace[1], v_ref)
    for k = 1:M
        get_traj!(prob.core, subspace[k+1], v)
        v = v - v_ref
        push!(V, deepcopy(v))
    end

    pr_mask = primals_mask(prob.core, probsize)
	R = length(pr_mask)
	W = [v[pr_mask] for v in V]
    m = sum(W) ./ M
    Wc = [w .- m for w in W]
    cov = sum([wc*wc' for wc in Wc])/(M-1)
    vecs = eigvecs(cov)
    vals = eigvals(cov)
    eigs = [(vecs[:,i], vals[i]) for i=1:R]
    sort!(eigs, by = x -> x[2], rev=true)
	vecs = [e[1] for e in eigs[1:M]]
    vals = [e[2] for e in eigs[1:M]]

    @show round.(log.(10, abs.(vals)), digits=3)

    plt = plot(legend=false)
    pdtraj_ref = subspace[1]
    plot!(prob.model, pdtraj_ref.pr, plt=plt)

    pdtraj_1 = deepcopy(pdtraj_ref)
    v1 = zeros(S)
    v1[pr_mask] = vecs[1]
    set_traj!(prob.core, pdtraj_1, v1)

    pdtraj_2 = deepcopy(pdtraj_ref)
    v2 = zeros(S)
    v2[pr_mask] = vecs[2]
    set_traj!(prob.core, pdtraj_2, v2)

    pdtraj_3 = deepcopy(pdtraj_ref)
    v3 = zeros(S)
    v3[pr_mask] = vecs[3]
    set_traj!(prob.core, pdtraj_3, v3)

    x = [[Algames.state(pdtraj_ref.pr[k])[px[i][1]] for k=2:N] for i=1:p]
    y = [[Algames.state(pdtraj_ref.pr[k])[px[i][2]] for k=2:N] for i=1:p]
    u1 = [[vals[1]*Algames.state(pdtraj_1.pr[k])[px[i][1]]  for k=2:N] for i=1:p]
    v1 = [[vals[1]*Algames.state(pdtraj_1.pr[k])[px[i][2]]  for k=2:N] for i=1:p]
    u2 = [[vals[2]*Algames.state(pdtraj_2.pr[k])[px[i][1]]  for k=2:N] for i=1:p]
    v2 = [[vals[2]*Algames.state(pdtraj_2.pr[k])[px[i][2]]  for k=2:N] for i=1:p]
    u3 = [[vals[3]*Algames.state(pdtraj_3.pr[k])[px[i][1]]  for k=2:N] for i=1:p]
    v3 = [[vals[3]*Algames.state(pdtraj_3.pr[k])[px[i][2]]  for k=2:N] for i=1:p]

    for i = 1:p
        quiver!(x[i],y[i],quiver=(β*u1[i],β*v1[i]), linewidth=1.5, color=:cyan)
        quiver!(x[i],y[i],quiver=(β*u2[i],β*v2[i]), linewidth=1.5, color=:red)
        # quiver!(x[i],y[i],quiver=(β*u3[i],β*v3[i]), linewidth=1.5, color=:black)
    end

    display(plt)
    return vals, pdtraj_ref, pdtraj_1, pdtraj_2
end

function primals_mask(core::NewtonCore, probsize::ProblemSize)
    N = probsize.N
    inds = Vector{Int}()
    stamp = HStamp()
    for k = 1:N-1
        # States
        stampify!(stamp, :x, 1, k+1)
        ind = horizontal_idx(core, stamp)
        inds = vcat(inds, Vector(ind))
        # Controls
        for i = 1:p
            stampify!(stamp, :u, i, k)
            ind = horizontal_idx(core, stamp)
            inds = vcat(inds, Vector(ind))
        end
    end
    return inds
end

function display_eigvals(vals)
    plt = plot(legend=false)
    plot!(vals,
        linetype=:steppost,
        linewidth=3.0,
        title="PCA on the Nash Equilibrium subspace \n around a nominal NE trajectory.",
        xlabel="eigval_index",
        ylabel="eigval_magnitude")
    display(plt)
	for (i,val) in enumerate(vals)
		println("($i,$val)")
	end
    return nothing
end

function cost(game_obj::Algames.GameObjective, pdtraj::PrimalDualTraj)
	probsize = game_obj.probsize
	N = probsize.N
	p = probsize.p
	p = probsize.p
	c = zeros(p)
	for i = 1:p
		for obj in game_obj.obj[i]
			c[i] = Algames.cost(obj, pdtraj.pr)
		end
	end
	return c
end

# Newton Solver
function simple_projection(α::T, λ::AbstractVector, x::AbstractVector, v::AbstractVector,
	β::T, ascore::ActiveSetCore, prob::Algames.GameProblem, vmask::Vector{Int}) where {T}
	probsize = prob.probsize
	N = probsize.N
	n = probsize.n
	m = probsize.m
	p = probsize.p
	S = probsize.S
	Sh = S + p*(p-1)*(N-1)
	Sv = S + Int(p*(p-1)*(N-1)/2)
	nx = Sh
	nλ = Sv
	ix = Vector(1:nx)
	iλ = Vector(nx .+ (1:nλ))
	jmask = vcat(ix[vmask], iλ[vmask])

	xp = deepcopy(x + α*v)

	for l = 1:1
		set_primal_dual!(ascore, prob, xp)
		#################
		prob.opts.shift = 0
		prob.opts.α_dual = 0.0
		prob.opts.αx_dual = 0.05*α*ones(p)
		prob.opts.dual_reset = false
		prob.opts.ρ_0 = prob.pen.ρ[1]
		prob.opts.ρ_max = prob.pen.ρ[1]
		prob.opts.inner_print = false
		prob.opts.outer_print = false
		prob.opts.inner_iter = 20
		prob.opts.outer_iter = 60
		prob.opts.reg_0 = 1e-6
		ϵ = 1e-6
		prob.opts.ϵ_opt = ϵ
		prob.opts.ϵ_dyn = ϵ
		prob.opts.ϵ_con = ϵ
		prob.opts.ϵ_sta = ϵ
		newton_solve!(prob)
		xp = get_primal_dual(ascore, prob)
		println("f_no_change", scn(f(x, α, x, v, β)))
		println("f_solution ", scn(f(xp, α, x, v, β)))
		#################
		record!(prob.stats, prob.core, prob.model,
			prob.game_con, prob.pdtraj, prob.opts.outer_iter)
		opt = mean(abs.(prob.stats.opt_vio[end].vio))
		sta = mean(abs.(prob.stats.sta_vio[end].vio))
		dyn = mean(abs.(prob.stats.dyn_vio[end].vio))
		con = mean(abs.(prob.stats.con_vio[end].vio))
		println("opt = "*scn(opt))
		println("sta = "*scn(sta))
		println("dyn = "*scn(dyn))
		println("con = "*scn(con))
	end
	return xp, λ
end

# Utils
function get_primal_dual(ascore::ActiveSetCore, prob::Algames.GameProblem)
	probsize = prob.probsize
	N = probsize.N
	n = probsize.n
	m = probsize.m
	p = probsize.p
	S = probsize.S
	Sh = S + p*(p-1)*(N-1)

	y = zeros(S)
	x = zeros(Sh)
	cs = CStamp()
	get_traj!(prob.core, prob.pdtraj, y)
	x[1:S] = y
	for k = 1:N-1
		for i = 1:p
			for j = 1:p
				conval, v = get_collision_conval(prob.game_con, i,j)
				if v
					for (l,k) in enumerate(conval.inds)
						stampify!(cs, :h, :col, i, j, k)
						x[horizontal_idx(ascore, cs)] = conval.λ[l]
					end
				end
			end
		end
	end
	return x
end

function set_primal_dual!(ascore::ActiveSetCore, prob::Algames.GameProblem, x::AbstractVector)
	probsize = prob.probsize
	N = probsize.N
	n = probsize.n
	m = probsize.m
	p = probsize.p
	S = probsize.S

	set_traj!(prob.core, prob.pdtraj, x[1:S])

	cs = CStamp()
	for k = 1:N-1
		for i = 1:p
			for j = 1:p
				conval, v = get_collision_conval(prob.game_con, i,j)
				if v
					for (l,k) in enumerate(conval.inds)
						stampify!(cs, :h, :col, i, j, k)
						conval.λ[l] = x[horizontal_idx(ascore, cs)]
					end
				end
			end
		end
	end
	return nothing
end

function f(xp::AbstractVector, α::T, x::AbstractVector, v::AbstractVector, β::T) where {T}
	Δ = (xp - (x + α*v))
	return 0.5*β*Δ'*Δ
end

function Algames.scn(a::Number; digits::Int=1)
	@assert digits >= 0
    # a = m x 10^e
    if a == 0
        e = 0
        m = 0.0
    else
        e = Int(floor(log(abs(a))/log(10)))
        m = a*exp(-e*log(10))
    end
    m = round(m, digits=digits)
    if digits == 0
        m = Int(floor(m))
		strm = string(m)
	else
		strm = string(m)
		is_neg = m < 0.
		strm = strm*"0"^abs(2+digits+is_neg-length(strm))
    end
    sgn = a >= 0 ? " " : ""
    sgne = e >= 0 ? "+" : ""
    return "$sgn$(strm)e$sgne$e"
end


text:
# # [a + b for a in 1:5, b in 1:3]
# function fill_phi(f, ts, Gs, u)
# 	f_u = x -> f(x..., u)
# 	@cast ϕ[i, j] := f_u((ts[i], ts[j], Gs[i, :], Gs[j, :]))
# 	# mapreduce((tsi, Gsi) ->
# 	# 	mapreduce((tsj, Gsj) ->
# 	# 		f(tsi, tsj, Gsi, Gsj, u),
# 	# 	vcat, ts, Gs),
# 	# hcat, ts, Gs)
# 	# reduce(hcat,
# 	# 			[[f(ts[j], ts[i], Gs[j, :], Gs[i, :], u)
# 	# 		for i in 1:s]
# 	# 	for	j in 1:s])
# 	# [f(ts[j], ts[], Gs[j, :], Gs[i, :], u) for i in 1:s, j in 1:s]
# end

"""
makes phi matrix for calculating the mean function
i.e. the eqn below eqn 11
"""
@views function get_phi_E(t::AbstractFloat, c::Int, d::Int,
						 gp::GaussianProcess)
	i_low = sum(1:c - 1)
	i_high = i_low + c

	Gs = gp.dpars.G[d, i_low + 1:i_high, :]

# 	reduce(hcat,
# 				[[gp.base_kernel(t, t, pi, ppi, gp.dpars.u)
# 			for ppi in G_pars_sub]
# 		for	pi in G_pars_sub])3
	ts = fill(t, c)
	gp.base_kernel(ts, Gs, gp.dpars.u)
end

# function fill_phi_zip_adj(f, ts, Gs, u)
# 	big_grad = 	reduce(hcat,
# 							[[collect(gradient(f, ti, tpi, pi, ppi, u))
# 						for (tpi, ppi) in zip(ts, Gs)]
# 					for	(ti, pi) in zip(ts, Gs)])
# 	big_grad
# end

    """
makes phi matrix for calculating the cross covariance
function i.e. the bit inside () just above example 2
"""
@views function get_phi_cov(t::AbstractFloat, tp::AbstractFloat, c::Int, cp::Int,
							d::Int, dp::Int, gp::GaussianProcess)
	i_low = sum(1:c - 1)
	i_high = i_low + c

	ip_low = sum(1:cp - 1)
	ip_high = ip_low + cp

	Gs = [gp.dpars.G[d, i_low + 1:i_high, :] ; gp.dpars.G[dp, ip_low + 1:ip_high, :]]
	ts = [fill(t, c) ; fill(tp, cp)]
	gp.base_kernel(ts, Gs, gp.dpars.u)
end

    """
calcultes terms inside the sum for prop one kan 2008, as it was faster this way
"""
function kan_rv_prod_inner(phi::AbstractArray, v::NTuple)
	vl = collect(v)
	h = 0.5 .- vl
	hKh =  ( 0.5 * dot(h, phi * h))^(size(phi)[1] / 2)
	(-1.0)^(sum(vl)) * hKh
end

"""
calculates the expectation of products of Gaussian RVs, with
exponents s and covariance matrix phi, using the identity in proposition 1
of Kan 2008.
"""
function kan_rv_prod(phi::AbstractArray)
	st = size(phi)[1]
	mapreduce(v -> kan_rv_prod_inner(phi, v), +, Iterators.product(fill(0:1, st)...)) / factorial(st ÷ 2)
end

"""
this is a custom adjoint for Zygote autodiff
https://github.com/FluxML/Zygote.jl/issues/292
is a bit of a hack and needs properly testing, but seems to work
and provided a 10x speedup
"""
function kan_rv_prod_adj(phi::Array{<:AbstractFloat,2})
    st = size(phi)[1]
    g = v -> gradient(kan_rv_prod_inner, phi, v)[1]
	sum(g, Iterators.product(fill(0:1, st)...)) / factorial(st ÷ 2)
end

@adjoint function kan_rv_prod(phi::Array{<:AbstractFloat,2})
    kan_rv_prod(phi), x -> (x * kan_rv_prod_adj(phi),)
end

const hs_fs = Dict{Int, Tuple{Matrix{Float64}, Vector{Float64}}}()
function get_hs_fs(st)
	get!(hs_fs, st) do
		print(st)
        hs = reduce(hcat, [0.5 .- digits(n, base=2, pad=st) for n in 0:2^st - 1])
        fs = map(v -> 2iseven(sum(v))-1, [digits(n, base=2, pad=st) for n in 0:2^st-1]) ./ factorial(st ÷ 2)
        (hs, fs)
    end
end

function kan_rv_prod_t(phi)
	st = size(phi, 1)
	hs, fs = get_hs_fs(st)
	@tullio ds[k] := hs[i,k] * phi[i,j] * hs[j,k] nograd=hs
	@tullio out := fs[k] * ds[k]^(st/2) nograd=fs
end

"""
implements eqn 10
"""
function full_E(t::Real, d::Int, gp::GaussianProcess)
	val = 0.
	for c in 1:gp.C
		if c % 2 == 0
			phi = get_phi_E(t, c, d, gp)
			# println(kan_rv_prod(phi))
			val += kan_rv_prod(phi)
		end
	end
	val
end

"""
implements eqn 12
"""
function full_cov(t::Real, tp::Real, d::Int, dp::Int, gp::GaussianProcess)
	val = 0.
        for c in 1:gp.C
		for cp in 1:gp.C
			if (c + cp) % 2 == 0
				# println(c,cp)
				phi = get_phi_cov(t, tp, c, cp, d, dp, gp)
				# println(size(phi))
				val += kan_rv_prod(phi)
			end
		end
	end
	val
end

"""
implements equation just above section 3.2
"""
function kernel(t::Real, tp::Real, d::Int, dp::Int, gp::GaussianProcess)
	cov = full_cov(t, tp, d, dp, gp)
    E = full_E(t, d, gp)
	Ep = full_E(tp, dp, gp)

	cov - E*Ep
end


text:
using Test

@testset "test two" begin
    @test true
    @test true
    @test true
end


text:
#!/usr/bin/env julia

# Gentlemen, start your engines
println("\e[34mInitialising\e[0m")

# Required packages
# -----------------------------------------------------------------------------

using EzXML

# Function definitions
# -----------------------------------------------------------------------------

# Core function for collecting XML data per the input spec
function collect_data(root::EzXML.Node, spec)::Dict{String,Any}
  local ns::String = namespace(root)
  local obj = Dict{String,Any}()
  for s in spec
    f = findall(s.xpath, root, ["x" => ns])  # Hardcoded as x!
    if length(f) === 0
      obj[s.name] = nothing
      continue
    end
    if length(f) > 1
      println(
        "\e[33mWARN |\e[0m  Spec. \"$(s.name)\" entries Gt1, using first"
      )
    end
    obj[s.name] = nodecontent(f[1])
  end
  return obj
end

# Abstracted reducer function to collapse the data into a TSV string array
function reduce_data_tsv(
    acc::Array{String,1},
    crt::Dict{String,Any},
    headers::Array{String,1}
  )::Array{String,1}
  local out = Array{String,1}()
  for h in headers
    push!(out, get(crt, h, ""))
  end
  return push!(acc, join(out, "\t"))
end

# Operations
# -----------------------------------------------------------------------------

# Prepare variables to use throughout
println("\e[34mAssigning common variables\e[0m")

# Paths in and out
path = (
  ip_dir = "../../hmdb/data/hmdb_metabolites",
  op = "../data/outRequest04AnalyteList.tsv"
)

# Headers for the output TSV
headers_tsv = [
  "RT",
  "Formula",
  "Name",
  "CCS [M+H]+",
  "CCS [M+Na]+",
  "CCS [M-H]-",
  "KEGG",
  "CAS",
  "PubChem",
  "ChemSpider",
  "HMDB",
  "BioCyc",
  "Metlin",
  "LipidMaps",
  "UserID",
  "InChI"
]

# Define the name-xpath specification pairs
pf = "//x:metabolite"  # Shorthand
spec = [
  (name = "Formula", xpath = "$pf/x:chemical_formula"),
  (name = "Name", xpath = "$pf/x:name"),
  (name = "KEGG", xpath = "$pf/x:kegg_id"),
  (name = "CAS", xpath = "$pf/x:cas_registry_number"),
  (name = "PubChem", xpath = "$pf/x:pubchem_compound_id"),
  (name = "ChemSpider", xpath = "$pf/x:chemspider_id"),
  (name = "HMDB", xpath = "$pf/x:accession"),
  (name = "BioCyc", xpath = "$pf/x:biocyc_id"),
  (name = "Metlin", xpath = "$pf/x:metlin_id"),
  (name = "InChI", xpath = "$pf/x:inchi"),
]

# Here we go then
println("\e[36mCollecting data\e[0m")

# Get it all
data = map(
  f -> collect_data(root(readxml("$(path.ip_dir)/$f")), spec),
  readdir(path.ip_dir)
)

# That wasn't too bad
println("\e[36mReducing output string\e[0m")

# Get it ready
out = reduce(
  (a, c) -> reduce_data_tsv(a, c, headers_tsv),
  data;
  init = ["$(join(headers_tsv, "\t"))"]
)

# Send it here
println("\e[32mWriting file:  $(path.op)\e[0m")

# Scribble
open(path.op, "w") do file
  write(file, "$(join(out, "\n"))\n")
end

# I am properly amazed at how fast and effortless that was!  Took less than a
# minute to cross 114,000+ entries.  Ruby scripts to do the same, which
# themselves were pretty quick considering previous attempts, took somewhere
# around 15 minutes (which is entirely to be expected with this volume of
# data).  I am very impressed with the capabilities of Julia!


text:
module REFPROPwrap
"""
Initializing by calling the SETUPdll with required fluid name

currently support only pure fluids with one component
"""
function refprop_init(fluid::String)
    ncmax = Ref{Cint}(2)
    refpropcharlength  = Ref{Cint}(255)
    filepathlength  = Ref{Cint}(255)
    lengthofreference = Ref{Cint}(3)
    errormessagelength =  Ref{Cint}(255)
    hfld_length = Ref{Cint}(refpropcharlength[]*ncmax[])
    hf = Array{UInt8}(undef,refpropcharlength[]*ncmax[])
    hrf = Array{UInt8}(undef, lengthofreference[])
    herr = Array{UInt8}(undef, errormessagelength[])
    hfmix = Array{UInt8}(undef,refpropcharlength[])
    hf_name = "/home/aaditya/Desktop/julia/REFPROP/FLUIDS/"*fluid*".FLD\0"
    hfmix_name = "/home/aaditya/Desktop/julia/REFPROP/FLUIDS/HMX.BNC"
    hrf_name = "DEF"
    herr_name = "Ok"
    hf[1:length(hf_name)] = collect(hf_name)
    hfmix[1:length(hfmix_name)] = collect(hfmix_name)
    hrf[1:length(hrf_name)] = collect(hrf_name)
    herr[1:length(herr_name)] = collect(herr_name)
    i = Ref{Cint}(1)
    ierr = Ref{Cint}(0)
    ccall((:SETUPdll,"/home/aaditya/Desktop/julia/librefprop.so"),Cvoid,(Ref{Cint},Ptr{UInt8},Ptr{UInt8},Ptr{UInt8},Ref{Cint},Ptr{UInt8},Ref{Cint},Ref{Cint},Ref{Cint},Ref{Cint}), i, hf, hfmix, hrf, ierr, herr, hfld_length, refpropcharlength, lengthofreference, errormessagelength)
    return String(herr)
end #refprop_init
"""
sample function to call the properties using pressure and enthalpy as the inputs
"""
function refprop_call(pre::Float64, enthal::Float64)
    #initializing the arrays and variables
    ncmax = Ref{Cint}(2)
    ierr = Ref{Cint}(0)
    herr = Array{UInt8}(undef, 255)

    x = Array{Cdouble}(zeros(ncmax[])); xliq = Array{Cdouble}(zeros(ncmax[])); xvap = Array{Cdouble}(zeros(ncmax[])); f = Array{Cdouble}(zeros(ncmax[]));

    wm = Ref{Cdouble}(0)
    vs = Ref{Cdouble}(0)
    ccall((:WMOLdll,"/home/aaditya/Desktop/julia/librefprop.so"), Cvoid, (Ptr{Cdouble}, Ref{Cdouble}), x, wm)
    p = Ref{Cdouble}(pre)
    e = Ref{Cdouble}(0)
    h = Ref{Cdouble}(enthal*wm[]/1000.0)
    s = Ref{Cdouble}(0)
    cv = Ref{Cdouble}(0)
    cp = Ref{Cdouble}(0)
    hjt = Ref{Cdouble}(0)
    temp = Ref{Cdouble}(0)
    rhop = Ref{Cdouble}(0) #(rho/wm[])
    rhol = Ref{Cdouble}(0)
    rhov = Ref{Cdouble}(0)
    qual = Ref{Cdouble}(0)

    #ccall((:THERMdll,"/home/aaditya/Desktop/julia/librefprop.so"),Cvoid,(Ref{Cdouble}, Ref{Cdouble},Ptr{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble}), temp, rhop, x, p, e, h, s, cv, cp, vs, hjt);
    ccall((:PHFLSHdll,"/home/aaditya/Desktop/julia/librefprop.so"),Cvoid,(Ref{Cdouble}, Ref{Cdouble}, Ptr{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cint}, Ptr{UInt8}), p,h,x,temp,rhop,rhol,rhov,xliq,xvap,qual,e,s,cv,cp,vs,ierr,herr);

    return p[], h[]/wm[]*1000.0, temp[], rhop[]*wm[], rhol[]*wm[], rhov[]*wm[], qual[],  e[]/wm[]*1000.0, s[]/wm[]*1000.0, cv[]/wm[]*1000.0, cp[]/wm[]*1000.0, vs[]
end #refprop_call

export refprop_inti, refprop_call

end # module