using BlockArrays, SparseArrays, Setfield

# structure to describe a (Time) mesh using the time steps t_{i+1} - t_{i}. If the time steps are constant, we do not record them but, instead, we save the number of time steps
struct TimeMesh{T}
	ds::T
end

TimeMesh(M::Int64) = TimeMesh{Int64}(M)

@inline canAdapt(ms::TimeMesh{Ti}) where Ti = !(Ti == Int64)
Base.length(ms::TimeMesh{Ti}) where Ti = length(ms.ds)
Base.length(ms::TimeMesh{Ti}) where {Ti <: Int} = ms.ds

# access the time steps
@inline getTimeStep(ms, i::Int) = ms.ds[i]
@inline getTimeStep(ms::TimeMesh{Ti}, i::Int) where {Ti <: Int} = 1.0 / ms.ds

Base.collect(ms::TimeMesh) = ms.ds
Base.collect(ms::TimeMesh{Ti}) where {Ti <: Int} = repeat([getTimeStep(ms, 1)], ms.ds)
####################################################################################################
# method using the Trapezoidal rule (Order 2 in time) and discretisation of the periodic orbit.
"""

This composite type implements Finite Differences based on a Trapezoidal rule to locate periodic orbits. More details (maths, notations, linear systems) can be found [here](https://bifurcationkit.github.io/BifurcationKitDocs.jl/dev/periodicOrbitTrapeze/).

## Arguments
- `F` vector field specified as a function of two arguments `F(x,p)`
- `J` is the jacobian of `F` at `(x, p)`. It can assume three forms.
    1. Either `J` is a function and `J(x, p)` returns a `::AbstractMatrix`. In this case, the default arguments of `contParams::ContinuationPar` will make `continuation` work.
    2. Or `J` is a function and `J(x, p)` returns a function taking one argument `dx` and returning `dr` of the same type as `dx`. In our notation, `dr = J * dx`. In this case, the default parameters of `contParams::ContinuationPar` will not work and you have to use a Matrix Free linear solver, for example `GMRESIterativeSolvers`,
    3. Or `J` is a function and `J(x, p)` returns a variable `j` which can assume any type. Then, you must implement a linear solver `ls` as a composite type, subtype of `AbstractLinearSolver` which is called like `ls(j, rhs)` and which returns the solution of the jacobian linear system. See for example `examples/SH2d-fronts-cuda.jl`. This linear solver is passed to `NewtonPar(linsolver = ls)` which itself passed to `ContinuationPar`. Similarly, you have to implement an eigensolver `eig` as a composite type, subtype of `AbstractEigenSolver`.
- `Jᵗ = nothing` jacobian transpose of `F` (optional), useful for continuation of Fold of periodic orbits. it should not be passed in case the jacobian is a (sparse) matrix as it is computed internally, and it would be computed twice in that case.
- `d2F = nothing` second derivative of F (optional), useful for continuation of Fold of periodic orbits. It has the definition `d2F(x,p,dx1,dx2)`.
- `M::Int` number of time slices
- `ϕ` used to set a section for the phase constraint equation, of size N*M
- `xπ` used in the section for the phase constraint equation, of size N*M
- `linsolver: = DefaultLS()` linear solver for each time slice, i.e. to solve `J⋅sol = rhs`. This is only needed for the computation of the Floquet multipliers.
- `isinplace::Bool` whether `F` and `J` are inplace functions (Experimental). In this case, the functions `F` and `J` must have the following definitions `(o, x, p) ->  F(o, x, p)` and `(o, x, p, dx) -> J(o, x, p, dx)`.
- `ongpu::Bool` whether the computation takes place on the gpu (Experimental)
- `massmatrix` a mass matrix. You can pass for example a sparse matrix. Default: identity matrix.

The scheme is as follows. We first consider a partition of ``[0,1]`` given by ``0<s_0<\\cdots<s_m=1`` and one looks for `T = x[end]` such that

 ``M_a\\cdot\\left(x_{i} - x_{i-1}\\right) - \\frac{T\\cdot h_i}{2} \\left(F(x_{i}) + F(x_{i-1})\\right) = 0,\\ i=1,\\cdots,m-1``

with ``u_{0} := u_{m-1}`` and the periodicity condition ``u_{m} - u_{1} = 0`` and

where ``h_1 = s_i-s_{i-1}``. ``M_a`` is a mass matrix. Finally, the phase of the periodic orbit is constrained by using a section (but you could use your own)

 ``\\sum_i\\langle x_{i} - x_{\\pi,i}, \\phi_{i}\\rangle=0.``

# Orbit guess
You will see below that you can evaluate the residual of the functional (and other things) by calling `pb(orbitguess, p)` on an orbit guess `orbitguess`. Note that `orbitguess` must be of size M * N + 1 where N is the number of unknowns in the state space and `orbitguess[M*N+1]` is an estimate of the period ``T`` of the limit cycle. More precisely, using the above notations, `orbitguess` must be ``orbitguess = [x_{1},x_{2},\\cdots,x_{M}, T]``.

Note that you can generate this guess from a function solution using `generateSolution`.

# Functional
 A functional, hereby called `G`, encodes this problem. The following methods are available

- `pb(orbitguess, p)` evaluates the functional G on `orbitguess`
- `pb(orbitguess, p, du)` evaluates the jacobian `dG(orbitguess).du` functional at `orbitguess` on `du`
- `pb(Val(:JacFullSparse), orbitguess, p)` return the sparse matrix of the jacobian `dG(orbitguess)` at `orbitguess` without the constraints. It is called `A_γ` in the docs.
- `pb(Val(:JacFullSparseInplace), J, orbitguess, p)`. Same as `pb(Val(:JacFullSparse), orbitguess, p)` but overwrites `J` inplace. Note that the sparsity pattern must be the same independantly of the values of the parameters or of `orbitguess`. In this case, this is significantly faster than `pb(Val(:JacFullSparse), orbitguess, p)`.
- `pb(Val(:JacCyclicSparse), orbitguess, p)` return the sparse cyclic matrix Jc (see the docs) of the jacobian `dG(orbitguess)` at `orbitguess`
- `pb(Val(:BlockDiagSparse), orbitguess, p)` return the diagonal of the sparse matrix of the jacobian `dG(orbitguess)` at `orbitguess`. This allows to design Jacobi preconditioner. Use `blockdiag`.

!!! note "GPU call"
    For these methods to work on the GPU, for example with `CuArrays` in mode `allowscalar(false)`, we face the issue that the function `extractPeriodFDTrap` won't be well defined because it is a scalar operation. Note that you must pass the option `ongpu = true` for the functional to be evaluated efficiently on the gpu.
"""
@with_kw_noshow struct PeriodicOrbitTrapProblem{TF, TJ, TJt, Td2F, Td3F, vectype, Tls <: AbstractLinearSolver, Tmesh, Tmass} <: AbstractPOFDProblem
	# Function F(x, par)
	F::TF = nothing

	# Jacobian of F w.r.t. x
	J::TJ = nothing

	# Jacobian transpose of F w.r.t. x. This is mainly used for matrix-free computation of Folds of limit cycles
	Jᵗ::TJt = nothing

	# Hessian of F w.r.t. x, useful for continuation of Fold of periodic orbits
	d2F::Td2F = nothing

	# 3rd differential of F w.r.t. x, useful for branch switching from branch of periodic orbits
	d3F::Td3F = nothing

	# variables to define a Section for the phase constraint equation
	ϕ::vectype = nothing
	xπ::vectype = nothing

	# discretisation of the time interval
	M::Int = 0
	mesh::Tmesh = TimeMesh(M)

	# dimension of the problem in case of an AbstractVector
	N::Int = 0

	# linear solver for each slice, i.e. to solve J⋅sol = rhs. This is mainly used for the computation of the Floquet coefficients
	linsolver::Tls = DefaultLS()

	# whether F and J are inplace functions
	isinplace::Bool = false

	# whether the computation takes place on the gpu
	ongpu::Bool = false

	# whether the time discretisation is adaptive
	adaptmesh::Bool = false

	# whether the problem is nonautonomous
	isautonomous::Bool = true

	# mass matrix
	massmatrix::Tmass = nothing
end

function Base.show(io::IO, pb::PeriodicOrbitTrapProblem)
	println(io, "┌─ Trapezoid problem")
	println(io, "├─ time slices  : ", pb.M)
	println(io, "├─ dimension    : ", pb.N)
	println(io, "└─ inplace      : ", pb.isinplace)
end

@inline getTimeStep(pb::AbstractPOFDProblem, i::Int) = getTimeStep(pb.mesh, i)
getTimes(pb::AbstractPOFDProblem) = cumsum(collect(pb.mesh))
@inline hasmassmatrix(pb::PeriodicOrbitTrapProblem) = ~isnothing(pb.massmatrix)
@inline function getMassMatrix(pb::PeriodicOrbitTrapProblem, returnArray = false)
	if returnArray == false
		return hasmassmatrix(pb) ? pb.massmatrix : spdiagm( 0 => ones(pb.N))
	else
		return hasmassmatrix(pb) ? pb.massmatrix : LinearAlgebra.I(pb.N)
	end
end

# for a dummy constructor, useful for specifying the "algorithm" to look for periodic orbits,
# just call PeriodicOrbitTrapProblem()

function PeriodicOrbitTrapProblem(F, J, d2F, ϕ::vectype, xπ::vectype, m::Union{Int, vecmesh}, ls::AbstractLinearSolver = DefaultLS(); isinplace = false, ongpu = false, adaptmesh = false, massmatrix = nothing) where {vectype, vecmesh <: AbstractVector}
	_length = ϕ isa AbstractVector ? length(ϕ) : 0
	M = m isa Number ? m : length(m) + 1

	return PeriodicOrbitTrapProblem(F = F, J = J, d2F = d2F, ϕ = ϕ, xπ = xπ, M = M, mesh = TimeMesh(m), N = _length ÷ M, linsolver = ls, isinplace = isinplace, ongpu = ongpu, adaptmesh = adaptmesh, massmatrix = massmatrix)
end

PeriodicOrbitTrapProblem(F, J, ϕ::vectype, xπ::vectype, m::Union{Int, vecmesh}, ls::AbstractLinearSolver = DefaultLS(); isinplace = false, ongpu = false, adaptmesh = false, massmatrix = nothing) where {vectype, vecmesh <: AbstractVector} = PeriodicOrbitTrapProblem(F, J, nothing, ϕ, xπ, m, ls; isinplace = isinplace, ongpu = ongpu, adaptmesh = adaptmesh, massmatrix = massmatrix)

function PeriodicOrbitTrapProblem(F, J, d2F, ϕ::vectype, xπ::vectype, m::Union{Int, vecmesh}, N::Int, ls::AbstractLinearSolver = DefaultLS(); isinplace = false, ongpu = false, adaptmesh = false, massmatrix = nothing) where {vectype, vecmesh <: AbstractVector}
	M = m isa Number ? m : length(m) + 1
	# we use 0 * ϕ to create a copy filled with zeros, this is useful to keep the types
	prob = PeriodicOrbitTrapProblem(F = F, J = J, d2F = d2F, ϕ = similar(ϕ, N*M), xπ = similar(xπ, N*M), M = M, mesh = TimeMesh(m), N = N, linsolver = ls, isinplace = isinplace, ongpu = ongpu, adaptmesh = adaptmesh, massmatrix = massmatrix)

	prob.xπ .= 0
	prob.ϕ .= 0

	prob.xπ[1:length(xπ)] .= xπ
	prob.ϕ[1:length(ϕ)] .= ϕ
	return prob
end

PeriodicOrbitTrapProblem(F, J, ϕ::vectype, xπ::vectype, m::Union{Int, vecmesh}, N::Int, ls::AbstractLinearSolver = DefaultLS(); isinplace = false, ongpu = false, adaptmesh = false, massmatrix = nothing) where {vectype, vecmesh <: AbstractVector} = PeriodicOrbitTrapProblem(F, J, nothing, ϕ, xπ, m, N, ls; isinplace = isinplace, ongpu = ongpu, adaptmesh = adaptmesh, massmatrix = massmatrix)

PeriodicOrbitTrapProblem(F, J, m::Union{Int, vecmesh}, N::Int, ls::AbstractLinearSolver = DefaultLS(); isinplace = false, ongpu = false, adaptmesh = false, massmatrix = nothing) where {vectype, vecmesh <: AbstractVector} = PeriodicOrbitTrapProblem(F, J, nothing, zeros(N*(m isa Number ? m : length(m) + 1)), zeros(N*(m isa Number ? m : length(m) + 1)), m, N, ls; isinplace = isinplace, ongpu = ongpu, adaptmesh = adaptmesh, massmatrix = massmatrix)

# these functions extract the last component of the periodic orbit guess
@inline extractPeriodFDTrap(pb::PeriodicOrbitTrapProblem, x::AbstractVector) = onGpu(pb) ? x[end:end] : x[end]
# @inline extractPeriodFDTrap(pb, x::BorderedArray)  = x.T

# these functions extract the time slices components
getTimeSlices(x::AbstractVector, N, M) = @views reshape(x[1:end-1], N, M)
# getTimeSlices(x::BorderedArray,  N, M) = x.u
getTimeSlices(pb::PeriodicOrbitTrapProblem, x) = getTimeSlices(x, pb.N, pb.M)

# do not type h::Number because this will imped using CUDA
function POTrapScheme!(pb::AbstractPOFDProblem, dest, u1, u2, du1, du2, par, h, tmp, linear::Bool = true; applyf::Bool = true)
	# this function implements the basic implicit scheme used for the time integration
	# because this function is called in a cyclic manner, we save in the variable tmp the value of F(u2) in order to avoid recomputing it in a subsequent call
	# basically tmp is F(u2)
	if linear
		dest .= tmp
		if applyf
			# tmp <- pb.F(u1, par)
			applyF(pb, tmp, u1, par) #TODO this line does not almost seem to be type stable in code_wartype, gives @_11::Union{Nothing, Tuple{Int64,Int64}}
		else
			applyJ(pb, tmp, u1, par, du1)
		end
		if hasmassmatrix(pb)
			dest .= pb.massmatrix * (du1 .- du2) .- h .* (dest .+ tmp)
		else
			dest .= @. (du1 - du2) - h * (dest + tmp)
		end
	else
		dest .-= h .* tmp
		# tmp <- pb.F(u1, par)
		applyF(pb, tmp, u1, par)
		dest .-= h .* tmp
	end
end
POTrapScheme!(pb::AbstractPOFDProblem, dest, u1, u2, par, h, tmp, linear::Bool = true; applyf::Bool = true) = POTrapScheme!(pb::AbstractPOFDProblem, dest, u1, u2, u1, u2, par, h, tmp, linear; applyf = applyf)

"""
This function implements the functional for finding periodic orbits based on finite differences using the Trapezoidal rule. It works for inplace / out of place vector fields `pb.F`
"""
function POTrapFunctional!(pb::AbstractPOFDProblem, out, u, par)
		M, N = size(pb)
		T = extractPeriodFDTrap(pb, u)

		uc = getTimeSlices(pb, u)
		outc = getTimeSlices(pb, out)

		# outc[:, M] plays the role of tmp until it is used just after the for-loop
		@views applyF(pb, outc[:, M], uc[:, M-1], par)

		h = T * getTimeStep(pb, 1)
		# https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-column-major
		# fastest is to do out[:, i] = x
		@views POTrapScheme!(pb, outc[:, 1], uc[:, 1], uc[:, M-1], par, h/2, outc[:, M])

		for ii in 2:M-1
			h = T * getTimeStep(pb, ii)
			# this function avoids computing F(uc[:, ii]) twice
			@views POTrapScheme!(pb, outc[:, ii], uc[:, ii], uc[:, ii-1], par, h/2, outc[:, M])
		end

		# closure condition ensuring a periodic orbit
		outc[:, M] .= @views uc[:, M] .- uc[:, 1]

		# this is for CuArrays.jl to work in the mode allowscalar(false)
		if onGpu(pb)
			return @views vcat(out[1:end-1], dot(u[1:end-1], pb.ϕ) - dot(pb.xπ, pb.ϕ)) # this is the phase condition
		else
			out[end] = @views dot(u[1:end-1], pb.ϕ) - dot(pb.xπ, pb.ϕ) #dot(u0c[:, 1] .- pb.xπ, pb.ϕ)
			return out
		end
end

"""
Matrix free expression of the Jacobian of the problem for computing periodic obits when evaluated at `u` and applied to `du`.
"""
function POTrapFunctionalJac!(pb::AbstractPOFDProblem, out, u, par, du)
	M, N = size(pb)
	T  = extractPeriodFDTrap(pb, u)
	dT = extractPeriodFDTrap(pb, du)

	uc = getTimeSlices(pb, u)
	outc = getTimeSlices(pb, out)
	duc = getTimeSlices(pb, du)

	# compute the cyclic part
	@views Jc(pb, outc, u[1:end-1-N], par, T, du[1:end-N-1], outc[:, M])

	# outc[:, M] plays the role of tmp until it is used just after the for-loop
	tmp = @view outc[:, M]

	# we now compute the partial derivative w.r.t. the period T
	@views applyF(pb, tmp, uc[:, M-1], par)

	h = dT * getTimeStep(pb, 1)
	@views POTrapScheme!(pb, outc[:, 1], uc[:, 1], uc[:, M-1], par, h/2, tmp, false)
	for ii in 2:M-1
		h = dT * getTimeStep(pb, ii)
		@views POTrapScheme!(pb, outc[:, ii], uc[:, ii], uc[:, ii-1], par, h/2, tmp, false)
	end

	# closure condition ensuring a periodic orbit
	outc[:, M] .= @views duc[:, M] .- duc[:, 1]

	# this is for CuArrays.jl to work in the mode allowscalar(false)
	if onGpu(pb)
		return @views vcat(out[1:end-1], dot(du[1:end-1], pb.ϕ))
	else
		out[end] = @views dot(du[1:end-1], pb.ϕ)
		return out
	end
end

(pb::PeriodicOrbitTrapProblem)(u::AbstractVector, par) = POTrapFunctional!(pb, similar(u), u, par)
(pb::PeriodicOrbitTrapProblem)(u::AbstractVector, par, du) = POTrapFunctionalJac!(pb, similar(du), u, par, du)

####################################################################################################
# Matrix free expression of matrices related to the Jacobian Matrix of the PO functional
"""
Function to compute the Matrix-Free version of Aγ, see docs for its expression.
"""
function Aγ!(pb::PeriodicOrbitTrapProblem, outc, u0::AbstractVector, par, du::AbstractVector; γ = 1)
	# u0 of size N * M + 1
	# du of size N * M
	M, N = size(pb)
	T = extractPeriodFDTrap(pb, u0)
	u0c = getTimeSlices(pb, u0)

	# compute the cyclic part
	@views Jc(pb, outc, u0[1:end-1-N], par, T, du[1:end-N], outc[:, M])

	# closure condition ensuring a periodic orbit
	duc = reshape(du, N, M)
	outc[:, M] .= @views duc[:, M] .- γ .* duc[:, 1]
	return nothing
end

"""
Function to compute the Matrix-Free version of the cyclic matrix Jc, see docs for its expression.
"""
function Jc(pb::PeriodicOrbitTrapProblem, outc::AbstractMatrix, u0::AbstractVector, par, T, du::AbstractVector, tmp)
	# tmp plays the role of buffer array
	# u0 of size N * (M - 1)
	# du of size N * (M - 1)
	# outc of size N * M
	M, N = size(pb)

	u0c = reshape(u0, N, M-1)
	duc = reshape(du, N, M-1)

	@views applyJ(pb, tmp, u0c[:, M-1], par, duc[:, M-1])

	h = T * getTimeStep(pb, 1)
	@views POTrapScheme!(pb, outc[:, 1], u0c[:, 1], u0c[:, M-1],
										 duc[:, 1], duc[:, M-1], par, h/2, tmp, true; applyf = false)

	for ii in 2:M-1
		h = T * getTimeStep(pb, ii)
		@views POTrapScheme!(pb, outc[:, ii], u0c[:, ii], u0c[:, ii-1],
											  duc[:, ii], duc[:, ii-1], par, h/2, tmp, true; applyf = false)
	end

	# we also return a Vector version of outc
	return vec(outc)
end

function Jc(pb::PeriodicOrbitTrapProblem, u0::AbstractVector, par, du::AbstractVector)
	M, N = size(pb)
	T = extractPeriodFDTrap(pb, u0)

	out  = similar(du)
	outc = reshape(out, N, M-1)
	tmp  = similar(view(outc, :, 1))
	return @views Jc(pb, outc, u0[1:end-1-N], par, T, du, tmp)
end
####################################################################################################
"""
Matrix by blocks expression of the Jacobian for the PO functional computed at the space-time guess: `u0`
"""
function jacobianPOTrapBlock(pb::PeriodicOrbitTrapProblem, u0::AbstractVector, par; γ = 1)
	# extraction of various constants
	M, N = size(pb)

	Aγ = BlockArray(spzeros(M * N, M * N), N * ones(Int64, M),  N * ones(Int64, M))
	cylicPOTrapBlock!(pb, u0, par, Aγ)

	In = spdiagm( 0 => ones(N))
	Aγ[Block(M, 1)] = -γ * In
	Aγ[Block(M, M)] = In
	return Aγ
end

"""
This function populates Jc with the cyclic matrix using the different Jacobians
"""
function cylicPOTrapBlock!(pb::PeriodicOrbitTrapProblem, u0::AbstractVector, par, Jc::BlockArray)
	# extraction of various constants
	M, N = size(pb)
	T = extractPeriodFDTrap(pb, u0)

	In = getMassMatrix(pb)

	u0c = getTimeSlices(pb, u0)
	outc = similar(u0c)

	tmpJ = @views pb.J(u0c[:, 1], par)

	h = T * getTimeStep(pb, 1)
	Jn = In - (h/2) .* tmpJ
	Jc[Block(1, 1)] = Jn

	# we could do a Jn .= -I .- ... but we want to allow the sparsity pattern to vary
	Jn = @views -In - (h/2) .* pb.J(u0c[:, M-1], par)
	Jc[Block(1, M-1)] = Jn

	for ii in 2:M-1
		h = T * getTimeStep(pb, ii)
		Jn = -In - (h/2) .* tmpJ
		Jc[Block(ii, ii-1)] = Jn

		tmpJ = @views pb.J(u0c[:, ii], par)

		Jn = In - (h/2) .* tmpJ
		Jc[Block(ii, ii)] = Jn
	end
	return Jc
end

function cylicPOTrapBlock(pb::PeriodicOrbitTrapProblem, u0::AbstractVector, par)
	# extraction of various constants
	M, N = size(pb)
	Jc = BlockArray(spzeros((M - 1) * N, (M - 1) * N), N * ones(Int64, M-1),  N * ones(Int64, M-1))
	cylicPOTrapBlock!(pb, u0, par, Jc)
end

cylicPOTrapSparse(pb::PeriodicOrbitTrapProblem, orbitguess0, par) = blockToSparse(cylicPOTrapBlock(pb, orbitguess0, par))

"""
This method returns the jacobian of the functional G encoded in PeriodicOrbitTrapProblem using a Sparse representation.
"""
function (pb::PeriodicOrbitTrapProblem)(::Val{:JacFullSparse}, u0::AbstractVector, par; γ = 1, δ = convert(eltype(u0), 1e-9))
	# extraction of various constants
	M, N = size(pb)
	T = extractPeriodFDTrap(pb, u0)
	AγBlock = jacobianPOTrapBlock(pb, u0, par; γ = γ)

	# we now set up the last line / column
	@views ∂TGpo = (pb(vcat(u0[1:end-1], T + δ), par) .- pb(u0, par)) ./ δ

	# this is "bad" for performance. Get converted to SparseMatrix at the next line
	Aγ = blockToSparse(AγBlock) # most of the computing time is here!!
	@views Aγ = hcat(Aγ, ∂TGpo[1:end-1])
	Aγ = vcat(Aγ, spzeros(1, N * M + 1))

	Aγ[N*M+1, 1:length(pb.ϕ)] .=  pb.ϕ
	Aγ[N*M+1, N*M+1] = ∂TGpo[end]
	return Aγ
end

"""
This method returns the jacobian of the functional G encoded in PeriodicOrbitTrapProblem using an inplace update. In case where the passed matrix J0 is a sparse one, it updates J0 inplace assuming that the sparsity pattern of J0 and dG(orbitguess0) are the same.
"""
@views function (pb::PeriodicOrbitTrapProblem)(::Val{:JacFullSparseInplace}, J0::Tj, u0::AbstractVector, par; γ = 1, δ = convert(eltype(u0), 1e-9)) where Tj
		M, N = size(pb)
		T = extractPeriodFDTrap(pb, u0)

		In = getMassMatrix(pb, ~(Tj <: SparseMatrixCSC))

		u0c = getTimeSlices(pb, u0)
		outc = similar(u0c)

		tmpJ = pb.J(u0c[:, 1], par)

		h = T * getTimeStep(pb, 1)
		Jn = In - (h/2) .* tmpJ
		# setblock!(Jc, Jn, 1, 1)
		J0[1:N, 1:N] .= Jn

		Jn .= -In .- (h/2) .* pb.J(u0c[:, M-1], par)
		# setblock!(Jc, Jn, 1, M-1)
		J0[1:N, (M-2)*N+1:(M-1)*N] .= Jn

		for ii in 2:M-1
			h = T * getTimeStep(pb, ii)
			@. Jn = -In - h/2 * tmpJ
			# the next lines cost the most
			# setblock!(Jc, Jn, ii, ii-1)
			J0[(ii-1)*N+1:(ii)*N, (ii-2)*N+1:(ii-1)*N] .= Jn

			tmpJ .= pb.J(u0c[:, ii], par)

			@. Jn = In - h/2 * tmpJ
			# setblock!(Jc, Jn, ii, ii)
			J0[(ii-1)*N+1:(ii)*N, (ii-1)*N+1:(ii)*N] .= Jn
		end

		# setblock!(Aγ, -γ * In, M, 1)
		# useless to update:
			# J0[(M-1)*N+1:(M)*N, (1-1)*N+1:(1)*N] .= -In
		# setblock!(Aγ,  In,     M, M)
		# useless to update:
			# J0[(M-1)*N+1:(M)*N, (M-1)*N+1:(M)*N] .= In

		# we now set up the last line / column
		∂TGpo = (pb(vcat(u0[1:end-1], T + δ), par) .- pb(u0, par)) ./ δ
		J0[:, end] .=  ∂TGpo

		# this following does not depend on u0, so it does not change. However we update it in case the caller updated the section somewhere else
		J0[N*M+1, 1:length(pb.ϕ)] .=  pb.ϕ

		return J0
end

@views function (pb::PeriodicOrbitTrapProblem)(::Val{:JacFullSparseInplace}, J0, u0::AbstractVector, par, indx; γ = 1, δ = convert(eltype(u0), 1e-9), updateborder = true)
	M, N = size(pb)
	T = extractPeriodFDTrap(pb, u0)

	In = getMassMatrix(pb)

	u0c = getTimeSlices(pb, u0)
	outc = similar(u0c)

	tmpJ = pb.J(u0c[:, 1], par)

	h = T * getTimeStep(pb, 1)
	Jn = In - tmpJ * (h/2)

	# setblock!(Jc, Jn, 1, 1)
	J0.nzval[indx[1, 1]] .= Jn.nzval

	Jn .= -In .- pb.J(u0c[:, M-1], par) .* (h/2)
	# setblock!(Jc, Jn, 1, M-1)
	J0.nzval[indx[1, M-1]] .= Jn.nzval

	for ii in 2:M-1
		h = T * getTimeStep(pb, ii)
		@. Jn = -In - tmpJ * (h/2)
		# the next lines cost the most
		# setblock!(Jc, Jn, ii, ii-1)
		J0.nzval[indx[ii, ii-1]] .= Jn.nzval

		tmpJ .= pb.J(u0c[:, ii], par)# * (h/2)

		@. Jn = In -  tmpJ * (h/2)
		# setblock!(Jc, Jn, ii, ii)
		J0.nzval[indx[ii,ii]] .= Jn.nzval
	end

	# setblock!(Aγ, -γ * In, M, 1)
	# useless to update:
		# J0[(M-1)*N+1:(M)*N, (1-1)*N+1:(1)*N] .= -In
	# setblock!(Aγ,  In,     M, M)
	# useless to update:
		# J0[(M-1)*N+1:(M)*N, (M-1)*N+1:(M)*N] .= In

	if updateborder
		# we now set up the last line / column
		∂TGpo = (pb(vcat(u0[1:end-1], T + δ), par) .- pb(u0, par)) ./ δ
		J0[:, end] .=  ∂TGpo

		# this following does not depend on u0, so it does not change. However we update it in case the caller updated the section somewhere else
		J0[N*M+1, 1:length(pb.ϕ)] .=  pb.ϕ
	end

	return J0
end

function (pb::PeriodicOrbitTrapProblem)(::Val{:JacCyclicSparse}, u0::AbstractVector, par, γ = 1)
	# extraction of various constants
	N = pb.N
	AγBlock = jacobianPOTrapBlock(pb, u0, par; γ = γ)

	# this is bad for performance. Get converted to SparseMatrix at the next line
	Aγ = blockToSparse(AγBlock) # most of the computing time is here!!
	# the following line is bad but still less costly than the previous one
	return Aγ[1:end-N, 1:end-N]
end

function (pb::PeriodicOrbitTrapProblem)(::Val{:BlockDiagSparse}, u0::AbstractVector, par)
	# extraction of various constants
	M, N = size(pb)
	T = extractPeriodFDTrap(pb, u0)

	A_diagBlock = BlockArray(spzeros(M * N, M * N), N * ones(Int64, M),  N * ones(Int64, M))

	In = getMassMatrix(pb)

	u0c = reshape(u0[1:end-1], N, M)
	outc = similar(u0c)

	h = T * getTimeStep(pb, 1)
	@views Jn = In - h/2 .* pb.J(u0c[:, 1], par)
	A_diagBlock[Block(1, 1)] = Jn

	for ii in 2:M-1
		h = T * getTimeStep(pb, ii)
		@views Jn = In - h/2 .* pb.J(u0c[:, ii], par)
		A_diagBlock[Block(ii, ii)]= Jn
	end
	A_diagBlock[Block(M, M)]= In

	A_diag_sp = blockToSparse(A_diagBlock) # most of the computing time is here!!
	return A_diag_sp
end
####################################################################################################
# Utils
"""
$(SIGNATURES)

Compute the full periodic orbit associated to `x`. Mainly for plotting purposes.
"""
@views function getPeriodicOrbit(prob::AbstractPOFDProblem, u::AbstractVector, p)
	T = getPeriod(prob, u, p)
	M, N = size(prob)
	uv = u[1:end-1]
	uc = reshape(uv, N, M)
	return SolPeriodicOrbit(t = cumsum(T .* collect(prob.mesh)), u = uc)
end

"""
$(SIGNATURES)

Compute the period of the periodic orbit associated to `x`.
"""
@inline getPeriod(prob::PeriodicOrbitTrapProblem, x, p) = extractPeriodFDTrap(prob, x)

"""
$(SIGNATURES)

Compute `norm(du/dt)`
"""
@views function getTimeDiff(pb::PeriodicOrbitTrapProblem, u)
	M, N = size(pb)
	T = extractPeriodFDTrap(pb, u)
	uc = reshape(u[1:end-1], N, M)
	return [norm(uc[:,ii+1].-uc[:,ii]) * T/M for ii in 1:M-1]
end

"""
$(SIGNATURES)

Compute the amplitude of the periodic orbit associated to `x`. The keyword argument `ratio = 1` is used as follows. If `length(x) = 1 + ratio * n`, the call returns the amplitude over `x[1:n]`.
"""
@views function getAmplitude(prob::PeriodicOrbitTrapProblem, x::AbstractVector, p; ratio = 1)
	n = div(length(x)-1, ratio)
	_max = maximum(x[1:n])
	_min = minimum(x[1:n])
	return maximum(_max .- _min)
end

"""
$(SIGNATURES)

Compute the maximum of the periodic orbit associated to `x`. The keyword argument `ratio = 1` is used as follows. If `length(x) = 1 + ratio * n`, the call returns the amplitude over `x[1:n]`.
"""
@views function getMaximum(prob::PeriodicOrbitTrapProblem, x::AbstractVector, p; ratio = 1)
	n = div(length(x)-1, ratio)
	return maximum(x[1:n])
end

# this function updates the section during the continuation run
@views function updateSection!(prob::PeriodicOrbitTrapProblem, x, par; stride = 0)
	M, N = size(prob)
	xc = getTimeSlices(prob, x)
	T = extractPeriodFDTrap(prob, x)

	# update the reference point
	prob.xπ .= x[1:end-1]

	# update the normals
	for ii in 0:M-1
		# ii2 = (ii+1)<= M ? ii+1 : ii+1-M
		applyF(prob, prob.ϕ[ii*N+1:ii*N+N], xc[:, ii+1], par)
		prob.ϕ[ii*N+1:ii*N+N] ./= M
	end

	return true
end
####################################################################################################
# Linear solvers for the jacobian of the functional G implemented by PeriodicOrbitTrapProblem
# composite type to encode the Aγ Operator and its associated cyclic matrix
abstract type AbstractPOTrapAγOperator end

# Matrix Free implementation of the operator Aγ
@with_kw mutable struct AγOperatorMatrixFree{Tvec, Tpb, Tpar} <: AbstractPOTrapAγOperator
	orbitguess::Tvec = zeros(1)				# point at which Aγ is evaluated, of size N * M + 1
	prob::Tpb = nothing						# PO functional, used when is_matrix_free = true
	par::Tpar = nothing						# parameters,    used when is_matrix_free = true
end

# implementation of Aγ which catches the LU decomposition of the cyclic matrix
@with_kw mutable struct AγOperatorLU{Tjc, Tpb} <: AbstractPOTrapAγOperator
	N::Int64 = 0							# dimension of time slice
	Jc::Tjc	= lu(spdiagm(0 => ones(1)))	    # lu factorisation of the cyclic matrix
	prob::Tpb = nothing						# PO functional
end

@with_kw struct AγOperatorSparseInplace{Tjc, Tjcf, Tind, Tpb} <: AbstractPOTrapAγOperator
	Jc::Tjc	=  nothing		# cyclic matrix
	Jcfact::Tjcf = nothing	# factorisation of Jc
	indx::Tind = nothing	# indices associated to the sparsity of Jc
	prob::Tpb = nothing		# PO functional
end

# functions to update the cyclic matrix
function (A::AγOperatorMatrixFree)(orbitguess::AbstractVector, par)
	copyto!(A.orbitguess, orbitguess)
	# update par for Matrix-Free
	A.par = par
	return A
end

function (A::AγOperatorLU)(orbitguess::AbstractVector, par)
	# we store the lu decomposition of the newly computed cyclic matrix
	A.Jc = SparseArrays.lu(cylicPOTrapSparse(A.prob, orbitguess, par))
	A
end

function (A::AγOperatorSparseInplace)(orbitguess::AbstractVector, par)
	# compute the cyclic matrix
	A.prob(Val(:JacFullSparseInplace), A.Jc, orbitguess, par, A.indx; updateborder = false)
	# update the Lu decomposition
	lu!(A.Jcfact, A.Jc)
	return A
end

# linear solvers designed specifically for AbstractPOTrapAγOperator
# this function is called whenever one wants to invert Aγ
@with_kw struct AγLinearSolver{Tls} <: AbstractLinearSolver
	# Linear solver to invert the cyclic matrix Jc contained in Aγ
	linsolver::Tls = DefaultLS()
end

@views function _combineSolutionAγLinearSolver(rhs, xbar, N)
	x = similar(rhs)
	x[1:end-N] .= xbar
	x[end-N+1:end] .= x[1:N] .+ rhs[end-N+1:end]
	return x
end

@views function (ls::AγLinearSolver)(A::AγOperatorMatrixFree, rhs)
	# dimension of a time slice
	N = A.prob.N
	# we invert the cyclic part Jc of Aγ
	xbar, flag, numiter = ls.linsolver(dx -> Jc(A.prob, A.orbitguess, A.par, dx), rhs[1:end - N])
	!flag && @warn "Matrix Free solver for Aγ did not converge"
	return _combineSolutionAγLinearSolver(rhs, xbar, N), flag, numiter
end

@views function (ls::AγLinearSolver)(A::AγOperatorLU, rhs)
	# dimension of a time slice
	N = A.N
	xbar, flag, numiter = ls.linsolver(A.Jc, rhs[1:end - N])
	!flag && @warn "Sparse solver for Aγ did not converge"
	return _combineSolutionAγLinearSolver(rhs, xbar, N), flag, numiter
end

@views function (ls::AγLinearSolver)(A::AγOperatorSparseInplace, rhs)
	# dimension of a time slice
	N = A.prob.N
	# we invert the cyclic part Jc of Aγ
	xbar, flag, numiter = ls.linsolver(A.Jcfact, rhs[1:end - N])
	!flag && @warn "Sparse solver for Aγ did not converge"
	return _combineSolutionAγLinearSolver(rhs, xbar, N), flag, numiter
end

####################################################################################################
# The following structure encodes the jacobian of a PeriodicOrbitTrapProblem which eases the use of PeriodicOrbitTrapBLS. It is made so that accessing the cyclic matrix Jc or Aγ is easier. It is combined with a specific linear solver. It is also a convenient structure for the computation of Floquet multipliers. Therefore, it is only used in the method continuationPOTrap
@with_kw struct POTrapJacobianBordered{T∂, Tag <: AbstractPOTrapAγOperator}
	∂TGpo::T∂ = nothing		# derivative of the PO functional G w.r.t. T
	Aγ::Tag					# Aγ Operator involved in the Jacobian of the PO functional
end

# this function is called whenever the jacobian of G has to be updated
function (J::POTrapJacobianBordered)(u0::AbstractVector, par; δ = convert(eltype(u0), 1e-9))
	T = extractPeriodFDTrap(J.Aγ.prob, u0)
	# we compute the derivative of the problem w.r.t. the period TODO: remove this or improve!!
	# TODO REMOVE CE vcat!
	@views J.∂TGpo .= (J.Aγ.prob(vcat(u0[1:end-1], T + δ), par) .- J.Aγ.prob(u0, par)) ./ δ

	J.Aγ(u0, par) # update Aγ

	# return J, needed to properly call the linear solver.
	return J
end

####################################################################################################
# linear solver for the PO functional, akin to a bordered linear solver
@with_kw struct PeriodicOrbitTrapBLS{Tl} <: AbstractLinearSolver
	linsolverbls::Tl = BorderingBLS(solver = AγLinearSolver(), checkPrecision = false)	# linear solver
end

# Linear solver associated to POTrapJacobianBordered
function (ls::PeriodicOrbitTrapBLS)(J::POTrapJacobianBordered, rhs)
	# we solve the bordered linear system as follows
	dX, dl, flag, liniter = @views ls.linsolverbls(J.Aγ, J.∂TGpo[1:end-1],
	 										J.Aγ.prob.ϕ, J.∂TGpo[end],
										   rhs[1:end-1], rhs[end])
	return vcat(dX, dl), flag, sum(liniter)
end

# One could think that by implementing (ls::PeriodicOrbitTrapBLS)(J::POTrapJacobianBLS, rhs1, rhs2), we could speed up the computation of the linear Bordered system arising in the continuation process. However, we can note that this speed up would be observed only if a factorization of J.Aγ is available like an LU one. When such factorization is available, it is automatically stored as such in J.Aγ and so no speed up would be gained by implementing (ls::PeriodicOrbitTrapBLS)(J::POTrapJacobianBLS, rhs1, rhs2)

####################################################################################################
const DocStrjacobianPOTrap = """
- `jacobianPO = :FullLU`. Specify the choice of the linear algorithm, which must belong to `[:FullLU, :FullSparseInplace, :BorderedLU, :FullMatrixFree, :BorderedMatrixFree, :FullSparseInplace]`. This is used to select a way of inverting the jacobian `dG` of the functional G.
    - For `:FullLU`, we use the default linear solver based on a sparse matrix representation of `dG`. This matrix is assembled at each newton iteration. This is the default algorithm.
    - For `:FullSparseInplace`, this is the same as for `:FullLU` but the sparse matrix `dG` is updated inplace. This method allocates much less. In some cases, this is significantly faster than using `:FullLU`. Note that this method can only be used if the sparsity pattern of the jacobian is always the same.
    - For `:Dense`, same as above but the matrix `dG` is dense. It is also updated inplace. This option is useful to study ODE of small dimension.
    - For `:BorderedLU`, we take advantage of the bordered shape of the linear solver and use a LU decomposition to invert `dG` using a bordered linear solver.
    - For `:BorderedSparseInplace`, this is the same as for `:BorderedLU` but the cyclic matrix `dG` is updated inplace. This method allocates much less. In some cases, this is significantly faster than using `:BorderedLU`. Note that this method can only be used if the sparsity pattern of the jacobian is always the same.
    - For `:FullMatrixFree`, a matrix free linear solver is used for `dG`: note that a preconditioner is very likely required here because of the cyclic shape of `dG` which affects negatively the convergence properties of GMRES.
    - For `:BorderedMatrixFree`, a matrix free linear solver is used but for `Jc` only (see docs): it means that `options.linsolver` is used to invert `Jc`. These two Matrix-Free options thus expose different part of the jacobian `dG` in order to use specific preconditioners. For example, an ILU preconditioner on `Jc` could remove the constraints in `dG` and lead to poor convergence. Of course, for these last two methods, a preconditioner is likely to be required.
"""
##########################
# newton wrappers
function _newton(probPO::PeriodicOrbitTrapProblem, orbitguess, par, options::NewtonPar, jacobianPO::Symbol = :FullLU; defOp::Union{Nothing, DeflationOperator{T, Tf, vectype}} = nothing, kwargs...) where {T, Tf, vectype}
	# this hack is for the test to work with CUDA
	@assert sum(extractPeriodFDTrap(probPO, orbitguess)) >= 0 "The guess for the period should be positive"
	@assert jacobianPO in (:Dense, :FullLU, :BorderedLU, :FullMatrixFree, :BorderedMatrixFree, :FullSparseInplace, :BorderedSparseInplace) "This jacobian is not defined. Please choose another one."
	M, N = size(probPO)

	if jacobianPO in (:Dense, :FullLU, :FullMatrixFree, :FullSparseInplace)
		if jacobianPO == :FullLU
			jac = (x, p) -> probPO(Val(:JacFullSparse), x, p)
		elseif jacobianPO == :FullSparseInplace
			# sparse matrix to hold the jacobian
			_J =  probPO(Val(:JacFullSparse), orbitguess, par)
			_indx = getBlocks(_J, N, M)
			# inplace modification of the jacobian _J
			jac = (x, p) -> probPO(Val(:JacFullSparseInplace), _J, x, p, _indx)
		elseif jacobianPO == :Dense
			_J =  probPO(Val(:JacFullSparse), orbitguess, par) |> Array
			jac = (x, p) -> probPO(Val(:JacFullSparseInplace), _J, x, p)
		else
		 	jac = (x, p) -> ( dx -> probPO(x, p, dx))
		end

		if isnothing(defOp)
			return newton(probPO, jac, orbitguess, par, options; kwargs...)
		else
			return newton(probPO, jac, orbitguess, par, options, defOp; kwargs...)
		end
	else # bordered linear solvers
		if jacobianPO == :BorderedLU
			Aγ = AγOperatorLU(N = N, Jc = lu(spdiagm( 0 => ones(N * (M - 1)) )), prob = probPO)
			# linear solver
			lspo = PeriodicOrbitTrapBLS()
		elseif jacobianPO == :BorderedSparseInplace
			_J =  probPO(Val(:JacCyclicSparse), orbitguess, par)
			_indx = getBlocks(_J, N, M-1)
			# inplace modification of the jacobian _J
			Aγ = AγOperatorSparseInplace(Jc = _J,  Jcfact = lu(_J), prob = probPO, indx = _indx)
			lspo = PeriodicOrbitTrapBLS()

		else	# :BorderedMatrixFree
			Aγ = AγOperatorMatrixFree(prob = probPO, orbitguess = zeros(N * M + 1), par = par)
			# linear solver
			lspo = PeriodicOrbitTrapBLS(BorderingBLS(solver = AγLinearSolver(options.linsolver), checkPrecision = false))
		end

		jacPO = POTrapJacobianBordered(zeros(N * M + 1), Aγ)

		if isnothing(defOp)
			return newton(probPO, jacPO, orbitguess, par, (@set options.linsolver = lspo); kwargs...)
		else
			return newton(probPO, jacPO, orbitguess, par, (@set options.linsolver = lspo), defOp; kwargs...)
		end
	end
end

"""
$(SIGNATURES)

This is the Krylov-Newton Solver for computing a periodic orbit using a functional G based on Finite Differences and a Trapezoidal rule.

# Arguments:
- `prob` a problem of type [`PeriodicOrbitTrapProblem`](@ref) encoding the functional G
- `orbitguess` a guess for the periodic orbit where `orbitguess[end]` is an estimate of the period of the orbit. It should be a vector of size `N * M + 1` where `M` is the number of time slices, `N` is the dimension of the phase space. This must be compatible with the numbers `N, M` in `prob`.
- `par` parameters to be passed to the functional
- `options` same as for the regular `newton` method
$DocStrjacobianPOTrap

# Output:
- solution
- history of residuals
- flag of convergence
- number of iterations
"""
newton(probPO::PeriodicOrbitTrapProblem, orbitguess, par, options::NewtonPar; jacobianPO::Symbol = :FullLU, kwargs...) = _newton(probPO, orbitguess, par, options, jacobianPO; defOp = nothing, kwargs...)

"""
	$(SIGNATURES)

This function is similar to `newton(probPO, orbitguess, options, jacobianPO; kwargs...)` except that it uses deflation in order to find periodic orbits different from the ones stored in `defOp`. We refer to the mentioned method for a full description of the arguments. The current method can be used in the vicinity of a Hopf bifurcation to prevent the Newton-Krylov algorithm from converging to the equilibrium point.
"""
newton(probPO::PeriodicOrbitTrapProblem, orbitguess::vectype, par, options::NewtonPar, defOp::DeflationOperator{Tp, Tdot, T, vectype}; jacobianPO::Symbol = :FullLU, kwargs...) where {Tp, Tdot, T, vectype} = _newton(probPO, orbitguess, par, options, jacobianPO; defOp = defOp, kwargs...)

####################################################################################################
# continuation wrapper
"""
	$(SIGNATURES)

This is the continuation routine for computing a periodic orbit using a functional G based on Finite Differences and a Trapezoidal rule.

# Arguments
- `prob::PeriodicOrbitTrapProblem` encodes the functional G
- `orbitguess` a guess for the periodic orbit where `orbitguess[end]` is an estimate of the period of the orbit. It could be a vector of size `N * M + 1` where `M` is the number of time slices, `N` is the dimension of the phase space. This must be compatible with the numbers `N, M` in `prob`.
- `p0` set of parameters passed to the vector field
- `contParams` same as for the regular [`continuation`](@ref) method
- `linearAlgo` same as in [`continuation`](@ref)
- `updateSectionEveryStep = 0` updates the section every `updateSectionEveryStep` step during continuation
$DocStrjacobianPOTrap

Note that by default, the method prints the period of the periodic orbit as function of the parameter. This can be changed by providing your `recordFromSolution` argument.
"""
function continuationPOTrap(prob::PeriodicOrbitTrapProblem, orbitguess, par, lens::Lens, contParams::ContinuationPar, linearAlgo::AbstractBorderedLinearSolver; jacobianPO = :FullLU, updateSectionEveryStep = 0, kwargs...)
	# this hack is for the test to work with CUDA
	@assert sum(extractPeriodFDTrap(prob, orbitguess)) >= 0 "The guess for the period should be positive"
	@assert jacobianPO in (:Dense, :DenseAD, :FullLU, :FullMatrixFree, :BorderedLU, :BorderedMatrixFree, :FullSparseInplace, :BorderedSparseInplace) "This jacobian is not defined. Please chose another one."

	M, N = size(prob)
	options = contParams.newtonOptions

	if computeEigenElements(contParams)
		contParams = @set contParams.newtonOptions.eigsolver =
		 FloquetQaD(contParams.newtonOptions.eigsolver)
	end

	# change the user provided finalise function by passing prob in its parameters
	_finsol = modifyPOFinalise(prob, kwargs, updateSectionEveryStep)
	_recordsol = modifyPORecord(prob, kwargs, par, lens)
	_plotsol = modifyPOPlot(prob, kwargs)

	if jacobianPO in (:Dense, :denseAD, :FullLU, :FullMatrixFree, :FullSparseInplace)

		if jacobianPO == :FullLU
			jac = (x, p) -> FloquetWrapper(prob, prob(Val(:JacFullSparse), x, p), x, p)
		elseif jacobianPO == :FullSparseInplace
			# sparse matrix to hold the jacobian
			_J =  prob(Val(:JacFullSparse), orbitguess, par)
			_indx = getBlocks(_J, N, M)
			# inplace modification of the jacobian _J
			jac = (x, p) -> (prob(Val(:JacFullSparseInplace), _J, x, p, _indx); FloquetWrapper(prob, _J, x, p));
		elseif jacobianPO == :Dense
			_J =  prob(Val(:JacFullSparse), orbitguess, par) |> Array
			jac = (x, p) -> (prob(Val(:JacFullSparseInplace), _J, x, p); FloquetWrapper(prob, _J, x, p));
		else
		 	jac = (x, p) -> FloquetWrapper(prob, x, p)
		end

		# we have to change the Bordered linearsolver to cope with our type FloquetWrapper
		linearAlgo = @set linearAlgo.solver = FloquetWrapperLS(linearAlgo.solver)

		br, z, τ = continuation(
			prob, jac,
			orbitguess, par, lens,
			(@set contParams.newtonOptions.linsolver = FloquetWrapperLS(options.linsolver)), linearAlgo; kwargs...,
			recordFromSolution = _recordsol,
			finaliseSolution = _finsol,
			plotSolution = _plotsol)
	else
		if jacobianPO == :BorderedLU
			Aγ = AγOperatorLU(N = N, Jc = lu(spdiagm( 0 => ones(N * (M - 1)) )), prob = prob)
			# linear solver
			lspo = PeriodicOrbitTrapBLS()
		elseif jacobianPO == :BorderedSparseInplace
			_J =  prob(Val(:JacCyclicSparse), orbitguess, par)
			_indx = getBlocks(_J, N, M-1)
			# inplace modification of the jacobian _J
			Aγ = AγOperatorSparseInplace(Jc = _J,  Jcfact = lu(_J), prob = prob, indx = _indx)
			lspo = PeriodicOrbitTrapBLS()

		else	# :BorderedMatrixFree
			Aγ = AγOperatorMatrixFree(prob = prob, orbitguess = zeros(N * M + 1), par = par)
			# linear solver
			lspo = PeriodicOrbitTrapBLS(BorderingBLS(solver = AγLinearSolver(options.linsolver), checkPrecision = false))
		end

		jacBD = POTrapJacobianBordered(zeros(N * M + 1), Aγ)
		jacPO = (x, p) -> FloquetWrapper(prob, jacBD(x, p), x, p)

		# we change the linear solver
		contParams = @set contParams.newtonOptions.linsolver = FloquetWrapperLS(lspo)

		# we have to change the Bordered linearsolver to cope with our type FloquetWrapper
		linearAlgo = @set linearAlgo.solver = contParams.newtonOptions.linsolver

		br, z, τ = continuation(prob, jacPO, orbitguess, par, lens,
			contParams, linearAlgo;
			kwargs...,
			recordFromSolution = _recordsol,
			finaliseSolution = _finsol,
			plotSolution = _plotsol)
	end
	return setproperties(br; type = :PeriodicOrbit, functional = prob), z, τ
end

"""
$(SIGNATURES)

This is the continuation routine for computing a periodic orbit using a functional G based on Finite Differences and a Trapezoidal rule.

# Arguments
- `prob::PeriodicOrbitTrapProblem` encodes the functional G
- `orbitguess` a guess for the periodic orbit where `orbitguess[end]` is an estimate of the period of the orbit. It could be a vector of size `N * M + 1` where `M` is the number of time slices, `N` is the dimension of the phase space. This must be compatible with the numbers `N, M` in `prob`.
- `p0` set of parameters passed to the vector field
- `contParams` same as for the regular [`continuation`](@ref) method
- `linearAlgo` same as in [`continuation`](@ref)
$DocStrjacobianPOTrap
- `updateSectionEveryStep = 1` updates the section every when `mod(step, updateSectionEveryStep) == 1` during continuation

Note that by default, the method prints the period of the periodic orbit as function of the parameter. This can be changed by providing your `recordFromSolution` argument.
"""
function continuation(prob::PeriodicOrbitTrapProblem, orbitguess, par, lens::Lens, _contParams::ContinuationPar; jacobianPO = :FullLU, recordFromSolution = (u, p) -> (period = u[end],), linearAlgo = nothing, updateSectionEveryStep = 0, kwargs...)
	_linearAlgo = isnothing(linearAlgo) ?  BorderingBLS(_contParams.newtonOptions.linsolver) : linearAlgo
	return continuationPOTrap(prob, orbitguess, par, lens, _contParams, _linearAlgo; jacobianPO = jacobianPO, recordFromSolution = recordFromSolution, updateSectionEveryStep = updateSectionEveryStep, kwargs...)
end

####################################################################################################
# function needed for automatic Branch switching from Hopf bifurcation point
function reMake(prob::PeriodicOrbitTrapProblem, F, dF, par, hopfpt, ζr::AbstractVector, orbitguess_a, period; kwargs...)
	M = length(orbitguess_a)
	N = length(ζr)

	# append period at the end of the initial guess
	orbitguess_v = reduce(vcat, orbitguess_a)
	orbitguess = vcat(vec(orbitguess_v), period) |> vec

	# update the problem
	probPO = setproperties(prob, N = N, F = F, J = dF, ϕ = zeros(N*M), xπ = zeros(N*M))

	probPO.ϕ[1:N] .= ζr
	probPO.xπ[1:N] .= hopfpt.x0

	return probPO, orbitguess
end

####################################################################################################
# Branch switching from Bifs of PO
# """
# $(SIGNATURES)
#
# Branch switching at a Bifurcation point of periodic orbits specified by a [`PeriodicOrbitTrapProblem`](@ref). This is still experimental. A deflated Newton-Krylov solver is used to improve the branch switching capabilities.
#
# # Arguments
# - `br` branch of periodic orbits computed with a [`PeriodicOrbitTrapProblem`](@ref)
# - `ind_bif` index of the branch point
# - `_contParams` parameters to be used by a regular [`continuation`](@ref)
#
# # Optional arguments
# - `Jᵗ = (x, p) -> transpose(d_xF(x, p))` jacobian adjoint, it should be implemented in an efficient manner. For matrix-free methods, `transpose` is not readily available and the user must provide a dedicated method. In the case of sparse based jacobian, `Jᵗ` should not be passed as it is computed internally more efficiently, i.e. it avoid recomputing the jacobian as it would be if you pass `Jᵗ = (x, p) -> transpose(dF(x, p))`
# - `δ` used internally to compute derivatives w.r.t the parameter `p`.
# - `δp = 0.1` used to specify a particular guess for the parameter in the branch which is otherwise determined by `contParams.ds`. This allows to use a step larger than `contParams.dsmax`.
# - `ampfactor = 1` factor which alter the amplitude of the bifurcated solution. Useful to magnify the bifurcated solution when the bifurcated branch is very steep.
# - `usedeflation = true` whether to use nonlinear deflation (see [Deflated problems](@ref)) to help finding the guess on the bifurcated branch
# - `jacobianPO = :BorderedLU` linear solver used for the Newton-Krylov solver when applied to [`PeriodicOrbitTrapProblem`](@ref).
# - `recordFromSolution = (u, p) -> u[end]`, print method used in the bifurcation diagram, by default this prints the period of the periodic orbit.
# - `linearAlgo = BorderingBLS()`, same as for [`continuation`](@ref)
# - `kwargs` keywords arguments used for a call to the regular [`continuation`](@ref)
# - `updateSectionEveryStep = 1` updates the section every when `mod(step, updateSectionEveryStep) == 1` during continuation
# """
# function continuationPOTrapFromBif(br::AbstractBranchResult, ind_bif::Int,
# 	_contParams::ContinuationPar;
# 	Jᵗ = nothing,
# 	δ = 1e-8, δp = 0.1,
# 	ampfactor = 1,
# 	usedeflation = true,
# 	jacobianPO = :BorderedLU,
# 	recordFromSolution = (u,p) -> (period = u[end],),
# 	linearAlgo = nothing,
# 	updateSectionEveryStep = 1,
# 	kwargs...)
#
# 	@assert br.functional isa PeriodicOrbitTrapProblem
# 	@assert abs(br.specialpoint[ind_bif].δ[1]) == 1 "Only simple bifurcation points are handled"
#
# 	verbose = get(kwargs, :verbosity, 0) > 0
# 	_linearAlgo = isnothing(linearAlgo) ?  BorderingBLS(_contParams.newtonOptions.linsolver) : linearAlgo
#
# 	bifpt = br.specialpoint[ind_bif]
#
# 	# let us compute the kernel
# 	λ = (br.eig[bifpt.idx].eigenvals[bifpt.ind_ev])
# 	verbose && print("--> computing nullspace...")
# 	ζ = geteigenvector(br.contparams.newtonOptions.eigsolver, br.eig[bifpt.idx].eigenvec, bifpt.ind_ev)
# 	# we normalize it by the sup norm because it could be too small/big in L2 norm
# 	ζ ./= norm(ζ, Inf)
# 	verbose && println("Done!")
#
# 	pb = br.functional
#
# 	# compute the full eigenvector
# 	ζ_a = MonodromyQaD(Val(:ExtractEigenVector), pb, bifpt.x, setParam(br, bifpt.param), real.(ζ))
# 	ζs = reduce(vcat, ζ_a)
#
# 	## predictor
# 	pbnew, orbitguess = predictor(pb, bifpt, ampfactor, ζs, bifpt.type)
# 	newp = bifpt.param + δp
#
# 	pbnew(orbitguess, setParam(br, newp))[end] |> abs > 1 && @warn "PO Trap constraint not satisfied"
#
# 	if usedeflation
# 		verbose && println("\n--> Attempt branch switching\n--> Compute point on the current branch...")
# 		optn = _contParams.newtonOptions
# 		# find point on the first branch
# 		sol0, _, flag, _ = newton(pbnew, bifpt.x, setParam(br, newp), optn; jacobianPO = jacobianPO, kwargs...)
#
# 		# find the bifurcated branch using deflation
# 		deflationOp = DeflationOperator(2, (x,y) -> dot(x[1:end-1], y[1:end-1]), 1.0, [sol0]; autodiff = true)
# 		verbose && println("\n--> Compute point on bifurcated branch...")
# 		solbif, _, flag, _ = newton(pbnew, orbitguess, setParam(br, newp), (@set optn.maxIter = 10*optn.maxIter), deflationOp; jacobianPO = jacobianPO, kwargs...)
# 		@assert flag "Deflated newton did not converge"
# 		orbitguess .= solbif
# 	end
#
# 	# TODO
# 	# we have to adjust the phase constraint.
# 	# Right now, it can be quite large.
#
# 	# perform continuation
# 	branch, u, τ = continuation(pbnew, orbitguess, setParam(br, newp), br.lens, _contParams;
# 		jacobianPO = jacobianPO,
# 		recordFromSolution = recordFromSolution,
# 		linearAlgo = _linearAlgo, kwargs...)
#
# 	# create a branch
# 	bppo = Pitchfork(bifpt.x, bifpt.param, setParam(br, bifpt.param), br.lens, ζ, ζ, nothing, :nothing)
#
# 	return Branch(setproperties(branch; type = :PeriodicOrbit, functional = br.functional), bppo), u, τ
# end

# predictor function close to bifurcations of PO
function predictor(pb::PeriodicOrbitTrapProblem, bifpt, ampfactor, ζ, bptype::Symbol)
	@assert bptype in (:bp, :pd)
	if bptype == :bp
		orbitguess = copy(bifpt.x)
		orbitguess[1:end-1] .+= ampfactor .*  ζ
	elseif bptype == :pd
		M, N = size(pb)
		orbitguess0 = copy(bifpt.x)[1:end-1]
		orbitguess0c = getTimeSlices(pb, copy(bifpt.x))
		ζc = reshape(ζ, N, M)
		orbitguess_c = orbitguess0c .+ ampfactor .*  ζc
		orbitguess_c = hcat(orbitguess_c, orbitguess0c .- ampfactor .*  ζc)
		orbitguess = vec(orbitguess_c[:,1:2:end])

		@debug size(orbitguess0) size(orbitguess) size(ζ) size(orbitguess_c)
		orbitguess = vcat(orbitguess, 2bifpt.x[end])
	end
	return pb, orbitguess
end


#=
  Tug of war
   From Hakaru example documentation/tugofwar_rejection.hk

  Compare with true_skill.jl

  cf ~/webppl/tug_of_war.wppl
     ~/blog/tug_of_war.blog
=#
using Turing, StatsPlots, DataFrames

#=
# For num_people = 5
strength (latent):
strength[1] > strength[2]: 0.67203
strength[1] > strength[3]: 0.641555
strength[1] > strength[4]: 0.76379
strength[1] > strength[5]: 0.73293
strength[2] > strength[3]: 0.60609
strength[2] > strength[4]: 0.69859
strength[2] > strength[5]: 0.688345
strength[3] > strength[4]: 0.54991
strength[3] > strength[5]: 0.640705
strength[4] > strength[5]: 0.49402
pulls (observed):
pulls[1] > pulls[2]: 1.0
pulls[1] > pulls[3]: 1.0
pulls[1] > pulls[4]: 1.0
pulls[1] > pulls[5]: 1.0
pulls[2] > pulls[3]: 1.0
pulls[2] > pulls[4]: 1.0
pulls[2] > pulls[5]: 1.0
pulls[3] > pulls[4]: 0.99546
pulls[3] > pulls[5]: 0.99546
pulls[4] > pulls[5]: 1.0

=#
@model function tug_of_war(n)

    strength = TArray{Real}(undef, n)
    pulls = TArray{Real}(undef, n)
    # strength = Vector{Real}(undef, n)
    # pulls = Vector{Real}(undef, n)

    for p in 1:n
        strength[p] ~ Uniform(0,1) # truncated(Normal(0,1),0,10) # latent
        pulls[p]    ~ truncated(Normal(strength[p],1),0,10) # observed # truncated(Normal(strength[p],1),0,10) # observed
    end

    # The Hakary/BLOG model had "pulls(a) < pulls(b) "
    # which is a little strange so I change to return the winner.
    function winner(a,b)
        return pulls[a] > pulls[b] ? a : b
    end

    # match1 = winner(alice, bob)
    # match2 = winner(bob, carol)
    # match3 = winner(alice, carol)

    # match1 = winner(1, 2)
    # match2 = winner(2, 3)
    # match3 = winner(1, 3)

    # A little more general approach
    for p1 in 1:n-1, p2 in p1+1:n
        true ~ Dirac(winner(p1,p2) == p1)
    end

end

num_people = 3
model = tug_of_war(num_people)

num_chains = 4
num_samples = num_people * 10_000 # 10_000

# High rhat
chains = sample(model,  MH(), MCMCThreads(), num_samples, num_chains)

# PG(20) doesn't differentiate the people's pull/strength
# chains = sample(model,  PG(20), MCMCThreads(), num_samples, num_chains)

# chains = sample(model,  HMC(0.1, 5), MCMCThreads(), num_samples, num_chains)

# IS() doesn't differentiate the people's pull/strength
# chains = sample(model,  IS(), MCMCThreads(), num_samples, num_chains)

# High rhat. Alot of rejected proposals
# chains = sample(model,  HMC(0.1, 5), MCMCThreads(), num_samples, num_chains)

# chains = sample(model,  HMCDA(200, 0.65, 0.3), MCMCThreads(), num_samples, num_chains)

# chains = sample(model,  NUTS(1000, 0.65), MCMCThreads(), num_samples, num_chains)

# Too slow
# chains = sample(model,  SMC(), MCMCThreads(), num_samples, num_chains)

# Good: slightly high rhat
# chains = sample(model,  NUTS(1000,0.65), MCMCThreads(), num_samples, num_chains)

display(chains)
display(gelmandiag(chains))
# display(plot(chains))
df = DataFrame(chains)

# Show the latent strength
println("strength (latent):")
for p1 in 1:num_people-1, p2 in p1+1:num_people
    println("strength[$p1] > strength[$p2]: ", mean(df[!,"strength[$p1]"] .>df[!,"strength[$p2]"]))
end

println("pulls (observed):")
for p1 in 1:num_people-1, p2 in p1+1:num_people
    println("pulls[$p1] > pulls[$p2]: ", mean(df[!,"pulls[$p1]"] .>df[!,"pulls[$p2]"]))
end


using Distributions, Random
using Test, LinearAlgebra, PDMats

p  = 3
n1 = p + abs(10randn())
n2 = p + 1 + abs(10randn())
B  = rand( InverseWishart(p + 2, Matrix(1.0I, p, p)) )
b  = Matrix((n2 / n1) * I, 1, 1)

G = MatrixFDist(n1, n2, B)
F = MatrixFDist(n1, n2, b)
f = FDist(n1, n2)

@testset "MatrixFDist promotion during construction" begin
    R = MatrixFDist(Float32(n1), BigFloat(n2), Matrix{Float64}(B))
    @test partype(R) == BigFloat
end

@testset "MatrixFDist construction errors" begin
    @test_throws ArgumentError MatrixFDist(1, n2, B)
    @test_throws ArgumentError MatrixFDist(n1, 1, B)
    @test_throws ArgumentError mean(MatrixFDist(n1, p, B))
end

@testset "MatrixFDist params" begin
    df1, df2, PDB = params(G)
    @test df1 == n1
    @test df2 == n2
    @test PDB.mat == B
end

@testset "MatrixFDist dim" begin
    @test dim(G) == p
    @test dim(F) == 1
end

@testset "MatrixFDist size" begin
    @test size(G) == (p, p)
    @test size(F) == (1, 1)
end

@testset "MatrixFDist rank" begin
    @test rank(G) == p
    @test rank(F) == 1
    @test rank(G) == rank(rand(G))
    @test rank(F) == rank(rand(F))
end

@testset "MatrixFDist insupport" begin
    @test insupport(G, rand(G))
    @test insupport(F, rand(F))

    @test !insupport(G, rand(G) + rand(G) * im)
    @test !insupport(G, randn(p, p + 1))
    @test !insupport(G, randn(p, p))
end

@testset "MatrixFDist logpdf" begin
    Z = rand(F)
    z = Z[1, 1]
    @test logpdf(F, Z) ≈ logpdf(f, z)
end

@testset "MatrixFDist cov and var" begin
    @test vec(var(G)) ≈ diag(cov(G))
    @test reshape(cov(G), p, p, p, p) ≈ cov(G, Val(false))
    @test cov(F)[1] ≈ var(f)
    @test_throws ArgumentError cov(MatrixFDist(n1, p + 1, B))
end

@testset "MatrixFDist sample moments" begin
    @test isapprox(mean(rand(G, 100000)), mean(G) , atol = 0.1)
    @test isapprox(cov(hcat(vec.(rand(G, 100000))...)'), cov(G) , atol = 0.1)
end

@testset "MatrixFDist conversion" for elty in (Float32, Float64, BigFloat)
    Gel1 = convert(MatrixFDist{elty}, G)
    Gel2 = convert(MatrixFDist{elty}, G.W, n2, G.logc0)

    @test Gel1 isa MatrixFDist{elty, Wishart{elty, PDMat{elty,Array{elty,2}}}}
    @test Gel2 isa MatrixFDist{elty, Wishart{elty, PDMat{elty,Array{elty,2}}}}
    @test partype(Gel1) == elty
    @test partype(Gel2) == elty
end


import IntervalArithmetic: diam, isinterior

export branch_and_prune, Bisection, Newton

diam(x::Root) = diam(x.interval)

Base.size(x::Interval) = (1,)

isinterior{N}(X::IntervalBox{N}, Y::IntervalBox{N}) = all(isinterior.(X, Y))

"""
    branch_and_prune(f, X, contractor, tol=1e-3)

Generic branch and prune routine for finding isolated roots of a function
`f:R^n → R^n` in a box.

Inputs:
- `f`: function whose roots will be found
- `X`: `Interval` or `IntervalBox`
- `contractor`: function that, when applied to the function `f`, determines
    the status of a given box `X`. It returns the new box and a symbol indicating
    the status. Current possible values are `Bisection` and `Newton`.

"""
function branch_and_prune(f, X, contractor, tol=1e-3)

    input_dim = length(X)
    output_dim = length(X)

    # @show input_dim
    # @show output_dim

    # if !(input_dim == output_dim)
    #     throw(ArgumentError("Input dimension ($input_dim) and output dimension ($output_dim) must be the same."))
    # end

    contract = contractor(Val{input_dim}, f)

    # main algorithm:

    working = [X]
    outputs = Root{typeof(X)}[]

    sizehint!(outputs, 100)
    sizehint!(working, 1000)

    while !isempty(working)
        # @show working
        X = pop!(working)

        status, output = contract(X)

        if status == :empty
            continue

        elseif status == :unique
            push!(outputs, Root(output, :unique))

        elseif diam(output) < tol
            push!(outputs, Root(output, :unknown))

        else  # branch
            X1, X2 = bisect(X)

            push!(working, X1, X2)
        end

    end

    return outputs
end

branch_and_prune(f, X::Root, contractor, tol=1e-3) =
    branch_and_prune(f, X.interval, contractor, tol)

function branch_and_prune(f, V::Vector{Root{T}}, contractor, tol=1e-3) where {T}
    reduce(append!, Root{T}[], [branch_and_prune(f, X.interval, contractor, tol) for X in V])
end

function branch_and_prune(f, V::Vector{T}, contractor, tol=1e-3) where {T}
    reduce(append!, Root{T}[], [branch_and_prune(f, X, contractor, tol) for X in V])
end

export recursively_branch_and_prune

function recursively_branch_and_prune(h, X, contractor=BisectionContractor, final_tol=1e-14)
    tol = 2
    roots = branch_and_prune(h, X, IntervalRootFinding.BisectionContractor, tol)

    while tol > 1e-14
       tol /= 2
       roots = branch_and_prune(h, roots, IntervalRootFinding.BisectionContractor, tol)
    end

    return roots
end

"""
If the input interval is complex, treat `f` as a complex function, currently of one complex variable `z`.
"""
function branch_and_prune{T}(f, Xc::Complex{Interval{T}}, contractor, tol=1e-3)

    g = realify(f)
    Y = IntervalBox(reim(Xc))

    roots = branch_and_prune(g, Y, contractor, tol)

    # @show roots

    return [Root(Complex(root.interval...), root.status) for root in roots]
end

contains_zero{T}(X::Interval{T}) = zero(T) ∈ X
contains_zero(X::SVector) = all(contains_zero.(X))
contains_zero(X::IntervalBox) = all(contains_zero(X[i]) for i in 1:length(X))

# contractors:

abstract type Contractor{F} end

export Bisection, Newton

struct Bisection{F} <: Contractor{F}
    dimension::Int
    f::F
end

Bisection{n}(::Type{Val{n}}, f) = Bisection(n, f)

function (contractor::Bisection)(X)
    image = contractor.f(X)

    if !(contains_zero(image))
        return :empty, X
    end

    return :unknown, X
end

"""
Generic refine operation for Krawczyk and Newton.
This assumes that it is already known that `X` contains a unique root.
Call using e.g. `op = X -> N(f, f_prime, X)`
"""
function refine(op, X)

    tolerance = 1e-16

    while diam(X) > tolerance  # avoid problem with tiny floating-point numbers if 0 is a root
        NX = op(X) ∩ X
        NX == X && break  # reached limit of precision
        X = NX
    end

    return X
end

struct Newton{F,FP,O} <: Contractor{F}
    dimension::Int
    f::F
    fp::FP
    op::O
end

function Newton(::Type{Val{1}}, f::Function)
    f_prime = x -> ForwardDiff.derivative(f, x)
    Newton(1, f, f_prime, N)
end

function Newton{n}(::Type{Val{n}}, f::Function)
    f_prime = x -> ForwardDiff.jacobian(f, x)
    Newton(n, f, f_prime, N)
end

function (C::Newton)(X)

    # use Bisection contractor for this:
    if !(contains_zero(IntervalBox(C.f(X))))
        return :empty, X
    end

    # given that have the Jacobian, can also do mean value form

    NX = C.op(C.f, C.fp, X) ∩ X

    isempty(NX) && return :empty, X

    if NX ⪽ X  # isinterior; know there's a unique root inside
        NX =  refine(X -> C.op(C.f, C.fp, X), NX)
        return :unique, NX
    end

    return :unknown, NX
end

"""
    roots(f, X, contractor, tol=1e-3)

Generic branch and prune routine for finding isolated roots of a function
`f:R^n → R^n` in a box `X`, or a vector of boxes.

Inputs:
- `f`: function whose roots will be found
- `X`: `Interval` or `IntervalBox`
- `contractor`: function that, when applied to the function `f`, determines
    the status of a given box `X`. It returns the new box and a symbol indicating
    the status. Current possible values are `Bisection` and `Newton`.

"""
roots{C<:Contractor}(f, X, contractor::Type{C}, tol::Float64=1e-3) = branch_and_prune(f, X, contractor, tol)

roots(f, X, tol::Float64=1e-3) = branch_and_prune(f, X, Newton, tol)


# # Baffled piston in a half-sphere domain

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

# ## Description

# Baffled piston in a half-sphere domain with Absorbing Boundary Condition (ABC).
# This is a steady-state simulation.

# ![](baffled_piston.png)

# ## Goals

# - Demonstrate the use of an algorithm to run the simulation.

##
# ## Definitions

# The finite element code relies on the basic functionality implemented in this
# package.
using FinEtools
# The acoustics functionality is brought in:
using FinEtoolsAcoustics

# We shall need some facilities from the linear algebra package
using LinearAlgebra

# The input quantities are provided including the units:
rho = 1.21*phun("kg/m^3");# mass density
c  = 343.0*phun("m/s");# sound speed
bulk =  c^2*rho;
omega =  7500*phun("rev/s");      # frequency of the piston
a_piston =  -1.0*phun("mm/s")     # amplitude of the piston acceleration
R = 50.0*phun("mm");# radius of the piston
Ro = 150.0*phun("mm"); # radius of the enclosure
nref = 4;#number of refinements of the sphere around the piston
nlayers = 35;                     # number of layers of elements surrounding the piston
tolerance = R/(2^nref)/100

# Hexahedral mesh
fens,fes  =  H8sphere(R,nref);
bfes  =  meshboundary(fes)
File  =   "baffledabc_boundary.vtk"
vtkexportmesh(File, connasarray(bfes), fens.xyz, FinEtools.MeshExportModule.VTK.Q4)
@async run(`"paraview.exe" $File`)

l = selectelem(fens, bfes, facing = true, direction = [1.0 1.0  1.0], dotmin= 0.001)
ex(xyz, layer) = (R+layer/nlayers*(Ro-R))*xyz/norm(xyz)
fens1, fes1  =  H8extrudeQ4(fens, subset(bfes,l), nlayers, ex);
fens, newfes1, fes2 =  mergemeshes(fens1, fes1, fens, fes, tolerance)
fes = cat(newfes1,fes2)

# Piston surface mesh
bfes  =  meshboundary(fes)
l1 = selectelem(fens, bfes, facing = true, direction = [-1.0 0.0 0.0])
l2 = selectelem(fens, bfes, distance = R, from = [0.0 0.0 0.0], inflate = tolerance)
piston_fes = subset(bfes,intersect(l1,l2));

# Outer spherical boundary
louter = selectelem(fens, bfes, facing = true, direction = [1.0 1.0  1.0], dotmin= 0.001)
outer_fes = subset(bfes,louter);

# The simulation is driven by setting up and then executing an algorithm. For
# this purpose we store the data in a "model data dictionary".

# Region of the fluid. Define the finite element machine for the fluid region.
material = MatAcoustFluid(bulk, rho)
region1 =  FDataDict("femm"=>FEMMAcoust(IntegDomain(fes, GaussRule(3, 2)), material))

# Define the finite element machine for the ABC (absorbing boundary condition).
abc1  =  FDataDict("femm"=>FEMMAcoustSurf(IntegDomain(outer_fes, GaussRule(2, 2)), material))

# The normal flux is prescribed on the surface of the piston in terms of the
# known acceleration.
flux1  =  FDataDict("femm"=>FEMMAcoustSurf(IntegDomain(piston_fes, GaussRule(2, 2)), material),  "normal_flux"=> -rho*a_piston+0.0im);

# Make model data dictionary. It completely defines the parameters of the
# problem.
modeldata =  FDataDict("fens"=>fens, "omega"=>omega, "regions"=>[region1], "flux_bcs"=>[flux1], "ABCs"=>[abc1])

# Call the solver. The model data is returned enriched of the solution
# parameters, such as the pressure field `P`.
modeldata = FinEtoolsAcoustics.AlgoAcoustModule.steadystate(modeldata)

# Extract geometry field and the pressure field for postprocessing.
geom = modeldata["geom"]
P = modeldata["P"]

# Write out the postprocessing data is a VTK file. Dump both the magnitude and
# the components of the solution (complex field).
File  =   "baffledabc.vtk"
vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.VTK.H8; scalars = [("absP", abs.(P.values)), ("realP", real.(P.values)), ("imagP", imag.(P.values))])
# If the `paraview` program is installed, run it.
@async run(`"paraview.exe" $File`)

true


using Bridge
using Bridge.Models: ℝ
import Bridge: b, σ, B, β, a, constdiff

struct SinDiffusion{T} <: ContinuousTimeProcess{ℝ{1, T}}
    a::T
    b::T
    c::T
    σ::T
end

b(t, x, P::SinDiffusion) = ℝ{1}(P.a + P.b*sin.(P.c * x))
σ(t, x, P::SinDiffusion) = ℝ{1}(P.σ)
constdiff(::SinDiffusion) = true
clone(::SinDiffusion, θ) = SinDiffusion(θ...)
params(P::SinDiffusion) = [P.a, P.b, P.c, P.σ]

struct SinDiffusionAux{T,S} <: ContinuousTimeProcess{ℝ{1, T}}
    a::T
    b::T
    c::T
    σ::T
    t::Float64
    u::S
    T::Float64
    v::S
end

#B(t, P::SinDiffusionAux) = @SMatrix [0.0]
#β(t, P::SinDiffusionAux) = ℝ{1}(0.0)

#B(t, P::SinDiffusionAux) = @SMatrix [(P.b*P.c*cos.(P.c * P.v)) * t/P.T + (P.b*P.c*cos.(P.c * P.u)) * (1-t/P.T)]
#β(t, P::SinDiffusionAux) = ℝ{1}((P.a + P.b*sin.(P.c * P.v) - P.b*P.c*cos.(P.c * P.v) * P.v) * t/P.T + (P.a + P.b*sin.(P.c * P.u) - P.b*P.c*cos.(P.c * P.u) * P.u) * (1-t/P.T))

#B(t, P::SinDiffusionAux) = @SMatrix [P.b*P.c*cos.(P.c * P.v)]
#β(t, P::SinDiffusionAux) = ℝ{1}(P.a + P.b*sin.(P.c * P.v) - P.b*P.c*cos.(P.c * P.v) * P.v)

B(t, P::SinDiffusionAux) = @SMatrix [0 + (t/P.T)/5]#+ (t/P.T)/5
β(t, P::SinDiffusionAux) = ℝ{1}((P.v-P.u)/P.T)

σ(t, P::SinDiffusionAux) = ℝ{1}(P.σ)
dependsOnParams(::SinDiffusionAux) = (4,)
constdiff(::SinDiffusionAux) = true
b(t, x, P::SinDiffusionAux) = B(t, P) * x + β(t, P)
a(t, P::SinDiffusionAux) = σ(t, P) * σ(t, P)'
clone(P::SinDiffusionAux, θ) = SinDiffusionAux(θ..., P.t, P.u, P.T, P.v)


using LinearAlgebraicRepresentation
LAR = LinearAlgebraicRepresentation
using ViewerGL
GL = ViewerGL

function randomarrangement2d(n::Int = 100)
	V, EV = LAR.randomcuboids(n, .4)

	model = Lar.Model(GL.normalize2(V,flag=true))
	Lar.addModelCells!(model, 1, EV)
	model = Lar.Arrangement.planar_arrangement(model)
	FVs = Lar.triangulate2D(model)
	EVs = Lar.FV2EVs(model) # polygonal face fragments

	GL.VIEW(GL.GLExplode(model.G, EVs, 1.2,1.2,1.2,1,1));
	GL.VIEW(GL.GLExplode(model.G, FVs, 1.2,1.2,1.2,99,1));
	GL.VIEW(GL.GLExplode(model.G, FVs, 1,1,1,99,1));
end


module StatisticalVerificationOfIsingEnergies

using Distributions
using Cumulants
using StatsBase

export estimate_ground_state_energy, squared_error, bootstrap_get_pvalue,
       bootstrap_hists_of_mins, estiamte_temperature

"""
    estimate_ground_state_energy(energies::Vector{Float64}, α::Float64)

Returns Float64, estimated ground state energy.

```jldoctest

julia> energies = [0.1, 1., 1.2, 2., 2.2, 3., 3.1, 3.5];

julia> estimate_ground_state_energy(energies, 0.19)
9.745023110079497

```
"""
function estimate_ground_state_energy(energies::Vector{Float64}, α::Float64)
    μ = mean(energies)
    σ = std(energies)
    η = skewness(energies)

    μ-(α+2)/(α+1)*σ/η
end

"""
    squared_error(energies::Vector{Float64}, α::Float64)

Returns Float64, estimated standard deviation of the estimation of the ground state energy

```jldoctest

julia> energies = [1. ,1.5, 2., 3., 4.];

julia> squared_error(energies, 0.19)
8.633622351519547
```
"""
function squared_error(energies::Vector{Float64}, α::Float64)
    N = length(energies)
    B = (α+2)/(α+1)
    c = [Array(x)[1] for x in cumulants(reshape(energies, :, 1), 6, 1)]
    δc2 = sqrt((c[4] + 2*c[2]^2)/N)
    δc3 = sqrt((c[6] + 9*c[2]*c[4]+ 9*c[3]^2 + 6*c[2]^3)/N)
    e2 = B*2*(c[2]/c[3])*δc2
    e3 = B*(c[2]/c[3])^2*δc3
    (e2^2+e3^2)^0.5
end

"""
    bootstrap_hists_of_mins(energies::Vector{Float64}, α::Float64, s::Int, l::Int=length(x))

Returns s-long vector of Float64, estimated ground state energies from l-long samples of dootstrap resampling

julia> energies = [1. ,1.5, 2., 3., 4.];

julia> Random.seed!(1234);

```jldoctest
julia> h = bootstrap_hists_of_mins(energies, 0.19, 3)
3-element Vector{Float64}:
 10.327274688631618
 -4.407048476374387
 -4.721357273543485
```
"""
function bootstrap_hists_of_mins(energies::Vector{Float64}, α::Float64, s::Int, l::Int=length(energies))
    ret = zeros(s)
    Threads.@threads for i ∈ 1:s
        ret[i] = estimate_ground_state_energy(rand(energies, l), α)
    end
    ret
end

"""
    bootstrap_get_pvalue(energies::Vector{Float64}, α::Float64, s::Int=1000)

returns Float64, p-value that energies contains the ground state if the Ising sampler,
s is the number ob bootstrap samples.

```jldoctest
julia> Random.seed!(1234);

julia> energies = [0.1, 1., 1.2, 2., 2.2, 3., 3.1, 3.5];

julia> bootstrap_get_pvalue(energies, 0.19)
0.655
```
"""
function bootstrap_get_pvalue(energies::Vector{Float64}, α::Float64, s::Int=1000)
    y = bootstrap_hists_of_mins(energies, α, s)
    f = ecdf(y)
    Hmin = minimum(energies)
    1.0 - f(Hmin)
end

"""
    estiamte_temperature(energies::Vector{Float64}, ground_energy::Float64)

returns Float64, estimated temperature.

```jldoctest
julia> energies = [1. ,1.5, 2., 3., 4.];

julia> estiamte_temperature(energies, 0.)
0.8958957699200034
```
"""
function estiamte_temperature(energies::Vector{Float64}, ground_energy::Float64)
    EH = ground_energy - mean(energies)
    EH/(std(energies)*skewness(energies)*EH-var(energies))
end

end


using EmpiricalBayes
using Test
using Distributions

## test whether speed-up for normal distributions indeed works
# i.e. whether it is faster and numerically works
marginal_grid = collect(range(-6, stop=6, length=1001));

μs = [0.0, 1.0]

for μ in μs
    @show μ
    # same distribution, but once as Normal type, and once as mixture type..
    dist1= MixtureModel([ Normal(μ,.2), Normal(μ,.2)])
    dist2= Normal(μ,.2)

    normal_conv_prob1 = NormalConvolutionProblem(dist1, marginal_grid)
    normal_conv_prob2 = NormalConvolutionProblem(dist2, marginal_grid)

    @test normal_conv_prob1.marginal ≈ normal_conv_prob2.marginal

    t1 = @elapsed NormalConvolutionProblem(dist1, marginal_grid)
    t2 = @elapsed NormalConvolutionProblem(dist2, marginal_grid)

    @test t1 > t2
end


function delta2Jl!(
    delta :: Matrix, AField :: Matrix, BField :: Matrix, Lap :: Matrix,
    Out :: Matrix, decay :: Real)

    d1, d2 = size(Out)
    for j in 1:d2
      for i in 1:d1
        Out[i,j] = delta[i,j] / (1 + AField[i,j]) - decay * BField[i,j] + Lap[i,j]
      end
    end
end


#
# GENERALIZED ARGMAX FUNCTION
# Implements the EVT-I business...

errmsg = "Please specify a log sum behavior. Usage is
`generalized_argmax!(y,x,sig, SumNormalize)` or
`generalized_argmax!(y,x,sig,SumRaw)`"

generalized_argmax!(y,x,sig)=error(errmsg)

function generalized_argmax!(y, x, sig, logsumspec)
    # `logsumspec` controls whether we subtract by greatest element of x or not
    # when computing stuff

    expsum = vexpsum(x, sig, logsumspec)  # Note: might be a number or a tuple
    # (depending on `logsumspec`)

    @inbounds for i in eachindex(x)
        y[i] = donorm(x[i], expsum, sig, logsumspec)
    end

    return genlog(expsum, sig, logsumspec)
end

struct SumNormalize end
struct SumRaw end

genlog(ss,sig,::Type{SumNormalize})=ss[2]+sig*log(ss[1])
genlog(ss,sig,::Type{SumRaw})=sig*log(ss)

vexpsum(x, sig, ::Type{SumNormalize})=vexpsum_norm(x,sig)
vexpsum(x, sig, ::Type{SumRaw})=vexpsum_raw(x,sig)

function vexpsum_norm(x::AbstractArray, sig)
    s = zero(eltype(x))
    maxel = maximum(x)
    @inbounds for i in eachindex(x)
        xnorm = x[i]-maxel
        s += exp(xnorm/sig)
    end
    return s, maxel
end

function vexpsum_raw(x::AbstractArray, sig)
    s = zero(eltype(x))
    @inbounds for i in eachindex(x)
        s += exp(x[i]/sig)
    end
    return s
end

donorm(xx, ss, sig, ::Type{SumNormalize})=exp((xx - ss[2])/sig)/ss[1]
donorm(xx, ss, sig, ::Type{SumRaw})=exp(xx/sig)/ss


#Table 6.2
#Run uility functions at bottom
#Change the setting of Lambda if necessary
using NMF, LinearAlgebra, MultiscaleGraphSignalTransforms, ScikitLearn
@sk_import linear_model:Lasso

########################### utility functions ###########################
function create_scaling_vectors(GProws::GraphPart)
    GP = GProws
    ### creating scaling vectors
    temp = GP.tag.== 0;
    scaling_index = findall(temp);# index of the scaling vectors .i.e. vectors with tag == 0

    num_scaling = length(scaling_index);
    scaling_vectors = fill(0., (size(GP.tag,1), num_scaling));

    for i = 1:num_scaling
        ind = scaling_index[i]
        dvec = fill(0., (size(GP.tag,1),1))
        dvec[ind[1],1] = 1.;
        BS = bs_level(GP, ind[2] - 1)
        scaling_vectors[:,i] = ghwt_synthesis(dvec, GP, BS)
    end
    return scaling_vectors
end

function create_PHI(scaling_vectors_rows, scaling_vectors_cols, m, n)
    PHI_right = repeat(scaling_vectors_cols[:,1:n], outer = (1, m))
    PHI_left = repeat(scaling_vectors_rows[:,1:m], inner = (1, n))
    return PHI_left, PHI_right
end

function ghwtinit(matrix, k)
    GProws, GPcols = partition_tree_matrixDhillon(matrix)
    dmatrix = ghwt_analysis_2d(matrix, GProws, GPcols)
    scaling_vectors_rows = create_scaling_vectors(GProws)
    scaling_vectors_cols = create_scaling_vectors(GPcols)

    PHI_left, PHI_right = create_PHI(scaling_vectors_rows, scaling_vectors_cols, 10, 10)
    scaling_vectors = fill(0., (length(matrix[:]), size(PHI_left,2)))
    for i = 1:size(PHI_left,2)
        scaling_vectors[:,i] = (PHI_left[:,i]*PHI_right[:,i]')[:]
    end
    ##########################
    ####################################### perform lasso

    f = matrix[:]
    #Lambda = range(0.00001,stop = 0.0001, length = 10)
    Lambda = range(0.001,stop = 0.01, length = 10)# m=125,n=25,k=5,sigma = 0.5,
    #Lambda = range(0.001, stop = 0.01, length = 10)
    B = fill(0., (size(scaling_vectors,2), length(Lambda))) ### initialize coefficients
    for i = 1:length(Lambda)
        lasso_out = Lasso(alpha = Lambda[i], positive = true, fit_intercept = false)
        ScikitLearn.fit!(lasso_out, scaling_vectors, f)
        B[:,i] = lasso_out.coef_
    end

    j = 1
    if sum((B[:,j].!=0)) < k
        return false, false
    else
        while j <10 && sum(B[:,j+1].!=0) >= k
            j += 1
        end
    end
    index = findall(B[:,j].!=0)[1:k]
    rescaleW = repeat(sqrt.(B[index,j]'),size(PHI_left,1),1)
    rescaleH = repeat(sqrt.(B[index,j]'),size(PHI_right,1),1)
    Wtilda = PHI_left[:,index].*rescaleW
    Htilda = Array{Float64,2}((PHI_right[:,index].*rescaleH)')
    return Wtilda, Htilda
end

function improved_NNDSVD(matrix, k)
    m,n = size(matrix)
    U,S,V = svd(matrix)
    W = zeros(m,k)
    H = zeros(k,n)
    Y = U*sqrt.(Diagonal(S))
    Z = sqrt.(Diagonal(S))*transpose(V)

    W[:,1] = abs.(Y[:,1])
    H[1,:] = abs.(Z[1,:])
    j = 2
    for i = 2:k
        if mod(i,2) == 0
            W[:,i] = Y[:,j]
            H[i,:] = Z[i,:]
        else
            W[:,i] = .-Y[:,j]
            H[i,:] = .-Z[i,:]
            j += 1
        end
        W[W[:,i] .< 0,i] .= 0
        H[i,H[i,:] .< 0] .= 0
    end
    return W, H
end
########################### end of utility functions ###########################

### Now the start of the actual NMF experiments.

k = 5
m = 125
n = 25
sigma = 0.5 #noise parameter
#methodid = 1 corresponds to ALSPGRAD.
#methodid = 2 corresponds to HALS.
methodid = 2

num_pos_eghwt_svd = 0
num_neg_eghwt_svd = 0
num_pos_eghwt_rand = 0
num_neg_eghwt_rand = 0
num_pos_svd_rand = 0
num_neg_svd_rand = 0
num_pos_svd_newsvd = 0
num_neg_svd_newsvd = 0
num_pos_eghwt_newsvd = 0
num_neg_eghwt_newsvd = 0
for i = 1:50
    Wtrue = rand(m,k)
    Htrue = rand(k,n)
    X = Wtrue*Htrue + sigma*rand(m,n)

    #X = rand(m,n)

    ################
    ######################################

    Wtilda, Htilda = ghwtinit(X, k)
    if Wtilda == false
        continue
    end
    Wsvd, Hsvd = NMF.nndsvd(X, k)
    Wrand, Hrand = NMF.randinit(X, k)
    Wnewsvd, Hnewsvd = improved_NNDSVD(X, k)

    iters = 1000000
    method = [NMF.ALSPGrad{Float64}(maxiter=iters, tolg=1.0e-6), NMF.CoordinateDescent{Float64}(maxiter=iters, α=0.5, l₁ratio=0.5)]

    a = NMF.solve!(method[methodid], X, copy(Wtilda), copy(Htilda))
    b = NMF.solve!(method[methodid], X, copy(Wsvd), copy(Hsvd))
    c = NMF.solve!(method[methodid], X, copy(Wrand), copy(Hrand))
    d = NMF.solve!(method[methodid], X, copy(Wnewsvd), copy(Hnewsvd))

    #println("initial norm ",norm(X - Wtilda*Htilda)," ",norm(X-Wsvd*Hsvd)," ",norm(X - Wrand*Hrand)," ", norm(X - Wnewsvd*Hnewsvd))
    #println(" ",a.niters," ",b.niters," ",c.niters," ",d.niters)
    #println("final norm",norm(X - a.W*a.H)," ",norm(X-b.W*b.H)," ",norm(X - c.W*c.H)," ", norm(X - d.W*d.H))
    if a.niters < b.niters
        global num_pos_eghwt_svd += 1
    else
        global num_neg_eghwt_svd += 1
    end
    if a.niters < c.niters
        global num_pos_eghwt_rand += 1
    else
        global num_neg_eghwt_rand += 1
    end
    if b.niters < c.niters
        global num_pos_svd_rand += 1
    else
        global num_neg_svd_rand += 1
    end
    if a.niters < d.niters
        global num_pos_eghwt_newsvd += 1
    else
        global num_neg_eghwt_newsvd += 1
    end
    if b.niters < d.niters
        global num_pos_svd_newsvd += 1
    else
        global num_neg_svd_newsvd += 1
    end
end

println("compare eGHWT with NNDSVD ", num_pos_eghwt_svd/(num_pos_eghwt_svd + num_neg_eghwt_svd))
println("compare eGHWT with Random ", num_pos_eghwt_rand/(num_pos_eghwt_rand + num_neg_eghwt_rand))
println("compare eGHWT with NNSVD-LRC ", num_pos_eghwt_newsvd/(num_pos_eghwt_newsvd + num_neg_eghwt_newsvd))


using CSV
using Plots
using Unitful
using UnitfulRecipes
using Measurements

# We're going to overengineer a bit.

struct Calibration
    wavelength::Vector{typeof(1.0u"nm")}
    counterAvg::Vector{Measurement}
    gradient
    yintercept
end

"""Load calibration data from path string and construct Calibration object."""
function Calibration(path::String)
    headers = ["order", "wavelength", "energy", "min", "max", "avg", "err"] # Custom column headers
    calData = CSV.File(path, header=headers, skipto=2) # skip to first row of data
    wavelength = Quantity.(calData.:wavelength, u"nm")
    counterAvg = calData.:avg .± calData.:err

    # Let's use linear least squares to find the fit
    # Can fit to wavelength = m1 * counts + m2
    # This makes our data kernel:
    kernel = [counterAvg ones(length(counterAvg))]
    # Generalised inverse
    genInv = inv(transpose(kernel) * kernel) * transpose(kernel)
    parameters = genInv * wavelength
    #display(parameters)

    Calibration(wavelength, counterAvg, parameters...)
end

# Expects to be passed an object of type Calibration
@userplot CalibrationPlot
@recipe function f(i::CalibrationPlot)
    cal = i.args[1]
    rawX = cal.counterAvg
    rawY = cal.wavelength

    legend := :bottomright
    xguide := "Counts"
    yguide := "Wavelength"
    grid := true
    minorgrid := true
    seriescolor := :auto

    @series begin
        seriestype := :scatter
        label := "Measurements"
        rawX, rawY
    end

    @series begin
        seriestype := :line
        label := "Fit"
        linestyle := :dash
        y = cal.gradient * rawX .+ cal.yintercept
        rawX, y
    end

end

cal = Calibration("data_calibration_v2.csv")
display(calibrationplot(cal))

λ(count, cal::Calibration=cal) = cal.gradient * count + cal.yintercept
E(λ) = uconvert(u"eV", u"h"* u"c0" / λ)

#------------------------------------------------

data = CSV.File("data_sample_v2.csv")

mutable struct OpticalData
    counter::Vector{Int}
    wavelength::Vector{typeof(1.0u"nm")}
    energy::Vector{typeof((1.0 ± 0.1)u"eV")}
    RawSi::Vector{typeof(1.0u"mV")}
    RawInP::Vector{typeof(1.0u"mV")}
    ScaledSi::Vector{<:Real}
    ScaledInP::Vector{<:Real}
end

"""Load data from path string and construct OpticalData object."""
function OpticalData(path::String)
    headers = ["counter", "wavelength", "energy", "spacer", "withoutInP", "withoutInPScaled", "withInP", "spacer2", "withInPScaled"] # Custom column headers
    data = CSV.File(path, header=headers, skipto=2) # skip to first row of data

    counter = data.:counter
    wavelength = Quantity.(data.:wavelength, u"nm")
    energy = E.(λ.(counter))
    RawSi = Quantity.(data.:withoutInP, u"mV")
    RawInP = Quantity.(data.:withInP, u"mV")
    ScaledSi = RawSi ./ maximum(RawSi)
    ScaledInP = RawInP ./ maximum(RawInP)

    OpticalData(counter, wavelength, energy, RawSi, RawInP, ScaledSi, ScaledInP)
end

data = OpticalData("data_sample_v2.csv")

# Expects to be passed x=energy, y=relative intensity and an optional label or any other plot options
@userplot OpticalDataPlot
@recipe function f(i::OpticalDataPlot)
    x, y = i.args

    legend := :topright
    xguide := "Incident Energy"
    yguide := "Relative Intensity"
    grid := true
    minorgrid := true

    x, y
end

SiPlot = opticaldataplot(data.energy, data.ScaledSi, label="Si")
InPPlot = opticaldataplot(data.energy, data.ScaledInP, label="InP", xlims=(-Inf, 1.5))
display(SiPlot)
display(InPPlot)

@userplot DirectnessCheck
@recipe function f(i::DirectnessCheck)
    energy, intensity = i.args
    invLog = log.(data.RawSi ./ intensity)

    legend := :topleft
    xguide := "Incident Energy"
    yguide := "f(αx)"
    grid := true
    minorgrid := true

    @series begin
        label := "Direct"
        energy, invLog .^ 2
    end

    @series begin
        label := "Indirect"
        energy, sqrt.(invLog)
    end

    @series begin
        # Filter out infinity and NaN rows
        arr = [energy invLog]
        replace!(arr, Inf=>NaN)
        nanrows = any(isnan.(arr), dims=2)
        arr = arr[.!vec(nanrows), :]
        arr[:,2] = arr[:,2] .^ 2

        line = arr[arr[:,1] .> 1.34u"eV", :]
        kernel = [ustrip.(line[:,1]) ones(length(line[:,1]))]
        genInv = inv(transpose(kernel) * kernel) * transpose(kernel)
        parameters = genInv * line[:,2]
        display(parameters)

        fline(x) = parameters[1].val * x + parameters[2].val
        x = LinRange(1.35, 1.385, 2)
        y = fline.(x)

        linestyle := :dash
        label := "Fit"

        x, y
    end

end

LinearPlot = directnesscheck(data.energy, data.RawInP, xlims=(1.1, 1.4))
display(LinearPlot)

# -------------------------------------------------


abstract type GeometricEntity end

struct Segment <: GeometricEntity
    l::Symbol
    r::Symbol
end
show(io::IO, AB::Segment) = print(io, AB.l, AB.r)
const segment = Segment

struct Circle <: GeometricEntity
    O::Symbol
    A::Symbol
end
show(io::IO, γ::Circle) = print(io, "circle center ", γ.O, " radius ", γ.O, γ.A)
const circle = Circle


module Ex5_11
using FaultDetectionTools
using DescriptorSystems
using LinearAlgebra
using Polynomials
using Test

# Example 5.11 - Solution of an AFDIP
println("Example 5.11")

# define s as an improper transfer function
s = rtf('s');
# define Gu(s), Gf(s) and Gw(s)
Gu = [(s+1)/(s+2); (s+2)/(s+3)];     # enter Gu(s)
Gw = [1/(s+2); 0];                   # enter Gw(s)
Gf = [(s+1)/(s+2) 0; 0 1];           # enter Gf(s)
p = 2; mu = 1; mw = 1; mf = 2;       # enter dimensions
S = eye(mf) .> 0

# Procedure AFDI

# Step 1): choose the left nullspace as Q1 = [I -Gu] and
# form Rf1 = Q1*[Gf;0] and Rw1 = Q1*[Gw;0]
Q1 = dss([eye(p) -Gu]); Rf1 = dss(Gf); Rw1 = dss(Gw);
atol = 1.e-7;

# Step 2): determine Q[i] and corresponding Rf[i] and Rw[i]

# initialize overall filter Q and fault detection system Rf
nb = size(S,1);      # number of necessary filters
Qt = similar(Vector{typeof(Q1)},nb)
Rfwt = similar(Vector{typeof(Q1)},nb)
scale = [1, -1]  # set appropriate scaling to match example

for i = 1:nb
    # perform Procedure AFD or EFD to compute Q[i]
    indd = Vector(1:mf)[S[i,:] .== false]
    Qi1 = glnull(Rf1[:,indd])[1]
    # initialize Q[i], Rf[i] and Rw[i]
    Qi = Qi1*Q1; Rfi = Qi1*Rf1; Rwi = Qi1*Rw1;

    if norm(evalfr(Rwi,rand())) > 0.0001
       # compute the quasi-co-outer-co-inner factorization
       Rwi, Rwo, = goifac(Rwi;atol)
       # update
       Qi = Rwo\Qi; Rfi = Rwo\Rfi
    end
    # update the solution if [Q[i], Rf[i], Rw[i]] is improper or unstable
    Qi_Rfi_Rwi, M = glcf([Qi Rfi Rwi], sdeg = -3, smarg = -3, mindeg = true)

    # adjust denominator M to unit infinity norm to match example
    Mnorm = ghinfnorm(M)[1]
    Qt[i] = (scale[i]/Mnorm)*Qi_Rfi_Rwi[:,1:p+mu]
    Rfwt[i] = (scale[i]/Mnorm)*Qi_Rfi_Rwi[:,p+mu+1:end]

    println("Q[$i] = $(dss2rm(Qt[i], atol = 1.e-7))")
    println("Rf[$i] = $(dss2rm(Rfwt[i][:,1:mf], atol = 1.e-7))")
    println("Rw[$i] = $(dss2rm(Rfwt[i][:,mf+1:end], atol = 1.e-7))")
end
Q = FDIFilter(gminreal.(Qt; atol), p, mu)
R = FDIFilterIF(gminreal.(Rfwt; atol), 0, 0, mf, mw)
println("Q = ")
display(Q)
println("R = ")
display(R)
println("gap = $(fdif2ngap(R,S)[1])")

end # module


"""
    Gaussian(θ, γ)

Gaussian layer, with location parameters `θ` and scale parameters `γ`.
"""
struct Gaussian{N,Aθ,Aγ} <: AbstractLayer{N}
    θ::Aθ
    γ::Aγ
    function Gaussian(θ::AbstractArray, γ::AbstractArray)
        @assert size(θ) == size(γ)
        return new{ndims(θ), typeof(θ), typeof(γ)}(θ, γ)
    end
end

Gaussian(::Type{T}, n::Int...) where {T} = Gaussian(zeros(T, n...), ones(T, n...))
Gaussian(n::Int...) = Gaussian(Float64, n...)

Base.repeat(l::Gaussian, n::Int...) = Gaussian(repeat(l.θ, n...), repeat(l.γ, n...))

energies(layer::Gaussian, x::AbstractArray) = gauss_energy.(layer.θ, layer.γ, x)

function free_energies(l::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    return @. -(l.θ .+ inputs)^2 / abs(2l.γ) + log(abs(l.γ)/π/2) / 2
end

function sample_from_inputs(layer::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    μ = mean_from_inputs(layer, inputs)
    σ = sqrt.(var_from_inputs(layer, inputs))
    z = randn!(similar(μ))
    return μ .+ σ .* z
end

mean_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0) = (l.θ .+ inputs) ./ abs.(l.γ)
var_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0) = inv.(abs.(l.γ .+ zero(inputs)))
std_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0) = sqrt.(var_from_inputs(l, inputs))
mode_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0) = mean_from_inputs(l, inputs)

function meanvar_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    return mean_from_inputs(l, inputs), var_from_inputs(l, inputs)
end

function mean_abs_from_inputs(layer::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    μ = mean_from_inputs(layer, inputs)
    ν = var_from_inputs(layer, inputs)
    return @. √(2ν/π) * exp(-μ^2 / (2ν)) + μ * erf(μ / √(2ν))
end

gauss_energy(θ::Real, γ::Real, x::Real) = (abs(γ) * x / 2 - θ) * x

function ∂free_energies(layer::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    θ = layer.θ .+ inputs
    abs_γ = abs.(layer.γ)
    return (
        θ = -θ ./ abs_γ,
        γ = @. sign(layer.γ) * (abs_γ + θ^2) / abs_γ^2 / 2
    )
end

struct GaussStats{A<:AbstractArray}
    x1::A
    x2::A
    function GaussStats(layer::AbstractLayer, data::AbstractArray; wts = nothing)
        x1 = batchmean(layer, data; wts)
        x2 = batchmean(layer, data.^2; wts)
        return new{typeof(x1)}(x1, x2)
    end
end

Base.size(stats::GaussStats) = size(stats.x1)
suffstats(layer::Gaussian, x::AbstractArray; wts = nothing) = GaussStats(layer, x; wts)

function ∂energy(layer::Gaussian, stats::GaussStats)
    @assert size(layer) == size(stats)
    return (; θ = -stats.x1, γ = sign.(layer.γ) .* stats.x2 / 2)
end


using Test
using StaticArrays
using FiniteVolumes
using FiniteVolumes.CartesianMeshes: Half, dx, _direction

@testset "Mesh" begin
    @test dx(CartesianMesh(10)) == SVector(0.1)
    @test dx(CartesianMesh(10, 10)) == SVector(0.1, 0.1)
    @test dx(CartesianMesh(SVector(1.0, 0.0), SVector(2.0, 2.0), SVector(10, 20))) == SVector(0.1, 0.1)

    @test FiniteVolumes.nb_cells(CartesianMesh(10)) == 10
    @test FiniteVolumes.nb_cells(CartesianMesh(5, 5)) == 25
    @test FiniteVolumes.nb_cells(CartesianMesh(2, 5)) == 10

    @test CartesianIndex(3, 3) in FiniteVolumes.all_cells(CartesianMesh(10, 10))
    @test FiniteVolumes.all_cells(CartesianMesh(2)) |> length == 2
    @test FiniteVolumes.all_cells(CartesianMesh(2, 2)) |> length == 4

    @test FiniteVolumes.inner_faces(CartesianMesh(2)) |> collect |> length == 1
    @test FiniteVolumes.inner_faces(PeriodicCartesianMesh(10, 10)) |> collect |> length == 200

    for m in [CartesianMesh(10), PeriodicCartesianMesh(3), CartesianMesh(5, 2), PeriodicCartesianMesh(10, 4)]
        @test FiniteVolumes.inner_faces(m) |> collect |> length == FiniteVolumes.nb_inner_faces(m)
        @test FiniteVolumes.boundary_faces(m) |> collect |> length == FiniteVolumes.nb_boundary_faces(m)
    end

    @test FiniteVolumes.boundary_faces(CartesianMesh(2)) |> collect |> length == 2
    # @test FiniteVolumes.boundary_faces(CartesianMesh(5, 5)) |> collect |> length == 20

    @test _direction((Half(3), Half(2))) == 1
    @test _direction((Half(2), Half(3))) == 2

    @test FiniteVolumes.cells_next_to_inner_face(PeriodicCartesianMesh(10), (Half(19),)) == (CartesianIndex(9), CartesianIndex(10))
    @test FiniteVolumes.cells_next_to_inner_face(PeriodicCartesianMesh(10), (Half(19),)) == (CartesianIndex(9), CartesianIndex(10))
    @test FiniteVolumes.cells_next_to_inner_face(PeriodicCartesianMesh(10), (Half(1),)) == (CartesianIndex(10), CartesianIndex(1))

    @test FiniteVolumes.cells_next_to_inner_face(CartesianMesh(10, 10), (Half(3), Half(2))) == (CartesianIndex(1, 1), CartesianIndex(2, 1))
    @test FiniteVolumes.cells_next_to_inner_face(CartesianMesh(10, 10), (Half(2), Half(3))) == (CartesianIndex(1, 1), CartesianIndex(1, 2))
    # @test FiniteVolumes.cell_next_to_boundary_face(CartesianMesh(10, 10), (Half(1), Half(2))) == CartesianIndex(1, 1)
    # @test FiniteVolumes.cell_next_to_boundary_face(CartesianMesh(10, 10), (Half(21), Half(2))) == CartesianIndex(10, 1)

    @test FiniteVolumes.cell_volume(CartesianMesh(10), CartesianIndex(3)) == 0.1
    @test FiniteVolumes.cell_volume(CartesianMesh(2, 5), CartesianIndex(1, 1)) == 0.1

    @test FiniteVolumes.face_area(CartesianMesh(5, 2), (Half(2), Half(1))) == 0.2
    @test FiniteVolumes.face_area(CartesianMesh(5, 2), (Half(1), Half(2))) == 0.5

    @test FiniteVolumes.cell_center(CartesianMesh(5, 5), CartesianIndex(1, 1)) ≈ [0.1, 0.1]
    @test FiniteVolumes.cell_center(CartesianMesh(5, 5), CartesianIndex(2, 1)) ≈ [0.3, 0.1]
    @test FiniteVolumes.face_center(CartesianMesh(5, 5), (Half(3), Half(2))) == [0.2, 0.1]
    @test FiniteVolumes.face_center(CartesianMesh(5, 5), (Half(2), Half(3))) == [0.1, 0.2]

    # @test FiniteVolumes.cell_corners(CartesianMesh(5, 5), ).bottom_left == @SVector [0.0, 0.0]
    # @test FiniteVolumes.cell_corners(CartesianMesh(5, 5), ).top_right == @SVector [1.0, 1.0]
end

# @btime cells_next_to_inner_face(mesh, i_face) setup=(mesh=CartesianMesh(10); i_face=(Half(19),))
# @btime cells_next_to_inner_face(mesh, i_face) setup=(mesh=PeriodicCartesianMesh(10); i_face=(Half(19),))
# @btime cell_center(mesh, (3,)) setup=(mesh=CartesianMesh(10))
# @btime face_center(mesh, (Half(3),)) setup=(mesh=CartesianMesh(10))


# Toy problem to test the solver

using luteos

P = P1() # Polynomial order of solution

mesh   = Mesh2D( "square", P, N = 11)
master = Master2D( P )

mat = Material( E = 1, ν = 0.0, dim = mesh.dim )

# source function
function funcS( p::Array{Float64} )
  return fill(1.0, size(p,1), 2)
end

function funcB( p::Array{Float64} )
  return fill(0.0, size(p,1), 2)
end

prob = Elas( "Example 1", mat, funcS, [1,1,1,1], true, [funcB, funcB, funcB, funcB] )

(uhath, uh, σh ) = hdgSolve( master, mesh, prob )

# write solution
# writeTecplot( "bla.dat", prob, mesh, uh, σh )


"""

"""
function potentialflow(ctx, A, V, W)
    n = getproperty(ctx, :n)

    Φ = scalar(:Φ, n)

    U = gradient(neu, ctx, A, V, 0, Φ, 0)
    ΔΦ = divergence(neu, ctx, A, U, W)

    sym = linearize(star, ctx, ΔΦ, Φ)

    # linearity of gradient and divergence
    num = Dict{Int,Vector{Float64}}()
    for (key, val) in sym
        num[key] = getproperty.(val, Ref(:val))
    end

    Δ = spdiagm(num...)

    for i in axes(Δ, 1)
        Δ[i, i] = abs(Δ[i, i]) ≤ eps(Δ[i, i]) ? 1. : Δ[i, i]
    end

    U = gradient(neu, ctx, A, V, 0, zeros(prod(n)), 0)
    b = divergence(neu, ctx, A, U, W)

    phi = -Δ \ b

    phi, gradient(neu, ctx, A, V, 0, phi, 0)
end


export DataSet
export CosSin
export FaouSchratz
export BaoDong

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

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

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

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

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

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

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

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

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

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

    end

end

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

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

    sigma    = 1
    llambda  = -4

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

end

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

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

    sigma   = 1
    llambda = -1

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

end

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

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

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

    sigma   = 1
    llambda = -1

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

end


@testset "Isomorphisms" begin
  K = GF(7)
  E = EllipticCurve(K, [1, 2, 3, 4, 5])
  P = rand(E)
  f = @inferred negation_map(E)
  @test f(P) == -P

  K = GF(2,4)
  a = gen(K)
  E1 = elliptic_curve_from_j_invariant(K(0))
  E2 = elliptic_curve_from_j_invariant(a)

  L = GF(3,2)
  b = gen(L)
  E3 = elliptic_curve_from_j_invariant(L(0))
  E4 = elliptic_curve_from_j_invariant(L(2))
  E5 = elliptic_curve_from_j_invariant(12)

  F1, phi1 = transform_rstu(E1, [1, a, 0, 1])
  F2, phi2 = transform_rstu(E2, [0, 1, 1, a])
  F3, phi3 = transform_rstu(E3, [1, b^2, b, 1])
  F4, phi4 = transform_rstu(E4, [0, 1, b^5, b - 1])
  F5, phi5 = transform_rstu(E5, [2, 0, 1, 3])

  @test @inferred is_isomorphic(E1, F1)
  @test @inferred is_isomorphic(E2, F2)
  @test @inferred is_isomorphic(E3, F3)
  @test @inferred is_isomorphic(E4, F4)
  @test @inferred is_isomorphic(E5, F5)

  psi1 = @inferred isomorphism(E1, F1)
  psi2 = @inferred isomorphism(E2, F2)
  psi3 = @inferred isomorphism(E3, F3)
  psi4 = @inferred isomorphism(E4, F4)
  psi5 = @inferred isomorphism(E5, F5)

  @test isomorphism_data(psi2 * inv(phi2)) == isomorphism_data(identity_map(E2)) || isomorphism_data(psi2 * inv(phi2)) == isomorphism_data(negation_map(E2))
  @test (psi4 * inv(phi4)) == identity_map(E4) || (psi4 * inv(phi4)) == negation_map(E4)
  @test (psi5 * inv(phi5)) == identity_map(E5) || (psi5 * inv(phi5)) == negation_map(E5)

  # automorphism group
  E6 = EllipticCurve([1, 2])
  K, = Hecke.rationals_as_number_field()
  E7 = EllipticCurve(K, [1, 2])
  Es = [E1, E2, E3, E4, E5, E6, E7]
  ids = [(24, 3), (2, 1), (12, 1), (2, 1), (2, 1), (2, 1), (2, 1)]
  for i in 1:length(Es)
    E = Es[i]
    G, GtoAut = @inferred automorphism_group(E)
    @test find_small_group(G)[1] == ids[i]
    for g in G
      a = GtoAut(g)
      @test domain(a) == codomain(a) == E
      for h in G
        @test GtoAut(g) * GtoAut(h) == GtoAut(g * h)
      end
      @test GtoAut\(GtoAut(g)) == g
    end
  end
end