abstract type LinearModel end

"""
    ESNtrain(esn::AbstractReservoirComputer, beta::Float64[, train_data])

Return the trained output layer using Ridge Regression.
"""
function ESNtrain(esn::AbstractReservoirComputer, beta::Float64; train_data::AbstractArray{Float64} = esn.train_data)

    states_new = nla(esn.nla_type, esn.states)
    W_out = (train_data*states_new')*inv(add_reg(states_new*states_new', beta))

    return W_out
end

function add_reg(X::AbstractArray{Float64}, beta::Float64)
    n = size(X, 1)
    for i=1:n
        X[i,i] += beta
    end
    return X
end

#MLJ Ridge
"""
    Ridge(lambda, solver::MLJLinearModels.Solver)

Return a LinearModel object for the training of the model using Ridge Regression with a MLJLinearModels method.
"""
struct Ridge{T<: AbstractFloat} <: LinearModel
    lambda::T
    solver::MLJLinearModels.Solver
end
#Ridge(lambda, solver) = Ridge{Float64}(lambda, solver)
"""
    ESNtrain(lm::LinearModel, esn::AbstractReservoirComputer[, train_data])

Return the trained output layer using an MLJLinearModels method built into a LinearSolver struct.
"""
ESNtrain(ridge::Ridge{T}, esn::AbstractReservoirComputer; train_data::AbstractArray{Float64} = esn.train_data) where T<: AbstractFloat = _ridge(esn, ridge; train_data = esn.train_data)

function _ridge(esn::AbstractReservoirComputer, ridge::Ridge; train_data::AbstractArray{Float64} = esn.train_data)

    states_new = nla(esn.nla_type, esn.states)
    W_out = zeros(Float64, size(train_data, 1), size(states_new, 1))
    for i=1:size(train_data, 1)
        r = RidgeRegression(lambda = ridge.lambda, fit_intercept = false)
        W_out[i,:] = MLJLinearModels.fit(r, states_new', train_data[i,:], solver = ridge.solver)
    end

    return W_out
end

#MLJ Lasso
"""
    Lasso(lambda, solver::MLJLinearModels.Solver)

Return a LinearModel object for the training of the model using Lasso with a MLJLinearModels method.
"""
struct Lasso{T<: AbstractFloat} <: LinearModel
    lambda::T
    solver::MLJLinearModels.Solver
end
#Lasso(lambda, solver) = Lasso{Float64}(lambda, solver)
ESNtrain(lasso::Lasso{T}, esn::AbstractReservoirComputer; train_data::AbstractArray{Float64} = esn.train_data) where T<: AbstractFloat = _lasso(esn, lasso; train_data = esn.train_data)

function _lasso(esn::AbstractReservoirComputer, lasso::Lasso; train_data::AbstractArray{Float64} = esn.train_data)

    states_new = nla(esn.nla_type, esn.states)
    W_out = zeros(Float64, size(train_data, 1), size(states_new, 1))
    for i=1:size(train_data, 1)
        l = LassoRegression(lambda = lasso.lambda, fit_intercept = false)
        W_out[i,:] = MLJLinearModels.fit(l, states_new', train_data[i,:], solver = lasso.solver)
    end

    return W_out
end

#MLJ ElastNet
"""
    ElastNet(lambda, gamma, solver::MLJLinearModels.Solver)

Return a LinearModel object for the training of the model using Elastic Net with a MLJLinearModels method.
"""
struct ElastNet{T<: AbstractFloat} <: LinearModel
    lambda::T
    gamma::T
    solver::MLJLinearModels.Solver
end
#ElastNet(lambda, gamma, solver) = ElastNet{Float64}(lambda, gamma, solver)
ESNtrain(elastnet::ElastNet{T}, esn::AbstractReservoirComputer; train_data::AbstractArray{Float64} = esn.train_data) where T<: AbstractFloat = _elastnet(esn, elastnet; train_data = esn.train_data)

function _elastnet(esn::AbstractReservoirComputer, elastnet::ElastNet; train_data::AbstractArray{Float64} = esn.train_data)

    states_new = nla(esn.nla_type, esn.states)
    W_out = zeros(Float64, size(train_data, 1), size(states_new, 1))
    for i=1:size(train_data, 1)
        en = ElasticNetRegression(lambda = elastnet.lambda, gamma = elastnet.gamma, fit_intercept = false)
        W_out[i,:] = MLJLinearModels.fit(en, states_new', train_data[i,:], solver = elastnet.solver)
    end

    return W_out
end

#MLJ Huber -> RobustHuber avoids conflict
"""
    ElastNet(delta, lambda, gamma, solver::MLJLinearModels.Solver)

Return a LinearModel object for the training of the model using the Huber function with a MLJLinearModels method.
"""
struct RobustHuber{T<: AbstractFloat} <: LinearModel
    delta::T
    lambda::T
    gamma::T
    solver::MLJLinearModels.Solver
end
#RobustHuber(delta, lambda, gamma, solver) = RobustHuber{Float64}(delta, lambda, gamma, solver)
ESNtrain(huber::RobustHuber{T}, esn::AbstractReservoirComputer; train_data::AbstractArray{Float64} = esn.train_data) where T<: AbstractFloat = _huber(esn, huber; train_data = esn.train_data)

function _huber(esn::AbstractReservoirComputer, huber::RobustHuber; train_data::AbstractArray{Float64} = esn.train_data)

    states_new = nla(esn.nla_type, esn.states)
    W_out = zeros(Float64, size(train_data, 1), size(states_new, 1))
    for i=1:size(train_data, 1)
        h = HuberRegression(delta = huber.delta,
            lambda = huber.lambda,
            gamma = huber.gamma,
            fit_intercept = false)
        W_out[i,:] = MLJLinearModels.fit(h, states_new', train_data[i,:], solver = huber.solver)
    end

    return W_out
end

"""
    HESNtrain(ridge::Ridge{T}, hesn::AbstractReservoirComputer; train_data::AbstractArray{Float64} = hesn.train_data)

Return the trained HESN output layer using an MLJLinearModels method built into a LinearSolver struct.
"""
HESNtrain(ridge::Ridge{T}, hesn::AbstractReservoirComputer; train_data::AbstractArray{Float64} = hesn.train_data) where T<: AbstractFloat = _PIridge(hesn, ridge; train_data = hesn.train_data)

function _PIridge(hesn::AbstractReservoirComputer, ridge::Ridge; train_data::AbstractArray{Float64} = hesn.train_data)

    states_new = nla(hesn.nla_type, hesn.states)
    states_new = vcat(states_new, hesn.physics_model_data[:, 2:end])
    W_out = zeros(Float64, size(train_data, 1), size(states_new, 1))
    for i=1:size(train_data, 1)
        r = RidgeRegression(lambda = ridge.lambda, fit_intercept = false)
        W_out[i,:] = MLJLinearModels.fit(r, states_new', train_data[i,:], solver = ridge.solver)
    end

    return W_out
end


"""
    icrs2g(dr, dd)

Transformation from ICRS to Galactic Coordinates.

### Given ###

- `dr`: ICRS right ascension (radians)
- `dd`: ICRS declination (radians)

### Returned ###

- `dl`: Galactic longitude (radians)
- `db`: Galactic latitude (radians)

### Notes ###

1. The IAU 1958 system of Galactic coordinates was defined with
   respect to the now obsolete reference system FK4 B1950.0.  When
   interpreting the system in a modern context, several factors have
   to be taken into account:

   . The inclusion in FK4 positions of the E-terms of aberration.

   . The distortion of the FK4 proper motion system by differential
     Galactic rotation.

   . The use of the B1950.0 equinox rather than the now-standard
     J2000.0.

   . The frame bias between ICRS and the J2000.0 mean place system.

   The Hipparcos Catalogue (Perryman & ESA 1997) provides a rotation
   matrix that transforms directly between ICRS and Galactic
   coordinates with the above factors taken into account.  The
   matrix is derived from three angles, namely the ICRS coordinates
   of the Galactic pole and the longitude of the ascending node of
   the galactic equator on the ICRS equator.  They are given in
   degrees to five decimal places and for canonical purposes are
   regarded as exact.  In the Hipparcos Catalogue the matrix
   elements are given to 10 decimal places (about 20 microarcsec).
   In the present ERFA function the matrix elements have been
   recomputed from the canonical three angles and are given to 30
   decimal places.

2. The inverse transformation is performed by the function eraG2icrs.

### Called ###

- `eraAnp`: normalize angle into range 0 to 2pi
- `eraAnpm`: normalize angle into range +/- pi
- `eraS2c`: spherical coordinates to unit vector
- `eraRxp`: product of r-matrix and p-vector
- `eraC2s`: p-vector to spherical

### Reference ###

- Perryman M.A.C. & ESA, 1997, ESA SP-1200, The Hipparcos and Tycho
    catalogues.  Astrometric and photometric star catalogues
    derived from the ESA Hipparcos Space Astrometry Mission.  ESA
    Publications Division, Noordwijk, Netherlands.

"""
function icrs2g(a, b)
    r1 = Ref(0.0)
    r2 = Ref(0.0)
    ccall((:eraIcrs2g, liberfa), Cvoid,
            (Cdouble, Cdouble, Ref{Cdouble}, Ref{Cdouble}),
            a, b, r1, r2)
    r1[], r2[]
end


module TensorNetworkAD
using Zygote, BackwardsLinalg
using OMEinsum

export trg, num_grad
export ctmrg
export optimiseipeps
export hamiltonian, model_tensor, mag_tensor
export Ising, TFIsing, Heisenberg

include("hamiltonianmodels.jl")

include("trg.jl")
include("fixedpoint.jl")
include("ctmrg.jl")
include("ipeps.jl")
include("autodiff.jl")
include("variationalipeps.jl")
include("exampletensors.jl")

end # module


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

module Rotations3D

using StaticArrays
using LinearAlgebra: norm, rank, svd, Diagonal, dot
using ACE.SphericalHarmonics: index_y
using Combinatorics: permutations

export ClebschGordan, Rot3DCoeffs, ri_basis, rpi_basis

"""
`ClebschGordan: ` storing precomputed Clebsch-Gordan coefficients; see
`?clebschgordan` for the convention that is use.
"""
struct ClebschGordan{T}
	vals::Dict{Tuple{Int, Int, Int, Int, Int, Int}, T}
end

"""
`ERot3DCoeffs: ` storing recursively precomputed matrix-valued coefficients for a
rotation-equivariant basis. L = 0 equiv invariant, L = 1 equiv equivariant
"""
struct Rot3DCoeffs{T,L}
   vals::Vector{Dict}
   cg::ClebschGordan{T}
end

# -----------------------------------
# iterating over an m collection
# -----------------------------------

_mvec(::CartesianIndex{0}) = SVector(Int(0))

_mvec(mpre::CartesianIndex) = SVector(Tuple(mpre)..., - sum(Tuple(mpre)))

struct MRange{N, T2}
   ll::SVector{N, Int}
   cartrg::T2
   L::Int
end

Base.length(mr::MRange) = sum(_->1, _mrange(mr.ll,mr.L))

_mrange(ll, L) = MRange(ll, Iterators.Stateful(
               filter((x) -> abs(sum(x))<= L, Tuple.(CartesianIndices(ntuple(i -> -ll[i]:ll[i], length(ll)))))
                     ),L)

function Base.iterate(mr::MRange, args...)
	while true
   		if isempty(mr.cartrg)
      		return nothing
   		end
   		mpre = popfirst!(mr.cartrg)
   		return SVector(mpre), nothing
	end
end

# ----------------------------------------------------------------------
#     ClebschGordan code
# ----------------------------------------------------------------------

cg_conditions(j1,m1, j2,m2, J,M) =
	cg_l_condition(j1, j2, J)   &&
	cg_m_condition(m1, m2, M)   &&
	(abs(m1) <= j1) && (abs(m2) <= j2) && (abs(M) <= J)

cg_l_condition(j1, j2, J) = (abs(j1-j2) <= J <= j1 + j2)

cg_m_condition(m1, m2, M) = (M == m1 + m2)

"""
`clebschgordan(j1, m1, j2, m2, J, M, T=Float64)` :

A reference implementation of Clebsch-Gordon coefficients based on

https://hal.inria.fr/hal-01851097/document
Equation (4-6)

This heavily uses BigInt and BigFloat and should therefore not be employed
for performance critical tasks, but only precomputation.

The ordering of parameters corresponds to the following convention:
```
clebschgordan(j1, m1, j2, m2, J, M) = C_{j1m1j2m2}^{JM}
```
where
```
   D_{m1k1}^{l1} D_{m2k2}^{l2}}
	=
	∑_j  C_{l1m1l2m2}^{j(m1+m2)} C_{l1k1l2k2}^{j2(k1+k2)} D_{(m1+m2)(k1+k2)}^{j}
```
"""
function clebschgordan(j1, m1, j2, m2, J, M, T=Float64)
	if !cg_conditions(j1, m1, j2, m2, J, M)
		return zero(T)
	end

   N = (2*J+1) *
       factorial(big(j1+m1)) * factorial(big(j1-m1)) *
       factorial(big(j2+m2)) * factorial(big(j2-m2)) *
       factorial(big(J+M)) * factorial(big(J-M)) /
       factorial(big( j1+j2-J)) /
       factorial(big( j1-j2+J)) /
       factorial(big(-j1+j2+J)) /
       factorial(big(j1+j2+J+1))

   G = big(0)
   # 0 ≦ k ≦ j1+j2-J
   # 0 ≤ j1-m1-k ≤ j1-j2+J   <=>   j2-J-m1 ≤ k ≤ j1-m1
   # 0 ≤ j2+m2-k ≤ -j1+j2+J  <=>   j1-J+m2 ≤ k ≤ j2+m2
   lb = (0, j2-J-m1, j1-J+m2)
   ub = (j1+j2-J, j1-m1, j2+m2)
   for k in maximum(lb):minimum(ub)
      bk = big(k)
      G += (-1)^k *
           binomial(big( j1+j2-J), big(k)) *
           binomial(big( j1-j2+J), big(j1-m1-k)) *
           binomial(big(-j1+j2+J), big(j2+m2-k))
   end

   return T(sqrt(N) * G)
end

ClebschGordan(T=Float64) =
	ClebschGordan{T}(Dict{Tuple{Int,Int,Int,Int,Int,Int}, T}())

_cg_key(j1, m1, j2, m2, J, M) = (j1, m1, j2, m2, J, M)
	# Int.((index_y(j1,m1), index_y(j2,m2), index_y(J,M)))

function (cg::ClebschGordan{T})(j1, m1, j2, m2, J, M) where {T}
	if !cg_conditions(j1,m1, j2,m2, J,M)
		return zero(T)
	end
	key = _cg_key(j1, m1, j2, m2, J, M)
	if haskey(cg.vals, key)
		return cg.vals[key]
	end
	val = clebschgordan(j1, m1, j2, m2, J, M, T)
	cg.vals[key] = val
	return val
end

# ----------------------------------------------------------------------
#     IRot3DCoeffs code: generalized cg coefficients
#
#  Note: in this section kk is a tuple of m-values, it is not
#        related to the k index in the 1-p basis (or radial basis)
# ----------------------------------------------------------------------

function get0(L::Int,T=Float64)
	if L== 0
		return T(0)
	elseif L == 1
		return @SArray zeros(Complex{T},3,3)
		#SMatrix{3, 3, Complex{T}, 9}(0, 0, 0, 0, 0, 0, 0, 0, 0)
	else
		#For L== 2 @SArray zeros(Complex{T},3,3,3)
		ErrorException("Only types for L in {0,1} implemented")
	end
end
dicttype(A::Rot3DCoeffs,N::Integer) = dicttype(A::Rot3DCoeffs,Val(N))

dicttype(A::Rot3DCoeffs{T,L},::Val{N}) where {T,L,N} =
   Dict{Tuple{SVector{N,Int}, SVector{N,Int}, SVector{N,Int}}, typeof(get0(L,T))}

Rot3DCoeffs(L,T=Float64) = Rot3DCoeffs{T,L}(Dict[], ClebschGordan(T))

function get_vals(A::Rot3DCoeffs, valN::Val{N}) where {N}
	if length(A.vals) < N
		for n = length(A.vals)+1:N
			push!(A.vals, dicttype(A,n)())
		end
	end
   	return A.vals[N]::dicttype(A,valN)
end

_key(ll::StaticVector{N}, mm::StaticVector{N}, kk::StaticVector{N}) where {N} =
      (SVector{N, Int}(ll), SVector{N, Int}(mm), SVector{N, Int}(kk))

  function (A::Rot3DCoeffs{T,L})(ll::StaticVector{N},
                              mm::StaticVector{N},
                              kk::StaticVector{N}) where {T,L, N}
     if       abs(sum(mm)) > L ||
              abs(sum(kk)) > L ||
              !all(abs.(mm) .<= ll) ||
              !all(abs.(kk) .<= ll)
        return get0(L,T)
     end
     vals = get_vals(A, Val(N))  # this should infer the type!
     key = _key(ll, mm, kk)
     if haskey(vals, key)
        val  = vals[key]
     else
        val = _compute_val(A, key...)
        vals[key] = val
     end
     return val
  end

  function (A::Rot3DCoeffs{T,0})(ll::StaticVector{1},
                              mm::StaticVector{1},
                              kk::StaticVector{1}) where {T}
     if ll[1] == mm[1] == kk[1] == 0
        return T(8 * pi^2)
     else
        return get0(0,T)
     end
  end

  function (A::Rot3DCoeffs{T,1})(ll::StaticVector{1},
                              mm::StaticVector{1},
                              kk::StaticVector{1}) where {T}
     if ll[1] == 1 && abs(mm[1]) <= 1 && abs(kk[1]) <= 1
        return  rmatrices[(mm[1],kk[1])]
     else
        return get0(1,T)
     end
  end

#function get0val(m,k,L)

#end
rmatrices = Dict(
    (-1,-1) => SMatrix{3, 3, ComplexF64, 9}(1/6, 1im/6, 0, -1im/6, 1/6, 0, 0, 0, 0),
    (-1,0) => SMatrix{3, 3, ComplexF64, 9}(0, 0, 0, 0, 0, 0, 1/(3*sqrt(2)), 1im/(3*sqrt(2)), 0),
    (-1,1) => SMatrix{3, 3, ComplexF64, 9}(-1/6, -1im/6, 0, -1im/6, 1/6, 0, 0, 0, 0),
    (0,-1) => SMatrix{3, 3, ComplexF64, 9}(0, 0, 1/(3*sqrt(2)), 0, 0, -1im/(3*sqrt(2)), 0, 0, 0),
    (0,0) => SMatrix{3, 3, ComplexF64, 9}(0, 0, 0, 0, 0, 0, 0, 0, 1/3),
    (0,1) => SMatrix{3, 3, ComplexF64, 9}(0, 0, -1/(3*sqrt(2)), 0, 0, -1im/(3*sqrt(2)), 0, 0, 0),
    (1,-1) => SMatrix{3, 3, ComplexF64, 9}(-1/6, 1im/6, 0, 1im/6, 1/6, 0, 0, 0, 0),
    (1,0) => SMatrix{3, 3, ComplexF64, 9}(0, 0, 0, 0, 0, 0, -1/(3*sqrt(2)), 1im/(3*sqrt(2)), 0),
    (1,1) => SMatrix{3, 3, ComplexF64, 9}(1/6, -1im/6, 0, 1im/6, 1/6, 0, 0, 0, 0)
    )

function (A::Rot3DCoeffs{T,L})(ll::StaticVector{1},
							 mm::StaticVector{1},
							 kk::StaticVector{1}) where {T,L}
	ErrorException("Not implemented for L = " + string(L))
end

function _compute_val(A::Rot3DCoeffs{T,L}, ll::StaticVector{N},
                                        mm::StaticVector{N},
                                        kk::StaticVector{N}) where {T,L,N}

   val = get0(L,T)
   llp = ll[1:N-2]
   mmp = mm[1:N-2]
   kkp = kk[1:N-2]
   for j = abs(ll[N-1]-ll[N]):(ll[N-1]+ll[N])
      if abs(kk[N-1]+kk[N]) > j || abs(mm[N-1]+mm[N]) > j
         continue
      end
		cgk = try
			A.cg(ll[N-1], kk[N-1], ll[N], kk[N], j, kk[N-1]+kk[N])
		catch
			@show (ll[N-1], kk[N-1], ll[N], kk[N], j, kk[N-1]+kk[N])
			get0(L,T)
		end
		cgm = A.cg(ll[N-1], mm[N-1], ll[N], mm[N], j, mm[N-1]+mm[N])
		if cgk * cgm  != 0
			val += cgk * cgm * A( SVector(llp..., j),
								       SVector(mmp..., mm[N-1]+mm[N]),
								       SVector(kkp..., kk[N-1]+kk[N]) )
		end
   end
   return val
end

function re_basis(A::Rot3DCoeffs{T,0}, ll::SVector{N, Int}; ordered=false) where {T,N}
	"""computes coefficients for a rotational 0-equinvariant (invariant) tensor basis"""
	CC = compute_Al(A, ll, Val(ordered))
	svdC = svd(CC)
	rk = rank(Diagonal(svdC.S))
	return svdC.U[:, 1:rk]'
end

function compute_Al(A::Rot3DCoeffs{T,0}, ll::SVector, ::Val{false}) where {T}
	len = length(_mrange(ll,0))
   CC = zeros(T, len, len)
   for (im, mm) in enumerate(_mrange(ll,0)), (ik, kk) in enumerate(_mrange(ll,0))
      CC[ik, im] = A(ll, mm, kk)
   end
   return CC
end

function re_basis(A::Rot3DCoeffs{T,L}, ll::SVector{N, Int}) where {T,L,N}
	"""return coefficients for a rotational 1-equinvariant (covariant) tensor basis"""
	#Mre = collect(Iterators.product(1:3, _mrange(ll,1)))
	GG = _gramian(A, ll)
   	S = svd(GG)
   	rk = rank(GG; rtol =  1e-7)
	Ure = S.U[:, 1:rk]'

	return Diagonal(1. ./ sqrt.(S.S[1:rk])) * Ure
end

imag_tol = 1e-15

function _gramian(A::Rot3DCoeffs{T,1}, ll::SVector) where {T}
	len = length(_mrange(ll,1))*3
	GG = zeros(ComplexF64, len, len) # = Gramian of rotational eqiuvariant span
	# kk = mprime
    for (i, (k,mm)) in enumerate(Iterators.product(1:3, _mrange(ll,1))), (ip, (kp,mmp)) in enumerate(Iterators.product(1:3, _mrange(ll,1)))
		for mu in _mrange(ll,1)
		  	GG[i,ip] += dot(A(ll, mm, mu)[:,k],A(ll, mmp, mu)[:,kp])
		end
	end
	@assert all(abs.(imag(GG)) .<= imag_tol)
    return real.(GG)
end

rpi_basis(A::Rot3DCoeffs, zz, nn, ll) =
			rpi_basis(A, SVector(zz...), SVector(nn...), SVector(ll...))

function rpi_basis(A::Rot3DCoeffs{T,L},
						 zz::SVector{N},
						 nn::SVector{N, Int},
						 ll::SVector{N, Int}) where {T,L,N}
	Ure = re_basis(A, ll)
	Mre = collect( _mrange(ll, L))   # rows...
	G = _gramian(zz, nn, ll, Ure, Mre)
	#print("G type :", typeof(G),"\n")
   	S = svd(G)
   	rk = rank(G; rtol =  1e-7)
	Urpi = S.U[:, 1:rk]'
	#print("Urpi type :", typeof(Urpi),"\n")
	return Diagonal(sqrt.(S.S[1:rk])) * Urpi * Ure, Mre
end

function rpi_basis(A::Rot3DCoeffs{T,1},
						 zz::SVector{N},
						 nn::SVector{N, Int},
						 ll::SVector{N, Int}) where {T,N}
	Mre = collect( _mrange(ll, 1))   # rows...
	G = _gramian(A, zz, nn, ll, Mre)
	#print("G type :", typeof(G),"\n")
   	S = svd(G)
   	rk = rank(G; rtol =  1e-7)
	Urepi = Diagonal(1/sqrt.(S.S[1:rk])) * S.U[:, 1:rk]'
	Utilde = zeros(SArray{Tuple{3},Complex{T},1,3}, rk, length(Mre))
	for α in 1:rk
		for (imu, mu) in enumerate(Mre)
			for (i, (k,mm)) in enumerate(Iterators.product(1:3, Mre))
				Utilde[α,imu] += Urepi[α,i] * A(ll,mu,mm)[:,k]
			end
		end
	end
	#print("Urpi type :", typeof(Urpi),"\n")
	return Utilde, Mre
end

function _gramian(A::Rot3DCoeffs{T,1}, zz, nn, ll::SVector{N}, Mre) where{T,N}
  	len = length(Mre)*3
	GG = zeros(Complex{T}, len, len)

	for (i1, (k1,mm1)) in enumerate(Iterators.product(1:3, _mrange(ll, 1))), (i2, (k2,mm2)) in enumerate(Iterators.product(1:3, _mrange(ll, 1)))
		for σ_ in permutations(1:N)
			σ = SVector{N,Int}(σ_)
	      	if ( zz[σ] != zz) || (nn[σ] != nn) || (ll[σ] != ll); continue; end
			for mu1 in  _mrange(ll, 1), mu2 in _mrange(ll, 1)
	        	if mu1[σ] == mu2
					GG[i1, i2] += dot(A(ll, mu2, mm2)[:,k2],A(ll[σ], mu1[σ], mm1[σ])[:,k1])
	         	end
	      	end
	   	end
	end
   	return GG
	#@assert all(abs.(imag(GG)) .<= imag_tol)
	#return real.(GG)
end

function _gramian(zz, nn, ll, Uri, Mri)
   N = length(nn)
   nri = size(Uri, 1)
   @assert size(Uri, 1) == nri
   G = zeros(Float64,nri, nri)

   for σ in permutations(1:N)
      if (zz[σ] != zz) || (nn[σ] != nn) || (ll[σ] != ll); continue; end
      for (iU1, mm1) in enumerate(Mri), (iU2, mm2) in enumerate(Mri)
         if mm1[σ] == mm2
            for i1 = 1:nri, i2 = 1:nri
               G[i1, i2] += conj(Uri[i1, iU1]) * Uri[i2, iU2]
            end
         end
      end
   end
   return G
   #@assert all(abs.(imag(G)) .<= imag_tol)
   #return real.(G)
end

end


@doc """
`CoB` is an abstract **c**hange **o**f **b**asis super-type
"""->
abstract CoB

# ------------------------------------------------------------
# Frequency to wavelets

@doc """
`Freq2Wave` is a change of basis type between frequency samples and wavelets.

There are sub types for wavelets with and without boundary correction.
To initialize a `Freq2Wave` type, run

	Freq2Wave(samples, wavename::String, J::Int, B; ...)

- `samples` are the sampling locations as a vector for 1D and M-by-2 matrix for 2D.
- `wave` is the name of the wavelet; see documentation for possibilities.
- `J` is the scale of the wavelet transform to reconstruct.
- `B` is the bandwidth of the samples; only needed if `samples` are non-uniform.
- Optional arguments (if needed) are passed to the functions computing Fourier transforms.
"""->
abstract Freq2Wave <: CoB

abstract Freq2Wave1D <: Freq2Wave
abstract Freq2Wave2D <: Freq2Wave
# The 2D Freq2BoundaryWave's need entries the 1D's do not, so don't
# use the dimension as a TypePar.

# No boundary correction
immutable Freq2NoBoundaryWave1D <: Freq2Wave1D
	# Sampling
	internal::Vector{Complex{Float64}}
	weights::Nullable{ Vector{Complex{Float64}} }

	# Reconstruction
	J::Int64
	wavename::AbstractString

	# Multiplication
	NFFT::NFFT.NFFTPlan{1,0,Float64}

	tmpMulVec::Vector{Complex{Float64}}
end

immutable Freq2NoBoundaryWave2D <: Freq2Wave2D
	internal::Vector{ Vector{Complex{Float64}} }
	weights::Nullable{ Vector{Complex{Float64}} }

	J::Int64
	wavename::AbstractString

	NFFT::NFFT.NFFTPlan{2,0,Float64}

	tmpMulVec::Vector{Complex{Float64}}
end

# Boundary correction
immutable Freq2BoundaryWave1D <: Freq2Wave1D
	internal::Vector{Complex{Float64}}
	weights::Nullable{Vector{Complex{Float64}}}

	J::Int64
	wavename::AbstractString

	NFFT::NFFT.NFFTPlan{1,0,Float64}

	left::Matrix{Complex{Float64}}
	right::Matrix{Complex{Float64}}

	tmpMulVec::Vector{Complex{Float64}}
end

immutable Freq2BoundaryWave2D <: Freq2Wave2D
	internal::Vector{ Vector{Complex{Float64}} }
	weights::Nullable{Vector{Complex{Float64}}}

	J::Int64
	wavename::AbstractString

	NFFT::NFFT.NFFTPlan{2,0,Float64}
	NFFTx::NFFT.NFFTPlan{2,1,Float64}
	NFFTy::NFFT.NFFTPlan{2,2,Float64}

	#= left::Dict{ Symbol, Matrix{Complex{Float64}} } =#
	left::Vector{ Matrix{Complex{Float64}} }
	right::Vector{ Matrix{Complex{Float64}} }

	tmpMulVec::Matrix{Complex{Float64}}
	tmpMulVecT::Matrix{Complex{Float64}} # Serves as the transpose of tmpMulVec
	tmpMulcVec::Vector{Complex{Float64}}
	weigthedVec::Vector{Complex{Float64}}
end

# ------------------------------------------------------------------------
# Constructors

function Freq2NoBoundaryWave1D(internal, weights, J, wavename, NFFT)
	tmpMulVec = Array{Complex{Float64}}( NFFT.M )
	Freq2NoBoundaryWave1D( internal, weights, J, wavename, NFFT, tmpMulVec )
end

function Freq2NoBoundaryWave2D(internal, weights, J, wavename, NFFT)
	tmpMulVec = Array{Complex{Float64}}( NFFT.M )
	Freq2NoBoundaryWave2D( internal, weights, J, wavename, NFFT, tmpMulVec )
end

function Freq2BoundaryWave1D(internal, weights, J, wavename, NFFT, left, right)
	tmpMulVec = Array{Complex{Float64}}( NFFT.M )
	Freq2BoundaryWave1D( internal, weights, J, wavename, NFFT, left, right, tmpMulVec )
end

function Freq2BoundaryWave2D(internal, weights, J, wavename, NFFT, left, right)
	tmpMulVec = similar(left[1])
	tmpMulVecT = tmpMulVec.'

	tmpMulcVec = Array{Complex{Float64}}( NFFT.M )
	weigthedVec = similar(tmpMulcVec)

	vm = van_moment(wavename)
	NFFTx = NFFTPlan( NFFT.x[1,:], 1, (NFFT.N[1],vm) )
	NFFTy = NFFTPlan( NFFT.x[2,:], 2, (vm,NFFT.N[2]) )

	Freq2BoundaryWave2D( internal, weights, J, wavename, NFFT,
	NFFTx, NFFTy, left, right, tmpMulVec, tmpMulVecT, tmpMulcVec, weigthedVec )
end

# ------------------------------------------------------------------------
# Load methods

include("freq2wave.jl")


module AeroMDAO

## Libraries
#==========================================================================================#

using LinearAlgebra
using StaticArrays
using CoordinateTransformations
using Rotations
using ForwardDiff: jacobian!
using DiffResults: JacobianResult, jacobian, value
using PrettyTables
using StructArrays

using Statistics: mean

using SplitApplyCombine: combinedimsview, combinedims
export combinedimsview, combinedims

using ComponentArrays

export ComponentVector, ComponentArray, valkeys

using Setfield
using LabelledArrays

## Methods to be extended in submodules
#==========================================================================================#

function collocation_point end

function velocity end

function solve_linear end

function solve_nonlinear end

function solve_linear! end

function solve_nonlinear! end

function solve_system end

function properties end

function surface_velocities end

function surface_coefficients end

export solve_linear, solve_nonlinear, solve_nonlinear!, collocation_point

## Math tools
#==========================================================================================#

include("Tools/MathTools.jl")
import .MathTools: forward_sum, forward_difference, forward_division, weighted_vector, vectarray, slope, splitat, adj3, columns, extend_yz, reflect_mapper, cosine_interp, structtolist, inverse_rotation, rotation, affine_2D, Point2D, Point3D, x, y, z, reshape_array, midpair_map, partition, uniform_spacing, linear_spacing, cosine_spacing, sine_spacing

# export forward_sum, forward_difference, forward_division, weighted_vector, vectarray, slope, splitat, adj3, columns, extend_yz, reflect_mapper, cosine_interp, structtolist, inverse_rotation, rotation, affine_2D, Point2D, Point3D, reshape_array, midpair_map, partition, uniform_spacing, linear_spacing, cosine_spacing, sine_spacing

## Non-dimensionalization
#==========================================================================================#

include("Tools/NonDimensional.jl")
import .NonDimensional: dynamic_pressure, force_coefficient, moment_coefficient, rate_coefficient, pressure_coefficient, aerodynamic_coefficients, reynolds_number, force, moment

export dynamic_pressure, force_coefficient, moment_coefficient, rate_coefficient, pressure_coefficient, reynolds_number, force, moment

## Panels
#==========================================================================================#

include("Geometry/PanelGeometry/PanelGeometry.jl")
import .PanelGeometry: AbstractPanel, AbstractPanel2D, Panel2D, WakePanel2D, AbstractPanel3D, Panel3D, p1, p2, p3, p4, zero, panel_length, transform_panel, transform_panel_points, panel_angle, panel_tangent, panel_normal, panel_location, panel_area, panel_coordinates, transform, midpoint, panel_points, wake_panel, wake_panels, reverse_panel, panel_velocity, panel_scalar, trailing_edge_panel, get_surface_values, panel_vector, distance, average_chord, average_width, wetted_area, make_panels, local_coordinate_system

export AbstractPanel, AbstractPanel2D, Panel2D, WakePanel2D, AbstractPanel3D, Panel3D, p1, p2, p3, p4, transform, panel_normal, midpoint, panel_location, panel_tangent, panel_points, distance, wake_panel, wake_panels, panel_area, reverse_panel, panel_length, transform_panel, panel_angle, panel_vector, panel_velocity, panel_scalar, trailing_edge_panel, get_surface_values, average_chord, average_width, wetted_area, make_panels, local_coordinate_system

## Aircraft geometry
#==========================================================================================#

include("Geometry/AircraftGeometry/AircraftGeometry.jl")
import .AircraftGeometry: AbstractAircraft, AbstractWing, AbstractFoil, Foil, arc_length, kulfan_CST, naca4, camber_CST, make_panels, read_foil, leading_edge_index, upper_surface, lower_surface, split_surface, coordinates_to_camber_thickness, camber_thickness_to_coordinates, camber_thickness, camber_thickness_to_coordinates, cosine_spacing, camber_thickness_to_CST, coordinates_to_CST, maximum_thickness_to_chord, translate, rotate, affine, scale, camber, camber_line, control_surface, # 2D setups
Fuselage, projected_area, length, cosine_spacing, # Fuselage
HalfWing, HalfWingSection, Wing, WingSection, affine_transformation, mean_aerodynamic_chord, span, aspect_ratio, projected_area, taper_ratio, leading_edge, trailing_edge, chop_leading_edge, chop_trailing_edge, chop_wing, chop_sections, chop_coordinates, chop_spanwise_sections, chop_chords, chop_spans, wing_bounds, make_panels, mesh_chords, mesh_wing, mesh_cambers, max_thickness_to_chord_ratio_sweeps, mean_aerodynamic_center, panel_wing, number_of_spanwise_panels, symmetric_spacing, coordinates, chord_coordinates, camber_coordinates, surface_coordinates, foils, chords, twists, spans, dihedrals, sweeps, position, orientation, WingMesh, chord_panels, camber_panels, normal_vectors, surface_panels, AbstractSpacing, Sine, Cosine, Uniform, properties

export AbstractAircraft, AbstractWing, AbstractFoil, Foil, arc_length, kulfan_CST, naca4, camber_CST, make_panels, read_foil, leading_edge_index, upper_surface, lower_surface, split_surface, coordinates_to_camber_thickness, camber_thickness_to_coordinates, camber_thickness, camber_thickness_to_coordinates, cosine_spacing, camber_thickness_to_CST, coordinates_to_CST, maximum_thickness_to_chord, translate, rotate, affine, scale, camber, camber_line, control_surface, # 2D setups
Fuselage, projected_area, length, cosine_spacing, # Fuselage
HalfWing, HalfWingSection, Wing, WingSection, affine_transformation, mean_aerodynamic_chord, span, aspect_ratio, projected_area, taper_ratio, leading_edge, trailing_edge, chop_leading_edge, chop_trailing_edge, chop_wing, chop_sections, chop_coordinates, chop_spanwise_sections, chop_chords, chop_spans, wing_bounds, make_panels, mesh_chords, mesh_wing, mesh_cambers, max_thickness_to_chord_ratio_sweeps, mean_aerodynamic_center, panel_wing, number_of_spanwise_panels, symmetric_spacing, coordinates, chord_coordinates, camber_coordinates, surface_coordinates, foils, chords, twists, spans, dihedrals, sweeps, position, orientation, WingMesh, chord_panels, camber_panels, normal_vectors, surface_panels, AbstractSpacing, Sine, Cosine, Uniform, properties

make_horseshoes(wing :: WingMesh) = StructArray(Horseshoe.(chord_panels(wing), normal_vectors(wing)))

make_vortex_rings(wing :: WingMesh) = StructArray(VortexRing.(camber_panels(wing)))

export make_horseshoes, make_vortex_rings

## Laplace
#==========================================================================================#

include("Tools/Laplace.jl")
import .Laplace: Uniform2D, potential, stream,  cartesian_to_freestream, freestream_to_cartesian

export Uniform2D, stream, vortex_stream_1, vortex_stream_2, source_stream, cartesian_to_freestream, freestream_to_cartesian

## Aerodynamic analyses
#==========================================================================================#

## Doublet-source panel method

include("Aerodynamics/DoubletSource/DoubletSource.jl")
import .DoubletSource: doublet_matrix, source_matrix, boundary_vector, wake_panels, source_strengths, surface_velocities, lift_coefficient

export doublet_matrix, source_matrix, boundary_vector, wake_panels, source_strengths, surface_velocities, lift_coefficient

## Linear-strength source and vorticity panel method

include("Aerodynamics/LinearVortexSource/LinearVortexSource.jl")
import .LinearVortexSource: total_velocity, source_velocity, vortex_velocity, vortex_influence_matrix, source_influence_matrix, neumann_boundary_condition, kutta_condition, two_point_matrix, linear_source_matrix, linear_vortex_matrix, constant_source_matrix, constant_source_boundary_condition

export total_velocity, source_velocity, vortex_velocity, vortex_influence_matrix, source_influence_matrix, neumann_boundary_condition, kutta_condition, two_point_matrix, linear_source_matrix, linear_vortex_matrix, constant_source_matrix, constant_source_boundary_condition

## Vortex lattice

include("Aerodynamics/VortexLattice/VortexLattice.jl")
import .VortexLattice: Horseshoe, VortexLatticeSystem, References, AbstractAxisSystem, Stability, Wind, Body, Geometry, streamlines, influence_coefficient, influence_matrix, boundary_condition, solve_system, transform, bound_leg, bound_leg_center, bound_leg_vector, r1, r2, Horseshoe, surface_velocity, surface_forces, surface_moments, nearfield_drag, geometry_to_wind_axes, geometry_to_stability_axes, stability_to_geometry_axes, wind_to_geometry_axes,  rate_coefficient, nearfield, farfield, farfield_forces, surface_velocities, surface_forces, surface_dynamics, surface_coefficients, nearfield_coefficients, farfield_coefficients, VortexRing, Freestream, velocity, body_frame_velocity, kinematic_viscosity, mach_number

export Horseshoe, VortexLatticeSystem, References, AbstractAxisSystem, Stability, Wind, Body, Geometry, streamlines, influence_coefficient, influence_matrix, boundary_condition, solve_system, transform, bound_leg, bound_leg_center, bound_leg_vector, r1, r2, Horseshoe, surface_velocity, surface_forces, surface_moments, nearfield_drag, geometry_to_wind_axes, geometry_to_stability_axes, stability_to_geometry_axes, wind_to_geometry_axes,  rate_coefficient, nearfield, farfield, farfield_forces, surface_velocities, surface_forces, surface_dynamics, surface_coefficients, nearfield_coefficients, farfield_coefficients, VortexRing, Freestream, velocity, body_frame_velocity, kinematic_viscosity, mach_number

## Profile drag estimation

include("Aerodynamics/profile_drag.jl")

export wetted_area_drag, profile_drag_coefficient, local_dissipation_drag

## Viscous airfoil analysis

include("Aerodynamics/ViscFoil/ViscFoil.jl")

import .ViscFoil: BoundaryLayer2D, solve_inviscid_doublets, solve_inviscid_vortices, defect_block, edge_velocities, solve_viscous_case

export BoundaryLayer2D, solve_inviscid_doublets, solve_inviscid_vortices, defect_block, edge_velocities, solve_viscous_case

## Cases

include("Aerodynamics/Cases/printing.jl")

export print_case, print_info, print_coefficients, print_derivatives

include("Aerodynamics/Cases/cases.jl")

export solve_case, spanwise_loading, triangle_connectivities, extrapolate_point_mesh

include("Aerodynamics/Cases/stability_cases.jl")

export solve_case_derivatives, longitudinal_stability_derivatives, longitudinal_stability_matrix, lateral_stability_derivatives, lateral_stability_matrix

include("Aerodynamics/Cases/foil_cases.jl")

## Structural analyses
#==========================================================================================#

include("Structures/Beams.jl")
import .Beams: Material, Tube, Beam, radii, area, moment_of_inertia, polar_moment_of_inertia, J_coeffs, Iyy_coeffs, Izz_coeffs, tube_stiffness_matrix, bending_stiffness_matrix, axial_stiffness_matrix, build_stiffness_matrix, solve_cantilever_beam, elastic_modulus, shear_modulus, yield_stress, density, principal_stress, torsional_stress, von_mises_stress

export Material, Tube, Beam, radii, area, moment_of_inertia, polar_moment_of_inertia, J_coeffs, Iyy_coeffs, Izz_coeffs, tube_stiffness_matrix, bending_stiffness_matrix, axial_stiffness_matrix, build_stiffness_matrix, solve_cantilever_beam, elastic_modulus, shear_modulus, yield_stress, density, principal_stress, torsional_stress, von_mises_stress

## Aerostructural analyses
#==========================================================================================#

include("Aerostructural/Aerostructural.jl")
import .Aerostructural: AerostructWing, make_beam_mesh, transform_stiffy, permute_stiffy, build_big_stiffy, adjacent_adder, section_moments, compute_loads, fem_load_vector, rotation_matrix, transfer_displacements, mesh_translation, mesh_rotation, new_horseshoes, solve_coupled_residual!, aerostruct_gauss_seidel

export AerostructWing, make_beam_mesh, transform_stiffy, permute_stiffy, build_big_stiffy, adjacent_adder, section_moments, compute_loads, fem_load_vector, rotation_matrix, transfer_displacements, mesh_translation, mesh_rotation, new_horseshoes, solve_coupled_residual!, aerostruct_gauss_seidel

## Post-processing
#==========================================================================================#

include("Tools/plot_tools.jl")

export plot_panels, plot_streamlines, plot_planform, plot_surface

end


# using Revise
# using CairoMakie, AbstractPlotting, ZigZagBoomerang, SparseArrays, LinearAlgebra

# Essence of the model section procedure.
# a = Dict((1,0)=>1, (0,0)=>2)
# push!(a, (2,0) => 3)
# (1,0) ∈ keys(a)
# key1 = (1,0)
# a[key1] = a[key1] + 1

# function ϕ(x, i, μ)
#     x[i] - μ[i]
# end
# κ = 1.0
# n = 100
# μ = zeros(n)
# x0 = randn(n)
# θ0 = rand((-1.0,1.0), n)
# T = 1000.0
# c = 10ones(n)
#@time trace0, _ = ZigZagBoomerang.spdmp(ϕ, 0.0, x0, θ0, T, c, ZigZag(sparse(1.0I,n,n), zeros(n)), μ)
# @time trace0, _ = ZigZagBoomerang.sspdmp(ϕ, 0.0, x0, θ0, T, c, κ, ZigZag(sparse(1.0I,n,n), zeros(n)), μ)

is_stuck(xi,vi) = xi == 0 && vi == 0 ? 1 : 0

"""
    model_trace(x0, v0, events)
given an initial position x0, v0 and a trace of events, assign to each
event, a vector k corresponding to the open sets as in Appendix 1.
`m0` is the initial model, `mm` is the trace of models at each event
"""
function model_trace(x0, v0, events)
    m0 = is_stuck.(x0, v0)
    m = copy(m0)
    mm = Vector{Int64}[]
    k = 1
    for event in events
        t, i, x, v = event
        m[i] = is_stuck(x, v)
        push!(mm, copy(m))
        k += 1
    end
    m0, mm
end

"""
    model_probs(m0, mm, events)
Given the initial model and the trace of models and trace of events,
computes the time spent by the sampler in each model explored. returns
a Ditcionary `probs` containing as `key` the model and as value its total time
"""
function model_probs(m0, mm, events)
    probs = Dict(m0=> events[1][1],)
    t = 0.0
    for i in eachindex(events)[2:end]
        tnew, i, x, v = events[i]
        δt = tnew - t
        t = tnew
        key = mm[i]
        if key ∈ keys(probs)
            probs[key] += δt
        else
            push!(probs, key => δt)
        end
    end
    probs
end

# m0, mm = model_trace(x0, θ0, trace0.events)
# model_probs(m0, mm, trace0.events)


println(sum(filter(x -> mod(x, 3) == 0 || mod(x, 5) == 0, 1:999)))


# # Getting started
# ExplainableAI.jl can be used on any classifier.
# In this first example, we will look at attributions on a LeNet5 model that was pretrained on MNIST.
#
# ### Loading the model
#md # !!! note
#md #
#md #     Outside of these docs, you should be able to load the model using
#md #     ```julia
#md #     using BSON: @load
#md #     @load "model.bson" model
#md #     ```

using ExplainableAI
using Flux
using BSON

model = BSON.load("../model.bson", @__MODULE__)[:model]

#md # !!! warning "Strip softmax"
#md #
#md #     For models with softmax activations on the output, it is necessary to call
#md #     ```julia
#md #     model = strip_softmax(model)
#md #     ```
#md #     before analyzing.

# ### Loading MNIST
# We use MLDatasets to load a single image from the MNIST dataset:
using MLDatasets
using ImageCore
using ImageIO
using ImageShow

index = 10
x, _ = MNIST(Float32, :test)[10]

convert2image(MNIST, x)

# By convention in Flux.jl, this input needs to be resized to WHCN format by adding a color channel and batch dimensions.
input = reshape(x, 28, 28, 1, :);

#md # !!! warning "Input format"
#md #
#md #     For any attribution of a model, ExplainableAI.jl assumes the batch dimension to be come last in the input.
#md #
#md #     For the purpose of heatmapping, the input is assumed to be in WHCN order
#md #     (width, height, channels, batch), which is Flux.jl's convention.

# ## Calling the analyzer
# We can now select an analyzer of our choice
# and call [`analyze`](@ref) to get an `Explanation`:
analyzer = LRP(model)
expl = analyze(input, analyzer);

# This `Explanation` bundles the following data:
# * `expl.attribution`: the analyzer's attribution
# * `expl.output`: the model output
# * `expl.neuron_selection`: the neuron index of used for the attribution
# * `expl.analyzer`: a symbol corresponding the used analyzer, e.g. `:LRP`

# Finally, we can visualize the `Explanation` through heatmapping:
heatmap(expl)

# Or get the same result by combining both analysis and heatmapping into one step:
heatmap(input, analyzer)

# ## Neuron selection
# By passing an additional index to our call to `analyze`, we can compute the attribution
# with respect to a specific output neuron.
# Let's see why the output wasn't interpreted as a 4 (output neuron at index 5)
heatmap(input, analyzer, 5)

# This heatmap shows us that the "upper loop" of the hand-drawn 9 has negative relevance
# with respect to the output neuron corresponding to digit 4!

#md # !!! note
#md #
#md #     The ouput neuron can also be specified when calling `analyze`:
#md #     ```julia
#md #     expl = analyze(img, analyzer, 5)
#md #     ```

# ## Input batches
# ExplainableAI also supports input batches:
batchsize = 100
xs, _ = MNIST(Float32, :test)[1:batchsize]
batch = reshape(xs, 28, 28, 1, :) # reshape to WHCN format
expl_batch = analyze(batch, analyzer);

# This will once again return a single `Explanation` `expl_batch` for the entire batch.
# Calling `heatmap` on `expl_batch` will detect the batch dimension and return a vector of heatmaps.
#
# Let's check if the digit at `index = 10` still matches.
hs = heatmap(expl_batch)
hs[index]

# JuliaImages' `mosaic` can be used to return a tiled view of all heatmaps:
mosaic(hs; nrow=10)

# We can also evaluate a batch with respect to a specific output neuron, e.g. for the digit zero at index `1`:
mosaic(heatmap(batch, analyzer, 1); nrow=10)

# ## Automatic heatmap presets
# Currently, the following analyzers are implemented:
#
# ```
# ├── Gradient
# ├── InputTimesGradient
# ├── SmoothGrad
# ├── IntegratedGradients
# └── LRP
#     ├── LRPZero
#     ├── LRPEpsilon
#     └── LRPGamma
# ```
#
# Let's try `InputTimesGradient`
analyzer = InputTimesGradient(model)
heatmap(input, analyzer)

# and `Gradient`
analyzer = Gradient(model)
heatmap(input, analyzer)

# As you can see, the function `heatmap` automatically applies common presets for each method.
#
# Since `InputTimesGradient` and LRP both compute attributions, their presets are similar.
# Gradient methods however are typically shown in grayscale.

# ## Custom heatmap settings
# We can partially or fully override presets by passing keyword arguments to `heatmap`:
using ColorSchemes
heatmap(expl; cs=ColorSchemes.jet)
#
heatmap(expl; reduce=:sum, rangescale=:extrema, cs=ColorSchemes.inferno)

# This also works with batches
mosaic(heatmap(expl_batch; rangescale=:extrema, cs=ColorSchemes.inferno); nrow=10)

# For the full list of keyword arguments, refer to the [`heatmap`](@ref) documentation.


# # Modal analysis of a free-floating steel circle

# Source code: [`circle_modal_conv_tut.jl`](circle_modal_conv_tut.jl)

# ## Description

# Vibration analysis of a free-floating steel ring. This is a
# benchmark from the NAFEMS Selected Benchmarks for Natural Frequency Analysis,
# publication: Test VM09: Circular Ring --  In-plane and Out-of-plane
# Vibration.

# The results can be compared with analytical expressions, but the main purpose
# is to compute data for extrapolation to the limit to predict the true natural
# frequencies.

# ## Reference frequencies

# There will be 6 rigid body modes (zero natural frequencies).

# The numerical results are due to the publication:
# NAFEMS Finite Element Methods & Standards, Abbassian, F., Dawswell, D. J., and
# Knowles, N. C. Selected Benchmarks for Natural Frequency Analysis, Test No.
# 6. Glasgow: NAFEMS, Nov., 1987.

# The reference values were analytically determined (Blevins, FORMULAS FOR
# DYNAMICS, ACOUSTICS AND VIBRATION, Table 4.16). Note that shear flexibility
# was neglected when computing the reference values.

# | Mode       |         Reference Value (Hz)  |  NAFEMS Target Value (Hz) |
# | -------   |     -------  |  ------- |
# | 7, 8 | (out of plane)   |        51.85          |         52.29  |
# | 9, 10 |  (in plane)       |       53.38         |          53.97  |
# | 11, 12 |  (out of plane)   |     148.8          |         149.7  |
# | 13, 14 |  (in plane)       |     151.0          |         152.4  |
# | 15, 16 |  (out of plane)   |     287.0          |         288.3  |
# |  17, 18 |  (in plane)      |      289.5         |          288.3  |

# ## Goals

# - Show convergence relative to reference values.
# - Compute data for extrapolation to the limit to predict the true natural
#   frequencies.
#

##
# ## Definition of the basic inputs

# Include the needed packages and modules.
using Arpack
using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsFlexBeams.MeshFrameMemberModule: frame_member
using FinEtoolsFlexBeams.CrossSectionModule: CrossSectionCircle
using FinEtoolsFlexBeams.FEMMCorotBeamModule
CB = FEMMCorotBeamModule
using FinEtoolsFlexBeams.FEMMCorotBeamModule: FEMMCorotBeam
using FinEtoolsFlexBeams.FESetCorotBeamModule: MASS_TYPE_CONSISTENT_NO_ROTATION_INERTIA,
    MASS_TYPE_CONSISTENT_WITH_ROTATION_INERTIA,
    MASS_TYPE_LUMPED_DIAGONAL_NO_ROTATION_INERTIA,
    MASS_TYPE_LUMPED_DIAGONAL_WITH_ROTATION_INERTIA

# The material parameters may be defined with the specification of the units.
# The elastic properties are:
E = 200.0 * phun("GPa")
nu = 0.3;

# The mass density is
rho = 8000 * phun("kg/m^3")
# Here are the cross-sectional dimensions and the length of the beam between supports.
radius = 1.0 * phun("m"); diameter = 0.1 * phun("m");

# We shall calculate these eigenvalues, but we are mostly interested in the
# first three  natural frequencies.
neigvs = 18;

# The mass shift needs to be applied since the structure is free-floating.
oshift = (2*pi*15)^2

# Here we get to choose the model: Bernoulli or Timoshenko
shear_correction_factor = 6/7 # Timoshenko
# shear_correction_factor = Inf # Bernoulli
cs = CrossSectionCircle(s -> diameter/2, s -> [1.0, 0.0, 0.0], shear_correction_factor)
@show cs.parameters(0.0)

# Here we can choose the mass- matrix type:

mtype = MASS_TYPE_CONSISTENT_NO_ROTATION_INERTIA
mtype = MASS_TYPE_CONSISTENT_WITH_ROTATION_INERTIA
mtype = MASS_TYPE_LUMPED_DIAGONAL_NO_ROTATION_INERTIA
mtype = MASS_TYPE_LUMPED_DIAGONAL_WITH_ROTATION_INERTIA

# Here are the formulas for the first two natural frequencies, obtained
# analytically with the shear flexibility  neglected. The parameters of the
# structure:
R = radius
Im = cs.parameters(0.0)[4]
m = rho * cs.parameters(0.0)[1]

# For instance the the first out of plane mode is listed in this table as
J = cs.parameters(0.0)[2]
G = E/2/(1+nu)
i = 2 # the first non-rigid body mode
@show i*(i^2-1)/(2*pi*R^2)*sqrt(E*Im/m/(i^2+E*Im/G/J))

# The first "ovaling" (in-plane) mode is:
i=2 # the first ovaling mode
@show i*(i^2-1)/(2*pi*R^2*(i^2+1)^(1/2))*sqrt(E*Im/m)

material = MatDeforElastIso(DeforModelRed3D, rho, E, nu, 0.0)

# We will generate this many elements  along the length of the circular ring.
results = let
    results = []
    for i in 1:3
        n = 10*2^i
        # beam elements along the member.
        tolerance = radius/n/1000;
        # Generate the mesh of a straight member.
        fens, fes = frame_member([0 0 0; 2*pi 0 0], n, cs)
        # Twist the straight member into a ring.
        for i in 1:count(fens)
            a = fens.xyz[i, 1]
            fens.xyz[i, :] .= (radius+radius*cos(a), radius*sin(a), 0)
        end
        # Merge the nodes of the bases, which involves renumbering the connectivity.
        fens, fes = mergenodes(fens, fes, tolerance, [1, n+1])
        # Generate the discrete model.
        geom0 = NodalField(fens.xyz)
        u0 = NodalField(zeros(size(fens.xyz, 1), 3))
        using FinEtoolsFlexBeams.RotUtilModule: initial_Rfield
        Rfield0 = initial_Rfield(fens)
        dchi = NodalField(zeros(size(fens.xyz, 1), 6))
        applyebc!(dchi)
        numberdofs!(dchi);
        femm = FEMMCorotBeam(IntegDomain(fes, GaussRule(1, 2)), material);
        K = CB.stiffness(femm, geom0, u0, Rfield0, dchi);
        M = CB.mass(femm, geom0, u0, Rfield0, dchi; mass_type = mtype);
        # Solve the free vibration problem.
        evals, evecs, nconv = eigs(K + oshift * M, M; nev=neigvs, which=:SM, ncv = 3*neigvs, maxiter = 2000);
        # Correct for the mass shift.
        evals = evals .- oshift;
        sigdig(n) = round(n * 10000) / 10000
        fs = real(sqrt.(complex(evals)))/(2*pi)
        println("Eigenvalues: $(sigdig.(fs)) [Hz]")
        push!(results, (fs[6+1], fs[6+3], fs[6+5]))
    end
    results # return these values from the block
end

@show results

##
# ## Richardson extrapolation

# Here we will use Richardson extrapolation from the three sets of data. This
# will allow us to predict the convergence rate and the true solution for each
# of the three frequencies (or rather the pairs of frequencies, 7 and 8, 9 and
# 10, and 11 and 12).

# We will immediately set up the convergence plots. We will extrapolate and then
# compute from that the normalized error to be plotted with respect to the
# refinement factors (which in this case are 4, 2, and 1). We use the
# refinement factor as a convenience: we will calculate the element size by
# dividing the circumference of the ring with a number of elements generated
# circumferentially.

using FinEtools.AlgoBaseModule: richextrapol

using Gnuplot
@gp  "set terminal windows 1 "  :-

# Modes 7 and 8
sols = [r[1] for r in results]
resextrap = richextrapol(sols, [4.0, 2.0, 1.0])
print("Predicted frequency 7 and 8: $(resextrap[1])\n")
errs = abs.(sols .- resextrap[1])./resextrap[1]
@gp  :- 2*pi*radius./[80, 160, 320] errs " lw 2 lc rgb 'red' with lp title 'Mode 7, 8' "  :-

# Modes 9 and 10
sols = [r[2] for r in results]
resextrap = richextrapol(sols, [4.0, 2.0, 1.0])
print("Predicted frequency 9 and 10: $(resextrap[1])\n")
errs = abs.(sols .- resextrap[1])./resextrap[1]
@gp  :- 2*pi*radius./[80, 160, 320] errs " lw 2 lc rgb 'green' with lp title 'Mode 9, 10' "  :-

# Modes 11 and 12
sols = [r[3] for r in results]
resextrap = richextrapol(sols, [4.0, 2.0, 1.0])
print("Predicted frequency 11 and 12: $(resextrap[1])\n")
errs = abs.(sols .- resextrap[1])./resextrap[1]
@gp  :- 2*pi*radius./[80, 160, 320] errs " lw 2 lc rgb 'blue' with lp title 'Mode 11, 12' "  :-

@gp  :- "set xrange [0.01:0.1]" "set logscale x" :-
@gp  :- "set logscale y" :-
@gp  :- "set xlabel 'Element size'" :-
@gp  :- "set ylabel 'Normalized error [ND]'" :-
@gp  :- "set title 'Beam: Convergence of modes 7, ..., 12'"


# Based on StackOverflow
function subscriptnumber(i::Int)
    if i < 0
        c = [Char(0x208B)]
    else
        c = []
    end
    for d in reverse(digits(abs(i)))
        push!(c, Char(0x2080+d))
    end
    return join(c)
end

function gather_plot_information(x::AbstractDataDrivenProblem{N, C, D}) where {N,C, D}

    X = getfield(x, :X)
    Y = getfield(x, :Y)
    DX = getfield(x, :DX)
    U = getfield(x, :U)

    t = x.t

    ylab = is_direct(x) ? "Sample ID" : "t"

    returns = []

    for (xi, si) in zip([X, DX, Y, U], ["x", "\U02202\U0209C"*"x", "y", "u"])
        isempty(xi) && continue
        push!(returns, (t, xi, si))
    end

    return returns, ylab
end

function gather_plot_information(x::AbstractDataDrivenSolution)
    p = get_problem(x)
    t = p.t

    t = is_discrete(p) ? t[2:end] : t
    ylab = is_direct(p) ? "Sample ID" : "t"

    Y = get_target(p)
    outsym = is_direct(p) ? "y" : (is_discrete(p) ? "x" : "\U02202\U0209C"*"x")
    est_sym = outsym * "\U00302"
    X, _, ts, u = get_oop_args(p)
    Ŷ = x.basis(X, parameters(x), ts, u)
    return [(t, Y, outsym), (t, Ŷ, est_sym)], ylab
end

## Problem

@recipe function probplot(p::AbstractDataDrivenProblem)
    plotins, ylab = gather_plot_information(p)

    layout := (length(plotins), 1)
    isdirec = is_direct(p)

    for (i, pins) in enumerate(plotins)
        yi, xi, lab = pins
        suff = isdirec ? "" : "(t)"
        @series begin
            if i > 1
                title := ""
            end
            label --> reduce(hcat, map(j->lab*subscriptnumber(j)*suff, 1:size(xi,1)))

            ylabel --> lab

            if length(plotins) == i
                xlabel --> ylab
            end

            subplot := i
            yi, permutedims(xi)
        end
    end

end

# Step 2...

## Solution !
# TODO Dispatch for SINDy to show a separate plot of the active
# terms in the candidates -> Need LatexStrings for that (most likely)
# Get the symbols of the states for plotting
# Add extreme values for error ?

@recipe function resplot(x::AbstractDataDrivenSolution; show_data::Bool = false)
    plotins, xlab = gather_plot_information(x)
    isdirec = is_direct(get_problem(x))
    layout := (2,1)

    suff = isdirec ? "" : "(t)"

    t, Y, lab = plotins[1]
    _ , Ŷ, lab_ = plotins[2]

    @series begin
        label --> reduce(hcat, map(j->lab_*subscriptnumber(j)*suff, 1:size(Y, 1)))
        ylabel --> lab_
        subplot := 1
        seriestype := :path
        t, permutedims(Ŷ)
    end

    if show_data
        @series begin
            label --> nothing
            ylabel --> lab
            subplot := 1
            seriestype --> :path
            primary := false
            linestyle --> :dash
            color --> :black
            t, permutedims(Y)
        end
    end

    elab = lab*suff*"-"*lab_*suff

    @series begin
        title := ""
        label --> reduce(hcat, map(j->"e"*subscriptnumber(j)*suff, 1:size(Y, 1)))
        ylabel --> elab
        subplot := 2
        xlabel --> xlab
        t, permutedims(Y-Ŷ)
    end

end


using StatsBase
using Test

n = 5000

# 1D integer counts

x = rand(1:5, n)
w = weights(rand(n))

c = counts(x, 5)
@test size(c) == (5,)
c0 = Int[count(v->v == i, x) for i in 1:5]
@test c == c0
@test counts(x .+ 1, 2:6) == c0
@test proportions(x, 1:5) ≈ (c0 ./ n)

c = counts(x)
@test size(c) == (5,)
c0 = Int[count(v->v == i, x) for i in 1:5]
@test c == c0
@test counts(x .+ 1, 2:6) == c0
@test proportions(x) ≈ (c0 ./ n)

c = counts(x, 5, w)
@test size(c) == (5,)
c0 = Float64[sum(w.values[x .== i]) for i in 1:5]
@test c                      ≈ c0
@test counts(x .+ 1, 2:6, w) ≈ c0
@test proportions(x, 1:5, w) ≈ (c0 ./ sum(w))

c = counts(x, w)
@test size(c) == (5,)
c0 = Float64[sum(w.values[x .== i]) for i in 1:5]
@test c                      ≈ c0
@test counts(x .+ 1, 2:6, w) ≈ c0
@test proportions(x, w)      ≈ (c0 ./ sum(w))

# 2D integer counts

x = rand(1:4, n)
y = rand(1:5, n)
w = weights(rand(n))

c = counts(x, y, (4, 5))
@test size(c) == (4, 5)
c0 = Int[count(t->t != 0,  (x .== i) .& (y .== j)) for i in 1:4, j in 1:5]
@test c == c0
@test counts(x .+ 2, y .+ 3, (3:6, 4:8)) == c0
@test proportions(x, y, (1:4, 1:5)) ≈ (c0 ./ n)

c = counts(x, y)
@test size(c) == (4, 5)
c0 = Int[count(t->t != 0, (x .== i) .& (y .== j)) for i in 1:4, j in 1:5]
@test c == c0
@test counts(x .+ 2, y .+ 3, (3:6, 4:8)) == c0
@test proportions(x, y,) ≈ (c0 ./ n)

c = counts(x, y, (4, 5), w)
@test size(c) == (4, 5)
c0 = Float64[sum(w.values[(x .== i) .& (y .== j)]) for i in 1:4, j in 1:5]
@test c                                     ≈ c0
@test counts(x .+ 2, y .+ 3, (3:6, 4:8), w) ≈ c0
@test proportions(x, y, (1:4, 1:5), w)      ≈ (c0 ./ sum(w))

c = counts(x, y, w)
@test size(c) == (4, 5)
c0 = Float64[sum(w.values[(x .== i) .& (y .== j)]) for i in 1:4, j in 1:5]
@test c                                     ≈ c0
@test counts(x .+ 2, y .+ 3, (3:6, 4:8), w) ≈ c0
@test proportions(x, y, w)                  ≈ (c0 ./ sum(w))

# count map

x = ["a", "b", "a", "a", "b", "c"]
w = [1.0, 1.5, 2.0, 2.5, 3.0, 3.5]

cm = countmap(x)
@test cm["a"] == 3
@test cm["b"] == 2
@test cm["c"] == 1

# iterator, non-radixsort
cm_missing = countmap(skipmissing(x))
cm_any_itr = countmap((i for i in x))
@test cm_missing == cm_any_itr == cm
@test cm_missing isa Dict{String, Int}
@test cm_any_itr isa Dict{Any, Int}

pm = proportionmap(x)
@test pm["a"] ≈ (1/2)
@test pm["b"] ≈ (1/3)
@test pm["c"] ≈ (1/6)

# testing the radixsort branch of countmap
xx = repeat([6, 1, 3, 1], outer=100_000)
cm = countmap(xx)
@test cm == Dict(1 => 200_000, 3 => 100_000, 6 => 100_000)

# with iterator
cm_missing = countmap(skipmissing(xx))
@test cm_missing isa Dict{Int, Int}
@test cm_missing == cm

cm_any_itr = countmap((i for i in xx))
@test cm_any_itr isa Dict{Any,Int} # no knowledge about type
@test cm_missing == cm

# testing the radixsort-based addcounts
xx = repeat([6, 1, 3, 1], outer=100_000)
cm = Dict{Int, Int}()
StatsBase.addcounts_radixsort!(cm,xx)
@test cm == Dict(1 => 200_000, 3 => 100_000, 6 => 100_000)
xx2 = repeat([7, 1, 3, 1], outer=100_000)
StatsBase.addcounts_radixsort!(cm,xx2)
@test cm == Dict(1 => 400_000, 3 => 200_000, 6 => 100_000, 7 => 100_000)
# with iterator
cm_missing = Dict{Int, Int}()
StatsBase.addcounts_radixsort!(cm_missing,skipmissing(xx))
@test cm_missing == Dict(1 => 200_000, 3 => 100_000, 6 => 100_000)
StatsBase.addcounts_radixsort!(cm_missing,skipmissing(xx2))
@test cm_missing == Dict(1 => 400_000, 3 => 200_000, 6 => 100_000, 7 => 100_000)

# testing the Dict-based addcounts
cm = Dict{Int, Int}()
cm_itr = Dict{Int, Int}()
StatsBase.addcounts_dict!(cm,xx)
StatsBase.addcounts_dict!(cm_itr,skipmissing(xx))
@test cm_itr == cm == Dict(1 => 200_000, 3 => 100_000, 6 => 100_000)
@test cm_itr isa Dict{Int, Int}

cm = countmap(x, weights(w))
@test cm["a"] == 5.5
@test cm["b"] == 4.5
@test cm["c"] == 3.5

@test cm == countmap(x, w)

pm = proportionmap(x, weights(w))
@test pm["a"] ≈ (5.5 / 13.5)
@test pm["b"] ≈ (4.5 / 13.5)
@test pm["c"] ≈ (3.5 / 13.5)

# testing small bits type
bx = [true, false, true, true, false]
cm_bx_missing = countmap(skipmissing(bx))
@test cm_bx_missing == countmap(bx) == Dict(true => 3, false => 2)
@test cm_bx_missing isa Dict{Bool, Int}

for T in [UInt8, UInt16, Int8, Int16]
    tx = T[typemin(T), 8, typemax(T), 19, 8]
    tx_missing = skipmissing(T[typemin(T), 8, typemax(T), 19, 8])
    cm_tx_missing = countmap(tx_missing)
    @test cm_tx_missing == countmap(tx) == Dict(typemin(T) => 1, typemax(T) => 1, 8 => 2, 19 => 1)
    @test cm_tx_missing isa Dict{T, Int}
end

@testset "views" begin
    X = view([1,1,1,2,2], 1:5)
    @test countmap(X) == countmap(copy(X))
end


module RationalSimplex

import LinearAlgebra: I

export simplex

"""
    simplex(c, obj, A, b, senses)

Solve the linear program

     {min/max}  dot(c, x)
    subject to  A x {>=|==|<=} b
                x >= 0,

where `x` is in the set of Rationals (so algorithm is exact).
`obj` should be either `:Min` or `:Max`.
`senses` should be a `Vector` of `Char`s: `'<','=','>'`.
"""
function simplex(c::Vector{T}, obj::Symbol, A::Matrix{T}, b::Vector{T},
                 senses::Vector{Char}) where {T <: Rational}
    # Any constraint that isn't an equality needs an auxiliary variable.
    num_auxiliary = count(sense -> sense != '=', senses)
    num_constraints, num_variables = size(A)
    num_variables_aux = num_variables + num_auxiliary
    # Extend constraint matrix and objective function, normalize objective.
    A_aux = zeros(T, num_constraints, num_variables_aux)
    A_aux[:, 1:num_variables] = A
    c_aux = zeros(T, num_variables_aux)
    c_aux[1:num_variables] = (obj == :Max) ? -c : c
    b_aux = copy(b)  # In case we modify signs.
    # Add the auxiliaries to the constraint matrix, and make sure b ≥ 0.
    offset = 1
    for (row, sense) in enumerate(senses)
        if sense == '<'
            A_aux[row, num_variables + offset] = one(T)
            offset += 1
        elseif sense == '>'
            A_aux[row, num_variables + offset] = -one(T)
            offset += 1
        end
        if b[row] < zero(T)
            A_aux[row, :] *= -one(T)
            b_aux[row] *= -one(T)
        end
    end
    # Solve problem with auxiliaries, but remove auxiliaries before returning.
    sense, x = simplex(c_aux, A_aux, b_aux)
    return sense, x[1:num_variables]
end

"""
    simplex(c, A, b)

Solve the linear program in standard computational form

           min  dot(c, x)
    subject to  A x == b [≥ 0]
                x >= 0,

where `x` is in the set of Rationals (so algorithm is exact).

The algorithm is the "two-phase primal revised simplex method".
In the first phase auxiliaries are created to get an initial basis.
We then eliminate the auxiliaries from the basis until it consists only of
actual variables. This is the "textbook" algorithm, and shouldn't be used for
anything that matters. It doesn't exploit sparsity at all. You could use it
with floating points but it won't work for anything except the most simple
problem due to accumulated errors and comparisons with zero.
"""
function simplex(c::Vector{T}, A::Matrix{T}, b::Vector{T}) where {T<:Rational}
    @assert all(b .> zero(T))

    # Set up data structures.
    num_constraints, num_variables = size(A)
    is_basic = zeros(Bool, num_variables + num_constraints)
    basic = zeros(Int, num_constraints)  # Indices of current basis.
    Binv = Matrix{T}(I, num_constraints, num_constraints)  # Basis inverse.
    cB = ones(T, num_constraints)  # Costs of basic variables.
    x = zeros(T, num_variables + num_constraints)  # Current solution.

    # Intialize phase one of RSM by setting basis = auxiliaries.
    for aux_index in 1:num_constraints
        basic[aux_index] = num_variables + aux_index
        is_basic[num_variables + aux_index] = true
        x[num_variables + aux_index] = b[aux_index]
    end
    phase_one = true

    # Begin simplex iterations.
    status = :Unknown
    while true
        # Calculate dual solution.
        π = vec(cB' * Binv)
        # Use it to calculate the reduced costs of the variables. Don't
        # calculate for auxiliaries - they can't re-enter the basis.
        rc = (phase_one ? zero(T) : c) .- vec(π' * A)
        @. rc *= !is_basic[1:num_variables]  # In basis, so don't consider.
        # Find variable to enter basis.
        min_rc, entering = findmin(rc)
        # If none have negative reduced cost, there are no variables that can
        # enter, and we are at optimality for this phase. This may or may not
        # be optimal for the actual problem.
        if min_rc >= zero(T)
            if phase_one
                phase_one = false
                # If any auxiliary still nonzero, we couldn't find a feasible
                # basis without auxiliaries.
                if any(x[num_variables+1:end] .> zero(T))
                    return :Infeasible, x[1:num_variables]
                end
                # Otherwise, start phase two with nice feasible basis.
                for i in 1:num_constraints
                    cB[i] = basic[i] > num_variables ? zero(T) : c[basic[i]]
                end
                continue
            else  # Phase two. We can't improve further, so we are optimal.
                return :Optimal, x[1:num_variables]
            end
        end
        # Calculate how the solution will change when our new variable enters
        # the basis and increases from zero.
        BinvAs = Binv * A[:, entering]
        # Perform a "ratio test" on each variable to determine which will
        # reach zero first.
        leaving = 0
        min_ratio = zero(T)
        for j in 1:num_constraints
            if BinvAs[j] > zero(T)
                ratio = x[basic[j]] / BinvAs[j]
                if ratio < min_ratio || leaving == 0
                    min_ratio = ratio
                    leaving = j
                end
            end
        end
        # If no variable will leave basis, then we have an unbounded problem.
        if leaving == 0
            return :Unbounded, x[1:num_variables]
        end
        # Update solution.
        for j in 1:num_constraints
            x[basic[j]] -= min_ratio * BinvAs[j]
        end
        x[entering] = min_ratio
        # Update basis inverse.
        # Our tableau is [ Binv b | Binv | BinvAs ] and we doing a pivot on the
        # `leaving` row of BinvAs.
        pivot_value = BinvAs[leaving]
        for basis_row in 1:num_constraints
            basis_row == leaving && continue  # All rows except `leaving` row.
            factor = BinvAs[basis_row] / pivot_value
            for basis_col in 1:num_constraints
                Binv[basis_row, basis_col] -= factor * Binv[leaving, basis_col]
            end
        end
        # Finally, the `leaving` row.
        for basis_col in 1:num_constraints
            Binv[leaving, basis_col] /= pivot_value
        end
        # Update variable status flags.
        is_basic[basic[leaving]] = false
        is_basic[entering] = true
        cB[leaving] = phase_one ? zero(T) : c[entering]
        basic[leaving] = entering
    end
end

end  # module RationalSimplex.


using Distributed, StatsBase, PyPlot

addprocs(4; exeflags="--project")
@everywhere begin

    using TransitionalMCMC, Distributions

    # Prior Bounds
    lb, ub  = -5, 5

    # Prior Density and sampler
    logprior(x) = logpdf(Uniform(lb, ub), x[1]) + logpdf(Uniform(lb, ub), x[2])
    priorRnd(Nsamples) = rand(Uniform(lb, ub), Nsamples, 2)

    # Log Likelihood
    function logLik(x)
        sleep(0.005)
        return -1 * ((x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7).^2)
    end

end

Nsamples = 200

samps, acc = tmcmc(logLik, logprior, priorRnd, Nsamples, 5, 2)

plt.scatter(samps[:,1], samps[:,2])

rmprocs(workers())


using LinearAlgebra

function build_AB(k, dQdy, U, Ψ, Dy, Dy2, y, params)

    k2 = k*k;
    Fr = params.Fr;
    Ny = length(U)-1;

    B = Dy2[2:end-1,2:end-1] - (k2 + Fr)*I;
    A = diagm(U[2:end-1]) * B + diagm(dQdy[2:end-1]);

    return A, B
end


"""
    MicrophysicsParameters

Module containing 1-moment bulk microphysics parameters.
"""
module MicrophysicsParameters
using CLIMA.ParametersType

@exportparameter MP_n_0          8e6 * 2         "Marshall-Palmer distribution n_0 coeff [1/m4]"
@exportparameter C_drag          0.55            "drag coefficient for rain drops [-]"

@exportparameter τ_cond_evap     10              "condensation/evaporation timescale [s]"

@exportparameter q_liq_threshold 5e-4            "autoconversion threshold [-]  ∈(0.5, 1) * 1e-3 "
@exportparameter τ_acnv          1e3             "autoconversion timescale [s]  ∈(1e3, 1e4) "

@exportparameter E_col           0.8             "collision efficiency [-]"

@exportparameter a_vent          1.5             "ventilation factor coefficient [-]"
@exportparameter b_vent          0.53            "ventilation factor coefficient [-]"
@exportparameter K_therm         2.4e-2          "thermal conductivity of air [J/m/s/K] "
@exportparameter D_vapor         2.26e-5         "diffusivity of water vapor [m2/s]"
@exportparameter ν_air           1.6e-5          "kinematic viscosity of air [m2/s]"
@exportparameter N_Sc            ν_air/D_vapor   "Schmidt number [-]"
end


module TestConsistency
    verbose = false

    using Base.Test

    import IntertemporalChoiceHeuristics:
        load,
        fit,
        models,
        simulate,
        make_closures,
        stderrs

    import Optim: minimizer

    import ForwardDiff: hessian

    import Distributions: ccdf, Chisq

    inputs, weights = load()
    n_trials = 1_000

    for model in models
        res = fit(model, inputs, weights)
        parameters = minimizer(res)
        m = model(parameters)

        for sim in 1:10
            sim_inputs, sim_weights = simulate(m, inputs, n_trials)

            sim_res = fit(model, sim_inputs, sim_weights)
            sim_parameters = minimizer(sim_res)
            sim_m = model(sim_parameters)

            p_value = NaN
            try
                se = stderrs(model, sim_inputs, sim_weights)
                z = (sim_parameters .- parameters) ./ se
                χ = sum(z.^2)
                p_value = ccdf(Chisq(length(parameters)), χ)
            end

            if verbose
                @printf(
                    "%s\t%.2f\t%.4f\n",
                    model,
                    norm(
                        100 * (parameters .- sim_parameters) ./ parameters,
                        Inf,
                    ),
                    p_value,
                )
            end
        end
    end
end


@everywhere using SharedArrays
@everywhere using Distributions
@everywhere using Random
@everywhere const EPS = 2.2204e-16

@everywhere function nmf(v,r)

    n, m = size(v)
    stopconv = 10000        # stopping criterion (can be adjusted)
    niter = 1000000         # maximum number of iterations (can be adjusted)

    cons = fill(false, (m,m))
    consold = fill(false, (m,m))
    inc = 0
    j = 0

    # initialize random w and h
    w = Random.rand(n, r)
    h = Random.rand(r, m)

    for i = 1:niter

        # divergence-reducing NMF iterations
        x1 = repeat(sum(w, dims=1)', 1, m)
        h = h.*(w'*(v./(w*h)))./x1
        x2 = repeat(sum(h, dims=2)', n, 1)
        w = w.*((v./(w*h))*h')./x2

        # test convergence every 10 iterations

        if i%10 == 0
            j = j + 1

            # adjust small values to avoid undeflow
            h = map(x->max(x, EPS), h)
            w = map(x->max(x, EPS), w)

            # construct connectivity matrix
            index = [i[1] for i in argmax(h, dims=1)]    # find largest factor
            mat1 = repeat(index, m, 1)                   # spread index down
            mat2 = repeat(index', 1, m)                  # spread index right
            cons = mat1==mat2

            if cons == consold             # connectivity matrix has not changed
                inc = inc + 1              # accumulate count
            else
                inc = 0                    # else restart count
            end

            if inc > stopconv
                break                      # assume convergence if connectivity stops changing
            end

            consold = cons

        end
    end

    return w, h
end

@everywhere function to_parallelize(i)

    if i%20 == 0
            @info "Iteration $i."
    end

    # Create a random seed for each iteration
    Random.seed!(i+123456)

    # Generating boostrapped genomes
    bootstrapGenomes = Array{Float64,2}(undef, size(parallelNormGenomes))
    for j = eachindex(parallelGenomesColsSum)
        bootstrapGenomes[:,j] = Random.rand(Multinomial(parallelGenomesColsSum[j], parallelNormGenomes[:,j]))
    end

    bootstrapGenomes = map(x->max(x, EPS), bootstrapGenomes)

    # Solving NMF for these boostrapped genomes
    W, H = nmf(bootstrapGenomes, parallelTotalProcesses)

    for j = 1 : parallelTotalProcesses
        total = sum( W[:, j] )
        W[:, j] = W[:, j] / total
        H[j, :] = H[j, :] * total
    end

    return i, bootstrapGenomes, W, H
end

function extractSignatures(genomes, processes, iterations)

    totalMutationTypes = size(genomes, 1)
    totalGenomes = size(genomes, 2)

    Wall = zeros( totalMutationTypes, processes * iterations )
    Hall = zeros( processes * iterations, totalGenomes )
    genomeErrors = zeros(totalMutationTypes, totalGenomes, iterations)
    genomesReconstructed = zeros(totalMutationTypes, totalGenomes, iterations)

    for (i, bootstrapGenomes, W, H) = pmap(to_parallelize, 1:iterations)
        processCount = (i-1)*processes + 1
        genomeErrors[:, :, i] = bootstrapGenomes -  W * H
        genomesReconstructed[:, :, i] = W * H
        Wall[ :, processCount : (processCount + processes - 1) ] = W
        Hall[ processCount : (processCount + processes - 1), : ] = H

    end

    return Wall, Hall, genomeErrors, genomesReconstructed
end


abstract type AbstractImplementation end
struct MatrixImplementation <: AbstractImplementation end
struct ConvolutionImplementation <: AbstractImplementation end

function optflow(first_img::AbstractArray{T, 2}, second_img::AbstractArray{T,2},  algorithm::Farneback{}) where T <: Gray
    displacement = Array{SVector{2, Float64}, 2}(undef,size(first_img))
    for i in eachindex(displacement)
            displacement[i] = SVector{2, Float64}(0.0, 0.0)
    end
    optflow!(first_img, second_img, displacement,  algorithm)
end

function optflow!(first_img::AbstractArray{T, 2}, second_img::AbstractArray{T,2}, displacement::Array{SVector{2, Float64}, 2}, algorithm::Farneback{}) where T <: Gray

    # Replace NaN with zero in both images.
    first_img = map(x -> isnan(x) ? zero(x) : x, first_img)
    second_img = map(x -> isnan(x) ? zero(x) : x, second_img)

    start_rows = first(first(axes(first_img)))
    start_cols = first(last(axes(first_img)))
    end_rows = last(first(axes(first_img)))
    end_cols = last(last(axes(first_img)))

    A1, B1, _ = polynomial_expansion(ConvolutionImplementation(), first_img, algorithm.expansion_window, algorithm.σ_expansion_window)
    A2, B2, _ = polynomial_expansion(ConvolutionImplementation(), second_img, algorithm.expansion_window, algorithm.σ_expansion_window)

    AtA = zero(similar(A1))
    AtB = zero(similar(B1))

    AtA_smoothed = zero(similar(A1))
    AtB_smoothed = zero(similar(B1))

    neighbourhood = 2*algorithm.estimation_window + 1
    σ, w =  (algorithm.σ_estimation_window,algorithm.σ_estimation_window), (neighbourhood,neighbourhood)
    kern = Kernel.gaussian(σ ,w)

    Aₗ = zeros(Float64,2,2)
    AtAₗ = zeros(Float64,2,2)
    AtBₗ = zeros(Float64,2,1)
    Adₗ = zeros(Float64,2,1)
    bₗ = zeros(Float64,2,1)
    dbₗ = zeros(Float64,2,1)
    dₗ = zeros(Float64,2,1)
    len =  algorithm.expansion_window
    for k = 1:algorithm.iterations
        # We don't utilise the raw polynomial expansion coefficients at the
        # border of the image; hence, we add and subtract `len`.
        for j = start_cols + len:end_cols - len
            for i = start_rows + len:end_rows - len
                d = displacement[i,j]

                mod_i = round(Int, i + d[2])
                mod_j = round(Int, j + d[1])
                mod_i = clamp(mod_i, 1, first(size(first_img)))
                mod_j = clamp(mod_j, 1, last(size(first_img)))

                # (A1[i,j,:,:] + A2[i,j,:,:]) / 2
                for p = 1:2
                    for q = 1:2
                        setindex!(Aₗ, A1[i,j,p,q] / 2 + A2[mod_i,mod_j,p,q] / 2, p, q)
                    end
                end
                # A*d
                setindex!(dₗ, d[1], 1)
                setindex!(dₗ, d[2], 2)
                mul!(Adₗ, Aₗ, dₗ)
                # (B2[i,j,:] - B1[i,j,:]) / 2
                for p = 1:2
                    setindex!(bₗ, (B2[mod_i,mod_j,p] - B1[i,j,p]) / 2 , p)
                end
                # dB = A*d - ((B2 - B1)/2)
                broadcast!(-, dbₗ, Adₗ, bₗ)
                # A'*A
                mul!(AtAₗ,Aₗ',Aₗ)
                # A'*dB
                mul!(AtBₗ,Aₗ',dbₗ)
                # Store result
                for p = 1:2
                    for q = 1:2
                        setindex!(AtA, AtAₗ[p,q], i, j, p, q)
                    end
                    setindex!(AtB, AtBₗ[p], i, j, p)
                end
            end
        end

        for p = 1:2
            for q = 1:2
                AtAview = @view AtA[:,:,p,q]
                AtAview_smoothed = @view AtA_smoothed[:,:,p,q]
                imfilter!(AtAview_smoothed, AtAview, kern, Pad(:replicate))
            end
            AtBview = @view AtB[:,:,p]
            AtBview_smoothed = @view AtB_smoothed[:,:,p]
            imfilter!(AtBview_smoothed, AtBview, kern, Pad(:replicate))
        end

        for j = start_cols:end_cols
            for i = start_rows:end_rows
                P = @view AtA_smoothed[i,j,:,:]
                Q = @view  AtB_smoothed[i,j,:]
                displacement[i,j] = SVector{2, Float64}(pinv2x2(P) * Q)
            end
        end
    end
    # Convert (x,y) coordinate convention to (row, col) convention.
    map!(x -> SVector(last(x), first(x)), displacement, displacement)
    return displacement
end

"""
```
A, B, C = polynomial_expansion(implementation, img, neighbourhood, σ)
```

Returns the polynomial coefficients of the approximation of the neighbourhood of a
point in a 2D function with a polynomial.
The expansion is: f(x) = x'Ax + B'x + C

# Options

Various options for the fields of this type are described in more detail
below.

## Choices for `implementation`

Selects implementation (`::MatrixImplementation` vs `::ConvolutionImplementation`).
The matrix implementation serves as a useful reference against which other
faster implementations can be validated.

## Choices for `img`

Grayscale (Float) image whose polynomial expansion is to be found.

## Choices for `window_size`

Determines the size of the pixel neighbourhood used to find polynomial expansion
for each pixel; larger values mean that the image will be approximated with
smoother surfaces, yielding more robust algorithm and more blurred motion field.
The total size equals `2*window_size + 1`.

## Choices for `σ`

Standard deviation of the Gaussian that is used to smooth the image for the purpose
of approximating it with a polynomial expansion.

## References

Farnebäck, G.: Polynomial Expansion for Orientation and Motion Estimation. PhD thesis, Linköping University,
Sweden, SE-581 83 Linköping, Sweden (2002) Dissertation No 790, ISBN 91-7373-475-6.

"""
function polynomial_expansion(implementation::MatrixImplementation, img::AbstractArray{T, 2}, window_size::Int, σ::Real) where T <: Gray
    len = window_size
    neighbourhood = 2*window_size + 1
    a = ImageFiltering.Kernel.gaussian((σ,σ), (neighbourhood,neighbourhood))[-len:len, -len:len]

    padded_img = parent(padarray(img, Fill(0, (len, len))))

    B0 = zeros(neighbourhood*neighbourhood, 6)
    for x = -len:len
        for y = -len:len
            B0[neighbourhood*(x + len) + (y + len) + 1, :] = [1, x, y, x*x, y*y, x*y]
        end
    end

    B = SMatrix{neighbourhood*neighbourhood, 6}(B0)
    Wa = SDiagonal{neighbourhood^2,Float64}( @view a[:])
    BtWa = B'*Wa

    # If we assume that the confidence is constant we can move several
    # matrix multiplications out of the main loop.
    Wc = SDiagonal(ones(1,neighbourhood^2)...)
    BtWaWc = BtWa*Wc
    G = BtWaWc*B
    G_inv = SMatrix{6,6,Float64}(pinv(convert(Array,G)))
    G_invBtWaWc = G_inv*BtWaWc
    poly_A = zeros(size(img)[1], size(img)[2], 2, 2)
    poly_B = zeros(size(img)[1], size(img)[2], 2)
    poly_C = zeros(size(img))
    f = @MVector zeros(neighbourhood^2)
    for i = (len + 1):(size(padded_img)[1] - len)
        for j = (len + 1):(size(padded_img)[2] - len)
            sub_img = @view padded_img[(i - len):(i + len), (j - len):(j + len)]
            f .= vec(sub_img)
            r = G_invBtWaWc*f

            poly_A[i-len, j-len, 1, 1] =  r[4]
            poly_A[i-len, j-len, 2, 1] =  r[6]/2
            poly_A[i-len, j-len, 1, 2] =  r[6]/2
            poly_A[i-len, j-len, 2, 2] =  r[5]

            poly_B[i-len, j-len, 1] = r[2]
            poly_B[i-len, j-len, 2] = r[3]

            poly_C[i-len, j-len] = r[1]
        end
    end
    return (poly_A, poly_B, poly_C)
end

function polynomial_expansion(imp::ConvolutionImplementation, img::AbstractArray{T, 2}, window_size::Int, σ::Real) where T <: Gray
    #len = floor(Int, neighbourhood/2)
    len = window_size
    neighbourhood = 2*window_size + 1

    # Define basis functions.
    P₂ = construct_basis(len, neighbourhood)

    # Construct Gaussian filter.
    g = collect(KernelFactors.gaussian(σ,neighbourhood))

    # Construct inverse metric using Gaussian "applicability".
    G⁻¹ = construct_inverse_metric(P₂, g)

    unity, x, x², y, xy, y² = calculate_correlations(img, len, g)

    # In-place version of the following code (thus reducing allocations):
    # unity_tmp = G⁻¹[1,1] * unity + G⁻¹[1,4] * x² + G⁻¹[1,5] * y²
    # x = G⁻¹[2,2] * x
    # y = G⁻¹[3,3] * y
    # x² = G⁻¹[4,4] * x² + G⁻¹[4,1] * unity
    # y² = G⁻¹[5,5] * y² + G⁻¹[5,1] * unity
    # xy = G⁻¹[6,6] * xy
    # unity = unity_tmp
    unity_tmp = similar(unity)
    broadcast!(+,unity_tmp, lmul!(G⁻¹[1,1], unity), lmul!(G⁻¹[1,4], x²), lmul!(G⁻¹[1,5] , y²))
    x =  lmul!(G⁻¹[2,2], x)
    y =  lmul!(G⁻¹[3,3], y)
    x² = broadcast!(+,x², lmul!(G⁻¹[4,4]/G⁻¹[1,4], x²), lmul!(G⁻¹[4,1]/G⁻¹[1,1], unity))
    y² = broadcast!(+,y², lmul!(G⁻¹[5,5]/G⁻¹[1,5], y² ), lmul!(G⁻¹[5,1]/G⁻¹[4,1], unity))
    xy = lmul!(G⁻¹[6,6], xy)
    unity = unity_tmp

    # # The result of each multiplication allocated memory which is undesirable.
    # # I've left this version because I am trying to track down an occasional
    # # NaN bug and I want to be sure it is not due to the `in-place` version of
    # # this code.
    # unity_tmp = G⁻¹[1,1] * unity + G⁻¹[1,4] * x² + G⁻¹[1,5] * y²
    # x = G⁻¹[2,2] * x
    # y = G⁻¹[3,3] * y
    # x² = G⁻¹[4,4] * x² + G⁻¹[4,1] * unity
    # y² = G⁻¹[5,5] * y² + G⁻¹[5,1] * unity
    # xy = G⁻¹[6,6] * xy
    # unity = unity_tmp

    poly_A = zeros(Float64,first(size(img)), last(size(img)), 2, 2)
    poly_B = zeros(Float64,first(size(img)), last(size(img)), 2)
    poly_C = unity
    for j =  1:last(size(img))
        for i = 1:first(size(img))
            @inbounds setindex!(poly_A, x²[i,j], i, j, 1, 1)
            @inbounds setindex!(poly_A, xy[i,j]/2, i, j, 1, 2)
            @inbounds setindex!(poly_A, poly_A[i,j,1,2], i, j, 2, 1)
            @inbounds setindex!(poly_A, y²[i,j], i, j, 2, 2)
            @inbounds setindex!(poly_B, x[i,j], i, j, 1)
            @inbounds setindex!(poly_B, y[i,j], i, j, 2)
        end
    end
    return (poly_A, poly_B, poly_C)
end

function construct_basis(len::Int, neighbourhood::Int)
    P₂ = zeros(neighbourhood*neighbourhood, 6)
    for x = -len:len
        for y = -len:len
            P₂[neighbourhood*(x + len) + (y + len) + 1, :] = [1, x, y, x*x, y*y, x*y]
        end
    end
    return P₂
end

function construct_inverse_metric(P₂::AbstractArray, g::AbstractArray)
    applicability = g*g'
    Wa = Matrix(Diagonal(@view applicability[:]))
    P₂_Wa_P₂ = P₂'*Wa*P₂
    G⁻¹ = pinv(P₂_Wa_P₂)
end

function calculate_correlations(img, len, g)
    # Define basis functions.
    l = collect(-len:len)
    lg = l.*g
    l²g = l.*lg

    # Correlations along the first dimensions.
    unity_first_dimension = imfilter(img, ( ImageFiltering.ReshapedOneD{2,0}(centered(g)),), Fill(0))
    y_first_dimension = imfilter(img, ( ImageFiltering.ReshapedOneD{2,0}(centered(lg)),),  Fill(0))
    y²_first_dimension = imfilter(img, ( ImageFiltering.ReshapedOneD{2,0}(centered(l²g)),), Fill(0))

    # Correlations along the second dimensions.
    # Figure 4.1 (Follow correlation structure) 1, y, y², x , xy, x²
    unity = imfilter(unity_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(g)),), Fill(0))
    x =  imfilter(unity_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(lg)),), Fill(0))
    x² =  imfilter(unity_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(l²g)),), Fill(0))

    y =  imfilter(y_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(g)),), Fill(0))
    xy = imfilter(y_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(lg)),), Fill(0))

    y² = imfilter(y²_first_dimension, ( ImageFiltering.ReshapedOneD{2,1}(centered(g)),), Fill(0))

    return unity, x, x², y, xy, y²
end


module QJuliaLanMR

using QJuliaInterface
using QJuliaEnums
using QJuliaBlas
using QJuliaReduce
using QJuliaSolvers

using LinearAlgebra
using MPI

norm2   = QJuliaReduce.gnorm2
rdot    = QJuliaReduce.reDotProduct
verbose = false

const ε = eps(Float64)

@inline function LanczosStep(Mat::Any, vk::AbstractArray, vkm1::AbstractArray, βold::Float64, σ::Float64, pk::AbstractArray)
    Mat(pk, vk)
    pk .=@. pk - βold*vkm1
    # Compute Lancsoz coefficient α
    α = rdot(vk, pk)
    pk .=@. pk - (α+σ)*vk
    # Compute Lancsoz coefficient β
    β = sqrt(rdot(pk, pk))
    if β  < ε ; error("Algorithm breakdown."); end
    vkm1 .=@. vk
    vk .=@. (1.0 / β)*pk
    return  α, β
end

function solver(x::AbstractArray, b::AbstractArray, Mat::Any, MatSloppy::Any, param::QJuliaSolvers.QJuliaSolverParam_qj)

    if verbose == true; println("Running Lanczos version of MR solver in maximum " , param.Nsteps, " step(s)."); end

    QJuliaReduce.set_blas_global_reduction(param.global_reduction)

    if MPI.Initialized() == false; error("MPI was not inititalized, copy source field to the solution."); end

    if (param.maxiter == 0) || (param.Nsteps == 0)
      if param.use_init_guess == false
        x .=@. 0.0
      end
      return
    end #if

    mixed = (param.dtype_sloppy != param.dtype)

    local r     = zero(typeof(x)(undef, length(x)))
    local Av    = zero(typeof(x)(undef, length(x)))
    # explicit Lanczos vectors
    local vk    = zero(typeof(x)(undef, length(x)))
    local vkm1  = zero(typeof(x)(undef, length(x)))
    # explicit Lanczos matrix
    # this is just (sub-)diagonal elements (but transposed when stored!):
    # #(0) β2      ... βk-1  βk
    # α1   α2   α3 ... αk-1  αk
    # β2   β3   β4 ... βk    βk+1
    #
    # R-matrix has the following structure
    # γ1 δ2 ϵ3  0  0  0  ...
    # 0  γ2 δ3 ϵ4  0  0  ...
    # 0  0  γ3 δ4  ϵ5 0  ...
    # 0 ................ γ_k

    local Tk    = zero(Matrix{Float64}(undef, (param.maxiter+1), 2))
    # some aux fields:
    local dk    = zero(typeof(x)(undef, length(x)))
    local dkm1  = zero(typeof(x)(undef, length(x)))

    bnorm = sqrt(norm2(b))  #Save norm of b
    rnorm = 0.0     #if zero source then we will exit immediately doing no work

    σ = param.shift

    if param.use_init_guess == true
      #r = b - Ax0
      Mat(r, x)
      r .=@. r - σ*x
      r .=@. b - r
      rnorm = sqrt(norm2(r))
    else
      rnorm = bnorm
      r .=@. b
    end

    if( rnorm == 0.0 ); return; end
    if verbose == true; println("MR: Initial residual = ", rnorm); end

    # Load init vectors
    vk   .=@. (1.0 / rnorm) * r
    # Givens parameters
    cs = -1.0; sn = 0.0
    # Matrix norm:
    Anorm = 0.0

    # Lancsoz coeffs
    α = 0.0; β = 0.0; βold = rnorm
    # aux parameters
    δk   = 0.0; ϵold = 0.0;
    #start iterations
    for k in 1:param.maxiter
      # Lanczos step:
      α, β = LanczosStep(Mat, vk, vkm1, βold, σ, Av); Tk[k,1] = α; Tk[k,2] = β
      # Givens QR
      # produce middle two enties in the last coloumn of Tk
      δ  = cs*δk + sn*α
      γk = sn*δk - cs*α
      # produce first two entries of T_{k+1}*e_{k+1}
      ϵ  =  sn*β
      δk = -cs*β
      # eliminate β
      γinv  = 1.0 / sqrt(γk*γk+β*β)
      sn    = β  * γinv
      cs    = γk * γinv
      # Solution update (using Av as tmp)
      Av   .=@. γinv*vkm1 - (γinv*δ)*dk - (γinv*ϵold)*dkm1
      dkm1 .=@. dk
      dk   .=@. Av
      #
      τk = rnorm*cs
      x .=@. x + τk*dk
      #residual norm and matrix norm:
      Arnorm = rnorm*sqrt(γk*γk+δk*δk)
      #
      rnorm  *= sn
      #
      Anorm = k == 1 ? sqrt(α*α+β*β) : max(Anorm, sqrt(βold*βold+α*α+β*β))

      if(verbose == true); println("Lanczos MR: ", k,"  residual norm = ", rnorm, "\t ( estimated matrix norm is ", Anorm,  ")"); end
      # Store current β and ϵ for the next cycle
      βold = β; ϵold = ϵ

    end #for itern

    Mat(r, x)

    r .=@. b - r
    rnorm = sqrt(norm2(r))

    param.true_res = rnorm / bnorm
    if verbose == true; println("MR: Relative residual: true = ", param.true_res); end

    QJuliaReduce.reset_blas_global_reduction()

    # Return computed matrix norm
    return Anorm

end #MR

end #QJuliaLanMR


using Pkg
Pkg.add(PackageSpec(name = "RDatasets", version = "0.7.5"))
Pkg.add(PackageSpec(name = "Flux", version = "0.12"))

# Import packages and set a seed
import Flux
using Random
Random.seed!(1);

# Load data and split it into training and test data
function rand_split_data(x, labels)
   nsamples = size(x, 2)
   testidxs = randperm(nsamples)[1:(round(Int, nsamples*0.3))]
   trainidxs = setdiff(1:nsamples, testidxs)
   x_train = x[:, trainidxs]
   x_test = x[:, testidxs]
   labels_train = labels[trainidxs]
   labels_test = labels[testidxs]
   y = Flux.onehotbatch(labels, unique(labels))
   y_train = y[:, trainidxs]
   y_test = y[:, testidxs]
   (training = (x = x_train, y = y_train),
         test = (x = x_test, y = y_test))
end

using RDatasets
import Tables
iris = dataset("datasets", "iris")
x = Tables.matrix(iris[:, 1:4])'
data = rand_split_data(x, iris[:, :Species])
trainingdata = data.training
testdata = data.test

# Load necessary features
import Flux
using Statistics

# Train the Flux model
model = Flux.Chain(
      Flux.Dense(4, 4, Flux.relu),
      Flux.Dense(4, 4, Flux.relu),
      Flux.Dense(4, 3),
      Flux.softmax)

loss(model, x, y) = Flux.crossentropy(model(x), y)
loss(model, data::NamedTuple) = loss(model, data.x, data.y)

function train_network!(model, x, y; epochs, callback)
   opt = Flux.ADAM()
   loss_(x, y) = loss(model, x, y)
   for i in 1:epochs
      Flux.train!(loss_, Flux.params(model), [(x, y)], opt)
      callback(i)
   end
end

epochs = 2500
train_losses = Vector{Float32}(undef, epochs);
test_losses = Vector{Float32}(undef, epochs);
train_network!(model, trainingdata.x, trainingdata.y, epochs = epochs,
      callback = (i) -> begin
         train_losses[i] = loss(model, trainingdata)
         test_losses[i] = loss(model, testdata)
      end)

# (Could plot losses here)

# Evaluate model
accuracy(model, data) =
      mean(Flux.onecold(model(data.x)) .== Flux.onecold(data.y))
accuracy(model, trainingdata)
accuracy(model, testdata)


function simulate!(sim::Simulation)

    T, p = size(sim.y)
    a = size(sim.m, 1)
    d = convert(Int, (a - 1)/2)
    y = zeros(T+d, p)
    F = ones(a)

    for t in 1:T
        if t > d
            # Lagged values
            F[2:end] = vec(sim.y[t-1:-1:t-d, :])
            # Simulate Ω
            sim.Ω[t, :, :] = rand(MatrixNormal(sim.m, sim.P, sim.S))
            # Simulate y
            sim.y[t, :] = rand(MvNormal(vec(F' * sim.Ω[t, :, :]), sim.S))
        else
            sim.Ω[t, :, :] = sim.m
        end
    end
    #return sim
end


using OrdinaryDiffEq

export ode_solver

"""
    ode_solver(f!, k, ts::Vector{Float64}, x0::Vector{Float64}, method; savedt=true)

    Solves differential equations!

    Input
    f!		set of differential equations
    k		reactions rates for the network
    ts		time grid vector
    x0		initial condition
    method	integration algorithm. If not specified: default is Tsit5()

	You can find a list of available solvers at https://diffeq.sciml.ai/stable/solvers/ode_solve/

	savedt	If true, the result will be saved after each time step.
			If false, an adaptive time step is used.
			For method==Euler, this has to be true and dt will be set to the
			discretization in the time grid vector `ts`

    Returns the result in a struct with time and mean copy
    numbers of molecules.

"""
function ode_solver(f!, k, ts::Vector{Float64}, x0::Vector{Float64}, method; savedt=true)

	tspan = (ts[1], ts[end])
	prob = ODEProblem(f!, x0, tspan, k)

	if method==Euler
		println("Euler algoritm needs choice of dt, setting savedt=true")
		delta_t = (ts[2]-ts[1])
		sol = solve(prob, Euler(), dt=delta_t, saveat=delta_t)
	elseif savedt
		sol = solve(prob, method(), saveat=(ts[2]-ts[1]))
	else
		sol = solve(prob, method())
	end

	col = []
	for i in 1:length(x0)
	    push!(col, map(A->A[i], sol.u))
	end

    xar = reshape(cat(col..., dims=1), length(sol.t), length(x0))
    return Result(sol.t, xar')
end

function ode_solver(f!, k, tspan::Tuple, x0)
	println("ode_solver: This is a old version! Please use a vector for tspan
	and a specified method for the integration algorithm (default: Tsit5()). ")
	ts = [tspan[1], tspan[end]]
	ode_solver(f!, k, ts, x0, Tsit5; savedt=false)
end

function ode_solver(f!, k, ts::Vector{Float64}, x0::Vector{Float64})
	println("No integration algorithm specified! Using Tsit5() as default!")
	ode_solver(f!, k, ts, x0, Tsit5)
end


using DiffEqBayesStan
using OrdinaryDiffEq
using ParameterizedFunctions
using RecursiveArrayTools
using StatsPlots, Statistics
using StanSample

ProjDir = @__DIR__

println("Four parameter case (using StanLanguage)")

f1 = @ode_def begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d
u0 = [1.0,1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob1 = ODEProblem(f1,u0,tspan,p)
sol = solve(prob1,Tsit5())
t = collect(range(1,stop=10,length=50))
randomized = VectorOfArray([(sol(t[i]) + .3randn(2)) for i in 1:length(t)])
data_array = convert(Array, randomized)'

stan_pes = "
functions {
  vector sho(real t,vector internal_var___u,
    real internal_var___p_1, real internal_var___p_2,
    real internal_var___p_3, real internal_var___p_4) {

    vector[2] internal_var___du;
    internal_var___du[1] = internal_var___p_1 * internal_var___u[1] -
      internal_var___p_2 * internal_var___u[1] * internal_var___u[2];
    internal_var___du[2] = internal_var___u[2] * -internal_var___p_3 +
      internal_var___p_4 * internal_var___u[1] * internal_var___u[2];

    return internal_var___du;
  }
}
data {
  vector[2] u0;
  int<lower=1> T;
  real internal_var___u[T,2];
  real t0;
  real ts[T];
}
parameters {
  row_vector<lower=0>[2] sigma1;
  real<lower=0.5,upper=2.0> theta_1;real<lower=0.5,upper=1.5> theta_2;real<lower=0.5,upper=4.0> theta_3;real<lower=0.5,upper=2.0> theta_4;
}
model{
  vector[2] u_hat[T];
  sigma1 ~ inv_gamma(4.0, 1.0);
  theta_1 ~ normal(1.1, 0.5) T[0.5,2.0];
  theta_2 ~ normal(1.5, 0.5) T[0.5,1.5];
  theta_3 ~ normal(2.0, 0.5) T[0.5,4.0];
  theta_4 ~ normal(1.5, 0.5) T[0.5,2.0];
  u_hat = ode_rk45_tol(sho, u0, t0, ts, 0.001, 1.0e-6, 100000,
    theta_1,theta_2,theta_3,theta_4);
  for (t in 1:T){
    internal_var___u[t,:] ~ normal(u_hat[t,1:2],sigma1);
    }
}";

pes4 = SampleModel("pes4", stan_pes, joinpath(ProjDir, "tmp"))
data = (u0=u0, T=size(t, 1), internal_var___u=data_array, t0=0.0, ts=t)
rc = stan_sample(pes4; data)

if success(rc)
  df = read_samples(pes4, :dataframe)

  println("\ntheta means:\n",
    [mean(df.theta_1), mean(df.theta_2), mean(df.theta_3), mean(df.theta_4)])
  println("\ntheta vars:\n",
    [var(df.theta_1), var(df.theta_2), var(df.theta_3), var(df.theta_4)])
  println()

  dfs = read_samples(pes4, :dataframes)
  for j in 1:4
    for i in 1:4
      if i == 1
        plot(dfs[i][:, "theta_$j"], lab="chain[$j]")
      else
        plot!(dfs[i][:, "theta_$j"], lab="chain[$j]")
      end
    end
    savefig(joinpath(ProjDir, "tmp", "theta_$j.png"))
  end

end


# ADD TESTS SET FOR INVERSE
#Test library imports
using Test

#Arithmetic library imports
using IntervalArithmetic
using IntervalContractors

#Preamble
setprecision(53)
setprecision(Interval, Float64)

# Using approximate checks for validaty update later?
import Base.isapprox
isapprox(x::Interval,y::Interval) = isapprox(x.lo,y.lo,atol=1E-4) && isapprox(x.hi,y.hi,atol=1E-4)

@testset "inv_rev_test" begin
    @test inv_rev(∅, entireinterval(Float64))[2] == ∅
    @test inv_rev(Interval(0.0, 1.0), ∅)[2] == ∅
    @test inv_rev(∅, Interval(0.0, 1.0))[2] == ∅
    @test inv_rev(Interval(-2.0, -1.0), entireinterval(Float64))[2] == Interval(-1.0, -0.5)
    @test inv_rev(Interval(1.0, 1.0), entireinterval(Float64))[2] == Interval(1.0, 1.0)
    @test inv_rev(entireinterval(Float64), entireinterval(Float64))[2] == Interval(-∞, ∞)
    @test inv_rev(Interval(-Inf, 0.0), entireinterval(Float64))[2] == Interval(-∞, 0.0)
end