text:
using Random
using LinearAlgebra

import CommutativeRings: GFImpl
const Polynomial = CommutativeRings.Polynomial

let rng = MersenneTwister(1), x
mat(p::Integer, n::Integer) = rand(rng, 0:p-1, n, n)

@testset "Galois Fields Implementation" begin

    @test GFImpl(7) <: ZZmod{7}
    @test GFImpl(3) == GFImpl(3,1)
    @test GFImpl(5,3) <: Quotient{<:UnivariatePolynomial{<:ZZmod{5},:α}}

    G7 = GFImpl(7)
    @test G7(3)^2 == G7(2)

    G53 = GFImpl(5,3)
    @test isone(inv(G53([1,0,1])) * G53([1,0,1]))

    x = monom(basetype(G53), 1)
    ma53(p, n, k) = Matrix{G53}(mat(p, n) .* x^k)
    A = ma53(4, 3, 0) + ma53(4, 3, 1) + ma53(4, 3, 2)
    @test inv(A) * A == I

    @test sprint(show, GFImpl(5,3)([1,2,3])) == "3°*α^2 + 2°*α + 1° mod(α^3 + 3°*α + 2°)"
end

@testset "Galois Fields" begin
    @test_throws ArgumentError GF(1)
    @test_throws ArgumentError GF(12)
    @test GF(32) == GF(2, 5)
    @test GF(625) == GF(5, 4)
end

@testset "Galois Fields $p^$r" for (p, r) in ((2,8), (7,2))
    @test GF(p) == GF(p,1)
    @test GF(p, r) <: GaloisField
    G = GF(p, r)
    @test Quotient(G) <: Quotient{<:UnivariatePolynomial{<:ZZmod{p},:α}}
    Q = Quotient(G)
    P = Polynomial(G)
    @test Q == basetype(G)
    @test P / modulus(G) == Q
    @test modulus(G) == modulus(Q)
    @test order(G) == order(Q) == p ^ r
    @test characteristic(G) == characteristic(Q) == p
    @test dimension(G) == dimension(Q) == r
    @test deg(modulus(G)) == r

    @test G[p^r-1] isa G
    @test G[0:p^r-1] isa Vector{G}

    g1, g2 = rand(rng, G, 2)
    while iszero(g1)
        g1 = rand(rng, G)
    end
    @test convert(G, convert(Q, g1)) == g1
    q1, q2 = Quotient.((g1, g2))
    @test G(q1) == g1
    @test G[1] == Q(1)
    @test G[p] != Q(p)
    @test q1 == Q(g1)
    @test g1 == G(q1)
    @test g1 == q1

    @test g1 * g2 == q1 * q2
    @test g1^17 == q1^17
    @test g1^-17 == q1^-17
    @test g1 + g2 == q1 + q2
    @test g1 - g2 == q1 - q2
    @test -g1 == -q1
    @test inv(g1) == inv(q1)

    @test iszero(G[0])
    @test isone(G[1])
    @test isunit(G[2])
    @test zero(G) == 0
    @test one(G) == 1
    @test G[10] / G[10] == 1
    @test one(G) * g1 == g1
    @test zero(G) * g1 == 0
    @test g1 * 1 == g1
    @test 1 * g2 == g2
    @test 2g1 == g1 + g1
    @test g2 ^ 2 == g2 * g2
    @test g1 + g1 == 2g1
    @test g1 + (-g2) == g1 - g2
    @test g1 * inv(g2) == g1 / g2
    @test (g1 - g1) - g2 == -g2
    @test_throws ArgumentError inv(G[0])
    @test_throws ArgumentError G[0] ^ -2

    @test sprint(show, g1) !== nothing
    @test sprint(show, q1) !== nothing

    @test G(one(ZZ/p)) === one(G)
    @test length(modulus(G).(G)) == length(G)

    @test value(g1) == value(q1)

    @test num_irreducibles(Polynomial(G), dimension(G)) < order(G)
    @test GF(p, r, nr=1) !== nothing
    @test_throws ArgumentError GF(p, r, nr = 10000000)
end

@testset "Galois Field Implementation - Homomorphisms" begin

    Z1 = GFImpl(5, 2)
    Z2 = GFImpl(5, 6)
    z1 = Z1(monom(basetype(Z1)))
    iso = homomorphism(Z1, Z2)
    z2 = iso(z1)
    @test z1 != z2
    @test iso(z1) == z2
    @test z1 + 1 == z1 + Z1(1)
    @test iso(Z1(0)) == Z2(0)
    @test iso(Z1(1)) == Z2(1)
    @test iso(z1^17 + 2z1^12 + 1) == z2^17 + 2z2^12 + 1

    G = GF(47)
    h = homomorphism(x->G(2x), Int, G)
    @test h(18) == G(36)
end

@testset "Galois field - irreducibles GF{$q,$r} ^ $s" for (q, r, s) in ((2, 3, 3), (3, 2, 3))
    G = GF(q, r)
    irr = irreducibles(G[:x], s)
    @test length(collect(irr)) == necklace(order(G), s)
    @test necklace(order(G), s) == num_irreducibles(G, s)
end

@testset "Galois field - user modulus" begin
    P = ZZ{Int}[:x]
    x = monom(P)
    p = x^3 + x + 1
    @test GF(5, mod = p) <: GaloisField
    @test_throws ArgumentError GF(11, mod = p) # p not irreducible in ZZ/11
end

@testset "isprimitive $G" for G in (GF(31), GF(2^3), GF(3^4))
    gG = (x for x in G)
    ord = order(G) - 1
    @test isprimitive.(gG) == ( order.(gG) .== ord )
    @test isfield(G)
end

end


text:
export rattle!, r_sum, r_dot

############################################################
###   Some useful utility functions
############################################################
###   Robust Summation and Dot Products
############################################################

"Robust summation. Uses `sum_kbn`."
r_sum(a) = sum_kbn(a)

## NOTE: if I see this correctly, then r_dot allocates a temporary
##       vector, which is likely quite a performance overhead.
##       probably, we want to re-implement this.
## TODO: new version without the intermediate allocation
"Robust inner product. Defined as `r_dot(a, b) = r_sum(a .* b)`"
r_dot(a, b) = r_sum(a[:] .* b[:])

"""
`rattle!(at::AbstractAtoms, r::Float64; rnn = 1.0)
  -> at`

randomly perturbs the atom positions

* `r`: magnitude of perturbation
* `rnn` : nearest-neighbour distance
"""
function rattle!(at::AbstractAtoms, r::Float64; rnn = 1.0)
   X = positions(at) |> mat
   X += r * rnn * 2.0/sqrt(3) * (rand(size(X)) - 0.5)
   return set_positions!(at, X)
end

ext: jl


text:
using ForwardDiff, SparseArrays
abstract type Hyperparameters end;

struct KernelHyperparameters <: Hyperparameters
    array
end

function kernel(sig::Number, l::Number, rcut::Number, ri::Vector, rj::Vector,
                ri1::Vector, rj1::Vector)
    rii1 = norm(ri .- ri1)
    rjj1 = norm(rj .- rj1)

    dij = abs(rii1-rjj1)

    if rcut < dij
        return 0.
    end
    return sig * exp(-(dij)^2/(2*l^2)) * (rcut - dij)^2
end

function kernel(sig::Number, l::Number, rcut::Number, ri::Vector, rj::Vector,
                ri1::Vector, rj1::Vector, ri2::Vector, rj2::Vector)

    rii1 = norm(ri .- ri1)
    rii2 = norm(ri .- ri2)
    ri1i2 = norm(ri1 .- ri2)

    rjj1 = norm(rj .- rj1)
    rjj2 = norm(rj .- rj2)
    rj1j2 = norm(rj1 .- rj2)

    dij = sqrt((rii1-rjj1)^2 + (rii2-rjj2)^2 + (ri1i2-rj1j2)^2)

    if rcut < dij
        return 0.
    end
    return sig * exp(-(dij)^2/(2*l^2)) * (rcut - dij)^2
end

function kgd(sig::Number, l::Number, rcut::Number, ri::Vector, rj::Vector, ri1::Vector, rj1::Vector)
    kgg(rri::Vector, rrj::Vector) = kernel(hyperparameters.array..., rcut, rri, rrj, ri1, rj1)
    kgd(rri::Vector, rrj::Vector) = ForwardDiff.gradient(x->kgg(x, rrj), rri)
    return kgd(ri, rj)
end

function kgd(sig::Number, l::Number, rcut::Number, ri::Vector, rj::Vector, ri1::Vector, rj1::Vector, ri2::Vector, rj2::Vector)
    kgg(rri::Vector, rrj::Vector) = kernel(hyperparameters.array..., rcut, rri, rrj, ri1, rj1, ri2, rj2)
    kgd(rri::Vector, rrj::Vector) = ForwardDiff.gradient(x->kgg(x, rrj), rri)
    return kgd(ri, rj)
end

function kgd(hyperparameters, rcut, ri::Vector, rj::Vector, ri1::Vector, ri2::Vector, rj1::Vector, rj2::Vector)
    kgg(rri::Vector, rrj::Vector) = kernel(hyperparameters.array..., rcut, rri, rrj, ri1, rj1, ri2, rj2)
    kgd(rri::Vector, rrj::Vector) = ForwardDiff.gradient(x->kgg(x, rrj), rri)
    return kgd(ri, rj)
end

function kdg(hyperparameters, rcut, ri::Vector, rj::Vector, ri1::Vector, ri2::Vector, rj1::Vector, rj2::Vector)
    kgg(rri::Vector, rrj::Vector) = kernel(hyperparameters.array..., rcut, rri, rrj, ri1, rj1, ri2, rj2)
    kgd(rri::Vector, rrj::Vector) = ForwardDiff.gradient(x->kgg(rri, x), rrj)
    return kgd(ri, rj)
end

function kdd(hyperparameters, rcut, ri::Vector, rj::Vector, ri1::Vector, ri2::Vector, rj1::Vector, rj2::Vector)
    kgg(rri::Vector, rrj::Vector) = kernel(hyperparameters.array..., rcut, rri, rrj, ri1, rj1, ri2, rj2)
    kgd(rri::Vector, rrj::Vector) = ForwardDiff.gradient(x->kgg(x, rrj), rri)
    kdd(rri::Vector, rrj::Vector) = ForwardDiff.jacobian(x->kgd(rri, x), rrj)
    return kdd(ri, rj)
end

function khd(hyperparameters, rcut, ri::Vector, rj::Vector, ri1::Vector, ri2::Vector, rj1::Vector, rj2::Vector)
    kgg(rri::Vector, rrj::Vector) = kernel(hyperparameters.array..., rcut, rri, rrj, ri1, rj1, ri2, rj2)
    kgd(rri::Vector, rrj::Vector) = ForwardDiff.gradient(x->kgg(x, rrj), rri)
    kdd(rri::Vector, rrj::Vector) = ForwardDiff.jacobian(x->kgd(rri, x), rrj)
    khd(rri::Vector, rrj::Vector) = ForwardDiff.jacobian(x->kdd(x, rrj), rri)
    return khd(ri, rj)
end

function kgd(descriptor::Taps.TwoThreeBody, db::Taps.TwoThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    K2 = kgd(descriptor, db.threebody, hyperparameters, rcut)
    K3 = kgd(descriptor, db.threebody, hyperparameters, rcut)
    return K2 + K3
end

function kdg(descriptor::Taps.TwoThreeBody, db::Taps.TwoThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    K2 = kdg(descriptor, db.threebody, hyperparameters, rcut)
    K3 = kdg(descriptor, db.threebody, hyperparameters, rcut)
    return K2 + K3
end

function kdd(descriptor::Taps.TwoThreeBody, db::Taps.TwoThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    K2 = kdd(descriptor, db.threebody, hyperparameters, rcut)
    K3 = kdd(descriptor, db.threebody, hyperparameters, rcut)
    return K2 + K3
end

function khd(descriptor::Taps.TwoThreeBody, db::Taps.TwoThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    K2 = khd(descriptor, db.threebody, hyperparameters, rcut)
    K3 = khd(descriptor, db.threebody, hyperparameters, rcut)
    return K2 + K3
end

function kgg(sparsetwothreebody::Taps.SparseAtomic{T, NN}, db::Taps.TwoThreeBodyDB,
             hyperparameters::KernelHyperparameters, rcut) where {T<:Taps.TwoThreeBody, NN}
    coords = sparsetwothreebody.coords
    N, A = size(coords)

    K = zeros(N*A, db.m)

    count = 1
    for i=1:N
        for a=1:A
            K[count, :] = kgg(coords[i, a], db, hyperparamters, rcut)
            count += 1
        end
    end
    return K
end

function kgg(descriptor::Taps.TwoThreeBody, db::Taps.TwoThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    K2 = kgg(descriptor, db.threebody, hyperparameters, rcut)
    K3 = kgg(descriptor, db.threebody, hyperparameters, rcut)
    return K2 + K3
end

function kgg(descriptor::Taps.TwoThreeBody, db::Taps.TwoBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    Ker = spzeros(maxidx)
    for cluster in chart

        permutations = dscpt[cluster]["permutations"]
        K = spzeros(maxidx)
        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]

                for permute in permutations
                    rj, rj1 = [rrj, rrj1][permute]
                    K[idx] += kernel(hyperparameters.array..., rcut, ri, rj, ri1, rj1)
                end
            end
        end
        Ker .+= K
    end
    return Ker
end

function kgg(descriptor::Taps.TwoThreeBody, db::Taps.ThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    Ker = spzeros(maxidx)
    for cluster in chart

        permutations = dscpt[cluster]["permutations"]
        K = spzeros(maxidx)
        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]
            ri2 = [dscpt[cluster]["xi2"][ci], dscpt[cluster]["yi2"][ci], dscpt[cluster]["zi2"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]
                rrj2 = [db.xi2[cluster][cj], db.yi2[cluster][cj], db.zi2[cluster][cj]]

                for permute in permutations
                    rj, rj1, rj2 = [rrj, rrj1, rrj2][permute]
                    K[idx] += kernel(hyperparameters.array..., rcut, ri, rj, ri1, rj1, ri2, rj2)
                end
            end
        end
        Ker .+= K
    end
    return Ker
end

function kgd(descriptor::Taps.TwoThreeBody, db::Taps.TwoBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    Kerx = spzeros(maxidx)
    Kery = spzeros(maxidx)
    Kerz = spzeros(maxidx)
    for cluster in chart

        permutations = dscpt[cluster]["permutations"]
        Kx = spzeros(maxidx)
        Ky = spzeros(maxidx)
        Kz = spzeros(maxidx)
        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]

                for permute in permutations
                    rj, rj1 = [rrj, rrj1][permute]
                    dK = kgd(hyperparameters, rcut, ri, rj, ri1, rj1)
                    Kx[idx] += dK[1]
                    Ky[idx] += dK[2]
                    Kz[idx] += dK[3]
                end
            end
        end
        Kerx .+= Kx
        Kery .+= Ky
        Kerz .+= Kz
    end
    return [Kerx, Kery, Kerz]
end

function kgd(descriptor::Taps.TwoThreeBody, db::Taps.ThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    Kerx = spzeros(maxidx)
    Kery = spzeros(maxidx)
    Kerz = spzeros(maxidx)

    for cluster in chart

        permutations = dscpt[cluster]["permutations"]
        Kx = spzeros(maxidx)
        Ky = spzeros(maxidx)
        Kz = spzeros(maxidx)
        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]
            ri2 = [dscpt[cluster]["xi2"][ci], dscpt[cluster]["yi2"][ci], dscpt[cluster]["zi2"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]
                rrj2 = [db.xi2[cluster][cj], db.yi2[cluster][cj], db.zi2[cluster][cj]]

                for permute in permutations
                    rj, rj1, rj2 = [rrj, rrj1, rrj2][permute]
                    dK = kgd(hyperparameters, rcut, ri, rj, ri1, rj1, ri2, rj2)
                    Kx[idx] += dK[1]
                    Ky[idx] += dK[2]
                    Kz[idx] += dK[3]
                end
            end
        end
        Kerx .+= Kx
        Kery .+= Ky
        Kerz .+= Kz
    end
    return [Kerx, Kery, Kerz]
end

function kdg(descriptor::Taps.TwoThreeBody, db::Taps.TwoBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    Kerx = spzeros(maxidx)
    Kery = spzeros(maxidx)
    Kerz = spzeros(maxidx)
    for cluster in chart

        permutations = dscpt[cluster]["permutations"]
        Kx = spzeros(maxidx)
        Ky = spzeros(maxidx)
        Kz = spzeros(maxidx)
        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]

                for permute in permutations
                    rj, rj1 = [rrj, rrj1][permute]
                    dK = kgd(hyperparameters, rcut, ri, rj, ri1, rj1)
                    Kx[idx] += dK[1]
                    Ky[idx] += dK[2]
                    Kz[idx] += dK[3]
                end
            end
        end
        Kerx .+= Kx
        Kery .+= Ky
        Kerz .+= Kz
    end
    return [Kerx, Kery, Kerz]
end

function kdg(descriptor::Taps.TwoThreeBody, db::Taps.ThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    Kerx = spzeros(maxidx)
    Kery = spzeros(maxidx)
    Kerz = spzeros(maxidx)

    for cluster in chart

        permutations = dscpt[cluster]["permutations"]
        Kx = spzeros(maxidx)
        Ky = spzeros(maxidx)
        Kz = spzeros(maxidx)
        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]
            ri2 = [dscpt[cluster]["xi2"][ci], dscpt[cluster]["yi2"][ci], dscpt[cluster]["zi2"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]
                rrj2 = [db.xi2[cluster][cj], db.yi2[cluster][cj], db.zi2[cluster][cj]]

                for permute in permutations
                    rj, rj1, rj2 = [rrj, rrj1, rrj2][permute]
                    dK = kdg(hyperparameters, rcut, ri, rj, ri1, rj1, ri2, rj2)
                    Kx[idx] += dK[1]
                    Ky[idx] += dK[2]
                    Kz[idx] += dK[3]
                end
            end
        end
        Kerx .+= Kx
        Kery .+= Ky
        Kerz .+= Kz
    end
    return [Kerx, Kery, Kerz]
end

function kdd(descriptor::Taps.TwoThreeBody, db::Taps.TwoBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    Kerxx = spzeros(maxidx)
    Kerxy = spzeros(maxidx)
    Kerxz = spzeros(maxidx)

    Keryy = spzeros(maxidx)
    Keryx = spzeros(maxidx)
    Keryz = spzeros(maxidx)

    Kerzx = spzeros(maxidx)
    Kerzy = spzeros(maxidx)
    Kerzz = spzeros(maxidx)
    for cluster in chart
        permutations = dscpt[cluster]["permutations"]
        Kxx = spzeros(maxidx)
        Kxy = spzeros(maxidx)
        Kxz = spzeros(maxidx)

        Kyy = spzeros(maxidx)
        Kyx = spzeros(maxidx)
        Kyz = spzeros(maxidx)

        Kzx = spzeros(maxidx)
        Kzy = spzeros(maxidx)
        Kzz = spzeros(maxidx)
        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]

                for permute in permutations
                    rj, rj1 = [rrj, rrj1][permute]
                    dK = kdd(hyperparameters, rcut, ri, rj, ri1, rj1)
                    Kxx[idx] += dK[1, 1]
                    Kxy[idx] += dK[1, 2]
                    Kxz[idx] += dK[1, 3]

                    Kyx[idx] += dK[2, 1]
                    Kyy[idx] += dK[2, 2]
                    Kyz[idx] += dK[2, 3]

                    Kzx[idx] += dK[3, 1]
                    Kzy[idx] += dK[3, 2]
                    Kzz[idx] += dK[3, 3]
                end
            end
        end
        Kerxx .+= Kxx
        Kerxy .+= Kxy
        Kerxz .+= Kxz

        Keryx .+= Kyx
        Keryy .+= Kyy
        Keryz .+= Kyz

        Kerzx .+= Kzx
        Kerzy .+= Kzy
        Kerzz .+= Kzz
    end
    return [Kerxx, Kerxy, Kerxz, Keryx, Keryy, Keryz, Kerzx, Kerzy, Kerzz]
end

function kdd(descriptor::Taps.TwoThreeBody, db::Taps.ThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    Kerxx = spzeros(maxidx)
    Kerxy = spzeros(maxidx)
    Kerxz = spzeros(maxidx)

    Keryy = spzeros(maxidx)
    Keryx = spzeros(maxidx)
    Keryz = spzeros(maxidx)

    Kerzx = spzeros(maxidx)
    Kerzy = spzeros(maxidx)
    Kerzz = spzeros(maxidx)

    for cluster in chart

        permutations = dscpt[cluster]["permutations"]
        Kxx = spzeros(maxidx)
        Kxy = spzeros(maxidx)
        Kxz = spzeros(maxidx)

        Kyy = spzeros(maxidx)
        Kyx = spzeros(maxidx)
        Kyz = spzeros(maxidx)

        Kzx = spzeros(maxidx)
        Kzy = spzeros(maxidx)
        Kzz = spzeros(maxidx)

        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]
            ri2 = [dscpt[cluster]["xi2"][ci], dscpt[cluster]["yi2"][ci], dscpt[cluster]["zi2"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]
                rrj2 = [db.xi2[cluster][cj], db.yi2[cluster][cj], db.zi2[cluster][cj]]

                for permute in permutations
                    rj, rj1, rj2 = [rrj, rrj1, rrj2][permute]
                    dK = kdd(hyperparameters, rcut, ri, rj, ri1, rj1, ri2, rj2)

                    Kxx[idx] += dK[1, 1]
                    Kxy[idx] += dK[1, 2]
                    Kxz[idx] += dK[1, 3]

                    Kyx[idx] += dK[2, 1]
                    Kyy[idx] += dK[2, 2]
                    Kyz[idx] += dK[2, 3]

                    Kzx[idx] += dK[3, 1]
                    Kzy[idx] += dK[3, 2]
                    Kzz[idx] += dK[3, 3]
                end
            end
        end
        Kerxx .+= Kxx
        Kerxy .+= Kxy
        Kerxz .+= Kxz

        Keryx .+= Kyx
        Keryy .+= Kyy
        Keryz .+= Kyz

        Kerzx .+= Kzx
        Kerzy .+= Kzy
        Kerzz .+= Kzz
    end
    return [Kerxx, Kerxy, Kerxz, Keryx, Keryy, Keryz, Kerzx, Kerzy, Kerzz]
end

function khd(descriptor::Taps.TwoThreeBody, db::Taps.TwoBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    # 3xJxK
    Kerijk = [spzeros(maxidx) for ii=1:27]

    for cluster in chart
        permutations = dscpt[cluster]["permutations"]
        Kijk = [spzeros(maxidx) for ii=1:27]

        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]

                for permute in permutations
                    rj, rj1 = [rrj, rrj1][permute]
                    dK = kdd(hyperparameters, rcut, ri, rj, ri1, rj1)

                    for ii=1:9
                        Kijk[ii][idx] += dK[ii, 1]
                        Kijk[ii+9][idx] += dK[ii, 2]
                        Kijk[ii+18][idx] += dK[ii, 3]
                    end
                end
            end
        end
        for ii=1:27
            Kerijk[ii] .+= Kijk[ii]
        end
    end
    return Kerijk
end

function khd(descriptor::Taps.TwoThreeBody, db::Taps.ThreeBodyDB, hyperparameters::KernelHyperparameters, rcut)
    dscpt = descriptor.descriptor
    chart = keys(db.idx)

    maxidx = 0
    for cluster in chart
        length(db.idx[cluster]) == 0 ? continue : nothing
        maxidx = max(maxidx, db.idx[cluster][end])
    end

    # 3xJxK
    Kerijk = [spzeros(maxidx) for ii=1:27]

    for cluster in chart

        permutations = dscpt[cluster]["permutations"]
        Kijk = [spzeros(maxidx) for i=1:27]

        for ci=1:length(dscpt[cluster]["xi"])
            ri = [dscpt[cluster]["xi"][ci], dscpt[cluster]["yi"][ci], dscpt[cluster]["zi"][ci]]
            ri1 = [dscpt[cluster]["xi1"][ci], dscpt[cluster]["yi1"][ci], dscpt[cluster]["zi1"][ci]]
            ri2 = [dscpt[cluster]["xi2"][ci], dscpt[cluster]["yi2"][ci], dscpt[cluster]["zi2"][ci]]

            for cj=1:length(db.idx[cluster])
                idx = db.idx[cluster][cj]

                rrj = [db.xi[cluster][cj], db.yi[cluster][cj], db.zi[cluster][cj]]
                rrj1 = [db.xi1[cluster][cj], db.yi1[cluster][cj], db.zi1[cluster][cj]]
                rrj2 = [db.xi2[cluster][cj], db.yi2[cluster][cj], db.zi2[cluster][cj]]

                for permute in permutations
                    rj, rj1, rj2 = [rrj, rrj1, rrj2][permute]
                    dK = khd(hyperparameters, rcut, ri, rj, ri1, rj1, ri2, rj2)

                    for ii=1:9
                        Kijk[ii][idx] += dK[ii, 1]
                        Kijk[ii+9][idx] += dK[ii, 2]
                        Kijk[ii+18][idx] += dK[ii, 3]
                    end
                end
            end
        end
        for ii=1:27
            Kerijk[ii] .+= Kijk[ii]
        end
    end
    return Kerijk
end

max_stars_repo_name: schinavro/taps

max_issues_repo_name: schinavro/taps

max_forks_repo_name: schinavro/taps


text:
using Test

@testset "Anderson accel. ($R)" for R in [Float32, Float64]

using LinearAlgebra
using ProximalAlgorithms: AndersonAcceleration, update!

n, k = 20, 10

Q = randn(R, n, k)
H = Q*Q' + I
l = randn(R, n)

f(x) = dot(x, H*x)/2 + dot(x, l)

x_star = -H\l
f_star = f(x_star)
norm_x_star = norm(x_star)

Lip = opnorm(H)
gamma = 1/Lip
x = zeros(n)
res = zeros(n)
x_prev = zeros(n)
res_prev = zeros(n)

acc = AndersonAcceleration(x, 5)

mul!(res, H, x)
res .+= l
res .*= gamma

d = acc * res

for it in 1:30
    # store iterate and residual for the operator update later
    res_prev .= res
    x_prev .= x

    # compute accelerated direction
    mul!(d, acc, res)

    # update iterate
    x .-= d

    # compute new residual
    mul!(res, H, x)
    res .+= l
    res .*= gamma

    # update operator
    update!(acc, x - x_prev, res - res_prev)
end

@test f(x) <= f_star + (1 + abs(f_star)) * 1e-6

end


text:
#####
##### include this separately for interactive tests
#####

####
#### packages and symbols
####

using DynamicHMC, Test, ArgCheck, DocStringExtensions, HypothesisTests, LinearAlgebra,
    MCMCDiagnostics, Parameters, Random, StatsBase, Statistics

import ForwardDiff, Random, TransformVariables

using DynamicHMC:
    # trees
    Directions, next_direction, biased_progressive_logprob2, adjacent_tree,
    sample_trajectory, InvalidTree,
    # Hamiltonian
    GaussianKineticEnergy, kinetic_energy, ∇kinetic_energy, rand_p, Hamiltonian,
    EvaluatedLogDensity, evaluate_ℓ, PhasePoint, logdensity, leapfrog, calculate_p♯,
    logdensity,
    # NUTS
    TrajectoryNUTS, rand_bool_logprob, GeneralizedTurnStatistic, AcceptanceStatistic,
    leaf_acceptance_statistic, acceptance_rate, TreeStatisticsNUTS, NUTS, sample_tree,
    # stepsize
    find_crossing_stepsize, bisect_stepsize, find_initial_stepsize, InitialStepsizeSearch,
    DualAveraging, initial_adaptation_state, adapt_stepsize, current_ϵ, final_ϵ,
    FixedStepsize,
    # mcmc
    mcmc_steps, mcmc_next_step, mcmc_keep_warmup, position_matrix, WarmupState

using DynamicHMC: logaddexp # FIXME remove when depending on LogExpFunctions

import DynamicHMC:
    # trees
    move, is_turning, combine_turn_statistics, is_divergent,
    combine_visited_statistics, calculate_logprob2, combine_proposals, leaf

using DynamicHMC.Diagnostics
using DynamicHMC.Diagnostics: ACCEPTANCE_QUANTILES

using LogDensityProblems: logdensity_and_gradient, dimension, LogDensityProblems

### uncomment code below to use latest LogDensityTestSuite
# if !isinteractive()             # on CI
#     @info "installing LogDensityTestSuite#master"
#     import Pkg
#     Pkg.API.add(Pkg.PackageSpec(; name = "LogDensityTestSuite", rev = "master"))
# end
using LogDensityTestSuite

####
#### utilities for testing
####

###
### general test environment
###

const RNG = Random.GLOBAL_RNG   # shorthand
Random.seed!(RNG, UInt32[0x23ef614d, 0x8332e05c, 0x3c574111, 0x121aa2f4])

"Tolerant testing in a CI environment."
const RELAX = (k = "CONTINUOUS_INTEGRATION"; haskey(ENV, k) && ENV[k] == "true")

###
### random values
###

"""
$(SIGNATURES)

Random positive definite matrix of size `n` x `n` (for testing).
"""
function rand_Σ(::Type{Symmetric}, n)
    A = randn(n, n)
    Symmetric(A'*A .+ 0.01)
end

rand_Σ(::Type{Diagonal}, n) = Diagonal(randn(n).^2 .+ 0.01)

rand_Σ(n::Int) = rand_Σ(Symmetric, n)

"""
$(SIGNATURES)

Simulated mean and covariance of `N` values from `f()`.
"""
function simulated_meancov(f, N)
    s = f()
    K = length(s)
    x = similar(s, (N, K))
    for i in 1:N
        x[i, :] = f()
    end
    m, C = mean_and_cov(x, 1)
    vec(m), C
end

@testset "simulated meancov" begin
    μ = [2, 1.2]
    D = [2.0, 0.7]
    m, C = simulated_meancov(()-> randn(2) .* D .+ μ, 10000)
    @test m ≈ μ atol = 0.05 rtol = 0.1
    @test C ≈ Diagonal(abs2.(D)) atol = 0.05 rtol = 0.1
end

###
### Multivariate normal ℓ for testing
###

"Random (positive) diagonal matrix."
rand_D(K) = Diagonal(abs.(randn(K)))

"Random Cholesky factor for correlation matrix."
function rand_C(K)
    t = TransformVariables.CorrCholeskyFactor(K)
    t(randn(TransformVariables.dimension(t)))
end

"Multivariate normal with `Σ = LL'`."
multivariate_normal(μ, L) = shift(μ, linear(L, StandardMultivariateNormal(length(μ))))

"Multivariate normal with diagonal `Σ` (constant `v` variance)."
multivariate_normal(μ, v::Real = 1) = multivariate_normal(μ, Diagonal(fill(v, length(μ))))

###
### Hamiltonian test helper functions
###

"""
$(SIGNATURES)

Return a reasonable estimate for the largest stable stepsize (which may not be
stable, but is a good starting point for finding that).

`q` is assumed to be normal with variance `Σ`. `κ` is the kinetic energy.

Using the transformation ``p̃ = W⁻¹ p``, the kinetic energy is

``p'M⁻¹p = p'W⁻¹'W⁻¹p/2=p̃'p̃/2``

Transforming to ``q̃=W'q``, the variance of which becomes ``W' Σ W``. Return the
square root of its smallest eigenvalue, following Neal (2011, p 136).

When ``Σ⁻¹=M=WW'``, this the variance of `q̃` is ``W' Σ W=W' W'⁻¹W⁻¹W=I``, and
thus decorrelates the density perfectly.
"""
find_stable_ϵ(κ::GaussianKineticEnergy, Σ) = √eigmin(κ.W'*Σ*κ.W)

"""
$(SIGNATURES)

A `NamedTuple` that contains

- a random `K`-element vector `μ`

- a random `K × K` covariance matrix `Σ`,

- a random Hamiltonian `H` with `ℓ` corresponding to a multivariate normal with `μ`, `Σ`,
  and a random Gaussian kinetic energy (unrelated to `ℓ`)

- a random phasepoint `z`.

Useful for testing.
"""
function rand_Hz(K)
    μ = randn(K)
    Σ = rand_Σ(K)
    L = cholesky(Σ).L
    κ = GaussianKineticEnergy(inv(rand_Σ(Diagonal, K)))
    H = Hamiltonian(κ, multivariate_normal(μ, L))
    z = PhasePoint(evaluate_ℓ(H.ℓ, μ .+ L * randn(K)), rand_p(RNG, κ))
    (μ = μ, Σ = Σ, H = H, z = z)
end


text:
using XGPaint
using Healpix
using HDF5
import ThreadsX
using JLD2

# ENV["SCRATCH"] = "/home/zequnl/scratch/"
@time halo_pos, halo_mass = read_halo_catalog_hdf5(
    joinpath(ENV["SCRATCH"],"websky_halos-light.hdf5"));

## Load halos from HDF5 files, establish a CIB model and cosmology
cosmo = get_cosmology(Float32, h = 0.7, OmegaM = 0.25)
radio_model = Radio_Sehgal2009{Float32}(a_0 = -1)
@time begin
    sources = generate_sources(radio_model, cosmo, halo_pos, halo_mass);
end;

@save joinpath(ENV["SCRATCH"], "CCAT/before_match/sources.h5") sources

freqs = ["93", "145", "220", "280", "350", "405", "860"]

m = Map{Float64,RingOrder}(radio_model.nside)

function generate_maps()
    scratch_dir = ENV["SCRATCH"]
    println("SCRATCH: ", scratch_dir)
    mkpath(joinpath(scratch_dir, "CCAT/before_match/"))
    mkpath(joinpath(scratch_dir, "CCAT/before_match/maps"))

    for freq in freqs
        @time begin

            flux_I, redshift_I, θ_I, ϕ_I, flux_II, redshift_II, θ_II, ϕ_II = paint!(
                m, parse(Float32, freq) * 1.0f9, radio_model, sources,
                return_fluxes=true)

            flux = vcat(flux_I, flux_II)
            redshift = vcat(redshift_I, redshift_II)
            theta = vcat(θ_I, θ_II)
            phi = vcat(ϕ_I, ϕ_II)

            h5open(joinpath(scratch_dir, "CCAT/before_match/catalog_$(freq).h5"), "w") do file
                write(file, "flux", flux)
                write(file, "redshift", redshift)
                write(file, "theta", theta)
                write(file, "phi", phi)
            end

            map_path = joinpath(scratch_dir, "CCAT/before_match/maps/radio$(freq).fits")
            Healpix.saveToFITS(m, "!$(map_path)")
        end
    end
end

generate_maps()
##

max_stars_repo_name: xzackli/XGPaint

max_issues_repo_name: xzackli/XGPaint

max_forks_repo_name: xzackli/XGPaint


text:
const _series_recipe_deps = Dict()

function series_recipe_dependencies(st::Symbol, deps::Symbol...)
    _series_recipe_deps[st] = deps
end

function seriestype_supported(st::Symbol)
    seriestype_supported(backend(), st)
end

# returns :no, :native, or :recipe depending on how it's supported
function seriestype_supported(pkg::AbstractBackend, st::Symbol)
    # is it natively supported
    if is_seriestype_supported(pkg, st)
        return :native
    end

    haskey(_series_recipe_deps, st) || return :no

    supported = true
    for dep in _series_recipe_deps[st]
        if seriestype_supported(pkg, dep) == :no
            supported = false
        end
    end
    supported ? :recipe : :no
end

macro deps(st, args...)
    :(PlotsBase.series_recipe_dependencies($(quot(st)), $(map(quot, args)...)))
end

# get a list of all seriestypes
function all_seriestypes()
    sts = Set{Symbol}(keys(_series_recipe_deps))
    for bsym in backends()
        btype = _backendType[bsym]
        sts = union(sts, Set{Symbol}(supported_seriestypes(btype())))
    end
    sort(collect(sts))
end

# ----------------------------------------------------------------------------------

num_series(x::AMat) = size(x,2)
num_series(x) = 1

RecipesBase.apply_recipe{T}(d::KW, ::Type{T}, plt::AbstractPlot) = throw(MethodError("Unmatched plot recipe: $T"))

# ---------------------------------------------------------------------------

# for seriestype `line`, need to sort by x values
@recipe function f(::Type{Val{:line}}, x, y, z)
    indices = sortperm(x)
    x := x[indices]
    y := y[indices]
    if typeof(z) <: AVec
        z := z[indices]
    end
    seriestype := :path
    ()
end
@deps line path

function hvline_limits(axis::Axis)
    vmin, vmax = axis_limits(axis)
    if vmin >= vmax
        if isfinite(vmin)
            vmax = vmin + 1
        else
            vmin, vmax = 0.0, 1.1
        end
    end
    vmin, vmax
end

@recipe function f(::Type{Val{:hline}}, x, y, z)
    xmin, xmax = hvline_limits(d[:subplot][:xaxis])
    n = length(y)
    newx = repmat(Float64[xmin, xmax, NaN], n)
    newy = vec(Float64[yi for i=1:3,yi=y])
    x := newx
    y := newy
    seriestype := :path
    ()
end
@deps hline path

@recipe function f(::Type{Val{:vline}}, x, y, z)
    ymin, ymax = hvline_limits(d[:subplot][:yaxis])
    n = length(y)
    newx = vec(Float64[yi for i=1:3,yi=y])
    newy = repmat(Float64[ymin, ymax, NaN], n)
    x := newx
    y := newy
    seriestype := :path
    ()
end
@deps vline path

# ---------------------------------------------------------------------------
# steps

function make_steps(x, y, st)
    n = length(x)
    n == 0 && return zeros(0),zeros(0)
    newx, newy = zeros(2n-1), zeros(2n-1)
    for i=1:n
        idx = 2i-1
        newx[idx] = x[i]
        newy[idx] = y[i]
        if i > 1
            newx[idx-1] = x[st == :steppre ? i-1 : i]
            newy[idx-1] = y[st == :steppre ? i   : i-1]
        end
    end
    newx, newy
end

# create a path from steps
@recipe function f(::Type{Val{:steppre}}, x, y, z)
    d[:x], d[:y] = make_steps(x, y, :steppre)
    seriestype := :path

    # create a secondary series for the markers
    if d[:markershape] != :none
        @series begin
            seriestype := :scatter
            x := x
            y := y
            label := ""
            primary := false
            ()
        end
        markershape := :none
    end
    ()
end
@deps steppre path scatter

# create a path from steps
@recipe function f(::Type{Val{:steppost}}, x, y, z)
    d[:x], d[:y] = make_steps(x, y, :steppost)
    seriestype := :path

    # create a secondary series for the markers
    if d[:markershape] != :none
        @series begin
            seriestype := :scatter
            x := x
            y := y
            label := ""
            primary := false
            ()
        end
        markershape := :none
    end
    ()
end
@deps steppost path scatter

# ---------------------------------------------------------------------------
# sticks

# create vertical line segments from fill
@recipe function f(::Type{Val{:sticks}}, x, y, z)
    n = length(x)
    fr = d[:fillrange]
    if fr == nothing
        yaxis = d[:subplot][:yaxis]
        fr = if yaxis[:scale] == :identity
            0.0
        else
            NaNMath.min(axis_limits(yaxis)[1], ignorenan_minimum(y))
        end
    end
    newx, newy = zeros(3n), zeros(3n)
    for i=1:n
        rng = 3i-2:3i
        newx[rng] = [x[i], x[i], NaN]
        newy[rng] = [_cycle(fr,i), y[i], NaN]
    end
    x := newx
    y := newy
    fillrange := nothing
    seriestype := :path

    # create a secondary series for the markers
    if d[:markershape] != :none
        @series begin
            seriestype := :scatter
            x := x
            y := y
            label := ""
            primary := false
            ()
        end
        markershape := :none
    end
    ()
end
@deps sticks path scatter

# ---------------------------------------------------------------------------
# bezier curves

# get the value of the curve point at position t
function bezier_value(pts::AVec, t::Real)
    val = 0.0
    n = length(pts)-1
    for (i,p) in enumerate(pts)
        val += p * binomial(n, i-1) * (1-t)^(n-i+1) * t^(i-1)
    end
    val
end

# create segmented bezier curves in place of line segments
@recipe function f(::Type{Val{:curves}}, x, y, z; npoints = 30)
    args = z != nothing ? (x,y,z) : (x,y)
    newx, newy = zeros(0), zeros(0)
    fr = d[:fillrange]
    newfr = fr != nothing ? zeros(0) : nothing
    newz = z != nothing ? zeros(0) : nothing
    # lz = d[:line_z]
    # newlz = lz != nothing ? zeros(0) : nothing

    # for each line segment (point series with no NaNs), convert it into a bezier curve
    # where the points are the control points of the curve
    for rng in iter_segments(args...)
        length(rng) < 2 && continue
        ts = linspace(0, 1, npoints)
        nanappend!(newx, map(t -> bezier_value(_cycle(x,rng), t), ts))
        nanappend!(newy, map(t -> bezier_value(_cycle(y,rng), t), ts))
        if z != nothing
            nanappend!(newz, map(t -> bezier_value(_cycle(z,rng), t), ts))
        end
        if fr != nothing
            nanappend!(newfr, map(t -> bezier_value(_cycle(fr,rng), t), ts))
        end
        # if lz != nothing
        #     lzrng = _cycle(lz, rng) # the line_z's for this segment
        #     push!(newlz, 0.0)
        #     append!(newlz, map(t -> lzrng[1+floor(Int, t * (length(rng)-1))], ts))
        # end
    end

    x := newx
    y := newy
    if z == nothing
        seriestype := :path
    else
        seriestype := :path3d
        z := newz
    end
    if fr != nothing
        fillrange := newfr
    end
    # if lz != nothing
    #     # line_z := newlz
    #     linecolor := (isa(d[:linecolor], ColorGradient) ? d[:linecolor] : cgrad())
    # end
    # PlotsBase.DD(d)
    ()
end
@deps curves path

# ---------------------------------------------------------------------------

# create a bar plot as a filled step function
@recipe function f(::Type{Val{:bar}}, x, y, z)
    procx, procy, xscale, yscale, baseline = _preprocess_barlike(d, x, y)
    nx, ny = length(procx), length(procy)
    axis = d[:subplot][isvertical(d) ? :xaxis : :yaxis]
    cv = [discrete_value!(axis, xi)[1] for xi=procx]
    procx = if nx == ny
        cv
    elseif nx == ny + 1
        0.5diff(cv) + cv[1:end-1]
    else
        error("bar recipe: x must be same length as y (centers), or one more than y (edges).\n\t\tlength(x)=$(length(x)), length(y)=$(length(y))")
    end

    # compute half-width of bars
    bw = d[:bar_width]
    hw = if bw == nothing
        0.5*_bar_width*ignorenan_mean(diff(procx))
    else
        Float64[0.5_cycle(bw,i) for i=1:length(procx)]
    end

    # make fillto a vector... default fills to 0
    fillto = d[:fillrange]
    if fillto == nothing
        fillto = 0
    end
    if (yscale in _logScales) && !all(_is_positive, fillto)
        fillto = map(x -> _is_positive(x) ? typeof(baseline)(x) : baseline, fillto)
    end

    # create the bar shapes by adding x/y segments
    xseg, yseg = Segments(), Segments()
    for i=1:ny
        yi = procy[i]
        if !isnan(yi)
            center = procx[i]
            hwi = _cycle(hw,i)
            fi = _cycle(fillto,i)
            push!(xseg, center-hwi, center-hwi, center+hwi, center+hwi, center-hwi)
            push!(yseg, yi, fi, fi, yi, yi)
        end
    end

    # widen limits out a bit
    expand_extrema!(axis, widen(ignorenan_extrema(xseg.pts)...))

    # switch back
    if !isvertical(d)
        xseg, yseg = yseg, xseg
    end

    # reset orientation
    orientation := default(:orientation)

    x := xseg.pts
    y := yseg.pts
    seriestype := :shape
    ()
end
@deps bar shape

# ---------------------------------------------------------------------------
# Histograms

_bin_centers(v::AVec) = (v[1:end-1] + v[2:end]) / 2

_is_positive(x) = (x > 0) && !(x ≈ 0)

_positive_else_nan{T}(::Type{T}, x::Real) = _is_positive(x) ? T(x) : T(NaN)

function _scale_adjusted_values{T<:AbstractFloat}(::Type{T}, V::AbstractVector, scale::Symbol)
    if scale in _logScales
        [_positive_else_nan(T, x) for x in V]
    else
        [T(x) for x in V]
    end
end

function _binbarlike_baseline{T<:Real}(min_value::T, scale::Symbol)
    if (scale in _logScales)
        !isnan(min_value) ? min_value / T(_logScaleBases[scale]^log10(2)) : T(1E-3)
    else
        zero(T)
    end
end

function _preprocess_binbarlike_weights{T<:AbstractFloat}(::Type{T}, w, wscale::Symbol)
    w_adj = _scale_adjusted_values(T, w, wscale)
    w_min = ignorenan_minimum(w_adj)
    w_max = ignorenan_maximum(w_adj)
    baseline = _binbarlike_baseline(w_min, wscale)
    w_adj, baseline
end

function _preprocess_barlike(d, x, y)
    xscale = get(d, :xscale, :identity)
    yscale = get(d, :yscale, :identity)
    weights, baseline = _preprocess_binbarlike_weights(float(eltype(y)), y, yscale)
    x, weights, xscale, yscale, baseline
end

function _preprocess_binlike(d, x, y)
    xscale = get(d, :xscale, :identity)
    yscale = get(d, :yscale, :identity)
    T = float(promote_type(eltype(x), eltype(y)))
    edge = T.(x)
    weights, baseline = _preprocess_binbarlike_weights(T, y, yscale)
    edge, weights, xscale, yscale, baseline
end

@recipe function f(::Type{Val{:barbins}}, x, y, z)
    edge, weights, xscale, yscale, baseline = _preprocess_binlike(d, x, y)
    if (d[:bar_width] == nothing)
        bar_width := diff(edge)
    end
    x := _bin_centers(edge)
    y := weights
    seriestype := :bar
    ()
end
@deps barbins bar

@recipe function f(::Type{Val{:scatterbins}}, x, y, z)
    edge, weights, xscale, yscale, baseline = _preprocess_binlike(d, x, y)
    xerror := diff(edge)/2
    x := _bin_centers(edge)
    y := weights
    seriestype := :scatter
    ()
end
@deps scatterbins scatter

function _stepbins_path(edge, weights, baseline::Real, xscale::Symbol, yscale::Symbol)
    log_scale_x = xscale in _logScales
    log_scale_y = yscale in _logScales

    nbins = length(linearindices(weights))
    if length(linearindices(edge)) != nbins + 1
        error("Edge vector must be 1 longer than weight vector")
    end

    x = eltype(edge)[]
    y = eltype(weights)[]

    it_e, it_w = start(edge), start(weights)
    a, it_e = next(edge, it_e)
    last_w = eltype(weights)(NaN)
    i = 1
    while (!done(edge, it_e) && !done(edge, it_e))
        b, it_e = next(edge, it_e)
        w, it_w = next(weights, it_w)

        if (log_scale_x && a ≈ 0)
            a = b/_logScaleBases[xscale]^3
        end

        if isnan(w)
            if !isnan(last_w)
                push!(x, a)
                push!(y, baseline)
            end
        else
            if isnan(last_w)
                push!(x, a)
                push!(y, baseline)
            end
            push!(x, a)
            push!(y, w)
            push!(x, b)
            push!(y, w)
        end

        a = b
        last_w = w
    end
    if (last_w != baseline)
        push!(x, a)
        push!(y, baseline)
    end

    (x, y)
end

@recipe function f(::Type{Val{:stepbins}}, x, y, z)
    axis = d[:subplot][PlotsBase.isvertical(d) ? :xaxis : :yaxis]

    edge, weights, xscale, yscale, baseline = _preprocess_binlike(d, x, y)

    xpts, ypts = _stepbins_path(edge, weights, baseline, xscale, yscale)
    if !isvertical(d)
        xpts, ypts = ypts, xpts
    end

    # create a secondary series for the markers
    if d[:markershape] != :none
        @series begin
            seriestype := :scatter
            x := _bin_centers(edge)
            y := weights
            fillrange := nothing
            label := ""
            primary := false
            ()
        end
        markershape := :none
        xerror := :none
        yerror := :none
    end

    x := xpts
    y := ypts
    seriestype := :path
    ()
end
PlotsBase.@deps stepbins path

wand_edges(x...) = (warn("Load the StatPlots package in order to use :wand bins. Defaulting to :auto", once = true); :auto)

function _auto_binning_nbins{N}(vs::NTuple{N,AbstractVector}, dim::Integer; mode::Symbol = :auto)
    _cl(x) = ceil(Int, NaNMath.max(x, one(x)))
    _iqr(v) = (q = quantile(v, 0.75) - quantile(v, 0.25); q > 0 ? q : oftype(q, 1))
    _span(v) = ignorenan_maximum(v) - ignorenan_minimum(v)

    n_samples = length(linearindices(first(vs)))

    # The nd estimator is the key to most automatic binning methods, and is modified for twodimensional histograms to include correlation
    nd = n_samples^(1/(2+N))
    nd = N == 2 ? nd / (1-cor(first(vs), last(vs))^2)^(3//8) : nd # the >2-dimensional case does not have a nice solution to correlations

    v = vs[dim]

    if mode == :auto
        mode = :fd
    end

    if mode == :sqrt  # Square-root choice
        _cl(sqrt(n_samples))
    elseif mode == :sturges  # Sturges' formula
        _cl(log2(n_samples) + 1)
    elseif mode == :rice  # Rice Rule
        _cl(2 * nd)
    elseif mode == :scott  # Scott's normal reference rule
        _cl(_span(v) / (3.5 * std(v) / nd))
    elseif mode == :fd  # Freedman–Diaconis rule
        _cl(_span(v) / (2 * _iqr(v) / nd))
    elseif mode == :wand
        wand_edges(v)  # this makes this function not type stable, but the type instability does not propagate
    else
        error("Unknown auto-binning mode $mode")
    end
end

_hist_edge{N}(vs::NTuple{N,AbstractVector}, dim::Integer, binning::Integer) = StatsBase.histrange(vs[dim], binning, :left)
_hist_edge{N}(vs::NTuple{N,AbstractVector}, dim::Integer, binning::Symbol) = _hist_edge(vs, dim, _auto_binning_nbins(vs, dim, mode = binning))
_hist_edge{N}(vs::NTuple{N,AbstractVector}, dim::Integer, binning::AbstractVector) = binning

_hist_edges{N}(vs::NTuple{N,AbstractVector}, binning::NTuple{N}) =
    map(dim -> _hist_edge(vs, dim, binning[dim]), (1:N...))

_hist_edges{N}(vs::NTuple{N,AbstractVector}, binning::Union{Integer, Symbol, AbstractVector}) =
    map(dim -> _hist_edge(vs, dim, binning), (1:N...))

_hist_norm_mode(mode::Symbol) = mode
_hist_norm_mode(mode::Bool) = mode ? :pdf : :none

function _make_hist{N}(vs::NTuple{N,AbstractVector}, binning; normed = false, weights = nothing)
    edges = _hist_edges(vs, binning)
    h = float( weights == nothing ?
        StatsBase.fit(StatsBase.Histogram, vs, edges, closed = :left) :
        StatsBase.fit(StatsBase.Histogram, vs, weights, edges, closed = :left)
    )
    normalize!(h, mode = _hist_norm_mode(normed))
end

@recipe function f(::Type{Val{:histogram}}, x, y, z)
    seriestype := length(y) > 1e6 ? :stephist : :barhist
    ()
end
@deps histogram barhist

@recipe function f(::Type{Val{:barhist}}, x, y, z)
    h = _make_hist((y,), d[:bins], normed = d[:normalize], weights = d[:weights])
    x := h.edges[1]
    y := h.weights
    seriestype := :barbins
    ()
end
@deps barhist barbins

@recipe function f(::Type{Val{:stephist}}, x, y, z)
    h = _make_hist((y,), d[:bins], normed = d[:normalize], weights = d[:weights])
    x := h.edges[1]
    y := h.weights
    seriestype := :stepbins
    ()
end
@deps stephist stepbins

@recipe function f(::Type{Val{:scatterhist}}, x, y, z)
    h = _make_hist((y,), d[:bins], normed = d[:normalize], weights = d[:weights])
    x := h.edges[1]
    y := h.weights
    seriestype := :scatterbins
    ()
end
@deps scatterhist scatterbins

@recipe function f{T, E}(h::StatsBase.Histogram{T, 1, E})
    seriestype --> :barbins

    st_map = Dict(
        :bar => :barbins, :scatter => :scatterbins, :step => :stepbins,
        :steppost => :stepbins # :step can be mapped to :steppost in pre-processing
    )
    seriestype := get(st_map, d[:seriestype], d[:seriestype])

    if d[:seriestype] == :scatterbins
        # Workaround, error bars currently not set correctly by scatterbins
        edge, weights, xscale, yscale, baseline = _preprocess_binlike(d, h.edges[1], h.weights)
        xerror --> diff(h.edges[1])/2
        seriestype := :scatter
        (PlotsBase._bin_centers(edge), weights)
    else
        (h.edges[1], h.weights)
    end
end

@recipe function f{H <: StatsBase.Histogram}(hv::AbstractVector{H})
    for h in hv
        @series begin
            h
        end
    end
end

# ---------------------------------------------------------------------------
# Histogram 2D

@recipe function f(::Type{Val{:bins2d}}, x, y, z)
    edge_x, edge_y, weights = x, y, z.surf

    float_weights = float(weights)
    if float_weights === weights
        float_weights = deepcopy(float_weights)
    end
    for (i, c) in enumerate(float_weights)
        if c == 0
            float_weights[i] = NaN
        end
    end

    x := PlotsBase._bin_centers(edge_x)
    y := PlotsBase._bin_centers(edge_y)
    z := Surface(float_weights)

    match_dimensions := true
    seriestype := :heatmap
    ()
end
PlotsBase.@deps bins2d heatmap

@recipe function f(::Type{Val{:histogram2d}}, x, y, z)
    h = _make_hist((x, y), d[:bins], normed = d[:normalize], weights = d[:weights])
    x := h.edges[1]
    y := h.edges[2]
    z := Surface(h.weights)
    seriestype := :bins2d
    ()
end
@deps histogram2d bins2d

@recipe function f{T, E}(h::StatsBase.Histogram{T, 2, E})
    seriestype --> :bins2d
    (h.edges[1], h.edges[2], Surface(h.weights))
end

# ---------------------------------------------------------------------------
# scatter 3d

@recipe function f(::Type{Val{:scatter3d}}, x, y, z)
    seriestype := :path3d
    if d[:markershape] == :none
        markershape := :circle
    end
    linewidth := 0
    linealpha := 0
    ()
end

# note: don't add dependencies because this really isn't a drop-in replacement

# ---------------------------------------------------------------------------
# contourf - filled contours

@recipe function f(::Type{Val{:contourf}}, x, y, z)
    fillrange := true
    seriestype := :contour
    ()
end

# ---------------------------------------------------------------------------
# Error Bars

function error_style!(d::KW)
    d[:seriestype] = :path
    d[:linecolor] = d[:markerstrokecolor]
    d[:linewidth] = d[:markerstrokewidth]
    d[:label] = ""
end

# if we're passed a tuple of vectors, convert to a vector of tuples
function error_zipit(ebar)
    if istuple(ebar)
        collect(zip(ebar...))
    else
        ebar
    end
end

function error_coords(xorig, yorig, ebar)
    # init empty x/y, and zip errors if passed Tuple{Vector,Vector}
    x, y = Array{float_extended_type(xorig)}(0), Array{Float64}(0)
    # for each point, create a line segment from the bottom to the top of the errorbar
    for i = 1:max(length(xorig), length(yorig))
        xi = _cycle(xorig, i)
        yi = _cycle(yorig, i)
        ebi = _cycle(ebar, i)
        nanappend!(x, [xi, xi])
        e1, e2 = if istuple(ebi)
            first(ebi), last(ebi)
        elseif isscalar(ebi)
            ebi, ebi
        else
            error("unexpected ebi type $(typeof(ebi)) for errorbar: $ebi")
        end
        nanappend!(y, [yi - e1, yi + e2])
    end
    x, y
end

# we will create a series of path segments, where each point represents one
# side of an errorbar
@recipe function f(::Type{Val{:yerror}}, x, y, z)
    error_style!(d)
    markershape := :hline
    d[:x], d[:y] = error_coords(d[:x], d[:y], error_zipit(d[:yerror]))
    ()
end
@deps yerror path

@recipe function f(::Type{Val{:xerror}}, x, y, z)
    error_style!(d)
    markershape := :vline
    d[:y], d[:x] = error_coords(d[:y], d[:x], error_zipit(d[:xerror]))
    ()
end
@deps xerror path

# TODO: move quiver to PlotRecipes

# ---------------------------------------------------------------------------
# quiver

# function apply_series_recipe(d::KW, ::Type{Val{:quiver}})
function quiver_using_arrows(d::KW)
    d[:label] = ""
    d[:seriestype] = :path
    if !isa(d[:arrow], Arrow)
        d[:arrow] = arrow()
    end

    velocity = error_zipit(d[:quiver])
    xorig, yorig = d[:x], d[:y]

    # for each point, we create an arrow of velocity vi, translated to the x/y coordinates
    x, y = zeros(0), zeros(0)
    for i = 1:max(length(xorig), length(yorig))
        # get the starting position
        xi = _cycle(xorig, i)
        yi = _cycle(yorig, i)

        # get the velocity
        vi = _cycle(velocity, i)
        vx, vy = if istuple(vi)
            first(vi), last(vi)
        elseif isscalar(vi)
            vi, vi
        elseif isa(vi,Function)
            vi(xi, yi)
        else
            error("unexpected vi type $(typeof(vi)) for quiver: $vi")
        end

        # add the points
        nanappend!(x, [xi, xi+vx, NaN])
        nanappend!(y, [yi, yi+vy, NaN])
    end

    d[:x], d[:y] = x, y
    # KW[d]
end

# function apply_series_recipe(d::KW, ::Type{Val{:quiver}})
function quiver_using_hack(d::KW)
    d[:label] = ""
    d[:seriestype] = :shape

    velocity = error_zipit(d[:quiver])
    xorig, yorig = d[:x], d[:y]

    # for each point, we create an arrow of velocity vi, translated to the x/y coordinates
    pts = P2[]
    for i = 1:max(length(xorig), length(yorig))

        # get the starting position
        xi = _cycle(xorig, i)
        yi = _cycle(yorig, i)
        p = P2(xi, yi)

        # get the velocity
        vi = _cycle(velocity, i)
        vx, vy = if istuple(vi)
            first(vi), last(vi)
        elseif isscalar(vi)
            vi, vi
        elseif isa(vi,Function)
            vi(xi, yi)
        else
            error("unexpected vi type $(typeof(vi)) for quiver: $vi")
        end
        v = P2(vx, vy)

        dist = norm(v)
        arrow_h = 0.1dist          # height of arrowhead
        arrow_w = 0.5arrow_h       # halfwidth of arrowhead
        U1 = v ./ dist             # vector of arrowhead height
        U2 = P2(-U1[2], U1[1])     # vector of arrowhead halfwidth
        U1 *= arrow_h
        U2 *= arrow_w

        ppv = p+v
        nanappend!(pts, P2[p, ppv-U1, ppv-U1+U2, ppv, ppv-U1-U2, ppv-U1])
    end

    d[:x], d[:y] = PlotsBase.unzip(pts[2:end])
    # KW[d]
end

# function apply_series_recipe(d::KW, ::Type{Val{:quiver}})
@recipe function f(::Type{Val{:quiver}}, x, y, z)
    if :arrow in supported_attrs()
        quiver_using_arrows(d)
    else
        quiver_using_hack(d)
    end
    ()
end
@deps quiver shape path

# -------------------------------------------------

# TODO: move OHLC to PlotRecipes finance.jl

"Represent Open High Low Close data (used in finance)"
type OHLC{T<:Real}
  open::T
  high::T
  low::T
  close::T
end
Base.convert(::Type{OHLC}, tup::Tuple) = OHLC(tup...)
# Base.tuple(ohlc::OHLC) = (ohlc.open, ohlc.high, ohlc.low, ohlc.close)

# get one OHLC path
function get_xy(o::OHLC, x, xdiff)
    xl, xm, xr = x-xdiff, x, x+xdiff
    ox = [xl, xm, NaN,
          xm, xm, NaN,
          xm, xr]
    oy = [o.open, o.open, NaN,
          o.low, o.high, NaN,
          o.close, o.close]
    ox, oy
end

# get the joined vector
function get_xy(v::AVec{OHLC}, x = 1:length(v))
    xdiff = 0.3ignorenan_mean(abs.(diff(x)))
    x_out, y_out = zeros(0), zeros(0)
    for (i,ohlc) in enumerate(v)
        ox,oy = get_xy(ohlc, x[i], xdiff)
        nanappend!(x_out, ox)
        nanappend!(y_out, oy)
    end
    x_out, y_out
end

# these are for passing in a vector of OHLC objects
# TODO: when I allow `@recipe f(::Type{T}, v::T) = ...` definitions to replace convertToAnyVector,
#       then I should replace these with one definition to convert to a vector of 4-tuples

# to squash ambiguity warnings...
@recipe f(x::AVec{Function}, v::AVec{OHLC}) = error()
@recipe f{R1<:Number,R2<:Number,R3<:Number,R4<:Number}(x::AVec{Function}, v::AVec{Tuple{R1,R2,R3,R4}}) = error()

# this must be OHLC?
@recipe f{R1<:Number,R2<:Number,R3<:Number,R4<:Number}(x::AVec, ohlc::AVec{Tuple{R1,R2,R3,R4}}) = x, OHLC[OHLC(t...) for t in ohlc]

@recipe function f(x::AVec, v::AVec{OHLC})
    seriestype := :path
    get_xy(v, x)
end

@recipe function f(v::AVec{OHLC})
    seriestype := :path
    get_xy(v)
end

# the series recipe, when passed vectors of 4-tuples

# -------------------------------------------------

# TODO: everything below here should be either changed to a
#       series recipe or moved to PlotRecipes

# "Sparsity plot... heatmap of non-zero values of a matrix"
# function spy{T<:Real}(z::AMat{T}; kw...)
#     mat = reshape(map(zi->float(zi!=0), z),1,:)
#     xn, yn = size(mat)
#     heatmap(mat; leg=false, yflip=true, aspect_ratio=:equal,
#         xlim=(0.5, xn+0.5), ylim=(0.5, yn+0.5),
#         kw...)
# end

# Only allow matrices through, and make it seriestype :spy so the backend can
# optionally handle it natively.

@userplot Spy

@recipe function f(g::Spy)
    @assert length(g.args) == 1 && typeof(g.args[1]) <: AbstractMatrix
    seriestype := :spy
    mat = g.args[1]
    if length(unique(mat[mat .!= 0])) < 2
        legend --> nothing
    end
    n,m = size(mat)
    PlotsBase.SliceIt, 1:m, 1:n, Surface(mat)
end

@recipe function f(::Type{Val{:spy}}, x,y,z)
    yflip := true
    aspect_ratio := 1
    rs, cs, zs = findnz(z.surf)
    xlim := ignorenan_extrema(cs)
    ylim := ignorenan_extrema(rs)
    if d[:markershape] == :none
        markershape := :circle
    end
    if d[:markersize] == default(:markersize)
        markersize := 1
    end
    markerstrokewidth := 0
    marker_z := zs
    label := ""
    x := cs
    y := rs
    z := nothing
    seriestype := :scatter
    grid --> false
    ()
end

# -------------------------------------------------

"Adds a+bx... straight line over the current plot"
function abline!(plt::Plot, a, b; kw...)
    plot!(plt, [ignorenan_extrema(plt)...], x -> b + a*x; kw...)
end

abline!(args...; kw...) = abline!(current(), args...; kw...)

# -------------------------------------------------
# Dates

dateformatter(dt) = string(Date(Dates.UTD(dt)))
datetimeformatter(dt) = string(DateTime(Dates.UTM(dt)))

@recipe f(::Type{Date}, dt::Date) = (dt -> Dates.value(dt), dateformatter)
@recipe f(::Type{DateTime}, dt::DateTime) = (dt -> Dates.value(dt), datetimeformatter)

# -------------------------------------------------
# Complex Numbers

@recipe function f{T<:Number}(A::Array{Complex{T}})
    xguide --> "Re(x)"
    yguide --> "Im(x)"
    real.(A), imag.(A)
end

# Splits a complex matrix to its real and complex parts
# Reals defaults solid, imaginary defaults dashed
# Label defaults are changed to match the real-imaginary reference / indexing
@recipe function f{T<:Real,T2}(x::AbstractArray{T},y::Array{Complex{T2}})
  ylabel --> "Re(y)"
  zlabel --> "Im(y)"
  x,real.(y),imag.(y)
end

# --------------------------------------------------
# Color Gradients

@userplot ShowLibrary
@recipe function f(cl::ShowLibrary)
    if !(length(cl.args) == 1 && isa(cl.args[1], Symbol))
        error("showlibrary takes the name of a color library as a Symbol")
    end

    library = PlotUtils.color_libraries[cl.args[1]]
    z = sqrt.((1:15)*reshape(1:20,1,:))

    seriestype := :heatmap
    ticks := nothing
    legend := false

    layout --> length(library.lib)

    i = 0
    for grad in sort(collect(keys(library.lib)))
        @series begin
            seriescolor := cgrad(grad, cl.args[1])
            title := string(grad)
            subplot := i += 1
            z
        end
    end
end

@userplot ShowGradient
@recipe function f(grad::ShowGradient)
    if !(length(grad.args) == 1 && isa(grad.args[1], Symbol))
        error("showgradient takes the name of a color gradient as a Symbol")
    end
    z = sqrt.((1:15)*reshape(1:20,1,:))
    seriestype := :heatmap
    ticks := nothing
    legend := false
    seriescolor := grad.args[1]
    title := string(grad.args[1])
    z
end


text:
using InternedStrings, SortingLab, DataFrames, StatsBase, BenchmarkTools
N = 200_000_000
K = 100

srand(1);
@time samplespace = InternedString.("id".*dec.(1:N÷K, 10));

function make_svec1(N, samplespace)
    rand(samplespace, N)
end

value(s) = s.value

function make_svec2(N, samplespace)
    @time vss = rand(value.(samplespace), N)
    @time InternedString.(vss, true)
end

function make_svec3(N, samplespace)
    @time vss = rand(value.(samplespace), N)
    @time svec = Vector{InternedString}(N)
    @time for i =1:N
        svec[i].value = vss[1]
    end
    svec
end

# @time make_svec(N, samplespace);
@time make_svec2(N, samplespace);

@time svec = make_svec3(N, samplespace);
too slow
using ProgressMeter
function newis(N, rid, samplespace)
    svec = Vector{InternedString}(N)
    p = Progress(N, 1)
    vss = value.(samplespace)
    for i = 1:N
        svec[i] = InternedString(vss[rid[i]], true)
        next!(p)
    end
    svec
end

@time make_svec1(N, samplespace);
@time make_svec2(N, samplespace); # winner
@time svec = make_svec3(N, samplespace); # doesn't even run


text:
"""
    Multiplier(x)

Stores a fixed `x` and multiplies by it, especially for testing
the gradient of higher order functions with respect to `x`.
```
julia> map(Multiplier(pi), [1 10 100 1000])
1×4 Matrix{Float64}:
 3.14159  31.4159  314.159  3141.59

julia> map(Multiplier(2), [1 10 100], [-1 3 -7])  # two arguments
1×3 Matrix{Int64}:
 -2  60  -1400

julia> map(Multiplier([1 2; 3 4]), ([5, 6], [7, 8]))  # x isa Matrix
([17, 39], [23, 53])
```
"""
struct Multiplier{T}
    x::T
end
(m::Multiplier)(y) = m.x * y
(m::Multiplier)(y, z) = (m.x * y) * z

function ChainRulesCore.rrule(m::Multiplier, y)
    Multiplier_pullback(dΩ) = (Tangent{typeof(m)}(; x = dΩ * y'), m.x' * dΩ)
    return m(y), Multiplier_pullback
end
function ChainRulesCore.rrule(m::Multiplier, y, z)
    Multiplier_pullback_3(dΩ) = (
        Tangent{typeof(m)}(; x = dΩ * (y * z)'),
        m.x' * dΩ * z',
        (m.x * y)' * dΩ,
    )
    return m(y, z), Multiplier_pullback_3
end

"""
    Counter()

Multiplies its input by number that increments on each call,
for testing execution order. Has a gradient `rrule` which
similarly increases by 10 each call, this is *not* the true gradient!
```
julia> map(Counter(), [1 1 1 10 10 10])
1×6 Matrix{Int64}:
 1  2  3  40  50  60
```
"""
mutable struct Counter
    n::Int
    Counter(n=0) = new(n)
end
Base.:(==)(a::Counter, b::Counter) = (a.n == b.n)
(c::Counter)(x) = x * (c.n += 1)
(c::Counter)(x, y) = (x + y) * (c.n += 1)

function ChainRulesCore.rrule(c::Counter, x)
    # False gradient, again just to record sequence of calls:
    Counter_back(dΩ) = (NoTangent(), dΩ * (c.n += 10))
    return c(x), Counter_back
end
function ChainRulesCore.rrule(c::Counter, x, y)
    Counter_back_2(dΩ) = (NoTangent(), dΩ * (c.n += 10), dΩ * (c.n += 10))
    return c(x, y), Counter_back_2
end

# NoRules - has no rules defined
struct NoRules; end

"A function that outputs a vector from a scalar for testing"
make_two_vec(x) = [x, x]
function ChainRulesCore.rrule(::typeof(make_two_vec), x)
    make_two_vec_pullback(ȳ) = (NoTangent(), sum(ȳ))
    return make_two_vec(x), make_two_vec_pullback
end

@testset "test_helpers.jl" begin

    @testset "Multiplier" begin
        # One argument
        test_rrule(Multiplier(4.0), 3.0)
        test_rrule(Multiplier(5.0 + 6im), 7.0 + 8im)
        test_rrule(Multiplier(rand(2,3)), rand(3,4))

        # Two arguments
        test_rrule(Multiplier(1.2), 3.4, 5.6)
        test_rrule(Multiplier(1.0 + 2im), 3.0 + 4im, 5.0 - 6im)
        test_rrule(Multiplier(rand(2,3)), rand(3,4), rand(4,5))
    end

    @testset "Counter" begin
        c = Counter()
        @test map(c, ones(5)) == 1:5
        @test c == Counter(5)
        y, back = rrule(c, 666)
        @test back(1) == (NoTangent(), 16)
        @test back(10) == (NoTangent(), 260)
        y2, back2 = rrule(c, 777, 888)
        @test back2(1) == (NoTangent(), 37, 47)
    end

    @testset "make_two_vec" begin
        test_rrule(make_two_vec, 1.5)
    end

end


text:
# Extracts all choice addresses from a choice map or trace
function all_addresses(choices::Gen.ChoiceMapNestedView,
                       exclude::Selection=EmptySelection())
    addrs = []
    choices = sort(collect(choices), by=first)
    for (k, v) in choices
        if v isa Gen.ChoiceMapNestedView
            append!(addrs, [Pair(k, a) for a in all_addresses(v, exclude[k])])
        elseif exclude[k] != AllSelection()
            push!(addrs, k)
        end
    end
    return addrs
end
all_addresses(choices::ChoiceMap, exclude::Selection=EmptySelection()) =
    all_addresses(nested_view(choices), exclude)
all_addresses(trace::Trace, exclude::Selection=EmptySelection()) =
    all_addresses(get_choices(trace), exclude)

# Single site Metropolis Hastings
function single_site_mh(T::Int, observations::ChoiceMap, n_iters::Int,
                        tracked_vars=[:β], obs_noise=5.0)
    trs, scores = Vector{Trace}(undef, n_iters) ,Vector{Float64}(undef, n_iters)
    data = DataFrame(fill(Float64, length(tracked_vars)), tracked_vars, n_iters)
    # Smart initialization of beta from the data
    init_choices, _, _ = propose(init_beta, (observations, min(T, 10)))
    observations = merge(observations, init_choices)
    trace, _ = generate(bayesian_sir, (T, obs_noise), observations)
    trs[1], scores[1] = trace, get_score(trace)
    data[1, tracked_vars] = [trace[v] for v in tracked_vars]
    # Extract addresses for all unobserved variables
    obs_addrs = Gen.select([:step => t => :obs for t in 1:T]...)
    latent_addrs = all_addresses(trace, obs_addrs)
    # Sequentially regenerate each address in the original trace
    for i in 2:n_iters
        @showprogress for addr in latent_addrs
            trace, _ = mh(trace, Gen.select(addr))
        end
        trs[i], scores[i] = trace, get_score(trace)
        data[i, tracked_vars] = [trace[v] for v in tracked_vars]
    end
    return trs, scores, data
end


text:
# Autogenerated wrapper script for Kokkos_jll for aarch64-apple-darwin
export libkokkoscontainers, libkokkoscore

using CompilerSupportLibraries_jll
JLLWrappers.@generate_wrapper_header("Kokkos")
JLLWrappers.@declare_library_product(libkokkoscontainers, "@rpath/libkokkoscontainers.3.5.dylib")
JLLWrappers.@declare_library_product(libkokkoscore, "@rpath/libkokkoscore.3.5.dylib")
function __init__()
    JLLWrappers.@generate_init_header(CompilerSupportLibraries_jll)
    JLLWrappers.@init_library_product(
        libkokkoscontainers,
        "lib/libkokkoscontainers.3.5.00.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libkokkoscore,
        "lib/libkokkoscore.3.5.00.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
put_real_eigenvalue_first(λ, T) = put_real_eigenvalue_first(λ, T, 0)
function put_real_eigenvalue_first(λ::Vector{Complex{BigFloat}}, T::Matrix{Complex{BigFloat}}, n_attempt::Int64)
    i_real = findfirst(λ .== real.(λ))
    if i_real == nothing
        return λ, T
    elseif i_real == 1
        return λ, T
    else
        n = length(λ)
        perm_mat = diagm(1=>ones(n-1), (1-n)=>ones(1))
        λ_mat = diagm(0=>λ)
        λ_perm_order = Int64.(collect(1:n)' * perm_mat)[:]
        λ_perm = λ[λ_perm_order]
        T_perm = T * perm_mat
        (n < n_attempt) && error("too many attempts")
        return put_real_eigenvalue_first(λ_perm, T_perm, n_attempt + 1)
    end
end

function big_eigen(M::Matrix{BigFloat})
    n, n2 = size(M)
    (n == n2) || error("matrix must be square")
    M_ = deepcopy(M)
    λ = eigvals(M_)
    T = zeros(Complex{BigFloat}, n, n)
    for k = 1:n
        sss = svd(M - I*λ[k])  # I'm aware this is massivly ineffecient
        T[:, k] = sss.Vt[end, :]'   # this does conjugate transpose
    end
    λ, T = put_real_eigenvalue_first(λ, T)
    return λ, T
end


text:
## test/parser.jl --- for the built in parser (still incomplete)

@testset "Atom/Residue parsing; residue unsorted" begin
    ## for the mean time use PDB data;
    ## return the result from frame extraction
    extracted = Trajectory(joinpath(dataloc, fname)) do traj
        fr = read(traj)
        extractframe(fr)
    end

    ## get the parsed StructureFrame from the same PDB
    parsed = read(joinpath(dataloc, fname), StructureFrame)

    # check if the atom numbers and lengths are the same
    @test length(extracted.at_pos) == length(parsed.at_pos)
    @test length(extracted.at_list) == length(parsed.at_list)
    @test length(extracted.res_list) == length(parsed.res_list)

    # check if the actual contents are the same
    # atom tests
    @test let (pl1, pl2) = (extracted.at_pos, parsed.at_pos)
        all(zip(pl1, pl2)) do (p1, p2)
            ## test approximate since there's something wrong with
            ## chemfiles extraction
            isapprox(p1, p2; rtol=1e-9)
        end
    end
    ## test skip this for the meantime since \' => p is not handled in extraction
    @test let (al1, al2) = (extracted.at_list, parsed.at_list)
        all(zip(al1, al2)) do (a1, a2)
            a1.name == a2.name
        end
    end skip=true
    # skip for the meantime since mass is not handled properly by parser
    @test let (al1, al2) = (extracted.at_list, parsed.at_list)
        all(zip(al1, al2)) do (a1, a2)
            a1.mass == a2.mass
        end
    end skip=true
    # skip for the meantime since charge is not handled properly by parser
    @test let (al1, al2) = (extracted.at_list, parsed.at_list)
        all(zip(al1, al2)) do (a1, a2)
            a1.charge == a2.charge
        end
    end skip=true

    ## residue tests
    @test let (rl1, rl2) = (extracted.res_list, parsed.res_list)
        all(zip(rl1, rl2)) do (res1, res2)
            res1.name == res2.name
        end
    end
    @test let (rl1, rl2) = (extracted.res_list, parsed.res_list)
        all(zip(rl1, rl2)) do (res1, res2)
            res1.chainid == res2.chainid
        end
    end
    @test let (rl1, rl2) = (extracted.res_list, parsed.res_list)
        all(zip(rl1, rl2)) do (res1, res2)
            !(res1.standard_pdb ⊻ res2.standard_pdb)
        end
    end
    ## test skip this for the meantime since \' => p is not handled in extraction
    @test let (rl1, rl2) = (extracted.res_list, parsed.res_list)
        all(zip(rl1, rl2)) do (res1, res2)
            res1.at_dict == res2.at_dict    ## なんか微妙だと思うけど;
        end
    end skip=true
end

@testset "Atom/residue parsing; residue sorted" begin
    unsorted = read(joinpath(dataloc, fname), StructureFrame; sort=false)
    sorted   = read(joinpath(dataloc, fname), StructureFrame; sort=true)

    ## check if the res list for "sorted" parsing is actually sorted
    @test let (; res_list) = sorted
        issorted(res_list;
                 lt=(x,y)->isless(x.standard_pdb, y.standard_pdb), rev=true)
    end

    ## TRIVIAL: check if the order of the atoms and positions are maintained
    @test let (pl1, pl2) = (unsorted.at_pos, sorted.at_pos)
        all(zip(pl1, pl2)) do (p1, p2)
            isapprox(p1, p2; rtol=1e-9)
        end
    end

    @test let (al1, al2) = (unsorted.at_list, sorted.at_list)
        all(zip(al1, al2)) do (a1, a2)
            a1.name == a2.name
        end
    end
end


text:
#!/usr/bin/env julia

include("../src/ReconstructionEvaluations.jl");

const RE = ReconstructionEvaluations
const CronUtils = RE.CronUtils
const Synaptor  = CronUtils.Synaptor
const Drivers   = CronUtils.Drivers
const NRI       = RE.NRI

using HDF5

seg_fname = ARGS[1]
output_prefix = ARGS[2]

append_score = false

#-------------------------
gt_seg = "data/pinky_golden_cube_021717.h5"
gt_edges = "data/pinky_golden_cube_021717_edges.csv"
gt_semmap = "data/pinky_golden_cube_021717_semmap.csv"

score_record = joinpath(dirname(output_prefix),"NRI_curve.h5")

dist_thr = 1000
res = [4,4,40]

overlap_chunk_shape = [1024,1024,128]
#-------------------------

#Defining output filename
score_h5 = "$(output_prefix)_scores.h5"

#==
Running Synaptor postprocessing
==#
println("RUNNING SYNAPTOR POSTPROCESSING")

cfg = Synaptor.make_gc_cfg(seg_fname, output_prefix)

edge_fname = Drivers.run_synaptor_cfg( cfg, output_prefix, dist_thr, res )

#==
Computing NRI scores
==#
println("")
println("COMPUTING NRI")

@time semmap = RE.load_semmap(gt_semmap)

@time full_nri, seg_nris, seg_nriw = NRI.nri(gt_edges, edge_fname)
@time full_pnri, pseg_nris, pseg_nriw = NRI.nri(edge_fname, gt_edges)#running a second time for proposed segments
println("NRI: $(full_nri)")
@time class_NRIs, class_weight = NRI.nri_by_class( seg_nris, seg_nriw, semmap )

println("Class NRI: $(class_NRIs)")
println("Class Weight: $(class_weight)")

seg_ids1, seg_nris = CronUtils.spvec_to_arrays(seg_nris)
seg_ids2, seg_nriw = CronUtils.spvec_to_arrays(seg_nriw)
seg_ids, seg_nris, seg_nriw = CronUtils.fill_array_indices(seg_ids1, seg_nris,
                                                           seg_ids2, seg_nriw)

pseg_ids1, pseg_nris = CronUtils.spvec_to_arrays(pseg_nris)
pseg_ids2, pseg_nriw = CronUtils.spvec_to_arrays(pseg_nriw)
pseg_ids, pseg_nris, pseg_nriw = CronUtils.fill_array_indices(pseg_ids1, pseg_nris,
                                                              pseg_ids2, pseg_nriw)

h5write( score_h5, "NRI"       , full_nri       )
h5write( score_h5, "seg_NRIs"  , seg_nris       )
h5write( score_h5, "seg_NRIws" , seg_nriw       )
h5write( score_h5, "seg_ids"   , seg_ids        )
h5write( score_h5, "ax_NRI"    , class_NRIs[2]  )
h5write( score_h5, "dend_NRI"  , class_NRIs[3]  )
h5write( score_h5, "ax_NRIw"   , class_weight[2])
h5write( score_h5, "dend_NRIw" , class_weight[3])
h5write( score_h5, "pseg_NRIs" , pseg_nris      )
h5write( score_h5, "pseg_NRIws", pseg_nriw      )
h5write( score_h5, "pseg_ids"  , pseg_ids       )

if append_score CronUtils.append_score( full_nri, score_record ) end

#==
Splits and Mergers
==#
println("")
println("FINDING SPLITS AND MERGERS")

asplits_fname = "$(output_prefix)_axon_splits.csv"
dsplits_fname = "$(output_prefix)_dend_splits.csv"
mergers_fname = "$(output_prefix)_mergers.csv"

@time om = Drivers.h5_file_om( gt_seg, seg_fname, overlap_chunk_shape )
@time splits, mergers = Drivers.splits_and_mergers(om)

axon_splits = filter( (k,v) -> semmap[k] == 2, splits )
dend_splits = filter( (k,v) -> semmap[k] == 3, splits )

num_axon_splits  = [ length(v) for v in values(axon_splits) ]
num_dend_splits  = [ length(v) for v in values(dend_splits) ]
num_mergers = [ length(v) for v in values(mergers) ]

RE.write_map_file( asplits_fname, axon_splits )
RE.write_map_file( dsplits_fname, dend_splits )
RE.write_map_file( mergers_fname, mergers )

h5write( score_h5, "num_axon_splits", num_axon_splits )
h5write( score_h5, "num_dend_splits", num_dend_splits )
h5write( score_h5, "num_mergers", num_mergers )


text:
using .StaticArrays

struct SArrayConstructor{S,N,L} end
struct MArrayConstructor{S,N,L} end
struct SizedArrayConstructor{S,N,M} end

(::SArrayConstructor{S,N,L})(data::NTuple{L,T}) where {S,T,N,L} = SArray{S,T,N,L}(data)
(::MArrayConstructor{S,N,L})(data::NTuple{L,T}) where {S,T,N,L} = MArray{S,T,N,L}(data)
(::SizedArrayConstructor{S,N,M})(data::AbstractArray{T,M}) where {S,T,N,M} =
    SizedArray{S,T,N,M}(data)

ConstructionBase.constructorof(sa::Type{<:SArray{S,<:Any,N,L}}) where {S,N,L} =
    SArrayConstructor{S,N,L}()
ConstructionBase.constructorof(sa::Type{<:MArray{S,<:Any,N,L}}) where {S,N,L} =
    MArrayConstructor{S,N,L}()
ConstructionBase.constructorof(sa::Type{<:SizedArray{S,<:Any,N,M}}) where {S,N,M} =
    SizedArrayConstructor{S,N,M}()


text:
macro def(name, definition)
    return quote
        macro $(esc(name))()
            esc($(Expr(:quote, definition)))
        end
    end
end

"""
`IterationLog` : store iteration information

### Usage Example:
```julia
ilog = IterationLog(:numE => Int, :numdE => Int, :res => Float64)
for n = 1:num_iter
   # do something with variables numE, numdE, res)
   push!(ilog, numE, numdE, res)
end
```
"""
type IterationLog
   keys::Tuple
   D::Dict
end

function IterationLog(args...)
   D = Dict()
   keys = tuple()
   for a in args
      if isa(a, Pair) && isa(a[2], DataType)
         keys = tuple(keys..., a[1])
         D[a[1]] = Vector{a[2]}()
      else
         keys = tuple(keys..., a)
         D[a] = Vector{Any}()
      end
   end
   return IterationLog(keys, D)
end

function Base.push!(log::IterationLog, args...)
   @assert length(args) == length(log.keys)
   for (k, a) in zip(log.keys, args)
      push!(log.D[k], a)
   end
   return log
end

"""
`DimerLog()` : generates an `IterationLog` for walker-type saddle search
methods
"""
DimerLog() = IterationLog(:numE => Int, :numdE => Int,
                          :res_trans => Float64, :res_rot => Float64)

"""
`PathLog()` : generates an `IterationLog` for string-type saddle search
methods
"""
PathLog() = IterationLog(:numE => Int, :numdE => Int, :maxres => Float64)

Base.getindex(l::IterationLog, idx) = l.D[idx]

# convenient access
numE(l::IterationLog) = l[:numE]
numdE(l::IterationLog) = l[:numdE]
res_trans(l::IterationLog) = l[:res_trans]
res_rot(l::IterationLog) = l[:res_rot]
maxres(l::IterationLog) = l[:maxres]

Base.dot{T}(x, A::UniformScaling{T}, y) = A.λ * dot(x,y)
Base.dot(x, A::AbstractMatrix, y) = dot(x, A*y)
Base.norm(P, x) = sqrt(dot(x, P*x))
dualnorm(P, f) = sqrt(dot(f, P \ f))

"""
An abstract linear operator representing `P + s * (Pv) ⊗ (Pv)`

Define `*` and `\`, the latter via Sherman-Morrison-Woodbury formula.
"""
type PreconSMW{T} <: AbstractMatrix{T}
   P       # an invertiable N x N matrix (probably spd)
   v       # a vector of length N
   Pv      # the vector P * v
   s::T    # see doc
   smw::T  # the SMW-factor
end

PreconSMW(P, v, s) = PreconSMW(P, v, P*v, s, s / (1.0 + s * dot(v, P, v)))

import Base: *, \, size
(*)(A::PreconSMW, x::AbstractVector) = A.P * x + (A.s * dot(A.Pv, x)) * A.Pv
(\)(A::PreconSMW, f::AbstractVector) = (A.P \ f) - ((A.smw * dot(A.v, f)) * A.v)


text:
using .Symbolics: jacobian, hessian

# ===========================================================================
# Functionalities for Jacobian matrices
# ===========================================================================

# ----------------------------------------
# Symbolic expressions for the Jacobian
# ----------------------------------------

function _jacobian_expression(f::AbstractVector{Num}, var::Vector{Num})
    Jsym = jacobian(f, var)
    return Jsym
end

function _jacobian_expression(sys::BBCS{Vector{Num}}, var)
    field = sys.f
    return _jacobian_expression(field, var)
end

# fallback for nonlinear systems: assumes in-place form f!(dx, x, p, t),
# *ignoring* p and t args and uses `var` to trace the associated symbolic function
function _jacobian_expression(sys::BBCS, var)
    n = length(var) # dimension
    expr = Vector{Num}(undef, n)
    field! = sys.f
    field!(expr, var, [], [])
    return _jacobian_expression(expr, var)
end

# fallback implementation: extract the vector field and use the symbolic variables
# `var` to trace the associated symbolic vector field
function _jacobian_expression(sys::AbstractContinuousSystem, var)
    F = field(VectorField(sys))
    expr = F(var)
    return _jacobian_expression(expr, var)
end

function _jacobian_expression(ivp::InitialValueProblem, var)
    return _jacobian_expression(ivp.s, var)
end

# ------------------------------------------
# Julia functions to evaluate the Jacobian
# ------------------------------------------

# given an ODE right-hand-side, compute the Jacobian function symbolically,
# then compile a Julia function for that symbolic exprsesion
function _jacobian_function(f::AbstractVector{Num}, var::Vector{Num})

    # compute the symbolic jacobian
    Jsym = _jacobian_expression(f, var)

    # use the out-of-place version
    Jexpr = build_function(Jsym, var)[1]

    # return the associated Julia function
    return eval(Jexpr)
end

function _jacobian_function(sys::BBCS{Vector{Num}}, var)
    field = sys.f
    return _jacobian_function(field, var)
end

# fallback for nonlinear systems: assumes in-place form f!(dx, x, p, t),
# *ignoring* p and t args and uses `var` to trace the associated symbolic function
function _jacobian_function(sys::BBCS, var)
    n = length(var) # dimension
    expr = Vector{Num}(undef, n)
    field! = sys.f
    field!(expr, var, [], [])
    return _jacobian_function(expr, var)
end

# fallback implementation: extract the vector field and use the symbolic variables
# `var` to trace the associated symbolic vector field
function _jacobian_function(sys::AbstractContinuousSystem, var)
    F = field(VectorField(sys))
    expr = F(var)
    return _jacobian_function(expr, var)
end

function _jacobian_function(ivp::InitialValueProblem, var)
    return _jacobian_function(ivp.s, var)
end

# ===========================================================================
# Functionalities for Hessian matrices
# ===========================================================================

# ----------------------------------------
# Symbolic expressions for the Hessian
# ----------------------------------------

# return a vector of Hessian symbolic expressions, one for each component function
function _hessian_expression(f::AbstractVector{Num}, var::Vector{Num})
    return [_hessian_expression(fi, var) for fi in f]
end

function _hessian_expression(f::AbstractVector{Num}, i::Int, var::Vector{Num})
    return _hessian_expression(f[i], var)
end

function _hessian_expression(fi::Num, var::Vector{Num})
    Hsym = hessian(fi, var)
    return Hsym
end

function _hessian_expression(sys::BBCS{Vector{Num}}, var)
    field = sys.f
    return _hessian_expression(field, var)
end
function _hessian_expression(sys::BBCS, var)
    n = length(var) # dimension
    expr = Vector{Num}(undef, n)
    field! = sys.f
    field!(expr, var, [], [])
    return _hessian_expression(expr, var)
end

function _hessian_expression(sys::AbstractContinuousSystem, var)
    F = field(VectorField(sys))
    expr = F(var)
    return _hessian_expression(expr, var)
end

function _hessian_expression(ivp::InitialValueProblem, var)
    return _hessian_expression(ivp.s, var)
end

# -----------------------------------------
# Julia functions to evaluate the Hessian
# -----------------------------------------

# return a vector of Hessian functions, one for each component function
function _hessian_function(f::AbstractVector{Num}, var::Vector{Num})
    return [_hessian_function(fi, var) for fi in f]
end

function _hessian_function(f::AbstractVector{Num}, i::Int, var::Vector{Num})
    return _hessian_function(f[i], var)
end

function _hessian_function(fi::Num, var::Vector{Num})

    # compute the symbolic jacobian
    Hsym = _hessian_expression(fi, var)

    # use the out-of-place version
    Hexpr = build_function(Hsym, var)[1]

    # return the associated Julia function
    return eval(Hexpr)
end

function _hessian_function(sys::BBCS{Vector{Num}}, var)
    field = sys.f
    return _hessian_function(field, var)
end

function _hessian_function(sys::BBCS, var)
    n = length(var) # dimension
    expr = Vector{Num}(undef, n)
    field! = sys.f
    field!(expr, var, [], [])
    return _hessian_function(expr, var)
end

function _hessian_function(sys::AbstractContinuousSystem, var)
    F = field(VectorField(sys))
    expr = F(var)
    return _hessian_function(expr, var)
end

function _hessian_function(ivp::InitialValueProblem, var)
    return _hessian_function(ivp.s, var)
end


text:
mutable struct MCVISimulator <: POMDPs.Simulator
    rng::AbstractRNG
    # init_state                  # In case we want to fix starting state
    times::Integer
    display::Bool
end
MCVISimulator() = MCVISimulator(MersenneTwister(420), 1, false)

function simulate(sim::MCVISimulator, pomdp::POMDPs.POMDP, policy::MCVIPolicy, updater::MCVIUpdater, initial_node::MCVINode, init_state=nothing)
    sum_reward::Reward = 0
    if init_state != nothing
        s = init_state
    else
        s = initialstate(pomdp, sim.rng)
    end
    for i in 1:sim.times
        n = copy(initial_node)
        disc::Float64 = 1
        sumr::Reward = 0
        while true
            a = action(policy, n)
            sprime, obs, r = gen(DDNOut(:sp,:o,:r), pomdp, s, a, sim.rng)
            o = @sprintf("%0.6f", obs)
            rtxt = @sprintf("%0.3f", r)
            sim.display && println("s: $s \t a: $a \t sp: $sprime \t o: $o \t r: $rtxt \t nid: $(n.id)")
            disc *= discount(pomdp)
            sumr += r*disc
            s = sprime
            if !hasnext(n)
                sim.display && println("===========")
                break
            end
            n = update(updater, n, n.act, obs)
        end
        sum_reward += sumr
    end
    sum_reward /= sim.times

    if isapprox(sum_reward, 0)
        @warn("Simulated reward close to 0")
    end
    return sum_reward
end

function simulate(sim::MCVISimulator, pomdp::POMDPs.POMDP, policy::POMDPs.Policy, updater::POMDPs.Updater, initial_belief::MCVIBelief)
    sum_reward::Reward = 0
    if sim.init_state != nothing
        s = sim.init_state
    else
        s = initialstate(pomdp, sim.rng)
    end
    for i in 1:sim.times
        b = initial_belief
        disc::Float64 = 1
        sumr::Reward = 0
        while true
            a = action(policy, b)
            if isterminal(pomdp, a)
                break
            end
            sprime, obs, r = gen(DDNOut(:sp,:o,:r), pomdp, s, a, sim.rng)
            disc *= discount(pomdp)
            sumr += r*disc
            s = sprime
            b = next(b, a, pomdp, sim.rng)
            b = next(b, obs, pomdp)
        end
        sum_reward += sumr
    end
    sum_reward /= sim.times

    if isapprox(sum_reward, 0)
        @warn("Simulated reward close to 0")
    end
    return sum_reward
end


text:
module Omega
using Reexport

@reexport using InferenceBase

@reexport using OmegaCore
@reexport using SoftPredicates

@reexport using OmegaDistributions
@reexport using OmegaSoftPredicates

@reexport using OmegaMH
@reexport using ReplicaExchange
# @reexport using InvolutiveMCMC
end

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
using Iris
using Test


text:
"""
    load_gpt_pretrain(path::AbstractString, sym = :all;
                      startsym="_start_",
                      delisym="_delimiter_",
                      clfsym="_classify_",
                      unksym="<unk>")

load gpt data/model from pretrain bson data. use `sym` to determine which data to load.
set `<xxx>sym`s for setting special symbols in vocabulary.

possible value: `:all`(default), `gpt_model`, `bpe`, `vocab`, `tokenizer`.
"""
function load_gpt_pretrain(path::AbstractString, sym = :all; kw...)
  @info "loading pretrain gpt model: $(basename(path)) $(sym == :all ? "" : sym)"
  @assert sym ∈ (:all, :gpt_model, :bpe, :vocab, :tokenizer) "sym only support :all, :gpt_model, :bpe, :vocab, :tokenizer"
  if isnpbson(path)

    bson = BSON.parse(path)

    if sym ∈ (:all, :bpe)
      bpe = BSON.raise_recursive(bson[:bpe])
      sym == :bpe && return bpe
    end

    tokenizer = BSON.raise_recursive(bson[:tokenizer])
    sym == :tokenizer && return tokenizer

    raw_vocab = BSON.raise_recursive(bson[:raw_vocab])
    vocab = build_vocab(raw_vocab; kw...)
    sym == :vocab && return vocab

    weights = BSON.raise_recursive(bson[:weights])
    gpt_model = load_gpt_from_npbson(weights, length(vocab))
    sym == :gpt_model && return gpt_model

    return gpt_model, bpe, vocab, tokenizer
  elseif isbson(path)
    bson = BSON.parse(path)
    if sym ∈ (:all, :bpe)
      bpe = BSON.raise_recursive(bson[:bpe])
      sym == :bpe && return bpe
    end
    tokenizer = BSON.raise_recursive(bson[:tokenizer])
    sym == :tokenizer && return tokenizer
    vocab = BSON.raise_recursive(bson[:vocab])
    sym == :vocab && return vocab
    gpt_model = BSON.raise_recursive(bson[:gpt_model])
    sym == :gpt_model && return gpt_model
    return gpt_model, bpe, vocab, tokenizer
  else
    error("""
          Unknown data type: Is it a gpt model?
          If it is the weight file get from finetune-transformer-lm, then run `npy2bson` on that file beforehand.
    """)
  end
end

# function load_gpt_pretrain(path::AbstractString)
#   if isnpbson(path)
#     @info "loading raw gpt np data: $path"
#     bson = BSON.load(path)
#     bson[:weights], bson[:bpe], bson[:raw_vocab], bson[:tokenizer]
#   elseif isbson(path)
#     @info "loading gpt model: $path"
#     bson = BSON.load(path)
#     bson[:gpt_model], bson[:bpe], bson[:vocab], bson[:tokenizer]
#   else
#     error("""
#           Unknown data type: Is it a gpt model?
#           If it is the weight file get from finetune-transformer-lm, then run npy2bson on that file beforehand.
#     """)
#   end
# end