text:
import Pkg
Pkg.activate(".")

import DFTK: index_G_vectors, eval_psp_projection_radial
import SpecialFunctions: erf
using ProgressMeter
using DFTK
using DoubleFloats
using GenericLinearAlgebra
using HDF5
using IntervalArithmetic
using LinearAlgebra
using Polynomials
using Printf
using Roots
using TimerOutputs

import FFTW
FFTW.set_num_threads(1)

const to = TimerOutput()

###
### Helper functions
###

@doc raw"""
Returns the result of the integral ``∫_τ^∞ t^k exp(-t^2 / σ) dt``
"""
function integral_monomial_gaussian(k::Integer, τ::T; σ=1, upper_bound_ok=false) where T
    @assert k ≥ -4
    @assert σ > 0
    τdσ = τ / T(sqrt(σ))  # τ divided by σ

    ## Even k
    if k == -4
        exp(-τdσ^2) / 3τ^3 - 2 / T(3σ) * integral_monomial_gaussian(-2, τ, σ=σ)
    elseif k == -2
        sqrt(T(π) / σ)        * (erf(τdσ) - 1) + exp(-τdσ^2) / τ
    elseif k == 0
        sqrt(T(π) * σ)    / 2 * (1 - erf(τdσ))
    elseif k == 2
        sqrt(T(π) * σ^3)  / 4 * (1 - erf(τdσ)) + (                             σ * τ / 2) * exp(-τdσ^2)
    elseif k == 4
        3sqrt(T(π) * σ^5) / 8 * (1 - erf(τdσ)) + (             σ * τ^3 / 2 +  3σ^2*τ / 4) * exp(-τdσ^2)
    elseif k == 6
        15sqrt(T(π) * σ^7)/16 * (1 - erf(τdσ)) + (σ*τ^5 / 2 + 5σ^2*τ^3 / 4 + 15σ^3*τ / 8) * exp(-τdσ^2)
    #
    ## Odd k
    elseif k == -3
        upper_bound_ok || error("The case k == -3 is not implemented.")
        # ∫_τ^∞ (1/t^3) * exp(-t^2 / σ) dt ≤ (1/τ^3) ∫_τ^∞ exp(-t^2 / σ) dt
        integral_monomial_gaussian(0, τ, σ=σ) / τ^3
    elseif k == -1
        upper_bound_ok || error("The case k == -1 is not implemented.")
        # ∫_τ^∞ (1/t) * exp(-t^2 / σ) dt ≤ (1/τ) ∫_τ^∞ exp(-t^2 / σ) dt
        integral_monomial_gaussian(0, τ, σ=σ) / τ
        # Better bound:
        # \intx^\infty t^{-1} e^{-t^2} dt is
        # = \int{x^2} s^{-1} e^{-s} ds (use the change of variable s=t^2) is e^{-x^2}/(2x^2)
    elseif k == 1
        (                                 σ) / 2 * exp(-τdσ^2)
    elseif k == 3
        (                   σ * τ^2 +   σ^2) / 2 * exp(-τdσ^2)
    elseif k == 5
        (         σ * τ^4 + 2σ^2*τ^2 + 2σ^3) / 2 * exp(-τdσ^2)
    elseif k == 7
        (σ*τ^6 + 3σ^2*τ^4 + 6σ^3*τ^2 + 6σ^4) / 2 * exp(-τdσ^2)
    #
    ## Fallback by recursion
    else
        σ * τ^(k - 1) / 2 * exp(-τdσ^2) + σ * (k - 1) / 2 * integral_monomial_gaussian(k - 2, τ, σ=σ)
    end
end

@doc raw"""
Returns the result of the integral ``∫_τ^∞ t^k P(t) exp(-t^2 / σ) dt``
where `P` is a polynomial in `t`
"""
function integral_polynomial_gaussian(P::Poly, τ::T; k=0, σ=1, upper_bound_ok=false) where T
    # -1 translates from index in coeff to power in t
    sum(coeff * integral_monomial_gaussian(i + k - 1, τ, σ=σ, upper_bound_ok=upper_bound_ok)
        for (i, coeff) in enumerate(P.a) if coeff != 0)
end

@doc raw"""
Computes an upper bound for the sum ``∑_{|G| > G0} |β G|^k Q(|β G|) exp(-|β G|^2 / σ)``
where all `G` are on a lattice defined by `recip_lattice`, `β` and `σ` are a positive
constants, `k` is integer and `Q` is a polynomial both chosen such that
``t^k Q(t) exp(-t / σ)`` is a decreasing function for `t > G0`.
"""
function bound_sum_polynomial_gaussian(polynomial::Poly{T}, recip_lattice, G0::T;
                                       β=1, σ=1, k=0, Gmin=G0) where {T}
    # Determine dimensionality: Note: Clashes with usual DFTK convention
    @assert !iszero(recip_lattice[:, end])
    m = size(recip_lattice, 1)
    @assert size(recip_lattice) == (m, m)

    # Diameter of recip_lattice unit cell.
    diameter = norm(recip_lattice * ones(m))
    @assert diameter ≥ 0  # We assume lattice to be positively orianted
    @assert Gmin > 0
    @assert G0 ≤ Gmin

    # The terms (depending on β, n and G) we sum over.
    term(β, n, G) = (β * G)^n * polynomial(β * G) * exp(-β*G^2)

    if m == 1
        a = abs(recip_lattice[1, 1])
        start = floor(G0 / a) * a
        @assert start ≤ Gmin
        intm = integral_polynomial_gaussian(polynomial, β * Gmin, k=k, σ=σ,
                                            upper_bound_ok=true) / T(β)
        return  2 / a * ((Gmin - start) * term(β, k, Gmin) + intm)
    elseif m == 2 || m == 3
        prefactor = (m == 2 ? 2T(π) : 4T(π)) / abs(det(recip_lattice))
        start = G0 - diameter

        intm = integral_polynomial_gaussian(polynomial, β * Gmin, σ=σ,
                                            k=k + m - 1, upper_bound_ok=true) / T(β)^m
        result = prefactor * ((Gmin - start) * term(β, k + m - 1, Gmin) + intm)

        for (j, aj) in enumerate(eachcol(recip_lattice))
            newlattice = []
            nRj = 1
            for (k, ak) in enumerate(eachcol(recip_lattice))
                k == j && continue
                angle = dot(aj, ak) / dot(aj, aj)

                nRj = nRj * (2 + ceil(abs(angle)))
                newlatticevector = Vector(ak - angle * aj)
                deleteat!(newlatticevector, j)
                push!(newlattice, newlatticevector)
            end
            newlattice = hcat(newlattice...)

            result += nRj * bound_sum_polynomial_gaussian(polynomial, newlattice,
                                                          G0 - diameter, σ=σ, β=β, k=k, Gmin=Gmin)
        end

        return result
    end
    error("Dimensionality > 3")
end

###
### Determination of minimal Ecut2
###

function is_cohen_bergstresser(model)
    (
           all(at isa ElementCohenBergstresser for (at, pos) in model.atoms)
        && typeof.(model.term_types) == [Kinetic, AtomicLocal]
    )
end

function is_linear_atomic(model)
    (
           all(at isa ElementPsp for (at, pos) in model.atoms)
        && typeof.(model.term_types) == [Kinetic, AtomicLocal, AtomicNonlocal,
                                         Ewald, PspCorrection]
    )
end

"""
Estimate a lower bound for the Ecut2, the energy cutoff on the bigger residual grid,
from the given PSP, reciprocal lattice and Ecut
"""
function psp_minimal_Ecut2(psp::PspHgh, recip_lattice, Ecut)
    T = eltype(recip_lattice)
    minimal_Ecuts = T[]

    # Compute the diameter of the first BZ
    diameter = norm(recip_lattice * Vec3(1, 1, 1))
    @assert diameter ≥ 0

    # Local part: We want sqrt(2 * Ecut2) - norm(G) ≥ qcut, where G in Ecut basis
    qcut = DFTK.qcut_psp_local(T, psp)
    push!(minimal_Ecuts, (qcut + sqrt(2Ecut))^2 / 2)

    # Nonlocal projectors: We want sqrt(2 * Ecut2) - norm(k) ≥ qcut,
    #                      where G outside of Ecut2 basis
    proj_idcs = unique((i, l) for (i, l, _) in DFTK.projector_indices(psp))
    for (i, l) in proj_idcs
        qcut = DFTK.qcut_psp_projection_radial(T, psp, i, l)
        push!(minimal_Ecuts, (qcut + diameter)^2 / 2)
    end

    maximum(minimal_Ecuts)
end

function minimal_Ecut2(basis)
    model = basis.model

    if is_cohen_bergstresser(model)
        element = first(model.atoms[1])
        Gmax = sqrt(maximum(keys(element.V_sym))) * (2π / element.lattice_constant)
        return basis.Ecut + sqrt(2basis.Ecut) * Gmax + Gmax^2 / 2
    elseif is_linear_atomic(model)
        maximum(psp_minimal_Ecut2(atom.psp, model.recip_lattice, basis.Ecut)
                for (atom, positions) in model.atoms)
    else
        error("Not implemented")
    end
end

###
### Bounds for \|P_{B_2}^\perp V P_{B_1} \|  (potential_perpb2_b1)
###

@doc raw"""
Estimates an upper bound for ``\|P_{B_2}^\perp Vloc P_{B_1} \|_{op}``

hamblock:  Hamiltonian k-block in small Ecut basis
Ecut2:     Large ecut used for ``B_2``
"""
function bound_potential_perpb2_b1_Vloc(hamblock, Ecut2)
    # Code duplication with bound_perpb2_residual_Vloc

    basis = hamblock.basis
    kpoint = hamblock.kpoint
    @assert any(t isa AtomicLocal for t in basis.model.term_types)
    T = eltype(basis)
    recip_lattice = basis.model.recip_lattice

    accu = zero(T)
    for (element, positions) in basis.model.atoms
        element isa ElementCohenBergstresser && continue
        @assert element isa ElementPsp
        @assert element.psp isa PspHgh
        psp = element.psp

        # We need a bound for Σ_{G, |G| > qcut} Q^2(t) exp(-t^2) / t^4
        # where t = psp.rloc * |G|.
        Qloc = DFTK.psp_local_polynomial(T, psp)
        R(q) = bound_sum_polynomial_gaussian(Qloc * Qloc, recip_lattice, q, β=psp.rloc, k=-4)

        for (i, G) in enumerate(G_vectors(kpoint))
            qcut = sqrt(2Ecut2) - norm(recip_lattice * G)
            accu += R(qcut) * length(positions) / basis.model.unit_cell_volume^2
        end
    end
    sqrt(accu)
end

@doc raw"""
Estimates an upper bound for ``\|P_{B_2}^\perp Vnlock P_{B_1} \|_{op}``

hamblock:  Hamiltonian k-block in small Ecut basis
Ecut2:     Large ecut used for ``B_2``
"""
function bound_potential_perpb2_b1_Vnlock(hamblock, Ecut2)
    # Highly related to bound_potential_perpb1_Vnloc, a lot of code duplication
    basis = hamblock.basis
    kpoint = hamblock.kpoint
    model = basis.model
    kpoint = hamblock.kpoint
    T = eltype(basis)
    recip_lattice = model.recip_lattice
    @assert any(t isa AtomicLocal for t in model.term_types)

    # Tabulated values for the L∞ norm of the spherical harmonics Ylm
    Ylm∞ = (sqrt(1 / 4T(π)), sqrt(3 / 4T(π)))  # My guess: (5 / 4π, 7 / 4π)

    # Norms of the G vectors in the Ecut basis (B1)
    Gnorms = [norm(model.recip_lattice * (G + kpoint.coordinate))
              for G in G_vectors(kpoint)]

    function projector_bound(psp, l)
        @assert psp isa PspHgh
        n_proj_l = size(psp.h[l + 1], 1)

        sum_proj_B1 = zeros(T, n_proj_l)
        sum_proj_B2perp = zeros(T, n_proj_l)

        for i in 1:n_proj_l
            # Bound for Σ_{G, |G| > qcut} Q^2(t) exp(-t^2) where t = psp.rp[l+1] * |G|.
            Qnloc = DFTK.psp_projection_radial_polynomial(T, psp, i, l)
            qcut = sqrt(2Ecut2) - norm(recip_lattice * kpoint.coordinate)
            Nli = bound_sum_polynomial_gaussian(Qnloc * Qnloc, recip_lattice,
                                                qcut, β=psp.rp[l + 1])

            psp_radials = eval_psp_projection_radial.(psp, i, l, Gnorms)
            sum_proj_B1[i] = norm(psp_radials) * Ylm∞[l + 1] / sqrt(model.unit_cell_volume)
            sum_proj_B2perp[i] = sqrt(Nli) * Ylm∞[l + 1] / sqrt(model.unit_cell_volume)
        end

        ret = sum_proj_B2perp' * abs.(psp.h[l + 1]) * sum_proj_B1
        @assert ndims(ret) == 0
        ret
    end

    @assert all(element isa ElementPsp for (element, _) in model.atoms)
    sum(maximum(projector_bound.(element.psp, 0:element.psp.lmax)) * length(positions)
        for (element, positions) in model.atoms)
end

@doc raw"""
Estimates an upper bound for ``\|P_{B_2}^\perp H P_{B_1} \|_{op}``

hamblock:  Hamiltonian k-block in small Ecut basis
Ecut2:     Large ecut used for ``B_2``
"""
function bound_ham_perpb2_b1(hamblock, Ecut2)
    T = eltype(hamblock.basis)
    model = hamblock.basis.model
    if is_cohen_bergstresser(model)
        @assert Ecut2 ≥ minimal_Ecut2(hamblock.basis)
        return (AtomicLocal=zero(T), )
    elseif is_linear_atomic(model)
        (AtomicLocal=bound_potential_perpb2_b1_Vloc(hamblock, Ecut2),
         AtomicNonlocal=bound_potential_perpb2_b1_Vnlock(hamblock, Ecut2))
    else
        error("Not implemented")
    end
end

###
### Bounds for \|P_{B_2}^\perp V ψnk\|  (perpb2_residual)
###

@doc raw"""
Estimates an upper bound for ``\|P_{B_2}^\perp V_loc ψnk\|`` using varying
Sobolev exponents `s`.

hamblock:  Hamiltonian k-block in small Ecut basis
ψnk:       Eigenvector for a particular band and k-point on small ecut
Ecut2:     Large ecut used for ``B_2``
s:         Sobolev exponent
"""
function bound_perpb2_residual_Vloc_sobolev(hamblock, ψnk::AbstractVector, Ecut2; s=2)
    basis = hamblock.basis
    kpoint = hamblock.kpoint
    @assert length(ψnk) == length(G_vectors(kpoint))
    @assert any(t isa AtomicLocal for t in basis.model.term_types)
    T = eltype(basis)
    recip_lattice = basis.model.recip_lattice

    accu = zero(T)
    for (element, positions) in basis.model.atoms
        element isa ElementCohenBergstresser && continue
        @assert element isa ElementPsp
        @assert element.psp isa PspHgh
        psp = element.psp

        # We need a bound for Σ_{G, |G| > qcut} Q^2(t) exp(-t^2) / t^4
        # where t = psp.rloc * |G|.
        Qloc = DFTK.psp_local_polynomial(T, psp)
        R(q) = bound_sum_polynomial_gaussian(Qloc * Qloc, recip_lattice, q, β=psp.rloc, k=-4)

        for (i, G) in enumerate(G_vectors(kpoint))
            Gnorm = norm(recip_lattice * G)
            qcut = sqrt(2Ecut2) - Gnorm
            accu += (1 + Gnorm^2)^s * abs2(ψnk[i]) * R(qcut) * length(positions)
        end
    end

    prefac = sum((1 + norm(recip_lattice * G)^2)^(-s) for G in G_vectors(kpoint))
    sqrt(prefac * accu / basis.model.unit_cell_volume^2)
end

@doc raw"""
Estimates an upper bound for ``\|P_{B_2}^\perp V_loc ψnk\|``

hamblock:  Hamiltonian k-block in small Ecut basis
ψnk:     Eigenvector for a particular band and k-point on small ecut
Ecut2:   Large ecut used for ``B_2``
"""
function bound_perpb2_residual_Vloc(hamblock, ψnk::AbstractVector, Ecut2; α=1)
    basis = hamblock.basis
    kpoint = hamblock.kpoint
    @assert length(ψnk) == length(G_vectors(kpoint))
    @assert any(t isa AtomicLocal for t in basis.model.term_types)
    T = eltype(basis)
    recip_lattice = basis.model.recip_lattice

    accu = zero(T)
    for (element, positions) in basis.model.atoms
        element isa ElementCohenBergstresser && continue
        @assert element isa ElementPsp
        @assert element.psp isa PspHgh
        psp = element.psp

        # We need a bound for Σ_{G, |G| > qcut} Q^2(t) exp(-t^2) / t^4
        # where t = psp.rloc * |G|.
        Qloc = DFTK.psp_local_polynomial(T, psp)
        R(q) = bound_sum_polynomial_gaussian(Qloc * Qloc, recip_lattice, q, β=psp.rloc, k=-4)

        for (i, G) in enumerate(G_vectors(kpoint))
            factor = α == 0 ? one(T) : abs(ψnk[i])^2α
            qcut = sqrt(2Ecut2) - norm(recip_lattice * G)
            accu += factor * R(qcut) * length(positions)
        end
    end

    if α == 0
        prefac = one(T)
    elseif α == 1
        prefac = length(ψnk)
    else
        prefac = sum(@. abs(ψnk)^(2-2α))
    end
    sqrt(prefac * accu) / basis.model.unit_cell_volume
end

@doc raw"""
Estimates an upper bound for ``\|P_{B_2}^\perp V_nlock ψnk\|``

hamblock:  Hamiltonian k-block in small Ecut basis
ψnk:       Eigenvector for a particular band and k-point on small Ecut
Ecut2:     Large ecut used for ``B_2``
"""
function bound_perpb2_residual_Vnlock(hamblock, ψnk::AbstractVector, Ecut2)
    basis = hamblock.basis
    kpoint = hamblock.kpoint
    T = eltype(basis)
    recip_lattice = basis.model.recip_lattice
    @assert length(ψnk) == length(G_vectors(kpoint))

    # Compute norm of vector f
    idx_nonlocal = only(findall(t -> t isa AtomicNonlocal, basis.model.term_types))
    Vlock = hamblock.operators[idx_nonlocal]
    fnorm = norm(Vlock.D * (Vlock.P' * ψnk))

    accu = zero(T)
    for (element, positions) in basis.model.atoms
        @assert element isa ElementPsp
        @assert element.psp isa PspHgh
        psp = element.psp

        # Ignore m ... sum over it done implicitly by the (2l+1)
        proj_idcs = unique((i, l) for (i, l, _) in DFTK.projector_indices(psp))
        for (i, l) in proj_idcs
            # Bound for Σ_{G, |G| > qcut} Q^2(t) exp(-t^2) where t = psp.rp[l+1] * |G|.
            Qnloc = DFTK.psp_projection_radial_polynomial(T, psp, i, l)
            qcut = sqrt(2Ecut2) - norm(recip_lattice * kpoint.coordinate)
            N = bound_sum_polynomial_gaussian(Qnloc * Qnloc, recip_lattice,
                                              qcut, β=psp.rp[l + 1])

            accu += (2l + 1) * length(positions) * N
        end
    end

    fnorm * sqrt(accu / 4T(π) / basis.model.unit_cell_volume)
end

@doc raw"""
Estimates an upper bound for ``\|P_{B_2}^\perp H ψnk\|`` if ``ψnk`` is given
on the `Ecut` basis.

hamblock:  Hamiltonian k-block in small `Ecut` basis
ψnk:       Eigenvector for a particular band and k-point on small `Ecut`
Ecut2:     Large ecut used for ``B_2``
"""
function bound_perpb2_residual(hamblock, ψnk::AbstractVector, Ecut2)
    T = eltype(hamblock.basis)
    model = hamblock.basis.model
    if is_cohen_bergstresser(model)
        @assert Ecut2 ≥ minimal_Ecut2(hamblock.basis)
        return (Kinetic=zero(T), AtomicLocal=zero(T))
    elseif is_linear_atomic(model)
        (Kinetic=zero(T),
         AtomicLocal=bound_perpb2_residual_Vloc(hamblock, ψnk, Ecut2),
         AtomicNonlocal=bound_perpb2_residual_Vnlock(hamblock, ψnk, Ecut2),
         Ewald=zero(T),
         PspCorrection=zero(T)
        )
    else
        error("Not implemented")
    end
end

###
### Apply P_{B_1^⟂ ∩ B_2} V ψnk
###

@doc raw"""
Get P_{B_1^⟂ ∩ B_2} op P_{B_1} ψk_Ecut2 as a dense matrix, where ψk_Ecut2 is a matrix
with columns in the Ecut basis.
"""
function apply_op_perpb1capb2(opblock, opblock_Ecut2, ψk_Ecut2::AbstractMatrix)
    basis2 = opblock_Ecut2.basis
    kpoint2 = opblock_Ecut2.kpoint
    kpoint = opblock.kpoint
    T = eltype(basis2)
    @assert kpoint2.coordinate == kpoint.coordinate
    @assert size(ψk_Ecut2, 1) == length(G_vectors(kpoint2))

    b1_in_b2 = indexin(G_vectors(kpoint), G_vectors(kpoint2))
    not_in_b1 = [i for i in 1:length(G_vectors(kpoint2)) if !(i in b1_in_b2)]
    @assert length(b1_in_b2) + length(not_in_b1) == length(G_vectors(kpoint2))

    # Compute the part inside B2 explicitly and select
    # only the indices inside B_1^⟂ ∩ B_2.
    (opblock_Ecut2 * ψk_Ecut2)[not_in_b1, :]
end

@doc raw"""
Get P_{B_1^⟂ ∩ B_2} H P_{B_1} ψk_Ecut2 as a dense matrix, where ψk_Ecut2 is a matrix
with columns in the Ecut2 basis.
"""
function apply_hamiltonian_perpb1capb2(hamblock, hamblock_Ecut2, ψk_Ecut2::AbstractMatrix)
    @assert hamblock_Ecut2.basis.Ecut > hamblock.basis.Ecut
    @assert hamblock_Ecut2.kpoint.coordinate == hamblock.kpoint.coordinate
    @assert hamblock.basis.model == hamblock_Ecut2.basis.model
    basis2 = hamblock_Ecut2.basis

    idx_local = only(findall(t -> t isa AtomicLocal, basis2.model.term_types))
    Vloc = hamblock.operators[idx_local]
    Vloc_Ecut2 = hamblock_Ecut2.operators[idx_local]

    if is_cohen_bergstresser(basis2.model)
        @assert basis2.Ecut ≥ minimal_Ecut2(hamblock.basis)
        apply_op_perpb1capb2(Vloc, Vloc_Ecut2, ψk_Ecut2)
    elseif is_linear_atomic(basis2.model)
        idx_nonlocal = only(findall(t -> t isa AtomicNonlocal,
                                    basis2.model.term_types))
        Vnlock = hamblock.operators[idx_nonlocal]
        Vnlock_Ecut2 = hamblock_Ecut2.operators[idx_nonlocal]

        (apply_op_perpb1capb2(Vloc, Vloc_Ecut2, ψk_Ecut2)
         + apply_op_perpb1capb2(Vnlock, Vnlock_Ecut2, ψk_Ecut2))
    else
        error("Not implemented")
    end
end

###
### Bounds for \|P_{B_2} V ψnk\|
###

@doc raw"""
Computes ``\|P_{B_1^⟂ ∩ B_2} op ψnk\|`` for an operator
"""
function bound_perpb1capb2_residual_operator(opblock, opblock_Ecut2, ψnk)
    basis2 = opblock_Ecut2.basis
    kpoint2 = opblock_Ecut2.kpoint
    kpoint = opblock.kpoint
    T = eltype(basis2)
    @assert kpoint2.coordinate == kpoint.coordinate
    @assert length(ψnk) == length(G_vectors(kpoint))

    error("This is currently broken, because the interface " *
          "of apply_op_perpb1capb2 updated")

    # Apply operator to psi and project onto B_1^⟂ ∩ B_2.
    opψnk = apply_op_perpb1capb2(opblock, opblock_Ecut2, reshape(Array(ψnk), :, 1))

    norm(opψnk)  # Compute L2-norm
end

@doc raw"""
Estimates an upper bound for ``\|P_{B_1}^\perp V_loc ψnk\|``

hamblock:        Hamiltonian k-block in small Ecut basis
hamblock_Ecut2:  Hamiltonian k-block in large Ecut basis
ψnk:             Eigenvector for a particular band and k-point on small Ecut
ψnk_Ecut2:       Eigenvector for a particular band and k-point on large Ecut
"""
function bound_perpb1_residual_Vloc(hamblock, hamblock_Ecut2, ψnk)
    basis2 = hamblock_Ecut2.basis
    idx_local = only(findall(t -> t isa AtomicLocal, basis2.model.term_types))
    Vloc = hamblock.operators[idx_local]
    Vloc_Ecut2 = hamblock_Ecut2.operators[idx_local]

    (   bound_perpb1capb2_residual_operator(Vloc, Vloc_Ecut2, ψnk)
      + bound_perpb2_residual_Vloc(hamblock, ψnk, basis2.Ecut))
end

@doc raw"""
Estimates an upper bound for ``\|P_{B_1}^\perp V_nlock ψnk\|``

hamblock:        Hamiltonian k-block in small Ecut basis
hamblock_Ecut2:  Hamiltonian k-block in large Ecut basis
ψnk:             Eigenvector for a particular band and k-point on small Ecut
ψnk_Ecut2:       Eigenvector for a particular band and k-point on large Ecut
"""
function bound_perpb1_residual_Vnlock(hamblock, hamblock_Ecut2, ψnk)
    basis2 = hamblock_Ecut2.basis
    idx_nonlocal = only(findall(t -> t isa AtomicNonlocal, basis2.model.term_types))
    Vnlock = hamblock.operators[idx_nonlocal]
    Vnlock_Ecut2 = hamblock_Ecut2.operators[idx_nonlocal]

    (   bound_perpb1capb2_residual_operator(Vnlock, Vnlock_Ecut2, ψnk)
      + bound_perpb2_residual_Vnlock(hamblock, ψnk, basis2.Ecut))
 end

@doc raw"""
Estimates an upper bound for ``\|P_{B_1}^\perp V ψnk\|`` if ``ψnk`` is given
on the `Ecut` basis.

hamblock:        Hamiltonian k-block in small Ecut basis
hamblock_Ecut2:  Hamiltonian k-block in large Ecut basis
ψnk:             Eigenvector for a particular band and k-point on small Ecut
ψnk_Ecut2:       Eigenvector for a particular band and k-point on large Ecut
"""
function bound_perpb1_residual(hamblock, hamblock_Ecut2, ψnk)
    @assert hamblock_Ecut2.basis.Ecut > hamblock.basis.Ecut
    @assert hamblock_Ecut2.kpoint.coordinate == hamblock.kpoint.coordinate
    @assert hamblock.basis.model == hamblock_Ecut2.basis.model
    basis2 = hamblock_Ecut2.basis

    if is_cohen_bergstresser(basis2.model)
        @assert basis2.Ecut ≥ minimal_Ecut2(hamblock.basis)
        (AtomicLocal=bound_perpb1_residual_Vloc(hamblock, hamblock_Ecut2, ψnk), )
    elseif is_linear_atomic(basis2.model)
        (AtomicLocal=bound_perpb1_residual_Vloc(hamblock, hamblock_Ecut2, ψnk),
         AtomicNonlocal=bound_perpb1_residual_Vnlock(hamblock, hamblock_Ecut2, ψnk))
    else
        error("Not implemented")
    end
end

###
### Bounds for \|V P_{B_1}^⟂\|_op (potential_perpb1)
###

function bound_potential_perpb1_Vloc(hamblock, hamblock_Ecut2)
    basis2 = hamblock_Ecut2.basis
    kpoint2 = hamblock_Ecut2.kpoint
    model = basis2.model
    Ecut2 = basis2.Ecut
    kpoint = hamblock.kpoint
    T = eltype(basis2)
    @assert kpoint.coordinate == kpoint2.coordinate

    accu = zero(T)
    for (element, positions) in model.atoms
        if element isa ElementPsp
            @assert element.psp isa PspHgh

            # We need a bound for Σ_{G, |G| > qcut} |Q(t)| exp(-t^2 / 2) / t^2
            # where t = psp.rloc * |G|. Since |Q| is a decaying function not going
            # through zero after |G| > qcut, we may replace |Q| by Q or -Q depending
            # on the sign at qcut.
            Qloc = DFTK.psp_local_polynomial(T, element.psp)
            qcut = sqrt(2Ecut2)
            rts = [real.(rt) for rt in roots(Qloc) if imag(rt) < 1e-10]
            @assert length(rts) == 0 || maximum(rts) < qcut
            sign(Qloc(qcut)) < 0 && (Qloc = -Qloc)
            R = bound_sum_polynomial_gaussian(Qloc, model.recip_lattice, qcut,
                                              β=element.psp.rloc, k=-2, σ=2)
        else
            @assert element isa ElementCohenBergstresser
            R = 0  # assume Ecut2 to be large enough
        end
        accu += length(positions) * R / model.unit_cell_volume
    end

    # For the remaining term (inside B2 ∩ B1^⟂) we use an estimate precomputed
    # with the function bound_potential_b2_Vloc_b2_infty at a very large value
    # of Ecut for the employed basis in the hamblock. This yields a value for
    # |P_B V P_B |_\intfy norm for a basis B which is guaranteed to be larger than
    # B2 and thus an upper bound to the term |P_{B2} V P_{B2} |_\intfy.
    #
    # NOTE: These values are specific to the lattice and atomic positions
    #       we employ here and are not transferable.
    @assert basis2.Ecut ≤ 5000  # Values precomputed at Ecut=5000
    if is_cohen_bergstresser(model)
        bound_Vinf = 0.6681312576847188
    elseif is_linear_atomic(model)
        bound_Vinf = 7.230556115103383
    else
        error("Not implemented")
    end

    accu + bound_Vinf
end

@doc raw"""
Compute ``\|P_{B} Vloc P_{B}\|_\infty`` for ``B`` being the basis used in `hamblock`.
"""
function bound_potential_b2_Vloc_b2_infty(hamblock)
    basis = hamblock.basis
    model = basis.model
    T = eltype(basis)
    idx_local = only(findall(t -> t isa AtomicLocal, model.term_types))

    # We compute the supremum of the real-space values we know
    # (shifting the potential implicitly, which we know does not change the gap)
    Vloc_extrema = extrema(hamblock.operators[idx_local].potential)
    extent = (Vloc_extrema[2] - Vloc_extrema[1]) / T(2)

    # Then we use a gradient correction for the fact that we do not know the values
    # inside the cell. If g is the gradient vector and δ the cell diameter than the
    # possible addition is <g|d>/2 ≤ ||g|| δ / 2 where d is any lattice vector.
    # In turn we estimate ||g||_2 ≤ ||g||_∞ = ||\hat{g}||_1 i.e. by the l1-norm of
    # the Fourier coefficients of the gradient, which is |G| times the Fourier coefficients
    # of the potential itself.

    # local_potential_fourier is the actual FT of the real-space potential
    potterm(el, r, Gcart) = Complex{T}(DFTK.local_potential_fourier(el, norm(Gcart))
                                       * cis(-dot(Gcart, model.lattice * r)))

    # sqrt(Ω) because of normalised basis used in DFTK
    pot(G) = sum(potterm(elem, r, model.recip_lattice * G) / sqrt(model.unit_cell_volume)
                 for (elem, positions) in model.atoms
                 for r in positions)

    # Another sqrt(Ω) from going to the l1-norm of the Fourier coefficients
    sqrtΩ = sqrt(model.unit_cell_volume)
    sumVderiv = sum(norm(model.recip_lattice * G) * abs(pot(G)) for G in G_vectors(basis))
    diameter = norm(model.lattice * 1 ./ T.(basis.fft_size))
    derivative_term = sumVderiv * diameter / T(2) / sqrtΩ

    extent + derivative_term
end

function bound_potential_perpb1_Vnloc(hamblock, hamblock_Ecut2)
    basis2 = hamblock_Ecut2.basis
    kpoint2 = hamblock_Ecut2.kpoint
    model = basis2.model
    Ecut2 = basis2.Ecut
    kpoint = hamblock.kpoint
    T = eltype(basis2)
    recip_lattice = model.recip_lattice
    @assert any(t isa AtomicLocal for t in model.term_types)
    @assert kpoint.coordinate == kpoint2.coordinate

    # Tabulated values for the L∞ norm of the spherical harmonics Ylm
    Ylm∞ = (sqrt(1 / 4T(π)), sqrt(3 / 4T(π)))  # My guess: (5 / 4π, 7 / 4π)

    # Norms of the G vectors in the Ecut2 basis, ignoring the DC
    Gnorms_Ecut2 = [norm(model.recip_lattice * (G + kpoint.coordinate))
                    for G in G_vectors(kpoint2)]

    # Norms of the G vectors only in Ecut2 basis, but not in Ecut basis
    Gs_complement = setdiff(G_vectors(kpoint2), G_vectors(kpoint))
    Gnorms_complement = [norm(model.recip_lattice * (G + kpoint.coordinate))
                         for G in Gs_complement]

    function norm_projectors(psp, i, l, Nli, Gnorms)
        psp_radials = eval_psp_projection_radial.(psp, i, l, Gnorms)
        sqrt(Nli + sum(abs2, psp_radials)) * Ylm∞[l + 1] / sqrt(model.unit_cell_volume)
    end

    function projector_bound(psp, l)
        @assert psp isa PspHgh
        n_proj_l = size(psp.h[l + 1], 1)

        sum_proj_Ecut2 = zeros(T, n_proj_l)
        sum_proj_complement = zeros(T, n_proj_l)

        for i in 1:n_proj_l
            # Bound for Σ_{G, |G| > qcut} Q^2(t) exp(-t^2) where t = psp.rp[l+1] * |G|.
            Qnloc = DFTK.psp_projection_radial_polynomial(T, psp, i, l)
            qcut = sqrt(2Ecut2) - norm(recip_lattice * kpoint2.coordinate)
            Nli = bound_sum_polynomial_gaussian(Qnloc * Qnloc, recip_lattice,
                                                qcut, β=psp.rp[l + 1])

            sum_proj_Ecut2[i] = norm_projectors(psp, i, l, Nli, Gnorms_Ecut2)
            sum_proj_complement[i] = norm_projectors(psp, i, l, Nli, Gnorms_complement)
        end

        sum_proj_Ecut2' * abs.(psp.h[l + 1]) * sum_proj_complement
    end

    @assert all(element isa ElementPsp for (element, _) in model.atoms)
    sum(maximum(projector_bound.(element.psp, 0:element.psp.lmax)) * length(positions)
        for (element, positions) in model.atoms)
end

@doc raw"""
Estimates a (rough) upper bound for ``\|V φ\|`` for any ``φ`` outside `Ecut`
and `V` are the potential terms of the hamiltonian `hamblock` given
on the `Ecut2` basis.

It is assumed that potential terms are decaying functions in Fourier
space beyond `Ecut2`.

hamblock:        Hamiltonian k-block in Ecut basis
hamblock_Ecut2:  Hamiltonian k-block in Ecut2 basis
"""
function bound_potential_perpb1(hamblock, hamblock_Ecut2)
    @assert hamblock_Ecut2.basis.Ecut > hamblock.basis.Ecut
    @assert hamblock_Ecut2.kpoint.coordinate == hamblock.kpoint.coordinate
    @assert hamblock.basis.model == hamblock_Ecut2.basis.model
    basis2 = hamblock_Ecut2.basis

    if is_cohen_bergstresser(basis2.model)
        @assert hamblock_Ecut2.basis.Ecut ≥ minimal_Ecut2(hamblock.basis)
        (AtomicLocal=bound_potential_perpb1_Vloc(hamblock, hamblock_Ecut2), )
    elseif is_linear_atomic(basis2.model)
        (AtomicLocal=bound_potential_perpb1_Vloc(hamblock, hamblock_Ecut2),
         AtomicNonlocal=bound_potential_perpb1_Vnloc(hamblock, hamblock_Ecut2))
    else
        error("Not implemented")
    end
end

###
### Estimating eigenvalue gaps
###

"""
Estimate the gap using a Schur complement.

hamblock:        Hamiltonian k-block in Ecut basis
hamblock_Ecut2:  Hamiltonian k-block in Ecut2 basis
λtilde:          Approximate eigenvalues of hamblock for which gaps to the respective
                 next eigenvalues are estimated.
ψk_Ecut2:        Eigenvectors for this k-point in Ecut2 basis
ngap:            Number of eigenvalues for which to estimate the gap
                 (Default: length(λtilde) - 1)
fast:            Use a slightly faster algorithm relying on (λtilde, ψk_Ecut2) being
                 the exact lower end of the spectrum of hamblock.
"""
function estimate_gaps_schur(hamblock, hamblock_Ecut2, λtilde, ψk_Ecut2; fast=false,
                             ngap=length(λtilde) - 1, λtol=1e-6, full_block=false)
    # Use a Schur complement to estimate the gap. If fast is true gap might not be exact.
    #
    # ham is the Hamiltonian in the Ecut basis we use.
    # This allows to decompose the (exact) shifted Hamiltonian into blocks:
    #     H = ( ham   W       )
    #         ( W†    hamperp )
    # where in a Fourier basis W only depends on the potential V. The eigenvalues
    # of this operator are
    #     λ1 ≤ λ2 ≤ ... ≤ λn
    # An upper bound to λ1 is λtilde[1]. To estimate the gap δ1 between λ2 and λ1
    # we need a lower bound to λ2. This can be estimated by that root
    # of the function
    #    f(μ) = ||Vφ||^2 / min(|λtilde .- μ|) + Ecut - ||Vφ|| - μ
    # which is closest to λtilde[2]. In this φ is normalised and has only support
    # beyond Ecut). This also assumes that the eigenvalues are all of single
    # multiplicity.
    #
    # Denoting this root as μ[1] an estimate for the gap of each eigenvalue
    # is the difference (μ[1] - λtilde[1])
    @assert hamblock_Ecut2.basis.Ecut > hamblock.basis.Ecut
    @assert hamblock_Ecut2.kpoint.coordinate == hamblock.kpoint.coordinate
    @assert ngap ≤ length(λtilde) - 1
    @assert length(λtilde) == size(ψk_Ecut2, 2)
    @assert length(G_vectors(hamblock_Ecut2.kpoint)) == size(ψk_Ecut2, 1)

    basis = hamblock.basis
    T = eltype(basis)
    Ecut = basis.Ecut
    nbas = size(hamblock, 2)
    gap = zero(λtilde)
    M = length(λtilde) - 1

    if full_block
        @timeit to "Full diagonalisation with evecs" begin
            # Don't trust anything. Fully diagonalise hamblock:
            λtilde, Xtest = eigen!(Hermitian(Array(hamblock)))
            @assert length(basis.kpoints) == 1  # Only implemented for 1-kpt case!
            ψk_Ecut2 = DFTK.interpolate_blochwave([Xtest], hamblock.basis,
                                                   hamblock_Ecut2.basis)[1]
            Xtest = nothing
            M = length(λtilde)
        end
    elseif !fast
        # Check approximation for eigenpairs is actually correct
        if T in (Float32, Float64)
            @timeit to "Partial diagonalisation" begin
                λtest = eigvals(Hermitian(Array(hamblock)), 1:length(λtilde) + 15)
                λtest = λtest[1:length(λtilde)]
            end
        else
            @timeit to "Full diagonalisation" begin
                λtest = eigvals(Hermitian(Array(hamblock)))[1:length(λtilde)]
            end
        end
        @assert maximum(abs, λtilde - λtest) < max(10λtol, 1e-16)
    end

    @timeit to "Compute terms estimating potential coupling of B_1 and B_1^⟂" begin
        # Compute upper bound for ||P_{B_1} V P_{B_1^⟂} ||_op  (just use ||V P_{B_1^⟂}||_op)
        Vop_b1_b1perp = sum(bound_potential_perpb1(hamblock, hamblock_Ecut2))

        # Compute upper bound for || P_{B_1^⟂} V P_{B_1^⟂} ||_op  (just use ||V P_{B_1^⟂}||_op)
        Vop_b1perp_b1perp = Vop_b1_b1perp

        # Compute upper bound for || P_B_2^⟂ V P_{B_1} ||_op
        Vop_b1_b2perp = sum(bound_ham_perpb2_b1(hamblock, hamblock_Ecut2.basis.Ecut))

        # Compute P_{B_1^⟂ ∩ B_2} H P_{B_1} ψk_Ecut2
        Hψk = apply_hamiltonian_perpb1capb2(hamblock, hamblock_Ecut2,
                                            @view ψk_Ecut2[:, 1:M])
    end

    for iλ in 1:ngap
        λn = λtilde[iλ]
        λnext = λtilde[iλ + 1]

        @timeit to "QR factorisation of Hψk_iλ" begin
            Hψk_iλ = @view Hψk[:, iλ+1:M]
            HψkR = qr(Hψk_iλ).R
            @assert size(HψkR, 1) == length(iλ+1:M)
            @assert size(HψkR, 2) == length(iλ+1:M)
            Λ = @view λtilde[iλ+1:M]
        end

        @timeit to "Bisection to find best μ" begin
            # The condition for the best μ is S_μ == 0 or (in our approximation)
            #   0 == Ecut - ||P_{B_1^⟂} V P_{B_1^⟂}||_op - μ
            #        - ||R (Λ - μI)^{-1} R'||_op
            #        - 2 ||R (Λ - μI)^{-1}||_op ||P_{B_1} V P_{B_2^⟂}||_op
            #        - ||P_{B_1} V P_{B_2^⟂}||_op^2 / (λ_{iλ+1} - μ)
            #        - ||P_{B_1} V P_{B_1^⟂}||_op^2 / (λ_{M+1} - μ)
            function Sμ(μn)
                accu = Ecut - Vop_b1perp_b1perp - μn
                accu += -opnorm(HψkR * Diagonal(1 ./ (Λ .- μn)) * HψkR')
                accu += -2opnorm(HψkR * Diagonal(1 ./ (Λ .- μn))) * Vop_b1_b2perp
                accu += -Vop_b1_b2perp^2 / (λnext - μn)
                if M != nbas
                    accu += -Vop_b1_b1perp^2 / (λtilde[M+1] - μn)
                end
                accu
            end

            # Use Bisection to find the best μ
            ε = max(λtol, 1000eps(T))
            μn = λn
            if sign(Sμ(λn + ε)) < 0
                # If this happens, we cannot find a suitable zero => zero gap
            elseif sign(Sμ(λnext - ε)) < 0
                try
                    μn = find_zero(Sμ, (λn + ε, λnext - ε), Bisection(), xatol=ε)
                    μn = max(λn, μn - ε)  # To be on the safe side
                catch e
                    e isa Roots.ConvergenceFailed || rethrow()
                end
            else  # sign(Sμ(λnext - ε)) ≥ 0 && sign(Sμ(λn + ε)) ≥ 0
                μn = λnext - ε  # Stick to λnext - ε
            end

            # Check that `μn` is in [λtilde[iλ], λtilde[iλ+1]]
            # and gives rise to an (almost) positive Sμ:
            if !(λn ≤ μn ≤ λnext) || !isfinite(Sμ(μn)) || Sμ(μn) < -ε
                gap[iλ] = 0.0
                μn = λtilde[iλ]
            else
                gap[iλ] = μn - λn
            end
        end
    end

    gap
end

###
### Model runners producing data
###

function run_model(model; Ecut=10, tol=1e-6, n_bands=8, only_gamma=false,
                   kline_density=10, kwargs...)
    basis = PlaneWaveBasis(model, Ecut)
    ham = Hamiltonian(basis)

    # Diagonalise (at the Gamma point) to get εF
    @timeit to "Diagonalise Gamma point" begin
        ham = Hamiltonian(basis)
        eigensolver(args...; kwargs...) = DFTK.lobpcg_hyper(args...; kwargs...,
                                                            display_progress=true)
        eigres = diagonalise_all_kblocks(eigensolver, ham, n_bands; tol=tol, kwargs...)
        εF = find_fermi_level(basis, eigres.λ)
    end

    ret_common = (basis=basis, εF=εF, Ecut=Ecut, ham=ham, tol=tol)
    only_gamma && return merge(eigres, ret_common)

    # Compute kpath and band data
    @timeit to "Compute bands" begin
        kpath = DFTK.high_symmetry_kpath(basis.model; kline_density=kline_density)
        ρ0 = guess_density(basis)
        data = DFTK.compute_bands(basis, ρ0, kpath.kcoords, n_bands, tol=tol)
    end

    merge(ret_common, data, (kpath=kpath, ))
end

function model_cohen_bergstresser(;T=Float64)
    Si = ElementCohenBergstresser(:Si)
    a = Si.lattice_constant
    lattice = a / 2 .* [[0 1 1.]; [1 0 1.]; [1 1 0.]]
    atoms = [Si => [ones(3)/8, -ones(3)/8]]
    Model(Matrix{T}(lattice); atoms=atoms, terms=[Kinetic(), AtomicLocal()])
end

function model_linear_silicon(;T=Float64)
    a = 10.263141334305942  # Silicon lattice constant in Bohr
    lattice = a / 2 .* [[0 1 1.]; [1 0 1.]; [1 1 0.]]
    Si = ElementPsp(:Si, psp=load_psp("hgh/lda/si-q4"))
    atoms = [Si => [ones(3)/8, -ones(3)/8]]
    model_atomic(lattice, atoms)
end

function run_cohen_bergstresser(;T=Float64, kwargs...)
    data = run_model(model_cohen_bergstresser(T=T); kwargs...)
    merge(data, (system="Cohen-Bergstresser", ))
end

function run_linear_silicon(;T=Float64, kwargs...)
    data = run_model(model_linear_silicon(T=T); kwargs...)
    merge(data, (system="Linear Silicon", ))
end

###
### Error estimator functions for data
###

"""
Estimate the discretisation error for the passed data.
Multiple error estimates are used (Bauer-Fike and Kato-Temple).
Ecut2 is the Ecut for the bigger grid used for the residual computation
If nothing, a value will be derived automatically.
If `fast` is true a faster, but not necessarily exact algorithm is chosen to
estimate the eigenvalue gap.
"""
function estimate_discretisation_error(data; Ecut2=nothing, fast=false,
                                       ngap=nothing, full_block=false)
    @timeit to "Estimate discretisation error" begin
        basis = data.basis
        minEcut2 = minimal_Ecut2(basis)

        if isnothing(Ecut2)
            if is_cohen_bergstresser(basis.model)
                Ecut2 = ceil(minEcut2)
            else
                # Make the Ecut2 a bit larger than the minimal value needed
                # TODO No clue if this is the best way to do it, but it works for us.
                # Maybe it makes more sense to add a constant onto the minimal_Ecut2?
                Ecut2 = ceil(4minEcut2)
            end
        end

        # Compute part of the residual in bigger grid
        @timeit to "Compute residual part on Ecut2" begin
            kcoords = [kpt.coordinate for kpt in basis.kpoints]
            basis_big = PlaneWaveBasis(basis.model, Ecut2, kcoords, basis.ksymops)
            X_big = DFTK.interpolate_blochwave(data.X, basis, basis_big)
            ham_big = Hamiltonian(basis_big)
            # TEMPORARY
            bound_potential_b2_Vloc_b2_infty(ham_big.blocks[1])
            residual_Ecut2 = [ham_big.blocks[ik] * X_big[ik] - X_big[ik] * Diagonal(data.λ[ik])
                              for ik in 1:length(basis.kpoints)]
            residual_norm_Ecut2 = [norm.(eachcol(resk)) for resk in residual_Ecut2]
        end

        # Compute upper bound for the part of the residual we miss
        @timeit to "Estimate residual part beyond Ecut2" begin
            ham = Hamiltonian(basis)
            residual_norm_terms = [[bound_perpb2_residual(ham.blocks[ik], ψnk, Ecut2)
                                    for ψnk in eachcol(data.X[ik])]
                                   for ik in 1:length(basis.kpoints)]
            n_bands = length(residual_norm_terms[1])
            residual_norm_terms = Dict(
                key => [[residual_norm_terms[ik][iband][key] for iband in 1:n_bands]
                        for ik in 1:length(basis.kpoints)]
                for key in keys(residual_norm_terms[1][1])
            )
            residual_norm_terms = merge((computed=residual_norm_Ecut2, ),
                                        (; residual_norm_terms...))
            residual_norm = sum(residual_norm_terms)
        end

        # Estimate eigenvalue gaps
        @timeit to "Estimate gap (fast = $fast)" begin
            # gap[i] = estimated gap between λ[i] and λ[i+1]
            # mingap[i] = min(gap[i-1], gap[i])

            ham = Hamiltonian(data.basis)
            gap = [estimate_gaps_schur(ham.blocks[ik], ham_big.blocks[ik],
                                       data.λ[ik], X_big[ik], fast=fast,
                                       λtol=data.tol, full_block=full_block)
                   for ik in 1:length(basis.kpoints)]

            mingap = deepcopy(gap)
            for ik = 1:length(basis.kpoints)
                mingap[ik][1] = gap[ik][1]
                for iband = 2:length(gap[ik])
                    mingap[ik][iband] = min(gap[ik][iband - 1], gap[ik][iband])
                end
            end
        end

        # Bauer-Fike error estimate
        bound_bauer_fike = residual_norm

        # Kato-Temple error estimate
        bound_kato_temple = [residual_norm[ik].^2 ./ mingap[ik]
                             for ik in 1:length(basis.kpoints)]

        (gap=gap, mingap=mingap, bound_kato_temple=bound_kato_temple, Ecut2=Ecut2,
         bound_bauer_fike=bound_bauer_fike, residual_norm=residual_norm,
         residual_Ecut2=residual_Ecut2, residual_norm_terms=residual_norm_terms,
         exact_bounds=!fast, minEcut2=minEcut2)
    end
end

function estimate_arithmetic_error(data; n_bands=nothing)
    @timeit to "Estimate arithmetic error" begin
        T = eltype(data.basis)
        IntT = Interval{T}
        modelT = data.basis.model

        if isnothing(n_bands)
            n_bands = length(data.λ[1])
        end

        println("Interval-Hamiltonian setup ... this will take some time ...")
        @timeit to  "Hamiltonian setup" begin
            # TODO Would need to copy temperature, spin_polarisation etc ...
            model = Model(Matrix{IntT}(modelT.lattice); atoms=modelT.atoms, terms=modelT.term_types)
            kcoords = [kpt.coordinate for kpt in data.basis.kpoints]
            basis = PlaneWaveBasis(model, data.Ecut, kcoords, data.basis.ksymops,
                                   data.basis.kweights)
            ham = Hamiltonian(basis)
        end
        println("Interval-Hamiltonian setup done!")

        @timeit to "Convert eigenpairs to intervals"  begin
            # Get eigenvalues / eigenvector coefficients as intervals
            eigvals = [Array{IntT}(λk[1:n_bands]) for λk in data.λ]
            eigvecs = [Array{Complex{IntT}}(ψk[:, 1:n_bands]) for ψk in data.X]
        end

        @timeit to  "Compute residuals in IntervalArithmetic" begin
            residual_norm = similar(eigvals)
            @showprogress 0.5 "Compute Interval residuals: " for ik in 1:length(basis.kpoints)
                # Form Ritz values via Rayleigh quotient in interval arithmetic
                Λks = eigvecs[ik]' * (ham.blocks[ik] * eigvecs[ik])
                residual_norm[ik] = norm.(eachcol(Λks - Diagonal(eigvals[ik])))
            end
        end

        # The error of the algorithm (e.g. not iterated till full numerical convergence)
        # is (roughly) the midpoint of the residual error. This is not fully correct
        # as the interval arithmetic might capture errors unevenly: E.g. the residual by
        # the numerical procedure could be 1e-10 and interval arithmetic could report the
        # interval [0 1e-8] in which case this routine would give 5e-7, whereas the
        # actual answer is 1e-10. This fine because we will only deal with cases
        # where the numerical error is *much* smaller than the residual error,
        # such that in practice only intervals like [1e-10 - 1e-14, 1e-10 + 1e-14]
        # will result here. But this should be changed for a more general implementation.
        error_algorithm = [mid.(residual_norm[ik]) for ik in 1:length(basis.kpoints)]

        # The difference between the supremum of the residual_norm and the
        # error_algorithm is the arithmetic error. This is also the radius of the interval.
        error_arithmetic = [radius.(residual_norm[ik]) for ik in 1:length(basis.kpoints)]

        if n_bands != length(data.λ[1])
            for ik in 1:length(basis.kpoints)
                append!(error_algorithm[ik], zeros(length(data.λ[1]) - n_bands))
                append!(error_arithmetic[ik], zeros(length(data.λ[1]) - n_bands))
            end
        end
        (error_algorithm=error_algorithm, error_arithmetic=error_arithmetic)
    end
end

###
### Storing results
###

function store_results(h5name, data; T=Float64, subdir=nothing, n_bands=nothing)
    isfile(h5name) || close(h5open(h5name, "w"))

    if isnothing(n_bands)  # Determine number of bands to store
        n_bands = length(data.λ[1])
    end

    h5open(h5name, "r+") do h5f
        if isnothing(subdir)
            h5gr = h5f
        elseif exists(h5f, subdir)
            h5gr = g_open(h5f, subdir)
        else
            h5gr = g_create(h5f, subdir)
        end

        if exists(h5gr, "System")
            @assert read(h5gr, "System") == data.system
        else
            h5gr["System"] = data.system
        end

        Ecut = round(data.Ecut, digits=3)
        if exists(h5gr, "Ecuts")
            Ecuts = append!(h5gr["Ecuts"][:], Ecut)
            o_delete(h5gr, "Ecuts")
            h5gr["Ecuts"] = Ecuts
        else
            h5gr["Ecuts"] = Float64[Ecut]
        end

        sEcut = (@sprintf "%.3f" Ecut)
        if hasproperty(data, :kpath)
            for label in keys(data.kpath.klabels)
                h5gr["$sEcut/klabels/$label"] = data.kpath.klabels[label]
            end
        end
        h5gr["$sEcut/n_kpoints"] = length(data.basis.kpoints)

        for symbol_float in (:Ecut2, :tol, :Ecut, :εF, :minEcut2)
            if hasproperty(data, symbol_float)
                h5gr["$sEcut/" * string(symbol_float)] =  T(getproperty(data, symbol_float))
            end
        end
        for symbol in (:exact_bounds, )
            if hasproperty(data, symbol)
                h5gr["$sEcut/" * string(symbol)] = getproperty(data, symbol)
            end
        end

        for ik in 1:length(data.basis.kpoints)
            exists(h5gr, "$sEcut/kpt_$ik") && o_delete(h5gr, "$sEcut/kpt_$ik")
            h5kpt = g_create(h5gr, "$sEcut/kpt_$ik")

            h5kpt["coordinate"] = Vector{T}(data.basis.kpoints[ik].coordinate)

            if hasproperty(data, :λ)
                h5kpt["eigenvalues"] = Array{T}(data.λ[ik][1:n_bands])
                h5kpt["eigenvalues_string"] = join(string.(data.λ[ik][1:n_bands]), ",")
            end

            for symbol in (:gap, :mingap, :bound_kato_temple, :bound_bauer_fike,
                           :residual_norm, :error_algorithm, :error_arithmetic)
                if hasproperty(data, symbol)
                    h5kpt[string(symbol)] = (
                        Array{T}(getproperty(data, symbol)[ik][1:n_bands])
                    )
                end
            end

            if hasproperty(data, :residual_norm_terms)
                for key in keys(data.residual_norm_terms)
                    h5kpt["residual_norm_term_" * string(key)] = (
                        Array{T}(getproperty(data.residual_norm_terms, key)[ik][1:n_bands])
                    )
                end
            end
        end
    end

    nothing
end

function load_results(h5name; has_subdirs=false)
    function parse_subdir(h5gr)
        system = read(h5gr, "System")
        Ecuts = read(h5gr, "Ecuts")

        results = []
        for Ecut in Ecuts
            sEcut = @sprintf "%.3f" Ecut
            result = Dict{Symbol, Any}()
            result[:n_kpoints] = read(h5gr, "$sEcut/n_kpoints")

            if "klabels" in names(h5gr["$sEcut"])
                result[:klabels] = Dict(label => h5gr["$sEcut/klabels/$label"][:]
                                        for label in names(h5gr["$sEcut/klabels"]))
            end
            for symbol in (:Ecut_big, :exact_bounds, :tol, :Ecut, :εF, :Ecut2, :minEcut2)
                if string(symbol) in names(h5gr["$sEcut"])
                    result[symbol] = read(h5gr, "$sEcut/" * string(symbol))
                end
            end

            ksymbols = Vector{Symbol}()
            for symbol in (:gap, :mingap, :bound_kato_temple, :bound_bauer_fike,
                           :residual_norm, :error_algorithm, :error_arithmetic,
                           :residual_norm_term_computed,
                           :residual_norm_term_AtomicLocal,
                           :residual_norm_term_AtomicNonlocal)
                if string(symbol) in names(h5gr["$sEcut/kpt_1"])
                    push!(ksymbols, symbol)
                    result[symbol] = Vector{Vector{Float64}}()
                end
            end

            result[:λ] = Vector{Vector{Double64}}()
            result[:kcoords] = Vector{Vector{Float64}}()
            for ik in 1:result[:n_kpoints]
                h5kpt = h5gr["$sEcut/kpt_$ik"]

                push!(result[:kcoords], read(h5kpt, "coordinate"))
                push!(result[:λ], parse.(Double64,
                                         split(read(h5kpt, "eigenvalues_string"), ",")))
                for symbol in ksymbols
                    push!(result[symbol], read(h5kpt, string(symbol)))
                end
            end
            push!(results, result)
        end
        (Ecuts=Ecuts, results=results, system=system)
    end

    h5open(h5name, "r") do h5f
        if has_subdirs
            results = Dict{String, Any}()
            for subdir in names(h5f)
                results[subdir] = parse_subdir(h5f[subdir])
            end
            return results
        else
            return parse_subdir(h5f)
        end
    end
end


text:
# helper
function runcommand(cmd)
    @info cmd
    run(cmd)
end

# `env` is the environment directory
function generate(env; execute=true, pluto=true)
    quote
        using Pkg
        Pkg.activate(temp=true)
        Pkg.add(name="Literate", rev="fe/pluto")
        using Literate

        const ENVDIR = $env
        const OUTDIR = joinpath(ENVDIR, "..")
        const INFILE = joinpath(OUTDIR, "notebook.jl")

        # generate pluto notebook:
        if $pluto
            Literate.notebook(INFILE, ENVDIR, flavor=Literate.PlutoFlavor())
            runcommand(`mv $ENVDIR/notebook.jl $OUTDIR/notebook.pluto.jl`)
        end

        Literate.notebook(INFILE, OUTDIR, execute=false)
        runcommand(
            `mv $OUTDIR/notebook.ipynb $OUTDIR/notebook.unexecuted.ipynb`)
        $execute && Literate.notebook(INFILE, OUTDIR, execute=true)

    end |> eval
end

# Pkg.add("Pluto")
# using Pluto
# Pluto.run(notebook=joinpath(OUTDIR, "notebook.pluto.jl"))

# Pkg.add("IJulia")
# Pkg.instantiate()
# using IJulia
# IJulia.notebook(dir=OUTDIR)

ext: jl


text:
using DataFrames
using Distributions
using NearestNeighbors
using StatsBase
using StaticArrays

using Base.Threads
using Dates: now, DateTime

import LightGraphs

function drop_unused_components!(data::BmmData; min_n_samples::Int=2)
    non_noise_ids = findall(data.assignment .> 0)
    n_mols_per_cell = num_of_molecules_per_cell(data) # Can't use comp.n_samples because it's updated only on the maximization step
    existed_ids = findall(n_mols_per_cell .>= min_n_samples)
    id_map = zeros(Int, length(data.components))
    id_map[existed_ids] .= 1:length(existed_ids)

    data.assignment[non_noise_ids] .= id_map[data.assignment[non_noise_ids]]
    data.components = data.components[existed_ids]
end

adjacent_component_ids(assignment::Array{Int, 1}, adjacent_points::Array{Int, 1})::Array{Int, 1} =
    [x for x in Set(assignment[adjacent_points]) if x > 0]

function assign_molecule!(data::BmmData, mol_id::Int; denses::Vector{Float64}, adj_classes::Vector{Int}, stochastic::Bool=true, noise_density_threshold::Float64=1e-100)
    if sum(denses) < noise_density_threshold
        assign!(data, mol_id, 0) # Noise class
    elseif !stochastic
        assign!(data, mol_id, adj_classes[findmax(denses)[2]])
    else
        assign!(data, mol_id, fsample(adj_classes, denses))
    end
end

"""
...
# Arguments
- `assignment::Array{Int, 1}`:
- `adjacent_points::Array{Int, 1}`:
- `adjacent_weights::Array{Float64, 1}`: weights here mean `1 / distance` between two adjacent points
"""
@inline function aggregate_adjacent_component_weights!(comp_ids::Vector{Int}, comp_weights::Vector{Float64}, component_weights::Dict{Int, Float64},
        assignment::Vector{Int}, adjacent_points::Vector{Int}, adjacent_weights::Vector{Float64})
    empty!(component_weights)
    empty!(comp_weights)
    empty!(comp_ids)
    zero_comp_weight = 0.0

    @inbounds @simd for i in 1:length(adjacent_weights)
        c_point = adjacent_points[i]
        c_id = assignment[c_point]
        cw = adjacent_weights[i]
        if c_id == 0
            zero_comp_weight += cw
        else
            component_weights[c_id] = get(component_weights, c_id, 0.0) + cw
        end
    end

    for (k, v) in component_weights
        push!(comp_ids, k)
        push!(comp_weights, v)
    end

    return zero_comp_weight
end

@inline function fill_adjacent_component_weights!(adj_classes::Vector{Int}, adj_weights::Vector{Float64}, data::BmmData, mol_id::Int;
        component_weights::Dict{Int, Float64}, adj_classes_global::Dict{Int, Vector{Int}})
    # Looks like it's impossible to optimize further, even with vectorization. It means that creating vectorized version of expect_dirichlet_spatial makes few sense
    zero_comp_weight = aggregate_adjacent_component_weights!(adj_classes, adj_weights, component_weights, data.assignment,
        data.adjacent_points[mol_id], data.adjacent_weights[mol_id])

    if mol_id in keys(adj_classes_global)
        n1 = length(adj_classes)
        append!(adj_classes, adj_classes_global[mol_id])
        append!(adj_weights, ones(length(adj_classes) - n1) .* data.real_edge_weight)
    end

    return zero_comp_weight
end

function adjust_densities_by_prior_segmentation!(denses::Vector{Float64}, segment_id::Int, largest_cell_id::Int,
        data::BmmData, adj_classes::Vector{Int}, adj_weights::Vector{Float64})
    seg_prior_pow = data.prior_seg_confidence * exp(3*data.prior_seg_confidence)
    seg_size = data.n_molecules_per_segment[segment_id]
    largest_cell_size = min(get(data.components[largest_cell_id].n_molecules_per_segment, segment_id, 0) + 1, seg_size)

    for j in eachindex(adj_weights)
        c_adj = adj_classes[j]
        if c_adj == largest_cell_id
            continue
        end

        main_seg = data.main_segment_per_cell[c_adj]
        n_cell_mols_per_seg = min(get(data.components[c_adj].n_molecules_per_segment, segment_id, 0) + 1, seg_size)

        if (segment_id == main_seg) || (main_seg == 0) # Part against over-segmentation
            denses[j] *= (1. - data.prior_seg_confidence)^0.5 * (n_cell_mols_per_seg / largest_cell_size)^data.prior_seg_confidence
        else # Part against under-segmentation and overlap
            denses[j] *= (1. - data.prior_seg_confidence)^0.5 * (1. - n_cell_mols_per_seg / seg_size)^seg_prior_pow
        end
    end
end

function expect_density_for_molecule!(denses::Vector{Float64}, data::BmmData{N}, mol_id::Int;
        zero_comp_weight::Float64, adj_classes::Vector{Int}, adj_weights::Vector{Float64}) where N
    x = SVector{N}(position_data(data)[:,mol_id]...)
    gene::Union{Int, Missing} = composition_data(data)[mol_id]
    confidence::Float64 = data.confidence[mol_id]
    mol_cluster::Int = get(data.cluster_per_molecule, mol_id, 0)
    segment_id::Int = get(data.segment_per_molecule, mol_id, 0)

    empty!(denses)
    largest_cell_id = 0
    largest_cell_size = 0
    for j in eachindex(adj_weights)
        c_adj = adj_classes[j]
        cc = data.components[c_adj]
        c_dens = confidence * exp(data.mrf_strength * adj_weights[j]) * pdf(cc, x, gene, use_smoothing=data.use_gene_smoothing)

        ## Only if molecule clustering is provided

        if (c_adj > 0) && (c_adj < length(data.cluster_per_cell)) && (data.cluster_per_cell[c_adj] != mol_cluster)
            c_dens *= data.cluster_penalty_mult
        end

        if (segment_id > 0) && (data.main_segment_per_cell[c_adj] in (segment_id, 0))
            # Find an adjacent cell with the largest number of molecules per cell after assigning the current point to it
            cur_mols_per_seg = min(get(cc.n_molecules_per_segment, segment_id, 0) + 1, data.n_molecules_per_segment[segment_id])
            if (cur_mols_per_seg > largest_cell_size) || ((cur_mols_per_seg == largest_cell_size) && (cc.n_samples > data.components[largest_cell_id].n_samples))
                largest_cell_size = cur_mols_per_seg
                largest_cell_id = c_adj
            end
        end

        push!(denses, c_dens)
    end

    if (segment_id > 0) & (largest_cell_id > 0)
        adjust_densities_by_prior_segmentation!(denses, segment_id, largest_cell_id, data, adj_classes, adj_weights)
    end

    if confidence < 1.0
        push!(denses, (1 - confidence) * exp(data.mrf_strength * fmax(data.real_edge_weight, zero_comp_weight)) * data.noise_density)
        push!(adj_classes, 0)
    end

    return largest_cell_id
end

function expect_dirichlet_spatial!(data::BmmData; stochastic::Bool=true)
    component_weights = Dict{Int, Float64}();
    adj_classes = Int[]
    adj_weights = Float64[]
    denses = Float64[]

    adj_classes_global = get_global_adjacent_classes(data)

    for i in 1:size(data.x, 1)
        zero_comp_weight = fill_adjacent_component_weights!(adj_classes, adj_weights, data, i;
            component_weights=component_weights, adj_classes_global=adj_classes_global)

        expect_density_for_molecule!(denses, data, i; adj_classes=adj_classes, adj_weights=adj_weights, zero_comp_weight=zero_comp_weight)

        assign_molecule!(data, i; denses=denses, adj_classes=adj_classes, stochastic=stochastic)
    end
end

function update_prior_probabilities!(components::Array{<:Component, 1}, new_component_weight::Float64)
    c_weights = [max(c.n_samples, new_component_weight) for c in components]
    prior_probs = rand(Distributions.Dirichlet(c_weights))
    # t_mmc = 150.0; # TODO: finish the idea
    # t_dist = Normal(t_mmc, t_mmc * 3)
    # prior_probs = pdf.(t_dist, c_weights)
    # prior_probs[c_weights .< t_mmc] .= c_weights[c_weights .< t_mmc] ./ t_mmc .* pdf(t_dist, t_mmc)

    for (c, p) in zip(components, prior_probs)
        c.prior_probability = p
    end
end

function maximize!(data::BmmData)
    ids_by_assignment = split_ids(data.assignment .+ 1, max_factor=length(data.components)+1)[2:end]

    @inbounds @views for i in 1:length(data.components)
        p_ids = ids_by_assignment[i]
        nuc_probs = isempty(data.nuclei_prob_per_molecule) ? nothing : data.nuclei_prob_per_molecule[p_ids]
        maximize!(data.components[i], position_data(data)[:, p_ids], composition_data(data)[p_ids];
            nuclei_probs=nuc_probs, min_nuclei_frac=data.min_nuclei_frac)
    end

    data.noise_density = estimate_noise_density_level(data)
    if isinf(data.noise_density)
        error("Infinity noise density")
    end

    if length(data.cluster_per_molecule) > 0
        data.cluster_per_cell = [isempty(ids) ? 0 : mode(data.cluster_per_molecule[ids]) for ids in ids_by_assignment]
    end
end

function noise_composition_density(data::BmmData)::Float64
#     mean([mean(c.composition_params.counts[c.composition_params.counts .> 0] ./ c.composition_params.sum_counts) for c in data.components]);
    acc = 0.0 # Equivalent to the above commented expression if sum_counts == sum(counts)
    n_comps = 0.0
    for (ci, c) in enumerate(data.components)
        if c.composition_params.sum_counts < 1e-3
            continue
        end

        inn_acc = 0
        for v in c.composition_params.counts
            inn_acc += (v > 0)
        end

        acc += 1.0 / inn_acc
        n_comps += 1.0
    end

    return acc / fmax(n_comps, 1.0)
end

function estimate_noise_density_level(data::BmmData)::Float64
    composition_density = noise_composition_density(data)

    std_vals = data.distribution_sampler.shape_prior.std_values;
    position_density = pdf(MultivariateNormal(zeros(length(std_vals)), diagm(0 => std_vals.^2)), 3 .* std_vals)

    return position_density * composition_density
end

append_empty_component!(data::BmmData) =
    push!(data.components, sample_distribution!(data))[end]

function append_empty_components!(data::BmmData, new_component_frac::Float64)
    for i in 1:round(Int, new_component_frac * length(data.components))
        append_empty_component!(data)
    end
end

function get_global_adjacent_classes(data::BmmData)::Dict{Int, Vector{Int}}
    adj_classes_global = Dict{Int, Vector{Int}}()
    for (cur_id, comp) in enumerate(data.components)
        if comp.n_samples > 0
            continue
        end

        nearest_id = knn(data.position_knn_tree, comp.position_params.μ, 1)[1][1]
        for t_id in data.adjacent_points[nearest_id]
            if t_id in keys(adj_classes_global)
                push!(adj_classes_global[t_id], cur_id)
            else
                adj_classes_global[t_id] = [cur_id]
            end
        end

        if nearest_id in keys(adj_classes_global)
            push!(adj_classes_global[nearest_id], cur_id)
        else
            adj_classes_global[nearest_id] = [cur_id]
        end
    end

    return adj_classes_global
end

function build_cell_graph(assignment::Vector{Int}, adjacent_points::Array{Vector{Int}, 1}, mol_ids::Vector{Int}, cell_id::Int; confidence::Union{Vector{Float64}, Nothing}=nothing, confidence_threshold::Float64=0.5)
    cur_adj_point = Vector{Int}[]
    if confidence === nothing
        cur_adj_point = filter.(p -> assignment[p] == cell_id, adjacent_points[mol_ids])
    else
        mol_ids = mol_ids[confidence[mol_ids] .> confidence_threshold]
        cur_adj_point = filter.(p -> (assignment[p] == cell_id) & (confidence[p] .> confidence_threshold), view(adjacent_points, mol_ids))
    end

    vert_id_per_mol_id = Dict(mi => vi for (vi, mi) in enumerate(mol_ids))

    for points in cur_adj_point
        for j in 1:length(points)
            points[j] = vert_id_per_mol_id[points[j]]
        end
    end

    neg = LightGraphs.SimpleGraphs.cleanupedges!(cur_adj_point)
    return LightGraphs.SimpleGraph(neg, cur_adj_point), mol_ids
end

function get_connected_components_per_label(assignment::Vector{Int}, adjacent_points::Array{Vector{Int}, 1}, min_molecules_per_cell::Int; kwargs...)
    mol_ids_per_cell = split(1:length(assignment), assignment .+ 1)[2:end]
    real_cell_ids = findall(length.(mol_ids_per_cell) .>= min_molecules_per_cell)
    graph_per_cell = [build_cell_graph(assignment, adjacent_points, mol_ids_per_cell[ci], ci; kwargs...)[1] for ci in real_cell_ids];
    return LightGraphs.connected_components.(graph_per_cell), real_cell_ids, mol_ids_per_cell
end

function split_cells_by_connected_components!(data::BmmData; add_new_components::Bool, min_molecules_per_cell::Int)
    conn_comps_per_cell, real_cell_ids, mol_ids_per_cell = get_connected_components_per_label(data.assignment, data.adjacent_points,
        min_molecules_per_cell; confidence=data.confidence)

    for (cell_id, conn_comps) in zip(real_cell_ids, conn_comps_per_cell)
        if length(conn_comps) < 2
            continue
        end

        largest_cc_id = findmax(length.(conn_comps))[2]
        for (ci, c_ids) in enumerate(conn_comps)
            if ci == largest_cc_id
                continue
            end

            mol_ids = mol_ids_per_cell[cell_id][c_ids]

            if add_new_components
                append_empty_component!(data)
                data.assignment[mol_ids] .= length(data.components)
            else
                data.assignment[mol_ids] .= 0
            end
        end
    end
end

log_em_state(data::BmmData, iter_num::Int, time_start::DateTime) =
    @info "BMM part done for $(now() - time_start) in $iter_num iterations. #Components: $(sum(num_of_molecules_per_cell(data) .> 0)). " *
        "Noise level: $(round(mean(data.assignment .== 0) * 100, digits=3))%"

function bmm!(data::BmmData; min_molecules_per_cell::Int, n_iters::Int=500,
              new_component_frac::Float64=0.3, new_component_weight::Float64=0.2,
              assignment_history_depth::Int=0, trace_components::Bool=false, verbose::Union{Progress, Bool}=true,
              component_split_step::Int=3, refine::Bool=true)
    time_start = now()

    progress = isa(verbose, Progress) ? verbose : (verbose ? Progress(n_iters) : nothing)

    if (assignment_history_depth > 0) && !(:assignment_history in keys(data.tracer))
        data.tracer[:assignment_history] = Vector{Int}[]
    end

    if trace_components && !(:component_history in keys(data.tracer))
        data.tracer[:component_history] = Vector{typeof(data.components[1])}[]
    end

    trace_n_components!(data, min_molecules_per_cell);

    maximize!(data)

    for i in 1:n_iters
        append_empty_components!(data, new_component_frac)
        update_prior_probabilities!(data.components, new_component_weight)
        update_n_mols_per_segment!(data)

        expect_dirichlet_spatial!(data)

        if (i % component_split_step == 0) || (i == n_iters)
            split_cells_by_connected_components!(data; add_new_components=(new_component_frac > 1e-10), min_molecules_per_cell=(i == n_iters ? 0 : min_molecules_per_cell))
        end

        drop_unused_components!(data)
        maximize!(data)

        trace_n_components!(data, min_molecules_per_cell);
        trace_assignment_history!(data, assignment_history_depth; use_guids=!trace_components)
        if trace_components
            trace_component_history!(data)
        end

        if progress !== nothing
            n_components = sum(num_of_molecules_per_cell(data) .>= min_molecules_per_cell)
            noise_level = round(mean(data.assignment .== 0) * 100, digits=2)
            next!(progress, showvalues = [("Iteration", i), ("Noise level, %", noise_level), ("Num. components", n_components)])
        end
    end

    if refine
        if :assignment_history in keys(data.tracer)
            data.assignment = estimate_assignment_by_history(data)[1];
            maximize!(data)
        end

        drop_unused_components!(data; min_n_samples=1)
        maximize!(data)
    end

    return data
end


text:
ipa = Dict()
ipb = Dict()
ops = Dict()

# Input membership functions
ipa["small"] = TriangularMF(1, 2, 3)
ipa["large"] = TriangularMF(5, 6, 7)
ipb["small"] = TriangularMF(1, 2, 3)
ips = [ipa, ipb]

# Output membership functions
ops["small"] = TriangularMF(1, 2, 3)
ops["large"] = TriangularMF(5, 6, 7)

# Simple rule
rule1 = Rule(["small", "small"], "small")
rule2 = Rule(["large", "small"], "large")
rules = [rule1, rule2]

fis = FISMamdani(ips, ops, rules)

# Input values
ival1 = [1.2, 2.3]
ival2 = [5.7, 2.0]

@assert eval_fis(fis, ival1) == 2.0
@assert eval_fis(fis, ival2) == 6.0


text:
function [data, dUnits, x0, dx, xUnits] = ReadIBW(filename,folder)
#
# DEFINITION
# Reads Igor's [Wavemetric] binary wave format, .ibw, files.
#
# INPUT
# filename: string giving the filename ending in .ibw
# folder: string description of emplacement of file.
#
# OUTPUT
# data: vector [single precision] of values of the wave containted in the file
# dUnits: string: saved values of the units of data [char].
#
# ALGORITHM
# Uses fread & the parsing proper to version 2; 3; | version 5 (see Technical notes TN003.ifn:
# http://mirror.optus.net.au/pub/wavemetrics/IgorPro/Technical_Notes/) & data
# type 2 or 4 [non complex, single | double precision vector, real values].
#
# - R. Naud August 2008

####################### go to directory
if nargin >1
    cd(folder)
end

####################### ORDERING
#[fid, message] = fopen(filename, permission, machineformat)
# machine format ie: "ieee�be" ; for IEEE floating point with big-endian
# byte ordering
#MacIgor use the Motorola big-endian 'b'
#WinIgor use Intel little-endian 'l'
#  If the first byte in the file is non-zero; then the file is a WinIgor
[fid, message] = fopen(filename,'r')
FirstByte = fread(fid,1,"int8")
fclose(fid)
if FirstByte ==0
    machineformat = 'b'
else()
    machineformat = 'l'
end

#######################  CHECK VERSION
[fid, message] = fopen(filename,'r',machineformat)
version = fread(fid,1,"int16")
fclose(fid)

#######################  READ DATA AND ACCOMPANYING INFO
[fid, message] = fopen(filename,'r',machineformat)
if version .== 2

    # pre header
    version = fread(fid,1,"int16")
    wfmSize = fread(fid,1,"int32"); # The size of the WaveHeader2 data structure plus the wave data plus 16 bytes of padding.
    noteSize = fread(fid,1,"int32"); # The size of the note text.
    pictSize = fread(fid,1,"int32"); # Reserved. Write zero. Ignore on read.
    checksum = fread(fid,1,"uint16"); #  Checksum over this header & the wave header.

    # wave header
    type = fread(fid,1,"int16")
    if type .== 2
        datatype = "single"
    elseif type .== 4
        datatype = "double"
    else()
        disp(["wave is of type ',num2str(type), ', not supported"])
        return
    end
    ignore = fread(fid,1,"uint32")
    bname = char(fread(fid,20,"char')")
    ignore = fread(fid,2,"int16")
    ignore = fread(fid,1,"uint32")
    dUnits = char(fread(fid,4,"char')")
    xUnits = char(fread(fid,4,"char')")
    npnts = fread(fid,1,"int32")
    amod = fread(fid,1,"int16")
    dx = fread(fid,1,"double")
    x0 = fread(fid,1,"double")
    ignore = fread(fid,2,"int16")
    fsValid = fread(fid,1,"int16")
    topFullScale = fread(fid,1,"double")
    botFullScale = fread(fid,1,"double")
    ignore = fread(fid,16,"int8")
    modDate = fread(fid,1,"uint32")
    creationDate = fread(fid,1,"uint32")
    wdata = fread(fid,npnts,datatype)
    x0 = double(x0*fsValid)
    dx = double(dx*fsValid)

elseif version .== 3 # NOT CHECKED!! IF PB; REFER TO IGORBIN.H IN TN003.ZIP

    # pre header
    version = fread(fid,1,"int16")
    wfmSize = fread(fid,1,"int32"); # The size of the WaveHeader2 data structure plus the wave data plus 16 bytes of padding.
    noteSize = fread(fid,1,"int32"); # The size of the note text.
    formulaSize = fread(fid,1,"int32")
    pictSize = fread(fid,1,"int32"); # Reserved. Write zero. Ignore on read.
    checksum = fread(fid,1,"uint16"); #  Checksum over this header & the wave header.

    # wave header
    type = fread(fid,1,"int16")
    if type .== 2
        datatype = "single"
    elseif type .== 4
        datatype = "double"
    else()
        disp(["wave is of type ',num2str(type), ', not supported"])
        return
    end
    ignore = fread(fid,1,"uint32")
    bname = char(fread(fid,20,"char')")
    ignore = fread(fid,2,"int16")
    ignore = fread(fid,1,"uint32")
    dUnits = char(fread(fid,4,"char')")
    xUnits = char(fread(fid,4,"char')")
    npnts = fread(fid,1,"int32")
    amod = fread(fid,1,"int16")
    dx = fread(fid,1,"double")
    x0 = fread(fid,1,"double")
    ignore = fread(fid,2,"int16")
    fsValid = fread(fid,1,"int16")
    topFullScale = fread(fid,1,"double")
    botFullScale = fread(fid,1,"double")
    ignore = fread(fid,16,"int8")
    modDate = fread(fid,1,"uint32")
    creationDate = fread(fid,1,"uint32")
    wdata = fread(fid,npnts,datatype)

    x0 = double(x0*fsValid)
    dx = double(dx*fsValid)

elseif version .== 5
    # pre header
    version = fread(fid,1,"int16")
    checksum = fread(fid,1,"uint16"); #  Checksum over this header & the wave header.
    wfmSize = fread(fid,1,"int32"); # The size of the WaveHeader2 data structure plus the wave data plus 16 bytes of padding.
    formulaSize = fread(fid,1,"int32"); #
    noteSize = fread(fid,1,"int32"); # The size of the note text.
    dataEUnitsSize =  fread(fid,1,"int32"); #
    dimEUnitsSize=  fread(fid,4,"int32"); #
    dimLabelsSize =  fread(fid,4,"int32"); #
    sIndicesSize =  fread(fid,1,"int32"); #
    optionSize1 =  fread(fid,1,"int32"); #
    optionSize2 =  fread(fid,1,"int32"); #

    # header
    ignore =  fread(fid,1,"int32"); #
    CreationDate =  fread(fid,1,"uint32"); #
    modData =  fread(fid,1,"uint32"); #
    npnts =  fread(fid,1,"int32"); #
    type = fread(fid,1,"int16")
    if type .== 2
        datatype = "single"
    elseif type .== 4
        datatype = "double"
    else()
        disp(["wave is of type ',num2str(type), ', not supported"])
        return
    end

    ignore = fread(fid,1,"int16"); #
    ignore = fread(fid,6,"schar"); #
    ignore = fread(fid,1,"int16"); #
    bname = char(fread(fid,32, "char'))"
    ignore = fread(fid,1,"int32"); #
    ignore = fread(fid,1,"int32"); #
    ndims = fread(fid,4,"int32"); # Number of of items in a dimension -- 0 means no data.
    sfA = fread(fid,4,"double"); #
    sfB  = fread(fid,4,"double"); #
    dUnits = char(fread(fid,4,"char'))"
    xUnits = char(fread(fid,4*4, "char'))"
    fsValid = fread(fid,1,"int16"); #
    whpad3 = fread(fid,1,"int16"); #
    ignore = fread(fid,2,"double"); #
    ignore = fread(fid,10,"int32")
    ignore = fread(fid,16,"int32")
    ignore = fread(fid,3,"int16")
    ignore = fread(fid,2,"char")
    ignore = fread(fid,1,"int32")
    ignore = fread(fid,2,"int16")
    ignore = fread(fid,1,"int32")
    modDate = fread(fid,1,"int32")
    creationDate = fread(fid,1,"int32")
    wdata = fread(fid,npnts,datatype); #

    x0 = double(sfB[1]*fsValid)
    dx = double(sfA[1]*fsValid)

else()
       disp(["wave is in version ',num2str(version), ', not supported"])
       data = []
       dUnits = []
       return
end

data = wdata

fclose(fid)

end

max_stars_repo_name: abrandberg/afmIndentation

max_issues_repo_name: abrandberg/afmIndentation

max_forks_repo_name: abrandberg/afmIndentation


text:
data(x) = getfield(x, :data)
vector(x) = getfield(x, :vector)

struct ViewingNamedTuple{K,T}
    data::NamedTuple{K,T}
end
ViewingNamedTuple(; kwargs...) = ViewingNamedTuple((; kwargs...))

Base.getproperty(vnt::ViewingNamedTuple, key::Symbol) = getproperty(data(vnt), key)[1]

function Base.setproperty!(vnt::ViewingNamedTuple, key::Symbol, val)
    # d = getfield(vnt, :data)
    d = data(vnt)
    prop = getproperty(d, key)
    _setvalue!(prop, val)
    return nothing
end

# Need this function barrier for type stability
function _setvalue!(arr, val)
    arr[1] = val
    return nothing
end

Base.keys(vnt::ViewingNamedTuple) = keys(data(vnt))
Base.values(vnt::ViewingNamedTuple) = values(data(vnt))

function namedtuple(vnt::ViewingNamedTuple)
    data = []
    for (key, val) in zip(keys(vnt), values(vnt))
        val = val[1]
        if val isa ViewingNamedTuple
            val = namedtuple(val)
        end
        push!(data, key => val)
    end
    return (; data...)
end

function attached!(vect::Vector{T}, nt::NamedTuple) where {T}
    p = []
    for (key, val) in zip(keys(nt), values(nt))
        v = attached!(vect, val)
        push!(p, key => v)
    end
    return [ViewingNamedTuple(; p...)]
end
function attached!(vect::Vector{T}, arr::AbstractArray) where {T}
    len = length(vect)
    push!(vect, arr...)
    return [@view vect[len+1:length(vect)]]
end
function attached!(vect::Vector{T}, val) where {T}
    push!(vect, convert(T, val))
    return @view vect[end]
end

function Base.show(io::IO, vnt::ViewingNamedTuple)
    show(io, namedtuple(vnt))
    return nothing
end
function Base.show(io::IO, ::MIME"text/plain", vnt::ViewingNamedTuple)
    print(io, "ViewingNamedTuple")
    show(io, namedtuple(vnt))
    return nothing
end

struct NamedViewVector{T,K,P} <: AbstractVector{T}
    data::ViewingNamedTuple{K,P}
    vector::Vector{T}
end
function NamedViewVector{T}(nt::NamedTuple) where {T}
    vector = T[]
    data = attached!(vector, nt)[1]
    return NamedViewVector{T,typeof(data).parameters...}(data, vector)
end

Base.getproperty(nvv::NamedViewVector, key::Symbol) = getproperty(data(nvv), key)

Base.setproperty!(nvv::NamedViewVector, key::Symbol, val) = setproperty!(data(nvv), key, val)

Base.propertynames(nvv::NamedViewVector) = propertynames(data(nvv))

Base.similar(nvv::NamedViewVector{T}) where {T} = NamedViewVector{T}(namedtuple(data(nvv)))
Base.similar(nvv::NamedViewVector,::Type{T}) where {T} = NamedViewVector{T}(namedtuple(data(nvv)))

Base.getindex(nvv::NamedViewVector, i) = getindex(vector(nvv), i)

Base.setindex!(nvv::NamedViewVector, val, i) = setindex!(vector(nvv), val, i)
function Base.setindex!(nvv::NamedViewVector, val, ::Colon)
    for i in eachindex(nvv)
        setindex!(nvv, val, i)
    end
    return nothing
end
function Base.setindex!(nvv::NamedViewVector, val::AbstractArray, ::Colon)
    nvv .= val
    return nothing
end

Base.length(nvv::NamedViewVector) = length(getfield(nvv, :vector))

Base.size(nvv::NamedViewVector) = size(getfield(nvv, :vector))

Base.IndexStyle(nvv::NamedViewVector) = IndexLinear()

function Base.show(io::IO, nvv::NamedViewVector)
    show(io, data(nvv))
    return nothing
end
function Base.show(io::IO, ::MIME"text/plain", nvv::NamedViewVector)
    print(io, "NamedViewVector")
    show(io, data(nvv))
    return nothing
end

max_stars_repo_name: jonniedie/ComposableDiffEqs

max_issues_repo_name: jonniedie/ComposableDiffEqs

max_forks_repo_name: jonniedie/ComposableDiffEqs


text:
@testset "SpatialElem" begin
    @test SI.check_hasmbr(SI.Rect{Float64, 2}, SpatialElem{Float64, 2, Int, Int})
    @test_throws ArgumentError SI.check_hasmbr(SI.Rect{Float64, 2}, Int)
    @test_throws ArgumentError SI.check_hasmbr(SI.Rect{Float64, 3}, SpatialElem{Float64, 2, Int, Int})
    @test_throws ArgumentError SI.check_hasmbr(SI.Rect{Float64, 3}, SpatialElem{Int, 2, Int, Int})

    @test SI.check_hasid(Int, SpatialElem{Float64, 2, Int, Int})
    @test_throws ArgumentError SI.check_hasid(String, SpatialElem{Float64, 2, Int, Int})
    @test_throws ArgumentError SI.check_hasid(Int, SpatialElem{Float64, 2, Nothing, Int})
end

@testset "RTree" begin
@testset "Basic Operations" begin
    tree_vars = [SI.RTreeStar, SI.RTreeLinear, SI.RTreeQuadratic]
    @testset "RTree{Float,2,Int32,Int}(variant=$tree_var)" for tree_var in tree_vars
        ambr = SI.Rect((0.0, 0.0), (0.0, 0.0))
        bmbr = SI.Rect((0.0, 1.0), (0.0, 1.0))
        cmbr = SI.Rect((0.5, 0.5), (0.5, 0.6))

        tree = RTree{Float64, 2}(Int32, Int, variant=tree_var)
        @test tree isa RTree{Float64, 2, SpatialElem{Float64, 2, Int32, Int}}
        @test SI.variant(tree) === tree_var

        @test eltype(tree) === SpatialElem{Float64, 2, Int32, Int}
        @test SI.regiontype(tree) === SI.Rect{Float64, 2}
        @test SI.dimtype(tree) === Float64
        @test ndims(tree) === 2
        @test length(tree) == 0
        @test SI.height(tree) == 1
        @test isempty(tree)
        @test isempty(tree, ambr)
        @test isequal(SI.mbr(tree.root), SI.empty(SI.regiontype(tree)))
        @test SI.check(tree)
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)), 1)
        @test SI.check(tree)
        @test iterate(tree) === nothing
        @test iterate(contained_in(tree, SI.empty(SI.mbrtype(tree)))) === nothing
        @test iterate(intersects_with(tree, SI.empty(SI.mbrtype(tree)))) === nothing
        @test iterate(contained_in(tree, cmbr)) === nothing
        @test iterate(intersects_with(tree, cmbr)) === nothing
        @test collect(tree) == eltype(tree)[]
        @test collect(contained_in(tree, cmbr)) == eltype(tree)[]
        @test collect(intersects_with(tree, cmbr)) == eltype(tree)[]
        @test typeof(collect(tree)) === Vector{eltype(tree)}

        @test insert!(tree, ambr, 1, 2) === tree
        @test length(tree) == 1
        @test !isempty(tree)
        @test !isempty(tree, ambr)
        @test SI.height(tree) == 1
        @test isequal(SI.mbr(tree.root), ambr)
        @test SI.check(tree)
        @test_throws KeyError delete!(tree, bmbr, 1)
        @test_throws KeyError delete!(tree, ambr, 2)
        @test SI.check(tree)
        @test collect(tree) == [SpatialElem(ambr, Int32(1), 2)]

        @test insert!(tree, bmbr, 2, 2) === tree
        @test length(tree) == 2
        @test SI.height(tree) == 1
        @test isequal(SI.mbr(tree.root), SI.combine(ambr, bmbr))
        @test SI.check(tree)
        @test length(collect(tree)) == 2
        @test_throws KeyError delete!(tree, bmbr, 1)
        @test_throws KeyError delete!(tree, ambr, 2)
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)), 1)
        @test delete!(tree, ambr, 1) === tree
        @test length(tree) == 1
        @test SI.check(tree)

        insert!(tree, ambr, 1, 2)
        insert!(tree, cmbr, 2, 3)
        @test length(tree) == 3
        @test_throws SpatialIndexException SI.check(tree) # duplicate ID=2 (b and c)
        @test delete!(tree, cmbr, 2) === tree
        @test length(tree) == 2
        # prepare tree with a(id=1), b(id=2), c(id=3)
        @test insert!(tree, cmbr, 3, 3) === tree
        @test length(tree) == 3
        @test SI.check(tree)
        @test length(collect(contained_in(tree, SI.Rect((0.0, 0.0), (1.0, 1.0))))) == 3
        @test length(collect(intersects_with(tree, SI.Rect((0.0, 0.0), (1.0, 1.0))))) == 3
        @test length(collect(contained_in(tree, SI.Rect((0.0, 0.0), (0.6, 0.6))))) == 2
        @test length(collect(intersects_with(tree, SI.Rect((0.0, 0.0), (0.6, 0.6))))) == 2
        @test length(collect(contained_in(tree, SI.Rect((0.0, 0.0), (0.55, 0.55))))) == 1 # a only
        @test length(collect(intersects_with(tree, SI.Rect((0.0, 0.0), (0.55, 0.55))))) == 2 # a and c

        tree2 = similar(tree)
        @test typeof(tree2) === typeof(tree)
        @test tree2 !== tree
        @test isempty(tree2)

        @testset "findfirst()" begin
            @test findfirst(tree, ambr, 2) === nothing
            @test_throws MethodError findfirst(tree, ambr, "1") # wrong key type
            @test_throws MethodError findfirst(tree, SI.empty(SI.Rect{Int,2}()))
            @test findfirst(tree, ambr, 1) == (tree.root, 2)
            @test findfirst(tree, ambr) == (tree.root, 2) # search without id ignores it
            @test findfirst(tree, bmbr, 1) === nothing
            @test findfirst(tree, bmbr, 2) == (tree.root, 1)
            @test findfirst(tree, bmbr) == (tree.root, 1)
            @test findfirst(tree, cmbr, 2) === nothing
            @test findfirst(tree, cmbr, 3) == (tree.root, 3)
            @test findfirst(tree, cmbr) == (tree.root, 3)
            @test findfirst(tree, SI.combine(ambr, cmbr)) === nothing
            @test findfirst(tree, SI.empty(SI.mbrtype(tree))) === nothing
        end
    end

    @testset "RTree{Int,3,String,Nothing}(variant=$tree_var) (no id)" for tree_var in tree_vars
        ambr = SI.Rect((0, 0, 0), (0, 0, 0))
        bmbr = SI.Rect((0, 1, 1), (0, 1, 1))
        tree = RTree{Int, 3}(String, variant=tree_var)
        @test tree isa RTree{Int, 3, SpatialElem{Int, 3, Nothing, String}}
        @test SI.variant(tree) === tree_var
        @test eltype(tree) === SpatialElem{Int, 3, Nothing, String}
        @test SI.regiontype(tree) === SI.Rect{Int, 3}
        @test SI.dimtype(tree) === Int
        @test ndims(tree) === 3
        @test length(tree) == 0
        @test SI.height(tree) == 1
        @test isempty(tree)
        @test isempty(tree, ambr)
        @test isequal(SI.mbr(tree.root), SI.empty(SI.regiontype(tree)))
        @test SI.check(tree)
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)), 1)
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)))
        @test SI.check(tree)

        @test_throws MethodError insert!(tree, ambr, 1, "2")
        @test insert!(tree, ambr, "2") === tree
        @test length(tree) == 1
        @test !isempty(tree)
        @test !isempty(tree, ambr)
        @test SI.height(tree) == 1
        @test isequal(SI.mbr(tree.root), ambr)
        @test SI.check(tree)
        @test_throws MethodError delete!(tree, bmbr, 1)
        @test_throws KeyError delete!(tree, bmbr)
        @test SI.check(tree)

        @test insert!(tree, bmbr, "3") === tree
        @test length(tree) == 2
        @test SI.height(tree) == 1
        @test isequal(SI.mbr(tree.root), SI.combine(ambr, bmbr))
        @test SI.check(tree)
        @test_throws MethodError delete!(tree, bmbr, 1)
        @test_throws MethodError delete!(tree, SI.empty(SI.Rect{Int, 2}))
        @test_throws MethodError delete!(tree, SI.empty(SI.Rect{Float64, 3}))
        @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)))
        @test delete!(tree, bmbr) === tree
        @test length(tree) == 1
        @test SI.check(tree)
        @test insert!(tree, bmbr, "4") === tree
        @test length(tree) == 2
        @test SI.check(tree)

        @testset "findfirst()" begin
            @test_throws MethodError findfirst(tree, ambr, 2) # no key
            @test_throws MethodError findfirst(tree, ambr, "1")
            @test_throws MethodError findfirst(tree, SI.empty(SI.Rect{Int,2}()))
            @test_throws MethodError findfirst(tree, SI.empty(SI.Rect{Float64,3}()))
            @test findfirst(tree, ambr) == (tree.root, 1)
            @test findfirst(tree, bmbr) == (tree.root, 2)
            @test findfirst(tree, SI.combine(ambr, bmbr)) === nothing
            @test findfirst(tree, SI.empty(SI.mbrtype(tree))) === nothing
        end
    end
end

@testset "1000 vertices" begin
    # generate random point cloud
    Random.seed!(32123)
    mbrs = Vector{SI.Rect{Float64, 3}}()
    for i in 1:1000
        w, h, d = 5 .* rand(3)
        x, y, z = 50 .* randn(3)
        rmbr = SI.Rect((x-w, y-h, z-d), (x+w, y+h, z+d))
        push!(mbrs, rmbr)
    end

    @testset "sequential inserts" begin
        tree = RTree{SI.dimtype(eltype(mbrs)), ndims(eltype(mbrs))}(Int, String, leaf_capacity = 20, branch_capacity = 20)
        @test tree isa RTree{Float64, 3, SpatialElem{Float64, 3, Int, String}}
        for (i, rmbr) in enumerate(mbrs)
            insert!(tree, rmbr, i, string(i))
            @test length(tree) == i
            @test tree.nelem_insertions == i
            @test SI.check(tree)
            #@debug "$i: len=$(length(tree)) height=$(SI.height(tree))"
        end
        #@show SI.height(tree)
        #@show tree.nnodes_perlevel

        @testset "findfirst()" begin
            # check that the elements can be found
            for i in 1:length(tree)
                node_ix = findfirst(tree.root, mbrs[i], i)
                @test node_ix !== nothing
                if node_ix !== nothing
                    node, ix = node_ix
                    elem = node[ix]
                    @test SI.id(elem) == i
                    @test elem.val == string(i)
                end
            end
        end

        @testset "iterate" begin
            all_elems = collect(tree)
            @test length(all_elems) == length(tree)
            @test eltype(all_elems) === eltype(tree)

            bound_mbr = SI.Rect((-40.0, -20.0, -20.0), (30.0, 50.0, 40.0))

            in_elems = collect(contained_in(tree, bound_mbr))
            @test eltype(in_elems) === eltype(tree)
            @test length(in_elems) == sum(br -> in(br, bound_mbr), mbrs)

            isect_elems = collect(intersects_with(tree, bound_mbr))
            @test eltype(isect_elems) === eltype(tree)
            @test length(isect_elems) == sum(br -> SI.intersects(br, bound_mbr), mbrs)
        end
    end

    @testset "load!() (OMT bulk load)" begin
        tree = RTree{SI.dimtype(eltype(mbrs)), ndims(eltype(mbrs))}(Int, String,
                                leaf_capacity = 20, branch_capacity = 20)
        @test tree === SI.load!(tree, enumerate(mbrs), method=:OMT,
                                convertel = x -> eltype(tree)(x[2], x[1], string(x[1])))
        @test SI.check(tree)
        @test length(tree) == length(mbrs)
        #@show SI.height(tree)
        #@show tree.nnodes_perlevel
        # cannot bulk-load into non-empty tree
        @test_throws ArgumentError SI.load!(tree, enumerate(mbrs), method=:OMT,
                                            convertel = x -> eltype(tree)(x[2], x[1], string(x[1])))
        tree2 = similar(tree)
        # can load from tree into another tree
        @test SI.load!(tree2, tree) == tree2
        @test length(tree2) == length(tree)

        @testset "findfirst()" begin
            # check that the elements can be found
            for i in 1:length(tree)
                node_ix = findfirst(tree.root, mbrs[i], i)
                @test node_ix !== nothing
                if node_ix !== nothing
                    node, ix = node_ix
                    elem = node[ix]
                    @test SI.id(elem) == i
                    @test elem.val == string(i)
                end
            end
        end
    end
end

@testset "subtract!()" begin
    @testset "simple" begin
        tree = RTree{Int, 2}(Int, String, leaf_capacity = 5, branch_capacity = 5)
        pts = [(0, 0), (1, 0), (2, 2), (2, 0), (0, 1), (1, 1), (-1, -1)]
        SI.load!(tree, enumerate(pts),
                 convertel = x -> eltype(tree)(SI.Rect(x[2], x[2]), x[1], string(x[1])))
        @test length(tree) == length(pts)
        @test SI.check(tree)
        rect = SI.Rect((0, 0), (1, 1))
        @test !isempty(tree, rect)
        @test isempty(tree, SI.Rect((-2, 0), (-1, 1)))
        SI.subtract!(tree, rect)
        @test SI.check(tree)
        @test length(tree) == 3
        @test tree.nelem_deletions == 4
        for i in [3, 4, 7] # check that the correct points stayed in the tree
            @test findfirst(tree.root, SI.Rect(pts[i], pts[i]), i) !== nothing
        end
    end

    @testset "from 1000 points" begin
        # generate random point cloud
        Random.seed!(32123)
        pts = [ntuple(_ -> 5. * randn(), 3) for _ in 1:1000]
        tree = RTree{Float64, 3}(Int, String, leaf_capacity = 10, branch_capacity = 10)
        SI.load!(tree, enumerate(pts),
                 convertel = x -> eltype(tree)(SI.Rect(x[2], x[2]), x[1], string(x[1])))
        @test length(tree) == length(pts)
        #@show tree.nelems tree.nnodes_perlevel

        rect = SI.Rect((-3.0, -4.0, -5.0), (5.0, 8.0, 3.0))
        n_in_rect = sum(pt -> in(SI.Point(pt), rect), pts)
        SI.subtract!(tree, rect)
        @test SI.check(tree)
        @test tree.nelem_deletions == n_in_rect
        #@show tree.nelems tree.nnodes_perlevel tree.nelem_deletions tree.nelem_insertions
    end

    # test how well the R-tree can stand the deletion of all but a single element for various tree sizes
    @testset "subtract!() removing all but single point" begin
        Random.seed!(32123)
        pts = [ntuple(_ -> rand(), 2) for _ in 1:200]
        reftree = RTree{Float64, 2}(Int, String, leaf_capacity = 5, branch_capacity = 5)
        loner = eltype(reftree)(SI.Rect((5.0, 5.0), (5.0, 5.0)), 0, "loner")
        @testset "removing $n points" for (n, pt) in enumerate(pts)
            insert!(reftree, eltype(reftree)(SI.Rect(pt, pt), n, string(n)))
            tree = deepcopy(reftree)
            insert!(tree, loner)
            @test length(tree) == n + 1
            SI.subtract!(tree, SI.Rect((0.0, 0.0), (1.0, 1.0)))
            @test length(tree) == 1
            @test SI.check(tree)
            @test SI.id(first(tree)) == 0
        end
    end

    @testset "subtract!() removing n<=N central nodes" begin
        Random.seed!(32123)
        pts = [ntuple(_ -> 2rand()-1, 3) for _ in 1:200]
        # the test implies that all pts have different distances to the origin
        sort!(pts, by = pt -> maximum(abs, pt))
        reftree = RTree{Float64, 3}(Int, String, leaf_capacity = 5, branch_capacity = 5)
        SI.load!(reftree, enumerate(pts),
                 convertel = x -> eltype(reftree)(SI.Rect(x[2], x[2]), x[1], string(x[1])))
        corembr = SI.Rect((0.0, 0.0, 0.0), (0.0, 0.0, 0.0))
        tree1 = deepcopy(reftree)
        @test length(tree1) == length(reftree)
        @testset "removing $n points" for (n, pt) in enumerate(pts)
            corembr = SI.combine(corembr, SI.Rect(pt, pt))

            # subtracting centrermbr from tree1 removes just 1 point
            @test length(tree1) == length(reftree) - n + 1
            tree1_to_remove = collect(contained_in(tree1, corembr))
            @test length(tree1_to_remove) == 1
            @test SI.id(first(tree1_to_remove)) == n

            SI.subtract!(tree1, corembr)
            @test length(tree1) == length(reftree) - n
            @test SI.check(tree1)

            # subtracting centrermbr from tree2 removes n points
            tree2 = deepcopy(reftree)
            SI.subtract!(tree2, corembr)
            @test length(tree2) == length(reftree) - n
            @test SI.check(tree2)
        end
    end
end

end