text:
@testset "insert/deleteat" begin
    @test @inferred(ArrayInterface.insert([1,2,3], 2, -2)) == [1, -2, 2, 3]
    @test @inferred(ArrayInterface.deleteat([1, 2, 3], 2)) == [1, 3]

    @test @inferred(ArrayInterface.deleteat([1, 2, 3], [1, 2])) == [3]
    @test @inferred(ArrayInterface.deleteat([1, 2, 3], [1, 3])) == [2]
    @test @inferred(ArrayInterface.deleteat([1, 2, 3], [2, 3])) == [1]

    @test @inferred(ArrayInterface.insert((2,3,4), 1, -2)) == (-2, 2, 3, 4)
    @test @inferred(ArrayInterface.insert((2,3,4), 2, -2)) == (2, -2, 3, 4)
    @test @inferred(ArrayInterface.insert((2,3,4), 3, -2)) == (2, 3, -2, 4)

    @test @inferred(ArrayInterface.deleteat((2, 3, 4), 1)) == (3, 4)
    @test @inferred(ArrayInterface.deleteat((2, 3, 4), 2)) == (2, 4)
    @test @inferred(ArrayInterface.deleteat((2, 3, 4), 3)) == (2, 3)
    @test ArrayInterface.deleteat((1, 2, 3), [1, 2]) == (3,)
end

ext: jl


text:
#####
##### Sampling: high-level interface and building blocks
#####

"Significant digits to display for reporting."
const REPORT_SIGDIGITS = 3

####
#### parts unaffected by warmup
####

"""
$(TYPEDEF)
A log density bundled with an RNG and options for sampling. Contains the parts of the
problem which are not changed during warmup.
# Fields
$(FIELDS)
"""
struct SamplingLogDensity{D,R,L<:AbstractProbabilityModel{D},O,S}
    "Random number generator."
    rng::R
    "Log density."
    ℓ::L
    """
    Algorithm used for sampling, also contains the relevant parameters that are not affected
    by adaptation. See eg [`NUTS`](@ref).
    """
    algorithm::O
    "Reporting warmup information and chain progress."
    reporter::S
end

####
#### warmup building blocks
####

###
### warmup state
###

"""
$(TYPEDEF)
Representation of a warmup state. Not part of the API.
# Fields
$(FIELDS)
"""
struct WarmupState{D,T,L,Tκ <: KineticEnergy, Tϵ <: Union{Real,Nothing}}
    z::PhasePoint{D,T,L}
    κ::Tκ
    ϵ::Tϵ
end

function Base.show(io::IO, warmup_state::WarmupState)
    @unpack κ, ϵ = warmup_state
    ϵ_display = ϵ ≡ nothing ? "unspecified" : "≈ $(round(ϵ; sigdigits = REPORT_SIGDIGITS))"
    print(io, "adapted sampling parameters: stepsize (ϵ) $(ϵ_display), $(κ)")
end

###
### warmup interface and stages
###

"""
$(SIGNATURES)
Return the *results* and the *next warmup state* after warming up/adapting according to
`warmup_stage`, starting from `warmup_state`.
Use `nothing` for a no-op.
"""
function warmup!(tree, chain, sampling_logdensity::SamplingLogDensity, warmup_stage::Nothing, warmup_state)
    nothing, warmup_state
end

random_position!(rng::AbstractRNG, q::PtrVector) = rand!(rng, q, -2.0, 2.0)

"""
$(SIGNATURES)
Helper function to create random starting positions in the `[-2,2]ⁿ` box.
"""
function random_position(rng::AbstractRNG, sptr::StackPointer, ::Static{P}) where {P}
    sptr, q = PtrVector{P,Float64}(sptr)
    sptr, random_position!(rng, q)
end

"Docstring for initial warmup arguments."
const DOC_INITIAL_WARMUP_ARGS =
"""
- `q`: initial position. *Default*: random (uniform [-2,2] for each coordinate).
- `κ`: kinetic energy specification. *Default*: Gaussian with identity matrix.
- `ϵ`: a scalar for initial stepsize, or `nothing` for heuristic finders.
"""

"""
$(SIGNATURES)
Create an initial warmup state from a random position.
# Keyword arguments
$(DOC_INITIAL_WARMUP_ARGS)
"""
function initialize_warmup_state(rng, sptr::StackPointer, ℓ::AbstractProbabilityModel{D}; q = nothing, κ = nothing, ϵ = nothing) where {D}
    if κ ≡ nothing
        sptr, κ′ = GaussianKineticEnergy(sptr, dimension(ℓ))
    else
        κ′ = κ
    end
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    tree = Tree{D,Float64}(sptr)
    # println("Allocated tree:")
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    treeptr, flag = undefined_z(tree)
    # @show treeptr
    LT = aligned_offset(tree)
    p = PtrVector{D,Float64}(treeptr)
    # println("About to generate a random position.")
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    q′ = q ≡ nothing ? random_position!(rng, PtrVector{D,Float64}(treeptr + LT)) : q
    ∇ℓq = PtrVector{D,Float64}(treeptr + 2LT)
    # println("About to evaluate the logdensity:")
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    Q = evaluate_ℓ!(tree.sptr, ∇ℓq, ℓ, q′)
    # println("Just evaluated the logdensity:")
    # @show κ′.W.diag, κ′.M⁻¹.diag
    # @show pointer(κ′.W.diag), pointer(κ′.M⁻¹.diag)
    tree, WarmupState(PhasePoint(Q, p, flag), κ′, ϵ)
end

"""
$(TYPEDEF)
Find a local optimum (using quasi-Newton methods).
It is recommended that this stage is applied so that the initial stepsize selection happens
in a region which is at least plausible.
"""
Base.@kwdef struct FindLocalOptimum{T}
    """
    Add `-0.5 * magnitude_penalty * sum(abs2, q)` to the log posterior **when finding the local
    optimum**. This can help avoid getting into high-density edge areas of the posterior
    which are otherwise not typical (eg multilevel models).
    """
    magnitude_penalty::T = 1e-4
    """
    Maximum number of iterations in the optimization algorithm. Recall that we don't need to
    find the mode, or even a local mode, just be in a reasonable region.
    """
    iterations::Int = 50
    # FIXME allow custom algorithm, tolerance, etc
end
@noinline ThrowOptimizationError(str) = throw(str)
function warmup!(
    tree::Tree, chain, tree_statistics, ϵs, sampling_logdensity::SamplingLogDensity{D},
    local_optimization::FindLocalOptimum, warmup_state
) where {D}
    @unpack rng, ℓ, reporter = sampling_logdensity
    @unpack magnitude_penalty, iterations = local_optimization
    @unpack z, κ, ϵ = warmup_state
    @unpack Q, p, flag = z
    @unpack q, ℓq, ∇ℓq = Q
    report(reporter, "finding initial optimum")
    for _ in 1:100
        ℓq = QuasiNewtonMethods.proptimize!(tree.sptr, ℓ, q, ∇ℓq, ℓq, magnitude_penalty, iterations)#+100)
        # @show q
        # @show ℓq
        # @show ∇ℓq
        isfinite(ℓq) && return WarmupState(PhasePoint(EvaluatedLogDensity(q, ℓq, ∇ℓq), p, flag), κ, ϵ)
        random_position!(rng, q)
        ℓq = evaluate_ℓ!(tree.sptr, ∇ℓq, ℓ, q).ℓq
        magnitude_penalty += magnitude_penalty
    end
    ThrowOptimizationError("Optimization failed to converge, returning $ℓq; Thread ID: $(Threads.threadid()).")
    # fg! = function(F, G, q)
        # ℓq, ∇ℓq = logdensity_and_gradient(ℓ, q)
        # if G ≠ nothing
            # @. G = -∇ℓq - q * magnitude_penalty
        # end
        # -ℓq - (0.5 * magnitude_penalty * sum(abs2, q))
    # end
    # objective = NLSolversBase.OnceDifferentiable(NLSolversBase.only_fg!(fg!), q)
    # opt = Optim.optimize(objective, q, Optim.LBFGS(),
                         # Optim.Options(; iterations = iterations))
    # sp, q = PtrVector{D,Float64}(sp)
    # ∇ℓq = PtrVector{D,Float64}(pointer(sp,Float64))
    # sptr is set ahead by proptimize! to store optim and gradient.
end
Base.length(::FindLocalOptimum) = 0
function warmup!(
    tree::Tree, chain, tree_statistics, ϵs,
    sampling_logdensity, stepsize_search::InitialStepsizeSearch, warmup_state
)
    @unpack rng, ℓ, reporter = sampling_logdensity
    @unpack z, κ, ϵ = warmup_state
    # @argcheck ϵ ≡ nothing "stepsize ϵ manually specified, won't perform initial search"
    rand_p!(rng, z.p, κ)
    ϵ = find_initial_stepsize(stepsize_search, local_acceptance_ratio(tree.sptr, Hamiltonian(κ, ℓ), z))
    report(reporter, "found initial stepsize",
           ϵ = round(ϵ; sigdigits = REPORT_SIGDIGITS))
    WarmupState(z, κ, ϵ)
end
Base.length(::InitialStepsizeSearch) = 0
"""
$(TYPEDEF)
Tune the step size `ϵ` during sampling, and the metric of the kinetic energy at the end of
the block. The method for the latter is determined by the type parameter `M`, which can be
1. `Diagonal` for diagonal metric (the default),
2. `Symmetric` for a dense metric,
3. `Nothing` for an unchanged metric.
# Results
A `NamedTuple` with the following fields:
- `chain`, a vector of position vectors
- `tree_statistics`, a vector of tree statistics for each sample
- `ϵs`, a vector of step sizes for each sample
# Fields
$(FIELDS)
"""
struct TuningNUTS{M,D}
    "Number of samples."
    N::Int
    "Dual averaging parameters."
    stepsize_adaptation::D
    """
    Regularization factor for normalizing variance. An estimated covariance matrix `Σ` is
    rescaled by `λ` towards ``σ²I``, where ``σ²`` is the median of the diagonal. The
    constructor has a reasonable default.
    """
    λ::Float64
    function TuningNUTS{M}(N::Integer, stepsize_adaptation::D,
                           λ = 5.0/N) where {M <: Union{Nothing,Diagonal,Symmetric},D}
        # @argcheck N ≥ 20        # variance estimator is kind of meaningless for few samples
        # @argcheck λ ≥ 0
        new{M,D}(N, stepsize_adaptation, λ)
    end
end

function Base.show(io::IO, tuning::TuningNUTS{M}) where {M}
    @unpack N, stepsize_adaptation, λ = tuning
    print(io, "Stepsize and metric tuner, $(N) samples, $(M) metric, regularization $(λ)")
end

Base.length(t::TuningNUTS) = t.N

# """
# $(SIGNATURES)
# Form a matrix from positions (`q`), with each column containing a position.
# """
# position_matrix(chain) = reduce(hcat, chain)

# """
# $(SIGNATURES)
# Estimate the inverse metric from the chain.
# In most cases, this should be regularized, see [`regularize_M⁻¹`](@ref).
# """
# sample_M⁻¹(::Type{Diagonal}, chain) = Diagonal(vec(var(position_matrix(chain); dims = 2)))

# sample_M⁻¹(::Type{Symmetric}, chain) = Symmetric(cov(position_matrix(chain); dims = 2))

# function store_sample!(_∇ℓq::PtrVector{P,T,L}, _q::PtrVector{P,T,L}, Q::EvaluatedLogDensity) where {P,T,L}
    # @unpack q, ∇ℓq = Q
    # @inbounds for l in 1:L
        # _q[l] = q[l]
        # _∇ℓq[l] = ∇ℓq[l]
    # end
    # EvaluatedLogDensity(
        # _q, Q.ℓq, _∇ℓq
    # )
# end

function warmup!(
    tree::Tree{D,T,L}, chain₊::AbstractMatrix{T},
    tree_statistics::AbstractVector{TreeStatisticsNUTS}, ϵs::AbstractVector{T},
    sampling_logdensity::SamplingLogDensity{D},
    tuning::TuningNUTS{M},
    warmup_state
) where {D,T,L,M}
    @unpack rng, ℓ, algorithm, reporter = sampling_logdensity
    @unpack z, κ, ϵ = warmup_state
    @unpack N, stepsize_adaptation, λ = tuning
    # L = VectorizationBase.align(D, T)
    chain_ptr = pointer(chain₊)
    chain = DynamicPtrMatrix{T}(chain_ptr, (D, N), L)
    # tree_statistics = Vector{TreeStatisticsNUTS}(undef, N)
    H = Hamiltonian(κ, ℓ)
    ϵ_state = initial_adaptation_state(stepsize_adaptation, ϵ)
    # ϵs = Vector{Float64}(undef, N)
    mcmc_reporter = make_mcmc_reporter(reporter, N; tuning = "stepsize and Diagonal{T,PtrVector{}} metric")
    # sp, ∇ℓq = PtrVector{D,T}(sp)
    for n in 1:N
        ϵ = current_ϵ(ϵ_state)
        # @assert ϵ > 1e-10 "Current ϵ: $ϵ; final: $(final_ϵ(ϵ_state))"
        if ϵ < 1e-10
            @show Threads.threadid(), z
            Q = evaluate_ℓ!(tree.sptr, z.Q.∇ℓq, sampling_logdensity.ℓ, z.Q.q)
            @show Threads.threadid(), Q
            throw(AssertionError("Current ϵ: $ϵ; final: $(final_ϵ(ϵ_state))"))
        end
        ϵs[n] = ϵ
        z, stats = sample_tree(rng, tree, algorithm, H, z, ϵ)
        copyto!( PtrVector{D,T}( chain_ptr ), z.Q.q ) # relocate to base of stack
        # @show z.Q.q
        chain_ptr += L*sizeof(T)
        tree_statistics[n] = stats
        ϵ_state = adapt_stepsize(stepsize_adaptation, ϵ_state, stats.acceptance_rate)
        reporting(mcmc_reporter) && report(mcmc_reporter, n; ϵ = round(ϵ; sigdigits = REPORT_SIGDIGITS))
    end
    # κ = GaussianKineticEnergy(regularize_M⁻¹(sample_M⁻¹(M, chain), λ))
    # sp, κ = GaussianKineticEnergy(sp, chain, λ, Val{D}())
    if M ≢ Nothing
        GaussianKineticEnergy!(κ, chain, λ)
        reporting(mcmc_reporter) && report(mcmc_reporter, "adaptation finished", adapted_kinetic_energy = κ)
    end
    # (chain = chain, tree_statistics = tree_statistics, ϵs = ϵs)
    WarmupState(z, κ, final_ϵ(ϵ_state))
end

function mcmc!(tree::Tree{D,T,L}, chain::AbstractMatrix, tree_statistics::AbstractVector{TreeStatisticsNUTS}, sampling_logdensity::SamplingLogDensity{D}, N, warmup_state) where {D,T,L}
    @unpack rng, ℓ, algorithm, reporter = sampling_logdensity
    @unpack z, κ, ϵ = warmup_state

    # tree_statistics = Vector{TreeStatisticsNUTS}(undef, N)
    H = Hamiltonian(κ, ℓ)
    mcmc_reporter = make_mcmc_reporter(reporter, N)
    chain_ptr = pointer(chain)
    for n in 1:N
        z, tree_statistics[n] = sample_tree(rng, tree, algorithm, H, z, ϵ)
        copyto!( PtrVector{D,T}( chain_ptr ), z.Q.q )
        chain_ptr += L*sizeof(T)

        report(mcmc_reporter, n)
    end
    # tree_statistics
end

"""
$(SIGNATURES)
Helper function for constructing the “middle” doubling warmup stages in
[`default_warmup_stages`](@ref).
"""
function _doubling_warmup_stages(M, stepsize_adaptation, middle_steps,
                                 doubling_stages::Val{D}) where {D}
    ntuple(d -> TuningNUTS{M}(middle_steps << (d - 1), stepsize_adaptation), Val(D))
end

"""
$(SIGNATURES)
A sequence of warmup stages:
1. find the local optimum using `local_optimization`,
2. select an initial stepsize using `stepsize_search` (default: based on a heuristic),
3. tuning stepsize with `init_steps` steps
4. tuning stepsize and covariance: first with `middle_steps` steps, then repeat with twice
   the steps `doubling_stages` times
5. tuning stepsize with `terminating_steps` steps.
`M` (`Diagonal`, the default or `Symmetric`) determines the type of the metric adapted from
the sample.
This is the suggested tuner of most applications.
Use `nothing` for `local_optimization` or `stepsize_adaptation` to skip the corresponding
step.
"""
function default_warmup_stages(;
                               local_optimization = FindLocalOptimum(),
                               stepsize_search = InitialStepsizeSearch(),
                               M::Type{<:Union{Diagonal,Symmetric}} = Diagonal,
                               stepsize_adaptation = DualAveraging(),
                               init_steps = 75, middle_steps = 25, doubling_stages::Val{DS} = Val(5),
                               terminating_steps = 50) where {DS}
    (local_optimization, stepsize_search,
     TuningNUTS{Nothing}(init_steps, stepsize_adaptation),
     _doubling_warmup_stages(M, stepsize_adaptation, middle_steps, Val(DS))...,
     TuningNUTS{Nothing}(terminating_steps, stepsize_adaptation))
end

"""
$(SIGNATURES)
A sequence of warmup stages for fixed stepsize:
1. find the local optimum using `local_optimization`,
2. tuning covariance: first with `middle_steps` steps, then repeat with twice
   the steps `doubling_stages` times
Very similar to [`default_warmup_stages`](@ref), but omits the warmup stages with just
stepsize tuning.
"""
function fixed_stepsize_warmup_stages(;
                                      local_optimization = FindLocalOptimum(),
                                      M::Type{<:Union{Diagonal,Symmetric}} = Diagonal,
                                      middle_steps = 25, doubling_stages::Val{DS} = 5) where {DS}
    (local_optimization,
     _doubling_warmup_stages(M, FixedStepsize(), middle_steps, Val(DS))...)
end

"""
$(SIGNATURES)
Helper function for implementing warmup.
!!! note
    Changes may imply documentation updates in [`mcmc_keep_warmup`](@ref).
"""
@generated function _warmup!(tree::Tree, chain, tree_statistics, ϵs, sampling_logdensity, stages::T, warmup_state_0) where {T}
    N = length(T.parameters)
    q = quote end
    for n in 1:N
        warmup_call_q = quote
            $(Symbol(:warmup_state_,n)) = warmup!(tree, chain, tree_statistics, ϵs, sampling_logdensity, stages[$n], $(Symbol(:warmup_state_,n-1)))
        end
        push!(q.args, warmup_call_q)
    end
    push!(q.args, Symbol(:warmup_state_,N))
    q
end


text:
@testset "Init pop plots" begin
    mpSim2 = MPsim()
    filePath = joinpath( "base", "baseConfig" )
    fileName = string( filePath, ".xlsx" )

    XLSX.openxlsx( fileName ) do xf
        MPIO.readInitPop( mpSim2, xf, filePath )
    end  # XLSX.openxlsx( fileName ) do xf

    excelInitPopPlot( mpSim2, "excelrequest/plotRequest" )
end  # @testset "Init pop plots"


text:
module CellDofsTests

using Test
using FillArrays
using Gridap.Helpers
using Gridap.Fields
using Gridap.ReferenceFEs
using Gridap.Arrays
using Gridap.Integration
using Gridap.CellData
using Gridap.TensorValues
using Gridap.Geometry
using LinearAlgebra

domain = (0,1,0,1,0,1)
cells = (2,2,2)
model = simplexify(CartesianDiscreteModel(domain,cells))

trian = Triangulation(model)

v = GenericCellField(get_cell_shapefuns(trian),trian,ReferenceDomain())

ctype_to_reffe, cell_to_ctype = compress_cell_data(get_cell_reffe(trian))
ctype_to_dofbasis = map(get_dof_basis,ctype_to_reffe)
cell_dofbasis = expand_cell_data(ctype_to_dofbasis,cell_to_ctype)

s = CellDof(cell_dofbasis,trian,ReferenceDomain())

a = s(v)
r = fill(Matrix(I,4,4),num_cells(trian))
test_array(a,r)

acell_to_cell = [1,3,5,3]
atrian = Triangulation(trian,acell_to_cell)
s_a = change_domain(s,atrian,DomainStyle(s))
v_a = change_domain(v,atrian,DomainStyle(v))

a = s_a(v_a)
r = fill(Matrix(I,4,4),num_cells(atrian))
test_array(a,r)

end # module


text:
using SFGReservoirs
using SFGAnalysis
using Optim
using Test
using StaticArrays
using DifferentialEquations
using BenchmarkTools

begin
      N = 4 # number of states
      p = [0.1, 0.2, 30.0, -1.0] # [k1, k2, x, t0]
      k = @MMatrix zeros(N,N)
      assignments_k = Dict(
            (1,2) => 1, # the entry k[1,2] should map to p[1]
            (4,1) => 2,
            (4,2) => 2,
      )
      assignments_k_matrix = @MMatrix zeros(Int, 4,4)
      assignments_k_matrix[1,2] = 1
      assignments_k_matrix[4,1] = 2
      assignments_k_matrix[4,2] = 2

      @btime SFGReservoirs.parameters_to_rate_matrix!(k, assignments_k, p)
      @btime SFGReservoirs.parameters_to_rate_matrix!(k, assignments_k_matrix, p)
end

@testset "Model Setup" begin
      N = 4 # number of states
      p = [0.1, 0.2, 30.0, -1.0] # [k1, k2, x, t0]
      k = @MMatrix zeros(N,N)
      assignments_k = Dict(
            (1,2) => 1, # the entry k[1,2] should map to p[1]
            (4,1) => 2,
            (4,2) => 2,
      )
      SFGReservoirs.parameters_to_rate_matrix!(k, assignments_k, p)
      @test k[1,2] == p[1]
      @test k[4,1] == p[2]
      @test k[4,2] == p[2]
      @test size(k) == (4,4)

      k2 = @MMatrix zeros(N,N)
      assignments_k_matrix = @MMatrix zeros(Int, 4,4)
      assignments_k_matrix[1,2] = 1
      assignments_k_matrix[4,1] = 2
      assignments_k_matrix[4,2] = 2
      SFGReservoirs.parameters_to_rate_matrix!(k2, assignments_k_matrix, p)
      @test k2 == k

      x = @MMatrix zeros(N,N)
      assignments_x = Dict(
            (1,2) => 3,
            (2,1) => 3,
      )
      SFGReservoirs.parameters_to_rate_matrix!(x, assignments_x, p)
end

@test SFGReservoirs.gauss(0, 0, 3) == 1
@test SFGReservoirs.gauss(100, 0, 3) ≈ 0
@test SFGReservoirs.gauss(4, 1, 3) ≈ 1/ℯ

begin
      du = @MArray zeros(3)
      u  = @MArray zeros(3)
      k  = @MMatrix rand(3,3)
      x  = @MMatrix zeros(3,3)
      t0 = 0.0
      σ  = 3.0
      x[1,2] = 0.1
      x[2,1] = 0.1
      p  = (k,x,t0,σ)
      t  = 0.0
      @btime model(du, u, p, t)
end

@testset "General Model" begin
      tstart = -30.0
      tend = 270.0
      tspan = (tstart, tend)

      # All in ground state and pumping hard. We shoud end up with
      # 1/2 in both ground and excited state.
      u0 = [1.0, 0, 0]
      N = length(u0)
      k = @MMatrix zeros(N,N)
      x = @MMatrix zeros(N,N)
      x[1,2] = 1.0
      x[2,1] = 1.0
      t0 = 0
      σ  = 3.0
      p = (k,x,t0,σ)
      prob = ODEProblem(model, u0, tspan, p)
      sol = solve(prob)
      @test isapprox(sol[end][1], 0.5, atol=1e-4) &&
            isapprox(sol[end][2], 0.5, atol=1e-4)
      @test sum(sol[end]) ≈ 1.0

      # All in 2nd state. let flow to 1st and 3nd state. Let the flow
      # to state 3 be twice as fast.
      # We should end up with 1/3 and 2/3 in states 1 and 3 respectively.
      u0 = [0, 1.0, 0]
      N = length(u0)
      k = @MMatrix zeros(N,N)
      x = @MMatrix zeros(N,N)
      k[2,1] = 1.0
      k[2,3] = 2.0
      t0 = 0
      σ  = 3.0
      p = (k,x,t0,σ)
      prob = ODEProblem(model, u0, tspan, p)
      sol = solve(prob)
      @test isapprox(sol[end][1], 1/3, atol=1e-4) &&
            isapprox(sol[end][3], 2/3, atol=1e-4)
      @test sum(sol[end]) ≈ 1.0

      # BenchmarkTools
      @btime solve($prob)
end

@testset "Population to Attenuation" begin
      # let's have a 4-state model
      #     1: ground state
      #     2: excited state (shows up in spectrum)
      #     3: excited state (shows up in spectrum)
      #     4: excited state (does not show up in spectrum and this state
      #                       does not contribute to SF intensity of states 2,3)
      # populations
      u = [0.7, 0.1, 0.05, 0.15]
      s1 = SFGReservoirs.State()
      s2 = SFGReservoirs.State(2, [1,3])
      s3 = SFGReservoirs.State(3, [1,2])
      states1 = (s2, s3)
      states2 = (s1, s2, s3)
      a1 = @MArray [0.0, 0.0]
      a2 = @MArray [0.0, 0.0, 0.0]
      SFGReservoirs.population_to_attenuation!(a1, u, states1)
      @test all(a1 .≈ [0.65, 0.75])
      SFGReservoirs.population_to_attenuation!(a2, u, states2)
      @test all(a2 .≈ [1.0, 0.65, 0.75])
end

begin
      u = [0.7, 0.1, 0.05, 0.15]
      s1 = SFGReservoirs.State()
      s2 = SFGReservoirs.State(2, [1,3])
      s3 = SFGReservoirs.State(3, [1,2])
      states1 = (s2, s3)
      states2 = (s1, s2, s3)
      a1 = @MArray [0.0, 0.0]
      a2 = @MArray [0.0, 0.0, 0.0]
      @btime SFGReservoirs.population_to_attenuation!(a1, u, states1)
      @btime SFGReservoirs.population_to_attenuation!(a2, u, states2)
end

# attenuated spectrum
begin
      x = collect(LinRange(2830, 3000, 512))
      a = [0.65, 1.1]
      y = zeros(ComplexF64, 512)
      A = [5, -3]
      ω = [2870, 2940]
      Γ = [8, 8]
      @btime SFGReservoirs.attenuated_sfspectrum!(y, x, a, A, ω, Γ)
      # plot(x, real(y))
      @btime sfspectrum!(y, x, A, ω, Γ)
      # plot!(x, real(y))
      nothing
end

begin
      u0 = [1.0, 0.0, 0.0, 0.0]
      delaytimes = range(-30.0, 120.0, length=121)  |> collect
      wavenumbers = range(2830, 3000.0, length=400) |> collect

      A = [0.2, 9, -3, -4, 9, -6, -1, -1]
      ω = [2845.0, 2876, 2886, 2905, 2938, 2947, 2961, 2966]
      Γ = [5.0, 8, 8, 10, 8, 8, 5, 5]
      spec_params = (A, ω, Γ)

      global p = [
            0.01,       # k21, k31
            0.08,       # k23
            0.02,       # k24, k34
            0.04,       # k32
            0.005,      # k41
            0.05,       # x12, x21
            0.01,       # t0
            7.6,        # σ_pulse
      ]

      n_states = length(u0)

      k  = @MMatrix zeros(n_states, n_states)
      # this maps the flow between the states
      # row: from state
      # column: to state
      ak = @SMatrix [
            0 0 0 0
            1 0 2 3
            1 4 0 3
            5 0 0 0
      ]

      x  = @MMatrix zeros(n_states, n_states)
      ax = @SMatrix [
            0 6 0 0
            6 0 0 0
            0 0 0 0
            0 0 0 0
      ]

      t0 = p[end-1]
      σ  = p[end]

      # states represented in the spectrum
      s1 = SFGReservoirs.State()
      s2 = SFGReservoirs.State(2, [1,3])
      s3 = SFGReservoirs.State(3, [1,2])
      statez = [s1, s2, s3, s3, s1, s3, s2, s2]

      SFGReservoirs.parameters_to_rate_matrix!(k, ak, p)
      SFGReservoirs.parameters_to_rate_matrix!(x, ax, p)

      tspan = extrema(delaytimes)
      params = (k,x,t0,σ)
      prob = ODEProblem(model, u0, tspan, params)
      # saving at the delaytimes is performance critical
      @btime solve(prob; saveat=delaytimes)
      sol = solve(prob; saveat=delaytimes)

      # display(plot(sol))

      # preallocations
      Y = zeros(ComplexF64, (length(wavenumbers), length(delaytimes)))
      a = zeros(length(statez), length(delaytimes))

      @btime SFGReservoirs.solution_to_attenuated_spectra!(
            Y, a, wavenumbers, delaytimes, spec_params, sol, statez
      )
      @test Y[:,1] == [sfspectrum(x, A, ω, Γ) for x in wavenumbers]
      # display(heatmap(real.(Y)))

      # make synthetic data
      global syn_data = real.(Y)
      # heatmap(syn_data) |> display
end

begin
      # this is the data we provide from the experiment
      delaytimes = range(-30.0, 120.0, length=121)  |> collect
      wavenumbers = range(2830, 3000.0, length=400) |> collect
      data = (delaytimes, wavenumbers, syn_data)

      # from the fitting procedure we already have good estimations for
      # the parameters that define the SF spectrum.
      A = [0.2, 9, -3, -4, 9, -6, -1, -1]
      ω = [2845.0, 2876, 2886, 2905, 2938, 2947, 2961, 2966]
      Γ = [5.0, 8, 8, 10, 8, 8, 5, 5]
      spec_params = (A, ω, Γ)

      # we have to provide a guess for the initial population distribution
      u0 = [1.0, 0.0, 0.0, 0.0]

      # this is the initial guess for the parameters
      p_initial = [
            0.01,       # k21, k31
            0.01,       # k23
            0.01,       # k24, k34
            0.01,       # k32
            0.01,      # k41
            0.01,       # x12, x21
            0.01,       # t0 (position fixed @end-1)
            10,        # σ_pulse (position fixed @end)
      ]

      # assign the energy flow directions to parameters in the initial guess
      # the values in the matrix are the indices of the initial guess parameters_to_rate_matrix
      # array
      ak = @SMatrix [
            0 0 0 0
            1 0 2 3
            1 4 0 3
            5 0 0 0
      ]
      # we do the same for the parameter that is defining the pump strength
      ax = @SMatrix [
            0 6 0 0
            6 0 0 0
            0 0 0 0
            0 0 0 0
      ]
      assignments = (ak, ax)

      # states represented in the spectrum
      s1 = State()
      s2 = State(2, [1,3])
      s3 = State(3, [1,2])
      state_array = [s1, s2, s3, s3, s1, s3, s2, s2]

     result = SFGReservoirs.fit(data, spec_params, u0, p, assignments, state_array)
     p_result = result.minimizer
     @test isapprox(p_result[end], p[end]; atol=1e-3)
end


text:
using Eigen_jll
using boost_jll
using libcxxwrap_julia_jll
using libint_jll

function build_jeri()
  eigen_root = Eigen_jll.artifact_dir
  boost_root = boost_jll.artifact_dir
  jlcxx_root = libcxxwrap_julia_jll.artifact_dir
  libint_root = libint_jll.artifact_dir

  if ispath("build") rm("build", recursive=true) end
  mkdir("build")
  cd("build")
  run(`cmake
    -DEIGEN_PATH=$eigen_root -DBOOST_PATH=$boost_root
    -DJLCXX_PATH=$jlcxx_root -DLIBINT_PATH=$libint_root ../`)
  run(`make`)
  run(`make install`)
end

build_jeri()


text:
module Util

using Dates
using Infiltrator
using DataFrames

export checkname

"""
    LSTNLEVEL

Höchster zulässiger Abstand beim Namensvergleich
"""
const LSTNLEVEL = 0.3

const RGXINP = r" ?\([^\)]*\)"

hasdata(r) = !isnothing(r) && !ismissing(r) && !(r in ("", "(?)", "?"))

"""
    checkname(name::AbstractString, nameref::AbstractString)

Vergleiche Namen mithilfe der Levenshtein-Distanz
"""
function checkname(name::Union{Missing, AbstractString},
                   nameref::Union{Missing, AbstractString})::Bool
    if ismissing(name) || ismissing(nameref) || "" in (name, nameref)
        return false
    elseif name == nameref # Performance
        return true
    else
        # entferne Klammerausdrücke
        name = replace(name, RGXINP => "")
        nameref = replace(nameref, RGXINP => "")

        ld = levenshtein(name, nameref)
        # if ld > 0 && ld < 0.5 && FLOG
        #     global iolog
        #     println(iolog, name, "\t", nameref, "\t", ld)
        # end
        # println(name, "\t", nameref, "\t", ld)
        return ld < LSTNLEVEL
    end
end

"""
    splitname(name::AbstractString)

Teile `name` auf und entferne leere Felder und Ordnungszahlen
"""
function splitname(name::AbstractString)
    rgx = r"[^\w]+"

    filtercard(a) = filter(p -> !occursin(r"^[IVX]+$", p), a)
    filterempty(a) = filter(p -> p != "", a)
    split(name, rgx) |> filterempty |> filtercard
end

"""
    levenshtein(u::AbstractString, v::AbstractString)

Berechne die Levenshtein-Distanz zwischen `u` und `v`. Dividiere durch die Länge der
kürzeren Zeichenkette.
"""
function levenshtein(u::AbstractString, v::AbstractString)
    cau = codeunits(u)
    cav = codeunits(v)
    nu = length(cau)
    nv = length(cav)

    cmove = 0.4
    crepl = 1.0

    Dm = fill(1.0, (nu + 1, nv + 1))

    Dm[1, 1] = 0
    Dm[2:nu + 1, 1] = range(cmove, step = cmove, length = nu)
    Dm[1, 2:nv + 1] = range(cmove, step = cmove, length = nv)

    for i in 2:size(Dm, 1), j in 2:size(Dm, 2)
        c = cau[i-1] == cav[j-1] ? 0 : crepl
        # println(i, ", ", j, ", ", c)
        Dm[i, j] = min(Dm[i-1, j-1] + c,
                       Dm[i, j-1] + cmove, # Einfügung
                       Dm[i-1, j] + cmove) # Löschung
    end
    return Dm[end] / min(nu, nv)
    # return Dm[end]
end

"""
    rowselect(df, v, col)

Gib eine View auf `df` zurück, wo die Werte in `col` mit `v` übereinstimmen (Zahl)
oder wo `v` in der Spalte `col` von `df` vorkommt.
"""
function rowselect(df::AbstractDataFrame, v::Number, col::Symbol)
    # Sieht nicht so aus, ist aber vergleichsweise schnell
    if eltype(df[!, col]) == String
        v = string(v)
    end
    ix = v .== df[!, col]
    return @view df[ix, :]
end

function rowselect(df::AbstractDataFrame, v::Union{AbstractString, Regex}, col::Symbol)
    ix = occursin.(v, @view df[!, col])
    return @view df[ix, :]
end

function rowselect(df::AbstractDataFrame, v, dix::Dict{Int, T}) where T
    return @view df[get(dix, v, Int[]), :]
end

"""
    makeindex(df, colname)

Erzeuge einen Index auf `df` ausgehend vom Feld `col`.
"""
function makeindex(df, colname)
    col = df[!, colname]
    coltype = eltype(col)
    di = Dict{coltype, Array{Int, 1}}()
    ciep = unique(col)
    ix = BitArray(undef, nrow(df))
    for iep in ciep
        ix = col .== iep
        di[iep] = findall(ix)
    end
    return di
end

"""
    makedict(path::AbstractString, delim='\t')
"""
function makedict(path::AbstractString, delim='\t')
    dss = Dict{String, String}()
    open(path, "r") do io
        while !eof(io)
            line = readline(io)
            match(r"^[[:space:]]*$", line) != nothing && continue
            p = split(line, "\t")
            dss[p[1]] = p[2]
        end
    end
    dss
end

end


text:
using Polynomials, Combinatorics, LinearAlgebra, Random, Statistics

"""
The moment (Prony) mapping
"""
𝔐(x,c) = k->sum([c[i]*x[i]^k for i=1:length(x)])

"""
The Fourier mapping
"""
𝔉(x,c) = k->sum([c[i]*exp(im*x[i]*k) for i=1:length(x)])

"""
Given measurement function and an approximation to the nodes,
    solve the linear least squares system.
"""
function solveLinear(f,Ω,x̃)
    V = [exp(im*x*k) for k=0:Ω,x=x̃]
    d = [f(k) for k=0:Ω]
    V\d
end

"""
The inverse moment mapping
"""
function 𝔓(mvec)
    d2 = length(mvec)
    if (isodd(d2))
        error("Length of moment vector should be even")
    end

    d = d2>>1

    # construct the Hankel
    H = [mvec[i+j+1] for i=0:d-1,j=0:d-1]
    rhs = -mvec[d+1:end]
    q = H\rhs
    p = Polynomials.Poly(vcat(q,1))
    xrec=roots(p)
    V = [x^k for k=0:d-1,x=xrec]
    crec = V\mvec[1:d]
    xrec,crec
end

"""
Worst case perturbation function
"""
𝔉ᵪ(x,c,ϵ,s::Array) = begin
    R = length(x)
    x₀ = x[s]
    c₀ = c[s]

    # move to origin
    μ = (minimum(x₀)+maximum(x₀))/2
    y₀ = x₀.-μ

    m₀=[𝔐(y₀,c₀)(j+0.) for j=0:(2*length(s)-1)]
    #m₁ = copy(m₀)
    m₀[end] = m₀[end]+ϵ
    y₁,c₁ = 𝔓(m₀)

    # move back
    x₁ = y₁ .+ μ

    xₙ = copy(vec(x))
    xₙ[s]=real.(x₁)

    cₙ = convert(typeof(c₁),copy(vec(c)))
    cₙ[s] = c₁
    𝔉(xₙ,cₙ)
end

𝔉ᵪ(x,c,ϵ,ℓ::Integer) = 𝔉ᵪ(x,c,ϵ,collect(1:ℓ))
𝔉ᵪ(x,c,ϵ) = 𝔉ᵪ(x,c,ϵ,length(x))
𝔉ᵪ(ℓ) = (x,c,ϵ)->𝔉ᵪ(x,c,ϵ,ℓ)

"""
Random perturbation function
"""
𝔉ᵣ(x,c,ϵ) = k->(𝔉(x,c)(k) + ϵ*((-1+2*rand(MersenneTwister(Int(floor(k)))))+im*(-1+2*rand(MersenneTwister(Int(floor(k)))))))

𝔉ᵣ(Ωₘ::Integer,seed::Integer) = begin
    rng = MersenneTwister(seed)
    seq = -1 .+ rand(rng,Ωₘ+1)

    (x,c,ϵ)->(k->(𝔉(x,c)(k)+ϵ*seq[Int(floor(k+1))]))
end

Hank(F,L,M) = [F(i+j) for i=1:L,j=1:M]

"""
The matrix pencil algorithm for recovery of complex exponentials from equispaced samples

F - function handle for sampling from the model
N - number of equispaced samples
n - the model order
L - the pencil parameter L (default = ceil(N/2))

Returns: array of n estimated frequencies, normalized to [0,2π]
"""
function MatrixPencil(F,N,n,L)
    # check we have enough samples
    if(N < 2*n)
        error("N should be at least 2n")
    end

    if (L < n+1)
        error("L should be at least n+1")
    end

    # construct the matrices H₁ and H₂
    M = N-L
    H = Hank(F,L,M)

    Rₗ = H[2:end,:] # lower part
    Rᵤ = H[1:end-1,:] # upper part

    # compute truncated SVDs
    S₁ = svd(Rᵤ;full=true)
    u₁ = S₁.U[:,1:n]
    Σ₁ = Diagonal(S₁.S[1:n])
    v₁ = S₁.V[:,1:n]

    S₂ = svd(Rₗ;full=true)
    u₂ = S₂.U[:,1:n]
    Σ₂ = Diagonal(S₂.S[1:n])
    v₂ = S₂.V[:,1:n]

    # compute the generalized eigenvalues of the corresponding *regular* pencil
    A = (u₂)'*u₁*Σ₁*(v₁)'*v₂
    B = Σ₂
    z = eigvals(B\A)
    mod.(2π.-angle.(z),2π)

end

MatrixPencil(F,N,n) = MatrixPencil(F,N,n,ceil(N/2))

"""
Amplitude function: all 1's
"""
all1s(ℓ,R) = ones(R)

"""
Amplitude function: i^j
"""
imag1s(ℓ,R) = [(im)^j for j=0:R-1]

"""
Amplitude function: plus minus 1
"""
opposite1s(ℓ,R) = [(-1)^j for j=0:R-1]

"""
Amplitude function: complex amplitudes
"""
mixed1s(ℓ,R) = [1+im*(-1)^j for j=0:R-1]

"""
Creates a clustered node configuration with a single cluster
and other nodes well-separated
(equidistributed in the remaining part of the circle)
"""
Ζ(ℓ,R,Δ,endpoint=π) = begin

    if (ℓ>R)
        error("ℓ should be at most R")
    end

    ρ = (ℓ-1)*Δ
    if (endpoint < ρ)
        error("Δ is too large")
    end

    δ = (endpoint-ρ)/(R-ℓ+1)

    r₁ = collect(range(0,step=Δ,length=ℓ))
    r₂ = collect(range(ρ+δ,step=δ,length=R-ℓ))
    [r₁;r₂].-π/2
end

distCirc(x₀::Number,x₁::Number) = min(mod(x₀-x₁,2π),mod(x₁-x₀,2π))

function housedorff_distance(test_points, candidates)
    n = length(test_points)
    m = length(candidates)
    @assert m==n "Err"
    t = reshape(test_points,(n,1))
    c = reshape(candidates,(n,1))
    return minimum([maximum(abs.([tt[i]-c[i] for i=1:n])) for tt in permutations(t)])
end

"""
Computes the minimal separation between the nodes in an array
"""
function min_sep(X::Array)
    # distance matrix
    D = [distCirc(x,y) for x in X, y in X]
    return minimum(D+diagm(Inf*ones(length(X))))
end

get_separation(X::Array) = [minimum(
            [distCirc(X[j],X[k]) for k in setdiff(1:length(X),j)])
        for j=1:length(X)]

get_separation(X::Array,Y::Array) = [minimum(
            [distCirc(X[j],Y[k]) for k in 1:length(Y)])
        for j=1:length(X)]

get_separation_ext(X::Array,Y::Array) = [findmin(
            [distCirc(X[j],Y[k]) for k in 1:length(Y)])
        for j=1:length(X)]

"""
Test the Matrix Pencil algo: maximal error
"""
testMP(ℓ,R,Δ,Ω,ϵ,clusterFun=Ζ,pertFun=𝔉ᵪ,amplitudeFun = all1s) = begin
    x = clusterFun(ℓ,R,Δ)
    c = amplitudeFun(ℓ,R)
    f = pertFun(x,c,ϵ)
    ϵ₀ = maximum(abs.([𝔉(x,c)(k+0.) - f(k+0.) for k=0:Ω]))
    x̃ = MatrixPencil(f,Ω,R)
    c̃ = solveLinear(f,Ω,x̃)
    Δx = housedorff_distance(exp.(im.*x),exp.(im.*x̃))/ϵ₀
    Δc = housedorff_distance(c,c̃)/ϵ₀

    # need to match the first two
    η = get_separation(x) # distance to closest neighbor
    σ = get_separation(x,x̃) # error for each node

    # success if all errors are smaller than corresponding separations
    succ = reduce((x,y)->x&&y,σ .< η/3)
    [Ω*Δx,Δc,Int(succ),ϵ₀]
end

"""
Test the Matrix Pencil algo: a specific node
"""
testMP_individual(ℓ,R,Δ,Ω,ϵ,node_index,clusterFun=Ζ,pertFun=𝔉ᵪ,amplitudeFun = all1s) = begin
    x = clusterFun(ℓ,R,Δ)
    c = amplitudeFun(ℓ,R)
    f = pertFun(x,c,ϵ)
    ϵ₀ = maximum(abs.([𝔉(x,c)(k+0.) - f(k+0.) for k=0:Ω]))
    x̃ = MatrixPencil(f,Ω,R)
    c̃ = solveLinear(f,Ω,x̃)

    # need to match the first two
    η = get_separation(x) # distance to closest neighbor
    σ̃ = get_separation_ext(x,x̃) # error for each node, including index
    Δx = σ̃[node_index][1]
    idx = σ̃[node_index][2]
    Δc = (abs(c̃[idx] - c[node_index]))

    succ = (σ̃[node_index][1] < η[node_index]/3)

    # success if all errors are smaller than corresponding separations
    [Ω*Δx/ϵ₀,Δc/ϵ₀,Int(succ),ϵ₀]
end

"""
Test the Matrix Pencil algorithm: accuracy for every node
"""
testMP_all(ℓ,R,Δ,Ω,ϵ,clusterFun=Ζ,pertFun=𝔉ᵪ,amplitudeFun = all1s) = begin
    x = clusterFun(ℓ,R,Δ)
    c = amplitudeFun(ℓ,R)
    f = pertFun(x,c,ϵ)
    ϵ₀ = maximum(abs.([𝔉(x,c)(k+0.) - f(k+0.) for k=0:Ω]))
    x̃ = MatrixPencil(f,Ω,R)
    c̃ = solveLinear(f,Ω,x̃)

    # need to match the first two
    η = get_separation(x) # distance to closest neighbor
    σ̃ = get_separation_ext(x,x̃) # error for each node, including index
    Δx = [sep[1] for sep=σ̃]
    Δc = [abs(c̃[σ̃[i][2]]-c[i]) for i=1:R]

    # success if all errors are smaller than corresponding separations
    succ = [Int(σ̃[i][1] < η[i]/3) for i=1:R]

    [Ω*Δx/ϵ₀,Δc/ϵ₀,succ,ϵ₀]
end

NrangeLog(low,high,n=100) = Int.(unique(floor.(10.0.^LinRange(low,high,n))))

"""
Run a single random experiment testing accuracy
"""
runSingle_Acc(ℓ,R,clusterFun, pertFun, amplitudeFun, ϵR,ΩR,ΔR) = begin
    ϵ = ϵR[rand(1:end)]
    Ω = ΩR[rand(1:end)]
    Δ = ΔR[rand(1:end)]

    res = testMP_all(ℓ,R,Δ,Ω,ϵ,clusterFun,pertFun,amplitudeFun)
    [inv(Ω*Δ),res...]
end

"""
Run a single random experiment testing accuracy
"""
runSingle(ℓ,R,clusterFun, pertFun, amplitudeFun, ϵR,ΩR,ΔR) = begin
    ϵ = ϵR[rand(1:end)]
    Ω = ΩR[rand(1:end)]
    Δ = ΔR[rand(1:end)]

    res = testMP(ℓ,R,Δ,Ω,ϵ,clusterFun,pertFun,amplitudeFun)

    [inv(Ω*Δ),res[3:4]...]
end

runSingle_ind(ℓ,R,clusterFun, pertFun, amplitudeFun, ni, ϵR,ΩR,ΔR) = begin
    ϵ = ϵR[rand(1:end)]
    Ω = ΩR[rand(1:end)]
    Δ = ΔR[rand(1:end)]

    res = testMP_individual(ℓ,R,Δ,Ω,ϵ,ni,clusterFun,pertFun,amplitudeFun)

    [inv(Ω*Δ),res[3:4]...]
end

cc(x) = (x==0) ? (:red) : (:blue)
ss(x) = (x==0) ? (:circle) : (:dtriangle)

max_stars_repo_name: batenkov-team/iai-code

max_issues_repo_name: batenkov-team/iai-code

max_forks_repo_name: batenkov-team/iai-code


text:
# Transformations
# ===============
#
# Methods that manipulate and change a biological sequence.
#
# This file is a part of BioJulia.
# License is MIT: https://github.com/BioJulia/BioSequences.jl/blob/master/LICENSE.md

"""
    empty!(seq::BioSequence)

Completely empty a biological sequence `seq` of nucleotides.
"""
Base.empty!(seq::BioSequence) = resize!(seq, 0)

"""
    push!(seq::BioSequence, x)

Append a biological symbol `x` to a biological sequence `seq`.
"""
function Base.push!(seq::BioSequence, x)
    x_ = convert(eltype(seq), x)
    resize!(seq, length(seq) + 1)
    unsafe_setindex!(seq, x_, lastindex(seq))
    return seq
end

"""
    pop!(seq::BioSequence)

Remove the symbol from the end of a biological sequence `seq` and return it.
Returns a variable of `eltype(seq)`.
"""
function Base.pop!(seq::BioSequence)
    if isempty(seq)
        throw(ArgumentError("sequence must be non-empty"))
    end
    @inbounds x = seq[end]
    deleteat!(seq, lastindex(seq))
    return x
end

"""
    insert!(seq::BioSequence, i, x)

Insert a biological symbol `x` into a biological sequence `seq`, at the given
index `i`.
"""
function Base.insert!(seq::BioSequence, i::Integer, x)
    checkbounds(seq, i)
    resize!(seq, length(seq) + 1)
    copyto!(seq, i + 1, seq, i, lastindex(seq) - i)
    unsafe_setindex!(seq, x, i)
    return seq
end

"""
    deleteat!(seq::BioSequence, range::UnitRange{<:Integer})

Deletes a defined `range` from a biological sequence `seq`.

Modifies the input sequence.
"""
function Base.deleteat!(seq::BioSequence, range::UnitRange{<:Integer})
    checkbounds(seq, range)
    copyto!(seq, range.start, seq, range.stop + 1, length(seq) - range.stop)
    resize!(seq, length(seq) - length(range))
    return seq
end

"""
    deleteat!(seq::BioSequence, i::Integer)

Delete a biological symbol at a single position `i` in a biological sequence
`seq`.

Modifies the input sequence.
"""
function Base.deleteat!(seq::BioSequence, i::Integer)
    checkbounds(seq, i)
    copyto!(seq, i, seq, i + 1, length(seq) - i)
    resize!(seq, length(seq) - 1)
    return seq
end

"""
    append!(seq, other)

Add a biological sequence `other` onto the end of biological sequence `seq`.
Modifies and returns `seq`.
"""
function Base.append!(seq::BioSequence, other::BioSequence)
    resize!(seq, length(seq) + length(other))
    copyto!(seq, lastindex(seq) - length(other) + 1, other, 1, length(other))
    return seq
end

"""
    popfirst!(seq)

Remove the symbol from the beginning of a biological sequence `seq` and return
it. Returns a variable of `eltype(seq)`.
"""
function Base.popfirst!(seq::BioSequence)
    if isempty(seq)
        throw(ArgumentError("sequence must be non-empty"))
    end
    @inbounds x = seq[1]
    deleteat!(seq, 1)
    return x
end

"""
    pushfirst!(seq, x)

Insert a biological symbol `x` at the beginning of a biological sequence `seq`.
"""
function Base.pushfirst!(seq::BioSequence, x)
    resize!(seq, length(seq) + 1)
    copyto!(seq, 2, seq, 1, length(seq) - 1)
    unsafe_setindex!(seq, x, 1)
    return seq
end

Base.filter(f::Function, seq::BioSequence) = filter!(f, copy(seq))
Base.map(f::Function, seq::BioSequence) = map!(f, copy(seq))
Base.reverse(seq::BioSequence) = reverse!(copy(seq))

"""
    complement(seq)

Make a complement sequence of `seq`.
"""
function BioSymbols.complement(seq::NucleotideSeq)
    return complement!(copy(seq))
end

"""
    reverse_complement!(seq)

Make a reversed complement sequence of `seq` in place.
"""
function reverse_complement!(seq::NucleotideSeq)
    return complement!(reverse!(seq))
end

"""
    reverse_complement(seq)

Make a reversed complement sequence of `seq`.
"""
function reverse_complement(seq::NucleotideSeq)
    return complement!(reverse(seq))
end

"""
    canonical!(seq::NucleotideSeq)

Transforms the `seq` into its canonical form, if it is not already canonical.
Modifies the input sequence inplace.

For any sequence, there is a reverse complement, which is the same sequence, but
on the complimentary strand of DNA:

```
------->
ATCGATCG
CGATCGAT
<-------
```

!!! note
    Using the [`reverse_complement`](@ref) of a DNA sequence will give give this
    reverse complement.

Of the two sequences, the *canonical* of the two sequences is the lesser of the
two i.e. `canonical_seq < other_seq`.

Using this function on a `seq` will ensure it is the canonical version.
"""
function canonical!(seq::NucleotideSeq)
    if !iscanonical(seq)
        reverse_complement!(seq)
    end
    return seq
end

"""
    canonical(seq::NucleotideSeq)

Create the canonical sequence of `seq`.

"""
canonical(seq::NucleotideSeq) = canonical!(copy(seq))

"Create a copy of a sequence with gap characters removed."
ungap(seq::BioSequence)  =  filter(!isgap, seq)

"Remove gap characters from an input sequence."
ungap!(seq::BioSequence) = filter!(!isgap, seq)

###
### Shuffle
###

function Random.shuffle(seq::BioSequence)
    return shuffle!(copy(seq))
end


text:
module Fast_Hertz

using JuLIP
using JuLIP.Potentials: pairs
import JuLIP: energy, forces, cutoff, rnn
using NeighbourLists
using StaticArrays
using LinearAlgebra

const SVecG{T} = SVector{3, T}
const SMatG{T} = SMatrix{3, 3, T, 9}

export Fast_Hertz_Potential

# @pot mutable struct Fast_HP <: AbstractCalculator
mutable struct Fast_HP <: AbstractCalculator
    k::Float64
    α::Float64
    rcut::Float64    # cutoff when calculating neighbourlist
    rtol2::Float64   # if max_displacement > rtol, recalculate the neighbourlist rtol2 = rtol^2
    rsum::Vector{Float64} # rsum = [2*rmin; rmin + rmax; 2*rmax]
    Lhalf::Float64
#    Force::Vector{JVecF}
end

cutoff(V::Fast_HP) = V.rcut
## define

function Fast_Hertz_Potential(L, r0::Float64, r1::Float64, α::Float64; k = 1.0, rcutfact = 2.5)
    rcut = rcutfact * r0 # rcut is rcutfact*2*R_{max}, used when searching possible pairs
    rtol = 0.25 * rcut    # rtol is the neighbourlist recomputing tolerance
    rtol2 = rtol * rtol
    rsum = [r1; (r0 + r1)/2; r0]
    Lhalf = L/2.0
    return Fast_HP(k, α, rcut, rtol2, rsum, Lhalf)
end

function energy(V::Fast_HP, at::AbstractAtoms)
    nlist_i, nlist_j, nlist_Z = update_neighbourlist!(V, at)
    bins = at.cell'
    inv_bins = inv(bins)
    return energy_inner!(nlist_i, nlist_j, nlist_Z, at, V, bins, inv_bins)
end

# compute the maximum displacement of all particles
@inline max_dr2(X::Vector{JVecF}, X_Old::Vector{JVecF}) =
                    maximum( [ sum( (X[i] - X_Old[i]) .* (X[i] - X_Old[i]) ) for i = 1:length(X) ] )

@inline PairType(i::Vector{Int32}, j::Vector{Int32}, Z::Vector{Int64}) = [ Z[i[k]] + Z[j[k]] - 1 for k = 1:length(i) ]

function update_neighbourlist!(V::Fast_HP, at::AbstractAtoms)
    if has_data(at, :nlist_X)
        X_Old = get_data(at, :nlist_X)::Vector{JVecF}

        X = at.X::Vector{JVecF}
        max_d = max_dr2(X, X_Old)
        if max_d < V.rtol2
            nlist_i = get_data(at, :nlist_i)::Vector{Int32}
            nlist_j = get_data(at, :nlist_j)::Vector{Int32}
            nlist_Z = get_data(at, :nlist_Z)::Vector{Int64}
            return nlist_i, nlist_j, nlist_Z
        end
    end
    nlist = neighbourlist(at, V.rcut)
    temp = (nlist.j - nlist.i) .> 0
    nlist_i = nlist.i[temp]
    nlist_j = nlist.j[temp]
    nlist_Z = PairType(nlist_i, nlist_j, at.Z)
    set_data!(at, :nlist_X, copy(at.X))
    set_data!(at, :nlist_i, nlist_i)
    set_data!(at, :nlist_j, nlist_j)
    set_data!(at, :nlist_Z, nlist_Z)
    return nlist_i, nlist_j, nlist_Z
end

@inline wrap_to_unit(inv_bins::SMatG{T}, x::SVecG{T}) where {T <: Real} = mod.(inv_bins * x, 1)
@inline adjust_dx(bins::SMatG{T}, dx::SVecG{T}) where {T <: Real} = bins * (dx - (dx .> 0.5))
@inline dx2dx(dx::SVecG{T}, inv_bins::SMatG{T}, bins::SMatG{T}) where {T <: Real} =
                                    adjust_dx(bins, wrap_to_unit(inv_bins, dx))

function energy_inner!(nlist_i, nlist_j, nlist_Z, X, V::Fast_HP, bins, inv_bins)
    k = V.k
    α = V.α
    E = 0.0
#    F = zeros(JVecF, length(at))
    for k = 1:length(nlist_i)
        i, j = nlist_i[k], nlist_j[k]
        dx = X[i] - X[j]
        r = norm(dx)
        if r > V.Lhalf
            dx = dx2dx(dx, inv_bins, bins)
            r = norm(dx)
        end
        nlist_s = V.rsum[nlist_Z[k]] - r
        if nlist_s > 0
            Eij = nlist_s^α
            E += Eij
            # Fij = α*Eij/nlist_s/r*dx

            # # E += nlist_s * nlist_s
            # # Fij = 2.0*nlist_s/r*dx

            # F[i] -= Fij
            # F[j] += Fij
        end
#        V.Force = F
    end
    # const = -Alpha*E0*temp.^(Alpha-1)./dij;

    # dV(1:Num) = dV(1:Num) + accumarray(i, const.*dx, [Num 1]);
    # dV(1+Num:end) = dV(1+Num:end) + accumarray(i, const.*dy, [Num 1]);
    # dV(1:Num) = dV(1:Num) - accumarray(j, const.*dx, [Num 1]);
    # dV(1+Num:end) = dV(1+Num:end) - accumarray(j, const.*dy, [Num 1]);

    return E
end

function forces(V::Fast_HP, at::AbstractAtoms)
    nlist_i, nlist_j, nlist_Z = update_neighbourlist!(V, at)
    bins = at.cell'
    inv_bins = inv(bins)
    F = zeros(JVecF, length(at))
    force_inner!(F, nlist_i, nlist_j, nlist_Z, at.X, V::Fast_HP, bins, inv_bins)
    return F
end

function force_inner!(F, nlist_i, nlist_j, nlist_Z, X, V::Fast_HP, bins, inv_bins)
    α = V.α
    for k = 1:length(nlist_i)
        i, j = nlist_i[k], nlist_j[k]
        dx = X[i] - X[j]
        r = norm(dx)
        if r > V.Lhalf
            dx = dx2dx(dx, inv_bins, bins)
            r = norm(dx)
        end
        nlist_s = V.rsum[nlist_Z[k]] - r
        if nlist_s > 0
            Fij = α*nlist_s^(α - 1)/r*dx

            #Fij = 2.0*nlist_s/r*dx

            F[i] += Fij
            F[j] -= Fij
        end
    end
    return F
end

end


text:
### Definitions of datastructures and basic constructors and operators
# Written by Boris Bonev, Jan. 2021

"""
    HssMatrix{T<:Number} <: AbstractMatrix{T}

Matrix type for storing hierarchically semi-separable matrices.
These matrices can be constructed via the smart contructor [`hss`](@ref) function or
the compression methods [`compress`](@ref), [`randcompress`](@ref) and [`randcompress_adaptive`](@ref)
"""
mutable struct HssMatrix{T<:Number} <: AbstractMatrix{T}
  # toggles for the type of node
  leafnode::Bool
  rootnode::Bool

  # fields for leaf nodes
  D ::Matrix{T}
  U ::Matrix{T}
  V ::Matrix{T}

  # fields for branch nodes
  A11 ::HssMatrix{T}
  A22 ::HssMatrix{T}
  B12 ::Matrix{T}
  B21 ::Matrix{T}

  sz1 ::Tuple{Int, Int}
  sz2 ::Tuple{Int, Int}

  R1 ::Matrix{T}
  W1 ::Matrix{T}
  R2 ::Matrix{T}
  W2 ::Matrix{T}

  # internal constructors for leaf nodes
  function HssMatrix(D::AbstractMatrix{T}, rootnode=true) where T
    m, n = size(D)
    return new{T}(true, rootnode, D, Matrix{T}(undef, m, 0), Matrix{T}(undef, n, 0))
  end
  function HssMatrix(D::AbstractMatrix{T}, U::AbstractMatrix{T}, V::AbstractMatrix{T}, rootnode=false) where T
    if size(D,1) != size(U,1) throw(ArgumentError("D and U must have same number of rows")) end
    if size(D,2) != size(V,1) throw(ArgumentError("D and V must have same number of columns")) end
    return new{T}(true, rootnode, D, U, V)
  end
  # internal constructors for branch nodes
  function HssMatrix(A11::HssMatrix{T}, A22::HssMatrix{T}, B12::AbstractMatrix{T}, B21::AbstractMatrix{T}, rootnode=true) where T
    kr1, kw1 = gensize(A11); kr2, kw2 = gensize(A22)
    hssA = new{T}(false, rootnode)
    hssA.A11 = A11; hssA.A22 = A22
    hssA.B12 = B12; hssA.B21 = B21
    hssA.sz1 = size(A11); hssA.sz2 = size(A22)
    hssA.R1 = Matrix{T}(undef,kr1,0); hssA.W1 = Matrix{T}(undef,kw1,0)
    hssA.R2 = Matrix{T}(undef,kr2,0); hssA.W2 = Matrix{T}(undef,kw2,0)
    return hssA
  end
  function HssMatrix(A11::HssMatrix{T}, A22::HssMatrix{T}, B12::AbstractMatrix{T}, B21::AbstractMatrix{T},
    R1::AbstractMatrix{T}, W1::AbstractMatrix{T}, R2::AbstractMatrix{T}, W2::AbstractMatrix{T}, rootnode=false) where T
    if size(R1,2) != size(R2,2) throw(DimensionMismatch("R1 and R2 must have same number of columns")) end
    if size(W1,2) != size(W2,2) throw(DimensionMismatch("W1 and W2 must have same number of rows")) end
    hssA = new{T}(false, rootnode)
    hssA.A11 = A11; hssA.A22 = A22
    hssA.B12 = B12; hssA.B21 = B21
    hssA.sz1 = size(A11); hssA.sz2 = size(A22)
    hssA.R1 = R1; hssA.W1 = W1
    hssA.R2 = R2; hssA.W2 = W2
    return hssA
  end
end

# custom constructors which are calling the compression algorithms
function hss(A::AbstractMatrix, opts=HssOptions(Float64); args...)
  opts = copy(opts; args...)
  chkopts!(opts)
  return hss(A, bisection_cluster(size(A,1), leafsize=opts.leafsize), bisection_cluster(size(A,2), leafsize=opts.leafsize); args...)
end
hss(A::Matrix, rcl::ClusterTree, ccl::ClusterTree; args...) = compress(A, rcl, ccl; args...)
function hss(A::AbstractSparseMatrix, rcl::ClusterTree, ccl::ClusterTree; args...)
  m, n = size(A)
  # estimate rank by assuming that the non-zero entries are clustered on the diagonal
  nl = nleaves(rcl)
  m0 = Int(ceil(m/nl))
  n0 = Int(ceil(n/nl))
  kest = max(nnz(A) - nl*m0*n0,0)
  return randcompress_adaptive(A, rcl, ccl; kest=kest, args...)
end
hss(A::LinearMap, rcl::ClusterTree, ccl::ClusterTree; args...) = randcompress_adaptive(A, rcl, ccl; args...)

@inline isleaf(hssA::HssMatrix) = hssA.leafnode # check whether making this inline speeds up things ?
@inline isbranch(hssA::HssMatrix) = !hssA.leafnode
@inline isroot(hssA::HssMatrix) = hssA.rootnode

@inline ishss(A::AbstractMatrix) = typeof(A) <: HssMatrix

size(hssA::HssMatrix) = isleaf(hssA) ? size(hssA.D) : hssA.sz1 .+ hssA.sz2
size(hssA::HssMatrix, dim::Integer) = size(hssA)[dim]

show(io::IO, hssA::HssMatrix) = print(io, "$(size(hssA,1))x$(size(hssA,2)) HssMatrix{$(eltype(hssA))}")

# implement sorted access to entries via recursion
# in the long run we might want to return an HssMatrix when we access via getindex
# TODO: add @boundscheck for the bound checking
getindex(hssA::HssMatrix, i::Int, j::Int) = _getidx(hssA, i, j)[1]
getindex(hssA::HssMatrix, i::Int, j::AbstractRange) = getindex(hssA, [i], j)[:]
getindex(hssA::HssMatrix, i::AbstractRange, j::Int) = getindex(hssA, i, [j])[:]
getindex(hssA::HssMatrix, i::AbstractRange, j::AbstractRange) = getindex(hssA, collect(i), collect(j))
getindex(hssA::HssMatrix, ::Colon, ::Colon) = full(hssA)
getindex(hssA::HssMatrix, i, ::Colon) = getindex(hssA, i, 1:size(hssA,2))
getindex(hssA::HssMatrix, ::Colon, j) = getindex(hssA, 1:size(hssA,1), j)
function getindex(hssA::HssMatrix{T}, i::Vector{Int}, j::Vector{Int}) where T
  m, n  = size(hssA)
  ip = sortperm(i); jp = sortperm(j)
  if (length(i) == 0 || length(j) == 0) return Matrix{T}(undef, length(i), length(j)) end
  return full(_getidx(hssA, i[ip], j[jp]))[invperm(ip), invperm(jp)]
end

# First construct a sub-HSS matrix and then call full()
function _getidx(hssA::HssMatrix, i::Vector{Int}, j::Vector{Int})
  if isleaf(hssA)
    HssMatrix(hssA.D[i,j], hssA.U[i,:], hssA.V[j,:])
  else
    m1, n1 = hssA.sz1
    i1 = i[i .<= m1]; j1 = j[j .<= n1]
    i2 = i[i .> m1] .- m1; j2 = j[j .> n1] .- n1
    A11 = _getidx(hssA.A11, i1, j1)
    A22 = _getidx(hssA.A22, i2, j2)
    return HssMatrix(A11, A22, hssA.B12, hssA.B21, hssA.R1, hssA.W1, hssA.R2, hssA.W2)
  end
end
# for individual indices, it should be faster to access them in the following way
function _getidx(hssA::HssMatrix, i::Int, j::Int)
  if isleaf(hssA)
    return hssA.D[i,j]
  else
    m1, n1 = hssA.sz1
    if i <= m1
      if j <= n1
        return _getidx(hssA.A11, i, j)
      else
        U1 = _getindex_colgenerator(hssA.A11, i)
        V2 = _getindex_rowgenerator(hssA.A22, j-n1)
        return dot(U1*hssA.B12, V2)
      end
    else
      if j <= n1
        U2 = _getindex_colgenerator(hssA.A22, i-m1)
        V1 = _getindex_rowgenerator(hssA.A11, j)
        return dot(U2*hssA.B21, V1)
      else
        return _getidx(hssA.A22, i-m1, j-n1)
      end
    end
  end
end

# perhaps, this should be deepcopy?
function copy(hssA::HssMatrix)
  if isleaf(hssA)
    return HssMatrix(copy(hssA.D), copy(hssA.U), copy(hssA.V), hssA.rootnode)
  else
    return HssMatrix(copy(hssA.A11), copy(hssA.A22), copy(hssA.B12), copy(hssA.B21), copy(hssA.R1), copy(hssA.W1), copy(hssA.R2), copy(hssA.W2), hssA.rootnode)
  end
end

# introduce lazy adjoint, then all this copying won't be necessary
function adjoint(hssA::HssMatrix)
  if isleaf(hssA)
    return HssMatrix(copy(adjoint(hssA.D)), copy(hssA.V), copy(hssA.U), hssA.rootnode)
  else
    return HssMatrix(adjoint(hssA.A11), adjoint(hssA.A22), copy(adjoint(hssA.B21)), copy(adjoint(hssA.B12)), copy(hssA.W1), copy(hssA.R1), copy(hssA.W2), copy(hssA.R2), hssA.rootnode)
  end
end

# transpose
function transpose(hssA::HssMatrix)
  if isleaf(hssA)
    return HssMatrix(copy(transpose(hssA.D)), copy(hssA.V), copy(hssA.U), hssA.rootnode)
  else
    return HssMatrix(transpose(hssA.A11), transpose(hssA.A22), copy(transpose(hssA.B21)), copy(transpose(hssA.B12)), copy(hssA.W1), copy(hssA.R1), copy(hssA.W2), copy(hssA.R2), hssA.rootnode)
  end
end

# Define Matlab-like convenience functions, which are used throughout the library
#blkdiagm(A::Matrix, B::Matrix) = [A zeros(size(A,1), size(B,2)); zeros(size(B,1), size(A,2)) B]
#blkdiagm(A::Matrix... ) = blkdiagm(A[1], blkdiagm(A[2:end]...))
#blkdiagm(A::Matrix{T}...) where T = cat(A[1:end]..., dims=(1,2))

function blkdiagm(As::Matrix{T}...) where T
  # pre-allocate memory
  m = 0; n = 0
  for A in As
    m = m + size(A,1)
    n = n + size(A,2)
  end
  #B = similar(As[1], m, n)
  B = zeros(T, m, n)
  ro = 0; co = 0
  @simd for A in As
    m, n = size(A)
    @simd for j in 1:n
      @simd for i in 1:m
        @inbounds B[ro+i, co+j] = A[i, j]
      end
    end
    ro = ro + m
    co = co + n
  end
  return B
end

## basic algebraic operations (taken and modified from LowRankApprox.jl)
for op in (:+,:-)
  @eval begin
    $op(hssA::HssMatrix) = isleaf(hssA) ? HssMatrix($op(hssA.D), hssA.U, hssA.V) : HssMatrix($op(hssA.A11), $op(hssA.A22), $op(hssA.B12), $op(hssA.B21), hssA.R1, hssA.W1, hssA.R2, hssA.W2)

    $op(a::Bool, hssA::HssMatrix{Bool}) = error("Not callable")
    $op(L::HssMatrix{Bool}, a::Bool) = error("Not callable")

    function $op(hssA::HssMatrix, hssB::HssMatrix)
      size(hssA) == size(hssB) || throw(DimensionMismatch("A has dimensions $(size(hssA)) but B has dimensions $(size(hssB))"))
      if isleaf(hssA) && isleaf(hssB)
        return HssMatrix($op(hssA.D, hssB.D), [hssA.U hssB.U], [hssA.V hssB.V], isroot(hssA) && isroot(hssB))
      elseif isbranch(hssA) && isbranch(hssB)
        hssA.sz1 == hssB.sz1 || throw(DimensionMismatch("A11 has dimensions $(hssA.sz1) but B11 has dimensions $(hssB.sz1)"))
        hssA.sz2 == hssB.sz2 || throw(DimensionMismatch("A22 has dimensions $(hssA.sz2) but B22 has dimensions $(hssA.sz2)"))
        return HssMatrix($op(hssA.A11, hssB.A11), $op(hssA.A22, hssB.A22), blkdiagm(hssA.B12, $op(hssB.B12)), blkdiagm(hssA.B21, $op(hssB.B21)),blkdiagm(hssA.R1, hssB.R1), blkdiagm(hssA.W1, hssB.W1), blkdiagm(hssA.R2, hssB.R2), blkdiagm(hssA.W2, hssB.W2), isroot(hssA) && isroot(hssB))
      else
        throw(ArgumentError("Cannot add leaves and branches."))
      end
    end
  end
end

# matrix division involving HSS matrices
\(hssA::HssMatrix, B::Matrix) = ulvfactsolve(hssA, B)
\(hssA::HssMatrix, hssB::HssMatrix) = ldiv!(hssA, copy(hssB))
/(A::Matrix, hssB::HssMatrix) = ulvfactsolve(hssB', collect(A'))'
/(hssA::HssMatrix, hssB::HssMatrix) = rdiv!(copy(hssA), hssB)

# Scalar multiplication
*(hssA::HssMatrix, a::Number) = *(a, hssA)
function *(a::Number, hssA::HssMatrix)
  if isleaf(hssA)
    return HssMatrix(a*hssA.D, hssA.U, hssA.V, isroot(hssA))
  else
    return HssMatrix(a*hssA.A11, a*hssA.A22, a*hssA.B12, a*hssA.B21, hssA.R1, hssA.W1, hssA.R2, hssA.W2, isroot(hssA))
  end
end

## Some more fundamental operations
# compute the HSS rank
hssrank(hssA::HssMatrix) = isleaf(hssA) ? 0 : max(hssrank(hssA.A11), hssrank(hssA.A22), size(hssA.B12)..., size(hssA.B21)...)

# compute the sizeof the underlying generators
function gensize(hssA::HssMatrix)
  if isleaf(hssA)
    return size(hssA.U,2), size(hssA.V,2)
  else
    (kr = size(hssA.R1,2)) == size(hssA.R2,2) || throw(DimensionMismatch("dimensions of column-translators do not match"))
    (kw = size(hssA.W1,2)) == size(hssA.W2,2) || throw(DimensionMismatch("dimensions of row-translators do not match"))
    return kr, kw
  end
end

# function that returns alternative HssNode acting as rootnode
# this makes multiplication etc. safe if we use subblocks of HSS matrices
rooted(hssA::HssMatrix) = isleaf(hssA) ? HssMatrix(hssA.D, true) : HssMatrix(hssA.A11, hssA.A22, hssA.B12, hssA.B21, true)

# return a full matrix (hopefully efficient implementation with pre-allocated memory)
Matrix(hssA::HssMatrix) = full(hssA)
function full(hssA::HssMatrix{T}) where T
  if isleaf(hssA)
    return hssA.D
  else
    m, n = size(hssA)
    k = hssrank(hssA)
    A = Matrix{T}(undef, m, n)
    U = Matrix{T}(undef, m, k)
    V = Matrix{T}(undef, n, k)
    _full!(hssA, A, U, V, 0, 0; rootnode=true)
  end
  return A
end
# function that expands hss matrix using pre-allocated memory
function _full!(hssA::HssMatrix, A::Matrix,  U::Matrix, V::Matrix, ro::Int, co::Int; rootnode=false)
  if isleaf(hssA)
    m, n = size(hssA)
    A[ro+1:ro+m, co+1:co+n] = hssA.D
    U[ro+1:ro+m, 1:size(hssA.U,2)] = hssA.U
    V[co+1:co+n, 1:size(hssA.V,2)] = hssA.V
  else
    m1, n1 = hssA.sz1; m2, n2 = hssA.sz2
    ru1, rv1 = gensize(hssA.A11)
    ru2, rv2 = gensize(hssA.A22)
    _full!(hssA.A11, A, U, V, ro, co)
    _full!(hssA.A22, A, U, V, ro+m1, co+n1)
    A[ro+1:ro+m1, co+n1+1:co+n1+n2] = U[ro+1:ro+m1, 1:ru1]*hssA.B12*V[co+n1+1:co+n1+n2, 1:rv2]'
    A[ro+m1+1:ro+m1+m2, co+1:co+n1] = U[ro+m1+1:ro+m1+m2, 1:ru2]*hssA.B21*V[co+1:co+n1, 1:rv1]'
    if !rootnode
      U[ro+1:ro+m1, 1:size(hssA.R1,2)] = U[ro+1:ro+m1, 1:ru1]*hssA.R1
      U[ro+m1+1:ro+m1+m2, 1:size(hssA.R2,2)] = U[ro+m1+1:ro+m1+m2, 1:ru2]*hssA.R2
      V[co+1:co+n1, 1:size(hssA.W1,2)] = V[co+1:co+n1, 1:rv1]*hssA.W1
      V[co+n1+1:co+n1+n2, 1:size(hssA.W2,2)] = V[co+n1+1:co+n1+n2, 1:rv2]*hssA.W2
    end
  end
end

# useful routine to check whether dimensions are compatible
checkdims(hssA::HssMatrix)= _checkdims(hssA, 1)[1]
function _checkdims(hssA::HssMatrix, i::Int)
  if isleaf(hssA)
    compatible = (size(hssA.D,1) == size(hssA.U,1)) && (size(hssA.D,2) == size(hssA.V,1))
    if !compatible println("dimensions don't match in node ", i) end
    return compatible, i+1
  else
    comp1, i = _checkdims(hssA.A11, i)
    comp2, i = _checkdims(hssA.A22, i)
    r1, w1 = gensize(hssA.A11); r2, w2 = gensize(hssA.A22)
    compatible = (r1 == size(hssA.R1,1)) && (r2 == size(hssA.R2,1)) && (w1 == size(hssA.W1,1)) && (w2 == size(hssA.W2,1))
    if !compatible println("dimensions don't match in node ", i) end
    return (compatible && comp1 && comp2), i+1
  end
end

# remove leaves on the bottom level
function prune_leaves!(hssA::HssMatrix)
  isleaf(hssA) && return hssA
  (isnothing(hssA.A11) || isnothing(hssA.A22)) && error("Something went wrong.")
  if isleaf(hssA.A11) && isleaf(hssA.A22)
    return HssMatrix(_hssleaf(hssA)..., hssA.rootnode)
  else
    hssA.A11 = prune_leaves!(hssA.A11)
    hssA.A22 = prune_leaves!(hssA.A22)
    return hssA
  end
end

# returns D, U and V. replaces _full function
function _hssleaf(hssA::HssMatrix)
  if isleaf(hssA)
    return hssA.D, hssA.U, hssA.V
  else
    A11, U1, V1 = _hssleaf(hssA.A11)
    A22, U2, V2 = _hssleaf(hssA.A22)
    return [A11 U1*hssA.B12*V2'; U2*hssA.B21*V1' A22], [U1*hssA.R1; U2*hssA.R2], [V1*hssA.W1; V2*hssA.W2]
  end
end

## write function that extracts the clustwer tree from an HSS matrix
cluster(hssA) = _cluster(hssA, 0, 0)
function _cluster(hssA::HssMatrix, ro::Int, co::Int)
  if isleaf(hssA)
    m, n = size(hssA)
    return ClusterTree(ro .+ (1:m)), ClusterTree(co .+ (1:n))
  else
    rcl1, ccl1 = _cluster(hssA.A11, ro, co)
    rcl2, ccl2 = _cluster(hssA.A22, rcl1.data[end], ccl1.data[end])
    return ClusterTree(ro+1:rcl2.data[end], rcl1, rcl2), ClusterTree(co+1:ccl2.data[end], ccl1, ccl2)
  end
end