text:
include("testfunctions.jl")

@testset "DQMC Parameters" begin
    P1 = MonteCarlo.DQMCParameters(beta=5.0) #defaults to delta_tau = 0.1
    @test all((P1.beta, P1.delta_tau, P1.slices) .== (5.0, 0.1, 50))

    P2 = MonteCarlo.DQMCParameters(beta=5.0, delta_tau=0.01)
    @test all((P2.beta, P2.delta_tau, P2.slices) .== (5.0, 0.01, 500))

    P3 = MonteCarlo.DQMCParameters(beta=50.0, slices=20)
    @test all((P3.beta, P3.delta_tau, P3.slices) .== (50.0, 2.5, 20))

    P4 = MonteCarlo.DQMCParameters(delta_tau=0.1, slices=50)
    @test all((P4.beta, P4.delta_tau, P4.slices) .== (5.0, 0.1, 50))

    @test parameters(P4) == (beta = 5.0, delta_tau = 0.1, thermalization = 100, sweeps = 100)
end

@testset "Magnitude Stats" begin
    m = MonteCarlo.MagnitudeStats()
    @test m.max == -Inf
    @test m.min == Inf
    @test m.sum == 0.0
    @test m.count == 0

    @test min(m) == 0.0
    @test max(m) == 0.0
    @test mean(m) == 0.0
    @test length(m) == 0

    push!(m, 1e-7)
    @test m.max == log10(1e-7)
    @test m.min == log10(1e-7)
    @test m.sum == log10(1e-7)
    @test m.count == 1

    push!(m, 1e-5)
    @test m.max == log10(1e-5)
    @test m.min == log10(1e-7)
    @test m.sum == log10(1e-7) + log10(1e-5)
    @test m.count == 2

    @test min(m) ≈ 1e-7
    @test max(m) ≈ 1e-5
    @test mean(m) ≈ 10^(0.5*(log10(1e-7) + log10(1e-5)))
    @test length(m) == 2
end

@testset "DQMC utilities" begin
    m = HubbardModelAttractive(8, 1);

    # Getters
    dqmc = DQMC(m; beta=5.0)

    @test MonteCarlo.beta(dqmc) == dqmc.parameters.beta
    @test MonteCarlo.nslices(dqmc) == dqmc.parameters.slices
    @test MonteCarlo.current_slice(dqmc) == dqmc.stack.current_slice
    @test MonteCarlo.configurations(dqmc) == dqmc.recorder
    @test MonteCarlo.parameters(dqmc) == merge(
        parameters(dqmc.parameters), parameters(dqmc.model)
    )

    io = IOBuffer()
    show(io, dqmc)
    @test String(take!(io)) == "Determinant quantum Monte Carlo simulation\nModel: attractive Hubbard model, 8 sites\nBeta: 5.0 (T ≈ 0.2)\nMeasurements: 0 (0 + 0)"

    # DQMC mandatory
    m = DummyModel(SquareLattice(2))
    @test_throws MethodError rand(DQMC, m, 4)
    @test_throws ErrorException MonteCarlo.nflavors(m)
    @test_throws MethodError MonteCarlo.hopping_matrix(dqmc, m)
    @test_throws MethodError MonteCarlo.interaction_matrix_exp!(dqmc, m, zeros(2,2), zeros(2,2), 1, 1.0)
    @test_throws MethodError MonteCarlo.propose_local(dqmc, m, 1, 1, zeros(2,2))
    @test_throws MethodError MonteCarlo.accept_local!(dqmc, m, 1, 1, zeros(2,2), 1.0, 1.0, nothing)

    MonteCarlo.nflavors(::DummyModel) = 2
    MonteCarlo.lattice(m::DummyModel) = m.lattice

    # DQMC optional
    @test MonteCarlo.greenseltype(DQMC, m) == ComplexF64
    @test MonteCarlo.interaction_matrix_type(DQMC, m) == Matrix{ComplexF64}
    @test MonteCarlo.init_interaction_matrix(m) == zeros(ComplexF64, 8, 8)
    @test_throws MethodError MonteCarlo.energy_boson(dqmc, m, MonteCarlo.conf(dqmc))
    @test parameters(m) == NamedTuple()

    # constructors
    mc = DQMC{
        typeof(dqmc.model), MonteCarlo.CheckerboardFalse, typeof(dqmc.conf),
        typeof(dqmc.recorder), typeof(dqmc.stack), typeof(dqmc.ut_stack),
        typeof(dqmc.scheduler)
    }(
        dqmc.model, dqmc.conf, dqmc.temp_conf, dqmc.last_sweep, dqmc.stack,
        dqmc.ut_stack, dqmc.scheduler, dqmc.parameters, dqmc.analysis,
        dqmc.recorder, dqmc.thermalization_measurements, dqmc.measurements
    )
    for field in fieldnames(DQMC)
        field == :lattice_iterator_cache && continue
        @test getfield(dqmc, field) == getfield(mc, field)
    end

    mc = DQMC(
        MonteCarlo.CheckerboardFalse, dqmc.model, dqmc.conf, dqmc.temp_conf,
        dqmc.last_sweep, dqmc.stack, dqmc.ut_stack, dqmc.scheduler,
        dqmc.parameters, dqmc.analysis, dqmc.recorder,
        dqmc.thermalization_measurements, dqmc.measurements
    )
    for field in fieldnames(DQMC)
        field == :lattice_iterator_cache && continue
        @test getfield(dqmc, field) == getfield(mc, field)
    end

    mc = DQMC(dqmc)
    for field in fieldnames(DQMC)
        field == :lattice_iterator_cache && continue
        @test getfield(dqmc, field) == getfield(mc, field)
    end

    mc = DQMC(dqmc, last_sweep = 9147, recorder = Discarder())
    for field in fieldnames(DQMC)
        if field == :last_sweep
            @test mc.last_sweep == 9147
        elseif field == :recorder
            @test mc.recorder isa Discarder
        elseif field == :lattice_iterator_cache
        else
            @test getfield(dqmc, field) == getfield(mc, field)
        end
    end
    @test mc isa DQMC{
        typeof(dqmc.model), MonteCarlo.CheckerboardFalse, typeof(dqmc.conf),
        Discarder, typeof(dqmc.stack), typeof(dqmc.ut_stack),
        typeof(dqmc.scheduler)
    }
end

@testset "GreensMatrix" begin
    m = HubbardModelAttractive(2, 2, mu=0.5)
    mc = DQMC(m, beta=1.0, safe_mult=10, thermalization=1, sweeps=1)
    MonteCarlo.initialize_stack(mc, mc.ut_stack)

    G = greens(mc)
    @test G isa GreensMatrix
    @test G.k == 0
    @test G.l == 0
    @test size(G) == size(G.val)

    G = greens(mc, 4)
    @test G isa GreensMatrix
    @test G.k == 4
    @test G.l == 4

    G = greens!(mc, 4)
    @test G isa GreensMatrix
    @test G.k == 4
    @test G.l == 4

    G = greens(mc, 2, 7)
    @test G isa GreensMatrix
    @test G.k == 2
    @test G.l == 7

    G = greens(mc, 5, 3)
    @test G isa GreensMatrix
    @test G.k == 5
    @test G.l == 3

    M = Float64[1 2; 3 4]
    G = GreensMatrix(0, 7, M)
    @test G[1, 1] == M[1, 1]
    @test G[1, 2] == M[1, 2]
    @test G[2, 1] == M[2, 1]
    @test G[2, 2] == M[2, 2]

    D = swapop(G)
    @test D isa MonteCarlo.Permuted
    @test D[1, 1] == -M[1, 1]
    @test D[2, 1] == -M[1, 2]
    @test D[1, 2] == -M[2, 1]
    @test D[2, 2] == -M[2, 2]
    @test size(D) == size(D.x.val)

    G = GreensMatrix(7, 7, M)
    D = swapop(G)
    @test D[1, 1] == 1 - M[1, 1]
    @test D[2, 1] == -M[1, 2]
    @test D[1, 2] == -M[2, 1]
    @test D[2, 2] == 1 - M[2, 2]
    @test swapop(D) == G

    G2 = GreensMatrix(7, 7, M)
    @test G == G2
    @test G ≈ G2

    G2 = GreensMatrix(6, 7, M)
    @test !(G == G2)
    @test !(G ≈ G2)

    G2 = GreensMatrix(7, 7, M .+ eps(10.0))
    @test !(G == G2)
    @test G ≈ G2

    G2 = GreensMatrix(6, 7, M .+ eps(10.0))
    @test !(G == G2)
    @test !(G ≈ G2)
end

@testset "DQMC stack" begin
    m = HubbardModelAttractive(8, 1);

    # constructors
    dqmc = DQMC(m; beta=5.0)

    # generic checkerboard
    sq = MonteCarlo.SquareLattice(4);
    @test MonteCarlo.build_checkerboard(sq) == ([1.0 3.0 5.0 7.0 9.0 11.0 13.0 15.0 1.0 2.0 4.0 6.0 9.0 10.0 12.0 14.0 2.0 3.0 4.0 5.0 8.0 10.0 11.0 16.0 6.0 7.0 8.0 12.0 13.0 14.0 15.0 16.0; 2.0 4.0 6.0 8.0 10.0 12.0 14.0 16.0 5.0 3.0 8.0 7.0 13.0 11.0 16.0 15.0 6.0 7.0 1.0 9.0 12.0 14.0 15.0 13.0 10.0 11.0 5.0 9.0 1.0 2.0 3.0 4.0; 1.0 5.0 9.0 13.0 17.0 21.0 25.0 29.0 2.0 3.0 8.0 11.0 18.0 19.0 24.0 27.0 4.0 6.0 7.0 10.0 16.0 20.0 22.0 31.0 12.0 14.0 15.0 23.0 26.0 28.0 30.0 32.0], UnitRange[1:8, 9:16, 17:24, 25:32], 4)

    m = HubbardModelAttractive(8, 2, mu=0.5)
    mc1 = DQMC(m, beta=5.0)
    mc2 = DQMC(m, beta=5.0, checkerboard=false)
    mc2.conf = deepcopy(mc1.conf)
    MonteCarlo.init_hopping_matrices(mc1, m)
    MonteCarlo.init_hopping_matrices(mc2, m)
    MonteCarlo.build_stack(mc1, mc1.stack)
    MonteCarlo.build_stack(mc2, mc2.stack)
    @test MonteCarlo.slice_matrix(mc1, m, 1, 1.) == MonteCarlo.slice_matrix(mc2, m, 1, 1.)

    mc = DQMC(m, beta=5.0, checkerboard=true, delta_tau=0.1)
    MonteCarlo.init_hopping_matrices(mc, m)
    hop_mat_exp_chkr = foldl(*,mc.stack.chkr_hop_half) * sqrt.(mc.stack.chkr_mu)
    r = MonteCarlo.effreldiff(mc.stack.hopping_matrix_exp,hop_mat_exp_chkr)
    r[findall(x -> x==zero(x), hop_mat_exp_chkr)] .= 0.
    @test maximum(MonteCarlo.absdiff(mc.stack.hopping_matrix_exp,hop_mat_exp_chkr)) <= mc.parameters.delta_tau

    # initial greens test
    mc = DQMC(m, beta=5.0, safe_mult=1)
    MonteCarlo.build_stack(mc, mc.stack)
    MonteCarlo.propagate(mc)
    # With this we effectively test calculate_greens without wrap_greens
    greens, = calculate_greens_and_logdet(mc, mc.stack.current_slice)
    MonteCarlo.wrap_greens!(mc, greens, mc.stack.current_slice+1, -1)
    @test greens ≈ mc.stack.greens
    # here with a single implied wrap
    greens, = calculate_greens_and_logdet(mc, mc.stack.current_slice-1)
    @test maximum(MonteCarlo.absdiff(greens, mc.stack.greens)) < 1e-12

    # wrap greens test
    for k in 0:9
        MonteCarlo.wrap_greens!(mc, mc.stack.greens, mc.stack.current_slice - k, -1)
    end
    greens, = calculate_greens_and_logdet(mc, mc.stack.current_slice-11)
    @test maximum(MonteCarlo.absdiff(greens, mc.stack.greens)) < 1e-9

    # Check greens reconstruction used in replay
    mc = DQMC(m, beta=5.0, safe_mult=5)
    # Make sure this works with any values
    for k in shuffle(0:MonteCarlo.nslices(mc))
        G1, _ = calculate_greens_and_logdet(mc, k)
        G2 = MonteCarlo.calculate_greens(mc, k)
        @test G1 ≈ G2
    end

    # check reverse_build_stack
    m = HubbardModelAttractive(8, 1);
    mc1 = DQMC(m; beta=5.0)
    MonteCarlo.init!(mc1)
    MonteCarlo.build_stack(mc1, mc1.stack)
    MonteCarlo.propagate(mc1)
    for t in 1:MonteCarlo.nslices(mc1)
        MonteCarlo.propagate(mc1)
    end

    m = HubbardModelAttractive(8, 1);
    mc2 = DQMC(m; beta=5.0)
    mc2.conf .= mc1.conf
    MonteCarlo.init!(mc2)
    MonteCarlo.reverse_build_stack(mc2, mc2.stack)
    MonteCarlo.propagate(mc2)

    @test mc1.stack.current_slice == mc2.stack.current_slice
    @test mc1.stack.direction == mc2.stack.direction
    @test mc1.stack.greens ≈ mc2.stack.greens
    for field in (:u_stack, :d_stack, :t_stack)
        for i in 1:mc1.stack.n_elements
            @test getfield(mc1.stack, field)[i] ≈ getfield(mc2.stack, field)[i]
        end
    end

    for field in (:Ul, :Ur, :Dl, :Dr, :Tl, :Tr)
        @test getfield(mc1.stack, field) ≈ getfield(mc2.stack, field)
    end

end

@testset "Unequal Time Stack" begin
    m = HubbardModelAttractive(6, 1);
    dqmc = DQMC(m; beta=15.0, safe_mult=5)
    MonteCarlo.initialize_stack(dqmc, dqmc.ut_stack)
    MonteCarlo.build_stack(dqmc, dqmc.stack)

    @testset "lazy build_stack forward/backward" begin
        for i in 2:length(dqmc.ut_stack.forward_u_stack)
            dqmc.ut_stack.forward_u_stack[i]  .= 0.0
            dqmc.ut_stack.forward_d_stack[i]  .= 0.0
            dqmc.ut_stack.forward_t_stack[i]  .= 0.0
        end
        for i in 1:length(dqmc.ut_stack.backward_d_stack)-1
            dqmc.ut_stack.backward_u_stack[i] .= 0.0
            dqmc.ut_stack.backward_d_stack[i] .= 0.0
            dqmc.ut_stack.backward_t_stack[i] .= 0.0
        end

        for upto in (4, 6)
            MonteCarlo.lazy_build_forward!(dqmc, dqmc.ut_stack, upto)
            for i in 1:upto
                @test dqmc.stack.u_stack[i] ≈ dqmc.ut_stack.forward_u_stack[i]
                @test dqmc.stack.d_stack[i] ≈ dqmc.ut_stack.forward_d_stack[i]
                @test dqmc.stack.t_stack[i] ≈ dqmc.ut_stack.forward_t_stack[i]
            end
            for i in upto+1:length(dqmc.ut_stack.forward_u_stack)
                @test all(dqmc.ut_stack.forward_u_stack[i] .== 0)
                @test all(dqmc.ut_stack.forward_d_stack[i] .== 0)
                @test all(dqmc.ut_stack.forward_t_stack[i] .== 0)
            end
        end

        while dqmc.stack.direction == -1
            MonteCarlo.propagate(dqmc)
        end

        for downto in (8, 6)
            MonteCarlo.lazy_build_backward!(dqmc, dqmc.ut_stack, downto)
            for i in length(dqmc.ut_stack.backward_d_stack):-1:downto
                @test dqmc.stack.u_stack[i] ≈ dqmc.ut_stack.backward_u_stack[i]
                @test dqmc.stack.d_stack[i] ≈ dqmc.ut_stack.backward_d_stack[i]
                @test dqmc.stack.t_stack[i] ≈ dqmc.ut_stack.backward_t_stack[i]
            end
            for i in downto-1:-1:1
                @test all(dqmc.ut_stack.backward_u_stack[i] .== 0)
                @test all(dqmc.ut_stack.backward_d_stack[i] .== 0)
                @test all(dqmc.ut_stack.backward_t_stack[i] .== 0)
            end
        end
    end

    @testset "build_stack forward/backward" begin
        MonteCarlo.build_stack(dqmc, dqmc.stack)
        MonteCarlo.build_stack(dqmc, dqmc.ut_stack)

        # test B(τ, 1) / B_l1 stacks
        @test dqmc.stack.u_stack ≈ dqmc.ut_stack.forward_u_stack
        @test dqmc.stack.d_stack ≈ dqmc.ut_stack.forward_d_stack
        @test dqmc.stack.t_stack ≈ dqmc.ut_stack.forward_t_stack

        # generate B(β, τ) / B_Nl stacks
        while dqmc.stack.direction == -1
            MonteCarlo.propagate(dqmc)
        end

        # test B(β, τ) / B_Nl stacks
        # Note: dqmc.stack doesn't generate the full stack here
        @test dqmc.stack.u_stack[2:end] ≈ dqmc.ut_stack.backward_u_stack[2:end]
        @test dqmc.stack.d_stack[2:end] ≈ dqmc.ut_stack.backward_d_stack[2:end]
        @test dqmc.stack.t_stack[2:end] ≈ dqmc.ut_stack.backward_t_stack[2:end]
    end

    while !(MonteCarlo.current_slice(dqmc) == 1 && dqmc.stack.direction == -1)
        MonteCarlo.propagate(dqmc)
    end

    @testset "equal time Greens function" begin
        # Check equal time greens functions from equal time and unequal time
        # calculation against each other
        for slice in 0:MonteCarlo.nslices(dqmc)
            G1 = deepcopy(MonteCarlo.calculate_greens(dqmc, slice))
            G2 = deepcopy(MonteCarlo.calculate_greens(dqmc, slice, slice))
            @test maximum(abs.(G1 .- G2)) < 1e-14
        end

        # Check G(t, 0) + G(0, beta - t) = 0
        for slice in 0:MonteCarlo.nslices(dqmc)-1
            G1 = MonteCarlo.greens(dqmc, slice, 0).val
            G2 = MonteCarlo.greens(dqmc, slice, MonteCarlo.nslices(dqmc)).val
            @test G1 ≈ -G2 atol = 1e-13
        end
    end

    # Check Iterators
    # For some reason Gkl/Gk0 errors jump in the last 10 steps
    # when using UnequalTimeStack

    # As in G(τ = Δτ * k, 0)
    # Note: G(0, l) = G(M-l, 0)
    Gk0s = [deepcopy(MonteCarlo.greens(dqmc, slice, 0).val) for slice in 0:MonteCarlo.nslices(dqmc)]

    @testset "GreensIterator" begin
        # Calculated from UnequalTimeStack (high precision)
        it = MonteCarlo.GreensIterator(dqmc, :, 0, dqmc.parameters.safe_mult)
        for (i, G) in enumerate(it)
            @test maximum(abs.(G.val .- Gk0s[i])) < 1e-14
        end

        # Calculated from mc.stack.greens using UDT decompositions (lower precision)
        it = MonteCarlo.GreensIterator(dqmc, :, 0, 4dqmc.parameters.safe_mult)
        for (i, G) in enumerate(it)
            @test maximum(abs.(G.val .- Gk0s[i])) < 1e-11
        end
    end

    Gkks = map(0:MonteCarlo.nslices(dqmc)) do slice
        g = MonteCarlo.calculate_greens(dqmc, slice)
        deepcopy(MonteCarlo._greens!(dqmc, dqmc.stack.greens_temp, g))
    end
    G0ks = [deepcopy(MonteCarlo.greens(dqmc, 0, slice).val) for slice in 0:MonteCarlo.nslices(dqmc)]
    MonteCarlo.calculate_greens(dqmc, 0) # restore mc.stack.greens

    @testset "CombinedGreensIterator" begin
        # high precision
        it = MonteCarlo.CombinedGreensIterator(dqmc, recalculate = dqmc.parameters.safe_mult)
        @test it isa CombinedGreensIterator
        @test it.recalculate == dqmc.parameters.safe_mult
        @test it.start == 1
        @test it.stop == MonteCarlo.nslices(dqmc)
        @test length(it) == it.stop - it.start + 1

        iter = MonteCarlo.init(dqmc, it)
        @test iter isa MonteCarlo._CombinedGreensIterator
        @test iter.mc == dqmc
        @test iter.spec == it
        @test length(iter) == length(it)

        for (i, (G0k, Gk0, Gkk)) in enumerate(iter)
            @test maximum(abs.(Gk0.val .- Gk0s[i+1])) < 1e-14
            @test maximum(abs.(G0k.val .- G0ks[i+1])) < 1e-14
            @test maximum(abs.(Gkk.val .- Gkks[i+1])) < 1e-14
        end

        # low precision
        it = MonteCarlo.CombinedGreensIterator(dqmc, recalculate = 4dqmc.parameters.safe_mult)
        for (i, (G0k, Gk0, Gkk)) in enumerate(MonteCarlo.init(dqmc, it))
            @test maximum(abs.(Gk0.val .- Gk0s[i+1])) < 1e-10
            @test maximum(abs.(G0k.val .- G0ks[i+1])) < 1e-10
            @test maximum(abs.(Gkk.val .- Gkks[i+1])) < 1e-10
        end
    end
end

@testset "Exact Greens comparison (Analytic, medium systems)" begin
    # These are theoretically the same but their implementation differs on
    # some level. To make sure both are correct it makes sense to check both here.
    models = (
        HubbardModelRepulsive(7, 2, U = 0.0, t = 1.0),
        HubbardModelAttractive(8, 2, U = 0.0, mu = 1.0, t = 1.0)
    )

    @info "Exact Greens comparison"
    for model in models, beta in (1.0, 10.0)
        @testset "$(typeof(model).name.name) β=$(beta)" begin
            Random.seed!(123)
            dqmc = DQMC(
                model, beta=beta, delta_tau = 0.1, safe_mult=5, recorder = Discarder(),
                thermalization = 1, sweeps = 2, measure_rate = 1
            )
            # @info "Running DQMC ($(typeof(model).name.name)) β=$(dqmc.parameters.beta)"

            dqmc[:G] = greens_measurement(dqmc, model)
            run!(dqmc, verbose=false)

            # error tolerance
            atol = 1e-13
            rtol = 1e-13
            N = length(lattice(model))

            # Direct calculation similar to what DQMC should be doing
            T = Matrix(MonteCarlo.hopping_matrix(dqmc, model))
            # Doing an eigenvalue decomposition makes this pretty stable
            vals, U = eigen(exp(-T))
            D = Diagonal(vals)^(dqmc.parameters.beta)

            # Don't believe "Quantum Monte Carlo Methods", this is the right
            # formula (believe dos Santos DQMC review instead)
            G = U * inv(I + D) * adjoint(U)

            @test check(mean(dqmc[:G]), G, atol, rtol)
        end
    end
end


text:
"
Maps `u` to an element (or array of elements) of [0,1]s by reflection.
"
function reflect(u::Float64)
    v = mod(u, 2.0)
    v > 1.0 ? 2.0-v : v
end

function reflect(u::Array{Float64, 1})
    Float64[reflect(x) for x in u]
end

"
Perform an iteration of slice sampling as in Murray and Graham.

## Arguments

- `u` Initial state vector of uniform variables
- `dosim` Function mapping a vector of uniforms to `sim, dist`, a simulation and a distance to the observations
- `epsilon` Acceptance threshold
- `w` Controls initial stepsize distribution

n.b. It's assumed that `dosim(u)[2] <= epsilon`.

## Returns

- A tuple `u, sim, z`: a new state vector, the corresponding simulation and the final stepsize.
"
function slicesample(u::Array{Float64, 1}, dosim::Function, epsilon::Float64; w=1.0)
    ##Initialisation
    p = length(u)
    v = randn(p)
    a = -rand()*w
    b = w + a
    ##Main loop
    while true
        z = a + rand() * (b-a)
        uprop = reflect(u+z*v)
        sim, dist = dosim(uprop) ##Profiling suggests this line doesn't cause type-stability problems
        if dist <= epsilon
            return uprop, sim, z
        end
        if z < 0
            a = z
        else
            b = z
        end
    end
end

ext: jl


text:
# Reexport SinkingParticles as they are useful outside too
#@reexport module SinkingParticles

using Unitful
using LinearAlgebra, SparseArrays
using OceanGrids

"""
    PFDO(grd; w_top)

Builds the particle-flux-divergence operator `PFDO` for a given particle sinking speed
(`w_top`).

Schematic of a grid cell:
```
     top  ┌─────────────────────────────────┐   ┬
          │      ↓ w_top          ↓ Φ_top   │   │
          │ (settling velovity)    (flux)   │   │
          │                                 │   │
          │            x                    │   δz
          │      (particle conc.)           │   │
          │                                 │   │
          │                                 │   │
  bottom  └─────────────────────────────────┘   ┴
```

- `δz` is the height of grid cell [m]
- `w_top` is the particle sinking speed at the top of the grid cell [m s⁻¹]
- `Φ_top` is the flux at the top of the grid cell [mol m⁻² s⁻¹]
- `x` is the particle concentration of the cell [mol m⁻³]

The PFDO is defined by

    PFDO * x = δΦ/δz ≈ dΦ/dz,

i.e., so that applied to `x` it approximates the flux divergence of `x`, dΦ/δz.
It is calculated as the matrix product of DIVO and FATO:

    PFDO = DIVO * FATO.

where the divergence operator `DIVO`, is defined by (`z` increasing downward)

    DIVO * ϕ_top = 1/δz * (ϕ_top[below] - ϕ_top) = δϕ/δz ≈ dΦ/δz.

and FATO, the flux-at-top operator, is defined by

    FATO * x = w_top * x[above] ≈ ϕ_top.

# Example

```julia-repl
julia> PFDO(grd, w_top=1.0) # 1.0 m/s (SI units assumed)
```
"""
function PFDO(grd; w_top)
    iwet = findall(vec(grd.wet3D))
    DIVO = DIVO(grd)
    Iabove = buildIabove(grd.wet3D, iwet)
    w_top = ustrip.(upreferred.(w_top))
    return PFDO(w_top, DIVO, Iabove)
end
"""
    PFDO(w, DIVO, Iabove)

Returns `DIVO * FATO(w, Iabove)`

This function is useful to avoid reconstructing `DIVO` and `Iabove` every time.
It should allow for faster runs.
"""
PFDO(w_top, DIVO, Iabove) = DIVO * FATO(w_top, Iabove)

"""
    PFDO(grd, δz, w_top, w_bot, frac_seafloor, cumfrac_seafloor, fsedremin, Iabove)

Returns the particle-flux-divergence operator for a given sinking speed as a function of depth.

This is a slightly different construction where I take in top and bottom settling velocities,
and where the bottom one can be modified to further allow a fraction of particles to sink through
(buried into) the sea floor.

Below is a detailed explanation of how this function computes the particle-flux divergence.
Take these 3 boxes on top of each other:

```
┌──────────────────┐
│ water            │
│←----c----→       │
├───────────┲━━━━━━┥ ←- Φ_top
│           ┃⣿⣿⣿⣿⣿⣿│
│           ┃⣿⣿⣿⣿⣿⣿│
│←-d-→ ←-a-→┃⣿⣿⣿⣿⣿⣿│
├─────┲━━━━━┹──────┤ ←- Φ_bot
│     ┃←----b-----→│
│     ┃⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿│
│     ┃⣿⣿⣿ land ⣿⣿⣿│
│     ┃⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿│
└━━━━━┹────────────┘
```

A part of each of these boxes is water, while the rest is land. For the middle box,
we will use the fractional area `a = frac_seafloor` of seafloor,
and the cumulated fractional area `b = cumfrac_seafloor` of seafloor, to determine the
flux of particles at the top `ϕ_top` and bottom `ϕ_bot`.

From the top box, only the particles in the water can enter the middle box. So the flux
at the top of the middle box, `ϕ_top`, is proportional to `c = 1 - Iabove * b`.
At the bottom of the middle box, we have to take care of the case of particles hitting
the sediments and potentially buried there, or kept in the water.
For the part going through the area `d`, the flux is proportional to `d = 1 - b`.
For the part going through `a` (hitting the sediments), the flux is proportional to
`a * (1 - f)` where `f` is the fraction of particles forcibly kept in the water.
(So if `f = 0`, the particles appear to move out of the middle box, but are not carried
into the bottom one, and thus are removed from the system / buried.)
"""
function PFDO(grd, δz, w_top, w_bot, frac_seafloor, cumfrac_seafloor, fsedremin, Iabove)
    fw_bot = @. (1.0 - cumfrac_seafloor + (1.0 - fsedremin) * frac_seafloor) * w_bot
    fw_top = (1.0 .- Iabove * cumfrac_seafloor) .* w_top
    return sparse(Diagonal(fw_bot ./ δz)) - sparse(Diagonal(fw_top ./ δz)) * Iabove
end

"""
    DIVO(grd)

Build the `DIVO` operator such that

    DIVO * ϕ_top = 1/δz * (ϕ_top - ϕ_top[below]) ≈ dΦ/δz.
"""
function DIVO(grd)
    Ibelow = buildIbelow(grd)
    iwet = indices_of_wet_boxes(grd)
    δz = ustrip.(grd.δz_3D[iwet])
    return sparse(Diagonal(1 ./ δz)) * (Ibelow - I) # divergence with positive downwards
end

"""
    FATO(w_top, Iabove)

Build the `FATO` operator for a particle sinking speed `w_top`

(`w_top` is the sinking speed at the top of each grid cell.)

The `FATO` operator is defined by

    FATO * x = w_top * x(above) ≈ ϕ_top.
"""
FATO(w_top::Vector, Iabove) = sparse(Diagonal(w_top)) * Iabove
FATO(w_top::Number, Iabove) = w_top * Iabove

export DIVO, PFDO, FATO
#end # module

"""
    transportoperator(grd, w)

Returns the transportoperator for the given settling velocity `w`.

The settling velocity can be provided as either a scalar
(e.g., `w = 100.0` in units of meters per second)
or as a function of depth
(e.g., `w(z) = 2z + 1`).

# Examples

Create the particle flux divergence with settling velocity of 100m/s

```julia-repl
julia> T = transportoperator(grd, 100.0)
```

Or with settling velocity function w(z) = 2z + 1

```julia-repl
julia> T = transportoperator(grd, z -> 2z + 1)
```

By default, the seafloor flux is set to zero, so that all the particles
that reach it are remineralized there. You can let particles go through
by setting `fsedremin=0.0`, via, e.g.,

```julia-repl
julia> T = transportoperator(grd, z -> 2z + 1; fsedremin=0.0)
```

For finer control and advanced use, see the particle-flux divergence
operator function, `PFDO`.
"""
transportoperator(grd, w_top; DIVop=DIVO(grd), Iabove=buildIabove(grd)) = PFDO(w_top, DIVop, Iabove)
function transportoperator(grd, w::Function;
              δz = ustrip.(grd.δz_3D[iswet(grd)]),
              Iabove = buildIabove(grd),
              fsedremin = 1.0,
              z_top = topdepthvec(grd),
              z_bot = bottomdepthvec(grd),
              frac_seafloor = float.(isseafloorvec(grd)),
              cumfrac_seafloor = zcumsum(frac_seafloor, grd))
    return PFDO(grd, δz, ustrip.(upreferred.(w.(z_top))), ustrip.(upreferred.(w.(z_bot))), frac_seafloor, cumfrac_seafloor, fsedremin, Iabove)
end

export transportoperator


text:
using BasicPOMCP
using Test
using POMDPs
using POMDPModels
using NBInclude
using D3Trees
using Random
using POMDPSimulators
using POMDPModelTools
using POMDPTesting

test_solver(POMCPSolver(), BabyPOMDP())

pomdp = BabyPOMDP()
solver = POMCPSolver(rng = MersenneTwister(1))
planner = solve(solver, pomdp)

tree = BasicPOMCP.POMCPTree(pomdp, solver.tree_queries)
node = BasicPOMCP.POMCPObsNode(tree, 1)

r = @inferred BasicPOMCP.simulate(planner, initialstate(pomdp, MersenneTwister(1)), node, 20)

sim = HistoryRecorder(max_steps=10)
simulate(sim, pomdp, planner, updater(pomdp))

solver = POMCPSolver(max_time=0.1, tree_queries=typemax(Int), rng = MersenneTwister(1))
planner = solve(solver, pomdp)
a, info = action_info(planner, initialstate_distribution(pomdp))
a, info = action_info(planner, initialstate_distribution(pomdp))
println("time below should be about 0.1 seconds")
etime = @elapsed a, info = action_info(planner, initialstate_distribution(pomdp))
@show etime
@test etime < 0.2
@show info[:search_time_us]

solver = POMCPSolver(max_time=0.1, tree_queries=typemax(Int), rng = MersenneTwister(1))
planner = solve(solver, pomdp)
a, info = action_info(planner, initialstate_distribution(pomdp), tree_in_info=true)

#d3t = D3Tree(planner, title="test")
d3t = D3Tree(info[:tree], title="test tree")
# inchrome(d3t)
show(stdout, MIME("text/plain"), d3t)

solver = POMCPSolver(max_time=0.1, tree_queries=typemax(Int), rng=MersenneTwister(1), tree_in_info=true)
planner = solve(solver, pomdp)
a, info = action_info(planner, initialstate_distribution(pomdp))

d3t = D3Tree(info[:tree], title="test tree (tree_in_info solver option)")

@nbinclude(joinpath(dirname(@__FILE__), "..", "notebooks", "Minimal_Example.ipynb"))

#d3t = D3Tree(planner, title="test")
# inchrome(d3t)

@testset "consistency" begin
    # test consistency when rng is specified
    pomdp = BabyPOMDP()
    solver = POMCPSolver(rng = MersenneTwister(1))
    planner = solve(solver, pomdp)
    hist1 = simulate(HistoryRecorder(max_steps=1000, rng=MersenneTwister(3)), pomdp, planner)

    solver = POMCPSolver(rng = MersenneTwister(1))
    planner = solve(solver, pomdp)
    hist2 = simulate(HistoryRecorder(max_steps=1000, rng=MersenneTwister(3)), pomdp, planner)

    @test discounted_reward(hist1) == discounted_reward(hist2)
end

@testset "requires" begin
    # REQUIREMENTS
    solver = POMCPSolver()
    pomdp = TigerPOMDP()

    println("============== @requirements_info with only solver:")
    @requirements_info solver
    println("============== @requirements_info with solver and pomdp:")
    @requirements_info solver pomdp
end


text:
"""
    visualize(mechanism, storage; vis, build, show_contact, animation, color, name)

    visualize mechanism using trajectory from storage

    mechanism: Mechanism
    storage: Storage
    vis: Visualizer
    build: flag to construct mechanism visuals (only needs to be built once)
    show_contact: flag to show contact locations on system
    color: RGBA
    name: unique identifier for mechanism
"""
function visualize(mechanism::Mechanism, storage::Storage{T,N}; vis::Visualizer=Visualizer(),
    build::Bool=true,
    show_joint=false,
    joint_radius=0.1,
    show_contact=false,
    animation=nothing,
    color=nothing,
    name::Symbol=:robot) where {T,N}

    storage = deepcopy(storage)
    bodies = mechanism.bodies
    origin = mechanism.origin

    # Build robot in the visualizer
    build && build_robot(mechanism,
        vis=vis,
        show_joint=show_joint,
        show_contact=show_contact,
        color=color,
        name=name)

    # Create animations
    framerate = Int64(round(1/mechanism.timestep))
    (animation == nothing) && (animation =
        MeshCat.Animation(Dict{MeshCat.SceneTrees.Path,MeshCat.AnimationClip}(), framerate))

    # Bodies and Contacts
    for (id, body) in enumerate(bodies)
        shape = body.shape
        visshape = convert_shape(shape)
        subvisshape = nothing
        showshape = false
        if visshape !== nothing
            subvisshape = vis[name][:bodies][Symbol(body.name, "__id_$id")]
            showshape = true
        end

        animate_node!(storage, id, shape, animation, subvisshape, showshape)

        if show_joint
            for (jd, joint) in enumerate(mechanism.joints)
                if joint.child_id == body.id
                    radius = 0.1
                    joint_shape = Sphere(radius,
                        position_offset=joint.translational.vertices[2],
                        color=RGBA(0.0, 0.0, 1.0, 0.5))
                    visshape = convert_shape(joint_shape)
                    subvisshape = nothing
                    showshape = false
                    if visshape !== nothing
                        subvisshape = vis[name][:joints][Symbol(joint.name, "__id_$(jd)")]
                        showshape = true
                    end
                    animate_node!(storage, id, joint_shape, animation, subvisshape, showshape)
                end
            end
        end

        if show_contact
            for (jd, contact) in enumerate(mechanism.contacts)
                if contact.parent_id == body.id
                    radius = abs(contact.model.collision.contact_radius)
                    (radius == 0.0) && (radius = 0.01)
                    contact_shape = Sphere(radius,
                        position_offset=contact.model.collision.contact_origin, #TODO: generalize for collision checking
                        axis_offset=one(Quaternion),
                        color=RGBA(1.0, 0.0, 0.0, 0.5))
                    visshape = convert_shape(contact_shape)
                    subvisshape = nothing
                    showshape = false
                    if visshape !== nothing
                        subvisshape = vis[name][:contacts][Symbol(contact.name, "__id_$(jd)")]
                        showshape = true
                    end
                    animate_node!(storage, id, contact_shape, animation, subvisshape, showshape)
                end
            end
        end
    end

    # Origin
    id = origin.id
    shape = origin.shape
    visshape = convert_shape(shape)
    if visshape !== nothing
        subvisshape = vis[name][:bodies][Symbol(:origin, "_id")]
        shapetransform = transform(szeros(T,3), one(Quaternion{T}), shape)
        settransform!(subvisshape, shapetransform)
    end

    setanimation!(vis, animation)
    return vis, animation
end

"""
    build_robot(mechanism; vis, show_contact, name, color)

    construct visuals for mechanism

    mechanism: Mechanism
    vis: Visualizer
    show_contact: flag to show contact locations on mechanism
    name: unique identifier
    color: RGBA
"""
function build_robot(mechanism::Mechanism;
    vis::Visualizer=Visualizer(),
    show_joint=false,
    joint_radius=0.1,
    show_contact=false,
    name::Symbol=:robot,
    color=nothing) where {T,N}

    bodies = mechanism.bodies
    origin = mechanism.origin
    set_background!(vis)
    set_light!(vis)
    set_floor!(vis)

    # Bodies and Contacts
    for (id, body) in enumerate(bodies)
        if color !== nothing
            shape = deepcopy(body.shape)
            set_color!(shape, color)
        else
            shape = body.shape
        end
        visshape = convert_shape(shape)
        subvisshape = nothing
        if visshape !== nothing
            subvisshape = vis[name][:bodies][Symbol(body.name, "__id_$id")]
            setobject!(subvisshape, visshape, shape,
                transparent=(show_joint || show_contact))
        end

        if show_joint
            for (jd, joint) in enumerate(mechanism.joints)
                if joint.child_id == body.id
                    radius = joint_radius
                    joint_shape = Sphere(radius,
                        position_offset=joint.translational.vertices[2],
                        color=RGBA(0.0, 0.0, 1.0, 0.5))
                    visshape = convert_shape(joint_shape)
                    subvisshape = nothing
                    if visshape !== nothing
                        subvisshape = vis[name][:joints][Symbol(joint.name, "__id_$(jd)")]
                        setobject!(subvisshape, visshape, joint_shape,
                            transparent=false)
                    end
                end
            end
        end

        if show_contact
            for (jd, contact) in enumerate(mechanism.contacts)
                if contact.parent_id == body.id
                    radius = abs(contact.model.collision.contact_radius)
                    (radius == 0.0) && (radius = 0.01)
                    contact_shape = Sphere(radius,
                        position_offset=(contact.model.collision.contact_origin),
                        axis_offset=one(Quaternion), color=RGBA(1.0, 0.0, 0.0, 0.5))
                    visshape = convert_shape(contact_shape)
                    subvisshape = nothing
                    if visshape !== nothing
                        subvisshape = vis[name][:contacts][Symbol(contact.name, "__id_$(jd)")]
                        setobject!(subvisshape,visshape,contact_shape,transparent=false)
                    end
                end
            end
        end
    end

    # Origin
    id = origin.id
    shape = origin.shape
    visshape = convert_shape(shape)
    if visshape !== nothing
        subvisshape = vis[name][:bodies][Symbol(:origin, "_id")]
        setobject!(subvisshape,visshape,shape,transparent=(show_joint || show_contact))
    end
    return vis
end

"""
    set_robot(vis, mechanism, z; show_contact, name)

    visualze mechanism configuration from maximal representation

    vis: Visualizer
    mechanism: Mechanism
    z: maximal state
    show_contact: flag to show contact locations on mechanism
    name: unique identifier
"""
function set_robot(vis::Visualizer, mechanism::Mechanism, z::Vector{T};
    show_joint::Bool=false,
    joint_radius=0.1,
    show_contact::Bool=true,
    name::Symbol=:robot) where {T,N}

    (length(z) == minimal_dimension(mechanism)) && (z = minimal_to_maximal(mechanism, z))
    bodies = mechanism.bodies
    origin = mechanism.origin

    # Bodies and Contacts
    for (id, body) in enumerate(bodies)
        x, _, q, _ = unpack_maximal_state(z, id)
        shape = body.shape
        visshape = convert_shape(shape)
        subvisshape = nothing
        showshape = false
        if visshape !== nothing
            subvisshape = vis[name][:bodies][Symbol(body.name, "__id_$id")]
            showshape = true
        end

        set_node!(x, q, id, shape, subvisshape, showshape)

        if show_joint
            for (jd, joint) in enumerate(mechanism.joints)
                if joint.child_id == body.id
                    radius = joint_radius
                    joint_shape = Sphere(radius,
                        position_offset=joint.translational.vertices[2],
                        color=RGBA(0.0, 0.0, 1.0, 0.5))
                    visshape = convert_shape(joint_shape)
                    subvisshape = nothing
                    showshape = false
                    if visshape !== nothing
                        subvisshape = vis[name][:joints][Symbol(joint.name, "__id_$(jd)")]
                        showshape = true
                    end
                    set_node!(x, q, id, joint_shape, subvisshape, showshape)
                end
            end
        end

        if show_contact
            for (jd, contact) in enumerate(mechanism.contacts)
                if contact.parent_id == body.id
                    radius = abs(contact.model.collision.contact_radius)
                    (radius == 0.0) && (radius = 0.01)
                    contact_shape = Sphere(radius,
                        position_offset=(contact.model.collision.contact_origin),
                        axis_offset=one(Quaternion), color=RGBA(1.0, 0.0, 0.0, 0.5))
                    visshape = convert_shape(contact_shape)
                    subvisshape = nothing
                    showshape = false
                    if visshape !== nothing
                        subvisshape = vis[name][:contacts][Symbol(contact.name, "__id_$(jd)")]
                        showshape = true
                    end
                    set_node!(x, q, id, contact_shape, subvisshape, showshape)
                end
            end
        end
    end

    # Origin
    id = origin.id
    shape = origin.shape
    visshape = convert_shape(shape)
    if visshape !== nothing
        subvisshape = vis[name][:bodies][Symbol(:origin, "_id")]
        shapetransform = transform(szeros(T,3), one(Quaternion{T}), shape)
        settransform!(subvisshape, shapetransform)
    end
    return vis
end

function transform(x, q, shape)
    scale_transform = MeshCat.LinearMap(diagm(shape.scale))
    x_transform = MeshCat.Translation(x + vector_rotate(shape.position_offset, q))
    q_transform = MeshCat.LinearMap(q * shape.axis_offset)
    return MeshCat.compose(x_transform, q_transform, scale_transform)
end

MeshCat.js_scaling(s::AbstractVector) = s
MeshCat.js_position(p::AbstractVector) = p

function set_node!(x, q, id, shape, shapevisualizer, showshape) where {T,N}
    if showshape
        # TODO currently setting props directly because MeshCat/Rotations doesn't convert scaled rotation properly.
        # If this changes, do similarily to origin
        setprop!(shapevisualizer, "scale", MeshCat.js_scaling(shape.scale))
        setprop!(shapevisualizer, "position", MeshCat.js_position(x + vector_rotate(shape.position_offset, q)))
        setprop!(shapevisualizer, "quaternion", MeshCat.js_quaternion(q * shape.axis_offset))
    end
    return
end

function animate_node!(storage::Storage{T,N}, id, shape, animation, shapevisualizer, showshape) where {T,N}
    for i=1:N
        x = storage.x[id][i]
        q = storage.q[id][i]
        atframe(animation, i) do
            set_node!(x, q, id, shape, shapevisualizer, showshape)
        end
    end
    return
end

function MeshCat.setobject!(subvisshape, visshape, shapes::Shapes;
    transparent=false)
    for (i, s) in enumerate(shapes.shape)
        v = subvisshape["node_$i"]
        setobject!(v, visshape[i], s, transparent=transparent)
        scale_transform = MeshCat.LinearMap(diagm(s.scale))
        x_transform = MeshCat.Translation(s.position_offset)
        q_transform = MeshCat.LinearMap(s.axis_offset)
        t = MeshCat.compose(x_transform, q_transform, scale_transform)
        settransform!(v, t)
    end
end

function MeshCat.setobject!(subvisshape, visshape, shape::Shape;
    transparent=false)
    setobject!(subvisshape, visshape, MeshPhongMaterial(color=(transparent ? RGBA(0.75, 0.75, 0.75, 0.5) : shape.color)))
end

function MeshCat.setobject!(subvisshape, visshape::Vector, shape::Capsule; transparent=false)
    setobject!(subvisshape["cylinder"], visshape[1], MeshPhongMaterial(color=(transparent ? RGBA(0.75, 0.75, 0.75, 0.5) : shape.color)))
    setobject!(subvisshape["cap1"], visshape[2], MeshPhongMaterial(color=(transparent ? RGBA(0.75, 0.75, 0.75, 0.5) : shape.color)))
    setobject!(subvisshape["cap2"], visshape[3], MeshPhongMaterial(color=(transparent ? RGBA(0.75, 0.75, 0.75, 0.5) : shape.color)))
end

function MeshCat.setobject!(subvisshape, visshape, shape::Mesh; transparent=false)
    if visshape.mtl_library == ""
        visshape = MeshFileGeometry(visshape.contents, visshape.format)
        setobject!(subvisshape, visshape, MeshPhongMaterial(color=shape.color))
    else
        setobject!(subvisshape, visshape)
    end
end


text:
# Precompile
# ==========


text:
module GExt
""" Perform extrapolation of density matrices using their geometric structure with
  least-squares method.

  nbas: number of basis functions
  nocc: number of occupied orbitals (half the number of electrons)
  nmat: number of elements used to perform the extrapolation plus one (the first element is
        supposed to contain the point we use as a reference, and we do not need its
        associated positions)
  nqm:  number of QM atoms
"""

using LinearAlgebra, DelimitedFiles, Printf
# To make sure the guesses we write are precise enough
Base.show(io::IO, f::Float64) = @printf(io, "%24.16e", f)

export gext

function gext(; ε=1e-5)
  """ Main function which does the extrapolation """
  cmat = get_cmat("cmat.dat")
  smat = get_smat("smat.dat")
  hpos = get_hpos("hpos.dat")

  # Make sure that we have the correct number of elements
  @assert size(smat, 2) == size(cmat, 3) == size(hpos, 3)
  nbas, nocc, nmat = size(cmat)

  # Perform orthonormalisation
  C = Array{Float64}(undef, nbas, nocc, nmat)
  for i in 1:nmat
    C[:, :, i] = sqrt(tril_to_sym(smat[:, i]))*cmat[:, :, i]
  end

  # The reference point
  γ_ref = C[:, :, 1]

  # Project the points on the tangent space at γ_ref
  Γs = logGrs(γ_ref, C)

  # Create a matrix of previous descriptors
  poly = [get_coulomb(hpos[:, :, imat]) for imat in 1:nmat-1]
  P = [poly[i][j] for i in eachindex(poly), j in eachindex(poly[1])]

  # Descriptor of the current point
  poly_test = get_coulomb(hpos[:, :, nmat])

  # Optional stablisation
  ds = hcat(poly_test, zeros(nmat-1))
  Ps = hcat(P, ε*Matrix(I, nmat-1, nmat-1))

  # Least-squares on the Coulomb descriptors
  α = ds'*pinv(Ps)

  # Extrapolation on the tangent space
  Γguess = sum(α[i]*Γs[:, :, i+1] for i in 1:nmat-1)
  # Retract the point back to the Grassmannian
  Cguess = expGr(γ_ref, Γguess)

  # Compute the idempotent guess
  pguess = 2*Cguess*Cguess'
  lguess = vec_tril(pguess)

  # Write the guess on file
  writedlm("guess", lguess)
end

function get_coulomb(hpos)
  """ Return the Coulomb descriptor, which is the lower triangular part of the symmetric
  Coulomb matrix """
  charges = hpos[1, :]
  positions = hpos[2:4, :]

  nqm = size(positions, 2)
  G = Matrix{Float64}(undef, nqm, nqm)
  for i in 1:nqm
    for j in 1:nqm
      if i == j
        G[i, i] = 0.5*charges[i]^(2.4)
      else
        G[i, j] = charges[i]*charges[j]/norm(positions[:, i] - positions[:, j])
      end
    end
  end

  return vec_tril(G)
end

#############################################
#### Exponential and logarithm functions ####
#############################################

function expGr(γ_ref, Γ)
  """ Perform the exponential of the tangent vector Γ at base point γ_ref """
  d = svd(Γ)
  Y = γ_ref * d.V * cos(Diagonal(d.S)) + (d.U) * sin(Diagonal(d.S))
  # Optional, to make sure that rounding errors do not prevent Y from being orthogonal
  return Matrix(qr(Y).Q)
end

function logGr(γ_ref, γ)
  """ Perform the logarithm of the point γ at base point γ_ref """
  Z = γ'*γ_ref
  W = Z \ (γ'-Z*γ_ref')
  d = svd(W, full = false)
  A = (d.V)*atan(Diagonal(svd(W, full = false).S))*(d.U)'
  return A
end

function logGrs(γ_ref, γs)
  """ Perform the logarithms on a list of points """
  nbas, nocc, nmat = size(cmat)

  Γs = Array{Float64}(undef, nbas, nocc, nmat)
  for iγ in 1:nmat
    Γs[:, :, iγ] = logGr(γ_ref, γs[:, :, iγ])
  end

  return Γs
end

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

function tril_to_sym(v)
  """ Convert a vector representing the lower triangular part of a symmetric matrix to
  a symmetric matrix """
  n = Int(floor(sqrt(2*length(v))))
  M = zeros(n, n)

  k = 0
  for i in 1:n
    for j in 1:i
      k += 1
      M[i, j] = v[k]
      M[j, i] = v[k]
    end
  end

  return M
end

function vec_tril(M)
  """ Extract the the lower triangular part of a square matrix """
  m, n = size(M)
  @assert m == n
  l = n*(n+1) ÷ 2
  v = zeros(l)

  k = 0
  for i in 1:n
    for j in 1:i
      v[k + j] = M[i, j]
    end
    k += i
  end

  return v
end

#############################################
####      Data processing functions      ####
#############################################

function get_cmat(file)
  """ Placeholder for a function that returns an array cmat of size nbas×nocc×nmat
  containing molecular orbitals matrices """
end

function get_smat(file)
  """ Placeholder for a function that returns an array smat of size nbas×(nbas+1)/2×nmat
  containing lower diagonal of overlap matrices """
end

function get_hpos(file)
  """ Placeholder for a function that returns an array smat of size 4×nqm×nmat
  containing charges and positions of the QM atoms for the previous positions as well as
  the current for which we want the extrapolation """
end

end

max_stars_repo_name: epolack/gext

max_issues_repo_name: epolack/gext

max_forks_repo_name: epolack/gext


text:
using DataFrames
using CSV
using PGFPlotsX

push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\pgfplotsset{compat=1.7}")
df = DataFrame(CSV.File("benchmark-results/part2_semi_implicit_vs_explicit_experiment_results.csv"))

beta(b) = df -> filter(:beta => ==(b), df)

@pgf common_opts = {
  title_style={align="center", font=raw"\large"},
  xlabel_style={align="center", font=raw"\Large",},
  ylabel=raw"$\mathit{Perf} = \frac{\mathit{Work}}{t_{\rm elapsed}}\ \left[GFlop/s\right]$",
  ylabel_style={align="center", font=raw"\Large"},
  # log_ticks_with_fixed_point,
  grid="major",
  legend_pos="outer north east",
}

ax1 = LogLogAxis(@pgf({
                      common_opts...,
                      legend_pos="south east",
                      title=raw"Navier-Stokes Solver Timestepping Schemes \\ $nx=2^{11}, ny=2^9, \mathit{ttot}=0.005s$",
                      xlabel = raw"Prandtl number \\ $Pr = \frac{\mathit{kinematic\ viscosity}}{\mathit{thermal\ diffusivity}}$",
                      ylabel="Timesteps"}),
          PlotInc(Table(beta(0.0)(df).Pr, beta(0.0)(df).timed_iters)),
          PlotInc(Table(beta(0.5)(df).Pr, beta(0.5)(df).timed_iters)),
          PlotInc(Table(beta(1.0)(df).Pr, beta(1.0)(df).timed_iters)),
          Legend([raw"explicit ($\beta=0$)", raw"semi-implicit ($\beta = 0.5$)", raw"implicit ($\beta = 1.0$)"])
)

ax2 = LogLogAxis(@pgf({
                      common_opts...,
                      legend_pos="south east",
                      title=raw"Navier-Stokes Solver Timestepping Schemes \\ $nx=2^{11}, ny=2^9, \mathit{ttot}=0.005s$",
                      xlabel = raw"Prandtl number \\ $Pr = \frac{\mathit{kinematic\ viscosity}}{\mathit{thermal\ diffusivity}}$",
                      ylabel="Time [s]"}),
          PlotInc(Table(beta(0.0)(df).Pr, beta(0.0)(df).t_elapsed)),
          PlotInc(Table(beta(0.5)(df).Pr, beta(0.5)(df).t_elapsed)),
          PlotInc(Table(beta(1.0)(df).Pr, beta(1.0)(df).t_elapsed)),
          Legend([raw"explicit ($\beta=0$)", raw"semi-implicit ($\beta = 0.5$)", raw"implicit ($\beta = 1.0$)"])
)

gp = GroupPlot(@pgf({group_style={group_size="2 by 1", horizontal_sep="7cm", vertical_sep="1.8cm"}}))
push!(gp, ax1, ax2)
fig = TikzPicture(gp)

pgfsave("figs/navier_stokes_timestepping_schemes.png", fig)
pgfsave("figs/navier_stokes_timestepping_schemes.pdf", fig)


text:
"""
    BSplineBasis{k}

B-spline basis for splines of order `k`.

The basis is defined by a set of knots and by the B-spline order.

---

    BSplineBasis(k, breakpoints::AbstractVector; augment = Val(true))

Create B-spline basis of order `k` with the given breakpoints.

If `augment = Val(true)`, breakpoints will be "augmented" so that boundary knots
have multiplicity `k`. Note that, if they are passed as a regular `Vector`, the
input may be modified. See [`augment_knots!`](@ref) for details.

# Examples

```jldoctest BSplineBasis
julia> breaks = range(-1, 1, length = 21)
-1.0:0.1:1.0

julia> B = BSplineBasis(5, breaks)
24-element BSplineBasis of order 5, domain [-1.0, 1.0]
 knots: [-1.0, -1.0, -1.0, -1.0, -1.0, -0.9, -0.8, -0.7, -0.6, -0.5  …  0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.0, 1.0, 1.0, 1.0]
```

This is equivalent to the above:

```jldoctest BSplineBasis
julia> B = BSplineBasis(BSplineOrder(5), breaks)
24-element BSplineBasis of order 5, domain [-1.0, 1.0]
 knots: [-1.0, -1.0, -1.0, -1.0, -1.0, -0.9, -0.8, -0.7, -0.6, -0.5  …  0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.0, 1.0, 1.0, 1.0]
```

Note that first and last knots are repeated ``k = 5`` times.

If `augment = Val(false)`, input breakpoints are taken without modification as
the knots ``t_i`` of the B-spline basis. Note that the valid domain is reduced to
``[-0.6, 0.6]``. The domain is always defined as the range ``[t_k, t_{N + 1}]``,
where ``N`` is the length of the basis (below, ``N = 16``).

```jldoctest BSplineBasis
julia> Bn = BSplineBasis(5, breaks, augment = Val(false))
16-element BSplineBasis of order 5, domain [-0.6, 0.6]
 knots: -1.0:0.1:1.0
```

"""
struct BSplineBasis{k, T, Knots <: AbstractVector{T}} <: AbstractBSplineBasis{k,T}
    N :: Int    # number of B-splines
    t :: Knots  # knots (length = N + k)
    function BSplineBasis(
            ::BSplineOrder{k}, knots::AbstractVector{T};
            augment::Val{Augment} = Val(true),
        ) where {k,T,Augment}
        k :: Integer
        Augment :: Bool
        if k <= 0
            throw(ArgumentError("B-spline order must be k ≥ 1"))
        end
        # TODO when `knots` is a regular Vector, it would be nice to avoid
        # modifying the vector and to instead use an AugmentedKnots instance
        # (commented line below).
        # However, for now that seems to be sligthly slower when evaluating
        # splines, not sure why...
        t = Augment ? augment_knots!(knots, k) : knots
        # t = Augment ? AugmentedKnots{k}(knots) : knots
        N = length(t) - k
        Knots = typeof(t)
        new{k, T, Knots}(N, t)
    end
end

@inline BSplineBasis(k::Integer, args...; kwargs...) =
    BSplineBasis(BSplineOrder(k), args...; kwargs...)

Base.:(==)(A::BSplineBasis, B::BSplineBasis) =
    A === B ||
    order(A) == order(B) && knots(A) == knots(B)

"""
    getindex(B::AbstractBSplineBasis, i, [T = Float64])

Get ``i``-th basis function.

This is an alias for `BasisFunction(B, i, T)` (see [`BasisFunction`](@ref) for details).

The returned object can be evaluated at any point within the boundaries defined
by the basis.

# Examples

```jldoctest
julia> B = BSplineBasis(BSplineOrder(4), -1:0.1:1)
23-element BSplineBasis of order 4, domain [-1.0, 1.0]
 knots: [-1.0, -1.0, -1.0, -1.0, -0.9, -0.8, -0.7, -0.6, -0.5, -0.4  …  0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.0, 1.0, 1.0]

julia> B[6]
Basis function i = 6
  from 23-element BSplineBasis of order 4, domain [-1.0, 1.0]
  support: [-0.8, -0.4) (knots 6:10)

julia> B[6](-0.5)
0.16666666666666666

julia> B[6, Float32](-0.5)
0.16666667f0

julia> B[6](-0.5, Derivative(1))
-5.000000000000001
```
"""
@inline function Base.getindex(
        B::AbstractBSplineBasis, i::Integer, ::Type{T} = Float64) where {T}
    @boundscheck checkbounds(B, i)
    BasisFunction(B, i, T)
end

@inline function Base.checkbounds(B::AbstractBSplineBasis, I)
    checkbounds(eachindex(B), I)
end

@inline Base.eachindex(B::AbstractBSplineBasis) = Base.OneTo(length(B))

@inline function Base.iterate(B::AbstractBSplineBasis, i = 0)
    i == length(B) && return nothing
    i += 1
    B[i], i
end

function Base.show(io::IO, B::AbstractBSplineBasis)
    summary(io, B)
    let io = IOContext(io, :compact => true, :limit => true)
        print(io, "\n knots: ", knots(B))
    end
    nothing
end

Base.summary(io::IO, B::BSplineBasis) = summary_basis(io, B)

function summary_basis(io, B::AbstractBSplineBasis)
    a, b = boundaries(B)
    print(io, length(B), "-element ", nameof(typeof(B)))
    print(io, " of order ", order(B), ", domain [", a, ", ", b, "]")
    nothing
end

# Make BSplineBasis behave as scalar when broadcasting.
Broadcast.broadcastable(B::AbstractBSplineBasis) = Ref(B)

"""
    length(g::BSplineBasis)

Returns the number of B-splines composing a spline.
"""
Base.length(g::BSplineBasis) = g.N
Base.size(g::AbstractBSplineBasis) = (length(g), )
Base.parent(g::BSplineBasis) = g

"""
    boundaries(B::AbstractBSplineBasis)

Returns `(xmin, xmax)` tuple with the boundaries of the domain supported by the
basis.
"""
function boundaries(B::BSplineBasis)
    k = order(B)
    N = length(B)
    t = knots(B)
    t[k], t[N + 1]
end

"""
    knots(g::BSplineBasis)
    knots(g::Spline)

Returns the knots of the B-spline basis.
"""
knots(g::BSplineBasis) = g.t

"""
    order(::Type{BSplineBasis}) -> Int
    order(::Type{Spline}) -> Int
    order(::BSplineOrder) -> Int

Returns order of B-splines as an integer.
"""
order(::Type{<:BSplineBasis{k}}) where {k} = k
order(b::AbstractBSplineBasis) = order(typeof(b))
order(::BSplineOrder{k}) where {k} = k

max_stars_repo_name: jipolanco/BasisSplines

max_issues_repo_name: jipolanco/BasisSplines

max_forks_repo_name: jipolanco/BasisSplines


text:
using ARules

transactions = [["milk", "eggs", "bread"],
			    ["butter", "milk", "sugar", "flour", "eggs"],
			    ["bacon", "eggs", "milk", "beer"],
			    ["bread", "ham", "turkey"],
			    ["cheese", "ham", "bread", "ketchup"],
			    ["mustard", "hot dogs", "buns", "hamburger", "cheese", "beer"]]

# freq = frequent(transactions, 1, 8)
# display(freq)
#
# for x in freq[:itemset]
#     println(x)
# end

unq = unique_items(transactions)

@time node = ARules.frequent_item_tree(transactions, unq, 1, 8)

sup = gen_support_dict(node, 9)

for k in keys(sup)
    println(k)
end

rules = apriori(transactions, supp = 0.01, conf = 0.01, maxlen = 6)
showall(rules)


text:
using MacroTools
using MacroTools: @q, combinedef

named(arg) =
    isexpr(arg, :(::)) && length(arg.args) == 1 ? :($(gensym())::$(arg.args[1])) : arg

typeless(x) = MacroTools.postwalk(x -> isexpr(x, :(::), :kw) ? x.args[1] : x, x)
isvararg(x) = isexpr(x, :(::)) && namify(x.args[2]) == :Vararg

for n = 0:3
    gradtuple = Symbol(:gradtuple, n)
    @eval begin
        $gradtuple(x::Tuple) = ($(ntuple(_ -> :nothing, n)...), x...)
        $gradtuple(x::Nothing) = nothing
        $gradtuple(x) = error("Gradient $x should be a tuple")
    end
end

abstract type AContext end
function adjoint end
function _forward end

function gradm(ex, mut = false)
    @capture(
        shortdef(ex),
        (name_(args__) = body_) | (name_(args__) where {Ts__} = body_),
    ) || error("Need a function definition")
    kw = length(args) > 1 && isexpr(args[1], :parameters) ? esc(popfirst!(args)) : nothing
    isclosure = isexpr(name, :(::)) && length(name.args) > 1
    f, T =
        isexpr(name, :(::)) ?
        (length(name.args) == 1 ? (esc(gensym()), esc(name.args[1])) : esc.(name.args)) :
        (esc(gensym()), :(Core.Typeof($(esc(name)))))
    kT = :(Core.kwftype($T))
    Ts == nothing && (Ts = [])
    args = named.(args)
    argnames = Any[typeless(arg) for arg in args]
    !isempty(args) && isvararg(args[end]) && (argnames[end] = :($(argnames[end])...,))
    args = esc.(args)
    argnames = esc.(argnames)
    Ts = esc.(Ts)
    cx = :($(esc(:__context__))::AContext)
    fargs = kw == nothing ? [cx, :($f::$T), args...] : [kw, cx, :($f::$T), args...]
    gradtuple = isclosure ? gradtuple0 : gradtuple1
    gradtuplekw = isclosure ? gradtuple2 : gradtuple3
    adj = @q @inline ZygoteRules.adjoint($(fargs...)) where {$(Ts...)} = $(esc(body))
    quote
        $adj
        @inline function ZygoteRules._forward($cx, $f::$T, $(args...)) where {$(Ts...)}
            y, _back = adjoint(__context__, $f, $(argnames...))
            $(mut ? nothing : :(back(::Nothing) = nothing))
            back(Δ) = $gradtuple(_back(Δ))
            return y, back
        end
        @inline function ZygoteRules._forward(
            $cx,
            ::$kT,
            kw,
            $f::$T,
            $(args...),
        ) where {$(Ts...)}
            y, _back = adjoint(__context__, $f, $(argnames...); kw...)
            $(mut ? nothing : :(back(::Nothing) = nothing))
            back(Δ) = $gradtuplekw(_back(Δ))
            return y, back
        end
        nothing
    end
end

macro adjoint(ex)
    gradm(ex)
end

macro adjoint!(ex)
    gradm(ex, true)
end


text:
"""
This module defines types for Pauli Matrix
"""
module Pauli

abstract type PauliMatrix{T <: Complex} <: AbstractMatrix{T} end

struct X{T} <: PauliMatrix{T} end
struct Y{T} <: PauliMatrix{T} end
struct Z{T} <: PauliMatrix{T} end

Base.Matrix(x::PauliMatrix{T}) where T = Matrix{T}(x)
Base.Matrix{T}(::X) where T = T[0 1;1 0]
Base.Matrix{T}(::Y) where T = T[0 -im;im 0]
Base.Matrix{T}(::Z) where T = T[1 0;0 -1]

using SparseArrays

SparseArrays.SparseMatrixCSC(x::PauliMatrix{T}) where T = SparseMatrixCSC{T}(x)
SparseArrays.SparseMatrixCSC{T}(x::X) where T = sparse(Matrix{T}(x))
SparseArrays.SparseMatrixCSC{T}(x::Y) where T = sparse(Matrix{T}(x))
SparseArrays.SparseMatrixCSC{T}(x::Z) where T = sparse(Matrix{T}(x))

SparseArrays.nonzeros(::X{T}) where T = T[1, 1]
SparseArrays.nonzeros(::Y{T}) where T = T[im, -im]
SparseArrays.nonzeros(::Z{T}) where T = T[1, -1]

function Base.checkbounds(::Type{Bool}, x::PauliMatrix, ind1, ind2)
    1 <= ind1 <= 2 || throw(BoundsError(x, (ind1, ind2)))
    1 <= ind2 <= 2 || throw(BoundsError(x, (ind1, ind2)))
    true
end

function Base.getindex(x::X{T}, inds...) where T
    @boundscheck checkbounds(Bool, x, inds...)

    inds == (1, 2) || inds == (2, 1) ? one(T) : zero(T)
end

function Base.getindex(x::Y{T}, inds...) where T
    @boundscheck checkbounds(Bool, x, inds...)

    inds == (1, 2) || inds == (2, 1) ? T(-im) : zero(T)
end

function Base.getindex(x::Z{T}, inds...) where T
    @boundscheck checkbounds(Bool, x, inds...)

    inds == (1, 1) ? -one(T) :
    inds == (2, 2) ? one(T)  :
    zero(T)
end

Base.size(::PauliMatrix) = (2, 2)

Base.summary(io::IO, ::X{T}) where T = print(io, "Pauli X{$T}")
Base.summary(io::IO, ::Y{T}) where T = print(io, "Pauli Y{$T}")
Base.summary(io::IO, ::Z{T}) where T = print(io, "Pauli Z{$T}")

end

max_stars_repo_name: Roger-luo/QuHamiltonian

max_issues_repo_name: Roger-luo/QuHamiltonian

max_forks_repo_name: Roger-luo/QuHamiltonian


text:
using PAINTER
using Base.Test

include("paintertest.jl")


text:
using SafeTestsets
using Test

@safetestset "matching" begin
    include("matching.jl")
end
@safetestset "persistencecurves" begin
    include("persistencecurves.jl")
end
@safetestset "persistenceimages" begin
    include("persistenceimages.jl")
end
@safetestset "scitypes" begin
    include("scitypes.jl")
end
@safetestset "mlj" begin
    include("mlj.jl")
end
@safetestset "plotsrecipes" begin
    include("plotsrecipes.jl")
end
@safetestset "aqua" begin
    include("aqua.jl")
end
@safetestset "doctests" begin
    include("doctests.jl")
end


text:
export PlayerA, PlayerB, Outcome, FactorA, FactorB, ScoreA, ScoreB, Margin, No # various useful symbol abbreviations
export check_ratings, read_ratings, write_ratings
export check_results, read_results, write_results, summarize, summarizeF
export read_ratingstable, write_ratingstable

# useful symbols for working with ratings data frames
PlayerA = Symbol("Player A")
PlayerB = Symbol("Player B")
Outcome = Symbol("Outcome")
FactorA = Symbol("Factor A")
FactorB = Symbol("Factor B")
ScoreA = Symbol("Score A")
ScoreB = Symbol("Score B")
Margin = Symbol("Margin")
Result = Symbol("Result")
No = Symbol("No. of matches")
# don't need to abbreviate
#  :Players, :Ratings

#####################################
#  ranked lists
#####################################

#####################################
#  I/O of ratings dataframes
#####################################

# check you have a valid data frame for a set of ratings
function check_ratings( df::DataFrame )
    # check it has the required columns with correct datatype
    if !haskey(df, :Players)
       error("Input DataFrame must have a 'Players' column'")
    end
    if !haskey(df, :Ratings)
        error("Input DataFrame must have a 'Ratings' column")
    end
    if !( eltype(df[:Players]) <: Union{Missing, AbstractString} )
       error("Input DataFrame 'Players' column must have type String")

    end
    if !( eltype(df[:Ratings]) <: Union{Missing, AbstractFloat} )
        error("Input DataFrame must have a 'Ratings' column must have type Float64")
    end
    # do I care if there is other stuff???
end

# convert an input DataFrame into a RatingsList object
function convert(::Type{RatingsList}, df::DataFrame )
    n = size(df,1)
    r = RatingsList(df[:Players],
                    Dict( [ df[r,:Players] => df[r,:Ratings] for r=findall( .!ismissing.(df[:,:Ratings])) ] )
                    )
    return r
end

# convert a RatingsList object into a DataFrame for output
function convert(::Type{DataFrame}, r::RatingsList )
    m = length(r.players)
    x = Array{Union{Missing, Float64}, 1}(undef, m)
    for i=1:m
        if haskey(r.ratings, r.players[i])
            x[i] = r.ratings[ r.players[i] ]
        end
    end
    df = DataFrame(Players=r.players, Ratings=x)
    return df
end

# read in a CSV file of ratings
function read_ratings( file::String )
    df = CSV.read(file; comment="#")
    check_ratings( df )
    r = convert( RatingsList, df )
    return df, r
end

# write a RatingsList to a file
function write_ratings( file::String, r::RatingsList )
    # write comments into file??? -- perhaps by appending to a file???
    df = convert(DataFrame, r)
    CSV.write(file, df)
end

#####################################
#  competition lists
#####################################
function check_results( df::DataFrame )
    # check it has the required columns with correct datatype
    if !haskey(df, PlayerA)
       error("Input DataFrame must have a 'Player A' column'")
    end
    if !( eltype(df[PlayerA]) <: Union{Missing, AbstractString} )
       error("Input DataFrame 'Player A' column must have type String")
    end
    if !haskey(df, PlayerB)
       error("Input DataFrame must have a 'Player B' column'")
    end
    if !( eltype(df[PlayerB]) <: Union{Missing, AbstractString} )
       error("Input DataFrame 'Player B' column must have type String")
    end

    if haskey(df, Outcome)
        if !( eltype(df[Outcome]) <: Union{Missing, Integer} )
            error("Input DataFrame 'Outcome' column must have type Union{Missing, Integer}")
        end
        if maximum(skipmissing(df[Outcome])) > 1 ||  minimum(skipmissing(df[Outcome])) < -1
            error("Input DataFrame 'Outcome' column should be -1,0,1 (or missing)")
        end
    end

    if haskey(df, Result)
        if !( eltype(df[Result]) <: Union{Missing, Real} )
            error("Input DataFrame 'Result' column must have type Union{Missing, Real}")
        end
        if minimum(skipmissing(df[Result])) < 0
            error("Input DataFrame 'Result' column should be >= 0 (or missing)")
            # should really check it is an integer or half integer
        end
        if !haskey(df, No)
            error("Input DataFrame 'Result' column also requires 'No. of matches' column")
        end
        if !( eltype(df[No]) <: Union{Missing, Integer} )
            error("Input DataFrame 'No. of matches' column must have type Union{Missing, Integer}")
        end
        if minimum(skipmissing(df[No])) < 1
            error("Input DataFrame 'No. of matches' column should be >= 1 (or missing)")
        end
    end

    if haskey(df, ScoreA) && haskey(df, ScoreB)
        if !( eltype(df[ScoreA]) <: Union{Missing, Integer} )
            error("Input DataFrame 'Score A' column must have type Union{Missing, Integer}")
        end
        if !( eltype(df[ScoreB]) <: Union{Missing, Integer} )
            error("Input DataFrame 'Score B' column must have type Union{Missing, Integer}")
        end
    elseif haskey(df, ScoreA)
            error("Input DataFrame only has 'Score A' column, not 'Score B'")
    elseif haskey(df, ScoreB)
            error("Input DataFrame only has 'Score B' column, not 'Score A'")
    end

    if haskey(df, Margin) && !( eltype(df[Margin]) <: Union{Missing, Integer} )
        error("Input DataFrame 'Margin' column must have type Union{Missing, Integer}")
    end

    if haskey(df, FactorA) && haskey(df, FactorB)
        if !( eltype(df[FactorA]) <: Union{Missing, Integer} )
            error("Input DataFrame 'Factor A' column must have type Union{Missing, Integer}")
        end
        if !( eltype(df[FactorB]) <: Union{Missing, Integer} )
            error("Input DataFrame 'Factor B' column must have type Union{Missing, Integer}")
        end
    elseif haskey(df, FactorA)
            error("Input DataFrame only has 'Factor A' column, not 'Factor B'")
    elseif haskey(df, FactorB)
            error("Input DataFrame only has 'Factor B' column, not 'Factor A'")
    end

    # add columns that we can add
    if haskey(df, ScoreA) && haskey(df, ScoreB) && !haskey(df, Margin)
        df[Margin] = df[ScoreA] - df[ScoreB]
    end

    if haskey(df, Margin) && !haskey(df, Outcome)
        df[Outcome] = sign.(df[Margin])
    end

    # check for consistency between different types of inputs
    if !all(skipmissing( df[Margin] .== df[ScoreA] - df[ScoreB] ))
        error("Input DataFrame: Margins are inconsistent with Scores")
    end
    if !all(skipmissing( df[Outcome] .== sign.(df[Margin]) ))
        error("Input DataFrame: Outcomes are inconsistent with Margins")
    end

    # results format isn't compatible with outcomes (or scores and margins)
    # so don't try to translate from one to the other (at the moment)
    # but we could at least check that both aren't there
    # should also check that
    #   - Result <= No

end

# read in a CSV file of competitions
function read_results( file::String )
    df = CSV.read(file; comment="#")
    check_results( df )
    df[PlayerA] = coalesce.( df[PlayerA], "" ) # replace missing player names with blanks (for the moment)
    df[PlayerB] = coalesce.( df[PlayerB], "" ) # replace missing player names with blanks (for the moment)
    player_list = merge(+, countmap( df[PlayerA] ), countmap( df[PlayerB] )  )
    return df, player_list
end

function write_results( file::String, df::DataFrame)
    CSV.write(file, df)
end

# summarize a list of competitions
function summarize( df::DataFrame,  player_list::Dict{String,Int64})
    wins = copy(player_list)
    reset!(wins)
    score_diff = Dict{String,Int64}()
    for i=1:size(df,1)
        # ignores ties
        if df[i,Outcome] == 1
            increment!(wins, df[i,PlayerA], 1)
        elseif df[i,Outcome] == -1
            increment!(wins, df[i,PlayerB], 1)
        end
        diff = df[i,ScoreA] - df[i,ScoreB]
        increment!(score_diff, df[i,PlayerA], diff)
        increment!(score_diff, df[i,PlayerB], -diff)
    end

    new_df = DataFrame( Player=Array{String,1}(), Record=Array{String,1}(), ScoreDiff = Array{Int64,1}() )
    for p in keys(player_list)
        push!( new_df, [p, "$(wins[p])/$(player_list[p])", score_diff[p] ])
    end
    return new_df
end

# summarize a list of competitions with home/away factors
#   presumes that exactly one of FactorA and FactorB is 0, and one is 1 indicating home-ground
function summarizeF( df::DataFrame,  player_list::Dict{String,Int64})
    home_wins = copy(player_list)
    away_wins = copy(player_list)
    home_games = copy(player_list)
    away_games = copy(player_list)
    reset!(home_wins)
    reset!(away_wins)
    reset!(home_games)
    reset!(away_games)
    score_diff = Dict{String,Int64}()
    for i=1:size(df,1)
        # ignores ties
        if df[i,Outcome] == 1 && df[i,FactorA] == 1
            increment!(home_wins, df[i,PlayerA], 1)
            increment!(home_games, df[i,PlayerA], 1)
            increment!(away_games, df[i,PlayerB], 1)
        elseif df[i,Outcome] == 1 && df[i,FactorA] == 0
            increment!(away_wins, df[i,PlayerA], 1)
            increment!(away_games, df[i,PlayerA], 1)
            increment!(home_games, df[i,PlayerB], 1)
        elseif df[i,Outcome] == -1 && df[i,FactorB] == 1
            increment!(home_wins, df[i,PlayerB], 1)
            increment!(home_games, df[i,PlayerB], 1)
            increment!(away_games, df[i,PlayerA], 1)
        elseif df[i,Outcome] == -1 && df[i,FactorB] == 0
            increment!(away_wins, df[i,PlayerB], 1)
            increment!(away_games, df[i,PlayerB], 1)
            increment!(home_games, df[i,PlayerA], 1)
        end
        diff = df[i,ScoreA] - df[i,ScoreB]
        increment!(score_diff, df[i,PlayerA], diff)
        increment!(score_diff, df[i,PlayerB], -diff)
    end

    new_df = DataFrame( Player=Array{String,1}(), HomeRecord=Array{String,1}(), AwayRecord=Array{String,1}() )
    for p in keys(player_list)
        push!( new_df, [p, "$(home_wins[p])/$(home_games[p])", "$(away_wins[p])/$(away_games[p])" ])
    end
    return new_df
end

# is there any need to write out a list of competition results????

#####################################
#  RatingsTable
#####################################

# read in a CSV file of ratings
function read_ratingstable( file::String )
    df = CSV.read(file; comment="#")
    # haven't written routines to check the input ratings table yet
    players = String.(names(df))
    return RatingsTable(players, df)
end

# write a RatingsList to a file
function write_ratingstable( file::String, R::RatingsTable )
    # write comments into file??? -- perhaps by appending to a file???
    CSV.write(file, R.ratings)
end


text:
using Pkg
Pkg.activate("./")
using RK

using Plots

function OutputSolution()
    # Output numerical solution in the end
    nz = length(history(tspan[1]))
    if nz == 1
        println("y(", tspan[2], ") = ", sol[2][end])
    else
        for i = 1 : nz
            println("y", i, "(", tspan[2], ") = ", (sol[2]')[end - nz + i])
        end
    end
end

function OutputConvOrder(; dde=false, overlapping=false)
    # Check convergence order of examples
    # "dde=true" - if equation is DDE (with dicrete delays)
    n_start = 2
    n_end = 8
    err = zeros(n_end - n_start + 1)
    nsteps = zeros(n_end - n_start + 1)

    i = 1
    for steppow = 2 : n_end
        stepsize = (2.0)^(-steppow)

        if dde
            if overlapping
                sol = ide_solve(idefun, K, delays_int, history, tspan, stepsize, delays, overlapping=true)
            else
                sol = ide_solve(idefun, K, delays_int, history, tspan, stepsize, delays)
            end
        else
            sol = ide_solve(idefun, K, delays_int, history, tspan, stepsize)
        end

        err[i] = abs(analytic_sol - sol[2][end])
        nsteps[i] = stepsize

        i += 1
    end

    println()
    println("Convergence order: ", (log10(err[end-1]) - log10(err[end])) / (log10((2.0)^(-n_end+1)) - log10((2.0)^(-n_end))))

    X_ticks = [(10.0)^(-x) for x in 1:15]
    Y_ticks = [(10.0)^(-x) for x in 0:0.5:3]
    display(plot(err,
                nsteps,
                linewidth=3,
                xaxis=:log10,
                yaxis=:log10,
                xticks=X_ticks,
                yticks=Y_ticks,
                xlabel="ERROR",
                ylabel="STEPSIZE",
                legend=false,
                title="Convergence order"))
end

# Example 5 (2 discrete delays + State-dependent)
idefun(t,y,z,i) = -z[1]^((t + 1) / 2) * z[2] * y.^2 * (1 + exp(t^2) * t * i[1]) / exp(0.5)
K(t,s,y)        = y * exp(s - s * t)
delays(t,y)     = [ (log(y[1]))^2 / (t + 1) - 0.5;
                    (t - 1) / 4 ]
delays_int(t)   = t / 2 - 1
history(t)      = exp(-t)

tspan = [0 5]
stepsize = 1e-2

sol = ide_solve(idefun, K, delays_int, history, tspan, stepsize, delays, overlapping=true)

OutputSolution()

display(plot(sol,
            linewidth=3,
            xlabel="TIME",
            ylabel="SOLUTION",
            xticks=trunc(Int,tspan[1]):(tspan[2]-trunc(Int,tspan[1]))/10:tspan[2],
            title="Example 5 (2 discrete delays + State-dependent)",
            legend=:outertopright))

# Check convergence order
fun(t) = exp(-t)
analytic_sol = fun(tspan[end]);

# dde=true - if equation is DDE (with dicrete delays 'z')
# overlapping=true - if equation has overlapping in discrete delays
OutputConvOrder(dde=true, overlapping=true)

max_stars_repo_name: vitesempl/IDE-Delay-Runge-Kutta

max_issues_repo_name: vitesempl/IDE-Delay-Runge-Kutta

max_forks_repo_name: vitesempl/IDE-Delay-Runge-Kutta


text:
#Pkg.checkout("EcoEvo3D","model1")
using EcoEvo3D

function main(ngenana_)
	seed = 177;#seed for random numbers (to control the outputs)
	nreal = 5;#Number of realizations
	Gmax = 10000;#Maximum number of Generations
	J = 10000;#Total Number of individuals in the system
        nGenAna = ngenana_
	anaG = J*nGenAna;#Threshold to consider speciation (anagenesis and cladogenesis)
	v =  (0.000001 + rand()* 0.00001);#Cladogenesis speciation
	mr = (0.00001 + rand()* 0.0001);#Regional migration -
	ml = (0.002 + rand()* 0.02);#Local Migration

	distmatfile = "upstream_cost0.001_size8_m1.txt";#Name of the file that contains the location of the points with clownfishes.
	verticesdata = "VerticesData.txt";#Name of the file that contains the size and the height of the points with clownfishes.
	model = 1;

	EcoEvo3D.dynamic(seed,nreal,Gmax,J,v,mr,ml,anaG,distmatfile,verticesdata,model);
end

if(length(ARGS)==0)
    println();
    println("Please give as an argment to the program the 'number of generations for anagenesis speciation':");
    println("To run: julia testm1.jl <anaG>");
    println();
else
    main(float(ARGS[1]))
end


text:
type MakeBuilder <: AbstractBuilder
  args
  MakeBuilder(args=[]) = new(args)
end

function rules(package::AbstractPackage,::AbstractDebianizer, build::MakeBuilder; kwargs...)
  debian_file("rules",
    "#! /usr/bin/make -f",
    "include /usr/share/cdbs/1/rules/debhelper.mk",
    "include /usr/share/cdbs/1/class/makefile.mk",
    build.args...
  )
end

type CMakeBuilder <: AbstractBuilder
  flags
  args
  CMakeBuilder(flags="", args=[]) = new(flags, args)
end

function rules(package::AbstractPackage, ::AbstractDebianizer, build::CMakeBuilder; kwargs...)
  specials = String["include /usr/share/cdbs/1/class/cmake.mk"]
  if (!isa(build.flags, Nothing)) && length(build.flags) > 0
    push!(specials, "DEB_CMAKE_EXTRA_FLAGS=$(build.flags)")
  end
  debian_file("rules",
    "#! /usr/bin/make -f",
    "include /usr/share/cdbs/1/rules/debhelper.mk",
    specials..., build.args...
  )
end

type PythonBuilder <: AbstractBuilder
  args
  PythonBuilder(args=[]) = new(args)
end

function rules(package::AbstractPackage, ::AbstractDebianizer, build::PythonBuilder; kwargs...)
  debian_file("rules",
    "#! /usr/bin/make -f",
    "DEB_PYTHON_SYSTEM=pysupport",
    "include /usr/share/cdbs/1/rules/debhelper.mk",
    "include /usr/share/cdbs/1/class/python-distutils.mk",
    build.args...
  )
end

max_stars_repo_name: UCL/DebbyPacker

max_issues_repo_name: UCL/DebbyPacker

max_forks_repo_name: UCL/DebbyPacker


text:
using Enzyme
using Test
using FiniteDifferences
using ForwardDiff
using Statistics

# Test against FiniteDifferences
function test_scalar(f, x; rtol=1e-9, atol=1e-9, fdm=central_fdm(5, 1), kwargs...)
    ∂x, = autodiff(f, Active, Active(x))
    @test isapprox(∂x, fdm(f, x); rtol=rtol, atol=atol, kwargs...)
end

include("abi.jl")

@testset "Internal tests" begin
    f(x) = 1.0 + x
    thunk_a = Enzyme.Compiler.thunk(f, Active, Tuple{Active{Float64}})
    thunk_b = Enzyme.Compiler.thunk(f, Const, Tuple{Const{Float64}})
    thunk_c = Enzyme.Compiler.thunk(f, Active{Float64}, Tuple{Active{Float64}})
    @test thunk_a.adjoint !== thunk_b.adjoint
    @test thunk_c.adjoint === thunk_a.adjoint

    @test thunk_a(Active(2.0), 1.0) == (1.0,)
    @test thunk_a(Active(2.0), 2.0) == (2.0,)
    @test thunk_b(Const(2.0)) === ()

    forward, pullback = Enzyme.Compiler.thunk(f, Active, Tuple{Active{Float64}}, Val(true))
    # @test thunk_split.primal !== C_NULL
    # @test thunk_split.primal !== thunk_split.adjoint
    # @test thunk_a.adjoint !== thunk_split.adjoint
end

# @testset "Split Tape" begin
#     f(x) = x[1] * x[1]

#     thunk_split = Enzyme.Compiler.thunk(f, Tuple{Duplicated{Array{Float64,1}}}, Val(true))
#     @test thunk_split.primal !== C_NULL
#     @test thunk_split.primal !== thunk_split.adjoint
# end

@testset "Simple tests" begin
    f1(x) = 1.0 + x
    f2(x) = x*x
    @test autodiff(f1, Active, Active(1.0))[1] ≈ 1.0
    @test autodiff(f2, Active, Active(1.0))[1] ≈ 2.0
    @test autodiff(tanh, Active, Active(1.0))[1] ≈ 0.41997434161402606939
    @test autodiff(tanh, Active, Active(1.0f0))[1] ≈ Float32(0.41997434161402606939)
    test_scalar(f1, 1.0)
    test_scalar(f2, 1.0)
end

@testset "Duplicated" begin
    x = Ref(1.0)
    y = Ref(2.0)

    ∇x = Ref(0.0)
    ∇y = Ref(0.0)

    autodiff((a,b)->a[]*b[], Active, Duplicated(x, ∇x), Duplicated(y, ∇y))

    @test ∇y[] == 1.0
    @test ∇x[] == 2.0
end

@testset "Simple tests" begin
    g(x) = real((x + im)*(1 - im*x))
    @test first(autodiff(g, Active, Active(2.0))) ≈ 2.0
    @test first(autodiff(g, Active, Active(3.0))) ≈ 2.0
    test_scalar(g, 2.0)
    test_scalar(g, 3.0)
end

@testset "Taylor series tests" begin

# Taylor series for `-log(1-x)`
# eval at -log(1-1/2) = -log(1/2)
function euroad(f::T) where T
    g = zero(T)
    for i in 1:10^7
        g += f^i / i
    end
    return g
end

euroad′(x) = first(autodiff(euroad, Active, Active(x)))

@test euroad(0.5) ≈ -log(0.5) # -log(1-x)
@test euroad′(0.5) ≈ 2.0 # d/dx -log(1-x) = 1/(1-x)
test_scalar(euroad, 0.5)
end

@testset "Array tests" begin

    function arsum(f::Array{T}) where T
        g = zero(T)
        for elem in f
            g += elem
        end
        return g
    end

    inp = Float64[1.0, 2.0]
    dinp = Float64[0.0, 0.0]
    autodiff(arsum, Active, Duplicated(inp, dinp))
    @test inp ≈ Float64[1.0, 2.0]
    @test dinp ≈ Float64[1.0, 1.0]
end

@testset "Advanced array tests" begin
    function arsum2(f::Array{T}) where T
        return sum(f)
    end
    inp = Float64[1.0, 2.0]
    dinp = Float64[0.0, 0.0]
    autodiff(arsum2, Active, Duplicated(inp, dinp))
    @test inp ≈ Float64[1.0, 2.0]
    @test dinp ≈ Float64[1.0, 1.0]
end

@testset "Bithacks" begin
    function fneg(x::Float64)
        xptr = reinterpret(Int64, x)
        y = Int64(-9223372036854775808)
        out = y ⊻ xptr;
        return reinterpret(Float64, out)
    end
    @test autodiff(fneg, Active, Active(2.0))[1] ≈ -1.0
    function expor(x::Float64)
        xptr = reinterpret(Int64, x)
        y = UInt64(4607182418800017408)
        out = y | xptr;
        return reinterpret(Float64, out)
    end
    @test autodiff(expor, Active, Active(0.42))[1] ≈ 4.0
end

@testset "GC" begin
    function gc_alloc(x)  # Basically g(x) = x^2
        a = Array{Float64, 1}(undef, 10)
        for n in 1:length(a)
            a[n] = x^2
        end
        return mean(a)
    end
    @test autodiff(gc_alloc, Active, Active(5.0))[1] ≈ 10

    # TODO (after BLAS)
    # A = Vector[2.0, 3.0]
    # B = Vector[4.0, 5.0]
    # dB = Vector[0.0, 0.0]
    # f = (X, Y) -> sum(X .* Y)
    # Enzyme.autodiff(f, Active, A, Duplicated(B, dB))

    function gc_copy(x)  # Basically g(x) = x^2
        a = x * ones(10)
        for n in 1:length(a)
            a[n] = x^2
        end
        return mean(a)
    end
    # TODO(wsmoses): Illegal update analysis
    # @test Enzyme.autodiff(gc_copy, Active, Active(5.0))[1] ≈ 10
end

@testset "Compare against" begin
    x = 3.0
    fd = central_fdm(5, 1)(sin, x)

    @test fd ≈ ForwardDiff.derivative(sin, x)
    @test fd ≈ first(autodiff(sin, Active, Active(x)))

    x = 0.2 + sin(3.0)
    fd = central_fdm(5, 1)(asin, x)

    @test fd ≈ ForwardDiff.derivative(asin, x)
    @test fd ≈ first(autodiff(asin, Active, Active(x)))
    test_scalar(asin, x)

    function foo(x)
        a = sin(x)
        b = 0.2 + a
        c = asin(b)
        return c
    end

    x = 3.0
    fd = central_fdm(5, 1)(foo, x)

    @test fd ≈ ForwardDiff.derivative(foo, x)
    @test fd ≈ first(autodiff(foo, Active, Active(x)))
    test_scalar(foo, x)

    # Input type shouldn't matter
    x = 3
    @test fd ≈ ForwardDiff.derivative(foo, x)
    @test fd ≈ first(autodiff(foo, Active, Active(x)))

    f74(a, c) = a * √c
    @test √3 ≈ first(autodiff(f74, Active, Active(2), 3))
end

"""
    J(ν, z) := ∑ (−1)^k / Γ(k+1) / Γ(k+ν+1) * (z/2)^(ν+2k)
"""
function mybesselj(ν, z, atol=1e-8)
    k = 0
    s = (z/2)^ν / factorial(ν)
    out = s
    while abs(s) > atol
        k += 1
        s *= (-1) / k / (k+ν) * (z/2)^2
        out += s
    end
    out
end
mybesselj0(z) = mybesselj(0, z)
mybesselj1(z) = mybesselj(1, z)

@testset "Bessel" begin
    autodiff(mybesselj, Active, Const(0), Active(1.0))
    autodiff(mybesselj, Active, 0, Active(1.0))
    @testset "besselj0/besselj1" for x in (1.0, -1.0, 0.0, 0.5, 10, -17.1,) # 1.5 + 0.7im)
        test_scalar(mybesselj0, x, rtol=1e-5, atol=1e-5)
        test_scalar(mybesselj1, x, rtol=1e-5, atol=1e-5)
    end

end

## https://github.com/JuliaDiff/ChainRules.jl/tree/master/test/rulesets
if !Sys.iswindows()
    include("packages/specialfunctions.jl")
end

@testset "DiffTest" begin
    include("DiffTests.jl")

    n = rand()
    x, y = rand(5, 5), rand(26)
    A, B = rand(5, 5), rand(5, 5)

    # f returns Number
    @testset "Number to Number" for f in DiffTests.NUMBER_TO_NUMBER_FUNCS
        test_scalar(f, n)
    end

    # TODO(vchuravy/wsmoses): Enable these tests
    # for f in DiffTests.VECTOR_TO_NUMBER_FUNCS
    #     @test isa(f(y), Number)
    # end

    # for f in DiffTests.MATRIX_TO_NUMBER_FUNCS
    #     @test isa(f(x), Number)
    # end

    # for f in DiffTests.TERNARY_MATRIX_TO_NUMBER_FUNCS
    #     @test isa(f(A, B, x), Number)
    # end

    # # f returns Array

    # for f in DiffTests.NUMBER_TO_ARRAY_FUNCS
    #     @test isa(f(n), Array)
    # end

    # for f in DiffTests.ARRAY_TO_ARRAY_FUNCS
    #     @test isa(f(A), Array)
    #     @test isa(f(y), Array)
    # end

    # for f in DiffTests.MATRIX_TO_MATRIX_FUNCS
    #     @test isa(f(A), Array)
    # end

    # for f in DiffTests.BINARY_MATRIX_TO_MATRIX_FUNCS
    #     @test isa(f(A, B), Array)
    # end

    # # f! returns Nothing

    # for f! in DiffTests.INPLACE_ARRAY_TO_ARRAY_FUNCS
    #     @test isa(f!(y, x), Nothing)
    # end

    # for f! in DiffTests.INPLACE_NUMBER_TO_ARRAY_FUNCS
    #     @test isa(f!(y, n), Nothing)
    # end

end

@testset "IO" begin

    function printsq(x)
        println(x)
        x*x
    end

    autodiff(printsq, Active, Active(2.3))
end

@testset "hmlstm" begin
    sigm(x)  = @fastmath 1 / (1 + exp(-x))
    @fastmath function hmlstm_update_c_scalar(z, zb, c, f, i, g)
        if z == 1.0f0 # FLUSH
            return sigm(i) * tanh(g)
        elseif zb == 0.0f0 # COPY
            return c
        else # UPDATE
            return sigm(f) * c + sigm(i) * tanh(g)
        end
    end

    N = 64
    Z = round.(rand(Float32, N))
    Zb = round.(rand(Float32, N))
    C = rand(Float32, N, N)
    F = rand(Float32, N, N)
    I = rand(Float32, N, N)
    G = rand(Float32, N, N)

    function broadcast_hmlstm(out, Z, Zb, C, F, I, G)
        out .= hmlstm_update_c_scalar.(Z, Zb, C, F, I, G)
        return nothing
    end

    ∇C = zeros(Float32, N, N)
    ∇F = zeros(Float32, N, N)
    ∇I = zeros(Float32, N, N)
    ∇G = zeros(Float32, N, N)

    # TODO(wsmoses): Check after updating Enzyme_jll
    # autodiff(broadcast_hmlstm, Const,
    #          Const(zeros(Float32, N, N)), Const(Z), Const(Zb),
    #          Duplicated(C, ∇C), Duplicated(F, ∇F), Duplicated(I, ∇I), Duplicated(G, ∇G))
end

@testset "generic" begin
    genlatestsin(x)::Float64 = Base.invokelatest(sin, x)
    @test -0.4161468365471424 ≈ Enzyme.autodiff(genlatestsin, Active, Active(2.0))[1]
end

@testset "broadcast" begin
    A = rand(10); B = rand(10); R = similar(A)
    dA = zero(A); dB = zero(B); dR = fill!(similar(R), 1)

    function foo_bc!(R, A, B)
        R .= A .+ B
        return nothing
    end

    autodiff(foo_bc!, Const, Duplicated(R, dR), Duplicated(A, dA), Duplicated(B, dB))

    # works since aliasing is "simple"
    autodiff(foo_bc!, Const, Duplicated(R, dR), Duplicated(R, dR), Duplicated(B, dB))

    A = rand(10,10); B = rand(10, 10)
    dA = zero(A); dB = zero(B); dR = fill!(similar(A), 1)

    @test_throws ErrorException autodiff(foo_bc!, Const, Duplicated(A, dR), Duplicated(transpose(A), transpose(dA)), Duplicated(B, dB))
end

@testset "DuplicatedReturn" begin
    moo(x) = fill(x, 10, 10)

    @test_throws ErrorException autodiff(moo, Active(2.1))
end

@testset "GCPreserve" begin
    function f(x, y)
        GC.@preserve x y begin
            ccall(:memcpy, Cvoid,
                (Ptr{Float64},Ptr{Float64},Csize_t), x, y, 8)
        end
        nothing
    end
    autodiff(f, Duplicated([1.0], [0.0]), Duplicated([1.0], [0.0]))
end

@testset "GCPreserve2" begin
    function f!(a_out, a_in)
           a_out[1:end-1] .= a_in[2:end]
           return nothing
    end
    a_in = rand(4)
    a_out = a_in

    shadow_a_out = ones(4)
    shadow_a_in = shadow_a_out

    autodiff(f!, Const, Duplicated(a_out, shadow_a_out), Duplicated(a_in, shadow_a_in))

    @test shadow_a_in ≈ Float64[0.0, 1.0, 1.0, 2.0]
    @test shadow_a_out ≈ Float64[0.0, 1.0, 1.0, 2.0]
end

@testset "UndefVar" begin
    function f(x, y)
        if x
            undefinedfnthowmagic()
        end
        y
    end
    @test 1.0 ≈ autodiff(f, false, Active(2.14))[1]
    @test_throws Base.UndefVarError autodiff(f, true, Active(2.14))

    function foo(x, y)
        if x
            Threads.@threads for N in 1:5:20
                println("The number of this iteration is $N")
            end
        end
        y
    end
    @test 1.0 ≈ autodiff(foo, false, Active(2.14))[1]
end

@testset "Split GC" begin
    @noinline function bmat(x)
        data = [x]
        return data
    end

    function f(x::Float64)
        @inbounds return bmat(x)[1]
    end
    @test 1.0 ≈ autodiff(f, Active(0.1))[1]
end


text:
"""
    VanGenuchten(; n, Ks=1, l=0.5, α=1, θr=0, θs=1) <: UnsaturatedFlowModel
    VanGenuchten(; m, Ks=1, l=0.5, α=1, θr=0, θs=1) <: UnsaturatedFlowModel
    VanGenuchten(; n, k, l=0.5, α=1, θr=0, θs=1, ρ=1e3, μ=1e-3, g=9.81) <: UnsaturatedFlowModel
    VanGenuchten(; m, k, l=0.5, α=1, θr=0, θs=1, ρ=1e3, μ=1e-3, g=9.81) <: UnsaturatedFlowModel

Create a Van Genuchten porous model.

# Keyword arguments
- `n`, `m`: _n_ or _m_ parameter (the parameters are related by _m_ = 1-1/_n_).
- `Ks=1`: saturated hydraulic conductivity.
- `k`: intrinsic permeability.
- `l=0.5`: _l_ parameter.
- `α=1`  (`\\alpha<tab>`): _α_ parameter.
- `θr=0` (`\\theta<tab>r`): residual moisture content.
- `θs=1` (`\\theta<tab>s`): moisture content when saturated.
- `ρ=1e3` (`\\rho<tab>`): density of the fluid.
- `μ=1e-3` (`\\mu<tab>`): viscosity of the fluid.
- `g=9.81`: magnitude of the gravitational acceleration.

# References
VAN GENUCHTEN, M. Th. A closed-form equation for predicting the hydraulic conductivity of unsaturated soils.
Soil Science Society of America Journal, 1980, vol. 44, no 5, p. 892-898.
"""
struct VanGenuchten{_Tml,_Tα,_TKs,_Tθ} <: UnsaturatedFlowModel
    m::_Tml
    l::_Tml
    α::_Tα
    Ks::_TKs
    θr::_Tθ
    θs::_Tθ

    function VanGenuchten(; n=nothing, m=nothing, l=0.5, α=1, Ks=nothing, k=nothing,
                            θr=0, θs=1, ρ=1e3, μ=1e-3, g=9.81)
        if !isnothing(n)
            @argcheck isnothing(m) "must assign only one of Ks and k, got both"
            @argcheck n>1
            m = 1-1/n

        else
            @argcheck !isnothing(m) "either n or m must be assigned"
        end

        @argcheck 0<m<1
        @argcheck α>zero(α)
        @argcheck θr<θs

        Ks = _asKs(Ks=Ks, k=k, ρ=ρ, μ=μ, g=g)

        new{promote_type(typeof(m),typeof(l)),typeof(α),typeof(Ks),promote_type(typeof(θr),typeof(θs))}(m,l,α,Ks,θr,θs)
    end
end

function θh(pm::VanGenuchten, h)
    if h ≥ zero(h)
        return pm.θs
    end

    n = 1/(1 - pm.m)

    Se = 1 / (1 + (pm.α*(-h))^n)^pm.m

    return pm.θr + Se*(pm.θs - pm.θr)
end

function hθ(pm::VanGenuchten, θ)
    n = 1/(1 - pm.m)

    Se = (θ - pm.θr)/(pm.θs - pm.θr)

    return -(1/(Se^(1/pm.m)) - 1)^(1/n)/pm.α
end

function Cθ(pm::VanGenuchten, θ)
    Se = (θ - pm.θr)/(pm.θs - pm.θr)

    return pm.α*pm.m/(1 - pm.m)*(pm.θs - pm.θr)*Se^(1/pm.m)*(1 - Se^(1/pm.m))^pm.m
end

function Ch(pm::VanGenuchten, h)
    if h ≥ zero(h)
        return zero(pm.α)
    end

    return Cθ(pm, θh(pm, h))
end

function Kθ(pm::VanGenuchten, θ)
    Se = (θ - pm.θr)/(pm.θs - pm.θr)

    return pm.Ks*Se^pm.l*(1 - (1 - Se^(1/pm.m))^pm.m)^2
end

function Kh(pm::VanGenuchten, h)
    if h ≥ zero(h)
        return pm.Ks
    end

    return Kθ(pm, θh(pm, h))
end

function Dθ(pm::VanGenuchten, θ)
    Se = (θ - pm.θr)/(pm.θs - pm.θr)
    return (1-pm.m)*pm.Ks/(pm.α*pm.m*(pm.θs - pm.θr)) * Se^pm.l*Se^(-1/pm.m) * ((1-Se^(1/pm.m))^(-pm.m) + (1-Se^(1/pm.m))^pm.m - 2)
end


text:
using HMCExamples

HMCExamples.main_FVGQ_2_joint()


text:
module Dynamics

using TrajectoryOptimization: Model, UnconstrainedObjective, ConstrainedObjective, LQRObjective, Trajectory
using RigidBodyDynamics
using LinearAlgebra

export
    pendulum,
    doublependulum,
    cartpole,
    ballonbeam,
    acrobot

include("../dynamics/pendulum.jl")
include("../dynamics/doublependulum.jl")
include("../dynamics/acrobot.jl")
include("../dynamics/ballonbeam.jl")
include("../dynamics/cartpole.jl")
include("../dynamics/dubinscar.jl")
include("../dynamics/quadrotor.jl")
include("../dynamics/kuka.jl")
include("../dynamics/double_integrator.jl")

end


text:
export Index,
       IndexVal,
       adjoint,
       dag,
       dim,
       prime,
       noprime,
       addtags,
       settags,
       replacetags,
       removetags,
       hastags,
       id,
       isdefault,
       dir,
       plev,
       tags,
       ind,
       setprime,
       sim,
       val

const IDType = UInt64

"""
   Arrow
`enum` type that can take three values: `In`, `Out`, or `Neither`, representing a directionality
associated with an index, i.e. the index leg is directed into or out of a given tensor
"""
@enum Arrow In=-1 Out=1 Neither=0

"""
    -(dir::Arrow)
Reverse direction of a directed `Arrow`. Will throw an error for an input of `Neither`.
"""
function -(dir::Arrow)
  dir==Neither && throw(ArgumentError("Cannot reverse direction of Arrow direction 'Neither'"))
  return dir==In ? Out : In
end

"""
An `Index` represents a single tensor index with fixed dimension `dim`. Copies of an Index compare equal unless their
`tags` are different.

An Index carries a `TagSet`, a set of tags which are small strings that specify properties of the `Index` to help
distinguish it from other Indices. There is a special tag which is referred to as the integer tag or prime
level which can be incremented or decremented with special priming functions.

Internally, an `Index` has a fixed `id` number, which is how the ITensor library knows two indices are copies of a
single original `Index`. `Index` objects must have the same `id`, as well as the `tags` to compare equal.
"""
struct Index
  id::IDType
  dim::Int
  dir::Arrow
  tags::TagSet
end

Index() = Index(IDType(0),1,Neither,TagSet("0"))

"""
    Index(dim::Integer, tags="0")
Create an `Index` with a unique `id` and a tagset given by `tags`.

Example: create a two dimensional index with tag `l`:
    Index(2, "l")
"""
function Index(dim::Integer,tags="0")
  ts = TagSet(tags)
  # By default, an Index has a prime level of 0
  # A prime level less than 0 is interpreted as the
  # prime level not being set
  plev(ts) < 0 && (ts = setprime(ts,0))
  Index(rand(IDType),dim,Out,ts)
end

id(i::Index) = i.id
dim(i::Index) = i.dim
dir(i::Index) = i.dir
tags(i::Index) = i.tags
plev(i::Index) = plev(tags(i))

"""
==(i1::Index, i1::Index)

Compare indices for equality. First the id's are compared,
then the prime levels are compared, and finally that
tags are compared.
"""
function ==(i1::Index,i2::Index)
  return id(i1) == id(i2) && tags(i1) == tags(i2)
end

"""
    copy(i::Index)
Create a copy of index `i` with identical `id`, `dim`, `dir` and `tags`.
"""
copy(i::Index) = Index(id(i),dim(i),dir(i),copy(tags(i)))

"""
    sim(i::Index)
Similar to `copy(i::Index)` except `sim` will produce an `Index` with a new, unique `id` instead of the same `id`.
"""
sim(i::Index) = Index(rand(IDType),dim(i),dir(i),copy(tags(i)))

"""
    dag(i::Index)
Copy an index `i` and reverse it's direction
"""
dag(i::Index) = Index(id(i),dim(i),-dir(i),tags(i))

"""
    isdefault(i::Index)
Check if an `Index` `i` was created with the default options.
"""
isdefault(i::Index) = (i==Index())

hastags(i::Index, ts) = hastags(tags(i),ts)

function settags(i::Index, strts)
  ts = TagSet(strts)
  # By default, an Index has a prime level of 0
  plev(ts) < 0 && (ts = setprime(ts,0))
  Index(id(i),dim(i),dir(i),ts)
end

addtags(i::Index, ts) = settags(i, addtags(tags(i), ts))
removetags(i::Index, ts) = settags(i, removetags(tags(i), ts))
replacetags(i::Index, tsold, tsnew) = settags(i, replacetags(tags(i), tsold, tsnew))

prime(i::Index, plinc = 1) = settags(i, prime(tags(i), plinc))
setprime(i::Index, plev) = settags(i, setprime(tags(i), plev))
noprime(i::Index) = setprime(i, 0)

Base.adjoint(i::Index) = prime(i)

# To use the notation i^5 == prime(i,5)
^(i::Index,pl::Integer) = prime(i,pl)

# Iterating over Index I gives
# integers from 1...dim(I)
start(i::Index) = 1
next(i::Index,n::Int) = (n,n+1)
done(i::Index,n::Int) = (n > dim(i))
colon(n::Int,i::Index) = range(n,dim(i))

function show(io::IO,
              i::Index)
  idstr = "$(id(i) % 1000)"
  if length(tags(i)) > 0
    print(io,"($(dim(i))|id=$(idstr)|$(tagstring(tags(i))))$(primestring(tags(i)))")
  else
    print(io,"($(dim(i))|id=$(idstr))$(primestring(tags(i)))")
  end
end

struct IndexVal
  ind::Index
  val::Int
  function IndexVal(i::Index,n::Int)
    n>dim(i) && throw(ErrorException("Value $n greater than size of Index $i"))
    n<1 && throw(ErrorException("Index value must be >= 1 (was $n)"))
    return new(i,n)
  end
end
getindex(i::Index, j::Int) = IndexVal(i, j)
getindex(i::Index, c::Colon) = [IndexVal(i, j) for j in 1:dim(i)]

(i::Index)(n::Int) = IndexVal(i,n)

val(iv::IndexVal) = iv.val
ind(iv::IndexVal) = iv.ind

==(i::Index,iv::IndexVal) = (i==ind(iv))
==(iv::IndexVal,i::Index) = (i==iv)

plev(iv::IndexVal) = plev(ind(iv))
prime(iv::IndexVal,inc::Integer=1) = IndexVal(prime(ind(iv),inc),val(iv))
adjoint(iv::IndexVal) = IndexVal(adjoint(ind(iv)),val(iv))

show(io::IO,iv::IndexVal) = print(io,ind(iv),"=$(val(iv))")


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    MeshData(domain, values)

A `domain` together with a dictionary of data `values`. For each rank `r`
(or parametric dimension) there can exist a corresponding Tables.jl table
`values[r]`. The helper function [`meshdata`](@ref) is recommended instead
of the raw constructor of the type.
"""
struct MeshData{D<:Domain,V<:Dict} <: Data
  domain::D
  values::V
end

domain(data::MeshData) = getfield(data, :domain)

function values(data::MeshData, rank=nothing)
  domain = getfield(data, :domain)
  values = getfield(data, :values)
  r = isnothing(rank) ? paramdim(domain) : rank
  haskey(values, r) ? values[r] : nothing
end

constructor(::Type{D}) where {D<:MeshData} = MeshData

# ----------------
# HELPER FUNCTION
# ----------------

"""
    meshdata(domain, values)

Create spatial data from a `domain` implementing the
the [`Domain`](@ref) trait and a dictionary of data
`values` where `values[r]` holds a Tables.jl table
for the rank `r`.

## Example

```julia
# attach temperature and pressure to grid elements
meshdata(CartesianGrid(10,10),
  Dict(2 => (temperature=rand(100), pressure=rand(100)))
)
```
"""
meshdata(domain::Domain, values::Dict) = MeshData(domain, values)

"""
    meshdata(vertices, elements, values)

Create spatial data from a [`SimpleMesh`](@ref) with
`vertices` and `elements`, and a dictionary of data
`values`.

## Example

```julia
# vertices and elements
vertices = Point2[(0,0),(1,0),(1,1),(0,1)]
elements = connect.([(1,2,3),(3,4,1)])

# attach data to ranks 0 and 2
meshdata(vertices, elements,
  Dict(
    0 => (temperature=[1.0,2.0,3.0,4.0], pressure=[4.0,3.0,2.0,1.0]),
    2 => (quality=["A","B"], state=[true,false])
  )
)
```
"""
meshdata(vertices, elements, values) =
  meshdata(SimpleMesh(vertices, elements), values)

"""
    meshdata(domain; vtable, etable)

Create spatial data from a `domain`, a table `vtable`
with data for the vertices, and a table `etable` with
data for the elements.

## Example

```julia
meshdata(CartesianGrid(10,10),
  etable = (temperature=rand(100), pressure=rand(100))
)
```
"""
function meshdata(domain; vtable=nothing, etable=nothing)
  d = paramdim(domain)
  values = if !isnothing(vtable) && !isnothing(etable)
    Dict(0 => vtable, d => etable)
  elseif isnothing(vtable)
    Dict(d => etable)
  elseif isnothing(etable)
    Dict(0 => vtable)
  else
    throw(ArgumentError("missing data tables"))
  end
  meshdata(domain, values)
end

"""
    meshdata(vertices, elements; vtable, etable)

Create spatial data from a [`SimpleMesh`](@ref) with
`vertices` and `elements`, a table `vtable` with data
for the vertices, and a table `etable` with data for
the elements.

## Example

```julia
# vertices and elements
vertices = Point2[(0,0),(1,0),(1,1),(0,1)]
elements = connect.([(1,2,3),(3,4,1)])

# attach data to mesh
meshdata(vertices, elements,
  vtable = (temperature=[1.0,2.0,3.0,4.0], pressure=[4.0,3.0,2.0,1.0]),
  etable = (quality=["A","B"], state=[true,false])
)
```
"""
meshdata(vertices, elements; vtable=nothing, etable=nothing) =
  meshdata(SimpleMesh(vertices, elements); vtable=vtable, etable=etable)


text:
@testset "data" begin
    @test normalize([1,2],0.5,0.5) == [1.0,3.0]
end