text:
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


text:
# support functions

using InteractiveUtils: methodswith

function not_iterator_of_pairs(kv)
    return any(x->isempty(methodswith(typeof(kv), x, true)),
               [iterate]) ||
           any(x->!isa(x, Union{Tuple,Pair}), kv)
end

ext: jl

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

module ADPerfTest

export Test

"""
Julia automatic differentiation performance testing module for an objective
corresponding to given Input and Output types. Differentiates the objective
using some AD framework.
"""
struct Test{Input, Output}
    """
    An object containing the state of the `Test`. Is passed to every callback
    during the benchmarking.
    """
    context::Any
    """
        test.prepare!(context, input::Input)

    Converts the input data from the `Input` type in which it is provided by
    the calling benchmark runner into the format optimized for use with the
    tested AD framework.
    Stores it in the `context`.

    Optionally, performs other preparatory activities need by the tested AD framework.
    """
    prepare!::Function
    """
        test.calculate_objective!(context, times::Int)

    Repeatedly computes the objective function `times` times for the input
    stored in the context. Stores results in the context.
    """
    calculate_objective!::Function
    """
        test.calculate_jacobian!(context, times::Int)

    Repeatedly computes the jacobian of the objective function `times` times
    for the input stored in the context. Stores results in the context.
    """
    calculate_jacobian!::Function
    """
        test.output!(output::Output, context)

    Convertes outputs saved in the `context` into the format specified by the runner.
    """
    output!::Function
    """
    Creates a Test{Input, Output} object. Checks that all the callbacks have all
    the necessary methods.
    """
    function Test{Input, Output}(empty_context::Context, prepare!::Function, calculate_objective!::Function, calculate_jacobian!::Function, output!::Function) where {Input, Output, Context}
        if !Base.hasmethod(prepare!, Tuple{Context, Input})
            throw(ArgumentError("$(string(prepare!)) doesn't have a method matching $(string(prepare!))(::$(string(Context)), ::$(string(Input)))."))
        end
        if !Base.hasmethod(calculate_objective!, Tuple{Context, Int})
            throw(ArgumentError("$(string(calculate_objective!)) doesn't have a method matching $(string(calculate_objective!))(::$(string(Context)), ::$(string(Int)))."))
        end
        if !Base.hasmethod(calculate_jacobian!, Tuple{Context, Int})
            throw(ArgumentError("$(string(calculate_jacobian!)) doesn't have a method matching $(string(calculate_jacobian!))(::$(string(Context)), ::$(string(Int)))."))
        end
        if !Base.hasmethod(output!, Tuple{Output, Context})
            throw(ArgumentError("$(string(output!)) doesn't have a method matching $(string(output!))(::$(string(Output)), ::$(string(Context)))."))
        end
        new{Input, Output}(empty_context, prepare!, calculate_objective!, calculate_jacobian!, output!)
    end
end

end

max_stars_repo_name: dsyme/ADBench

max_issues_repo_name: dsyme/ADBench

max_forks_repo_name: dsyme/ADBench


text:
module GameSession

using ..Gameplay, ..Wikipedia, ..Wikipedia.Articles
using Random

mutable struct Game
  id::String
  articles::Vector{Article}
  history::Vector{Article}
  steps_taken::UInt8
  difficulty::UInt8

  Game(game_difficulty) = new(randstring(), newgame(game_difficulty), Article[], 0, game_difficulty)
end

const GAMES = Dict{String,Game}()

export newgamesession, gamesession, destroygamesession

function newgamesession(difficulty)
  game = Game(difficulty)
  GAMES[game.id] = game

  game
end

function gamesession(id)
  GAMES[id]
end

function destroygamesession(id)
  delete!(GAMES, id)
end

end

max_stars_repo_name: PacktPublishing/Julia-Projects

max_issues_repo_name: PacktPublishing/Julia-Projects

max_forks_repo_name: PacktPublishing/Julia-Projects


text:
function add_noise!(x, W, C, hₙ, ϵ)
    Cc = cholesky(C + ϵ*I)
    x.+= sqrt(2*hₙ)*Cc.L*W
end

function teks_update(u::Array{T,2}, y, Γ, C₀, G;
    parallel=false,
    h₀=1,
    iδ=100,
    λ=1,
    ϵ=convert(T, 1e-10)) where T

    J = size(u, 2)
    #calculate i-step sample covariance in parameter space
    ū = mean(u, dims=2)
    Cu = Symmetric((u .- ū)*(u .- ū)')/(J-1)
    # forward map and differencing operator
    if parallel
        g = reduce(hcat, pmap(G, eachcol(u)))
    else
        g = reduce(hcat, map(G, eachcol(u)))
    end
    ḡ = mean(g, dims=2)
    D = (g .- y)' * (Γ \ (g .- ḡ))
    #Regularization/prior, adaptive stepsize, gradient step
    ∇R = (C₀\ Cu)' # because both are symmetric you can do this...
    hₙ = h₀ / (norm(D) + 1/iδ)
    Du = u .- ((hₙ/J)*D*u')'
    #Implicit solve of partial step
    I∇R = I + hₙ*∇R
    u = I∇R \ Du
    # perturb parameters according to i-step sample covariance
    Wu = randn(T, size(u))
    # u .+= sqrt(2*hₙ)*Cuc.L*Wu
    add_noise!(u, Wu, Cu, hₙ, ϵ)
    return (u, hₙ, g)
end

function teks(u::Array{T,2}, y, Γ, C₀⁻¹, G, n_steps;
              parallel=false,
              h₀=1,
              iδ=100,
              λ=1,
              ϵ=convert(T, 1e-10),
              savechain=false) where T
    if savechain
        uchain = Array{T,2}[]
    end
    hₙchain = T[]
    for i = 1:n_steps
        u, hₙ, _ = teks_update(u, y, Γ, C₀⁻¹, G;
                            parallel=parallel,
                            h₀=h₀,
                            iδ=iδ,
                            λ=λ,
                            ϵ=ϵ)
        if savechain
            push!(uchain, u)
        end
        push!(hₙchain, hₙ)
    end
    if savechain
        return (uchain, hₙchain)
    else
        return (u, hₙchain)
    end
end

function implicit_hyper_solve(θΔθ, Cθ, hₙ, ∇nlpθ_fun; solver=NelderMead())
    θ, Δθ = θΔθ
    # ∇nlpθ_fun = gradient of negative log prior in respect to θ
    res = optimize(θ̂ -> sum((θ̂ .+ hₙ*Cθ*∇nlpθ_fun(θ̂) - Δθ).^2), θ[:], solver)
    return Optim.minimizer(res)
end

function whteks_update(ξ::Array{T,2}, θ::Array{T,2}, y, Γ, GT, ∇nlpθ_fun;
                        parallel=false,
                        h₀=1,
                        iδ=100,
                        λ=1,
                        ϵ=convert(T, 1e-10),
                        hyper_solver=NelderMead()) where T
    J = size(ξ, 2)
    #calculate i-step sample covariance in parameter space & update parameters
    ξ̄ = mean(ξ, dims=2)
    θ̄ = mean(θ, dims=2)
    Cξ = Symmetric((ξ .- ξ̄)*(ξ .- ξ̄)')/(J-1)
    Cθ = Symmetric((θ .- θ̄)*(θ .- θ̄)')/(J-1)
    Cξc = cholesky(Cξ + ϵ*I)
    Cθc = cholesky(Cθ + ϵ*I)
    # forward map and differencing operator
    if parallel
        g = reduce(hcat, pmap(GT, zip(eachcol(ξ), eachcol(θ))))
    else
        g = reduce(hcat, map(GT, zip(eachcol(ξ), eachcol(θ))))
    end
    ḡ = mean(g, dims=2)
    D = (g .- y)' * (Γ \ (g .- ḡ))
    #Regularization / prior and adaptive stepsize, gradient step
    ∇Rξ = Cξ
    hₙ = h₀ / (norm(D) + 1/iδ)
    Dξ = ξ .- ((hₙ/J)*D*ξ')'
    Dθ = θ .- ((hₙ/J)*D*θ')'
    #Implicit solve of partial step
    # I∇Rξ = cholesky(Symmetric(I + hₙ*∇Rξ))
    I∇Rξ = I + hₙ*∇Rξ
    ξ = I∇Rξ \ Dξ
    if parallel
        θ = reduce(hcat,
                    pmap(θDθ->implicit_hyper_solve(θDθ, Cθ, hₙ, ∇nlpθ_fun, solver=hyper_solver),
                    zip(eachcol(θ), eachcol(Dθ))))
    else
        θ = reduce(hcat,
                    map(θDθ->implicit_hyper_solve(θDθ, Cθ, hₙ, ∇nlpθ_fun, solver=hyper_solver),
                    zip(eachcol(θ), eachcol(Dθ))))
    end
    # perturb parameters according to i-step sample covariance
    Wξ = randn(T, size(ξ))
    Wθ = randn(T, size(θ))
    ξ .+= sqrt(2*hₙ)*Cξc.L*Wξ
    θ .+= sqrt(2*hₙ)*Cθc.L*Wθ
    return (ξ, θ, hₙ, g)
end

function whteks(ξ::Array{T,2}, θ::Array{T,2}, y, Γ⁻¹, GT, ∇nlpθ_fun, n_steps;
                parallel=false,
                h₀=1,
                iδ=100,
                λ=1,
                ϵ=convert(T, 1e-10),
                savechain=false,
                hyper_solver=NelderMead()) where T

    if savechain
        ξchain = Array{T,2}[]
        θchain = Array{T,2}[]
    end
    hₙchain = T[]
    J = size(ξ, 2)
    for i = 1:n_steps
        ξ, θ, hₙ, _ = whteks_update(ξ, θ, y, Γ⁻¹, GT, ∇nlpθ_fun;
                                parallel=parallel,
                                h₀=h₀,
                                iδ=iδ,
                                λ=λ,
                                ϵ=ϵ,
                                hyper_solver=hyper_solver)
        if savechain
            push!(ξchain, ξ)
            push!(θchain, θ)
        end
        push!(hₙchain, hₙ)
    end
    if savechain
        return (ξchain, θchain, hₙchain)
    else
        return (ξ, θ, hₙchain)
    end

end


text:
"constraint: `c^2 + d^2 <= a*b`"
function relaxation_complex_product(m, a, b, c, d)
    @assert (getlowerbound(a) >= 0 && getlowerbound(b) >= 0) || (getupperbound(a) <= 0 && getupperbound(b) <= 0)
    @constraint(m, c^2 + d^2 <= a*b)
end

"In the literature this constraints are called the Lifted Nonlinear Cuts (LNCs)"
function cut_complex_product_and_angle_difference(m, wf, wt, wr, wi, angmin, angmax)
    @assert angmin >= -pi/2 && angmin <= pi/2
    @assert angmax >= -pi/2 && angmax <= pi/2
    @assert angmin < angmax

    vfub = sqrt(getupperbound(wf))
    vflb = sqrt(getlowerbound(wf))
    vtub = sqrt(getupperbound(wt))
    vtlb = sqrt(getlowerbound(wt))
    tdub = angmax
    tdlb = angmin

    phi = (tdub + tdlb)/2
    d   = (tdub - tdlb)/2

    sf = vflb + vfub
    st = vtlb + vtub

    @constraint(m, sf*st*(cos(phi)*wr + sin(phi)*wi) - vtub*cos(d)*st*wf - vfub*cos(d)*sf*wt >=  vfub*vtub*cos(d)*(vflb*vtlb - vfub*vtub))
    @constraint(m, sf*st*(cos(phi)*wr + sin(phi)*wi) - vtlb*cos(d)*st*wf - vflb*cos(d)*sf*wt >= -vflb*vtlb*cos(d)*(vflb*vtlb - vfub*vtub))
end

"""
```
c^2 + d^2 <= a*b*getupperbound(z)
c^2 + d^2 <= getupperbound(a)*b*getupperbound(z)
c^2 + d^2 <= a*getupperbound(b)*z
```
"""
function relaxation_complex_product_on_off(m, a, b, c, d, z)
    @assert getlowerbound(c) <= 0 && getupperbound(c) >= 0
    @assert getlowerbound(d) <= 0 && getupperbound(d) >= 0
    # assume c and d are already linked to z in other constraints
    # and will be forced to 0 when z is 0

    a_ub = getupperbound(a)
    b_ub = getupperbound(b)
    z_ub = getupperbound(z)

    @constraint(m, c^2 + d^2 <= a*b*z_ub)
    @constraint(m, c^2 + d^2 <= a_ub*b*z)
    @constraint(m, c^2 + d^2 <= a*b_ub*z)
end

"`x - getupperbound(x)*(1-z) <= y <= x - getlowerbound(x)*(1-z)`"
function relaxation_equality_on_off(m, x, y, z)
    # assumes 0 is in the domain of y when z is 0

    x_ub = getupperbound(x)
    x_lb = getlowerbound(x)

    @constraint(m, y >= x - x_ub*(1-z))
    @constraint(m, y <= x - x_lb*(1-z))
end

"""
general relaxation of a square term

```
x^2 <= y <= (getupperbound(x)+getlowerbound(x))*x - getupperbound(x)*getlowerbound(x)
```
"""
function relaxation_sqr(m, x, y)
    @constraint(m, y >= x^2)
    @constraint(m, y <= (getupperbound(x)+getlowerbound(x))*x - getupperbound(x)*getlowerbound(x))
end

"general relaxation of a sine term, in -pi/2 to pi/2"
function relaxation_sin(m, x, y)
    ub = getupperbound(x)
    lb = getlowerbound(x)
    @assert lb >= -pi/2 && ub <= pi/2

    max_ad = max(abs(lb),abs(ub))

    if lb < 0 && ub > 0
        @constraint(m, y <= cos(max_ad/2)*(x - max_ad/2) + sin(max_ad/2))
        @constraint(m, y >= cos(max_ad/2)*(x + max_ad/2) - sin(max_ad/2))
    end
    if ub <= 0
        @constraint(m, y <= (sin(lb) - sin(ub))/(lb-ub)*(x - lb) + sin(lb))
        @constraint(m, y >= cos(max_ad/2)*(x + max_ad/2) - sin(max_ad/2))
    end
    if lb >= 0
        @constraint(m, y <= cos(max_ad/2)*(x - max_ad/2) + sin(max_ad/2))
        @constraint(m, y >= (sin(lb) - sin(ub))/(lb-ub)*(x - lb) + sin(lb))
    end
end

"general relaxation of a cosine term, in -pi/2 to pi/2"
function relaxation_cos(m, x, y)
    ub = getupperbound(x)
    lb = getlowerbound(x)
    @assert lb >= -pi/2 && ub <= pi/2

    max_ad = max(abs(lb),abs(ub))

    @constraint(m, y <= 1 - (1-cos(max_ad))/(max_ad*max_ad)*(x^2))
    @constraint(m, y >= (cos(lb) - cos(ub))/(lb-ub)*(x - lb) + cos(lb))
end

"general relaxation of a sine term, in -pi/2 to pi/2"
function relaxation_sin_on_off(m, x, y, z, M_x)
    ub = getupperbound(x)
    lb = getlowerbound(x)
    @assert lb >= -pi/2 && ub <= pi/2

    max_ad = max(abs(lb),abs(ub))

    @constraint(m, y <= z*sin(ub))
    @constraint(m, y >= z*sin(lb))

    @constraint(m,  y - cos(max_ad/2)*(x) <= z*(sin(max_ad/2) - cos(max_ad/2)*max_ad/2) + (1-z)*(cos(max_ad/2)*M_x))
    @constraint(m, -y + cos(max_ad/2)*(x) <= z*(sin(max_ad/2) + cos(max_ad/2)*max_ad/2) + (1-z)*(cos(max_ad/2)*M_x))

    @constraint(m, y <= z*(sin(max_ad/2) + cos(max_ad/2)*max_ad/2))
    @constraint(m, -y <= z*(sin(max_ad/2) + cos(max_ad/2)*max_ad/2))

    @constraint(m, cos(max_ad/2)*x <= z*(sin(max_ad/2) - cos(max_ad/2)*max_ad/2 + sin(max_ad)) + (1-z)*(cos(max_ad/2)*M_x))
    @constraint(m, -cos(max_ad/2)*x <= z*(sin(max_ad/2) - cos(max_ad/2)*max_ad/2 + sin(max_ad)) + (1-z)*(cos(max_ad/2)*M_x))
end

"general relaxation of a cosine term, in -pi/2 to pi/2"
function relaxation_cos_on_off(m, x, y, z, M_x)
    ub = getupperbound(x)
    lb = getlowerbound(x)
    @assert lb >= -pi/2 && ub <= pi/2

    max_ad = max(abs(lb),abs(ub))

    @constraint(m, y <= z)
    @constraint(m, y >= z*cos(max_ad))
    # can this be integrated?
    #@constraint(m, y >= (cos(lb) - cos(ub))/(lb-ub)*(x - lb) + cos(lb))

    @constraint(m, y <= z - (1-cos(max_ad))/(max_ad^2)*(x^2) + (1-z)*((1-cos(max_ad))/(max_ad^2)*(M_x^2)))
end

"""
general relaxation of binlinear term (McCormick)

```
z >= getlowerbound(x)*y + getlowerbound(y)*x - getlowerbound(x)*getlowerbound(y)
z >= getupperbound(x)*y + getupperbound(y)*x - getupperbound(x)*getupperbound(y)
z <= getlowerbound(x)*y + getupperbound(y)*x - getlowerbound(x)*getupperbound(y)
z <= getupperbound(x)*y + getlowerbound(y)*x - getupperbound(x)*getlowerbound(y)
```
"""
function relaxation_product(m, x, y, z)
    x_ub = getupperbound(x)
    x_lb = getlowerbound(x)
    y_ub = getupperbound(y)
    y_lb = getlowerbound(y)

    @constraint(m, z >= x_lb*y + y_lb*x - x_lb*y_lb)
    @constraint(m, z >= x_ub*y + y_ub*x - x_ub*y_ub)
    @constraint(m, z <= x_lb*y + y_ub*x - x_lb*y_ub)
    @constraint(m, z <= x_ub*y + y_lb*x - x_ub*y_lb)
end

"""
On/Off variant of binlinear term (McCormick)
requires that all variables (x,y,z) go to zero with ind
"""
function relaxation_product_on_off(m, x, y, z, ind)
    @assert getlowerbound(x) <= 0 && getupperbound(x) >= 0
    @assert getlowerbound(y) <= 0 && getupperbound(y) >= 0
    @assert getlowerbound(z) <= 0 && getupperbound(z) >= 0

    x_ub = getupperbound(x)
    x_lb = getlowerbound(x)
    y_ub = getupperbound(y)
    y_lb = getlowerbound(y)

    @constraint(m, z >= x_lb*y + y_lb*x - ind*x_lb*y_lb)
    @constraint(m, z >= x_ub*y + y_ub*x - ind*x_ub*y_ub)
    @constraint(m, z <= x_lb*y + y_ub*x - ind*x_lb*y_ub)
    @constraint(m, z <= x_ub*y + y_lb*x - ind*x_ub*y_lb)
end

"""
convex hull relaxation of trilinear term

```
w₁ = getlowerbound(x)*getlowerbound(y)*getlowerbound(z)
w₂ = getlowerbound(x)*getlowerbound(y)*getupperbound(z)
w₃ = getlowerbound(x)*getupperbound(y)*getlowerbound(z)
w₄ = getlowerbound(x)*getupperbound(y)*getupperbound(z)
w₅ = getupperbound(x)*getlowerbound(y)*getlowerbound(z)
w₆ = getupperbound(x)*getlowerbound(y)*getupperbound(z)
w₇ = getupperbound(x)*getupperbound(y)*getlowerbound(z)
w₈ = getupperbound(x)*getupperbound(y)*getupperbound(z)
w = λ₁*w₁ + λ₂*w₂ + λ₃*w₃ + λ₄*w₄ + λ₅*w₅ + λ₆*w₆ + λ₇*w₇ + λ₈*w₈
x = (λ₁ + λ₂ + λ₃ + λ₄)*getlowerbound(x) + (λ₅ + λ₆ + λ₇ + λ₈)*getupperbound(x)
y = (λ₁ + λ₂ + λ₅ + λ₆)*getlowerbound(x) + (λ₃ + λ₄ + λ₇ + λ₈)*getupperbound(x)
z = (λ₁ + λ₃ + λ₅ + λ₇)*getlowerbound(x) + (λ₂ + λ₄ + λ₆ + λ₈)*getupperbound(x)
λ₁ + λ₂ + λ₃ + λ₄ + λ₅ + λ₆ + λ₇ + λ₈ = 1
```
"""
function relaxation_trilinear(m, x, y, z, w, lambda)
    x_ub = getupperbound(x)
    x_lb = getlowerbound(x)
    y_ub = getupperbound(y)
    y_lb = getlowerbound(y)
    z_ub = getupperbound(z)
    z_lb = getlowerbound(z)

    @assert length(lambda) == 8

    w_val = [x_lb * y_lb * z_lb
             x_lb * y_lb * z_ub
             x_lb * y_ub * z_lb
             x_lb * y_ub * z_ub
             x_ub * y_lb * z_lb
             x_ub * y_lb * z_ub
             x_ub * y_ub * z_lb
             x_ub * y_ub * z_ub]

    @constraint(m, w == sum(w_val[i]*lambda[i] for i in 1:8))
    @constraint(m, x == (lambda[1] + lambda[2] + lambda[3] + lambda[4])*x_lb +
                        (lambda[5] + lambda[6] + lambda[7] + lambda[8])*x_ub)
    @constraint(m, y == (lambda[1] + lambda[2] + lambda[5] + lambda[6])*y_lb +
                        (lambda[3] + lambda[4] + lambda[7] + lambda[8])*y_ub)
    @constraint(m, z == (lambda[1] + lambda[3] + lambda[5] + lambda[7])*z_lb +
                        (lambda[2] + lambda[4] + lambda[6] + lambda[8])*z_ub)
    @constraint(m, sum(lambda) == 1)
end


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

using BAT
using Test

using Random, StatsBase, Distributions
using HypothesisTests

import NestedSamplers

@testset "test_ellipsoidal_nested_sampling" begin

    p = Uniform(-80,80)
    prior = BAT.NamedTupleDist(a=[p,p,p])

    dist = [
        MixtureModel(Normal, [(-50.0,2.5),(50.0,2.5)],[0.2,0.8]),
        MixtureModel(Normal, [(-20.0,2.5),(20.0,2.5)],[0.7,0.3]),
        MixtureModel(Normal, [(-40.0,5.0),(10.0,1.0)],[0.5,0.5])
    ]

    likelihood = params -> begin
        r1 = logpdf(dist[1],params.a[1])
        r2 = logpdf(dist[2],params.a[2])
        r3 = logpdf(dist[3],params.a[3])
        return LogDVal(r1+r2+r3)
    end

    posterior = PosteriorDensity(likelihood, prior)
    algorithm = EllipsoidalNestedSampling()
    r = bat_sample(posterior, algorithm);
    smpls = r.result

    @test logvalof(posterior).(smpls.v) ≈ smpls.logd

    iid = BAT.NamedTupleDist(a=dist)
    iidsamples, chains = bat_sample(iid, IIDSampling());
    @test ones(3) ≈ isapprox.(bat_compare(smpls,iidsamples).result.ks_p_values, 1.0, atol = 0.8)

    logz_expected = -log(prod(160))*3
    @test isapprox(r.logintegral.val, logz_expected, atol = 100 * r.logintegral.err)

end


text:
using AutomotiveDrivingModels
using AutoViz
using NearestNeighbors

import PyPlot

# export
#     DemoState,
#     DemoMotionPrimitives,
#     DemoIncreasedHorizonBehavior,
#     DemoBuildingBlocks,
#     DemoObjective,
#     DemoTailgateAvoidance,
#     DemoObserveActObjective,
#     DemoObstacleMap

function DemoState(track,scene,carcolors;targetid=1,zoom=5.0)
    x = scene.vehicles[1].state.posG.x
    y = scene.vehicles[1].state.posG.y
    θ = scene.vehicles[1].state.posG.θ
    v = scene.vehicles[1].state.v

    s,t,ϕ = kdProject(x,y,θ,track.tree,track.roadway,track.models[1])
    println("Vehicle 1: ")
    @show x
    @show y
    @show θ
    @show v
    @show s
    @show t
    @show ϕ

    render(scene, track.roadway, cam=CarFollowCamera(targetid, zoom), car_colors=carcolors, canvas_height=300)
end
function DemoMotionPrimitives(track,context)
    demoCar = HRHC(1,track.roadway,context,h=50,v_max=120.0,μ=30.0,a_step=12.0,a_range=[-1,0,1],k=2)

    v_idx = [50,150,220]
    δ_idx = [18,14,40]
    a_idx = [3,1,2]
    titles = ["V = 50", "V = 150", "V = 220"]
    h=12
    motion_map = demoCar.motion_map

    PyPlot.figure(figsize=[12,4])
    for i in 1:3
        PyPlot.subplot(1,3,i)
        PyPlot.axis("tight")
        PyPlot.axis("off")
        ΔXYθ = motion_map[v_idx[i]][:,:,h,:]
        trajectory = motion_map[v_idx[i]][a_idx[i],:,1:h,:]
        for j in 1:size(trajectory,1)
            PyPlot.plot(trajectory[j,:,2],trajectory[j,:,1],c="red")
        end
        PyPlot.plot(trajectory[δ_idx[i],:,2],trajectory[δ_idx[i],:,1],c="lime",linewidth=2)
        PyPlot.scatter(ΔXYθ[:,:,2],ΔXYθ[:,:,1],edgecolor="none")

        PyPlot.axis("equal")
        PyPlot.ylim([0,80])
        PyPlot.xlim([-40,40])

        PyPlot.title(titles[i])
    end
end
function DemoIncreasedHorizonBehavior(track,context)
    demoCar = HRHC(1,track.roadway,context,h=50,v_max=120.0,μ=30.0,a_step=12.0,a_range=[-1,0,1],k=2)
    v_idx = 220
    δ_idx = 40
    a_idx = 3
    h=demoCar.h
    motion_map = demoCar.motion_map
    PyPlot.figure()
    ΔXYθ = motion_map[v_idx][:,:,h,:]
    trajectory = motion_map[v_idx][a_idx,:,1:h,:]
    for j in 1:size(trajectory,1)
        PyPlot.plot(trajectory[j,:,2],trajectory[j,:,1],c="red")
    end
    PyPlot.plot(trajectory[δ_idx,:,2],trajectory[δ_idx,:,1],c="lime",linewidth=2)
    PyPlot.scatter(ΔXYθ[:,:,2],ΔXYθ[:,:,1],edgecolor="none")

    PyPlot.axis("tight")
    PyPlot.axis("off")
    PyPlot.axis("equal")
    PyPlot.title("Increased Horizon Behavior (V = 220)")
end
function DemoBuildingBlocks()
    t1 = linspace(-13,13,200)
    T_MAX = 10
    t_cost = exp(abs(t1/T_MAX).^6)

    ϕ1 = linspace(-Float64(π),Float64(π),200)
    ϕ_MAX = Float64(π)/2
    ϕ_cost = exp(abs(ϕ1/ϕ_MAX).^2)

    s1 = linspace(0,1,20)
    s_factor = 5
    s_cost = s_factor*(1-s1/maximum(s1))

    PyPlot.figure(figsize=[12,4])
    PyPlot.subplot(1,4,1)
    PyPlot.plot(t1/T_MAX,t_cost)
    PyPlot.ylim([0,50])
    PyPlot.title("cost(t)")
    PyPlot.subplot(1,4,2)
    PyPlot.plot(ϕ1/ϕ_MAX,ϕ_cost)
    PyPlot.ylim([0,50])
    PyPlot.title("cost(phi)")
    PyPlot.subplot(1,4,3)
    PyPlot.plot(s1,s_cost)
    PyPlot.ylim([0,50])
    PyPlot.title("cost(s)")
end
function DemoObjective(t_shift)
    θ = linspace(-Float64(π),Float64(π),360)
    r = linspace(1.0,20.0,20)
    x = [cos(i) for i in θ]*[i for i in r]'
    y = [sin(i) for i in θ]*[i for i in r]'
    R = [1 for i in θ]*[i for i in r]'
    Θ = [i for i in θ]*[1 for i in r]'
    s = x
    t = y - t_shift
    ϕ = Θ

    car_ID = 1
    s₀ = -21 # scene.vehicles[car_ID].state.posF.s

    T_MAX= 10
    ϕ_MAX = Float64(π)

    s_norm = s/maximum(s)
    t_norm = t/T_MAX
    ϕ_norm = ϕ/ϕ_MAX

    #costs
    t_cost = abs(t_norm).^6
    ϕ_cost = abs(ϕ_norm).^6
    s_factor = 1
    s_cost = s_factor*(1-s_norm)
    A = [1 .5; #  [ϕ t] [a1 a2] [ϕ]
        .5 0] #         [a2 a3] [t]
    tϕ_cost = A[1,1]*(ϕ_norm).^2 + (A[1,2]+A[2,1])*(ϕ_norm).*(t_norm) + A[2,2]*(t_norm).^2

    objective = calculateObjective(car_ID,s,s₀,t,ϕ,T_MAX)

    PyPlot.figure(figsize=[12,3])

    PyPlot.subplot(1,4,1)
    PyPlot.plot([-20,20],[T_MAX+t_shift, T_MAX+t_shift],c="black",linewidth=2)
    PyPlot.plot([-20,20],[-T_MAX+t_shift, -T_MAX+t_shift],c="black",linewidth=2)
    PyPlot.scatter(x,y,c=s_cost,edgecolor="none",s=4)
    PyPlot.axis("off")
    PyPlot.axis("equal")
    PyPlot.title("cost(s)")
    PyPlot.subplot(1,4,2)
    PyPlot.plot([-20,20],[T_MAX+t_shift, T_MAX+t_shift],c="black",linewidth=2)
    PyPlot.plot([-20,20],[-T_MAX+t_shift,-T_MAX+t_shift],c="black",linewidth=2)
    PyPlot.scatter(x,y,c=t_cost,edgecolor="none",s=4)
    PyPlot.axis("off")
    PyPlot.axis("equal")
    PyPlot.title("cost(t)")
    PyPlot.subplot(1,4,3)
    PyPlot.plot([-20,20],[T_MAX+t_shift, T_MAX+t_shift],c="black",linewidth=2)
    PyPlot.plot([-20,20],[-T_MAX+t_shift,-T_MAX+t_shift],c="black",linewidth=2)
    PyPlot.scatter(x,y,c=tϕ_cost,edgecolor="none",s=4)
    PyPlot.axis("off")
    PyPlot.axis("equal")
    PyPlot.title("cost(t,phi)")
    PyPlot.subplot(1,4,4)
    PyPlot.plot([-20,20],[T_MAX+t_shift, T_MAX+t_shift],c="black",linewidth=2)
    PyPlot.plot([-20,20],[-T_MAX+t_shift,-T_MAX+t_shift],c="black",linewidth=2)
    PyPlot.scatter(x,y,c=log(objective),edgecolor="none",s=4)
    PyPlot.scatter(x[indmin(objective)],y[indmin(objective)],c="white",s=30)
    PyPlot.axis("off")
    PyPlot.axis("equal")
    PyPlot.title("objective = cost(s,t,phi)")
end
function DemoTailgateAvoidance()
    θ = linspace(-Float64(π),Float64(π),360)
    r = linspace(1.0,20.0,20)
    x = [cos(i) for i in θ]*[i for i in r]'
    y = [sin(i) for i in θ]*[i for i in r]'
    R = [1 for i in θ]*[i for i in r]'
    ψ = [i for i in θ]*[1 for i in r]'

    #costs
    R_cost = 1./(2.5*R.*cos(ψ).^2 + 2*4.8*R.*sin(ψ).^2 + 1)
    # ψ_cost = 1./(R.*(sin(ψ).^2 + 1))
    ψ_cost = -cos(ψ).^3 + 1.1
    Rψ_cost = R_cost.*ψ_cost + 1

    PyPlot.figure(figsize=[12,4])

    PyPlot.subplot(1,3,1)
    PyPlot.scatter(x,y,c=R_cost,edgecolor="none",s=8)
    PyPlot.axis("equal")
    PyPlot.axis("off")
    PyPlot.title("cost(r)")
    PyPlot.subplot(1,3,2)
    PyPlot.scatter(x,y,c=ψ_cost,edgecolor="none",s=8)
    PyPlot.axis("equal")
    PyPlot.axis("off")
    PyPlot.title("cost(psi)")
    PyPlot.subplot(1,3,3)
    PyPlot.scatter(x,y,c=Rψ_cost,edgecolor="none",s=8)
    PyPlot.axis("equal")
    PyPlot.axis("off")
    PyPlot.title("cost(r,psi)")

    return x,y,Rψ_cost
end
function DemoObserveActObjective(track,scene;steps=50)
    lane_width = track.roadway.segments[1].lanes[1].width
    actions = Array(DriveAction, length(scene))
    for i in 1:steps
        k_level = 0 # needs to be updated into a loop
        for k_level in 0:maximum([model.k for (id,model) in track.models])
            for (i, veh) in enumerate(scene)
                model = track.models[veh.def.id]
                observe!(model, scene, track.roadway, veh.def.id, track.tree, track.obstacleMap, k_level)
                actions[i] = rand(model)
            end
        end
        for (veh, action) in zip(scene, actions)
            model = track.models[veh.def.id]
            context = action_context(model)
            veh.state = propagate(veh, action, context, track.roadway)
        end
    end

    PyPlot.figure(figsize=[12,8])

    lo = Inf
    hi = 0

    for (id,hrhc) in track.models
        lo = Int(min(lo,hrhc.curve_ind))
        hi = Int(max(hi,hrhc.curve_ind + Int(1+div(hrhc.v_range[end]*hrhc.Δt*hrhc.h,hrhc.Δs))))
        lane_width = track.roadway.segments[1].lanes[1].width

        x = zeros(hrhc.h,size(hrhc.successor_states,1),size(hrhc.successor_states,2))
        y = zeros(size(x))
        Θ = zeros(size(x))
        s = zeros(size(x))
        t = zeros(size(x))
        ϕ = zeros(size(x))
        objective = zeros(size(x))
        avoidanceCost = zeros(size(x))
        collisionFlag = zeros(size(x))

        s₀ = scene.vehicles[hrhc.car_ID].state.posF.s
        for i in 1:hrhc.h
            successor_states = getSuccessorStates(hrhc.motion_map[hrhc.v_cmd],hrhc.car_ID, i, scene)
            x[i,:,:] = copy(successor_states[:,:,1])
            y[i,:,:] = copy(successor_states[:,:,2])
            Θ[i,:,:] = copy(successor_states[:,:,3])
            s[i,:,:], t[i,:,:], ϕ[i,:,:] = loopProjectionKD(hrhc,scene,track.roadway,track.tree)
            objective[i,:,:] = calculateObjective(hrhc.car_ID,s[i,:,:],s₀,t[i,:,:],ϕ[i,:,:],hrhc.T_MAX)
            avoidanceCost[i,:,:] = tailgateAvoidance(hrhc, track.obstacleMap, track.tree, track.roadway, scene, hrhc.k)
            objective[i,:,:] = objective[i,:,:] + avoidanceCost[i,:,:]
            collisionFlag[i,:,:] = screenCollision(hrhc, track.obstacleMap, track.tree, track.roadway, scene, hrhc.k)
        end
        objective[collisionFlag .> 0] = Inf
        trajectory = track.obstacleMap[hrhc.k][hrhc.car_ID]

        PyPlot.scatter(x,y,c=log(objective),edgecolor="none")
        PyPlot.plot(trajectory[:,1],trajectory[:,2],color="red")
        PyPlot.axis("off")
        PyPlot.title("log objective functions of vehicles on track")

    end
    plotSplineRoadway(track.x[lo:hi],track.y[lo:hi],track.θ[lo:hi],lane_width)
end
function DemoObstacleMap(track,k_level)
    lo = Inf
    hi = 0
    lane_width = track.roadway.segments[1].lanes[1].width
    for (id,hrhc) in track.models
        lo = Int(min(lo,hrhc.curve_ind))
        hi = Int(max(hi,hrhc.curve_ind + Int(1+div(hrhc.v_range[end]*hrhc.Δt*hrhc.h,hrhc.Δs))))
        # lane_width = track.roadway.segments[1].lanes[1].width
    end
    plotSplineRoadway(track.x[lo:hi],track.y[lo:hi],track.θ[lo:hi],lane_width)

    for (car_id,trajectory) in track.obstacleMap[k_level]
        PyPlot.plot(trajectory[:,1],trajectory[:,2],c="r",linewidth=1)
        PyPlot.scatter(trajectory[end,1],trajectory[end,2],c="r",edgecolor="none",s=10)
    end
    PyPlot.axis("equal")
end


text:
function beep(sound::Tuple{Matrix{<:AbstractFloat}, <:AbstractFloat})::Nothing
    wavplay(sound[1], sound[2])
    return nothing
end

function beep(sound::Symbol)::Nothing
    @assert sound ∈ keys(beep_sounds) "No such sound. Check out: `keys(beep_sounds)`."

    sound_to_play = beep_sounds[sound]
    beep(sound_to_play)
end

function load_sound(fname::String)::Tuple{Matrix{<:AbstractFloat}, <:AbstractFloat}
    raw_wav = wavread(fname)
    return (raw_wav[1], raw_wav[2])
end


text:
using PrettyTables
using Statistics

function hasoutliers(data)
    Qi, Qf = quantile(data, [0.01, 0.99])
    return any( data .> Qf) || any(data .< Qi)
end

function raw_average(p = :fermi, engine = "MKL")
    if p == :fermi
        timings = get_fermi(engine)
        avg = [sum(timings[s,:])/10 for s = 1:22]
    else
        timings = get_psi4()
        avg = [sum(timings[s,:])/5 for s = 1:22]
    end
    return avg
end

function get_fermi(engine = "MKL")
    timings = zeros(22,10)
    for s in 1:22
        for i in 1:10
            rex = Regex("S$s - \\(T\\) alg $(engine == "MKL" ? 1 : 2)\\s+?$i run:\\s+([0-9]*\\.?[0-9]*)")
            timings[s,i] = get_linedata(rex, "fermi/output.dat")
        end
    end
    return timings
end

function get_psi4()
    Erex = r"FNOCC:\striples\s+:.+?([0-9]*\.?[0-9]*)w\s+[0-9]+\scalls"
    timings = zeros(22, 5)
    for s in 1:22
        for i in 1:5
            path = joinpath(@__DIR__, "psi4/R$i/S$s/timer.dat")
            timings[s,i] = get_linedata(Erex, path)
        end
    end
    return timings
end

function get_linedata(rex, fname)
    fpath = joinpath(@__DIR__, fname)
    out = 0.0
    f = false
    for l = eachline(fpath)
        m = match(rex, l)
        if m !== nothing
            out = m.captures[1] |> String |> x->parse(Float64,x)
            f = true
            break
        end
    end
    if !f
        error("Data not found for $(rex) in $(fname)")
    end
    return out
end

if abspath(PROGRAM_FILE) == @__FILE__
    fermi_mkl = get_mkl()
    fermi_octavian = get_octavian()
    psi4 = get_psi4()

    @pt :header = ["Fermi - MKL", "Fermi - Octavian", "Psi4"] [fermi_mkl fermi_octavian psi4]
end

max_stars_repo_name: FermiQC/Fermi

max_issues_repo_name: FermiQC/Fermi

max_forks_repo_name: FermiQC/Fermi


text:
abstract type AbstractResult end

mutable struct Result{T <: Configuration, R <: Number} <: AbstractResult
    technique::String
    start::T
    minimum::T
    cost_minimum::R
    iterations::Int
    current_iteration::Int
    cost_calls::Int
    is_final::Bool
    current_time::Float64

    function Result{T, R}(technique::String,
                          start::T,
                          minimum::T,
                          cost_minimum::R,
                          iterations::Int,
                          current_iteration::Int,
                          cost_calls::Int,
                          is_final::Bool) where {T <: Configuration, R <: Number}
        new(technique, start, minimum, cost_minimum,
            iterations, current_iteration, cost_calls, is_final)
    end
end

function Result(technique::String,
                start::T,
                minimum::T,
                cost_minimum::R,
                iterations::Int,
                current_iteration::Int,
                cost_calls::Int,
                is_final::Bool) where {T <: Configuration, R <: Number}
    Result{T, R}(technique, start, minimum, cost_minimum,
                 iterations, current_iteration, cost_calls, is_final)
end


text:
using ProblemParser
using Utils

using Underscores

const file = find_input(@__FILE__)
const input = parse(Split(",", Convert(Int16)), slurp(file))

function part1(input)
    sort!(input)
    median = input[length(input) ÷ 2]
    sum(abs(median - pos) for pos in input)
end

assertequal(part1(input), 349769)

triangular(n) = n*(n + 1) ÷ 2

function part2(input)
    mean = sum(input) ÷ length(input)
    sum(triangular(abs(mean - pos)) for pos in input)
end

assertequal(part2(input), 99540554)

max_stars_repo_name: CmdQ/AoC


text:
using DiffGPFlux
using Test

γ = rand(3)
X = rand(3, 10)

@testset "DiffGPFlux.jl" begin
    @test (K=rbf(γ, X, X); isposdef(K))   # verify postive definite
    @test_broken (K=linear(γ, X, X); issymmetric(K))   # numerical error will affect linear kernel !!!
    @test (K=rbf(γ, X, X); diag(K)==ones(size(K, 1)))  # for RBF kernel, diagonal elements should all be 1
    @test (K=rbf(γ, X, X); eltype(K)==eltype(γ))    # type of K should be the same as type of γ
end

max_stars_repo_name: HamletWantToCode/DiffGPFlux

max_issues_repo_name: HamletWantToCode/DiffGPFlux


text:
using Gridap
using GridapEmbedded
using LineSearches: Static

using GridapMHD
using GridapMHD: driver_inductionless_MHD

# Problem setting
ρ = 1.0
ν = 1.0
σ = 1.0

U0 = 10.0
B0 = 20.0
L = 1.0
Re = U0 * L / ν
Ha = B0 * L * sqrt(σ/(ρ*ν))

function g_u(x)
  if abs(x[1] + 0.05) < 1e-8
    y = x[2]; z = x[3]
    return VectorValue(1.0/0.0125^4 * (y-0.0125) * (y+0.0125) *
                      (z-0.0125) * (z+0.0125), 0.0, 0.0)
  else
    return VectorValue(0.0, 0.0, 0.0)
  end
end

function g_j(x)
  return VectorValue(0.0,0.0,0.0)
end

function f_u(x)
  return VectorValue(0.0,0.0,0.0)
end

function map1(coord)
  ncoord = streching(coord,  domain=(0.0,0.025,-0.05,-0.0375,-0.0125,0.0),factors=(3.0,3.0,3.0),dirs=(1,2,3))
  ncoord = streching(ncoord, domain=(0.025,0.05,-0.0375,-0.025,0.0,0.0125),factors=(1.25,1.25,1.25),dirs=(1,2,3))
  ncoord = streching(ncoord, domain=(0.175,0.2,-0.0125,0.0),factors=(1.25,3.0),dirs=(1,2))
  ncoord = streching(ncoord, domain=(0.0,0.0125),factors=(1.25,),dirs=(2,))
  ncoord = streching(ncoord, domain=(0.025,0.0375),factors=(3.0,),dirs=(2,))
  ncoord = streching(ncoord, domain=(0.0375,0.05),factors=(1.25,),dirs=(2,))
  ncoord
end

# Background mesh definition
partition = (10,12,4);
#partition = (60,24,6);
domain = (-0.05,0.2,-0.05,0.05,-0.0125,0.0125)
bgmodel=CartesianDiscreteModel(domain,partition,map=map1);

# Build the final domain using simple geomtries
inlet = cuboid(dx=0.05-1e-6,dy=0.025-1e-6,dz=0.025,x0=Point(-0.025,0.0,0.0));
out = cuboid(dx=0.2-1e-6,dy=0.1-1e-6,dz=0.025,x0=Point(0.1,0.0,0.0));

out2 = cuboid(dx=0.15+1e-6,dy=0.0125+1e-6,dz=0.025,x0=Point(0.125,0.01875,0.0));
out3 = cuboid(dx=0.15+1e-6,dy=0.0125+1e-6,dz=0.025,x0=Point(0.125,-0.01875,0.0));
out1 = setdiff(out,out3);
outlet = setdiff(out1,out2);

# Add the simple geomtries and "cut"
geo = union(inlet,outlet,name="csg");

# "Cut" the final geomtry over the Background mesh
cutgeo = cut(bgmodel,geo);
model = DiscreteModel(cutgeo,"csg");

# Filter to select all walls
function walls(coord)
  x = coord[1]
  y = coord[2]
  z = coord[3]
  tol = 1e-6
  if (-0.05+tol < x < -tol) & ((0.0125-tol < y) | (y < -0.0125+tol ))
    return true
  elseif (z > 0.0125 - tol) | ( z < -0.0125 + tol)
    return true
  elseif (-tol < x < tol)
    if (-0.05-tol < y < -0.0125+tol) | (0.0125-tol < y < 0.5+tol)
      return true
    end
  elseif (0.05-tol < x < 0.2-tol) & ((0.0125-tol < y < 0.025+tol) | (-0.025-tol < y < -0.0125+tol ))
      return true
  elseif (y < -0.05 + tol ) | ( y > 0.05 - tol)
    return true
  end
  false
end

# Filter to select inflow boundary
function inlet_boundary(coord)
  x = coord[1]
  y = coord[2]
  tol = 1e-6

  if (x<-0.05+tol)
    return true
  end
  false
end

# Filter to select outflow boundary
function outlet_boundary(coord)
  x = coord[1]
  y = coord[2]
  tol = 1e-6

  if (x>0.2-tol)
    return true
  end
  false
end

# Name inlet, outlet, and walls element sets.
add_entity!(model,walls,"walls")
add_entity!(model,inlet_boundary,"inlet")
add_entity!(model,outlet_boundary,"outlet")

model1 = simplexify(model)

# Call main MHD driver
xh, trian, dΩ = driver_inductionless_MHD(;
  Re=Re,
  Ha=Ha,
  model=model1,
  fluid_dirichlet_tags = ["inlet","walls"],
  fluid_neumann_tags = ["outlet"],
  magnetic_dirichlet_tags = ["walls"],
  magnetic_neumann_tags = ["inlet","outlet"],
  fluid_dirichlet_conditions = g_u ,
  magnetic_dirichlet_conditions = g_j,
  fluid_body_force = f_u,
  linesearch=Static(),
  constraint_presures = (false,false),
  precond_tau = 1e-10
)

uh, ph, jh, φh = xh

# Scale unknowns
uh = uh * U0
jh = jh * σ * B0 * U0

uh_r = restrict(uh,trian)
ph_r = restrict(ph,trian)
jh_r = restrict(jh,trian)
φh_r = restrict(φh,trian)
writevtk(trian,"results",cellfields=["u"=>uh_r,"p"=>ph_r,"j"=>jh_r,"phi"=>φh_r])


text:
module ReversePropagation

export gradient, forward_backward_contractor

import Symbolics: toexpr

using SymbolicUtils
using SymbolicUtils: Sym, Term
using SymbolicUtils.Rewriters

using Symbolics
using Symbolics: value,
                istree, operation, arguments,
                Assignment

using IntervalContractors

import Base: ∩
import Base: ∪

@register a ∩ b
@register a ∪ b

using OrderedCollections

using ChainRules

# struct Assignment
#     lhs
#     rhs
# end

# Base.show(io::IO, eq::Assignment) = print(io, lhs(eq), " := ", rhs(eq))

include("make_variable.jl")
include("chain_rules_interface.jl")
# include("cse.jl")
include("cse_new.jl")
include("reverse_diff.jl")
include("reverse_icp.jl")

end


text:
using Tar, CodecZlib, ZipFile, Downloads

const IHACRES_VERSION = v"0.4.1"
const BASE_URL = "https://github.com/ConnectedSystems/ihacres_nim/releases/download/v$(IHACRES_VERSION)/ihacres_nim_"
const HERE = @__DIR__

extract_path = joinpath(HERE, "usr", "lib")

mkpath(extract_path)

if Sys.iswindows()
    target = "windows.zip"
elseif Sys.islinux()
    target = "linux.tar.gz"
elseif Sys.isapple()
    target = "macos.tar.gz"
else
    throw(DomainError("Unsupported platform"))
end

target_url = BASE_URL * target

function untar(data::IO, exdir::String)
    if isdir(exdir) == false mkdir(exdir) end
    tar = GzipDecompressorStream(data)
    Tar.extract(tar, exdir; copy_symlinks=true)
    close(tar)
end

fn = Downloads.download(target_url, "ihacres_nim")
if endswith(target_url, "zip")
    dir_name = split(replace(target_url, ".zip" => ""), "/")[end]
    zarchive = ZipFile.Reader(fn)
    for f in zarchive.files
        if endswith(f.name, "/")
            # skip directories
            continue
        end

        extract_fn = replace(f.name, dir_name * "/" => "")
        file_loc = joinpath(extract_path, extract_fn)
        open(file_loc, "w") do fp
            write(fp, read(f))
        end
    end

    close(zarchive)

elseif endswith(target_url, "tar.gz")
    tmp_dir = mktempdir()

    open(fn, "r") do fp
        untar(fp, tmp_dir)
    end

    # Copy files to destination
    # Necessary as Julia's `mv` and `cp` functions
    # require empty directories (`force` option deletes an existing directory)
    for pfn in readdir(tmp_dir)
        open(joinpath(tmp_dir, pfn), "r") do src
            open(joinpath(extract_path, pfn), "w") do dst
                write(dst, read(src))
            end
        end
    end
end

rm(fn)


text:
# This file is part of the TaylorIntegration.jl package; MIT licensed

testfiles = (
    "one_ode.jl",
    "many_ode.jl",
    "complex.jl",
    "jettransport.jl",
    "lyapunov.jl",
    "bigfloats.jl",
    "common.jl",
    "rootfinding.jl",
    "taylorize.jl"
    )

for file in testfiles
    include(file)
end


text:
using Statistics
using Plots
using BenchmarkTools
using Profile
using FFTW
using Statistics
using Optim
using Images, FileIO, ImageIO
using Printf
using Revise
using Profile
using LinearAlgebra
using Random
using Distributions
using FITSIO
using SparseArrays

push!(LOAD_PATH, pwd()*"/main")
using DHC_2DUtils
#include("Deriv_Utils_New.jl")
#import Deriv_Utils_New
push!(LOAD_PATH, pwd()*"/scratch_NM")
using Deriv_Utils_New
using Data_Utils
using Visualization
using ReconFuncs

#Rewritten
logbool= false                                                #Calculate coefficients of the log of the image?
apdbool = true                                                #Apodize?
isobool = false                                               #Iso?

#Image you want to start with
Nx=64
sfdall = readsfd_fromsrc("scratch_NM/data/dust10000.fits", Nx, logbool=false)
input = sfdall[:, :, 5]                     #Less of a memory crunch if you just read in the one image you want to start with

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
dhc_args = Dict(:doS2=>false, :doS20=>true, :apodize=>apdbool, :iso=>isobool)
lambda = 0.0
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])

#Coefficient to optimize wrt
Jvalue = 1
Lvalue = 0
#Code currently requires you to access S1 through the S2 matrix diagonal
mask_nf = (filter_hash["J_L"][:, 1] .== Jvalue) .& (filter_hash["J_L"][:, 2] .== Lvalue)
mask_nf2 = Diagonal(mask_nf)
(Nf,) = size(filter_hash["filt_index"])
coeff_mask = falses(2+Nf+Nf^2)
coeff_mask[3+Nf:end] .= mask_nf2[:]

#Setting a target coefficient: currently I'm taking the 10*true target coefficient
#=

if logbool
    procdbn = log.(sfdall)
else
    procdbn = sfdall
end

dbncoeffs = get_dbn_coeffs(procdbn, filter_hash, dhc_args, coeff_mask = coeff_mask)
sigval = std(dbncoeffs, dims=1)
=#

if logbool
    s_targ_mean = 0.1*DHC_compute_wrapper(log.(input), filter_hash, norm=false; dhc_args...)[coeff_mask]             #target coefficient vector = twice the current wst value
    s_targ_invcov = Diagonal([0.1.^(-2)])
    res, recon_img = Deriv_Utils_New.image_recon_derivsum_regularized(log.(input), filter_hash, s_targ_mean, s_targ_invcov, falses(Nx, Nx), dhc_args, coeff_mask=coeff_mask, optim_settings=optim_settings, lambda=lambda)
    recon_img = exp.(recon_img)
else
    s_targ_mean = 0.1*DHC_compute_wrapper(input, filter_hash, norm=false; dhc_args...)[coeff_mask]             #target coefficient vector = twice the current wst value
    s_targ_invcov = Diagonal([0.0015.^(-2)])          #Might need to tune. inverse covariance matrix with sigma= true_coeff/10 for J=1 L=2
    res, recon_img = Deriv_Utils_New.image_recon_derivsum_regularized(input, filter_hash, s_targ_mean, s_targ_invcov, falses(Nx, Nx), dhc_args, optim_settings=optim_settings, coeff_mask=coeff_mask, lambda=lambda)
end

heatmap(input)
heatmap(recon_img)

#Iso
#Rewritten
logbool= false                                               #Calculate coefficients of the log of the image?
apdbool = true                                               #Apodize?
isobool = true                                               #Iso?

#Image you want to start with
Nx=64
sfdall = readsfd_fromsrc("scratch_NM/data/dust10000.fits", Nx, logbool=false)
input = sfdall[:, :, 5]                     #Less of a memory crunch if you just read in the one image you want to start with

filter_hash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1, Omega=true)
dhc_args = Dict(:doS2=>false, :doS20=>true, :apodize=>apdbool, :iso=>isobool)
lambda = 0.0
optim_settings = Dict([("iterations", 1000), ("norm", false), ("minmethod", ConjugateGradient())])

#Coefficient to optimize wrt
Jvalue = 1
Lvalue = 0
filter_hash["S2_iso_mat"]
#Code currently requires you to access S1 through the S2 matrix diagonal
mask_nf = (filter_hash["J_L"][:, 1] .== Jvalue) .& (filter_hash["J_L"][:, 2] .== Lvalue)
mask_nf2 = Diagonal(mask_nf)
(Nf,) = size(filter_hash["filt_index"])
coeff_mask = falses(2+Nf+Nf^2)
coeff_mask[3+Nf:end] .= mask_nf2[:]

if logbool
    s_targ_mean = 0.1*DHC_compute_wrapper(log.(input), filter_hash, norm=false; dhc_args...)[coeff_mask]             #target coefficient vector = twice the current wst value
    s_targ_invcov = Diagonal([0.1.^(-2)])
    res, recon_img = Deriv_Utils_New.image_recon_derivsum_regularized(log.(input), filter_hash, s_targ_mean, s_targ_invcov, falses(Nx, Nx), dhc_args, coeff_mask=coeff_mask, optim_settings=optim_settings, lambda=lambda)
    recon_img = exp.(recon_img)
else
    s_targ_mean = 0.1*DHC_compute_wrapper(input, filter_hash, norm=false; dhc_args...)[coeff_mask]             #target coefficient vector = twice the current wst value
    s_targ_invcov = Diagonal([0.0015.^(-2)])          #Might need to tune. inverse covariance matrix with sigma= true_coeff/10 for J=1 L=2
    res, recon_img = Deriv_Utils_New.image_recon_derivsum_regularized(input, filter_hash, s_targ_mean, s_targ_invcov, falses(Nx, Nx), dhc_args, optim_settings=optim_settings, coeff_mask=coeff_mask, lambda=lambda)
end

heatmap(input)
heatmap(recon_img)

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
using IntervalArithmetic
using IntervalRootFinding

@testset "Root search iterator interface" begin
    contractor = Newton(sin, cos)
    search = BreadthFirstSearch(-10..10, contractor, 1e-3)

    # First iteration of the search
    elem, state = iterate(search)

    # Optional iterator methods
    @test eltype(search) != Any
end

@testset "Search strategies" begin
    # DepthFristSearch and BreadthFristSearch should return the same results
    X = -5..5
    rts = roots(sin, cos, X, Newton, DepthFirstSearch)
    @test Set(rts) == Set(roots(sin, cos, X, Newton, BreadthFirstSearch))
end


text:
module FastSphericalHarmonics

using FastTransforms
using LinearAlgebra
using StaticArrays

include("points.jl")
include("scalar.jl")
# include("vector.jl")
include("spin.jl")

end


text:
function convertBI(path::String)
    data = []
    for line in open(readlines,path)
        line = chomp(line)
        if isempty(line)
            push!(data, "")
        else
            items = split(line, ' ')
            tag = items[3]
            if tag == "O"
            elseif tag[1] == 'B'
                tag = tag[3:end]
            elseif tag[1] == 'I'
                tag = "_"
            else
                tbrow("")
            end
            push!(data, "$(items[1])\t$(items[2])\t$tag")
        end
    end
    open("test.out", "w") do f
        foreach(x -> println(f,x), data)
    end
end

convertBI("C:/Users/hshindo/Dropbox/chunking/test.txt")

max_stars_repo_name: hshindo/Merlin-Examples

max_issues_repo_name: hshindo/Merlin-Examples

max_forks_repo_name: hshindo/Merlin-Examples


text:
__precompile__(false)

# LibTOAST.jl: interface to the TOAST++ library
# Copyright (C) 2019 Samuel Powell

module LibTOAST

using SparseArrays
using LinearAlgebra
using Libdl
using Compat
using Cxx

# Set verbosity, provide info function
verbose = true
_info(s) = verbose && @info(s)

# Load Toast++ libarary path
include(joinpath(dirname(@__FILE__), "..", "deps", "path.jl"))

function __init__()

  # Add header locations: toast, include, libfe, libmath
  addHeaderDir(_jl_toast_dir; kind = C_System)
  addHeaderDir(joinpath(_jl_toast_dir, "include"); kind = C_System)
  addHeaderDir(joinpath(_jl_toast_dir, "src", "libfe"); kind = C_System)
  addHeaderDir(joinpath(_jl_toast_dir, "src", "libmath"); kind = C_System)
  addHeaderDir(joinpath(_jl_toast_dir, "src", "libstoast"); kind = C_System)

  # Define header options accroding to library build options
  defineMacro("TOAST_THREAD")     # Enable threading
  defineMacro("FDOT")             # Enable flourescence (projection)

  # Include headers: felib (this includes all required headers)
  cxxinclude("mathlib.h")         # Matrices and vectors, tasks and verbosity
  cxxinclude("felib.h")           # Mesh related functions
  cxxinclude("stoastlib.h")
  # cxxinclude("source.h")          # Source and detector profiles

  # Import dynamic libraries: libsuperlu, libmath, libfe
  dlopen(_jl_toast_libsuperlu, Libdl.RTLD_GLOBAL)
  dlopen(_jl_toast_libmath, Libdl.RTLD_GLOBAL)
  dlopen(_jl_toast_libfe, Libdl.RTLD_GLOBAL)
  dlopen(_jl_toast_libstoast, Libdl.RTLD_GLOBAL)

  # Initialise Toast++ thread pool
  @cxx Task_Init(0)

end

include("util.jl")            # Utility functions
include("mesh.jl")            # Mesh types
include("sysmat.jl")          # System matrices
include("raster.jl")       # Raster mapping
include("coefftypes.jl")      # Coefficients
include("assembly.jl")        # System matrix assembly
include("regul.jl")  # Regularisation

end # module


text:
using DataArrays,DataFrames

v110 = rcopy(DataArray,reval("x <- 1:10"))
@test isa(v110,DataVector)
@test eltype(v110) == Cint

attenu = rcopy(DataFrame,:attenu)
@test isa(attenu,DataFrame)
@test size(attenu) == (182,5)

dist = attenu[:dist]
@test isa(dist,DataArray{Float64})

@test rcopy(DataArray,"c(NA,TRUE)").na == @data([NA,true]).na
@test rcopy(DataArray,"c(NA,1)").na == @data([NA,1.0]).na
@test rcopy(DataArray,"c(NA,1+0i)").na == @data([NA,1.0+0.0*im]).na
@test rcopy(DataArray,"c(NA,1L)").na == @data([NA,one(Int32)]).na
@test rcopy(DataArray,"c(NA,'NA')").na == @data([NA,"NA"]).na
@test_throws ErrorException rcopy(DataArray,"as.factor(c('a','a','c'))")
@test rcopy(PooledDataArray,"as.factor(c('a','a','c'))").pool == ["a","c"]

@test rcopy(DataArray,RObject(@data([NA,true]))).na == @data([NA,true]).na
@test rcopy(DataArray,RObject(@data([NA,1]))).na == @data([NA,1]).na
@test rcopy(DataArray,RObject(@data([NA,1.]))).na == @data([NA,1.]).na
@test rcopy(DataArray,RObject(@data([NA,1.+0*im]))).na == @data([NA,1.+0*im]).na
@test rcopy(DataArray,RObject(@data([NA,NA,"a","b"]))).na == @data([NA,NA,"a","b"]).na
pda = PooledDataArray(repeat(["a", "b"], inner = [5]))
@test rcopy(PooledDataArray,RObject(pda)).refs == repeat([1,2], inner = [5])

@test rcopy(rcall(:dim,RObject(attenu))) == [182,5]


text:
"""
    AbstractManifold{𝔽}

A type to represent a (Riemannian) manifold.
The [`AbstractManifold`](@ref) is a central type of this interface.
It allows to distinguish different implementations of functions like the [`exp`](@ref)onential and
[`log`](@ref)arithmic map for different manifolds.
Usually, the manifold is the first parameter in any of these functions within `ManifoldsBase.jl`.
Based on these, say “elementary” functions, as the two mentioned above, more general functions are built,
for example the [`shortest_geodesic`](@ref) and the [`geodesic`](@ref).
These should only be overwritten (reimplemented) if for a certain manifold specific, more efficient implementations are possible, that do not just call the elementary functions.

The [`AbstractManifold`] is parametrized by [`AbstractNumbers`](@ref) to distinguish for example
real (ℝ) and complex (ℂ) manifolds.

For subtypes the preferred order of parameters is: size and simple value parameters,
followed by the [`AbstractNumbers`](@ref) `field`, followed by data type parameters,
which might depend on the abstract number field type.
"""
abstract type AbstractManifold{𝔽} end

"""
    AbstractManifoldPoint

Type for a point on a manifold.
While an [`AbstractManifold`](@ref) does not necessarily require this
type, for example when it is implemented for `Vector`s or `Matrix` type elements, this type
can be used either

* for more complicated representations,
* semantic verification, or
* when dispatching on different representations of points on a manifold.

Since semantic verification and different representations usually might still only store a
matrix internally, it is possible to use [`@manifold_element_forwards`](@ref) and
[`@default_manifold_fallbacks`](@ref) to reduce implementation overhead.
"""
abstract type AbstractManifoldPoint end