text:
using PyPlot
using LinearAlgebra
using PolynomialBasis
using ImplicitDomainQuadrature
using Revise
using CutCell
include("../../../test/useful_routines.jl")

function bulk_modulus(l, m)
    return l + 2m / 3
end

function lame_lambda(k, m)
    return k - 2m / 3
end

function displacement(x)
    u1 = x[1]^2 + 2 * x[1] * x[2]
    u2 = x[2]^2 + 3x[1]
    return [u1, u2]
end

function body_force(lambda, mu)
    b1 = -2 * (lambda + 2mu)
    b2 = -(4lambda + 6mu)
    return [b1, b2]
end

function onboundary(x, L, W)
    return x[2] ≈ 0.0 || x[1] ≈ L || x[2] ≈ W || x[1] ≈ 0.0
end

function solve_for_displacement(
    basis,
    cellquads,
    interfacequads,
    facequads,
    cutmesh,
    levelset,
    levelsetcoeffs,
    stiffness,
    penalty,
    alpha,
)

    L, W = CutCell.widths(cutmesh)
    lambda, mu = CutCell.lame_coefficients(stiffness, +1)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition = CutCell.coherent_interface_condition(
        basis,
        interfacequads,
        stiffness,
        cutmesh,
        penalty,
    )
    displacementbc = CutCell.DisplacementCondition(
        x -> displacement(x),
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onboundary(x, L, W),
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(
        sysmatrix,
        interfacecondition,
        cutmesh,
    )
    CutCell.assemble_body_force_linear_form!(
        sysrhs,
        x -> body_force(lambda, mu),
        basis,
        cellquads,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(
        sysmatrix,
        sysrhs,
        displacementbc,
        cutmesh,
    )

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    nodaldisplacement = matrix \ rhs

    return nodaldisplacement
end

function spatial_closest_points(refclosestpoints, refclosestcellids, mesh)
    dim, npts = size(refclosestpoints)
    spclosestpoints = zeros(dim, npts)
    for i = 1:npts
        cellmap = CutCell.cell_map(mesh, refclosestcellids[i])
        spclosestpoints[:, i] .= cellmap(refclosestpoints[:, i])
    end
    return spclosestpoints
end

K1, K2 = 247.0, 247.0    # Pa
mu1, mu2 = 126.0, 126.0   # Pa
lambda1 = lame_lambda(K1, mu1)
lambda2 = lame_lambda(K2, mu2)
stiffness = CutCell.HookeStiffness(lambda1, mu1, lambda2, mu2)

theta0 = 0.0
transfstress = CutCell.plane_strain_transformation_stress(lambda1,mu1,theta0)

width = 1.0
displacementscale = 0.01 * width
penaltyfactor = 1e2

polyorder = 2
numqp = required_quadrature_order(polyorder) + 2
nelmts = 3
delta = 0.01
interfacepoint = [1/3+delta, 1/3]
interfaceangle = 45
interfacenormal = [cosd(interfaceangle),sind(interfaceangle)]

dx = width / nelmts
meanmoduli = 0.5 * (lambda1 + lambda2 + mu1 + mu2)
penalty = penaltyfactor * meanmoduli

basis = TensorProductBasis(2, polyorder)
mesh = CutCell.Mesh([0.0, 0.0], [width, width], [nelmts, nelmts], basis)
levelset = InterpolatingPolynomial(1, basis)
levelsetcoeffs = CutCell.levelset_coefficients(
    x -> plane_distance_function(x, interfacenormal, interfacepoint),
    mesh,
)

cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
interfacequads =
    CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
facequads = CutCell.FaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

nodaldisplacement = solve_for_displacement(
    basis,
    cellquads,
    interfacequads,
    facequads,
    cutmesh,
    levelset,
    levelsetcoeffs,
    stiffness,
    penalty,
    displacementscale,
)

refseedpoints, spatialseedpoints, seedcellids =
    CutCell.seed_zero_levelset_with_interfacequads(interfacequads, cutmesh)
interfacenormals = CutCell.collect_interface_normals(interfacequads,cutmesh)

spatialpoints = spatialseedpoints
referencepoints = refseedpoints
referencecellids = seedcellids

sortidx = sortperm(spatialpoints[2, :])
spatialpoints = spatialpoints[:, sortidx]
referencepoints = referencepoints[:, sortidx]
referencecellids = referencecellids[sortidx]
spycoords = spatialpoints[2, :]
interfacenormals = interfacenormals[:,sortidx]

productdisplacement = CutCell.displacement_at_reference_points(
    referencepoints,
    referencecellids,
    +1,
    basis,
    nodaldisplacement,
    cutmesh,
)

parentdisplacement = CutCell.displacement_at_reference_points(
    referencepoints,
    referencecellids,
    -1,
    basis,
    nodaldisplacement,
    cutmesh,
)

exactdisplacement = hcat(
    [
        displacement(spatialpoints[:, i]) for
        i = 1:size(spatialpoints)[2]
    ]...,
)

# fig, ax = PyPlot.subplots(2,1)
# ax[1].plot(spycoords, productdisplacement[1, :], label = "product u1")
# ax[1].plot(spycoords, parentdisplacement[1, :], label = "parent u1")
# # ax[1].plot(spycoords, exactdisplacement[1, :], "--", label = "exact")
# ax[1].grid()
# ax[1].legend()
# ax[2].plot(spycoords, productdisplacement[2, :], label = "product u2")
# ax[2].plot(spycoords, parentdisplacement[2, :], label = "parent u2")
# # ax[2].plot(spycoords, exactdisplacement[2, :], "--", label = "exact")
# ax[2].grid()
# ax[2].legend()
# fig.tight_layout()
# fig

productstress = CutCell.product_stress_at_reference_points(
    referencepoints,
    referencecellids,
    basis,
    stiffness,
    transfstress,
    theta0,
    nodaldisplacement,
    cutmesh,
)
parentstress = CutCell.parent_stress_at_reference_points(
    referencepoints,
    referencecellids,
    basis,
    stiffness,
    nodaldisplacement,
    cutmesh,
)
#
# exactstress = hcat([stress_field(lambda1,mu1,displacementscale,spatialpoints[:,i]) for i = 1:size(spatialpoints)[2]]...)

producttraction = CutCell.traction_force_at_points(productstress,interfacenormals)
parenttraction = CutCell.traction_force_at_points(parentstress,interfacenormals)

fig, ax = PyPlot.subplots(2, 1)
ax[1].plot(spycoords, producttraction[1, :], label = "product")
ax[1].plot(spycoords, parenttraction[1, :], label = "parent")
ax[1].set_title("t1")
ax[1].legend()
ax[1].grid()
ax[2].plot(spycoords, producttraction[2, :], label = "product")
ax[2].plot(spycoords, parenttraction[2, :], label = "parent")
ax[2].set_title("t2")
ax[2].legend()
ax[2].grid()
fig.tight_layout()
fig
#
#
# using Statistics
# difftraction = parenttraction - producttraction
#
# idx = findall(difftraction[1,:] .> 3.0)
# oscillating_cells = referencecellids[idx]
#
# fig, ax = PyPlot.subplots(2, 1)
# ax[1].plot(spycoords, difftraction[1, :])
# ax[1].set_title("diff t1")
# ax[1].legend()
# ax[1].grid()
# ax[2].plot(spycoords, difftraction[2, :])
# ax[2].set_title("diff t2")
# ax[2].legend()
# ax[2].grid()
# fig.tight_layout()
# fig

# fig, ax = PyPlot.subplots(3, 1)
# ax[1].plot(spycoords, productstress[1, :], label = "product")
# ax[1].plot(spycoords, parentstress[1, :], label = "parent")
# ax[1].plot(spycoords, exactstress[1,:], "--", label = "exact")
# ax[1].set_title("S11")
# ax[1].legend()
# ax[1].grid()
# ax[2].plot(spycoords, productstress[2, :], label = "product")
# ax[2].plot(spycoords, parentstress[2, :], label = "parent")
# ax[2].plot(spycoords, exactstress[2,:], "--", label = "exact")
# ax[2].set_title("S22")
# ax[2].legend()
# ax[2].grid()
# ax[3].plot(spycoords, productstress[3, :], label = "product")
# ax[3].plot(spycoords, parentstress[3, :], label = "parent")
# ax[3].plot(spycoords, exactstress[3,:], "--", label = "exact")
# ax[3].set_title("S12")
# ax[3].legend()
# ax[3].grid()
# # ax[3].set_ylim(2,2.5)
# fig.tight_layout()
# fig

# productpressure = CutCell.pressure_at_points(productstress)
# parentpressure = CutCell.pressure_at_points(parentstress)
#
# fig,ax = PyPlot.subplots()
# ax.plot(spycoords,productpressure,label="product")
# ax.plot(spycoords,parentpressure,label="parent")
# ax.legend()
# ax.grid()
# ax.set_title("Pressure on interface")
# fig


text:
using Test

include("insertion_sort.jl")
include("selection_sort.jl")
include("merge_sort.jl")
include("heap_sort.jl")
include("quick_sort.jl")
include("counting_sort.jl")

@testset "test sorting correctness" begin
    inputs = [
        [5 2 4 7 1 3 2 6],
        [4 1 3 2 16 9 10 14 8 7],
        [2 8 7 1 3 5 6 4],
        [2 5 3 0 2 3 0 3],
        [329 457 657 839 436 720 355],
    ]
    outputs = [
        [1 2 2 3 4 5 6 7],
        [1 2 3 4 7 8 9 10 14 16],
        [1 2 3 4 5 6 7 8],
        [0 0 2 2 3 3 3 5],
        [329 355 436 457 657 720 839],
    ]

    @testset "insertion sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            insertion_sort!(A)
            @test A == output
        end
    end

    @testset "selection sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            selection_sort!(A)
            @test A == output
        end
    end

    @testset "merge sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            merge_sort!(A, 1, length(A))
            @test A == output
        end
    end

    @testset "heap sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            heap_sort!(A)
            @test A == output
        end
    end

    @testset "quick sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            quick_sort!(A, 1, length(A))
            @test A == output

            A = copy(input)
            randomized_quick_sort!(A, 1, length(A))
            @test A == output
        end
    end

    @testset "counting sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            k = maximum(A)
            B = counting_sort(A, k)
            @test B == output
        end
    end

    @testset "radix sort" begin
        for (input, output) in zip(inputs, outputs)
            A = copy(input)
            radix_sort!(A)
            @test A == output
        end
    end
end

ext: jl

max_stars_repo_name: blallen/CodePractice

max_issues_repo_name: blallen/CodePractice

max_forks_repo_name: blallen/CodePractice


text:
# This code calculates the total energies for all atoms Z = 1..92
# in the LDA(SVWN) DFT approximation and compares them with
# the NIST data available at https://www.nist.gov/document/dftdata-targz

const ENIST = [-0.445671;     -2.834836;     -7.335195;    -14.447209;    -24.344198;
              -37.425749;    -54.025016;    -74.473077;    -99.099648;   -128.233481;
             -161.440060;   -199.139406;   -241.315573;   -288.198397;   -339.946219;
             -396.716081;   -458.664179;   -525.946195;   -598.200590;   -675.742283;
             -758.679275;   -847.277216;   -941.678904;  -1042.030238;  -1148.449372;
            -1261.093056;  -1380.091264;  -1505.580197;  -1637.785861;  -1776.573850;
            -1921.846456;  -2073.807332;  -2232.534978;  -2398.111440;  -2570.620700;
            -2750.147940;  -2936.337293;  -3129.453161;  -3329.520604;  -3536.737751;
            -3751.196175;  -3973.013235;  -4202.188857;  -4438.981228;  -4683.301031;
            -4935.368406;  -5195.031215;  -5462.390982;  -5737.309064;  -6019.953353;
            -6310.376268;  -6608.631413;  -6914.773092;  -7228.856107;  -7550.557710;
            -7880.111578;  -8217.575230;  -8563.360285;  -8917.664369;  -9280.311037;
            -9651.484134; -10031.259090; -10419.710775; -10816.653877; -11222.941975;
           -11637.869664; -12061.770549; -12494.718304; -12936.786494; -13388.048594;
           -13848.230375; -14317.493720; -14795.884202; -15283.448822; -15780.236024;
           -16286.295408; -16801.677471; -17326.576377; -17860.790943; -18404.274220;
           -18956.957627; -19518.993145; -20090.414449; -20671.256539; -21261.555215;
           -21861.346869; -22470.319655; -23088.688083; -23716.496952; -24353.832231;
           -25001.291382; -25658.417889];

using AtomEnergyLevels
using Test, Printf, Statistics

function test_all(N = length(ENIST))
  E = zeros(N)
  it = zeros(N)
  for at_number=1:N
    element = keys(atom)[at_number]
    @info "Calculating $element:"
    results = lda(at_number, -35:0.1:5, Α = 1e5)
    E[at_number] = results.energy.total
    it[at_number] = results.iterations
    @info @sprintf("Energy levels for the %s atom are:", element)
    for (quantum_numbers, ψ) in sort(collect(results.orbitals), by = x -> last(x).ϵᵢ)
      nᵣ, l = quantum_numbers
      n = nᵣ + l + 1
      @info @sprintf("\t%i%s\t(%4.1f)\t%14.6f", n, shells[l + 1], ψ.nᵢ, ψ.ϵᵢ)
    end
    @info @sprintf("ΔE = %12.6f", ENIST[at_number]-E[at_number])
  end
  return E, N, it
end

function summary(E, N, it)
  # Median absolute deviation
  MAD(X) = median(abs.(X .- median(X)))

  @info "== RESULTS SUMMARY =="
  for Z=1:N
    element = keys(atom)[Z]
    @info @sprintf("%2s: ΔE = %12.6f", element, ENIST[Z]-E[Z])
  end

  err, atn = findmax(abs.(ENIST[1:N] .- E))
  el = keys(atom)[atn]

  @info "== STATISTICS =="
  @info @sprintf("Median absolute deviation:\t%4.1e", MAD(ENIST[1:N] .- E))
  @info @sprintf("Maximum absolute deviation:\t%4.1e is for %2s", err, el)
  @info @sprintf("Average number of iterations:\t%i", mean(it))
end

summary(test_all(18)...)


text:
export minres_iterable, minres, minres!
using Printf
import LinearAlgebra: BLAS.axpy!, givensAlgorithm
import Base: iterate

mutable struct MINRESIterable{matT, solT, vecT <: DenseVector, smallVecT <: DenseVector, rotT <: Number, realT <: Real}
    A::matT
    skew_hermitian::Bool
    x::solT

    # Krylov basis vectors
    v_prev::vecT
    v_curr::vecT
    v_next::vecT

    # W = R * inv(V) is computed using 3-term recurrence
    w_prev::vecT
    w_curr::vecT
    w_next::vecT

    # Vector of size 4, holding the active column of the Hessenberg matrix
    # rhs is just two active values of the right-hand side.
    H::smallVecT
    rhs::smallVecT

    # Some Givens rotations
    c_prev::rotT
    s_prev::rotT
    c_curr::rotT
    s_curr::rotT

    # Bookkeeping
    mv_products::Int
    maxiter::Int
    tolerance::realT
    resnorm::realT
end

function minres_iterable!(x, A, b;
    initially_zero::Bool = false,
    skew_hermitian::Bool = false,
    tol = sqrt(eps(real(eltype(b)))),
    maxiter = size(A, 2)
)
    T = eltype(x)
    HessenbergT = skew_hermitian ? T : real(T)

    v_prev = similar(x)
    v_curr = similar(x)
    copyto!(v_curr, b)
    v_next = similar(x)
    w_prev = similar(x)
    w_curr = similar(x)
    w_next = similar(x)

    mv_products = 0

    # For nonzero x's, we must do an MV for the initial residual vec
    if !initially_zero
        # Use v_next to store Ax; v_next will soon be overwritten.
        mul!(v_next, A, x)
        axpy!(-one(T), v_next, v_curr)
        mv_products = 1
    end

    resnorm = norm(v_curr)
    reltol = resnorm * tol

    # Last active column of the Hessenberg matrix
    # and last two entries of the right-hand side
    H = zeros(HessenbergT, 4)
    rhs = [resnorm; zero(HessenbergT)]

    # Normalize the first Krylov basis vector
    rmul!(v_curr, inv(resnorm))

    # Givens rotations
    c_prev, s_prev = one(T), zero(T)
    c_curr, s_curr = one(T), zero(T)

    MINRESIterable(
        A, skew_hermitian, x,
        v_prev, v_curr, v_next,
        w_prev, w_curr, w_next,
        H, rhs,
        c_prev, s_prev, c_curr, s_curr,
        mv_products, maxiter, reltol, resnorm
    )
end

converged(m::MINRESIterable) = m.resnorm ≤ m.tolerance

start(::MINRESIterable) = 1

done(m::MINRESIterable, iteration::Int) = iteration > m.maxiter || converged(m)

function iterate(m::MINRESIterable, iteration::Int=start(m))
    if done(m, iteration) return nothing end

    # v_next = A * v_curr - H[2] * v_prev
    mul!(m.v_next, m.A, m.v_curr)

    iteration > 1 && axpy!(-m.H[2], m.v_prev, m.v_next)

    # Orthogonalize w.r.t. v_curr
    proj = dot(m.v_curr, m.v_next)
    m.H[3] = m.skew_hermitian ? proj : real(proj)
    axpy!(-proj, m.v_curr, m.v_next)

    # Normalize
    m.H[4] = norm(m.v_next)
    rmul!(m.v_next, inv(m.H[4]))

    # Rotation on H[1] and H[2]. Note that H[1] = 0 initially
    if iteration > 2
        m.H[1] = m.s_prev * m.H[2]
        m.H[2] = m.c_prev * m.H[2]
    end

    # Rotation on H[2] and H[3]
    if iteration > 1
        tmp = -conj(m.s_curr) * m.H[2] + m.c_curr * m.H[3]
        m.H[2] = m.c_curr * m.H[2] + m.s_curr * m.H[3]
        m.H[3] = tmp
    end

    # Next rotation
    c, s, m.H[3] = givensAlgorithm(m.H[3], m.H[4])

    # Apply as well to the right-hand side
    m.rhs[2] = -conj(s) * m.rhs[1]
    m.rhs[1] = c * m.rhs[1]

    # Update W = V * inv(R). Two axpy's can maybe be one MV.
    copyto!(m.w_next, m.v_curr)
    iteration > 1 && axpy!(-m.H[2], m.w_curr, m.w_next)
    iteration > 2 && axpy!(-m.H[1], m.w_prev, m.w_next)
    rmul!(m.w_next, inv(m.H[3]))

    # Update solution x
    axpy!(m.rhs[1], m.w_next, m.x)

    # Move on: next -> curr, curr -> prev
    m.v_prev, m.v_curr, m.v_next = m.v_curr, m.v_next, m.v_prev
    m.w_prev, m.w_curr, m.w_next = m.w_curr, m.w_next, m.w_prev
    m.c_prev, m.s_prev, m.c_curr, m.s_curr = m.c_curr, m.s_curr, c, s
    m.rhs[1] = m.rhs[2]

    # Due to symmetry of the tri-diagonal matrix
    m.H[2] = m.skew_hermitian ? -m.H[4] : m.H[4]

    # The approximate residual is cheaply available
    m.resnorm = abs(m.rhs[2])

    m.resnorm, iteration + 1
end

"""
    minres!(x, A, b; kwargs...) -> x, [history]

Solve Ax = b for (skew-)Hermitian matrices A using MINRES.

# Arguments

- `x`: initial guess, will be updated in-place;
- `A`: linear operator;
- `b`: right-hand side.

## Keywords

- `initially_zero::Bool = false`: if `true` assumes that `iszero(x)` so that one
  matrix-vector product can be saved when computing the initial
  residual vector;
- `skew_hermitian::Bool = false`: if `true` assumes that `A` is skew-symmetric or skew-Hermitian;
- `tol`: tolerance for stopping condition `|r_k| / |r_0| ≤ tol`. Note that the residual is computed only approximately;
- `maxiter::Int = size(A, 2)`: maximum number of iterations;
- `Pl`: left preconditioner;
- `Pr`: right preconditioner;
- `log::Bool = false`: keep track of the residual norm in each iteration;
- `verbose::Bool = false`: print convergence information during the iterations.

# Return values

**if `log` is `false`**

- `x`: approximate solution.

**if `log` is `true`**

- `x`: approximate solution;
- `history`: convergence history.
"""
function minres!(x, A, b;
    skew_hermitian::Bool = false,
    verbose::Bool = false,
    log::Bool = false,
    tol = sqrt(eps(real(eltype(b)))),
    maxiter::Int = size(A, 2),
    initially_zero::Bool = false
)
    history = ConvergenceHistory(partial = !log)
    history[:tol] = tol
    log && reserve!(history, :resnorm, maxiter)

    iterable = minres_iterable!(x, A, b;
        skew_hermitian = skew_hermitian,
        tol = tol,
        maxiter = maxiter,
        initially_zero = initially_zero
    )

    if log
        history.mvps = iterable.mv_products
    end

    for (iteration, resnorm) = enumerate(iterable)
        if log
            nextiter!(history, mvps = 1)
            push!(history, :resnorm, resnorm)
        end
        verbose && @printf("%3d\t%1.2e\n", iteration, resnorm)
    end

    verbose && println()
    log && setconv(history, converged(iterable))
    log && shrink!(history)

    log ? (iterable.x, history) : iterable.x
end

"""
    minres(A, b; kwargs...) -> x, [history]

Same as [`minres!`](@ref), but allocates a solution vector `x` initialized with zeros.
"""
minres(A, b; kwargs...) = minres!(zerox(A, b), A, b; initially_zero = true, kwargs...)


text:
using DataFrames

function molar_mass(MI)
	sql = "	SELECT 	molecule_id, local_id, molar_mass
			FROM	isotopologues
			WHERE 	(molecule_id, local_id)
			IN 		(VALUES " * join(["(" * join("?"^length(t), ',') * ")" for t in MI], ',') * ")"
	result = query_local_db(sql, [i for t in MI for i in t])

	out = Dict{Tuple{Int, Int}, Float64}()
	for row in result
		out[(row.molecule_id, row.local_id)] = row.molar_mass
	end

	return out
end

## Default data for molecules and isotopologues for new databases
const molecules = DataFrame(id=Int[], formula=String[], name=String[])
push!(molecules,
    (1, "H2O", "Water"),
    (2, "CO2", "Carbon Dioxide"),
    (3,	"O3", "Ozone"),
    (4,	"N2O", "Nitrogen oxide"),
    (5,	"CO", "Carbon Monoxide"),
    (6,	"CH4", "Methane"),
    (7, "O2", "Oxygen"),
    (8,	"NO", "Nitric Oxide"),
    (9,	"SO2", "Sulfur Dioxide"),
    (10, "NO2", "Nitrogen Dioxide"),
    (11, "NH3", "Ammonia"),
    (12, "HNO3", "Nitric Acid"),
    (13, "OH", "Hydroxyl"),
    (14, "HF", "Hydrogen Fluoride"),
    (15, "HCl", "Hydrogen Chloride"),
    (16, "HBr", "Hydrogen Bromide"),
    (17, "HI", "Hydrogen Iodide"),
    (18, "ClO", "Chlorine Monoxide"),
    (19, "OCS", "Carbonyl Sulfide"),
    (20, "H2CO", "Formaldehyde"),
    (21, "HOCl", "Hypochlorous Acid"),
    (22, "N2", "Nitrogen"),
    (23, "HCN", "Hydrogen Cyanide"),
    (24, "CH3Cl", "Methyl Chloride"),
    (25, "H2O2", "Hydrogen Peroxide"),
    (26, "C2H2", "Acetylene"),
    (27, "C2H6", "Ethane"),
    (28, "PH3", "Phosphine"),
    (29, "COF2", "Carbonyl Fluoride"),
    (30, "SF6", "Sulfur Hexafluoride"),
    (31, "H2S", "Hydrogen Sulfide"),
    (32, "HCOOH", "Formic Acid"),
    (33, "HO2", "Hydroperoxyl"),
    (34, "O", "Oxygen Atom"),
    (35, "ClONO2", "Chlorine Nitrate"),
    (36, "NO+", "Nitric Oxide Cation"),
    (37, "HOBr", "Hypobromous Acid"),
    (38, "C2H4", "Ethylene"),
    (39, "CH3OH", "Methanol"),
    (40, "CH3Br", "Methyl Bromide"),
    (41, "CH3CN", "Acetonitrile"),
    (42, "CF4", "PFC-14"),
    (43, "C4H2", "Diacetylene"),
    (44, "HC3N", "Cyanoacetylene"),
    (45, "H2", "Hydrogen"),
    (46, "CS", "Carbon Monosulfide"),
    (47, "SO3", "Sulfur trioxide"),
    (48, "C2N2", "Cyanogen"),
    (49, "COCl2", "Phosgene"),
	(53, "CS2", "Carbon disulfide"),
	(55, "NF3", "Nitrogen trifluoride")
)

const isotopologues = DataFrame(
    molecule_id=Int[],
    global_id=Int[],
    local_id=Int[],
    formula=String[],
    afgl_code=Int[],
    abundance=Float64[],
    molar_mass=Float64[],
    q_t0=Float64[],
    g_i=Int[]
    )

push!(isotopologues,
    # 1: H2O
	(1, 1, 1, "H216O", 161, 0.997317, 18.010565, 174.58, 1),
	(1, 2, 2, "H218O", 181, 0.002000, 20.014811, 176.05, 1),
	(1, 3, 3, "H217O", 171, 3.718840e-4, 19.01478, 1052.14, 6),
	(1, 4, 4, "HD16O", 162, 3.106930e-4, 19.01674, 864.74, 6),
	(1, 5, 5, "HD18O", 182, 6.230030e-7, 21.020985, 875.57, 6),
	(1, 6, 6, "HD17O", 172, 1.158530e-7, 20.020956, 5226.79, 36),
	(1, 129, 7, "D216O", 262, 2.419700e-8, 20.022915, 1027.80, 1),

    # 2: CO2
	(2, 7, 1, "12C16O2", 626, 0.984204, 43.98983, 286.09, 1),
	(2, 8, 2, "13C16O2", 636, 0.011057, 44.993185, 576.64, 2),
	(2, 9, 3, "16O12C18O", 628, 0.003947, 45.994076, 607.81, 1),
	(2, 10, 4, "16O12C17O", 627, 7.339890e-4, 44.994045, 3542.61, 6),
	(2, 11, 5, "16O13C18O", 638, 4.434460e-5, 46.997431, 1225.46, 2),
	(2, 12, 6, "16O13C17O", 637, 8.246230e-6, 45.9974, 7141.32, 12),
	(2, 13, 7, "12C18O2", 828, 3.957340e-6, 47.998322, 323.42, 1),
	(2, 14, 8, "17O12C18O", 827, 1.471800e-6, 46.998291, 3766.58, 6),
	(2, 121, 9, "12C17O2", 727, 1.368470e-7, 45.998262, 10971.57, 1),
	(2, 15, 0, "13C18O2", 838, 4.446000e-8, 49.001675, 652.24, 2),
	(2, 120, 11, "18O13C17O", 837, 1.653540e-8, 48.001646, 7595.04, 12),
	(2, 122, 12, "13C17O2", 737, 1.537500e-9, 47.0016182378, 22120.47, 2),

    # 3: O3
	(3, 16, 1, "16O3", 666, 0.992901, 47.984745, 3483.71, 1),
	(3, 17, 2, "16O16O18O", 668, 0.003982, 49.988991, 7465.68, 1),
	(3, 18, 3, "16O18O16O", 686, 0.001991, 49.988991, 3647.08, 1),
	(3, 19, 4, "16O16O17O", 667, 7.404750e-4, 48.98896, 43330.85, 6),
	(3, 20, 5, "16O17O16O", 676, 3.702370e-4, 48.98896, 21404.96, 6),

    # 4: N2O
	(4, 21, 1, "14N216O", 446, 0.990333, 44.001062, 4984.90, 9),
	(4, 22, 2, "14N15N16O", 456, 0.003641, 44.998096, 3362.01, 6),
	(4, 23, 3, "15N14N16O", 546, 0.003641, 44.998096, 3458.58, 6),
	(4, 24, 4, "14N218O", 448, 0.001986, 46.005308, 5314.74, 9),
	(4, 25, 5, "14N217O", 447, 3.692800e-4, 45.005278, 30971.79, 54),

    # 5: CO
	(5, 26, 1, "12C16O", 26, 0.986544, 27.994915, 107.42, 1),
	(5, 27, 2, "13C16O", 36, 0.011084, 28.99827, 224.69, 2),
	(5, 28, 3, "12C18O", 28, 0.001978, 29.999161, 112.77, 1),
	(5, 29, 4, "12C17O", 27, 3.678670e-4, 28.99913, 661.17, 6),
	(5, 30, 5, "13C18O", 38, 2.222500e-5, 31.002516, 236.44, 2),
	(5, 31, 6, "13C17O", 37, 4.132920e-6, 30.002485, 1384.66, 12),

    # 6: CH4
	(6, 32, 1, "12CH4", 211, 0.988274, 16.0313, 590.48, 1),
	(6, 33, 2, "13CH4", 311, 0.011103, 17.034655, 1180.82, 2),
	(6, 34, 3, "12CH3D", 212, 6.157510e-4, 17.037475, 4794.73, 3),
	(6, 35, 4, "13CH3D", 312, 6.917850e-6, 18.04083, 9599.16, 6),

    # 7: O2
	(7, 36, 1, "16O2", 66, 0.995262, 31.98983, 215.73, 1),
	(7, 37, 2, "16O18O", 68, 0.003991, 33.994076, 455.23, 1),
	(7, 38, 3, "16O17O", 67, 7.422350e-4, 32.994045, 2658.12, 6),

    # 8: NO
	(8, 39, 1, "14N16O", 46, 0.993974, 29.997989, 1142.13, 3),
	(8, 40, 2, "15N16O", 56, 0.003654, 30.995023, 789.26, 2),
	(8, 41, 3, "14N18O", 48, 0.001993, 32.002234, 1204.44, 3),

    # 9: SO2
	(9, 42, 1, "32S16O2", 626, 0.945678, 63.961901, 6340.30, 1),
	(9, 43, 2, "34S16O2", 646, 0.041950, 65.957695, 6368.98, 1),

    # 10: NO2
	(10, 44, 1, "14N16O2", 646, 0.991616, 45.992904, 13577.48, 3),
	(10, 130, 2, "15N16O2", 656, 0.003646, 46.989938, 9324.70, 2),

    # 11: NH3
	(11, 45, 1, "14NH3", 4111, 0.995872, 17.026549, 1725.22, 3),
	(11, 46, 2, "15NH3", 5111, 0.003661, 18.023583, 1153.30, 2),

    # 12: HNO3
	(12, 47, 1, "H14N16O3", 146, 0.989110, 62.995644, 2.14e5, 6),
	(12, 117, 2, "H15N16O3", 156, 0.003636, 63.99268, 1.43e5, 4),

    # 13: OH
	(13, 48, 1, "16OH", 61, 0.997473, 17.00274, 80.35, 2),
	(13, 49, 2, "18OH", 81, 0.002000, 19.006986, 80.88, 2),
	(13, 50, 3, "16OD", 62, 1.553710e-4, 18.008915, 209.32, 3),

    # 14: HF
	(14, 51, 1, "H19F", 19, 0.999844, 20.006229, 41.47, 4),
	(14, 110, 2, "D19F", 29, 1.557410e-4, 21.012404, 115.91, 6),

    # 15: HCl
	(15, 52, 1, "H35Cl", 15, 0.757587, 35.976678, 160.65, 8),
	(15, 53, 2, "H37Cl", 17, 0.242257, 37.973729, 160.89, 8),
	(15, 107, 3, "D35Cl", 25, 1.180050e-4, 36.982853, 462.78, 12),
	(15, 108, 4, "D37Cl", 27, 3.773500e-5, 38.979904, 464.13, 12),

    # 16: HBr
	(16, 54, 1, "H79Br", 19, 0.506781, 79.92616, 200.17, 8),
	(16, 55, 2, "H81Br", 11, 0.493063, 81.924115, 200.23, 8),
	(16, 111, 3, "D79Br", 29, 7.893840e-5, 80.932336, 586.40, 12),
	(16, 112, 4, "D81Br", 21, 7.680160e-5, 82.930289, 586.76, 12),

    # 17: HI
	(17, 56, 1, "H127I", 17, 0.999844, 127.912297, 388.99, 12),
	(17, 113, 2, "D127I", 27, 1.557410e-4, 128.918472, 1147.06, 18),

    # 18: ClO
	(18, 57, 1, "35Cl16O", 56, 0.755908, 50.963768, 3274.61, 4),
	(18, 58, 2, "37Cl16O", 76, 0.241720, 52.960819, 3332.29, 4),

    # 19: OCS
	(19, 59, 1, "16O12C32S", 622, 0.937395, 59.966986, 1221.01, 1),
	(19, 60, 2, "16O12C34S", 624, 0.041583, 61.96278, 1253.48, 1),
	(19, 61, 3, "16O13C32S", 632, 0.010531, 60.970341, 2484.15, 2),
	(19, 62, 4, "16O12C33S", 623, 0.007399, 60.966371, 4950.11, 4),
	(19, 63, 5, "18O12C32S", 822, 0.001880, 61.971231, 1313.78, 1),
	(19, 135, 6, "16O13C34S", 634, 4.675080e-4, 62.966136, 2546.53, 2),

    # 20: H2CO
	(20, 64, 1, "H212C16O", 126, 0.986237, 30.010565, 2844.53, 1),
	(20, 65, 2, "H213C16O", 136, 0.011080, 31.01392, 5837.69, 2),
	(20, 66, 3, "H212C18O", 128, 0.001978, 32.014811, 2986.44, 1),

    # 21: HOCl
	(21, 67, 1, "H16O35Cl", 165, 0.755790, 51.971593, 19274.79, 8),
	(21, 68, 2, "H16O37Cl", 167, 0.241683, 53.968644, 19616.20, 8),

    # 22: N2
	(22, 69, 1, "14N2", 44, 0.992687, 28.006148, 467.10, 1),
	(22, 118, 2, "14N15N", 45, 0.007478, 29.003182, 644.10, 6),

    # 23: HCN
	(23, 70, 1, "H12C14N", 124, 0.985114, 27.010899, 892.20, 6),
	(23, 71, 2, "H13C14N", 134, 0.011068, 28.014254, 1830.97, 12),
	(23, 72, 3, "H12C15N", 125, 0.003622, 28.007933, 615.28, 4),

    # 24: CH3Cl
	(24, 73, 1, "12CH335Cl", 215, 0.748937, 49.992328, 57916.12, 4),
	(24, 74, 2, "12CH337Cl", 217, 0.239491, 51.989379, 58833.90, 4),

    # 25: H2O2
	(25, 75, 1, "H216O2", 1661, 0.994952, 34.00548, 9847.99, 1),

    # 26: C2H2
	(26, 76, 1, "12C2H2", 1221, 0.977599, 26.01565, 412.45, 1),
	(26, 77, 2, "H12C13CH", 1231, 0.021966, 27.019005, 1656.18, 8),
	(26, 105, 3, "H12C12CD", 1222, 3.045500e-4, 27.021825, 1581.84, 6),

    # 27: C2H6
	(27, 78, 1, "12C2H6", 1221, 0.976990, 30.04695, 70882.52, 1),
	(27, 106, 2, "12CH313CH3", 1231, 0.021953, 31.050305, 36191.80, 2),

    # 28: PH3
	(28, 79, 1, "31PH3", 1111, 0.999533, 33.997238, 3249.44, 2),

    # 29: COF2
	(29, 80, 1, "12C16O19F2", 269, 0.986544, 65.991722, 70028.43, 1),
	(29, 119, 2, "13C16O19F2", 369, 0.011083, 66.995083, 1.40e5, 2),

    # 30: SF6
	(30, 126, 1, "32S19F6", 29, 0.950180, 145.962492, 1.62e6, 1),

    # 31: H2S
	(31, 81, 1, "H232S", 121, 0.949884, 33.987721, 505.79, 1),
	(31, 82, 2, "H234S", 141, 0.042137, 35.983515, 504.35, 1),
	(31, 83, 3, "H233S", 131, 0.007498, 34.987105, 2014.94, 4),

    # 32: HCOOH
	(32, 84, 1, "H12C16O16OH", 126, 0.983898, 46.00548, 39132.76, 4),

    # 33: HO2
	(33, 85, 1, "H16O2", 166, 0.995107, 32.997655, 4300.39, 2),

    # 34: O
	(34, 86, 1, "16O", 6, 0.997628, 15.994915, 6.72, 1),

    # 35: ClONO2
	(35, 127, 1, "35Cl16O14N16O2", 5646, 0.749570, 96.956672, 4.79e6, 12),
	(35, 128, 2, "37Cl16O14N16O2", 7646, 0.239694, 98.953723, 4.91e6, 12),

    # 36: NO+
    (36, 87, 1, "14N16O+", 46, 0.993974, 29.997989, 311.69, 3),

    # 37: HOBr
	(37, 88, 1, "H16O79Br", 169, 0.505579, 95.921076, 28339.38, 8),
	(37, 89, 2, "H16O81Br", 161, 0.491894, 97.919027, 28237.98, 8),

    # 38: C2H4
	(38, 90, 1, "12C2H4", 221, 0.977294, 28.0313, 11041.54, 1),
	(38, 91, 2, "12CH213CH2", 231, 0.021959, 29.034655, 45196.89, 2),

    # 39: CH3OH
	(39, 92, 1, "12CH316OH", 2161, 0.985930, 32.026215, 70569.92, 2),

    # 40: CH3Br
	(40, 93, 1, "12CH379Br", 219, 0.500995, 93.941811, 83051.98, 4),
	(40, 94, 2, "12CH381Br", 211, 0.487433, 95.939764, 83395.21, 4),

    # 41: CH3CN
	(41, 95, 1, "12CH312C14N", 2124, 0.973866, 41.026549, 88672.19, 3),

    # 42: CF4
	(42, 96, 1, "12C19F4", 29, 0.988890, 87.993616, 1.21e5, 1),

    # 43: C4H2
	(43, 116, 1, "12C4H2", 2211, 0.955998, 50.01565, 9818.97, 1),

    # 44: HC3N
	(44, 109, 1, "H12C314N", 1224, 0.963346, 51.010899, 24786.84, 6),

    # 45: H2
	(45, 103, 1, "H2", 11, 0.999688, 2.01565, 7.67, 1),
	(45, 115, 2, "HD", 12, 3.114320e-4, 3.021825, 29.87, 6),

    # 46: CS
	(46, 97, 1, "12C32S", 22, 0.939624, 43.971036, 253.62, 1),
	(46, 98, 2, "12C34S", 24, 0.041682, 45.966787, 257.77, 1),
	(46, 99, 3, "13C32S", 32, 0.010556, 44.974368, 537.50, 2),
	(46, 100, 4, "12C33S", 23, 0.007417, 44.970399, 1022.97, 4),

    # 47: SO3
	(47, 114, 1, "32S16O3", 26, 0.943400, 79.95682, 7783.30, 1),

    # 48: C2N2
	(48, 123, 1, "12C214N2", 4224, 0.970752, 52.006148, 15582.44, 1),

    # 49: COCl2
	(49, 124, 1, "12C16O35Cl2", 2655, 0.566392, 97.9326199796, 1.48e6, 1),
	(49, 125, 2, "12C16O35Cl37Cl", 2657, 0.362235, 99.9296698896, 3.04e6, 16),

    # 53: CS2
	(53, 131, 1, "12C32S2", 222, 0.892811, 75.94414, 1352.60, 1),
	(53, 132, 2, "32S12C34S", 224, 0.079260, 77.93994, 2798.00, 1),
	(53, 133, 3, "32S12C33S", 223, 0.014094, 76.943256, 1107.00, 4),
	(53, 134, 4, "13C32S2", 232, 0.010310, 76.947495, 2739.70, 2),

	# 55: NF3
	(55, 136, 1, "14N19F3", 4999, 0.996337, 70.998284, 3.46e5, 3)
)


text:
import Random
import Zygote
import RestrictedBoltzmannMachines as RBMs

using Test: @test, @testset, @inferred
using Statistics: mean, var, cov
using Random: bitrand
using LogExpFunctions: logistic
using QuadGK: quadgk
using RestrictedBoltzmannMachines: flatten, batch_size, batchmean, batchvar, batchcov
using RestrictedBoltzmannMachines: Binary, Spin, Potts, Gaussian, ReLU, dReLU, xReLU, pReLU
using RestrictedBoltzmannMachines: mean_from_inputs, var_from_inputs, meanvar_from_inputs,
    std_from_inputs, mean_abs_from_inputs, sample_from_inputs, mode_from_inputs
using RestrictedBoltzmannMachines: energy, free_energy, free_energies, energies

Random.seed!(2)

_layers = (
    Binary,
    Spin,
    Potts,
    Gaussian,
    ReLU,
    dReLU,
    pReLU,
    xReLU
)

random_layer(::Type{T}, N::Int...) where {T <: Union{Binary,Spin,Potts}} = T(randn(N...))
random_layer(::Type{T}, N::Int...) where {T <: Union{Gaussian,ReLU}} = T(randn(N...), randn(N...))
random_layer(::Type{dReLU}, N::Int...) = dReLU(randn(N...), randn(N...), randn(N...), randn(N...))
random_layer(::Type{xReLU}, N::Int...) = xReLU(randn(N...), randn(N...), randn(N...), randn(N...))
random_layer(::Type{pReLU}, N::Int...) = pReLU(randn(N...), randn(N...), randn(N...), 2rand(N...) .- 1)

@testset "testing $Layer" for Layer in _layers
    N = (3,2)
    layer = random_layer(Layer, N...)

    @test (@inferred size(layer)) == N
    for d in 1:length(N)
        @test (@inferred size(layer, d)) == N[d]
    end
    @test (@inferred length(layer)) == prod(N)
    @test (@inferred ndims(layer)) == length(N)
    @test size(@inferred repeat(layer,2,3,4)) == (((2, 3) .* N)..., 4)
    @test (@inferred batch_size(layer, rand(N...))) == ()
    @test (@inferred energy(layer, rand(N...))) isa Number
    @test (@inferred free_energy(layer)) isa Number
    @test (@inferred free_energy(layer, rand(N...))) isa Number
    @test size(@inferred mean_from_inputs(layer)) == size(layer)
    @test size(@inferred var_from_inputs(layer)) == size(layer)
    @test size(@inferred sample_from_inputs(layer)) == size(layer)
    @test free_energies(layer, 0) ≈ free_energies(layer)
    @test std_from_inputs(layer) ≈ sqrt.(var_from_inputs(layer))

    if layer isa RBMs.Potts
        @test size(@inferred free_energies(layer)) == (1, size(layer)[2:end]...)
    else
        @test size(@inferred free_energies(layer)) == size(layer)
    end

    @test size(@inferred sample_from_inputs(layer, 0)) == size(layer)

    for B in ((), (2,), (1,2))
        x = rand(N..., B...)
        @test (@inferred batch_size(layer, x)) == (B...,)
        @test (@inferred RBMs.batchdims(layer, x)) == (length(N) + 1):ndims(x)
        @test size(@inferred energy(layer, x)) == (B...,)
        @test size(@inferred free_energy(layer, x)) == (B...,)
        @test size(@inferred energies(layer, x)) == size(x)
        @test size(@inferred sample_from_inputs(layer, x)) == size(x)
        @test size(@inferred mean_from_inputs(layer, x)) == size(x)
        @test size(@inferred var_from_inputs(layer, x)) == size(x)
        @test @inferred(std_from_inputs(layer, x)) ≈ sqrt.(var_from_inputs(layer, x))
        @test all(energy(layer, mode_from_inputs(layer)) .≤ energy(layer, x))

        μ, ν = meanvar_from_inputs(layer, x)
        @test μ ≈ mean_from_inputs(layer, x)
        @test ν ≈ var_from_inputs(layer, x)

        if layer isa RBMs.Potts
            @test size(@inferred free_energies(layer, x)) == (1, size(x)[2:end]...)
        else
            @test size(@inferred free_energies(layer, x)) == size(x)
        end

        if B == ()
            @test @inferred(energy(layer, x)) ≈ sum(energies(layer, x))
            @test @inferred(free_energy(layer, x)) ≈ sum(free_energies(layer, x))
            @test @inferred(batchmean(layer, x)) ≈ x
            @test @inferred(batchvar(layer, x)) == zeros(N)
            @test @inferred(batchcov(layer, x)) == zeros(N..., N...)
        else
            @test @inferred(energy(layer, x)) ≈ reshape(sum(energies(layer, x); dims=1:ndims(layer)), B)
            @test @inferred(free_energy(layer, x)) ≈ reshape(sum(free_energies(layer, x); dims=1:ndims(layer)), B)
            @test @inferred(batchmean(layer, x)) ≈ reshape(mean(x; dims=(ndims(layer) + 1):ndims(x)), N)
            @test @inferred(batchvar(layer, x)) ≈ reshape(var(x; dims=(ndims(layer) + 1):ndims(x), corrected=false), N)
            @test @inferred(batchcov(layer, x)) ≈ reshape(cov(flatten(layer, x); dims=2, corrected=false), N..., N...)
        end

        μ = @inferred mean_from_inputs(layer, x)
        @test only(Zygote.gradient(j -> sum(free_energies(layer, j)), x)) ≈ -μ
    end

    ∂F = @inferred RBMs.∂free_energy(layer)
    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    for (ω, ∂ω) in pairs(∂F)
        @test ∂ω ≈ getproperty(only(gs), ω)
    end

    samples = @inferred sample_from_inputs(layer, zeros(size(layer)..., 10^6))
    @test @inferred(mean_from_inputs(layer)) ≈ reshape(mean(samples; dims=3), size(layer)) rtol=0.1 atol=0.01
    @test @inferred(var_from_inputs(layer)) ≈ reshape(var(samples; dims=ndims(samples)), size(layer)) rtol=0.1
    @test @inferred(mean_abs_from_inputs(layer)) ≈ reshape(mean(abs.(samples); dims=ndims(samples)), size(layer)) rtol=0.1

    ∂F = @inferred RBMs.∂free_energy(layer)
    ∂E = @inferred RBMs.∂energy(layer, samples)
    @test length(∂F) == length(∂E)
    @test propertynames(∂F) == propertynames(∂E)
    for (∂f, ∂e) in zip(∂F, ∂E)
        @test ∂f ≈ ∂e rtol=0.1
    end

    gs = Zygote.gradient(layer) do layer
        sum(energies(layer, samples)) / size(samples)[end]
    end
    for (ω, ∂ω) in pairs(∂E)
        @test ∂ω ≈ getproperty(only(gs), ω)
    end
end

@testset "discrete layers ($Layer)" for Layer in (Binary, Spin, RBMs.Potts)
    N = (3, 4, 5)
    B = 13
    layer = Layer(randn(N...))
    x = bitrand(N..., B)
    @test energies(layer, x) ≈ -layer.θ .* x
    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    @test RBMs.∂free_energy(layer).θ ≈ only(gs).θ ≈ -mean_from_inputs(layer)
end

@testset "Binary" begin
    @testset "binary_rand" begin
        for θ in -5:5, u in 0.0:0.1:1.0
            @test (@. u * (1 + exp(-θ)) < 1) == @inferred RBMs.binary_rand(θ, u)
        end
        θ = randn(1000)
        u = rand(1000)
        @test RBMs.binary_rand.(θ, u) == (@. u * (1 + exp(-θ)) < 1)
    end

    layer = Binary(randn(7, 4, 5))
    var_from_inputs(layer) ≈ @. logistic(layer.θ) * logistic(-layer.θ)
    @test free_energies(layer) ≈ -log.(sum(exp.(layer.θ .* h) for h in 0:1))
    @test sort(unique(sample_from_inputs(layer))) == [0, 1]

    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    ∂ = RBMs.∂free_energy(layer)
    @test ∂.θ ≈ only(gs).θ ≈ -mean_from_inputs(layer)
    @test RBMs.grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
    @test RBMs.grad2var(layer, ∂) ≈ var_from_inputs(layer)
end

@testset "Spin" begin
    layer = Spin(randn(7, 4, 5))
    @test free_energies(layer) ≈ -log.(sum(exp.(layer.θ .* h) for h in (-1, 1)))
    @test sort(unique(sample_from_inputs(layer))) == [-1, 1]

    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    ∂ = RBMs.∂free_energy(layer)
    @test ∂.θ ≈ only(gs).θ ≈ -mean_from_inputs(layer)
    @test RBMs.grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
    @test RBMs.grad2var(layer, ∂) ≈ var_from_inputs(layer)
end

@testset "Potts" begin
    q = 3
    N = (4, 5)
    layer = RBMs.Potts(randn(q, N...))
    @test free_energies(layer) ≈ -log.(sum(exp.(layer.θ[h:h,:,:,:]) for h in 1:q))
    @test all(sum(mean_from_inputs(layer); dims=1) .≈ 1)
    # samples are proper one-hot
    @test sort(unique(sample_from_inputs(layer))) == [0, 1]
    @test all(sum(sample_from_inputs(layer); dims=1) .== 1)

    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    ∂ = RBMs.∂free_energy(layer)
    @test ∂.θ ≈ only(gs).θ ≈ -mean_from_inputs(layer)
    @test RBMs.grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
    @test RBMs.grad2var(layer, ∂) ≈ var_from_inputs(layer)
end

@testset "Gaussian" begin
    N = (3, 4, 5)
    B = 7

    # bound γ away from zero to avoid numerical issues with quadgk, but allow negative
    layer = Gaussian(randn(N...), (rand(N...) .+ 1) .* (rand(Bool, N...) .- 0.5))

    x = randn(size(layer)..., B)
    @test energies(layer, x) ≈ @. abs(layer.γ) * x^2 / 2 - layer.θ * x

    function quad_free(θ::Real, γ::Real)
        Z, ϵ = quadgk(h -> exp(-RBMs.gauss_energy(θ, γ, h)), -Inf,  Inf)
        return -log(Z)
    end

    @test free_energies(layer) ≈ quad_free.(layer.θ, layer.γ) rtol=1e-6

    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    μ = mean_from_inputs(layer)
    ν = var_from_inputs(layer)
    μ2 = @. ν + μ^2
    ∂ = RBMs.∂free_energy(layer)
    @test ∂.θ ≈ only(gs).θ ≈ -μ
    @test ∂.γ ≈ only(gs).γ ≈ sign.(layer.γ) .* μ2/2
    @test RBMs.grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
    @test RBMs.grad2var(layer, ∂) ≈ var_from_inputs(layer)
end

@testset "ReLU" begin
    N = (3, 4, 5)
    B = 7

    # bound γ away from zero to avoid numerical issues with quadgk, but allow negative
    layer = ReLU(randn(N...), (rand(N...) .+ 1) .* (rand(Bool, N...) .- 0.5))

    x = abs.(randn(size(layer)..., B))
    @test energies(layer, x) ≈ energies(Gaussian(layer.θ, layer.γ), x)

    function quad_free(θ::Real, γ::Real)
        Z, ϵ = quadgk(h -> exp(-RBMs.relu_energy(θ, γ, h)), 0,  Inf)
        return -log(Z)
    end
    @test free_energies(layer) ≈ @. quad_free(layer.θ, layer.γ)

    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    μ = mean_from_inputs(layer)
    ν = var_from_inputs(layer)
    μ2 = @. ν + μ^2
    ∂ = RBMs.∂free_energy(layer)
    @test ∂.θ ≈ only(gs).θ ≈ -μ
    @test ∂.γ ≈ only(gs).γ ≈ sign.(layer.γ) .* μ2/2
    @test RBMs.grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
    @test RBMs.grad2var(layer, ∂) ≈ var_from_inputs(layer)
end

@testset "pReLU / xReLU / dReLU convert" begin
    N = (10, 7)
    B = 13
    x = randn(N..., B)

    drelu = dReLU(randn(N...), randn(N...), randn(N...), randn(N...))
    prelu = @inferred pReLU(drelu)
    xrelu = @inferred xReLU(drelu)
    @test drelu.θp ≈ dReLU(prelu).θp ≈ dReLU(xrelu).θp
    @test drelu.θn ≈ dReLU(prelu).θn ≈ dReLU(xrelu).θn
    @test abs.(drelu.γp) ≈ abs.(dReLU(prelu).γp)
    @test abs.(drelu.γn) ≈ abs.(dReLU(prelu).γn)
    @test abs.(drelu.γp) ≈ abs.(dReLU(xrelu).γp)
    @test abs.(drelu.γn) ≈ abs.(dReLU(xrelu).γn)
    @test energies(drelu, x) ≈ energies(prelu, x) ≈ energies(xrelu, x)
    @test free_energies(drelu) ≈ free_energies(prelu) ≈ free_energies(xrelu)
    @test mode_from_inputs(drelu) ≈ mode_from_inputs(prelu) ≈ mode_from_inputs(xrelu)
    @test mean_from_inputs(drelu) ≈ mean_from_inputs(prelu) ≈ mean_from_inputs(xrelu)
    @test mean_abs_from_inputs(drelu) ≈ mean_abs_from_inputs(prelu) ≈ mean_abs_from_inputs(xrelu)
    @test var_from_inputs(drelu) ≈ var_from_inputs(prelu) ≈ var_from_inputs(xrelu)

    prelu = pReLU(randn(N...), randn(N...), randn(N...), 2rand(N...) .- 1)
    drelu = @inferred dReLU(prelu)
    xrelu = @inferred xReLU(prelu)
    @test prelu.θ ≈ pReLU(drelu).θ ≈ pReLU(xrelu).θ
    @test abs.(prelu.γ) ≈ abs.(pReLU(drelu).γ) ≈ abs.(pReLU(xrelu).γ)
    @test prelu.Δ ≈ pReLU(drelu).Δ ≈ pReLU(xrelu).Δ
    @test prelu.η ≈ pReLU(drelu).η ≈ pReLU(xrelu).η
    @test energies(drelu, x) ≈ energies(prelu, x) ≈ energies(xrelu, x)
    @test free_energies(drelu) ≈ free_energies(prelu) ≈ free_energies(xrelu)
    @test mode_from_inputs(drelu) ≈ mode_from_inputs(prelu) ≈ mode_from_inputs(xrelu)
    @test mean_from_inputs(drelu) ≈ mean_from_inputs(prelu) ≈ mean_from_inputs(xrelu)
    @test mean_abs_from_inputs(drelu) ≈ mean_abs_from_inputs(prelu) ≈ mean_abs_from_inputs(xrelu)
    @test var_from_inputs(drelu) ≈ var_from_inputs(prelu) ≈ var_from_inputs(xrelu)

    xrelu = xReLU(randn(N...), randn(N...), randn(N...), randn(N...))
    drelu = @inferred dReLU(xrelu)
    prelu = @inferred pReLU(xrelu)
    @test xrelu.θ ≈ xReLU(drelu).θ ≈ xReLU(prelu).θ
    @test xrelu.Δ ≈ xReLU(drelu).Δ ≈ xReLU(prelu).Δ
    @test xrelu.ξ ≈ xReLU(drelu).ξ ≈ xReLU(prelu).ξ
    @test abs.(xrelu.γ) ≈ abs.(xReLU(prelu).γ)
    @test energies(drelu, x) ≈ energies(prelu, x) ≈ energies(xrelu, x)
    @test free_energies(drelu) ≈ free_energies(prelu) ≈ free_energies(xrelu)
    @test mode_from_inputs(drelu) ≈ mode_from_inputs(prelu) ≈ mode_from_inputs(xrelu)
    @test mean_from_inputs(drelu) ≈ mean_from_inputs(prelu) ≈ mean_from_inputs(xrelu)
    @test mean_abs_from_inputs(drelu) ≈ mean_abs_from_inputs(prelu) ≈ mean_abs_from_inputs(xrelu)
    @test var_from_inputs(drelu) ≈ var_from_inputs(prelu) ≈ var_from_inputs(xrelu)

    gauss = Gaussian(randn(N...), randn(N...))
    drelu = @inferred dReLU(gauss)
    prelu = @inferred pReLU(gauss)
    xrelu = @inferred xReLU(gauss)
    @test (
        energies(gauss, x) ≈ energies(drelu, x) ≈
        energies(prelu, x) ≈ energies(xrelu, x)
    )
    @test (
        free_energies(gauss) ≈ free_energies(drelu) ≈
        free_energies(prelu) ≈ free_energies(xrelu)
    )
    @test (
        mode_from_inputs(gauss) ≈ mode_from_inputs(drelu) ≈
        mode_from_inputs(prelu) ≈ mode_from_inputs(xrelu)
    )
    @test (
        mean_from_inputs(gauss) ≈ mean_from_inputs(drelu) ≈
        mean_from_inputs(prelu) ≈ mean_from_inputs(xrelu)
    )
    @test (
        mean_abs_from_inputs(gauss) ≈ mean_abs_from_inputs(drelu) ≈
        mean_abs_from_inputs(prelu) ≈ mean_abs_from_inputs(xrelu)
    )
    @test (
        var_from_inputs(gauss) ≈ var_from_inputs(drelu) ≈
        var_from_inputs(prelu) ≈ var_from_inputs(xrelu)
    )

    drelu = dReLU(randn(1), [0.0], randn(1), [Inf])
    relu = ReLU(drelu.θp, drelu.γp)
    x = rand(1, 100)
    @test energies(relu, x) ≈ energies(drelu, x)
    @test free_energies(relu) ≈ free_energies(drelu)
    @test mode_from_inputs(relu) ≈ mode_from_inputs(drelu)
    #@test mean_from_inputs(relu) ≈ mean_from_inputs(drelu)
    #@test mean_abs_from_inputs(relu)  ≈ mean_abs_from_inputs(drelu)
    #@test var_from_inputs(relu) ≈ var_from_inputs(drelu)
end

@testset "dReLU" begin
    N = (3, 5)
    B = 13
    x = randn(N..., B)

    # bound γ away from zero to avoid issues with quadgk
    layer = dReLU(
        randn(N...), randn(N...),
        (rand(N...) .+ 1) .* (rand(Bool, N...) .- 0.5),
        (rand(N...) .+ 1) .* (rand(Bool, N...) .- 0.5)
    )

    Ep = energy(ReLU( layer.θp, layer.γp), max.( x, 0))
    En = energy(ReLU(-layer.θn, layer.γn), max.(-x, 0))
    @test energy(layer, x) ≈ Ep + En
    @test iszero(energy(layer, zero(x)))

    function quad_free(θp::Real, θn::Real, γp::Real, γn::Real)
        Z, ϵ = quadgk(h -> exp(-RBMs.drelu_energy(θp, θn, γp, γn, h)), -Inf, Inf)
        return -log(Z)
    end
    @test free_energies(layer) ≈ quad_free.(layer.θp, layer.θn, layer.γp, layer.γn)

    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    ∂ = RBMs.∂free_energy(layer)
    @test ∂.θp ≈ only(gs).θp
    @test ∂.θn ≈ only(gs).θn
    @test ∂.γp ≈ only(gs).γp
    @test ∂.γn ≈ only(gs).γn
    μ = mean_from_inputs(layer)
    ν = var_from_inputs(layer)
    μ2 = @. ν + μ^2
    @test ∂.θp + ∂.θn ≈ -μ
    @test ∂.γp .* sign.(layer.γp) + ∂.γn .* sign.(layer.γn) ≈ μ2/2
    @test RBMs.grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
    @test RBMs.grad2var(layer, ∂) ≈ var_from_inputs(layer)

    # check law of total variance
    inputs = randn(size(layer)..., 1000)
    ∂ = RBMs.∂free_energy(layer, inputs)
    h_ave = mean_from_inputs(layer, inputs)
    h_var = var_from_inputs(layer, inputs)
    μ = batchmean(layer, h_ave)
    ν_int = batchmean(layer, h_var)
    ν_ext = batchvar(layer, h_ave; mean = μ)
    ν = ν_int + ν_ext # law of total variance
    @test RBMs.grad2ave(layer, ∂) ≈ μ
    @test RBMs.grad2var(layer, ∂) ≈ ν
    μ1, ν1 = RBMs.total_meanvar_from_inputs(layer, inputs)
    @test μ1 ≈ μ ≈ RBMs.total_mean_from_inputs(layer, inputs)
    @test ν1 ≈ ν ≈ RBMs.total_var_from_inputs(layer, inputs)
end

@testset "pReLU" begin
    N = (3, 5, 7)
    layer = pReLU(randn(N...), randn(N...), randn(N...), 2rand(N...) .- 1)
    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    ∂ = RBMs.∂free_energy(layer)
    @test ∂.θ ≈ only(gs).θ ≈ -mean_from_inputs(layer)
    @test ∂.γ ≈ only(gs).γ
    @test ∂.Δ ≈ only(gs).Δ
    @test ∂.η ≈ only(gs).η
    @test RBMs.grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
    @test RBMs.grad2var(layer, ∂) ≈ var_from_inputs(layer)
end

@testset "xReLU" begin
    N = (3, 5, 7)
    layer = xReLU(randn(N...), randn(N...), randn(N...), randn(N...))
    gs = Zygote.gradient(layer) do layer
        sum(free_energies(layer))
    end
    ∂ = RBMs.∂free_energy(layer)
    @test ∂.θ ≈ only(gs).θ ≈ -mean_from_inputs(layer)
    @test ∂.γ ≈ only(gs).γ
    @test ∂.Δ ≈ only(gs).Δ
    @test ∂.ξ ≈ only(gs).ξ
    @test RBMs.grad2ave(layer, ∂) ≈ mean_from_inputs(layer)
    @test RBMs.grad2var(layer, ∂) ≈ var_from_inputs(layer)
end


text:
using AtomsBase
using InteratomicPotentials
using StaticArrays
using Unitful
using UnitfulAtomic

position = @SVector [1.0, 0.0, 0.0]
element = :Ar
atom1 = Atom(element, position * u"Å")

ϵ = 1.0
σ = 0.25
rcutoff = 2.0
lj = LennardJones(ϵ, σ, rcutoff)

@test isa(atom1, Atom)
@test isa(lj, EmpiricalPotential)

atom2 = Atom(element, 0.25 * position * u"Å")
atom3 = Atom(element, 0.5 * position * u"Å")
atom4 = Atom(element, 0.75 * position * u"Å")
atom5 = Atom(element, 1.5 * position * u"Å")
atom6 = Atom(element, 2.0 * position * u"Å")
atoms = [atom1, atom2, atom3, atom4, atom5, atom6]
box = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
bcs = [DirichletZero(), Periodic(), Periodic()]
system = FlexibleSystem(atoms, box * u"Å", bcs)
n = length(system)
nnlist = neighborlist(system, rcutoff)
@test isa(nnlist, NeighborList)

@test isa(potential_energy(system, lj), AbstractFloat)
@test isa(force(system, lj), AbstractVector{<:SVector{3,<:AbstractFloat}})

@test isa(virial(system, lj), AbstractFloat)
@test isa(virial_stress(system, lj), SVector{6,<:AbstractFloat})


text:
function get_pointer(str::ASCIIString)
    convert(Ptr{Int8}, pointer(str))
end

function get_pointer(kstr::KStr)
    convert(Ptr{KStr}, pointer_from_objref(kstr))
end

strptr(kstr::KStr) = strptr(kstr.s,1,kstr.l)
strptr(p::Ptr{Int8},st::Integer,ed::Integer) = stringfrompointer(convert(Ptr{UInt8},p),Int64(st),Int64(ed))
strptr(p::Ptr{UInt8},st::Integer,ed::Integer) = stringfrompointer(p,Int64(st),Int64(ed))
function stringfrompointer(p::Ptr{UInt8},st::Int64,ed::Int64)
    len = ed-st+1
    @assert len>0
    data = Array{UInt8,1}(len)
    for i in st:ed
        data[i-st+1] = unsafe_load(p,i)
    end
    ASCIIString(data)
end


text:
# Polynomial mean function

"""
    MeanPoly <: Mean

Polynomial mean function
```math
m(x) = ∑ᵢⱼ βᵢⱼxᵢʲ
```
with polynomial coefficients ``βᵢⱼ`` of shape ``d × D`` where ``d`` is the dimension of
observations and ``D`` is the degree of the polynomial.
"""
mutable struct MeanPoly <: Mean
    "Polynomial coefficients"
    β::Matrix{Float64}
    "Priors for mean parameters"
    priors::Array

    """
        MeanPoly(β::Matrix{Float64})

    Create `MeanPoly` with polynomial coefficients `β`.
    """
    MeanPoly(β::Matrix{Float64}) = new(β, [])
end

function mean(mPoly::MeanPoly, x::AbstractVector)
    dim = length(x)
    deg = size(mPoly.β, 2)
    dim == size(mPoly.β, 1) || throw(ArgumentError("Observations and mean function have inconsistent dimensions"))
    return sum(dot(x.^j, mPoly.β[:,j]) for j in 1:deg)
end

get_params(mPoly::MeanPoly) = vec(mPoly.β)
get_param_names(mPoly::MeanPoly) = get_param_names(vec(mPoly.β), :β)
num_params(mPoly::MeanPoly) = length(mPoly.β)

function set_params!(mPoly::MeanPoly, hyp::AbstractVector)
    length(hyp) == num_params(mPoly) || throw(ArgumentError("Polynomial mean function has $(num_param) parameters"))
    copyto!(mPoly.β, hyp)
end

function grad_mean(mPoly::MeanPoly, x::AbstractVector)
    dim = length(x)
    deg = size(mPoly.β, 2)
    dM_theta = Array{Float64}(undef, dim, deg)
    for i in 1:dim
        for j in 1:deg
            dM_theta[i,j] = x[i].^j
        end
    end
    return vec(dM_theta)
end


text:
# Workaround for GR warnings
ENV["GKSwstype"] = "100"

using Documenter, GeoStats

istravis = "TRAVIS" ∈ keys(ENV)

makedocs(
  format = Documenter.HTML(assets=["assets/style.css"], prettyurls=istravis),
  sitename = "GeoStats.jl",
  authors = "Júlio Hoffimann",
  pages = [
    "Home" => "index.md",
    "User guide" => "userguide.md",
    "Tutorials" => "tutorials.md",
    "Reference guide" => [
      "Data" => "data.md",
      "Domains" => "domains.md",
      "Problems" => "problems.md",
      "Variography" => [
        "empirical_variograms.md",
        "theoretical_variograms.md",
        "fitting_variograms.md"
      ],
      "Kriging estimators" => "estimators.md",
      "Solver comparisons" => "comparisons.md",
      "Spatial operations" => "operations.md",
      "Plotting" => "plotting.md"
    ],
    "Developer guide" => "devguide.md",
    "Contributing" => "contributing.md",
    "About" => [
      "Community" => "about/community.md",
      "License" => "about/license.md",
      "Citing" => "about/citing.md"
    ],
    "Index" => "links.md"
  ]
)

deploydocs(repo="github.com/juliohm/GeoStats.jl.git")


text:
# make sure not to underplot below 0
# we
function density_bands_to_ribbons(ys, bands_width)
    lower_lims = ys .- max.(ys .- bands_width, 0)
    upper_lims = copy(lower_lims)
    upper_lims .= bands_width
    (lower_lims, upper_lims)
end

function limits_to_ribbons(ys, bands::AbstractVector{<:Tuple})
    bands_left =  ys .- first.(bands)
    bands_right = last.(bands) .- ys
    (bands_left, bands_right)
end

@userplot ConfIntPlot

@recipe function f(c::ConfIntPlot)

      if length(c.args) == 2
            targets, cis = c.args
            estimated_targets = mean.(cis)
      elseif length(c.args) == 3
             targets, cis, estimated_targets = c.args
      else
             error("Expected 2 or 3 arguments")
      end

      xs = response.(targets)

      left_ci = first.(cis)
      right_ci = last.(cis)

      lower_lims = estimated_targets .- left_ci
      upper_lims = right_ci .- estimated_targets

      ribbons --> (lower_lims, upper_lims)
      fillalpha --> 0.4

      xlab --> L"\mu"
      ylab --> pretty_label(targets[1])

      xs, estimated_targets
end

# to plot results from real data analyses.

@recipe function f(resvec::Vector{<:CalibratedMinimaxEstimator})
	targets = [res.target for res in resvec]
	target_grid = response.(targets)

	calib_estimates = estimate.(targets, resvec)
	cis = confint.(targets, resvec)
	cis_ribbon  = limits_to_ribbons(calib_estimates, cis)
	pilots = [res.pilot for res in resvec]

	@series begin
		seriestype  -->  :path
		color --> "#550133"
		linestyle --> :dot
		label -->  L"\bar{\theta}(x)"#" #L"Pilot estimator $\bar{\theta}(x)$"
		target_grid, pilots
	end

	color --> "#018AC4"
	fillalpha --> 0.39
	ribbon --> cis_ribbon
	xguide --> L"x"

	yguide --> pretty_label(targets[1])
	ylims --> extrema(targets[1])
	legend --> :topleft
	label --> L"\bar{\theta}(x) + \hat{\Delta}(x)" #L"Calibrated estimator $\bar{\theta}(x) + \hat{\Delta}(x)$"

	target_grid, calib_estimates
end


text:
###########################################################
# Waveguide eigenvalue problem (WEP)
# Sum of products of matrices and functions (SPMF)
"""
Waveguide eigenvalue problem (WEP)
Sum of products of matrices and functions (SPMF)
"""
function assemble_waveguide_spmf_fd(nx::Integer, nz::Integer, hx, Dxx::SparseMatrixCSC, Dzz::SparseMatrixCSC, Dz::SparseMatrixCSC, C1::SparseMatrixCSC, C2T::SparseMatrixCSC, K::Union{Matrix{ComplexF64},Matrix{Float64}}, Km, Kp, pre_Schur_fact::Bool)
    Ix = sparse(ComplexF64(1)I, nx, nx)
    Iz = sparse(ComplexF64(1)I, nz, nz)
    Q0 = kron(Ix, Dzz) + kron(Dxx, Iz) + sparse(Diagonal(vec(K)))
    Q1 = kron(Ix, 2*Dz)
    Q2 = kron(Ix, Iz)

    A = Vector{SparseMatrixCSC}(undef, 3+2*nz)
    A[1] = hvcat((2,2), Q0, C1, C2T, spzeros(ComplexF64,2*nz, 2*nz) )
    A[2] = hvcat((2,2), Q1, spzeros(ComplexF64,nx*nz, 2*nz), spzeros(ComplexF64,2*nz, nx*nz), spzeros(ComplexF64,2*nz, 2*nz) )
    A[3] = hvcat((2,2), Q2, spzeros(ComplexF64,nx*nz, 2*nz), spzeros(ComplexF64,2*nz, nx*nz), spzeros(ComplexF64,2*nz, 2*nz) )

    f = Vector{Function}(undef, 3+2*nz)
    f[1] = λ -> one(λ)
    f[2] = λ -> λ
    f[3] = λ -> λ^2

    R, Rinv = generate_R_matvecs(nz)
    S = generate_S_function(nz, hx, Km, Kp)

    for j = 1:nz
        f[j+3] = λ -> S(λ, j)
        e_j = zeros(nz)
        e_j[j] = 1
        E_j = [R(e_j); spzeros(ComplexF64,nz)]
        E_j = E_j * (E_j/nz)'
        A[j+3] =  hvcat((2,2), spzeros(ComplexF64,nx*nz,nx*nz), spzeros(ComplexF64,nx*nz, 2*nz), spzeros(ComplexF64,2*nz, nx*nz), E_j)
    end
    for j = 1:nz
        f[j+nz+3] = λ -> S(λ, nz+j)
        e_j = zeros(nz)
        e_j[j] = 1
        E_j = [spzeros(ComplexF64,nz); R(e_j)]
        E_j = E_j * (E_j/nz)'
        A[j+nz+3] =  hvcat((2,2), spzeros(ComplexF64,nx*nz,nx*nz), spzeros(ComplexF64,nx*nz, 2*nz), spzeros(ComplexF64,2*nz, nx*nz), E_j)
    end
    return SPMF_NEP(A,f;Schur_fact=pre_Schur_fact)
end

###########################################################
# Generate R-matrix
# Part of defining the P-matrix, see above, Jarlebring-(1.6) and Ringh-(2.8) and Remark 1
# OBS: R' = nz * Rinv, as noted in Ringh between (2.8) and Remark 1
function generate_R_matvecs(nz)
    # The scaled FFT-matrix R
    p = (nz-1)/2
    bb = exp.(-2im*pi*((1:nz).-1)*(-p)/nz)  # scaling to do after FFT
    function R(X) # Note! Only works for vectors or one-dim matrices
        return reverse(bb .* fft(vec(X)), dims = 1)
    end
    bbinv = 1 ./ bb # scaling to do before inverse FFT
    function Rinv(X)
        return ifft(bbinv .* reverse(vec(X), dims = 1))
    end
    return R, Rinv
end

###########################################################
# Generate S-function for matrix argument
# Part of defining the P-matrix, see above, Jarlebring-(2.4) and Ringh-(2.8)(2.3)
function generate_S_function(nz, hx, Km, Kp)
    # Constants from the problem
    p = (nz-1)/2
    d0 = -3/(2*hx)
    b = 4*pi*1im * (-p:p)
    cM = Km^2 .- 4*pi^2 * ((-p:p).^2)
    cP = Kp^2 .- 4*pi^2 * ((-p:p).^2)

    # Note: λ should be scalar or matrix (not vector)
    betaM = function(λ, j)
        return λ^2 + b[j]*λ + cM[j]*one(λ)
    end
    betaP = function(λ, j)
        return λ^2 + b[j]*λ + cP[j]*one(λ)
    end

    sM = function(λ, j)
        return  1im*sqrt_schur_pos_imag(betaM(λ, j)) + d0*one(λ)
    end
    sP = function(λ, j)
        return  1im*sqrt_schur_pos_imag(betaP(λ, j)) + d0*one(λ)
    end

    S = function(λ, j)
        if j <= nz
            return sM(λ,j)
        elseif j <= 2*nz
            return sP(λ,(j-nz))
        else
            error("The chosen j = ", j, "but the setup nz = ", nz, ". Hence j>2nz which is illegal.")
        end
    end

    return S
end

###########################################################
"""
    sqrt_schur_pos_imag(A::AbstractMatrix)
 Computes the matrix square root on the 'correct branch',
 that is, with positivt imaginary part. Similar to Schur method
 in Algorithm 6.3 in Higham matrix functions.
"""
sqrt_schur_pos_imag(A::Number) = sqrt_pos_imag(A)
function sqrt_schur_pos_imag(A::AbstractMatrix)
    n = size(A,1);
    AA = Matrix{ComplexF64}(A);
    (T, Q, ) = schur(AA)
    U = zeros(ComplexF64,n,n);
    for i = 1:n
        U[i,i] = sqrt_pos_imag(T[i,i])
    end
    private_inner_loops_sqrt!(n, U, T)
    return Q*U*Q'
end
#Helper function executing the inner loop (more Juliaesque)
function private_inner_loops_sqrt!(n, U, T)
    temp = zero(ComplexF64);
    for j = 2:n
        for i = (j-1):-1:1
            temp *= zero(ComplexF64);
            for k = (i+1):(j-1)
                temp += U[i,k]*U[k,j]
            end
            U[i,j] = (T[i,j] - temp)/(U[i,i]+U[j,j])
        end
    end
end

"""
    sqrt_pos_imag(a::ComplexF64) and sqrt_pos_imag(a::Float64)
 Helper function: Computes the scalar square root on the 'correct branch',
 that is, with positivt imaginary part.
"""
function sqrt_pos_imag(a::ComplexF64)
    imag_sign = sign(imag(a))
    if imag_sign == 0 #Real value in complex type
        sqrt(a)
    else
        sign(imag(a))*sqrt(a)
    end
end
function sqrt_pos_imag(a::Float64)
    return sqrt(a)
end

    function Pinv(nep,λ,x)
        return vec(  [R(nep,Rinv(nep,x[1:Int(end/2)]) ./ sM(nep,λ));
                      R(nep,Rinv(nep,x[Int(end/2)+1:end]) ./ sP(nep,λ))  ]  )
    end

    function R(nep,X) # Note! Only works for vectors or one-dim matrices
        return reverse(nep.bb .* fft(vec(X)), dims = 1)
    end

    function Rinv(nep,X)
        return ifft(nep.bbinv .* reverse(vec(X), dims = 1))
    end

    function betaM(nep,λ)
        return λ^2 .+ nep.b*λ + nep.cM
    end

    function betaP(nep,λ)
        return λ^2 .+ nep.b*λ + nep.cP
    end

    function sM(nep,λ)
        bbeta = betaM(nep,λ)
        return 1im*sign.(imag(bbeta)).*sqrt.(bbeta) .+ nep.d0
    end

    function sP(nep,λ)
        bbeta = betaP(nep,λ)
        return 1im*sign.(imag(bbeta)).*sqrt.(bbeta) .+ nep.d0
    end

###########################################################
# Waveguide eigenvalue problem - WEP
# A more optimized (native) implementation of the WEP with FD discretization
"""
    WEP_FD

Waveguide eigenvalue problem

A more optimized implementation of the WEP for FD-discretization.\\
Closer to what is done in the article:
''E. Ringh, and G. Mele, and J. Karlsson, and E. Jarlebring,
Sylvester-based preconditioning for the waveguide eigenvalue problem,
Linear Algebra and its Applications''
"""
    struct WEP_FD <: WEP
        nx::Int
        nz::Int
        hx::Float64
        hz::Float64
        Dxx::SparseMatrixCSC{Float64,Int}
        Dzz::SparseMatrixCSC{Float64,Int}
        Dz::SparseMatrixCSC{Float64,Int}
        Iz::SparseMatrixCSC{Float64,Int}
        C1::SparseMatrixCSC{Float64,Int}
        C2T::SparseMatrixCSC{Float64,Int}
        k_bar::ComplexF64
        K::Matrix{ComplexF64}
        p::Integer
        d0::Float64
        d1::Float64
        d2::Float64
        b::Vector{ComplexF64}
        cM::Vector{ComplexF64}
        cP::Vector{ComplexF64}
        bb::Vector{ComplexF64}
        bbinv::Vector{ComplexF64}

        function WEP_FD(nx, nz, hx, hz, Dxx, Dzz, Dz, C1, C2T, K, Km, Kp)
            n = nx*nz + 2*nz
            k_bar = mean(K)
            K_scaled = K-k_bar*ones(ComplexF64,nz,nx)

            p = (nz-1)/2

            d0 = -3/(2*hx)
            d1 = 2/hx
            d2 = -1/(2*hx)

            b = 4*pi*1im * (-p:p)
            cM = Km^2 .- 4*pi^2 * ((-p:p).^2)
            cP = Kp^2 .- 4*pi^2 * ((-p:p).^2)

            bb = exp.(-2im*pi*((1:nz).-1)*(-p)/nz)  # scaling to do after FFT
            bbinv = 1 ./ bb # scaling to do before inverse FFT

            Iz = sparse(ComplexF64(1)I, nz, nz)

            return new(nx, nz, hx, hz, Dxx, Dzz, Dz, Iz, C1, C2T, k_bar, K_scaled, p, d0, d1, d2, b, cM, cP, bb, bbinv)
        end
    end

    function A(nep::WEP_FD, λ, d=0)
        nz = nep.nz
        if(d == 0)
            return nep.Dzz + 2*λ*nep.Dz + (λ^2 + nep.k_bar)*nep.Iz
        elseif(d == 1)
            return 2*nep.Dz + 2*λ*nep.Iz
        elseif(d == 2)
            return 2*nep.Iz
        else
            return spzeros(ComplexF64, nz, nz)
        end
    end

    function B(nep::WEP_FD, λ, d=0)
        if(d == 0)
            return nep.Dxx
        else
            return spzeros(ComplexF64, nx, nx)
        end
    end

    #Inverses of the boundary operators, Ringh - (2.8)
    #To be used in the Schur-complement- and SMW-context.
    function P_inv_m(nep::WEP_FD, λ, v)
        nz = nep.nz
        coeffs = zeros(ComplexF64, nz)
            aa = 1.0
        for j = 1:nz
            bb = nep.b[j]
            cc = nep.cM[j]
            coeffs[j] = 1im*sqrt_derivative(aa, bb, cc, 0, λ) + nep.d0
        end
        return R(nep,Rinv(nep,v) ./ coeffs)
    end

    function P_inv_p(nep::WEP_FD, λ, v)
        nz = nep.nz
        coeffs = zeros(ComplexF64, nz)
            aa = 1.0
        for j = 1:nz
            bb = nep.b[j]
            cc = nep.cP[j]
            coeffs[j] = 1im*sqrt_derivative(aa, bb, cc, 0, λ) + nep.d0
        end
        return R(nep,Rinv(nep,v) ./ coeffs)
    end

    function size(nep::WEP_FD)
        n = nep.nx*nep.nz + 2*nep.nz
        return (n,n)
    end
    function size(nep::WEP_FD, dim)
        n = nep.nx*nep.nz + 2*nep.nz
        return n
    end

    function issparse(nep::WEP_FD)
        return false
    end

"""
    compute_Mlincomb(nep::WEP_FD, λ::Number, V; a=ones(ComplexF64,size(V,2)))

Specialized for Waveguide Eigenvalue Problem discretized with Finite Difference\\\\
 Computes the linear combination of derivatives\\
 ``Σ_i a_i M^{(i)}(λ) v_i``
"""
    function compute_Mlincomb(nep::WEP_FD, λ::Number, V::AbstractVecOrMat,
                              a::Vector=ones(ComplexF64,size(V,2)))
        na = size(a,1)
        nv = size(V,1)
        mv = size(V,2)
        n_nep = size(nep,1)
        if(na != mv)
            error("Incompatible sizes: Number of coefficients = ", na, ", number of vectors = ", mv, ".")
        end
        if(nv != n_nep)
            error("Incompatible sizes: Length of vectors = ", nv, ", size of NEP = ", n_nep, ".")
        end
        nx = nep.nx
        nz = nep.nz
        max_d = na - 1 #Start on 0:th derivative

        V1 = view(V, 1:nx*nz, :)
        V1_mat = reshape(V1, nz, nx, na)
        V2 = view(V, nx*nz+1:n_nep, :)

        # Compute the top part (nx*nz)
        y1_mat::Matrix{ComplexF64} = (A(nep,λ) * V1_mat[:,:,1] + V1_mat[:,:,1] * B(nep,λ)  +  nep.K .* V1_mat[:,:,1])*a[1]
        for d = 1:min(max_d,3)
            y1_mat[:,:] += A(nep,λ,d) * V1_mat[:,:,d+1] * a[d+1];
        end
        y1::Vector{ComplexF64} = y1_mat[:]
        y1[:] += nep.C1 * V2[:,1] * a[1]

        # Compute the bottom part (2*nz)
        D::Matrix{ComplexF64} = zeros(ComplexF64, 2*nz, na)
        cMP = vcat(nep.cM, nep.cP)
        for j = 1:2*nz
            aa = 1
            bb = nep.b[rem(j-1,nz)+1]
            cc = cMP[j]
            der_coeff = 1im*sqrt_derivative(aa, bb, cc, max_d, λ)
            for jj = 1:na
                D[j, jj] = der_coeff[jj]
            end
        end

        #Multiplication with diagonal matrix optimized by working "elementwise" Jarlebring-(4.6)
        y2_temp::Vector{ComplexF64} =
            (D[:,1] .+ nep.d0) .* [Rinv(nep,V2[1:nz,1]);
                                   Rinv(nep,V2[nz+1:2*nz,1])]*a[1]
        for jj = 2:na
            #Multiplication with diagonal matrix optimized by working "elementwise" Jarlebring-(4.6)
            y2_temp[:] += D[:,jj] .* [Rinv(nep,V2[1:nz,jj]);
                                      Rinv(nep,V2[nz+1:2*nz,jj])] *a[jj]
        end
        y2::Vector{ComplexF64} = [R(nep,y2_temp[1:nz,1]);
                                  R(nep,y2_temp[nz+1:2*nz,1])]
        y2[:] += nep.C2T * V1[:,1]*a[1] #Action of C2T. OBS: Add last because of implcit storage in R*D_i*R^{-1}*v_i

        return vcat(y1, y2)
    end

"""
    compute_Mder(nep::WEP,λ::Number,i::Integer=0)
"""
    function compute_Mder(nep::WEP,λ::Number,i::Integer=0)
        error("The WEP does not implement this function. If this was called in a situation where you want to solve linear systems please look at `WEPLinSolverCreator`")
    end

###########################################################
# Linear Solvers for WEP

    # Matrix vector operations for the Schur complement (to be used in GMRES call)
    # Matrix-vector product according to Ringh (2.13) and (3.3)
    struct SchurMatVec
        nep::WEP_FD
        λ::ComplexF64
        SchurMatVec(nep::WEP_FD, λ::Union{ComplexF64,Float64}) = new(nep, λ)
    end

    function *(M::SchurMatVec,v::AbstractVector)
        λ = M.λ
        nep = M.nep

        X = reshape(v, nep.nz, nep.nx)
        return vec(  vec( A(nep,λ)*X + X*B(nep,λ) + nep.K.*X ) - nep.C1 * Pinv(nep, λ, nep.C2T*v)  )
    end

    function (M::SchurMatVec)(v::AbstractVector) #Overload the ()-function so that a SchurMatVec struct can act and behave like a function
        return M*v
    end

    function size(M::SchurMatVec, dim=-1)
        n = M.nep.nx*M.nep.nz
        if (dim==-1)
            return (n,n)
        else
            return n
        end
    end

    function eltype(M::SchurMatVec)
        return ComplexF64
    end

    # GMRES Solver
    struct WEPGMRESLinSolver<:LinSolver
        schur_comp::LinearMap{ComplexF64}
        kwargs
        gmres_log::Bool
        nep::WEP_FD
        λ::ComplexF64

        function WEPGMRESLinSolver(nep::WEP_FD,λ::Union{ComplexF64,Float64},kwargs)
            f = SchurMatVec(nep, λ)
            schur_comp = LinearMap{ComplexF64}(f, nep.nx*nep.nz, ismutating=false, issymmetric=false, ishermitian=false);
            gmres_log = false
            for elem in kwargs
                gmres_log |= ((elem[1] == :log) && elem[2])
            end
            return new(schur_comp, kwargs, gmres_log,nep,λ)
        end
    end

    function WEP_inner_lin_solve(solver::WEPGMRESLinSolver, rhs::Vector, tol)
        if( solver.gmres_log )
            q, convhist = gmres(solver.schur_comp, rhs; tol=tol, solver.kwargs...)
        else
            q = gmres(solver.schur_comp, rhs; tol=tol, solver.kwargs...)
        end
        return q
    end

    # Direct Backslash solver
    struct WEPBackslashLinSolver<:LinSolver
        schur_comp::SparseMatrixCSC{ComplexF64,Int}
        nep::WEP_FD
        λ::ComplexF64

        function WEPBackslashLinSolver(nep::WEP_FD, λ::Union{ComplexF64,Float64}, kwargs=())
            schur_comp = construct_WEP_schur_complement(nep, λ)
            return new(schur_comp, nep, λ)
        end
    end

    function WEP_inner_lin_solve(solver::WEPBackslashLinSolver, rhs::Vector, tol)
        return solver.schur_comp \ rhs
    end

    # Direct pre-factorized solver
    struct WEPFactorizedLinSolver<:LinSolver
        schur_comp_fact
        nep::WEP_FD
        λ::ComplexF64

        function WEPFactorizedLinSolver(nep::WEP_FD, λ::Union{ComplexF64,Float64}, kwargs=())
            schur_comp_fact = factorize(construct_WEP_schur_complement(nep, λ))
            return new(schur_comp_fact, nep, λ)
        end
    end

    function WEP_inner_lin_solve(solver::WEPFactorizedLinSolver, rhs::Vector, tol)
        return solver.schur_comp_fact \ rhs
    end

"""
    struct WEPLinSolverCreator <: LinSolverCreator
    function WEPLinSolverCreator(;solver_type=:backslash,kwargs=())

The a linsolver creator for the waveguide eigenvalue problem. The
kwarg `solver_type` is one of `:backslash`, `:factorized`, `:gmres`.
The `kwargs` keyword argument is passed to the solver.
"""
    struct WEPLinSolverCreator <: LinSolverCreator
        solver_type::Symbol
        kwargs;
        function WEPLinSolverCreator(;solver_type=:factorized,kwargs=())
            return new(solver_type,kwargs);
        end
    end

    function create_linsolver(creator::WEPLinSolverCreator,nep,λ)
        if (!(nep isa WEP_FD))
            t=typeof(nep)
            error("WEPLinSolver can only be used in combination with WEPs: type(nep)=$t");
        end
        if (creator.solver_type==:backslash)
            return WEPBackslashLinSolver(nep, λ, creator.kwargs)
        elseif (creator.solver_type==:gmres)
            return WEPGMRESLinSolver(nep, λ, creator.kwargs)
        elseif (creator.solver_type==:factorized)
            return WEPFactorizedLinSolver(nep, λ, creator.kwargs)
        else
            s=creator.solver_type
            error("Unknown type of solver_type in linsolvercreator:$s");
        end
    end

    # Helper functions for WEP LinSolvers. To avoid code repetition.
    # Assembls the full Schur-complement, used in both Backslash and LU solvers
    function construct_WEP_schur_complement(nep::WEP_FD, λ::Union{ComplexF64,Float64})
        nz = nep.nz
        nx = nep.nx
        Inz = sparse(ComplexF64(1)I, nz, nz)
        Inx = sparse(ComplexF64(1)I, nx, nx)

        Pinv_minus = Matrix{ComplexF64}(undef, nz, nz)
        Pinv_plus = Matrix{ComplexF64}(undef, nz, nz)
        e = zeros(ComplexF64,nz)
        for i = 1:nz
            e[:] .= 0
            e[i] = 1
            Pinv_minus[:,i] = P_inv_m(nep, λ, e)
            Pinv_plus[:,i] = P_inv_p(nep, λ, e)
        end

        E = spzeros(nx,nx)
        E[1,1] = nep.d1/(nep.hx^2)
        E[1,2] = nep.d2/(nep.hx^2)
        EE = spzeros(nx,nx)
        EE[nx,nx] = nep.d1/(nep.hx^2)
        EE[nx,nx-1] = nep.d2/(nep.hx^2)

        # Kronecker product form of Ringh - Proposition 3.1
        return kron(copy(B(nep,λ)'), Inz) + kron(Inx, A(nep,λ)) + sparse(Diagonal(nep.K[:])) - kron(E, Pinv_minus) - kron(EE, Pinv_plus)
    end

    # lin_solve function to wrapp all the WEP linear solvers.
    # Since Schur-complement transformations are the same.
    # Does transforming between that and the full system.
    # Ringh - Proposition 2.1, see also Algorithm 2, step 10-11.
    function lin_solve(solver::Union{WEPBackslashLinSolver,WEPGMRESLinSolver,WEPFactorizedLinSolver}, x::Vector; tol=eps(Float64))
    # Ringh - Proposition 2.1
        λ = solver.λ
        nep = solver.nep

        x_int = x[1:(nep.nx*nep.nz)]
        x_ext = x[((nep.nx*nep.nz)+1):((nep.nx*nep.nz) + 2*nep.nz)]
        rhs =  vec(  x_int - nep.C1*Pinv(nep, λ, x_ext))

        q = WEP_inner_lin_solve(solver, rhs, tol)

        return [q; vec(Pinv(nep, λ, -nep.C2T * q + x_ext))]

    end

###########################################################
#Square root of second degree polynomial (Gegenbauer polynomials)
#Jarlebring - Appendix C
"""
    sqrt_derivative(a,b,c, d=0, x=0)

Computes all d derivatives of sqrt(a*z^2 + b*z + c)
in the point z = x.
Returns a d+1 vector with all numerical values
"""
function sqrt_derivative(a,b,c, d=0, x=0)
    if(d<0)
        error("Cannot take negative derivative. d = ", d)
    end

    aa = a
    bb = b + 2*a*x
    cc = c + a*x^2 + b*x

    derivatives = zeros(ComplexF64,d+1)

    yi = sqrt_pos_imag(cc)
    derivatives[1] = yi
    if( d==0 )
        return derivatives[1] #OBS: If only function is sought, return it as an integer and not array
    end

    yip1 = bb/(2*sqrt_pos_imag(cc))
    fact = Float64(1)
    derivatives[2] = yip1 * fact
    if( d==1 )
        return derivatives
    end

    yip2 = zero(ComplexF64)
    for i = 2:d
        m = i - 2
        yip2 = - (2*aa*(m-1)*yi  +  bb*(1+2*m)*yip1) / (2*cc*(2+m))
        fact *= i

        yi = yip1
        yip1 = yip2

        derivatives[i+1] = yip2 * fact
    end
    return derivatives
end


text:
import Base: exp, log, transpose
import Base: .+, +, .-, -, .*, *

for op in (:exp, :log)
    @eval begin
        @generated function $op{T}(x::CuArray{T})
            op = $op
            f = CuFunction("""
            __global__ void f($T *x, $T *y, int length) {
                int idx = blockIdx.x * blockDim.x + threadIdx.x;
                if (idx < length) {
                    y[idx] = $op(x[idx]);
                }
            }""")
            quote
                y = similar(x)
                $f(x.ptr, y.ptr, length(x), dx=length(x))
                y
            end
        end
    end
end

@generated function transpose2{T}(x::CuMatrix{T})

    quote
        y = similar(x, size(x,2), size(x,1))
        $f(y.ptr, x.ptr)
    end
end

for op in (:+, :-)
    @eval begin
        @generated function $op{T}(x1::CuArray{T}, x2::CuArray{T})
            op = $op
            f = CuFunction("""
            __global__ void f($T *x1, $T *x2, $T *y, int length) {
                int idx = blockIdx.x * blockDim.x + threadIdx.x;
                if (idx < length) {
                    y[idx] = x1[idx] $op x2[idx];
                }
            }""")
            quote
                size(x1) == size(x2) || throw(DimensionMismatch())
                y = similar(x1)
                $f(x1.ptr, x2.ptr, y.ptr, length(y), dx=length(y))
                y
            end
        end
        @generated function $op{T}(x1::Number, x2::CuArray{T})
            op = $op
            f = CuFunction("""
            __global__ void f($T x1, $T *x2, $T *y, int length) {
                int idx = blockIdx.x * blockDim.x + threadIdx.x;
                if (idx < length) {
                    y[idx] = x1 $op x2[idx];
                }
            }""")
            quote
                y = similar(x2)
                $f(T(x1), x2.ptr, y.ptr, length(y), dx=length(y))
                y
            end
        end
        @generated function $op{T}(x1::CuArray{T}, x2::Number)
            op = $op
            f = CuFunction("""
            __global__ void f($T *x1, $T x2, $T *y, int length) {
                int idx = blockIdx.x * blockDim.x + threadIdx.x;
                if (idx < length) {
                    y[idx] = x1[idx] $op x2;
                }
            }""")
            quote
                y = similar(x1)
                $f(x1.ptr, x2, y.ptr, length(y), dx=length(y))
                y
            end
        end
    end
end

@generated function -{T}(x::CuArray{T})
    f = CuFunction("""
    __global__ void f($T *x, $T *y, int length) {
        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx < length) {
            y[idx] = -x[idx];
        }
    }""")
    quote
        y = similar(x)
        $f(x.ptr, y.ptr, length(y), dx=length(y))
        y
    end
end

for op in (:.+, :.-, :.*)
    @eval begin
        function $op{T,N}(x1::CuArray{T,N}, x2::CuArray{T,N})
            dims = ntuple(N) do i
                size(x1,i) == size(x2,i) && return size(x1,i)
                size(x1,i) == 1 && return size(x2,i)
                size(x2,i) == 1 && return size(x1,i)
                throw("Cannot be broadcasted.")
            end
            y = CuArray{T}(dims)
            broadcast!($op, y, x1, x2)
        end
    end
end

*(x1::CuMatrix, x2::CuVector) = BLAS.gemv(x1, x2)
*(x1::CuMatrix, x2::CuMatrix) = BLAS.gemm(x1, x2)


text:
################################################################################
#
#  Minkowski matrix
#
################################################################################

function minkowski_matrix_parallel(O::NfOrd, abs_tol::Int = 64)
  if isdefined(O, :minkowski_matrix) && O.minkowski_matrix[2] > abs_tol
    A = deepcopy(O.minkowski_matrix[1])
  else
    T = Vector{Vector{arb}}(undef, degree(O))
    B = O.basis_nf
    @Threads.threads for i in 1:degree(O)
      T[i] = minkowski_map(B[i], abs_tol)
    end
    p = maximum(Int[ prec(parent(T[i][j])) for i in 1:degree(O), j in 1:degree(O) ])
    M = zero_matrix(ArbField(p, cached = false), degree(O), degree(O))
    for i in 1:degree(O)
      for j in 1:degree(O)
        M[i, j] = T[i][j]
      end
    end
    O.minkowski_matrix = (M, abs_tol)
    A = deepcopy(M)
  end
  return A
end

@doc Markdown.doc"""
    minkowski_gram_mat_scaled(O::NfOrd, prec::Int = 64) -> fmpz_mat

Let $c$ be the Minkowski matrix as computed by `minkowski_matrix` with precision $p$.
This function computes $d = round(c 2^p)$ and returns $round(d d^t/2^p)$.
"""
function minkowski_gram_mat_scaled_parallel(O::NfOrd, prec::Int = 64)
  if isdefined(O, :minkowski_gram_mat_scaled) && O.minkowski_gram_mat_scaled[2] >= prec
    A = deepcopy(O.minkowski_gram_mat_scaled[1])
    Hecke.shift!(A, prec - O.minkowski_gram_mat_scaled[2])
  else
    c = minkowski_matrix_parallel(O, prec)
    d = zero_matrix(FlintZZ, degree(O), degree(O))
    A = zero_matrix(FlintZZ, degree(O), degree(O))
    round_scale!(d, c, prec)
    ccall((:fmpz_mat_gram, libflint), Nothing, (Ref{fmpz_mat}, Ref{fmpz_mat}), A, d)
    Hecke.shift!(A, -prec)
    O.minkowski_gram_mat_scaled = (A, prec)
    A = deepcopy(A)
  end
  # to ensure pos. definitenes, we add n to the diag.
  for i=1:degree(O)
    A[i, i] += nrows(A)
  end
  return A
end

function parallel_lll_precomputation(M::NfOrd, prec::Int, nblocks::Int = 4)
  n = degree(M)
  K = nf(M)
  dimension_blocks = div(n, nblocks)
  blocks = Vector{Int}[]
  for i = 1:nblocks-1
    int = (dimension_blocks*(i-1)+1):dimension_blocks*i
    push!(blocks, collect(int))
  end
  int = (dimension_blocks*(nblocks-1)+1):n
  push!(blocks, collect(int))
  new_prec = prec
  to_do = Hecke.subsets(blocks, 2)
  @vprint :LLL 3 "Subsets computed\n"
  #I need first to create all the blocks to do simultaneously.
  done = falses(length(to_do))
  while !all(done)
    blocks_selection = Vector{Int}[]
    block = 1
    @vprint :LLL 3 "Starting block selection\n"
    while block < length(to_do)+1
      while block < length(to_do)+1 && (done[block] || !Hecke._has_trivial_intersection(to_do[block], blocks_selection))
        block += 1
      end
      if block == length(to_do)+1
        break
      end
      done[block] = true
      indices = vcat(to_do[block][1], to_do[block][2])
      push!(blocks_selection, indices)
      block += 1
    end
    @vprint :LLL 3 "Blocks selection finished\n"
    g = identity_matrix(FlintZZ, n)
    while true
      try
        Hecke.minkowski_gram_mat_scaled(M, prec)
	    break
      catch e
        prec *= 2
      end
    end
    @vprint :LLL 3 "Start computation \n"
    @Threads.threads for i = 1:length(blocks_selection)
      indices = blocks_selection[i]
      g1 = Hecke._lll_sublattice(M, indices, prec = prec)[2]
      Hecke._copy_matrix_into_matrix(g, indices, indices, g1)
    end
    On = NfOrd(K, g*basis_matrix(M, copy = false))
    On.ismaximal = M.ismaximal
    if isdefined(M, :index)
      On.index = M.index
    end
    if isdefined(M, :disc)
      On.disc = M.disc
    end
    if isdefined(M, :gen_index)
      On.gen_index = M.gen_index
    end
    M = On
  end
  @vprint :LLL 3 "Precomputation finished with precision $(prec)\n"
  return prec, M
end


text:
# This file was generated, do not modify it. # hide
coerce!(df0, Textual => Multiclass);

max_stars_repo_name: alan-turing-institute/MLJTutorials

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
"""
    PlotConfig(;kwargs...)

Configuration options to be sent to the frontend to control aspects of how the plot is rendered. The acceptable keyword arguments are:

- `scrollZoom`: Determines whether mouse wheel or two-finger scroll zooms is enable. Turned on by default for gl3d, geo and mapbox subplots (as these subplot types do not have zoombox via pan), but turned off by default for cartesian subplots. Set `scrollZoom` to *false* to disable scrolling for all subplots.
- `editable`: Determines whether the graph is editable or not. Sets all pieces of `edits` unless a separate `edits` config item overrides individual parts.
- `staticPlot`: Determines whether the graphs are interactive or not. If *false*, no interactivity, for export or image generation.
- `toImageButtonOptions`: Statically override options for toImage modebar button allowed keys are format, filename, width, height
- `displayModeBar`: Determines the mode bar display mode. If *true*, the mode bar is always visible. If *false*, the mode bar is always hidden. If *hover*, the mode bar is visible while the mouse cursor is on the graph container.
- `modeBarButtonsToRemove`: Remove mode bar buttons by name
- `modeBarButtonsToAdd`: Add mode bar button using config objects. To enable predefined modebar buttons e.g. shape drawing, hover and spikelines, simply provide their string name(s). This could include: *v1hovermode*, *hoverclosest*, *hovercompare*, *togglehover*, *togglespikelines*, *drawline*, *drawopenpath*, *drawclosedpath*, *drawcircle*, *drawrect* and *eraseshape*. Please note that these predefined buttons will only be shown if they are compatible with all trace types used in a graph.
- `modeBarButtons`: Define fully custom mode bar buttons as nested array where the outer arrays represents button groups, and the inner arrays have buttons config objects or names of default buttons.
- `showLink`: Determines whether a link to Chart Studio Cloud is displayed at the bottom right corner of resulting graphs. Use with `sendData` and `linkText`.
- `plotlyServerURL`: When set it determines base URL for the 'Edit in Chart Studio' `showEditInChartStudio`/`showSendToCloud` mode bar button and the showLink/sendData on-graph link. To enable sending your data to Chart Studio Cloud, you need to set both `plotlyServerURL` to 'https://chart-studio.plotly.com' and also set `showSendToCloud` to true.
- `linkText`: Sets the text appearing in the `showLink` link.
- `showEditInChartStudio`: Same as `showSendToCloud`, but use a pencil icon instead of a floppy-disk. Note that if both `showSendToCloud` and `showEditInChartStudio` are turned, only `showEditInChartStudio` will be honored.
- `locale`: Which localization should we use? Should be a string like 'en' or 'en-US'.
- `displaylogo`: Determines whether or not the plotly logo is displayed on the end of the mode bar.
- `responsive`: Determines whether to change the layout size when window is resized.
- `doubleClickDelay`: Sets the delay for registering a double-click in ms. This is the time interval (in ms) between first mousedown and 2nd mouseup to constitute a double-click. This setting propagates to all on-subplot double clicks (except for geo and mapbox) and on-legend double clicks.
"""
@with_kw mutable struct PlotConfig
    scrollZoom::Union{Nothing,Bool} = true
    editable::Union{Nothing,Bool} = false
    staticPlot::Union{Nothing,Bool} = false
    toImageButtonOptions::Union{Nothing,Dict} = nothing
    displayModeBar::Union{Nothing,Bool} = nothing
    modeBarButtonsToRemove::Union{Nothing,Array} = nothing
    modeBarButtonsToAdd::Union{Nothing,Array} = nothing
    modeBarButtons::Union{Nothing,Array} = nothing
    showLink::Union{Nothing,Bool} = false
    plotlyServerURL::Union{Nothing,String} = nothing
    linkText::Union{Nothing,String} = nothing
    showEditInChartStudio::Union{Nothing,Bool} = nothing
    locale::Union{Nothing,String} = nothing
    displaylogo::Union{Nothing,Bool} = nothing
    responsive::Union{Nothing,Bool} = true
    doubleClickDelay::Union{Nothing,Int} = nothing
end

function JSON.lower(pc::PlotConfig)
    out = Dict{Symbol,Any}()
    for fn in fieldnames(PlotConfig)
        field = getfield(pc, fn)
        if !isnothing(field)
            out[fn] = field
        end
    end
    out
end

function Base.get(pc::PlotConfig, field::Symbol, default::Any)
    if hasfield(pc, field)
        val = getfield(pc, field)
        if !isnothing(val)
            return val
        end
    end
    return default
end