function isotropic_turbulence(spectrum, p)

    # --> Get the simulation parameters.
    @unpack Lx, Ly, nx, ny = p

    # --> Wavenumbers.
    α = fftfreq(nx, nx/Lx) * 2π
    β = fftfreq(ny, ny/Ly) * 2π

    # --> Random Fourier modes with unit energy.
    ω = randn(ny, nx)
    ω .-= mean(ω)

    ω = fft(ω)
    ω ./= abs.(ω)

    # --> Apply the isotropic enstrophy spectrum
    for j = 1:nx, i = 1:ny
        # --> Wavenumber
        k = √(α[j]^2 + β[i]^2)

        # --> Rescale the Fourier mode.
        ω[i, j] = i == j == 1 ? zero(ω[i, j]) : ω[i, j] * √spectrum(k)
    end

    # --> Zero-out the last wavenumber.
    ω[:, nx÷2+1] .= 0
    ω[ny÷2+1, :] .= 0

    # --> Scale the vorticity field to the desired Reynolds number.
    ω = scale_by_re(ω, p)

    return ω
end

function scale_by_re(ω, p)

    # --> Simulation parameters.
    @unpack Lx, Ly, nx, ny = p

    # --> Rescale the vorticity field to enforce the desired Re.
    u, v = compute_velocity(ω, p)

    E = mean(u.^2 + v.^2)

    ω /= √E * 2π

    return ω
end

function isolated_vortex(x, y ; x₀=0.0, y₀=0.0, Γ=1.0, r=1.0)

    # --> Center the mesh around the (x₀, y₀) position.
    Lx = maximum(abs.(x))
    Ly = maximum(abs.(y))
    x, y = collect(x .- x₀), collect(y .- y₀)

    # --> Enforce periodicity.
    x[x .> Lx] .= x[x .> Lx] .- 2Lx
    x[x .< -Lx] .= x[x .< -Lx] .+ 2Lx

    y[y .> Ly] .= y[y .> Ly] .- 2Ly
    y[y .< -Ly] .= y[y .< -Ly] .+ 2Ly

    # --> Isolated vortex vorticity field.
    r² = (x').^2 .+ y.^2
    ω = Γ/(π*r^2) .* exp.(-r²./r^2)

    # --> Map to spectral space.
    ω = fft(ω)

    return ω
end

function dipole(p)

    # --> Unpack parameters.
    @unpack Lx, Ly, nx, ny = p

    # --> Get mesh.
    x, y = mesh(p)

    # --> Dipole.
    ω = isolated_vortex(x, y ; Γ=1.0, r=0.5, y₀=0.5)
    ω += isolated_vortex(x, y ; Γ=-1.0, r=0.5, y₀=-0.5)

    # --> Remove the mean.
    ω[1, 1] = 0.0

    # --> Scale vorticity field to the desired Reynolds number.
    ω = scale_by_re(ω, p)

    return ω
end

function taylor_green(p)

    # --> Construct the mesh.
    x, y = mesh(p)

    # --> Taylor-Green vortices.
    ω = -2 .* cos.(x') .* cos.(y)

    # --> Map to spectral space.
    ω = fft(ω)

    return ω
end


# TODO: Fix image birghtness when writing images out to file

using DelimitedFiles
using LinearAlgebra
using Images
using ImageIO

# Function to output a set of points as an image
function create_image_array(points, ranges, res; point_offset = 10)
    output_array = zeros(res)

    range_diffs = [ranges[1,2] - ranges[1,1], ranges[2,2] - ranges[2,1]]

    # bin the pixels
    for i = point_offset:size(points)[1]

        yval = floor(Int, (points[i,1] - ranges[2,1]) /
                          (range_diffs[2]/ res[2]))
        xval = floor(Int, (points[i,2] - ranges[1,1]) /
                          (range_diffs[1] / res[1]))

        if (xval < res[1] && yval < res[2] && xval > 0 && yval > 0)
            output_array[xval,yval] += 1.0
        end
    end

    if maximum(output_array) > 0
        return normalize(output_array)
    else
        return zeros(size(output_array))
    end

end

function write_image(points, ranges, res, filename; point_offset = 10)
    img_array = create_image_array(points, ranges, res;
                                   point_offset = point_offset)
    write_image(img_array, filename)

end

function write_image(img_array, filename)
    save(filename, reverse(img_array / maximum(img_array),dims=1))
end

# Create a set of affine transforms for function system
function create_sierpinski_set(A, B, C)
    return [[[0.5 0 A[1]*0.5; 0 0.5 A[2]*0.5; 0 0 1], .33],
            [[0.5 0 B[1]*0.5; 0 0.5 B[2]*0.5; 0 0 1], 0.33],
            [[0.5 0 C[1]*0.5; 0 0.5 C[2]*0.5; 0 0 1], 0.34]]
end

function create_barnsley_set()
    return [[[0.0 0.0 0.0; 0.0 0.16 0.0; 0.0 0.0 1.0], 0.01],
            [[0.85 0.04 0.0; -0.04 0.85 1.60; 0.0 0.0 1.0], 0.85],
            [[0.20 -0.26 0.0; 0.23 0.22 1.60; 0.0 0.0 1.0], 0.07],
            [[-0.15 0.28 0.0; 0.26 0.24 0.44; 0.0 0.0 1.0], 0.07]]
end

function select_array(hutchinson)

    rnd = rand()
    if rnd < hutchinson[1][2]
        return hutchinson[1][1]
    end

    rnd_sum = hutchinson[1][2]

    for i = 2:length(hutchinson)
        if rnd > rnd_sum && rnd < rnd_sum + hutchinson[i][2]
            return hutchinson[i][1]
        end
        rnd_sum += hutchinson[i][2]
    end

    return hutchinson[end][2]

end

# This is a function to simulate a "chaos game"
function chaos_game(n::Int, hutchinson)

    # Initializing the output array and the initial point
    output_points = zeros(n,2)
    point = [rand(), rand(), 1]

    for i = 1:n
        output_points[i,:] .= point[1:2]
        point = select_array(hutchinson)*point
    end

    return output_points

end

function main()
    #hutchinson = create_sierpinski_set([0.0, 0.0], [0.5, sqrt(0.75)], [1.0, 0.0])
    hutchinson = create_barnsley_set()
    output_points = chaos_game(100000000, hutchinson)
    #println(output_points)
    writedlm("out.dat", output_points)
    return output_points
end


# ENV["CUDA_VISIBLE_DEVICES"] = ""
using Revise
using ADCME
using ADSeismic
using PyPlot
using DelimitedFiles
matplotlib.use("agg")
o = placeholder(ones(1))
param = ElasticPropagatorParams(NX=150, NY=150, NSTEP=5, DELTAT=1e-4, DELTAX=1.0, DELTAY=1.0,
                                vp_ref = 3300.0, USE_PML_XMIN=true, USE_PML_XMAX=true, USE_PML_YMIN=true, USE_PML_YMAX=true)

source = Ricker(param, 15.0, 100.0, 1e6)
srci = repeat([param.NX ÷ 2], 5)
srcj = repeat([param.NY ÷ 2], 5)
srctype = [0;1;2;3;4] # vx, vy, sigmaxx, sigmayy, sigmaxy
srcv = reshape(source, :, 1)
src = ElasticSource(srci, srcj, srctype, srcv)

vp = 3000. * ones(param.NX+2, param.NY+2)
vs = vp / 1.732
rho = 2800. * ones(param.NX+2, param.NY+2)
vp0 = vp
vp_ = Variable(vp[:]) + o[1]^2
vp = reshape(vp_, (152,152))
λ, μ, ρ = compute_lame_parameters(param.NX, param.NY, vp, vs, rho)
model = ElasticPropagatorSolver(param, src, ρ, λ, μ)

sess = Session(); init(sess)

# loss = sum(model.vx)
loss = sum(model.vx^2) + sum(model.vy^2) + sum(model.sigmaxx^2) + sum(model.sigmaxy^2) + sum(model.sigmayy^2)
grad = gradients(loss, o)

# vx = run(sess, model.vx)
# p = visualize_wavefield(vx, param)
# saveanim(p, "elastic-wavefield.gif")

function test(x)

  l = run(sess, loss, o=>x)
  g = run(sess, grad, o=>x)

  return l, g

end

close("all")
ADCME.test_gradients(test, ones(1))
savefig("gradtest_sigmaxy_gpu.png")


using Base.Test
using StatsFuns
using PMTK4Stats

a = 3
b = 9
icdf(p::Real) = betainvcdf(a, b, p)
lb, ub = hdiFromIcdf(icdf)

@test_approx_eq_eps .04 lb 1e-2
@test_approx_eq_eps .48 ub 1e-2


using Measurements
r_ini = [14.0,0.1,
        0.14,0.4,
        5.0,0.2,
        20.0,4.0]

model = bimolecular_diss_bell_activated_state_w_cons!
# u₀ = [1.0,0.0]
u₀ = [] # don't provide initial states
u₀_type = "equilibrium" # calculate initial states based on equilibrium
# u₀ = "detailed_balance"
# u₀ = "force"
# u₀ = []
tspan = (0.0,maximum(γϵ.t))
# k₋ᵒf_γϵ = 14.1
k₋ᵒf_γϵ = 12.2
γϵ_ft_M = M(model,u₀,u₀_type,tspan,[],[k₋ᵒf_γϵ],0.0,Rodas5)
# γϵ_ft_M = M(model,u₀,u₀_type,tspan,[],[0.5],0.0,Rodas5)

pinit = [   0.01,
        0.1,0.6,
        3.,0.2,
        150.0,4.0]
opt =  [        (0.,10.),
        (0.,1.),(0.,10.),
        (0.,10.),(0.,10.),
        (0.,500.),(0.,10.)]

@time mle_loss(pinit,γϵ,γϵ_ft_M)
# loss_in = (p) -> mle_loss(p,γϵ,γϵ_ft_M)
# hes = ForwardDiff.hessian(loss_in,rates_γϵ)
# σ_γϵ = hessian2σ(hes)

@time rates_γϵ, l_γϵ, hes_γϵ = optimization(γϵ_ft_M,γϵ,opt)
ks_γϵ = rates_γϵ
σ_γϵ = hessian2σ(hes_γϵ)
k_γϵ = rates_γϵ .± σ_γϵ

plotly()
plot_ft_fit(γϵ_ft_M,ks_γϵ,γϵ)
# plot_ft_fit(γϵ_ft_M,ks_γϵ,γϵ,dF=30)
plot(); Color = :blue
plot_mean_ft_fit(γϵ_ft_M,k_γϵ,γϵ,γϵ_binned;species="γϵ")
# sols_γϵ = ftsolve(γϵ_ft_M,k_γϵ) #change solver to Tsit5
# plot(Fs,mean_lifetime.(sols_γϵ),ylabel="<t>",xlabel="F",title="γϵ")
# plot_F_vs_t(clean_dissociation_data(γϵ_binned))
u₀γϵ = initial_u(ks_γϵ,γϵ_ft_M)
________________________________________________________________________________________________
## δϵ fitting with 2 species model
model = bimolecular_diss_bell_activated_state_w_cons!
u₀ = []
u₀_type = "equilibrium"
tspan = (0.0,maximum(δϵ.t))
# k₋ᵒf_δϵ = 11.1
k₋ᵒf_δϵ = 13.8
δϵ_ft_M = M(model,u₀,u₀_type,tspan,[],[k₋ᵒf_δϵ],0.0,Rodas5)
pinit = [   0.1,
        0.1,0.5,
        0.1,0.5,
        10.0,5.0]
opt =  [        (0.,10.),
        (0.,1.),(0.,10.),
        (0.,10.),(0.,10.),
        (0.,100.),(0.,10.)]

@time mle_loss(pinit,δϵ,δϵ_ft_M)

@time rates_δϵ, l_δϵ, hes_δϵ = optimization(δϵ_ft_M,δϵ,opt)
ks_δϵ = rates_δϵ
σ_δϵ = hessian2σ(hes_δϵ)
k_δϵ = rates_δϵ .± σ_δϵ
u₀δϵ = initial_u(ks_δϵ,δϵ_ft_M)

plot_ft_fit(δϵ_ft_M,ks_δϵ,δϵ)
plot();
Color = :red
plot_mean_ft_fit(δϵ_ft_M,k_δϵ,δϵ,δϵ_binned;species="δϵ")
# plot(Fs,mean_lifetime.(sols_δϵ),ylabel="<t>",xlabel="F",title="δϵ")
# plot_F_vs_t(clean_dissociation_data(δϵ_binned))

## mix fitting with 2 species model

model = cd3_tri_dissociation_bell!
# u₀_type = "equilibrium"
u₀_type = []
# u₀ = vcat(u₀γϵ./2, u₀δϵ./2, 0.0,0.0)
u₀ = vcat(0.454*u₀γϵ,0.290*u₀δϵ, 0.256*[1.,0.])
# u₀ = vcat(0.454*u₀γϵ,0.290*u₀δϵ, 0.256*[0.,0.])/(1-0.256)
sum(u₀)
# u₀ = vcat(u₀γϵ/3, u₀δϵ/3)
tspan = (0.0,maximum(mix.t))
k₋ᵒf_mix = 0.6
cons = [k₋ᵒf_γϵ;
        ks_γϵ;
        k₋ᵒf_δϵ;
        ks_δϵ;
        k₋ᵒf_mix ]
mix_ft_M = M(model,u₀,u₀_type,tspan,[],cons,0.0,Rodas5)

pinit = [   0.1,
        0.1,0.5,
        0.1,0.5,
        10.0,5.0]
# pinit = zeros(7)
opt =  [        (0.,10.),
        (0.,1.),(0.,1.),
        (0.,10.),(0.,1.),
        (0.,100.),(0.,10.)]

@time mle_loss(pinit,mix,mix_ft_M)
loss_in = (p) -> mle_loss(p,mix,mix_ft_M)
hes_mix = ForwardDiff.hessian(loss_in,pinit)
# hes = Zygote.hessian(loss_in,pinit)

@time rates_mix, l_mix, hes_mix = optimization(mix_ft_M,mix,opt)
ks_mix = rates_mix
σ_mix = hessian2σ(hes_mix)
k_mix = rates_mix .± σ_mix

ts = 0:dt:maximum(mix.t)
Fs = 0:1:maximum(mix.F)
plot_ft_fit(mix_ft_M,ks_mix,mix)
# sols_mix = plot_ft_fit(mix_ft_M,pinit,mix)

plot_ft_fit(mix_ft_M,ks_mix,mix)
plot(); Color = :green
plot_mean_ft_fit(mix_ft_M,k_mix,mix,mix_binned;species="mix")
# plot(Fs,mean_lifetime.(sols_mix),ylabel="<t>",xlabel="F",title="mix")
# plot_F_vs_t(clean_dissociation_data(mix_binned))

## plot rates

plt = plot_bell_equations(vcat(k₋ᵒf_γϵ, ks_γϵ))
title!("γϵ")
plt = plot_bell_equations(vcat(k₋ᵒf_δϵ, ks_δϵ))
title!("δϵ")
plt = plot_bell_equations(vcat(k₋ᵒf_mix, ks_mix))
title!("mix")

pyplot()
plot(); Color = :blue
plot_mean_ft_fit(γϵ_ft_M,k_γϵ,γϵ,γϵ_binned;species="γϵ")
Color = :red
plot_mean_ft_fit(δϵ_ft_M,k_δϵ,δϵ,δϵ_binned;species="δϵ")
Color = :green
plot_mean_ft_fit(mix_ft_M,k_mix,mix,mix_binned;species="mix")

plot(); Color = :blue
extrapolate_fit(γϵ_ft_M,ks_γϵ,γϵ,150.0;species="γϵ",dF=10.0)
Color = :red
extrapolate_fit(δϵ_ft_M,ks_δϵ,δϵ,150.0;species="δϵ",dF=10.0)
Color = :green
extrapolate_fit(mix_ft_M,ks_mix,mix,150.0;species="mix",dF=10.0)

mean_t_γϵ = mean_ft_fit(γϵ_ft_M,k_γϵ,γϵ,35;species=[],dt=0.1,dF=1.0)
mean_t_δϵ = mean_ft_fit(δϵ_ft_M,k_δϵ,δϵ,35;species=[],dt=0.1,dF=1.0)
mean_t_mix = mean_ft_fit(mix_ft_M,k_mix,mix,35;species=[],dt=0.1,dF=1.0)

mean_t_sum_bi = 0.610*mean_t_γϵ + 0.390*mean_t_δϵ
plot!(mean_t_sum_bi)

pyplot()
lftm_coop = (mean_t_mix -mean_t_sum_bi) ./ mean_t_sum_bi
plot(0:35, 100 .* lftm_coop,errorbar=:ribbon, xlabel="F",ylabel="lifetime cooperativity (%)")

CSV.write("results/γϵ_fit.csv", measurements2Df(0:35,mean_t_γϵ))
CSV.write("results/δϵ_fit.csv", measurements2Df(0:35,mean_t_δϵ))
CSV.write("results/mix_fit.csv", measurements2Df(0:35,mean_t_mix))
CSV.write("results/lftm_coop.csv", measurements2Df(0:35,lftm_coop))


@testset "CTMCSt1" begin
    Q = @expr [-x x 0;
        1 -2 1;
        0 1 -1]
    v = @expr [1, 0, 0]

    ma = prob(Q)

    x = 2.0
    ex = prob(Float64[-x x 0; 1 -2 1; 0 1 -1])
    @bind x = x
    res = seval(ma)
    @test isapprox(ex, res)
end

@testset "CTMCSt2" begin
    Q = @expr [-x x 0;
        1 -2 1;
        0 1 -1]
    v = @expr [1, 0, 0]
    csc = SparseCSC(Q)
    ma = prob(csc)

    x = 2.0
    ex = prob(Float64[-x x 0; 1 -2 1; 0 1 -1])
    @bind x = x
    res = seval(ma)
    @test isapprox(ex, res)
end

@testset "CTMCSt3" begin
    Q = @expr [-x x 0;
        1 -2 1;
        0 1 -1]
    v = @expr [1, 0, 0]

    ma = dot(prob(Q), v)
    println(ma)
    x = 2.0
    ex = sum(prob(Float64[-x x 0; 1 -2 1; 0 1 -1]) .* [1,0,0])
    @bind x = x
    res = seval(ma)
    println(res, ex)
    @test isapprox(ex, res)
end

@testset "CTMCStsen1" begin
    Q = @expr [
        -x x 0;
        1 -2 1;
        0 1 -1]
    v = @expr [1, 0, 0]
    csc = SparseCSC(Q)
    ma = prob(csc)

    x0 = 2.0
    test = SymbolicEnv()
    @bind test x = x0
    h = 0.0001

    test0 = SymbolicEnv()
    @bind test0 x = x0 + h
    test1 = SymbolicEnv()
    @bind test1 x = x0 - h

    vv0 = seval(ma, test0)
    vv1 = seval(ma, test1)
    ex = (vv0 - vv1) / (2*h)
    @test isapprox(seval(ma, :x, test), ex)
end

@testset "CTMCStsen2" begin
    Q = @expr [
        -x x 0;
        1 -2 1;
        0 1 -1]
    v = @expr [1, 0, 0]
    csc = SparseCSC(Q)
    ma = dot(prob(csc), v)

    x0 = 2.0
    test = SymbolicEnv()
    @bind test x = x0
    h = 0.0001

    test0 = SymbolicEnv()
    @bind test0 x = x0 + h
    test1 = SymbolicEnv()
    @bind test1 x = x0 - h

    vv0 = seval(ma, test0)
    vv1 = seval(ma, test1)
    ex = (vv0 - vv1) / (2*h)
    @test isapprox(seval(ma, :x, test), ex)
end

@testset "CTMCsensen1" begin
    Q = @expr [
        -x x 0;
        1 -2 1;
        0 1 -1]
    v = @expr [1, 0, 0]
    csc = SparseCSC(Q)
    ma = prob(csc)

    x0 = 2.0
    test = SymbolicEnv()
    @bind test x = x0
    h = 0.0001

    test0 = SymbolicEnv()
    @bind test0 x = x0 + h
    test1 = SymbolicEnv()
    @bind test1 x = x0 - h

    vv0 = seval(ma, test0)
    vv1 = seval(ma, test)
    vv2 = seval(ma, test1)
    ex = (vv0 -2*vv1 + vv2) / (h^2)
    @test isapprox(seval(ma, (:x, :x), test), ex, atol = 1.0e-5)
end

@testset "CTMCsensen2" begin
    Q = @expr [
        -x x 0;
        1 -2 1;
        0 1 -1]
    v = @expr [1, 0, 0]
    csc = SparseCSC(Q)
    ma = dot(prob(csc), v)

    x0 = 2.0
    test = SymbolicEnv()
    @bind test x = x0
    h = 0.0001

    test0 = SymbolicEnv()
    @bind test0 x = x0 + h
    test1 = SymbolicEnv()
    @bind test1 x = x0 - h

    vv0 = seval(ma, test0)
    vv1 = seval(ma, test)
    vv2 = seval(ma, test1)
    ex = (vv0 -2*vv1 + vv2) / (h^2)
    @test isapprox(seval(ma, (:x, :x), test), ex, atol = 1.0e-5)
end

@testset "CTMCStsen3" begin
    Q = @expr [
        -x x 0;
        1 -2 1;
        0 1 -1]
    v = @expr [1, 0, 0]
    csc = Q
    ma = dot(prob(csc), v)

    x0 = 2.0
    test = SymbolicEnv()
    @bind test x = x0
    h = 0.0001

    test0 = SymbolicEnv()
    @bind test0 x = x0 + h
    test1 = SymbolicEnv()
    @bind test1 x = x0 - h

    vv0 = seval(ma, test0)
    vv1 = seval(ma, test1)
    ex = (vv0 - vv1) / (2*h)
    @test isapprox(seval(ma, :x, test), ex)
end


using Makie

"""
convert_to_cartesian(p)
Converts a point in the geographic coordinate system to a point in the
cartesian one with the given point in radians.
"""
function convert_to_cartesian(p)
    [cos(p[2]) * cos(p[1]),
     cos(p[2]) * sin(p[1]),
     sin(p[2])]
end

"""
convert_to_geographic(p)
Converts a point in the cartesian coordinate system to a point in the
geographic one with the given point.
"""
function convert_to_geographic(p)
    r = sqrt(p[1]^2 + p[2]^2 + p[3]^2)
    if p[1] > 0
          ϕ = atan(p[2] / p[1])
    elseif p[2] > 0
          ϕ = atan(p[2] / p[1]) + pi
    else
          ϕ = atan(p[2] / p[1]) - pi
    end
    θ = asin(p[3] / r)
    [ϕ, θ]
end

"""
λmap(p)
Sends a point on a 3-sphere to a point in the plane x₄=0 with the given point
in the form of 2 complex numbers representing a unit quaternion. This is the
stereographic projection of a 3-sphere. S³ ↦ R³
"""
function λmap(p)
    [real(p[1]), imag(p[1]), real(p[2])] ./ (1 - imag(p[2]))
end

"""
λ⁻¹map(p)
Sends a point on the plane back to a point on a unit sphere with the given
point. This is the inverse stereographic projection of a 3-sphere.
"""
function λ⁻¹map(p)
    x₁ = 2p[1] / (1 + p[1]^2 + p[2]^2 + p[3]^2)
    x₂ = 2p[2] / (1 + p[1]^2 + p[2]^2 + p[3]^2)
    x₃ = 2p[3] / (1 + p[1]^2 + p[2]^2 + p[3]^2)
    x₄ = (-1 + p[1]^2 + p[2]^2 + p[3]^2) / (1 + p[1]^2 + p[2]^2 + p[3]^2)
    [Complex(x₁, x₂), Complex(x₃, x₄)]
end

"""
S¹action(α, p)
Performs a group action corresponding to moving along the circumference of a
circle with the given angle and the point in the form of 2 complex numbers
representing a unit quaternion.
"""
function S¹action(α, p)
    [exp(im * α) * p[1], exp(im * α) * p[2]]
end

"""
get_manifold(points, segments, start, finish)

Calculates a grid of points in R³ for constructing a surface in a specific way
with the given points in the base space, the number of segments, and the start
and finish angles that determine where to cut.
"""
function get_manifold(points, segments, start, finish)
    samples = size(points, 1)
    manifold = Array{Float64}(undef, segments, samples, 3)
    for i in 1:samples
            ϕ, θ = points[i, :]
        f = θ / (pi / 2)
        α = -(finish - start) / (segments - 1) / f
        p = S¹action(-start / f - pi/2, λ⁻¹map(convert_to_cartesian([ϕ, θ])))
        for j in 1:segments
            x₁, x₂, x₃ = λmap(p)
            manifold[j, i, :] = [x₁, x₂, x₃]
            p = S¹action(α, p)
        end
    end
    manifold
end

"""
build_surface(scene, points, color; transparency, shading)

Builds a surface with the given scene, points, color, transparency and shading.
"""
function build_surface(scene,
                       points,
                       color;
                       transparency = false,
                       shading = true)
    surface!(scene,
             @lift($points[:, :, 1]),
             @lift($points[:, :, 2]),
             @lift($points[:, :, 3]),
             color = color,
             transparency = transparency,
             shading = shading)
end

# The scene object that contains other visual objects
universe = Scene(backgroundcolor = :orange, show_axis=false, resolution = (360, 360))
# Use a slider for rotating the base space in an interactive way
sg, og = textslider(0:0.05:2pi, "g", start = 0)
segments = 100
samples = 100
# Construct a manifold for each country in the dictionary
latitudes = [pi/12, pi/6, pi/4, pi/3, 5pi/12]
for l in 1:length(latitudes)
    latitude = latitudes[l]
    # Sample a circle of points
    points = Array{Float64}(undef, samples, 2)
    for i in 1:samples
        longitude = i * 2pi / (samples - 1)
        points[i, :] = [longitude, latitude]
    end
    if (-1) ^ l > 0
        color = fill(RGBAf0(0.0, 0.0, 1.0, 1.0), segments, samples)
    else
        color = fill(RGBAf0(0.0, 0.0, 0.75, 1.0), segments, samples)
    end
    for i in range(1, stop = samples, step = 10)
        color[:, i] = fill(RGBAf0(rand(), rand(), rand(), 1.0), segments)
    end
    # Rotate the points by adding to the longitudes
    rotated = @lift begin
        R = similar(points)
        for i in 1:samples
                  ϕ, θ = points[i, :]
            R[i, :] = [ϕ + $og, θ]
        end
        R
    end
    start = -latitude
    finish = latitude
    manifold = @lift(get_manifold($rotated,
                                  segments,
                                  start,
                                  finish))
    build_surface(universe, manifold, color)
end

# Instantiate a horizontal box for holding the visuals and the controls
#scene = hbox(universe,
#             vbox(sg),
#             parent = Scene(resolution = (360, 360)))

# update eye position
eye_position, lookat, upvector = Vec3f0(-2, 0, 2), Vec3f0(0), Vec3f0(0, 0, 1.0)
update_cam!(universe, eye_position, lookat)
universe.center = false # prevent scene from recentering on display
rotate_cam!(universe, 0.0, 0.0, pi/2)

record(universe, "gallery/londontsai.gif") do io
    frames = 90
    for i in 1:frames
        og[] = i*2pi/frames # animate scene
        rotate_cam!(universe, 0.0, 2pi/frames, 0.0)
        recordframe!(io) # record a new frame
    end
end


# MobilityTheories.jl

"""
    polaronmobility(Trange, ε_Inf, ε_S, freq, effectivemass; verbose::Bool=false)

Solves the Feynman polaron problem variationally with finite temperature Osaka energies. From the resulting v, and w parameters, calculates polaron structure (wave function size, etc.).  Uses FHIP, Kadanoff (Boltzmann relaxation time) and Hellwarth direct contour integration to predict a temperature-dependent mobility for the material system.

# Arguments
- `Trange::range`: temperature range.
- `ε_Inf`: reduced optical dielectric constant.
- `ε_S`: reduced static dielectric constant.
- `freq`: characteristic dielectric phonon frequency (THz).
- `effectivemass`: bare-band effective mass (mₑ).

Returns a structure of type `Polaron`, containing arrays of useful
information.  Also prints a lot of information to the standard out - which
may be more useful if you're just inquiring as to a particular data point,
rather than plotting a temperature-dependent parameter.

As an example, to calculate the electron polaron in MAPI at 300 K:
# Examples
```jldoctest
polaronmobility(300, 4.5, 24.1, 2.25, 0.12)
```
"""
function polaronmobility(Trange, ε_Inf, ε_S, freq, effectivemass; verbose::Bool=false)
    println("\n\nPolaron mobility for system ε_Inf=$ε_Inf, ε_S=$ε_S, freq=$freq,
                 effectivemass=$effectivemass; with Trange $Trange ...")

    # Internally we we use 'mb' for the 'band mass' in SI units, of the effecitve-mass of the electron
    mb=effectivemass*MassElectron
    ω = (2*pi)*freq*1e12 # angular-frequency, of the phonon mode

    α=frohlichalpha(ε_Inf, ε_S,  freq,    effectivemass)
    #α=2.395939683378253 # Hard coded; from MAPI params, 4.5, 24.1, 2.25THz, 0.12me

    v=7.1 # starting guess for v,w variational parameters
    w=6.5

    @printf("Polaron mobility input parameters: ε_Inf=%f ε_S=%f freq=%g α=%f \n",ε_Inf, ε_S, freq, α )
    @printf("Derived params in SI: ω =%g mb=%g \n",ω ,mb)

    # Empty struct for storing data
    # A slightly better way of doing this ф_ф ...
    p=Polaron()
    # populate data structure with (athermal) parameters supplied...
    append!(p.α,α) # appending so as not to mess with type immutability
    append!(p.mb,mb)
    append!(p.ω,ω)

    # We define βred as the subsuming the energy of the phonon; i.e. kbT c.f. ħω
    for T in Trange
        β=1/(kB*T)
        βred=ħ*ω*β
        append!(p.βred,βred)
        @printf("T: %f β: %.3g βred: %.3g ħω  = %.3g meV\t",T,β,βred, 1E3* ħ*ω  / q)

        if T==0
            v,w=feynmanvw(α, v=v,w=w)
        else
            v,w=feynmanvw(α,βred, v=v,w=w)
        end

        @printf("\n Polaron Parameters:  v= %.4f  w= %.4f",v,w)

        # From 1962 Feynman, definition of v and w in terms of the coupled Mass and spring-constant
        # See Page 1007, just after equation (18)
        # Units of M appear to be 'electron masses'
        # Unsure of units for k, spring coupling constant
        k=(v^2-w^2)
        M=(v^2-w^2)/w^2
        append!(p.k,k)
        append!(p.M,M)

        @printf("  ||   M=%f  k=%f\t",M,k)
        @printf("\n Bare-band effective mass: %f Polaron effective mass: %f Polaron mass enhancement: %f%%",effectivemass,effectivemass*(1+M),M*100)

        @printf("\n Polaron frequency (SI) v= %.2g Hz  w= %.2g Hz",
                v*ω /(2*pi),w*ω /(2*pi))

		# (46) in Feynman1955
		meSmallAlpha(α )=α /6 + 0.025*α ^2
		# (47) In Feynman1955
		meLargeAlpha(α )=16*α ^4 / (81*π ^4)
		#meLargeAlpha(α )=202*(α /10)^4
        if (verbose) # asymptotic solutions - not that interesting when you have the actual ones!
		    @printf("\n Feynman1955(46,47): meSmallAlpha(α)= %.3f meLargeAlpha(α)= %.3f",
                    meSmallAlpha(α),meLargeAlpha(α))
            @printf("\n Feynman1962: Approximate ~ Large alpha limit, v/w = %.2f  =~approx~= alpha^2 = %.2f ",
                    v/w,α ^2)
        end

        # POLARON SIZE
        @printf("\n Polaron size (rf), following Schultz1959. (s.d. of Gaussian polaron ψ )")
        # Schultz1959 - rather nicely he actually specifies everything down into units!
        # just before (2.4) in Schultz1959
        mu=((v^2-w^2)/v^2)
        # (2.4)
        rf=sqrt(3/(2*mu*v))
        # (2.4) SI scaling inferred from units in (2.5a) and Table II
        rfsi=rf*sqrt(2*me*ω )
        @printf("\n\t Schultz1959(2.4): rf= %g (int units) = %g m [SI]",rf,rfsi )
        append!(p.rfsi,rfsi)

        if (verbose)
        scale=sqrt(2*mb*ω) # Note we're using mb;
        #band effective-mass in SI units (i.e. meff*melectron)

        rfa=(3/(0.44*α ))^0.5 # As given in Schultz1959(2.5a), but that 0.44 is actually 4/9
        @printf("\n\t Schultz1959(2.5a) with 0.44: Feynman α→0 expansion: rfa= %g (int units) = %g m [SI]",rfa,scale*rfa )
        rfa=(3/((4/9)*α ))^0.5 # Rederived from Feynman1955, 8-8-2017; Yellow 2017.B Notebook pp.33-34
        @printf("\n\t Schultz1959(2.5a) with 4/9 re-derivation: Feynman α→0 expansion: rfa= %g (int units) = %g m [SI]",rfa,scale*rfa )
        append!(p.rfsmallalpha,scale*rfa)

        rfb=3*(pi/2)^0.5 * α
        @printf("\n\t Schultz1959(2.5b): Feynman α→∞ expansion: rf= %g (int units) = %g m [SI]",rfb,scale*rfb )

        # Schultz1959 - Between (5.7) and (5.8) - resonance of Feynman SHM system
        phononfreq=sqrt(k/M)
        @printf("\n\t Schultz1959(5.7-5.8): fixed-e: phononfreq= %g (int units) = %g [SI, Hz] = %g [meV]",
            phononfreq,phononfreq*ω /(2*pi), phononfreq*hbar*ω *1000/q)

        phononfreq=sqrt(k/mu) # reduced mass
        @printf("\n\t Schultz1959(5.7-5.8): reducd mass: phononfreq= %g (int units) = %g [SI, Hz] = %g [meV]",
            phononfreq,phononfreq*ω /(2*pi), phononfreq*hbar*ω *1000/q)

        @printf("\n\t Schultz1959: electronfreq= %g (int units) = %g [SI, Hz] = %g [meV]",
            sqrt(k/1),sqrt(k/1)*ω /(2*pi), sqrt(k/1)*hbar*ω *1000/q)
        @printf("\n\t Schultz1959: combinedfreq= %g (int units) = %g [SI, Hz] = %g [meV]",
            sqrt(k/(1+M)),sqrt(k/(1+M))*ω /(2*pi), sqrt(k/(1+M))*hbar*ω *1000/q)

        # Devreese1972: 10.1103/PhysRevB.5.2367
        # p.2371, RHS.
        @printf("\n Devreese1972: (Large Alpha) Franck-Condon frequency = %.2f", 4*α ^2/(9*pi))
        end

        # F(v,w,β,α)=-(A(v,w,β)+B(v,w,β,α)+C(v,w,β)) #(62a) - Hellwarth 1999
        if (T>0)
            @printf("\n Polaron Free Energy: A= %f B= %f C= %f F= %f",A(v,w,βred),B(v,w,βred,α),C(v,w,βred),F(v,w,βred,α))
            @printf("\t = %f meV",1000.0 * F(v,w,βred,α) * ħ*ω  / q)
            append!(p.A,A(v,w,βred))
            append!(p.B,B(v,w,βred,α))
            append!(p.C,C(v,w,βred))
            append!(p.F,F(v,w,βred,α))
        else # Athermal case; Enthalpy
            @printf("\n Polaron Enthalpy: F= %f = %f meV \n",F(v,w,α), 1_000*F(v,w,α) * ħ*ω/q)

            return # return early, as if T=0, all mobility theories = infinite / fall over
        end

        # FHIP
        #    - low-T mobility, final result of Feynman1962
        # [1.60] in Devreese2016 page 36; 6th Edition of Frohlich polaron notes (ArXiv)
        # I believe here β is in SI (expanded) units
        @printf("\n Polaron Mobility theories:")
        μ=(w/v)^3 * (3*q)/(4*mb*ħ*ω^2*α*β) * exp(ħ*ω*β) * exp((v^2-w^2)/(w^2*v))
        @printf("\n\tμ(FHIP)= %f m^2/Vs \t= %.2f cm^2/Vs",μ,μ*100^2)
        append!(p.T,T)
        append!(p.FHIPμ,μ*100^2)

        # Kadanoff
        #     - low-T mobility, constructed around Boltzmann equation.
        #     - Adds factor of 3/(2*beta) c.f. FHIP, correcting phonon emission behaviour
        # [1.61] in Devreese2016 - Kadanoff's Boltzmann eqn derived mob
        μ=(w/v)^3 * (q)/(2*mb*ω*α) * exp(ħ*ω*β) * exp((v^2-w^2)/(w^2*v))
        @printf("\n\tμ(Kadanoff,via Devreese2016)= %f m^2/Vs \t= %.2f cm^2/Vs",μ,μ*100^2)

        append!(p.Kμ,μ*100^2)

        ######
        # OK, now deep-diving into Kadanoff1963 itself to extract Boltzmann equation components
        # Particularly right-hand-side of page 1367
        #
        # From 1963 Kadanoff, (9), define eqm. number of phonons (just from T and phonon omega)
        Nbar=(exp(βred) - 1)^-1
        @printf("\n\t\tEqm. Phonon. pop. Nbar: %f ",Nbar)
        if (verbose)
            @printf("\n\texp(Bred): %f exp(-Bred): %f exp(Bred)-1: %f",exp(βred),exp(-βred),exp(βred)-1)
        end
        Nbar=exp(-βred)
        #Note - this is only way to get Kadanoff1963 to be self-consistent with
        #FHIP, and later statements (Devreese) of the Kadanoff mobility.
        #It suggests that Kadanoff used the wrong identy for Nbar in 23(b) for
        #the Gamma0 function, and should have used a version with the -1 to
        #account for Bose / phonon statistics

        #myv=sqrt(k*(1+M)/M) # cross-check maths between different papers
        #@printf("\nv: %f myv: %f\n",v,myv)

        # Between 23 and 24 in Kadanoff 1963, for small momenta skip intergration --> Gamma0
        Gamma0=2*α * Nbar * (M+1)^(1/2) * exp(-M/v)
        Gamma0*=ω  #* (ω *hbar) # Kadanoff 1963 uses hbar=omega=mb=1 units
            # Factor of omega to get it as a rate relative to phonon frequency
            # Factor of omega*hbar to get it as a rate per energy window
        μ=q/( mb*(M+1) * Gamma0 ) #(25) Kadanoff 1963, with SI effective mass

        if (verbose) # these are cross-checks
            @printf("\n\tμ(Kadanoff1963 [Eqn. 25]) = %f m^2/Vs \t = %.2f cm^2/Vs",μ,μ*100^2)
            @printf("\n\t\t Eqm. Phonon. pop. Nbar: %f ",Nbar)
        end

        @printf("\n\t\tGamma0 = %g rad/s = %g /s ",
                Gamma0, Gamma0/(2*pi))
        @printf(" \n\t\tTau=1/Gamma0 = %g s = %f ps",
            2*pi/Gamma0, 2*pi*1E12/Gamma0)
        Eloss=hbar*ω * Gamma0/(2*pi) # Simply Energy * Rate
        @printf("\n\t\tEnergy Loss = %g J/s = %g meV/ps",Eloss,Eloss * 1E3 / (q*1E12) )
        append!(p.Tau, 2*pi*1E12/Gamma0) # Boosted into ps ?

        # Hellwarth1999 - directly do contour integration in Feynman1962, for
        # finite temperature DC mobility
        # Hellwarth1999 Eqn (2) and (1) - These are going back to the general
        # (pre low-T limit) formulas in Feynman1962.  to evaluate these, you
        # need to do the explicit contour integration to get the polaron
        # self-energy
        R=(v^2-w^2)/(w^2*v) # inline, page 300 just after Eqn (2)

        #b=R*βred/sinh(b*βred*v/2) # This self-references b! What on Earth?
        # OK! I now understand that there is a typo in Hellwarth1999 and
        # Biaggio1997. They've introduced a spurious b on the R.H.S. compared to
        # the original, Feynman1962:
        b=R*βred/sinh(βred*v/2) # Feynman1962 version; page 1010, Eqn (47b)

        a=sqrt( (βred/2)^2 + R*βred*coth(βred*v/2))
        k(u,a,b,v) = (u^2+a^2-b*cos(v*u))^(-3/2)*cos(u) # integrand in (2)
        K=quadgk(u->k(u,a,b,v),0,Inf)[1] # numerical quadrature integration of (2)

        #Right-hand-side of Eqn 1 in Hellwarth 1999 // Eqn (4) in Baggio1997
        RHS= α/(3*sqrt(π)) * βred^(5/2) / sinh(βred/2) * (v^3/w^3) * K
        μ=RHS^-1 * (q)/(ω*mb)
        @printf("\n\tμ(Hellwarth1999)= %f m^2/Vs \t= %.2f cm^2/Vs",μ,μ*100^2)
        append!(p.Hμ,μ*100^2)

        if (verbose)
        # Hellwarth1999/Biaggio1997, b=0 version... 'Setting b=0 makes less than 0.1% error'
        # So let's test this
        R=(v^2-w^2)/(w^2*v) # inline, page 300 just after Eqn (2)
        b=0
        a=sqrt( (βred/2)^2 + R*βred*coth(βred*v/2))
        #k(u,a,b,v) = (u^2+a^2-b*cos(v*u))^(-3/2)*cos(u) # integrand in (2)
        K=quadgk(u->k(u,a,b,v),0,Inf)[1] # numerical quadrature integration of (2)

        #Right-hand-side of Eqn 1 in Hellwarth 1999 // Eqn (4) in Baggio1997
        RHS= α/(3*sqrt(π)) * βred^(5/2) / sinh(βred/2) * (v^3/w^3) * K
        μ=RHS^-1 * (q)/(ω*mb)
        @printf("\n\tμ(Hellwarth1999,b=0)= %f m^2/Vs \t= %.2f cm^2/Vs",μ,μ*100^2)
        @printf("\n\tError due to b=0; %f",(100^2*μ-p.Hμ[length(p.Hμ)])/(100^2*μ))
        #append!(Hμs,μ*100^2)
        end
        @printf("\n") # blank line at end of spiel.

        # Recycle previous variation results (v,w) as next guess
        initial=[v,w] # Caution! Might cause weird sticking in local minima

        append!(p.v,v)
        append!(p.w,w)

    end

    return(p)
end

"""
    make_polaron(ϵ_optic, ϵ_static, phonon_freq, m_eff; temp = 300.0, efield_freq = 0.0, volume = nothing, ir_activity = nothing, N_params = 1, rtol = 1e-3, verbose = false, threads = false)

Solves the Feynman polaron problem variationally with finite temperature
Osaka energies. From the resulting v and w parameters, calculates polaron
structure (wave function size, etc.).

Uses FHIP to predict a temperature- and frequency-dependent mobility, complex conductivity and impedance for the material system.

Can evaluate polaron properties for materials with multiple phonon branches using infrared activities and phonon branch frequencies.

# Arguments
- `ϵ_optic`: reduced optical dielectric constant.
- `ϵ_static`: reduced static dielectric constant.
- `phonon_freq`: vector of characteristic dielectric phonon frequencies (THz).
- `m_eff`: bare-band effective-mass (mₑ).
- `Trange`: temperature value or range.
- `Ω`: electric field frequency value or range (THz).
- `volume`: unit cell volume for the material (m³). `nothing` for one phonon mode.
- `ir_activity`: vector of infrared activities. `nothing` for one phonon  mode.
- `N_params`: number of variational parameters to minimise the polaron energy.
- `rtol`: relative tolerance for the accuracy of any quadrature integrations.
- 'verbose': `true` to print progress. `false` to print nothing.
- 'threads'; `true` activates multithreading through use of @tullio macro. Tullio utilises einsum notation.

Returns a structure of type NewPolaron, containing arrays of useful
information.  Also prints a lot of information to the standard out - which
may be more useful if you're just inquiring as to a particular data point,
rather than plotting a temperature-dependent parameter.

As an example, to calculate the electron polaron in MAPI, at temperatures 0:100:400 K and electric field frequencies 0.0:0.1:5.0 THz, and inclusive of 15 optical phonon modes:

# Examples
```jldoctest
make_polaron(
    4.5,
    24.1,
    [4.016471586720514, 3.887605410774121, 3.5313112232401513, 2.755392921480459, 2.4380741812443247, 2.2490917637719408, 2.079632190634424, 2.0336707697261187, 1.5673011873879714, 1.0188379384951798, 1.0022960504442775, 0.9970130778462072, 0.9201781906386209, 0.800604081794174, 0.5738689505255512],
    0.12;
    temp = 0.0:100.0:400.0,
    efield_freq = 0.0:0.1:5.0,
    volume = (6.29e-10)^3,
    ir_activity = [0.08168931020200264, 0.006311654262282101, 0.05353548710183397, 0.021303020776321225, 0.23162784335484837, 0.2622203718355982, 0.23382298607799906, 0.0623239656843172, 0.0367465760261409, 0.0126328938653956, 0.006817361620021601, 0.0103757951973341, 0.01095811116040592, 0.0016830270365341532, 0.00646428491253749],
    N_params = 1)
```

"""
function make_polaron(ϵ_optic, ϵ_static, phonon_freq, m_eff, Trange, Ω; volume = nothing, ir_activity = nothing, rtol = 1e-4, N_params = 1, verbose = false, threads = false)

    # Number of phonon modes.
    N_modes = length(phonon_freq)

    if verbose
        println("Calculating Fröhlich alpha...")
    end

    # Convert THz phonon frequencies to 2π THz.
    ω = 2π .* phonon_freq

    if N_modes == 1
        # One phonon mode.

        # Calculate Frohlich alpha coupling parameters.
        α = frohlichalpha(ϵ_optic, ϵ_static, phonon_freq, m_eff)
    else
        # Multiple phonon modes

        # Calculate contribution to the ionic delectric constant for each phonon mode.
        ϵ_ionic = [ϵ_ionic_mode(i, j, volume) for (i, j) in zip(phonon_freq, ir_activity)]

        # Calculate contribution to Frohlich alpha for each phonon mode.
        α = [multi_frohlichalpha(ϵ_optic, i, sum(ϵ_ionic), j, m_eff) for (i, j) in zip(ϵ_ionic, phonon_freq)]
    end

    if verbose
        show(IOContext(stdout, :limit => true), round.(α, digits = 3))
        print("\n\n")
        println("Calculating thermodynamic betas...")
    end

    # Calculate reduced thermodynamic betas for each phonon mode.
    # If temperature is zero, return Inf.
    @tullio threads = threads betas[m, i] := Trange[i] == 0.0 ? Inf64 : ħ * ω[m] / (kB * Trange[i]) * 1e12 (i in eachindex(Trange))

    if verbose
        show(IOContext(stdout, :limit => true), round.(betas, digits = 3))
        print("\n\n")
        println("Calculating variational parameters...")
        global count = 0
        global processes = length(Trange)
    end

    # Calculate variational parameters for each temperature from multiple phonon frequencies.
    @tullio threads = threads params[i] := Trange[i] == 0.0 ? var_params(α; v = 5.6, w = 3.4, ω = ω, rtol = rtol, N = N_params, verbose = verbose) : var_params(α, betas[:, i]; v = 5.6, w = 3.4, ω = ω, rtol = rtol, T = Trange[i], N = N_params, verbose = verbose) (i in eachindex(Trange))

    # Separate tuples of variational parameters into a list of 'v' and 'w' parameters for each temperature.
    @tullio threads = threads v_params[i] := params[i][1] (i in eachindex(Trange))
    @tullio threads = threads w_params[i] := params[i][2] (i in eachindex(Trange))

    if verbose
        println("\nv: ")
        show(IOContext(stdout, :limit => true), round.(v_params, digits = 3))
        print("\n\n")
        println("w: ")
        show(IOContext(stdout, :limit => true), round.(w_params, digits = 3))
        print("\n\n")
        println("Calculating spring constants...")
    end

    # Calculate fictitious spring constants for each temperature.
    @tullio threads = threads spring_constants[i] := v_params[i]^2 - w_params[i]^2

    if verbose
        show(IOContext(stdout, :limit => true), round.(spring_constants, digits = 3))
        print("\n\n")
        println("Calculating fictitious masses...")
    end

    # Calculate fictitious masses for each temperature.
    @tullio threads = threads masses[i] := spring_constants[i] / w_params[i]^2

    if verbose
        show(IOContext(stdout, :limit => true), round.(masses, digits = 3))
        print("\n\n")
        println("Calculating free energies...")
        global count = 0
        global processes = length(Trange)
    end

    # Calculate ground-state free energies for each temperature.
    @tullio threads = threads energies[i] := Trange[i] == 0.0 ? multi_F(v_params[i], w_params[i], α; ω = ω, rtol = rtol, verbose = verbose) * 1000 * ħ / eV * 1e12 : multi_F(v_params[i], w_params[i], α, betas[:, i]; ω = ω, rtol = rtol, T = Trange[i],  verbose = verbose) * 1000 * ħ / eV * 1e12 (i in eachindex(Trange))

    if verbose
        show(IOContext(stdout, :limit => true), round.(energies, digits = 3))
        print("\n\n")
        println("Calculating complex impedances...")
        global count = 0
        global processes = length(Trange) * length(Ω)
    end

    # Calculate complex impedances for each frequency and temperature. Returns a matrix.
    @tullio threads = threads impedances[f, i] := Ω[f] == Trange[i] == 0.0 ? 0.0 + 1im * 0.0 : polaron_complex_impedence(Ω[f], betas[:, i], α, v_params[i], w_params[i]; ω = ω, rtol = rtol, T = Trange[i], verbose = verbose) / eV^2 * 1e12 * me * m_eff * volume * 100^3 (i in eachindex(Trange), f in eachindex(Ω))

    if verbose
        show(IOContext(stdout, :limit => true), round.(impedances, digits = 3))
        print("\n\n")
        println("Calculating complex conductivities...")
    end

    # Calculate complex conductivities for each frequency and temperature. Returns a matrix.
    @tullio threads = threads conductivities[f, i] := Ω[f] == Trange[i] == 0.0 ? Inf64 + 1im * 0.0 : 1 / impedances[f, i] (i in eachindex(Trange), f in eachindex(Ω))

    if verbose
        show(IOContext(stdout, :limit => true), round.(conductivities, digits = 3))
        print("\n\n")
        println("Calculating mobilities...")
        global count = 0
        global processes = length(Trange)
    end

    # Calculates the dc mobility for each temperature.
    @tullio threads = threads mobilities[i] := Trange[i] == 0.0 ? Inf64 : polaron_mobility(betas[:, i], α, v_params[i], w_params[i]; ω = ω, rtol = rtol, T = Trange[i], verbose = verbose) * eV / (1e12 * me * m_eff) * 100^2 (i in eachindex(Trange))

    if verbose
        show(IOContext(stdout, :limit => true), round.(mobilities, digits = 3))
        print("\n")
    end

    # Return Polaron mutable struct with evaluated data.
    return NewPolaron(α, Trange, betas, phonon_freq, v_params, w_params, spring_constants, masses, energies, Ω, impedances, conductivities, mobilities)
end

"""
    make_polaron(α, Trange, Ω; ω = 1.0, rtol = 1e-4, verbose = false, threads = false)

Same as above but from a model system with specified alpha values rather than from material properties. Here we only have one phonon mode with a normalised frequency `ω = 1.0`.
"""
function make_polaron(α, Trange, Ω; ω = 1.0, rtol = 1e-4, verbose = false, threads = false)

    if verbose
        println("\nCalculating thermodynamic betas...")
    end

    # Calculate reduced thermodynamic betas for each phonon mode.
    # If temperature is zero, return Inf.
    @tullio threads = threads betas[i] := Trange[i] == 0.0 ? Inf64 : ω / Trange[i]

    if verbose
        show(IOContext(stdout, :limit => true), round.(betas, digits = 3))
        print("\n\n")
        println("Calculating variational parameters...")
        global count = 0
        global processes = length(Trange) * length(α)
    end

    # Calculate variational parameters for each alpha parameter and temperature. Returns a Matrix of tuples.
    @tullio threads = threads params[j, i] := Trange[i] == 0.0 ? var_params(α[j]; v = 5.6, w = 3.4, ω = ω, rtol = rtol, verbose = verbose) : var_params(α[j], betas[i]; v = 5.6, w = 3.4, ω = ω, rtol = rtol, T = Trange[i], verbose = verbose) (i in eachindex(Trange), j in eachindex(α))

    # Separate tuples of variational parameters into a Matrices of 'v' and 'w' parameters for each alpha parameter and temperature.
    @tullio threads = threads v_params[j, i] := params[j, i][1]
    @tullio threads = threads w_params[j, i] := params[j, i][2]

    if verbose
        println("\nv: ")
        show(IOContext(stdout, :limit => true), round.(v_params, digits = 3))
        print("\n\n")
        println("w: ")
        show(IOContext(stdout, :limit => true), round.(w_params, digits = 3))
        print("\n\n")
        println("Calculating spring constants...")
    end

    # Calculate fictitious spring constants for each alpha parameter and temperature. Returns a Matrix.
    @tullio threads = threads spring_constants[j, i] := v_params[j, i]^2 - w_params[j, i]^2

    if verbose
        show(IOContext(stdout, :limit => true), round.(spring_constants, digits = 3))
        print("\n\n")
        println("Calculating fictitious masses...")
    end

    # Calculate fictitious masses for each alpha parameter and temperature. Returns a Matrix.
    @tullio threads = threads masses[j, i] := spring_constants[j, i] / w_params[j, i]^2

    if verbose
        show(IOContext(stdout, :limit => true), round.(masses, digits = 3))
        print("\n\n")
        println("Calculating free energies...")
        global count = 0
        global processes = length(T) * length(α)
    end

    # Calculate ground-state free energies for each alpha parameter and temperature. Returns a Matrix.
    @tullio threads = threads energies[j, i] := Trange[i] == 0.0 ? multi_F(v_params[j, i], w_params[j, i], α[j]; ω = ω, rtol = rtol, verbose = verbose) : multi_F(v_params[j, i], w_params[j, i], α[j], betas[i]; ω = ω, rtol = rtol, T = Trange[i],  verbose = verbose) (i in eachindex(Trange), j in eachindex(α))

    if verbose
        show(IOContext(stdout, :limit => true), round.(energies, digits = 3))
        print("\n\n")
        println("Calculating complex impedances...")
        global count = 0
        global processes = length(T) * length(Ω) * length(α)
    end

    # Calculate complex impedances for each alpha parameter, frequency and temperature. Returns a 3D Array.
    @tullio threads = threads impedances[k, j, i] := Ω[j] == Trange[i] == 0.0 ? 0.0 + 1im * 0.0 : polaron_complex_impedence(Ω[j], [betas[i]], α[k], v_params[k, i], w_params[k, i]; ω = ω, rtol = rtol, T = Trange[i],  verbose = verbose) (i in eachindex(Trange), j in eachindex(Ω), k in eachindex(α))

    if verbose
        show(IOContext(stdout, :limit => true), round.(impedances, digits = 3))
        print("\n\n")
        println("Calculating complex conductivities...")
    end

    # Calculate complex conductivities for each alpha parameter, frequency and temperature. Returns a 3D array.
    @tullio threads = threads conductivities[k, j, i] := Ω[j] == Trange[i] == 0.0 ? Inf64 + 1im * 0.0 : 1 / impedances[k, j, i] (i in eachindex(Trange), j in eachindex(Ω), k in eachindex(α))

    if verbose
        show(IOContext(stdout, :limit => true), round.(conductivities, digits = 3))
        print("\n\n")
        println("Calculating mobilities...")
        global count = 0
        global processes = length(T) * length(α)
    end

    # Calculates the dc mobility for each alpha parameter and each temperature.
    @tullio threads = threads mobilities[j, i] := Trange[i] == 0.0 ? Inf64 : polaron_mobility(betas[i], α[j], v_params[j, i], w_params[j, i]; ω = ω, rtol = rtol, T = Trange[i], verbose = verbose) (i in eachindex(Trange), j in eachindex(α))

    if verbose
        show(IOContext(stdout, :limit => true), round.(mobilities, digits = 3))
        print("\n")
    end

    # Return Polaron mutable struct with evaluated data.
    return NewPolaron(α, Trange, betas, ω, v_params, w_params, spring_constants, masses, energies, Ω, impedances, conductivities, mobilities)
end

"""
    save_polaron(p::NewPolaron, prefix)

Saves data from 'polaron' into file "prefix".
This is a .jdl file for storing the polaron data whilst preserving types. Allows for saving multidimensional arrays that sometimes arise in the polaron data.
Each parameter in the NewPolaron type is saved as a dictionary entry. E.g. NewPolaron.α is saved under JLD.load("prefix.jld")["alpha"].
"""
function save_polaron(polaron::NewPolaron, prefix)

    println("Saving polaron data to $prefix.jld ...")

    JLD.save("$prefix.jld",
    "alpha", polaron.α,
    "temperature", polaron.T,
    "beta", polaron.β,
    "phonon freq", polaron.ω,
    "v", polaron.v,
    "w", polaron.w,
    "spring", polaron.κ,
    "mass", polaron.M,
    "energy", polaron.F,
    "efield freq", polaron.Ω,
    "impedance", polaron.Z,
    "conductivity", polaron.σ,
    "mobility", polaron.μ
    )

    println("... Polaron data saved.")
end

"""
    load_polaron(p::NewPolaron, prefix)

Loads data from file "polaron_file_path" into a NewPolaron type.
"""
function load_polaron(polaron_file_path)

    println("Loading polaron data from $polaron_file_path ...")

    data = JLD.load("$polaron_file_path")

    polaron = NewPolaron(
        data["alpha"],
        data["temperature"],
        data["beta"],
        data["phonon freq"],
        data["v"],
        data["w"],
        data["spring"],
        data["mass"],
        data["energy"],
        data["efield freq"],
        data["impedance"],
        data["conductivity"],
        data["mobility"]
    )

    println("... Polaron loaded.")

    return polaron
end

"""
    savepolaron(fileprefix, p::Polaron)

Saves data from polaron 'p' into file "fileprefix".
This is a simple space-delimited text file, with each entry a separate temperature, for plotting with Gnuplot or similar.

Structure of file is written to the header:
# Ts, βreds, Kμs, Hμs, FHIPμs, vs, ws, ks, Ms, As, Bs, Cs, Fs, Taus, rfsis
# 1    2     3    4     5      6   7   8   9  10  11  12  13    14     15
"""
function savepolaron(fileprefix, p::Polaron)
    println("Saving data to $fileprefix.dat ...")
    f=open("$fileprefix.dat","w")

    @printf(f,"# %s\n",fileprefix) # put name / material at header
    @printf(f,"# Params in SI: ω =%g mb=%g \n",p.ω[1] ,p.mb[1])
    @printf(f,"# Alpha parameter: α = %f  \n",p.α[1] )

    @printf(f,"# Ts, βreds, Kμs, Hμs, FHIPμs, vs, ws, ks, Ms, As, Bs, Cs, Fs, Taus, rfsis\n")
    @printf(f,"#  1    2     3    4     5      6   7   8   9  10  11  12  13    14     15\n") # columns for GNUPLOT etc.

    for i in 1:length(p.T)
        @printf(f,"%d %03f %g %g %g %g %g %g %g %g %g %g %g %g %g \n",
        p.T[i], p.βred[i], p.Kμ[i], p.Hμ[i], p.FHIPμ[i],
        p.v[i], p.w[i],
        p.k[i], p.M[i], p.A[i], p.B[i], p.C[i], p.F[i],
        p.Tau[i], p.rfsi[i])
    end
    close(f)
end


"""
    ExpMethodHigham2005(A::AbstractMatrix);
    ExpMethodHigham2005(b::Bool=true);

Computes the matrix exponential using the algorithm Higham, N. J. (2005). "The scaling and squaring method for the matrix exponential revisited." SIAM J. Matrix Anal. Appl.Vol. 26, No. 4, pp. 1179–1193" based on generated code. If a matrix is specified, balancing is determined automatically.
"""
struct ExpMethodHigham2005
    do_balancing::Bool
end
ExpMethodHigham2005(A::AbstractMatrix)=ExpMethodHigham2005(A isa StridedMatrix)
ExpMethodHigham2005()=ExpMethodHigham2005(true)

function alloc_mem(A,::ExpMethodHigham2005)
    return [similar(A) for i=1:5];
end

# Import the generated code
for i=1:13
    include("exp_generated/exp_$i.jl")
end

function getmem(cache,k) # Called from generated code
    return cache[k-1];
end
function ldiv_for_generated!(C,A,B) # C=A\B. Called from generated code
    F= lu!(A);
    ldiv!(F,B); # Result stored in B
    if (pointer_from_objref(C) != pointer_from_objref(B)) # Aliasing allowed
        copyto!(C,B)
    end
    return C
end

# Inplace add of a UniformScaling object (support julia 1.6.2)
@inline function inplace_add!(A,B::UniformScaling) # Called from generated code
    s = B.λ
    @inbounds for i in diagind(A)
        A[i] += s
    end
end
function exponential!(A,method::ExpMethodHigham2005,cache=alloc_mem(A,method))
    n = LinearAlgebra.checksquare(A)
    nA = opnorm(A,1);

    # Maybe to balancing
    if method.do_balancing
        ilo, ihi, scale = LAPACK.gebal!('B', A)    # modifies A
    end

    # Select how many multiplications to use
    rhov=    [0.015; 0.25; 0.95; 2.1; 5.4];
    # Number of memslots needed (beside A)
    memslots=[3   ;    4;    5;   5;   5];
    for s=1:7 # Only 8 since exp(5.4*2^8)=Inf
        push!(rhov, rhov[end]*2);
        push!(memslots,memslots[end]);
    end
    i = findfirst(nA .< rhov)

    # Make the call to the appropriate exp_gen! function
    X = Base.Cartesian.@nif 13 d -> begin
        nA < rhov[d]
    end d -> begin # if condition
        exp_gen!(cache, A, Val(d))
    end d -> begin # fallback (d == 13)
        exp_gen!(cache, A, Val(d))
    end

    # Undo the balancing
    if method.do_balancing
        for j = ilo:ihi
            scj = scale[j]
            for i = 1:n
                X[j,i] *= scj
            end
            for i = 1:n
                X[i,j] /= scj
            end
        end

        if ilo > 1       # apply lower permutations in reverse order
            for j in (ilo-1):-1:1; LinearAlgebra.rcswap!(j, Int(scale[j]), X) end
        end
        if ihi < n       # apply upper permutations in forward order
            for j in (ihi+1):n;    LinearAlgebra.rcswap!(j, Int(scale[j]), X) end
        end
    end

    return X

end


@testset "Distance" begin
    a = randn(Float,10)
    b = randn(Float,10)
    c1 = [a b a b]
    c2 = [a b]
    c3 = [a b b b]

    @test VectorDataUtils.sql2 <: VectorDataUtils.Distance
    @test VectorDataUtils.l2 <: VectorDataUtils.Distance
    @test VectorDataUtils.l1 <: VectorDataUtils.Distance

    @test VectorDataUtils.evaluate(VectorDataUtils.sql2(),a,b) ≈ VectorDataUtils.sql2dist(a,b)
    @test VectorDataUtils.evaluate(VectorDataUtils.l2(),a,b) ≈ VectorDataUtils.l2dist(a,b)
    @test VectorDataUtils.evaluate(VectorDataUtils.l1(),a,b) ≈ VectorDataUtils.l1dist(a,b)

    distances = (VectorDataUtils.sql2(), VectorDataUtils.l2(), VectorDataUtils.l1())
    for dm ∈ distances

        d = VectorDataUtils.colwise(dm,c1,c3)
        @test length(d) == size(c1,2)
        @test d[1] == 0.0
        @test d[2] == 0.0
        @test d[3] == VectorDataUtils.evaluate(dm,a,b)

        d = VectorDataUtils.colwise(dm,a,c1)
        @test length(d) == size(c1,2)
        @test d[1] == 0.0
        @test d[2] == VectorDataUtils.evaluate(dm,a,b)
        @test d[3] == 0.0
        @test d[4] == d[2]

        d = VectorDataUtils.pairwise(dm,a,c1)
        @test length(d) == size(c1,2)
        @test d[1] == 0.0
        @test d[2] == VectorDataUtils.evaluate(dm, a, b)

        d = VectorDataUtils.pairwise(dm,c1,a)
        @test length(d) == size(c1,2)
        @test d[3] == 0.0
        @test d[4] == VectorDataUtils.evaluate(dm, a, b)

        dd = VectorDataUtils.pairwise(dm, c2, c1)
        @test size(dd,1) == size(c2,2)
        @test size(dd,2) == size(c1,2)
        @test dd[1,1] == 0.0
        @test dd[2,2] == 0.0
        @test dd[1,2] == VectorDataUtils.evaluate(dm, a, b)
        @test dd[2,1] == VectorDataUtils.evaluate(dm, b, a)
    end
end


#
# Proof of concept using artificial data.
# Illustrate usefulness of the PCCF on artificially generated signal.
#

include("./lib/metadata.jl")
include("./2bayes_detector_replica.jl")
include("./lib/SignalGenerator.jl")
include("./lib/roc.jl")
include("./lib/pccf.jl")

using SignalGenerator
using BayesChpOrig
using Metadata
using Roc
using Pccf

function run_sig_gen()
    ret = genSig(5, 100.0, 1.0, 2.0, 0.5, 1)
    sig = ret["sig"]
    chps= ret["changes"]
    writedlm(outdir*"testsig.txt", sig)
    writedlm(outdir*"testsigchps.txt", chps)
    print("run1")
end

function run_proof3(outputfolder)
    println("... run_proof3 > output folder = ", outputfolder)
    # genSig(n, mu, sigma, changeStep, sigStd)
    mu    = 100.0
    sigma = 10.0
    srand(1234)
    ret   = genSig(10, mu, sigma, 1.3, 0.75)
    sig   = ret["sig"]
    n     = length(sig)
    chps  = ret["changes"]
    probs = scale01(pccf(n, mu, sigma))
    # pccf_matrix_max_vals, pccf_matrix = pccf_bayes_detector(sig, 200 * probs, 200, 200 + 30, 0.15) # one outlier less
    # start INTERESTING EDGE CASE
    #pccf_matrix_max_vals, pccf_matrix = pccf_bayes_detector(sig, 200 * probs, 200, 200 + 150, 0.3) #150-max,min-110.
    # -- end INTERESTING EDGE CASE
    # start One outlier is removed case
    pccf_matrix_max_vals, pccf_matrix = pccf_bayes_detector(sig, probs, 200+100, 200, 0.3)
    # -- end One outlier is removed case
    matrix_max_vals, matrix           = bayes_detector(sig, lambda = 200)
    detections                        = extract_changes(matrix_max_vals)
    pccf_detections                   = extract_changes(pccf_matrix_max_vals)
    #
    writedlm(outputfolder * "pccf.txt", probs)
    writedlm(outputfolder * "changes.txt", chps)
    writedlm(outputfolder * "signal.txt", sig)
    writedlm(outputfolder * "detections.txt", detections)
    writedlm(outputfolder * "pccf_detections.txt", pccf_detections)
    if 1 == 0
        max_delay = 30
        tpr = tpRate(detections, chps, max_delay)
        fpr = faRateScaled(detections, chps, max_delay, length(sig), tn_price_scale=0.1)
        println("... test_bd > TPrate = ", tpr, " / FArate = ", fpr)
        ##writedlm(outdir*"concept_proof/matrix.txt", matrix, ",")
        ##writedlm(outdir*"concept_proof/matrix_max_vals.txt", matrix_max_vals, ",")
        # println("... test_bd > changes:", chps)
        # println("... test_bd > detections:", detections)
    end
end

run_proof3(outdir * "concept_proof3/")


using FuzzySets, Test

step_size = 0.005
levels = collect(0:step_size:1)

A₁ = FuzzyNumber(levels, number=1)
A₂ = FuzzyNumber(levels, number=2)
B₁ = FuzzyNumber(levels, number=0)
B₂ = FuzzyNumber(levels, number=0)

A⃗ = FuzzyVector([A₁, A₂])
@test length(A⃗) == 2
@test A⃗[1] == A₁
@test A⃗[2] == A₂
@test iterate(A⃗, 1) == A₁
@test iterate(A⃗, 2) == A₂
@test isnothing(iterate(A⃗, 3))

B⃗ = FuzzyVector([B₁])
@test length(B⃗) == 1
@test B⃗[1] == B₁
B⃗ = B⃗ ∪ [B₂]
@test length(B⃗) == 2
@test B⃗[1] == B₁
@test B⃗[2] == B₂

@test A⃗ == A⃗
@test A⃗ == FuzzyVector([A₁, A₂])
@test B⃗ != FuzzyVector([A₁, A₂])

@test cut(A⃗, 0) == [Interval(0.5, 1.5), Interval(1.5, 2.5)]
@test cut(A⃗, 1) == [Interval(1), Interval(2)]

C₁ = FuzzyNumber(levels, number=-1)
C₂ = FuzzyNumber(levels, number=3)
C⃗ = FuzzyVector([C₁, C₂])

patterns = [A⃗, C⃗]
@test cut(patterns, 1, 1) == [Interval(1), Interval(-1)]
@test cut(patterns, 2, 1) == [Interval(2), Interval(3)]


struct BinnedMarginalDensity # this encodes the estimated density..
    marginal::Vector{Float64} # measure after truncation & binning
    marginal_grid::Vector{Float64}
    marginal_h::Float64
end

function BinnedMarginalDensity(d::NormalConvolutionProblem)
    BinnedMarginalDensity(d.marginal, d.marginal_grid, d.marginal_h)
end

function BinnedMarginalDensity(M_bd::Float64, marginal_grid, marginal_h)
    BinnedMarginalDensity(fill(M_bd, size(marginal_grid)), marginal_grid, marginal_h);
end


using DeltaSugar, Test, Flux

function reverse_grad_tester(d0,d1,d2)
    return all(-(d0.weight - d1.weight) .≈ (d0.weight - d2.weight)) &&
        all(-(d0.bias - d1.bias) .≈ (d0.bias - d2.bias))
end

loss( x1, x2, m ) = sum( abs2, ( m( x1 ) - x2 ) .^ 2 )
mloss( x1, x2, m ) = sum( abs2, ( m( x1 ) - x2 ) .^ 2 ) / length(x2)

@testset "TenseLayer" begin
    a = randn( 10, 6 )
    b = randn( 10, 3 )
    model = TenseLayer( 6, 3 )
    @test all(size( model( a ) ) .== ( 10, 3 ) )

    #ensure the gradient moves
    w0 = deepcopy(model.weight)
    train_light!(a, b, model, Flux.params(model), loss, Descent(0.01))
    @test sum(abs, w0 - model.weight) > 0.0
end

@testset "TenseLayer - Reverse Gradient Check" begin
    a,b = randn( 10, 6 ), randn( 10, 3 )
    original      = TenseLayer( 6, 3 )
    model_forward = deepcopy(original)
    model_reverse = RevGrad( deepcopy(original) )
    @test all(size( model_forward( a ) ) .== ( 10, 3 ) )
    @test all(size( model_reverse( a ) ) .== ( 10, 3 ) )
    #ensure the gradient moves
    ps = Flux.params( model_forward )
    train_light!(a,b, model_forward, ps, loss, Descent(1.))
    ps = Flux.params( model_reverse )
    train_light!(a,b, model_reverse, ps, loss, Descent(1.))
    @test reverse_grad_tester(original, model_forward, model_reverse)
    #now a 2 layer model check...
    a, b = randn( 10, 6 ), randn( 10 )
    layer1, layer2 = TenseLayer( 6, 3 ), TenseLayer( 3, 1 )
    model_forward = Chain( deepcopy(layer1), deepcopy(layer2) )
    model_reverse = Chain( RevGrad(layer1), RevGrad(layer2) )
    #ensure the gradient moves
    ps = Flux.params( model_forward )
    train_light!(a,b, model_forward, ps, loss, Descent(1.))
    ps = Flux.params( model_reverse )
    train_light!(a,b, model_reverse, ps, loss, Descent(1.))
    @test reverse_grad_tester(layer1, model_forward[1], model_reverse[1])
    @test reverse_grad_tester(layer2, model_forward[2], model_reverse[2])

    #test that the chain operator works
    layer1, layer2 = TenseLayer( 6, 3 ), TenseLayer( 3, 1 )
    model_forward = Chain( deepcopy(layer1), deepcopy(layer2) )
    model_reverse = RevGrad( model_forward )
    ps = Flux.params( model_forward )
    train_light!(a,b, model_forward, ps, loss, Descent(1.))
    ps = Flux.params( model_reverse )
    train_light!(a,b, model_reverse, ps, loss, Descent(1.))
    @test reverse_grad_tester(layer1, model_forward[1], model_reverse[1])
    @test reverse_grad_tester(layer2, model_forward[2], model_reverse[2])
end

@testset "DenseLayer" begin
    a = randn( 6, 10 )
    b = randn( 3, 10 )
    model = DenseLayer( 6, 3 )
    @test all(size( model( a ) ) .== ( 3, 10 ) )
    #ensure the gradient moves
    w0 = deepcopy(model.weight)
    train_light!(a, b, model, Flux.params(model), loss, Descent(0.01))
    @test sum(abs, w0 - model.weight) > 0.0
end

@testset "DenseLayer - Reverse Gradient Check" begin
    a = randn( 6, 10 )
    b = randn( 3, 10 )
    original      = DenseLayer( 6, 3 )
    model_forward = deepcopy(original)
    model_reverse = RevGrad( deepcopy(original) )
    @test all(size( model_forward( a ) ) .== ( 3, 10 ) )
    @test all(size( model_reverse( a ) ) .== ( 3, 10 ) )
    #ensure the gradient moves
    ps = Flux.params( model_forward )
    train_light!(a,b, model_forward, ps, loss, Descent(1.))
    ps = Flux.params( model_reverse )
    train_light!(a,b, model_reverse, ps, loss, Descent(1.))
    @test reverse_grad_tester(original,model_forward,model_reverse)

    #now a 2 layer model check...
    a, b = randn( 6, 10 ), randn( 1, 10 )
    layer1, layer2 = DenseLayer( 6, 3 ), DenseLayer( 3, 1 )
    model_forward = Chain( deepcopy(layer1), deepcopy(layer2) )
    model_reverse = Chain( RevGrad(layer1), RevGrad(layer2) )
    #ensure the gradient moves
    ps = Flux.params( model_forward )
    train_light!(a,b, model_forward, ps, loss, Descent(1.))
    ps = Flux.params( model_reverse )
    train_light!(a,b, model_reverse, ps, loss, Descent(1.))
    @test reverse_grad_tester(layer1, model_forward[1], model_reverse[1])
    @test reverse_grad_tester(layer2, model_forward[2], model_reverse[2])

    #test that the chain operator works
    layer1, layer2 = DenseLayer( 6, 3 ), DenseLayer( 3, 1 )
    model_forward = Chain( deepcopy(layer1), deepcopy(layer2) )
    model_reverse = RevGrad( model_forward )
    ps = Flux.params( model_forward )
    train_light!(a,b, model_forward, ps, loss, Descent(1.))
    ps = Flux.params( model_reverse )
    train_light!(a,b, model_reverse, ps, loss, Descent(1.))
    @test reverse_grad_tester(layer1, model_forward[1], model_reverse[1])
    @test reverse_grad_tester(layer2, model_forward[2], model_reverse[2])
end

@testset "Gradient Clip Test" begin
    a = randn( 6, 10 )
    b = randn( 3, 10 )
    ϵ = 1e-2
    model = GradientClamp( DenseLayer( 6, 3 ), ϵ )
    @test all(size( model( a ) ) .== ( 3, 10 ) )
    #ensure the gradient is clipped
    ps = Flux.Params(Flux.params(model))
    gs = Flux.gradient(ps) do
        loss(a, b, model)
    end
    @test all( [ reduce(max, g) <= ϵ for (_,g) in gs.grads ])
    @test all( [ reduce(min, g) >= -ϵ for (_,g) in gs.grads ])

    ϵ = 100.0
    model = GradientClamp( DenseLayer( 6, 3 ), ϵ )
    @test all(size( model( a ) ) .== ( 3, 10 ) )
    #ensure the gradient is clipped
    ps = Flux.Params(Flux.params(model))
    gs = Flux.gradient(ps) do
        loss(a, b, model)
    end
    #println.([ g for (_,g) in gs.grads ])
    @test all( [ reduce(max, g) <= ϵ for (_,g) in gs.grads ])
    @test all( [ reduce(min, g) >= -ϵ for (_,g) in gs.grads ])
end

@testset "Ridge Test" begin
    a = randn( 10, 20 )
    b = randn( 1, 20 )
    a[5,:]  = 0.5 * b[:]
    model = L2( DenseLayer( 10, 1 ), 0.02 )
    weight_i = deepcopy( model.weight )
    @test all(size( model( a ) ) .== ( 1, 20 ) )
    ps = Flux.params( model )
    for _ in 1:100
        train_light!( a, b, model, ps, mloss, Descent( 0.01 ) )
    end
    @test sum(abs2, model.weight ) < sum(abs2, weight_i )
end

@testset "FF Chain Sugar Tests" begin
    z = PlainChain( 5,4,3,2,1 ; σ = identity )
    @test all(size(z(randn(5,10))) .== (1,10))

    x = TChain( 5,4,3,2,1 ; σ = identity )
    @test all(size(x(randn(10,5))) .== (10,1))
end

@testset "LASSO Test" begin
    """
    This one won't always work, c'est la vie...
    """
    a = randn( 10, 20 )
    b = randn( 1, 20 )
    a[5,:]  = 0.5 * b[:]
    model = L1( DenseLayer( 10, 1 ), 0.00001 )
    #model.weight[1,5] .+= 0.1
    @test all(size( model( a ) ) .== ( 1, 20 ) )
    ps = Flux.params( model )
    for _ in 1:10000
        train_light!( a, b, model, ps, mloss, Descent( 0.001 ) )
    end
    #println( model.weight )
    @test argmax( model.weight[:] ) == 5
end


using Test
using Random

include("yahtzee.jl")

@testset "test roll_dice" begin
    # fix random seed
    Random.seed!(1234);
    # known sequence for given seed
    exp_first_throw = [1, 3, 5, 1, 6]
    exp_second_throw = [1, 6, 4, 5, 4]
    exp_third_throw = [1, 2, 5, 3, 2]

    n_dice = 5
    sides=(1,2,3,4,5,6)
    throw = roll_dice(n_dice, sides)
    @test throw == exp_first_throw
    throw = roll_dice(n_dice, sides)
    @test throw == exp_second_throw
    throw = roll_dice(n_dice, sides)
    @test throw == exp_third_throw
end

@testset "testing yahtzee" begin
    # try a million throws and see if we get close
    # to the statistical average of 4.6%
    # (https://en.wikipedia.org/wiki/Yahtzee#Probabilities)
    n_yahtzees = 0
    n_tests = 1_000_000
    for i in 1:n_tests
        collected_side, n_collected = yahtzee()
        if n_collected == 5
            n_yahtzees += 1
        end
    end
    @test n_yahtzees/n_tests ≈ 0.046 atol=0.001
end


#=
# Learning rate decay

Effect of decaying the learning rate during training to achieve convergence.
=#

using Statistics: mean, std, var
using Random: bitrand
using LinearAlgebra: dot
using ValueHistories: MVHistory
import Makie
import CairoMakie
import Flux
import MLDatasets
using RestrictedBoltzmannMachines: BinaryRBM, initialize!, pcd!, log_pseudolikelihood, minibatch_count
using RestrictedBoltzmannMachines: free_energy, sample_v_from_v, mean_h_from_v, default_optimizer

#=
Load MNIST dataset. We select only 0 digits and binarize pixel intensities.
=#

Float = Float32
train_x, train_y = MLDatasets.MNIST.traindata()
train_x = Array{Float64}(train_x[:,:, train_y .== 0] .> 0.5)
nsamples = size(train_x)[end]
nothing #hide

# Some hyper-parameters

nh = 100 # number of hidden units
epochs = 1000
batchsize = 128
batchcount = minibatch_count(train_x; batchsize)
nothing #hide

#=
Consider first an RBM that we train without decaying the learning rate.
=#

rbm_nodecay = initialize!(BinaryRBM(Float, (28,28), nh), train_x)
history_nodecay = MVHistory()
@time pcd!(
    rbm_nodecay, train_x; epochs, batchsize,
    optim = default_optimizer(nsamples, batchsize, epochs; decay_after=1),
    callback = function callback(; rbm, epoch, batch_idx, _...)
        if batch_idx == batchcount && epoch % 5 == 0
            push!(history_nodecay, :lpl, mean(log_pseudolikelihood(rbm, train_x)))
        end
    end
)
nothing #hide

#=
Now train an RBM with lr decay after half training (this is the default
behavior).
=#

rbm_decaylr = initialize!(BinaryRBM(Float, (28,28), nh), train_x)
history_decaylr = MVHistory()
@time pcd!(
    rbm_decaylr, train_x; epochs, batchsize,
    callback = function callback(; rbm, epoch, batch_idx, _...)
        if batch_idx == batchcount && epoch % 5 == 0
            push!(history_decaylr, :lpl, mean(log_pseudolikelihood(rbm, train_x)))
        end
    end
)
nothing #hide

#=
Compare the results
=#

fig = Makie.Figure(resolution=(600,400))
ax = Makie.Axis(fig[1,1])
Makie.lines!(ax, get(history_nodecay, :lpl)..., label="normal")
Makie.lines!(ax, get(history_decaylr, :lpl)..., label="decay")
Makie.axislegend(ax, position=:rb)
fig

#=
Check convergence by computing the moment-matching conditions.
First generate MC data from the RBMs.
=#

nsteps = 1000
nsamples = 5000
F_nodecay = zeros(nsamples, nsteps)
F_decaylr = zeros(nsamples, nsteps)
samples_v_nodecay = bitrand(28,28,nsamples)
samples_v_decaylr = bitrand(28,28,nsamples)
F_nodecay[:,1] .= free_energy(rbm_nodecay, samples_v_nodecay)
F_decaylr[:,1] .= free_energy(rbm_decaylr, samples_v_decaylr)
@time for step in 2:nsteps
    samples_v_nodecay .= sample_v_from_v(rbm_nodecay, samples_v_nodecay)
    samples_v_decaylr .= sample_v_from_v(rbm_decaylr, samples_v_decaylr)
    F_nodecay[:,step] .= free_energy(rbm_nodecay, samples_v_nodecay)
    F_decaylr[:,step] .= free_energy(rbm_decaylr, samples_v_decaylr)
end
nothing #hide

# Check equilibration of sampling

# Samples without lr decay

fig = Makie.Figure(resolution=(400,300))
ax = Makie.Axis(fig[1,1])
F_nodecay_μ = vec(mean(F_nodecay; dims=1))
F_nodecay_σ = vec(std(F_nodecay; dims=1))
Makie.band!(ax, 1:nsteps, F_nodecay_μ - F_nodecay_σ/2, F_nodecay_μ + F_nodecay_σ/2)
Makie.lines!(ax, 1:nsteps, F_nodecay_μ, label="no decay")
fig

# Samples with lr decay

fig = Makie.Figure(resolution=(400,300))
ax = Makie.Axis(fig[1,1])
F_decaylr_μ = vec(mean(F_decaylr; dims=1))
F_decaylr_σ = vec(std(F_decaylr; dims=1))
Makie.band!(ax, 1:nsteps, F_decaylr_μ - F_decaylr_σ/2, F_decaylr_μ + F_decaylr_σ/2)
Makie.lines!(ax, 1:nsteps, F_decaylr_μ, label="decay lr")
fig

# Now make the plots. Average digit shapes.

fig = Makie.Figure(resolution=(900, 300))
ax = Makie.Axis(fig[1,1], title="data", yreversed=true)
Makie.heatmap!(ax, mean(train_x, dims=3)[:,:,1])
Makie.hidedecorations!(ax)
ax = Makie.Axis(fig[1,2], title="const. lr", yreversed=true)
Makie.heatmap!(ax, mean(samples_v_nodecay, dims=3)[:,:,1])
Makie.hidedecorations!(ax)
ax = Makie.Axis(fig[1,3], title="lr decay", yreversed=true)
Makie.heatmap!(ax, mean(samples_v_decaylr, dims=3)[:,:,1])
Makie.hidedecorations!(ax)
fig

#=
Moment matching conditions, first row for RBM with constant learning rate,
second row for RBM with learning rate decay.
=#

h_data_nodecay = mean_h_from_v(rbm_nodecay, train_x)
h_data_decaylr = mean_h_from_v(rbm_decaylr, train_x)
h_model_nodecay = mean_h_from_v(rbm_nodecay, samples_v_nodecay)
h_model_decaylr = mean_h_from_v(rbm_decaylr, samples_v_decaylr)
nothing #hide

fig = Makie.Figure(resolution=(900, 600))

ax = Makie.Axis(fig[1,1], xlabel="<v>_data", ylabel="<v>_model", limits=(0,1,0,1))
Makie.scatter!(ax, vec(mean(train_x; dims=3)), vec(mean(samples_v_nodecay; dims=3)))
Makie.abline!(ax, 0, 1; color=:red)

ax = Makie.Axis(fig[1,2], xlabel="<h>_data", ylabel="<h>_model", limits=(0,1,0,1))
Makie.scatter!(ax, vec(mean(h_data_nodecay; dims=2)), vec(mean(h_model_nodecay; dims=2)))
Makie.abline!(ax, 0, 1; color=:red)

ax = Makie.Axis(fig[1,3], xlabel="<vh>_data", ylabel="<vh>_model", limits=(0,1,0,1))
Makie.scatter!(ax,
    vec([dot(train_x[i,j,:], h_data_nodecay[μ,:]) / size(train_x,3) for i=1:28, j=1:28, μ=1:nh]),
    vec([dot(samples_v_nodecay[i,j,:], h_model_nodecay[μ,:]) / size(samples_v_nodecay,3) for i=1:28, j=1:28, μ=1:nh])
)
Makie.abline!(ax, 0, 1; color=:red)

ax = Makie.Axis(fig[2,1], xlabel="<v>_data", ylabel="<v>_model", limits=(0,1,0,1))
Makie.scatter!(ax, vec(mean(train_x; dims=3)), vec(mean(samples_v_decaylr; dims=3)))
Makie.abline!(ax, 0, 1; color=:red)

ax = Makie.Axis(fig[2,2], xlabel="<h>_data", ylabel="<h>_model", limits=(0,1,0,1))
Makie.scatter!(ax, vec(mean(h_data_decaylr; dims=2)), vec(mean(h_model_decaylr; dims=2)))
Makie.abline!(ax, 0, 1; color=:red)

ax = Makie.Axis(fig[2,3], xlabel="<vh>_data", ylabel="<vh>_model", limits=(0,1,0,1))
Makie.scatter!(ax,
    vec([dot(train_x[i,j,:], h_data_decaylr[μ,:]) / size(train_x,3) for i=1:28, j=1:28, μ=1:nh]),
    vec([dot(samples_v_decaylr[i,j,:], h_model_decaylr[μ,:]) / size(samples_v_decaylr,3) for i=1:28, j=1:28, μ=1:nh])
)
Makie.abline!(ax, 0, 1; color=:red)

fig


@testset "DFS" begin

  g5 = SimpleDiGraph(4)
  add_edge!(g5, 1, 2); add_edge!(g5, 2, 3); add_edge!(g5, 1, 3); add_edge!(g5, 3, 4)
  gx = cycle_digraph(3)

  @testset "dfs_tree" begin
    for g in testdigraphs(g5)
      z = @inferred(dfs_tree(g, 1))
      @test ne(z) == 3 && nv(z) == 4
      @test !has_edge(z, 1, 3)
      @test !is_cyclic(g)
    end
  end

  @testset "topological_sort_by_dfs" begin
    for g in testdigraphs(g5)
      @test @inferred(topological_sort_by_dfs(g)) == [1, 2, 3, 4]
    end

    for g in testdigraphs(gx)
      @test @inferred(is_cyclic(g))
      @test_throws ErrorException topological_sort_by_dfs(g)
    end
  end

  @testset "is_cyclic" begin
    for g in testgraphs(path_graph(2))
      @test @inferred(is_cyclic(g))
      @test @inferred(!is_cyclic(zero(g)))
    end
  end

end


using Decimals
using Base.Test

# Normalization
@test Decimal(1, 151100, -4) == Decimal(1, 1511, -2)
@test Decimal(0, 100100, -5) == Decimal(0, 1001, -3)
@test norm(Decimal(1, 151100, -4)) == Decimal(1, 1511, -2)
@test norm(Decimal(0, 100100, -5)) == Decimal(0, 1001, -3)
@test decimal("3.0") == Decimal(0, 3, 0)
@test decimal("3.0") == Decimal(0, 30, -1)
@test decimal("3.1400") == Decimal(0, 314, -2)
@test decimal("1234") == Decimal(0, 1234, 0)


# Day 13: Transparent Origami
function load(::Day{13}, path)
    positions = Set{CartesianIndex}()
    folds = Vector{Tuple{Char,Int}}()
    regex = r"fold along ([xy])=(\d+)"
    for line in readlines(path)
        isempty(line) && continue
        instruction = match(regex, line)
        if isnothing(instruction)
            x, y = parse.(Int, split(line, ","))
            push!(positions, CartesianIndex(x, y))
        else
            push!(folds, (first(instruction[1]), parse(Int, instruction[2])))
        end
    end
    return (positions, folds)
end

# Part 1 – single fold
@inline function foldx(I::CartesianIndex, pos)
    x, y = Tuple(I)
    return CartesianIndex(x - 2 * abs(pos - x), y)
end
@inline function foldy(I::CartesianIndex, pos)
    x, y = Tuple(I)
    return CartesianIndex(x, y - 2 * abs(pos - y))
end

function fold!(positions, fold)
    pfold = last(fold)
    if first(fold) == 'x'
        folded = [
            foldx(pop!(positions, p), pfold) for p in filter(p -> p[1] > pfold, positions)
        ]
    else
        folded = [
            foldy(pop!(positions, p), pfold) for p in filter(p -> p[2] > pfold, positions)
        ]
    end
    union!(positions, folded)
    return nothing
end

function solve1(::Day{13}, data)
    positions, folds = data
    fold!(positions, first(folds))
    return length(positions)
end

# Part 2 – show the final code
function show_code(positions::Set{CartesianIndex})
    pos = Tuple.(positions)
    h = maximum(first(t) + 1 for t in pos)
    w = maximum(last(t) + 1 for t in pos)
    img = zeros(Bool, h, w)
    for p in positions
        img[p + CartesianIndex(1, 1)] = 1
    end
    return transpose(img)
end

function solve2(::Day{13}, data)
    positions, folds = data
    for f in folds
        fold!(positions, f)
    end
    # @show show_code(positions) # worth plugging through ImageCore's `Gray`
    return 0
end

testresult1(::Day{13}) = 17

# No test for part 2, just pass:
testresult2(::Day{13}) = 0
test2(::Day{13}) = 0