text:
const epsilon = 10^-3

@testset "Configuration objects" begin
    @test_throws ArgumentError MTConfig{Float64}(100; nfft = 99)
    @test_throws ArgumentError MTConfig{Float64}(100; fs=-1)
    @test_throws DimensionMismatch MTConfig{Float64}(100; window = rand(1, 200))
    @test_throws ArgumentError MTConfig{Complex{Float64}}(100; onesided=true)

    @testset "`MTSpectrogramConfig` with `n_samples`=$(n_samples), `n_samples_per_window`=$(n_samples_per_window), `n_overlap`=$(n_overlap) " for n_samples in (20:20:100), n_samples_per_window in (20:20:100), n_overlap in (0:20:(n_samples_per_window-1))
        config = MTSpectrogramConfig{Float64}(n_samples, n_samples_per_window, n_overlap)
        @test length(config.time) == length(arraysplit(1:n_samples, n_samples_per_window, n_overlap))
    end

    @testset "`MTConfig` inner constructor errors" begin
        n_samples = 100
        nfft = 200
        ntapers = 5
        fs = 1
        window = rand(n_samples, ntapers)
        for T in (Float32, Float64, Complex{Float32}, Complex{Float64})
            fft_input_tmp = Vector{T}(undef, nfft)
            onesided = T <: Real
            fft_flags = FFTW.ESTIMATE
            freqs = onesided ? rfftfreq(nfft, fs) : fftfreq(nfft, fs)
            fft_output_tmp = Vector{fftouttype(T)}(undef, length(freqs))
            r = fs*ntapers*ones(ntapers)
            plan = onesided ? plan_rfft(fft_input_tmp; flags=fft_flags) :
                plan_fft(fft_input_tmp; flags=fft_flags)
            # Test that the current configuation is valid so we know if it errors later
            # it's because we changed it, not that it was always broken
            @test MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r) isa MTConfig
            # now do a series of changes (in let-blocks to introduce new local bindings)
            # and check that they are each invalid
            let n_samples = 201
                @test_throws ArgumentError MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r)
            end
            let n_samples = -1
                @test_throws ArgumentError MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r)
            end
            let ntapers = -1
                @test_throws ArgumentError MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r)
            end
            let fs = -1
                @test_throws ArgumentError MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r)
            end
            let fft_input_tmp = Vector{T}(undef, 2*nfft)
                @test_throws DimensionMismatch MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r)
            end
            let nfft = 2*nfft, fft_input_tmp = Vector{T}(undef, 2*nfft)
                @test_throws DimensionMismatch MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r)
            end
            let window = rand(2*ntapers, n_samples)
                @test_throws DimensionMismatch MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r)
            end
            let r = 1.0
                @test_throws DimensionMismatch MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r)
            end
            if T <: Complex
                let onesided=true, freqs=rfftfreq(nfft, fs), fft_output_tmp=Vector{fftouttype(T)}(undef, length(freqs)), plan=plan_fft(fft_input_tmp; flags=fft_flags)
                    @test_throws ArgumentError MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r)
                end
            end
            # let's check that our original config is still valid, meaning we haven't
            # accidentally broken the config by not properly scoping our changes
            # (and therefore invalidating the tests that follow).
            @test MTConfig{T}(n_samples, nfft, ntapers, freqs, fs, plan, fft_input_tmp, fft_output_tmp, window, onesided, r) isa MTConfig
        end
    end

end

@testset "`dpss_config`" begin
    dpss_config = DSP.Periodograms.dpss_config(Float64, 1024; keep_only_large_evals=false, weight_by_evals=false)
    mt_config = MTConfig{Float64}(1024)

    @test dpss_config.window ≈ mt_config.window
    @test dpss_config.ntapers == mt_config.ntapers
    @test dpss_config.r ≈ mt_config.r

    dpss_config = DSP.Periodograms.dpss_config(Float64, 1024; keep_only_large_evals=false, weight_by_evals=true)
    @test dpss_config.window ≈ mt_config.window
    @test dpss_config.ntapers == mt_config.ntapers
    # should be different since we're weighting by eigenvalues now
    @test sum(abs, dpss_config.r - mt_config.r) > 0.1
end

avg_coh(x) = dropdims(mean(coherence(x); dims=3); dims=3)

@testset "Coherence (synthetic data)" begin
    fs = 1000.0
    n_samples = 1024
    t = (0:1023) ./ fs
    sin_1 = sin.(2 * π * 12.0 * t)  # 12 Hz sinusoid signal
    sin_2 = sin.(2 * π * 12.0 * t .+ π)
    noise = rand(1024) * 2 .- 1

    same_signal = Matrix{Float64}(undef, 2, n_samples)
    same_signal[1, :] = sin_1
    same_signal[2, :] = sin_1
    coh = avg_coh(mt_coherence(same_signal; demean=true, fs=fs, freq_range = (10, 15)))
    same_signal_coherence = coh[2,1]

    # test in-place gets the same result
    config = MTCoherenceConfig{eltype(same_signal)}(size(same_signal)...; demean=true, fs=fs, freq_range = (10, 15))
    out = allocate_output(config)
    coh2 = avg_coh(mt_coherence!(out, same_signal, config))
    @test coh ≈ coh2

    # check that the output object has the right frequency information and properties
    coh_object = mt_coherence!(out, same_signal, config)
    @test freq(coh_object) == coh_object.freq == config.cs_config.freq
    @test all(x -> 10 <= x <= 15, freq(coh_object))
    # Check symmetries
    for i in 1:2, j in 1:2, f in eachindex(freq(coh_object))
        @test coherence(coh_object)[i, j, f] == coherence(coh_object)[j, i, f]
        if i == j
            @test coherence(coh_object)[i,j, f] == 1
        end
    end

    # check that manually demeaning with `demean=false` gives the same result
    same_signal_demeaned = same_signal .- mean(same_signal; dims = 2)
    coh3 = avg_coh(mt_coherence(same_signal_demeaned; demean = false))
    @test coh3 ≈ coh2

    @test_throws DimensionMismatch mt_coherence!(zeros(size(out,1), 2*size(out, 2), size(out, 3)), same_signal, config)
    @test_throws DimensionMismatch mt_coherence!(out, vcat(same_signal, same_signal), config)

    @test abs(same_signal_coherence - 1) < epsilon

    phase_shift = Matrix{Float64}(undef, 2, n_samples)
    phase_shift[1, :] = sin_1
    phase_shift[2, :] = sin_2
    coh = avg_coh(mt_coherence(phase_shift; fs=fs, freq_range = (10, 15)))
    phase_shift_coherence = coh[2,1]
    @test abs(phase_shift_coherence - 1) < epsilon
    @test coh[1,1] ≈ 1
    @test coh[2,2] ≈ 1

    # test in-place gets the same result
    config = MTCoherenceConfig{eltype(phase_shift)}(size(phase_shift)...; fs=fs, freq_range = (10, 15))
    out = allocate_output(config)
    coh2 = avg_coh(mt_coherence!(out, phase_shift, config))
    @test coh ≈ coh2

    # test out-of-place with config the same result
    config = MTCoherenceConfig{eltype(phase_shift)}(size(phase_shift)...; fs=fs, freq_range = (10, 15))
    coh3 = avg_coh(mt_coherence(phase_shift, config))
    @test coh ≈ coh3

    # test in-place without config the same result
    coh4 = avg_coh(mt_coherence!(out, phase_shift; fs=fs, freq_range = (10, 15)))
    @test coh ≈ coh4

    # Construct the config via an `MTConfig`
    mt_config = MTConfig{eltype(phase_shift)}(size(phase_shift, 2); fs=fs)
    config = MTCoherenceConfig(size(phase_shift, 1), mt_config; freq_range = (10, 15))
    coh5 = avg_coh(mt_coherence(phase_shift, config))
    @test coh ≈ coh5

    # Construct the config via an `MTCrossSpectraConfig`
    config = MTCoherenceConfig(MTCrossSpectraConfig(size(phase_shift, 1), mt_config; freq_range = (10, 15)))
    coh5 = avg_coh(mt_coherence(phase_shift, config))
    @test coh ≈ coh5

    different_signal = Matrix{Float64}(undef, 2, n_samples)
    different_signal[1, :] = sin_1
    different_signal[2, :] = noise
    coh = avg_coh(mt_coherence(different_signal; fs=fs, freq_range = (10, 15)))
    different_signal_coherence = coh[2,1]
    # .8 is arbitrary, but represents a high coherence, so if a sine wave is
    # that coherent with noise, there is a problem
    @test different_signal_coherence < 0.8

    # test in-place gets the same result
    config = MTCoherenceConfig{eltype(different_signal)}(size(different_signal)...; fs=fs, freq_range = (10, 15))
    out = allocate_output(config)
    coh2 = avg_coh(mt_coherence!(out, different_signal, config))
    @test coh ≈ coh2

    less_noisy = Matrix{Float64}(undef, 2, n_samples)
    less_noisy[1, :] = sin_1
    less_noisy[2, :] .= sin_1 .+ noise
    coh = avg_coh(mt_coherence(less_noisy; fs=fs, freq_range = (10, 15)))
    less_noisy_coherence = coh[2, 1]

    # test in-place gets the same result
    config = MTCoherenceConfig{eltype(less_noisy)}(size(less_noisy)...; fs=fs, freq_range = (10, 15))
    out = allocate_output(config)
    coh2 = avg_coh(mt_coherence!(out, less_noisy, config))
    @test coh ≈ coh2

    more_noisy = Matrix{Float64}(undef, 2, n_samples)
    more_noisy[1, :] = sin_1
    more_noisy[2, :] .= sin_1 .+ 3 * noise
    coh = avg_coh(mt_coherence(more_noisy; fs=fs, freq_range = (10, 15)))
    more_noisy_coherence = coh[2, 1]
    @test less_noisy_coherence < same_signal_coherence
    @test more_noisy_coherence < less_noisy_coherence
    @test different_signal_coherence < more_noisy_coherence

    # test in-place gets the same result
    config = MTCoherenceConfig{eltype(more_noisy)}(size(more_noisy)...; fs=fs, freq_range = (10, 15))
    out = allocate_output(config)
    coh2 = avg_coh(mt_coherence!(out, more_noisy, config))
    @test coh ≈ coh2

    several_signals = Matrix{Float64}(undef, 3, n_samples)
    several_signals[1, :] = sin_1
    several_signals[2, :] = sin_2
    several_signals[3, :] = noise
    several_signals_coherences = avg_coh(mt_coherence(several_signals; fs=fs, freq_range = (10, 15)))
    @test length(several_signals_coherences) == 9
    @test several_signals_coherences[2, 1] ≈ phase_shift_coherence
    @test several_signals_coherences[3, 1] ≈ different_signal_coherence

    # test in-place gets the same result
    config = MTCoherenceConfig{Float64}(size(several_signals)...; fs=fs, freq_range = (10, 15))
    out = allocate_output(config)
    @test eltype(out) == Float64
    coh2 = avg_coh(mt_coherence!(out, several_signals, config))
    @test coh2 ≈ several_signals_coherences

    # Float32 output:
    config = MTCoherenceConfig{Float32}(size(several_signals)...; fs=fs, freq_range = (10, 15))
    out = allocate_output(config)
    @test eltype(out) == Float32
    coh3 = avg_coh(mt_coherence!(out, several_signals, config))
    @test coh3 ≈ several_signals_coherences

    # Float32 input and output:
    coh4 = avg_coh(mt_coherence!(out, Float32.(several_signals), config))
    @test coh4 ≈ several_signals_coherences
end

@testset "`mt_coherence` reference test" begin
    fs = 1000.0
    n_samples = 1024
    t = (0:1023) ./ fs
    sin_1 = sin.(2 * π * 12.0 * t)  # 12 Hz sinusoid signal
    sin_2 = sin.(2 * π * 12.0 * t .+ π)
    noise = vec(readdlm(joinpath(@__DIR__, "data", "noise.txt"))) # generated by `rand(1024) * 2 .- 1`
    more_noisy = Array{Float64}(undef, 1, 2, n_samples)
    more_noisy[1, 1, :] = sin_1
    more_noisy[1, 2, :] .= sin_1 .+ 3 * noise

    ## to generate the reference result:
    # using PyMNE
    # mne_coherence_matrix, _ = PyMNE.connectivity.spectral_connectivity(more_noisy, method="coh",sfreq=fs,mode="multitaper",fmin=10,fmax=15,verbose=false)
    # coh = dropdims(mean(mne_coherence_matrix; dims=3); dims=3)[2, 1]
    coh = 0.982356762670818

    mt_config = DSP.Periodograms.dpss_config(Float64, n_samples; fs=fs, keep_only_large_evals=true, weight_by_evals=true)
    config = MTCoherenceConfig(2, mt_config; freq_range = (10,15), demean=true)
    result = avg_coh(mt_coherence(dropdims(more_noisy;dims=1), config))
    @test result[2, 1] ≈ coh
end

@testset "`mt_cross_power_spectra`" begin
    fs = 1000.0
    n_samples = 1024
    t = (0:1023) ./ fs
    data = Array{Float64, 3}(undef, 1, 2, n_samples)
    data[1,1,:] = sin.(2 * π * 12.0 * t) # 12 Hz sinusoid signal
    data[1,2,:] = sin.(2 * π * 12.0 * t .+ π)
    ## Script to generate reference data:
    # using PyMNE, DelimitedFiles
    # result = PyMNE.time_frequency.csd_array_multitaper(data, fs, n_fft = nextpow(2, n_samples), low_bias=true, adaptive=false)
    # csd_array_multitaper_frequencies = result.frequencies
    # csd_array_multitaper_values = reduce((x,y) -> cat(x,y; dims=3), [ r.get_data() for r in result])
    # writedlm("data/csd_array_multitaper_frequencies.txt", csd_array_multitaper_frequencies)
    # writedlm("data/csd_array_multitaper_values_re.txt", real(csd_array_multitaper_values))
    # writedlm("data/csd_array_multitaper_values_im.txt", imag(csd_array_multitaper_values))
    csd_array_multitaper_frequencies = readdlm(joinpath(@__DIR__, "data", "csd_array_multitaper_frequencies.txt"))
    csd_array_multitaper_values_re = reshape(readdlm(joinpath(@__DIR__, "data", "csd_array_multitaper_values_re.txt")), (2,2,512))
    csd_array_multitaper_values_im = reshape(readdlm(joinpath(@__DIR__, "data", "csd_array_multitaper_values_im.txt")), (2,2,512))
    csd_array_multitaper_values = csd_array_multitaper_values_re + im*csd_array_multitaper_values_im

    signal = dropdims(data; dims=1)
    mt_config = DSP.Periodograms.dpss_config(Float64, n_samples; fs=fs, keep_only_large_evals=true, weight_by_evals=true)
    config = MTCrossSpectraConfig(2, mt_config; demean=true)
    result = mt_cross_power_spectra(signal, config)
    @test signal isa Matrix{Float64}
    @test freq(result)[2:end] ≈ csd_array_multitaper_frequencies
    @test power(result)[:,:,2:end] ≈ csd_array_multitaper_values

    # Test in-place. Full precision:
    config = MTCrossSpectraConfig(size(signal, 1), mt_config; demean=true)
    out = allocate_output(config)
    @test eltype(out) == Complex{Float64}
    result2 = mt_cross_power_spectra!(out, signal, config)
    @test freq(result) ≈ freq(result2)
    @test power(result) ≈ power(result2)

    # Float32 output:
    mt_config32 = DSP.Periodograms.dpss_config(Float32, n_samples; fs=fs, keep_only_large_evals=true, weight_by_evals=true)
    config = MTCrossSpectraConfig(size(signal, 1), mt_config32; demean=true)
    out = allocate_output(config)
    @test eltype(out) == Complex{Float32}
    result2 = mt_cross_power_spectra!(out, signal, config)
    @test freq(result2) ≈ freq(result)
    @test power(result2) ≈ power(result)

    # Float32 input and output:
    result3 = mt_cross_power_spectra!(out, Float32.(signal), config)
    @test freq(result3) ≈ freq(result)
    @test power(result3) ≈ power(result)

    @test_throws DimensionMismatch mt_cross_power_spectra!(similar(out, size(out, 1) + 1, size(out, 2), size(out, 3)), signal, config)
    @test_throws DimensionMismatch mt_cross_power_spectra!(out, vcat(signal, signal), config)

end

@testset "`mt_cross_power_spectra` agrees with `mt_pgram`" begin
    fs = 1000.0
    n_samples = 1024
    t = (0:1023) ./ fs
    noise = vec(readdlm(joinpath(@__DIR__, "data", "noise.txt")))
    signal = sin.(2 * π * 12.0 * t) .+ 3 * noise
    cs = mt_cross_power_spectra(reshape(signal, 1, :); fs=fs)
    p = mt_pgram(signal; fs=fs)

    @test freq(cs) ≈ freq(p)
    @test dropdims(power(cs); dims=(1,2)) ≈ power(p)

    # out-of-place with config
    config = MTCrossSpectraConfig{Float64}(1, length(signal); fs=fs)
    cs = mt_cross_power_spectra(reshape(signal, 1, :), config)
    @test freq(cs) ≈ freq(p)
    @test dropdims(power(cs); dims=(1,2)) ≈ power(p)

    # in-place without config
    out = allocate_output(config)
    cs = mt_cross_power_spectra!(out, reshape(signal, 1, :); fs=fs)
    @test freq(cs) ≈ freq(p)
    @test dropdims(power(cs); dims=(1,2)) ≈ power(p)

    # rm once two-sided FFTs supported
    @test_throws ArgumentError mt_cross_power_spectra(reshape(complex.(signal), 1, :); fs=fs)
end


text:
module Input

using Fermi
using Fermi.HartreeFock.RHF
using Lints

export run #shorthand for read->exec
export read
export exec

"""
    run(fname::String)

read and execute the input file specified by fname.
"""
function run(fname::String = "input.dat")
    cont = read(fname)
    exec(cont)
end

"""
    read(fname::String)

read an input file and convert to a Dict. exec takes the output of this function as input.
"""
function read(fname)
    if isfile(fname)
        lines = readlines(fname)
    else
        error("Input file not found $fname")
    end
    molecule,options,command = parse(lines)
    cont = Dict("molecule"=>molecule,"options"=>options,"command"=>command)
end

"""
    exec(cont::Dict)

execute the computation specified in cont.
cont is composed of:
    Dict("molecule"=>molecule::String,
         "options"=>options::Dict,
         "command"=>command::function)
"""
function exec(cont)
    open("/tmp/molfile.xyz","w") do molfile
        natom = length(split(strip(cont["molecule"]),"\n"))
        write(molfile,"$natom\n\n")
        write(molfile,cont["molecule"])
    end
    Lints.libint2_init()
    mol = Lints.Molecule("/tmp/molfile.xyz")
    bas = Lints.BasisSet(cont["options"]["basis"],mol)
    rhfwfn = Fermi.HartreeFock.RHF.RHFWavefunction(bas,mol,5,5)
    com = cont["command"]
    E = com(rhfwfn)
    Lints.libint2_finalize()
    return E
end

function parse(lines;delim="--")
    chunks = split(join(lines,"\n"),delim)
    for chunk in chunks
        expr = Meta.parse(chunk)
        eval(expr)
    end
    molecule,options,command
end

end # module

ext: jl


text:
using Pkg
Pkg.activate(joinpath(@__DIR__, ".."))
Pkg.instantiate()

# ## Setup
using Dojo

# ## Simulation
timestep = 0.01
gravity = 0.0
mech = get_mechanism(:dzhanibekov,
        timestep=timestep,
        gravity=gravity);

# ## Simulate
initialize!(mech, :dzhanibekov,
    angular_velocity=[8.0; 0.01; 0.0]);
storage = simulate!(mech, 8.0,
    record=true,
    verbose=true)

# ## Visualizers
vis = Visualizer()
render(vis)
visualize(mech, storage,
    vis=vis);


text:
# a log of things that I put into the julia REPL
# useful things will be later absorbed into the package

using TensorKit
using TensorOperations
using TensorKitAD
using KrylovKit
using Zygote
using KrylovKit
using LinearAlgebra
using Plots
using Optim
using ChainRules
using ChainRulesCore
using BenchmarkTools
using Tullio
using LoopVectorization

using Revise
using statmech_tm_solver

lop = lieb_liniger_h_tangent_map(ψ, 0, 0, 1.0, 100, 1)
lop(V)

A = TensorMap(rand, ComplexF64, ℂ^64, ℂ^64)
B = TensorMap(rand, ComplexF64, ℂ^64, ℂ^64)
Wvec = 0.01*(rand(64) + im .* rand(64))

L = 1
M = zero(A)
@benchmark @tullio M.data[i, k] = A.data[i, j] * B.data[j, k] * statmech_tm_solver.theta3(L, Wvec[i], Wvec[j], Wvec[k])
C = zeros(ComplexF64, (64, 64, 64))
@benchmark @tullio C[i, j, k] = statmech_tm_solver.theta3(L, Wvec[i], Wvec[j], Wvec[k])
@tullio C[i, j, k] = statmech_tm_solver.theta3(L, Wvec[i], Wvec[j], Wvec[k])

@benchmark statmech_tm_solver.elem_mult_f2(A, B, (ix, im, iy)->statmech_tm_solver.theta3(L, Wvec[ix], Wvec[im], Wvec[iy]))
@benchmark @tullio M.data[i, k] = A.data[i, j] * B.data[j, k] * C[i, j, k]

L = 1
χ = 12
ψ = cmps(rand, χ, 1);
V = TensorMap(rand, ComplexF64, ℂ^χ, ℂ^χ);
lop = tangent_map(ψ, L);
@benchmark lop(V)

hop = lieb_liniger_h_tangent_map(ψ, 0, 0, 1.0, 100, 1);
@benchmark hop(V)

#########################################################
f1 = p -> (sqrt((p^2 - μ)^2 - 4 * ν^2) - (p^2 - μ)) / 2

Eexact = map(f1, (-100000:100000) .* 2*pi/L ) .* (1 / L) |> sum

@show (E - Eexact) / Eexact


text:
using Test, Interpolations, OffsetArrays

@testset "ExtrapNon1" begin

    f(x) = sin((x-3)*2pi/9 - 1)
    xinds = -3:6
    A = OffsetArray(Float64[f(x) for x in xinds], xinds)
    itpg = interpolate(A, BSpline(Linear()))

    schemes = (
        Flat,
        Line,
        Reflect,
        Periodic
    )

    for etp in map(E -> extrapolate(itpg, E()), schemes), x in xinds
        @test parent(etp) === itpg
        @test @inferred(getindex(etp, x)) ≈ A[x]
    end

    g(y) = (y/100)^3
    yinds = 2:5
    A = OffsetArray(Float64[f(x)*g(y) for x in xinds, y in yinds], xinds, yinds)
    itp2 = interpolate(A, BSpline(Linear()))

    for (etp2,E) in map(E -> (extrapolate(itp2, E()), E), schemes)
        @test parent(etp2) === itp2
        E == Periodic && continue  # g isn't periodic
        for y in yinds, x in xinds
            @test @inferred(getindex(etp2, x, y)) ≈ A[x, y]
        end
    end

end


text:
"""
Nodal admittance matrix (Ybus) is an N x N matrix describing a power system with N buses. It represents the nodal admittance of the buses in a power system.

The Ybus Struct is indexed using the Bus Numbers, no need for them to be sequential
"""
struct Ybus{Ax, L <: NTuple{2, Dict}} <: PowerNetworkMatrix{ComplexF64}
    data::SparseArrays.SparseMatrixCSC{ComplexF64, Int}
    axes::Ax
    lookup::L
end

function _ybus!(
    ybus::SparseArrays.SparseMatrixCSC{ComplexF64, Int},
    b::ACBranch,
    num_bus::Dict{Int, Int},
)
    arc = get_arc(b)
    bus_from_no = num_bus[arc.from.number]
    bus_to_no = num_bus[arc.to.number]
    Y_l = (1 / (get_r(b) + get_x(b) * 1im))
    Y11 = Y_l + (1im * get_b(b).from)
    ybus[bus_from_no, bus_from_no] += Y11
    Y12 = -Y_l
    ybus[bus_from_no, bus_to_no] += Y12
    #Y21 = Y12
    ybus[bus_to_no, bus_from_no] += Y12
    Y22 = Y_l + (1im * get_b(b).to)
    ybus[bus_to_no, bus_to_no] += Y22
    return
end

function _ybus!(
    ybus::SparseArrays.SparseMatrixCSC{ComplexF64, Int},
    b::Transformer2W,
    num_bus::Dict{Int, Int},
)
    arc = get_arc(b)
    bus_from_no = num_bus[arc.from.number]
    bus_to_no = num_bus[arc.to.number]
    Y_t = 1 / (get_r(b) + get_x(b) * 1im)
    Y11 = Y_t
    b = get_primary_shunt(b)
    ybus[bus_from_no, bus_from_no] += Y11
    ybus[bus_from_no, bus_to_no] += -Y_t
    ybus[bus_to_no, bus_from_no] += -Y_t
    ybus[bus_to_no, bus_to_no] += Y_t + (1im * b)
    return
end

function _ybus!(
    ybus::SparseArrays.SparseMatrixCSC{ComplexF64, Int},
    b::TapTransformer,
    num_bus::Dict{Int, Int},
)
    arc = get_arc(b)
    bus_from_no = num_bus[arc.from.number]
    bus_to_no = num_bus[arc.to.number]

    Y_t = 1 / (get_r(b) + get_x(b) * 1im)
    c = 1 / get_tap(b)
    b = get_primary_shunt(b)

    Y11 = (Y_t * c^2)
    ybus[bus_from_no, bus_from_no] += Y11
    Y12 = (-Y_t * c)
    ybus[bus_from_no, bus_to_no] += Y12
    #Y21 = Y12
    ybus[bus_to_no, bus_from_no] += Y12
    Y22 = Y_t
    ybus[bus_to_no, bus_to_no] += Y22 + (1im * b)
    return
end

function _ybus!(
    ybus::SparseArrays.SparseMatrixCSC{ComplexF64, Int},
    b::PhaseShiftingTransformer,
    num_bus::Dict{Int, Int},
)
    arc = get_arc(b)
    bus_from_no = num_bus[arc.from.number]
    bus_to_no = num_bus[arc.to.number]
    Y_t = 1 / (get_r(b) + get_x(b) * 1im)
    tap = (get_tap(b) * exp(get_α(b) * 1im))
    c_tap = (get_tap(b) * exp(-1 * get_α(b) * 1im))
    b = get_primary_shunt(b)
    Y11 = (Y_t / abs(tap)^2)
    ybus[bus_from_no, bus_from_no] += Y11
    Y12 = (-Y_t / c_tap)
    ybus[bus_from_no, bus_to_no] += Y12
    Y21 = (-Y_t / tap)
    ybus[bus_to_no, bus_from_no] += Y21
    Y22 = Y_t
    ybus[bus_to_no, bus_to_no] += Y22 + (1im * b)
    return
end

function _ybus!(
    ybus::SparseArrays.SparseMatrixCSC{ComplexF64, Int},
    fa::FixedAdmittance,
    num_bus::Dict{Int, Int},
)
    bus = get_bus(fa)
    bus_no = num_bus[get_number(bus)]
    ybus[bus_no, bus_no] += fa.Y
    return
end

function _buildybus(branches, nodes, fixed_admittances)
    buscount = length(nodes)
    num_bus = Dict{Int, Int}()

    for (ix, b) in enumerate(nodes)
        num_bus[get_number(b)] = ix
    end
    ybus = SparseArrays.spzeros(ComplexF64, buscount, buscount)
    for (ix, b) in enumerate(branches)
        if get_name(b) == "init"
            throw(DataFormatError("The data in Branch is invalid"))
        end

        _ybus!(ybus, b, num_bus)
    end
    for fa in fixed_admittances
        _ybus!(ybus, fa, num_bus)
    end
    return SparseArrays.dropzeros!(ybus)
end

function _goderya(ybus::SparseArrays.SparseMatrixCSC{ComplexF64, Int})
    node_count = size(ybus)[1]
    max_I = node_count^2
    I, J, val = SparseArrays.findnz(ybus)
    T = SparseArrays.sparse(I, J, ones(length(val)))
    T_ = T * T
    for n in 1:(node_count - 1)
        I, _, _ = SparseArrays.findnz(T_)
        if length(I) == max_I
            @info "The System has no islands"
            break
        elseif length(I) < max_I
            temp = T_ * T
            I_temp, _, _ = SparseArrays.findnz(temp)
            if all(I_temp == I)
                throw(DataFormatError("The system contains islands"))
            end
            T_ = temp
        else
            @assert false
        end
        @assert n < node_count - 1
    end
    return
end

"""
Builds a Ybus from a collection of buses and branches. The return is a Ybus Array indexed with the bus numbers and the branch names.

# Keyword arguments
- `check_connectivity::Bool`: Checks connectivity of the network using Goderya's algorithm
"""
function Ybus(branches, nodes; check_connectivity::Bool = true)
    #Get axis names
    bus_ax = [get_number(bus) for bus in nodes]
    sort!(bus_ax)
    axes = (bus_ax, bus_ax)
    look_up = (_make_ax_ref(bus_ax), _make_ax_ref(bus_ax))
    ybus = _buildybus(branches, nodes, Vector{FixedAdmittance}())
    if check_connectivity
        _goderya(ybus)
    end
    return Ybus(ybus, axes, look_up)
end

"""
Builds a Ybus from the system. The return is a Ybus Array indexed with the bus numbers and the branch names.

# Keyword arguments
- `check_connectivity::Bool`: Checks connectivity of the network using Goderya's algorithm
"""

function Ybus(sys::System; check_connectivity::Bool = true)
    branches = get_components(ACBranch, sys)
    nodes = sort!(collect(get_components(Bus, sys)), by = x -> x.number)
    fixed_admittances = get_components(FixedAdmittance, sys)
    # Get axis names
    bus_ax = sort!([get_number(bus) for bus in nodes])
    axes = (bus_ax, bus_ax)
    look_up = (_make_ax_ref(bus_ax), _make_ax_ref(bus_ax))
    ybus = _buildybus(branches, nodes, fixed_admittances)
    if check_connectivity
        _goderya(ybus)
    end
    return Ybus(ybus, axes, look_up)
end


text:
function test_f0_otypes()
    srand(98765)
    dummy_input = rand(Float64, 4096)

    for (otype_int, otype_str) in zip(0:2, SPTK.supported_otypes)
        f01 = swipe(dummy_input, 16000, 80, otype=otype_int)
        f02 = swipe(dummy_input, 16000, 80, otype=otype_str)
        @test_approx_eq f01 f02
    end

    srand(98765)
    dummy_input = rand(Float32, 4096)

    for (otype_int, otype_str) in zip(0:2, SPTK.supported_otypes)
        f01 = rapt(dummy_input, 16000, 80, otype=otype_int)
        f02 = rapt(dummy_input, 16000, 80, otype=otype_str)
        @test_approx_eq f01 f02
    end
end

function test_swipe()
    srand(98765)
    dummy_input = rand(1024)

    println("-- test_swipe")
    for otype in [0, 1, 2]
        for fs in [16000]
            for hopsize in [40, 80, 160, 320]
                f0 = swipe(dummy_input, fs, hopsize, otype=otype)
                @test all(isfinite(f0))
                otype == 0 && @test all(f0 .>= 0)
            end
        end
    end

    fs = 16000

    # invalid otype
    @test_throws ArgumentError swipe(dummy_input, fs, 80, otype=-1)
    @test_throws ArgumentError swipe(dummy_input, fs, 80, otype=-3)

    # invalid min/max
    @test_throws ArgumentError swipe(dummy_input, fs, 80, min=60.0, max=60.0)
    swipe(dummy_input, 16000, 80, min=60.0, max=7999.0)
    @test_throws ArgumentError swipe(dummy_input, fs, 80, min=60.0, max=8000.0)
end

function test_rapt()
    srand(98765)
    dummy_input = rand(Float32, 1024)

    println("-- test_rapt")
    for otype in [0, 1, 2]
        for fs in [16000]
            for hopsize in [40, 80, 160, 320]
                println(" where otype = $otype, fs = $fs and hopsize = $hopsize")
                f0 = rapt(dummy_input, fs, hopsize, otype=otype)
                @test all(isfinite(f0))
                otype == 0 && @test all(f0 .>= 0)
            end
        end
    end

    fs = 16000

    # invalid otype
    @test_throws ArgumentError rapt(dummy_input, fs, 80, otype=-1)
    @test_throws ArgumentError rapt(dummy_input, fs, 80, otype=-3)

    # invalid min/max
    @test_throws ArgumentError rapt(dummy_input, fs, 80, min=60.0, max=60.0)
    @test_throws ArgumentError rapt(dummy_input, fs, 80, min=60.0, max=8000.0)

    # valid frame period (corner case)
    # TODO: should be 1599 instead of 1600
    rapt(rand(Float32, 10000), fs, 1599)

    # invalid frame_period
    @test_throws ArgumentError rapt(dummy_input, fs, 1601)

    # invalid input length (too small)
    @test_throws ErrorException rapt(dummy_input[1:100], fs, 80)
end

println("test f0 estimation")
test_f0_otypes()
test_swipe()
test_rapt()


text:
@recipe function f(::Type{Val{:tswindows}}, x, y, z; windows=5, N=min(length(x)÷10, 1000))
    L = length(y)
    xlims, ylims = (extrema(x), extrema(y))
    scale = (ylims[2] - ylims[1])*0.1
    ylims = (ylims[1] - scale, ylims[2] + scale)
    legend --> false
    ylims := ylims
    #seriescolor --> :black
    #grid --> false

    # We want to take a long time series, choose some windows of length N and plot these adjacently
    dL = (L÷(2*windows))
    idxs = [2*a*dL+dL-Int(floor(N/2)):1:2*a*dL+dL+Int(floor(N/2)-1) for a ∈ (1:windows).-1]
    dispxs = LinRange(0, max(x...), length(vcat(idxs...)))
    nn = -N+1
    for s in 1:length(idxs)
        nn += N
        # if nn-N < 1
        #     annx = (dispxs[nn+N])/2
        # elseif nn+N > length(dispxs)
        #     annx = (dispxs[nn] + dispxs[end])/2
        #     println(dispxs[nn])
        # else
        #     annx = (dispxs[nn+N] + dispxs[nn])/2
        # end
        # anum = (round((x[idxs[s][end]] + x[idxs[s][1]])/2, sigdigits=3))
        # if dispxs[end] > 99
        #     anum = Int(anum)
        # end
        # annotations := [(annx, min(y...)-2*scale, text("$anum", :black, :centre, 8))]
        if s < length(idxs)
            barx = (max(dispxs[nn:nn+N-1]...) + min(dispxs[(nn+N):(nn+2*N-1)]...))/2
            @series begin
                seriestype := :path
                linewidth := 3
                linecolor := :red
                x := [barx, barx]
                y := [min(y...)-scale, max(y...)+scale]
            end
        end
        @series begin
            seriestype := :path
            #framestyle --> :box
            #xlims --> xlims
            #ylims --> ylims
            seriescolor := :black
            x := dispxs[nn:nn+N-1]
            y := y[idxs[s]]
        end
    end
end

# ------------------------------------------------------------------------------------------------ #
#                              Function for plotting array of rasters                              #
# ------------------------------------------------------------------------------------------------ #

@userplot RasterArray
@recipe function f(g::RasterArray)
    @assert length(g.args) == 3 && g.args[3] isa AbstractMatrix
    seriestype := :rasterarray
    mat = g.args[3]
    g.args[1], g.args[2], Surface(mat)
end

@recipe function f(::Type{Val{:rasterarray}}, x, y, z)
    # Currently only works with gr()
    x, y, T = x, y, z.surf
    T = normalise(T, standardise, 2)
    if size(T, 1) > size(T, 2)
        @warn "Seems like a lot of time series, sure you didn't mean to transpose?"
    end
    uy = unique(y)
    ux = unique(x)
    yax = (sort∘unique)(y) # Must have equally spaced y's and x's
    xax = (sort∘unique)(x)
    P = Array{Matrix{Float64}}(undef, (length(yax), length(xax)))
    for 😄 ∈ CartesianIndices(P)
        (i, j) = Tuple(😄)
        P[😄] = T[(x.==xax[j]) .& (y.==yax[i]), :]
    end
    P .= reshape.(P, 1, size(T, 2))
    if size(P, 2) > 1
        P = [hcat(P[i, :]...) for i ∈ 1:size(P, 1)]
    end
    if size(P, 1) > 1
        P = vcat(P...)
    end
    seriestype := :heatmap
    xticks --> (LinRange(0, size(P, 2), length(ux)+1)[1:end-1].+size(P, 2)/length(ux)/2, ux)
    yticks --> (1:size(P, 1), uy)
    xlims := (0.5, size(P, 2)+0.5)
    ylims := (0.5, size(P, 1)+0.5)
    y := collect(1:size(P, 1))
    x := collect(1:size(P, 2))
    z := Surface(P)
    ()
end

# @userplot TSRasterArray
# @recipe function f(g::TSRasterArray)
#     @assert length(g.args) == 3 && typeof(g.args[3]) <: AbstractMatrix
#     seriestype := :tsrasterarray
#     mat = g.args[3]
#     g.args[1], g.args[2], Surface(mat)
# end
# @shorthands tsrasterarray
# @userplot TSRasterArray
# @recipe function f(g::TSRasterArray)
#     #x, y, z = plotattributes[:x], plotattributes[:y], plotattributes[:y].surf
#     x = g.args[1]
#     y = g.args[2]
#     z = g.args[3]
#     @series begin
#         seriestype := rasterarray
#         (x, y, z)
#     end
# end


text:
using MAT, JLD2

# Read .mat file and write to .jld2 file

matf = matread("f0_jar.mat");
f0 = matf["f0"];          # 3D fieldmap we wish to shim (Hz)
fov = matf["fov"];
mask = matf["mask"];
@save "f0_jar.jld2" f0 fov mask


text:
@testset "QuadTreeTests" begin
    @testset "Neighbour search NE" begin
      bb = GeometryTypes.SimpleRectangle(2.0, 3.0, 4.0, 4.0)
      q = QuadTreeMeshes.QuadTree{Int}(bb)

      # subdivide twice
      QuadTreeMeshes.subdivide!(q, 1, subdividefunc)
      firstElement = get_neighbour_from_index(q, 1, 2)
      QuadTreeMeshes.subdivide!(q, get(firstElement), subdividefunc)

      # and plot it
      #filename = "subdivide_plot_4.svg"
      #Plots.plot(q)
      #Plots.savefig(filename)

      @test get(QuadTreeMeshes.find_neighbour(q, 3, QuadTreeMeshes.south)) == 5
      @test get(QuadTreeMeshes.find_neighbour(q, 3, QuadTreeMeshes.west)) == 2
      @test isnull(QuadTreeMeshes.find_neighbour(q, 3, QuadTreeMeshes.east))
      @test isnull(QuadTreeMeshes.find_neighbour(q, 3, QuadTreeMeshes.north))
      @test get(QuadTreeMeshes.find_neighbour(q, 2, QuadTreeMeshes.east)) == 3
      @test get(QuadTreeMeshes.find_neighbour(q, 5, QuadTreeMeshes.north)) == 3

      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.west)) == 2
      @test isnull(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.north))
      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.east)) == 7
      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.south)) == 8

      @test get(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.west)) == 6
      @test isnull(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.north))
      @test isnull(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.east))
      @test get(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.south)) == 9
      #println("$q")
    end

    @testset "Neighbour search SE" begin
      bb = GeometryTypes.SimpleRectangle(2.0, 3.0, 4.0, 4.0)
      q = QuadTreeMeshes.QuadTree{Int}(bb)

      # subdivide twice
      QuadTreeMeshes.subdivide!(q, 1, subdividefunc)
      firstElement = get_neighbour_from_index(q, 1, 4)
      QuadTreeMeshes.subdivide!(q, get(firstElement), subdividefunc)

      # and plot it
      #filename = "subdivide_plot_3.svg"
      #Plots.plot(q)
      #Plots.savefig(filename)

      @test isnull(QuadTreeMeshes.find_neighbour(q, 5, QuadTreeMeshes.south))
      @test get(QuadTreeMeshes.find_neighbour(q, 5, QuadTreeMeshes.west)) == 4
      @test isnull(QuadTreeMeshes.find_neighbour(q, 5, QuadTreeMeshes.east))
      @test get(QuadTreeMeshes.find_neighbour(q, 5, QuadTreeMeshes.north)) == 3
      @test get(QuadTreeMeshes.find_neighbour(q, 4, QuadTreeMeshes.east)) == 5
      @test get(QuadTreeMeshes.find_neighbour(q, 3, QuadTreeMeshes.south)) == 5

      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.west)) == 4
      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.north)) == 3
      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.east)) == 7
      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.south)) == 8

      @test get(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.west)) == 6
      @test get(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.north)) == 3
      @test isnull(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.east))
      @test get(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.south)) == 9
      #println("$q")
    end

    @testset "Neighbour search SW" begin
      bb = GeometryTypes.SimpleRectangle(2.0, 3.0, 4.0, 4.0)
      q = QuadTreeMeshes.QuadTree{Int}(bb)

      # subdivide twice
      QuadTreeMeshes.subdivide!(q, 1, subdividefunc)
      firstElement = get_neighbour_from_index(q, 1, 3)
      QuadTreeMeshes.subdivide!(q, get(firstElement), subdividefunc)

      # and plot it
      #filename = "subdivide_plot_2.svg"
      #Plots.plot(q)
      #Plots.savefig(filename)

      @test isnull(QuadTreeMeshes.find_neighbour(q, 4, QuadTreeMeshes.south))
      @test get(QuadTreeMeshes.find_neighbour(q, 4, QuadTreeMeshes.east)) == 5
      @test isnull(QuadTreeMeshes.find_neighbour(q, 4, QuadTreeMeshes.west))
      @test get(QuadTreeMeshes.find_neighbour(q, 4, QuadTreeMeshes.north)) == 2
      @test get(QuadTreeMeshes.find_neighbour(q, 5, QuadTreeMeshes.west)) == 4
      @test get(QuadTreeMeshes.find_neighbour(q, 5, QuadTreeMeshes.north)) == 3

      @test isnull(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.west))
      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.north)) == 2
      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.east)) == 7
      @test get(QuadTreeMeshes.find_neighbour(q, 6, QuadTreeMeshes.south)) == 8

      @test get(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.west)) == 6
      @test get(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.north)) == 2
      @test get(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.east)) == 5
      @test get(QuadTreeMeshes.find_neighbour(q, 7, QuadTreeMeshes.south)) == 9
      #println("$q")
    end

    @testset "Neighbour search NW" begin
      bb = GeometryTypes.SimpleRectangle(2.0, 3.0, 4.0, 4.0)
      q = QuadTreeMeshes.QuadTree{Int}(bb)

      # subdivide twice
      QuadTreeMeshes.subdivide!(q, 1, subdividefunc)
      firstElement = get_neighbour_from_index(q, 1, 1)
      QuadTreeMeshes.subdivide!(q, get(firstElement), subdividefunc)

      # and plot it
      #filename = "subdivide_plot_1.svg"
      #Plots.plot(q)
      #Plots.savefig(filename)

      @test isnull(QuadTreeMeshes.find_neighbour(q, 2, QuadTreeMeshes.north))
      @test get(QuadTreeMeshes.find_neighbour(q, 2, QuadTreeMeshes.east)) == 3
      @test isnull(QuadTreeMeshes.find_neighbour(q, 2, QuadTreeMeshes.west))
      @test get(QuadTreeMeshes.find_neighbour(q, 2, QuadTreeMeshes.south)) == 4
      @test get(QuadTreeMeshes.find_neighbour(q, 3, QuadTreeMeshes.west)) == 2
      @test get(QuadTreeMeshes.find_neighbour(q, 3, QuadTreeMeshes.south)) == 5
      @test isnull(QuadTreeMeshes.find_neighbour(q, 3, QuadTreeMeshes.north))

      #println("$q")
    end

    @testset "Random Line Subdivision" begin
    # seed rng and store seed in file
    #local seed = 338139087
    local seed = rand(UInt32)
    local fs = open("RandomLineSubdivision.seed", "w")
    write(fs, string(seed))
    close(fs)
    srand(seed)

    bb = GeometryTypes.SimpleRectangle(2.0, 3.0, 4.0, 4.0)
    q = QuadTreeMeshes.QuadTree{Int}(bb)

    # pick random line inside bounding box
    linepos = rand(Float64, 4) * 2.0
    r1 = GeometryTypes.Point(linepos[1] + 2.0, linepos[2] + 3.0)
    r2 = GeometryTypes.Point(linepos[3] + 2.0, linepos[4] + 3.0)
    ls = GeometryTypes.LineSegment(r1, r2)
    QuadTreeMeshes.refine_line_to_level(q, ls, 6, subdividefunc)

    # and plot it
    filename = "random_line_subdivision.svg"
    Plots.plot(q)
    Plots.plot!([r1[1], r2[1]], [r1[2], r2[2]])
    Plots.savefig(filename)

    # go through all elements and check that they either don't intersect
    # or are in the list
    for (elIndex, el) in enumerate(q.elements)
      if (!QuadTreeMeshes.has_child(q, elIndex))
        bb = QuadTreeMeshes.get_element_bounding_box(q, el)
        intersecting = line_intersects_rectangle(bb, ls)
        if intersecting
          @test q.elements[elIndex].level == 6
        end
      end
    end
  end

  @testset "Random Line Intersection" begin
    # seed rng and store seed in file
    #local seed = 4196620933
    local seed = rand(UInt32)
    local fs = open("RandomLineIntersection.seed", "w")
    write(fs, string(seed))
    close(fs)
    srand(seed)

    bb = GeometryTypes.SimpleRectangle(2.0, 3.0, 4.0, 4.0)
    q = QuadTreeMeshes.QuadTree{Int}(bb)

    # subdivide three random children
    QuadTreeMeshes.subdivide!(q, 1, subdividefunc)
    subdivide_random!(q)
    subdivide_random!(q)
    subdivide_random!(q)

    # pick random line inside bounding box
    linepos = rand(Float64, 4) * 2.0
    r1 = GeometryTypes.Point(linepos[1] + 2.0, linepos[2] + 3.0)
    r2 = GeometryTypes.Point(linepos[3] + 2.0, linepos[4] + 3.0)
    ls = GeometryTypes.LineSegment(r1, r2)
    intersectingLeaves = QuadTreeMeshes.query_line(q, ls)

    # and plot it
    #filename = "random_line_intersection.svg"
    #Plots.plot(q)
    #Plots.plot!([r1[1], r2[1]], [r1[2], r2[2]])
    #Plots.savefig(filename)

    # go through all elements and check that they either don't intersect
    # or are in the list
    for (elIndex, el) in enumerate(q.elements)
      if (!QuadTreeMeshes.has_child(q, elIndex))
        bb = QuadTreeMeshes.get_element_bounding_box(q, el)
        intersecting = line_intersects_rectangle(bb, ls)
        interPos = findfirst(intersectingLeaves, elIndex)
        if intersecting
          @test interPos > 0
          if (interPos > 0)
            deleteat!(intersectingLeaves, interPos)
          end
        else
          @test interPos == 0
        end
      end
    end
    @test size(intersectingLeaves)[1] == 0
  end

  @testset "Subdivide Level 2" begin for u in 1:4, v in 1:4 begin
    bb = GeometryTypes.SimpleRectangle(2.0, 3.0, 4.0, 4.0)
    q = QuadTreeMeshes.QuadTree{Int}(bb)

    # subdivide twice
    QuadTreeMeshes.subdivide!(q, 1, subdividefunc)
    firstElement = get_neighbour_from_index(q, 1, u)
    QuadTreeMeshes.subdivide!(q, get(firstElement), subdividefunc)
    secondElement = get_neighbour_from_index(q, get(firstElement), v)
    QuadTreeMeshes.subdivide!(q, get(secondElement), subdividefunc)

    # and plot it
    filename = "subdivide_plot_$u$v.svg"
    Plots.plot(q)
    Plots.savefig(filename)

    # check that all levels are correctly subdivided
    # (i.e. it is balanced)
    check_subdivision_levels(q)

  end
  end
  end

  @testset "Constructor" begin

  # construct quadtree object and check that types are correct
  bb = GeometryTypes.SimpleRectangle(2.0, 3.0, 4.0, 4.0)
  q = QuadTreeMeshes.QuadTree{Int}(bb)

  @test length(q.elements) == 1
  firstEl = q.elements[1]
  @test isnull(firstEl.parent)
  @test isnull(firstEl.northWest)
  @test isnull(firstEl.northEast)
  @test isnull(firstEl.southEast)
  @test isnull(firstEl.southWest)
  @test firstEl.level == 1
  @test firstEl.index == 1
  testBB = QuadTreeMeshes.get_element_bounding_box(q, firstEl)
  @test testBB == bb
  @test typeof(q.values) == Array{Array{Int, 1}, 1}
  @test length(q.values) == 1
  end
end


text:
using DataFrames, CSV
using PyPlot
using Printf
using SpecialFunctions
include("utilities.jl")

filename = @sprintf("../data/zeta_test_data.csv")
data1 = CSV.read(filename; delim=",", type=String)

data1[!,:s] = parse.(Complex{Float64}, data1[!,:s] )
data1[!,:z] = parse.(Complex{Float64}, data1[!,:z] )
data1[!,:zeta] = parse.(Complex{Float64}, data1[!,:zeta] )

m = size(data1,1)
s = data1[!,:s]
z = complex(data1[!,:z])

S1 = zeros(Complex{Float64}, m)
for i = 1:m
    S1[i] = SpecialFunctions.zeta( 1-s[i], 0.5 + log(z[i])/(2*pi*im) )
end

error1 = S1 - data1[!,:zeta]
rel_error = abs.( error1 ) ./ abs.( data1[!,:zeta] )
maximum(rel_error)


text:
using JuMP, Gurobi

function solomon_exact(vrptw::VRPTW_Instance)
    n_nodes = size(vrptw.travel_time, 1)
    n_requests = length(vrptw.service_time)
    n_vehicles = n_requests
    @show n_nodes, n_requests
    @assert n_nodes - 2 == n_requests

    V = 1:n_vehicles
    N = 1:n_nodes
    C = 1:n_requests
    depot0 = n_requests + 1
    depot_dummy = n_requests + 2

    d = vrptw.load
    q = vrptw.capacity
    a = vrptw.early_time
    b = vrptw.late_time
    cost = copy(vrptw.travel_time)
    for i in N, j in N
        if cost[i, j] == Inf
            cost[i, j] = 0.0
        end
    end

    t = copy(vrptw.travel_time)
    for i in N, j in N
        if i in C
            t[i, j] = t[i, j] + vrptw.service_time[i]
        else
            t[i, j] = t[i, j]
        end
        if t[i, j] == Inf
            t[i, j] = 0.0
        end
    end

    a = [a; 0; 0]
    b = [b; 0; vrptw.max_travel_time]
    bigM = vrptw.max_travel_time * 2
    bigQ = q + maximum(d)

    m = nothing
    x = nothing
    s = nothing

    m = Model(Gurobi.Optimizer)
    @variable(m, x[N, N], Bin)
    @variable(m, s[N] >= 0)
    @variable(m, qq[N] >= 0)

    @objective(m, Min, sum(cost[i,j] * x[i,j] for i in N, j in N))

    @constraint(m, [i in C], sum(x[i,j] for j in N) == 1 )

    @constraint(m, [i in N, j in C], qq[j] >= qq[i] + d[j] - bigQ * (1 - x[i,j]))
    @constraint(m, [i in N], 0 <= qq[i] <= q)

    @constraint(m, sum(x[depot0,j] for j in N) <= n_vehicles )
    @constraint(m, [h in C], sum(x[i,h] for i in N) - sum(x[h,j] for j in N) == 0 )
    @constraint(m, sum(x[i,depot_dummy] for i in N) == sum(x[depot0,j] for j in N) )

    @constraint(m, [i in N, j in N], s[j] >= s[i] + t[i,j] - bigM * (1 - x[i,j]))
    @constraint(m, [i in N], a[i] <= s[i] <= b[i])

    @constraint(m, [i in N], x[i, depot0] == 0)
    @constraint(m, [i in N], x[depot_dummy, i] == 0)
    @constraint(m, [i in N], x[i, i] == 0)

    start_time = time()
    optimize!(m)
    duration = time() - start_time

    @show raw_status(m)
    @show objective_value(m)

    x_val = JuMP.value.(x).data
    first_nodes = findall(x->x>0.9, x_val[depot0, :])
    route_costs = Dict()
    for i in first_nodes
        next = findfirst(x->x>0.9, x_val[i, :])
        route = [depot0, i, next]
        cc = cost[depot0, i] + cost[i, next]
        while next != depot_dummy
            current = next
            next = findfirst(x->x>0.9, x_val[current, :])
            push!(route, next)
            cc += cost[current, next]
            # @show route
        end
        route_costs[route] = cc
    end
    display(route_costs)
    @show sum([route_costs[s] for s in keys(route_costs)])
    println(route_costs)
end


text:
#
println("Dcp) Test of the Photoionization and PlasmaShift module with ASF from an internally generated initial- and final-state multiplet.")
#
setDefaults("print summary: open", "zzz-PhotoIonization-Plasma.sum")
setDefaults("method: continuum, Galerkin")
setDefaults("method: normalization, pure sine")   ## setDefaults("method: normalization, pure Coulomb")    setDefaults("method: normalization, pure sine")

wa = Atomic.Computation("xx",  Nuclear.Model(36.); grid=JAC.Radial.Grid("grid: by given parameters"; rnt = 2.0e-6, h = 5.0e-2, hp = 1.0e-2, NoPoints = 600),
                        initialConfigs=[Configuration("1s^2 2s^2 2p^6")],
                        finalConfigs  =[Configuration("1s^2 2s^2 2p^5"), Configuration("1s^2 2s 2p^6") ],
                        process = JAC.PhotoInPlasma,
                        processSettings=PlasmaShift.PhotoSettings(PlasmaShift.DebyeHueckel, 0.25, 0., 0, [E1], [JAC.UseCoulomb],
                                                                  [3000., 3200.], true, false, Tuple{Int64,Int64}[]) )

wb = perform(wa)
setDefaults("print summary: close", "")


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMDynamics.jl/blob/master/LICENSE

using JuliaFEM
using JuliaFEM.Preprocess
add_elements! = JuliaFEM.add_elements!
using FEMDynamics
using Logging
Logging.configure(level=DEBUG)

## Generating structured block mesh

# TODO: this should be in some package!
function quad_mesh(nel_x=10, nel_y=10; lx=1.0, ly=1.0)
    nnodes_x = nel_x+1
    nnodes_y = nel_y+1
    nnode = nnodes_x*nnodes_y
    nodemap = reshape(1:nnode, nnodes_x, nnodes_y)
    nodes_1 = vec(nodemap[1:nnodes_x-1, 1:nnodes_y-1])
    nodes_2 = vec(nodemap[2:nnodes_x, 1:nnodes_y-1])
    nodes_3 = vec(nodemap[2:nnodes_x, 2:nnodes_y])
    nodes_4 = vec(nodemap[1:nnodes_x-1, 2:nnodes_y])

    mesh = Mesh()

    nid = 1
    for y in linspace(0, ly, nnodes_y)
        for x in linspace(0, lx, nnodes_x)
            add_node!(mesh, nid, [x, y])
            add_node_to_node_set!(mesh, :NALL, nid)
            nid += 1
        end
    end

    elid = 1
    for c in zip(nodes_1, nodes_2, nodes_3, nodes_4)
        add_element!(mesh, elid, :Quad4, collect(c))
        add_element_to_element_set!(mesh, :EALL, elid)
        elid += 1
    end

    return mesh
end

nel_x = 10
nel_y = 3
mesh = quad_mesh(nel_x, nel_y; lx=5000.0, ly=90.0)
for nid=1:nel_x+1:(nel_x+1)*(nel_y+1)
    add_node_to_node_set!(mesh, :LEFT, nid)
end

block_elements = create_elements(mesh, "EALL")
update!(block_elements, "youngs modulus", 210.0e3)
update!(block_elements, "poissons ratio", 0.3)
update!(block_elements, "density", 7.80e-7)
update!(block_elements, "displacement load 2", 7.80e-7*9.81)
block = Problem(Elasticity, "10x1 block", 2)
block.properties.formulation = :plane_stress
add_elements!(block, block_elements)

fixed = Problem(Dirichlet, "fixed", 2, "displacement")
fixed_nodes = mesh.node_sets[:LEFT]
fixed_elements = [Element(Poi1, [j]) for j in fixed_nodes]
update!(fixed_elements, "geometry", mesh.nodes)
for j=1:2
    update!(fixed_elements, "displacement $j", 0.0)
end
add_elements!(fixed, fixed_elements)

xdmf = Xdmf("block_resulst"; overwrite=true)

# analysis = Analysis(Nonlinear, "modal analysis of block under self-weight")
# add_results_writer!(analysis, xdmf)
# add_problems!(analysis, [block, fixed])
# run!(analysis)

# analysis = Analysis(Modal, "modal analysis of block under self-weight")
# add_results_writer!(analysis, xdmf)
# add_problems!(analysis, [block, fixed])
# run!(analysis)

analysis = Analysis(LinearImplicit, "transient analysis of block under self-weight")
add_results_writer!(analysis, xdmf)
add_problems!(analysis, [block, fixed])
run!(analysis)

close(xdmf.hdf)

max_el_id = maximum(keys(mesh.elements))
element = first(filter(e -> e.id == max_el_id, block_elements))
disp = element.fields["displacement"].data
t = [d.first for d in disp]
u = [d.second[end][2] for d in disp]

# using Plots
# plot(t[1:10:end], u[1:10:end])


text:
using Revise
using SeniorProjectAssignments
using DataFrames
using CSV

using SeniorProjectAssignments: match

using Random
Random.seed!(151)

survey_dataframe = CSV.read("survey_011.csv", DataFrame, header=1, datarow=4)
project_dataframe = CSV.read("projects_011.csv", DataFrame)

projects = process_projects(project_dataframe, min=10, max=12)
@show pnames = [p.id for p in projects]

students, groups = process_survey(survey_dataframe, pnames)

CSV.write("processed_students.csv", convert(DataFrame, students))

@show groups

sdf = convert(DataFrame, students)

force = ["Stephen Peng"=>"Avalanche Drone: Snowpack Measurement"]

# push!(force, ["Michelle

m = match(students, projects, groups, force=force)
sort!(m, :project)
CSV.write("output_011.csv", m)
display(m)
display([r=>sum(skipmissing(m[:,:rank]).==r) for r in sort(unique(m[:,:rank]))])


text:
import ..SnakeAI
import ..SnakeAI: play_step!, reset!

const WIDTH = SnakeAI.WIDTH
const HEIGHT = SnakeAI.HEIGHT
const BACKGROUND = colorant"black"
const FRAME_BUFFER = 8     # Updates UI at every n frames

game = SnakeAI.Game()
step = 1

function draw()
    # Drawing the snake
    for p in game.snake.body
        outside = Rect(p.x, p.y, SnakeAI.BLOCK_SIZE, SnakeAI.BLOCK_SIZE)
        inside = Rect(p.x + 2, p.y + 2, SnakeAI.BLOCK_SIZE - 4, SnakeAI.BLOCK_SIZE - 4)
        draw(outside, colorant"#00262a", fill=true)
        draw(inside, colorant"#00717f", fill=true)

    end

    # Drawing the food
    f = Rect(game.food.x, game.food.y, SnakeAI.BLOCK_SIZE, SnakeAI.BLOCK_SIZE)
    draw(f, colorant"#7F0071", fill=true)
end

function update(g::GameZero.Game)
    global game, step

    # Movement keys
    if g.keyboard.UP && game.snake.last_move != SnakeAI.DOWN
        game.direction = SnakeAI.UP
    elseif g.keyboard.DOWN && game.snake.last_move != SnakeAI.UP
        game.direction = SnakeAI.DOWN
    elseif g.keyboard.LEFT && game.snake.last_move != SnakeAI.RIGHT
        game.direction = SnakeAI.LEFT
    elseif g.keyboard.RIGHT && game.snake.last_move != SnakeAI.LEFT
        game.direction = SnakeAI.RIGHT
    end

    if step == FRAME_BUFFER
        step = 1
        _, game_over, score = play_step!(game)

        if game_over
            println(
                """
                ==================
                GAME OVER!
                Final Score: $(score)
                ==================\n""")
            # Reset the game
            reset!(game)
        end
    else
        step += 1
    end
end


text:
mutable struct Trace{Tspl<:AbstractSampler, Tvi<:AbstractVarInfo, Tmodel<:Model}
    model::Tmodel
    spl::Tspl
    vi::Tvi
    task::Task

    function Trace{SampleFromPrior}(model::Model, spl::AbstractSampler, vi::AbstractVarInfo)
        return new{SampleFromPrior,typeof(vi),typeof(model)}(model, SampleFromPrior(), vi)
    end
    function Trace{S}(model::Model, spl::S, vi::AbstractVarInfo) where S<:Sampler
        return new{S,typeof(vi),typeof(model)}(model, spl, vi)
    end
end

function Base.copy(trace::Trace)
    vi = deepcopy(trace.vi)
    res = Trace{typeof(trace.spl)}(trace.model, trace.spl, vi)
    res.task = copy(trace.task)
    return res
end

# NOTE: this function is called by `forkr`
function Trace(f::Function, m::Model, spl::AbstractSampler, vi::AbstractVarInfo)
    res = Trace{typeof(spl)}(m, spl, deepcopy(vi));
    # CTask(()->f());
    res.task = CTask( () -> begin res=f(); produce(Val{:done}); res; end )
    if res.task.storage === nothing
        res.task.storage = IdDict()
    end
    res.task.storage[:turing_trace] = res # create a backward reference in task_local_storage
    return res
end
function Trace(m::Model, spl::AbstractSampler, vi::AbstractVarInfo)
    res = Trace{typeof(spl)}(m, spl, deepcopy(vi));
    # CTask(()->f());
    reset_num_produce!(res.vi)
    res.task = CTask( () -> begin vi_new=m(vi, spl); produce(Val{:done}); vi_new; end )
    if res.task.storage === nothing
        res.task.storage = IdDict()
    end
    res.task.storage[:turing_trace] = res # create a backward reference in task_local_storage
    return res
end

# step to the next observe statement, return log likelihood
Libtask.consume(t::Trace) = (increment_num_produce!(t.vi); consume(t.task))

# Task copying version of fork for Trace.
function fork(trace :: Trace, is_ref :: Bool = false)
    newtrace = copy(trace)
    is_ref && set_retained_vns_del_by_spl!(newtrace.vi, newtrace.spl)
    newtrace.task.storage[:turing_trace] = newtrace
    return newtrace
end

# PG requires keeping all randomness for the reference particle
# Create new task and copy randomness
function forkr(trace :: Trace)
    newtrace = Trace(trace.task.code, trace.model, trace.spl, deepcopy(trace.vi))
    newtrace.spl = trace.spl
    reset_num_produce!(newtrace.vi)
    return newtrace
end

current_trace() = current_task().storage[:turing_trace]

const Particle = Trace

"""
Data structure for particle filters
- effectiveSampleSize(pc :: ParticleContainer)
- normalise!(pc::ParticleContainer)
- consume(pc::ParticleContainer): return incremental likelihood
"""
mutable struct ParticleContainer{T<:Particle, F}
    model::F
    vals::Vector{T}
    # logarithmic weights (Trace) or incremental log-likelihoods (ParticleContainer)
    logWs::Vector{Float64}
    # log model evidence
    logE::Float64
    # helpful for rejuvenation steps, e.g. in SMC2
    n_consume::Int
end

ParticleContainer(model, particles::Vector{<:Particle}) =
    ParticleContainer(model, particles, zeros(length(particles)), 0.0, 0)

Base.collect(pc::ParticleContainer) = pc.vals
Base.length(pc::ParticleContainer) = length(pc.vals)
Base.@propagate_inbounds Base.getindex(pc::ParticleContainer, i::Int) = pc.vals[i]

# registers a new x-particle in the container
function Base.push!(pc::ParticleContainer, p::Particle)
    push!(pc.vals, p)
    push!(pc.logWs, 0.0)
    pc
end

# clones a theta-particle
function Base.copy(pc::ParticleContainer)
    # fork particles
    vals = eltype(pc.vals)[fork(p) for p in pc.vals]

    # copy weights
    logWs = copy(pc.logWs)

    ParticleContainer(pc.model, vals, logWs, pc.logE, pc.n_consume)
end

# run particle filter for one step, return incremental likelihood
function Libtask.consume(pc :: ParticleContainer)
    # normalisation factor: 1/N
    z1 = logZ(pc)
    n = length(pc)

    particles = collect(pc)
    num_done = 0
    for i=1:n
        p = particles[i]
        score = Libtask.consume(p)
        if score isa Real
            score += getlogp(p.vi)
            resetlogp!(p.vi)
            increase_logweight(pc, i, Float64(score))
        elseif score == Val{:done}
            num_done += 1
        else
            error("[consume]: error in running particle filter.")
        end
    end

    if num_done == n
        res = Val{:done}
    elseif num_done != 0
        error("[consume]: mis-aligned execution traces, num_particles= $(n), num_done=$(num_done).")
    else
        # update incremental likelihoods
        z2 = logZ(pc)
        res = increase_logevidence(pc, z2 - z1)
        pc.n_consume += 1
        # res = increase_loglikelihood(pc, z2 - z1)
    end

    res
end

# compute the normalized weights
getweights(pc::ParticleContainer) = softmax(pc.logWs)

# compute the log-likelihood estimate, ignoring constant term ``- \log num_particles``
logZ(pc::ParticleContainer) = logsumexp(pc.logWs)

# compute the effective sample size ``1 / ∑ wᵢ²``, where ``wᵢ```are the normalized weights
function effectiveSampleSize(pc :: ParticleContainer)
    Ws = getweights(pc)
    return inv(sum(abs2, Ws))
end

increase_logweight(pc :: ParticleContainer, t :: Int, logw :: Float64) =
    (pc.logWs[t]  += logw)

increase_logevidence(pc :: ParticleContainer, logw :: Float64) =
    (pc.logE += logw)

function resample!(
    pc :: ParticleContainer,
    randcat :: Function = Turing.Inference.resample_systematic,
    ref :: Union{Particle, Nothing} = nothing;
    weights = getweights(pc)
)
    # check that weights are not NaN
    @assert !any(isnan, weights)

    # sample ancestor indices
    n = length(pc)
    nresamples = ref === nothing ? n : n - 1
    indx = randcat(weights, nresamples)

    # count number of children for each particle
    num_children = zeros(Int, n)
    @inbounds for i in indx
        num_children[i] += 1
    end

    # fork particles
    particles = collect(pc)
    children = similar(particles)
    j = 0
    @inbounds for i in 1:n
        ni = num_children[i]

        if ni > 0
            # fork first child
            pi = particles[i]
            isref = pi === ref
            p = isref ? fork(pi, isref) : pi
            children[j += 1] = p

            # fork additional children
            for _ in 2:ni
                children[j += 1] = fork(p, isref)
            end
        end
    end

    if ref !== nothing
        # Insert the retained particle. This is based on the replaying trick for efficiency
        # reasons. If we implement PG using task copying, we need to store Nx * T particles!
        @inbounds children[n] = ref
    end

    # replace particles and log weights in the container with new particles and weights
    pc.vals = children
    pc.logWs = zeros(n)

    pc
end

struct ResampleWithESSThreshold{R, T<:Real}
    resampler::R
    threshold::T
end

function ResampleWithESSThreshold()
    ResampleWithESSThreshold(Turing.Inference.resample_systematic)
end
ResampleWithESSThreshold(resampler) = ResampleWithESSThreshold(resampler, 0.5)

function resample!(
    pc::ParticleContainer,
    resampler::ResampleWithESSThreshold,
    ref::Union{Particle,Nothing} = nothing;
    weights = getweights(pc)
)
    # Compute the effective sample size ``1 / ∑ wᵢ²`` with normalized weights ``wᵢ``
    ess = inv(sum(abs2, weights))

    if ess ≤ resampler.threshold * length(pc)
        resample!(pc, resampler.resampler, ref; weights = weights)
    end

    pc
end


text:
# REPL experiments

#= --- how to use HTTP ---

julia> using HTTP

julia> url = "https://hacker-news.firebaseio.com/v0/topstories.json";

julia> response = HTTP.request("GET", url);

julia> typeof(response)
HTTP.Messages.Response

julia> response.body
4357-element Array{UInt8,1}:
 0x5b
 0x32
 0x31
 0x36
 0x37
 0x36
    ⋮

julia> String(response.body)
"[21676252,21676933,21677389,21674752,21675456,21676027,21671304,
21676606,21674599,21675498,21676923,21674729,21669787,21669530,21676543,
21675894,21676384,21677533,21677556,21674610,21667223,21675030,
21675228,21675280,21671799,..."
=#

#= --- putting it altogether ---

julia> using HTTP, JSON3

julia> url = "https://hacker-news.firebaseio.com/v0/topstories.json";

julia> response = HTTP.request("GET", url);

julia> JSON3.read(String(response.body))
484-element JSON3.Array{Int64,Base.CodeUnits{UInt8,String},Array{UInt64,1}}:
 21676252
 21676933
 21677389
 21674752
 21675456
 21676027
        ⋮
=#

#= --- getting a story ---

julia> fetch_story(21676252)
Story("kkm", 383, 535, 1575218518, 21676252,
"The world needs more search engines",
[21676507, 21676549, 21676457, 21676352, 21677053, 21676790, 21676616,
21676315, 21676385, 21677210  …  21676442, 21676505, 21676858, 21676598,
21676736, 21677029, 21676919, 21676853, 21676397, 21676703],
"https://www.0x65.dev/blog/2019-12-01/the-world-needs-cliqz-the-world-needs-more-search-engines.html")
=#

#= --- average score ---
julia> average_score()
2019-12-01T12:44:05.653 Found 484 stories
2019-12-01T12:44:06.453 Fetched 10 story details
2019-12-01T12:44:06.453 Average score = 131.3
131.3
=#


text:
using Avalon

include("model.jl")
include("imagefolder.jl")

# this should point to a real folder containing subfolders with images
# each subfolder name will be taken as class name, while contained images
# will be used for training
# images can have any size, but will be resized to (100, 100) during loading
const IMAGE_DATASET_PATH = "~/data/food-101-sample"
# number of subfolders with images
const NUM_CLASSES = 5

# TODO:
# 4. load pretrained net
# 5. make wrappers like resnet18(), etc.

################################################################################
#                         TRAINING FROM SCRATCH                                #
################################################################################

# note that this method doesn't do validation or testing, but is only useful to make sure
# value of loss function is reduced over time
function Avalon.fit!(m::ResNet, dataset::ImageFolder, loss_fn;
              n_epochs=10, batch_size=100, opt=Adam(;lr=1e-3), device=CPU(), report_every=1)
    f = (m, x, y) -> loss_fn(m(x), y)
    num_batches = length(dataset) // batch_size
    for epoch in 1:n_epochs
        epoch_loss = 0
        for (i, (x, y)) in enumerate(Avalon.batchiter(dataset, sz=batch_size))
            x = device(copy(x))
            y = device(copy(y))
            loss, g = grad(f, m, x, y)
            update!(opt, m, g[2])
            epoch_loss += loss
            println("  batch loss = $loss")
        end
        if epoch % report_every == 0
            println("Epoch $epoch: avg_loss=$(epoch_loss / num_batches)")
        end
    end
    return m
end

function train_from_scretch()
    device = best_available_device()
    dataset = ImageFolder(expanduser(IMAGE_DATASET_PATH), transform=img -> imresize(img, (100, 100)))
    block = Bottleneck
    block_nums = [2, 2, 2, 2]
    stride = 1
    m = ResNet(block, block_nums, zero_init_residual=true, num_classes=NUM_CLASSES) |> device
    loss_fn = CrossEntropyLoss()
    fit!(m, dataset, loss_fn; device=device, batch_size=5)
end

################################################################################
#                         USING PRETRAINED MODEL                               #
################################################################################

# Will be added once ONNX import/export is implemented


text:
module OrthogonalPolynomials

# 0. Our mission
# Dear codes, should you choose to accept it, is to learn and implement the following
# mathematical in Julia, whilst learning the magic of multiple dispatch as a design
# paradigm. Ready? Steady? Go!
# If you are a beginner and this code just looks like squiggles, worry not!
# There is a video tutorial online at my youtube channel, BrainRPG.
# Happy Hacking!
# - Miguel Raz Guzman Macedo
# http://people.math.sfu.ca/~cbm/aands/page_789.jpg

# 1. First design - Laguerre only
# Note the _names with underscore are used as convention for package internals.
_d(n, α=0) = binomial(n+α, n)
_b(n, m) = n-m+1
_c(α, m) = m*(α+m)
_k(n, α=0) = [-_b(n,i)*inv(_c(α, i)) for i in 1:n]
_f(x) = x
_α = 0

_a(x, n, ks = _k(n,α), i=0) = i == n ? :(1) : return :(muladd( $(ks[i+1]*_f(x)), $((_a)(x, n, ks, i+1)) , 1))

# 2. Macro function as a for loop
macro _a(x,n,ks=k(n,α))
    ex = 1
    for i in n-1 : -1 : 0
        ex = :(muladd( $(ks[i+1]*_f(x)), $ex, 1))
    end
    return :($ex)
end

# 2. Second design, now with dispatching on types.
# 2.1 FIRST, we will define all the infrastructure for dispatching for Laguerre.
# 2.2 THEN, we will extend that design by adding dispatches for the Hermite case.
# 2.3 LAST, we will extend for all the other cases.

# 2.1 FIRST, the Laguerre case.
abstract type OP end
struct Laguerre{α,n} <: OP end
params(a :: T) where T = Tuple(T.parameters) # Credit to @yingboma
Laguerre(α, n :: Int) = Laguerre{α, n}()
Laguerre(n :: Int) = n >= 0 ? Laguerre{0,n}() : throw("The degree n must be non-negative.")
# Generalized binomial, useful too in the Jacobi case.
function Base.binomial(α :: T,k :: Int) where T<:AbstractFloat
    res = 1
    for i in 0:k-1
        res = res * (α - i)
    end
    return res/factorial(k)
end

d(p :: T, x) where T<:Laguerre = d(p)
function d(p :: T) where T<:Laguerre
    α,n = params(p)
    return binomial(α + n, n)
end

function b(p :: T, m) where T<:Laguerre
    α,n = params(p)
    return n-m+1
end

function c(p :: T, m) where T<:Laguerre
    α, n = params(p)
    return m*(α + m)
end

f(p :: T, x) where T<:Laguerre = x

# This now works for all Orthogonal Polynomials!
function k(p :: T, x) where T<:OP
    n = params(p)[end]
    return [-b(p, i)*inv(c(p, i)) for i in 1:n]
end

# Notice this works for all types Orthogonal Polynomials!
# Note: It doesn't return a specialized function yet - try to run it with @code_llvm
# We will make this specialized in the Third stage of design, after we add all the dispatches
# for all the polynomials.
function a(p :: T, x, m = params(p)[end]) where T<:OP
    res = oneunit(x)
    ks = k(p, x)

    for i in m-1: -1 : 0
        res = muladd(res, ks[i+1]*f(p, x),1)
    end

    return res * d(p, x)
end

# function a(p :: T, x, i=

# 2.2 SECOND - add a little abstract type machinery to handle dispatches for even/odd cases.
abstract type EvenOP <: OP end
abstract type OddOP <: OP end
struct Hermite{n} <: OP end
struct HermiteEven{n} <: EvenOP end
struct HermiteOdd{n} <: OddOP end

Hermite{n}() where n = iseven(n) ? HermiteEven{n ÷ 2}() : HermiteOdd{(n-1) ÷ 2}()
Hermite(n) =  n>= 0 ? Hermite{n}() : throw("The degree n must be non-negative.")

#Heven
function d(p :: T, x) where T<:HermiteEven
    n = params(p)[1]
    return (-1) ^ n * (factorial(2n)/factorial(n))
end

function b(p :: T, m) where T<:HermiteEven
    n = params(p)[1]
    return  2*(n - m + 1)
end

# This now covers half the cases!
c(p :: T, m) where T<:EvenOP = m*(2m - 1)

# This catches almost every case!
f(p :: T, x) where T<:OP = x^2

#Hodd
function d(p :: T, x) where T<:HermiteOdd
    n = params(p)[1]
    return (-1) ^ n * (factorial(2n +1)/factorial(n))*2x
end

function b(p :: T, m) where T<:HermiteOdd
    n = params(p)[1]
    return  2*(n - m + 1)
end

# This catches almost every case left over! We are only missing the Jacobi case.
c(p :: T, m) where T<:OddOP = m * (2m + 1)

# 2.3 THIRD, we now add dispatches to all the other cases.
struct Legendre{n} <: OP end
struct LegendreEven{n} <: EvenOP end
struct LegendreOdd{n} <: OddOP end
Legendre{n}() where n = iseven(n) ? LegendreEven{n ÷ 2}() : LegendreOdd{(n-1) ÷ 2}()
Legendre(n :: Int) = n>=0 ? Legendre{n}() : throw("The degree n must be non-negative.")

# PEven
# Note! We only have to define two functions per case here, because the other cases already got caught by
# some adequate planning and dispatching.
function d(p :: T, x) where T<:LegendreEven
    n = params(p)[1]
    return (-1/4)^n * binomial(2n,n)
end

function b(p :: T, m) where T<:LegendreEven
    n = params(p)[1]
    return (n - m + 1) * (2n + 2m -1)
end

# POdd
function d(p :: T, x) where T<:LegendreOdd
    n = params(p)[1]
    return (-1/4)^n * binomial(2n+1,n) * (n+1) * x
end

function b(p :: T, m) where T<:LegendreOdd
    n = params(p)[1]
    return (n - m + 1) * (2n + 2m + 1)
end

# TEven
struct ChebyshevFirstKind{n} <: OP end
struct ChebyshevFirstKindEven{n} <: EvenOP end
struct ChebyshevFirstKindOdd{n} <: OddOP end
ChebyshevFirstKind{n}() where n = iseven(n) ? ChebyshevFirstKindEven{n ÷ 2}() : ChebyshevFirstKindOdd{(n-1) ÷ 2}()
ChebyshevFirstKind(n :: Int) = n>=0 ? ChebyshevFirstKind{n}() : throw("The degree n must be non-negative")

function d(p :: T, x) where T<:ChebyshevFirstKindEven
    n = params(p)[1]
    return (-1)^n
end

function b(p :: T, m) where T<:ChebyshevFirstKindEven
    n = params(p)[1]
    return 2(n-m+1)*(n+m-1)
end

# TOdd
function d(p :: T, x) where T<: ChebyshevFirstKindOdd
    n = params(p)[1]
    return (-1^n)*(2n+1)*x
end

function b(p ::T, m) where T<: ChebyshevFirstKindOdd
    n = params(p)[1]
    return  2*(n-m+1)*(n+m)
end

# UEven
struct ChebyshevSecondKind{n} <: OP end
struct ChebyshevSecondKindEven{n} <: EvenOP end
struct ChebyshevSecondKindOdd{n} <: OddOP end
ChebyshevSecondKind{n}() where n = iseven(n) ? ChebyshevSecondKindEven{n ÷ 2}() : ChebyshevSecondKindOdd{(n-1) ÷ 2}()
ChebyshevSecondKind(n :: Int) = n>=0 ? ChebyshevSecondKind{n}() : throw("The degree n must be non-negative")

function d(p :: T, x) where T<:ChebyshevSecondKindEven
    n = params(p)[1]
    return (-1)^n
end

function b(p :: T, m) where T<:ChebyshevSecondKindEven
    n = params(p)[1]
    return 2*(n-m+1)*(n+m)
end

# UOdd
function d(p :: T, x) where T<:ChebyshevSecondKindOdd
    n = params(p)[1]
    return (-1)^n * 2(n+1)*x
end

function b(p :: T, m) where T<:ChebyshevSecondKindOdd
    n = params(p)[1]
    return 2*(n-m+1)*(n+m)
end

# CEven
struct Gegenbauer{α, n} <: OP end
struct GegenbauerEven{α, n} <: EvenOP end
struct GegenbauerOdd{α, n} <: OddOP end
function Gegenbauer(α, n :: Int)
     n >= 0 && α >= 0 ? nothing : throw("The degree n must be non-negative")
     iseven(n) ? GegenbauerEven{α, n ÷ 2}() : GegenbauerOdd{α, (n-1) ÷ 2}()
end
Gegenbauer(n :: Int) = Gegenbauer(0, n)

# We need to define a pochhammer symbol here!
function poch(a,::Val{N}) where N
            res = a; for i in Base.OneTo(N-1)
                res = res * (a+i)
            end
        return res
end
poch(a, ::Val{0}) = one(a)
poch(a, ::Val{1}) = a

function d(p :: T, x = 0) where T<:GegenbauerEven
    α, n  = params(p)
    return (-1)^n * (poch(α,Val(n))/factorial(n))
end

function b(p :: T, m) where T<:GegenbauerEven
    α, n = params(p)
    return 2*(n-m+1)*(α+n+m-1)
end

# COdd
function d(p :: T, x) where T<:GegenbauerOdd
    α, n = params(p)
    return (-1)^n * (poch(α,Val(n+1))/factorial(n))*2x
end

function b(p :: T, m) where T<:GegenbauerOdd
    α, n = params(p)
    return  2*(n-m+1)*(α + n + m)
end

# Jacobi
struct Jacobi{α, β, n} <: OP end
Jacobi(α, β, n :: Int) = α >= 0 && β >= 0 && n >= 0 ? Jacobi{α, β, n}() : throw("The parameters must be non-negative")

function d(p :: T, x) where T<:Jacobi
    α, β, n = params(p)
    return binomial(n + α, n)
end

function b(p :: T, m) where T<:Jacobi
    α, β, n = params(p)
    return (n - m + 1) * (α + β + n + m)
end

function c(p :: T, m) where T<:Jacobi
    α, β, n = params(p)
    return 2m * (α + m)
end

function f(p :: T, x) where T<:Jacobi
    return 1 - x
end

# Pass: Laguerre, ChebyFE, ChebySE,
# Check Jacobi-binomial, HermiteEven, HermiteOdd, LegendreE, LegendreO, ChebySO, ChebyFO, GegenbauerE, GegenbauerO
# Strat, look at the a(j,2,0) and k(j,.5) methods
export d,b,c,k,f,a
export Jacobi                                                                   # J(α,β,n)
export GegenbauerEven, GegenbauerOdd, Gegenbauer                                # C(α,n)
export ChebyshevSecondKindEven, ChebyshevSecondKindOdd, ChebyshevSecondKind     # U(n)
export ChebyshevFirstKindEven, ChebyshevFirstKindOdd, ChebyshevFirstKind        # T(n)
export Legendre                                                                 # L(n)
export HermiteOdd, HermiteEven, Hermite                                         # H(n)
export Laguerre

end # module


text:
@testset "WorldClim Weather" begin

    using SimpleSDMDataSources: rastername, zipurl, zipname, zippath

    raster_file = joinpath(ENV["ECODATASOURCES_PATH"], "WorldClim/Weather/prec/wc2.1_2.5m_prec_2001-01.tif")
    @test rasterpath(WorldClim{Weather}, :prec, Date(2001, 1)) == raster_file
    @test rastername(WorldClim{Weather}, :prec, Date(2001, 1)) == "wc2.1_2.5m_prec_2001-01.tif"

    zip_file = joinpath(ENV["ECODATASOURCES_PATH"], "WorldClim/Weather/zips/wc2.1_2.5m_prec_2010-2018.zip")
    @test zippath(WorldClim{Weather}, :prec, Date(2010)) == zip_file
    @test zipurl(WorldClim{Weather}, :prec, Date(2010)) ==
        "https://biogeo.ucdavis.edu/data/worldclim/v2.1/hist/wc2.1_2.5m_prec_2010-2018.zip"
    @test zipname(WorldClim{Weather}, :prec, Date(2010)) ==
        "wc2.1_2.5m_prec_2010-2018.zip"

    # These files are 3GB each. Probably too big to test in CI.
    # Not sure what to do about that.
    # download_raster(WorldClim{Weather}, :prec; dates=Date(2001):Month(1):Date(2001, 12))
end


text:
@testset "Basic 3D Orientation Tests" begin
	w = (0f0,0f0); x = (1,1,2,2); y = (2,2,1,1); z = (1,2,1,2);
	p = (2f0,2f0,2f0); q = (3,4,3); r = (3f0,1f0,1f0,1f0); s = (2,3,4); t = (-4f0,-1f0,-1f0,-7f0);

	@test plane(x,y,z) == (-3, 3, 3, -3)
	@test dot(q,s) == 30
	@test toPoint3H(q) == (1,q[1],q[2],q[3])
	@test abovePl(t,plane(x,y,z))
	@test !belowPl(toPoint3H(p),plane(x,y,z))
	@test plane(p, q, s) == plane(toPoint3H(p), toPoint3H(q), toPoint3H(s))

#	pts = [Point3H(1f0, a, b, c) for a in [2f0 4f0], b in [2f0 4f0], c in [2f0 4f0]]

	a = (1f0,3f0,3f0,3f0); b = (1f0,2f0,2f0,2f0); c = (1f0,2f0,2f0,4f0); d = (1f0,4f0,2f0,2f0); e = (1f0,2f0,4f0,2f0); f = (1f0,2f0,4f0,4f0); g = (1f0,4f0,4f0,2f0);
	h = (1f0); i = (1f0);

#	@test CH([a;b;c;d;e;f;g])

end;

max_stars_repo_name: snoeyink/CGinJulia

max_issues_repo_name: snoeyink/CGinJulia

max_forks_repo_name: snoeyink/CGinJulia