using BlockArrays, SparseArrays, Setfield
import Base: length

# 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)
length(ms::TimeMesh{Ti}) where Ti = length(ms.ds)
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

####################################################################################################
# method using the Trapezoidal rule (Order 2 in time) and discretisation of the periodic orbit.
"""
	pb = PeriodicOrbitTrapProblem(F, J, ϕ, xπ, M::Int)
This composite type implements Finite Differences based on a Trapezoidal rule to locate periodic orbits. The arguments are as follows
- `F(x,p)` vector field
- `J` jacobian of `F`. Can be matrix based `J(x,p)` or Matrix-Free
- `Jt = nothing` jacobian tranpose 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)`.`
- `ϕ` used to set a section for the phase constraint equation
- `xπ` used in the section for the phase constraint equation
- `M::Int` number of time slices
- `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)

You can then call `pb(orbitguess, p)` to compute the functional on a `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 of the limit cycle.

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

 ``\\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}``. Finally, the phase of the periodic orbit is constrained by using a section

 ``\\langle x[1] - x_\\pi, \\phi\\rangle=0.``

 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. One may have to redefine it like `extractPeriodFDTrap(x::CuArray) = x[end:end]` or something else. Also, note that you must pass the option `ongpu = true` for the functional to be evaluated efficiently on the gpu.
"""
@with_kw struct PeriodicOrbitTrapProblem{TF, TJ, TJt, Td2F, Td3F, vectype, Tls <: AbstractLinearSolver, Tm} <: AbstractPOFDProblem
	# Function F(x, par)
	F::TF = nothing

	# Jacobian of F wrt x
	J::TJ = nothing

	# Jacobian transpose of F wrt x. This is mainly used for matrix-free computation of Folds of limit cycles
	Jt::TJt = nothing

	# Hessian of F wrt x, useful for continuation of Fold of periodic orbits
	d2F::Td2F = nothing

	# 3rd differential of F wrt 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::Tm = 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
end

isInplace(pb::PeriodicOrbitTrapProblem) = pb.isinplace
onGpu(pb::PeriodicOrbitTrapProblem) = pb.ongpu
hasHessian(pb::PeriodicOrbitTrapProblem) = pb.d2F == nothing
@inline getTimeStep(pb::PeriodicOrbitTrapProblem, i::Int) = getTimeStep(pb.mesh, i)

function applyF(pb::PeriodicOrbitTrapProblem, dest, x, p)
	if isInplace(pb)
		pb.F(dest, x, p)
	else
		dest .= pb.F(x, p)
	end
	dest
end

function applyJ(pb::PeriodicOrbitTrapProblem, dest, x, p, dx)
	if isInplace(pb)
		pb.J(dest, x, p, dx)
	else
		dest .= apply(pb.J(x, p), dx)
	end
	dest
end

# 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) 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, linsolver = ls, isinplace = isinplace, ongpu = ongpu, adaptmesh = adaptmesh)
end

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

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

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

function POTrapScheme!(pb::PeriodicOrbitTrapProblem, dest, u1, u2, par, h::Number, tmp, linear::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
		# 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}}
		dest .= u1 .- u2 .- h .* (dest .+ tmp)
	else
		dest .-= h .* tmp
		# tmp <- pb.F(u1, par)
		applyF(pb, tmp, u1, par)
		dest .-= h .* tmp
	end
end

function POTrapSchemeJac!(pb::PeriodicOrbitTrapProblem, dest, u1, u2, du1, du2, par, h::Number, tmp)
	# this function implements the basic implicit scheme used for the time integration
	# useful for the matrix-free jacobian
	# basically tmp is dF(u2).du2 (see above for explanation)
	dest .= tmp
	# tmp <- apply(pb.J(u1, par), du1)
	applyJ(pb, tmp, u1, par, du1)
	dest .= du1 .- du2 .- h .* (dest .+ tmp)
end

"""
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::PeriodicOrbitTrapProblem, out, u0, par)
	M = pb.M
	N = pb.N
	T = extractPeriodFDTrap(u0)

	u0c  = extractTimeSlices(u0, N, M)
	outc = extractTimeSlices(out, N, M)

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

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

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

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

	# this is for CuArrays.jl to work in the mode allowscalar(false)
	if onGpu(pb)
		return @views vcat(out[1:end-1], dot(u0c[:, 1], pb.ϕ) - dot(pb.xπ, pb.ϕ)) # this is the phase condition
	else
		out[end] = @views dot(u0c[:, 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 `u0` and applied to `du`.
"""
function POTrapFunctionalJac!(pb::PeriodicOrbitTrapProblem, out, u0, par, du)
	M = pb.M
	N = pb.N
	T  = extractPeriodFDTrap(u0)
	dT = extractPeriodFDTrap(du)

	u0c = extractTimeSlices(u0, N, M)
	outc = extractTimeSlices(out, N, M)
	duc = extractTimeSlices(du, N, M)

	# compute the cyclic part
	@views Jc(pb, outc, u0[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
	# the .+ is for the GPU
	# out .+= @views (pb(vcat(u0[1:end-1], T .+ δ)) .- pb(u0)) ./ δ .* dT
	@views applyF(pb, tmp, u0c[:, M-1], par)

	h = dT * getTimeStep(pb, 1)
	@views POTrapScheme!(pb, outc[:, 1], u0c[:, 1], u0c[:, M-1], par, h/2, tmp, false)
	for ii in 2:M-1
		h = dT * getTimeStep(pb, ii)
		@views POTrapScheme!(pb, outc[:, ii], u0c[:, ii], u0c[:, 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(duc[:, 1], pb.ϕ))
	else
		out[end] = @views dot(duc[:, 1], pb.ϕ)
		return out
	end
end

function (pb::PeriodicOrbitTrapProblem)(u0::AbstractVector, par)
	out = similar(u0)
	POTrapFunctional!(pb, out, u0, par)
end

function (pb::PeriodicOrbitTrapProblem)(u0::AbstractVector, par, du)
	out = similar(du)
	POTrapFunctionalJac!(pb, out, u0, par, du)
end

####################################################################################################
# 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 Agamma!(pb::PeriodicOrbitTrapProblem, outc, u0::AbstractVector, par, du::AbstractVector)
	# u0 of size N * M + 1
	# du of size N * M
	M = pb.M
	N = pb.N
	T = extractPeriodFDTrap(u0)
	u0c = extractTimeSlices(u0, N, M)

	# 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 = pb.M
	N = pb.N

	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 POTrapSchemeJac!(pb, outc[:, 1], u0c[:, 1], u0c[:, M-1],
											duc[:, 1], duc[:, M-1], par, h/2, tmp)

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

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

function Jc(pb::PeriodicOrbitTrapProblem, u0::AbstractVector, par, du::AbstractVector)
	M = pb.M
	N = pb.N

	T = extractPeriodFDTrap(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 jacobianPOFDBlock(pb::PeriodicOrbitTrapProblem, u0::AbstractVector, par; γ = 1.0)
	# extraction of various constants
	M = pb.M
	N = pb.N

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

	In = spdiagm( 0 => ones(N))
	setblock!(Aγ, -γ * In, M, 1)
	setblock!(Aγ,  In,     M, M)
	return Aγ
end

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

	In = spdiagm( 0 => ones(N))
	On = spzeros(N, N)

	u0c = extractTimeSlices(u0, N, M)
	outc = similar(u0c)

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

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

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

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

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

		Jn = In - h/2 .* tmpJ
		setblock!(Jc, Jn, ii, ii)
	end
	return Jc
end

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

cylicPOTrapSparse(pb::PeriodicOrbitTrapProblem, orbitguess0, par) = blockToSparse(cylicPOFDBlock(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.0, δ = 1e-9)
	# extraction of various constants
	M = pb.M
	N = pb.N
	T = extractPeriodFDTrap(u0)
	AγBlock = jacobianPOFDBlock(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:N] .=  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 a Sparse representation and inplace update.
"""
@views function (pb::PeriodicOrbitTrapProblem)(::Val{:JacFullSparseInplace}, J0, u0::AbstractVector, par; γ = 1.0, δ = 1e-9)
		# update J0 inplace assuming that the sparsity pattern of J0 and dG(orbitguess0) are the same
		M = pb.M
		N = pb.N
		T = extractPeriodFDTrap(u0)

		In = spdiagm( 0 => ones(N))
		On = spzeros(N, N)

		u0c = extractTimeSlices(u0, N, M)
		outc = similar(u0c)

		tmpJ = @views 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
		# J0[N*M+1, 1:N] .=  pb.ϕ

		return J0
end

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

	In = spdiagm( 0 => ones(N))
	On = spzeros(N, N)

	u0c = extractTimeSlices(u0, N, M)
	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

	# 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
	# J0[N*M+1, 1:N] .=  pb.ϕ

	return J0
end

function (pb::PeriodicOrbitTrapProblem)(::Val{:JacCyclicSparse}, u0::AbstractVector, par, γ = 1.0)
	# extraction of various constants
	N = pb.N
	AγBlock = jacobianPOFDBlock(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!!
	return Aγ[1:end-N, 1:end-N]
end

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

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

	In = spdiagm( 0 => ones(N))

	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)
	setblock!(A_diagBlock, Jn, 1, 1)

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

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

function getTimeDiff(pb::PeriodicOrbitTrapProblem, u0)
	M = pb.M
	N = pb.N
	T = extractPeriodFDTrap(u0)
	u0c = reshape(u0[1:end-1], N, M)

	res = Float64[]
	for ii in 1:M-1
		push!(res, norm(u0c[:,ii+1].-u0c[:,ii]) * T/M)
	end
	return res
end
####################################################################################################
# The following struct encodes a jacobian of PeriodicOrbitTrapProblem which is a convenient composite type for the computation of Floquet multipliers. Therefore, it is only used in the method continuationPOTrap
mutable struct PeriodicOrbitTrapJacobianFull{Tpb, Tj, vectype, Tpar}
	pb::Tpb								# PeriodicOrbitTrapProblem
	J::Tj								# jacobian of the problem
	orbitguess0::vectype				# point at which the jacobian is computed
	par::Tpar							# parameter passed to F, J...
end

# computation of the jacobian, nothing to be done
(pojacfull::PeriodicOrbitTrapJacobianFull)(x, p) = return pojacfull

# linear solver for the PO functional, akin to a bordered linear solver
@with_kw mutable struct PeriodicOrbitTrapLS{Tl} <: AbstractLinearSolver
	linsolver::Tl = DefaultLS()			# linear solver
end

# linear solver for the jacobian
(pols::PeriodicOrbitTrapLS)(pojacfull::PeriodicOrbitTrapJacobianFull, rhs) = pols.linsolver(pojacfull.J, rhs)

(pols::PeriodicOrbitTrapLS)(pojacfull::PeriodicOrbitTrapJacobianFull, rhs1, rhs2) = pols.linsolver(pojacfull.J, rhs1, rhs2)

####################################################################################################
# Linear solvers of the linearized version of the functional G implemented by PeriodicOrbitTrapProblem

# composite type to encode the Aγ Operator and its associated cyclic matrix
@with_kw mutable struct AγOperator{Tvec, Tjc, T, Tpb, Tpar}
	N::Int64 = 0				    		# dimension of a time slice
	orbitguess::Tvec = zeros(1)				# point at which Aγ is evaluated, of size N * M + 1
	Jc::Tjc	= lu(spdiagm(0 => ones(1)))	    # lu factorisation of the cyclic matrix
	is_matrix_free::Bool = false	    	# whether we consider a sparse matrix representation or a Matrix Free one
	γ::T = 1.0				    			# factor γ can be used to compute Floquet multipliers
	prob::Tpb = nothing						# PO functional, used when is_matrix_free = true
	par::Tpar = nothing						# parameter passed to vector field F(x, par)
end

ismatrixfree(A::AγOperator) = A.is_matrix_free

# function to update the cyclic matrix
function (A::AγOperator)(pb::PeriodicOrbitTrapProblem, orbitguess::AbstractVector, par)
	if ismatrixfree(A) == false
		# we store the lu decomposition of the newly computed cyclic matrix
		A.Jc = SparseArrays.lu(cylicPOTrapSparse(pb, orbitguess, par))
	else
		copyto!(A.orbitguess, orbitguess)
		# update par for Matrix-Free
		A.par = par
	end
end

# linear solver designed specifically to deal with AγOperator
@with_kw struct AγLinearSolver{Tls} <: AbstractLinearSolver
	# Linear solver to invert the cyclic matrix Jc contained in Aγ
	linsolver::Tls = DefaultLS()
end

# this function is called whenever one wants to invert Aγ
function (ls::AγLinearSolver)(A::AγOperator, rhs)
	N = A.N
	if ismatrixfree(A) == false
		# we invert the cyclic part Jc of Aγ
		xbar, flag, numiter = @views ls.linsolver(A.Jc, rhs[1:end - N])
		!flag && @warn "Sparse solver for Aγ did not converge"
	else
		# we invert the cyclic part Jc of Aγ
		xbar, flag, numiter = @views 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"
	end
	x = similar(rhs)
	x[1:end-N] .= xbar
	x[end-N+1:end] .= @views A.γ .* x[1:N] .+ rhs[end-N+1:end]
	return x ,flag, numiter
end

####################################################################################################
# The following structure encodes a jacobian of PeriodicOrbitTrapProblem which eases the use of PeriodicOrbitTrapBLS. It is made so that accessing to 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 mutable struct PeriodicOrbitTrapJacobianBordered{Tpb <: PeriodicOrbitTrapProblem, T∂, vectype, Tpar}
	pb::Tpb								# PeriodicOrbitTrapProblem
	∂TGpo::T∂	   = nothing			# derivative of the PO functional G wrt T
	Aγ::AγOperator = AγOperator()		# Aγ Operator involved in the Jacobian of the PO functional
	orbitguess0::vectype = nothing		# point at which the jacobian is computed
	par::Tpar							# parameter passed to F, J...
end

# this function is called whenever the jacobian of G has to be updated
function (J::PeriodicOrbitTrapJacobianBordered)(orbitguess0::AbstractVector, par; δ = 1e-9)
	# u0 must be an orbit guess
	@views J.orbitguess0 .= orbitguess0[1:length(J.orbitguess0)]

	# we compute the derivative of the problem wrt the period TODO: remove this or improve!!
	T = extractPeriodFDTrap(orbitguess0)
	# TODO REMOVE CE vcat!
	@views J.∂TGpo .= (J.pb(vcat(orbitguess0[1:end-1], T + δ), par) .- J.pb(orbitguess0, par)) ./ δ

	# update Aγ
	J.Aγ(J.pb, orbitguess0, par)

	# 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 mutable struct PeriodicOrbitTrapBLS{Tl} <: AbstractLinearSolver
	linsolverbls::Tl = BorderingBLS(AγLinearSolver())	# linear solver
end

# Linear solver associated to PeriodicOrbitTrapJacobianBordered
function (ls::PeriodicOrbitTrapBLS)(J::PeriodicOrbitTrapJacobianBordered, rhs)
	N = J.pb.N

	# TODO REMOVE THIS HACK
	ϕ = zeros(length(rhs)-1)
	ϕ[1:N] .= J.pb.ϕ

	# we solve the bordered linear system as follows
	dX, dl, flag, liniter = @views ls.linsolverbls(J.Aγ, J.∂TGpo[1:end-1],
	 										 		  ϕ, 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::PeriodicOrbitTrapJacobianBordered, 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::PeriodicOrbitTrapJacobianBordered, rhs1, rhs2)

####################################################################################################
# newton wrappers
function _newton(probPO::PeriodicOrbitTrapProblem, orbitguess, par, options::NewtonPar, linearPO::Symbol = :BorderedLU; defOp::Union{Nothing, DeflationOperator{T, Tf, vectype}} = nothing, kwargs...) where {T, Tf, vectype}
	@assert linearPO in [:FullLU, :BorderedLU, :FullMatrixFree, :BorderedMatrixFree, :FullSparseInplace]
	N = probPO.N
	M = probPO.M

	if linearPO in [:FullLU, :FullMatrixFree, :FullSparseInplace]
		@assert orbitguess isa AbstractVector
		@assert length(orbitguess) == N * M + 1 "Error with size of the orbitguess"

		if linearPO == :FullLU
			jac = (x, p) -> probPO(Val(:JacFullSparse), x, p)
		elseif linearPO == :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)
		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
		@assert orbitguess isa AbstractVector
		@assert length(orbitguess) == N * M + 1 "Error with size of the orbitguess"

		if linearPO == :BorderedLU
			Aγ = AγOperator(is_matrix_free = false, N = probPO.N, Jc = lu(spdiagm( 0 => ones(N * (M - 1)) )) )
			# linear solver
			lspo = PeriodicOrbitTrapBLS()
		else	# :BorderedMatrixFree
			Aγ = AγOperator(is_matrix_free = true, prob = probPO, N = probPO.N, orbitguess = zeros(N * M + 1), Jc = lu(spdiagm( 0 => ones(N * (M - 1)) )), par = par)
			# linear solver
			lspo = PeriodicOrbitTrapBLS(BorderingBLS(AγLinearSolver(options.linsolver)))
		end

		# create the jacobian
		jacPO = PeriodicOrbitTrapJacobianBordered(probPO, zeros(N * M + 1), Aγ, zeros(N * M + 1), par)

		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

"""
	newton(prob::PeriodicOrbitTrapProblem, orbitguess, options::NewtonPar, linearPO = :BorderedLU; kwargs...)

This is the Newton-Krylov 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`.
- `options` same as for the regular `newton` method
- `linearPO = :BorderedLU`. 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.
    - 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 `: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. This is the default algorithm.
    - 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.

# Output:
- solution
- history of residuals
- flag of convergence
- number of iterations
"""

newton(probPO::PeriodicOrbitTrapProblem, orbitguess, par, options::NewtonPar; linearPO::Symbol = :BorderedLU, kwargs...) = _newton(probPO, orbitguess, par, options, linearPO; defOp = nothing, kwargs...)

"""
	newton(probPO::PeriodicOrbitTrapProblem, orbitguess, options::NewtonPar, defOp::DeflationOperator{T, Tf, vectype}, linearPO = :BorderedLU; kwargs...) where {T, Tf, vectype}

This function is similar to `newton(probPO, orbitguess, options, linearPO; 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{T, Tf, vectype}; linearPO::Symbol = :BorderedLU, kwargs...) where {T, Tf, vectype} = _newton(probPO, orbitguess, par, options, linearPO; defOp = defOp, kwargs...)

####################################################################################################
# continuation wrapper

"""
	continuationPOTrap(probPO::PeriodicOrbitTrapProblem, orbitguess, par, lens::Lens, _contParams::ContinuationPar, linearAlgo::AbstractBorderedLinearSolver; linearPO = :BorderedLU, printSolution = (u,p) -> u[end], kwargs...)

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)
- `linearPO = :BorderedLU`. Same as `newton` when applied to `PeriodicOrbitTrapProblem`. More precisely:
    - For `:FullLU`, we use the default linear solver on a sparse matrix representation of `dG`. This matrix is assembled at each newton iteration.
    - 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 `:BorderedLU`, we take advantage of the bordered shape of the linear solver and use LU decomposition to invert `dG` using a bordered linear solver. This is the default algorithm.
    - 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.

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 `printSolution` argument.
"""
function continuationPOTrap(probPO::PeriodicOrbitTrapProblem, orbitguess, par, lens::Lens, _contParams::ContinuationPar, linearAlgo::AbstractBorderedLinearSolver; linearPO = :BorderedLU, printSolution = (u,p) -> u[end], kwargs...)
	@assert linearPO in [:FullLU, :FullMatrixFree, :BorderedLU, :BorderedMatrixFree, :FullSparseInplace]
	contParams = check(_contParams)

	N = probPO.N
	M = probPO.M
	options = contParams.newtonOptions

	if contParams.computeEigenValues
		contParams = @set contParams.newtonOptions.eigsolver = FloquetQaDTrap(contParams.newtonOptions.eigsolver)
	end

	if linearPO in [:FullLU, :FullMatrixFree, :FullSparseInplace]
		@assert length(orbitguess) == N * M + 1 "Error with size of the orbitguess"

		if linearPO == :FullLU
			jac = (x, p) -> probPO(Val(:JacFullSparse), x, p)
		elseif linearPO == :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)
		else
		 	jac = (x, p) ->   ( dx -> probPO(x, p, dx))
		end

		lspo = PeriodicOrbitTrapLS(options.linsolver)

		return continuation(
			probPO,
			(x, p) -> PeriodicOrbitTrapJacobianFull(probPO, jac(x, p), x, p),
			orbitguess, par, lens,
			(@set contParams.newtonOptions.linsolver = lspo);
			printSolution = printSolution,
			kwargs...)
	else
		@assert orbitguess isa AbstractVector
		@assert length(orbitguess) == N * M + 1 "Error with size of the orbitguess"

		if linearPO == :BorderedLU
			Aγ = AγOperator(is_matrix_free = false, N = N,
					Jc = lu(spdiagm( 0 => ones(N * (M - 1)) )) )
			# linear solver
			lspo = PeriodicOrbitTrapBLS()
		else
			Aγ = AγOperator(is_matrix_free = true, prob = probPO, N = N,
					orbitguess = zeros(N * M + 1),
					Jc = lu(spdiagm( 0 => ones(N * (M - 1)) )), par = par)
			# linear solver
			lspo = PeriodicOrbitTrapBLS(BorderingBLS(AγLinearSolver(options.linsolver)))
		end

		# create the jacobian
		jacPO = PeriodicOrbitTrapJacobianBordered(probPO, zeros(N * M + 1), Aγ, zeros(N * M + 1), par)

		return continuation(probPO, jacPO, orbitguess, par, lens,
			(@set contParams.newtonOptions.linsolver = lspo);
			printSolution = printSolution,
			kwargs...)
	end
end

"""
	continuation(probPO, orbitguess, p0::Real, _contParams::ContinuationPar; linearPO = :BorderedLU, printSolution = (u,p) -> u[end], linearAlgo = BorderingBLS(), kwargs...)

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)
- `linearPO = :BorderedLU`. Same as `newton` when applied to `PeriodicOrbitTrapProblem`. More precisely:
    - For `:FullLU`, we use the default linear solver on a sparse matrix representation of `dG`. This matrix is assembled at each newton iteration.
    - 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 `:BorderedLU`, we take advantage of the bordered shape of the linear solver and use LU decomposition to invert `dG` using a bordered linear solver. This is the default algorithm.
    - 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.

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 `printSolution` argument.
"""
function continuation(probPO::PeriodicOrbitTrapProblem, orbitguess, par, lens::Lens, _contParams::ContinuationPar; linearPO = :BorderedLU, printSolution = (u,p) -> u[end], linearAlgo = BorderingBLS(), kwargs...)
	_linearAlgo = @set linearAlgo.solver = _contParams.newtonOptions.linsolver
	return continuationPOTrap(probPO, orbitguess, par, lens, _contParams, _linearAlgo; linearPO = linearPO, printSolution = printSolution, kwargs...)
end

####################################################################################################
# functions needed Branch switching from Hopf bifurcation point
function update(prob::PeriodicOrbitTrapProblem, F, dF, hopfpt, ζr::AbstractVector, M, orbitguess_a, period)
	# 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 = length(ζr), M = M, F = F, J = dF, ϕ = ζr, xπ = hopfpt.x0)
	return probPO, orbitguess
end

####################################################################################################
# Branch switching from BP of PO
"""
$(SIGNATURES)

Branch switching at a Branch 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
- `Jt = (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, `Jt` 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 `Jt = (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
- `linearPO = :BorderedLU` linear solver used for the Newton-Krylov solver when applied to [`PeriodicOrbitTrapProblem`](@ref).
- `printSolution = (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)
"""
function continuationPOTrapBPFromPO(br::ContResult, ind_bif::Int, _contParams::ContinuationPar ; Jt = nothing, δ = 1e-8, δp = 0.1, ampfactor = 1, usedeflation = true, linearPO = :BorderedLU, printSolution = (u,p) -> u[end], linearAlgo = BorderingBLS(), kwargs...)
	verbose = get(kwargs, :verbosity, 0) > 0

	@assert br.functional isa PeriodicOrbitTrapProblem
	@assert abs(br.bifpoint[ind_bif].δ[1]) == 1

	bifpt = br.bifpoint[ind_bif]

	# let us compute the kernel
	λ = (br.eig[bifpt.idx].eigenvals[bifpt.ind_bif])
	verbose && print("--> computing kernel...")
	ζ0 = geteigenvector(br.contparams.newtonOptions.eigsolver, br.eig[bifpt.idx].eigenvec, bifpt.ind_bif)
	verbose && println("Done!")
	# we normalize it by the sup norm because it could be too small/big in L2 norm
	ζ0 ./= norm(ζ0, Inf)

	pb = br.functional

	ζ_a = MonodromyQaDFD(Val(:ExtractEigenVector), pb, bifpt.x, set(br.params, br.param_lens, bifpt.param), real.(ζ0))
	ζ = reduce(vcat, ζ_a)
	# @show λ typeof(reduce(vcat,ζ_a))

	orbitguess = copy(bifpt.x)
	orbitguess[1:end-1] .+= ampfactor .*  ζ
	# orbitguess[end] = 3.4

	newp = bifpt.param + δp

	pb(orbitguess, set(br.params, br.param_lens, 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(pb, bifpt.x, set(br.params, br.param_lens, newp), optn, linearPO; kwargs...)

		# find the bifurcated branch using deflation
		deflationOp = DeflationOperator(2.0, (x,y) -> dot(x[1:end-1], y[1:end-1]), 1.0, [sol0])
		verbose && println("\n--> Compute point on bifurcated branch...")
		solbif, _, flag, _ = newton(pb, orbitguess, set(br.params, br.param_lens, newp), (@set optn.maxIter = 10*optn.maxIter), deflationOp, linearPO; kwargs...)
		@assert flag "Deflated newton did not converge"
		orbitguess .= solbif
	end

	# TODO
	# we have to adjust the constraint. Right now, it can be quite big.

	# perform continuation
	branch, u, tau = continuation(br.functional, orbitguess, set(br.params, br.param_lens, newp), br.param_lens, _contParams; linearPO = linearPO, printSolution = printSolution, linearAlgo = linearAlgo, kwargs...)

	return setproperties(branch; type = :PeriodicOrbit, functional = br.functional), u, tau
end


module TestVecnorm
using ParallelAccelerator

@acc norm_t(y) = norm(y,1)
@acc norm_t2(y) = norm(y,2)

function test()
    y = [1.,2.,3.]
    z = norm_t(y)
    return z
end

function test2()
    y = [1.,2.,3.]
    z = norm_t2(y)
    return z
end

end

using Base.Test
println("testing vecnorm...")
@test_approx_eq TestVecnorm.test() 6.0
@test_approx_eq TestVecnorm.test2() 3.7416573867739413
println("Done testing vecnorm.")


using Sundials
using Base.Test

examples_path = joinpath(dirname(@__FILE__), "..", "examples")

# run cvode example
println("== start cvode Roberts example (simplified)")
let
include(joinpath(examples_path, "cvode_Roberts_simplified.jl"))
println("result at t=$(t[end]):")
println(res[:,end], "\n")
end

println("== start cvode Roberts example")
let
include(joinpath(examples_path, "cvode_Roberts_dns.jl"))
end

# run ida examples
println("== start ida_Roberts example (simplified)")
let
include(joinpath(examples_path, "ida_Roberts_simplified.jl"))
end

println("== start ida_Roberts example")
let
include(joinpath(examples_path, "ida_Roberts_dns.jl"))
println("result at t=$(t[end]):")
println(yout[:,end], "\n")
end

println("== start ida_Heat2D example")
let
include(joinpath(examples_path, "ida_Heat2D.jl"))
println("result at t=$(t[end]):")
println(yout[:,end], "\n")
end

#= requires Grid package
println("== start ida_Cable example")
let
include(joinpath(examples_path, "ida_Cable.jl"))
end
=#

# run kinsol example
println("== start kinsol example (simplified)")
let
include(joinpath(examples_path, "kinsol_mkin_simplified.jl"))
println("solution:")
println(res)
residual = ones(2)
sysfn(res, residual)
println("residual:")
println(residual, "\n")
end

println("== start kinsol example")
let
include(joinpath(examples_path, "kinsol_mkinTest.jl"))
@test abs(minimum(residual)) < 1e-5
end


using Symbolics: scalarize
@variables x[1:3] u[1:2]
@parameters p[1:2] t

g = [x[1] * p[1] + p[2] * x[2]; x[2] * u[1]; u[2] * x[3] + exp(-t)]

b = Basis(g, x, parameters = p, controls = u, iv = t)

# Vector call
x0 = randn(3)
p0 = randn(2)
t0 = 0.0
u0 = randn(2)

true_res(x, p, t, u) =
    [sum(x[1:2] .* p); x[2] .* u[1]; u[2] .* x[3] .+ exp.(-t)]
true_res_(x, p, t) =
    scalarize.(hcat([true_res(x[:, i], p, t[i], zeros(2)) for i = 1:100]...))
true_res_(x, p, t, u) =
    scalarize.(hcat([true_res(x[:, i], p, t[i], u[:, i]) for i = 1:100]...))

@test isequal(b(x0), DataDrivenDiffEq.get_f(b)(x0, p, t, zeros(2)))
@test isequal(b(x0, p0), DataDrivenDiffEq.get_f(b)(x0, p0, t, zeros(2)))
@test isequal(b(x0, p0, t0), DataDrivenDiffEq.get_f(b)(x0, p0, t0, zeros(2)))
@test isequal(b(x0, p0, t0, u0), DataDrivenDiffEq.get_f(b)(x0, p0, t0, u0))

# Array call
x0 = randn(3, 100)
p0 = randn(2)
t0 = randn(100)
u0 = randn(2, 100)

# These first two fail, since exp(-t) != exp(getindex(t,1))
@test isequal(b(x0), true_res_(x0, p, [t for i = 1:100]))
@test isequal(b(x0, p0), true_res_(x0, p0, [t for i = 1:100]))
@test isequal(b(x0, p0, t0), true_res_(x0, p0, t0))
@test isequal(b(x0, p0, t0, u0), true_res_(x0, p0, t0, u0))

@parameters w[1:2] t
@variables u[1:3](t)

h = [u[1]; u[2]; cos(w[1] * u[2] + w[2] * u[3]); 5*u[3] + u[2]]
h_not_unique = [u[1]; u[1]; u[1]^1; h; 1]
basis = Basis(h_not_unique, u, parameters = w, iv = t)
basis_2 =
    Basis(h_not_unique, u, parameters = w, iv = t, linear_independent = true)
# Check getters
@test isequal(states(basis), u)
@test isequal(parameters(basis), w)
@test isequal(ModelingToolkit.get_iv(basis), t)
@test isequal(controls(basis), [])

# Check free parameter calculation
@test free_parameters(basis, operations = [+]) == 7
@test free_parameters(basis_2) == 6
@test free_parameters(basis, operations = [+, cos]) == 8
@test free_parameters(basis_2, operations = [+, cos]) == 7
@test DataDrivenDiffEq.count_operation(
    (1 + cos(u[2]) * sin(u[1]))^3,
    [+, cos, ^, *],
) == 4

# Check array functionalities
basis_2 = unique(basis)
@test isequal(basis, basis_2)
@test size(basis) == (5,)
@test basis([1.0; 2.0; π], [0.0; 1.0]) ≈ [1.0; 2.0; -1.0; 5*π + 2.0; 1.0]
@test size(basis) == size(basis_2)
push!(basis_2, sin(u[2]))
@test size(basis_2)[1] == length(h) + 2
basis_3 = merge(basis, basis_2)
@test size(basis_3) == (6,)
@test isequal(states(basis_3), states(basis_2))
@test isequal(parameters(basis_3), parameters(basis_2))
merge!(basis_3, basis)
@test basis_3 == basis_2
push!(basis, 5*u[3] + u[2])
unique!(basis) # Does not remove
@test size(basis) == (5,)

# Further callables
@variables a
g = [u[1]; u[3]; a]
basis = Basis(g, [u; a])
@test basis([1; 2; 3; 4]) == [1; 3; 4]
g = [u[1]; u[3]; u[2]]
basis = Basis(g, u, parameters = [])
X = ones(Float64, 3, 10)
X[1, :] .= 3 * X[1, :]
X[3, :] .= 5 * X[3, :]
# Check the array evaluation
@test basis(X) ≈ [1.0 0.0 0.0; 0.0 0.0 1.0; 0.0 1.0 0.0] * X
Y = similar(X)
basis(Y, X)
@test Y ≈ [1.0 0.0 0.0; 0.0 0.0 1.0; 0.0 1.0 0.0] * X
f = DataDrivenDiffEq.jacobian(basis)
@test f([1; 1; 1], [0.0; 0.0], 0.0) ≈ [1.0 0.0 0.0; 0.0 0.0 1.0; 0.0 1.0 0.0]
@test_nowarn [xi for xi in basis]
@test_nowarn basis[2:end];

@variables u[1:3] t

g = [u[2]; -sin(u[1]) * exp(-t); u[2] + u[3]]
basis = Basis(g, u, iv = t)

f_(u, p, t) = [u[3]; u[2] * u[1]; p[1] * sin(u[1]) * u[2]; p[2] * t]
b = Basis(f_, u, parameters = w, iv = t)
@test f_([1; 2; 3], [2; 0], 3.0) ≈ b([1; 2; 3], [2; 0], 3.0)


@testset "1631.path-with-minimum-effort.jl" begin
    @test minimum_effort_path([[1,2,2],[3,8,2],[5,3,5]]) == 2
    @test minimum_effort_path([[1,2,3],[3,8,4],[5,3,5]]) == 1
    @test minimum_effort_path([[1,2,1,1,1],[1,2,1,2,1],[1,2,1,2,1],[1,2,1,2,1],[1,1,1,2,1]]) == 0
    @test minimum_effort_path([[1,10,6,7,9,10,4,9]]) == 9
end


# Import modules
using Distributions  # probability distributions and associated functions
using LinearAlgebra
using StatsPlots

using Random
using JLD2

using Plots
# CES
using EnsembleKalmanProcesses.EnsembleKalmanProcessModule
using EnsembleKalmanProcesses.Observations
using EnsembleKalmanProcesses.ParameterDistributionStorage

rng_seed = 4137
Random.seed!(rng_seed)

# Output figure save directory
homedir = pwd()
println(homedir)
figure_save_directory = homedir * "/output/"
data_save_directory = homedir * "/output/"
if ~isdir(figure_save_directory)
    mkdir(figure_save_directory)
end
if ~isdir(data_save_directory)
    mkdir(data_save_directory)
end

###
###  model settings
###

###  Generate (artificial) truth samples
###  Note: The observables y are related to the parameters θ by:
###        y = G(θ) + η
###
# G(u) = [u[1] ; u[2]]
G(u) = [sqrt(u[1] - u[2]^2) ; u[2]]
θ1_true, θ2_true = 2.0, 1.0
prior_mean = [θ1_true + 1.0; θ2_true + 0.5]
prior_cov = Array(Diagonal(fill(0.1^2,  length(prior_mean))))
# prior_cov = Array(Diagonal(fill(1.0^2,  length(prior_mean))))

params_true = [θ1_true, θ2_true]
param_names = ["θ1", "θ2"]
n_param = length(param_names)
params_true = reshape(params_true, (n_param, 1))

println(n_param)
println(params_true)

# prior_stds = [2.0, 0.5*A_true]

prior_distributions = [Parameterized(Normal(prior_mean[1], sqrt(prior_cov[1,1]))), Parameterized(Normal(prior_mean[2], sqrt(prior_cov[2,2])))]
constraints = [[no_constraint()], [no_constraint()]]

priors = ParameterDistribution(prior_distributions, constraints, param_names)

###
###  Define the data from which we want to learn the parameters
###
data_names = ["y0", "y1"]

function run_G_ensemble(params)
    N_ensemble = size(params, 2)
    return hcat([G(params[:, i]) for i in 1:N_ensemble]...)
end
# Lorenz forward
# Input: params: [N_params, N_ens]
# Output: gt: [N_data, N_ens]
# Dropdims of the output since the forward model is only being run with N_ens=1
# corresponding to the truth construction
n_samples = 100
G_t = G(params_true)
y_t = zeros(length(G_t), n_samples)
Γy = Array(Diagonal(fill(0.1^2,  length(G_t))))

# Add noise
for i in 1:n_samples
    y_t[:, i] = G_t .+ rand(MvNormal(zeros(size(G_t)), Γy))
end

# Construct observation object
truth = Observations.Obs(y_t, Γy, data_names)
truth_sample = truth.mean


export CG_generic

"""
A non-linear conjugate gradient algorithm.
Can use differents conjugate gradient formulae:
  - Fletcher & Reeves
  - Hager & Zhang
  - Pollack & Ribière
  - Hestenes & Stiefel

  Inputs:
  - An AbstractNLPModel (our problem)
  - An NLPStopping (stopping fonctions/criterion of our problem)
  - (opt) verbose = true will show function value, norm of gradient and the
    inner product of the gradient and the descent direction
  - (opt) hessian_rep, different way to implement the Hessian. The default
    way is a LinearOperator

  Outputs:
  - An NLPAtX, so all the information of the last iterate
  - A boolean true if we reahced an optimal solution, false otherwise
"""
function CG_generic(nlp        :: AbstractNLPModel,
                    nlp_stop   :: NLPStopping;
                    verbose    :: Bool = false,
                    linesearch :: Function = armijo_ls,
                    CG_formula :: Function = formula_HZ,
                    scaling    :: Bool = true,
                    kwargs...)
    nlp_at_x = nlp_stop.current_state
    x = copy(nlp.meta.x0)
    n = nlp.meta.nvar

    xt = Array{Float64}(undef, n)
    ∇ft = Array{Float64}(undef, n)

    f = obj(nlp, x)
    ∇f = grad(nlp, x)

    iter = 0

    OK = update_and_start!(nlp_stop, x = x, fx = f, gx = ∇f, g0 = ∇f)
    ∇fNorm = norm(nlp_at_x.gx)

    verbose && @printf("%4s  %8s  %7s  %8s \n", " iter", "f", "‖∇f‖", "α")
    verbose && @printf("%5d  %9.2e  %8.1e", iter, nlp_at_x.fx, ∇fNorm)

    β = 0.0
    # d = zeros(∇f)
    d = zero(nlp_at_x.gx)
    scale = 1.0

    h = LineModel(nlp, x, d * scale)

    while !OK
        d = - ∇f + β * d
        slope = ∇f⋅d
        if slope > 0.0  # restart with negative gradient
          #stalled_ascent_dir = true
          d = - ∇f
          slope =  ∇f⋅d
        end

        h = redirect!(h, x, d * scale)
        ls_at_t = LSAtT(0.0, h₀ = nlp_at_x.fx, g₀ = slope)
        stop_ls = LS_Stopping(h, (x, y) -> armijo(x, y; kwargs...), ls_at_t)
        verbose && println(" ")
        ls_at_t, good_step_size = linesearch(h, stop_ls)

        ls_at_t.x *= scale
        xt = nlp_at_x.x + ls_at_t.x * d
        ft = obj(nlp, xt);
        ∇ft = grad(nlp, xt);

        ∇ft = grad!(nlp, xt, ∇ft)
        # Move on.
        s = xt - x
        y = ∇ft - ∇f
        β = 0.0
        if (∇ft⋅∇f) < 0.2 * (∇ft⋅∇ft)   # Powell restart
            β = CG_formula(∇f, ∇ft, s, d)
        end
        if scaling
            scale = (y⋅s) / (y⋅y)
        end
        if scale <= 0.0
            #println(" scale = ",scale)
            #println(" ∇f⋅s = ",∇f⋅s,  " ∇ft⋅s = ",∇ft⋅s)
            scale = 1.0
        end
        x = xt
        f = ft
        dxt = xt - x
        BLAS.blascopy!(n, ∇ft, 1, ∇f, 1)

        ∇fNorm = norm(∇ft, Inf)
        iter = iter + 1
        verbose && @printf("%4d  %9.2e  %8.1e %8.2e", iter, nlp_at_x.fx, ∇fNorm, ls_at_t.x)

        OK = update_and_stop!(nlp_stop, x = xt, fx = ft, gx = ∇ft)
    end
    verbose && @printf("\n")

    return nlp_stop, nlp_stop.meta.optimal
end


using Match

function countletters()
end

function countletters(n)
	assert(0 <= n <= 1000)
	return @match n begin
		0 => length("")
   		1 => length("one")
   		2 => length("two")
   		3 => length("three")
   		4 => length("four")
   		5 => length("five")
   		6 => length("six")
	   	7 => length("seven")
   		8 => length("eight")
   		9 => length("nine")
   		10 => length("ten")
   		11 => length("eleven")
   		12 => length("twelve")
   		13 => length("thirteen")
   		14 => length("fourteen")
   		15 => length("fifteen")
   		16 => length("sixteen")
   		17 => length("seventeen")
   		18 => length("eighteen")
   		19 => length("nineteen")
		1000 => length("onethousand")
   		_    =>
   			if 20 <= n <= 29
   				length("twenty") + countletters(n%10)
   			elseif 30 <= n <= 39
   				length("thirty") + countletters(n%10)
   			elseif 40 <= n <= 49
   				length("forty") + countletters(n%10)
   			elseif 50 <= n <= 59
   				length("fifty") + countletters(n%10)
   			elseif 60 <= n <= 69
   				length("sixty") + countletters(n%10)
   			elseif 70 <= n <= 79
   				length("seventy") + countletters(n%10)
   			elseif 80 <= n <= 89
   				length("eighty") + countletters(n%10)
   			elseif 90 <= n <= 99
   				length("ninety") + countletters(n%10)
   			elseif 100 <= n <= 999
   				countletters(div(n,100)) + length("hundred") + ((n % 100) > 0 ? length("and") + countletters(n % 100) : 0)
   			end
	end
end

function main()
	assert(countletters(342) == 23) #"three hundred and forty two"
	assert(countletters(115) == 20) #"one hundred and fifteen"
	assert(sum(map(countletters,1:5)) == 19)
	assert(sum(map(countletters,1:9)) == 36)
	assert(sum(map(countletters,10:19)) == 70)
	assert(sum(map(countletters,20:99)) == 748)
	assert(sum(map(countletters,100:999)) == 20259)
	N = 1000
	s = sum(map(countletters,1:N))
	println(s)
end

main()


#populate array
# %%

# Atom Editorで起動 uber-junoがすでに導入されているとする。
# ここら辺にカーソルを移動
# option + enter を押す。
# 次の`# %%` が出てくるまでの部分まで実行される
println("ここから処理が始まるよ")
arr = [1,2,3]
@show arr

arr = Float64[]
push!(arr,1)
@show arr

arr = Int64[]
push!(arr,1)
@show arr

println("ここで処理が終わるよ")

# %%
@show collect(1:7)
@show size(collect(1:7),1)

arr = Float64[]
sizehint!(arr,6)
@show size(arr)
for i in 1:10
    push!(arr,i)
end

@show arr

# %%

arr = collect(1:10)
@show arr[3:7]

# %%
@show println(Array{Any}(undef,4))
@show fill!(arr,42)

# %% concatenate arrays

a₁ = [-1,-2,-3]
a₂ = [4,5,6]

append!(a₁,a₂)
@show a₁

# you can't do
#append!(a₁,9)

# %%

arr = collect(1:8)
@show arr
p=pop!(arr)
@show p
@show arr

p=popfirst!(arr)
@show p
@show arr
pushfirst!(arr,-1)
@show arr

# %%
# splice

arr = collect(1:8)

@show arr
splice!(arr,3)
@show arr

# %%

#=
Checking wheather an array contains an element
=#
arr = collect(1:8)

@show　4 in arr
@show -1 in arr
@show in(4,arr)
@show in(-1,arr)

#=
julia> import LinearAlgebra:dot

julia> f = x -> dot(x,x); g = x -> sum(x .* x);

julia> arr =collect(1:1000);

julia> @benchmark f(arr)
BenchmarkTools.Trial:
  memory estimate:  16 bytes
  allocs estimate:  1
  --------------
  minimum time:     264.665 ns (0.00% GC)
  median time:      268.938 ns (0.00% GC)
  mean time:        326.445 ns (10.09% GC)
  maximum time:     276.418 μs (99.82% GC)
  --------------
  samples:          10000
  evals/sample:     328

julia> @benchmark g(arr)
BenchmarkTools.Trial:
  memory estimate:  7.95 KiB
  allocs estimate:  2
  --------------
  minimum time:     715.370 ns (0.00% GC)
  median time:      1.134 μs (0.00% GC)
  mean time:        1.772 μs (37.87% GC)
  maximum time:     897.657 μs (99.85% GC)
  --------------
  samples:          10000
  evals/sample:     108

julia> f(arr) == g(arr)

=#


## Quite a few of these are borrowed from Flux.jl

# Misc
nfan() = 1, 1 # fan_in, fan_out
nfan(n) = 1, n # A vector is treated as a n×1 matrix
nfan(n_out, n_in) = n_in, n_out # In case of Dense kernels: arranged as matrices
nfan(dims::Tuple) = nfan(dims...)
nfan(dims...) = prod(dims[1:(end - 2)]) .* (dims[end - 1], dims[end]) # In case of convolution kernels

# Neural Network Initialization
## NOTE: Would be great if these could be moved into its own package and NN frameworks
##       could just import it.
"""
    zeros32(rng::AbstractRNG, size...) = zeros(Float32, size...)

Return an `Array{Float32}` of zeros of the given `size`. (`rng` is ignored)
"""
zeros32(rng::AbstractRNG, args...; kwargs...) = zeros(rng, Float32, args...; kwargs...)

"""
    ones32(rng::AbstractRNG, size...) = ones(Float32, size...)

Return an `Array{Float32}` of ones of the given `size`. (`rng` is ignored)
"""
ones32(rng::AbstractRNG, args...; kwargs...) = ones(rng, Float32, args...; kwargs...)

Base.zeros(rng::AbstractRNG, args...; kwargs...) = zeros(args...; kwargs...)
Base.ones(rng::AbstractRNG, args...; kwargs...) = ones(args...; kwargs...)

"""
    glorot_uniform(rng::AbstractRNG, size...; gain = 1)

Return an `Array{Float32}` of the given `size` containing random numbers drawn from a uniform distribution on the interval ``[-x, x]``, where `x = gain * sqrt(6 / (fan_in + fan_out))`. This method is described in [1] and also known as Xavier initialization.

# References

[1] Glorot, Xavier, and Yoshua Bengio. "Understanding the difficulty of training deep feedforward neural networks." _Proceedings of the thirteenth international conference on artificial intelligence and statistics_. 2010.
"""
function glorot_uniform(rng::AbstractRNG, dims::Integer...; gain::Real=1)
    scale = Float32(gain) * sqrt(24.0f0 / sum(nfan(dims...)))
    return (rand(rng, Float32, dims...) .- 0.5f0) .* scale
end

"""
    glorot_normal(rng::AbstractRNG, size...; gain = 1)

Return an `Array{Float32}` of the given `size` containing random numbers drawn from a normal distribution with standard deviation `gain * sqrt(2 / (fan_in + fan_out))`. This method is described in [1] and also known as Xavier initialization.

# References

[1] Glorot, Xavier, and Yoshua Bengio. "Understanding the difficulty of training deep feedforward neural networks." _Proceedings of the thirteenth international conference on artificial intelligence and statistics_. 2010.
"""
function glorot_normal(rng::AbstractRNG, dims::Integer...; gain::Real=1)
    std = Float32(gain) * sqrt(2.0f0 / sum(nfan(dims...)))
    return randn(rng, Float32, dims...) .* std
end

# PRNG Handling
replicate(rng::AbstractRNG) = copy(rng)
replicate(rng::CUDA.RNG) = deepcopy(rng)

# Training Check
@inline istraining() = false
@inline istraining(::Val{training}) where {training} = training
@inline istraining(st::NamedTuple) = istraining(st.training)

# Linear Algebra
@inline _norm(x; dims=Colon()) = sqrt.(sum(abs2, x; dims=dims))
@inline function _norm_except(x::AbstractArray{T, N}, except_dim=N) where {T, N}
    _norm(x; dims=filter(i -> i != except_dim, 1:N))
end

# Convolution
function convfilter(rng::AbstractRNG, filter::NTuple{N, Integer},
                    ch::Pair{<:Integer, <:Integer};
                    init=glorot_uniform, groups=1) where {N}
    cin, cout = ch
    @assert cin % groups==0 "Input channel dimension must be divisible by groups."
    @assert cout % groups==0 "Output channel dimension must be divisible by groups."
    return init(rng, filter..., cin ÷ groups, cout)
end

expand(N, i::Tuple) = i
expand(N, i::Integer) = ntuple(_ -> i, N)

_maybetuple_string(pad) = string(pad)
_maybetuple_string(pad::Tuple) = all(==(pad[1]), pad) ? string(pad[1]) : string(pad)

# Padding
struct SamePad end

function calc_padding(lt, pad, k::NTuple{N, T}, dilation, stride) where {T, N}
    expand(Val(2 * N), pad)
end

function calc_padding(lt, ::SamePad, k::NTuple{N, T}, dilation, stride) where {N, T}
    # Ref: "A guide to convolution arithmetic for deep learning" https://arxiv.org/abs/1603.07285
    # Effective kernel size, including dilation
    k_eff = @. k + (k - 1) * (dilation - 1)
    # How much total padding needs to be applied?
    pad_amt = @. k_eff - 1
    # In case amount of padding is odd we need to apply different amounts to each side.
    return Tuple(mapfoldl(i -> [cld(i, 2), fld(i, 2)], vcat, pad_amt))
end

# Handling ComponentArrays
## NOTE: We should probably upsteam some of these
function Base.zero(c::ComponentArray{T, N, <:CuArray{T}}) where {T, N}
    ComponentArray(zero(getdata(c)), getaxes(c))
end

Base.vec(c::ComponentArray{T, N, <:CuArray{T}}) where {T, N} = getdata(c)

function Base.:-(x::ComponentArray{T, N, <:CuArray{T}}) where {T, N}
    ComponentArray(-getdata(x), getaxes(x))
end

function Base.similar(c::ComponentArray{T, N, <:CuArray{T}},
                      l::Vararg{Union{Integer, AbstractUnitRange}}) where {T, N}
    return similar(getdata(c), l)
end

function Functors.functor(::Type{<:ComponentArray}, c)
    return NamedTuple{propertynames(c)}(getproperty.((c,), propertynames(c))),
           ComponentArray
end

# Updating a monolithic vector is way faster than chunks of smaller ones
# Also helps in distributed settings
Optimisers.setup(opt, ps::ComponentArray) = Optimisers.setup(opt, getdata(ps))

function Optimisers.update(opt_state, ps::ComponentArray, gs::ComponentArray)
    opt_state, ps_new = Optimisers.update(opt_state, getdata(ps), getdata(gs))
    return opt_state, ComponentArray(ps_new, getaxes(ps))
end

function Optimisers.update!(st, ps::ComponentArray, gs::ComponentArray)
    st, ps_ = Optimisers.update!(st, NamedTuple(ps), NamedTuple(gs))
    return st, ComponentArray(ps_)
end

function ComponentArrays.make_carray_args(nt::NamedTuple)
    data, ax = ComponentArrays.make_carray_args(Vector, nt)
    data = length(data) == 0 ? Float32[] :
           (length(data) == 1 ? [data[1]] : reduce(vcat, data))
    return (data, ax)
end

## For being able to print empty ComponentArrays
function ComponentArrays.last_index(f::FlatAxis)
    nt = ComponentArrays.indexmap(f)
    length(nt) == 0 && return 0
    return ComponentArrays.last_index(last(nt))
end

ComponentArrays.recursive_length(nt::NamedTuple{(), Tuple{}}) = 0

# Return Nothing if field not present
function safe_getproperty(x::Union{ComponentArray, NamedTuple}, k::Symbol)
    k ∈ propertynames(x) && return getproperty(x, k)
    return nothing
end

# Getting typename
get_typename(::T) where {T} = Base.typename(T).wrapper

# For Normalization
@inline @generated safe_copy(x::T) where {T} = hasmethod(copy, (T,)) ? :(copy(x)) : :x

@inline @generated safe_vec(x::T) where {T} = hasmethod(vec, (T,)) ? :(vec(x)) : :x

@inline @inbounds function get_reshape_dims(sx::NTuple{N, <:Int},
                                            ly::Int)::typeof(sx) where {N}
    if ly == sx[N - 1]
        return ntuple(i -> i == N - 1 ? ly : 1, N)
    elseif N > 2 && ly == sx[N - 1] * sx[N - 2]
        return ntuple(i -> i == (N - 1) || i == (N - 2) ? sx[i] : 1, N)
    else
        error("Invalid Dimensions")
    end
end

@inline reshape_into_proper_shape(x::Nothing, y)::Nothing = x
@inline reshape_into_proper_shape(x, y)::typeof(y) = reshape(x,
                                                             get_reshape_dims(size(y),
                                                                              length(x)))

# RNN Utilities
@inline gate(h::Int, n::Int) = (1:h) .+ h * (n - 1)
@inline gate(x::AbstractVector, h::Int, n::Int) = view(x, gate(h, n))
@inline gate(x::AbstractMatrix, h::Int, n::Int) = view(x, gate(h, n), :)

"""
    multigate(x::AbstractArray, ::Val{N})

Split up `x` into `N` equally sized chunks (along dimension `1`).
"""
@inline multigate(x::AbstractArray, ::Val{N}) where {N} = gate.((x,), size(x, 1) ÷ N, 1:N)

# Val utilities
get_known(::Val{T}) where {T} = T


"""
```
corners = imcorner(img; [method])
corners = imcorner(img, threshold, percentile; [method])
```

Performs corner detection using one of the following methods -

    1. harris
    2. shi_tomasi
    3. kitchen_rosenfeld

The parameters of the individual methods are described in their documentation. The
maxima values of the resultant responses are taken as corners. If a threshold is
specified, the values of the responses are thresholded to give the corner pixels.
The threshold is assumed to be a percentile value unless `percentile` is set to false.
"""
function imcorner(img::AbstractArray; method::Function = harris, args...)
    responses = method(img; args...)
    corners = similar(img, Bool)
    fill!(corners, false)
    maxima = map(CartesianIndex, findlocalmaxima(responses))
    for m in maxima corners[m] = true end
    corners
end

function imcorner(img::AbstractArray, threshold; method::Function = harris, args...)
    responses = method(img; args...)
    map(i -> i > threshold, responses)
end

function imcorner(img::AbstractArray, thresholdp::Percentile; method::Function = harris, args...)
    responses = method(img; args...)
    threshold = StatsBase.percentile(vec(responses), thresholdp.p)
    map(i -> i > threshold, responses)
end

"""
```
corners = imcorner_subpixel(img; [method])
         -> Vector{HomogeneousPoint{Float64,3}}
corners = imcorner_subpixel(img, threshold, percentile; [method])
         -> Vector{HomogeneousPoint{Float64,3}}
```

Same as [`imcorner`](@ref), but estimates corners to sub-pixel precision.

Sub-pixel precision is achieved by interpolating the corner response values using
the 4-connected neighbourhood of a maximum response value.
See [`corner2subpixel`](@ref) for more details of the interpolation scheme.

"""
function imcorner_subpixel(img::AbstractArray; method::Function = harris, args...)
    responses = method(img; args...)
    corner_indicator = similar(img, Bool)
    fill!(corner_indicator, false)
    maxima = map(CartesianIndex, findlocalmaxima(responses))
    for m in maxima corner_indicator[m] = true end
    corners = corner2subpixel(responses,corner_indicator)
end

function imcorner_subpixel(img::AbstractArray, threshold; method::Function = harris, args...)
    responses = method(img; args...)
    corner_indicator = map(i -> i > threshold, responses)
    corners = corner2subpixel(responses,corner_indicator)
end

function imcorner_subpixel(img::AbstractArray, thresholdp::Percentile; method::Function = harris, args...)
    responses = method(img; args...)
    threshold = StatsBase.percentile(vec(responses), thresholdp.p)
    corner_indicator = map(i -> i > threshold, responses)
    corners = corner2subpixel(responses,corner_indicator)
end

"""
```
corners = corner2subpixel(responses::AbstractMatrix,corner_indicator::AbstractMatrix{Bool})
        -> Vector{HomogeneousPoint{Float64,3}}
```
Refines integer corner coordinates to sub-pixel precision.

The function takes as input a matrix representing corner responses
and a boolean indicator matrix denoting the integer coordinates of a corner
in the image. The output is a vector of type [`HomogeneousPoint`](@ref)
storing the sub-pixel coordinates of the corners.

The algorithm computes a correction factor which is added to the original
integer coordinates. In particular, a univariate quadratic polynomial is fit
separately to the ``x``-coordinates and ``y``-coordinates of a corner and its immediate
east/west, and north/south neighbours. The fit is achieved using a local
coordinate system for each corner, where the origin of the coordinate system is
a given corner, and its immediate neighbours are assigned coordinates of  minus
one and plus one.

The corner and its two neighbours form a system of three equations. For example,
let  ``x_1 = -1``,  ``x_2 = 0`` and  ``x_3 = 1`` denote the local ``x`` coordinates
of the west, center and east pixels and let the vector ``\\mathbf{b} = [r_1, r_2, r_3]``
denote the corresponding corner response values. With

```math
    \\mathbf{A} =
        \\begin{bmatrix}
            x_1^2 & x_1  & 1  \\\\
            x_2^2 & x_2  & 1 \\\\
            x_3^2 & x_3  & 1 \\\\
        \\end{bmatrix},
```
the coefficients of the quadratic polynomial can be found by solving the
system of equations ``\\mathbf{b} = \\mathbf{A}\\mathbf{x}``.
The result is given by ``x = \\mathbf{A}^{-1}\\mathbf{b}``.

The vertex of the quadratic polynomial yields a sub-pixel estimate of the
true corner position. For example, for a univariate quadratic polynomial
``px^2 + qx + r``, the ``x``-coordinate of the vertex is ``\\frac{-q}{2p}``.
Hence, the refined sub-pixel coordinate is equal to:
 ``c +  \\frac{-q}{2p}``, where ``c`` is the integer coordinate.

!!! note
    Corners on the boundary of the image are not refined to sub-pixel precision.

"""
function corner2subpixel(responses::AbstractMatrix, corner_indicator::AbstractMatrix{Bool})
    row_range, col_range = axes(corner_indicator)
    idxs = findall(!iszero, corner_indicator) # findnz
    row, col = (getindex.(idxs,1),getindex.(idxs,2))
    ncorners = length(row)
    corners = fill(HomogeneousPoint((0.0,0.0,0.0)),ncorners)
    invA = @SMatrix [0.5 -1.0 0.5; -0.5 0.0 0.5; 0.0 1.0 -0.0]
    for k = 1:ncorners
        # Corners on the perimeter of the image will not be interpolated.
        if  (row[k] == first(row_range) || row[k] == last(row_range) ||
             col[k] == first(col_range) || col[k] == last(col_range))
            y = convert(Float64,row[k])
            x = convert(Float64,col[k])
            corners[k] = HomogeneousPoint((x,y,1.0))
        else
            center, north, south, east, west =
                                unsafe_neighbourhood_4(responses,row[k],col[k])
            # Solve for the coefficients of the quadratic equation.
            a, b, c = invA* @SVector [west, center, east]
            p, q, r = invA* @SVector [north, center, south]
            # Solve for the first coordinate of the vertex.
            u = -b/(2.0a)
            v = -q/(2.0p)
            corners[k] = HomogeneousPoint((col[k]+u,row[k]+v,1.0))
        end
    end
    return corners
end

"""
```
unsafe_neighbourhood_4(matrix::AbstractMatrix,r::Int,c::Int)
```

Returns the value of a matrix at given coordinates together with the values
of the north, south, east and west neighbours.

This function does not perform bounds checking. It is up to the user to ensure
that the function is not called with indices that are on the boundary of the
matrix.

"""
function unsafe_neighbourhood_4(matrix::AbstractMatrix,r::Int,c::Int)
    center = matrix[r,c]
    north = matrix[r-1,c]
    south = matrix[r+1,c]
    east = matrix[r,c+1]
    west = matrix[r,c-1]
    return center, north, south, east, west
end

"""
```
harris_response = harris(img; [k], [border], [weights])
```

Performs Harris corner detection. The covariances can be taken using either a mean
weighted filter or a gamma kernel.
"""
function harris(img::AbstractArray; k::Float64 = 0.04, args...)
    cov_xx, cov_xy, cov_yy = gradcovs(img, args...)
    corner = map((xx, yy, xy) -> xx * yy - xy ^ 2 - k * (xx + yy) ^ 2, cov_xx, cov_yy, cov_xy)
    corner
end

"""
```
shi_tomasi_response = shi_tomasi(img; [border], [weights])
```

Performs Shi Tomasi corner detection. The covariances can be taken using either a mean
weighted filter or a gamma kernel.
"""
function shi_tomasi(img::AbstractArray; border::AbstractString = "replicate", args...)
    cov_xx, cov_xy, cov_yy = gradcovs(img, border; args...)
    corner = map((xx, yy, xy) -> ((xx + yy) - (sqrt((xx - yy) ^ 2 + 4 * xy ^ 2))) / 2, cov_xx, cov_yy, cov_xy)
    corner
end

"""
```
kitchen_rosenfeld_response = kitchen_rosenfeld(img; [border])
```

Performs Kitchen Rosenfeld corner detection. The covariances can be taken using either a mean
weighted filter or a gamma kernel.
"""
function kitchen_rosenfeld(img::AbstractArray; border::AbstractString = "replicate")
    meth = KernelFactors.sobel
    (grad_x, grad_y) = imgradients(img, meth, border)
    (grad_xx, grad_xy) = imgradients(grad_x, meth, border)
    (grad_yx, grad_yy) = imgradients(grad_y, meth, border)
    map(kr, grad_x, grad_y, grad_xx, grad_xy, grad_yy)
end

function kr(x::T, y::T, xx::T, xy::T, yy::T) where T<:Real
    num = xx*y*y + yy*x*x - 2*xy*x*y
    denom = x*x + y*y
    ifelse(denom == 0, zero(num)/one(denom), -num/denom)
end

function kr(x::Real, y::Real, xx::Real, xy::Real, yy::Real)
    xp, yp, xxp, xyp, yyp = promote(x, y, xx, xy, yy)
    kr(xp, yp, xxp, xyp, yyp)
end

kr(x::RealLike, y::RealLike, xx::RealLike, xy::RealLike, yy::RealLike) =
    kr(gray(x), gray(y), gray(xx), gray(xy), gray(yy))

function kr(x::AbstractRGB, y::AbstractRGB, xx::AbstractRGB, xy::AbstractRGB, yy::AbstractRGB)
    krrgb = RGB(kr(red(x), red(y), red(xx), red(xy), red(yy)),
                kr(green(x), green(y), green(xx), green(xy), green(yy)),
                kr(blue(x),  blue(y),  blue(xx),  blue(xy),  blue(yy)))
    gray(convert(Gray, krrgb))
end

"""
    fastcorners(img, n, threshold) -> corners

Performs FAST Corner Detection. `n` is the number of contiguous pixels
which need to be greater (lesser) than intensity + threshold (intensity - threshold)
for a pixel to be marked as a corner. The default value for n is 12.
"""
function fastcorners(img::AbstractArray{T}, n::Int = 12, threshold::Float64 = 0.15) where T
    img_padded = padarray(img, Fill(0, (3,3)))
    corner = falses(size(img))
    R = CartesianIndices(size(img))
    idx = map(CartesianIndex{2}, [(0, 3), (1, 3), (2, 2), (3, 1), (3, 0), (3, -1), (2, -2), (1, -3),
            (0, -3), (-1, -3), (-2, -2), (-3, -1), (-3, 0), (-3, 1), (-2, 2), (-1, 3)])

    idxidx = [1, 5, 9, 13]
    for I in R
        bright_threshold = img_padded[I] + threshold
        dark_threshold = img_padded[I] - threshold
        if n >= 12
            sum_bright = 0
            sum_dark = 0
            for k in idxidx
                pixel = img_padded[I + idx[k]]
                if pixel > bright_threshold
                    sum_bright += 1
                elseif pixel < dark_threshold
                    sum_dark += 1
                end
            end
            if sum_bright < 3 && sum_dark < 3
                continue
            end
        end
        consecutive_bright = 0
        consecutive_dark = 0

        for i in 1:15 + n
            k = mod1(i, 16)
            pixel = img_padded[I + idx[k]]
            if pixel > bright_threshold
                consecutive_dark = 0
                consecutive_bright += 1
            elseif pixel < dark_threshold
                consecutive_bright = 0
                consecutive_dark += 1
            end

            if consecutive_dark == n || consecutive_bright == n
                corner[I] = true
                break
            end
        end
    end
    corner
end

function gradcovs(img::AbstractArray, border::AbstractString = "replicate"; weights::Function = meancovs, args...)
    (grad_x, grad_y) = imgradients(img, KernelFactors.sobel, border)

    cov_xx = dotc.(grad_x, grad_x)
    cov_xy = dotc.(grad_x, grad_y)
    cov_yy = dotc.(grad_y, grad_y)

    weights(cov_xx, cov_xy, cov_yy, args...)
end

function meancovs(cov_xx, cov_xy, cov_yy, blockSize::Int = 3)

    box_filter_kernel = centered((1 / (blockSize * blockSize)) * ones(blockSize, blockSize))

    filt_cov_xx = imfilter(cov_xx, box_filter_kernel)
    filt_cov_xy = imfilter(cov_xy, box_filter_kernel)
    filt_cov_yy = imfilter(cov_yy, box_filter_kernel)

    filt_cov_xx, filt_cov_xy, filt_cov_yy
end

function gammacovs(cov_xx, cov_xy, cov_yy, gamma::Float64 = 1.4)
    kernel = KernelFactors.gaussian((gamma, gamma))

    filt_cov_xx = imfilter(cov_xx, kernel)
    filt_cov_xy = imfilter(cov_xy, kernel)
    filt_cov_yy = imfilter(cov_yy, kernel)

    filt_cov_xx, filt_cov_xy, filt_cov_yy
end


#=
Construct a 100 x 100 grid and run discreteVI on it. Also run global approx VI where any integral point in the grid
can be sampled. Compare the values at the integer grid points. Confirm that the max difference reduces with increasing
number of samples.
=#

# Feature vector conversion
# For point (x,y), fv is (1 x y xy x^2 y^2)
function GlobalApproximationValueIteration.convert_featurevector(::Type{SVector{10, Float64}}, s::GWPos, mdp::SimpleGridWorld)

    x = s[1]
    y = s[2]

    v = SVector{10, Float64}(1, x, y, x*y, x^2, y^2, x^3, x^2*y, x*y^2, y^3)

    return v
end

# Sample a specific integral point in the grid
function GlobalApproximationValueIteration.sample_state(mdp::SimpleGridWorld, rng::RNG=Random.GLOBAL_RNG) where {RNG <: AbstractRNG}

    x = rand(rng, 1:mdp.size[1])
    y = rand(rng, 1:mdp.size[2])

    return GWPos(x, y)
end

function test_absolute_error()

    rng = MersenneTwister(1234)

    # Set solver and grid world parameters
    MAX_ITERS = 500
    NUM_SAMPLES = 1000

    SIZE_X = 5
    SIZE_Y = 5
    REWARD_COV_PROB = 0.4

    rewards = Dict{GWPos, Float64}()
    for x = 1:SIZE_X
        for y = 1:SIZE_Y
            if rand(rng) < REWARD_COV_PROB
                rewards[GWPos(x, y)] = 1
            end
        end
    end

    # Create the MDP for a typical grid world
    mdp = SimpleGridWorld(size=(SIZE_X, SIZE_Y), rewards=rewards)

    # Create the linear function approximation with 10 weight parameters, initialized to zero
    lin_gfa = LinearGlobalFunctionApproximator(zeros(10))

    # Initialize the global approximation solver with the linear approximator and solve the MDP to obtain the policy
    gfa_solver = GlobalApproximationValueIterationSolver(lin_gfa, num_samples=NUM_SAMPLES, max_iterations=MAX_ITERS, verbose=true, fv_type=SVector{10, Float64})
    gfa_policy = solve(gfa_solver, mdp)

    error_arr = Vector{Float64}(undef, 0)

    for state in states(mdp)

        full_val = value(policy, state)
        approx_val = value(gfa_policy, state)
        abs_diff = abs(full_val - approx_val)

        push!(error_arr, abs_diff)
    end

    @show mean(error_arr)
    @show maximum(error_arr)

    return (mean(error_arr) < 0.04 && maximum(error_arr) < 0.35)
end

function test_relative_error()

    rng = MersenneTwister(2378)

    # Attempt to approximate globally with N samples and M iterations. As N increases, the average error should decrease
    MAX_ITERS = 500
    NUM_SAMPLES_LOW = 30
    NUM_SAMPLES_HI = 1000

    # Grid probabilities
    SIZE_X = 5
    SIZE_Y = 5
    REWARD_COV_PROB = 0.4

    # Generate a block of reward states from 40,40 to 60,60
    rewards = Dict{GWPos, Float64}()
    for x = 1:SIZE_X
        for y = 1:SIZE_Y
            if rand(rng) < REWARD_COV_PROB
                rewards[GWPos(x, y)] = 1
            end
        end
    end

    # Create MDP
    mdp = SimpleGridWorld(size=(SIZE_X, SIZE_Y), rewards=rewards)

    lin_gfa_1 = LinearGlobalFunctionApproximator(zeros(10))
    lin_gfa_2 = LinearGlobalFunctionApproximator(zeros(10))

    solver_low = GlobalApproximationValueIterationSolver(lin_gfa_1; num_samples=NUM_SAMPLES_LOW, max_iterations=MAX_ITERS, verbose=true, fv_type=SVector{10, Float64})
    solver_hi = GlobalApproximationValueIterationSolver(lin_gfa_2; num_samples=NUM_SAMPLES_HI, max_iterations=MAX_ITERS, verbose=true, fv_type=SVector{10, Float64})

    policy_low = solve(solver_low, mdp)
    policy_hi = solve(solver_hi, mdp)

    # Now solve with dicrete VI
    solver = ValueIterationSolver(max_iterations=1000, verbose=true)
    policy = solve(solver, mdp)

    err_arr_low = Vector{Float64}(undef, 0)
    err_arr_hi = Vector{Float64}(undef, 0)

    for state in states(mdp)

        full_val = value(policy, state)

        approx_val_low = value(policy_low, state)
        approx_val_hi = value(policy_hi, state)

        push!(err_arr_low, abs(full_val - approx_val_low))
        push!(err_arr_hi, abs(full_val - approx_val_hi))

    end

    @show mean(err_arr_low), mean(err_arr_hi)
    @show maximum(err_arr_low), maximum(err_arr_hi)

    return (mean(err_arr_low) > mean(err_arr_hi) && maximum(err_arr_low) > maximum(err_arr_hi))
end

# @test test_absolute_error() == true
@test test_relative_error() == true


using LoopVectorization, Test

function multiassign!(y, x)
  @assert length(y)+3 == length(x)
  @inbounds for i ∈ eachindex(y)
    x₁, ((x₂,x₃), (x₄,x₅)) = x[i], (sincos(x[i+1]), (x[i+2], x[i+3]))
    y[i] = x₁ * x₄ - x₂ * x₃
  end
  y
end
multiassign(x) = multiassign!(similar(x, length(x)-3), x)
function multiassign_turbo!(y, x)
  @assert length(y)+3 == length(x)
  @turbo for i ∈ eachindex(y)
    x₁, ((x₂,x₃), (x₄,x₅)) = x[i], (sincos(x[i+1]), (x[i+2], x[i+3]))
    y[i] = x₁ * x₄ - x₂ * x₃
  end
  y
end
multiassign_turbo(x) = multiassign_turbo!(similar(x, length(x)-3), x)

@testset "Multiple assignments" begin
  @show @__LINE__
  x = rand(111);
  @test multiassign(x) ≈ multiassign_turbo(x)
end


struct VariableStd <: DataItem
    name :: String
    value :: Vector{<:Real}
    offset :: Real
    scale :: Real
end

"""
    VariableStd(name::String, z::Vector{<:Real})::VariableStd

Construct a VariableStd type from the standardized vector `z` with the name `name`.

"""
function VariableStd(name::String, z::Vector{<:Real})::VariableStd

    m = mean(z)
    s = std(z)

    @assert isapprox(m, 0.0, atol = sqrt(eps())) "the mean should be equal to zero. Use the type Variable instead."
    @assert isapprox(s, 1.0, atol = sqrt(eps())) "the standard deviation should be equal to one. Use the type Variable instead."

    return VariableStd(name,z,0,1)
end

"""
    reconstruct(vstd::VariableStd)

Reconstruct the original Variable from the standardized one.

See also

See also [`Variable`](@ref), [`VariableStd`](@ref)  and [`standardize`](@ref).
"""
function reconstruct(vstd::VariableStd)::Variable

    z = vstd.value

    x = vstd.scale*z .+ vstd.offset

    v = Variable(vstd.name, x)

    return v

end


"""\
 C version: Dmitri Laikov
 F77 version: Christoph van Wuellen, http://www.ccl.net
 Python version: Richard P. Muller, 2002.
 Julia version: John C. Snyder, 2015.

 This subroutine is part of a set of subroutines that generate
 Lebedev grids [1-6] for integration on a sphere. The original
 C-code [1] was kindly provided by Dr. Dmitri N. Laikov and
 translated into fortran by Dr. Christoph van Wuellen.
 This subroutine was translated from C to fortran77 by hand.

 Users of this code are asked to include reference [1] in their
 publications, and in the user- and programmers-manuals
 describing their codes.

   [1] V.I. Lebedev, and D.N. Laikov
       'A quadrature formula for the sphere of the 131st
        algebraic order of accuracy'
       Doklady Mathematics, Vol. 59, No. 3, 1999, pp. 477-481.

   [2] V.I. Lebedev
       'A quadrature formula for the sphere of 59th algebraic
        order of accuracy'
       Russian Acad. Sci. Dokl. Math., Vol. 50, 1995, pp. 283-286.

   [3] V.I. Lebedev, and A.L. Skorokhodov
       'Quadrature formulas of orders 41, 47, and 53 for the sphere'
       Russian Acad. Sci. Dokl. Math., Vol. 45, 1992, pp. 587-592.

   [4] V.I. Lebedev
       'Spherical quadrature formulas exact to orders 25-29'
       Siberian Mathematical Journal, Vol. 18, 1977, pp. 99-107.

   [5] V.I. Lebedev
       'Quadratures on a sphere'
       Computational Mathematics and Mathematical Physics, Vol. 16,
       1976, pp. 10-24.

   [6] V.I. Lebedev
       'Values of the nodes and weights of ninth to seventeenth
        order Gauss-Markov quadrature formulae invariant under the
        octahedron group with inversion'
       Computational Mathematics and Mathematical Physics, Vol. 15,
       1975, pp. 44-51.
"""

function genOh_a00(v)
    "(0,0,a) etc. (6 points)"
    a=1.0
    hcat([a,0,0,v],[-a,0,0,v],[0,a,0,v],[0,-a,0,v],[0,0,a,v],[0,0,-a,v])'
end

function genOh_aa0(v)
    "(0,a,a) etc, a=1/sqrt(2) (12 points)"
    a=sqrt(0.5)
    hcat([0,a,a,v],[0,-a,a,v],[0,a,-a,v],[0,-a,-a,v],
         [a,0,a,v],[-a,0,a,v],[a,0,-a,v],[-a,0,-a,v],
         [a,a,0,v],[-a,a,0,v],[a,-a,0,v],[-a,-a,0,v])'
 end

function genOh_aaa(v)
    "(a,a,a) etc, a=1/sqrt(3) (8 points)"
    a = sqrt(1./3.)
    hcat([a,a,a,v],[-a,a,a,v],[a,-a,a,v],[-a,-a,a,v],
         [a,a,-a,v],[-a,a,-a,v],[a,-a,-a,v],[-a,-a,-a,v])'
 end

function genOh_aab(v,a)
    "(a,a,b) etc, b=sqrt(1-2 a^2), a input (24 points)"
    b = sqrt(1.0 - 2.0*a*a)
    hcat([a,a,b,v],[-a,a,b,v],[a,-a,b,v],[-a,-a,b,v],
         [a,a,-b,v],[-a,a,-b,v],[a,-a,-b,v],[-a,-a,-b,v],
         [a,b,a,v],[-a,b,a,v],[a,-b,a,v],[-a,-b,a,v],
         [a,b,-a,v],[-a,b,-a,v],[a,-b,-a,v],[-a,-b,-a,v],
         [b,a,a,v],[-b,a,a,v],[b,-a,a,v],[-b,-a,a,v],
         [b,a,-a,v],[-b,a,-a,v],[b,-a,-a,v],[-b,-a,-a,v])'
 end

function genOh_ab0(v,a)
    "(a,b,0) etc, b=sqrt(1-a^2), a input (24 points)"
    b=sqrt(1.0-a*a)
    hcat([a,b,0,v],[-a,b,0,v],[a,-b,0,v],[-a,-b,0,v],
         [b,a,0,v],[-b,a,0,v],[b,-a,0,v],[-b,-a,0,v],
         [a,0,b,v],[-a,0,b,v],[a,0,-b,v],[-a,0,-b,v],
         [b,0,a,v],[-b,0,a,v],[b,0,-a,v],[-b,0,-a,v],
         [0,a,b,v],[0,-a,b,v],[0,a,-b,v],[0,-a,-b,v],
         [0,b,a,v],[0,-b,a,v],[0,b,-a,v],[0,-b,-a,v])'
end

function genOh_abc(v,a,b)
    "(a,b,c) etc, c=sqrt(1-a^2-b^2), a,b input  (48 points)"
    c=sqrt(1.0 - a*a - b*b)
    hcat([a,b,c,v],[-a,b,c,v],[a,-b,c,v],[-a,-b,c,v],
         [a,b,-c,v],[-a,b,-c,v],[a,-b,-c,v],[-a,-b,-c,v],
         [a,c,b,v],[-a,c,b,v],[a,-c,b,v],[-a,-c,b,v],
         [a,c,-b,v],[-a,c,-b,v],[a,-c,-b,v],[-a,-c,-b,v],
         [b,a,c,v],[-b,a,c,v],[b,-a,c,v],[-b,-a,c,v],
         [b,a,-c,v],[-b,a,-c,v],[b,-a,-c,v],[-b,-a,-c,v],
         [b,c,a,v],[-b,c,a,v],[b,-c,a,v],[-b,-c,a,v],
         [b,c,-a,v],[-b,c,-a,v],[b,-c,-a,v],[-b,-c,-a,v],
         [c,a,b,v],[-c,a,b,v],[c,-a,b,v],[-c,-a,b,v],
         [c,a,-b,v],[-c,a,-b,v],[c,-a,-b,v],[-c,-a,-b,v],
         [c,b,a,v],[-c,b,a,v],[c,-b,a,v],[-c,-b,a,v],
         [c,b,-a,v],[-c,b,-a,v],[c,-b,-a,v],[-c,-b,-a,v])'
end

function leb6()
    genOh_a00(0.1666666666666667)
end

function leb14()
    [genOh_a00(0.06666666666666667);
     genOh_aaa(0.07500000000000000)]
end

function leb26()
    [genOh_a00(0.04761904761904762);
     genOh_aa0(0.03809523809523810);
     genOh_aaa(0.03214285714285714)]
end

function leb38()
    [genOh_a00(0.009523809523809524);
     genOh_aaa(0.3214285714285714E-1);
     genOh_ab0(0.2857142857142857E-1,0.4597008433809831E+0)]
end

function leb50()
    [genOh_a00(0.1269841269841270E-1);
     genOh_aa0(0.2257495590828924E-1);
     genOh_aaa(0.2109375000000000E-1);
     genOh_aab(0.2017333553791887E-1,0.3015113445777636E+0)]
end

function leb74()
    [genOh_a00(0.5130671797338464E-3);
     genOh_aa0(0.1660406956574204E-1);
     genOh_aaa(-0.2958603896103896E-1);
     genOh_aab(0.2657620708215946E-1,0.4803844614152614E+0);
     genOh_ab0(0.1652217099371571E-1,0.3207726489807764E+0)]
end

function leb86()
    [genOh_a00(0.1154401154401154E-1);
     genOh_aaa(0.1194390908585628E-1);
     genOh_aab(0.1111055571060340E-1,0.3696028464541502E+0);
     genOh_aab(0.1187650129453714E-1,0.6943540066026664E+0);
     genOh_ab0(0.1181230374690448E-1,0.3742430390903412E+0)]
end

function leb110()
    [genOh_a00(0.3828270494937162E-2);
     genOh_aaa(0.9793737512487512E-2);
     genOh_aab(0.8211737283191111E-2,0.1851156353447362E+0);
     genOh_aab(0.9942814891178103E-2,0.6904210483822922E+0);
     genOh_aab(0.9595471336070963E-2,0.3956894730559419E+0);
     genOh_ab0(0.9694996361663028E-2,0.4783690288121502E+0)]
end

function leb146()
    [genOh_a00(0.5996313688621381E-3);
     genOh_aa0(0.7372999718620756E-2);
     genOh_aaa(0.7210515360144488E-2);
     genOh_aab(0.7116355493117555E-2,0.6764410400114264E+0);
     genOh_aab(0.6753829486314477E-2,0.4174961227965453E+0);
     genOh_aab(0.7574394159054034E-2,0.1574676672039082E+0);
     genOh_abc(0.6991087353303262E-2,0.1403553811713183E+0,
                     0.4493328323269557E+0)]
end

function leb170()
    [genOh_a00(0.5544842902037365E-2);
     genOh_aa0(0.6071332770670752E-2);
     genOh_aaa(0.6383674773515093E-2);
     genOh_aab(0.5183387587747790E-2,0.2551252621114134E+0);
     genOh_aab(0.6317929009813725E-2,0.6743601460362766E+0);
     genOh_aab(0.6201670006589077E-2,0.4318910696719410E+0);
     genOh_ab0(0.5477143385137348E-2,0.2613931360335988E+0);
     genOh_abc(0.5968383987681156E-2,0.4990453161796037E+0,
                     0.1446630744325115E+0)]
end

function leb194()
    [genOh_a00(0.1782340447244611E-2);
     genOh_aa0(0.5716905949977102E-2);
     genOh_aaa(0.5573383178848738E-2);
     genOh_aab(0.5608704082587997E-2,0.6712973442695226E+0);
     genOh_aab(0.5158237711805383E-2,0.2892465627575439E+0);
     genOh_aab(0.5518771467273614E-2,0.4446933178717437E+0);
     genOh_aab(0.4106777028169394E-2,0.1299335447650067E+0);
     genOh_ab0(0.5051846064614808E-2,0.3457702197611283E+0);
     genOh_abc(0.5530248916233094E-2,0.1590417105383530E+0,
                     0.8360360154824589E+0)]
end

leb = Dict(
    6   => leb6,
    14  => leb14,
    26  => leb26,
    38  => leb38,
    50  => leb50,
    74  => leb74,
    86  => leb86,
    110 => leb110,
    146 => leb146,
    170 => leb170,
    194 => leb194
)

# generate lebedev grids
function generate_lebedev()

    function lebgrid(n::Integer)
        data = leb[n]()'
        w = vec(4pi*data[end,:])
        p = data[1:3,:]
        p = Point3{Float64}[tuple(p[:,i]...) for i in 1:size(p,2)]
        AngularGrid(p,w)
    end

    lebedev = Dict{Int,AngularGridf}([k=>lebgrid(k) for k in keys(leb)])

    open(joinpath(Pkg.dir("Element"), "deps/lebedev.jls"), "w") do io serialize(io, lebedev) end

end


"""Sums the element of a 2-dimensional matrix."""
function sum_matrix(a)
    sum = 0
    n = array_length(a)
    for i = 0 : (n - 1)
        m = array_length(a[i])
        for j = 0 : (m - 1)
            sum = sum + a[i][j]
        end
    end
    return sum
end

"""Makes matrix of dimension d with level i having a size of l[i] and initial value of val."""
function make_matrix(d :: Int64, l :: Array, val :: Any) :: Any
    if d > 0
        n = array_length(l)
        a = newarray(l[n - d], make_matrix(d - 1, l, val))
        for i = 0 : (l[n - d] - 1)
            b = make_matrix(d - 1, l, val)
            a[i] = b
        end
        return a
    end
    # if d = 0, we simply return the value
    return val
end


# Tests for Thompson Sampling agents

noOfArms = 5

@testset "Thompson Sampling" begin

@testset "TS" begin
    agent = TS( noOfArms )

    @test agent.noOfArms == noOfArms
    @test agent.noOfSteps == 0
    @test agent.lastPlayedArm == 0

    @test info_str( agent ) == "Thompson Sampling"
    @test info_str( agent, false ) == "Thompson Sampling"
    @test info_str( agent, true ) == "Thompson Sampling"

    for t = 1:10
        get_arm_index( agent ); update_reward!( agent, rand() );
    end
    @test agent.noOfSteps == 10
    @test sum(agent.cummSuccess+agent.cummFailure) == 10
    @test sum(agent.cummSuccess+agent.cummFailure+agent.α0+agent.β0) == 10+2noOfArms

    reset!( agent )
    @test agent == TS(noOfArms)

    # With different prior
    agent = TS( [ (1,2), (2,3), (3,4) ] )
    @test agent.noOfArms == 3
    for t = 1:10
        get_arm_index( agent ); update_reward!( agent, rand() )
    end
    @test sum(agent.cummSuccess+agent.cummFailure+agent.α0+agent.β0) == 10+1+2+2+3+3+4
    @test [ agent.α0+agent.cummSuccess, agent.β0+agent.cummFailure ] ==
            [ [agent.samplingDist[i].α for i = 1:3], [agent.samplingDist[i].β for i = 1:3]  ]
    reset!( agent )
    @test agent == TS([(1,2),(2,3),(3,4)])
end

@testset "DynamicTS" begin
    agent = DynamicTS( noOfArms, 10 )
    @test agent.noOfArms == noOfArms
    @test agent.noOfSteps == 0
    @test agent.lastPlayedArm == 0

    @test info_str( agent ) == "Dynamic Thompson Sampling (C = 10.000)"
    @test info_str( agent, false ) == "Dynamic Thompson Sampling (C = 10.000)"
    @test info_str( agent, true ) == "Dynamic Thompson Sampling (\$C = 10.000\$)"

    for t = 1:10
        get_arm_index( agent ); update_reward!( agent, rand() )
    end
    @test agent.noOfSteps == 10

    reset!( agent )
    @test agent == DynamicTS(noOfArms,10)
end

@testset "OTS" begin

end

@testset "TSNormal" begin

end

@testset "Discounted TS" begin

end

@testset "Discounted OTS" begin

end

@testset "RestartTS" begin

end

@testset "TSGaussPrior" begin

end

end


using E2ga
using Test

@testset "promote" begin
    a = E2ga.Geometric2(2.0, 5, 11, 17)
    @test typeof(a.α) == Float64
    @test typeof(a.x) == Float64
    @test typeof(a.y) == Float64
    @test typeof(a.β) == Float64
    @test "$(a)" == "2.0 + 5.0e1 + 11.0e2 + 17.0I"
end

@testset "scalar" begin
    a = E2ga.Geometric2(1)
    @test a.α == 1
    @test a.x == 0
    @test a.y == 0
    @test a.β == 0
end

@testset "+" begin
    a = E2ga.Geometric2(1, 3, 7, 13)
    b = E2ga.Geometric2(2, 5, 11, 17)
    c = a + b
    @test c.α == 3
    @test c.x == 8
    @test c.y == 18
    @test c.β == 30
end

@testset "-" begin
    a = E2ga.Geometric2(2, 5, 11, 17)
    b = E2ga.Geometric2(1, 3, 7, 13)
    c = a - b
    @test c.α == 1
    @test c.x == 2
    @test c.y == 4
    @test c.β == 4
end

@testset "*" begin
    a = E2ga.Geometric2(2, 5, 11, 17)
    b = E2ga.Geometric2(1, 3, 7, 13)
    c = a * b
    @test c.α == (2 * 1) + (5 * 3) + (11 * 7) - (17 * 13)
    @test c.x == (2 * 3) + (5 * 1) - (11 * 13) + (17 * 7)
    @test c.y == (2 * 7) + (5 * 13) + (11 * 1) - (17 * 3)
    @test c.β == (2 * 13) + (5 * 7) - (11 * 3) + (17 * 1)
end

@testset "one" begin
    one = E2ga.Geometric2(1)
    @test one.α == 1
    @test one.x == 0
    @test one.y == 0
    @test one.β == 0
end

@testset "e1" begin
    e1 = E2ga.Geometric2(0, 1, 0, 0)
    @test e1.α == 0
    @test e1.x == 1
    @test e1.y == 0
    @test e1.β == 0
end

@testset "e2" begin
    e2 = E2ga.Geometric2(0, 0, 1, 0)
    @test e2.α == 0
    @test e2.x == 0
    @test e2.y == 1
    @test e2.β == 0
end

@testset "I" begin
    I = E2ga.Geometric2(0, 0, 0, 1)
    @test I.α == 0
    @test I.x == 0
    @test I.y == 0
    @test I.β == 1
end

@testset "show" begin
    a = E2ga.Geometric2(2, 5, 11, 17)
    @test "$(a)" == "2 + 5e1 + 11e2 + 17I"
end


using Random
export RandomShuffle

"""
    RandomShuffle() <: Surrogate

A random constrained surrogate, generated by shifting values around.

Random shuffle surrogates preserve the mean, variance and amplitude
distribution of the original signal. Properties not preserved are *any
temporal information*, such as the power spectrum and hence linear
correlations.

The null hypothesis this method can test for is whether the data
are uncorrelated noise, possibly measured via a nonlinear function.
Specifically, random shuffle surrogate can test
the null hypothesis that the original signal is produced by independent and
identically distributed random variables[^Theiler1991, ^Lancaster2018].

*Beware: random shuffle surrogates do not cover the case of correlated noise*[^Lancaster2018].

[^Theiler1991]: J. Theiler, S. Eubank, A. Longtin, B. Galdrikian, J. Farmer, Testing for nonlinearity in time series: The method of surrogate data, Physica D 58 (1–4) (1992) 77–94.
"""
struct RandomShuffle <: Surrogate end

function surrogenerator(x::AbstractVector, rf::RandomShuffle, rng = Random.default_rng())
    n = length(x)
    idxs = collect(1:n)

    init = (
        permutation = zeros(Int, n),
        idxs = idxs,
    )

    return SurrogateGenerator(rf, x, similar(x), init, rng)
end

function (sg::SurrogateGenerator{<:RandomShuffle})()
    x, s, rng = sg.x, sg.s, sg.rng
    permutation, idxs  = getfield.(Ref(sg.init), (:permutation, :idxs))
    sample!(rng, idxs, permutation; replace = false)
    s .= x[permutation]
    return s
end