text:
using Flux
using Random
using EcologicalNetworks
using StatsPlots
using MultivariateStats
using ProgressMeter
using BSON: @save, @load
using Statistics
using TSne
using DataFrames
using CSV: CSV
using ParallelKMeans
using JSON

# We will need some folders to store the results
for f in ["figures", "artifacts"]
    isdir(f) || mkdir(f)
end

# We set a chill little seed for the random number generator
Random.seed!(420)

# Default theme for the plots
theme(:bright)

# Get the data from the Web of Life database, and merge them into a metaweb M
ids = map(i -> i.ID, filter(i -> contains(i.Reference, "Hadfield"), web_of_life()))
B = convert.(BipartiteNetwork, web_of_life.(ids))
M = reduce(∪, B)

# We convert the metaweb into a unipartite network (easier for PCA)
# and create a mirrored version K
N = convert(UnipartiteNetwork, M)
K = EcologicalNetworks.mirror(N)

# To get the latent traits, we use a PPCA on the mirrored edges matrix,
# as this works with sparse data, and captures in/out interactions
pc = fit(PPCA, Float64.(Array(K.edges)))
pr = MultivariateStats.transform(pc, Float64.(Array(K.edges)))

# Plot of the features for axis 1 and 2
scatter(
    pr[:, 1],
    pr[:, 2];
    frame=:origin,
    alpha=0.6,
    lab=false,
    aspectratio=1,
    dpi=600,
    size=(400, 400),
)
xaxis!("PC1")
yaxis!("PC2")
savefig("figures/features.png")

# We only retain nf features for the host, and nf features for the parasite
nf = 15

# These objects will store some temporary data
cooc = zeros(Bool, prod(size(M)))                  # H/P co-occurrence?
labels = zeros(Bool, prod(size(M)))                # H/P interaction?
features = zeros(Float64, (2 * nf, prod(size(M))))   # H/P latent traits

# We then move through the metaweb step by step
cursor = 0
for i in species(M; dims=1), j in species(M; dims=2)
    global cursor += 1
    # Interaction in the metaweb?
    labels[cursor] = M[i, j]
    # Values in the PCA space
    p_i = findfirst(i .== species(N))
    p_j = findfirst(j .== species(N))
    features[1:nf, cursor] .= pr[1:nf, p_i]
    features[(nf + 1):end, cursor] .= pr[1:nf, p_j]
    # Co-occurrence?
    for b in B
        if i in species(b)
            if j in species(b)
                cooc[cursor] = true
            end
        end
    end
end

# We only work on the species pairs that DO co-occur
kept = findall(cooc)
x = Float32.(copy(features[:, kept])) # Latent traits
y = Matrix(hcat(labels[kept])')       # Interaction bit

# This will setup the model for training,
# with 20% of data leftover for validation
training_size = convert(Int64, floor(size(x, 2) * 0.8))
train = sort(sample(1:size(x, 2), training_size; replace=false))
test = filter(i -> !(i in train), 1:size(x, 2))

# The final datasets we will use are as follows:
data = (x[:, train], y[:, train])
data_test = (x[:, test], y[:, test])

# The model is specified as a series of chained layers
m = Chain(
    # The first layer uses relu
    Dense(2nf, ceil(Int64, 2.4nf), relu),
    # The first dropout rate is 0.8 as we're using a lot of features
    Dropout(0.8),
    # All other layers are sigmoid with 0.6 dropout rate
    Dense(ceil(Int64, 2.4nf), ceil(Int64, 1.5nf), σ),
    Dropout(0.6),
    Dense(ceil(Int64, 1.5nf), ceil(Int64, 0.8nf), σ),
    Dropout(0.6),
    # The last layer has a single bit! P(parasite → host)
    Dense(ceil(Int64, 0.8nf), 1, σ),
)

# This step is a little long, but it's basically a plot of the untrained network
include("plotnetwork.jl")
savefig("figures/network-untrained.png")

# We use MSE (averaged over all interactions) as the loss function
loss(x, y) = Flux.mse(m(x), y)
ps = Flux.params(m)
opt = ADAM() # ADAM is the optimizer, used with default parameters

# To avoid overfitting, we will present the training data in batches
# n_batches is the number of batches
# batch_size is the number of interactions per batch
# This gives us a good guarantee that the model won't revisit the same interactions
# in the same order, and so this should prevent overfitting
n_batches, batch_size = 50000, 64

# Because we run a lot of epocs, we only save the results (loss)
# every mat_at timesteps - 500 is reasonable in terms of performance/dataviz
# tradeoff. Note that we always save after the first batch.
mat_at = 500
epc = mat_at:mat_at:n_batches
epc = vcat(1, epc...)

# These two arrays are storing the loss values
trainlossvalue = zeros(Float64, n_batches)
testlossvalue = zeros(Float64, n_batches)

# This is the main training loop
@showprogress for i in 1:n_batches
    # We pick a random batch out of the training set
    ord = sample(train, batch_size; replace=false)
    data_batch = (x[:, ord], y[:, ord])
    # If the training batch is too unbalanced, we draw another one
    while sum(data_batch[2]) < ceil(Int64, 0.25batch_size)
        ord = sample(train, batch_size; replace=false)
        data_batch = (x[:, ord], y[:, ord])
    end
    # This trains the model
    Flux.train!(loss, ps, [data_batch], opt)
    # We only save the loss at the correct interval
    if i in epc
        trainlossvalue[i] = loss(data...)
        testlossvalue[i] = loss(data_test...)
    end
end

# This is a simple plot of the loss values over time, to aid with the diagnosis of
# overfitting -- if the two curves diverge a lot, especially with the testing
# loss going back up, the model is mad sus
plot(epc, trainlossvalue[epc]; lab="Training", dpi=600, frame=:box, size=(400, 400))
plot!(epc, testlossvalue[epc]; lab="Testing")
xaxis!("Epoch")
yaxis!("Loss (MSE)")

# This saves the figure to the correct folder
savefig("figures/loss.png")

# The next bit is about thresholding - we will specifically figure out which
# value of the interaction bit is high enough to say "there is an interaction".
# We will pick the value maximizing Youden's J.

# We get the predictions and observations for the testing dataset
predictions = vec(m(data_test[1]))
obs = vec(data_test[2])

# And we pick thresholds in the [0,1] range
thresholds = range(0.0, 1.0; length=500)

# All this is going to be the components of the adjacency matrix at a given threshold
tp = zeros(Float64, length(thresholds))
fp = zeros(Float64, length(thresholds))
tn = zeros(Float64, length(thresholds))
fn = zeros(Float64, length(thresholds))

# Main loop to get the four components
for (i, thr) in enumerate(thresholds)
    pred = vec(predictions .>= thr)
    tp[i] = sum(pred .& obs)
    tn[i] = sum(.!(pred) .& (.!obs))
    fp[i] = sum(pred .& (.!obs))
    fn[i] = sum(.!(pred) .& obs)
end

# Total number of cases
n = tp .+ fp .+ tn .+ fn

# Diagnostic measures
tpr = tp ./ (tp .+ fn)
fpr = fp ./ (fp .+ tn)
tnr = tn ./ (tn .+ fp)
fnr = fn ./ (fn .+ tp)
acc = (tp .+ tn) ./ (n)
racc = ((tn .+ fp) .* (tn .+ fn) .+ (fn .+ tp) .* (fp .+ tp)) ./ (n .* n)
bacc = ((tp ./ (tp .+ fn)) .+ (tn ./ (fp .+ tn))) ./ 2.0
J = (tp ./ (tp .+ fn)) + (tn ./ (tn .+ fp)) .- 1.0
κ = (acc .- racc) ./ (1.0 .- racc)
threat = tp ./ (tp .+ fn .+ fp)
fomrate = fn ./ (fn .+ tn)
fdirate = fp ./ (fp .+ tp)
ppv = tp ./ (tp .+ fp)
npv = tn ./ (tn .+ fn)

# This bit is here to get the AUC
dx = [reverse(fpr)[i] - reverse(fpr)[i - 1] for i in 2:length(fpr)]
dy = [reverse(tpr)[i] + reverse(tpr)[i - 1] for i in 2:length(tpr)]
AUC = sum(dx .* (dy ./ 2.0))

# Final thresholding results - we pick the value maximizing Youden's J
thr_index = last(findmax(J))
thr_final = thresholds[thr_index]

# Save the validation measures to a plot
validation = Dict{String,Float64}()
validation["ROC-AUC"] = AUC
validation["Threat score"] = threat[thr_index]
validation["Youden's J"] = J[thr_index]
validation["True Positive Rate"] = tpr[thr_index]
validation["True Negative Rate"] = tnr[thr_index]
validation["False Positive Rate"] = fpr[thr_index]
validation["False Negative Rate"] = fnr[thr_index]
validation["Kappa"] = κ[thr_index]
validation["Accuracy"] = acc[thr_index]
validation["Accuracy (random)"] = racc[thr_index]
validation["Accuracy (balanced)"] = bacc[thr_index]
validation["False Discovery Rate"] = fdirate[thr_index]
validation["False Omission Rate"] = fomrate[thr_index]
validation["Positive Predictive Value"] = ppv[thr_index]
validation["Negative Predictive Value"] = npv[thr_index]

open("artifacts/validation.json", "w") do f
    JSON.print(f, validation, 4)
end

# ROC plot
plot(fpr, tpr; aspectratio=1, frame=:box, lab="", dpi=600, size=(400, 400))
scatter!([fpr[thr_index]], [tpr[thr_index]]; lab="", c=:black)
plot!([0, 1], [0, 1]; c=:grey, ls=:dash, lab="")
xaxis!("False positive rate", (0, 1))
yaxis!("True positive rate", (0, 1))

# We also save this one to a file
savefig("figures/roc-auc.png")

# Precision-Recall plot
plot(tpr, ppv; aspectratio=1, frame=:box, lab="", dpi=600, size=(400, 400))
scatter!([tpr[thr_index]], [ppv[thr_index]]; lab="", c=:black)
plot!([0, 1], [1, 0]; c=:grey, ls=:dash, lab="")
xaxis!("True positive rate", (0, 1))
yaxis!("Positive predictive value", (0, 1))

# We also save this one to a file
savefig("figures/precision-recall.png")

# We save the network itself as a bson object
@save "artifacts/netpred.bson" m

# And finally we plot it one more time
include("plotnetwork.jl")
savefig("figures/network-trained.png")

# Predictions and co-occurrences
pr = vec(m(features))    # Raw predictions
imp = pr .>= thr_final   # Larger than threshold
nwi = imp .& .!labels    # New

# We can now move on to the actual imputation
P = copy(M)

# This will store the interactions
new_interactions = DataFrame(; parasite=String[], host=String[], cooc=Bool[], p=Float64[])

# We will fill the network, but also add the interaction values to a DataFrame
cursor = 0
for i in eachindex(species(P; dims=1)), j in eachindex(species(P; dims=2))
    global cursor += 1
    if imp[cursor]
        P[i, j] = true
        if !M[i, j]
            push!(
                new_interactions,
                (species(P; dims=1)[i], species(P; dims=2)[j], cooc[cursor], pr[cursor]),
            )
        end
    end
end

# We save the new interactions file to disk
CSV.write("artifacts/predictions.csv", sort(new_interactions, :p; rev=true))

# We can do simple diagnostic plots, like change in degree
dh = [(degree(N)[s], degree(P)[s]) for s in species(M; dims=2)]
dp = [(degree(N)[s], degree(P)[s]) for s in species(M; dims=1)]

# We can plot this to check that rich do not get richer
scatter(
    dh;
    frame=:box,
    dpi=600,
    legend=:bottomright,
    aspectratio=1,
    label="Hosts",
    ms=6,
    alpha=0.6,
)
scatter!(dp; label="Parasites", ms=5, size=(400, 400), alpha=0.6)
plot!([1, 200], [1, 200]; lab="", c=:grey, ls=:dot)
xaxis!(:log10, "Degree (measured)", (1, 200))
yaxis!(:log10, "Degree (imputed)", (1, 200))

# We save this
savefig("figures/degree-change.png")

# We can check the distribution of specificity
density(
    collect(values(specificity(M)));
    frame=:box,
    dpi=600,
    lab="Empirical",
    fill=(0, 0.2),
    size=(400, 400),
)
density!(collect(values(specificity(P))); lab="Imputed", fill=(0, 0.2))
vline!([0.5]; lab="", c=:grey, ls=:dash)
xaxis!("Specificity", (0, 1))
yaxis!("Density", (0, 10))

# And save it as well
savefig("figures/specificity.png")

# For a final figure, we will produce a tSNE embedding
emb_post = tsne(
    convert.(
        Float64, Array(EcologicalNetworks.mirror(convert(UnipartiteNetwork, P)).edges)
    ),
    3,
    nf,
    15000,
    5;
    pca_init=true,
)

# k-means
elbow = [
    ParallelKMeans.kmeans(
        emb_post', i; n_threads=4, tol=1e-6, max_iters=500, verbose=false
    ).totalcost for i in 2:20
];
# use scatter(elbow) to check

# Cluster based on the above color
clusters = kmeans(emb_post', 7).assignments
# wng = palette(distinguishable_colors(length(unique(clusters))))
wng = palette([
    colorant"#e69f00",
    colorant"#56b4e9",
    colorant"#009e73",
    colorant"#f0e442",
    colorant"#0072b2",
    colorant"#d55e00",
    colorant"#cc79a7",
])

# Positions of the hosts and parasites
idx_para = indexin(species(M; dims=1), species(K))
idx_host = indexin(species(M; dims=2), species(K))

# This next bit is just fugly code to get the edges for the plot
old_interacting_pairs = findall(Array(M.edges) .> 0)
new_interacting_pairs = findall(Array(P.edges) .- Array(M.edges) .> 0)
nx = Float64[]
ny = Float64[]
ox = Float64[]
oy = Float64[]
for (i, ip) in enumerate(new_interacting_pairs)
    cp, ch = ip.I
    append!(nx, [emb_post[idx_para[cp], 1], emb_post[idx_host[ch], 1], NaN])
    append!(ny, [emb_post[idx_para[cp], 2], emb_post[idx_host[ch], 2], NaN])
end
for (i, ip) in enumerate(old_interacting_pairs)
    cp, ch = ip.I
    append!(ox, [emb_post[idx_para[cp], 1], emb_post[idx_host[ch], 1], NaN])
    append!(oy, [emb_post[idx_para[cp], 2], emb_post[idx_host[ch], 2], NaN])
end

# We start by plotting the edges
plot(
    ox,
    oy;
    frame=:none,
    lab="",
    legend=false,
    c=:lightgrey,
    dpi=600,
    lw=0.5,
    size=(400, 400),
)

# Then we add the nodes
scatter!(
    emb_post[idx_host, 1],
    emb_post[idx_host, 2];
    m=:square,
    marker_z=clusters[idx_host],
    msw=0.0,
    ms=3,
    c=wng,
)
scatter!(
    emb_post[idx_para, 1], emb_post[idx_para, 2]; ms=3, marker_z=clusters[idx_para], c=wng
)

# And we save
savefig("figures/tsne-original.png")

# We do the same plot but with the
plot(
    ox,
    oy;
    frame=:none,
    lab="",
    legend=false,
    c=:lightgrey,
    dpi=600,
    lw=0.5,
    size=(400, 400),
)
plot!(nx, ny; lab="", c=:black, alpha=0.2, lw=0.6)

# Then we add the nodes
scatter!(
    emb_post[idx_host, 1],
    emb_post[idx_host, 2];
    m=:square,
    marker_z=clusters[idx_host],
    msw=0.0,
    ms=3,
    c=wng,
)
scatter!(
    emb_post[idx_para, 1], emb_post[idx_para, 2]; ms=3, marker_z=clusters[idx_para], c=wng
)

# And we save
savefig("figures/tsne-imputed.png")

max_stars_repo_name: PoisotLab/NetworkLearningIllustration

max_issues_repo_name: PoisotLab/NetworkLearningIllustration

max_forks_repo_name: PoisotLab/NetworkLearningIllustration


text:
"""
    sample_locus(locus::Dict, n::Int, ploidy::Signed)
Internal function used by `simulate` to take a `Dict` of alleles => frequencies of a locus and return
`n` number of genotypes (n_alleles = `ploidy`) by using weighted sampling of the
allele-frequency pairs.

**Example**
```julia
d = Dict(133 => 0.125,135 => 0.5625,143 => 0.25,137 => 0.0625)

julia> sample_locus(d, 3, 2)
5-element Array{Tuple{Int16,Int16},1}:
 (133, 135)
 (135, 135)
 (143, 137)

julia> sample_locus(d, 3, 3)
5-element Array{Tuple{Int16,Int16,Int16},1}:
 (135, 135, 133)
 (143, 135, 133)
 (137, 135, 135)
```
"""
function sample_locus(locus::Dict, n::Int, ploidy::Signed)
    isempty(locus) && return fill(missing, n)
    k,v = collect(keys(locus)), collect(values(locus))
    Tuple.(sort.([sample(k, Weights(v), ploidy) for i in 1:n]))
end

"""
    simulate(data::PopData; n::Int)
Simulate `n` number of individuals per population using per-population
allele frequencies derived from a `PopData` object. Returns a new `PopData` object with `n` * `n_populations` samples.

    simulate(data::PopData; scale::Int)
Simulate individuals per population in the same proportions they appear in the PopData
using per-population allele frequencies. Simulation volume can be multiplied using `scale`,
i.e. if you want to keep the same proportions but generate twice the number of samples, `scale`
would be `2`. Returns a new `PopData` object with `n_samples` * `scale` samples.

**Example**
```julia
julia> cats = @nanycats;

julia> sims = simulate(cats, n = 100)
PopData{Diploid, 9 Microsatellite Loci}
  Samples: 1700
  Populations: 17

julia> sims_prop = simulate(cats, scale = 3)
  PopData{Diploid, 9 Microsatellite Loci}
    Samples: 711
    Populations: 17
```
"""
function simulate(data::PopData; n::Union{Dict{String,Int}, Int}=0, scale::Int = 0)
    n == scale && throw(ArgumentError("Please use one of n (flat) or scale (proportional) keywords for simulations. See ?simulate for more info."))
    ((n isa Dict) | (n != 0)) & !iszero(scale) && throw(ArgumentError("Must use only one of n (flat) or scale (proportional) keywords for simulations. See ?simulate for more info."))
    length(data.metadata.ploidy) > 1 && error("Simulations do not work on mixed-ploidy data (yet)")
    geno_out =
        if n != 0
            n isa Int ? _simulateflat(data, n) : _simulatearbitrary(data, n)
        else
            _simulatescale(data, scale)
        end
    transform!(
        geno_out,
        :name => (i -> PooledArray(i, compress = true)) => :name,
        :population => (i -> PooledArray(i, compress = true)) => :population,
        :locus => (i -> PooledArray(i, compress = true)) => :locus,
        :genotype
    )
    PopData(geno_out)
end

function _simulatearbitrary(data::PopData, n::Dict{String, Int})
    ploidy = data.metadata.ploidy
    pops = collect(keys(n))
    popcounts =  collect(values(n))
    nloci = data.metadata.loci
    nsamples = sum(popcounts)
    # instantiate output df
    simnames = repeat(["sim_" * "$i" for i in 1:nsamples], inner = nloci)
    popnames = reduce(vcat, [fill(i, nloci * j) for (i,j) in zip(pops,popcounts)])
    locinames = repeat(unique(data.genodata.locus), outer = nsamples)
    geno_out = DataFrame(:name => simnames, :population => popnames, :locus => locinames, :genotype => similar(data.genodata.genotype, length(locinames)))
    # generate allele freqs per population
    gdf = groupby(data.genodata[data.genodata.population .∈ Ref(pops),:], [:population, :locus])
    freqs = DataFrames.combine(
        gdf,
        :genotype => allelefreq => :frq,
    )
    freqs[:, :n] .= getindex.(Ref(n), freqs.population)
    transform!(freqs, [:frq, :n] => ((i,j) -> sample_locus.(i, j, ploidy)) => :frq)
    # populate out df
    out_gdf = groupby(geno_out, :population)
    geno_gdf = groupby(freqs, :population)
    for pop in pops
        out_gdf[(population = pop,)][:,:genotype] = reduce(hcat, geno_gdf[(population = pop,)].frq) |> permutedims |> vec
    end
    return geno_out
end

function _simulateflat(data::PopData, n::Int)
    ploidy = data.metadata.ploidy
    pops = unique(data.sampleinfo.population)
    npops = data.metadata.populations
    nloci = data.metadata.loci

    # instantiate output df
    simnames = repeat(["sim_" * "$i" for i in 1:(n*npops)], inner = nloci)
    popnames = repeat(pops, inner = (nloci * n))
    locinames = repeat(unique(data.genodata.locus), outer = (n * npops))
    geno_out = DataFrame(:name => simnames, :population => popnames, :locus => locinames, :genotype => similar(data.genodata.genotype, length(locinames)))

    # generate allele freqs per population
    gdf = groupby(data.genodata, [:population, :locus])
    freqs = DataFrames.combine(
        gdf,
        :genotype => allelefreq => :frq
    )
    # create new genotypes
    transform!(freqs, :frq => (i -> sample_locus.(i,n,ploidy)) => :frq)

    # populate out df
    out_gdf = groupby(geno_out, :population)
    geno_gdf = groupby(freqs, :population)
    for pop in pops
        out_gdf[(population = pop,)][:,:genotype] = reduce(hcat, geno_gdf[(population = pop,)].frq) |> permutedims |> vec
    end
    return geno_out
end

function _simulatescale(data::PopData, n::Int)
    ploidy = data.metadata.ploidy
    pops = unique(data.sampleinfo.population)
    popcounts =  [count(i -> i == j, data.sampleinfo.population) for j in pops]
    nloci = data.metadata.loci

    # instantiate output df
    simnames = repeat(["sim_" * "$i" for i in 1:(data.metadata.samples * n)], inner = nloci)
    popnames = reduce(vcat, [fill(i, nloci * j * n) for (i,j) in zip(pops,popcounts)])
    locinames = repeat(unique(data.genodata.locus), outer = (n * data.metadata.samples))
    geno_out = DataFrame(:name => simnames, :population => popnames, :locus => locinames, :genotype => similar(data.genodata.genotype, length(locinames)))
    # generate allele freqs per population
    gdf = groupby(data.genodata, [:population, :locus])
    freqs = DataFrames.combine(
        gdf,
        :genotype => allelefreq => :frq
    )
    freqs[:, :n] .= repeat(popcounts, inner = nloci)
    transform!(freqs, [:frq, :n] => ((i,j) -> sample_locus.(i, j*n, ploidy)) => :frq)
    # populate out df
    out_gdf = groupby(geno_out, :population)
    geno_gdf = groupby(freqs, :population)
    for pop in pops
        out_gdf[(population = pop,)][:,:genotype] = reduce(hcat, geno_gdf[(population = pop,)].frq) |> permutedims |> vec
    end
    return geno_out
end

ext: jl


text:
__precompile__()
module AppleAccelerate

using Libdl

try
    global const libacc = dlopen("/System/Library/Frameworks/Accelerate.framework/Accelerate")
catch
    error("Accelerate framework not found.")
end

get_fptr(s) = dlsym(libacc, s)

include("Array.jl")
include("DSP.jl")
include("Util.jl")

end # module


text:
n_k = [80; 100; 120]
k = length(n_k)
n = sum(n_k)
p = 3

y, X = sampledata(n_k, p)
M = MOD.classmeans(Val{:row}, X, y)
H = MOD.centerclasses!(Val{:row}, copy(X), M, y)
Σ = H'H/(n-1)
Σ_k =
priors = ones(k)./k

info("Testing ", MOD.classwhiteners!)
for order in (:row, :col)
    for T in FloatingPointTypes
        isrowmajor = order == :row
        H_tst = isrowmajor ? copy(convert(Array{T}, H)) : transpose(convert(Array{T}, H))
        Σ_tst = isrowmajor ? copy(convert(Array{T}, Σ)) : transpose(convert(Array{T}, Σ))
        Σ_k_tst = isrowmajor ? [MOD.covmatrix(Val{:row}, H_tst[y .== i,:]) for i = 1:y.k] :
                               [MOD.covmatrix(Val{:col}, H_tst[:,y .== i]) for i = 1:y.k]

        for γ in (Nullable{T}(), Nullable(convert(T,0.4)), Nullable(one(T)))
            for λ in (Nullable{T}(), Nullable(convert(T,0.4)), Nullable(one(T)))
                W_k = isnull(λ) ? MOD.classwhiteners!(Val{order}, H_tst, y, γ) :
                                  MOD.classwhiteners!(Val{order}, H_tst, y, γ, get(λ))
                for i in eachindex(W_k)
                    S = isnull(λ) ? Σ_k_tst[i] : (1-get(λ))*Σ_k_tst[i] + get(λ)*Σ_tst
                    S = isnull(γ) ? S      : (1-get(γ))*S + get(γ)*trace(S)/p*I
                    if order == :row
                        @test_approx_eq eye(T,p) W_k[i]'*S*W_k[i]
                    else
                        @test_approx_eq eye(T,p) W_k[i]*S*(W_k[i]')
                    end
                end
            end
        end
    end
end

info("Testing ", MOD.qda!)
for order in (:row, :col)
    for T in FloatingPointTypes
        isrowmajor = order == :row
        X_tst = isrowmajor ? copy(convert(Matrix{T}, X)) : transpose(convert(Matrix{T}, X))
        M_tst = isrowmajor ? copy(convert(Matrix{T}, M)) : transpose(convert(Matrix{T}, M))
        H_tst = isrowmajor ? copy(convert(Matrix{T}, H)) : transpose(convert(Matrix{T}, H))

        for γ in (Nullable{T}(), Nullable(convert(T,0.5)))
            for λ in (Nullable{T}(), Nullable(convert(T,0.5)))
                W_k_tst = MOD.qda!(Val{order}, copy(X_tst), copy(M_tst), y, γ, λ)
                W_k_tmp = isnull(λ) ? MOD.classwhiteners!(Val{order}, H_tst, y, γ) :
                                      MOD.classwhiteners!(Val{order}, H_tst, y, γ, get(λ))
                for i in eachindex(W_k_tmp)
                    @test_approx_eq W_k_tst[i] W_k_tmp[i]
                end
            end
        end
    end
end

info("Testing ", MOD.qda)
for order in (:row, :col)
    for T in FloatingPointTypes
        isrowmajor = order == :row
        X_tst = isrowmajor ? copy(convert(Matrix{T}, X)) : transpose(convert(Matrix{T}, X))
        M_tst = isrowmajor ? copy(convert(Matrix{T}, M)) : transpose(convert(Matrix{T}, M))

        for γ in (Nullable{T}(), zero(T), Nullable(convert(T, 0.5)), convert(T, 0.5), one(T))
            for λ in (Nullable{T}(), zero(T), Nullable(convert(T, 0.5)), convert(T, 0.5), one(T))
                model = MOD.qda(X_tst, y, order=Val{order}, gamma=γ, lambda=λ)
                γ_tst = isa(γ, Nullable) ? γ : Nullable(γ)
                λ_tst = isa(λ, Nullable) ? λ : Nullable(λ)
                W_k_tst = MOD.qda!(Val{order}, copy(X_tst), copy(M_tst), y, γ_tst, λ_tst)
                for i in eachindex(model.W_k)
                    @test_approx_eq model.W_k[i] W_k_tst[i]
                end
                @test_approx_eq model.M M_tst
            end
        end
    end
end

info("Testing ", MOD.discriminants_qda)
for T in FloatingPointTypes
    X_tmp = copy(convert(Matrix{T}, X))
    priors_tmp = convert(Vector{T}, priors)

    model1 = qda(X_tmp,  y, order=Val{:row})
    model2 = qda(X_tmp', y, order=Val{:col})
    D = hcat(-[MOD.dotvectors(Val{:row}, (X_tmp .- M[i,:])*model1.W_k[i])/2
                for i in eachindex(priors_tmp)]...) .+ log(priors_tmp)'

    @test_approx_eq D  MOD.discriminants(model1, X_tmp)
    @test_approx_eq D' MOD.discriminants(model2, X_tmp')
end

info("Testing ", MOD.classify)
for order in (:row, :col)
    for T in FloatingPointTypes
        isrowmajor = order == :row
        X_tst = isrowmajor ? copy(convert(Matrix{T}, X)) : transpose(convert(Matrix{T}, X))
        priors_tst = convert(Vector{T}, priors)

        model = qda(X_tst, y, order=Val{order})
        D1 = vec(mapslices(indmax, MOD.discriminants(model, X_tst), isrowmajor ? 2 : 1))
        @test D1 == MOD.classify(model, X_tst)
    end
end

info("Testing ", MOD.ModelQDA)
show(DevNull, qda(X, y))


text:
"""
 Take the L2 Tikhonov regularised inverse of a matrix `m`.

 The regularisation parameter is the larger of `delta` (the optional argument that defaults
 to zero) and square root of machine epsilon. The inverse is done using an SVD.
 """
 function reginv(a::AbstractMatrix, δ = defaulttol(a))
     U, S, V = svd(a)
     Sinv = inv.(hypot.(S, δ))
     return  V * Diagonal(Sinv) * U'
 end

 function posreginv(a::AbstractMatrix, δ = defaulttol(a))
     # assumes a is positive definite (or close to)
     D, V = eigen(Hermitian(a))
     Dinv = inv.(hypot.(D, δ))
     return  V * Diagonal(Dinv) * V'
 end


text:
facts("Fitness") do

  context("hat_compare floats") do
    @fact hat_compare(1.0, 2.0) => -1
    @fact hat_compare(-1.0, 1.0) => -1

    @fact hat_compare(2.0, 1.0) => 1
    @fact hat_compare(-2.0, -3.0) => 1

    @fact hat_compare(1.0, 1.0) => 0
    @fact hat_compare(0.0, 0.0) => 0
    @fact hat_compare(-1.0, -1.0) => 0
  end

  context("is_minimizing in float fitness schemes") do
    mins = BlackBoxOptim.FloatFitness(true)
    @fact BlackBoxOptim.is_minimizing(mins) => true

    maxs = BlackBoxOptim.FloatFitness(false)
    @fact BlackBoxOptim.is_minimizing(maxs) => false
  end

  context("hat_compare floats in a minimizing fitness scheme") do
    scheme = BlackBoxOptim.FloatFitness(true)

    @fact hat_compare(1.0, 2.0, scheme) => -1
    @fact hat_compare(2.0, 1.0, scheme) => 1
    @fact hat_compare(1.0, 1.0, scheme) => 0
  end

  context("hat_compare floats in a maximizing fitness scheme") do
    scheme = BlackBoxOptim.FloatFitness(false)

    @fact hat_compare(1.0, 2.0, scheme) => 1
    @fact hat_compare(2.0, 1.0, scheme) => -1
    @fact hat_compare(1.0, 1.0, scheme) => 0
  end

  context("hat_compare fitnesses of size 1 in a minimizing FitnessScheme") do
    scheme = float_vector_scheme_min()

    @fact hat_compare([-1.0], [1.0], scheme) => -1
    @fact hat_compare([0.0], [0.3], scheme) => -1
    @fact hat_compare([11.3], [354.65], scheme) => -1

    @fact hat_compare([-1.0], [-1.0], scheme) => 0
    @fact hat_compare([0.0], [0.0], scheme) => 0
    @fact hat_compare([0.2], [0.2], scheme) => 0

    @fact hat_compare([1.0], [0.0], scheme) => 1
    @fact hat_compare([0.0], [-0.4], scheme) => 1
    @fact hat_compare([-0.65], [-34.2], scheme) => 1
  end

  context("hat_compare fitnesses of size 1 in a maximizing FitnessScheme") do
    scheme = float_vector_scheme_max()

    @fact hat_compare([-1.0], [1.0], scheme) => 1
    @fact hat_compare([0.0], [0.3], scheme) => 1
    @fact hat_compare([11.3], [354.65], scheme) => 1

    @fact hat_compare([-1.0], [-1.0], scheme) => 0
    @fact hat_compare([0.0], [0.0], scheme) => 0
    @fact hat_compare([0.2], [0.2], scheme) => 0

    @fact hat_compare([1.0], [0.0], scheme) => -1
    @fact hat_compare([0.0], [-0.4], scheme) => -1
    @fact hat_compare([-0.65], [-34.2], scheme) => -1
  end

  context("hat_compare fitnesses of size > 1 in a minimizing FitnessScheme") do
    scheme = float_vector_scheme_min()

    @fact hat_compare([-1.0, 0.0], [1.0, 0.0], scheme) => -1
    @fact hat_compare([0.0, -1.0], [0.3, 1.0], scheme) => -1
    @fact hat_compare([11.3, 100.0], [354.65, 10.0], scheme) => -1

    @fact hat_compare([-1.0, 2.0], [-1.0, 2.0], scheme) => 0
    @fact hat_compare([0.0, 0.0], [0.0, 0.0], scheme) => 0
    @fact hat_compare([0.2, -0.4], [0.2, -0.4], scheme) => 0

    @fact hat_compare([1.0, 0.0], [0.0, -1.0], scheme) => 1
    @fact hat_compare([0.0, 34.0], [-0.4, 20.0], scheme) => 1
    @fact hat_compare([-0.65, 1.0], [1.0, -34.2], scheme) => 1
  end

  context("hat_compare fitnesses of size > 1 in a minimizing FitnessScheme") do
    scheme = float_vector_scheme_max()

    @fact hat_compare([-1.0, 0.0], [1.0, 0.0], scheme) => 1
    @fact hat_compare([0.0, -1.0], [0.3, 1.0], scheme) => 1
    @fact hat_compare([11.3, 100.0], [354.65, 10.0], scheme) => 1

    @fact hat_compare([-1.0, 2.0], [-1.0, 2.0], scheme) => 0
    @fact hat_compare([0.0, 0.0], [0.0, 0.0], scheme) => 0
    @fact hat_compare([0.2, -0.4], [0.2, -0.4], scheme) => 0

    @fact hat_compare([1.0, 0.0], [0.0, -1.0], scheme) => -1
    @fact hat_compare([0.0, 34.0], [-0.4, 20.0], scheme) => -1
    @fact hat_compare([-0.65, 1.0], [1.0, -34.2], scheme) => -1
  end

  context("is_better/is_worse/same_fitness in a minimizing FitnessScheme") do
    scheme = float_vector_scheme_min()

    @fact is_better([-1.0, 0.0], [1.0, 0.0], scheme) => true
    @fact is_better([0.0, 0.0], [1.0, 0.0], scheme) => true

    @fact is_better([1.0, 0.0], [-1.0, 0.0], scheme) => false
    @fact is_better([1.0, 0.0], [0.0, 0.0], scheme) => false

    @fact is_worse([-1.0, 0.0], [1.0, 0.0], scheme) => false
    @fact is_worse([0.0, 0.0], [1.0, 0.0], scheme) => false

    @fact is_worse([1.0, 0.0], [-1.0, 0.0], scheme) => true
    @fact is_worse([1.0, 0.0], [0.0, 0.0], scheme) => true

    @fact same_fitness([-1.0, 0.0], [1.0, 0.0], scheme) => false
    @fact same_fitness([0.0, 0.0], [1.0, 0.0], scheme) => false
    @fact same_fitness([0.0], [1.0], scheme) => false

    @fact same_fitness([-1.0, 0.0], [-1.0, 0.0], scheme) => true
    @fact same_fitness([0.0, 0.0], [0.0, 0.0], scheme) => true
    @fact same_fitness([0.0], [0.0], scheme) => true
  end

  context("is_better/is_worse/same_fitness in a maximizing FitnessScheme") do
    scheme = float_vector_scheme_max()

    @fact is_better([-1.0, 0.0], [1.0, 0.0], scheme) => false
    @fact is_better([0.0, 0.0], [1.0, 0.0], scheme) => false

    @fact is_better([1.0, 0.0], [-1.0, 0.0], scheme) => true
    @fact is_better([1.0, 0.0], [0.0, 0.0], scheme) => true

    @fact is_worse([-1.0, 0.0], [1.0, 0.0], scheme) => true
    @fact is_worse([0.0, 0.0], [1.0, 0.0], scheme) => true

    @fact is_worse([1.0, 0.0], [-1.0, 0.0], scheme) => false
    @fact is_worse([1.0, 0.0], [0.0, 0.0], scheme) => false

    @fact same_fitness([-1.0, 0.0], [1.0, 0.0], scheme) => false
    @fact same_fitness([0.0, 0.0], [1.0, 0.0], scheme) => false
    @fact same_fitness([0.0], [1.0], scheme) => false

    @fact same_fitness([-1.0, 0.0], [-1.0, 0.0], scheme) => true
    @fact same_fitness([0.0, 0.0], [0.0, 0.0], scheme) => true
    @fact same_fitness([0.0], [0.0], scheme) => true
  end

end


text:
# ---
# title: 587. Erect the Fence
# id: problem587
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Geometry
# link: <https://leetcode.com/problems/erect-the-fence/description/>
# hidden: true
# ---
#
# There are some trees, where each tree is represented by (x,y) coordinate in a
# two-dimensional garden. Your job is to fence the entire garden using the
# **minimum length** of rope as it is expensive. The garden is well fenced only
# if all the trees are enclosed. Your task is to help find the coordinates of
# trees which are exactly located on the fence perimeter.
#
#
#
# **Example 1:**
#
#
#
#     Input: [[1,1],[2,2],[2,0],[2,4],[3,3],[4,2]]
#     Output: [[1,1],[2,0],[4,2],[3,3],[2,4]]
#     Explanation:
#     ![](https://assets.leetcode.com/uploads/2018/10/12/erect_the_fence_1.png)
#
#
# **Example 2:**
#
#
#
#     Input: [[1,2],[2,2],[4,2]]
#     Output: [[1,2],[2,2],[4,2]]
#     Explanation:
#     ![](https://assets.leetcode.com/uploads/2018/10/12/erect_the_fence_2.png)
#     Even you only have trees in a line, you need to use rope to enclose them.
#
#
#
#
# **Note:**
#
#   1. All trees should be enclosed together. You cannot cut the rope to enclose trees that will separate them in more than one group.
#   2. All input integers will range from 0 to 100.
#   3. The garden has at least one tree.
#   4. All coordinates are distinct.
#   5. Input points have **NO** order. No order required for output.
#   6. input types have been changed on April 15, 2019. Please reset to default code definition to get new method signature.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
thisdir = dirname(@__FILE__());
addvectorfile = thisdir * "/addVectors.okl";
function test_okl_ofl_vectoradd(backend::TestBackend)
    linfo=get_info_string(backend);
    entries = 5

    device = OCCA.Device(linfo);

    a  = Float32[1 - i for i in 1:entries]
    b  = Float32[i     for i in 1:entries]
    ab = Array(Float32,(length(a),));

    correctvals = [1.0 for i in 1:entries];

    o_a  = OCCA.malloc(device, a);
    o_b  = OCCA.malloc(device, b);
    o_ab = OCCA.malloc(device, ab);

    addvectors = OCCA.buildkernel(device,
                                            addvectorfile,
                                            "addVectors")

    OCCA.runkernel!(addvectors,
                   entries,
                   o_a, o_b, o_ab)

    OCCA.memcpy!(ab, o_ab)

    return (norm(correctvals-ab)/norm(correctvals))<1e-5
end


text:
### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# ╔═╡ bfe42bab-98e8-4460-a8ba-12968297d227
begin
	using PlutoUI, LinearAlgebra
	PlutoUI.TableOfContents(aside=true)
end

# ╔═╡ 43d24d37-df93-4ee6-8e48-89f9ca6f0f92
md"""
# Singular Value Decomposition - Algorithms and Error Analysis

We study only algorithms for real matrices, which are most commonly used in the applications described in this course.

For more details, see
[A. Kaylor Cline and I. Dhillon, Computation of the Singular Value Decomposition, pp. 58.1-58.13](https://www.routledge.com/Handbook-of-Linear-Algebra/Hogben/p/book/9781138199897) and the references therein.

__Prerequisites__

The reader should be familiar with facts about the singular value decomposition and perturbation theory and algorithms for the symmetric eigenvalue decomposition.

__Competences__

The reader should be able to apply an adequate algorithm to a given problem, and to assess the accuracy of the solution.
"""

# ╔═╡ 64ca893a-a1ee-4172-bd20-cf4aca1c3d41
md"""
# Basics

## Definitions

The __singular value decomposition__ (SVD) of $A\in\mathbb{R}^{m\times n}$ is $A=U\Sigma V^T$, where $U\in\mathbb{R}^{m\times m}$ is orthogonal, $U^TU=UU^T=I_m$, $V\in\mathbb{R}^{n\times n}$ is orthogonal, $V^TV=VV^T=I_n$, and $\Sigma \in\mathbb{R}^{m\times n}$ is diagonal with singular values $\sigma_1,\ldots,\sigma_{\min\{m,n\}}$ on the diagonal.

If $m>n$, the __thin SVD__ of $A$ is $A=U_{1:m,1:n} \Sigma_{1:n,1:n} V^T$.

## Facts

1. Algorithms for computing SVD of $A$ are modifications of algorithms for the symmetric eigenvalue decomposition of the matrices $AA^T$, $A^TA$ and $\begin{bmatrix} 0 & A\\ A^T & 0 \end{bmatrix}$.

2. Most commonly used approach is the three-step algorithm:
    1. Reduce $A$ to bidiagonal matrix $B$ by orthogonal transformations, $X^TAY=B$.
    2. Compute the SVD of $B$ with QR iterations, $B=W\Sigma Z^T$.
    3. Multiply $U=XW$ and $V=YZ$.

3. If $m\geq n$, the overall operation count for this algorithm is $O(mn^2)$ operations.

4. __Error bounds.__ Let $U\Sigma V^T$ and $\tilde U \tilde \Sigma \tilde V^T$ be the exact and the computed SVDs of $A$, respectively. The algorithms generally compute the SVD with errors bounded by

$$
|\sigma_i-\tilde \sigma_i|\leq \phi \epsilon\|A\|_2,
\qquad
\|u_i-\tilde u_i\|_2, \| v_i-\tilde v_i\|_2 \leq \psi\epsilon \frac{\|A\|_2}
{\min_{j\neq i}
|\sigma_i-\tilde \sigma_j|},$$

where $\epsilon$ is machine precision and $\phi$ and $\psi$ are slowly growing polynomial functions of $n$ which depend upon the algorithm used (typically $O(n)$ or $O(n^2)$). These bounds are obtained by combining perturbation bounds with the floating-point error analysis of the algorithms.
"""

# ╔═╡ 38b90f14-4c98-4a9c-a54a-5fabacb252ea
md"""
#  Bidiagonalization

## Facts

1. The reduction of $A$ to bidiagonal matrix can be performed by applying $\min\{m-1,n\}$ Householder reflections $H_L$ from the left and $n-2$ Householder reflections $H_R$ from the right. In the first step, $H_L$ is chosen to annihilate all elements of the first column below the diagonal, and $H_R$ is chosen to annihilate all elements of the first row right of the first super-diagonal. Applying this procedure recursively yields the bidiagonal matrix.

2.  $H_L$ and $H_R$ do not depend on the normalization of the respective Householder vectors $v_L$ and $v_R$. With the normalization $[v_L]_1=[V_R]_1=1$, the vectors $v_L$ are stored in the lower-triangular part of $A$, and the vectors $v_R$ are stored in the upper-triangular part of $A$ above the super-diagonal.

3. The matrices $H_L$ and $H_R$ are not formed explicitly - given $v_L$ and $v_R$, $A$ is overwritten with $H_L A H_R$ in $O(mn)$ operations by using matrix-vector multiplication and rank-one updates.

4. Instead of performing rank-one updates, $p$ transformations can be accumulated, and then applied. This __block algorithm__ is rich in matrix-matrix multiplications (roughly one half of the operations is performed using BLAS 3 routines), but it requires extra workspace.

5. If the matrices $X$ and $Y$ are needed explicitly, they can be computed from the stored Householder vectors. In order to minimize the operation count, the computation starts from the smallest matrix and the size is gradually increased.

6. The backward error bounds for the bidiagonalization are as follows: The computed matrix $\tilde B$ is equal to the matrix which would be obtained by exact bidiagonalization of some perturbed matrix $A+\Delta A$, where $\|\Delta A\|_2 \leq \psi \varepsilon \|A\|_2$ and $\psi$ is a slowly increasing function of $n$. The computed matrices $\tilde X$ and $\tilde Y$ satisfy $\tilde X=X+\Delta X$ and $\tilde Y=Y+\Delta Y$, where $\|\Delta X \|_2,\|\Delta Y\|_2\leq \phi \varepsilon$ and $\phi$ is a slowly increasing function of $n$.

7. The bidiagonal reduction is implemented in the [LAPACK](http://www.netlib.org/lapack) subroutine [DGEBRD](http://www.netlib.org/lapack/explore-html/dd/d9a/group__double_g_ecomputational.html#ga9c735b94f840f927f8085fd23f3ee2e6). The computation of $X$ and $Y$ is implemented in the subroutine [DORGBR](http://www.netlib.org/lapack/lapack-3.1.1/html/dorgtr.f.html), which is not yet wrapped in Julia.

8. Bidiagonalization can also be performed using Givens rotations. Givens rotations act more selectively than Householder reflectors, and are useful if $A$ has some special structure, for example, if $A$ is a banded matrix. Error bounds for function `BidiagG()` are the same as above, but with slightly different functions $\psi$ and $\phi$.
"""

# ╔═╡ 875fce0a-1b1b-4332-a150-6ec5f20aa7dc
md"
## Example
"

# ╔═╡ 3c8c847e-978d-4d43-99ea-773d915e03ea
begin
	m=8
	n=5
	import Random
	Random.seed!(421)
	A=rand(-9.0:9,m,n)
end

# ╔═╡ fa7b3f29-d9bc-4a64-90ae-16a6a3f8e3a6
# Householder matrix from vector
function H(x)
    v=copy(x)
    v[1]+=sign(x[1])*norm(x)
    # display(v/v[1])
    I-(2/(v⋅v))*v*v'
end

# ╔═╡ 78d6a3a1-9fd5-4eec-94dd-2baf31cf312e
H₁=H(A[:,1])

# ╔═╡ 2447bb48-334b-4554-b7ca-5b4f74f68be9
# Orthogonality
norm(H₁'*H₁-I)

# ╔═╡ ed9bfaff-ed9f-442d-b5a5-c0435e22d2ef
B₁=H₁*A

# ╔═╡ 8856d77a-f2fb-4d7e-aca3-bde649508c8a
H¹=cat(1,H(B₁[1,2:5]),dims=(1,2))

# ╔═╡ 5d742283-df40-416b-8067-1a21f89f6e71
H₁*A*H¹

# ╔═╡ fd73c193-a5bb-4eb9-a0ce-aca44b7ef7ee
# ?LAPACK.gebrd!

# ╔═╡ 783deb44-2d8f-4a40-80d5-cc65bcf9018c
# We need copy()
O=LAPACK.gebrd!(copy(A))

# ╔═╡ cfba6cc8-451c-4716-95a2-9a3f49009253
O[4]

# ╔═╡ 2f92ce80-72d0-4424-949e-a4f07065102c
B=Bidiagonal(O[2],O[3][1:end-1],'U')

# ╔═╡ e132f765-d232-40e3-9d41-1ec07528dbae
[svdvals(A) svdvals(B)]

# ╔═╡ 583548a3-b3ad-493f-8061-fae21138208a
# Extract X
function BidiagX(H::Matrix)
    m,n=size(H)
    T=typeof(H[1,1])
    X = Matrix{T}(I,m,n)
    v = Array{T}(undef,m)
    for j = n : -1 : 1
        v[j] = one(T)
        v[j+1:m] = H[j+1:m, j]
        γ = -2 / (v[j:m]⋅v[j:m])
        w = γ * X[j:m, j:n]'*v[j:m]
        X[j:m, j:n] = X[j:m, j:n] + v[j:m]*w'
    end
    X
end

# ╔═╡ 75d77c3b-649e-4d20-a1cf-5cc4d1335896
# Extract Y
function BidiagY(H::AbstractMatrix)
    n,m=size(H)
    T=typeof(H[1,1])
    Y = Matrix{T}(I,n,n)
    v = Array{T}(undef,n)
    for j = n-2 : -1 : 1
        v[j+1] = one(T)
        v[j+2:n] = H[j+2:n, j]
        γ = -2 / (v[j+1:n]⋅v[j+1:n])
        w = γ * Y[j+1:n, j+1:n]'*v[j+1:n]
        Y[j+1:n, j+1:n] = Y[j+1:n, j+1:n] + v[j+1:n]*w'
    end
    Y
end

# ╔═╡ 28a70c10-3900-4a22-9c93-41abe15b5ce8
X=BidiagX(O[1])

# ╔═╡ a0065053-51a6-4351-b7e8-da5e37b3eb94
Y=BidiagY(O[1]')

# ╔═╡ 583eec81-0a99-4ccc-9e7e-0625a942afbc
# Orthogonality
norm(X'*X-I), norm(Y'*Y-I)

# ╔═╡ 0b4010cc-5956-4a52-879f-dbc3dd6cc519
# Residual error
norm(A*Y-X*B)

# ╔═╡ 7d6da96f-c272-46e8-979c-14cc2a92d071
# Bidiagonalization using Givens rotations
function BidiagG(A1::Matrix)
    A=deepcopy(A1)
    m,n=size(A)
    T=typeof(A[1,1])
    X=Matrix{T}(I,m,m)
    Y=Matrix{T}(I,n,n)
    for j = 1 : n
        for i = j+1 : m
            G,r=givens(A,j,i,j)
            # Use the faster in-place variant
            # A=G*A
            lmul!(G,A)
            # X=G*X
            lmul!(G,X)
            # display(A)
        end
        for i=j+2:n
            G,r=givens(A',j+1,i,j)
            # A*=adjoint(G)
            rmul!(A,adjoint(G))
            # Y*=adjoint(G)
            rmul!(Y,adjoint(G))
            # display(A)
        end
    end
    X',Bidiagonal(diag(A),diag(A,1),'U'), Y
end

# ╔═╡ 5496e8da-6899-4792-8c20-46a5f8843354
X₁, B¹, Y₁ = BidiagG(A)

# ╔═╡ 1caedf0a-23d5-460e-8e61-f533f7f731a0
# Orthogonality
norm(X₁'*X₁-I), norm(Y₁'*Y₁-I)

# ╔═╡ e5ccbecf-5bf3-4c0e-a036-2f8fab7d0a3f
# Diagonalization
X₁'*A*Y₁

# ╔═╡ c5820b28-6e62-4b87-af13-6cbf3a3369a5
# X, Y and B are not unique
B

# ╔═╡ c2552861-ffc9-45e1-aac4-e44b30f28ee3
B¹

# ╔═╡ 28ee391d-2ae4-4d35-9832-973071e55681
md"""
# Bidiagonal QR method

Let $B$ be a real upper-bidiagonal matrix of order $n$ and let $B=W\Sigma Z^T$ be its SVD.

All metods for computing the SVD of bidiagonal matrix are derived from the methods for computing the EVD of the tridiagonal matrix $T=B^T B$.

## Facts

1. The shift $\mu$ is the eigenvalue of the $2\times 2$ matrix $T_{n-1:n,n-1:n}$ which is closer to $T_{n,n}$. The first Givens rotation from the right is the one which annihilates the element $(1,2)$ of the shifted $2\times 2$ matrix $T_{1:2,1:2}-\mu I$. Applying this rotation to $B$ creates the bulge at the element $B_{2,1}$. This bulge is subsequently chased out by applying adequate Givens rotations alternating from the left and from the right. This is the __Golub-Kahan algorithm__.

2. The computed SVD satisfes error bounds from the Fact 4 above.

3. The special variant of zero-shift QR algorithm (the __Demmel-Kahan algorithm__) computes the singular values with high relative accuracy.

4. The tridiagonal divide-and-conquer method, bisection and inverse iteration, and MRRR method can also be adapted for bidiagonal matrices.

5. Zero shift QR algorithm for bidiagonal matrices is implemented in the LAPACK routine [DBDSQR](http://www.netlib.org/lapack/explore-html/db/dcc/dbdsqr_8f.html). It is also used in the function `svdvals()`. Divide-and-conquer algorithm for bidiagonal matrices is implemented in the LAPACK routine [DBDSDC](http://www.netlib.org/lapack/explore-html/d9/d08/dbdsdc_8f.html). However, this algorithm also calls zero-shift QR to compute singular values.
"""

# ╔═╡ e64fc85e-2861-4e42-9789-ee12b6f9950f
md"
## Examples
"

# ╔═╡ a0d540ea-d116-499a-bcad-d6bfe3e52470
W,σ,Z=svd(B)

# ╔═╡ 2b900980-6ccb-441c-9ef5-e7dc68ff843e
@which svd(B)

# ╔═╡ 90970845-d921-4b88-9681-d344451d4c3c
σ₁=svdvals(B)

# ╔═╡ 860a3476-b257-4384-b255-4a8bd68728c9
@which svdvals!(B)

# ╔═╡ 2c1943df-95d8-43bb-bf6a-8fe13424a2c9
σ-σ₁

# ╔═╡ 1a79b153-6560-45d7-8113-e39101d17460
# ?LAPACK.bdsqr!

# ╔═╡ 7fafeb92-568a-4270-b4fb-4d44c655bcfb
B.dv

# ╔═╡ 99109bd8-e2ca-435c-a242-5772ea4432c0
one(B)

# ╔═╡ ef17e375-670b-4b7d-878a-bc2c06760651
begin
	V₀=Matrix(one(B))
	U₀=Matrix(one(B))
	C₀=Matrix(one(B))
	σ₂,Z₂,W₂,C = LAPACK.bdsqr!('U',copy(B.dv),copy(B.ev),V₀,U₀,C₀)
end

# ╔═╡ 613c0e7b-7288-4dac-a6be-731cf4f8be63
W₂'*B*Z₂'

# ╔═╡ e8119506-bd02-47e2-890c-c088fd883532
#?LAPACK.bdsdc!

# ╔═╡ ecf74ee7-97e3-4749-b63d-2cf200e4451e
σ₃,ee,W₃,Z₃,rest=LAPACK.bdsdc!('U','I',copy(B.dv),copy(B.ev))

# ╔═╡ cbed834b-4935-431c-96f6-3630534ba2eb
W₃'*B*Z₃'

# ╔═╡ ffa18a0d-c576-473a-9851-f7bd70a843ba
md"""
Functions `svd()`, `LAPACK.bdsqr!()` and `LAPACK.bdsdc!()` use the same algorithm to compute singular values.
"""

# ╔═╡ beec8bad-9621-4c66-88d3-4f48dbcf5566
[σ₃-σ₂ σ₃-σ]

# ╔═╡ 569d2241-b9a8-4501-8b17-9527e1777cdc
md"
Let us compute some timings. We observe $O(n^2)$ operations.
"

# ╔═╡ cfd21881-5cfe-430e-b60a-a783feac2cdb
begin
	n₀=1000
	Bₙ=Bidiagonal(rand(n₀),rand(n₀-1),'U')
	B₂ₙ=Bidiagonal(rand(2*n₀),rand(2*n₀-1),'U')
	println(" ")
	@time svdvals(Bₙ)
	@time svdvals(B₂ₙ)
	@time LAPACK.bdsdc!('U','N',copy(Bₙ.dv),copy(Bₙ.ev))
	@time svd(Bₙ)
	@time svd(B₂ₙ)
	1
end

# ╔═╡ 5ef6758d-51f9-49c5-a736-413f8be3e734
md"""
# QR method

Final algorithm is obtained by combining bidiagonalization and bidiagonal SVD methods. Standard method is implemented in the LAPACK routine [DGESVD](http://www.netlib.org/lapack/explore-html/d8/d2d/dgesvd_8f.html). Divide-and-conquer method is implemented in the LAPACK routine [DGESDD](http://www.netlib.org/lapack/explore-html/db/db4/dgesdd_8f.html).

The functions `svd()`, `svdvals()`, and `svdvecs()`  use `DGESDD`. Wrappers for `DGESVD` and `DGESDD` give more control about output of eigenvectors.
"""

# ╔═╡ 416c8bca-215c-491d-b1f1-1abf0793b65b
# The built-in algorithm
U,σA,V=svd(A)

# ╔═╡ 17d90d71-0776-498d-8c09-9d651ee3ea1f
Sv=svd(A)

# ╔═╡ c8e602f7-9ccf-4753-9aa5-a07df141cf40
Sv.Vt

# ╔═╡ 4ef07f50-e9aa-4075-960a-dce2ee9f9d9e
begin
	# With our building blocks
	U₁=X*W
	V₁=Y*Z
	U₁'*A*V₁
end

# ╔═╡ bb4e96f8-e61a-486c-b385-72dd2542e7fc
#?LAPACK.gesvd!

# ╔═╡ 7c5cea63-5273-46df-ae5e-58a7b4312360
# DGESVD
LAPACK.gesvd!('A','A',copy(A))

# ╔═╡ dd07def6-6ee3-4f48-ba11-340674fbc22d
#?LAPACK.gesdd!

# ╔═╡ 7e8aebb8-7621-4fd2-afeb-bdc71b091cd1
LAPACK.gesdd!('N',copy(A))

# ╔═╡ 08f390ca-5122-4a38-b7c0-631e7d9ea726
md"""
Let us perform some timings. We observe $O(n^3)$ operations.
"""

# ╔═╡ f425bd0a-ebde-42a5-8585-270b7e86f27f
begin
	n₁=1000
	Aₙ=rand(n₁,n₁)
	A₂ₙ=rand(2*n₁,2*n₁)
	println(" ")
	@time Uₙ,σₙ,Vₙ=svd(Aₙ)
	@time svd(A₂ₙ)
	@time LAPACK.gesvd!('A','A',copy(Aₙ))
	@time LAPACK.gesdd!('A',copy(Aₙ))
	@time LAPACK.gesdd!('A',copy(A₂ₙ))
	1
end

# ╔═╡ ed98fe9c-1ee8-4333-8ace-958223df6552
# Residual
norm(Aₙ*Vₙ-Uₙ*Diagonal(σₙ))

# ╔═╡ Cell order:
# ╟─bfe42bab-98e8-4460-a8ba-12968297d227
# ╟─43d24d37-df93-4ee6-8e48-89f9ca6f0f92
# ╟─64ca893a-a1ee-4172-bd20-cf4aca1c3d41
# ╟─38b90f14-4c98-4a9c-a54a-5fabacb252ea
# ╟─875fce0a-1b1b-4332-a150-6ec5f20aa7dc
# ╠═3c8c847e-978d-4d43-99ea-773d915e03ea
# ╠═fa7b3f29-d9bc-4a64-90ae-16a6a3f8e3a6
# ╠═78d6a3a1-9fd5-4eec-94dd-2baf31cf312e
# ╠═2447bb48-334b-4554-b7ca-5b4f74f68be9
# ╠═ed9bfaff-ed9f-442d-b5a5-c0435e22d2ef
# ╠═8856d77a-f2fb-4d7e-aca3-bde649508c8a
# ╠═5d742283-df40-416b-8067-1a21f89f6e71
# ╠═fd73c193-a5bb-4eb9-a0ce-aca44b7ef7ee
# ╠═783deb44-2d8f-4a40-80d5-cc65bcf9018c
# ╠═cfba6cc8-451c-4716-95a2-9a3f49009253
# ╠═2f92ce80-72d0-4424-949e-a4f07065102c
# ╠═e132f765-d232-40e3-9d41-1ec07528dbae
# ╠═583548a3-b3ad-493f-8061-fae21138208a
# ╠═75d77c3b-649e-4d20-a1cf-5cc4d1335896
# ╠═28a70c10-3900-4a22-9c93-41abe15b5ce8
# ╠═a0065053-51a6-4351-b7e8-da5e37b3eb94
# ╠═583eec81-0a99-4ccc-9e7e-0625a942afbc
# ╠═0b4010cc-5956-4a52-879f-dbc3dd6cc519
# ╠═7d6da96f-c272-46e8-979c-14cc2a92d071
# ╠═5496e8da-6899-4792-8c20-46a5f8843354
# ╠═1caedf0a-23d5-460e-8e61-f533f7f731a0
# ╠═e5ccbecf-5bf3-4c0e-a036-2f8fab7d0a3f
# ╠═c5820b28-6e62-4b87-af13-6cbf3a3369a5
# ╠═c2552861-ffc9-45e1-aac4-e44b30f28ee3
# ╟─28ee391d-2ae4-4d35-9832-973071e55681
# ╟─e64fc85e-2861-4e42-9789-ee12b6f9950f
# ╠═a0d540ea-d116-499a-bcad-d6bfe3e52470
# ╠═2b900980-6ccb-441c-9ef5-e7dc68ff843e
# ╠═90970845-d921-4b88-9681-d344451d4c3c
# ╠═860a3476-b257-4384-b255-4a8bd68728c9
# ╠═2c1943df-95d8-43bb-bf6a-8fe13424a2c9
# ╠═1a79b153-6560-45d7-8113-e39101d17460
# ╠═7fafeb92-568a-4270-b4fb-4d44c655bcfb
# ╠═99109bd8-e2ca-435c-a242-5772ea4432c0
# ╠═ef17e375-670b-4b7d-878a-bc2c06760651
# ╠═613c0e7b-7288-4dac-a6be-731cf4f8be63
# ╠═e8119506-bd02-47e2-890c-c088fd883532
# ╠═ecf74ee7-97e3-4749-b63d-2cf200e4451e
# ╠═cbed834b-4935-431c-96f6-3630534ba2eb
# ╟─ffa18a0d-c576-473a-9851-f7bd70a843ba
# ╠═beec8bad-9621-4c66-88d3-4f48dbcf5566
# ╟─569d2241-b9a8-4501-8b17-9527e1777cdc
# ╠═cfd21881-5cfe-430e-b60a-a783feac2cdb
# ╟─5ef6758d-51f9-49c5-a736-413f8be3e734
# ╠═416c8bca-215c-491d-b1f1-1abf0793b65b
# ╠═17d90d71-0776-498d-8c09-9d651ee3ea1f
# ╠═c8e602f7-9ccf-4753-9aa5-a07df141cf40
# ╠═4ef07f50-e9aa-4075-960a-dce2ee9f9d9e
# ╠═bb4e96f8-e61a-486c-b385-72dd2542e7fc
# ╠═7c5cea63-5273-46df-ae5e-58a7b4312360
# ╠═dd07def6-6ee3-4f48-ba11-340674fbc22d
# ╠═7e8aebb8-7621-4fd2-afeb-bdc71b091cd1
# ╟─08f390ca-5122-4a38-b7c0-631e7d9ea726
# ╠═f425bd0a-ebde-42a5-8585-270b7e86f27f
# ╠═ed98fe9c-1ee8-4333-8ace-958223df6552

max_stars_repo_path: src/Module B - Eigenvalue and Singular Value Decompositions/L6a Singular Value Decomposition - Algorithms and Error Analysis.jl

max_stars_repo_name: ivanslapnicar/GIAN-Applied-NLA-Course

max_issues_repo_name: ivanslapnicar/GIAN-Applied-NLA-Course

max_forks_repo_name: ivanslapnicar/GIAN-Applied-NLA-Course


text:
# This shows basic data analysis using the Card data
# on returns to education

##
using CSV, DataFrames, DataFramesMeta, DelimitedFiles, StatsPlots
cd(@__DIR__)
wait_for_key(prompt) = (print(stdout, prompt); read(stdin, 1); nothing)
card = CSV.read("../Data/card.csv", DataFrame) # how to read CSV data into DataFrame
size(card)
names(card)

## see the first 6 lines
show(first(card,6))
# see some statistics using DataFrames functionality
show(describe(card))

## add some variables
card[!,:lnwage] = round.(log.(card[!,:wage]), digits=4) # round to save storage space
card[!,:expsq] = (card[!,:exper].^2)/100
card[!,:agesq] = card[!,:age].^2

## select the ones we want
cooked = @select(card, :lnwage, :educ, :exper, :expsq, :black, :south, :smsa, :age, :agesq, :nearc4 )
display(first(cooked,6))

## write as CSV
CSV.write("cooked.csv", cooked)

## write as plain ASCII, using limited precision to reduce file size
data = Matrix{Float32}(cooked)
writedlm("cooked.txt", data)

## a simple way of looking at geographical and racial discrimination
# select north and white
nw = @subset(card, :south .== 0, :black .== 0)
@df nw scatter(:educ, :wage,label="North-White")
# select south and black
sb = @subset(card, :south .== 1, :black .== 1)
@df sb scatter!(:educ, :wage, label="South-Black",title="Wage vs. Educ",xlabel="Educ", ylabel="Wage")

## the following is similar, but maybe better, as it plots all points
@df card scatter(:educ, :wage, zcolor=:south+:black)

## density plots by groups
@df card density(:wage, group=(:black,:south), fill=true, fillalpha=0.3,legend=:topright, title="wage density by south and\n black categorical variables")

## contours and marginal densities
@df card marginalkde(:wage, :educ)

## violin and box plots
@df card boxplot(:wage) # puts is beside
@df card violin!(:wage)

## violin plot of wage, by education
@df card violin(:educ, :wage)
# add a boxplot on top
@df card boxplot!(:educ, :wage)

max_stars_repo_name: mcreel/EconometricsNotes

max_issues_repo_name: mcreel/EconometricsNotes

max_forks_repo_name: mcreel/EconometricsNotes


text:
function getenv(var::AbstractString)
    val = ccall((:getenv, "libc"), Cstring, (Cstring,), var)
    if val == C_NULL
        error("getenv: undefined variable: ", var)
    end
    unsafe_string(val)
end

rows = parse(Int32, getenv("INPUTMAT_ROWS"))
columns = parse(Int32, getenv("INPUTMAT_COLS"))

function linspace(start::Float64, finish::Float64, n::Integer)
    linspace = [start]
    distance = (finish - start) / Float64(n-1)
    next = start
    while abs(next - finish) > 0.000000000001
        next = next + distance
        linspace = append!(linspace, next)
    end
    return linspace
end

X = linspace(-2.5, 1.0, rows)
Y = linspace(-1.25, 1.25, columns)
Z = [complex(x,y) for x in X, y in Y]

max_stars_repo_name: holypriest/undergrad-thesis

max_issues_repo_name: holypriest/undergrad-thesis

max_forks_repo_name: holypriest/undergrad-thesis


text:
module Hexagons

import Base: convert, length, collect, iterate

export CoordAxial, CoordCubic,
       HexFlatTop, HexPointyTop,
       to_coord_form,
       center, vertices,
       hex_containing

# Various ways to index hexagons in a grid
# ----------------------------------------

##
# Hex coordinates:
# Each hexagon has coordinates in the hexagonal grid system.
# Conversions between hex coordinates do not depend on the
# mapping between hexagon space and Cartesian space.
##

"""Abstract hexagon coordinates. A subtype of `AbstractCoord` locates a hexagon
within a grid of hexagons, but does not contain enough information to map a hexagon
to non-hexagonal (e.g. Cartesian) space. For example, `CoordAxial(0, 0)` can be thought
of as the center of a grid of hexagons. But without knowing hexagon orientation (see
`Hexagon`), the Cartesian `(x,y)` coordinates of the vertices of the hexagon at
`CoordAxial(0, 0)` are not defined.

See https://www.redblobgames.com/grids/hexagons/#coordinates for some details."""
abstract type AbstractCoord end

struct CoordAxial{T <: Real} <: AbstractCoord
    q::T
    r::T
end
# semantic based equality and hash
Base.:(==)(c1::CoordAxial, c2::CoordAxial) = c1.q == c2.q && c1.r == c2.r
Base.hash(c::CoordAxial, h::UInt64) = hash(c.q, hash(c.r, hash(:CoordAxial, h)))

"""Cubic hexagon coordinates. Components must sum to zero.

See [`AbstractCoord`](@ref) for more information on hexagon coordinates."""
struct CoordCubic{T <: Real} <: AbstractCoord
    x::T
    y::T
    z::T

    function CoordCubic(x, y, z)
        T = promote_type(typeof(x), typeof(y), typeof(z))
        maxabs = T <: Integer ? 0 : sqrt(eps(T))
        if abs(x + y + z) > maxabs
            throw(DomainError(
                (x,y,z),
                "Components of cubic hexagon coordinates must sum to 0."
            ))
        end
        return new{T}(x, y, z)
    end
end
# use axial form for == and hash for non axial types
Base.:(==)(c1::AbstractCoord, c2::AbstractCoord) = CoordAxial(c1) == CoordAxial(c2)
Base.hash(c::AbstractCoord, h::UInt64) = hash(CoordAxial(c), h)

const _COORD_TYPES = (CoordAxial, CoordCubic)

# Convert between hexagon coordinate systems

function convert(::Type{CoordAxial}, coord::CoordCubic)
    CoordAxial(coord.x, coord.z)
end
CoordAxial(c::AbstractCoord) = convert(CoordAxial, c)

function convert(::Type{CoordCubic}, coord::CoordAxial)
    CoordCubic(coord.q, -coord.q - coord.r, coord.r)
end
CoordCubic(c::AbstractCoord) = convert(CoordCubic, c)

##
# Hexagons:
# A hexagon has an orientation (either flat or pointy top, encoded by the type)
# and coordinates in hex space.
# Orientation plus hex coordinates (plus an implicit unit-edge-length assumption) allows
# for mapping between hexagon space and Cartesian space.
# to-do: support variable size?
##

abstract type Hexagon end

"""
    HexPointyTop{T <: AbstractCoord} <: Hexagon
A hexagon oriented so that in the Cartesian plane, the hexagon's top and bottom
are points, and left/right sides are parallel to the y-axis.
"""
struct HexPointyTop{T <: AbstractCoord} <: Hexagon
    coords::T
end
# Hexagon equality doesn't depend on internal coordinate representation,
# just the actual hexagon represented.
function Base.:(==)(hex1::HexPointyTop, hex2::HexPointyTop)
    return CoordAxial(hex1.coords) == CoordAxial(hex2.coords)
end
Base.hash(hex::HexPointyTop, h::UInt64) = hash(hex.coords, hash(:HexPointyTop, h))

"""
    HexFlatTop{T <: AbstractCoord} <: Hexagon
A hexagon oriented so that in the Cartesian plane, the hexagon's top and bottom
are parallel to the x-axis.
"""
struct HexFlatTop{T <: AbstractCoord} <: Hexagon
    coords::T
end
function Base.:(==)(hex1::HexFlatTop, hex2::HexFlatTop)
    return CoordAxial(hex1.coords) == CoordAxial(hex2.coords)
end
Base.hash(hex::HexFlatTop, h::UInt64) = hash(hex.coords, hash(:HexFlatTop, h))

# Convenience constructors without manual instantiation of coord objects
"""
    HexFlatTop(q,r)
A flat top hexagon with axial hexagon coordinates `(q,r)`.
"""
HexFlatTop(q, r) = HexFlatTop(CoordAxial(q, r))
"""
    HexFlatTop(x,y,z)
A flat top hexagon with cubic coordinates `(x,y,z)`.
"""
HexFlatTop(x, y, z) = HexFlatTop(CoordCubic(x, y, z))
"""
    HexPointyTop(q,r)
A pointy top hexagon with axial hexagon coordinates `(q,r)`.
"""
HexPointyTop(q, r) = HexPointyTop(CoordAxial(q, r))
"""
    HexPointyTop(x,y,z)
A pointy top hexagon with cubic coordinates `(x,y,z)`.
"""
HexPointyTop(x, y, z) = HexPointyTop(CoordCubic(x, y, z))

##
# Hexagon coordinate Conversions
##

"""
    to_coord_form(coord_type, hex)

Return a hexagon equal to `hex` with internal coordinate representation
of type `coord_type`.
"""
function to_coord_form end
for hextype in (HexPointyTop, HexFlatTop), coordtype in _COORD_TYPES
    @eval begin
        function to_coord_form(::Type{$coordtype}, h::$hextype)
            return ($hextype)(convert($coordtype, h.coords))
        end
    end
end

##
# Hexagon to Cartesian mappings
##

"""
    center(hex::Hexagon)

Return the Cartesian coordiantes of a hexagon's center as a tuple (x,y).
"""
function center(hex::HexPointyTop)
    c_ax = CoordAxial(hex.coords)
    return (
        sqrt(3) * (c_ax.q + c_ax.r / 2),
        3/2 * c_ax.r
    )
end
function center(hex::HexFlatTop)
    c_ax = CoordAxial(hex.coords)
    return (
        3/2 * c_ax.q,
        sqrt(3) * (c_ax.q / 2 + c_ax.r)
    )
end

"""
    vertices(hex::Hexagon)

Return the Cartesian coordinates of the vertices of `hex` as a list of `(x,y)` tuples.
The returned vertices are ordered counter-clockwise from the hexagon center, starting
from the positive x-axis.
"""
function vertices(hex::HexFlatTop)
    thetas = (i * pi / 6 for i in 0:2:10)
    center_x, center_y = center(hex)
    return [(center_x + cos(theta), center_y + sin(theta)) for theta in thetas]
end
function vertices(hex::HexPointyTop)
    thetas = (i * pi / 6 for i in 1:2:11)
    center_x, center_y = center(hex)
    return [(center_x + cos(theta), center_y + sin(theta)) for theta in thetas]
end

##
# Cartesian to hexagon
##

"""
    _nearest_integer_coord(c::AbstractCoord)

Find the integer hex coordinates closest to `c`.
"""
_nearest_integer_coord(c::CoordCubic{T} where {T <: Integer}) = c
function _nearest_integer_coord(c::CoordCubic{T}) where {T}
    rx, ry, rz = round(Int, c.x), round(Int, c.y), round(Int, c.z)
    x_diff, y_diff, z_diff = abs(rx - c.x), abs(ry - c.y), abs(rz - c.z)

    if x_diff > y_diff && x_diff > z_diff
        rx = -ry - rz
    elseif y_diff > z_diff
        ry = -rx - rz
    else
        rz = -rx - ry
    end

    CoordCubic(rx, ry, rz)
end
_nearest_integer_coord(coord::CoordAxial) = coord |> CoordCubic |> _nearest_integer_coord |> CoordAxial

"""
    hex_containing(HexFlatTop, x, y)

Return the HexFlatTop containing Cartesian point `(x,y)`.
"""
function hex_containing(::Type{HexFlatTop}, x, y)
    q = 2 / 3 * x
    r = -x / 3 + sqrt(3) / 3 * y
    return HexFlatTop(_nearest_integer_coord(CoordAxial(q,r)))
end

"""
    hex_containing(HexPointyTop, x, y)

Return the HexPointyTop containing Cartesian point `(x,y)`.
"""
function hex_containing(::Type{HexPointyTop}, x, y)
    q = sqrt(3) / 3 * x - y / 3
    r = 2 / 3 * y
    return HexPointyTop(_nearest_integer_coord(CoordAxial(q,r)))
end

##
# Utilities in hexagon space
##

"""
    hex_distance(hexA::T, hexB::T) where {T <: Hexagon}

The distance in hexagon-space between two hexagons. Distances are only defined between
hexagons of the same type
"""
hex_distance(hexA::HexPointyTop, hexB::HexPointyTop) = hex_distance(hexA.coords, hexB.coords)

"""
    hex_distance(c1::AbstractCoord, c2::AbstractCoord)

The distance in hexagon-space between two hexagon coordinates.
"""
function hex_distance(c1::AbstractCoord, c2::AbstractCoord)
    return hex_distance(CoordCubic(c1), CoordCubic(c2))
end

# see RedBlobGames reference for derivation
function hex_distance(c1::CoordCubic, c2::CoordCubic)
    return (abs(c1.x - c2.x) + abs(c1.y - c2.y) + abs(c1.z - c2.z)) / 2
end

end # module Hexagons


text:
function gather_kpts_scfres(scfres::NamedTuple)
    # TODO Rename this to gather_kpts once scfres gets its proper type

    # Need gathering over k-Points:
    kpt_properties = (:ψ, :occupation, :eigenvalues)
    scfdict = Dict{Symbol, Any}()
    for symbol in kpt_properties
        scfdict[symbol] = gather_kpts(getproperty(scfres, symbol), scfres.basis)
    end
    scfdict[:basis] = gather_kpts(scfres.basis)

    if mpi_master()
        (; (symbol => get(scfdict, symbol, getproperty(scfres, symbol))
            for symbol in propertynames(scfres))...)
    else
        nothing
    end
end

"""
    load_scfres(filename)

Load back an `scfres`, which has previously been stored with [`save_scfres`](@ref).
Note the warning in [`save_scfres`](@ref).
"""
function load_scfres end

"""
    save_scfres(filename, scfres)

Save an `scfres` obtained from `self_consistent_field` to a file.
The format is determined from the file extension. Currently the following
file extensions are recognized and supported:

- **jld2**: A JLD2 file. Stores the complete state and can be used
  (with [`load_scfres`](@ref)) to restart an SCF from a checkpoint or
  post-process an SCF solution. See [Saving SCF results on disk and SCF checkpoints](@ref)
  for details.
- **vts**: A VTK file for visualisation e.g. in [paraview](https://www.paraview.org/).
  Stores the density, spin density and some metadata (energy, Fermi level, occupation etc.).
  Supports these keyword arguments:
    * `save_ψ`: Save the real-space representation of the orbitals as well
      (may lead to larger files).
    * `extra_data`: `Dict{String,Array}` with additional data on the 3D real-space
      grid to store into the VTK file.

!!! warning "No compatibility guarantees"
    No guarantees are made with respect to this function at this point.
    It may change incompatibly between DFTK versions or stop working / be removed
    in the future.
"""
function save_scfres(filename::AbstractString, scfres::NamedTuple; kwargs...)
    _, ext = splitext(filename)
    ext = Symbol(ext[2:end])

    # Whitelist valid extensions
    !(ext in (:jld2, :vts)) && error("Extension '$ext' not supported by DFTK.")

    # Gather scfres data on master MPI process
    scfres = gather_kpts_scfres(scfres)

    # On master MPI process dispatch to individual functions based on extension
    ret = isnothing(scfres) ? nothing : save_scfres_master(filename, scfres, Val(ext); kwargs...)

    # Return from this function synchronously on all processes
    # (to ensure that data has actually been written upon return)
    MPI.Barrier(MPI.COMM_WORLD)
    ret
end

function save_scfres_master(filename::AbstractString, ::NamedTuple, ::Any; kwargs...)
    error("The extension $(last(splitext(filename))) is currently not available. " *
          "A required package (e.g. JLD2 or WriteVTK) is not yet loaded.")
end


text:
function checklinkid(net::FastNet,kid,task)
    if kid<1 || kid>net.K
        throw(ArgumentError("$task, but $kid is not a valid link id"))
    end
end

function checklinkstate(net::FastNet,cls,task)
    if cls<1 || cls>net.L-1
        throw(ArgumentError("$task, but $cls is not a valid link state"))
    end
end

function checklinkexists(net::FastNet,kid,task)
    if net.kstate[kid]==net.L
        throw(ArgumentError("$task, but link $kid is not in the network (it has been deleted or hasn't been created yet)"))
    end
end

function checknonempty(net::FastNet,task)
    if countlinks_f(net)===0
        throw(ArgumentError("$task, but the network does not contain any links"))
    end
end

function checknonempty(net::FastNet,cls,task)
    if countlinks_f(net,cls)===0
        throw(ArgumentError("$task, but the network does not contain any links of state $cls"))
    end
end

function checknodeid(net::FastNet,nid,task)
    if nid<1 || nid>net.N
        throw(ArgumentError("$task, but $nid is not a valid node id"))
    end
end

function checknodestate(net::FastNet,cls,task)
    s=0
    try
        s=convert(Int,cls)
    catch e
        throw(ArgumentError("$task, but the provided node state does not seem to be an integer value"))
    end
    if s<1 || s>net.C
        throw(ArgumentError("$task, but $s is not a valid node state"))
    end
    s
end

function checknodeexists(net::FastNet,nid,task)
    if net.nstate[nid]==net.C
        throw(ArgumentError("$task, but node $nid is not in the network (it has been deleted or hasn't been created yet)"))
    end
end

function checknonnull(net::FastNet,task)
    if countnodes_f(net)===0
        throw(ArgumentError("$task, but the network does not contain any nodes"))
    end
end

function checknonnull(net::FastNet,cls,task)
    if countnodes_f(net,cls)===0
        throw(ArgumentError("$task, but the network does not contain any nodes of state $cls"))
    end
end

"""
    shownodes(net)

Print information on all nodes in FastNet *net*.

This function is mainly intended for testing/debugging. You might want to think twice before
calling it for a large network. Having 10 million nodes printed to your REPL is much less fun than
it sounds.

See also [showlinks](#Fastnet.showlinks)

# Examples
```jldoctest
julia> using Fastnet

julia> net=FastNet(1000,2000,2,[]);

julia> makenodes!(net,5,1)

julia> makenodes!(net,5,2)

julia> shownodes(net)
id      state
1       1
2       1
3       1
4       1
5       1
6       2
7       2
8       2
9       2
10      2
"""
function shownodes(net::FastNet)
    println("id\tstate")
    @inbounds begin
        for i=1:net.N
            if (net.nstate[i]<net.C)
                println(i,'\t',net.nstate[i])
            end
        end
    end
end

"""
    showlinks(net)

Print information on all links in FastNet *net*.

This function is mainly intended for testing/debugging. Use it only for networks with few links.

See also [shownodes](#Fastnet.shownodes)

# Examples
```jldoctest
julia> using Fastnet

julia> net=FastNet(1000,2000,2,[LinkType(1,1),LinkType(1,2),LinkType(2,2)]);

julia> makenodes!(net,5,1)

julia> makenodes!(net,5,2)

julia> makelink!(net,node(net,1,1),node(net,1,2));
1

julia> makelink!(net,node(net,1,1),node(net,2,1))
2

julia> makelink!(net,node(net,2,1),node(net,2,2))
3

julia> makelink!(net,node(net,2,2),node(net,1,2))
4

julia> showlinks(net)
id      src     dest    state
1       1       2       1
2       1       6       2
3       6       7       3
4       7       2       2

"""
function showlinks(net::FastNet)
    println("id\tsrc\tdest\tstate")
    @inbounds begin
        for i=1:net.K
            if (net.kstate[i]<net.L)
                println(i,'\t',net.ksrc[i],'\t',net.kdst[i],'\t',net.kstate[i])
            end
        end
    end
end

function check_repositoryconsistency(net::FastNet)
    print("  Checking repository consistency ... ")
    error=false
    for nid=1:net.N                # check nodes
        pos=net.npos[nid]
        id=net.nid[pos]
        if id!=nid
            error=true
            print("\n     + Node $nid is listed at position $pos but this position refers to node $id")
        end
    end
    for kid=1:net.K                # check links
        pos=net.kpos[kid]
        id=net.kid[pos]
        if id!=kid
            error=true
            print("\n     + Link $kid is listed at position $pos but this position refers to link $id")
        end
    end

    if error
        println("\n    THERE WERE ERRORS\n")
        false
    else
        println("OK")
        true
    end
end

function check_nodeaccounting(net::FastNet)
    print("  Checking node accounting ... ")
    error=false
    numnodes=0
    for c=1:net.C
        if net.ncstart[c]!=numnodes                        # Check that the state starts at the right place
            error=true
            print("\n     + Node state $c starts at position $(net.ncstart[c]) but I found $numnodes nodes prior to this state")
        end
        numnodes+=net.nclen[c]                             # Adding up nodes in states to see
    end
    if numnodes!=net.N
        error=true
        print("\n     + Only $numnodes out of $(net.N) are accounted for in node states")
    end
    if error
        println("\n    THERE WERE ERRORS\n")
        false
    else
        println("OK")
        true
    end
end

function check_linkaccounting(net::FastNet)
    print("  Checking link accounting ... ")
    error=false
    numlinks=0
    for c=1:net.L
        if net.kcstart[c]!=numlinks                        # Check that the state starts at the right place
            error=true
            print("\n     + Link state $c starts at position $(net.kcstart[c]) but I found $numlinks links prior to this state")
        end
        numlinks+=net.kclen[c]                             # Adding up nodes in states to see
    end
    if numlinks!=net.K
        error=true
        print("\n     + Only $numnlinks out of $(net.K) are accounted for in link states")
    end
    if error
        println("\n    THERE WERE ERRORS\n")
        false
    else
        println("OK")
        true
    end
end

function check_endpointconsistency(net::FastNet)
    print("  Checking endpoint consistency ... ")
    error=false
    for nid=1:net.N         #check from node perspective
        c=net.nstate[nid]               # Check outgoing
        cur=net.nout[nid]
        if cur!=0 && c==net.C           # Deleted nodes should not have links
            error=true
            print("\n     + Deleted node $nid has an outgoing link (id: $cur)")
        end
        while (cur!=0)                     # Check source is correct
            src=net.ksrc[cur]
            if src!=nid
                error=true
                print("\n     + Node $nid has an outgoing link (id: $cur) which has node $src as source.")
            end
            if net.kstate[cur]==net.L
                error=true
                print("\n     + Outgoing link $cur from node $nid looks like it has been deleted.")
            end
            cur=net.knexts[cur]
        end
        cur=net.nin[nid]                        # Check incoming
        if cur!=0 && c==net.C                       # Deleted nnodes should not have links
            error=true
            print("\n     + Deleted node $nid has an incoming link (id: $cur)")
        end
        while (cur!=0)                             # Check link destination is correct
            dst=net.kdst[cur]
            if dst!=nid
                error=true
                print("\n     + Node $nid has an incoming link (id: $cur) which has node $dst as destination.")
            end
            if net.kstate[cur]==net.L
                error=true
                print("\n     + Incoming link $cur to node $nid looks like it has been deleted.")
            end
            cur=net.knextd[cur]
        end
    end

    for kid=1:net.K                             # Now checking from the link perspective
        if net.kstate[kid]<net.L                 # no need checking deleted links
            src=net.ksrc[kid]
            dst=net.kdst[kid]
            if net.nstate[src]==net.C             # Check that nodes actually exist
                error=true
                print("\n     + Link $kid starts at deleted node $src")
            end
            if net.nstate[dst]==net.C
                error=true
                print("\n     + Link $kid ends at deleted node $dst")
            end
            cur=net.nout[src]                   # Check that it is listed as outgoing from source
            while cur!=0 && cur!=kid
                cur=net.knexts[cur]
            end
            if cur!=kid
                error=true
                print("\n     + Link $kid says it starts at $src but that node does not list it as an outgoing link.")
            end
            cur=net.nin[dst]                   # Check that it is listed as incoming at dst
            while cur!=0 && cur!=kid
                cur=net.knextd[cur]
            end
            if cur!=kid
                error=true
                print("\n     + Link $kid says it ends at $dst but that node does not list it as an incoming link.")
            end
        end
    end

    if error
        println("\n    THERE WERE ERRORS\n")
        false
    else
        println("OK")
        true
    end
end

function check_nodestateification(net::FastNet)
    print("  Checking node stateification ... ")
    error=false
    for c=1:net.C                   # Check that state array is consistent with position in repo
        for r=1:net.nclen[c]
            nid=node(net,c,r)
            nc=net.nstate[nid]
            if nc<1 || nc>net.C
                error=true
                print("\n     + The state of node $nid is listed as $nc which is out of range.")
            end
            if nc!=c
                error=true
                print("\n     + Node $nid thinks it is in state $(net.nstate[nid]) but it's actually in state $c")
            end
        end
    end
    if error
        println("\n    THERE WERE ERRORS\n")
        false
    else
        println("OK")
        true
    end
end

function check_linkstateification(net::FastNet)
    print("  Checking link stateification ... ")
    error=false
    for c=1:net.L
        for r=1:net.kclen[c]
            kid=link_f(net,c,r)           # Check that state array is consistent with position in repo
            kc=net.kstate[kid]
            if kc!=c
                error=true
                print("\n     + Link $kid thinks it is in state $kc but it's actually in state $c")
            end
            if kc<1 || kc>net.L
                error=true
                print("\n     + The state of link $kid is listed as $kc which is out of range.")
            end
            if (c<net.L)                  # if the link is not deleted also check if state is consistent with linked node states
                src=net.ksrc[kid]
                dst=net.kdst[kid]
                sc=net.nstate[src]
                dc=net.nstate[dst]
                if sc<1 || dc<1 || sc>net.C || dc>net.C
                    error=true
                    print("\n     + I can't check the stateification of link $kid as one of it's endpoints has a state that is out of range.")
                else
                    cc=net.ttable[sc,dc]
                    if cc!=c
                        error=true
                        print("\n     + Link $kid is in state $c but it connects node $src (state $sc) to $dst (state $dc) which means it schould be in state $cc")
                    end
                end
            end
        end
    end
    if error
        println("\n    THERE WERE ERRORS\n")
        false
    else
        println("OK")
        true
    end
end

"""
    healthcheck(net)

Perform an internal consistencey check on a FastNet *net*.

To achieve the desired performance Fastnet engages in a certain amount of double bookeeping.
In an ideal world the FastNet structures should always stay internally consistent. However,
inconsistencies could arise from a number of sources including software bugs, CPU and memeory
errors. This function checks the internal data stored in FastNet for consistency to make sure
that everything is alright.

The return value is true if all chacks have been passed, false otherwise.

See also [link](#Fastnet.link)

# Examples
```jldoctest
julia> using Fastnet

julia> net=FastNet(100,200,10,[])
Network of 0 nodes and 0 links

julia> healthcheck(net);
  Checking repository consistency ... OK
  Checking node accounting ... OK
  Checking link accounting ... OK
  Checking endpoint consistency ... OK
  Checking node stateification ... OK
  Checking link stateification ... OK
```
"""
function healthcheck(net::FastNet)
    ok=true
    ok=ok && check_repositoryconsistency(net)
    ok=ok && check_nodeaccounting(net)
    ok=ok && check_linkaccounting(net)
    ok=ok && check_endpointconsistency(net)
    ok=ok && check_nodestateification(net)
    ok=ok && check_linkstateification(net)
end


text:
#####
##### Import dict: SingleFile
#####

export import_dict

function import_dict(D::Dict, outstyle::SingleFile, ft::AbstractFilterTypes=EntireStruct())
  is_leaf(s, prop) = __is_leaf(s, prop, ft)
  if !isempty([triggers(is_leaf(nothing, v)) for v in values(D)])
    open(outstyle.filename,"r") do io
      for (k,v) in D
        if triggers(is_leaf(nothing, v))
          val_s = strip(last(split(readline(io), "=")), ' ')
          D[k] = get_val_from_var(val_s)
        end
      end
    end
  end
end

#####
##### Import dict: FolderStructure
#####

function import_dict(D::Dict, outstyle::FolderStructure, ft::AbstractFilterTypes=EntireStruct())
  get_path(prop_chain) = joinpath(split(prop_chain, ".")[1:end-1]...)
  get_prop_chain_root(k) = join(split(k, ".")[1:end-1]...)
  get_varname(prop_chain) = split(prop_chain, ".")[end]
  is_leaf(s, prop) = __is_leaf(s, prop, ft)
  mkpath(outstyle.root_folder)
  D_paths = Dict([get_path(k) => [] for k in keys(D)])
  for (k,v) in D
    kp = get_path(k)
    var = get_varname(k)
    push!(D_paths[kp], (var, get_prop_chain_root(k), D[k]))
  end

  for (path,arr) in D_paths
    _fullpath = joinpath(outstyle.root_folder, path)
    folder_needed = [triggers(is_leaf(nothing, val)) for (var, pcr, val) in arr]
    if any(folder_needed)
      mkpath(_fullpath)
      open(joinpath(_fullpath, outstyle.filename), "r") do io
        for (var, prop_chain_root, val) in arr
          if triggers(is_leaf(nothing, val))
            D[prop_chain_root*"."*k] = get_val_from_var(strip(last(split(readline(io), "=")), ' '))
          end
        end
      end
    end
  end
end


text:
using ForecastAccuracy, Test, Random

include("utils.jl")

@testset "Exact forecast" begin
    for m in ForecastAccuracy.METRICS
        test_exact_forecast(m)
    end
end


text:
###############################################################################
#
# Big leaf canopy model
#
###############################################################################
const _K_D         = 0.7
const _ALPAR       = 0.8
const _ALNIR       = 0.2
const _ALPAR_SQ    = sqrt(_ALPAR)
const _ALNIR_SQ    = sqrt(_ALNIR)
const _ALPAR_SQ_KD = _ALPAR_SQ * _K_D
const _ALNIR_SQ_KD = _ALNIR_SQ * _K_D

"""
    big_leaf_partition(
                lai::FT,
                zenith::FT,
                r_all::FT,
                r_dir::FT = FT(0.8)
    ) where {FT<:AbstractFloat}

Partition the big-leaf canopy into sunlit and shaded layers, given
- `lai` Leaf area index
- `zenith` Zenith angle in degree
- `r_all` Total radiation in `[W m⁻²]`
- `r_dir` Direct radiation partition in `r_all`

The function returns
- `ratio` ratio of sunlit leaves out of all leaves
- `q_slm` Mean sunlit layer PAR
- `q_shm` Mean shaded layer PAR
- `e_sl` Mean sunlit layer absorbed total energy
- `e_sh` Mean shaded layer absorbed total energy
"""
function big_leaf_partition(
            lai::FT,
            zenith::FT,
            r_all::FT,
            r_dir::FT = FT(0.8)
) where {FT <:AbstractFloat}
    # 1. assume 50%-50% PAR and NIR
    par_tot = r_all / 2 / FT(0.235);
    q_ob    = par_tot * r_dir;
    q_od    = par_tot * (1 - r_dir);

    # 2. calculate the LAI
    shape = FT(1.0);
    shapa = FT(2.0);

    # 3. calculate the mean sunlit layer PAR
    # For vertical leaves k_be = 2.0 * tand(zenith) / pi
    k_be  = sqrt( shape^2+tand(zenith)^2 ) /
              ( shape + FT(1.774) * (shape+FT(1.182))^FT(-0.733) );
    q_btp = q_ob * exp( -FT(_ALPAR_SQ) * k_be * lai );
    q_btn = q_ob * exp( -FT(_ALNIR_SQ) * k_be * lai );
    q_b   = q_ob * exp( -k_be*lai );
    q_scp = q_btp - q_b;
    q_scn = q_btn - q_b;

    # 4. mean par in sunlit and shade layers
    q_dm  = q_od * ( 1 - exp(-FT(_ALPAR_SQ_KD)) ) / ( FT(_ALPAR_SQ_KD) * lai );
    r_dm  = q_od * ( 1 - exp(-FT(_ALNIR_SQ_KD)) ) / ( FT(_ALNIR_SQ_KD) * lai );
    q_shm = q_dm + q_scp/shapa;
    r_shm = r_dm + q_scn/shapa;
    q_slm = k_be*q_ob + q_dm + q_scp/shapa;
    r_slm = k_be*q_ob + r_dm + q_scn/shapa;

    # 5. calculate lasi_sl and lai_sh
    lai_sl = ( 1 - exp(-k_be*lai) ) / k_be;
    lai_sh = lai - lai_sl;
    ratio  = lai_sl / lai;

    # 6. calculate the radiation based on PAR and NIR
    ep_sl = q_slm * FT(_ALPAR) * FT(0.235);
    ep_sh = q_shm * FT(_ALPAR) * FT(0.235);
    en_sl = r_slm * FT(_ALNIR) * FT(0.235);
    en_sh = r_shm * FT(_ALNIR) * FT(0.235);

    # e_sl and e_sh are leaf-area-based mean radiation
    e_sl = ep_sl + en_sl;
    e_sh = ep_sh + en_sh;

    # return values
    return ratio, q_slm, q_shm, e_sl, e_sh
end

max_stars_repo_name: bbuman/Land

max_issues_repo_name: bbuman/Land

max_forks_repo_name: bbuman/Land


text:
using Base.Test
using QuantumOptics

# Set up operators
spinbasis = SpinBasis(1//2)

sx = sigmax(spinbasis)
sy = sigmay(spinbasis)

sx_dense = full(sx)
sy_dense = full(sy)

@test typeof(sx_dense) == DenseOperator
@test typeof(sparse(sx_dense)) == SparseOperator
@test sparse(sx_dense) == sx

b = FockBasis(3)
I = identity(b)
I_dense = dense_identity(b)

s = tensor(sx, sy)
s_dense = tensor(sx_dense, sy_dense)

@test typeof(I) == SparseOperator
@test typeof(I_dense) == DenseOperator
@test_approx_eq 0. norm((I_dense-full(I)).data)
@test_approx_eq 0. norm((s_dense - full(s)).data)

@test I == identity(destroy(b))

type A <: Operator
end

a = A()

@test_throws ArgumentError sparse(a)


text:
center = 5
sigma = 2

mf = GaussianMF(center, sigma)

# Evaluation tests
@assert mf.eval(center + sigma) == exp(-0.5)
@assert mf.eval(center) == 1
@assert mf.eval(center + sigma) == mf.eval(center - sigma)

# Mean finding tests
@assert mf.mean_at(1) == mf.mean_at(0.6) == mf.mean_at(0.3) == mf.mean_at(0) == center


text:
using Langevin
using KitBase.ProgressMeter: @showprogress

function ev!(KS, sol, flux, dt)
    @inbounds for i in eachindex(flux.fw)
        for j in axes(sol.w[1], 2)
            fw = @view flux.fw[i][:, j]
            fh = @view flux.fh[i][:, :, j]
            fb = @view flux.fb[i][:, :, j]

            flux_kfvs!(
                fw,
                fh,
                fb,
                sol.h[i-1][:, :, j] .+ 0.5 .* KS.pSpace.dx[i-1] .* sol.∇h[i-1][:, :, j],
                sol.b[i-1][:, :, j] .+ 0.5 .* KS.pSpace.dx[i-1] .* sol.∇b[i-1][:, :, j],
                sol.h[i][:, :, j] .- 0.5 .* KS.pSpace.dx[i] .* sol.∇h[i][:, :, j],
                sol.b[i][:, :, j] .- 0.5 .* KS.pSpace.dx[i] .* sol.∇b[i][:, :, j],
                KS.vSpace.u,
                KS.vSpace.v,
                KS.vSpace.weights,
                dt,
                1.0, # interface length
                sol.∇h[i-1][:, :, j],
                sol.∇b[i-1][:, :, j],
                sol.∇h[i][:, :, j],
                sol.∇b[i][:, :, j],
            )
        end
    end
end

function up!(KS, uq, sol, flux, dt, residual)
    w_old = deepcopy(sol.w)

    @inbounds Threads.@threads for i = 1:KS.pSpace.nx
        @. sol.w[i] += (flux.fw[i] - flux.fw[i+1]) / KS.pSpace.dx[i]
        sol.prim[i] .= uq_conserve_prim(sol.w[i], KS.gas.γ, uq)
    end

    τ = uq_vhs_collision_time(sol, KS.gas.μᵣ, KS.gas.ω, uq)
    H = [
        uq_maxwellian(KS.vSpace.u, KS.vSpace.v, sol.prim[i], uq) for
        i in eachindex(sol.prim)
    ]
    B = deepcopy(H)
    for i = 1:KS.pSpace.nx
        for j in axes(B[1], 3)
            B[i][:, :, j] .= H[i][:, :, j] .* KS.gas.K ./ (2.0 .* sol.prim[i][end, j])
        end
    end

    @inbounds Threads.@threads for i = 1:KS.pSpace.nx
        for j in axes(sol.w[1], 2)
            @. sol.h[i][:, :, j] =
                (
                    sol.h[i][:, :, j] +
                    (flux.fh[i][:, :, j] - flux.fh[i+1][:, :, j]) / KS.pSpace.dx[i] +
                    dt / τ[i][j] * H[i][:, :, j]
                ) / (1.0 + dt / τ[i][j])
            @. sol.b[i][:, :, j] =
                (
                    sol.b[i][:, :, j] +
                    (flux.fb[i][:, :, j] - flux.fb[i+1][:, :, j]) / KS.pSpace.dx[i] +
                    dt / τ[i][j] * B[i][:, :, j]
                ) / (1.0 + dt / τ[i][j])
        end
    end

    # record residuals
    sumRes = zeros(axes(sol.w[1], 1))
    sumAvg = zeros(axes(sol.w[1], 1))
    for j in axes(sumRes, 1)
        for i = 1:KS.pSpace.nx
            sumRes[j] += sum((sol.w[i][j, :] .- w_old[i][j, :]) .^ 2)
            sumAvg[j] += sum(abs.(sol.w[i][j, :]))
        end
    end
    @. residual = sqrt(sumRes * KS.pSpace.nx) / (sumAvg + 1.e-7)

    return nothing
end

set = Setup(case = "layer", space = "1d2f2v", maxTime = 0.5539)
ps = PSpace1D(-1, 1, 1000, 1)
vs = VSpace2D(-4.5, 4.5, 32, -4.5, 4.5, 64)
gas = Gas(Kn = 0.005, K = 1)

ib = begin
    primL = [1.0, 0.0, 1.0, 1.0]
    primR = [1.0, 0.0, -1.0, 2.0]
    wL = prim_conserve(primL, gas.γ)
    wR = prim_conserve(primR, gas.γ)

    p = (
        x0 = ps.x0,
        x1 = ps.x1,
        u = vs.u,
        γ = gas.γ,
        K = gas.K,
        wL = wL,
        wR = wR,
        primL = primL,
        primR = primR,
    )

    fw = function (x, p)
        if x <= (p.x0 + p.x1) / 2
            return p.wL
        else
            return p.wR
        end
    end

    bc = function (x, p)
        if x <= (p.x0 + p.x1) / 2
            return p.primL
        else
            return p.primR
        end
    end

    ff = function (x, p)
        w = ifelse(x <= (p.x0 + p.x1) / 2, p.wL, p.wR)
        prim = conserve_prim(w, p.γ)
        h = maxwellian(p.u, p.v, prim)
        b = @. h * p.K / 2 / prim[end]
        return h, b
    end

    IB2F{typeof(bc)}(fw, ff, bc, p)
end

ks = SolverSet(set, ps, vs, gas, ib)
uq = UQ1D(5, 10, 0.9, 1.1, "uniform", "collocation")
sol, flux = Langevin.init_collo_sol(ks, uq)

# add uncertainty
for i in eachindex(ks.pSpace.x)
    if i <= ks.pSpace.nx ÷ 2
        for j in axes(sol.w[1], 2)
            sol.prim[i][3, j] *= uq.pceSample[j]
            sol.w[i] .= uq_prim_conserve(sol.prim[i], ks.gas.γ, uq)
            sol.h[i], sol.b[i] =
                uq_maxwellian(ks.vSpace.u, ks.vSpace.v, sol.prim[i], uq, ks.gas.K)
        end
    end
end

simTime = 0.0
iter = 0

dt = 0.2 * ks.pSpace.dx[1]
nt = floor(ks.set.maxTime / dt + 1) |> Int

#simTime, iter = solve!(ks, uq, sol, flux, simTime, dt, nt)

reconstruct!(ks, sol)
ev!(ks, sol, flux, dt)
up!(ks, uq, sol, flux, dt, zeros(4, uq.op.quad.Nquad))


text:
using .IntervalConstraintProgramming

"""
returns the unrolled expression for \$|a ⋅x - b|\$
\$a\$ and \$x\$ must be vectors of the same length and \$b\$ is scalar.
The absolue value in the equation is taken elementwise.
"""
function oettli_lhs(a, b, x)
    ex = :( $(a[1])*$(x[1]))
    for i = 2:length(x)
        ex = :( $ex + $(a[i])*$(x[i]))
    end
    return :(abs($ex - $b))
end

"""
returns the unrolled expression for \$a ⋅|x| + b\$
\$a\$ and \$x\$ must be vectors of the same length and \$b\$ is scalar.
The absolue value in the equation is taken elementwise.
"""
function oettli_rhs(a, b, x)
    ex = :( $(a[1])*abs($(x[1])))
    for i = 2:length(x)
        ex = :( $ex + $(a[i])*abs($(x[i])))
    end
    return :($ex + $b)
end

"""
Returns the separator for the constraint `|a_c ⋅x - b_c| - a_r ⋅|x| - b_r <= 0`.

`a` and `x` must be vectors of the same length and `b` is scalar.

The absolue values in the equation are taken elementwise.

`a_c` and `a_r` are vectors containing midpoints and radii of the intervals in `a`. Similar `b_c` and `b_r`.
"""
function oettli_eq(a, b, x)
    ac = mid.(a)
    ar = IntervalArithmetic.radius.(a)
    bc = mid(b)
    br = IntervalArithmetic.radius(b)
    lhs = oettli_lhs(ac, bc, x)
    rhs = oettli_rhs(ar, br, x)
    ex = :(@constraint $lhs - $rhs <= 0)
    @eval $ex

end

function (op::NonLinearOettliPrager)(A, b, X::IntervalBox)
    vars = ntuple(i -> Symbol(:x, i), length(b))
    separators = [oettli_eq(A[i,:], b[i], vars) for i in 1:length(b)]
    S = reduce(∩, separators)
    return Base.invokelatest(pave, S, X, op.tol)
end

(op::NonLinearOettliPrager)(A, b, X=enclose(A, b)) = op(A, b, IntervalBox(X))

_default_precondition(_, ::NonLinearOettliPrager) = NoPrecondition()


text:
module SlurmAutoDetectWorkers

import Distributed
import ExistingProcessManagers

_foo(x) = x + x

end


text:
# ---------------------- #
# Jupyter notebook setup #
# ---------------------- #

mutable struct JupyterDisplay <: AbstractPlotlyDisplay
    divid::Base.Random.UUID
    displayed::Bool
    cond::Condition  # for getting data back from js
end

const JupyterPlot = SyncPlot{JupyterDisplay}

JupyterDisplay(p::Plot) = begin
    _ijulia_eval_comm[] = Comm(:plotlyjs_eval)
    JupyterDisplay(p.divid, false, Condition())
end
JupyterPlot(p::Plot) = JupyterPlot(p, JupyterDisplay(p))

fork(jp::JupyterPlot) = JupyterPlot(fork(jp.plot))

const _jupyter_js_loaded = [false]
js_loaded(::JupyterDisplay) = _jupyter_js_loaded[1]
js_loaded(::Type{JupyterDisplay}) = _jupyter_js_loaded[1]

function html_body(p::JupyterPlot)
    """

    """
end

function init_notebook(force=false)
    # TODO: figure out a way to ask the notebook if the js is currently loaded
    #       and active.

    # if we're in IJulia call setup the notebook js interop
    if _isijulia()
        # borrowed from https://github.com/plotly/plotly.py/blob/2594076e29584ede2d09f2aa40a8a195b3f3fc66/plotly/offline/offline.py#L64-L71
        # and https://github.com/JuliaLang/Interact.jl/blob/cc5f4cfd34687000bc6bc70f0513eaded1a7c950/src/IJulia/setup.jl#L15
        if !js_loaded(JupyterDisplay) || force
            _ijulia_js = readstring(joinpath(dirname(@__FILE__), "ijulia.js"))

            # three script tags for loading ijulia setup, and plotly
            display("text/html", """

Plotly javascript loaded.

To load again call
init_notebook(true)

             """)
            _jupyter_js_loaded[1] = true
        end
    end
end

# ---------------- #
# Helper functions #
# ---------------- #

_call_js(jd::JupyterDisplay, code) =
    send_comm(_ijulia_eval_comm[], Dict("code" => code))

## API Methods for JupyterDisplay
_the_div_js(jd::JupyterDisplay) = "document.getElementById('$(jd.divid)')"
_the_div_js(jp::JupyterPlot) = _the_div_js(jp.view)

# function _img_data(jp::JupyterPlot, format::String)
#     _formats = ["png", "jpeg", "webp", "svg"]
#     if !(format in _formats)
#         error("Unsupported format $format, must be one of $_formats")
#     end
#
#     if format == "svg"
#         return svg_data(jp)
#     end
#
#     code =  """
#     ev = Plotly.Snapshot.toImage($(_the_div_js(jp)), {format: '$(format)'});
#     new Promise(function(resolve) {ev.once("success", resolve)});
#     """
#     _call_js_return(jp.view, code)
# end
#
# function svg_data(jp::JupyterPlot, format="png")
#     code =  "Plotly.Snapshot.toSVG($(_the_div_js(jp)), '$(format)')"
#     _call_js_return(jp.view, code)
# end

function _call_plotlyjs(jd::JupyterDisplay, func::AbstractString, args...)
    arg_str = length(args) > 0 ? string(",", join(map(json, args), ", ")) : ""
    code = "Plotly.$func($(_the_div_js(jd)) $arg_str)"
    jd.displayed && _call_js(jd, code)
    nothing
end

# ---------------------- #
# Javascript API methods #
# ---------------------- #

relayout!(jd::JupyterDisplay, update::Associative=Dict(); kwargs...) =
    _call_plotlyjs(jd, "relayout", merge(update, prep_kwargs(kwargs)))

restyle!(jd::JupyterDisplay, ind::Int, update::Associative=Dict(); kwargs...) =
    _call_plotlyjs(jd, "restyle", merge(update, prep_kwargs(kwargs)), ind-1)

function restyle!(jd::JupyterDisplay, inds::AbstractVector{Int},
                  update::Associative=Dict();  kwargs...)
    _call_plotlyjs(jd, "restyle", merge(update, prep_kwargs(kwargs)), inds-1)
end

restyle!(jd::JupyterDisplay, update::Associative=Dict(); kwargs...) =
    _call_plotlyjs(jd, "restyle", merge(update, prep_kwargs(kwargs)))

addtraces!(jd::JupyterDisplay, traces::AbstractTrace...) =
    _call_plotlyjs(jd, "addTraces", traces)

addtraces!(jd::JupyterDisplay, where::Int, traces::AbstractTrace...) =
    _call_plotlyjs(jd, "addTraces", traces, where-1)

deletetraces!(jd::JupyterDisplay, traces::Int...) =
    _call_plotlyjs(jd, "deleteTraces", collect(traces)-1)

movetraces!(jd::JupyterDisplay, to_end::Int...) =
    _call_plotlyjs(jd, "moveTraces", collect(to_end)-1)

movetraces!(jd::JupyterDisplay, src::AbstractVector{Int}, dest::AbstractVector{Int}) =
    _call_plotlyjs(jd, "moveTraces", src-1, dest-1)

redraw!(jd::JupyterDisplay) = _call_plotlyjs(jd, "redraw")
purge!(jd::JupyterDisplay) = _call_plotlyjs(jd, "purge")

to_image(jd::JupyterDisplay; kwargs...) =
    _call_plotlyjs(jd, "to_image", Dict(kwargs))

download_image(jd::JupyterDisplay; kwargs...) =
    _call_plotlyjs(jd, "download_image", Dict(kwargs))

# unexported (by plotly.js) api methods
extendtraces!(jd::JupyterDisplay, update::Associative=Dict(),
              indices::Vector{Int}=[1], maxpoints=-1;) =
    _call_plotlyjs(jd, "extendTraces", update, indices-1, maxpoints)

prependtraces!(jd::JupyterDisplay, update::Associative=Dict(),
               indices::Vector{Int}=[1], maxpoints=-1;) =
    _call_plotlyjs(jd, "prependTraces", update, indices-1, maxpoints)


text:
using SparseArrays

# Dummy functions to allow not specifying these characteristics
function F_dummy(ẋ, xc, xd, parms, t)
	fill!(ẋ, 0)
	nothing
end

# Dummy flow to be used in rejection algorithm
function Phi_dummy(out, xc, xd, parms, t)
	# vector field used for the continuous variable
	# trivial dynamics
	out[1,:] .= xc
	out[2,:] .= xc
	nothing
end

mutable struct PDMPJumpTime{Tc <: Real, Td}
	tstop_extended::Tc
	lastjumptime::Tc
	njumps::Td

	# fields required for the rejection method
	lambda_star::Tc					# bound on the total rate
	ppf::Vector{Tc}
	reject::Bool					# boolean to know whether to reject or not the step
	fictitous_jumps::Td
end

struct PDMPCaracteristics{TF, TR, TJ, vecc, vecd, vecrate, Tparms}
	F::TF						# vector field for ODE between jumps
	R::TR			    		# rate function for jumps
	pdmpjump::TJ

	xc::vecc					# current continuous variable
	xd::vecd					# current discrete variable

	xc0::vecc					# initial continuous variable
	xd0::vecd					# initial discrete variable
	ratecache::vecrate			# to hold the rate vector for inplace computations. Also used to initialise rate as this can be an issue for StaticArrays.jl
	parms::Tparms				# container to hold parameters to be passed to F, R, Delta

	function PDMPCaracteristics(F, R, Delta, nu::Tnu, xc0::vecc, xd0::vecd, parms::Tparms) where {Tc, Td, Tparms, Tnu <: AbstractMatrix{Td},
						vecc <: AbstractVector{Tc},
						vecd <: AbstractVector{Td}}
		jump = Jump(nu, Delta)
		rate = dualcache(get_rate_prototype(jump, Tc))
		ratefunction = VariableRate(R)
		return new{typeof(F), typeof(ratefunction), typeof(jump), vecc, vecd, typeof(rate), Tparms}(F, ratefunction, jump, copy(xc0), copy(xd0), copy(xc0), copy(xd0), rate, parms)
	end

	function PDMPCaracteristics(F, R::TR, Delta, nu::Tnu, xc0::vecc, xd0::vecd, parms::Tparms) where {Tc, Td, Tparms, Tnu <: AbstractMatrix{Td},
						vecc <: AbstractVector{Tc},
						vecd <: AbstractVector{Td},
						TR <: AbstractRate}
		jump = Jump(nu, Delta)
		rate = dualcache(get_rate_prototype(jump, Tc))
		return new{typeof(F), typeof(R), typeof(jump), vecc, vecd, typeof(rate), Tparms}(F, R, jump, copy(xc0), copy(xd0), copy(xc0), copy(xd0), rate, parms)
	end
end

function PDMPCaracteristics(F, R, nu::Tnu, xc0::vecc, xd0::vecd, parms::Tparms) where {Tc, Td, Tparms, Tnu <: AbstractMatrix{Td}, vecc <: AbstractVector{Tc}, vecd <: AbstractVector{Td}}
	return PDMPCaracteristics(F, R, Delta_dummy, nu, xc0, xd0, parms)
end

function init!(pb::PDMPCaracteristics)
	pb.xc .= pb.xc0
	pb.xd .= pb.xd0
	init!(pb.R)
end

struct PDMPProblem{Tc, Td, vectype_xc <: AbstractVector{Tc},
						vectype_xd <: AbstractVector{Td},
						Tcar}
	tspan::Vector{Tc}				    			# final simulation time interval, we use an array to be able to mutate it
	simjptimes::PDMPJumpTime{Tc, Td}				# space to save result
	time::Vector{Tc}
	Xc::VectorOfArray{Tc, 2, Array{vectype_xc, 1}}	# continuous variable history
	Xd::VectorOfArray{Td, 2, Array{vectype_xd, 1}}	# discrete variable history
	# variables for debugging
	rate_hist::Vector{Tc}							# to save the rates for debugging purposes
	caract::Tcar									# struct for characteristics of the PDMP
end

pushTime!(pb::PDMPProblem, t) = push!(pb.time, t)
pushXc!(pb::PDMPProblem, xc) = push!(pb.Xc, xc)
pushXd!(pb::PDMPProblem, xd) = push!(pb.Xd, xd)

function init!(pb::PDMPProblem)
	init!(pb.caract)
	pb.simjptimes.tstop_extended = -log(rand())
	pb.simjptimes.lastjumptime = pb.tspan[1]
	pb.simjptimes.njumps = 0
	pb.simjptimes.fictitous_jumps = 0
	resize!(pb.time, 1)
	resize!(pb.rate_hist, 1)
	resize!(pb.Xc.u, 1)
	resize!(pb.Xd.u, 1)
end

# callable struct used in the iterator interface
function (prob::PDMPProblem)(u, t, integrator)
	t == prob.simjptimes.tstop_extended
end

function PDMPProblem(F::TF, R::TR, DX::TD, nu::Tnu,
				xc0::vecc, xd0::vecd, parms::Tp,
				tspan) where {Tc, Td, Tnu <: AbstractMatrix{Td}, Tp, TF ,TR ,TD, vecc <: AbstractVector{Tc}, vecd <:  AbstractVector{Td}}
	ti, tf = tspan
	caract = PDMPCaracteristics(F, R, DX, nu, xc0, xd0, parms)
	return PDMPProblem{Tc, Td, vecc, vecd, typeof(caract)}(
			[ti, tf],
			PDMPJumpTime{Tc, Td}(Tc(0), ti, 0, Tc(0), Vector{Tc}([0, 0]), false, 0),
			[ti],
			VectorOfArray([copy(xc0)]), VectorOfArray([copy(xd0)]),
			Tc[],
			caract)
end

function PDMPProblem(F, R, nu::Tnu, xc0::vecc, xd0::vecd, parms,
				tspan; kwargs...) where {Tc, Td, Tnu <: AbstractMatrix{Td}, vecc <: AbstractVector{Tc}, vecd <:  AbstractVector{Td}}
	return PDMPProblem(F, R, Delta_dummy, nu, xc0, xd0, parms, tspan; kwargs...)
end

function PDMPProblem(F, R, Delta, reaction_number::Int64, xc0::vecc, xd0::vecd, parms,
				tspan; kwargs...) where {Tc, Td, vecc <: AbstractVector{Tc}, vecd <:  AbstractVector{Td}}
	return PDMPProblem(F, R, Delta, spzeros(Int64, reaction_number, length(xd0)), xc0, xd0, parms, tspan; kwargs...)
end