text:
"""
    train!(model::AbstractGP, X::AbstractMatrix, y::AbstractVector;obsdim = 1, iterations::Int=10,callback=nothing,conv=0)
    train!(model::AbstractGP, X::AbstractVector, y::AbstractVector;iterations::Int=20,callback=nothing,conv=0)

Function to train the given GP `model`.

**Keyword Arguments**

there are options to change the number of max iterations,
- `iterations::Int` : Number of iterations (not necessarily epochs!)for training
- `callback::Function` : Callback function called at every iteration. Should be of type `function(model,iter) ...  end`
- `conv::Function` : Convergence function to be called every iteration, should return a scalar and take the same arguments as `callback`
"""
function train!(
    m::OnlineSVGP,
    X::AbstractMatrix,
    y::AbstractArray;
    iterations::Int=20,
    callback=nothing,
    convergence=nothing,
    obsdim::Int=1,
)
    return train!(
        m, KernelFunctions.vec_of_vecs(X; obsdim=obsdim), y; iterations=iterations, callback=callback
    )
end

function train!(
    m::OnlineSVGP,
    X::AbstractVector,
    y::AbstractArray;
    iterations::Int=20,
    callback::Union{Nothing,Function}=nothing,
    conv::Union{Nothing,Function}=nothing,
)
    X, _ = wrap_X(X)
    y, _nLatent, m.likelihood = check_data!(y, likelihood(m))

    wrap_data!(data(m), X, y) # Set the data in the model

    _nLatent == nLatent(m) || "Data should always contains the same number of outputs"
    iterations > 0 || "Number of iterations should be positive"
    setnMinibatch!(inference(m), nSamples(data(m)))
    setnSamples!(inference(m), nSamples(data(m)))
    # setMBIndices!(inference(m), collect(1:nMinibatch(inference(m))))

    if nIter(m) == 1 # The first time data is seen, initialize all parameters
        init_onlinemodel(m)
        m.likelihood = init_likelihood(
            likelihood(m), inference(m), nLatent(m), nSamples(data(m))
        )
    else
        save_old_parameters!(m)
        m.likelihood = init_likelihood(likelihood(m), inference(m), nLatent(m), nSamples(m))
        updateZ!(m)
    end

    # model.evol_conv = [] #Array to check on the evolution of convergence
    local_iter = 1
    conv = Inf

    while true # Loop until one condition is matched
        try # Allow for keyboard interruption without losing the model
            if local_iter == 1
                compute_old_matrices!(m)
                local_updates!(likelihood(m), yview(m), mean_f(m), var_f(m))
                ∇E_μs = ∇E_μ(m)
                ∇E_Σs = ∇E_Σ(m) # They need to be computed before recomputing the matrices
                compute_kernel_matrices!(m)
                natural_gradient!.(
                    ∇E_μs, ∇E_Σs, getρ(inference(m)), get_opt(inference(m)), Zviews(m), m.f
                )
                global_update!(m)
            else
                update_parameters!(m) #Update all the variational parameters
            end
            set_trained!(m, true)
            if !isnothing(callback)
                callback(m, nIter(m)) #Use a callback method if given by user
            end
            if (nIter(m) % m.atfrequency == 0) && nIter(m) >= 3
                update_hyperparameters!(m) #Update the hyperparameters
            end
            if verbose(m) > 2 || (verbose(m) > 1 && local_iter % 10 == 0)
                print("Iteration : $(nIter(m)), ")
                print("ELBO is : $(objective(m))")
                print("\n")
                println("number of points : $(dim(m[1]))")
            end
            ### Print out informations about the convergence
            local_iter += 1
            m.inference.nIter += 1
            (local_iter <= iterations) || break # Verify if the number of maximum iterations has been reached
        # (iter < model.nEpochs && conv > model.ϵ) || break; #Verify if any condition has been broken
        catch e
            if e isa InterruptException
                @warn "Training interrupted by user at iteration $local_iter"
                break
            else
                rethrow(e)
            end
        end
    end
    # if model.verbose > 0
    # println("Training ended after $local_iter iterations. Total number of iterations $(model.inference.nIter)")
    # end
    compute_kernel_matrices!(m) #Compute final version of the matrices for prediction
    return set_trained!(m, true)
end

# Update all variational parameters of the online sparse
# variational GP Model
function update_parameters!(model::OnlineSVGP)
    compute_kernel_matrices!(model) #Recompute the matrices if necessary (always for the stochastic case, or when hyperparameters have been updated)
    variational_updates!(model)
    return nothing
end

function InducingPoints.updateZ!(m::OnlineSVGP)
    for gp in m.f
        gp.Z = InducingPoints.updateZ(gp.Z, gp.Zalg, input(m); kernel=kernel(gp))
        gp.post.dim = length(Zview(gp))
    end
    setHPupdated!(inference(m), true)
    return nothing
end

function save_old_parameters!(m::OnlineSVGP)
    for gp in m.f
        save_old_gp!(gp)
    end
end

function save_old_gp!(gp::OnlineVarLatent{T}) where {T}
    gp.Zₐ = deepcopy(gp.Z)
    gp.Z = InducingPoints.remove_point(Random.GLOBAL_RNG, gp.Z, gp.Zalg, Matrix(pr_cov(gp)))# Matrix(pr_cov(gp)))
    gp.invDₐ = Symmetric(-2.0 * nat2(gp) - inv(pr_cov(gp))) # Compute Σ⁻¹ₐ - K⁻¹ₐ
    gp.prevη₁ = copy(nat1(gp))
    gp.prev𝓛ₐ =
        (-logdet(cov(gp)) + logdet(pr_cov(gp)) - dot(mean(gp), nat1(gp))) / 2
    return nothing
end

function init_onlinemodel(m::OnlineSVGP{T}) where {T<:Real}
    m.f = ntuple(length(m.f)) do i
        init_online_gp!(m.f[i], m)
    end
    # for gp in m.f
    #     init_online_gp!(gp, m)
    # end
    setρ!(inference(m), one(T))
    return setHPupdated!(inference(m), false)
end

function init_online_gp!(gp::OnlineVarLatent{T}, m::OnlineSVGP, jitt::T=T(jitt)) where {T}
    Z = InducingPoints.initZ(gp.Zalg, input(m); kernel=kernel(gp))
    k = length(Z)
    Zₐ = deepcopy(Z)
    post = OnlineVarPosterior{T}(k)
    prior = GPPrior(
        kernel(gp), pr_mean(gp), cholesky(kernelmatrix(kernel(gp), Z) + jitt * I)
    )

    Kab = zeros(T, k, k)
    κₐ = Matrix{T}(I(k))
    K̃ₐ = zero(Kab)

    Knm = kernelmatrix(kernel(gp), input(m), Z)
    κ = Knm / (kernelmatrix(kernel(gp), Z) + jitt * I)
    K̃ = kernelmatrix_diag(kernel(gp), input(m)) .+ jitt - diag_ABt(κ, Knm)
    all(K̃ .> 0) || error("K̃ has negative values")

    invDₐ = Symmetric(Matrix{T}(I(k)))
    prev𝓛ₐ = zero(T)
    prevη₁ = zeros(T, k)
    return OnlineVarLatent(
        prior,
        post,
        Z,
        gp.Zalg,
        Knm,
        κ,
        K̃,
        gp.Zupdated,
        gp.opt,
        gp.Zopt,
        Zₐ,
        Kab,
        κₐ,
        K̃ₐ,
        invDₐ,
        prev𝓛ₐ,
        prevη₁
    )
    # return nothing
end

function compute_old_matrices!(m::OnlineSVGP{T}) where {T}
    for gp in m.f
        compute_old_matrices!(gp, xview(m), T(jitt))
    end
end

function compute_old_matrices!(gp::OnlineVarLatent, X::AbstractVector, jitt::Real)
    gp.prior.K = cholesky(kernelmatrix(kernel(gp), gp.Zₐ) + jitt * I)
    gp.Knm = kernelmatrix(kernel(gp), X, gp.Zₐ)
    gp.κ = gp.Knm / pr_cov(gp)
    gp.K̃ = kernelmatrix_diag(kernel(gp), X) .+ jitt - diag_ABt(gp.κ, gp.Knm)
    all(gp.K̃ .> 0) || error("K̃ has negative values")
    return nothing
end


text:
"""
    steadystate_iterative!(rho0, H, J, solver, args...; [log=false], kwargs...) -> rho[, log]

Compute the steady state density matrix of a Hamiltonian and a set of jump
operators by solving `L rho = 0` via an iterative method provided as argument.

# Arguments
* `rho0`: Initial guess.
* `H`: dense Hamiltonian.
* `J`: array of dense jump operators.
* `solver::Symbol`: name of a desired iterative method from scope or `IterativeSolvers.jl`.
* `args...`: further arguments are passed on to the iterative solver.
* `kwargs...`: further keyword arguments are passed on to the iterative solver.

See also: [`steadystate_bicg`](@ref)
"""
function steadystate_iterative!(ρ0::DenseOperator{B,B}, H::DenseOperator{B,B}, J::Vector{O}, solver::Symbol, args...; log::Bool=false, tol::Float64 = sqrt(eps(real(ComplexF64))), kwargs...) where {B<:Basis,O<:DenseOperator{B,B}}
    method! = isdefined(Main,solver) ? getfield(Main,solver) : getfield(@__MODULE__,solver)
    jump_ops = map(x->x.data, J)
    ρ0_mat = ρ0.data

    # Size of the Hilbert space
    M = size(H.data,1)
    # Non-Hermitian Hamiltonian
    iHnh = -im*H.data
    for Jᵢ=jump_ops
        iHnh .+= -0.5adjoint(Jᵢ)*Jᵢ
    end
    Jρ_cache = similar(iHnh)

    # In-place update of y = Lx where L and x are respectively the vectorized
    # Liouvillian and the vectorized density matrix. y[1] is set to the trace
    # of the density matrix so as to enforce a trace one non-trivial solution.
    function mvecmul!(y::AbstractVector, x::AbstractVector)
        y .= zero(eltype(y));
        ym = @views reshape(y[2:end], M, M)
        ρ  = @views reshape(x[2:end], M, M)

        BLAS.gemm!('N', 'N', one(eltype(y)), iHnh, ρ, one(eltype(y)), ym)
        BLAS.gemm!('N', 'C', one(eltype(y)), ρ, iHnh, one(eltype(y)), ym)
        for Jᵢ=jump_ops
            BLAS.gemm!('N','N', one(eltype(y)), Jᵢ, ρ, zero(eltype(y)), Jρ_cache)
            BLAS.gemm!('N','C', one(eltype(y)), Jρ_cache, Jᵢ, one(eltype(y)), ym)
        end
        y[1] = tr(ρ)

        return y
    end
    # Solution x must satisfy L.x = y with y[1] = tr(x) = 1 and y[j≠1] = 0.
    x0 = similar(ρ0_mat, M^2+1)
    x0[1] = zero(x0[1])
    x0[2:end] .= reshape(ρ0_mat, M^2)

    y = similar(ρ0_mat, M^2+1)
    y .= zero(y)
    y[1] = one(y[1])

    # Define the linear map lm: ρ ↦ L(ρ)
    lm = LinearMap{eltype(iHnh)}(mvecmul!, length(y)::Int, length(y)::Int; ismutating=true, issymmetric=false, ishermitian=false, isposdef=false)

    # Perform the stabilized biconjugate gradient procedure and devectorize ρ
    res0_norm = norm(mvecmul!(similar(y),x0) .- y)
    tol /= res0_norm + eps(Float64)
    if !log
        ρ0.data .= @views reshape(method!(x0,lm,y,args...;tol=tol,kwargs...)[2:end],(M,M))
        return ρ0
    else
        R, history = method!(x0,lm,y,args...;log=log,tol=tol,kwargs...)
        ρ0.data .= @views reshape(R[2:end],(M,M))
        return ρ0, history
    end
end

function steadystate_iterative!(ρ0::SparseOperator{B,B}, H::SparseOperator{B,B}, J::Vector{O}, solver::Symbol, args...; log::Bool=false, tol::Float64 = sqrt(eps(real(ComplexF64))), kwargs...) where {B<:Basis,O<:SparseOperator{B,B}}
    method! = isdefined(Main,solver) ? getfield(Main,solver) : getfield(@__MODULE__,solver)
    # Size of the Hilbert space
    M::Int = size(H.data,1)
    # Non-Hermitian Hamiltonian
    iHnh::SparseMatrixCSC{ComplexF64,Int} = -im*H.data
    for i in 1:length(J)
        iHnh .+= -0.5adjoint(J[i].data)*J[i].data
    end
    Jρ_cache::SparseMatrixCSC{ComplexF64,Int} = similar(iHnh)

    # In-place update of y = Lx where L and x are respectively the vectorized
    # Liouvillian and the vectorized density matrix. y[1] is set to the trace
    # of the density matrix so as to enforce a trace one non-trivial solution.
    function mvecmul!(y::AbstractVector, x::AbstractVector)
        y .= zero(eltype(y));
        ym::SparseMatrixCSC{eltype(y),Int} = @views reshape(y[2:end],M,M)
        ρ::SparseMatrixCSC{eltype(x),Int}  = @views reshape(x[2:end],M,M)

        ym .= iHnh * ρ .+ ρ * adjoint(iHnh)
        #ym .+= adjoint(ym)
        @inbounds @views for i in 1:length(J)
            Jρ_cache = J[i].data * ρ
            ym .+= Jρ_cache * adjoint(J[i].data)
        end
        @views y[2:end] .= reshape(ym,length(ym))
        y[1] = tr(ρ)

        return y
    end
    # Solution x must satisfy L.x = y with y[1] = tr(x) = 1 and y[j≠1] = 0.
    x0::SparseVector{ComplexF64,Int} = [zero(ComplexF64); reshape(ρ0.data,M^2)]
    y::SparseVector{ComplexF64,Int}  = [one(ComplexF64); zeros(ComplexF64,M^2)]

    # Define the linear map lm: ρ ↦ L(ρ)
    lm = LinearMap{ComplexF64}(mvecmul!, length(y)::Int, length(y)::Int; ismutating=true, issymmetric=false, ishermitian=false, isposdef=false)

    # Perform the stabilized biconjugate gradient procedure and devectorize ρ
    res0_norm::Float64 = norm(mvecmul!(similar(y),x0) .- y)
    tol /= res0_norm + eps(Float64)
    if !log
        ρ0.data .= reshape(method!(x0,lm,y,args...;tol=tol,kwargs...)[2:end],(M,M))
        return ρ0
    else
        R::Vector{ComplexF64}, history = method!(x0,lm,y,args...;log=log,tol=tol,kwargs...)
        ρ0.data .= reshape(R[2:end],(M,M))
        return ρ0, history
    end
end

ext: jl


text:
header("parse tree")
str = """
<?xml version="1.0"?><!-- https://github.com/daemontus/pnml-parser -->
<pnml xmlns="http://www.pnml.org/version-2009/grammar/pnml">
  <net id="small-net" type="http://www.pnml.org/version-2009/grammar/ptnet">
    <name> <text>P/T Net with one place</text> </name>
    <page id="page0">
      <place id="place1">
        <name> <text>Some place</text> </name>
        <initialMarking> <text>100</text> </initialMarking>
      </place>
      <transition id="transition1">
        <name> <text>Some transition </text> </name>
      </transition>
      <arc source="transition1" target="place1" id="arc1">
        <inscription> <text>12 </text> </inscription>
      </arc>
      <arc source="place1" target="transition1" id="arc2">
      </arc>
    </page>
  </net>
</pnml>
"""
doc = EzXML.parsexml(str) # shared by testsets

#if PRINT_PNML
#    EzXML.prettyprint(doc);
#    println()
#end

@testset "parse tree" begin
    pnml = root(doc)
    @test EzXML.nodename(pnml) == "pnml"
    @test EzXML.namespace(pnml) == "http://www.pnml.org/version-2009/grammar/pnml"

    reg = PNML.IDRegistry()
    # Manually decend tree parsing leaf-enough elements because this is a test!
    foreach(PNML.allchildren("net", pnml)) do net
        @test nodename(net) == "net"

        nn = parse_name(PNML.firstchild("name", net), PnmlCore(); reg)
        @test isa(nn, PNML.Name)
        @test nn.text == "P/T Net with one place"
        @test nn.graphics === nothing
        @test nn.tools === nothing || isempty(nn.tools)

        nd = PNML.allchildren("declaration", net)
        @test_call PNML.allchildren("declaration", net)
        @test isempty(nd)
        @test isempty(parse_node.(nd; reg)) # Empty elements are leaf-enough.

        nt = PNML.allchildren("toolspecific", net)
        @test isempty(nt)
        @test isempty(parse_node.(nt; reg))

        pages = PNML.allchildren("page", net)
        @test !isempty(pages)

        foreach(pages) do page
            @test nodename(page) == "page"

            @test !isempty(PNML.allchildren("place", page))
            foreach(PNML.allchildren("place", page)) do p
                @test nodename(p) == "place"
                i = parse_node(PNML.firstchild("initialMarking", p); reg)
                @test_call PNML.firstchild("initialMarking", p)
                @test typeof(i) <: PNML.PTMarking
                @test typeof(i.value) <: Number
                @test i.value >= 0
                @test xmlnode(i) isa Maybe{EzXML.Node}
            end

            @test !isempty(PNML.allchildren("transition", page))
            foreach(PNML.allchildren("transition", page)) do t
                @test nodename(t) == "transition"
                cond = PNML.firstchild("condition", t)
                @test cond === nothing
                #i = parse_node(PNML.firstchild("condition", t); reg)
            end

            @test !isempty(PNML.allchildren("arc", page))
            foreach(PNML.allchildren("arc", page)) do a
                @test nodename(a) == "arc"
                ins = PNML.firstchild("inscription", a)
                if ins !== nothing
                    i = parse_node(ins; reg)
                    @test typeof(i) <: PNML.PTInscription
                    @test typeof(i.value) <: Number
                    @test i.value > 0
                    @test xmlnode(i) isa Maybe{EzXML.Node}
                end
            end
        end
    end
    PNML.reset_registry!(reg)
end

@testset "parse node level" begin

    # Do a full parse and maybe print the generated data structure.
    reg = PNML.IDRegistry()
    pnml_ir = parse_pnml(root(doc); reg)
    @test typeof(pnml_ir) <: PNML.PnmlModel

    foreach(nets(pnml_ir)) do net
        @test net isa PNML.PnmlNet
        @test net.id isa Symbol

        foreach(net.pages) do page
            @test page isa PNML.Page
            @test pid(page) isa Symbol
            foreach(page.places) do place
                @test place isa PNML.Place
                @test pid(place) isa Symbol
            end
            foreach(page.transitions) do transition
                @test transition isa PNML.Transition
                @test pid(transition) isa Symbol
            end
            foreach(page.arcs) do arc
                @test arc isa PNML.Arc
                @test pid(arc) isa Symbol
            end
            foreach(PNML.declarations(page)) do decl
                @test decl isa PNML.Declaration
                @test decl[:text] !== nothing || decl[:structure] !== nothing
            end
        end

        foreach(PNML.declarations(net)) do decl
            @test decl isa PNML.Declaration
            @test decl[:text] !== nothing || decl[:structure] !== nothing
        end
    end

    PNML.reset_registry!(reg)
end


text:
export DataSet
export CosSin
export FaouSchratz
export BaoDong

"""
Class with initial data Relativistic Klein-Gordon equation
"""
struct DataSet

    nx        :: Int64
    xmin      :: Float64
    xmax      :: Float64
    epsilon   :: Float64
    kx        :: Vector{Float64}
    T         :: Float64
    Tfinal    :: Float64
    sigma     :: Int64
    llambda   :: Int64
    x         :: Array{Float64,1}
    u         :: Array{ComplexF64,1}
    v         :: Array{ComplexF64,1}
    dx        :: Float64

    function DataSet( xmin, xmax, nx, epsilon, T, Tfinal,
                      phi, gamma, sigma, llambda)

        kx  = zeros(Float64, nx)
        kx .= 2 * pi / (xmax - xmin) * vcat(0:nx÷2-1,-nx÷2:-1)

        x   = zeros(Float64, nx)
        x  .= range(xmin, stop=xmax, length=nx+1)[1:end-1]
        dx  = (xmax - xmin) / nx

        ϕ   = zeros(ComplexF64, nx)
        γ   = zeros(ComplexF64, nx)

        ϕ  .= phi(x)
        γ  .= gamma(x)

        u = zeros(ComplexF64, nx)
        v = zeros(ComplexF64, nx)

        u .= ϕ .- 1im * ifft((1 .+ epsilon * kx.^2) .^ (-1/2) .* fft(γ,1), 1)
        v .= conj.(ϕ) .- 1im * ifft((1 .+ epsilon * kx.^2) .^ (-1/2) .* fft(conj.(γ),1),1)

        new(nx, xmin, xmax, epsilon, kx, T, Tfinal, sigma, llambda, x, u, v, dx)

    end

end

function BaoDong( xmin, xmax, nx, epsilon, T, Tfinal)

    phi(x)   = 2 / (exp.(x .^ 2) + exp.(-x .^ 2))
    gamma(x) = 0

    sigma    = 1
    llambda  = -4

    DataSet(xmin, xmax, nx, epsilon, T, Tfinal, phi, gamma, sigma, llambda)

end

function CosSin( xmin, xmax, nx, epsilon, T, Tfinal)

    phi(x)   = (1 + 1im) .* cos.(x)
    gamma(x) = (1 - 1im) .* sin.(x)

    sigma   = 1
    llambda = -1

    DataSet(xmin, xmax, nx, epsilon, T, Tfinal, phi, gamma, sigma, llambda)

end

function FaouSchratz( xmin, xmax, nx, epsilon, T, Tfinal)

    # Faou-Schratz 6.2
    #ϕ  .= (2 + 1im) / sqrt(5) * cos.(x)
    #γ  .= (1 + 1im) / sqrt(2) * sin.(x) .+ 0.5 * cos.(x)

    # Faou-Schratz 6.3
    phi(x)   = cos.(x)
    gamma(x) = 1 / 4 * sin.(x) .+ 0.5 * cos.(x)

    sigma   = 1
    llambda = -1

    DataSet(xmin, xmax, nx, epsilon, T, Tfinal, phi, gamma, sigma, llambda)

end


text:
module FESetCorotBeamModule

using FinEtools
import FinEtools: cat, subset
using FinEtools.MatrixUtilityModule: complete_lt!
using ..CrossSectionModule: AbstractCrossSectionType
using LinearAlgebra: norm, Transpose, mul!

mutable struct FESetL2CorotBeam{CT} <: AbstractFESet1Manifold{2}
    conn::Array{NTuple{2, FInt}, 1};
    label::FIntVec;
    crosssection::CT
    A::FFltVec
    I1::FFltVec
    I2::FFltVec
    I3::FFltVec
    J::FFltVec
    A2s::FFltVec
    A3s::FFltVec
    x1x2_vector::Vector{FFltVec}
    dimensions::Vector{FFltVec}
end

function FESetL2CorotBeam(conn::FIntMat, crosssection::CT) where {CT}
    par = crosssection.parameters(0.0)
    N = size(conn, 1)
    _A = fill(par.A, N)
    _I1 = fill(par.I1, N)
    _I2 = fill(par.I2, N)
    _I3 = fill(par.I3, N)
    _J = fill(par.J, N)
    _A2s = fill(par.A2s, N)
    _A3s = fill(par.A3s, N)
    _x1x2_vector = [par.x1x2_vector for i in 1:N]
    _dimensions = [par.dimensions for i in 1:N]
    self = FESetL2CorotBeam(NTuple{2, FInt}[], FInt[], crosssection, _A, _I1, _I2, _I3, _J, _A2s, _A3s, _x1x2_vector, _dimensions)
    self = fromarray!(self, conn)
    setlabel!(self, 0)
    return self
end

function FESetL2CorotBeam(conn::FIntMat, crosssection::CT, _A, _I1, _I2, _I3, _J, _A2s, _A3s, _x1x2_vector, _dimensions) where {CT}
    dummy = FESetL2(conn)
    setlabel!(dummy, 0)
    return FESetL2CorotBeam(dummy.conn, dummy.label, crosssection, _A, _I1, _I2, _I3, _J, _A2s, _A3s, _x1x2_vector, _dimensions)
end

"""
    Type of the mass matrix formulation.
"""
const MASS_TYPE_CONSISTENT_NO_ROTATION_INERTIA=0;
const MASS_TYPE_CONSISTENT_WITH_ROTATION_INERTIA=1;
const MASS_TYPE_LUMPED_DIAGONAL_NO_ROTATION_INERTIA=2;
const MASS_TYPE_LUMPED_DIAGONAL_WITH_ROTATION_INERTIA=3;

"""
    cat(self::T,  other::T) where {T<:FESetL2CorotBeam}

Concatenate two sets of beam elements.
"""
function cat(self::T,  other::T) where {T<:FESetL2CorotBeam}
    @assert self.crosssection === other.crosssection "Cannot concatenate sets with distinct cross-sections"
    result = deepcopy(self)
    result.conn = vcat(self.conn, other.conn);
    setlabel!(result, vcat(self.label, other.label))
    result.A = vcat(self.A, other.A);
    result.I1 = vcat(self.I1, other.I1);
    result.I2 = vcat(self.I2, other.I2);
    result.I3 = vcat(self.I3, other.I3);
    result.J = vcat(self.J, other.J);
    result.A2s = vcat(self.A2s, other.A2s);
    result.A3s = vcat(self.A3s, other.A3s);
    result.x1x2_vector = vcat(self.x1x2_vector, other.x1x2_vector);
    result.dimensions = vcat(self.dimensions, other.dimensions);
    return result
end

"""
    subset(self::T, L::FIntVec) where {T<:FESetL2CorotBeam}

Subset of a beam-element set
"""
function subset(self::T, L::FIntVec) where {T<:FESetL2CorotBeam}
    result = deepcopy(self)
    result.conn = deepcopy(self.conn[L])
    result.label = deepcopy(self.label[L])
    result.A = self.A[L]
    result.I1 = self.I1[L]
    result.I2 = self.I2[L]
    result.I3 = self.I3[L]
    result.J = self.J[L]
    result.A2s = self.A2s[L]
    result.A3s = self.A3s[L]
    result.x1x2_vector = self.x1x2_vector[L]
    result.dimensions = self.dimensions[L]
    return  result
end

function initial_local_frame!(F0, x0, x1x2_vector)
    # This is the element frame in the configuration t=0
    F0[:,1] = (x0[2,:]-x0[1,:]);
    L0 = norm(@view F0[:,1]);
    F0[:,1] /= L0;
    #     F0(:,3)=skewmat(F0(:,1))*x1x2_vector;
    F0[:,3] .= (-F0[3,1]*x1x2_vector[2]+F0[2,1]*x1x2_vector[3],
                 F0[3,1]*x1x2_vector[1]-F0[1,1]*x1x2_vector[3],
                -F0[2,1]*x1x2_vector[1]+F0[1,1]*x1x2_vector[2]);
    F0[:,3] /= norm(@view F0[:,3]);
    #     F0(:,2)=skewmat(F0(:,3))*F0(:,1);
    F0[:,2] .= (-F0[3,3]*F0[2,1]+F0[2,3]*F0[3,1],
                 F0[3,3]*F0[1,1]-F0[1,3]*F0[3,1],
                -F0[2,3]*F0[1,1]+F0[1,3]*F0[2,1]);
    return L0, F0
end

function current_local_frame!(Ft, xt, FtI, FtJ)
    Ft = fill(0.0, 3, 3);
    # Compute the element frame in configuration t
    Ft[:,1] = (xt[2,:]-xt[1,:]);
    Lt = norm(@view Ft[:,1]);
    Ft[:,1] /= Lt;
    x1x2_vectort = FtI[:,2]+FtJ[:,2];
    #     Ft(:,3)=skewmat(Ft(:,1))*x1x2_vectort; Ft(:,3)=Ft(:,3)/norm(Ft(:,3));
    #     Ft(:,2)=skewmat(Ft(:,3))*Ft(:,1);
    #     Ft(:,3)=skewmat(Ft(:,1))*x1x2_vectort; # In the interest of speed,
    #     replace with below explicit rewrite
    Ft[:,3] .= (-Ft[3,1]*x1x2_vectort[2]+Ft[2,1]*x1x2_vectort[3],
                Ft[3,1]*x1x2_vectort[1]-Ft[1,1]*x1x2_vectort[3],
               -Ft[2,1]*x1x2_vectort[1]+Ft[1,1]*x1x2_vectort[2]);
    Ft[:,3] /= norm(@view Ft[:,3]);
    #     Ft(:,2)=skewmat(Ft(:,3))*Ft(:,1); # In the interest of speed,
    #     replace with below explicit rewrite
    Ft[:,2] .= (-Ft[3,3]*Ft[2,1]+Ft[2,3]*Ft[3,1],
                Ft[3,3]*Ft[1,1]-Ft[1,3]*Ft[3,1],
               -Ft[2,3]*Ft[1,1]+Ft[1,3]*Ft[2,1]);
    return Lt, Ft
end

"""
    local_frame_and_def!(Ft, dN, F0, FtI, FtJ, x0, x1x2_vector, xt, RI, RJ)

Compute the current length of the element, the current element frame,
and the natural deformations

# Arguments
x0= array of node coordinates, one node per row, in initial configuration
x1x2_vector= vector that lies in the x1-x2 local element coordinate
     plane, in initial configuration
xt= array of node coordinates, one node per row, in initial configuration
RI,RJ=nodal rotation (orthogonal) matrix

# Outputs
`Lt`= current length of the element,
`Ft`= current element frame (orthogonal rotation matrix) whose columns are unit
     vectors: they are centered halfway between the current locations of
     the nodes, vector 1 points from node I to node J, vector 3 is
     orthogonal to the sum of the nodal cross-section frame vectors 2
`dN`= vector of natural deformations; dN(1)= total change in length
     between configurations 0 and t; dN(2)= symmetric bending;
     dN(3)= anti-symmetric bending;    dN(4)= symmetric bending
     dN(5)= anti-symmetric bending; dN(6)=total axial torsion angle.

"""
function local_frame_and_def!(Ft, dN, F0, FtI, FtJ, x0, x1x2_vector, xt, RI, RJ)
    # This is the element frame in the configuration t=0
    L0, F0 = initial_local_frame!(F0, x0, x1x2_vector)

    # The nodal cross-section frames are rotated by the nodal rotation matrices
    mul!(FtI, RI, F0);
    mul!(FtJ, RJ, F0);

    # The element frame in the configuration t
    Lt, Ft = current_local_frame!(Ft, xt, FtI, FtJ)

    # Components of FtI,FtJ in the element frame
    LFtI = Ft'*FtI;
    LFtJ = Ft'*FtJ;

    # C  COMPUTE NET DEFORMATIONS- DELTA L,TH1IJ,TH2I,TH2J,TH3I,TH3J
    dN = fill(0.0, 6);
    # Total change in length between configurations 0 and t
    dN[1] = Lt-L0;
    # Total axial TORSION ANGLE
    dN[6] = (LFtJ[3,2]/LFtJ[2,2]-LFtI[3,2]/LFtI[2,2] -LFtJ[2,3]/LFtJ[3,3]+LFtI[2,3]/LFtI[3,3])/2;
    # Total rotation  angles of nodal cross-sections relative to the element frame
    TH2I =-LFtI[3,1]/LFtI[1,1];
    TH2J =-LFtJ[3,1]/LFtJ[1,1];
    TH3I = LFtI[2,1]/LFtI[1,1];
    TH3J = LFtJ[2,1]/LFtJ[1,1];
    dN[2] = TH3I-TH3J; # symmetric bending
    dN[3] = TH3I+TH3J; # anti-symmetric bending
    dN[4] = -TH2I+TH2J; # symmetric bending
    dN[5] = -TH2I-TH2J; # anti-symmetric bending
    return Lt, Ft, dN
end

"""
    local_cartesian_to_natural!(aN, L)

Compute transformation from local Cartesian displacements to natural deformations
of a beam element.

Matrix defined in Eq (4.8) of COMPUTER  METHODS  IN APPLIED  MECHANICS  AND  ENGINEERING   14  (1978)  401-451
ON LARGE  DISPLACEMENT-SMALL   STRAIN  ANALYSIS   OF  STRUCTURES
WITH  ROTATIONAL   DEGREES  OF  FREEDOM.
J.H.  ARGYRIS,   P.C.  DUNNE  and  D.W. SCHARPF

# Arguments
`L`= current length of the element

# Outputs
`aN`= transformation matrix to take Cartesian (local) displacement increments in the
     element frame and to produce increments of natural deformations;
     see `local_frame_and_def!` for the definition of the natural deformations
"""
function local_cartesian_to_natural!(aN, L)
    fill!(aN, 0.0)
    aN[1, 1] = -1; aN[1, 7] = +1
    aN[2, 6] = +1; aN[2, 12] = -1
    aN[3, 2] = 2/L; aN[3, 6] = +1; aN[3, 8] = -2/L; aN[3, 12] = +1
    aN[4, 5] = -1; aN[4, 11] = +1
    aN[5, 3] = 2/L; aN[5, 5] = -1; aN[5, 9] = -2/L; aN[5, 11] = -1
    aN[6, 4] = -1; aN[6, 10] = +1
    # aN=[[ -1,   0,   0,  0,  0, 0, 1,    0,    0, 0,  0,  0]
    #     [  0,   0,   0,  0,  0, 1, 0,    0,    0, 0,  0, -1]
    #     [  0, 2/L,   0,  0,  0, 1, 0, -2/L,    0, 0,  0,  1]
    #     [  0,   0,   0,  0, -1, 0, 0,    0,    0, 0,  1,  0]
    #     [  0,   0, 2/L,  0, -1, 0, 0,    0, -2/L, 0, -1,  0]
    #     [  0,   0,   0, -1,  0, 0, 0,    0,    0, 1,  0,  0]];
        return aN
end

"""
    local_geometric_stiffness!(SM, A, I2, I3, PN, L)

Compute the local geometric stiffness matrix.

# Arguments
`A`= cross-sectional area,
`I2`, `I3`=central moment of inertia of the cross-section about the x2 and x3
coordinate axis,
`PN`= vector of natural forces; see natural_forces() for definitions
`L`= current length of the element,

# Outputs
`SM` = local geometric stiffness matrix, 12 x 12

This geometric stiffness matrix this consistent with relationship between
the natural deformations and the natural forces that assumes there is
only a linear constitutive link: no non-constitutive effects (bowing
etc.) are included. This form of the geometric matrix was derived by
Krenk.
@BOOK{Krenk:2009,
  AUTHOR =       {S. Krenk},
  TITLE =        {Non-linear Modeling and Analysis of Solids and Structures },
  PUBLISHER =    {Cambridge University Press},
  YEAR =         {2009},
  isbn =         {9780521830546}
}
"""
function local_geometric_stiffness!(SM, A, I2, I3, PN, L)
    N = PN[1];
    S_2 = -2*PN[3]/L;
    S_3 = -2*PN[5]/L;
    M_1 = PN[6];
    M_2I = PN[4]+PN[5];
    M_2J = PN[4]-PN[5];
    M_3I = -(PN[2]+PN[3]);
    M_3J = -(PN[2]-PN[3]);
    SM[1:3, 1:3] = [0   -S_2/L   -S_3/L;
                    -S_2/L   N/L   0;
                    -S_3/L   0   N/L];
    SM[1:3, 4:6] = [0   0   0;
                    -M_2I/L   M_1/L   0;
                    -M_3I/L   0   M_1/L];
    SM[1:3, 7:9] = [0   +S_2/L   +S_3/L;
                    +S_2/L   -N/L   0;
                    +S_3/L   0   -N/L];
    SM[1:3, 10:12] = [0   0   0;
                    M_2J/L   -M_1/L   0;
                    M_3J/L   0   -M_1/L];
    SM[4:6, 4:6] = [0   M_3I/2   -M_2I/2;
                    M_3I/2   0   0;
                    -M_2I/2   0   0];
    SM[4:6, 7:9] = [0   M_2I/L   M_3I/L;
                    0   -M_1/L   0;
                    0   0   -M_1/L];
    SM[4:6, 10:12] = [0   0   0;
                    0   0   M_1/2;
                    0   -M_1/2   0];
    SM[7:9, 7:9] = [0   -S_2/L   -S_3/L;
                    -S_2/L   N/L   0;
                    -S_3/L   0   N/L];
    SM[7:9, 10:12] = [0   0   0;
                    -M_2J/L   M_1/L   0;
                    -M_3J/L   0   M_1/L];
    SM[10:12, 10:12] = [0   -M_3J/2   M_2J/2;
                        -M_3J/2   0   0;
                        M_2J/2   0   0];
    complete_lt!(SM)
    return SM
end

function local_mass_CONSISTENT_NO_ROTATION_INERTIA!(MM, A, I1, I2, I3, rho, L)
    # C
    # C  CONSISTENT MASS MATRIX including ROTATIONAL MASSES
    # C  Formulation of the (3.38), (3.39) equation from Dykstra's thesis
    fill!(MM, 0.0)
    c1 = (rho*A*L)::Float64
    MM[1, 1] = c1 * (1/3)
    MM[1, 7] = c1 * (1/6)
    MM[2, 2] = c1 * (13/35)
    MM[2, 6] = c1 * (11*L/210)
    MM[2, 8] = c1 * (9/70)
    MM[2, 12] = c1 * (-13*L/420)
    MM[3, 3] = c1 * (13/35)
    MM[3, 5] = c1 * (-11*L/210)
    MM[3, 9] = c1 * (9/70)
    MM[3, 11] = c1 * (13*L/420)
    MM[4, 4] = c1 * (I1/3/A)
    MM[4, 10] = c1 * (I1/6/A)
    MM[5, 5] = c1 * (L^2/105)
    MM[5, 9] = c1 * (-13*L/420)
    MM[5, 11] = c1 * (-L^2/140)
    MM[6, 6] = c1 * (L^2/105)
    MM[6, 8] = c1 * (13*L/420)
    MM[6, 12] = c1 * (-L^2/140)
    MM[7, 7] = c1 * (1/3)
    MM[8, 8] = c1 * (13/35)
    MM[8, 12] = c1 * (-11*L/210)
    MM[9, 9] = c1 * (13/35)
    MM[9, 11] = c1 * (11*L/210)
    MM[10, 10] = c1 * (I1/3/A)
    MM[11, 11] = c1 * (L^2/105)
    MM[12, 12] = c1 * (L^2/105)
    # 1       2        3        4       5         6       7        8       9       10      11       12
    complete_lt!(MM)
    return MM
end

function local_mass_CONSISTENT_WITH_ROTATION_INERTIA!(MM, A, I1, I2, I3, rho, L)
    # C
    # C  CONSISTENT MASS MATRIX including ROTATIONAL MASSES
    # C  Formulation of the (3.38), (3.39) equation from Dykstra's thesis
    fill!(MM, 0.0)
    c1 = (rho*A*L)::Float64
    MM[1, 1] = c1 * (1/3)
    MM[1, 7] = c1 * (1/6)
    MM[2, 2] = c1 * (13/35)
    MM[2, 6] = c1 * (11*L/210)
    MM[2, 8] = c1 * (9/70)
    MM[2, 12] = c1 * (-13*L/420)
    MM[3, 3] = c1 * (13/35)
    MM[3, 5] = c1 * (-11*L/210)
    MM[3, 9] = c1 * (9/70)
    MM[3, 11] = c1 * (13*L/420)
    MM[4, 4] = c1 * (I1/3/A)
    MM[4, 10] = c1 * (I1/6/A)
    MM[5, 5] = c1 * (L^2/105)
    MM[5, 9] = c1 * (-13*L/420)
    MM[5, 11] = c1 * (-L^2/140)
    MM[6, 6] = c1 * (L^2/105)
    MM[6, 8] = c1 * (13*L/420)
    MM[6, 12] = c1 * (-L^2/140)
    MM[7, 7] = c1 * (1/3)
    MM[8, 8] = c1 * (13/35)
    MM[8, 12] = c1 * (-11*L/210)
    MM[9, 9] = c1 * (13/35)
    MM[9, 11] = c1 * (11*L/210)
    MM[10, 10] = c1 * (I1/3/A)
    MM[11, 11] = c1 * (L^2/105)
    MM[12, 12] = c1 * (L^2/105)
    c2 = (rho/L)::Float64
    MM[2, 2] += c2 * (6/5*I2)
    MM[2, 6] += c2 * (L/10*I2)
    MM[2, 8] += c2 * (-6/5*I2)
    MM[2, 12] += c2 * (L/10*I2)
    MM[3, 3] += c2 * (6/5*I3)
    MM[3, 5] += c2 * (-L/10*I3)
    MM[3, 9] += c2 * (-6/5*I3)
    MM[3, 11] += c2 * (-L/10*I3)
    MM[5, 5] += c2 * (2*L^2/15*I3)
    MM[5, 9] += c2 * (L/10*I3)
    MM[5, 11] += c2 * (-L^2/30*I3)
    MM[6, 6] += c2 * (2*L^2/15*I2)
    MM[6, 8] += c2 * (-L/10*I2)
    MM[6, 12] += c2 * (-L^2/30*I2)
    MM[8, 8] += c2 * (6/5*I2)
    MM[8, 12] += c2 * (-L/10*I2)
    MM[9, 9] += c2 * (6/5*I3)
    MM[9, 11] += c2 * (L/10*I3)
    MM[11, 11] += c2 * (2*L^2/15*I3)
    MM[12, 12] += c2 * (2*L^2/15*I2)
    # 1       2        3        4       5         6       7        8       9       10      11       12
    complete_lt!(MM)
    return MM
end

function local_mass_LUMPED_DIAGONAL_WITH_ROTATION_INERTIA!(MM, A, I1, I2, I3, rho, L)
    # C
    # C  LUMPED DIAGONAL MASS MATRIX WITH ROTATIONAL MASSES
    # C
    HLM  = A*rho*L/2.;
    HLI1 = rho*I1* L/2.;
    HLI2 = rho*I2* L/2.;
    HLI3 = rho*I3* L/2.;
    CA = HLM;
    CB = HLI1;
    CC = HLI2;
    CD = HLI3;
    fill!(MM, 0.0);
    MM[1,1]    = MM[1,1]    + CA;
    MM[2,2]    = MM[2,2]    + CA;
    MM[3,3]    = MM[3,3]    + CA;
    MM[4,4]    = MM[4,4]    + CB;
    MM[5,5]    = MM[5,5]    + CC;
    MM[6,6]    = MM[6,6]    + CD;
    MM[7,7]    = MM[7,7]    + CA;
    MM[8,8]    = MM[8,8]    + CA;
    MM[9,9]    = MM[9,9]    + CA;
    MM[10,10]  = MM[10,10]  + CB;
    MM[11,11]  = MM[11,11]  + CC;
    MM[12,12]  = MM[12,12]  + CD;
    return MM
end

function local_mass_LUMPED_DIAGONAL_NO_ROTATION_INERTIA!(MM, A, I1, I2, I3, rho, L)
    # C
    # C  LUMPED DIAGONAL ISOTROPIC MASS MATRIX WITHOUT ROTATIONAL MASSES
    # C
    HLM  = A*rho*L/2.;
    CA = HLM;
    CB = 0.0;
    CC = 0.0;
    CD = 0.0;
    fill!(MM, 0.0);
    MM[1,1]    = MM[1,1]    + CA;
    MM[2,2]    = MM[2,2]    + CA;
    MM[3,3]    = MM[3,3]    + CA;
    MM[7,7]    = MM[7,7]    + CA;
    MM[8,8]    = MM[8,8]    + CA;
    MM[9,9]    = MM[9,9]    + CA;
    return MM
end

"""
    local_mass!(MM, A, I1, I2, I3, rho, L, mass_type)

Mass matrix of the beam.

# Arguments
`A`= cross-sectional area,
`I1`=central moment of inertia of the cross-section about the x1 axis,
`I2`, `I3`=central moment of inertia of the cross-section about the x2 and x3
coordinate axis,
`rho`=mass density,
`L`= initial length of the element,

# Outputs
`MM` = local mass matrix, 12 x 12
In the element frame the mass matrix is constant.
"""
function local_mass!(MM, A, I1, I2, I3, rho, L, mass_type)
    if (mass_type == MASS_TYPE_CONSISTENT_WITH_ROTATION_INERTIA)
        return local_mass_CONSISTENT_WITH_ROTATION_INERTIA!(MM, A, I1, I2, I3, rho, L)
    elseif (mass_type == MASS_TYPE_CONSISTENT_NO_ROTATION_INERTIA)
        return local_mass_CONSISTENT_NO_ROTATION_INERTIA!(MM, A, I1, I2, I3, rho, L)
    elseif (mass_type == MASS_TYPE_LUMPED_DIAGONAL_WITH_ROTATION_INERTIA)
        return local_mass_LUMPED_DIAGONAL_WITH_ROTATION_INERTIA!(MM, A, I1, I2, I3, rho, L)
    elseif (mass_type == MASS_TYPE_LUMPED_DIAGONAL_NO_ROTATION_INERTIA)
        return local_mass_LUMPED_DIAGONAL_NO_ROTATION_INERTIA!(MM, A, I1, I2, I3, rho, L)
    end
end

function local_mass_original!(MM, A, I1, I2, I3, rho, L, mass_type)
    if (mass_type == MASS_TYPE_CONSISTENT_WITH_ROTATION_INERTIA)
        # C
        # C  CONSISTENT MASS MATRIX including ROTATIONAL MASSES
        # C  Formulation of the (3.38), (3.39) equation from Dykstra's thesis
        MM .= (rho*A*L)*[
            1/3      0        0        0        0        0       1/6       0        0        0        0        0
            0     13/35       0        0        0    11*L/210     0      9/70       0        0        0    -13*L/420
            0        0     13/35       0    -11*L/210    0        0        0       9/70      0    13*L/420     0
            0        0        0     I1/3/A      0        0        0        0        0      I1/6/A     0        0
            0        0   -11*L/210     0     L^2/105     0        0        0    -13*L/420    0    -L^2/140     0
            0    11*L/210     0        0        0     L^2/105     0    13*L/420     0        0        0    -L^2/140
            1/6      0        0        0        0        0       1/3       0        0        0        0        0
            0      9/70       0        0        0     13*L/420    0      13/35      0        0        0   -11*L/210
            0        0       9/70    0     -13*L/420     0        0        0      13/35      0    11*L/210     0
            0        0        0     I1/6/A      0        0        0        0        0      I1/3/A     0        0
            0        0    13*L/420     0    -L^2/140     0        0        0     11*L/210     0     L^2/105    0
            0   -13*L/420     0        0        0    -L^2/140     0   -11*L/210     0         0        0    L^2/105];
        MM .+= (rho/L)*[
            0       0        0        0       0         0       0       0        0       0       0        0
            0    6/5*I2      0        0       0     L/10*I2     0    -6/5*I2     0       0       0    L/10*I2
            0       0     6/5*I3      0   -L/10*I3      0       0       0     -6/5*I3    0   -L/10*I3     0
            0       0        0        0       0         0       0       0        0       0       0        0
            0       0    -L/10*I3     0  2*L^2/15*I3    0       0       0     L/10*I3    0   -L^2/30*I3   0
            0    L/10*I2     0        0       0   2*L^2/15*I2   0    -L/10*I2    0       0       0    -L^2/30*I2
            0       0        0        0       0         0       0       0        0       0       0        0
            0   -6/5*I2      0        0       0     -L/10*I2    0    6/5*I2      0       0       0     -L/10*I2
            0       0     -6/5*I3     0    L/10*I3      0       0       0      6/5*I3    0    L/10*I3     0
            0       0        0        0       0         0       0       0        0       0       0        0
            0       0     -L/10*I3    0  -L^2/30*I3     0       0       0     L/10*I3    0   2*L^2/15*I3  0
            0   L/10*I2      0        0       0    -L^2/30*I2   0   -L/10*I2     0       0       0   2*L^2/15*I2];
    elseif (mass_type == MASS_TYPE_CONSISTENT_NO_ROTATION_INERTIA)
        # C
        # C  CONSISTENT MASS MATRIX excluding ROTATIONAL MASSES
        # C  Formulation of the (3.38) equation from Dykstra's thesis, no rotational inertia
        MM .= (rho*A*L)*[
            1/3      0        0        0        0        0       1/6       0        0        0        0        0
            0     13/35       0        0        0    11*L/210     0      9/70       0        0        0    -13*L/420
            0        0     13/35       0    -11*L/210    0        0        0       9/70      0    13*L/420     0
            0        0        0     I1/3/A      0        0        0        0        0      I1/6/A     0        0
            0        0   -11*L/210     0     L^2/105     0        0        0    -13*L/420    0    -L^2/140     0
            0    11*L/210     0        0        0     L^2/105     0    13*L/420     0        0        0    -L^2/140
            1/6      0        0        0        0        0       1/3       0        0        0        0        0
            0      9/70       0        0        0     13*L/420    0      13/35      0        0        0   -11*L/210
            0        0       9/70    0     -13*L/420     0        0        0      13/35      0    11*L/210     0
            0        0        0     I1/6/A      0        0        0        0        0      I1/3/A     0        0
            0        0    13*L/420     0    -L^2/140     0        0        0     11*L/210     0     L^2/105    0
            0   -13*L/420     0        0        0    -L^2/140     0   -11*L/210     0         0        0    L^2/105];
    elseif (mass_type == MASS_TYPE_LUMPED_DIAGONAL_WITH_ROTATION_INERTIA)
        # C
        # C  LUMPED DIAGONAL MASS MATRIX WITH ROTATIONAL MASSES
        # C
        HLM  = A*rho*L/2.;
        HLI1 = rho*I1* L/2.;
        HLI2 = rho*I2* L/2.;
        HLI3 = rho*I3* L/2.;
        CA = HLM;
        CB = HLI1;
        CC = HLI2;
        CD = HLI3;
        fill!(MM, 0.0);
        MM[1,1]    = MM[1,1]    + CA;
        MM[2,2]    = MM[2,2]    + CA;
        MM[3,3]    = MM[3,3]    + CA;
        MM[4,4]    = MM[4,4]    + CB;
        MM[5,5]    = MM[5,5]    + CC;
        MM[6,6]    = MM[6,6]    + CD;
        MM[7,7]    = MM[7,7]    + CA;
        MM[8,8]    = MM[8,8]    + CA;
        MM[9,9]    = MM[9,9]    + CA;
        MM[10,10]  = MM[10,10]  + CB;
        MM[11,11]  = MM[11,11]  + CC;
        MM[12,12]  = MM[12,12]  + CD;
    elseif (mass_type == MASS_TYPE_LUMPED_DIAGONAL_NO_ROTATION_INERTIA)
        # C
        # C  LUMPED DIAGONAL ISOTROPIC MASS MATRIX WITHOUT ROTATIONAL MASSES
        # C
        HLM  = A*rho*L/2.;
        CA = HLM;
        CB = 0.0;
        CC = 0.0;
        CD = 0.0;
        fill!(MM, 0.0);
        MM[1,1]    = MM[1,1]    + CA;
        MM[2,2]    = MM[2,2]    + CA;
        MM[3,3]    = MM[3,3]    + CA;
        MM[7,7]    = MM[7,7]    + CA;
        MM[8,8]    = MM[8,8]    + CA;
        MM[9,9]    = MM[9,9]    + CA;
    end
    return MM
end

"""
    local_stiffness!(SM, E, G, A, I2, I3, J, A2s, A3s, L, aN, DN)

Compute the local elastic stiffness matrix.

# Arguments
`E`, `G`= Young's and shear modulus,
`A`= cross-sectional area,
`I2`, `I3`=central moment of inertia of the cross-section about the x2 and x3
coordinate axis,
`J`=St Venant torsion constant,
`L`= current length of the element,

# Outputs
`SM` = local stiffness matrix, 12 x 12
"""
function local_stiffness!(SM, E, G, A, I2, I3, J, A2s, A3s, L, aN, DN)
    local_cartesian_to_natural!(aN, L);
    natural_stiffness!(DN, E, G, A, I2, I3, J, A2s, A3s, L);
    SM .= aN'*DN*aN;
    return SM
end

"""
    natural_forces!(PN, E, G, A, I2, I3, J, A2s, A3s, L, dN, DN)

Compute the natural forces from the natural deformations.

# Argument
`E`, `G`= Young's and shear modulus,
`A`= cross-sectional area,
`I2`, `I3`=central moment of inertia of the cross-section about the x2 and x3
coordinate axis,
`J`=St Venant torsion constant,
`L`= current length of the element,
`dN`= column vector of natural deformations; see local_frames()

# Outputs
`PN` = column vector of natural forces;
     `PN[1]`= axial force;
     `PN[2]`= symmetric bending moment in the plane x1-x2;
     `PN[3]`= anti-symmetric bending bending moment in the plane x1-x2;
     `PN[4]`= symmetric bending bending moment in the plane x1-x3;
     `PN[5]`= anti-symmetric bending bending moment in the plane x1-x3;
     `PN[6]`= axial torque.
"""
function natural_forces!(PN, E, G, A, I2, I3, J, A2s, A3s, L, dN, DN)
    natural_stiffness!(DN, E, G, A, I2, I3, J, A2s, A3s, L);
    #     Natural forces
    mul!(PN, DN, dN);
    # Note that the non-constitutive stiffness due to pre-existing internal forces is currently omitted
end

"""
    natural_stiffness_Bernoulli!(DN, E, G, A, I2, I3, J, A2s, A3s, L)

Compute the natural stiffness matrix.

function DN = natural_stiffness(self, E, G, A, I2, I3, J, L)

# Arguments
- `E`, `G`= Young's and shear modulus,
- `A`= cross-sectional area,
- `I2`, `I3`= central moment of inertia of the cross-section about the x2 and x3
  coordinate axis,
- `J`= St Venant torsion constant,
- `A2s` = effective area for shear in the direction of x2 (ignored)
- `A3s` = effective area for shear in the direction of x3 (ignored)
- `L`= current length of the element,

# Outputs
- `DN` = 6 x 6 natural stiffness matrix
"""
function natural_stiffness_Bernoulli!(DN, E, G, A, I2, I3, J, A2s, A3s, L)
    fill!(DN, 0.0)
    DN[1, 1] = E*A/L
    DN[2, 2] = E*I3/L
    DN[3, 3] = 3*E*I3/L
    DN[4, 4] = E*I2/L
    DN[5, 5] = 3*E*I2/L
    DN[6, 6] = G*J/L
    return DN
end

"""
    natural_stiffness_Timoshenko!(DN, E, G, A, I2, I3, J, A2s, A3s, L)

Compute the natural stiffness matrix.

function DN = natural_stiffness(self, E, G, A, I2, I3, J, L)

# Arguments
- `E`, `G`= Young's and shear modulus,
- `A`= cross-sectional area,
- `I2`, `I3`= central moment of inertia of the cross-section about the x2 and x3
  coordinate axis,
- `J`= St Venant torsion constant,
- `A2s` = effective area for shear in the direction of x2
- `A3s` = effective area for shear in the direction of x3
- `L`= current length of the element,

# Outputs
- `DN` = 6 x 6 natural stiffness matrix
"""
function natural_stiffness_Timoshenko!(DN, E, G, A, I2, I3, J, A2s, A3s, L)
    fill!(DN, 0.0)
    DN[1, 1] = E*A/L
    DN[2, 2] = E*I3/L
    Phi3 = 12 * E * I3 / (G * A2s * L^2)
    DN[3, 3] = 3*E*I3/L/(1+Phi3)
    DN[4, 4] = E*I2/L
    Phi2 = 12 * E * I2 / (G * A3s * L^2)
    DN[5, 5] = 3*E*I2/L/(1+Phi3)
    DN[6, 6] = G*J/L
    return DN
end

function natural_stiffness!(DN, E, G, A, I2, I3, J, A2s, A3s, L)
    if A2s == Inf || A3s == Inf
        return natural_stiffness_Bernoulli!(DN, E, G, A, I2, I3, J, A2s, A3s, L)
    else
        return natural_stiffness_Timoshenko!(DN, E, G, A, I2, I3, J, A2s, A3s, L)
    end
end

"""
    local_forces!(FL, PN, L, aN)

Compute forces through which the element acts on the nodes in the
local coordinate system.

# Arguments
`PN` = column vector of natural forces;
`L`= current length of the element,
`aN`= transformation matrix to take Cartesian (local) displacement increments in the
     element frame and to produce increments of natural deformations;
     see `local_frame_and_def!` for the definition of the natural deformations

# Outputs
FL = vector of forces acting on the nodes in the local coordinate system
"""
function local_forces!(FL, PN, L, aN)
    local_cartesian_to_natural!(aN, L);
    mul!(FL, Transpose(aN), PN);
    return FL
end

end # module


text:
@testset "Ideal manifold" begin
    f_0 = 1575420e3
    c₀ = 299_792_458
    λ = c₀ / f_0
    ant_pos = SVector(λ / 4 * SVector(1, 1, 0), λ / 4 * SVector(-1, 1, 0), λ / 4 * SVector(1, -1, 0), λ / 4 * SVector(-1, -1 , 0))
    manifold = @inferred IdealManifold(f_0, ant_pos)

    @test typeof(manifold) <: AbstractManifold{4}
    @test @inferred(get_steer_vec(manifold, SVector(0,0,1))) ≈ [1, 1, 1, 1]
    @test @inferred(get_steer_vec(manifold, SVector(0,1,0))) ≈ [1im, 1im, -1im, -1im]
    @test @inferred(get_steer_vec(manifold, SVector(1,0,0))) ≈ [1im, -1im, 1im, -1im]

    @test @inferred(get_steer_vec(manifold, Spherical(1.0, 0.0, 0.0))) ≈ [1im, -1im, 1im, -1im]
    @test @inferred(get_steer_vec(manifold, Spherical(1.0, 0.0, 0.0), RotXYZ(0.0,0.0,π/2))) ≈ [1im, 1im, -1im, -1im]
    @test @inferred(get_steer_vec(manifold, SVector(1,0,0), RotXYZ(0.0,0.0,π/2))) ≈ [1im, 1im, -1im, -1im]

    manifold = @inferred IdealManifold(f_0, λ / 4 * SVector(1, 1, 0), λ / 4 * SVector(-1, 1, 0), λ / 4 * SVector(1, -1, 0), λ / 4 * SVector(-1, -1 , 0))

    @test typeof(manifold) <: AbstractManifold{4}
    @test @inferred(get_steer_vec(manifold, SVector(0,0,1))) ≈ [1, 1, 1, 1]
    @test @inferred(get_steer_vec(manifold, SVector(0,1,0))) ≈ [1im, 1im, -1im, -1im]
    @test @inferred(get_steer_vec(manifold, SVector(1,0,0))) ≈ [1im, -1im, 1im, -1im]

    @test @inferred(get_steer_vec(manifold, Spherical(1.0, 0.0, 0.0))) ≈ [1im, -1im, 1im, -1im]
    @test @inferred(get_steer_vec(manifold, Spherical(1.0, 0.0, 0.0), RotXYZ(0.0,0.0,π/2))) ≈ [1im, 1im, -1im, -1im]
    @test @inferred(get_steer_vec(manifold, SVector(1,0,0), RotXYZ(0.0,0.0,π/2))) ≈ [1im, 1im, -1im, -1im]

    manifold = @inferred IdealManifold(f_0, λ / 4 * @SMatrix([1 -1 1 -1; 1 1 -1 -1; 0 0 0 0]))

    @test typeof(manifold) <: AbstractManifold{4}
    @test @inferred(get_steer_vec(manifold, SVector(0,0,1))) ≈ [1, 1, 1, 1]
    @test @inferred(get_steer_vec(manifold, SVector(0,1,0))) ≈ [1im, 1im, -1im, -1im]
    @test @inferred(get_steer_vec(manifold, SVector(1,0,0))) ≈ [1im, -1im, 1im, -1im]

    @test @inferred(get_steer_vec(manifold, Spherical(1.0, 0.0, 0.0))) ≈ [1im, -1im, 1im, -1im]
    @test @inferred(get_steer_vec(manifold, Spherical(1.0, 0.0, 0.0), RotXYZ(0.0,0.0,π/2))) ≈ [1im, 1im, -1im, -1im]
    @test @inferred(get_steer_vec(manifold, SVector(1,0,0), RotXYZ(0.0,0.0,π/2))) ≈ [1im, 1im, -1im, -1im]

    manifold = @inferred IdealManifold()
    @test typeof(manifold) <: AbstractManifold{1}
    @test @inferred(get_steer_vec(manifold, SVector(1,0,0))) == 1.0
    @test @inferred(get_steer_vec(manifold, SVector(1,0,0), RotXYZ(0.0,0.0,0.0))) == 1.0

    num_ant_x = 2
    dist_x = 0.5
    manifold = IdealManifold(1, num_ant_x, dist_x; c0 = 1)
    positions = manifold.scaled_antenna_positions
    @test get_num_ants(manifold) == num_ant_x^2
    @test get_steer_vec(manifold, [0, 0, 1]) ≈ ones(num_ant_x^2)
    @test get_steer_vec(manifold, [1, 0, 0]) ≈ sign.(positions[1,:]).*im
    @test get_steer_vec(manifold, [0, 1, 0]) ≈ sign.(positions[2,:]).*im
    @test all(abs.(rem.(positions[1,:] .- π/2, 1.0π, RoundNearest)) .≈ 0)
    @test all(abs.(rem.(positions[2,:] .- π/2, 1.0π, RoundNearest)) .≈ 0)

    num_ant_y = 3
    dist_y = 1
    manifold = IdealManifold(1, num_ant_x, num_ant_y, dist_x, dist_y; c0 = 1)
    positions = manifold.scaled_antenna_positions
    @test get_num_ants(manifold) == num_ant_x * num_ant_y
    @test get_steer_vec(manifold, [0, 0, 1]) ≈ ones(num_ant_x*num_ant_y)
    @test get_steer_vec(manifold, [1, 0, 0]) ≈ sign.(positions[1,:]).*im
    @test get_steer_vec(manifold, [0, 1, 0]) ≈ ones(num_ant_x*num_ant_y)
    @test all(rem.(positions[1,:] .- π/2, 1.0π, RoundNearest) .≈ 0)
    @test all(rem.(positions[2,:], 2.0π, RoundNearest) .≈ 0)
end


text:
using REPL

lbs = String[]
tc = Tuple{String,String}[]
for (a,b) in REPL.REPLCompletions.latex_symbols
	push!(tc, (a[2:end],b))
	bs=Symbol(b)
	if isdefined(Main,bs)
		be=string(eval(bs))
		push!(lbs, b^2*"\n$be"^(be≠b))
	end
end

lbs = sort(unique(lbs))
tc = sort(tc, by=x->(length(x[1]),x))
tcf = first.(tc)
tcc = join(tcf) |> unique |> sort |> join |> x->replace(x, ["\$()*+-./?[\\]^{|}"...]=>x->'\\'*x) |> x-> "/\\\\([$x]+)\$/"
tcu = filter(i->!isnothing(match(r"^\^.$",i)), tcf) .|> last |> sort |> join |> x->replace(x, ["\$()*+-./?[\\]^{|}"...]=>x->'\\'*x) |> x-> "/\\\\\\^([$x]+)\$/"
tcl = filter(i->!isnothing(match(r"^_.$",i)), tcf)  .|> last |> sort |> join |> x->replace(x, ["\$()*+-./?[\\]^{|}"...]=>x->'\\'*x) |> x-> "/\\\\_([$x]+)\$/"

slbs = '['*join(repr.(lbs),',')*']'
stc = '{'*join(tc.|>x->join(repr.(x),':'),',')*'}'

out = """
, lbs = $slbs
, tc = $stc
, tcc = $tcc
, tcl = $tcl
, tcu = $tcu
"""
try
	clipboard(out)
catch
	print(out)
end

max_stars_repo_name: razetime/lb

max_issues_repo_name: razetime/lb

max_forks_repo_name: razetime/lb


text:
#/usr/bin/env julia

using Wut
using Base.Test

# Run tests

tic()

@time @test include("ScalarEncoderTests.jl")
@time @test include("RandomDistributedScalarEncoderTests.jl")

toc()


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::NoDensityTransform, ::AnyDensityLike) = DensityIdentityTransform()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::NoDensityTransform, ::AbstractPosteriorDensity) = DensityIdentityTransform()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::NoDensityTransform, ::DensityWithDiff{<:Any,<:AbstractPosteriorDensity}) = DensityIdentityTransform()

bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToUniform, ::AbstractPosteriorDensity) = PriorSubstitution()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToUniform, ::DensityWithDiff{<:Any,<:AbstractPosteriorDensity}) = PriorSubstitution()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToUniform, ::DistributionDensity) = PriorSubstitution()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToUniform, ::RenormalizedDensity{<:DistributionDensity}) = PriorSubstitution()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToUniform, ::StandardUniformDensity) = DensityIdentityTransform()

bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToGaussian, ::AbstractPosteriorDensity) = PriorSubstitution()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToGaussian, ::DensityWithDiff{<:Any,<:AbstractPosteriorDensity}) = PriorSubstitution()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToGaussian, ::DistributionDensity) = PriorSubstitution()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToGaussian, ::RenormalizedDensity{<:DistributionDensity}) = PriorSubstitution()
bat_default(::typeof(bat_transform), ::Val{:algorithm}, ::PriorToGaussian, ::StandardNormalDensity) = DensityIdentityTransform()

# ToDo: Add ToUnitBounded and ToUnbounded


text:
"""
    bridges(g)

Compute the [bridges](https://en.m.wikipedia.org/wiki/Bridge_(graph_theory))
of a connected graph `g` and return an array containing all bridges, i.e edges
whose deletion increases the number of connected components of the graph.
# Examples
```jldoctest
julia> using LightGraphs

julia> bridges(star_graph(5))
8-element Array{LightGraphs.SimpleGraphs.SimpleEdge{Int64},1}:
 Edge 1 => 2
 Edge 1 => 3
 Edge 1 => 4
 Edge 1 => 5

julia> bridges(path_graph(5))
8-element Array{LightGraphs.SimpleGraphs.SimpleEdge{Int64},1}:
 Edge 4 => 5
 Edge 3 => 4
 Edge 2 => 3
 Edge 1 => 2
```
"""
function bridges end
@traitfn function bridges(g::AG::(!IsDirected)) where {T, AG<:AbstractGraph{T}}
    s = Vector{Tuple{T, T, T}}()
    low = zeros(T, nv(g)) #keeps track of the earliest accesible time of a vertex in DFS-stack, effect of having back-edges is considered here
    pre = zeros(T, nv(g)) #checks the entry time of a vertex in the DFS-stack, pre[u] = 0 if a vertex isn't visited; non-zero, otherwise
    bridges = Edge{T}[]   #keeps record of the bridge-edges

    # We iterate over all vertices, and if they have already been visited (pre != 0), we don't start a DFS from that vertex.
    # The purpose is to create a DFS forest.
    @inbounds for u in vertices(g)
        pre[u] != 0 && continue
        v = u #currently visiting vertex
        wi::T = zero(T) #index of children of v
        w::T = zero(T) #children of v
        cnt::T = one(T) # keeps record of the time
        first_time = true

        #start of DFS
        while !isempty(s) || first_time
            first_time = false
            if  wi < 1 #initialisation for vertex v
                pre[v] = cnt
                cnt += 1
                low[v] = pre[v]
                v_neighbors = outneighbors(g, v)
                wi = 1
            else
                wi, u, v = pop!(s) # the stack states, explained later
                v_neighbors = outneighbors(g, v)
                w = v_neighbors[wi]
                low[v] = min(low[v], low[w]) # condition check for (v, w) being a tree-edge
                if low[w] > pre[v]
                    edge = v < w ? Edge(v, w) : Edge(w, v)
                    push!(bridges, edge)
                end
                wi += 1
            end

            # here, we're iterating of all the childen of vertex v, if unvisited, we start a DFS from that child, else we update the low[v] as the edge is a back-edge.
            while wi <= length(v_neighbors)
                w = v_neighbors[wi]
                # If this is true , this indicates the vertex is still unvisited, then we push this on the stack.
                # Pushing onto the stack is analogous to visiting the vertex and starting DFS from that vertex.
                if pre[w] == 0
                    push!(s, (wi, u, v)) # the stack states are (index of child, currently visiting vertex, parent vertex of the child)
                    #updates the value for stimulating DFS from top of the stack
                    wi = 0
                    u = v
                    v = w
                    break
                elseif w != u # (v, w) is a back-edge
                    low[v] = min(low[v], pre[w]) # condition for back-edges
                end
                wi += 1
            end
            wi < 1 && continue
        end

    end

    return bridges
end


text:
using LoopVectorization: @avx

export  ExpandedFunction,
        active_dim,
        alleval,
        getbasis,
        evaluate_basis!,
        evaluate_basis,
        repeated_evaluate_basis,
        grad_xk_basis!,
        grad_xk_basis,
        repeated_grad_xk_basis!,
        repeated_grad_xk_basis,
        repeated_hess_xk_basis!,
        repeated_hess_xk_basis,
        grad_x_basis!,
        grad_x_basis,
        hess_x_basis!,
        hess_x_basis,
        evaluate, grad_x, hess_x,
        grad_xd, hess_xd,
        grad_x_grad_xd,
        hess_x_grad_xd,
        reduced_grad_x_grad_xd!,
        reduced_grad_x_grad_xd,
        reduced_hess_x_grad_xd!,
        reduced_hess_x_grad_xd,
        grad_coeff, hess_coeff,
        grad_coeff_grad_xd,
        hess_coeff_grad_xd

"""
$(TYPEDEF)

`ExpandedFunction` decomposes a multi-dimensional function f:R^{Nx} → R onto
a basis of `MultiFunction` ψ_α where c_α are scalar coefficients
for each MultiFunction:
f(x1, x2, ..., xNx) = ∑_α c_α ψ_α(x1, x2, ..., xNx), where
`Nψ` is the number of `MultiFunction`s used, and
`Nx` is the dimension of the input vector `x`.

## Fields
$(TYPEDFIELDS)

"""
struct ExpandedFunction
    m::Int64
    Nψ::Int64
    Nx::Int64
    B::MultiBasis
    idx::Array{Int64,2}
    dim::Array{Int64, 1} # contains the active dimensions, i.e. columns of idx not equal to zeros
    coeff::Array{Float64,1}
    function ExpandedFunction(B::MultiBasis, idx::Array{Int64,2}, coeff::Array{Float64,1})
            Nψ = size(idx,1)
            Nx = B.Nx
            @assert Nψ == size(coeff, 1) "The dimension of the basis functions don't
                                            match the number of coefficients"

            @assert size(idx,2) == Nx "Size of the array of multi-indices idx is wrong"
        return new(B.B.m, Nψ, Nx, B, idx, active_dim(idx, B), coeff)
    end
end

"""
$(TYPEDSIGNATURES)

Returns the kind of basis of the `ExpandedFunction` `f`.
"""
getbasis(f::ExpandedFunction) = getbasis(f.B)

"""
$(TYPEDSIGNATURES)

Evaluates the `ExpandedFunction` `f` at `x`.
"""
function (f::ExpandedFunction)(x::Array{T,1}) where {T<:Real}
    out = 0.0
    @inbounds for i=1:f.Nψ
        fi = MultiFunction(f.B, f.idx[i,:])
        out += f.coeff[i]*fi(x)
    end
    return out
end

"""
$(TYPEDSIGNATURES)

Returns the active dimensions of the  set of multi-indices `idx` for the `Basis` `B`.
"""
function active_dim(idx::Array{Int64,2}, B::T) where {T<:Basis}
    # Nx should always be an active dimension (we need to ensures
    # that we have a strictly increasing function in the last component)
    dim = Int64[]
    Nx = size(idx,2)
    if iszerofeatureactive(B) == false
        @inbounds for i=1:Nx-1
            if !all(view(idx,:,i) .== 0)
                push!(dim, i)
            end
        end
        push!(dim, Nx)
    else
        dim = collect(1:size(idx,2))
    end

    return dim
end

"""
$(TYPEDSIGNATURES)

Returns the active dimensions of the  set of multi-indices `idx` for the  `MultiBasis` `B`.
"""
active_dim(idx::Array{Int64,2}, B::MultiBasis) = active_dim(idx, B.B)

"""
$(TYPEDSIGNATURES)

Returns the active dimensions of the `ExpandedFunction` `f`.
"""
active_dim(f::ExpandedFunction) = f.dim

# alleval computes the evaluation, gradient and hessian of the function
# use it for validatio since it is slower than the other array-based variants
"""
$(TYPEDSIGNATURES)

Returns the evaluation, gradient and hessian of the `ExpandedFunction` `f` for the ensemble matrix `X`.
Note that this routine is not designed for speed, but for validation only,
as all the derivatives are computedto machine precision with `ForwardDiff.jl`
"""
function alleval(f::ExpandedFunction, X)
        Nx, Ne = size(X)
        Nψ = f.Nψ
        ψ = zeros(Ne, Nψ)
       dψ = zeros(Ne, Nψ, Nx)
      d2ψ = zeros(Ne, Nψ, Nx, Nx)
   result = DiffResults.HessianResult(zeros(Nx))

    for i=1:Nψ
        fi = MultiFunction(f.B, f.idx[i,:])
        for j=1:Ne
            result = ForwardDiff.hessian!(result, fi, X[:,j])
            ψ[j,i] = DiffResults.value(result)
            dψ[j,i,:,:] .= DiffResults.gradient(result)
            d2ψ[j,i,:,:] .= DiffResults.hessian(result)
        end
    end
    return ψ, dψ, d2ψ
end

"""
$(TYPEDSIGNATURES)

Evaluates in-place the basis of `ExpandedFunction` `f` for the ensemble matrix `X` along the dimensions `dims` for the set of multi-indices of features `idx`
"""
function evaluate_basis!(ψ, f::ExpandedFunction, X, dims::Union{Array{Int64,1},UnitRange{Int64}}, idx::Array{Int64,2})
    Nψreduced = size(idx,1)
    NxX, Ne = size(X)
    @assert NxX == f.Nx "Wrong dimension of the input sample X"
    @assert size(ψ) == (Ne, Nψreduced) "Wrong dimension of the ψ"

    maxdim = maximum(idx)
    # ψvander = zeros(Ne, maxdim)
    fill!(ψ, 1.0)
    if maxdim+1<= Nψreduced
        ψtmp = zero(ψ)
    else
        ψtmp = zeros(Ne, maxdim+1)
    end

    # The maximal size of ψtmp assumes that the set of index is downward closed
    # such that Nψreduced is always smaller of equal to maxj+1
    @inbounds for j in intersect(dims, f.dim)
        idxj = view(idx,:,j)
        maxj = maximum(idxj)
        Xj = view(X,j,:)
        ψj = ψtmp[:,1:maxj+1]

        vander!(ψj, f.B.B, maxj, 0, Xj)

        @avx for l = 1:Nψreduced
            for k=1:Ne
                ψ[k,l] *= ψj[k, idxj[l] + 1]
            end
        end
        # @avx  ψ .*= view(ψj,:, midxj .+ 1)#view(ψvanderj, :, midxj .+ 1)#
    end
    return ψ
end

# In-place versions
evaluate_basis!(ψ, f::ExpandedFunction, X, dims::Union{Array{Int64,1},UnitRange{Int64}}) =
              evaluate_basis!(ψ, f, X, dims, f.idx)

evaluate_basis!(ψ, f::ExpandedFunction, X, idx::Array{Int64,2}) =
            evaluate_basis!(ψ, f, X, f.dim, idx)
            # evaluate_basis!(ψ, f, X, 1:f.Nx, idx)

evaluate_basis!(ψ, f::ExpandedFunction, X) =
            evaluate_basis!(ψ, f, X, f.dim, f.idx)
            # evaluate_basis!(ψ, f, X, 1:f.Nx, f.idx)

# Versions with allocations
"""
$(TYPEDSIGNATURES)

Evaluates the basis of `ExpandedFunction` `f` for the ensemble matrix `X` along the dimensions `dims` for the set of multi-indices of features `idx`
"""
evaluate_basis(f::ExpandedFunction, X, dims::Union{Array{Int64,1},UnitRange{Int64}}, idx::Array{Int64,2}) =
              evaluate_basis!(zeros(size(X,2),size(idx,1)), f, X, dims, idx)

evaluate_basis(f::ExpandedFunction, X, dims::Union{Array{Int64,1},UnitRange{Int64}}) =
              evaluate_basis!(zeros(size(X,2),size(f.idx,1)), f, X, dims, f.idx)

evaluate_basis(f::ExpandedFunction, X, idx::Array{Int64,2}) =
            evaluate_basis!(zeros(size(X,2),size(idx,1)), f, X, f.dim, idx)
            # evaluate_basis!(zeros(size(X,2),size(idx,1)), f, X, 1:f.Nx, idx)

evaluate_basis(f::ExpandedFunction, X) =
            evaluate_basis!(zeros(size(X,2),size(f.idx,1)), f, X, f.dim, f.idx)
            # evaluate_basis!(zeros(size(X,2),size(f.idx,1)), f, X, 1:f.Nx, f.idx)

"""
$(TYPEDSIGNATURES)

Evaluates the basis of `ExpandedFunction` `f` for the last component
"""
function repeated_evaluate_basis(f::ExpandedFunction, x, idx::Array{Int64,2})
    # Compute the last component
    midxj = idx[:,f.Nx]
    maxj = maximum(midxj)
    ψj = vander(f.B.B, maxj, 0, x)
    return ψj[:, midxj .+ 1]
end

repeated_evaluate_basis(f::ExpandedFunction, x) = repeated_evaluate_basis(f, x, f.idx)

# function grad_xk_basis!(dkψ, f::ExpandedFunction, X::Array{Float64,2}, k::Int64, grad_dim::Union{Int64, Array{Int64,1}}, dims::Union{Int64, UnitRange{Int64}, Array{Int64,1}}, idx::Array{Int64,2})
"""
$(TYPEDSIGNATURES)

Evaluates in-place the k-th (k>0) derivative of the basis features of `f` with respect to the `grad_dim` components of the states.
The i-th column of the output contains ∂^k...∂^k ∏_{j ∈ dims} ψ^i_{j}(x_1:n) ∂x_{grad_dim[1]}^k ... ∂x_{grad_dim[end]}^k
for the different columns of the ensemble matrix `X`, where ψ^i is the i-th feature.
"""
function grad_xk_basis!(dkψ, f::ExpandedFunction, X, k::Int64, grad_dim::Union{Int64, Array{Int64,1}}, dims::Union{Int64, UnitRange{Int64}, Array{Int64,1}}, idx::Array{Int64,2})
    m = f.m
    Nx = f.Nx

    NxX, Ne = size(X)
    Nψreduced = size(idx,1)

    @assert NxX == Nx "Wrong dimension of the input sample X"
    @assert size(dkψ) == (Ne, Nψreduced) "Wrong dimension of ψ"
    # ∂ᵏf/∂x_{grad_dim} = ψ
    @assert k>0  "The derivative order k must be >0, if k=0, you better use evaluate_basis!"

    T = typeof(grad_dim)
    if T <:Array{Int64,1}
        @assert all(1 .<= grad_dim)
        @assert all(grad_dim .<= Nx)
    elseif T <: Int64
        @assert 1 <= grad_dim <= Nx
    end

    # Check if we are taking derivative with respect
    # to non active dimension
    if any([!(gdim ∈ f.dim) for gdim in grad_dim])
        fill!(dkψ, 0.0)
    else
        fill!(dkψ, 1.0)

        for j in intersect(dims, f.dim)
            midxj = idx[:,j]
            maxj = maximum(midxj)
            Xj = view(X,j,:)
            if j in grad_dim # Compute the kth derivative along grad_dim
                dkψj = vander(f.B.B, maxj, k, Xj)

            else # Simple evaluation
                dkψj = vander(f.B.B, maxj, 0, Xj)
            end
            dkψ .*= dkψj[:, midxj .+ 1]
        end
    end
    return dkψ
end

# In-place versions of grad_xk_basis!

grad_xk_basis!(dkψ, f::ExpandedFunction, X, k::Int64, grad_dim::Union{Int64, Array{Int64,1}}, dims::Union{Int64, UnitRange{Int64}, Array{Int64,1}}) =
              grad_xk_basis!(dkψ, f, X, k, grad_dim, dims, f.idx)

grad_xk_basis!(dkψ, f::ExpandedFunction, X, k::Int64, grad_dim::Union{Int64, Array{Int64,1}}, idx::Array{Int64,2}) =
              grad_xk_basis!(dkψ, f, X, k, grad_dim, 1:f.Nx, idx)

grad_xk_basis!(dkψ, f::ExpandedFunction, X, k::Int64, grad_dim::Union{Int64, Array{Int64,1}}) =
              grad_xk_basis!(dkψ, f, X, k, grad_dim, 1:f.Nx, f.idx)

# Versions with allocations
"""
$(TYPEDSIGNATURES)

Evaluates in-place the k-th (k>0) derivative of the basis features of `f` with respect to the `grad_dim` components of the states.
The i-th column of the output contains ∂^k...∂^k ∏_{j ∈ dims} ψ^i_{j}(x_{1:n}) ∂x_{grad_dim[1]}^k ... ∂x_{grad_dim[end]}^k
for the different columns of the ensemble matrix `X`, where ψ^i is the i-th feature.
"""
grad_xk_basis(f::ExpandedFunction, X, k::Int64, grad_dim::Union{Int64, Array{Int64,1}}, dims::Union{Int64, UnitRange{Int64}, Array{Int64,1}}, idx::Array{Int64,2}) =
             grad_xk_basis!(zeros(size(X,2), size(idx,1)), f, X, k, grad_dim, dims, f.idx)

grad_xk_basis(f::ExpandedFunction, X, k::Int64, grad_dim::Union{Int64, Array{Int64,1}}, dims::Union{Int64, UnitRange{Int64}, Array{Int64,1}}) =
             grad_xk_basis!(zeros(size(X,2), size(f.idx,1)), f, X, k, grad_dim, dims, f.idx)

grad_xk_basis(f::ExpandedFunction, X, k::Int64, grad_dim::Union{Int64, Array{Int64,1}}, idx::Array{Int64,2}) =
             grad_xk_basis!(zeros(size(X,2), size(idx,1)), f, X, k, grad_dim, 1:f.Nx, idx)

grad_xk_basis(f::ExpandedFunction, X, k::Int64, grad_dim::Union{Int64, Array{Int64,1}}) =
             grad_xk_basis!(zeros(size(X,2), size(f.idx,1)), f, X, k, grad_dim, 1:f.Nx, f.idx)

 """
 $(TYPEDSIGNATURES)

 Evaluates in-place the gradient of the basis features of `f` with respect to the different state components
 The i-th column of the output contains ∂ψ^{i}(x_1:n)∂x for the different columns of the ensemble matrix `X`,
 where ψ^{i} is the i-th feature.
 """
function grad_x_basis!(dψ::Array{Float64,3}, f::ExpandedFunction, X, idx::Array{Int64,2})
    m = f.m
    Nx = f.Nx
    # Compute the k-th order deriviative of an expanded function along the direction grad_dim
    Nψreduced = size(idx,1)
    fill!(dψ, 1.0)
    Ne, Nψr1, Nxψ = size(dψ)
    @assert Nψr1 == size(idx,1) "Wrong dimension of dψ"

    @inbounds for j=1:Nx
        dψj = view(dψ,:,:,j)
        grad_xk_basis!(dψj, f, X, 1, j, idx)
    end
    return dψ
end

# In place version
grad_x_basis!(dψ::Array{Float64,3}, f::ExpandedFunction, X::Array{Float64,2}) = grad_x_basis!(dψ, f, X, f.idx)

# Version with allocations
"""
$(TYPEDSIGNATURES)

Evaluates in-place the gradient of the basis features of `f` with respect to the different state components
The i-th column of the output contains ∂ψ^{i}(x_1:n)∂x for the different columns of the ensemble matrix `X`,
where ψ^{i} is the i-th feature.
"""
grad_x_basis(f::ExpandedFunction, X::Array{Float64,2}, idx::Array{Int64,2}) = grad_x_basis!(zeros(size(X,2), size(idx,1), f.Nx), f, X, idx)

grad_x_basis(f::ExpandedFunction, X::Array{Float64,2}) = grad_x_basis!(zeros(size(X,2), size(f.idx,1), f.Nx), f, X, f.idx)

"""
$(TYPEDSIGNATURES)

Evaluates in-place the hessian of the basis features of `f` with respect to the different state components
The i-th column of the output contains ∂^2ψ^{i}(x_1:n)∂x^2 for the different columns of the ensemble matrix `X`,
where ψ^{i} is the i-th feature.
"""
function hess_x_basis!(d2ψ::Array{Float64,4}, f::ExpandedFunction, X::Array{Float64,2}, idx::Array{Int64,2})
    m = f.m
    Nx = f.Nx
    # Compute the k=th order deriviative of an expanded function along the direction grad_dim
    Nψreduced = size(idx,1)
    fill!(d2ψ, 1)
    Ne, Nψr1, Nxψ1, Nxψ2 = size(d2ψ)
    @assert Ne == size(X,2)
    @assert Nxψ1 == Nxψ2 "Wrong dimension of d2ψ"
    @assert Nψr1 == size(idx,1) "Wrong dimension of d2ψ"

    # Fill diagonal components
    @inbounds for j=1:Nx
        d2ψj = view(d2ψ,:,:,j,j)
        grad_xk_basis!(d2ψj, f, X, 2, j, idx)
    end

    # Fill off-diagonal and exploit symmetry (Schwartz theorem)
    @inbounds for i=1:Nx
                for j=i+1:Nx
                    d2ψij = view(d2ψ,:,:,i,j)
                    grad_xk_basis!(d2ψij, f, X, 1, [i;j], idx)
                    d2ψ[:,:,j,i] .= d2ψ[:,:,i,j]
                end
    end
    return d2ψ
end

# In place version
hess_x_basis!(d2ψ::Array{Float64,4}, f::ExpandedFunction, X::Array{Float64,2}) = hess_x_basis!(d2ψ, f, X, f.idx)

# Version with allocations
"""
$(TYPEDSIGNATURES)

Evaluates the hessian of the basis features of `f` with respect to the different state components
The i-th column of the output contains ∂^2ψ^{i}(x_1:n)∂x^2 for the different columns of the ensemble matrix `X`,
where ψ^{i} is the i-th feature.
"""
hess_x_basis(f::ExpandedFunction, X::Array{Float64,2}, idx::Array{Int64,2}) = hess_x_basis!(zeros(size(X,2), size(idx,1), f.Nx, f.Nx), f, X, idx)

hess_x_basis(f::ExpandedFunction, X::Array{Float64,2})  = hess_x_basis!(zeros(size(X,2), size(f.idx,1), f.Nx, f.Nx), f, X, f.idx)

"""
$(TYPEDSIGNATURES)

Evaluates in-place the function `f` for the ensemble matrix `X`.
"""
function evaluate!(ψ, f::ExpandedFunction, X::Array{Float64,2})
    ψ .= evaluate_basis(f, X)*f.coeff
    return ψ
end

"""
$(TYPEDSIGNATURES)

Evaluates the function `f` for the ensemble matrix `X`.
"""
evaluate(f::ExpandedFunction, X::Array{Float64,2}) = evaluate!(zeros(size(X,2)), f, X)

"""
$(TYPEDSIGNATURES)

Evaluates the gradient of `f` with respect to the different state components
The i-th column of the output contains ∂f/∂x_i for the different columns of the ensemble matrix `X`.
"""
function grad_x(f::ExpandedFunction, X::Array{Float64,2})
    NxX, Ne = size(X)
    @assert NxX == f.Nx "Wrong dimension of the input"
    dψ = zeros(Ne, f.Nx)
    dψ_basis = grad_x_basis(f, X)
    @tensor dψ[a,b] = dψ_basis[a,c,b] * f.coeff[c]
    return dψ
end

"""
$(TYPEDSIGNATURES)

Evaluates the hessian of `f` with respect to the different state components
The (i,j,k) entry of the output contains ∂^2f(X[:,i])/∂x_j∂x_k.
"""
function hess_x(f::ExpandedFunction, X::Array{Float64,2})
    Nx = f.Nx
    NxX, Ne = size(X)
    @assert NxX == Nx "Wrong dimension of the input"
    d2ψ = zeros(Ne, Nx, Nx)
    d2ψ_basis = hess_x_basis(f, X)

    @tensor d2ψ[a,b,c] = d2ψ_basis[a,d,b,c] * f.coeff[d]
    return d2ψ
end

"""
$(TYPEDSIGNATURES)

Evaluates the gradient of `f` with respect to the last state component.
"""
function grad_xd(f::ExpandedFunction, X::Array{Float64,2})
    dψxd = grad_xk_basis(f, X, 1, f.Nx)
    return dψxd*f.coeff
end

"""
$(TYPEDSIGNATURES)

Evaluates the hessian of `f` with respect to the last state component.
"""
function hess_xd(f::ExpandedFunction, X::Array{Float64,2})
    d2ψxd = grad_xk_basis(f, X, 2, f.Nx)
    return d2ψxd*f.coeff
end

"""
$(TYPEDSIGNATURES)

Computes in-place the gradient with respect to the last state component of the basis of the last univariate function of each feature with multi-indices `idx` at `x`.
"""
function repeated_grad_xk_basis!(out, cache, f::ExpandedFunction, x, idx::Array{Int64,2})
    # Compute the k-th order deriviative of an expanded function along the last state component
    Ne = size(x,1)
    Nx = f.Nx

    # @assert size(out,1) = (N, size(idx, 1)) "Wrong dimension of the output vector"
    # ∂ᵏf/∂x_{grad_dim} = ψ
    k = 1
    grad_dim = Nx
    dims = Nx

    midxj = idx[:, Nx]
    maxj = maximum(midxj)
    # dkψj = zeros(Ne, maxj+1)
    vander!(cache, f.B.B, maxj, k, x)
    Nψreduced = size(idx, 1)
    @avx for l = 1:Nψreduced
        for k=1:Ne
            out[k,l] = cache[k, midxj[l] + 1]
        end
    end

    return out#dkψj[:, midxj .+ 1]
end

repeated_grad_xk_basis!(out, cache, f::ExpandedFunction, x) = repeated_grad_xk_basis!(out, cache, f, x, f.idx)

# repeated_grad_xk_basis(f::ExpandedFunction{m, Nψ, Nx}, x::Array{Float64,1}) where {m, Nψ, Nx} =
"""
$(TYPEDSIGNATURES)

Computes the gradient with respect to the last state component of the basis of the last univariate function of each feature with multi-indices `idx` at `x`.
"""
repeated_grad_xk_basis(f::ExpandedFunction, x, idx::Array{Int64,2}) =
    repeated_grad_xk_basis!(zeros(size(x,1),size(idx,1)), zeros(size(x,1), maximum(idx[:,f.Nx])+1), f, x, idx)
repeated_grad_xk_basis(f::ExpandedFunction, x) = repeated_grad_xk_basis(f, x, f.idx)

"""
$(TYPEDSIGNATURES)

Computes in-place the hessian with respect to the last state component of the basis of the last univariate function of each feature with multi-indices `idx` at `x`.
"""
function repeated_hess_xk_basis!(out, cache, f::ExpandedFunction, x, idx::Array{Int64,2})
    # Compute the Hessian of the basis functions with respect to the last component
    Ne = size(x,1)
    Nx = f.Nx

    # @assert size(out,1) = (N, size(idx, 1)) "Wrong dimension of the output vector"
    # ∂ᵏf/∂x_{grad_dim} = ψ
    k = 2
    grad_dim = Nx
    dims = Nx

    midxj = idx[:, Nx]
    maxj = maximum(midxj)
    #   Compute the kth derivative along grad_dim
    # dkψj = zeros(Ne, maxj+1)
    vander!(cache, f.B.B, maxj, k, x)
    Nψreduced = size(idx, 1)
    @avx for l = 1:Nψreduced
        for k=1:Ne
            out[k,l] = cache[k, midxj[l] + 1]
        end
    end

    return out#dkψj[:, midxj .+ 1]
end

repeated_hess_xk_basis!(out, cache, f::ExpandedFunction, x) = repeated_hess_xk_basis!(out, cache, f, x, f.idx)

repeated_hess_xk_basis(f::ExpandedFunction, x, idx::Array{Int64,2}) =
    repeated_hess_xk_basis!(zeros(size(x,1),size(idx,1)), zeros(size(x,1), maximum(idx[:,f.Nx])+1), f, x, idx)

"""
$(TYPEDSIGNATURES)

Computes the hessian (with respect to the last state component) of the basis of the last univariate function of each feature with multi-indices `idx` at `x`.
"""
repeated_hess_xk_basis(f::ExpandedFunction, x) = repeated_hess_xk_basis(f, x, f.idx)

## Compute ∂_i (∂_k f(x_{1:k}))

# function active_dimik(idx::Array{Int64,2})
#     dim = Int64[]
#     Nx = size(idx,2)
#     boolik = sum(idx .* idx[:,end]; dims = 1)[1,:] .> 0
#     return (1:Nx)[boolik]
# end
#
# active_dimik(f::ExpandedFunction) = active_dimik(f.idx)

# function grad_x_grad_xd(f::ExpandedFunction, X::Array{Float64,2}, idx::Array{Int64,2})

function grad_x_grad_xd(f::ExpandedFunction, X, idx::Array{Int64,2})
    NxX, Ne = size(X)
    m = f.m
    Nx = f.Nx
    Nψ = f.Nψ
    @assert NxX == Nx "Wrong dimension of the input"
    dxdxkψ = zeros(Ne, Nx)

    if f.dim == Int64[]
        return dxdxkψ
    end

    if f.Nx == f.dim[end]
        dxdxkψ_basis = zeros(Ne, Nψ)
        # dxdxkψ_basis = zeros(Ne, Nψ, Nx)
        @inbounds for i ∈ f.dim[f.dim .< f.Nx]
            # Reduce further the computation, we have a non-zero output only if
            # there is a feature such that idx[:,i]*idx[:,Nx]>0
            if any([line[i]*line[f.Nx] for line in eachslice(idx; dims = 1)] .> 0) || iszerofeatureactive(f.B.B)
                fill!(dxdxkψ_basis, 0.0)
                grad_xk_basis!(dxdxkψ_basis, f, X, 1, [i;Nx], idx)
                dxidxkψ = view(dxdxkψ,:,i)
                mul!(dxidxkψ, dxdxkψ_basis, f.coeff)
            end
        end

        d2xkψ = view(dxdxkψ,:,f.Nx)
        fill!(dxdxkψ_basis, 0.0)
        grad_xk_basis!(dxdxkψ_basis, f, X, 2, f.Nx)
        mul!(d2xkψ, dxdxkψ_basis, f.coeff)
        # dxdxkψ = zeros(Ne, Nx)
        # @tensor dxdxkψ[a,b] = dxdxkψ_basis[a,c,b] * f.coeff[c]
    end

    return dxdxkψ
end

"""
$(TYPEDSIGNATURES)

Computes ∂_i (∂_k f(x_{1:k}))
"""
grad_x_grad_xd(f::ExpandedFunction, X) = grad_x_grad_xd(f, X, f.idx)

# This version outputs an object of dimension (Ne, f.dim)
"""
$(TYPEDSIGNATURES)

Computes in-place ∂_i (∂_k f(x_{1:k})). In this routine, gradients are only computed with respect to the active dimensions of `f`.
"""
function reduced_grad_x_grad_xd!(dxdxkψ, f::ExpandedFunction, X, idx::Array{Int64,2})
    NxX, Ne = size(X)
    m = f.m
    Nx = f.Nx
    Nψ = f.Nψ

    dim = f.dim
    dimoff = dim[ dim .< Nx]
    @assert NxX == Nx "Wrong dimension of the input"
    @assert size(dxdxkψ) == (Ne, length(f.dim))

    if f.dim == Int64[]
        fill!(dxdxkψ, 0.0)
        return dxdxkψ
    end

    if f.Nx == f.dim[end]
        dxdxkψ_basis = zeros(Ne, Nψ)
        # dxdxkψ_basis = zeros(Ne, Nψ, Nx)
        @inbounds for i=1:length(dimoff)
            # Reduce further the computation, we have a non-zero output only if
            # there is a feature such that idx[:,i]*idx[:,Nx]>0
            if any([line[dim[i]]*line[f.Nx] for line in eachslice(idx; dims = 1)] .> 0) || iszerofeatureactive(f.B.B)
                fill!(dxdxkψ_basis, 0.0)
                grad_xk_basis!(dxdxkψ_basis, f, X, 1, [dim[i]; Nx], idx)
                dxidxkψ = view(dxdxkψ,:,i)
                mul!(dxidxkψ, dxdxkψ_basis, f.coeff)
            end
        end

        d2xkψ = view(dxdxkψ,:,length(f.dim))
        fill!(dxdxkψ_basis, 0.0)
        grad_xk_basis!(dxdxkψ_basis, f, X, 2, f.Nx)
        mul!(d2xkψ, dxdxkψ_basis, f.coeff)
        # dxdxkψ = zeros(Ne, Nx)
        # @tensor dxdxkψ[a,b] = dxdxkψ_basis[a,c,b] * f.coeff[c]
    end
    return dxdxkψ
end

reduced_grad_x_grad_xd!(dxdxkψ, f::ExpandedFunction, X) = reduced_grad_x_grad_xd!(dxdxkψ, f, X, f.idx)

# This version outputs an object of dimension (Ne, f.dim)
"""
$(TYPEDSIGNATURES)

Computes ∂_i (∂_k f(x_{1:k})). In this routine, gradients are only computed with respect to the active dimensions of `f`.
"""
reduced_grad_x_grad_xd(f::ExpandedFunction, X) = reduced_grad_x_grad_xd!(zeros(size(X,2), length(f.dim)), f, X, f.idx)

# Compute ∂_i ∂_j (∂_k f(x_{1:k}))

# function hess_x_grad_xd(f::ExpandedFunction, X::Array{Float64,2}, idx::Array{Int64,2})
# This version outputs an object of dimension (Ne, f.dim)
"""
$(TYPEDSIGNATURES)

Computes ∂_i ∂_j (∂_k f(x_{1:k}))
"""
function hess_x_grad_xd(f::ExpandedFunction, X, idx::Array{Int64,2})
    NxX, Ne = size(X)
    m = f.m
    Nx = f.Nx
    Nψ = f.Nψ
    @assert NxX == Nx "Wrong dimension of the input"
    # dxidxjdxkψ_basis = zeros(Ne, Nψ, Nx, Nx)
    d2xdxkψ = zeros(Ne, Nx, Nx)

    d2xdxkψ_basis = zeros(Ne, Nψ)

    # Store the derivative of the basis with respect to the last component
    dxkψ  = repeated_grad_xk_basis(f, X[Nx,:])
    d2xkψ = repeated_hess_xk_basis(f, X[Nx,:])

    @inbounds for i ∈ f.dim
        for j ∈ f.dim[f.dim .>= i]
            # Reduce further the computation, we have a non-zero output only if
            # there is a feature such that idx[:,i]*idx[:,j]*idx[:,Nx]>0 or if the first feature is active
            if any([line[i]*line[j]*line[f.Nx] for line in eachslice(f.idx; dims = 1)] .> 0) || iszerofeatureactive(f.B.B)
                fill!(d2xdxkψ_basis, 0.0)
                dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                #  Case i = j = k
                if i==Nx && j==Nx
                    grad_xk_basis!(d2xdxkψ_basis, f, X, 3, Nx)
                    dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                    mul!(dxidxjdxkψ, d2xdxkψ_basis, f.coeff)
                elseif i==j && j!=Nx
                    grad_xk_basis!(d2xdxkψ_basis, f, X, 2, i, f.dim[f.dim .< f.Nx], idx)
                    d2xdxkψ_basis .*= dxkψ
                    dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                    mul!(dxidxjdxkψ, d2xdxkψ_basis, f.coeff)
                elseif i!=Nx && j==Nx #(use symmetry as well)
                    grad_xk_basis!(d2xdxkψ_basis, f, X, 1, i, f.dim[f.dim .< f.Nx], idx)
                    d2xdxkψ_basis .*= d2xkψ
                    dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                    dxjdxidxkψ = view(d2xdxkψ,:,j,i)
                    mul!(dxidxjdxkψ, d2xdxkψ_basis, f.coeff)
                    dxjdxidxkψ .= dxidxjdxkψ
                else # the rest of the cases (use symmetry as well)
                    grad_xk_basis!(d2xdxkψ_basis, f, X, 1, [i;j], f.dim[f.dim .< f.Nx], idx)
                    d2xdxkψ_basis .*= dxkψ
                    dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                    dxjdxidxkψ = view(d2xdxkψ,:,j,i)
                    mul!(dxidxjdxkψ, d2xdxkψ_basis, f.coeff)
                    dxjdxidxkψ .= dxidxjdxkψ
                end
            end
        end
    end

    # dxidxjdxkψ = zeros(Ne, Nx, Nx)
    # @tensor dxidxjdxkψ[a,b,c] = dxidxjdxkψ_basis[a,d,b, c] * f.coeff[d]

    return d2xdxkψ
end

# hess_x_grad_xd(f::ExpandedFunction, X::Array{Float64,2}) = hess_x_grad_xd(f, X, f.idx)
"""
$(TYPEDSIGNATURES)

Computes ∂_i ∂_j (∂_k f(x_{1:k}))
"""
hess_x_grad_xd(f::ExpandedFunction, X) = hess_x_grad_xd(f, X, f.idx)

# function hess_x_grad_xd(f::ExpandedFunction, X::Array{Float64,2}, idx::Array{Int64,2})
"""
$(TYPEDSIGNATURES)

Computes in-place ∂_i ∂_j (∂_k f(x_{1:k})). In this routine, gradients are only computed with respect to the active dimensions of `f`.
"""
function reduced_hess_x_grad_xd!(d2xdxkψ, f::ExpandedFunction, X, idx::Array{Int64,2})
    NxX, Ne = size(X)
    m = f.m
    Nx = f.Nx
    Nψ = f.Nψ
    dim = f.dim
    @assert NxX == Nx "Wrong dimension of the input"
    # dxidxjdxkψ_basis = zeros(Ne, Nψ, Nx, Nx)
    @assert size(d2xdxkψ) == (Ne, length(dim), length(dim))

    d2xdxkψ_basis = zeros(Ne, Nψ)

    # Store the derivative of the basis with respect to the last component
    dxkψ  = repeated_grad_xk_basis(f, view(X, Nx,:))
    d2xkψ = repeated_hess_xk_basis(f, view(X, Nx,:))

    @inbounds for i=1:length(dim)
        for j = i:length(dim)
            # Reduce further the computation, we have a non-zero output only if
            # there is a feature such that idx[:,i]*idx[:,j]*idx[:,Nx]>0
            if any([line[dim[i]]*line[dim[j]]*line[f.Nx] for line in eachslice(f.idx; dims = 1)] .> 0) || iszerofeatureactive(f.B.B)
                fill!(d2xdxkψ_basis, 0.0)
                dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                #  Case i = j = k
                if dim[i]==Nx && dim[j]==Nx
                    grad_xk_basis!(d2xdxkψ_basis, f, X, 3, Nx)
                    dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                    mul!(dxidxjdxkψ, d2xdxkψ_basis, f.coeff)
                elseif dim[i]==dim[j] && dim[j]!=Nx
                    grad_xk_basis!(d2xdxkψ_basis, f, X, 2, dim[i], f.dim[f.dim .< f.Nx], idx)
                    d2xdxkψ_basis .*= dxkψ
                    dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                    mul!(dxidxjdxkψ, d2xdxkψ_basis, f.coeff)
                elseif dim[i]!=Nx && dim[j]==Nx #(use symmetry as well)
                    grad_xk_basis!(d2xdxkψ_basis, f, X, 1, dim[i], f.dim[f.dim .< f.Nx], idx)
                    d2xdxkψ_basis .*= d2xkψ
                    dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                    dxjdxidxkψ = view(d2xdxkψ,:,j,i)
                    mul!(dxidxjdxkψ, d2xdxkψ_basis, f.coeff)
                    dxjdxidxkψ .= dxidxjdxkψ
                else # the rest of the cases (use symmetry as well)
                    grad_xk_basis!(d2xdxkψ_basis, f, X, 1, [dim[i]; dim[j]], f.dim[f.dim .< f.Nx], idx)
                    d2xdxkψ_basis .*= dxkψ
                    dxidxjdxkψ = view(d2xdxkψ,:,i,j)
                    dxjdxidxkψ = view(d2xdxkψ,:,j,i)
                    mul!(dxidxjdxkψ, d2xdxkψ_basis, f.coeff)
                    dxjdxidxkψ .= dxidxjdxkψ
                end
            end
        end
    end

    return d2xdxkψ
end

"""
$(TYPEDSIGNATURES)

Computes ∂_i ∂_j (∂_k f(x_{1:k})). In this routine, gradients are only computed with respect to the active dimensions of `f`.
"""
reduced_hess_x_grad_xd(f::ExpandedFunction, X) = reduced_hess_x_grad_xd!(zeros(size(X, 2), length(f.dim), length(f.dim)), f, X, f.idx)

# Derivative with respect to the some coefficients
"""
$(TYPEDSIGNATURES)

Computes the gradient f(x_{1:k})) with respect to the components `coeff_idx` of `getcoeff(f)`.
"""
function grad_coeff(f::ExpandedFunction, X::Array{Float64,2}, coeff_idx::Array{Int64, 1})
    Nψ = f.Nψ
    # Verify that all the index
    @assert all([0 <=idx <= Nψ for idx in coeff_idx]) "idx is larger than Nψ"
    return evaluate_basis(f, X, f.idx[coeff_idx,:])
end

# Derivative with respect to the coefficients
"""
$(TYPEDSIGNATURES)

Computes the gradient f(x_{1:k})) with respect to the entire vector of coefficient `getcoeff(f)`.
"""
function grad_coeff(f::ExpandedFunction, X::Array{Float64,2})
    return evaluate_basis(f, X)
end

# Hessian with respect to the some coefficients
"""
$(TYPEDSIGNATURES)

Computes the hessian f(x_{1:k})) with respect to the components `coeff_idx` of `getcoeff(f)`.
"""
function hess_coeff(f::ExpandedFunction, X::Array{Float64,2}, coeff_idx::Array{Int64, 1})
    # Verify that all the index
    Nψ = f.Nψ
    Ne = size(X,2)
    Nψreduced = size(coeff_idx,1)
    @assert all([0 <=idx <= Nψ for idx in coeff_idx]) "idx is larger than Nψ"
    return zeros(Ne, Nψreduced, Nψreduced)
end

# Hessian with respect to the coefficients
"""
$(TYPEDSIGNATURES)

Computes the hessian f(x_{1:k})) with respect to the entire vector of coefficient `getcoeff(f)`.
"""
function hess_coeff(f::ExpandedFunction, X::Array{Float64,2})
    Nψ = f.Nψ
    return zeros(size(X,2), Nψ, Nψ)
end

# Gradient with respect to the coefficients of the gradient with respect to xd
"""
$(TYPEDSIGNATURES)

Computes ∂_c ∂_xₖ f.
"""
function grad_coeff_grad_xd(f::ExpandedFunction, X::Array{Float64,2})
    return grad_xk_basis(f, X, 1, f.Nx)
end
# function grad_xk_basis(f::ExpandedFunction{m, Nψ, Nx}, grad_dim::Union{Int64, Array{Int64,1}}, k::Int64, X::Array{Float64,2}, idx::Array{Int64,2}) where {m, Nψ, Nx, Ne}

"""
$(TYPEDSIGNATURES)

Computes ∂_c ∂_xₖ f for the components `coeff_idx` of the coefficient vector of `f`.
"""
function grad_coeff_grad_xd(f::ExpandedFunction, X::Array{Float64,2}, coeff_idx::Array{Int64,1})
    return grad_xk_basis(f, X, 1, f.Nx, f.idx[coeff_idx,:])
end

# Hessian with respect to the coefficients of the gradient with respect to xd
"""
$(TYPEDSIGNATURES)

Computes ∂^2_c ∂_xₖ f.
"""
function hess_coeff_grad_xd(f::ExpandedFunction, X::Array{Float64,2})
    Ne = size(X,2)
    Nψ = f.Nψ
    return zeros(Ne, Nψ, Nψ)
end

"""
$(TYPEDSIGNATURES)

Computes ∂_c ∂_xₖ f for the components `coeff_idx` of the coefficient vector of `f`.
"""
function hess_coeff_grad_xd(f::ExpandedFunction, X::Array{Float64,2}, coeff_idx::Array{Int64,1})
    # Verify that all the index
    Ne = size(X,2)
    Nψ = f.Nψ
    Nψreduced = size(coeff_idx,1)
    @assert all([0 <=idx <= Nψ for idx in coeff_idx]) "idx is larger than Nψ"
    return zeros(Ne, Nψreduced, Nψreduced)
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule bottom_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("bottom")
JLLWrappers.@generate_main_file("bottom", UUID("77dc8835-e407-5831-b71c-802d10921d58"))
end  # module bottom_jll


text:
"""
    locw(Xtrain, Ytrain, X; listnn, listw = nothing, fun, verbose = false, kwargs...)
Compute predictions for a given kNN model.
* `Xtrain` : Training X-data.
* `Ytrain` : Training Y-data.
* `X` : X-data (m observations) to predict.
* `listnn` : List of m vectors of indexes.
* `listw` : List of m vectors of weights.
* `fun` : Function computing the model on the m neighborhoods.
* `verbose` : If true, fitting information are printed.
* `kwargs` : Keywords arguments to pass in function fun.

Each component i of `listnn` and `listw` contains the indexes and weights, respectively,
of the nearest neighbors of x_i in Xtrain. The sizes of the neighborhood for i = 1,...,m can be different.

All the arguments in kwargs must have length = 1 (not collections).
"""
function locw(Xtrain, Ytrain, X ;
    listnn, listw = nothing, fun, verbose = false, kwargs...)
    m = size(X, 1)
    q = size(Ytrain, 2)
    #pred = zeros(m, q)
    pred = similar(Ytrain, m, q)
    @inbounds for i = 1:m
        verbose ? print(i, " ") : nothing
        s = listnn[i]
        length(s) == 1 ? s = (s:s) : nothing
        zYtrain = Ytrain[s, :]
        # For discrimination,
        # case where all the neighbors are of same class
        if q == 1 && length(unique(zYtrain)) == 1
            pred[i, :] .= zYtrain[1]
        # End
        else
            if isnothing(listw)
                fm = fun(Xtrain[s, :],  zYtrain; kwargs...)
            else
                fm = fun(Xtrain[s, :], zYtrain, listw[i] ; kwargs...)
            end
            pred[i, :] = predict(fm, X[i:i, :]).pred
        end
    end
    verbose ? println() : nothing
    (pred = pred,)
end

"""
    locwlv(Xtrain, Ytrain, X ; listnn, listw = nothing, fun, nlv, verbose = false, kwargs...)
Same as [`locw`](@ref) but specific (and much faster) to LV-based models (e.g. PLSR).
* `nlv` : Nb. or collection of nb. of latent variables (LVs).

Argument pars must not contain nlv.
"""
function locwlv(Xtrain, Ytrain, X;
    listnn, listw = nothing, fun, nlv, verbose = true, kwargs...)
    m = size(X, 1)
    q = size(Ytrain, 2)
    nlv = max(minimum(nlv), 0):maximum(nlv)
    le_nlv = length(nlv)
    #zpred = zeros(m, q, le_nlv)
    #zpred = Array{Float64}(undef, m, q, le_nlv)
    zpred = similar(Ytrain, m, q, le_nlv)
    @inbounds for i = 1:m
        verbose ? print(i, " ") : nothing
        s = listnn[i]
        length(s) == 1 ? s = (s:s) : nothing
        zYtrain = Ytrain[s, :]
        # For discrimination,
        # case where all the neighbors are of same class
        if q == 1 && length(unique(zYtrain)) == 1
            for a = 1:le_nlv
                zpred[i, :, a] .= zYtrain[1]
            end
        # End
        else
            if isnothing(listw)
                fm = fun(Xtrain[s, :],  zYtrain ; nlv = maximum(nlv), kwargs...)
            else
                fm = fun(Xtrain[s, :], zYtrain, listw[i] ; nlv = maximum(nlv), kwargs...)
            end
            @inbounds for a = 1:le_nlv
                zpred[i, :, a] = predict(fm, X[i:i, :] ; nlv = nlv[a]).pred
            end
        end
    end
    verbose ? println() : nothing
    pred = list(le_nlv, Matrix{Float64})
    for a = 1:le_nlv
        pred[a] = zpred[:, :, a]
    end
    le_nlv == 1 ? pred = pred[1] : nothing
    (pred = pred, )
end


text:
### Rendering function for the Mars Exploration Problem

using Cairo, Colors

const WALLE_IMG = read_from_png("wall_e.png")
const PLANT_IMG = read_from_png("plant.png")
const EVE_IMG = read_from_png("eve.png")

mutable struct GridDraw
    x0::Int64
    y0::Int64
    cell_width::Float64
    cell_height::Float64
    n_cells::Int64
    line_width::Float64
    function GridDraw(;x0=50, y0=50, cell_width=50., cell_height=50., n_cells=10, line_width=1.)
        return new(x0, y0, cell_width, cell_height, n_cells, line_width)
    end
end

function render_line{T<:Real}(
    ctx        :: CairoContext,
    pts        :: Matrix{T}, # 2×n
    color      :: RGB{Float64},
    line_width :: Real = 1.0,
    line_cap   :: Integer=Cairo.CAIRO_LINE_CAP_ROUND, # CAIRO_LINE_CAP_BUTT, CAIRO_LINE_CAP_ROUND, CAIRO_LINE_CAP_SQUARE
    )

    line_width = user_to_device_distance!(ctx, [line_width,0])[1]

    save(ctx)
    set_source_rgb(ctx, color.:r, color.:g, color.:b)
    set_line_width(ctx,line_width)
    set_line_cap(ctx, line_cap)

    move_to(ctx, pts[1,1], pts[2,1])
    for i = 2 : size(pts,2)
        line_to(ctx, pts[1,i], pts[2,i])
    end
    stroke(ctx)
    restore(ctx)
end

function set_context_and_surface(canvas_width::Int=600, canvas_height::Int=100)
    c = CairoRGBSurface(canvas_width, canvas_height)
    ctx = creategc(c)
    ctx, c
end

function render_background(ctx::CairoContext, color::RGB{Float64})
    save(ctx)

    set_source_rgb(ctx, color.:r, color.:g, color.:b)

    rectangle(ctx, 0., 0., ctx.surface.width, ctx.surface.height) # background
    fill(ctx)
    restore(ctx)
end

function render_grid(ctx::CairoContext,
                     grid::GridDraw = GridDraw())

    #draw upper and lower line
    save(ctx)
    upper_pts = [grid.x0 grid.y0+grid.cell_height/2;
                 grid.x0 + grid.n_cells*grid.cell_width grid.y0 + grid.cell_height/2]'
    render_line(ctx, upper_pts, RGB(0.,0.,0.), grid.line_width)

    lower_pts = [grid.x0  grid.y0 - grid.cell_height/2;
                 grid.x0 + grid.n_cells*grid.cell_width grid.y0 - grid.cell_height/2]'
    render_line(ctx, lower_pts, RGB(0.,0.,0.), grid.line_width)

    # draw vertical separation
    for i = 0:grid.n_cells
        pts = [grid.x0+ i*grid.cell_width grid.y0 - grid.cell_height/2;
               grid.x0 + i*grid.cell_width grid.y0 + grid.cell_height/2]'
        render_line(ctx, pts, RGB(0.,0.,0.), grid.line_width)
    end
    restore(ctx)
end

function render_rewards(ctx::CairoContext, grid::GridDraw = GridDraw())
    render_agent(ctx, 10, grid, PLANT_IMG)
    render_agent(ctx, 1, grid, EVE_IMG)
end

function render_agent(ctx::CairoContext, cell::Int64,
                      grid::GridDraw = GridDraw(),
                      img::CairoSurface = WALLE_IMG)
    save(ctx)
    w = img.width
    h = img.height
    translate(ctx, grid.x0 +(cell - 0.5)*grid.cell_width, grid.y0)
    scale(ctx, grid.cell_width/w, grid.cell_height/h)
    translate(ctx, -0.5*w, -0.5*h)
    set_source_surface(ctx, img, 0, 0)
    paint(ctx)
    restore(ctx)
end

function paint_cell(ctx::CairoContext, cell::Int64, color::RGB{Float64},
                    grid::GridDraw = GridDraw())
    save(ctx)
    set_source_rgb(ctx, color.:r, color.:g, color.:b)
    # bottom left corner of the cell
    x = grid.x0+ (cell-1)*grid.cell_width
    y = grid.y0 - grid.cell_height/2
    rectangle(ctx, x, y, grid.cell_width, grid.cell_height) # background
    fill(ctx)
    restore(ctx)
end

function render_action_value(ctx::CairoContext, action::Symbol, value::Float64, cell::Int64,
                       grid::GridDraw = GridDraw())
    save(ctx)
    select_font_face(ctx, "Sans", Cairo.FONT_SLANT_NORMAL,
                 Cairo.FONT_WEIGHT_NORMAL)
    font_size = 20.
    set_font_size(ctx, font_size)
    move_to(ctx, grid.x0 +(cell - 0.5)*grid.cell_width - 10, grid.y0)
    if action == :left
        str = @sprintf("← ")
    elseif action == :right
        str = @sprintf("→ ")
    end
    show_text(ctx, str)
    move_to(ctx, grid.x0 +(cell - 0.5)*grid.cell_width - 25, grid.y0 + font_size)
    val_str = @sprintf("%2.2f", value)
    show_text(ctx, val_str)
    restore(ctx)
end

function render_state(s::Int64)
    ctx, c = set_context_and_surface()
    bg_col = RGB(1., 1., 1.)
    render_background(ctx, bg_col)
    render_grid(ctx)
    render_rewards(ctx)
    render_agent(ctx, s)
    c
end

function colorval(policy::ValuePolicy, val, cmap)
    N = length(cmap)
    max_val = max(policy.mdp.r_left, policy.mdp.r_right)
    val_res = max_val/N
    val_ind = Int(floor(val/val_res)) + 1
    return cmap[val_ind]
end

function render_cell(cell::Int64)
    ctx, c = set_context_and_surface()
    bg_col = RGB(1., 1., 1.)
    render_background(ctx, bg_col)
    paint_cell(ctx, cell, RGB(1.,0.,0.))
    render_grid(ctx)
    render_action_value(ctx, :left, 1., cell)
    c
end

function render_policy(policy::ValuePolicy)
    ctx, c = set_context_and_surface()
    bg_col = RGB(1., 1., 1.)
    render_background(ctx, bg_col)

    cmap = colormap("Greens", 1000)
    for (i, s) in enumerate(ordered_states(policy.mdp))
        val = maximum(policy.value_table[i, :])
        a = action(policy, s)
        color = colorval(policy, val, cmap)
        paint_cell(ctx, s, color)
        render_action_value(ctx, a, val, s)
    end
    render_grid(ctx)
    c
end