text:
using Random, MLJBase

function gen_reg(; n=100, p=5, seed=143)
    Random.seed!(143)
    X = randn(n, p)
    y = randn(n)
    return MLJBase.table(X), y
end

function gen_classif(; n=100, p=5, seed=145, classes=["A", "B"])
    Random.seed!(seed)
    X = randn(n, p)
    # gen [1, 2, 1, 3,..]
    M = exp.(randn(n, length(classes)))
    Mn = M ./ sum(M, dims=2)
    z = multi_rand(Mn)
    # associate labels
    y = [classes[zᵢ] for zᵢ in z]
    return MLJBase.table(X), MLJBase.categorical(y)
end

function gen_dummy_classif_binary(; n=50, p=5, seed=1566)
    # create clouds of points that are super separated
    cloud1 = randn(n, p)  .+ 5.0
    cloud2 = randn(n, p)  .- 5.0
    test1  = randn(10, p) .+ 5.0
    test2  = randn(10, p) .- 5.0
    X = MLJBase.table(vcat(cloud1, cloud2))
    Xt = MLJBase.table(vcat(test1, test2))
    y1 = fill("A", n)
    y2 = fill("B", n)
    yt1 = fill("A", 10)
    yt2 = fill("B", 10)
    y = MLJBase.categorical(vcat(y1, y2))
    yt = MLJBase.categorical(vcat(yt1, yt2))
    return X, Xt, y, yt
end

function gen_dummy_classif(; n=50, p=5, seed=1566)
    # create clouds of points that are super separated
    cloud1 = randn(n, p)  .+ 5.0
    cloud2 = randn(n, p)
    cloud3 = randn(n, p)  .- 5.0
    test1  = randn(10, p) .+ 5.0
    test2  = randn(10, p)
    test3  = randn(10, p) .- 5.0
    X = MLJBase.table(vcat(cloud1, cloud2, cloud3))
    Xt = MLJBase.table(vcat(test1, test2, test3))
    y1 = fill("A", n)
    y2 = fill("B", n)
    y3 = fill("C", n)
    yt1 = fill("A", 10)
    yt2 = fill("B", 10)
    yt3 = fill("C", 10)
    y = MLJBase.categorical(vcat(y1, y2, y3))
    yt = MLJBase.categorical(vcat(yt1, yt2, yt3))
    return X, Xt, y, yt
end

# simple function to sample multinomial
function multi_rand(Mp)
    n, c = size(Mp)
    be   = reshape(rand(length(Mp)), n, c)
    y    = zeros(Int, n)
    @inbounds for i in eachindex(y)
        rp = 1.0
        for k in 1:c-1
            if (be[i, k] < Mp[i, k] / rp)
                y[i] = k
                break
            end
            rp -= Mp[i, k]
        end
    end
    y[y .== 0] .= c
    return y
end

function simple_test_reg(m, X, y)
    f, _, _ = fit(m, 1, X, y)
    p = predict(m, f, X)
    @test norm(p .- y) / norm(y) < 1
    m, f
end

function test_dummy_classif(m; seed=5154, binary=false, thresh=0.75)
    if binary
        X, Xt, y, yt = gen_dummy_classif_binary(seed=seed)
    else
        X, Xt, y, yt = gen_dummy_classif(seed=seed)
    end
    f, _, _ = fit(m, 1, X, y)
    p = typeof(m) <: Probabilistic ? predict_mode(m, f, Xt) : predict(m, f, Xt)
    @test sum(p .== yt) / length(yt) ≥ thresh
end

function simple_test_classif(m, X, y; dummybinary=false, nodummy=false, thresh=0.75)
    f, _, _ = fit(m, 1, X, y)
    p = predict(m, f, X)
    @test eltype(p) == eltype(y)
    @test Set(unique(p)) == Set(unique(y))
    nodummy || test_dummy_classif(m; binary=dummybinary, thresh=thresh)
    m, f
end

function simple_test_classif_prob(m, X, y; dummybinary=false, nodummy=false, thresh=0.75)
    f, _, _ = fit(m, 1, X, y)
    p = predict_mode(m, f, X)
    @test eltype(p) == eltype(y)
    @test Set(unique(p)) == Set(unique(y))
    p = predict(m, f, X)
    @test eltype(p) <: UnivariateFinite
    nodummy || test_dummy_classif(m; binary=dummybinary, thresh=thresh)
    m, f
end


text:
@testset "Prolog syntax conversion" begin

pl_input = @prolog """
    member(X, [X | Y]).
    member(X, [Y | YS]) :- member(X, YS).
    vertebrate(A) :- fish(A) ; amphibian(A) ; reptile(A) ; bird(A) ; mammal(A).
    bird(A) :- dinosaur(A), \\+reptile(A).
    reptile(A) :- member(A, .(stegosaurus, .(triceratops, []))).
    dinosaur(A) :- member(A, [archaeopteryx, stegosaurus, triceratops]).
"""

julog_clauses = @julog [
    member(X, [X | Y]) <<= true,
    member(X, [Y | YS]) <<= member(X, YS),
    vertebrate(A) <<= fish(A),
    vertebrate(A) <<= amphibian(A),
    vertebrate(A) <<= reptile(A),
    vertebrate(A) <<= bird(A),
    vertebrate(A) <<= mammal(A),
    bird(A) <<= dinosaur(A) & !reptile(A),
    reptile(A) <<= member(A, cons(stegosaurus, cons(triceratops, []))),
    dinosaur(A) <<= member(A, [archaeopteryx, stegosaurus, triceratops])
]

@test pl_input == julog_clauses

@test resolve(@prolog("bird(archaeopteryx)."), pl_input)[1] == true

pl_output = """
member(X, [X | Y]).
member(X, [Y | YS]) :- member(X, YS).
vertebrate(A) :- fish(A).
vertebrate(A) :- amphibian(A).
vertebrate(A) :- reptile(A).
vertebrate(A) :- bird(A).
vertebrate(A) :- mammal(A).
bird(A) :- dinosaur(A), \\+(reptile(A)).
reptile(A) :- member(A, [stegosaurus, triceratops]).
dinosaur(A) :- member(A, [archaeopteryx, stegosaurus, triceratops]).
"""

@test write_prolog(pl_input) == pl_output

end

ext: jl


text:
export dmrg!

"""
    dmrg!(psi::Mps{T}, H::Mpo{T}, dmrg_opts::DMRGOpts) where {T<:Number}

Minimize energy of psi with respect to H using a DMRG algorithm.

Warning: the Hamiltonian is assumed to be Hermitian. This function will silently
fail if that is not the case.

Output:
    - `E::Vector{Float64}`: energy of psi after every sweep.
    - `var::Vector{Float64}`: variance of psi after every sweep.
"""
function dmrg!(psi::Mps{T}, H::Mpo{T}, dmrg_opts::DMRGOpts) where {T<:Number}
    # Create left and right environments.
    Le = fill(ones(T, 1, 1, 1), psi.L)
    Re = fill(ones(T, 1, 1, 1), psi.L)
    # Initialize left environment.
    for i=2:psi.L
        Le[i] = prop_right3(Le[i-1], psi.M[i-1], H.W[i-1], psi.M[i-1])
    end

    # Compute energy of state and variance at each sweep. Sweep 1 is just
    # the computation of the energy and the variance.
    E = Float64[]
    var = Float64[]
    # Noise term for DMRG3S.
    α = dmrg_opts.α
    # Cache to store temporaries.
    cache = Cache(Vector{AbstractArray{T}}())
    for s=1:dmrg_opts.nsweeps
        # Do left and right sweeps.
        time = @elapsed begin
        if dmrg_opts.algorithm == "DMRG1"
            Es = do_sweep_1s!(psi, H, Le, Re, -1, s, dmrg_opts)
            Es = do_sweep_1s!(psi, H, Le, Re, +1, s, dmrg_opts)
        elseif dmrg_opts.algorithm == "DMRG2"
            Es = do_sweep_2s!(psi, H, Le, Re, -1, s, cache, dmrg_opts)
            Es = do_sweep_2s!(psi, H, Le, Re, +1, s, cache, dmrg_opts)
        elseif dmrg_opts.algorithm == "DMRG3S"
            Es = do_sweep_3s!(psi, H, Le, Re, -1, s, cache, α, dmrg_opts)
            Es = do_sweep_3s!(psi, H, Le, Re, +1, s, cache, α, dmrg_opts)
        end
        end

        # Update energy and variance of psi after the sweep.
        push!(E, Es)
        push!(var, real(m_variance(H, psi)))

        # Print debug information: energy, variance and their variation.
        if dmrg_opts.show_trace > 0
            println("Done sweep $s, max bond dimension: $(max_bond_dim(psi))")
            @printf("    E: %.6e, ΔE: %.2e\n",
                    E[s], E[s]-(s > 1 ? E[s-1] : 1e5))
            @printf("    var: %.6e, Δvar: %.2e\n",
                    var[s], var[s]-(s > 1 ? var[s-1] : 1e5))
            println("    Elapsed time: $time s")
        end

        # Check for convergence of variance.
        if s > 1 && abs(var[s] - var[s-1]) < dmrg_opts.dmrg_tol
            println("The variance of psi has converged. Finished DMRG.")
            break
        end
    end
    return E, var
end

#
# 1-SITE DMRG.
#

"""
    update_lr_envs_1s!(psi::Mps{T},
                       H::Mpo{T},
                       Le::Vector{Array{T, 3}},
                       Re::Vector{Array{T, 3}},
                       sense::Int,
                       s::Int,
                       i::Int,
                       Mi::Vector{T},
                       dmrg_opts::DMRGOpts) where {T<:Number}

Update the left and right environments after the local Hamiltonian is minimized.
"""
function update_lr_envs_1s!(psi::Mps{T},
                            H::Mpo{T},
                            Le::Vector{Array{T, 3}},
                            Re::Vector{Array{T, 3}},
                            sense::Int,
                            s::Int,
                            i::Int,
                            Mi::Vector{T},
                            dmrg_opts::DMRGOpts) where {T<:Number}
    if sense == +1
        Mi = reshape(Mi, (size(Le[i], 1), psi.d, size(Re[i], 1)))
        psi.M[i], R = factorize_svd_right(Mi, cutoff=dmrg_opts.cutoff[s])

        # Update left and right environments at Le[i+1] and Re[i] and psi.
        psi.M[i] = psi.M[i]
        if i < psi.L
            Le[i+1] = prop_right3(Le[i], psi.M[i], H.W[i], psi.M[i])
            Re[i] = absorb_Re(Re[i], R)
        end
    else
        Mi = reshape(Mi, (size(Le[i], 1), psi.d, size(Re[i], 1)))
        L, psi.M[i] = factorize_svd_left(Mi, cutoff=dmrg_opts.cutoff[s])

        # Update left and right environments at Le[i] and Re[i-1] and psi.
        if i > 1
            Re[i-1] = prop_left3(psi.M[i], H.W[i], psi.M[i], Re[i])
            Le[i] = absorb_Le(Le[i], L)
        end
    end
    return psi
end

"""
    do_sweep_1s!(psi::Mps{T},
                 H::Mpo{T},
                 Le::Vector{Array{T, 3}},
                 Re::Vector{Array{T, 3}},
                 sense::Int,
                 s::Int,
                 dmrg_opts::DMRGOpts) where {T<:Number}

Do a sweep to locally minimize the energy of `psi` at 1 site per step. The
direction of the sweep is given by `sense = +1, -1`.
"""
function do_sweep_1s!(psi::Mps{T},
                      H::Mpo{T},
                      Le::Vector{Array{T, 3}},
                      Re::Vector{Array{T, 3}},
                      sense::Int,
                      s::Int,
                      dmrg_opts::DMRGOpts) where {T<:Number}

    # Manually increase the bond dimension.
    enlarge_bond_dimension!(psi, dmrg_opts.maxm[s])
    # Update left environment (right environment is updated in the first sweep).
    for i=2:psi.L
        Le[i] = prop_right3(Le[i-1], psi.M[i-1], H.W[i-1], psi.M[i-1])
    end

    sense == 1 || sense == -1 || throw("`Sense` must be either `+1` or `-1`.")
    # Energy before the sweep starts. Depends on the sweep sense, assume the
    # last sweep had the opposite sense.
    E = sense == +1 ?
        sum(real(Re[1].*Le[2])) : sum(real(Re[psi.L-1].*Le[psi.L]))

    # Order of sites to do the sweep.
    sweep_sites = sense == +1 ? (1:psi.L) : reverse(1:psi.L)
    for i in sweep_sites
        # Compute local minimum.
        Hi = build_local_hamiltonian(Le[i], H.W[i], Re[i])

        Hi = Hermitian(Hi)
        f(x::Vector{T}) = Hi*x
        v0 = rand(size(Hi, 1))
        v0 ./= norm(v0)
        evals, evecs, info = eigsolve(
            f, v0, 1, :SR, KrylovKit.Lanczos(maxiter=dmrg_opts.lanczos_iters)
        )
        E = real(evals[1])
        Mi = vec(evecs[1])

        # Update left and right environments.
        update_lr_envs_1s!(psi, H, Le, Re, sense, s, i, Mi, dmrg_opts)

        # Useful debug information.
        if dmrg_opts.show_trace > 1
            println("site: $i, size(Hi): $(size(Hi)), Ei: $(E)")
        end
    end
    return E
end

#
# 2-SITE DMRG.
#

"""
    update_lr_envs_2s!(psi::Mps{T},
                       H::Mpo{T},
                       Le::Vector{Array{T, 3}},
                       Re::Vector{Array{T, 3}},
                       sense::Int,
                       s::Int,
                       i::Int,
                       Mi::Vector{T},
                       dmrg_opts::DMRGOpts) where {T<:Number}

Update the left and right environments after the local Hamiltonian is minimized
with 2-site algorithm.
"""
function update_lr_envs_2s!(psi::Mps{T},
                            H::Mpo{T},
                            Le::Vector{Array{T, 3}},
                            Re::Vector{Array{T, 3}},
                            sense::Int,
                            s::Int,
                            i::Int,
                            Mi::Vector{T},
                            dmrg_opts::DMRGOpts) where {T<:Number}

    # Decompose the Mi tensor spanning sites i and i+1 with SVD.
    Mi = reshape(Mi, size(psi.M[i], 1), psi.d, psi.d, size(psi.M[i+1], 3))
    dimcutoff = bond_dimension_with_m(psi.L, i+1, dmrg_opts.maxm[s], psi.d)
    psi.M[i], psi.M[i+1] = sense == 1 ?
        factorize_svd_right(Mi, dimcutoff=dimcutoff) :
        factorize_svd_left(Mi, dimcutoff=dimcutoff)

    # Update environments.
    Le[i+1] = prop_right3(Le[i], psi.M[i], H.W[i], psi.M[i])
    Re[i] = prop_left3(psi.M[i+1], H.W[i+1], psi.M[i+1], Re[i+1])
    return
end

"""
    do_sweep_2s!(psi::Mps{T},
                 H::Mpo{T},
                 Le::Vector{Array{T, 3}},
                 Re::Vector{Array{T, 3}},
                 sense::Int,
                 s::Int,
                 cache::Cache{T},
                 dmrg_opts::DMRGOpts) where {T<:Number}

Do a sweep to locally minimize the energy of `psi` at 2 sites per step. The
direction of the sweep is given by `sense = +1, -1`.
"""
function do_sweep_2s!(psi::Mps{T},
                      H::Mpo{T},
                      Le::Vector{Array{T, 3}},
                      Re::Vector{Array{T, 3}},
                      sense::Int,
                      s::Int,
                      cache::Cache{T},
                      dmrg_opts::DMRGOpts) where {T<:Number}

    sense == 1 || sense == -1 || throw("`Sense` must be either `+1` or `-1`.")
    # Energy before the sweep starts. Depends on the sweep sense, assume the
    # last sweep had the opposite sense.
    E = sense == +1 ?
        sum(real(Re[1].*Le[2])) : sum(real(Re[psi.L-1].*Le[psi.L]))

    # Order of sites to do the sweep.
    sweep_sites = sense == +1 ? (1:psi.L-1) : reverse(2:psi.L-1)
    for i in sweep_sites
        # Compute local minimum.
        Hi = build_local_hamiltonian_2(Le[i], H.W[i], H.W[i+1], Re[i+1], cache)
        # Build an initial vector for eigs using the previous M tensors.
        M1 = reshape(deepcopy(psi.M[i]),
                     (size(psi.M[i], 1)*size(psi.M[i], 2), size(psi.M[i], 3)))
        M2 = reshape(deepcopy(psi.M[i+1]),
                     (size(psi.M[i+1], 1), size(psi.M[i+1], 2)*size(psi.M[i+1], 3)))
        v0 = vec(M1*M2)

        Hi = Hermitian(Hi)
        f(x::Vector{T}) = Hi*x
        evals, evecs, _ = eigsolve(
            f, v0, 1, :SR, KrylovKit.Lanczos(maxiter=dmrg_opts.lanczos_iters)
        )
        E = real(evals[1])
        Mi = vec(evecs[1])

        # Update left and right environments.
        update_lr_envs_2s!(psi, H, Le, Re, sense, s, i, Mi, dmrg_opts)

        # Useful debug information.
        if dmrg_opts.show_trace > 1
            println("Site: $i, size(Hi): $(size(Hi)), Ei: $(E)")
        end
    end
    return E
end

#
# STRICTLY-SINGLE-SITE DMRG3S.
#

"""
    update_lr_envs_3s!(psi::Mps{T},
                       H::Mpo{T},
                       Le::Vector{Array{T, 3}},
                       Re::Vector{Array{T, 3}},
                       sense::Int,
                       s::Int,
                       i::Int,
                       Mi::Vector{T},
                       dmrg_opts::DMRGOpts) where {T<:Number}

Update the left and right environments after the local Hamiltonian is minimized.
"""
function update_lr_envs_3s!(psi::Mps{T},
                            H::Mpo{T},
                            Le::Vector{Array{T, 3}},
                            Re::Vector{Array{T, 3}},
                            sense::Int,
                            s::Int,
                            i::Int,
                            Mi::Vector{T},
                            α::Float64,
                            dmrg_opts::DMRGOpts) where {T<:Number}
    if sense == +1
        # Subspace expansion.
        Mi = reshape(Mi, (size(Le[i], 1), psi.d, size(Re[i], 1)))
        P = prop_right_subexp(Le[i], H.W[i], Mi)
        P = reshape(P, (size(Le[i], 1)*psi.d, size(Re[i], 2)*size(Re[i], 3)))
        # Add subpsace expansion to local mimimum.
        Mi = reshape(Mi, (size(Le[i], 1)*psi.d, size(Re[i], 1)))
        MP = hcat(Mi, α*P)

        # Svd.
        cutoff = dmrg_opts.cutoff[s]
        dimcutoff = bond_dimension_with_m(psi.L, i+1, dmrg_opts.maxm[s], psi.d)
        psi.M[i], SV = factorize_svd_right(reshape(MP, 1, size(MP)...),
                                           cutoff=cutoff, dimcutoff=dimcutoff)
        # Trim the dimension of SV to match with psi.M[i+1].
        SV = SV[:, 1:size(Re[i], 1)]
        # Unbind the bond dimension in psi.M[i].
        psi.M[i] = reshape(psi.M[i], size(Le[i], 1), psi.d, size(psi.M[i], 3))

        # Update left and right environments at Le[i+1] and Re[i] and psi.
        if i < psi.L
            Le[i+1] = prop_right3(Le[i], psi.M[i], H.W[i], psi.M[i])
            Re[i] = absorb_Re(Re[i], SV)
            # Update this tensor to provide a good initial vector later in eigs.
            psi.M[i+1] = absorb_fromleft(SV, psi.M[i+1])
        end
    else
        # Subspace expansion.
        Mi = reshape(Mi, (size(Le[i], 1), psi.d, size(Re[i], 1)))
        P = prop_left_subexp(H.W[i], Mi, Re[i])
        P = reshape(P, (psi.d*size(Re[i], 1), size(Le[i], 2)*size(Le[i], 3)))
        # Add subpsace expansion to local mimimum.
        Mi = reshape(Mi, (size(Le[i], 1), psi.d*size(Re[i], 1)))
        MP = vcat(Mi, α*transpose(P))

        # Svd.
        cutoff = dmrg_opts.cutoff[s]
        dimcutoff = bond_dimension_with_m(psi.L, i, dmrg_opts.maxm[s], psi.d)
        US, psi.M[i] = factorize_svd_left(reshape(MP, size(MP)..., 1),
                                          cutoff=cutoff, dimcutoff=dimcutoff)
        # Trim the dimension of US to match with psi.M[i-1].
        US = US[1:size(Le[i], 1), :]
        # Unbind the bond dimension in psi.M[i].
        psi.M[i] = reshape(psi.M[i], size(psi.M[i], 1), psi.d, size(Re[i], 1))

        # Update left and right environments at Le[i+1] and Re[i] and psi.
        if i > 1
            Re[i-1] = prop_left3(psi.M[i], H.W[i], psi.M[i], Re[i])
            Le[i] = absorb_Le(Le[i], US)
            # Update this tensor to provide a good initial vector later in eigs.
            psi.M[i-1] = absorb_fromright(psi.M[i-1], US)
        end
    end
    return psi
end

"""
    do_sweep_3s!(psi::Mps{T},
                 H::Mpo{T},
                 Le::Vector{Array{T, 3}},
                 Re::Vector{Array{T, 3}},
                 sense::Int,
                 s::Int,
                 cache::Cache{T},
                 dmrg_opts::DMRGOpts) where {T<:Number}

Do a sweep to locally minimize the energy of `psi` at 1 site per step. The
direction of the sweep is given by `sense = +1, -1`.
"""
function do_sweep_3s!(psi::Mps{T},
                      H::Mpo{T},
                      Le::Vector{Array{T, 3}},
                      Re::Vector{Array{T, 3}},
                      sense::Int,
                      s::Int,
                      cache::Cache{T},
                      α::Float64,
                      dmrg_opts::DMRGOpts) where {T<:Number}

    sense == 1 || sense == -1 || throw("`Sense` must be either `+1` or `-1`.")
    # Energy before the sweep starts. Depends on the sweep sense, assume the
    # last sweep had the opposite sense.
    E = sense == +1 ?
        sum(real(Re[1].*Le[2])) : sum(real(Re[psi.L-1].*Le[psi.L]))

    # Order of sites to do the sweep.
    sweep_sites = sense == +1 ? (1:psi.L-1) : reverse(2:psi.L)
    for i in sweep_sites
        # Compute local minimum.
        Hi = build_local_hamiltonian(Le[i], H.W[i], Re[i], cache)
        v0 = vec(deepcopy(psi.M[i]))
        Hi = Hermitian(Hi)
        f(x::Vector{T}) = Hi*x
        evals, evecs, _ = eigsolve(
            f, v0, 1, :SR, KrylovKit.Lanczos(maxiter=dmrg_opts.lanczos_iters)
        )
        E1 = real(evals[1])
        delta_E1 = E1-E
        Mi = vec(evecs[1])

        # Update left and right environments.
        update_lr_envs_3s!(psi, H, Le, Re, sense, s, i, Mi, α, dmrg_opts)

        # Compute new energy and update α.
        E = sense == +1 ? sum(real(Re[i].*Le[i+1])) : sum(real(Re[i-1].*Le[i]))
        delta_E = E-E1
        if -delta_E/delta_E1 > 0.3/log(dmrg_opts.maxm[s])
            α /= 2.
        elseif 0. < -delta_E/delta_E1 < min(1/dmrg_opts.maxm[s], 0.1)
            α *= 2.
        end

        # Useful debug information.
        if dmrg_opts.show_trace > 1
            println("site: $i, size(Hi): $(size(Hi)), Ei: $(E), α: $(dmrg_opts.α)")
        end
        if dmrg_opts.show_trace > 2
            @printf("ΔE_0: %.3e, ΔE_0/ΔE_T: %.3e\n", delta_E1, delta_E/delta_E1)
        end
    end
    return E
end


text:
"""Script to compute the memory capacity of a AssemblyCalculus.BrainArea

Script projects random stimuli into the brain area and compares assembly that
form in the region. This is an attempt to answer the question, how many
suitably unique assemblies can a region generate.
"""

using AssemblyCalculus
using AssemblyCalculus: overlap, forget_parents!
using Graphs
using Plots

NUM_BRAINS = 4
MAX_ASSEMBLIES = 500
MAX_ITERS = 100
TOL = 3
N = 10_000
K = 100
EDGE_PROBS = [0.005f0, 0.01f0, 0.05f0, 0.1f0]
β = 0.1f0

assemblies = [Assembly{Float32}[] for i in 1:NUM_BRAINS]
max_distances = [Int[] for i in 1:NUM_BRAINS]
iters_taken = [Int[] for i in 1:NUM_BRAINS]

for brain_num in 1:NUM_BRAINS
    cutoff = (2EDGE_PROBS[brain_num] / π) ^ 0.5
    g, _ = euclidean_graph(N, 2, cutoff=cutoff)
    println("Mean degree: $(length(edges(g)) / N)")
    brain = Brain(DiGraph(g), [N], [K], [β])
    for assem_num in 1:MAX_ASSEMBLIES
        stim = rand_stim(brain[1])
        assems, _, dists = simulate!([stim], tol=TOL, max_iters=MAX_ITERS)
        a = assems[1]
        forget_parents!(a)
        iters = size(dists)[2]
        push!(iters_taken[brain_num], iters)
        if assem_num != 1
            if dists[iters] <= TOL
                common_neurons = map(x -> overlap(a, x), assemblies[brain_num])
                push!(max_distances[brain_num], maximum(common_neurons))
            else
                push!(max_distances[brain_num], -1)
            end
        end
        push!(assemblies[brain_num], a)
    end
end

p1 = plot(legend=:bottomright)
colors = [:teal, :pink, :green, :orange]
for (i, ep) in enumerate(EDGE_PROBS)
    max_ds = max_distances[i]
    mean_max_iters = cumsum(max_ds) ./ collect(1:length(max_ds))
    plot!(p1, mean_max_iters, label="p = $ep", linewidth=3, color=colors[i])
end
xlabel!(p1, "Number of Assemblies Created")
ylabel!(p1, "Average max overlap \n with previous assemblies")

p2 = plot(legend=:bottomright)
for (i, ep) in enumerate(EDGE_PROBS)
    plot!(p2, max_distances[i], label="p = $ep", linewidth=3, color=colors[i])
end
xlabel!(p2, "Number of Assemblies Created")
ylabel!(p2, "Max overlap with \n previous assemblies")

l = @layout [a b]
plot(p1, p2, layout=l, size=(1200, 400))
suptitle!("Assembly capacity for different edge probabilities \n(Random Geometric topology)")


text:
module Filetimes

export filetime, datetime, EPOCH

using Dates

const EPOCH =  Dates.DateTime(1601, 1, 1)

"""
    filetime()

Returns the current time as the number of 100 nanosecond intervals
since January 1, 1601.

"""
function filetime()
    filetime(now())
end

"""
    filetime(d::Dates.DateTime)

Returns the number of 100 nanosecond intervals since January 1, 1601
for the given DateTime.

"""
function filetime(d::Dates.DateTime)
    Dates.value(d - EPOCH) * 10000
end

"""
    datetime(f::Integer)

Returns a DateTime object given a filetime f, being the number of 100
nanosecond intervals since January 1, 1601.

N.B. Possible loss of precision.

"""
function datetime(f::Integer)
    EPOCH + Dates.Microsecond(f/10)
end

"""
    filetime(s::AbstractString)

Returns the number of 100 nanosecond intervals since January 1, 1601,
for the given ISO8601 string.

    filetime("2014-09-02T08:20:07")

"""
function filetime(s::AbstractString)
    filetime(DateTime(s))
end

end # module


text:
# using LinearAlgebra
# using MeshCat
# using GeometryBasics
# using Colors
# using CoordinateTransformations

vis = Visualizer()
open(vis)

struct SphereContact{T}
    position::Vector{T}
    radius::T
end

p1 = [-0.11, 0.0, 0.0]
p2 = [0.11, 0.0, 0.0]
s1 = SphereContact(p1, 0.1)
s2 = SphereContact(p2, 0.1)
m1 = GeometryBasics.Sphere(Point(0.0, 0.0, 0.0), s1.radius)
m2 = GeometryBasics.Sphere(Point(0.0, 0.0, 0.0), s2.radius)
setobject!(vis[:s1], m1, MeshPhongMaterial(color=Colors.RGBA(1.0, 0.0, 0.0, 1.0)))
setobject!(vis[:s2], m2, MeshPhongMaterial(color=Colors.RGBA(1.0, 0.0, 0.0, 1.0)))
settransform!(vis[:s1], Translation(s1.position))
settransform!(vis[:s2], Translation(s2.position))

function sdf_sphere_sphere(s1::SphereContact, s2::SphereContact)
    position_error = s2.position - s1.position
    total_radius = s1.radius + s2.radius
    norm(position_error) - total_radius # should be >= 0
end

sdf_sphere_sphere(s1, s2)

function sdf_sq_sphere_sphere(s1::SphereContact, s2::SphereContact)
    position_error = s2.position - s1.position
    total_radius = s1.radius + s2.radius
    dot(position_error, position_error) - total_radius^2 # should be >= 0
end

sdf_sq_sphere_sphere(s1, s2)

function direction_sphere_sphere(s1::SphereContact, s2::SphereContact)
    # from sphere 1 to sphere 2
    position_error = s2.position - s1.position
    magnitude = norm(position_error)

    if magnitude > 0.0
        direction = position_error ./ magnitude
    else
        direction = [1.0, 0.0, 0.0]
    end

    return direction
end

direction_sphere_sphere(s1, s2)

function closest_points_sphere_sphere(s1::SphereContact, s2::SphereContact)
    # direction
    direction = direction_sphere_sphere(s1::SphereContact, s2::SphereContact)

    # closest point on sphere 1
    p1 = s1.position + direction * s1.radius

    # closest point on sphere 2
    p2 = s2.position - direction * s2.radius

    return p1, p2
end

c1, c2 = closest_points_sphere_sphere(s1, s2)

cs1 = GeometryBasics.Sphere(Point(0.0, 0.0, 0.0), 0.01)
cs2 = GeometryBasics.Sphere(Point(0.0, 0.0, 0.0), 0.01)
setobject!(vis[:cs1], cs1, MeshPhongMaterial(color=Colors.RGBA(0.0, 0.0, 0.0, 1.0)))
setobject!(vis[:cs2], cs2, MeshPhongMaterial(color=Colors.RGBA(0.0, 0.0, 0.0, 1.0)))
settransform!(vis[:cs1], Translation(c1))
settransform!(vis[:cs2], Translation(c2))


text:
module ObservationData

using ImageCore, ImageAxes, AxisIndices, NamedIndicesMeta

export
    # @defdim output
    obsdim,
    has_obsdim,
    obs_axis,
    obs_axis_type,
    obs_keys,
    obs_indices,
    to_obs_format,
    nobs,
    obs_format,
    is_observation

Base.@pure is_observation(x::Symbol) = (x === :obs) | (x === :observations) | (x === :samples)

@defdim obs is_observation

end


text:
function ecRCCSDpT{T}() where T <: AbstractFloat
    @output "Calling CASCI module...\n"
    # Call CASCI
    cas = Fermi.ConfigurationInteraction.CASCI{T}()
    ecRCCSDpT{T}(cas)
end

function ecRCCSDpT{T}(cas::Fermi.ConfigurationInteraction.CASCI) where T <: AbstractFloat

    @output "Processing CAS data...\n"
    refdet, Casdata = process_cas(cas)
    refwfn = cas.ref

    ecccsd = ecRCCSD{T}(refdet, refwfn, Casdata, CTF())
    drop_occ = Fermi.CurrentOptions["drop_occ"]
    ecRCCSDpT{T}(ecccsd, refwfn.ints, refdet, Casdata["C3dets"], refwfn.ndocc, drop_occ)
end

function ecRCCSDpT{T}(ecccsd::RCCSD, ints::IntegralHelper, ref::Determinant, casT3::Array{Determinant,1}, ndocc::Int, fcn::Int) where T <: AbstractFloat

    @output "\n   • Perturbative Triples Started\n\n"
    @output "T3 within EC active space are going to be skipped\n\n"

    T1 = ecccsd.T1.data
    T2 = ecccsd.T2.data

    Vvvvo = permutedims(ints["OVVV"], (4,2,3,1))
    Vvooo = permutedims(ints["OOOV"], (4,2,1,3))
    Vvovo = permutedims(ints["OOVV"], (3,1,4,2))

    o,v = size(T1)
    Et::T = 0.0

    fo = diag(ints["FOO"])
    fv = diag(ints["FVV"])

    # Pre-allocate Intermediate arrays
    W  = Array{T}(undef, v,v,v)
    V  = Array{T}(undef, v,v,v)

    # Pre-allocate Views
    Vvvvo_4k =    Array{T}(undef,v,v,v)
    Vvooo_2k_3j = Array{T}(undef,v,o)
    Vvooo_2k_3i = Array{T}(undef,v,o)
    Vvooo_2i_3k = Array{T}(undef,v,o)
    Vvooo_2j_3k = Array{T}(undef,v,o)
    T2_1k_2j    = Array{T}(undef,v,v)
    T2_1k_2i    = Array{T}(undef,v,v)
    T2_1i_2k    = Array{T}(undef,v,v)
    T2_1j_2k    = Array{T}(undef,v,v)
    T2_1k       = Array{T}(undef,o,v,v)
    T1_1k       = Array{T}(undef,v)
    Vvovo_2j_4k = Array{T}(undef,v,v)
    Vvovo_2i_4k = Array{T}(undef,v,v)
    Vvvvo_4i    = Array{T}(undef,v,v,v)
    T2_1i       = Array{T}(undef,o,v,v)
    T1_1i       = Array{T}(undef,v)
    Vvvvo_4j    = Array{T}(undef,v,v,v)
    Vvooo_2i_3j = Array{T}(undef,v,o)
    Vvooo_2j_3i = Array{T}(undef,v,o)
    T2_1i_2j    = Array{T}(undef,v,v)
    T2_1j_2i    = Array{T}(undef,v,v)
    T2_1j       = Array{T}(undef,o,v,v)
    Vvovo_2i_4j = Array{T}(undef,v,v)
    T1_1j       = Array{T}(undef,v)
    @output "Computing energy contribution from occupied orbitals:\n"
    for i in 1:o
        iα = ref.α ⊻ 1 << (i+fcn-1)
        @output "→ Orbital {} of {}\n" i o
        Vvvvo_4i .= view(Vvvvo, :,:,:,i)
        T2_1i    .= view(T2, i,:,:,:)
        T1_1i    .= view(T1, i, :)
        for j in 1:i
            jβ = ref.β ⊻ 1 << (j+fcn-1)
            Vvvvo_4j    .= view(Vvvvo, :,:,:,j)
            Vvooo_2i_3j .= view(Vvooo,:,i,j,:)
            Vvooo_2j_3i .= view(Vvooo,:,j,i,:)
            T2_1i_2j    .= view(T2, i,j,:,:)
            T2_1j_2i    .= view(T2, j,i,:,:)
            T2_1j       .= view(T2, j,:,:,:)
            Vvovo_2i_4j .= view(Vvovo, :,i,:,j)
            T1_1j       .= view(T1, j, :)
            δij = Int(i == j)
            for k in 1:j
                kiα = iα ⊻ 1 << (k+fcn-1)
                Vvvvo_4k    .= view(Vvvvo, :,:,:,k)
                Vvooo_2k_3j .= view(Vvooo,:,k,j,:)
                Vvooo_2k_3i .= view(Vvooo,:,k,i,:)
                Vvooo_2i_3k .= view(Vvooo,:,i,k,:)
                Vvooo_2j_3k .= view(Vvooo,:,j,k,:)

                T2_1k_2j .= view(T2, k,j,:,:)
                T2_1k_2i .= view(T2, k,i,:,:)
                T2_1i_2k .= view(T2, i,k,:,:)
                T2_1j_2k .= view(T2, j,k,:,:)

                T2_1k .= view(T2, k,:,:,:)
                T1_1k .= view(T1, k, :)

                Vvovo_2j_4k .= view(Vvovo, :,j,:,k)
                Vvovo_2i_4k .= view(Vvovo, :,i,:,k)

                # Build intermediates W and V for the set i j k
                @tensoropt begin
                    W[a,b,c]  = (Vvvvo_4i[b,d,a]*T2_1k_2j[c,d] - Vvooo_2k_3j[c,l]*T2_1i[l,a,b]  # ijk abc
                              +  Vvvvo_4i[c,d,a]*T2_1j_2k[b,d] - Vvooo_2j_3k[b,l]*T2_1i[l,a,c]  # ikj acb
                              +  Vvvvo_4k[a,d,c]*T2_1j_2i[b,d] - Vvooo_2j_3i[b,l]*T2_1k[l,c,a]  # kij cab
                              +  Vvvvo_4k[b,d,c]*T2_1i_2j[a,d] - Vvooo_2i_3j[a,l]*T2_1k[l,c,b]  # kji cba
                              +  Vvvvo_4j[c,d,b]*T2_1i_2k[a,d] - Vvooo_2i_3k[a,l]*T2_1j[l,b,c]  # jki bca
                              +  Vvvvo_4j[a,d,b]*T2_1k_2i[c,d] - Vvooo_2k_3i[c,l]*T2_1j[l,b,a]) # jik bac

                    V[a,b,c]  = W[a,b,c] + Vvovo_2j_4k[b,c]*T1_1i[a] + Vvovo_2i_4k[a,c]*T1_1j[b] + Vvovo_2i_4j[a,b]*T1_1k[c]
                end

                # Compute Energy contribution
                δjk = Int(j == k)
                for a in 1:v
                    akiα = kiα | 1 << (a+ndocc-1)
                    for b in 1:a
                        bjβ = jβ | 1 << (b+ndocc-1)
                        δab = Int(a == b)
                        for c in 1:b
                            cakiα = akiα | 1 << (c+ndocc-1)
                            _det = Determinant(cakiα, bjβ)
                            if _det in casT3
                                continue
                            end
                            Dd = fo[i] + fo[j] + fo[k] - fv[a] - fv[b] - fv[c]
                            δbc = Int(b == c)
                            @inbounds X = (W[a,b,c]*V[a,b,c] + W[a,c,b]*V[a,c,b] + W[b,a,c]*V[b,a,c] + W[b,c,a]*V[b,c,a] + W[c,a,b]*V[c,a,b] + W[c,b,a]*V[c,b,a])
                            @inbounds Y = (V[a,b,c] + V[b,c,a] + V[c,a,b])
                            @inbounds Z = (V[a,c,b] + V[b,a,c] + V[c,b,a])
                            E = (Y - 2*Z)*(W[a,b,c] + W[b,c,a] + W[c,a,b]) + (Z - 2*Y)*(W[a,c,b]+W[b,a,c]+W[c,b,a]) + 3*X
                            Et += E*(2-δij-δjk)/(Dd*(1+δab+δbc))
                        end
                    end
                end
            end
        end
    end
    @output "Final (T) contribution:   {:15.10f}\n" Et
    @output "ecCCSD(T) energy:         {:15.10f}\n" Et+ecccsd.CorrelationEnergy
    return RCCSDpT{T}(ecccsd, Et)
end


text:
using Base: Broadcast

struct DefaultJaxArrayStyle{N} <: JaxAbstractArrayStyle{N} end
DefaultJaxArrayStyle(::Val{N}) where {N} = DefaultJaxArrayStyle{N}()
DefaultJaxArrayStyle{M}(::Val{N}) where {N,M} = DefaultJaxArrayStyle{N}()
Broadcast.BroadcastStyle(::Type{<:AbstractJaxArray{T,N}}) where {T,N} =
    DefaultJaxArrayStyle{N}()

_BroadcastStyle(::Type{<:AbstractArray{T,N}}) where {T,N} =
    DefaultJaxArrayStyle{N}()
Base.convert(
    ::Type{<:DefaultJaxArrayStyle},
    a::Broadcast.DefaultArrayStyle{M},
) where {M} = DefaultJaxArrayStyle{M}()

# When broadcasting with Jax replace Julia functions with jax functions
Base.broadcasted(::JAS, f, args...) where {JAS<:JaxAbstractArrayStyle} =
    Broadcast.Broadcasted{JAS}(jaxfunc(f), args, nothing)

_pyconvert(o::PyObject) = convert(PyAny, o)
_pyconvert(o::AbstractJaxArray) = o

# Outermost layer of materialization. Convert the result of recursive inner
# _materialize calls
Broadcast.materialize(bc::Broadcast.Broadcasted{<:JaxAbstractArrayStyle}) =
    _pyconvert(_materialize(Broadcast.instantiate(bc)))

# Standard things, go back to Base
_materialize(bc::Broadcast.Broadcasted) = Base.materialize(bc)

# Numbers and other objects, just return them
_materialize(bc) = bc

# Jax stuff, go to python
_materialize(bc::Broadcast.Broadcasted{<:JaxAbstractArrayStyle}) =
    __pymaterialize(Broadcast.instantiate(bc))

__pymaterialize(bc::Broadcast.Broadcasted) = bc.f(map(_materialize, bc.args)...)

# Jax correctly preserves 0-dim objects
Base.broadcast_preserving_zero_d(f, As::AbstractJaxArray...) =
    Broadcast.materialize(Base.broadcasted(f, As...))

## inplace
Base.copyto!(dest::AbstractJaxArray, bc::Broadcast.Broadcasted{Nothing}) =
    _copyto!(dest, bc)

Base.copyto!(
    dest::AbstractJaxArray,
    bc::Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}},
) = _copyto!(dest, bc)

function _copyto!(dest::AbstractJaxArray{T,N}, bc::Broadcast.Broadcasted) where {T,N}
    obj = __pymaterialize(bc)

    indices = N != 0 ? jax.ops.index[:] : nothing

    o = pycall(jax.ops.index_update, PyObject, dest, indices, obj)

    # Swap the inner buffer because jax does not really do in-place mutation
    dest.o = o
    return dest
end

#=
## wrappers
using LinearAlgebra
for (W, ctor) in Adapt.wrappers
  @eval begin
    Broadcast.BroadcastStyle(T::Type{<:$W}) where {AT<:AbstractJaxArray} = _BroadcastStyle(T)
    #backend(::Type{<:$W}) where {AT<:AbstractGPUArray} = backend(AT)
  end
end

function _materialize(wa::LinearAlgebra.Transpose{T,A}) where {T,N,A<:AbstractJaxArray{T,N}}
    if N==1
        numpy.reshape(parent(wa), (-1,1))
    elseif N==2
        numpy.transpose(parent(wa))
    else
        error("Transpose not supported")
    end
end

_materialize(wa::LinearAlgebra.Adjoint) = conj(_materialize(Transpose(parent(wa))))

#=_materialize(bc::Broadcast.Broadcasted{<:JaxAbstractArrayStyle, T, typeof(_pmul), Tuple{V1,V2}}) where {T, V1<:AbstractJaxVector, V2<:AdjOrTrans} = begin
    @info "Jad._materialize{JaxAbstractArrayStyle} t" typeof(bc) bc.f bc.axes typeof(bc.args)
    arg1 = _materialize(bc.args[1])
    arg2 = _materialize(bc.args[2])
    #args = map(_materialize, bc.args...)
    numpy.outer(arg2, arg1)
end

_materialize(bc::Broadcast.Broadcasted{<:JaxAbstractArrayStyle, T, typeof(_pmul), Tuple{V1,V2}}) where {T, V1<:AdjOrTrans, V2<:AbstractJaxVector} = begin
    @info "Jad._materialize{JaxAbstractArrayStyle} t" typeof(bc) bc.f bc.axes typeof(bc.args)
    arg1 = _materialize(bc.args[1])
    arg2 = _materialize(bc.args[2])
    #args = map(_materialize, bc.args...)
    numpy.outer(arg1, arg2)
end
=#
=#


text:
# IndexMap is defined here because there is a boostrapping problem.
#  * IndexMap requires `Utilities.CleverDicts` and `Utilities.DoubleDicts`, so
#    if it were to be defined in MOI proper, it must be included after
#    Utilities.
#  * However, Utilities requires IndexMap, so it must be defined before
#    Utilities.jl is included.
# To work around this issue, we define `IndexMap` here.

struct IndexMap <: AbstractDict{MOI.Index,MOI.Index}
    var_map::CleverDicts.CleverDict{
        MOI.VariableIndex,
        MOI.VariableIndex,
        typeof(CleverDicts.key_to_index),
        typeof(CleverDicts.index_to_key),
    }
    con_map::DoubleDicts.IndexDoubleDict
end

"""
    IndexMap()

The dictionary-like object returned by [`VecMathOptInterface.copy_to`](@ref).
"""
function IndexMap()
    var_map = CleverDicts.CleverDict{MOI.VariableIndex,MOI.VariableIndex}()
    con_map = DoubleDicts.IndexDoubleDict()
    return IndexMap(var_map, con_map)
end

function _identity_constraints_map(
    model,
    map::MOIU.DoubleDicts.IndexDoubleDictInner{F,S},
) where {F,S}
    for c in MOI.get(model, MOI.ListOfConstraintIndices{F,S}())
        map[c] = c
    end
    return
end

"""
    identity_index_map(model::MOI.ModelLike)

Return an [`IndexMap`](@ref) that maps all variable and constraint indices of
`model` to themselves.
"""
function identity_index_map(model::MOI.ModelLike)
    variables = MOI.get(model, MOI.ListOfVariableIndices())
    map = IndexMap()
    for x in variables
        map[x] = x
    end
    for (F, S) in MOI.get(model, MOI.ListOfConstraintTypesPresent())
        _identity_constraints_map(model, map.con_map[F, S])
    end
    return map
end

Base.getindex(map::IndexMap, key::MOI.VariableIndex) = map.var_map[key]

function Base.getindex(map::IndexMap, key::MOI.ConstraintIndex{F,S}) where {F,S}
    return map.con_map[key]::MOI.ConstraintIndex{F,S}
end

function Base.setindex!(
    map::IndexMap,
    value::MOI.VariableIndex,
    key::MOI.VariableIndex,
)
    return map.var_map[key] = value
end

function Base.setindex!(
    map::IndexMap,
    value::MOI.ConstraintIndex{F,S},
    key::MOI.ConstraintIndex{F,S},
) where {F,S}
    return map.con_map[key] = value
end

Base.delete!(map::IndexMap, x::MOI.VariableIndex) = delete!(map.var_map, x)

Base.delete!(map::IndexMap, c::MOI.ConstraintIndex) = delete!(map.con_map, c)

Base.haskey(map::IndexMap, c::MOI.ConstraintIndex) = haskey(map.con_map, c)

Base.haskey(map::IndexMap, x::MOI.VariableIndex) = haskey(map.var_map, x)

function Base.keys(map::IndexMap)
    return Iterators.flatten((keys(map.var_map), keys(map.con_map)))
end

Base.length(map::IndexMap) = length(map.var_map) + length(map.con_map)

function Base.iterate(map::IndexMap, args...)
    return iterate(Base.Iterators.flatten((map.var_map, map.con_map)), args...)
end


text:
using HilbertSchmidtIndependenceCriterion
using Test

# data
X = [-0.2433   -1.1410    0.7903   -0.5062    0.5668    0.6861    0.3256   -1.2743    0.3403   -0.4115;
   -0.7473    0.7157    0.0778    0.5060   -0.8018   -0.4555    1.2192   -0.9725   -0.3975   -0.4692;
   -1.8649   -1.0257   -0.8888    0.8538   -0.2300    1.3616   -1.2759    0.6357   -1.2806    0.8666;
   -0.9981    0.6345   -0.2449   -1.1629   -1.5696    1.2884   -2.4034   -0.9240    0.6828    1.0895;
    0.6304    0.0738    0.1950    0.5096   -0.5894   -0.6975   -0.7219   -0.4880   -1.7749    0.1688;
    0.9231    0.1652    0.7928   -1.5919    0.2999   -0.5722   -0.8011    0.3885    1.8926    1.5672;
    0.9512    0.5355    0.1816   -1.5473    1.6256   -0.6757    0.2539   -0.3323    0.6048   -0.6128;
   -0.9314   -1.5384    2.2852   -0.5943   -0.4123   -1.6746   -0.5080   -1.6629    0.3894   -0.4103;
   -1.0362    0.7568    0.3759    0.8777   -0.8261    0.0162   -1.6744   -1.8138   -0.1012   -0.6914;
   -1.3328    0.2287   -0.1785   -0.3840    1.0539   -1.2816    0.0528    0.5623    0.2620   -0.1908]

Y = [-0.3025    0.3002   -0.1137   -0.1182    0.0001    0.0003   -0.0349    0.0622    0.1032    0.2134;
    0.1211    0.2137   -0.0572    0.1728   -0.1178    0.0028    0.1929    0.0682   -0.2034   -0.4904;
   -0.0240    0.2245   -0.0098   -0.2657   -0.0709   -0.0686   -0.1034    0.1989    0.0252    0.0654;
   -0.1064   -0.2296   -0.1406   -0.0826   -0.2258    0.0583   -0.1604    0.1040   -0.2303   -0.2311;
    0.1787   -0.0681    0.0634    0.0154    0.0965   -0.2212   -0.5266    0.0783   -0.1328    0.0395;
    0.1153    0.2439   -0.0669    0.4550    0.0372   -0.2356    0.0208    0.0617   -0.0346    0.3055;
   -0.0782   -0.2572    0.2275    0.1107    0.1034    0.1801    0.4776   -0.0141    0.0969   -0.2561;
    0.0732   -0.0995    0.2098    0.0772    0.4193   -0.2355    0.0905    0.2314   -0.3197   -0.1537;
    0.2838   -0.0845    0.0613   -0.2866   -0.1755   -0.0226    0.3473    0.2075   -0.1599   -0.0624;
    0.1457   -0.1459    0.0878    0.0085   -0.1861   -0.1096    0.2240   -0.0176    0.0995   -0.0136]

# Common tests
println("# Run Common Function Tests")
@test estimateKernelSize(X) ≈ 2 atol=1

# rbf dot product truth
rbftruth = [1.0000    0.2920    0.0945    0.0811    0.1872    0.1310    0.2986    0.2776    0.1904    0.2337;
    0.2920    1.0000    0.0786    0.0543    0.3356    0.0464    0.1449    0.1207    0.2955    0.2940;
    0.0945    0.0786    1.0000    0.1136    0.1459    0.0164    0.0145    0.0165    0.1107    0.1081;
    0.0811    0.0543    0.1136    1.0000    0.0746    0.1040    0.0283    0.0300    0.2275    0.0472;
    0.1872    0.3356    0.1459    0.0746    1.0000    0.0651    0.1239    0.1046    0.2800    0.1773;
    0.1310    0.0464    0.0164    0.1040    0.0651    1.0000    0.2751    0.0660    0.0352    0.1519;
    0.2986    0.1449    0.0145    0.0283    0.1239    0.2751    1.0000    0.0746    0.0577    0.3403;
    0.2776    0.1207    0.0165    0.0300    0.1046    0.0660    0.0746    1.0000    0.1383    0.1189;
    0.1904    0.2955    0.1107    0.2275    0.2800    0.0352    0.0577    0.1383    1.0000    0.1272;
    0.2337    0.2940    0.1081    0.0472    0.1773    0.1519    0.3403    0.1189    0.1272    1.0000]

@test sum(rbfDotProduct(X, X, 2.0) - rbftruth) ≈ 0.0 atol=1e-4

# Gamma HSIC tests
println("# Run Gamma HSIC Test")

@test gammaHSIC(X, Y)[2] ≈ 0.4 atol=1e-1
@test gammaHSIC(X, Y)[1] < gammaHSIC(X, Y)[2]

println("# Finished Tests")


text:
using Dates

# Date
today = Dates.today()
println(today)

# DateTime
now = Dates.now()
println(now)

# format
date_format = DateFormat("YYYY/MM/dd HH:mm")
println(Dates.format(now, date_format))

sleep(1)
dt2 = Dates.now()
diff = dt2 - now
println(( diff, typeof(diff) ))

println(now - DateTime(1970, 1, 1, 0, 0))

max_stars_repo_name: mkanenobu/trashbox

max_issues_repo_name: mkanenobu/trashbox

max_forks_repo_name: mkanenobu/trashbox


text:
# This does not work ATM
module TestWithSequential

include("../with_sequential.jl")
using KernelAbstractions

test_with_sequential(tests, get_executors(CPU()))

end


text:
using Statistics
using DataAPI

abstract type FitResult end
"""
A measured set of correlated k-ratios.
"""
struct BasicFitResult <: FitResult
    label::Label
    kratios::UncertainValues
end

function findlabel(ffr::FitResult, cxr::CharXRay)
    lbls = labels(ffr.kratios)
    fa = findall(lbl -> (lbl isa CharXRayLabel) && (cxr in lbl.xrays), lbls)
    @assert length(fa) ≠ 0 "No label found for $cxr."
    @assert length(fa) == 1 "Multiple labels found for $cxr."
    return lbls[fa[1]]
end

Base.show(io::IO, ffr::FitResult) = print(io, "$(ffr.label)")
NeXLUncertainties.NeXLUncertainties.value(ffr::FitResult, label::ReferenceLabel) =
    NeXLUncertainties.value(ffr.kratios, label)
NeXLUncertainties.σ(ffr::FitResult, label::ReferenceLabel) = σ(ffr.kratios, label)
NeXLUncertainties.uncertainty(ffr::FitResult, label::ReferenceLabel) =
    uncertainty(ffr.kratios, label)
NeXLUncertainties.getindex(ffr::FitResult, label::ReferenceLabel) =
    getindex(ffr.kratios, label)
Base.keys(ffr::FitResult) = keys(ffr.kratios)
NeXLUncertainties.labels(ffr::FitResult) = labels(ffr.kratios)
function kratio(cxr::CharXRay, ffr::FitResult)
    lbl = findlabel(ffr, cxr)
    return uv(NeXLUncertainties.value(ffr, lbl), σ(ffr, lbl))
end
unknown(ffr::FitResult)::Label = ffr.label
Base.values(ffrs::Vector{<:FitResult}, lbl::ReferenceLabel) =
    [NeXLUncertainties.value(ffr.kratios, lbl, 0.0) for ffr in ffrs]
σs(ffrs::Vector{<:FitResult}, lbl::ReferenceLabel) =
    [σ(ffr.kratios, lbl, 0.0) for ffr in ffrs]

"""
    kratios(ffr::FitResult)::Vector{KRatio}

The k-ratios associated with each `CharXRayLabel` as a vector 'KRatio' objects.
"""
function kratios(ffr::FitResult)::Vector{KRatio}
    lbls = filter(labels(ffr.kratios)) do l
        (l isa CharXRayLabel) && haskey(properties(l)[:Composition], element(l))
    end
    return map(lbls) do lbl
        props = properties(lbl)
        KRatio(lbl.xrays, properties(unknown(ffr)), props, props[:Composition], ffr.kratios[lbl])
    end
end

"""
    kratio(ffr::FitResult, cxr::CharXRay)::KRatio

Extract the k-ratio associated with the specified CharXRay (zero if not measured)
"""
function kratio(ffr::FitResult, cxr::CharXRay)
    lbls = labels(ffr.kratios)
    i = findfirst(lbl->(lbl isa CharXRayLabel) && (cxr in lbl.xrays), lbls)
    @assert !isnothing(i) "There is no k-ratio associated with $cxr."
    lbl, props = lbls[i], properties(lbls[i])
    KRatio(lbl.xrays, unknown(ffr).spectrum.properties, props, props[:Composition], ffr.kratios[lbl])
  end

"""
    extractStandards(ffr::FitResult, elm::Element, mat::Material)::Vector{KRatio}

Extract a `Vector{KRatio}` for `elm::Element` from a `ffr::FilterFitResult` measured from `mat::Material`.
"""
function extractStandards(ffr::FitResult, elm::Element, mat::Material)::Vector{KRatio}
    @assert haskey(mat, elm) "The element $elm is not contained within $mat."
    lbls = filter(labels(ffr.kratios)) do l
        (l isa CharXRayLabel) && (elm == element(l))
    end
    krs = map(lbls) do lbl
        props = properties(lbl)
        kr = KRatio(lbl.xrays, properties(unknown(ffr)), props, props[:Composition], ffr.kratios[lbl])
        kr.unkProps[:Composition] = mat
        kr
    end
    length(krs) == 0 && @info  "There are no k-ratios associated with the $elm in $ffr."
    return krs
end

"""
    extractStandards(ffr::FitResult, cxrs::AbstractVector{CharXRay}, mat::Material)::Vector{KRatio}

Extract a `KRatio` for the `CharXRay` from a `FilterFitResult` associated with
the `Material`.
"""
function extractStandards(ffr::FitResult, cxrs::AbstractVector{CharXRay}, mat::Material)::Vector{KRatio}
    lbls = filter(labels(ffr.kratios)) do lbl
        (lbl isa CharXRayLabel) && any(cxr->cxr in xrays(lbl), cxrs)
    end
    krs = map(lbls) do lbl
        props = properties(lbl)
        kr = KRatio(lbl.xrays, properties(unknown(ffr)), props, props[:Composition], ffr.kratios[lbl])
        kr.unkProps[:Composition] = mat
        kr
    end
    length(krs) == 0 && @info "There are no k-ratios associated with any of $cxrs in $ffr."
    return krs
end

"""
    heterogeneity(ffrs::Vector{FilterFitResult}, lbl::ReferenceLabel)

Computes the ratio of the standard deviation of the measured values over the mean calculated uncertainty
from the fit.  A value near 1 means the sample appears homogeneous and a value greater than 1 means the sample
appears heterogeneous.
"""
heterogeneity(ffrs::Vector{<:FitResult}, lbl::ReferenceLabel) =
    std(values(ffrs, lbl)) / mean(σs(ffrs, lbl))

"""
    NeXLUncertainties.asa(::Type{DataFrame}, ffrs::AbstractVector{<:FitResult}; charOnly = true, withUnc = false, format = :normal # :pivot or :long)

Return generic `FitResult` as a DataFrame.

Format:
  * :normal - One row per spectrum, one column per k-ratio
  * :pivot - One row per ROI, one column per spectrum (optional: column for 1σ uncertainty on k-ratio)
  * :long - One row per spectrum, feature, measured k-ratio

"""
function NeXLUncertainties.asa(
    ::Type{DataFrame},
    ffrs::AbstractVector{<:FitResult};
    charOnly = true,
    withUnc = false,
    format = :normal # :pivot or :long
)::DataFrame
    lbls = sort(unique(mapreduce(labels, append!, ffrs)))
    lbls = charOnly ? filter(lbl -> lbl isa CharXRayLabel, lbls) : lbls
    if format==:pivot
        res = DataFrame(ROI = lbls)
        for ffr in ffrs
            vals = [ value(ffr.kratios, lbl, missing) for lbl in lbls ]
            insertcols!(res, ncol(res) + 1, Symbol(repr(ffr.label)) => vals)
            if withUnc
                unc = [σ(ffr.kratios, lbl, missing) for lbl in lbls]
                insertcols!(res, ncol(res) + 1, Symbol('Δ' * repr(ffr.label)) => unc)
            end
        end
    elseif format==:long
        if withUnc
            res = DataFrame(Spectrum=String[], ROI = String[], k = Float64[], dk = Float64[])
            for ffr in ffrs, lbl in lbls
                push!(res, ( repr(ffr.label), repr(lbl), value(ffr.kratios, lbl, missing), σ(ffr.kratios, lbl, missing) ))
            end
        else
            res = DataFrame(Spectrum=String[], ROI = String[], k = Float64[])
            for ffr in ffrs, lbl in lbls
                push!(res, ( repr(ffr.label), repr(lbl), value(ffr.kratios, lbl, missing)))
            end
        end
    else
        rowLbls = [repr(ffr.label) for ffr in ffrs]
        res = DataFrame(Symbol("Spectra") => rowLbls)
        for lbl in lbls
            vals = [value(ffr.kratios, lbl, missing) for ffr in ffrs]
            insertcols!(res, ncol(res) + 1, Symbol(repr(lbl)) => vals)
            if withUnc
                unc = [σ(ffr.kratios, lbl, missing) for ffr in ffrs]
                insertcols!(res, ncol(res) + 1, Symbol('Δ' * repr(lbl)) => unc)
            end
        end
    end
    return res
end

"""
    NeXLUncertainties.asa(::Type{DataFrame}, fr::FitResult; withUnc = false)

Summarize the ROI and k-ratio data within a `FitResult` structure as a `DataFrame`.
"""
function NeXLUncertainties.asa(::Type{DataFrame}, fr::FitResult; withUnc = false)
    rois = labels(fr)
    res =  DataFrame(
        ROI = rois,
        K = map(l -> value(fr.kratios, l), rois),
    )
    withUnc && insertcols!(res, :dK => map(l -> σ(fr.kratios, l), rois) )
    return res
end

function DataAPI.describe(ffrs::Vector{<:FitResult})::DataFrame
    df = asa(DataFrame, ffrs)[:, 2:end]
    desc = describe(df, :mean, :std, :min, :q25, :median, :q75, :max)
    lbls = filter(
        lbl -> lbl isa CharXRayLabel,
        sort(collect(Set(mapreduce(labels, append!, ffrs)))),
    )
    insertcols!(desc, 4, :hetero => heterogeneity.(Ref(ffrs), lbls))
    return desc
end

"""
    FilterFitResult

Represents the result of fitting a FilteredUnknownW to a FilteredUnknown.

Struct elements

    label::UnknownLabel  # Identifies the unknown
    kratios::UncertainValues # Labeled with ReferenceLabel objects
    roi::UnitRange{Int} # Range of channels fit
    raw::Vector{Float64} # Raw spectrum data
    residual::Vector{Float64} # Residual spectrum
    peakback::Dict{ReferenceLabel,NTuple{3,Float64}} # peak counts, background counts and counts/(nAs)

Use asa(DataFrame, ffr::FilterFitResult) to summarize in tabular form.
"""
struct FilterFitResult{T <: AbstractFloat} <: FitResult
    label::UnknownLabel
    kratios::UncertainValues
    roi::UnitRange{Int}
    raw::Vector{T}
    residual::Vector{T}
    peakback::Dict{ReferenceLabel,NTuple{3,T}}
end

"""
    residual(ffr::FilterFitResult)::Spectrum

A Spectrum containing the histogram representing the unknown spectrum
minus the fitted characteristic peaks shapes times the best fit coefficient.
"""
function residual(ffr::FilterFitResult)
    props = copy(properties(ffr.label.spectrum))
    props[:Name] = "Residual[$(props[:Name])]"
    return Spectrum(ffr.label.spectrum.energy, ffr.residual, props)
end

"""
    spectrum(ffr::FilterFitResult)::Spectrum

Returns the original unknown spectrum.
"""
spectrum(ffr::FilterFitResult) = ffr.label.spectrum

"""
    characteristiccounts(ffr::FiterFitResult, strip)

Number of spectrum counts that were accounted for by the fitted elements with the `strip` Element(s) removed.
"""
characteristiccounts(ffr::FilterFitResult, strip) =
    sum(element(ref) in strip ? 0.0 : v[1] - v[2] for (ref, v) in ffr.peakback) # sum(ffr.raw[ffr.roi]-ffr.residual[ffr.roi])

"""
    peaktobackground(ffr::FilterFitResult, backwidth::Float64=10.0)::Float64

The peak-to-background ratio as determined from the raw and residual spectra integrated over
the fit region-of-interest and scaled to `backwidth` eV of continuum (nominally 10 eV).
"""
function peaktobackground(
    ffr::FilterFitResult{T},
    klabel::ReferenceLabel,
    backwidth::T = convert(T, 10.0),
) where { T <: AbstractFloat }
    unk = spectrum(unknown(ffr))
    peak, back, _ = ffr.peakback[klabel]
    return (peak * (energy(klabel.roi.stop, unk) - energy(klabel.roi.start, unk))) /
           (backwidth * back)
end

"""
    NeXLUncertainties.asa(
        ::Type{DataFrame},
        ffr::FilterFitResult;
        charOnly::Bool=true,
        material=nothing,
        mc = XPP, fc=ReedFluorescence,
        columns = () # Selected from ( :roi, :peakback, :counts, :dose )
    )::DataFrame

Tabulate details about each region-of-interest in the 'FilterFitResult' in a 'DataFrame'.
  * If charOnly then only display characteristic X-ray data (not escapes etc.)
  * If `material` is a Material then the computed k-ratio (KCalc) will also be tabulated along with kmeas/kcalc (KoKCalc).
  * columns - Select optional column outputs (see below)

Columns:

  * Spectrum : UnknownLabel - Identifies the fit spectrum
  * Feature  : Label - Identifies the fit feature (Vector{CharXRay} or other)
  * Reference: String - Name of reference against which :Spectrum was fit over :Feature
  * K : Float64 - The multiplicative fit constant
  * dK : Float64 - The 1σ uncertainty in :K
  * :Start : Int - Start index for fit channels (:roi ∈ columns)
  * :Stop : Int - Stop index for fit channels  (:roi ∈ columns)
  * :Peak : Float64 - Total counts in characteristic peak (:peakback ∈ columns)
  * :Back : Float64 - Total counts in background under the characteristic peak (:peakback ∈ columns)
  * :PtoB : Float64 - Peak-to-Background assuming 10 eV/channel (:peakback ∈ columns)
  * :KCalc : Float64 - Computed k-ratio assuming a composition. (Requires `material` argument to be specified.)
  * :KoKcalc : Float64 - Ratio of measured/computed k-ratio.  (Requires `material` argument to be specified.)
  * :LiveTime : Float64 - Acquisiton live time (s) (:dose ∈ columns)
  * :ProbeCurrent : Float64 - Probe current (nA) (:dose ∈ columns)
  * :DeadPct : Float64 - Dead time in ProbeCurrent (:dose ∈ columns)
  * :Counts : Float64 - Total counts in characteristic peak (:counts ∈ columns)
  * :RefCountsPernAs : Float64 - Estimated counts in :Reference in :Feature per unit dose.  (:counts ∈ columns)
  * :CountsPernAs : Float64 - Estimated counts in :Spectrum in :Feature per unit dose.  (:counts ∈ columns)
"""
function NeXLUncertainties.asa(
    ::Type{DataFrame},
    ffr::FilterFitResult;
    charOnly::Bool = true,
    material::Union{Material,Nothing} = nothing,
    columns::Tuple = (), # ( :roi, :peakback, :counts, :dose)
    mc = XPP, fc=ReedFluorescence,
)::DataFrame
    sl = charOnly ? #
        filter(lbl->lbl isa CharXRayLabel, NeXLUncertainties.sortedlabels(ffr.kratios)) : #
        NeXLUncertainties.sortedlabels(ffr.kratios)
    unkspec = ffr.label.spectrum
    # Build the base result DataFrame
    res = DataFrame(
        Spectrum = [ ffr.label for _ in sl ],
        Feature = sl,
        Reference = [properties(lbl)[:Name] for lbl in sl],
        K = [ NeXLUncertainties.value(ffr.kratios, lbl) for lbl in sl],
        dK = [ NeXLUncertainties.σ(ffr.kratios, lbl) for lbl in sl]
    )
    # Add optional columns
    if :roi ∈ columns
        insertcols!(res, 4, :Start => [ lbl.roi.start for lbl in sl])
        insertcols!(res, 5, :Stop => [ lbl.roi.stop for lbl in sl])
    end
    if :peakback ∈ columns
        insertcols!(res, :Peak => [ (ffr.peakback[lbl][1]-ffr.peakback[lbl][2]) for lbl in sl] )
        insertcols!(res, :Back => [ ffr.peakback[lbl][2] for lbl in sl] )
        insertcols!(res, :PtoB => [ peaktobackground(ffr, lbl) for lbl in sl] )
    end
    if material isa Material
        function kcalc(lbl)
            stdspec = lbl.spectrum
            zc = zafcorrection(
                mc,
                fc,
                NullCoating,
                material,
                stdspec[:Composition],
                lbl.xrays,
                unkspec[:BeamEnergy],
            )
            k(zc..., unkspec[:TakeOffAngle], stdspec[:TakeOffAngle])
        end
        insertcols!(res, :KCalc => kcalc.(sl))
        insertcols!(res, :KoKcalc => [ r[:K]/r[:KCalc] for r in eachrow(res) ])
    end
    if :dose ∈ columns
        dt(spec) = 100.0*(spec[:RealTime]-spec[:LiveTime])/spec[:RealTime]
        insertcols!(res, 2, :LiveTime => [ get(unkspec, :LiveTime, missing) for _ in sl ])
        insertcols!(res, 3, :ProbeCurrent => [ get(unkspec, :ProbeCurrent, missing) for _ in sl ])
        insertcols!(res, 4, :DeadPct => [ dt(unkspec) for lbl in sl ])

    end
    if :counts ∈ columns
        insertcols!(res, :Counts => [ ffr.peakback[lbl][1]-ffr.peakback[lbl][2] for lbl in sl] )
        insertcols!(res, :RefCountsPernAs => [ ffr.peakback[lbl][3] for lbl in sl] )
        insertcols!(res, :CountsPernAs => [ (ffr.peakback[lbl][1]-ffr.peakback[lbl][2]) / dose(unkspec) for lbl in sl ])
    end
    return res
end

"""
    filteredresidual(fit::FilterFitResult, unk::FilteredUnknown, ffs::AbstractVector{FilteredReference})::Vector{Float64}

Computes the difference between the best fit and the unknown filtered spectral data.
"""
function filteredresidual(
    fit::FilterFitResult{T},
    unk::FilteredUnknown,
    ffs::AbstractVector{FilteredReference{T}},
) where { T <: AbstractFloat }
    return unk.filtered - mapreduce(+, ffs) do ff
        (value(fit, ff.label) * (ff.scale / unk.scale)) * extract(ff, unk.ffroi)
    end
end

"""
    NeXLMatrixCorrection.estimatecoating(fr::FitResult, substrate::Material, coating::Material, cxr::CharXRay, mc::Type{<:MatrixCorrection}=XPP)::Film

Estimate the mass-thickness of `coating` on bulk `substrate` using the X-ray `cxr` for which there is KRatio in `fr`.  The result is
assigned to the properties of `fr` so that it can be account for in the matrix correction process.

This method is intended for use on standards where the substrate composition is known a priori.
"""
function NeXLMatrixCorrection.estimatecoating(fr::FitResult, substrate::Material, coating::Material, cxr::CharXRay, mc::Type{<:MatrixCorrection}=XPP)::Film
    properties(fr.label.spectrum)[:Coating] = estimatecoating(substrate, coating, kratio(fr, cxr), mc)
end

max_stars_repo_name: NicholasWMRitchie/NeXLSpectrum

max_issues_repo_name: NicholasWMRitchie/NeXLSpectrum

max_forks_repo_name: NicholasWMRitchie/NeXLSpectrum


text:
# TODO finish this
# function add_Heerden_genes_mods(metnet::MetNet; verbose = false)

#     verbose && print("Modifying reations")
#     for (Hd_id, model_ids) in to_inactivate_map
#         # Closing formate tranportation
#         if Hd_id == "Formate transporter"
#             METS.set_L!(0.0, Hd_mets_map["FORM"])
#             METS.set_V!(0.0, Hd_mets_map["FORM"])
#             verbose && print(Hd_id, "blocked")
#         else
#             # Blocking inactivated reactions
#             for rxn_id in model_ids
#                 RXNS.set_pub!(0.0, rxn_id)
#                 RXNS.set_nub!(0.0, rxn_id)
#                 verbose && println(rxn_id, "blocked")
#             end
#         end
#     end

# end