using Stats, Distributions, Turing
using Gadfly

# include("ASCIIPlot.jl")
import Gadfly.ElementOrFunction

# First add a method to the basic Gadfly.plot function for QQPair types (generated by Distributions.qqbuild())
Gadfly.plot(qq::QQPair, elements::ElementOrFunction...) = Gadfly.plot(x=qq.qx, y=qq.qy, Geom.point, Theme(highlight_width=0px), elements...)

# Now some shorthand functions
qqplot(x, y, elements::ElementOrFunction...) = Gadfly.plot(qqbuild(x, y), elements...)
qqnorm(x, elements::ElementOrFunction...) = qqplot(Normal(), x, Guide.xlabel("Theoretical Normal quantiles"), Guide.ylabel("Observed quantiles"), elements...)

NSamples = 5000

@model gdemo_fw() = begin
  s ~ InverseGamma(2,3)
  # s = 1
  m ~ Normal(0,sqrt.(s))
  y ~ MvNormal([m; m; m], [sqrt.(s) 0 0; 0 sqrt.(s) 0; 0 0 sqrt.(s)])
end

@model gdemo_bk(x) = begin
  # Backward Step 1: theta ~ theta | x
  s ~ InverseGamma(2,3)
  # s = 1
  m ~ Normal(0,sqrt.(s))
  x ~ MvNormal([m; m; m], [sqrt.(s) 0 0; 0 sqrt.(s) 0; 0 0 sqrt.(s)])
  # Backward Step 2: x ~ x | theta
  y ~ MvNormal([m; m; m], [sqrt.(s) 0 0; 0 sqrt.(s) 0; 0 0 sqrt.(s)])
end

fw = IS(NSamples)
# bk = Gibbs(10, PG(10,10, :s, :y), HMC(1, 0.25, 5, :m));
bk_sample_n = 20
bk = NUTS(bk_sample_n, 100, 0.65);

s = sample(gdemo_fw(), fw);
# describe(s)

N = div(NSamples, bk_sample_n)

x = [s[:y][1]...]
s_bk = Array{Turing.Chain}(undef, N)

simple_logger = Base.CoreLogging.SimpleLogger(stderr, Base.CoreLogging.Debug)
with_logger(simple_logger) do
  i = 1
  while i <= N
    s_bk[i] = sample(gdemo_bk(x), bk);
    x = [s_bk[i][:y][end]...];
    i += 1
  end
end

s2 = vcat(s_bk...);
# describe(s2)

# qqplot(s[:m], s2[:m])
# qqplot(s[:s], s2[:s])

qqm = qqbuild(s[:m], s2[:m])

using UnicodePlots

qqm = qqbuild(s[:m], s2[:m])
show(scatterplot(qqm.qx, qqm.qy, title = "QQ plot for m", canvas = DotCanvas))
show(scatterplot(qqm.qx[51:end-50], qqm.qy[51:end-50], title = "QQ plot for m (removing first and last 50 quantiles):", canvas = DotCanvas))
show(scatterplot(qqm.qx, qqm.qy, title = "QQ plot for m"))
show(scatterplot(qqm.qx[51:end-50], qqm.qy[51:end-50], title = "QQ plot for m (removing first and last 50 quantiles):"))

qqs = qqbuild(s[:s], s2[:s])
show(scatterplot(qqs.qx, qqs.qy, title = "QQ plot for s"))
show(scatterplot(qqs.qx[51:end-50], qqs.qy[51:end-50], title = "QQ plot for s (removing first and last 50 quantiles):"))

X = qqm.qx
y = qqm.qy
slope = (1 / (transpose(X) * X)[1] * transpose(X) * y)[1]

print("  slopeₛ = $slope ≈ 1 (ϵ = 0.1)")
ans1 = abs(slope - 1.0) <= 0.1
if ans1
  printstyled(" ✓\n", color=:green)
else
  printstyled(" X\n", color=:red)
  printstyled("    slope = $slope, diff = $(slope - 1.0)\n", color=:red)
end

# qqs = qqbuild(s[:s], s2[:s])


println("Using the Task's Test Data")
test = """12 127 28 42 39 113 42 18 44 118 44 37 113 124 37 48 127 36 29
   31 125 139 131 115 105 132 104 123 35 113 122 42 117 119 58 109 23 105
   63 27 44 105 99 41 128 121 116 125 32 61 37 127 29 113 121 58 114 126
   53 114 96 25 109 7 31 141 46 13 27 43 117 116 27 7 68 40 31 115 124 42
   128 52 71 118 117 38 27 106 33 117 116 111 40 119 47 105 57 122 109
   124 115 43 120 43 27 27 18 28 48 125 107 114 34 133 45 120 30 127 31
   116 146"""
test = map(parseint, split(test, r"\s"))
println(stemleaf(test))

println("Test with Reals and Negative Zero Stem")
test = [-23.678758, -12.45, -3.4, 4.43, 5.5, 5.678, 16.87, 24.7, 56.8]
println(stemleaf(test))

println("Test with Leaf Size Scaling")
test = int(500*randn(20))
println("Using:  ", test)
println(stemleaf(test, 10))


using Test, LinearAlgebra, DomainSets, StaticArrays
import DomainSets: elements, TypeFactory

include("test_spaces.jl")
include("test_maps.jl")
include("test_generic_domain.jl")
include("test_specific_domains.jl")


using Printf
using Random
using Statistics
using LinearAlgebra

using Flux
using DifferentialEquations
using DiffEqFlux

using JLD2
using Plots

using Flux: @epochs

#####
##### Load data from JLD2 file
#####

file = jldopen("../DataGeneration/rayleigh_taylor_instability_3d_horizontal_averages.jld2")

Is = keys(file["timeseries/t"])

Nz = file["grid/Nz"]
Lz = file["grid/Lz"]
Nt = length(Is)

t = zeros(Nt)
b = b_data = zeros(Nt, Nz)

for (i, I) in enumerate(Is)
    t[i] = file["timeseries/t/$I"]
    b[i, :] .= file["timeseries/b/$I"]
end

#####
##### Plot animation of b̅(z,t) from data
#####

z = file["grid/zC"]

@info "Animating buoyancy profile..."
anim = @animate for n=1:Nt
    t_str = @sprintf("%2.2f", t[n])
    plot(b[n, :], z, linewidth=2,
         xlim=(-1, 1), ylim=(-0.5, 0.5), label="",
         xlabel="buoyancy", ylabel="Depth (z)",
         title="Rayleigh-Taylor instability: t=$t_str", show=false)
end
mp4(anim, "rayleigh_taylor_instability_buoyancy.mp4", fps=15)

#####
##### Coarse grain data to 16 vertical levels (plus halo regions)
#####

function coarse_grain(data, resolution)
    @assert length(data) % resolution == 0
    s = length(data) / resolution

    data_cs = zeros(resolution)
    for i in 1:resolution
        t = data[Int((i-1)*s+1):Int(i*s)]
        data_cs[i] = mean(t)
    end

    return data_cs
end

coarse_resolution = cr = 16

b_cs = zeros(cr, Nt)
z_cs = coarse_grain(collect(z), cr)

for n=1:Nt
    b_cs[:, n] .= coarse_grain(b[n, :], cr)
end

#####
##### Plot coarse buoyancy profile
#####

@info "Animating coarse buoyancy profile..."
anim = @animate for n=1:Nt
    t_str = @sprintf("%2.2f", t[n])
    plot(b_cs[:, n], z_cs, linewidth=2,
         xlim=(-1, 1), ylim=(-1/2, 1/2), label="",
         xlabel="buoyancy", ylabel="Depth (z)",
         title="Rayleigh-Taylor instability: t=$t_str", show=false)
end
mp4(anim, "rayleigh_taylor_instability_buoyancy_coarse.mp4", fps=15)

#####
##### Create training data
#####

bₙ    = zeros(cr, Nt-1)
bₙ₊₁  = zeros(cr, Nt-1)

for i in 1:Nt-1
       bₙ[:, i] .=  b_cs[:,   i]
     bₙ₊₁[:, i] .=  b_cs[:, i+1]
end

N = 100  # Number of training data pairs.

training_data = [(bₙ[:, i], bₙ₊₁[:, i]) for i in 1:N]

#####
##### Create neural network
#####

# Complete black box right-hand-side.
dbdt_NN = Chain(Dense( cr, 2cr, tanh),
                Dense(2cr, 4cr, tanh),
                Dense(4cr, 4cr, tanh),
                Dense(4cr, 2cr, tanh),
                Dense(2cr,  cr))

NN_params = Flux.params(dbdt_NN)

#####
##### Define loss function
#####

tspan = (0.0, 0.1)  # 10 minutes
neural_pde_prediction(b₀) = neural_ode(dbdt_NN, b₀, tspan, Tsit5(), reltol=1e-4, save_start=false, saveat=tspan[2])

loss_function(bₙ, bₙ₊₁) = sum(abs2, bₙ₊₁ .- neural_pde_prediction(bₙ))

#####
##### Choose optimization algorithm
#####

opt = ADAM(1e-3)

#####
##### Callback function to observe training.
#####

function cb()
    train_loss = sum([loss_function(bₙ[:, i], bₙ₊₁[:, i]) for i in 1:N])

    # nn_pred = neural_ode(dTdt_NN, bₙ[:, 1], (t[1], t[N]), Tsit5(), saveat=t[1:N], reltol=1e-4) |> Flux.data
    # test_loss = sum(abs2, b_cs[:, 1:N] .- nn_pred)

    println("train_loss = $train_loss")
    return train_loss
end

cb()

#####
##### Train!
#####

epochs = 50
best_loss = Inf
last_improvement = 0

for epoch_idx in 1:epochs
    global best_loss, last_improvement

    @info "Epoch $epoch_idx"
    Flux.train!(loss_function, NN_params, training_data, opt, cb=cb) # cb=Flux.throttle(cb, 10))

    loss = cb()

    if loss <= best_loss
        @info("Record low loss! Saving neural network out to dbdt_NN.bson")
        BSON.@save "dbdt_NN.bson" dbdt_NN
        best_loss = loss
        last_improvement = epoch_idx
    end

    # If we haven't seen improvement in 2 epochs, drop our learning rate:
    if epoch_idx - last_improvement >= 2 && opt.eta > 1e-8
        opt.eta /= 10.0
        @warn("Haven't improved in a while, dropping learning rate to $(opt.eta)")

        # After dropping learning rate, give it a few epochs to improve
        last_improvement = epoch_idx
    end
end

#####
##### Run the neural PDE forward to see how well it performs just by itself.
#####

nn_pred = neural_ode(dbdt_NN, bₙ[:, 1], (t[1], t[end]), Tsit5(), saveat=t, reltol=1e-4) |> Flux.data

z_cs = coarse_grain(z, cr)

anim = @animate for n=1:Nt
    t_str = @sprintf("%2.2f", t[n])
    plot(b_cs[:, n], z_cs, linewidth=2,
         xlim=(-1, 1), ylim=(-0.5, 0.5), label="Data",
         xlabel="buoyancy", ylabel="Depth (z)",
         title="Rayleigh-Taylor instability: t=$t_str",
         legend=:bottomright, show=false)
    plot!(nn_pred[:, n], z_cs, linewidth=2, label="Neural DE", show=false)
end

mp4(anim, "rayleigh_taylor_instability_neural_PDE.mp4", fps=15)


using PtFEM

data = Dict(
  # GenericSolid(ndim, nst, nels, nn, nip, finite_element(nod, nodof), axisymmentrix)
  :struc_el => GenericSolid(2, 3, 6, 35, 9, Quadrilateral(9, 2), false),
  :properties => [1.0e6 0.3 0.0;],
  :g_coord => [
    0.0  0.0; 1.5  0.0; 3.0  0.0; 4.5  0.0; 6.0  0.0; 0.0 -1.5;
    1.5 -1.5; 3.0 -1.5; 4.5 -1.5; 6.0 -1.5; 0.0 -3.0; 1.5 -3.0;
    3.0 -3.0; 4.5 -3.0; 6.0 -3.0; 0.0 -4.5; 1.5 -4.5; 3.0 -4.5;
    4.5 -4.5; 6.0 -4.5; 0.0 -6.0; 1.5 -6.0; 3.0 -6.0; 4.5 -6.0;
    6.0 -6.0; 0.0 -7.5; 1.5 -7.5; 3.0 -7.5; 4.5 -7.5; 6.0 -7.5;
    0.0 -9.0; 1.5 -9.0; 3.0 -9.0; 4.5 -9.0; 6.0 -9.0],
  :g_num => [
    11,  6,  1,  2,  3,  8, 13, 12,  7,
    21, 16, 11, 12, 13, 18, 23, 22, 17,
    31, 26, 21, 22, 23, 28, 33, 32, 27,
    13,  8,  3,  4,  5, 10, 15, 14,  9,
    23, 18, 13, 14, 15, 20, 25, 24, 19,
    33, 28, 23, 24, 25, 30, 35, 34, 29
  ],
  :support => [
    ( 1, [0 1]), ( 5, [0 1]), ( 6, [0 1]), (10, [0 1]), (11, [0 1]), (15, [0 1]),
    (16, [0 1]), (20, [0 1]), (21, [0 1]), (25, [0 1]), (26, [0 1]), (30, [0 1]),
    (31, [0 0]), (32, [0 0]), (33, [0 0]), (34, [0 0]), (35, [0 0])
  ],
  :loaded_nodes => [
    ( 1, [0.0 -0.5]),
    ( 2, [0.0 -2.0]),
    ( 3, [0.0 -0.5])
  ]
)

data |> display
println()

@time m, dis_df, fm_df = p54(data)
println()

dis_df |> display
println()

println("\nThe integration point (nip = $(data[:struc_el].nip)) stresses are:\n")

fm_df |> display
println()


function packtablegen(lengths, rarity_pct)
    weights = rarity_pct ./ lengths
    ans = Vector{ Tuple{Int,Int} }(undef, sum(lengths) )
    wvec = zeros(Float64, sum(lengths) )
    i = 0
    for (rarity,len) in enumerate( lengths)
        for j in 1:len
            i += 1
            ans[i] = (rarity, j )
            wvec[i] = weights[rarity]
        end
    end
    ans, StatsBase.weights(wvec)
end


using Plots, LaTeXStrings
pyplot()

using Random
nsteps = 20
Y = zeros(Int, nsteps + 1)
T = zeros(Float64, nsteps + 1)
for i = 1:nsteps
    T[i+1] = T[i] + randexp()
    Y[i+1] = Y[i] + 1
end

# plot(T,Y,linetype=:steppost, xlabel=L"t", ylabel=L"Y(t)", fmt=:svg, legend=false)
# plot(T,Y,linetype=:steppost, xlabel=L"t", ylabel=L"Y(t)", fmt=:pdf, legend=false)
# fname = string(pwd(),"/figs/unitpoisson.pdf")
# savefig(fname)

using DiffEqBiological, Latexify
rn = @reaction_network begin
   (k₊,k₋), A + B ↔ C
    α, 0 --> C
    β, A --> 0
    γ, B --> 0
end k₊ k₋ α β γ
latexify(rn; env=:chemical)

latexify(rn)

using OrdinaryDiffEq
u0 = [0., 50., 100.]
p  = [1., 1., 1., 1., 1.]
tspan = (0.,30.)
oprob = ODEProblem(rn, u0, tspan, p)
osol = solve(oprob, Tsit5())
p1 = plot(osol, fmt=:svg, linestyle=:dash, xlabel=L"t")
xlabel!(p1,L"t")

using DiffEqJump
u0 = [0, 50, 100]
dprob = DiscreteProblem(rn, u0, tspan, p)
jprob = JumpProblem(dprob, Direct(), rn)
jsol = solve(jprob, SSAStepper())
# p2 = plot(osol, fmt=:svg, linestyle=:dash, xlabel=L"t")
# plot!(p2, jsol, fmt=:svg, xlabel=L"t")
# plot(p1,p2, size=(800,400))

# fmt = :pdf
# p1 = plot(osol, fmt=fmt, linestyle=:dash, xlabel=L"t")
# p2 = plot(osol, fmt=fmt, linestyle=:dash, xlabel=L"t")
# plot!(p2, jsol, fmt=fmt, xlabel=L"t")
# plot(p1,p2, size=(800,400))
# fname = string(pwd(),"/abtoc.pdf")
# savefig(fname)


@testset "Erode / dilate" begin
        A = zeros(4,4,3)
        A[2,2,1] = 0.8
        A[4,4,2] = 0.6
        Ae = erode(A)
        @test Ae == zeros(size(A))
        Ad = dilate(A, 1:2)
        Ar = [0.8 0.8 0.8 0;
              0.8 0.8 0.8 0;
              0.8 0.8 0.8 0;
              0 0 0 0]
        Ag = [0 0 0 0;
              0 0 0 0;
              0 0 0.6 0.6;
              0 0 0.6 0.6]
        @test Ad == cat(Ar, Ag, zeros(4,4), dims=3)
        Ae = erode(Ad, 1:2)
        Ar = [0.8 0.8 0 0;
              0.8 0.8 0 0;
              0 0 0 0;
              0 0 0 0]
        Ag = [0 0 0 0;
              0 0 0 0;
              0 0 0 0;
              0 0 0 0.6]
        @test Ae == cat(Ar, Ag, zeros(4,4), dims=3)
        # issue Images.jl #311
        @test dilate(trues(3)) == trues(3)
    end

    @testset "Opening / closing" begin
        A = zeros(4,4,3)
        A[2,2,1] = 0.8
        A[4,4,2] = 0.6
        Ao = opening(A)
        @test Ao == zeros(size(A))
        A = zeros(10,10)
        A[4:7,4:7] .= 1
        B = copy(A)
        A[5,5] = 0
        Ac = closing(A)
        @test Ac == B
    end

    @testset "Morphological Top-hat" begin
        A = zeros(13, 13)
        A[2:3, 2:3] .= 1
        Ae = copy(A)
        A[5:9, 5:9] .= 1
        Ao = tophat(A)
        @test Ao == Ae
        Aoo = tophat(Ae)
        @test Aoo == Ae
    end

    @testset "Morphological Bottom-hat" begin
        A = ones(13, 13)
        A[2:3, 2:3] .= 0
        Ae = 1 .- copy(A)
        A[5:9, 5:9] .= 0
        Ao = bothat(A)
        @test Ao == Ae
    end

    @testset "Morphological Gradient" begin
        A = zeros(13, 13)
        A[5:9, 5:9] .= 1
        Ao = morphogradient(A)
        Ae = zeros(13, 13)
        Ae[4:10, 4:10] .= 1
        Ae[6:8, 6:8] .= 0
        @test Ao == Ae
        Aee = dilate(A) - erode(A)
        @test Aee == Ae
    end

    @testset "Morphological Laplacian" begin
        A = zeros(13, 13)
        A[5:9, 5:9] .= 1
        Ao = morpholaplace(A)
        Ae = zeros(13, 13)
        Ae[4:10, 4:10] .= 1
        Ae[5:9, 5:9] .= -1
        Ae[6:8, 6:8] .= 0
        @test Ao == Ae
        Aee = dilate(A) + erode(A) - 2A
        @test Aee == Ae
    end


# bbox is a "bounding box" representing outer limits of latitude & longitude: [bottomlat leftlon; toplat rightlon]
# rasterdensity is number of pixels per degree, so rasterdensity = 1/res with res = raster resolution measured in degrees
# TO DO: make these consistent with lon, lat convention in rest of code, use resolution instead of rasterdensity
function bbox2ranges(bbox, rasterdensity)
    latindexes, lonindexes = bbox2indexes(bbox, rasterdensity)
    latindex = latindexes[1] : latindexes[2]
    lonindex = lonindexes[1] : lonindexes[2]
    return latindex, lonindex
end

function bbox2indexes(bbox, rasterdensity)
    latindexes = round.(Int, reverse(rasterdensity*(90 .- bbox[:,1]), dims=1) + [1 0]')
    lonindexes = round.(Int, rasterdensity*(bbox[:,2] .- (-180)) + [1 0]')
    return latindexes, lonindexes
end

function roundbbox(bbox, rasterdensity)
    newbbox = zeros(2,2)
    newbbox[1,:] = floor.(bbox[1,:]*rasterdensity)/rasterdensity
    newbbox[2,:] = ceil.(bbox[2,:]*rasterdensity)/rasterdensity
    return newbbox
end

filterrange(range, lowhigh) = range[(range .>= lowhigh[1]) .& (range .<= lowhigh[2])]

function getlats(bbox, rasterdensity, halfshift)
	bottomlat, toplat = bbox[:,1]
	res = 1/rasterdensity
	shift = halfshift ? res/2 : 0.0
	return bottomlat + shift : res : toplat - shift
	# return toplat - shift : -res : bottomlat + shift
end

function getlons(bbox, rasterdensity, halfshift)
	leftlon, rightlon = bbox[:,2]
	res = 1/rasterdensity
	shift = halfshift ? res/2 : 0.0
	return leftlon + shift : res : rightlon - shift
end

# Boolean function, true for a leap year (for Gregorian calendar, assumes year >= 1583).
# The Gregorian calendar has 97 leap years every 400 years:
#     Every year divisible by 4 is a leap year.
#     However, every year divisible by 100 is not a leap year.
#     However, every year divisible by 400 is a leap year after all.
# So, 1700, 1800, 1900, 2100, and 2200 are not leap years, but 1600, 2000, and 2400 are leap years.
leapyear(year) = divisible(year,4) && (!divisible(year,100) || divisible(year,400))
divisible(x,y) = x % y == 0

function meshgrid2d(x,y)
	xx = repeat(x', length(y), 1)
	yy = repeat(y, 1, length(x))
	return xx, yy
end

# Borrowed from Matlab, see fspecial('disk', rad)
function diskfilterkernel(rad)
	crad  = ceil(Int, rad-0.5)
	x,y = meshgrid2d(-crad:crad,-crad:crad)
	maxxy = max.(abs.(x),abs.(y))
	minxy = min.(abs.(x),abs.(y))
	m1 = (rad^2 .< (maxxy .+ 0.5).^2 + (minxy .- 0.5).^2) .* (minxy .- 0.5) +
	  (rad^2 .>= (maxxy .+ 0.5).^2 + (minxy .- 0.5).^2) .* sqrt.(max.(0, rad^2 .- (maxxy .+ 0.5).^2))
	m2 = (rad^2 .> (maxxy .- 0.5).^2 + (minxy .+ 0.5).^2) .* (minxy .+ 0.5) +
	  (rad^2 .<= (maxxy .- 0.5).^2 + (minxy .+ 0.5).^2) .* sqrt.(max.(rad^2 .- (maxxy .- 0.5).^2))
	sgrid = (rad^2 * (0.5*(asin.(m2/rad) .- asin.(m1/rad)) + 0.25*(sin.(2*asin.(m2/rad)) - sin.(2*asin.(m1/rad)))) +
	     		- (maxxy .- 0.5).*(m2-m1) + (m1 .- minxy .+ 0.5)) .*
	      ((((rad^2 .< (maxxy .+ 0.5).^2 + (minxy .+ 0.5).^2) .&
	      	(rad^2 .> (maxxy .- 0.5).^2 + (minxy .- 0.5).^2)) .|
	      	((minxy .== 0) .& (maxxy .- 0.5 .< rad) .& (maxxy .+ 0.5 .>= rad))))
	sgrid = sgrid .+ ((maxxy .+ 0.5).^2 + (minxy .+ 0.5).^2 .< rad^2)
	sgrid[crad+1,crad+1] = min.(pi*rad^2, pi/2)
	if ((crad>0) && (rad > crad-0.5) && (rad^2 < (crad-0.5)^2+0.25))
		m1a  = sqrt(rad^2 - (crad - 0.5).^2)
		m1n = m1a/rad
		sg0 = 2*(rad^2*(0.5*asin(m1n) + 0.25*sin(2*asin(m1n)))-m1a*(crad-0.5))
		sgrid[2*crad+1,crad+1] = sg0
		sgrid[crad+1,2*crad+1] = sg0
		sgrid[crad+1,1]        = sg0
		sgrid[1,crad+1]        = sg0
		sgrid[2*crad,crad+1]   = sgrid[2*crad,crad+1] - sg0
		sgrid[crad+1,2*crad]   = sgrid[crad+1,2*crad] - sg0
		sgrid[crad+1,2]        = sgrid[crad+1,2]      - sg0
		sgrid[2,crad+1]        = sgrid[2,crad+1]      - sg0
	end
	sgrid[crad+1,crad+1] = min(sgrid[crad+1,crad+1], 1)
	h = centered(sgrid/sum(sgrid))
end

meandrop(x; dims=dims) = dropdims(mean(x, dims=dims), dims=dims)
sumdrop(x; dims=dims) = dropdims(sum(x, dims=dims), dims=dims)

shifthalfcell(data, kernel=centered([.25 .25; .25 .25])) = imfilter(data, kernel)

function shiftallcells!(windCF)
    yearlength = size(windCF, 1)
    updateprogress = Progress(yearlength, 1)
    kernel = centered([.25 .25; .25 .25])
    for i=1:yearlength
        windCF[i,:,:] = shifthalfcell(windCF[i,:,:], kernel)
        next!(updateprogress)
    end
end

function max!(x::AbstractArray, val)
    @inbounds for i in eachindex(x)
        x[i] = max(x[i], val)
    end
	x
end

selfmap!(f,x) = map!(f,x,x)

# Apply a function fn to the data matrix by chunks to avoid memory issues
function gridsplit(data::AbstractArray, fn::Function, resulttype::DataType; nmax=9000, overlap=250)
    rows, cols = size(data)
    nparts = length(1:nmax:rows) * length(1:nmax:cols)
    nparts > 1 && println("\nSplitting data matrix into $nparts parts to avoid memory issues...")
    result = zeros(resulttype, size(data))
    part = 0
    for r = 1:nmax:rows, c = 1:nmax:cols
        part += 1
        nparts > 1 && println("\nPart $part/$nparts:")
        rowrange_in = max(r-overlap, 1):min(r-1+nmax+overlap, rows)
        colrange_in = max(c-overlap, 1):min(c-1+nmax+overlap, cols)
        rowrange_out = r:min(r-1+nmax, rows)
        colrange_out = c:min(c-1+nmax, cols)
        rowdelta = rowrange_out[1] - rowrange_in[1] + 1
        coldelta = colrange_out[1] - colrange_in[1] + 1
        rowrange_delta = rowdelta:(rowdelta + length(rowrange_out) - 1)
        colrange_delta = coldelta:(coldelta + length(colrange_out) - 1)
        result[rowrange_out,colrange_out] = fn(data[rowrange_in,colrange_in])[rowrange_delta,colrange_delta]
    end
    return result
end

# function gridsurface(x, y; color=color, nmax=9000, args...)
#     rows, cols = size(color)
#     rowchunks, colchunks = ceil(Int, rows/nmax), ceil(Int, cols/nmax)
#     rowsplit, colsplit = ceil(Int, rows/rowchunks), ceil(Int, cols/colchunks)
#     nparts = rowchunks * colchunks
#     nparts > 1 && println("\nSplitting surface plat into $nparts parts to avoid Makie GPU limits...")
#     part = 0
#     scene = Scene()
#     for r = 1:rowsplit:rows, c = 1:colsplit:cols
#         part += 1
#         nparts > 1 && println("\nPart $part/$nparts:")
#         rowrange = r:min(r-1+rowsplit, rows)
#         colrange = c:min(c-1+colsplit, cols)
#         surface!(x[rowrange,colrange], y[rowrange,colrange]; color=color[rowrange,colrange], args...)
#     end
#     return scene
# end

row2lon(row::Int, res) = (row - 0.5) * res - 180
col2lat(col::Int, res) = 90 - (col - 0.5) * res
lon2row(lon::Real, res) = floor(Int, mod(180+lon, 360) / res) + 1
lat2col(lat::Real, res) = floor(Int, min(180-res/2, 90-lat) / res) + 1

function rowcol2lonlat(rowcol::Tuple{Int,Int}, res)
    row, col = rowcol
    @assert 1 <= row <= 360/res
    @assert 1 <= col <= 180/res
    return row2lon(row,res), col2lat(col,res)
end

function lonlat2rowcol(lonlat::Tuple{<:Real,<:Real}, res)
    lon, lat = lonlat
    @assert -90 <= lat <= 90
    return lon2row(lon,res), lat2col(lat, res)
end

# convert indexes of datasets with 0.01 degree resolution to indexes of ERA5 resolution (0.28125 degrees)
function eraranges(lonrange, latrange, res, erares)
    topleft = rowcol2lonlat((lonrange[1],latrange[1]), res)
    bottomright = rowcol2lonlat((lonrange[end],latrange[end]), res)
    row1, col1 = lonlat2rowcol(topleft, erares)
    row2, col2 = lonlat2rowcol(bottomright, erares)
    nlons = round(Int, 360/erares)
    eralonranges = row2 >= row1 ? [row1:row2] : [row1:nlons, 1:row2]
    eralatrange = col1:col2
    return eralonranges, eralatrange
end

function eraindexlookup(lons, lats, eralonranges, eralatrange)
    eralonrange = length(eralonranges) == 1 ? eralonranges[1] : [eralonranges[1]; eralonranges[2]]
    lonindexlookup = zeros(Int,1280)
    latindexlookup = zeros(Int,640)
    for (i,r) in enumerate(eralonrange)
        lonindexlookup[r] = i
    end
    for (i,r) in enumerate(eralatrange)
        latindexlookup[r] = i
    end
    eralonindexes = lonindexlookup[lon2row.(lons,0.28125)]
    eralatindexes = latindexlookup[lat2col.(lats,0.28125)]
    return eralonindexes, eralatindexes
end

# Automatically detect the "bounding box" of nonzero data in a matrix.
# Returns a tuple of indexes of the box containing data, (lon,lat) or (row,col).
function getbboxranges(regions::AbstractMatrix, padding::Int=0)
    data = (regions .> 0) .& (regions .!= NOREGION)
    lonrange = dataindexes_lon(vec(any(data, dims=2)), padding)     # all longitudes with region data
    latrange = dataindexes_lat(vec(any(data, dims=1)), padding)     # all latitudes with region data
    return lonrange, latrange
end

# Given a vector of Bool indicating nonzero data along latitudes, return a range
# of indexes indicating the area where there is data.
# Optionally add some padding elements to the data area (to ensure that there is
# some offshore territory around the regions).
function dataindexes_lat(latdata::AbstractVector, padding::Int=0)
    len = length(latdata)
    first, last = extrema(findall(latdata))         # first & last indexes of the data region
    return max(1, first-padding):min(len, last+padding)   # lat indexes of region elements
end

# Given a vector of Bool indicating nonzero data along longitudes, return a vector
# of indexes indicating the area where there is data. Do this by eliminating the
# longest contiguous sequence of zero data, considering wraparound (at lon=180).
# Optionally add some padding elements to the data area (to ensure that there is
# some offshore territory around the regions).
function dataindexes_lon(londata::AbstractVector, padding::Int=0)
    len = length(londata)
    seq = longest_circular_sequence(londata, false) # first & last indexes of the longest contiguous sequence of NONregion elements
    first, last = seq[2]+1, seq[1]-1                # the rest are region elements (including "holes" in the sequence)
    last = last >= first ? last : last + len        # make sure last > first so the loop below works
    return [mod1(i,len) for i = first-padding:last+padding]   # lon indexes of region elements
end

# Return the start and end indexes of the longest contiguous sequence of v such that v[i] == x,
# allowing the sequence to wrap around the end.
function longest_circular_sequence(v::AbstractVector, x)
    longest = Int[]
    current = Int[]
    haswrapped = false
    i = 1
    len = length(v)
    sequencelength(seq) = isempty(seq) ? 0 : mod(seq[2]-seq[1], len) + 1
    while true
        if v[i] == x
            if isempty(current)
                current = [i,i]
            else
                current[2] = i
            end
        else
            longest = sequencelength(current) > sequencelength(longest) ? current : longest
            haswrapped && return longest
            current = Int[]
        end
        if !haswrapped && i==len
            haswrapped = true
            i = 1
            sequencelength(current) == len && return current
        else
            i += 1
        end
    end
end

function coordmap(len::Int, croprange)
    coords = zeros(Int, len)
    indexes = (1:len)[croprange]
    for (i,n) = enumerate(indexes)
        coords[n] = i
    end
    return coords
end

# area of a grid cell in km2 for a given latitude (in degrees) and raster resolution
rastercellarea(lat, res) = cosd(lat) * (2*6371*π/(360/res))^2

function eralonlat(options, lonrange, latrange)
    @unpack res, erares = options

    lons = (-180+res/2:res:180-res/2)[lonrange]         # longitude values (pixel center)
    lats = (90-res/2:-res:-90+res/2)[latrange]          # latitude values (pixel center)
    cellarea = rastercellarea.(lats, res)

    eralonranges, eralatrange = eraranges(lonrange, latrange, res, erares)
    eralonrange = length(eralonranges) == 1 ? eralonranges[1] : [eralonranges[1]; eralonranges[2]]
    eralons = (-180+erares/2:erares:180-erares/2)[eralonrange]     # longitude values (pixel center)
    eralats = (90-erares/2:-erares:-90+erares/2)[eralatrange]      # latitude values (pixel center)

    lonmap = coordmap(round(Int, 360/res), lonrange)      # map full longitude indexes to cropped indexes
    latmap = coordmap(round(Int, 180/res), latrange)      # ditto latitude

    return eralons, eralats, lonmap, latmap, cellarea
end

function uncrop(croppedarray, lonrange, latrange, res)
    nlon, nlat = round(Int, 360/res), round(Int, 180/res)
    full = zeros(eltype(croppedarray), nlon, nlat)
    full[lonrange, latrange] = croppedarray
    return full
end

function drawmap(mapdata)
    skip = ceil(Int, maximum(size(mapdata))/3000)
    mirrormap = reverse(mapdata[1:skip:end,1:skip:end], dims=2)
    # display(heatmap(mirrormap, size=(1200, 900), c=[cgrad(:viridis)[x] for x in 0.0:0.2:1.0]))
    # display(heatmap(mirrormap, size=(1200, 900), c=cgrad(:viridis, [0, 0.5, 1])))
    heatmap(mirrormap)
    # display(countmap(mirrormap[:]))
end

# I can't believe Makie can't do this more conveniently.
function plotlines(x,y)
    scene = Scene()
    for i = 1:size(y,2)
        lines!(scene, x, y[:,i], color=ColorSchemes.Set1_9[i])
    end
    scene
end

function drawregionmap(regionname)
    plotly()
    regions, offshoreregions, regionlist, lonrange, latrange = loadregions(regionname)
    skip = ceil(Int, maximum(size(mapdata))/3000)
    reg = reverse(regions[1:skip:end,1:skip:end]', dims=1)
    display(heatmap(reg .+ (reg.>0).*20, size=(1200, 900)))
    reg = reverse(offshoreregions[1:4:end,1:4:end]', dims=1)
    display(heatmap(reg .+ (reg.>0).*20, size=(1200, 900)))
end

function resize_categorical(regions, regionlist, lonrange, latrange; skipNOREGION=false)
    res = 0.01          # resolution of auxiliary datasets [degrees per pixel]
    erares = 0.28125    # resolution of ERA5 datasets [degrees per pixel]

    eralons, eralats, lonmap, latmap, cellarea = eralonlat(Dict(:res=>res, :erares=>erares), lonrange, latrange)
    numreg = length(regionlist)
    nlons, nlats = length(eralons), length(eralats)

    RegType = eltype(regions)
    smallregions = zeros(RegType, nlons, nlats)
    count0 = Dict(i => 0 for i = 0:numreg)
    count0[NOREGION] = 0

    for (j, eralat) in enumerate(eralats)
        colrange = latmap[lat2col(eralat+erares/2, res):lat2col(eralat-erares/2, res)-1]
        for (i, eralon) in enumerate(eralons)
            # get all high resolution row and column indexes within this ERA5 cell
            rowrange = lonmap[lon2row(eralon-erares/2, res):lon2row(eralon+erares/2, res)-1]
            count = copy(count0)
            foundnonzero = false
            for c in colrange, r in rowrange
                (c == 0 || r == 0) && continue
                reg = regions[r,c]
                skipNOREGION && (reg == 0 || reg == NOREGION) && continue
                count[reg] += 1
                foundnonzero = true
            end
            smallregions[i,j] = foundnonzero ? findmax(count)[2] : 0
        end
    end

    return smallregions
end

function matlab2elin(; gisregion="Europe8", year=2018)
    filenamesuffix = ""
    _, _, regionlist, _, _ = loadregions(gisregion)

    # CF_pvrooftop, capacity_pvrooftop
    region = string.(regionlist)
    tech = ["PVplantA", "PVplantB", "PVroof", "CSPA", "CSPB", "WindonshoreA", "WindonshoreB", "Windoffshore"]
    classname = ["PVPA", "PVPB", "PVR", "CSPA", "CSPB", "WONA", "WONB", "WOFF"]
    capvar = ["capacity_pvplantA", "capacity_pvplantB", "capacity_pvrooftop", "capacity_cspplantA", "capacity_cspplantB",
                "capacity_onshoreA", "capacity_onshoreB", "capacity_offshore"]
    cfvar = ["CFtime_pvplantA", "CFtime_pvplantB", "CFtime_pvrooftop", "CFtime_cspplantA", "CFtime_cspplantB",
                "CFtime_windonshoreA", "CFtime_windonshoreB", "CFtime_windoffshore"]

    winddata = matread(in_datafolder("output", "GISdata_wind$(year)_$gisregion$filenamesuffix.mat"))
    solardata = matread(in_datafolder("output", "GISdata_solar$(year)_$gisregion$filenamesuffix.mat"))
    data = merge(winddata, solardata)

    # read number of classes from wind & solar GIS output
    nwindclasses = [size(winddata[varname], 2) for varname in capvar[6:8]]
    nsolarclasses = [size(solardata[varname], 2) for varname in capvar[1:5]]
    nclasses = [nsolarclasses; nwindclasses]

    for t = 1:length(tech)
        open(in_datafolder("output", "capacity_$(tech[t]).inc"), "w") do f
            for (r,reg) in enumerate(region)
                for c = 1:nclasses[t]
                    val = data[capvar[t]][r,c]
                    !isnan(val) && val > 0 && @printf(f, "%-3s . %s%-2d %12.6f\n", reg, classname[t], c, val)
                end
            end
        end
        open(in_datafolder("output", "cf_$(tech[t]).inc"), "w") do f
            for (r,reg) in enumerate(region)
                for c = 1:nclasses[t]
                    for h = 1:8760
                        val = data[cfvar[t]][h,r,c]
                        !isnan(val) && val > 0 && @printf(f, "%-3s . %s%-2d . h%04d %10.6f\n", reg, classname[t], c, h, val)
                    end
                end
            end
        end
    end
end

function dms2deg(dms)
    sg = sign(dms)
    dms = abs(dms)
    d = sg*floor(dms/10000)
    dms = dms-10000*d
    m = floor(dms/100)
    return d + m/60 + (dms-100*m)/3600
end

# Capital recovery factor for investments
CRF(r, T) = r / (1 - 1/(1+r)^T)

in_datafolder(names...) = joinpath(getconfig("datafolder"), names...)


using Test
using HOI4CombatSimulator
using Statistics

macro rep(n::Int, exp::Expr)
    quote
        resp_array = Array{Any, 1}(undef, $n)
        for i in 1:$n
            resp_array[i] = $exp
        end
        m = length(resp_array)
        n = length(resp_array[1])
        mat = zeros(m, n)
        for i in 1:m
            for j in 1:n
                mat[i, j] = resp_array[i][j]
            end
        end
        mat
    end
end

@testset "HOI4CombatSimulator" begin
    @testset "get_unit_list" begin
        unit_list = HOI4CombatSimulator.get_unit_list()
        @test length(unit_list) > 0
    end

    @testset "wiki_example" begin
        HP_loss_Org_loss_mat = @rep 1000 HOI4CombatSimulator.fire(
            evade=500, HPPercent=0.99, ShooterHardAtk=0,
            ShooterSoftAtk=1000, VictimHardness=0,
            ShooterPierce=20, ShooterArmor=10,
            VictimPierce=1, VictimArmor=0
        ) # (n, 2)
        HP_loss_arr = HP_loss_Org_loss_mat[:, 1]
        Org_loss_arr = HP_loss_Org_loss_mat[:, 2]
        diff = abs(mean(Org_loss_arr) - 3.9)
        tol = (std(Org_loss_arr) / sqrt(1000))*5 # 5 sigma
        @test  diff < tol
    end

    @testset "conv" begin
        x = HOI4CombatSimulator.Inference.fast_conv(ones(2)/2, 2)
        @test maximum(abs.(x .- [0.25, 0.5, 0.25])) < 1e-6
    end

    @testset "DiceSumDistribution" begin
        d6x2 = HOI4CombatSimulator.Inference.dice_sum_distribution(6, 2)
        @test d6x2.support[1] == 2
        @test d6x2.support[end] == 12
        @test length(d6x2.support) == 11
        @test abs(d6x2.p[1] - 1/36) < 1e-6

        d6x4 = HOI4CombatSimulator.Inference.dice_sum_distribution(6, 4)
        @test d6x4.p[end] == (1/6)^4
        @test abs(sum(d6x4.p)- 1.) < 1e-6
    end

    @testset "prob_fire_loss" begin
        HP_Org_loss_mat = HOI4CombatSimulator.Inference.prob_fire(104.0, 160.0, 6)
        @test size(HP_Org_loss_mat, 1) > 1
        @test size(HP_Org_loss_mat, 2) > 1
        # TODO: detailed test
    end
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 9396fcac-b0b6-11eb-3a60-9f2ce25df953
begin
	using Pkg
	Pkg.activate(".")
	Pkg.instantiate()
end

# ╔═╡ 4ecf1267-aac9-4e8f-a522-5595b1205f43
begin
	using CSV, RCall, DataFrames, MixedModels
	using LinearAlgebra, Statistics, StatsBase
	using CategoricalArrays, Arrow
	using AlgebraOfGraphics
	using AlgebraOfGraphics: linear
	using CairoMakie  # for Scatter
end

# ╔═╡ eebef932-1d63-41e5-998f-9748379c43af
md"""
# Mixed Models Tutorial: Conditional Modes

Ths script uses a subset of data reported in Fühner, Golle, Granacher, & Kliegl (2021).
Physical fitness in third grade of primary school:
A mixed model analysis of 108,295 children and 515 schools.

To circumvent delays associated with model fitting we work with a reduced data set and less complex models than those in the reference publication. All the data to reproduce the models in the publication are used here, too; the script requires only a few changes to specify the more complex models in the paper.

The script is structured in three main sections:

1. **Setup** with reading and examing the data, plotting the main results, and specifying the contrasts for the fixed factor `Test`.
2. **Extraction of conditional modes**
3. **Caterpillar plot**

## 1. Setup
### 1.0 Packages and functions
"""

# ╔═╡ faafd359-fb38-4f49-9dfd-19cb4f0c5a54
begin
	function viewdf(x)
		DataFrame(Arrow.Table(take!(Arrow.write(IOBuffer(), x))))
	end
end

# ╔═╡ ee85abd9-0172-44d3-b03a-c6a780d33c72
md"""
### 1.1 Readme for 'EmotikonSubset.rds'

1. Cohort: 9 levels; 2011-2019
2. School: 46 levels
3. Child: 11566 levels; all children were between 6.0 and 6.99 years at legal keydate (30 September) of school enrollement.
4. Sex: 5893 girls, 5673 boys
5. age: test date - middle of month of birthdate (ranges between 7.9 and 9.2)
6. Test: 5 levels
     + Endurance (`Run`):  6 minute endurance run [m]; to nearest 9m in 9x18m field
     + Coordination (`Star_r`): star coordination run [m/s]; 9x9m field, 4 x diagonal = 50.912 m
     + Speed(`S20_r`): 20-meters sprint [m/s]
     + Muscle power low (`SLJ`): standing long jump [cm]
     + Muscle power up (`BPT`): 1-kg medicine ball push test [m]
 7. score - see units

### 1.2 Preprocessing
#### Read data
"""

# ╔═╡ 39443fb0-64ec-4a87-b072-bc8ad6fa9cf4
dat = rcopy(R"readRDS('./data/EmotikonSubset.rds')");

# ╔═╡ 0051ec83-7c30-4d28-9dfa-4c6f5d0259fa
md"""
#### Transformations
"""

# ╔═╡ 7440c439-9d4c-494f-9ac6-18c9fb2fe144
begin
	transform!(dat, :age, :age => (x -> x .- 8.5) => :a1); # centered age (linear)
	transform!(dat,  :a1, :a1  => (x -> x.^2) => :a2);     # centered age (quadr.)
	select!(groupby(dat,  :Test), :, :score => zscore => :zScore); # z-score
end;

# ╔═╡ 16c3cdaa-49fa-46d5-a844-03094329fe4c
viewdf(dat)

# ╔═╡ 64cc1f8e-f831-4a53-976f-dc7600b5634d
# ... by Test and Sex
begin
	dat2 = combine(groupby(dat, [:Test, :Sex]),
                             :score => mean, :score  => std,
                             :zScore => mean, :zScore => std)
	viewdf(dat2)
end

# ╔═╡ c6c6f056-b8b9-4190-ac14-b900bafa04df
md"""
## 2 LMMs and extraction of conditional modes
### 2.1 _SeqDiffCoding_ of `Test`

_SeqDiffCoding_ was used in the publication. This specification tests pairwise
differences between the five neighboring levels of `Test`, that is:

+ H1: `Star_r` - `Run` (2-1)
+ H2: `S20_r` - `Star_r` (3-2)
+ H3: `SLJ` - `S20_r` (4-3)
+ H4: `BPT` - `SLJ` (5-4)

Various options for contrast coding are the topic of the *MixedModelsTutorial_contrasts.jl*
notbebook.
"""

# ╔═╡ c5326753-a03b-4739-a82e-90ffa7c1ebdb
begin
	contr = merge(
        Dict(nm => Grouping() for nm in (:School, :Child, :Cohort)),
		Dict(:Sex => EffectsCoding(; levels=["Girls", "Boys"])),
	    Dict(:Test => SeqDiffCoding(; levels=["Run", "Star_r", "S20_r", "SLJ", "BPT"])),
        Dict(:TestHC => HelmertCoding(; levels=["S20_r", "SLJ", "Star_r", "Run", "BPT"])),
	   )
end;

# ╔═╡ e3cf1ee6-b64d-4356-af23-1dd5c7a0fec6
md"""
#### 2.2 LMMs  `m11` for the reduced data

We fit LMM `m11_Cohort` (only nine levels, six VCs for Test; no CPs) and `m11_School` (46 levels; VCs and CPs for test and linear effect of age)
"""

# ╔═╡ 0dd0d060-81b4-4cc0-b306-dda7589adbd2
f11_Cohort =  @formula zScore ~ 1 + Test * a1 * Sex +  zerocorr(1 + Test | Cohort)

# ╔═╡ f68f6884-7c37-4868-b0b7-01ad450ee383
m11_Cohort = fit(MixedModel, f11_Cohort, dat)

# ╔═╡ ef1dfbf4-700c-4093-9824-01368c511531
VarCorr(m11_Cohort)

# ╔═╡ ace9a2f8-075a-49b5-94eb-12be4ffb6568
f11_School =  @formula zScore ~ 1 + Test * a1 * Sex +  (1 + Test + a1 | School)

# ╔═╡ 167fbd71-2a7b-4afa-8ffb-815dcda63129
m11_School = fit(MixedModel, f11_School, dat)

# ╔═╡ 68ee8abb-04e6-45ee-8f0a-033e7d2ef57b
VarCorr(m11_School)

# ╔═╡ dd8b5c05-c1c8-4778-b179-1ad39f50f8d3
f12 = @formula zScore ~ 1 + Test * a1 * Sex +
                       (1 + Test + a1 | School) + zerocorr(1 + Test | Cohort)

# ╔═╡ 25376bb3-4673-4c89-8600-6e3da85971b4
m12 = fit(MixedModel, f12, dat)

# ╔═╡ e7cdd374-57c3-4d3c-9b1b-b6a211733b81
VarCorr(m12)

# ╔═╡ 5ffbc2ca-02df-4dd7-8e3b-f692d7d12adb

# ╔═╡ 1a2586a6-ff9d-4067-a5af-6c0ee6929e81
md"""
### 2.2 Conditional modes
"""

# ╔═╡ d2c0979e-ed52-4b56-b3ec-58db1609843e
cms = raneftables(m11_School)

# ╔═╡ 776cec2c-c566-4f40-81f5-12546c6a4da7
cms_Chrt = DataFrame(cms.School);

# ╔═╡ cf3f5312-6702-4069-95aa-862dc54d9a34
cvs = condVar(m11_School)

# ╔═╡ 10ef55d6-6048-4353-92e5-e4621bfb130e
md"""
### 3 MixedModelsMakie functions

Information on random effects conditional modes/means, variances, etc.
"""

# ╔═╡ cbe2197e-12c7-4d55-8817-d25d7b7cc5bd
struct RanefInfo{T<:AbstractFloat}
    cnames::Vector{String}
    levels::Vector
    ranef::Matrix{T}
    stddev::Matrix{T}
end

# ╔═╡ 32acfdf0-ecb6-4acb-904f-d7b963771d6d
md"""

ranefinfo(m::LinearMixedModel)

Return a `NamedTuple{fnames(m), NTuple(k, RanefInfo)}` from model `m`
"""

# ╔═╡ 90ce17cc-d497-4605-aa62-be69343e03eb
function ranefinfo(m::LinearMixedModel{T}) where {T}
    fn = fnames(m)
    val = sizehint!(RanefInfo[], length(fn))
    for (re, eff, cv) in zip(m.reterms, ranef(m), condVar(m))
        push!(
            val,
            RanefInfo(
                re.cnames,
                re.levels,
                Matrix(adjoint(eff)),
                Matrix(adjoint(dropdims(sqrt.(mapslices(diag, cv, dims=1:2)), dims=2))),
            )
        )
    end
    NamedTuple{fn}((val...,))
end

# ╔═╡ cd4d8409-cb5e-4f38-bd63-c54b8f09e5ff
function caterpillar!(f::Figure, r::RanefInfo; orderby=1)
    rr = r.ranef
    vv = view(rr, :, orderby)
    ord = sortperm(vv)
    y = axes(rr, 1)
    cn = r.cnames
    axs = [Axis(f[1, j]) for j in axes(rr, 2)]
    linkyaxes!(axs...)
    for (j, ax) in enumerate(axs)
        xvals = view(rr, ord, j)
        scatter!(ax, xvals, y, color=(:red, 0.2))
        errorbars!(ax, xvals, y, view(r.stddev, ord, j), direction=:x)
        ax.xlabel = cn[j]
        ax.yticks = y
        j > 1 && hideydecorations!(ax, grid=false)
    end
    axs[1].yticks = (y, r.levels[ord])
    f
end

# ╔═╡ bd10bfa8-9e22-4927-b83e-eb80f328c6a3
md"""
    caterpillar(m::LinearMixedModel, gf::Symbol)

Returns a `Figure` of a "caterpillar plot" of the random-effects means and prediction intervals

A "caterpillar plot" is a horizontal error-bar plot of conditional means and standard deviations
of the random effects.
"""

# ╔═╡ 727059df-d7a5-4714-8dcb-6568d9119510
function caterpillar(m::LinearMixedModel, gf::Symbol=first(fnames(m)))
    caterpillar!(Figure(resolution=(1000,800)), ranefinfo(m)[gf])
end

# ╔═╡ 5a62529e-cc10-4e2c-8d75-125951f2dc34
md"""
  caterpillar!(f::Figure, r::RanefInfo; orderby=1)

Add Axes of a caterpillar plot from `r` to `f`.

The order of the levels on the vertical axes is increasing `orderby` column
of `r.ranef`, usually the `(Intercept)` random effects.
"""

# ╔═╡ a9043efa-17e5-498f-a744-7c0690446263
md"""
## 4. Caterpillar plots
### 4.1 LMM `m11_Cohort`
"""

# ╔═╡ 07a00ce2-4658-49ff-a384-4b5230691a7b
caterpillar(m11_Cohort)

# ╔═╡ 53ed436c-82b9-4331-bef4-0e034c233e7f
md"""
### 4.2 LMM `m11_School`
"""

# ╔═╡ 44b78423-e28c-4787-921e-e4ba4173a31c
caterpillar(m11_School)

# ╔═╡ 8e15a0b8-b445-4bec-8016-7b6eacf3c889
md"""
### 4.3 LMM `m12`
"""

# ╔═╡ 90dc96b9-3cc6-47d3-84d0-6c3eb8a05a5d
last(fnames(m12))

# ╔═╡ 12518738-2269-4019-9572-bfe6c9dc224f
caterpillar(m12)

# ╔═╡ Cell order:
# ╠═9396fcac-b0b6-11eb-3a60-9f2ce25df953
# ╟─eebef932-1d63-41e5-998f-9748379c43af
# ╠═4ecf1267-aac9-4e8f-a522-5595b1205f43
# ╠═faafd359-fb38-4f49-9dfd-19cb4f0c5a54
# ╟─ee85abd9-0172-44d3-b03a-c6a780d33c72
# ╠═39443fb0-64ec-4a87-b072-bc8ad6fa9cf4
# ╟─0051ec83-7c30-4d28-9dfa-4c6f5d0259fa
# ╠═7440c439-9d4c-494f-9ac6-18c9fb2fe144
# ╠═16c3cdaa-49fa-46d5-a844-03094329fe4c
# ╠═64cc1f8e-f831-4a53-976f-dc7600b5634d
# ╟─c6c6f056-b8b9-4190-ac14-b900bafa04df
# ╠═c5326753-a03b-4739-a82e-90ffa7c1ebdb
# ╟─e3cf1ee6-b64d-4356-af23-1dd5c7a0fec6
# ╠═0dd0d060-81b4-4cc0-b306-dda7589adbd2
# ╠═f68f6884-7c37-4868-b0b7-01ad450ee383
# ╠═ef1dfbf4-700c-4093-9824-01368c511531
# ╠═ace9a2f8-075a-49b5-94eb-12be4ffb6568
# ╠═167fbd71-2a7b-4afa-8ffb-815dcda63129
# ╠═68ee8abb-04e6-45ee-8f0a-033e7d2ef57b
# ╠═dd8b5c05-c1c8-4778-b179-1ad39f50f8d3
# ╠═25376bb3-4673-4c89-8600-6e3da85971b4
# ╠═e7cdd374-57c3-4d3c-9b1b-b6a211733b81
# ╠═5ffbc2ca-02df-4dd7-8e3b-f692d7d12adb
# ╠═1a2586a6-ff9d-4067-a5af-6c0ee6929e81
# ╠═d2c0979e-ed52-4b56-b3ec-58db1609843e
# ╠═776cec2c-c566-4f40-81f5-12546c6a4da7
# ╠═cf3f5312-6702-4069-95aa-862dc54d9a34
# ╟─10ef55d6-6048-4353-92e5-e4621bfb130e
# ╠═cbe2197e-12c7-4d55-8817-d25d7b7cc5bd
# ╟─32acfdf0-ecb6-4acb-904f-d7b963771d6d
# ╠═90ce17cc-d497-4605-aa62-be69343e03eb
# ╠═cd4d8409-cb5e-4f38-bd63-c54b8f09e5ff
# ╟─bd10bfa8-9e22-4927-b83e-eb80f328c6a3
# ╠═727059df-d7a5-4714-8dcb-6568d9119510
# ╟─5a62529e-cc10-4e2c-8d75-125951f2dc34
# ╟─a9043efa-17e5-498f-a744-7c0690446263
# ╠═07a00ce2-4658-49ff-a384-4b5230691a7b
# ╟─53ed436c-82b9-4331-bef4-0e034c233e7f
# ╠═44b78423-e28c-4787-921e-e4ba4173a31c
# ╟─8e15a0b8-b445-4bec-8016-7b6eacf3c889
# ╠═90dc96b9-3cc6-47d3-84d0-6c3eb8a05a5d
# ╠═12518738-2269-4019-9572-bfe6c9dc224f


# Binary heap (non-mutable)

#################################################
#
#   core implementation
#
#################################################

function _heap_bubble_up!{Comp,T}(comp::Comp, valtree::Array{T}, i::Int)
    i0::Int = i
    @inbounds v = valtree[i]

    while i > 1  # nd is not root
        p = i >> 1
        @inbounds vp = valtree[p]

        if compare(comp, v, vp)
            # move parent downward
            @inbounds valtree[i] = vp
            i = p
        else
            break
        end
    end

    if i != i0
        @inbounds valtree[i] = v
    end
end

function _heap_bubble_down!{Comp,T}(comp::Comp, valtree::Array{T}, i::Int)
    @inbounds v::T = valtree[i]
    swapped = true
    n = length(valtree)
    last_parent = n >> 1

    while swapped && i <= last_parent
        lc = i << 1
        if lc < n   # contains both left and right children
            rc = lc + 1
            @inbounds lv = valtree[lc]
            @inbounds rv = valtree[rc]
            if compare(comp, rv, lv)
                if compare(comp, rv, v)
                    @inbounds valtree[i] = rv
                    i = rc
                else
                    swapped = false
                end
            else
                if compare(comp, lv, v)
                    @inbounds valtree[i] = lv
                    i = lc
                else
                    swapped = false
                end
            end
        else        # contains only left child
            @inbounds lv = valtree[lc]
            if compare(comp, lv, v)
                @inbounds valtree[i] = lv
                i = lc
            else
                swapped = false
            end
        end
    end

    valtree[i] = v
end

function _binary_heap_pop!{Comp,T}(comp::Comp, valtree::Array{T})
    # extract root
    v = valtree[1]

    if length(valtree) == 1
        empty!(valtree)
    else
        valtree[1] = pop!(valtree)
        if length(valtree) > 1
            _heap_bubble_down!(comp, valtree, 1)
        end
    end
    v
end

function _make_binary_heap{Comp,T}(comp::Comp, ty::Type{T}, xs)
    n = length(xs)
    valtree = copy(xs)
    for i = 2 : n
        _heap_bubble_up!(comp, valtree, i)
    end
    valtree
end

#################################################
#
#   heap type and constructors
#
#################################################

type BinaryHeap{T,Comp} <: AbstractHeap{T}
    comparer::Comp
    valtree::Array{T}

    function BinaryHeap(comp::Comp)
        new(comp, Array(T,0))
    end

    function BinaryHeap(comp::Comp, xs)  # xs is an iterable collection of values
        valtree = _make_binary_heap(comp, T, xs)
        new(comp, valtree)
    end
end

function binary_minheap{T}(ty::Type{T})
    BinaryHeap{T,LessThan}(LessThan())
end

binary_maxheap{T}(ty::Type{T}) = BinaryHeap{T,GreaterThan}(GreaterThan())
binary_minheap{T}(xs::AbstractVector{T}) = BinaryHeap{T,LessThan}(LessThan(), xs)
binary_maxheap{T}(xs::AbstractVector{T}) = BinaryHeap{T,GreaterThan}(GreaterThan(), xs)

#################################################
#
#   interfaces
#
#################################################

length(h::BinaryHeap) = length(h.valtree)

isempty(h::BinaryHeap) = isempty(h.valtree)

function push!{T}(h::BinaryHeap{T}, v::T)
    valtree = h.valtree
    push!(valtree, v)
    _heap_bubble_up!(h.comparer, valtree, length(valtree))
end

top(h::BinaryHeap) = h.valtree[1]

pop!{T}(h::BinaryHeap{T}) = _binary_heap_pop!(h.comparer, h.valtree)


struct Graph{T}
    adjacency_lists :: Dict{T, Set{T}}
    edge_weights :: Dict{Tuple{T, T}, Float64}
end

function Graph{T}() where T
    return Graph{T}(Dict(), Dict())
end

function Graph()
    return Graph{Int}()
end

"""
    add_edge!(G, u, v, w)

Add edge (u, v) with weight w to graph G.
"""
function add_edge!(G, u, v, w)
    G.edge_weights[(u, v)] = w
    if !(haskey(G.adjacency_lists, u))
        G.adjacency_lists[u] = Set()
    end
    push!(G.adjacency_lists[u], v)
end

"""
    get_edge_weight(G, u, v)

Return edge (u, v) weight from graph G.
"""
function get_edge_weight(G, u, v)
    return G.edge_weights[(u, v)]
end

"""
    get_adjacent_nodes(G, u)

Return a set of adjacent nodes to u in graph G.
"""
function get_adjacent_nodes(G::Graph{T}, u) where T
    return get(G.adjacency_lists, u, Set{T}())
end

"""
    generate_random_graph(nvertices, edge_density, min_weight, max_weight)

Generate random graph for test purposes.
"""
function generate_random_graph(nvertices, edge_density, min_weight, max_weight)
    G = Graph()
    for u in 1:nvertices
        nadj = round(Int, nvertices * edge_density)
        while length(get_adjacent_nodes(G, u)) < nadj
            v = rand(1:nvertices)
            u == v && continue
            w = rand(min_weight:max_weight)
            add_edge!(G, u, v, w)
        end
    end
    return G
end


connectivity(network::WattsStrogatzNetwork) = network.k

"""
	numshortcuts(network::WattsStrogatzNetwork)

Return the number of connections that were rewired.
"""
numshortcuts(network::WattsStrogatzNetwork) = network._numShortcuts

#---------- Calculators ----------
calcNumConnections!(network::WattsStrogatzNetwork) =
	network._props.numConnections = Int(network.k * network.N / (isdirected(network) ? 1 : 2))

function generateWSAdjMat(
	N::Integer,
	k::Integer,
	β::Real;
	directed::Bool,
	seed::Integer
)
	mat = adjMat(RegularNetwork(N, k), sparse=true)
	(β <= 0) && (return mat, 0)

	numShortcuts = 0
	rng = Random.MersenneTwister(seed)
	edges = allEdges(mat, directed)

	for edge in edges
		if rand(rng) < β
			# Find a valid reconnection
		    i = rand(rng, 1:N); j = rand(rng, 1:N)
			while mat[i,j] || i == j
				i = rand(rng, 1:N)
		    	j = rand(rng, 1:N)
			end

			mat[i,j] = 1
			mat[edge[1], edge[2]] = 0

			if !directed
				mat[j,i] = 1
				mat[edge[2], edge[1]] = 0
			end

			numShortcuts += 1
		end
	end

	return SparseArrays.dropzeros(mat), numShortcuts
end

function generateWSAdjMat(
	N::Integer,
	k::Integer,
	numShortcuts::Integer;
	directed::Bool,
	seed::Integer
)
	mat = adjMat(RegularNetwork(N, k), sparse=true)
	(numShortcuts <= 0) && (return mat)

	rng = Random.MersenneTwister(seed)
	edges = allEdges(mat, directed)

	# Connections that will be rewired
	rewiredEdges = Array{Tuple{Int,Int}}(undef, numShortcuts)
	for i in 1:numShortcuts
		rewiredEdges[i] = rand(rng, edges)
		edges = filter(x -> x != rewiredEdges[i], edges)
	end

	for edge in rewiredEdges
		# Find a valid reconnection
	    i = rand(rng, 1:N); j = rand(rng, 1:N)
		while mat[i,j] || i == j
			i = rand(rng, 1:N)
	    	j = rand(rng, 1:N)
		end

		mat[i,j] = 1
		mat[edge[1], edge[2]] = 0

		if !directed
			mat[j,i] = 1
			mat[edge[2], edge[1]] = 0
		end
	end

	return SparseArrays.dropzeros(mat)
end


# using Memoization
using Printf

function run_euler(p)
    # GC.gc(false)
    t = @timed p()
    # GC.gc(true)
  @printf "%s: %d: %2.8fs\n" p t.value t.time
end

# @memoize function fib(n)
#    n < 3 ? 1 : fib(n-1) + fib(n-2)
# end

function fib_nonmem(n)
    n < 3 ? 1 : fib(n-1) + fib(n-2)
end

# For Fibonacci with BigInt
# Note: @memoize is too slow, probably because of BigInts
# (or that I have not used @memoize correctly)
fibmemBigInt_dict = Dict{Int,BigInt}()
function fibmemBigInt(n)
    get!(fibmemBigInt_dict, n) do
        n < 3 ? 1 : fibmemBigInt(BigInt(n-1)) + fibmemBigInt(BigInt(n-2))
    end
end

function isPrime(n::Int)
    if n < 2
        return false
    end
    if n === 2 || n === 3
        return true
    end

    if n % 2 === 0
        return false
    end

    m = 1+ceil(Int,sqrt(n))
    for i in 3:2:m
        if n % i === 0
            return false;
        end
    end
    return true;
end

# @memoize function isPrimeCached(n)
#    return isPrime(n)
# end

# return all primes <= n
function primes(n::Int)
    primes = [2]
    for i in 3:2:n
        if isPrime(i)
            push!(primes,i)
        end
    end
    return primes;
end

#
# returns the prime divisors of n.
#

function prime_divisors(n::Int)
    if n === 1
        return []
    end
    divisors = []
    while (n > 1)
        m = round(Int,ceil(sqrt(n)))
        for i in 2:m
            if n % i == 0 && isPrime(i)
                push!(divisors,i)
                n /= i
                continue
            end
        end
    end
    return divisors
end

# Find the primes <= n
function sieve(n::Int)
    a = Array(range(1,length=n))
    a[1] = 0;
    a[2] = 2;
    for i in 2:ceil(Int,sqrt(n))
        for j in i*i:i:n
            a[j] = 0
        end
    end
    return filter(i->a[i]>0,1:n);

end

# Prime factors of n
function factors(n::Int)
    if n == 1
        return [1]
    end
    f = []
    while n % 2 == 0
        push!(f,2)
        n /= 2
    end
    t = 3
    while n > 1 && t < ceil(sqrt(n))
        while n % t == 0
            push!(f,t)
            n /= t
        end
        t += 2
    end
    if n > 1
        push!(f,round(Int,n))
    end
    return f
end

# Is n a palindromic number?
function palindromic_number(n::Int)
    s = string(n)
    len = length(s)
    for i in 1:ceil(Int,len/2)
        if s[i] != s[len-i+1]
            return false
        end
    end
    return true
end

# is the list a palindrom?
function palindromic_list(s)
    # return s == reverse(s)
    len = length(s)
    for i in 1:ceil(Int,len/2)
        if s[i] != s[len-i+1]
            return false
        end
    end
    return true
end

# All divisors of a number, except 1 and n
function all_divisors(n::Int)
    divisors = [];
    m = round(Int,n / 2)
    for i in 2:m
        if n % i == 0
            push!(divisors,i)
        end
    end
    return divisors;

end

# // All divisors of n, including 1 and n
function all_divisors2(n::Int)
    return hcat([1], all_divisors(n),[n])
end

# All divisors of n, including 1 (but not n)
function all_divisors3(n::Int)
    divisors = [1]
    m = round(Int,n / 2)
    for i in 2:m
        if n % i == 0
            push!(divisors,i)
        end
    end
    return divisors
end

# convert a list of elements in a to a hash with
#   { element: number of occurrences of element, ...}
function collect_list(a)
    m = Dict()
    a.|>(e->
        if !haskey(m,e)
            m[e] = 1
        else
            m[e] += 1
        end

    )
    return m
end

# next permutation of array p
function next_permutation(p)
    perm = p
    n = length(perm)
    k = n - 1
    while perm[k] > perm[k+1] && k >= 0
        k = k - 1
    end
    if k > 0
        j = n
        while perm[k] > perm[j]
            j = j - 1
        end
        perm[k],perm[j] = perm[j], perm[k]
        r = n
        s = k + 1
        while r > s
            perm[r],perm[s] = perm[s],perm[r]
            r = r - 1
            s = s + 1
        end
    end
    return perm
end

#
# Generate all permutations
# This is much faster than all_permutations2
#
function all_permutations(s)
    perms = []
    push!(perms,s)
    ss = copy(s)
    p = next_permutation(copy(s));;
    push!(perms,p)
    while true
        p = next_permutation(p)
        # Note: I have to copy p here
        push!(perms,copy(p))
        if p == reverse(s)
            break
        end
    end
    return perms;
end

# Inspired by an answer from
# http://stackoverflow.com/questions/4240080/generating-all-permutations-of-a-given-string
# This is too slow. Use all_permutations instead
function all_permutations2(s)
    perms = [];
    permutation_tmp([], s, perms);
    return perms;
end

function permutation_tmp(prefix,s, perms)
    n = length(s)
    if n == 0
        push!(perms,prefix)
    else
        for i in 1:n
            permutation_tmp(vcat(prefix,s[i]), vcat(s[1:i-1],s[i+1:n]), perms)
        end
    end
end

# Convert a number to its (integer) digits
# Note: this should be replaced with
#  digits(n)
function num_to_list(n::Int)
    # return split(string(n),"").|>i->parse(Int,i)
    return digits(n)
end

# length of a number (i.e. the number of digits)
function nlen(n::Int)
    return floor(Int,log10(n))+1
end

# check is s is a pandigital number  1..9
function is_pandigital(s)
    return length(s) === 9 && !occursin("0",s) && length(Set(s)) == 9
end


using Base: String, Int
import Base.+

########## Command ##########

@enum Direction forward down up

struct Command{Direction}

    coeff::Int

end

function Command(entry::String)

    direction_string, coeff_string = split(entry)

    direction_string == "forward" ? direction = forward :
    direction_string == "down" ? direction = down :
    direction_string == "up" ? direction = up :
    error("wrong command direction")

    Command{direction}(parse(Int, coeff_string))

end

coeff(c::Command) = c.coeff

########## Coordinates ##########

struct Coordinates

    x::Int
    y::Int

end

Coordinates() = Coordinates(0, 0)

x(c::Coordinates) = c.x
y(c::Coordinates) = c.y

+(a::Coordinates, b::Coordinates) = Coordinates(x(a) + x(b), y(a) + y(b))

########## Submarines_type 1 ##########

mutable struct Submarine_type1

    position::Coordinates

end

Submarine_type1() = Submarine_type1(Coordinates())

position(s::Submarine_type1) = s.position

instruct!(submarine::Submarine_type1, command::Command{forward}) = (submarine.position += Coordinates(coeff(command), 0))
instruct!(submarine::Submarine_type1, command::Command{down}) = (submarine.position += Coordinates(0, coeff(command)))
instruct!(submarine::Submarine_type1, command::Command{up}) = (submarine.position += Coordinates(0, -coeff(command)))

########## Submarines_type 2 ##########

mutable struct Submarine_type2

    position::Coordinates
    aim::Int

end

Submarine_type2() = Submarine_type2(Coordinates(), 0)

position(s::Submarine_type2) = s.position
aim(s::Submarine_type2) = s.aim

instruct!(submarine::Submarine_type2, command::Command{forward}) = (submarine.position += Coordinates(coeff(command), aim(submarine) * coeff(command)))
instruct!(submarine::Submarine_type2, command::Command{down}) = (submarine.aim += coeff(command))
instruct!(submarine::Submarine_type2, command::Command{up}) = (submarine.aim -= coeff(command))

########## Script ##########

function main()

    # open file and store content into a string
    input = open("inputs//day2.txt", "r") do io
        read(io, String)
    end

    # split the string into a vector and parse it as Commands
    data = Command.(String.(split(input, "\r\n")))

    # instantiate submarine
    submarine1 = Submarine_type1()
    submarine2 = Submarine_type2()

    # send all instructions
    for command in data
        instruct!(submarine1, command)
        instruct!(submarine2, command)
    end

    println("part1: ", x(position(submarine1)) * y(position(submarine1)))
    println("part2: ", x(position(submarine2)) * y(position(submarine2)))

end

main()


using MyJuliaGEMM
using LinearAlgebra, Random, Test

Random.seed!(44)

sizes = (48, 288, 576)
for (k, m, n) in Iterators.product(sizes, sizes, sizes)
  A, B, C = rand(m, k), rand(k, n), rand(m, n)
  Cref = C + A * B
  mygemm!(C, A, B)
  @test Cref ≈ C
end


plot(linspace(-4,4,100), erf(linspace(-4,4,100)),"k-")

xlabel(L"x")
ylabel(L"\frac{2}{\sqrt{\pi}}\int^x_{0} dt\, e^{-t^2}")


@testset "Interactions" begin
    c1 = SVector(1.0, 1.0, 1.0)u"nm"
    c2 = SVector(1.3, 1.0, 1.0)u"nm"
    c3 = SVector(1.4, 1.0, 1.0)u"nm"
    a1 = Atom(charge=1.0, σ=0.3u"nm", ϵ=0.2u"kJ * mol^-1")
    box_size = SVector(2.0, 2.0, 2.0)u"nm"
    dr12 = vector(c1, c2, box_size)
    dr13 = vector(c1, c3, box_size)

    for inter in (LennardJones(), Mie(m=6, n=12))
        @test isapprox(
            force(inter, dr12, c1, c2, a1, a1, box_size),
            SVector(16.0, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
            atol=1e-9u"kJ * mol^-1 * nm^-1",
        )
        @test isapprox(
            force(inter, dr13, c1, c3, a1, a1, box_size),
            SVector(-1.375509739, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
            atol=1e-9u"kJ * mol^-1 * nm^-1",
        )
        @test isapprox(
            potential_energy(inter, dr12, c1, c2, a1, a1, box_size),
            0.0u"kJ * mol^-1",
            atol=1e-9u"kJ * mol^-1",
        )
        @test isapprox(
            potential_energy(inter, dr13, c1, c3, a1, a1, box_size),
            -0.1170417309u"kJ * mol^-1",
            atol=1e-9u"kJ * mol^-1",
        )
    end

    inter = SoftSphere()
    @test isapprox(
        force(inter, dr12, c1, c2, a1, a1, box_size),
        SVector(32.0, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        force(inter, dr13, c1, c3, a1, a1, box_size),
        SVector(0.7602324486, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        potential_energy(inter, dr12, c1, c2, a1, a1, box_size),
        0.8u"kJ * mol^-1",
        atol=1e-9u"kJ * mol^-1",
    )
    @test isapprox(
        potential_energy(inter, dr13, c1, c3, a1, a1, box_size),
        0.0253410816u"kJ * mol^-1",
        atol=1e-9u"kJ * mol^-1",
    )

    inter = Coulomb()
    @test isapprox(
        force(inter, dr12, c1, c2, a1, a1, box_size),
        SVector(1543.727311, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-5u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        force(inter, dr13, c1, c3, a1, a1, box_size),
        SVector(868.3466125, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-5u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        potential_energy(inter, dr12, c1, c2, a1, a1, box_size),
        463.1181933u"kJ * mol^-1",
        atol=1e-5u"kJ * mol^-1",
    )
    @test isapprox(
        potential_energy(inter, dr13, c1, c3, a1, a1, box_size),
        347.338645u"kJ * mol^-1",
        atol=1e-5u"kJ * mol^-1",
    )

    c1_grav = SVector(1.0, 1.0, 1.0)u"m"
    c2_grav = SVector(6.0, 1.0, 1.0)u"m"
    a1_grav, a2_grav = Atom(mass=1e6u"kg"), Atom(mass=1e5u"kg")
    box_size_grav = SVector(20.0, 20.0, 20.0)u"m"
    dr12_grav = vector(c1_grav, c2_grav, box_size_grav)
    inter = Gravity()
    @test isapprox(
        force(inter, dr12_grav, c1_grav, c2_grav, a1_grav, a2_grav, box_size_grav),
        SVector(-0.266972, 0.0, 0.0)u"kg * m * s^-2",
        atol=1e-9u"kg * m * s^-2",
    )
    @test isapprox(
        potential_energy(inter, dr12_grav, c1_grav, c2_grav,
                         a1_grav, a2_grav, box_size_grav),
        -1.33486u"kg * m^2 * s^-2",
        atol=1e-9u"kg * m^2 * s^-2",
    )

    b1 = HarmonicBond(b0=0.2u"nm", kb=300_000.0u"kJ * mol^-1 * nm^-2")
    b2 = HarmonicBond(b0=0.6u"nm", kb=100_000.0u"kJ * mol^-1 * nm^-2")
    fs = force(b1, c1, c2, box_size)
    @test isapprox(
        fs.f1,
        SVector(30000.0, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        fs.f2,
        SVector(-30000.0, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    fs = force(b2, c1, c3, box_size)
    @test isapprox(
        fs.f1,
        SVector(-20000.0, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        fs.f2,
        SVector(20000.0, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        potential_energy(b1, c1, c2, box_size),
        1500.0u"kJ * mol^-1",
        atol=1e-9u"kJ * mol^-1",
    )
    @test isapprox(
        potential_energy(b2, c1, c3, box_size),
        2000.0u"kJ * mol^-1",
        atol=1e-9u"kJ * mol^-1",
    )

    b1 = MorseBond(D=100.0u"kJ * mol^-1", α=10.0u"nm^-1", r0=0.2u"nm")
    b2 = MorseBond(D=200.0u"kJ * mol^-1", α=5.0u"nm^-1" , r0=0.6u"nm")
    fs = force(b1, c1, c2, box_size)
    @test isapprox(
        fs.f1,
        SVector(465.0883158697, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        fs.f2,
        SVector(-465.0883158697, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    fs = force(b2, c1, c3, box_size)
    @test isapprox(
        fs.f1,
        SVector(-9341.5485409432, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        fs.f2,
        SVector(9341.5485409432, 0.0, 0.0)u"kJ * mol^-1 * nm^-1",
        atol=1e-9u"kJ * mol^-1 * nm^-1",
    )
    @test isapprox(
        potential_energy(b1, c1, c2, box_size),
        39.9576400894u"kJ * mol^-1",
        atol=1e-9u"kJ * mol^-1",
    )
    @test isapprox(
        potential_energy(b2, c1, c3, box_size),
        590.4984884025u"kJ * mol^-1",
        atol=1e-9u"kJ * mol^-1",
    )
end

@testset "Spatial" begin
    @test vector1D(4.0, 6.0, 10.0) ==  2.0
    @test vector1D(1.0, 9.0, 10.0) == -2.0
    @test vector1D(6.0, 4.0, 10.0) == -2.0
    @test vector1D(9.0, 1.0, 10.0) ==  2.0

    @test vector1D(4.0u"nm", 6.0u"nm", 10.0u"nm") ==  2.0u"nm"
    @test vector1D(1.0u"m" , 9.0u"m" , 10.0u"m" ) == -2.0u"m"
    @test_throws Unitful.DimensionError vector1D(6.0u"nm", 4.0u"nm", 10.0)

    @test vector(SVector(4.0, 1.0, 6.0), SVector(6.0, 9.0, 4.0),
                    SVector(10.0, 10.0, 10.0)) == SVector(2.0, -2.0, -2.0)
    @test vector(SVector(4.0, 1.0, 1.0), SVector(6.0, 4.0, 3.0),
                    SVector(10.0, 5.0, 3.5)) == SVector(2.0, -2.0, -1.5)
    @test vector(SVector(4.0, 1.0), SVector(6.0, 9.0),
                    SVector(10.0, 10.0)) == SVector(2.0, -2.0)
    @test vector(SVector(4.0, 1.0, 6.0)u"nm", SVector(6.0, 9.0, 4.0)u"nm",
                    SVector(10.0, 10.0, 10.0)u"nm") == SVector(2.0, -2.0, -2.0)u"nm"

    @test wrap_coords(8.0 , 10.0) == 8.0
    @test wrap_coords(12.0, 10.0) == 2.0
    @test wrap_coords(-2.0, 10.0) == 8.0

    @test wrap_coords(8.0u"nm" , 10.0u"nm") == 8.0u"nm"
    @test wrap_coords(12.0u"m" , 10.0u"m" ) == 2.0u"m"
    @test_throws ErrorException wrap_coords(-2.0u"nm", 10.0)

    vels_units   = [maxwell_boltzmann(12.0u"u", 300.0u"K") for _ in 1:1_000]
    vels_nounits = [maxwell_boltzmann(12.0    , 300.0    ) for _ in 1:1_000]
    @test 0.35u"nm * ps^-1" < std(vels_units) < 0.55u"nm * ps^-1"
    @test 0.35 < std(vels_nounits) < 0.55
end

@testset "Neighbor lists" begin
    for neighbor_finder in (DistanceNeighborFinder, TreeNeighborFinder, CellListMapNeighborFinder)
        nf = neighbor_finder(nb_matrix=trues(3, 3), n_steps=10, dist_cutoff=2.0u"nm")
        s = System(
            atoms=[Atom(), Atom(), Atom()],
            coords=[SVector(1.0, 1.0, 1.0)u"nm", SVector(2.0, 2.0, 2.0)u"nm",
                    SVector(5.0, 5.0, 5.0)u"nm"],
            box_size=SVector(10.0, 10.0, 10.0)u"nm",
            neighbor_finder=nf,
        )
        neighbors = find_neighbors(s, s.neighbor_finder; parallel=false)
        @test neighbors.list == [(2, 1, false)] || neighbors.list == [(1, 2, false)]
        if run_parallel_tests
            neighbors = find_neighbors(s, s.neighbor_finder; parallel=true)
            @test neighbors.list == [(2, 1, false)] || neighbors.list == [(1, 2, false)]
        end
        show(devnull, nf)
    end

    # Test passing the box_size and coordinates as keyword arguments to CellListMapNeighborFinder
    coords = [SVector(1.0, 1.0, 1.0)u"nm", SVector(2.0, 2.0, 2.0)u"nm", SVector(5.0, 5.0, 5.0)u"nm"]
    box_size = SVector(10.0, 10.0, 10.0)u"nm"
    neighbor_finder=CellListMapNeighborFinder(
        nb_matrix=trues(3, 3), n_steps=10, x0=coords,
        unit_cell=box_size, dist_cutoff=2.0u"nm",
    )
    s = System(
        atoms=[Atom(), Atom(), Atom()],
        coords=coords, box_size=box_size,
        neighbor_finder=neighbor_finder,
    )
    neighbors = find_neighbors(s, s.neighbor_finder; parallel=false)
    @test neighbors.list == [(2, 1, false)] || neighbors.list == [(1, 2, false)]
    if run_parallel_tests
        neighbors = find_neighbors(s, s.neighbor_finder; parallel=true)
        @test neighbors.list == [(2, 1, false)] || neighbors.list == [(1, 2, false)]
    end
end

@testset "Analysis" begin
    pdb_path = joinpath(data_dir, "1ssu.pdb")
    struc = read(pdb_path, BioStructures.PDB)
    cm_1 = BioStructures.coordarray(struc[1], BioStructures.calphaselector)
    cm_2 = BioStructures.coordarray(struc[2], BioStructures.calphaselector)
    coords_1 = SVector{3, Float64}.(eachcol(cm_1)) / 10 * u"nm"
    coords_2 = SVector{3, Float64}.(eachcol(cm_2)) / 10 * u"nm"
    @test rmsd(coords_1, coords_2) ≈ 2.54859467758795u"Å"
    if run_gpu_tests
        @test rmsd(cu(coords_1), cu(coords_2)) ≈ 2.54859467758795u"Å"
    end

    bb_atoms = BioStructures.collectatoms(struc[1], BioStructures.backboneselector)
    coords = SVector{3, Float64}.(eachcol(BioStructures.coordarray(bb_atoms))) / 10 * u"nm"
    bb_to_mass = Dict("C" => 12.011u"u", "N" => 14.007u"u", "O" => 15.999u"u")
    atoms = [Atom(mass=bb_to_mass[BioStructures.element(bb_atoms[i])]) for i in 1:length(bb_atoms)]
    @test isapprox(radius_gyration(coords, atoms), 11.51225678195222u"Å", atol=1e-6u"nm")
end


using RecordedArrays
using RecordedArrays.ResizingTools
using Random
using LoopVectorization
using EvolutionaryModelingTools
using EvolutionaryModelingTools: sample

const T = 10.0
const r = 0.5
const K = 100
const μ = 0.01
const v = [10;]
const m = ones(1, 1)

# simulate with this package
## define reactions

# define growth with @reaction_eq, loop vectorization is disabled but inbounds, fastmath and  offset are enable
@reaction_eq growth r * (1 - μ) v[i] → 2v[i] avx=false inbounds=true fastmath=true offset=true

# define mutation with @reaction and @quickloop with given index
@reaction mutation begin
    @quickloop r * μ * v[i]
    begin
        push!(v, 1)
        n = length(v)
        resize!(m, (n, n))
        # rand competition matrix
        m[:, n] = rand(rng, n)
        m[n, 1:n-1] = rand(rng, 1, n - 1)
    end
end

# define competition with @reaction and @quickloop
@reaction competition begin
    @quickloop (r / K) * v[i] * m[i, j] * v[j] turbo # force enable loop vectorization
    begin
        i = ind[1]
        v[i] -= 1
        if v[i] == 0 # check extinction
            deleteat!(v, i)
            resize!(m, (not(i), not(i)))
        end
    end
end

# run functions for test and benchmark
function run_gillespie(rng=Random.GLOBAL_RNG, T=T, r=r, K=K, μ=μ, v=v, m=m)
    ps = (; r, K, μ, v, m = SimpleRDArray(m))
    return gillespie(rng, T, ps, (growth, mutation, competition))
end

function run_gillespie_record(rng=Random.GLOBAL_RNG, T=T, r=r, K=K, μ=μ, v=v, m=m)
    clock = ContinuousClock(T)
    ps = (;
        r,
        K,
        μ,
        v=recorded(DynamicEntry, clock, v),
        m=recorded(StaticEntry, clock, m),
    )
    return gillespie(rng, clock, ps, (growth, mutation, competition))
end

# simulate manually
function run_manually(rng=Random.GLOBAL_RNG, T=T, r=r, K=K, μ=μ, v=v, m=m)
    t = 0.0
    v′ = deepcopy(v)
    m′ = SimpleRDArray(m)
    while t <= T
        growth_rate = r * (1 - μ) * v′
        mutation_rate = r * μ * v′
        c = r / K # baseline competition coefficient
        competition_rate = @. c * v′ * m′ * v′'
        rates_sum = (sum(growth_rate), sum(mutation_rate), sum(competition_rate))
        rates_sum_acc = accumulate(+, rates_sum)
        rates_sum_sum = rates_sum_acc[end]
        t += -log(rand(rng)) / rates_sum_sum
        rn = rand(rng) * rates_sum_sum
        ind = searchsortedfirst(collect(rates_sum_acc), rn)
        if ind == 1
            i = sample(growth_rate, rn)
            v′[i] += 1
        elseif ind == 2
            i = sample(mutation_rate, rn - rates_sum_acc[1])
            push!(v′, 1)
            n = length(v′)
            resize!(m′, (n, n))
            # rand competition matrix
            m′[:, n] = rand(rng, n)
            m′[n, 1:n-1] = rand(rng, 1, n - 1)
        elseif ind == 3
            i = sample(competition_rate, rn - rates_sum_acc[2])[1]
            v′[i] -= 1
            if v′[i] == 0 # check extinction
                deleteat!(v′, i)
                resize!(m′, (not(i), not(i)))
            end
        end
        length(v′) == 0 && break
    end
    return (; v=v′, m=m′)
end
function run_manually_record(rng=Random.GLOBAL_RNG, T=T, r=r, K=K, μ=μ, v=v, m=m)
    clock = ContinuousClock(T)
    v′ = recorded(DynamicEntry, clock, v)
    m′ = recorded(StaticEntry, clock, m)
    for _ in clock
        growth_rate = r * (1 - μ) * v′
        mutation_rate = r * μ * v′
        c = r / K # baseline competition coefficient
        competition_rate = @. c * v′ * m′ * v′'
        rates_sum = (sum(growth_rate), sum(mutation_rate), sum(competition_rate))
        rates_sum_acc = accumulate(+, rates_sum)
        rates_sum_sum = rates_sum_acc[end]
        τ = -log(rand(rng)) / rates_sum_sum
        increase!(clock, τ)
        rn = rand(rng) * rates_sum_sum
        ind = searchsortedfirst(collect(rates_sum_acc), rn)
        if ind == 1
            i = sample(growth_rate, rn)
            v′[i] += 1
        elseif ind == 2
            i = sample(mutation_rate, rn - rates_sum_acc[1])
            push!(v′, 1)
            n = length(v′)
            resize!(m′, (n, n))
            # rand competition matrix
            m′[:, n] = rand(rng, n)
            m′[n, 1:n-1] = rand(rng, 1, n - 1)
        elseif ind == 3
            i = sample(competition_rate, rn - rates_sum_acc[2])[1]
            v′[i] -= 1
            if v′[i] == 0 # check extinction
                deleteat!(v′, i)
                resize!(m′, (not(i), not(i)))
            end
        end
        length(v′) == 0 && break
    end
    return (; v=v′, m=m′)
end


module Planck

using Unitful
using Unitful: AbstractQuantity
using Unitful: h, c0, k

export blackbody

# constants in SI units
const _h = ustrip(u"J*s", h)
const _c0 = ustrip(u"m/s", c0)
const _k = ustrip(u"J/K", k)

using Unitful: 𝐋, 𝐓

"""
    blackbody([OT], x, T)

Evaluate a blackbody[^1] at `x`, which is in meters by default. If `OT` is given, the output will be converted to that type, which is convenient for unit conversions. Temperature is assumed to be Kelvin.

!!! warning "Units"

    If you do not use [Unitful.jl](https://github.com/PainterQubits/Unitful.jl), the defaults will be [SI](https://en.wikipedia.org/wiki/International_System_of_Units)

# Examples

```jldoctest
julia> using Unitful

julia> blackbody(6e-7, 5850)
2.583616647617973e13

julia> blackbody(Float32, 6e-7, 5850)
2.5836166f13

julia> blackbody(u"erg/s/cm^2/nm/sr", 600u"nm", 5850u"K")
2.5836166476179734e7 erg nm⁻¹ cm⁻² s⁻¹ sr⁻¹
```

# References
[^1]: [Blackbody radiation](https://en.wikipedia.org/wiki/Planck%27s_law)
"""
blackbody(OT, x, T) = OT(blackbody(x, T))

function blackbody(ν::AbstractQuantity{V,inv(𝐓)}, T) where V
    2 * h * ν^3 / c0^2 / expm1(h * ν / (k * T))
end

function blackbody(λ::AbstractQuantity{V,𝐋}, T) where V
    2 * h * c0^2 / λ^5 / expm1(h * c0 / (λ * k * T))
end

# SI (meters, Kelvin)
function blackbody(λ, T)
    2 * _h * _c0^2 / λ^5 / expm1(_h * _c0 / (λ * _k * T)) # W / m^3
end

end


# Ambegaokar, V., & Troyer, M. (2010). Estimating errors reliably in Monte
# Carlo simulations of the Ehrenfest model. American Journal of Physics, 78(2),
# 150-157.

function stderr_bin(xs::AbstractVector)
    stderr_full = std(xs) / sqrt(length(xs))
    stderr_max = stderr_full

    while length(xs) >= 60
        xs = [0.5 * (xs[i] + xs[i+1]) for i in (1+length(xs)%2):2:length(xs)]
        stderr_max = max(stderr_max, std(xs) / sqrt(length(xs)))
    end

    autocorrelation_steps = 0.5 * ((stderr_max / stderr_full)^2 - 1)

    stderr_max, autocorrelation_steps
end

function aggregate(xs::AbstractVector)
    mean(xs), stderr_bin(xs)...
end


using QuantLib

function main()
  euribor3m = euribor_index(QuantLib.Time.TenorPeriod(Dates.Month(3)))

  todaysDate = Date(2006, 5, 23)
  set_eval_date!(settings, todaysDate)
  calendar = euribor3m.fixingCalendar
  fixingDays = euribor3m.fixingDays
  settlementDate = QuantLib.Time.advance(Dates.Day(fixingDays), calendar, todaysDate)

  threeMonthFraQuote = zeros(10)
  threeMonthFraQuote[2] = 0.030
  threeMonthFraQuote[3] = 0.031
  threeMonthFraQuote[4] = 0.032
  threeMonthFraQuote[7] = 0.033
  threeMonthFraQuote[10] = 0.034

  ## QUOTES ##
  h1x4 = Quote(threeMonthFraQuote[2])
  h2x5 = Quote(threeMonthFraQuote[3])
  h3x6 = Quote(threeMonthFraQuote[4])
  h6x9 = Quote(threeMonthFraQuote[7])
  h9x12 = Quote(threeMonthFraQuote[10])

  ## RATE HELPERS ##
  fraDayCount = euribor3m.dc
  convention = euribor3m.convention
  endOfMonth = euribor3m.endOfMonth

  fra1x4 = FraRateHelper(h1x4, 1, 4, fixingDays, calendar, convention, endOfMonth, fraDayCount)
  fra2x5 = FraRateHelper(h2x5, 2, 5, fixingDays, calendar, convention, endOfMonth, fraDayCount)
  fra3x6 = FraRateHelper(h3x6, 3, 6, fixingDays, calendar, convention, endOfMonth, fraDayCount)
  fra6x9 = FraRateHelper(h6x9, 6, 9, fixingDays, calendar, convention, endOfMonth, fraDayCount)
  fra9x12 = FraRateHelper(h9x12, 9, 12, fixingDays, calendar, convention, endOfMonth, fraDayCount)

  fraInstruments = FraRateHelper[fra1x4, fra2x5, fra3x6, fra6x9, fra9x12]
  termStructureDC = QuantLib.Time.ISDAActualActual()
  tol = 1.0e-15

  fraTermStructure = PiecewiseYieldCurve(settlementDate, fraInstruments, termStructureDC, QuantLib.Math.LogLinear(), Discount(), tol, IterativeBootstrap())

  euribor3m = update_termstructure(euribor3m, fraTermStructure)

  fraCalendar = euribor3m.fixingCalendar
  fraBusinessDayConvention = euribor3m.convention
  fraFwdType = LongPosition()
  fraNotional = 100.0
  monthsToStart = [1, 2, 3, 6, 9]
  fraTermMonths = 3
  fraValueDate = QuantLib.Time.advance(Dates.Month(monthsToStart[1]), fraCalendar, settlementDate, fraBusinessDayConvention)
  fraMaturityDate = QuantLib.Time.advance(Dates.Month(fraTermMonths), fraCalendar, fraValueDate, fraBusinessDayConvention)
  fraStrikeRate = threeMonthFraQuote[monthsToStart[1] + 1]
  myFRA = ForwardRateAgreement(fraValueDate, fraMaturityDate, fraFwdType, fraStrikeRate, fraNotional, euribor3m, fraTermStructure)

  println(forward_rate(myFRA).rate)
  println(spot_value(myFRA))
  println(forward_value(myFRA))
  println(implied_yield(myFRA, spot_value(myFRA), forward_value(myFRA), settlementDate, SimpleCompounding(), fraDayCount).rate)
end


immutable TrivialModel{T}
    α::T
    β::T
end

@testset "Collocation" begin
    fam = Chebyshev(10)
    θ = fit(fam, x->x^2)
    model = TrivialModel(2.0, 3.0)
    residual_function(model::TrivialModel, f, x) = f(x-model.α)-model.β
    cres = CollocationResidual(model, fam, residual_function)
    r = cres(θ)
    @test r ≈ (points(fam) - model.α).^2-model.β
end

immutable TrivialModel2{T}
    α::T
end

@testset "Collocation solver" begin
    (m::TrivialModel2)(x) = m.α*x
    model0 = TrivialModel2(2.0)
    residualf(model::TrivialModel2, f, x) = f(x)-model(x)
    res = CollocationResidual(model0, Chebyshev(10), residualf)
    f_sol, o = solve_collocation(res, zero)
    @test norm([f_sol(x)-model0(x) for x in linspace(domain(f_sol), 100)]) ≤ 1e-14
end


# Tests the vectorized version of cel:
#include("cel_bulirsch.jl")
# include("../src/cel_bulirsch.jl")
#if VERSION >= v"0.7"
#  using Test
#else
#  using Base.Test
#end

import Limbdark: cel_bulirsch

# Randomizer seed
using Random
Random.seed!(42)

# Function which tests cel_bulirsch both numerically, one-by-one and vectorized:
function test_cel!(kc,p,a,b)
k2 = 1.0-kc^2
# Evaluate all three at once:
ell2 = cel_bulirsch(k2,kc,p,a[1],a[2],a[3],b[1],b[2],b[3])
ell1 = zeros(3)
nphi = 100000
if p < 0
#  Make same transformations as in cel_bulirsch for negative p:
  p0 = sqrt((kc*kc-p)/(1-p))
  a0 = (a[1]-b[1])/(1-p)
  b0 = -(b[1]-a[1]*p)/(1-p)^2*(1-kc^2)/p0+a0*p0
  a[1] = a0
  b[1] = b0*p0
  p = p0*p0
end
# Numerically compute the integrals:
dphi = .5*pi/nphi
phi = linearspace(0.5*dphi,pi/2-.5*dphi,nphi)
ell3 = zeros(3); cphi2 = cos.(phi).^2; sphi2=sin.(phi).^2
den = dphi./sqrt.(cphi2+kc*kc*sphi2)
ell3[1] = sum((a[1]*cphi2+b[1]*sphi2)./(cphi2+p*sphi2).*den)
ell3[2] = sum((a[2]*cphi2+b[2]*sphi2).*den)
ell3[3] = sum((a[3]*cphi2+b[3]*sphi2).*den)

# First, evaluate one by one:
ell1[1] = cel_bulirsch(k2,kc,p,a[1],b[1])
ell1[2] = cel_bulirsch(k2,kc,1.0,a[2],b[2])
ell1[3] = cel_bulirsch(k2,kc,1.0,a[3],b[3])

return ell1,ell2,ell3
end

@testset "cel_bulirsch" begin
  # Test negative p values:
  p = -rand(); a=rand(3); b=rand(3)
  kc = rand()
  ell1,ell2,ell3 = test_cel!(kc,p,a,b)
  println("ell1: ",ell1)
  println("ell2: ",ell2)
  println("ell3: ",ell3)
  for i=1:3
    @test isapprox(ell1[i],ell2[i])
#    @test isapprox(ell1[i],ell3[i],atol=1e-6)
    @test isapprox(ell1[i],ell3[i])
  end
  # Test positive p values:
  p = rand(); a=rand(3); b=rand(3)
  kc = rand()
  ell1,ell2,ell3 = test_cel!(kc,p,a,b)
  println("ell1: ",ell1)
  println("ell2: ",ell2)
  println("ell3: ",ell3)
  for i=1:3
    @test isapprox(ell1[i],ell2[i])
#    @test isapprox(ell1[i],ell3[i],atol=1e-6)
    @test isapprox(ell1[i],ell3[i])
  end
  # Test small kc values:
  p = rand(); a=rand(3); b=rand(3)
  kc = 1e-4
  ell1,ell2,ell3 = test_cel!(kc,p,a,b)
  println("ell1: ",ell1)
  println("ell2: ",ell2)
  println("ell3: ",ell3)
  for i=1:3
    @test isapprox(ell1[i],ell2[i])
#    @test isapprox(ell1[i],ell3[i],atol=1e-6)
    @test isapprox(ell1[i],ell3[i])
  end
  # Test large kc values:
  p = rand(); a=rand(3); b=rand(3)
  kc = 1.0-1e-4
  ell1,ell2,ell3 = test_cel!(kc,p,a,b)
  println("ell1: ",ell1)
  println("ell2: ",ell2)
  println("ell3: ",ell3)
  for i=1:3
    @test isapprox(ell1[i],ell2[i])
#    @test isapprox(ell1[i],ell3[i],atol=1e-6)
    @test isapprox(ell1[i],ell3[i])
  end
end


module BernsteinExpansions

using Reexport,
      RecursiveArrayTools,
      MultivariatePolynomials

@reexport using IntervalArithmetic
const IntervalOrIntervalBox = Union{Interval, IntervalBox}

# Bernstein coefficients of univariate and multivariate monomials
include("fastpow.jl")
include("monomials.jl")

export univariate,
       multivariate

# Bernstein expansion of polynomials on hyperrectangular domains
include("forms.jl")
include("implicit.jl")
include("full.jl")

export ImplicitBernsteinForm,
       ExplicitBernsteinForm,
       FullBernsteinForm

# Range enclosure methods
include("enclosures.jl")

end # module


module particle

import Base.==
import ..operators

using ..bases, ..states, ..operators, ..operators_sparse, ..operators_lazy

export PositionBasis, MomentumBasis,
        gaussianstate,
        spacing, samplepoints,
        positionoperator, momentumoperator, laplace_x, laplace_p, FFTOperator

"""
Basis for a particle in real space.

For simplicity periodic boundaries are assumed which means that
the rightmost point defined by xmax is not included in the basis
but is defined to be the same as xmin.

Arguments
---------

xmin
    Left boundary.
xmax
    Right boundary.
N
    Number of discrete points in the basis.
"""
type PositionBasis <: Basis
    shape::Vector{Int}
    xmin::Float64
    xmax::Float64
    N::Int
    PositionBasis(xmin::Real, xmax::Real, N::Int) = new([N], xmin, xmax, N)
end

"""
Basis for a particle in momentum space.

For simplicity periodic boundaries are assumed which means that
pmax is not included in the basis but is defined to be the same as pmin.

Arguments
---------

pmin
    Left boundary.
pmax
    Right boundary.
N
    Number of discrete points in the basis.
"""
type MomentumBasis <: Basis
    shape::Vector{Int}
    pmin::Float64
    pmax::Float64
    N::Int
    MomentumBasis(pmin::Real, pmax::Real, N::Int) = new([N], pmin, pmax, N)
end

"""
Create a PositionBasis from a MomentumBasis.

Since periodic boundary conditions are assumed the values for :math:`x_{min}`
and :math:`x_{max}` are more or less arbitrary and are chosen to be
:math:`-\\pi/dp` and :math:`\\pi/dp` with :math:`dp=(p_{max}-p_{min})/N`.
"""
PositionBasis(b::MomentumBasis) = (dp = (b.pmax - b.pmin)/b.N; PositionBasis(-pi/dp, pi/dp, b.N))

"""
Create a MomentumBasis from a PositionBasis.

Since periodic boundary conditions are assumed the values for :math:`p_{min}`
and :math:`p_{max}` are more or less arbitrary and are chosen to be
:math:`-\\pi/dx` and :math:`\\pi/dx` with :math:`dx=(x_{max}-x_{min})/N`.
"""
MomentumBasis(b::PositionBasis) = (dx = (b.xmax - b.xmin)/b.N; MomentumBasis(-pi/dx, pi/dx, b.N))

==(b1::PositionBasis, b2::PositionBasis) = b1.xmin==b2.xmin && b1.xmax==b2.xmax && b1.N==b2.N
==(b1::MomentumBasis, b2::MomentumBasis) = b1.pmin==b2.pmin && b1.pmax==b2.pmax && b1.N==b2.N

"""
Create a Gaussian state around x0 and p0 with width sigma in real space.

This implementation is based on the definition

.. math::

    \\Psi(x) = \\frac{\\sqrt{\\Delta x}}{\\pi^{1/4}\\sqrt{\\sigma}}
                e^{i p_0 (x-x_0) - \\frac{(x-x_0)^2}{2 \\sigma^2}}

which has the properties

* :math:`|\\Psi(p)|^2 = 1`
* :math:`\\langle p \\rangle = p_0`
* :math:`\\langle x \\rangle = x_0`
* :math:`\\mathrm{Var}(x) = \\frac{\\sigma^2}{2}`
* :math:`\\mathrm{Var}(p) = \\frac{1}{2 \\sigma^2}`

and is connected to the MomentumBasis based implementation via a
Fourier-transformation

.. math::

    \\Psi(p) = \\frac{1}{\\sqrt{2\\pi}}
                \\int_{-\\infty}^{\\infty} e^{-ipx}\\Psi(x) \\mathrm{d}x

Due to the numerically necessary discretization an additional scaling
factor :math:`\\sqrt{\\Delta x}` is used so that
:math:`\\Psi_i = \\sqrt{\\Delta x} \\Psi(x_i)` and the resulting state
is normalized to

.. math::

    \\sum_i |\\Psi_i|^2 = 1

Arguments
---------
b
    PositionBasis.
x0
    Center of the Gaussian in position space.
p0
    Center of the Gaussian in momentum space.
sigma
    Width of the Gaussian.
"""
function gaussianstate(b::PositionBasis, x0::Real, p0::Real, sigma::Real)
    psi = Ket(b)
    dx = spacing(b)
    alpha = 1./(pi^(1/4)*sqrt(sigma))*sqrt(dx)
    x = b.xmin
    for i=1:b.N
        psi.data[i] = alpha*exp(1im*p0*(x-x0) - (x-x0)^2/(2*sigma^2))
        x += dx
    end
    return psi
end

"""
Create a Gaussian state around x0 and p0 with width sigma in momentum space.

This implementation is based on the definition

.. math::

    \\Psi(p) = \\frac{\\sqrt{\\sigma} \\sqrt{\\Delta x}}{\\pi^{1/4}}
                e^{-i x_0 p - \\frac{1}{2}(p-p_0)^2 \\sigma^2}

which has the properties

* :math:`|\\Psi(p)|^2 = 1`
* :math:`\\langle p \\rangle = p_0`
* :math:`\\langle x \\rangle = x_0`
* :math:`\\mathrm{Var}(x) = \\frac{\\sigma^2}{2}`
* :math:`\\mathrm{Var}(p) = \\frac{1}{2 \\sigma^2}`

and is connected to the MomentumBasis based implementation via a
Fourier-transformation

.. math::

    \\Psi(x) = \\frac{1}{\\sqrt{2\\pi}}
                    \\int_{-\\infty}^{\\infty} e^{ipx}\\Psi(p) \\mathrm{d}p

Due to the numerically necessary discretization an additional scaling
factor :math:`\\sqrt{\\Delta x}` is used so that
:math:`\\Psi_i = \\sqrt{\\Delta p} \\Psi(p_i)` and the resulting state
is normalized to

.. math::

    \\sum_i |\\Psi_i|^2 = 1

Arguments
---------
b
    MomentumBasis.
x0
    Center of the Gaussian in position space.
p0
    Center of the Gaussian in momentum space.
sigma
    Width of the Gaussian.
"""
function gaussianstate(b::MomentumBasis, x0::Real, p0::Real, sigma::Real)
    psi = Ket(b)
    dp = spacing(b)
    alpha = sqrt(sigma)/pi^(1/4)*sqrt(dp)
    p = b.pmin
    for i=1:b.N
        psi.data[i] = alpha*exp(-1im*x0*p - (p-p0)^2/2*sigma^2)
        p += dp
    end
    return psi
end

"""
Distance between two adjacent points of the real space basis.
"""
spacing(b::MomentumBasis) = (b.pmax - b.pmin)/b.N

"""
Momentum difference between two adjacent points of the momentum basis.
"""
spacing(b::PositionBasis) = (b.xmax - b.xmin)/b.N

"""
x values of the real space basis.
"""
samplepoints(b::PositionBasis) = (dx = spacing(b); Float64[b.xmin + i*dx for i=0:b.N-1])

"""
p values of the momentum basis.
"""
samplepoints(b::MomentumBasis) = (dp = spacing(b); Float64[b.pmin + i*dp for i=0:b.N-1])

"""
Position operator in real space.
"""
positionoperator(b::PositionBasis) = SparseOperator(b, spdiagm(complex(samplepoints(b)), 0, length(b), length(b)))

"""
Position operator in momentum space.
"""
function positionoperator(b::MomentumBasis)
    p_op = DenseOperator(b)
    u = 1im/(12*spacing(b))
    for i=1:b.N-2
        p_op.data[i,i+2] = -u
        p_op.data[i,i+1] = 8*u
        p_op.data[i+1,i] = -8*u
        p_op.data[i+2,i] = u
    end
    return SparseOperator(p_op)
end

"""
Momentum operator in momentum space.
"""
momentumoperator(b::MomentumBasis) = SparseOperator(b, spdiagm(complex(samplepoints(b)), 0, length(b), length(b)))

"""
Momentum operator in real space.
"""
function momentumoperator(b::PositionBasis)
    p_op = DenseOperator(b)
    u = -1im/(12*spacing(b))
    for i=1:b.N-2
        p_op.data[i,i+2] = -u
        p_op.data[i,i+1] = 8*u
        p_op.data[i+1,i] = -8*u
        p_op.data[i+2,i] = u
    end
    return SparseOperator(p_op)
end

"""
Second x-derivative in real space.
"""
function laplace_x(b::PositionBasis)
    x_op = DenseOperator(b)
    u = 1/spacing(b)^2
    for i=1:b.N-1
        x_op.data[i+1,i] = u
        x_op.data[i,i] = -2*u
        x_op.data[i,i+1] = u
    end
    x_op.data[b.N,b.N] = -2*u
    return SparseOperator(x_op)
end

"""
Second x-derivative in momentum space.
"""
laplace_x(b::MomentumBasis) = SparseOperator(b, spdiagm(complex(samplepoints(b)).^2, 0, length(b), length(b)))

"""
Second p-derivative in momentum space.
"""
function laplace_p(b::MomentumBasis)
    p_op = DenseOperator(b)
    u = 1/spacing(b)^2
    for i=1:b.N-1
        p_op.data[i+1,i] = u
        p_op.data[i,i] = -2*u
        p_op.data[i,i+1] = u
    end
    p_op.data[b.N,b.N] = -2*u
    return SparseOperator(p_op)
end

"""
Second p-derivative in real space.
"""
laplace_p(b::PositionBasis) = SparseOperator(b, spdiagm(complex(samplepoints(b)).^2, 0, length(b), length(b)))

type FFTOperator <: LazyOperator
    basis_l::Basis
    basis_r::Basis
    fft_l!
    fft_r!
    mul_before::Vector{Complex128}
    mul_after::Vector{Complex128}
end

"""
DenseOperator performing a fast Fourier transformation of a state.
"""
function FFTOperator(basis_l::MomentumBasis, basis_r::PositionBasis)
    Lx = (basis_r.xmax - basis_r.xmin)
    dp = spacing(basis_l)
    dx = spacing(basis_r)
    if basis_l.N != basis_r.N || abs(2*pi/dp - Lx)/Lx > 1e-12
        throw(IncompatibleBases())
    end
    x = zeros(Complex128, length(basis_r))
    mul_before = exp(-1im*basis_l.pmin*(samplepoints(basis_r)-basis_r.xmin))
    mul_after = exp(-1im*basis_r.xmin*samplepoints(basis_l))/sqrt(basis_r.N)
    FFTOperator(basis_l, basis_r, plan_bfft!(x), plan_fft!(x), mul_before, mul_after)
end

function FFTOperator(basis_l::PositionBasis, basis_r::MomentumBasis)
    Lx = (basis_l.xmax - basis_l.xmin)
    dp = spacing(basis_r)
    dx = spacing(basis_l)
    if basis_l.N != basis_r.N || abs(2*pi/dp - Lx)/Lx > 1e-12
        throw(IncompatibleBases())
    end
    x = zeros(Complex128, length(basis_r))
    mul_before = exp(1im*basis_l.xmin*(samplepoints(basis_r)-basis_r.pmin))
    mul_after = exp(1im*basis_r.pmin*samplepoints(basis_l))/sqrt(basis_r.N)
    FFTOperator(basis_l, basis_r, plan_fft!(x), plan_bfft!(x), mul_before, mul_after)
end

operators.dagger(op::FFTOperator) = FFTOperator(op.basis_r, op.basis_l)

function operators.gemv!{T<:Complex}(alpha::T, M::FFTOperator, b::Ket, beta::T, result::Ket)
    N::Int = M.basis_r.N
    if beta==Complex(0.)
        @inbounds for i=1:N
            result.data[i] = M.mul_before[i] * b.data[i]
        end
        M.fft_r! * result.data
        @inbounds for i=1:N
            result.data[i] *= M.mul_after[i] * alpha
        end
    else
        psi_ = Ket(M.basis_l, deepcopy(b.data))
        @inbounds for i=1:N
            psi_.data[i] *= M.mul_before[i]
        end
        M.fft_r! * psi_.data
        @inbounds for i=1:N
            result.data[i] = beta*result.data[i] + alpha * psi_.data[i] * M.mul_after[i]
        end
    end
    nothing
end

function operators.gemv!{T<:Complex}(alpha::T, b::Bra, M::FFTOperator, beta::T, result::Bra)
    N::Int = M.basis_l.N
    if beta==Complex(0.)
        @inbounds for i=1:N
            result.data[i] = conj(M.mul_after[i]) * conj(b.data[i])
        end
        M.fft_l! * result.data
        @inbounds for i=1:N
            result.data[i] = conj(result.data[i]) * M.mul_before[i] * alpha
        end
    else
        psi_ = Bra(M.basis_r, conj(b.data))
        @inbounds for i=1:N
            psi_.data[i] *= conj(M.mul_after[i])
        end
        M.fft_l! * psi_.data
        @inbounds for i=1:N
            result.data[i] = beta*result.data[i] + alpha * conj(psi_.data[i]) * M.mul_before[i]
        end
    end
    nothing
end

end # module


using ACEfit
using LinearAlgebra

@info("Test Solver on overdetermined system")
Nobs = 10_000
Nfeat = 100
A = randn(Nobs, Nfeat) / sqrt(Nobs)
y = randn(Nobs)
P = Diagonal(1.0 .+ rand(Nfeat))

@info(" ... QR")
solver = ACEfit.QR()
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... regularised QR, λ = 1.0")
solver = ACEfit.QR(λ = 1e0, P = P )
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... regularised QR, λ = 10.0")
solver = ACEfit.QR(λ = 1e1, P = P)
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... RRQR, rtol = 1e-15")
solver = ACEfit.RRQR(rtol = 1e-15, P = P)
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... RRQR, rtol = 0.5")
solver = ACEfit.RRQR(rtol = 0.5, P = P)
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... RRQR, rtol = 0.99")
solver = ACEfit.RRQR(rtol = 0.99, P = P)
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... LSQR")
solver = ACEfit.LSQR(damp=0, atol=1e-6)
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... SKLEARN_BRR")
solver = ACEfit.SKLEARN_BRR()
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... SKLEARN_ARD")
solver = ACEfit.SKLEARN_ARD()
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... Bayesian Linear")
solver = ACEfit.BL()
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)

@info(" ... Bayesian ARD")
solver = ACEfit.BARD()
θ = ACEfit.solve_llsq(solver, A, y)
@show norm(A * θ - y)
@show norm(θ)


module SparseFEM

# assemble the finite-difference laplacian
function fdlaplacian(N)
    # create a 1D laplacian and a sparse identity
    fdl1 = spdiagm((ones(N-1),-2*ones(N),ones(N-1)), [-1,0,1])
    # laplace operator on the full grid
    return kron(speye(N), fdl1) + kron(fdl1, speye(N))
end

# get the list of boundary dof-indices
function get_free(N)
    L = zeros(Int, N, N)
    L[2:N-1, 2:N-1] = 1
    return find(L)
end

# timing of assembly, slice and solve
function perf_sparse_fem(N)
    Ifree = get_free(N)
    # assembly
    A = fdlaplacian(N)
    # boundary condition
    B = A[Ifree, Ifree]
    # solver
    return lufact(B)
end

end # module


using GeneralizedLeastSquares: coef, LinRegQR, LinRegCholesky, LinRegBunchKaufman, LinRegSweep
using LinearAlgebra
using Test

n, p = 10^6, 50

x = randn(n, p)

y = x * (1:50) + randn(n)

w = Diagonal(rand(n))

β̂ = x \ y

β̂2 = (x' * w * x) \ (x' * w * y)

@testset "OLS" begin
    @testset "LinRegQR" begin
        model = LinRegQR(x, y)
        @test coef(model) ≈ β̂
    end
    @testset "LinRegCholesky" begin
        model = LinRegCholesky(x, y)
        @test coef(model) ≈ β̂
    end
    @testset "LinRegBunchKaufman" begin
        model = LinRegBunchKaufman(x, y)
        @test coef(model) ≈ β̂
    end
    @testset "LinRegSweep" begin
        model = LinRegSweep(x, y)
        @test coef(model) ≈ β̂
    end
end
@testset "GLS" begin
    @testset "LinRegQR" begin
        wsqrt = sqrt(w)
        x2 = wsqrt * x
        y2 = wsqrt * y
        model = LinRegQR(x2, y2)
        @test coef(model) ≈ β̂2
    end
    @testset "LinRegCholesky" begin
        model = LinRegCholesky(x, y, w)
        @test coef(model) ≈ β̂2
    end
    @testset "LinRegBunchKaufman" begin
        model = LinRegBunchKaufman(x, y, w)
        @test coef(model) ≈ β̂2
    end
    @testset "LinRegSweep" begin
        model = LinRegSweep(x, y, w)
        @test coef(model) ≈ β̂2
    end
end


using ModelingToolkit, GalacticOptim, Optim

@variables t x(t) v(t) u(t)
@parameters p[1:2]
D = Differential(t)

loss = (4-x)^2 + 2v^2 + u^2
eqs = [
    D(x) ~ v - p[2]*x
    D(v) ~ p[1]*u^3 + v
]

sys = ControlSystem(loss,eqs,t,[x,v],[u],p)
dt = 0.1
tspan = (0.0,1.0)
sys = runge_kutta_discretize(sys,dt,tspan)

u0 = rand(length(states(sys))) # guess for the state values
prob = OptimizationProblem(sys,u0,[0.1,0.1],grad=true)
sol = solve(prob,BFGS())

# issue #819
@testset "Combined system name collisions" begin
    eqs_short = [
        D(x) ~ - p[2]*x
        D(v) ~ p[1]*u^3
    ]
    sys1 = ControlSystem(loss,eqs_short, t, [x, v], [u], p, name = :sys1)
    sys2 = ControlSystem(loss,eqs_short, t, [x, v], [u], p, name = :sys1)
    @test_throws ArgumentError ControlSystem(loss, [sys2.v ~ sys1.v], t, [], [], [], systems = [sys1, sys2])
end


function reflection_coefficient(ω::T, amps::AverageWave{T}, medium::Medium, specie::Specie;
        θin::T = 0.0, scheme::Symbol = :trapezoidal) where T <: AbstractFloat

    k = ω/medium.c
    M = amps.hankel_order
    σ = k .* integration_scheme(amps.x; scheme=scheme) # multiple by k to be the same as the non-dimensional version

    R = T(2)*specie.num_density/(cos(θin)*k^2)*sum(
        im^T(m)*exp(im*k*amps.x[j]*cos(θin) - im*θin*m)*amps.amplitudes[j,m+M+1,1]*σ[j]
    for m=-M:M, j in eachindex(amps.x))

    return R
end


# N-gons

#Makie

#Pkg.add("Makie")

using CairoMakie, Makie.GeometryBasics
CairoMakie.activate!()
let
    n = 20
    function ngonShape(h, k, r, n)
        Polygon([Point2f(h .+ r * sin.(m * 2π / n), k .+ r * cos.(m * 2π / n)) for m in 1:n])
    end
    polysCentric = [ngonShape(0, 0, 3 / i^1.5, i) for i in 3:n]
    polysCircular = [ngonShape(√2 / 2 * sin(θ), √2 / 2 * cos(θ), 0.15 / √idx, idx + 2)
                     for (idx, θ) in enumerate(LinRange(0, 2π * (1 - 1 / (n - 2)), n - 2))]
    cmap = to_colormap(:gist_ncar)[3:end]

    with_theme(theme_light()) do
        fig, ax, = poly(polysCentric; color = 1:n-2, colormap = cmap,
            axis = (; aspect = DataAspect()), figure = (; resolution = (600, 400)))
        poly!(polysCircular; color = 1:n-2, colormap = cmap)
        hidedecorations!(ax; grid = false)
        hidespines!(ax)
        display(fig)
    end
end