using DiffEqBase

####################################################################################################
# Poincare shooting based on Sánchez, J., M. Net, B. Garcı́a-Archilla, and C. Simó. “Newton–Krylov Continuation of Periodic Orbits for Navier–Stokes Flows.” Journal of Computational Physics 201, no. 1 (November 20, 2004): 13–33. https://doi.org/10.1016/j.jcp.2004.04.018.

function sectionHyp!(out, x, normals, centers)
	for ii = 1:length(normals)
		out[ii] = dot(normals[ii], x) - dot(normals[ii], centers[ii])
	end
	out
end

# this composite type encodes a set of hyperplanes which are used as Poincaré sections
struct HyperplaneSections{Tn, Tc, Ti, Tnb, Tcb}
	M::Int64		# number of hyperplanes
	normals::Tn 	# normals to define hyperplanes
	centers::Tc 	# representative point on each hyperplane
	indices::Ti 	# indices to be removed in the operator Ek

	normals_bar::Tnb
	centers_bar::Tcb

	function HyperplaneSections(normals, centers)
		M = length(normals)
		indices = zeros(Int64, M)
		for ii in 1:M
			indices[ii] = argmax(abs.(normals[ii]))
		end
		nbar = [R(normals[ii], indices[ii]) for ii=1:M]
		cbar = [R(centers[ii], indices[ii]) for ii=1:M]

		return new{typeof(normals), typeof(centers), typeof(indices), typeof(nbar), typeof(cbar)}(M, normals, centers, indices, nbar, cbar)
	end

	HyperplaneSections(M = 0) = new{Nothing, Nothing, Nothing, Nothing, Nothing}(M, nothing, nothing, nothing, nothing, nothing)
end

(hyp::HyperplaneSections)(out, u) = sectionHyp!(out, u, hyp.normals, hyp.centers)

"""
	update!(hyp::HyperplaneSections, normals, centers)

Update the hyperplanes saved in `hyp`.
"""
function update!(hyp::HyperplaneSections, normals, centers)
	M = hyp.M
	@assert length(normals) == M "Wrong number of normals"
	@assert length(centers) == M "Wrong number of centers"
	hyp.normals .= normals
	hyp.centers .= centers
	for ii in 1:M
		hyp.indices[ii] = argmax(abs.(normals[ii]))
		R!(hyp.normals_bar[ii], normals[ii], hyp.indices[ii])
		R!(hyp.centers_bar[ii], centers[ii], hyp.indices[ii])
	end
	return hyp
end

# Operateur Rk from the paper above
function R!(out, x::AbstractVector, k::Int)
	@views out[1:k-1] .= x[1:k-1]
	@views out[k:end] .= x[k+1:end]
	return out
end

R!(hyp::HyperplaneSections, out, x::AbstractVector, ii::Int) = R!(out, x, hyp.indices[ii])

function R(x::AbstractVector, k::Int)
	out = similar(x, length(x) - 1)
	R!(out, x, k)
end

function R(hyp::HyperplaneSections, x::AbstractVector, ii::Int)
	out = similar(x, length(x) - 1)
	R!(hyp, out, x, ii)
end

# differential of R
dR!(hyp::HyperplaneSections, out, dx::AbstractVector, ii::Int) = R!(hyp, out, dx, ii)

# Operateur Ek from the paper above
function E!(hyp::HyperplaneSections, out, xbar::AbstractVector, ii::Int)
	@assert length(xbar) == length(hyp.normals[1]) - 1 "Wrong size for the projector / expansion operators, length(xbar) = $(length(xbar)) and length(normal) = $(length(hyp.normals[1]))"
	k = hyp.indices[ii]
	nbar = hyp.normals_bar[ii]
	xcbar = hyp.centers_bar[ii]
	coord_k = hyp.centers[ii][k] - (dot(nbar, xbar) - dot(nbar, xcbar)) / hyp.normals[ii][k]

	@views out[1:k-1] .= xbar[1:k-1]
	@views out[k+1:end] .= xbar[k:end]
	out[k] = coord_k
	return out
end

function E(hyp::HyperplaneSections, xbar::AbstractVector, ii::Int)
	out = similar(xbar, length(xbar) + 1)
	E!(hyp, out, xbar, ii)
end

# differential of E!
function dE!(hyp::HyperplaneSections, out, dxbar::AbstractVector, ii::Int)
	k = hyp.indices[ii]
	nbar = hyp.normals_bar[ii]
	xcbar = hyp.centers_bar[ii]
	coord_k = - dot(nbar, dxbar) / hyp.normals[ii][k]

	@views out[1:k-1]   .= dxbar[1:k-1]
	@views out[k+1:end] .= dxbar[k:end]
	out[k] = coord_k
	return out
end

function dE(hyp::HyperplaneSections, dxbar::AbstractVector, ii::Int)
	out = similar(dxbar, length(dxbar) + 1)
	dE!(hyp, out, dxbar, ii)
end
####################################################################################################
"""

`pb = PoincareShootingProblem(flow::Flow, M, sections; δ = 1e-8, interp_points = 50, isparallel = false)`

This composite type implements the Poincaré Shooting method to locate periodic orbits by relying on Poincaré return maps. The arguments are as follows
- `flow::Flow`: implements the flow of the Cauchy problem though the structure [`Flow`](@ref).
- `M`: the number of Poincaré sections. If `M==1`, then the simple shooting is implemented and the multiple one otherwise.
- `sections`: function or callable strict which implements a Poincaré section condition. The evaluation `sections(x)` must return a scalar number when `M==1`. Otherwise, one must implement a function `section(out, x)` which populates `out` with the `M` sections.
- `δ = 1e-8` used to compute the jacobian of the fonctional by finite differences. If set to `0`, an analytical expression of the jacobian is used instead.
- `interp_points = 50` number of interpolation point used to define the callback (to compute the hitting of the hyperplan section)
- `isparallel = false` whether the shooting are computed in parallel (threading). Only available through the use of Flows defined by `EnsembleProblem`.

## Simplified constructors
- A simpler way is to create a functional is

`pb = PoincareShootingProblem(F, p, prob::ODEProblem, alg, section; kwargs...)`

for simple shooting or

`pb = PoincareShootingProblem(F, p, prob::Union{ODEProblem, EnsembleProblem}, alg, M::Int, section; kwargs...)`

for multiple shooting . Here `F(x,p)` is the vector field, `p` is a parameter (to be passed to the vector and the flow), `prob` is an `Union{ODEProblem, EnsembleProblem}` which is used to create a flow using the ODE solver `alg` (for example `Tsit5()`). Finally, the arguments `kwargs` are passed to the ODE solver defining the flow. We refere to `DifferentialEquations.jl` for more information.

- Another convenient call is

`pb = PoincareShootingProblem(F, p, prob::Union{ODEProblem, EnsembleProblem}, alg, normals::AbstractVector, centers::AbstractVector; δ = 1e-8, kwargs...)`

where `normals` (resp. `centers`) is a list of normals (resp. centers) which defines a list of hyperplanes ``\\Sigma_i``. These hyperplanes are used to define partial Poincaré return maps.

## Computing the functionals
A functional, hereby called `G` encodes this shooting problem. You can then call `pb(orbitguess, par)` to apply the functional to a guess. Note that `orbitguess::AbstractVector` must be of size M * N where N is the number of unknowns in the state space and `M` is the number of Poincaré maps. Another accepted `guess` is such that `guess[i]` is the state of the orbit on the `i`th section. This last form allows for non-vector state space which can be convenient for 2d problems for example.

- `pb(orbitguess, par)` evaluates the functional G on `orbitguess`
- `pb(orbitguess, par, du)` evaluates the jacobian `dG(orbitguess).du` functional at `orbitguess` on `du`

!!! tip "Tip"
    You can use the function `getPeriod(pb, sol, par)` to get the period of the solution `sol` for the problem with parameters `par`.
"""
@with_kw struct PoincareShootingProblem{Tf, Tsection <: HyperplaneSections} <: AbstractShootingProblem
	M::Int64 = 0								# number of Poincaré sections
	flow::Tf = Flow()							# should be a Flow{TF, Tf, Td}
	section::Tsection = HyperplaneSections(M)	# Poincaré sections
	δ::Float64 = 0e-8							# Numerical value used for the Matrix-Free Jacobian by finite differences. If set to 0, analytical jacobian is used
	isparallel::Bool = false					# whether we use DE in Ensemble mode for multiple shooting
end

@inline isParallel(psh::PoincareShootingProblem) = psh.isparallel

function PoincareShootingProblem(F, p,
			prob::ODEProblem, alg,
			hyp::HyperplaneSections;
			δ = 1e-8, interp_points = 50, isparallel = false, kwargs...)
	pSection(out, u, t, integrator) = (hyp(out, u); out .= out .* (integrator.iter > 1))
	affect!(integrator, idx) = terminate!(integrator)
	# we put nothing option to have an upcrossing
	cb = VectorContinuousCallback(pSection, affect!, hyp.M; interp_points = interp_points, affect_neg! = nothing)
	# change the ODEProblem -> EnsembleProblem for the parallel case
	_M = hyp.M
	isparallel = _M==1 ? false : isparallel
	_pb = isparallel ? EnsembleProblem(prob) : prob
	return PoincareShootingProblem(flow = Flow(F, p, _pb, alg; callback = cb, kwargs...), M = hyp.M, section = hyp, δ = δ, isparallel = isparallel)
end

# this is the "simplest" constructor to use in automatic branching from Hopf
# this is a Hack to pass the arguments to construct a Flow. Indeed, we need to provide the
# appropriate callback for Poincare Shooting to work
PoincareShootingProblem(M::Int, par, prob::ODEProblem, alg; isparallel = false, kwargs...) = PoincareShootingProblem(M = M, flow = (par = par, prob = prob, alg = alg), isparallel = isparallel)

PoincareShootingProblem(M::Int, par, prob1::ODEProblem, alg1, prob2::ODEProblem, alg2; isparallel = false, kwargs...) = PoincareShootingProblem(M = M, flow = (par = par, prob1 = prob1, alg1 = alg1, prob2 = prob2, alg2 = alg2), isparallel = isparallel)

function PoincareShootingProblem(F, p,
			prob::ODEProblem, alg,
			normals::AbstractVector, centers::AbstractVector;
			δ = 1e-8, interp_points = 50, isparallel = false, kwargs...)
	return PoincareShootingProblem(F, p,
					prob, alg,
					HyperplaneSections(normals, centers);
					δ = δ, interp_points = interp_points, isparallel = isparallel, kwargs...)
end

function PoincareShootingProblem(F, p,
				prob1::ODEProblem, alg1,
				prob2::ODEProblem, alg2,
				hyp::HyperplaneSections;
				δ = 1e-8, interp_points = 50, isparallel = false, kwargs...)
	pSection(out, u, t, integrator) = (hyp(out, u); out .= out .* (integrator.iter > 1))
	affect!(integrator, idx) = terminate!(integrator)
	# we put nothing option to have an upcrossing
	cb = VectorContinuousCallback(pSection, affect!, hyp.M; interp_points = interp_points, affect_neg! = nothing)
	# change the ODEProblem -> EnsembleProblem for the parallel case
	_M = hyp.M
	isparallel = _M==1 ? false : isparallel
	_pb1 = isparallel ? EnsembleProblem(prob1) : prob1
	_pb2 = isparallel ? EnsembleProblem(prob2) : prob2
	return PoincareShootingProblem(flow = Flow(F, p, _pb1, alg1, _pb2, alg2; callback = cb, kwargs...), M = hyp.M, section = hyp, δ = δ, isparallel = isparallel)
end

function PoincareShootingProblem(F, p,
				prob1::ODEProblem, alg1,
				prob2::ODEProblem, alg2,
				normals::AbstractVector, centers::AbstractVector;
				δ = 1e-8, interp_points = 50, isparallel = false, kwargs...)
	return PoincareShootingProblem(F, p,
					prob1, alg2, prob2, alg2,
					HyperplaneSections(normals, centers);
					δ = δ, interp_points = interp_points, isparallel = isparallel, kwargs...)
end

"""
	update!(pb::PoincareShootingProblem, centers_bar; _norm = norm)

This function updates the normals and centers of the hyperplanes defining the Poincaré sections.
"""
function update!(pb::PoincareShootingProblem, centers_bar; _norm = norm)
	centers = [E(pb.section, centers_bar[ii], ii) for ii in eachindex(centers_bar)]
	normals = [pb.flow.F(c) for c in centers]
	for ii in eachindex(normals)
		normals[ii] ./= _norm(normals[ii])
	end
	update!(pb.section, normals, centers)
end

function getPeriod(psh::PoincareShootingProblem, x_bar, par)

	M = getM(psh)
	Nm1 = div(length(x_bar), M)

	# reshape the period orbit guess
	x_barc = reshape(x_bar, Nm1, M)

	# variable to hold the computed result
	xc = similar(x_bar, Nm1 + 1, M)
	outc = similar(xc)

	Th = eltype(x_barc)
	period = Th(0)

	# we extend the state space to be able to call the flow, so we fill xc
	if ~isParallel(psh)
		for ii in 1:M
			E!(psh.section, view(xc, :, ii), view(x_barc, :, ii), ii)
			# We need the callback to be active here!!!
			period += @views psh.flow(Val(:TimeSol), xc[:, ii], par, Inf64).t
		end
	else
		for ii in 1:M
			E!(psh.section, view(xc, :, ii), view(x_barc, :, ii), ii)
		end
		solOde =  psh.flow(Val(:TimeSol), xc, par, repeat([Inf64],M))
		for ii in 1:M
			period += solOde[ii].t
		end
	end
	return period
end

function _getMax(psh::PoincareShootingProblem, x_bar::AbstractVector, par; ratio = 1)
	# this function extracts the amplitude of the cycle
	M = getM(psh)
	Nm1 = div(length(x_bar), M)

	# reshape the period orbit guess
	x_barc = reshape(x_bar, Nm1, M)
	xc = similar(x_bar, Nm1 + 1, M)

	Th = eltype(x_bar)
	mx = Th(0)

	if ~isParallel(psh)
		for ii in 1:M
			E!(psh.section, view(xc, :, ii), view(x_barc, :, ii), ii)
			# We need the callback to be active here!!!
			sol = @views psh.flow(Val(:Full), xc[:, ii], par, Inf64)
			mx = max(mx, maximum(sol[1:div(Nm1, ratio), :]))
		end
	else
		for ii in 1:M
			E!(psh.section, view(xc, :, ii), view(x_barc, :, ii), ii)
		end
		solOde =  psh.flow(Val(:Full), xc, par, repeat([Inf64], M) )
		for ii in 1:M
			mx = max(mx, maximum(solOde[ii].u[1:div(Nm1, ratio), :]))
		end
	end

	return mx
end

# Poincaré (multiple) shooting with hyperplanes parametrization
function (psh::PoincareShootingProblem)(x_bar::AbstractVector, par; verbose = false)
	M = getM(psh)
	Nm1 = div(length(x_bar), M)

	# reshape the period orbit guess
	x_barc = reshape(x_bar, Nm1, M)

	# TODO the following declaration of xc allocates. It would be better to make it inplace
	xc = similar(x_bar, Nm1 + 1, M)

	# variable to hold the result of the computations
	outc = similar(xc)

	# we extend the state space to be able to call the flow, so we fill xc
	#TODO create the projections on the fly
	for ii in 1:M
		E!(psh.section, view(xc, :, ii), view(x_barc, :, ii), ii)
	end

	if ~isParallel(psh)
		for ii in 1:M
			im1 = (ii == 1 ? M : ii - 1)
			# We need the callback to be active here!!!
			@views outc[:, ii] .= xc[:, ii] .- psh.flow(xc[:, im1], par, Inf64)
		end
	else
		solOde = psh.flow(xc, par, repeat([Inf64],M))
		for ii in 1:M
			im1 = (ii == 1 ? M : ii - 1)
			# We need the callback to be active here!!!
			@views outc[:, ii] .= xc[:, ii] .- solOde[im1][2]
		end
	end

	# build the array to be returned
	out_bar = similar(x_bar)
	out_barc = reshape(out_bar, Nm1, M)
	for i in 1:M
		R!(psh.section, view(out_barc, :, i), view(outc, :, i), i)
	end
	return out_bar
end

function diffPoincareMap(psh::PoincareShootingProblem, x, par, dx, ii::Int)
	normal = psh.section.normals[ii]
	abs(dot(normal, dx)) > 1e-12 && @warn "Vector does not belong to hyperplane!  dot(normal, dx) = $(abs(dot(normal, dx))) and $(dot(dx, dx))"
	# compute the Poincare map from x
	tΣ, solΣ = psh.flow(Val(:TimeSol), x, par, Inf64)
	z = psh.flow.F(solΣ, par)
	# solution of the variational equation at time tΣ
	# We need the callback to be INACTIVE here!!!
	y = psh.flow(x, par, dx, tΣ; callback = nothing).du
	out = y .- (dot(normal, y) / dot(normal, z)) .* z
end

# jacobian of the shooting functional
function (psh::PoincareShootingProblem)(x_bar::AbstractVector, par, dx_bar::AbstractVector)
	δ = psh.δ
	if δ > 0
		# mostly for debugging purposes
		return (psh(x_bar .+  δ .* dx_bar, par) .- psh(x_bar, par)) ./ δ
	end

	# otherwise analytical Jacobian
	M = getM(psh)
	Nm1 = div(length(x_bar), M)

	# reshape the period orbit guess
	x_barc  = reshape( x_bar, Nm1, M)
	dx_barc = reshape(dx_bar, Nm1, M)

	# variable to hold the computed result
	xc  = similar( x_bar, Nm1 + 1, M)
	dxc = similar(dx_bar, Nm1 + 1, M)
	outc = similar(xc)

	# we extend the state space to be able to call the flow, so we fill xc
	for ii in 1:M
		 E!(psh.section,  view(xc, :, ii),  view(x_barc, :, ii), ii)
		dE!(psh.section, view(dxc, :, ii), view(dx_barc, :, ii), ii)
	end

	if ~isParallel(psh)
		for ii in 1:M
			im1 = (ii == 1 ? M : ii - 1)
			@views outc[:, ii] .= dxc[:, ii] .- diffPoincareMap(psh, xc[:, im1], par, dxc[:, im1], im1)
		end
	else
		@assert 1==0 "Analytical Jacobian for parallel Poincare Shooting not implemented yet. Please use the option δ > 0."
	end

	# build the array to be returned
	out_bar = similar(x_bar)
	out_barc = reshape(out_bar, Nm1, M)
	for ii in 1:M
		dR!(psh.section, view(out_barc, :, ii), view(outc, :, ii), ii)
	end
	return out_bar
end

####################################################################################################
# functions needed Branch switching from Hopf bifurcation point

function update(prob::PoincareShootingProblem, F, dF, hopfpt, ζr, M, centers, period)
	# make the section
	normals = [F(u, hopfpt.params) for u in centers]
	for n in normals; n ./= norm(n); end

	@assert ~(prob.flow isa Flow) "Somehow, this method was not called as it should. prob.flow should be a Named Tuple, prob should be constructed with the simple constructor, not yielding a Flow for its flow field."

	# update the problem
	if length(prob.flow) == 3
	probPSh = PoincareShootingProblem(F, prob.flow.par, prob.flow.prob, prob.flow.alg, normals, centers)
	else
		probPSh = PoincareShootingProblem(F, prob.flow.par, prob.flow.prob1, prob.flow.alg1, prob.flow.prob2, prob.flow.alg2, normals, centers)
	end

	# create initial guess. We have to pass it through the projection R
	hyper = probPSh.section
	orbitguess_bar = zeros(length(centers[1])-1, M)
	for ii=1:length(normals)
		orbitguess_bar[:, ii] .= R(hyper, centers[ii], ii)
	end

	return probPSh, vec(orbitguess_bar)
end


# # Sensitivity Analysis of SVM

#md # [![](https://img.shields.io/badge/GitHub-100000?style=for-the-badge&logo=github&logoColor=white)](@__REPO_ROOT_URL__/examples/sensitivity-analysis-svm.jl)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/sensitivity-analysis-svm.ipynb)

# This notebook illustrates sensitivity analysis of data points in a [Support Vector Machine](https://en.wikipedia.org/wiki/Support-vector_machine) (inspired from [@matbesancon](http://github.com/matbesancon)'s [SimpleSVMs](http://github.com/matbesancon/SimpleSVMs.jl).)

# For reference, Section 10.1 of https://online.stat.psu.edu/stat508/book/export/html/792 gives an intuitive explanation of what it means to have a sensitive hyperplane or data point. The general form of the SVM training problem is given below (without regularization):

# ```math
# \begin{split}
# \begin{array} {ll}
# \mbox{minimize} & \sum_{i=1}^{N} \xi_{i} \\
# \mbox{s.t.} & \xi_{i} \ge 0 \quad i=1..N  \\
#             & y_{i} (w^T X_{i} + b) \ge 1 - \xi[i]\\
# \end{array}
# \end{split}
# ```
# where
# - `X`, `y` are the `N` data points
# - `ξ` is the soft-margin loss.

# ## Define and solve the SVM

# Import the libraries.

using SCS, DiffOpt, LinearAlgebra, JuMP
import Random, Plots

# Construct separable, non-trivial data points.

N = 100
D = 2
Random.seed!(62)
X = vcat(randn(N ÷ 2, D), randn(N ÷ 2, D) .+ [4.0, 1.5]')
y = append!(ones(N ÷ 2), -ones(N ÷ 2));

# Let's initialize a special model that can understand sensitivities

model = Model(() -> diff_optimizer(SCS.Optimizer))

# Add the variables

@variable(model, l[1:N])
@variable(model, w[1:D])
@variable(model, b);

# Add the constraints.

@constraint(
    model,
    1.0 * l ∈ MOI.Nonnegatives(N),
)
@constraint(
    model,
    cons,
    y .* (X * w .+ b) + l .- 1 ∈ MOI.Nonnegatives(N),
);

# Define the objective and solve

@objective(
    model,
    Min,
    sum(l),
)

optimize!(model) # solve

# We can visualize the separating hyperplane.

loss = objective_value(model)
wv = value.(w)
bv = value(b)

svm_x = [0.0, 3.0]
svm_y = (-bv .- wv[1] * svm_x )/wv[2]

p = Plots.scatter(X[:,1], X[:,2], color = [yi > 0 ? :red : :blue for yi in y], label = "")
Plots.yaxis!(p, (-2, 4.5))
Plots.plot!(p, svm_x, svm_y, label = "loss = $(round(loss, digits=2))", width=3)

# # Experiments
# Now that we've solved the SVM, we can compute the sensitivity of optimal values -- the separating hyperplane in our case -- with respect to perturbations of the problem data -- the data points -- using DiffOpt. For illustration, we've explored two questions:

# - How does a change in labels of the data points (`y=1` to `y=-1`, and vice versa) affect the position of the hyperplane? This is achieved by finding the gradient of `w`, `b` with respect to `y[i]`, the classification label of the ith data point.
# - How does a change in coordinates of the data points, `X`, affects the position of the hyperplane? This is achieved by finding gradient of `w`, `b` with respect to `X[i]`, 2D coordinates of the data points.

# Note that finding the optimal SVM can be modelled as a conic optimization problem:

# ```math
# \begin{align*}
# & \min_{x \in \mathbb{R}^n} & c^T x \\
# & \text{s.t.}               & A x + s = b  \\
# &                           & b \in \mathbb{R}^m  \\
# &                           & s \in \mathcal{K}
# \end{align*}
# ```

# where
# ```math
# \begin{align*}
# c &= [l_1 - 1, l_2 -1, ... l_N -1, 0, 0, ... 0 \text{(D+1 times)}] \\\\

# A &=
# \begin{bmatrix}
#  -l_1 &    0 & ... &    0 &            0 & ... & 0 & 0  \\
#     0 & -l_2 & ... &    0 &            0 & ... & 0 & 0  \\
#     : &    : & ... &    : &            0 & ... & 0 & 0  \\
#     0 &    0 & ... & -l_N &            0 & ... & 0 & 0  \\
#     0 &    0 & ... &    0 & -y_1 X_{1,1} & ... & -y_1 X_{1,N} & -y_1  \\
#     0 &    0 & ... &    0 & -y_2 X_{2,1} & ... & -y_1 X_{2,N} & -y_2  \\
#     : &    : & ... &    : &           :  & ... &          :   & :   \\
#     0 &    0 & ... &    0 & -y_N X_{N,1} & ... & -y_N X_{N,N} & -y_N  \\
# \end{bmatrix} \\\\

# b &= [0, 0, ... 0 \text{(N times)}, l_1 - 1, l_2 -1, ... l_N -1] \\\\

# \mathcal{K} &= \text{Set of Nonnegative cones}
# \end{align*}
# ```

# ## Experiment 1: Gradient of hyperplane wrt the data point labels

# Construct perturbations in data point labels `y` without changing the data point coordinates `X`.

∇ = Float64[]
dy = zeros(N);

# begin differentiating
for Xi in 1:N
    dy[Xi] = 1.0  # set

    MOI.set(
        model,
        DiffOpt.ForwardInConstraint(),
        cons,
        MOI.Utilities.vectorize(dy .* MOI.SingleVariable(b)),
    )

    DiffOpt.forward(model)

    dw = MOI.get.(
        model,
        DiffOpt.ForwardOutVariablePrimal(),
        w
    )
    db = MOI.get(
        model,
        DiffOpt.ForwardOutVariablePrimal(),
        b
    )
    push!(∇, norm(dw) + norm(db))

    dy[Xi] = 0.0  # reset the change made above
end
normalize!(∇);

# Visualize point sensitivities with respect to separating hyperplane. Note that the gradients are normalized.

p2 = Plots.scatter(
    X[:,1], X[:,2],
    color = [yi > 0 ? :red : :blue for yi in y], label = "",
    markersize = ∇ * 20,
)
Plots.yaxis!(p2, (-2, 4.5))
Plots.plot!(p2, svm_x, svm_y, label = "loss = $(round(loss, digits=2))", width=3)

# ## Experiment 2: Gradient of hyperplane wrt the data point coordinates

# Similar to previous example, construct perturbations in data points coordinates `X`.

∇ = Float64[]
dX = zeros(N, D);

# begin differentiating
for Xi in 1:N
    dX[Xi, :] = ones(D)  # set

    for i in 1:D
        MOI.set(
            model,
            DiffOpt.ForwardInConstraint(),
            cons,
            MOI.Utilities.vectorize(dX[:,i] .* MOI.SingleVariable(w[i])),
        )
    end

    DiffOpt.forward(model)

    dw = MOI.get.(
        model,
        DiffOpt.ForwardOutVariablePrimal(),
        w
    )
    db = MOI.get(
        model,
        DiffOpt.ForwardOutVariablePrimal(),
        b
    )
    push!(∇, norm(dw) + norm(db))

    dX[Xi, :] = zeros(D)  # reset the change made ago
end
normalize!(∇);

# We can visualize point sensitivity with respect to the separating hyperplane. Note that the gradients are normalized.

p3 = Plots.scatter(
    X[:,1], X[:,2],
    color = [yi > 0 ? :red : :blue for yi in y], label = "",
    markersize = ∇ * 20,
)
Plots.yaxis!(p3, (-2, 4.5))
Plots.plot!(p3, svm_x, svm_y, label = "loss = $(round(loss, digits=2))", width=3)


# Adapted from:
# SOSDEMO2 --- Lyapunov Function Search
# Section 3.2 of SOSTOOLS User's Manual

facts("SOSDEMO2") do
for solver in sdp_solvers
context("With solver $(typeof(solver))") do
  @polyvar x[1:3]

  # Constructing the vector field dx/dt = f
  f = [-x[1]^3-x[1]*x[3]^2,
      -x[2]-x[1]^2*x[2],
      -x[3]+3*x[1]^2*x[3]-3*x[3]/(x[3]^2+1)]

  m = Model(solver = solver)

  # The Lyapunov function V(x):
  Z = x.^2
  @polyvariable m V Z

  @polyconstraint m V >= sum(x.^2)

  # dV/dx*(x[3]^2+1)*f <= 0
  P = dot(differentiate(V, x), f)*(x[3]^2+1)
  @polyconstraint m P <= 0

  status = solve(m)

  @fact status --> :Optimal

  @fact removemonomials(getvalue(V), Z) --> zero(Polynomial{true, Float64})
end; end; end


#=

# 1: Static container

```@raw html
	<img src='../assets/static_container.png' width="50%" height="50%" alt='missing' /><br>
```

This is a hello world example of SPH.
Simulates motionless fluid in uniform gravitational field.
Things that may be tested by this benchmark are:
* walls,
* zero point stability of time integrator,
* stability of free surface.

Let us begin by declaring a module and importing some stuff.
=#

module static_container

using Printf
using SmoothedParticles

#=
Declare constant parameters
=#

const dr = 3.0e-3       #average particle distance
const h = 1.8*dr        #size of kernel support
const rho0 = 1000.0     #fluid density
const m = rho0*dr^2     #particle mass
const c = 40.0          #numerical speed of sound
const g = -VECY         #gravitational acceleration
const mu = 8.4e-4       #dynamic viscosity of water

const water_depth = 0.14
const box_height = 0.18
const box_width = 0.14
const wall_width = 2.5*dr

##temporal parameters
const dt = 0.2*h/c
const dt_frame = 0.1
const t_end = 0.5

##particle types
const FLUID = 0.
const WALL = 1.

#=
Declare variables to be stored in a Particle
=#

mutable struct Particle <: AbstractParticle
	x::RealVector #position
	v::RealVector #velocity
	a::RealVector #acceleration
	rho::Float64 #density
	type::Float64 #particle type
	Particle(x::RealVector, type::Float64) = new(
		x,
		VEC0,
		VEC0,
		rho0,
		type
	)
end

##dependance of pressure on density
function pressure(p::Particle)
	return c^2*(p.rho - rho0)
end

##fluid identier
function isfluid(p::Particle)::Float64
	return Float64(p.type == FLUID)
end

#=
Define geometry and make particles
=#
function make_system()
	grid = Grid(dr, :square)
	box = Rectangle(0., 0., box_width, box_height)
	fluid = Rectangle(0., 0., box_width, water_depth)
	walls = BoundaryLayer(box, grid, wall_width)
	sys = ParticleSystem(Particle, box + walls, h)
	generate_particles!(sys, grid, fluid, x -> Particle(x, FLUID))
	generate_particles!(sys, grid, walls, x -> Particle(x,  WALL))
	return sys
end

#=
Define particle interactions
=#
@inbounds function balance_of_mass!(p::Particle, q::Particle, r::Float64)
	p.rho += dt*dot(p.x-q.x, p.v-q.v)*m*rDwendland2(h,r)
end

@inbounds function internal_force!(p::Particle, q::Particle, r::Float64)
	if p.type == FLUID
		ker = m*rDwendland2(h,r)
		#pressure force
		p.a += -ker*(pressure(p)/p.rho^2 + pressure(q)/q.rho^2)*(p.x - q.x)
		#viscous force
		p.a += ker*2*mu/(p.rho*q.rho)*(p.v - q.v)
	end
end

function update!(p::Particle)
	p.v += dt*isfluid(p)*(p.a + g)
	p.x += dt*isfluid(p)*p.v
	p.a = VEC0
end

#=
Put everything into a time loop
=#
function main()
	sys = make_system()
	@show sys.key_max
	out = new_pvd_file("results/static_container")
	for k = 0 : Int64(round(t_end/dt))
		if (k %  Int64(round(dt_frame/dt)) == 0)
			println("# of particles = ", length(sys.particles))
			@printf("t = %.6e\n", k*dt)
			save_frame!(out, sys, :rho, :type, :v)
		end
		create_cell_list!(sys)
		apply!(sys, balance_of_mass!)
		apply!(sys, internal_force!)
		apply!(sys, update!)
	end
	save_pvd_file(out)
end ##function main()

end ##module


@doc doc"""
    MisclassLoss <: SupervisedLoss

Misclassification loss that assigns `1` for misclassified
examples and `0` otherwise. It is a generalization of
`ZeroOneLoss` for more than two classes.
"""
struct MisclassLoss <: SupervisedLoss end

agreement(target, output) = target == output

value(::MisclassLoss, agreement::Bool) = agreement ? 0 : 1
deriv(::MisclassLoss, agreement::Bool) = 0
deriv2(::MisclassLoss, agreement::Bool) = 0
value_deriv(::MisclassLoss, agreement::Bool) = agreement ? (0, 0) : (1, 0)

value(loss::MisclassLoss, target::Number, output::Number) = value(loss, agreement(target, output))
deriv(loss::MisclassLoss, target::Number, output::Number) = deriv(loss, agreement(target, output))
deriv2(loss::MisclassLoss, target::Number, output::Number) = deriv2(loss, agreement(target, output))

isminimizable(::MisclassLoss) = false
isdifferentiable(::MisclassLoss) = false
isdifferentiable(::MisclassLoss, at) = at != 0
istwicedifferentiable(::MisclassLoss) = false
istwicedifferentiable(::MisclassLoss, at) = at != 0
isnemitski(::MisclassLoss) = false
islipschitzcont(::MisclassLoss) = false
isconvex(::MisclassLoss) = false
isclasscalibrated(::MisclassLoss) = false
isclipable(::MisclassLoss) = false

# ===============================================================

@doc doc"""
    PoissonLoss <: SupervisedLoss

Loss under a Poisson noise distribution (KL-divergence)

``L(target, output) = exp(output) - target*output``
"""
struct PoissonLoss <: SupervisedLoss end

value(loss::PoissonLoss, target::Number, output::Number) = exp(output) - target*output
deriv(loss::PoissonLoss, target::Number, output::Number) = exp(output) - target
deriv2(loss::PoissonLoss, target::Number, output::Number) = exp(output)
function value_deriv(loss::PoissonLoss, target::Number, output::Number)
    exp_o = exp(output)
    return (exp_o-(target*output), exp_o-target)
end

isdifferentiable(::PoissonLoss) = true
isdifferentiable(::PoissonLoss, y, t) = true
istwicedifferentiable(::PoissonLoss) = true
istwicedifferentiable(::PoissonLoss, y, t) = true
islipschitzcont(::PoissonLoss) = false
isconvex(::PoissonLoss) = true
# TODO: isstrictlyconvex(::PoissonLoss) = ?
isstronglyconvex(::PoissonLoss) = false

# ===============================================================

@doc doc"""
    CrossEntropyLoss <: SupervisedLoss

Cross-entropy loss also known as log loss and logistic loss is defined as:

``L(target, output) = - target*ln(output) - (1-target)*ln(1-output)``
"""
struct CrossEntropyLoss <: SupervisedLoss end
const LogitProbLoss = CrossEntropyLoss

function value(loss::CrossEntropyLoss, target::Number, output::Number)
    target >= 0 && target <=1 || error("target must be in [0,1]")
    output >= 0 && output <=1 || error("output must be in [0,1]")
    if target == 0
        -log(1 - output)
    elseif target == 1
        -log(output)
    else
        -(target * log(output) + (1-target) * log(1-output))
    end
end
deriv(loss::CrossEntropyLoss, target::Number, output::Number) = (1-target) / (1-output) - target / output
deriv2(loss::CrossEntropyLoss, target::Number, output::Number) = (1-target) / (1-output)^2 + target / output^2
value_deriv(loss::CrossEntropyLoss, target::Number, output::Number) = (value(loss,target,output), deriv(loss,target,output))

isdifferentiable(::CrossEntropyLoss) = true
isdifferentiable(::CrossEntropyLoss, y, t) = t != 0 && t != 1
istwicedifferentiable(::CrossEntropyLoss) = true
istwicedifferentiable(::CrossEntropyLoss, y, t) = t != 0 && t != 1
isconvex(::CrossEntropyLoss) = true

# ===============================================================
# L(target, output) = sign(agreement) < 0 ? 1 : 0


using ToeplitzMatrices, StatsBase, Test, LinearAlgebra

using Base: copyto!

ns = 101
nl = 2000

xs = randn(ns, 5)
xl = randn(nl, 5)

@testset("Toeplitz: $st",
    for (As, Al, st) in ((Toeplitz(0.9.^(0:ns-1), 0.4.^(0:ns-1)),
                            Toeplitz(0.9.^(0:nl-1), 0.4.^(0:nl-1)),
                                "Real general square"),
                         (Toeplitz(complex(0.9.^(0:ns-1)), complex(0.4.^(0:ns-1))),
                            Toeplitz(complex(0.9.^(0:nl-1)), complex(0.4.^(0:nl-1))),
                                "Complex general square"),
                         (Circulant(0.9.^(0:ns - 1)),
                            Circulant(0.9.^(0:nl - 1)),
                                "Real circulant"),
                         (Circulant(complex(0.9.^(0:ns - 1))),
                            Circulant(complex(0.9.^(0:nl - 1))),
                                "Complex circulant"),
                         (TriangularToeplitz(0.9.^(0:ns - 1), :U),
                            TriangularToeplitz(0.9.^(0:nl - 1), :U),
                                "Real upper triangular"),
                         (TriangularToeplitz(complex(0.9.^(0:ns - 1)), :U),
                            TriangularToeplitz(complex(0.9.^(0:nl - 1)), :U),
                                "Complex upper triangular"),
                         (TriangularToeplitz(0.9.^(0:ns - 1), :L),
                            TriangularToeplitz(0.9.^(0:nl - 1), :L),
                                "Real lower triangular"),
                         (TriangularToeplitz(complex(0.9.^(0:ns - 1)), :L),
                            TriangularToeplitz(complex(0.9.^(0:nl - 1)), :L),
                                "Complex lower triangular"))

        @test As * xs ≈ Matrix(As) * xs
        @test Al * xl ≈ Matrix(Al) * xl
        @test [As[n] for n in 1:length(As)] == vec(As)
        @test [Al[n] for n in 1:length(Al)] == vec(Al)
        @test ldiv!(As, LinearAlgebra.copy_oftype(xs, eltype(As))) ≈ Matrix(As) \ xs
        @test ldiv!(Al, LinearAlgebra.copy_oftype(xl, eltype(Al))) ≈ Matrix(Al) \ xl
    end
)

@testset "Real general rectangular" begin
    Ar1 = Toeplitz(0.9.^(0:nl-1), 0.4.^(0:ns-1))
    Ar2 = Toeplitz(0.9.^(0:ns-1), 0.4.^(0:nl-1))
    @test Ar1 * xs ≈ Matrix(Ar1) * xs
    @test Ar2 * xl ≈ Matrix(Ar2) * xl
end

@testset "Complex general rectangular" begin
    Ar1 = Toeplitz(complex(0.9.^(0:nl-1)), complex(0.4.^(0:ns-1)))
    Ar2 = Toeplitz(complex(0.9.^(0:ns-1)), complex(0.4.^(0:nl-1)))
    @test Ar1 * xs ≈ Matrix(Ar1) * xs
    @test Ar2 * xl ≈ Matrix(Ar2) * xl
end

@testset "Symmetric Toeplitz" begin
    As = SymmetricToeplitz(0.9.^(0:ns-1))
    Ab = SymmetricToeplitz(abs.(randn(ns)))
    Al = SymmetricToeplitz(0.9.^(0:nl-1))
    @test As * xs ≈ Matrix(As) * xs
    @test Ab * xs ≈ Matrix(Ab) * xs
    @test Al * xl ≈ Matrix(Al) * xl
    @test ldiv!(As, copy(xs)) ≈ Matrix(As) \ xs
    @test ldiv!(Ab, copy(xs)) ≈ Matrix(Ab) \ xs
    @test ldiv!(Al, copy(xl)) ≈ Matrix(Al) \ xl
    @test StatsBase.levinson(As, xs) ≈ Matrix(As) \ xs
    @test StatsBase.levinson(Ab, xs) ≈ Matrix(Ab) \ xs
end

@testset "Hankel" begin
    @testset "Real square" begin
        H = Hankel([1.0,2,3,4,5],[5.0,6,7,8,9])
        @test Matrix(H) == [1 2 3 4 5;
                            2 3 4 5 6;
                            3 4 5 6 7;
                            4 5 6 7 8;
                            5 6 7 8 9]

        @test convert(Hankel{Float64}, H) == H
        @test convert(AbstractMatrix{Float64}, H) == H
        @test convert(AbstractArray{Float64}, H) == H

        @test H[2,2] == 3
        @test H[7]  == 3
        @test diag(H) == [1,3,5,7,9]

        x = ones(5)
        @test Matrix(H)*x ≈ H*x

        Hs = Hankel(0.9.^(ns-1:-1:0), 0.4.^(0:ns-1))
        Hl = Hankel(0.9.^(nl-1:-1:0), 0.4.^(0:nl-1))
        @test Hs * xs[:,1] ≈ Matrix(Hs) * xs[:,1]
        @test Hs * xs ≈ Matrix(Hs) * xs
        @test Hl * xl ≈ Matrix(Hl) * xl
    end

    @testset "Complex square" begin
        H = Hankel(complex([1.0,2,3,4,5]), complex([5.0,6,7,8,0]))
        x = ones(5)
        @test Matrix(H)*x ≈ H*x

        Hs = Hankel(complex(0.9.^(ns-1:-1:0)), complex(0.4.^(0:ns-1)))
        Hl = Hankel(complex(0.9.^(nl-1:-1:0)), complex(0.4.^(0:nl-1)))
        @test Hs * xs[:,1] ≈ Matrix(Hs) * xs[:,1]
        @test Hs * xs ≈ Matrix(Hs) * xs
        @test Hl * xl ≈ Matrix(Hl) * xl
    end

    @testset "Real rectangular" begin
        Hs = Hankel(0.9.^(ns-1:-1:0), 0.4.^(0:nl-1))
        Hl = Hankel(0.9.^(nl-1:-1:0), 0.4.^(0:ns-1))
        @test Hs * xl[:,1] ≈ Matrix(Hs) * xl[:,1]
        @test Hs * xl ≈ Matrix(Hs) * xl
        @test Hl * xs ≈ Matrix(Hl) * xs
    end

    @testset "Complex rectangular" begin
        Hs = Hankel(complex(0.9.^(ns-1:-1:0)), complex(0.4.^(0:nl-1)))
        Hl = Hankel(complex(0.9.^(nl-1:-1:0)), complex(0.4.^(0:ns-1)))
        @test Hs * xl[:,1] ≈ Matrix(Hs) * xl[:,1]
        @test Hs * xl ≈ Matrix(Hs) * xl
        @test Hl * xs ≈ Matrix(Hl) * xs
    end

    @testset "Convert" begin
        H = Hankel([1.0,2,3,4,5],[5.0,6,7,8,0])
        @test Hankel(H) == Hankel{Float64}(H) == H
        @test convert(Hankel,H) == convert(Hankel{Float64},H) ==
                convert(AbstractArray,H) == convert(AbstractArray{Float64},H) == H

        A = [1.0 2; 3 4]
        @test Hankel(A) == [1 3; 3 4]
        T = Toeplitz([1.0,2,3,4,5],[1.0,6,7,8,0])
        @test Hankel(T) == Hankel([1.0,2,3,4,5],[5.0,4,3,2,1])
        @test Hankel(T) ≠ ToeplitzMatrices._Hankel(T)
    end
end

@testset "Convert" begin
    T = Toeplitz(ones(2),ones(2))

    @test isa(convert(Matrix{ComplexF64},T),Matrix{ComplexF64})
    @test isa(convert(AbstractMatrix{ComplexF64},T),Toeplitz{ComplexF64})
    @test isa(convert(AbstractArray{ComplexF64},T),Toeplitz{ComplexF64})
    @test isa(convert(ToeplitzMatrices.AbstractToeplitz{ComplexF64},T),Toeplitz{ComplexF64})
    @test isa(convert(ToeplitzMatrices.Toeplitz{ComplexF64},T),Toeplitz{ComplexF64})

    T = SymmetricToeplitz(ones(2))

    @test isa(convert(Matrix{Float32},T),Matrix{Float32})
    @test isa(convert(AbstractMatrix{Float32},T),SymmetricToeplitz{Float32})
    @test isa(convert(AbstractArray{Float32},T),SymmetricToeplitz{Float32})
    @test isa(convert(ToeplitzMatrices.AbstractToeplitz{Float32},T),SymmetricToeplitz{Float32})
    @test isa(convert(ToeplitzMatrices.SymmetricToeplitz{Float32},T),SymmetricToeplitz{Float32})

    T = Circulant(ones(2))

    @test isa(convert(Matrix{ComplexF64},T),Matrix{ComplexF64})
    @test isa(convert(AbstractMatrix{ComplexF64},T),Circulant{ComplexF64})
    @test isa(convert(AbstractArray{ComplexF64},T),Circulant{ComplexF64})
    @test isa(convert(ToeplitzMatrices.AbstractToeplitz{ComplexF64},T),Circulant{ComplexF64})
    @test isa(convert(ToeplitzMatrices.Circulant{ComplexF64},T),Circulant{ComplexF64})

    T = TriangularToeplitz(ones(2),:U)

    @test isa(convert(Matrix{ComplexF64},T),Matrix{ComplexF64})
    @test isa(convert(AbstractMatrix{ComplexF64},T),TriangularToeplitz{ComplexF64})
    @test isa(convert(AbstractArray{ComplexF64},T),TriangularToeplitz{ComplexF64})
    @test isa(convert(ToeplitzMatrices.AbstractToeplitz{ComplexF64},T),TriangularToeplitz{ComplexF64})
    @test isa(convert(ToeplitzMatrices.TriangularToeplitz{ComplexF64},T),TriangularToeplitz{ComplexF64})

    T = Hankel(ones(2),ones(2))

    @test isa(convert(Matrix{ComplexF64},T),Matrix{ComplexF64})
    @test isa(convert(AbstractMatrix{ComplexF64},T),Hankel{ComplexF64})
    @test isa(convert(AbstractArray{ComplexF64},T),Hankel{ComplexF64})
    @test isa(convert(ToeplitzMatrices.Hankel{ComplexF64},T),Hankel{ComplexF64})

    @test Circulant(1:5) == Circulant(Vector(1.0:5))

end

@testset "Constructors" begin
    A = ones(10, 10)
    @test Matrix(Toeplitz(A)) == Matrix(Toeplitz{Float64}(A)) == A
    @test Matrix(SymmetricToeplitz(A)) == Matrix(SymmetricToeplitz{Float64}(A)) == A
    @test Matrix(Circulant(A)) == Matrix(Circulant{Float64}(A)) == A
    @test Matrix(Hankel(A)) == Matrix(Hankel{Float64}(A)) == A

    A = [1.0 2.0;
         3.0 4.0]

    @test Toeplitz(A) == Toeplitz([1.,3.], [1.,2.])
    @test Toeplitz{Float64}(A) == Toeplitz([1.,3.], [1.,2.])
    @test Matrix(SymmetricToeplitz(A)) == Matrix(SymmetricToeplitz{Float64}(A)) ==
                Matrix(Toeplitz(Symmetric(A))) == Matrix(Symmetric(Toeplitz(A))) == [1. 2.; 2. 1.]
    @test Matrix(Circulant(A)) == [1 3; 3 1]

    @test TriangularToeplitz(A, :U) == TriangularToeplitz{Float64}(A, :U) == Toeplitz(UpperTriangular(A)) == UpperTriangular(Toeplitz(A))
    @test TriangularToeplitz(A, :L) == TriangularToeplitz{Float64}(A, :L) == Toeplitz(LowerTriangular(A)) == LowerTriangular(Toeplitz(A))

    @test Matrix(Hankel(A)) == Matrix(Hankel{Float64}(A)) == [1.0 3; 3 4]

    # Constructors should be projections
    @test Toeplitz(Toeplitz(A)) == Toeplitz(A)
    @test SymmetricToeplitz(SymmetricToeplitz(A)) == SymmetricToeplitz(A)
    @test Circulant(Circulant(A)) == Circulant(A)
    @test TriangularToeplitz(TriangularToeplitz(A, :U), :U) == TriangularToeplitz(A, :U)
    @test TriangularToeplitz(TriangularToeplitz(A, :L), :L) == TriangularToeplitz(A, :L)
    @test Hankel(Hankel(A)) == Hankel(A)
end

@testset "Circulant mathematics" begin
    C1 = Circulant(rand(5))
    C2 = Circulant(rand(5))
    C3 = Circulant{ComplexF64}(rand(5))
    C4 = Circulant(ones(5))
    C5 = Circulant{ComplexF64}(ones(5))
    M1 = Matrix(C1)
    M2 = Matrix(C2)
    M3 = Matrix(C3)
    M4 = Matrix(C4)
    M5 = Matrix(C5)

    C = C1*C2
    @test C isa Circulant
    @test C ≈ M1*M2

    C = C1-C2
    @test C isa Circulant
    @test C ≈ M1-M2

    C = C1+C2
    @test C isa Circulant
    @test C ≈ M1+M2

    C = 2C1
    @test C isa Circulant
    @test C ≈ 2M1

    C = C1*2
    @test C isa Circulant
    @test C ≈ M1*2

    C = -C1
    @test C isa Circulant
    @test C ≈ -M1

    C = inv(C1)
    @test C isa Circulant
    @test C ≈ inv(M1)
    C = inv(C3)
    @test C isa Circulant
    @test C ≈ inv(M3)

    C = pinv(C1)
    @test C isa Circulant
    @test C ≈ pinv(M1)
    C = pinv(C3)
    @test C isa Circulant
    @test C ≈ pinv(M3)
    C = pinv(C4)
    @test C isa Circulant
    @test C ≈ pinv(M4)
    C = pinv(C5)
    @test C isa Circulant
    @test C ≈ pinv(M5)

    C = sqrt(C1)
    @test C isa Circulant
    @test C*C ≈ C1
    C = sqrt(C3)
    @test C isa Circulant
    @test C*C ≈ C3

    C = copy(C1)
    @test C isa Circulant
    C2 = similar(C1)
    copyto!(C2, C1)
    @test C1 ≈ C2

    v1 = eigvals(C1)
    v2 = eigvals(M1)
    for v1i in v1
        @test minimum(abs.(v1i .- v2)) < sqrt(eps(Float64))
    end
end

@testset "Cholesky" begin
    T = SymmetricToeplitz(exp.(-0.5 .* range(0, stop=5, length=100)))
    @test cholesky(T).U ≈ cholesky(Matrix(T)).U
    @test cholesky(T).L ≈ cholesky(Matrix(T)).L
end


using Combinatorics
using MLStyle

function parse_line(l)
    @match match(r"(\w+) would (lose|gain) (\d+) .* (\w+)\.", l).captures begin
        [u, "lose", n, v] => (u, v, -parse(Int, n))
        [u, "gain", n, v] => (u, v,  parse(Int, n))
        _ => nothing
    end
end

const data = parse_line.(readlines())
const g = begin
    gg = Dict{String, Dict{String, Int}}()
    for (u, v, n) in data
        get!(gg, u, Dict())[v] = n
    end
    gg
end

function happiness(names)
    n = length(names)
    sum(1:n) do i
        get(get(g, names[i], Dict()), names[mod1(i+1, n)], 0) + get(get(g, names[i], Dict()), names[mod1(i-1, n)], 0)
    end
end

# Part 1
g |> keys |> collect |> permutations .|> happiness |> maximum |> println

# Part 2
[collect(keys(g)); "Me"] |> permutations .|> happiness |> maximum |> println


export
    orthocomp,
    pdet_maker,
    bin2

function orthocomp(m)
    U, S, V = PolyLib.smith_normal_form(m)
    return lll(PolyLib.inverse(V)[:, size(m)[1]+1:end])[1]
end

nonneg(v) = all(v.>=0) ? true : false

function unequal_sample_maker(p)
    perm = sortperm(p, rev=true)
    p0 = p[perm]
    for i=1:length(p0)-1
        p0[i+1] += p0[i]
    end

    function sample()
        r = rand()
        for i=1:length(p0)
            if r < p0[i]
                return perm[i]
            end
        end
    end
end

function pdet_maker(eta, pdc)
    a = 1 - pdc
    b = 1 - eta.*a
    function pdet0(i)
        return a .* b.^i
    end
    return pdet0
end

function bin2(arr)
    n1 = length(arr)
    n2 = div(n1, 2)
    arr2 = zeros(n2)
    for i = 1:n2
        arr2[i] = (arr[2i-1] + arr[2i]) / 2
    end
    return arr2
end


# Example 2

using Base.Iterators
using Arbitrary

# Define your own type
struct Point{T}
    x::T
    y::T
end

# Define a method for Arbitrary.arbitrary
function Arbitrary.arbitrary(::Type{Point{T}}, ast::ArbState) where {T}
    xs = Iterators.Stateful(arbitrary(T, ast))
    flatten([Point{T}[Point(T(0), T(0)),
                      Point(T(0), T(1)),
                      Point(T(1), T(0)),
                      Point(T(-1), T(-1))],
             Generate{Point{T}}(
                 () -> Point(popfirst!(xs), popfirst!(xs)))])
end

collect(take(arbitrary(Point{Int}), 20))


using Luxor
using OffsetArrays

function make_drawing(width, height, img_path, bkg_color, origin_p)
    d = Drawing(width, height, img_path)
    background(bkg_color)
    origin(origin_p)
    return d
end

width = 500
height = 500
path = "voronoi.png"
color = "black"

my_draw = make_drawing(width, height, path, color, O)

slope(x₁, y₁, x₂, y₂) = (y₂ - y₁) / (x₂ - x₁)
slope(P1, P2) = (P2.y - P1.y) / (P2.x - P1.x)
midpoint(x₁, y₁, x₂, y₂) = ((x₁ + x₂) / 2, (y₁ + y₂) / 2)
midpoint(P1, P2) = ((P1.x + P2.x) / 2, (P2.y + P2.y) / 2)
point_slope(x, x₁, y₁, m) = (x, m * (x - x₁) + y₁)

mat = zeros(width, height)

p1 = (x = 125, y = 125)
p2 = (x = 375, y = 375)

sp = slope(p1, p2)
mp = midpoint(p1, p2)

d(x1, y1, x2, y2, x, y) = (x - x1) * (y2 - y1) - (y - y1) * (x2 - x1)

h1_plane = []
h2_plane = []
h3_plane = []

for x in 1:500, y in 1:500
	if d(500, 0, mp[1], mp[2], x, y) < 0
		push!(h1_plane, Point(x, y))
	elseif d(500, 0, mp[1], mp[2], x, y) > 0
		push!(h2_plane, Point(x, y))
	else
		push!(h3_plane, Point(x, y))
	end
end

# for p in vals
	# mat[Int(p.x), Int(p.y)] = 1
# end

# function half_plane(p1, p2, width, height)

    # mp = midpoint(p1.x, p1.y, p2.x, p2.y) |> Point
    # sp = slope(p1.x, p1.y, p2.x, p2.y)
    # first_point = point_slope(-250, mp.x, mp.y, sp == 0 ? 0 : -1 / sp) |> Point
    # second_point = point_slope(250, mp.x, mp.y, sp == 0 ? 0 : -1 / sp) |> Point

    # half_1 = [first_point, second_point, Point(width / 2, height / 2)]
    # half_2 = [first_point, second_point, Point(-width / 2, height / 2)]

    # return half_1, half_2
# end

setcolor(1, 0, 0, 1)
for c in h1_plane
	circle(c, 1; action = :stroke)
end

setcolor(0, 0, 1, 1)
for c in h2_plane
	circle(c, 1; action = :stroke)
end

setcolor(1, 1, 0, 1)
for c in h3_plane
	circle(c, 1; action = :stroke)
end

setcolor(1, 0, 0, 1)
circle(p1.x, p1.y, 15, action = :fill)

setcolor(0, 0, 1, 1)
circle(p2.x, p2.y, 15, action = :fill)

finish()


### A Pluto.jl notebook ###
# v0.17.4

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ bce7c57b-9250-496a-a943-275301ea8c6c
begin
    import Pkg
    Pkg.add(url="https://github.com/Pocket-titan/DarkMode")
    import DarkMode
    DarkMode.enable()
    # OR DarkMode.Toolbox(theme="default")
end

# ╔═╡ dacc8004-b9f0-11eb-04d8-a391fd316be6
using PlutoUI, WAV, UnicodePlots

# ╔═╡ afdee77e-ccdb-42ad-9402-e0dd214fc7aa
md"![](https://i.imgur.com/k2qKXDd.png)"

# ╔═╡ db3d1596-f84c-4e0d-a3cf-23b6cdd7b8b4
@bind A Slider(.1:.1:1,default=.8,show_value=true)

# ╔═╡ 671513e6-1d23-49e9-802c-b50d079c0b94
@bind phi Slider(0:pi/6:2pi,default=pi/2,show_value=true)

# ╔═╡ 828db6b2-6c0f-4a81-a804-0a1ae2ed1dc0
@bind f0 Slider(0:1000:10000,default=1000,show_value=true)

# ╔═╡ 01cb6294-a0d3-4e10-a919-7da531c36997
@bind fs Slider(0:1000:50000,default=44100,show_value=true)

# ╔═╡ b3b71013-5de2-4f7c-99a7-98f92de69ec7
begin

t = -0.002:1/fs:0.002
y = A*cos.(2*pi*f0*t .+ phi)

#y1 = squarewave.(f1*t1) * 1
#y1 = (sin.(2pi * f1* t1) .*0.1)
#y1 = trianglewave.(f1* t1) * 1
# y1 = (sin.(2pi * f1* t1) .*t1) .+ (sin.(2pi * f1*2* t1 .+pi) .*t1/2)
wavplay(y, fs)
end

# ╔═╡ 956dbaf0-f6f7-495f-acdd-03792be4eade
lineplot(t, y)

# ╔═╡ 13b2dac6-6f57-4922-9931-24ff501f7e9b
md"""
## Complex Numbers
```
(a + ib)  a,b: real numbers; i: imaginary unit
```
- Rectangular form: a+ib
- Polar form: A=sqroot(a^2+b^2);phi=a*tan(2(b/a))
"""

# ╔═╡ 5db2175f-d377-49b5-bfb3-f9403857c48b
md"""
## Euler's Formula

Fundamental to understand discrete fourier transform

$$e^{i\phi} = cos(\phi) + isin(\phi)$$

$$cos(\phi) = \frac{e^{i\phi}+ e^{-i\phi}}{2}$$
$$sin(\phi) = \frac{e^{i\phi}- e^{-i\phi}}{2i}$$
"""

# ╔═╡ d4e6f5af-2f22-4501-adf9-fadb96d9c436
md"""
![](https://i.imgur.com/HkbIUh8.png)

We plot the complex wave  in a 2d graph by plotting real and imaginary parts separately
"""

# ╔═╡ f705e1df-9ef3-4eae-9edb-1b0f932569a2
md"""
## Scalar(Dot Product of sequences)
$$x=[1,2,3] , y = [1,2,3]$$
$$dot = 1*1 + 2*2 + 3*3 = 14$$

Scalar product of two orthogonal sequences is 0.

$$x = [2,2], y = [2,-2]$$
$$dot = 2*2 + 2*-2 = 0$$

"""

# ╔═╡ 178a84f2-54e4-42f2-8ee2-165db0dc3618
md"""
![](https://i.imgur.com/iixirS3.png)
"""

# ╔═╡ 23b9bb3e-74f4-43f6-b3dc-8e2edcc5700d
@bind A1 Slider(.1:.1:1,default=.8,show_value=true)

# ╔═╡ 47ce6005-a5c9-4684-8632-33e675433bd9
@bind A2 Slider(.1:.1:1,default=.8,show_value=true)

# ╔═╡ 17258789-78dd-446e-bc57-36faf3fe664b
@bind phi1 Slider(0:pi/6:2pi,default=pi/2,show_value=true)

# ╔═╡ 0b7b17a5-d2c5-4af5-b1c9-376fdaffdd58
@bind phi2 Slider(0:pi/6:2pi,default=pi/2,show_value=true)

# ╔═╡ 51f34b02-d31f-435d-b816-7369da8764ce
@bind f01 Slider(0:1000:10000,default=1000,show_value=true)

# ╔═╡ c8a44a7c-ea12-4d7c-a2ab-c2b792f12a0e
@bind f02 Slider(0:1000:10000,default=1000,show_value=true)

# ╔═╡ 234f7169-6f24-4508-a35a-11664de63f48
@bind fst Slider(0:1000:50000,default=44100,show_value=true)

# ╔═╡ 6777d7e3-de89-45a8-a07f-bb22428b0418
begin
#tr = -1:1/fst:1	 #use this if want to hear something and comment next line
tr = -0.002:1/fs:0.002
y1 = A1*sin.(2*pi*f01*tr .+ phi1)
y2 = A2*cos.(2*pi*f02*tr .+ phi2)
addresult =y1 .+ y2
subresult =y1 .- y2
mulresult =y1 .* y2
divresult =y1 ./ y2
wavplay(subresult, fst)
end

# ╔═╡ 1a532407-0157-4ced-b657-7b0b3f664142
lineplot(tr, y1)

# ╔═╡ 8c0abdfe-ad72-4307-97c7-3358a9889a7e
lineplot(tr, y2)

# ╔═╡ 96048a2b-31dc-4377-bcd5-64528a5127d1
lineplot(tr, addresult)

# ╔═╡ 82b9d306-07a7-41d3-a0a6-dfd1ee3153fc
lineplot(tr, subresult)

# ╔═╡ c2b815e6-062d-4225-8ed6-221b16fbec7f
lineplot(tr, mulresult)

# ╔═╡ 53989efc-c028-4f84-abbd-fffc7727f2da
lineplot(tr, divresult)

# ╔═╡ a01ff431-a00f-4a36-9aed-520209e53e8d
md"""
Important Resources:

- https://ccrma.stanford.edu/~jos/mdft/
- https://github.com/MTG/sms-tools

"""

# ╔═╡ Cell order:
# ╠═bce7c57b-9250-496a-a943-275301ea8c6c
# ╠═dacc8004-b9f0-11eb-04d8-a391fd316be6
# ╠═afdee77e-ccdb-42ad-9402-e0dd214fc7aa
# ╠═db3d1596-f84c-4e0d-a3cf-23b6cdd7b8b4
# ╠═671513e6-1d23-49e9-802c-b50d079c0b94
# ╠═828db6b2-6c0f-4a81-a804-0a1ae2ed1dc0
# ╠═01cb6294-a0d3-4e10-a919-7da531c36997
# ╠═b3b71013-5de2-4f7c-99a7-98f92de69ec7
# ╠═956dbaf0-f6f7-495f-acdd-03792be4eade
# ╠═13b2dac6-6f57-4922-9931-24ff501f7e9b
# ╠═5db2175f-d377-49b5-bfb3-f9403857c48b
# ╠═d4e6f5af-2f22-4501-adf9-fadb96d9c436
# ╠═f705e1df-9ef3-4eae-9edb-1b0f932569a2
# ╠═178a84f2-54e4-42f2-8ee2-165db0dc3618
# ╠═23b9bb3e-74f4-43f6-b3dc-8e2edcc5700d
# ╠═47ce6005-a5c9-4684-8632-33e675433bd9
# ╠═17258789-78dd-446e-bc57-36faf3fe664b
# ╠═0b7b17a5-d2c5-4af5-b1c9-376fdaffdd58
# ╠═51f34b02-d31f-435d-b816-7369da8764ce
# ╠═c8a44a7c-ea12-4d7c-a2ab-c2b792f12a0e
# ╠═234f7169-6f24-4508-a35a-11664de63f48
# ╠═6777d7e3-de89-45a8-a07f-bb22428b0418
# ╠═1a532407-0157-4ced-b657-7b0b3f664142
# ╠═8c0abdfe-ad72-4307-97c7-3358a9889a7e
# ╠═96048a2b-31dc-4377-bcd5-64528a5127d1
# ╠═82b9d306-07a7-41d3-a0a6-dfd1ee3153fc
# ╠═c2b815e6-062d-4225-8ed6-221b16fbec7f
# ╠═53989efc-c028-4f84-abbd-fffc7727f2da
# ╟─a01ff431-a00f-4a36-9aed-520209e53e8d


module Representation
    export Organism, Population
    using Random, Distributions

    # Struct that represents a single organism
    mutable struct Organism
        gauss_kernels::Array{Tuple{Float64, Float64, Float64}, 1}
        σ::Array{Float64, 1}
        α::Array{Float64, 1}
        fitness::Union{Float64, Missing}
        hit_ratio::Union{Int64, Missing}
    end
    # Struct that represents a population as a list of organisms
    mutable struct Population
        members::Array{Organism, 1}
    end

    # Generate N organisms, therefore a population
    function generate_population(population_size, n_kernels, initial_σ)
        organism_list = Organism[]
        for i in range(1, population_size)
            kernels = Tuple{Float64, Float64, Float64}[]
            for k in range(1, n_kernels)
                weight, c, γ = generate_kernel_params(initial_σ)
                push!(kernels, tuple(weight, c, γ))
            end
            σ = rand(Normal(0.0, 1.0), 1)[1]
            this_organism = Organism(kernels,
                                    repeat([σ], (n_kernels * length(kernels[1]))),
                                    repeat([0.0], (n_kernels * length(kernels[1]))),
                                    missing,
                                    missing)
            push!(organism_list, this_organism)
        end
        #c = rand(range(0.817, step=0.001, length=floor(Int, (1.0-0.817)/0.001) + 1))
        population = Population(organism_list)
        return population
    end

    # Function to generate a single kernel
    function generate_kernel_params(σ)
        normal_d = Normal(0.0, σ)
        weight = rand(normal_d, 1)[1]
        c = rand(normal_d, 1)[1]
        γ = rand(normal_d, 1)[1]
        return weight, c, γ
    end

end  # module Representation


using LoopVectorization: @avx

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

"""
$(TYPEDEF)

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

## Fields
$(TYPEDFIELDS)

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

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

"""
$(TYPEDSIGNATURES)

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

"""
$(TYPEDSIGNATURES)

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

"""
$(TYPEDSIGNATURES)

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

    return dim
end

"""
$(TYPEDSIGNATURES)

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

"""
$(TYPEDSIGNATURES)

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

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

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

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

"""
$(TYPEDSIGNATURES)

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

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

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

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

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

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

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

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

# Versions with allocations
"""
$(TYPEDSIGNATURES)

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

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

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

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

"""
$(TYPEDSIGNATURES)

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

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

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

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

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

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

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

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

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

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

# In-place versions of grad_xk_basis!

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

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

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

# Versions with allocations
"""
$(TYPEDSIGNATURES)

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

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

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

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

 """
 $(TYPEDSIGNATURES)

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

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

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

# Version with allocations
"""
$(TYPEDSIGNATURES)

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

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

"""
$(TYPEDSIGNATURES)

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

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

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

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

# Version with allocations
"""
$(TYPEDSIGNATURES)

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

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

"""
$(TYPEDSIGNATURES)

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

"""
$(TYPEDSIGNATURES)

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

"""
$(TYPEDSIGNATURES)

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

"""
$(TYPEDSIGNATURES)

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

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

"""
$(TYPEDSIGNATURES)

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

"""
$(TYPEDSIGNATURES)

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

"""
$(TYPEDSIGNATURES)

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

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

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

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

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

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

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

"""
$(TYPEDSIGNATURES)

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

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

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

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

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

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

"""
$(TYPEDSIGNATURES)

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

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

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

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

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

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

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

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

    return dxdxkψ
end

"""
$(TYPEDSIGNATURES)

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

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

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

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

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

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

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

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

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

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

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

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

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

    d2xdxkψ_basis = zeros(Ne, Nψ)

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

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

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

    return d2xdxkψ
end

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

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

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

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

    d2xdxkψ_basis = zeros(Ne, Nψ)

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

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

    return d2xdxkψ
end

"""
$(TYPEDSIGNATURES)

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

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

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

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

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

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

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

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

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

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

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

"""
$(TYPEDSIGNATURES)

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

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

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

"""
$(TYPEDSIGNATURES)

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


"""

#### References

An incremental extreme learning machine for online sequential learning problems.

Lu Guo, Jing-hua Hao, and Min Liu.

Neurocomputing, 2014 vol. 128 pp. 50-58.

http://linkinghub.elsevier.com/retrieve/pii/S0925231213010059

"""
mutable struct LSIELM{TN,TV} <: AbstractSLFN where TN <: AbstractNodeInput where TV AbstractArray{Float64}
    p::Int  # Number of training points
    q::Int  # Dimensionality of function domain
    s::Int  # number of neurons
    neuron_type::TN
    c::Float64  # regularization parameter
    μx::Vector{Float64}
    σx::Vector{Float64}
    Wt::Matrix{Float64}  # transpose of W matrix
    d::Vector{Float64}

    M::Matrix{Float64}  # internal
    p_tot::Int
    v::TV

    function LSIELM(p::Int, q::Int, s::Int, neuron_type::TN, c::Float64,
                    μx, σx) where TN <: AbstractNodeInput
        # NOTE: Wt and d are constant throughout all learning chunks
        Wt = 2*rand(q, s) .- 1
        d = rand(s)
        M = zeros(Float64, s, s)
        new{TN,typeof(d)}(p, q, s, neuron_type, c, μx, σx, Wt, d, M, 0)
    end
end

function LSIELM(x::AbstractArray, y::AbstractArray;
                c::Float64=1/(50*eps()),
                neuron_type::AbstractNodeInput=Linear(Tanh()),
                s::Int=size(x, 1))
    q = size(x, 2)  # dimensionality of function domain
    p = size(x, 1)  # number of training points
    s = min(p, s)   # Can't have more neurons than training points
    xn, μx, σx = standardize(x[:, :])
    out = LSIELM(p, q, s, neuron_type, c, μx, σx)
    fit!(out, xn, y)
end

## API methods
function fit!(elm::LSIELM, x::AbstractArray, y::AbstractArray)
    # check if we are boosting or not
    S = hidden_out(elm, x)
    if elm.p_tot == 0  # this is the initialization phase
        # elm.v = S \ y
        elm.M = inv(1/elm.c .+ S'S)
        elm.v = elm.M*S'y
    else
        K = I - elm.M * S' * inv(S*elm.M*S' + I)*S
        elm.v = K*(elm.v + elm.M*S'y)
        elm.M = K*elm.M
    end

    elm.p_tot += size(x, 1)
    elm

end

function Base.show(io::IO, elm::LSIELM{TA}) where TA
    s =
    """
    LSIELM with
      - $(TA) Activation function
      - $(elm.q) input dimension(s)
      - $(elm.s) neuron(s)
      - $(elm.p) initial training point(s)
      - $(elm.p_tot) total training points
    """
    print(io, s)
end

function test_me(::Type{LSIELM})
    x = range(0, stop=1, length=20)
    y = sin.(6*x[:])
    oselm = LSIELM(x, y)

    xt = range(0, stop=1, length=300)
    yt = sin.(6*xt)

    @show maximum(abs.(oselm(xt) - yt))

    for i in 1:50
        x2 = rand(15)
        y2 = sin.(6*x2)
        fit!(oselm, x2, y2)
        @printf "%i\t%2.4e\n" i maximum(abs.(oselm(xt) - yt))
    end
end


using Test
using Random
using StatsBase
using Printf
using JITrench

include("test_utils.jl")

operators = [+, -, /, *]

N_TEST_COUNT = 100

@testset "OperatorTest" begin
    @testset "ForwardTest" begin
        for i in 1:N_TEST_COUNT
            for op in operators
                x1 = rand()
                x2 = rand()
                @test isAbout(op(Variable(x1), Variable(x2)).values, op(x1, x2), )
            end
        end
    end
    @testset "backwardTest" begin
        for i in 1:N_TEST_COUNT
            for op in operators[1:end-1]
                x1 = rand()
                x2 = rand()
                @test isAbout(backward_diff(op, [x1, x2]), numerical_diff(op, [x1, x2]))
            end
        end
    end
end

funcitons = [
    sin,
    cos,
    tan,
    log,
    exp
]

@testset "MathFuncTest" begin
    @testset "ForwardTest" begin
        for i in 1:N_TEST_COUNT
            for op in funcitons
                x = rand()
                @test isAbout(op(x), op(Variable(x).values))
            end
        end
    end
    @testset "backwardTest" begin
        for i in 1:N_TEST_COUNT
            for op in funcitons
                x = rand()
                @test isAbout(backward_diff(op, [x]), numerical_diff(op, [x]))
            end
        end
    end
end

activations = Dict(
    JITrench._sigmoid => JITrench.sigmoid
)

@testset "ActivationTest" begin
    @testset "ForwardTest" begin
        for i in 1:N_TEST_COUNT
            for (func, jt_func) in activations
                x = rand()
                isAbout(func(x), jt_func(Variable(x)).values)
            end
        end
    end
    @testset "backwardTest" begin
        for i in 1:N_TEST_COUNT
            for (func, jt_func) in activations
                x = rand()
                @test isAbout(backward_diff(jt_func, [x]), numerical_diff(func, [x]))
            end
        end
    end
end

@testset "ArrOperateTest" begin
    max_size = 120
    @testset "forward" begin
        @testset "reshape" begin
            for _ in 1:N_TEST_COUNT
                in_shape, out_shape = generate_shape(rand(1:max_size))
                x = rand(in_shape...)
                y_arr = reshape(x, out_shape)
                y_var = reshape(Variable(x), out_shape)
                @test y_arr == y_var.values
            end
        end
        @testset "transpose" begin
            for _ in 1:N_TEST_COUNT
                in_shape, out_shape = generate_shape(rand(1:max_size), max_dim=2)
                x = rand(in_shape...)
                y_arr = transpose(x)
                y_var = transpose(Variable(x))
                @test y_arr == y_var.values
            end
        end
    end
    @testset "backward" begin
        @testset "reshape" begin
            for _ in 1:N_TEST_COUNT
                in_shape, out_shape = generate_shape(rand(1:max_size))
                x = Variable(rand(in_shape...))
                y = reshape(x, out_shape)
                backward!(y, retain_grad=true)
                @test y.grad.values == ones(size(y.values))
                @test x.grad.values == ones(size(x.values))
            end
        end
        @testset "transpose" begin
            for _ in 1:N_TEST_COUNT
                in_shape, out_shape = generate_shape(rand(1:max_size), min_dim=2, max_dim=2)
                x = Variable(rand(in_shape...))
                y = JITrench.transpose(x)
                backward!(y, retain_grad=true)
                @test y.grad.values == ones(size(y.values))
                @test x.grad.values == ones(size(x.values))
            end
        end
    end
end


module AtomicOpt

using LinearAlgebra
using LinearMaps
using Printf
using Random
using SparseArrays
using Arpack
using IterativeSolvers
using Distributed
using FFTW
using LBFGSB
using BlockArrays

import Base: show, vec, size
import LinearAlgebra: dot, rank

export AbstractAtomicSet, AbstractAtom, AbstractFace
export MappedAtomicSet, MappedAtom, MappedFace
export SumAtomicSet, SumAtom, SumFace
export OneBall, OneBallAtom, OneBallFace
export NucBall, NucBallAtom, NucBallFace
export BlkNucBall, BlkNucBallAtom, BlkNucBallFace
# export PosSimplex, PosSimplexAtom, PosSimplexFace
# export TwoBall, TwoBallAtom, TwoBallFace
export TraceBall, TraceBallAtom, TraceBallFace
export face, face!, gauge, rank, support, expose, expose!
export face_project!
export level_set, conditional_graident, coordinate_descent
export Solution, constructPrimal
export MaskOP, TMaskOP
export dual_obj_gap

"Abstract atomic set."
abstract type AbstractAtomicSet end

"Abstract atoms."
abstract type AbstractAtom end

"Abstract face."
abstract type AbstractFace end

"Abstract operators."
abstract type AbstractOperator end

"Abstract Linear Operator"
LinearOp = Union{LinearMap, AbstractMatrix}
AbstractLinearOp = Union{LinearOp, AbstractOperator}

include("./BasicSets/OneBall.jl")
include("./BasicSets/NucBall.jl")
include("./BasicSets/BlkNucBall.jl")
include("./BasicSets/TraceBall.jl")
# include("src/BasicSets/PosSimplex.jl")
# include("src/BasicSets/TwoBall.jl")
include("./SetOperations/mapped.jl")
include("./SetOperations/sum.jl")
include("./SetOperations/operators.jl")
include("./SetOperations/facialprojection.jl")
include("./SetOperations/utils.jl")
include("./Optimization/solution.jl")
include("./Optimization/align.jl")
include("./Optimization/levelset.jl")
include("./Optimization/boxls.jl")
include("./Optimization/conditionalGradient.jl")
include("./Optimization/liftedCoordinateDescent.jl")

dot(a1::AbstractAtom, a2::Vector) = dot(vec(a1), a2)
dot(a1::Vector, a2::AbstractAtom) = dot(vec(a2), a1)
dot(a1::AbstractAtom, a2::AbstractAtom) = dot(vec(a1), vec(a2))

function Base.show(io::IO, A::AbstractAtomicSet)
    println(io, "atomic set  : ", atom_name(A))
    println(io, "description : ", atom_description(A))
    print(  io, "parameters  : ", atom_parameters(A))
end

function Base.show(io::IO, F::AbstractFace)
    println(io, "face of atomic set  : ", face_name(F))
    print(  io, "parameters          : ", face_parameters(F))
end

end # module


θ = rand() * -2π + π
a = U1(θ)
b = U1(rand())
c = U1(rand())

## Properties of the Group ##
@test isapprox(arg(a), θ)
@test isapprox(len(a), 1)

## Properties of the Unary Operators ##

@test a == +a
@test a == -(-a)

## Properties of Multiplication ##

@test a * (U1(π / 2) / a) == U1(π / 2) # Identity
@test a * b == b * a # Commutativity
@test (a * b) * c == a * (b * c) # Associativity


using Flux
using Flux:@adjoint

function kan_rv_prod_inner(phi::Array{AbstractFloat,2}, v::NTuple)
	vl = collect(v)
	h = 0.5 .- vl
	hKh =  ( 0.5 * dot(h, phi * h))^(size(phi)[1] / 2)
	(-1.0)^(sum(vl)) * hKh
end
"""
calculates the expectation of products of Gaussian RVs, with
exponents s and covariance matrix phi, using the identity in proposition 1
of Kan 2008.
"""
function kan_rv_prod(phi)
	st = size(phi)[1]
	mapreduce(v -> kan_rv_prod_inner(phi, v), +, Iterators.product(fill(0:1, st)...)) / factorial(st ÷ 2)
end

function kan_rv_prod2(phi)
    kan_rv_prod(phi)
end

function kan_rv_prod2_adj(phi)
    st = size(phi)[1]
    g = v -> gradient(kan_rv_prod_inner, phi, v)[1]
    sum(g, Iterators.product(fill(0:1, st)...)) / factorial(st ÷ 2)
    # gradient(kan_rv_prod, phi)[1]
end

@adjoint function kan_rv_prod2(phi)
    kan_rv_prod2(phi), x -> (x * kan_rv_prod2_adj(phi),)
end

println(kan_rv_prod2_adj(3 * ones(5, 5)))
# gradient(kan_rv_prod_inner, ones(2, 2), (1, 0))
# print(gradient(kan_rv_prod, ones((2, 2)))[1] - gradient(kan_rv_prod2, ones((2, 2)))[1])
"""
@btime VolterraGP.kan_rv_prod(ones((5, 5)))

@btime gradient(kan_rv_prod, ones((2, 2)))

@btime gradient(kan_rv_prod2, ones((2, 2)))

"""
# function sum2(op, arr)
# 	return sum(op, arr)
# end

# function sum2adj(Δ, op, arr)
# 	n = length(arr)
# 	g = x -> Δ * gradient(op, x)[1]
# 	return (nothing, map(g, arr))
# end

# @adjoint function sum2(op, arr)
# 	return sum2(op, arr), Δ -> sum2adj(Δ, op, arr)
# end

# function t1(arr, op)
# 	g = sum2(op, arr)
# 	return g
# end

# function t0(arr, op)
# 	g = sum(op, arr)
# 	return g
# end

# function testPerf()
	# nb = 100000
# 	arr = rand(nb)
# 	op = x -> x + x

# 	res = t1(arr, op)
# 	println(res)

# 	df = x -> gradient(x -> t1(x, op), x)
# 	res1 = df(arr)[1]

# 	df0 = x -> gradient(x -> t0(x, op), x)
# 	res0 = df0(arr)[1]

# 	diff = sum((res0 - res1).^2)
# 	println("diff : ")
# 	println(diff)

# 	println("Benchmarking t1")
# 	@time for i = 1:100
# 		t1(arr, op)
# 	end

# 	println("Benchmarking gradient")
# 	@time for i = 1:100
# 		df(arr)
# 	end

# end

# testPerf()


module SfMatter

using Unitful

import PhysicalConstants.CODATA2018: k_B, g_n, R
import ..SfGeometry: volume

export Material, density, yield_strength, bouyancy, compression_energy

struct Material
	density::Unitful.Density
	yield_strength::Union{Missing, Unitful.Pressure}
end

Material(density::Unitful.Density) = Material(density, missing)

# empty definition for overriding
mass(::Nothing) = nothing

density(x) = mass(x) / volume(x) |> u"kg/m^3"
density(x::Unitful.Density) = x
density(m::Material) = m.density

yield_strength(m::Material) = m.yield_strength

"""
	bouyancy(ρ_body::Unitful.Density, ρ_fluid::Unitful.Density; g = g_n)

Acceleration due to bouyancy effects experienced by an object with density `ρ_body` in a fluid with density `ρ_fluid`, undergoing acceleration `g` (which defaults to Earth gravity).
"""
bouyancy(ρ_body::Unitful.Density, ρ_fluid::Unitful.Density; g = g_n) = g * ρ_fluid / ρ_body

"""
	bouyancy(ρ_body::Unitful.Density, ρ_fluid::Unitful.Density; g = g_n)

Acceleration due to bouyancy effects experienced by an object made of `body` in a fluid composed of `fluid`, undergoing acceleration `g` (which defaults to Earth gravity).
"""
bouyancy(body::Material, fluid::Material; g = g_n) = g * density(fluid) / density(body)

"""
    bouyancy(v_body::Unitful.Volume, ρ_fluid::Unitful.Density; g = g_n)

Bouyancy force felt by an object of volume `v_body` in a fluid of density `ρ_fluid`, undergoing acceleration `g` (which defaults to Earth gravity).
"""
bouyancy(v_body::Unitful.Volume, ρ_fluid::Unitful.Density; g = g_n) = g * ρ_fluid * v_body

"""
    bouyancy(v_body::Unitful.Volume, fluid; g = g_n)

Bouyancy force felt by an object of volume `v_body` in a fluid composed of `fluid`, undergoing acceleration `g` (which defaults to Earth gravity).
"""
bouyancy(v_body::Unitful.Volume, fluid; g = g_n) = g * density(fluid) * v_body

"""
    compression_energy(n, T, v_initial, v_final)

Work done to isothermally compress `n` moles of a substance from a volume of `v_initial` to `v_final`, maintaining a temperature of `T`.

If `v_final` is larger than `v_initial`, the result will be negative.

Ref: https://physics.stackexchange.com/a/37676/
"""
compression_energy(n, T, v_initial, v_final) = n * R * T * log(v_initial/ v_final)

end


using Pkg
Pkg.activate(joinpath(@__DIR__,".."))
using LinearAlgebra
using BenchmarkTools, Test
using Cthulhu
using Plots ; theme(:dark)
using SmoothedSpectralAbscissa ; const SSA = SmoothedSpectralAbscissa

function rand_gauss(n::Integer)
  return randn(n,n) ./ sqrt(n)
end
function rand_nonnormal(n::Integer,(ud::Real)=1.01)
  mat = rand_gauss(n)
  sh = schur(mat)
  upd = diagm(0=>fill(1.0,n),1=>fill(ud,n-1))
  return sh.vectors*upd*sh.Schur*inv(upd)*sh.vectors'
end

##

sizes = [25,50,75,100,125,150]

ssa_times = similar(sizes,Float64)
ssa_allocs = similar(ssa_times)

for (k,n) in enumerate(sizes)
  alloc = SSA.SSAAlloc(n)
  A=rand_nonnormal(n,1.0)
  be = @benchmark SSA.ssa!($A,nothing,$alloc)
  ssa_times[k] = median(be.times)
  ssa_allocs[k] = be.allocs
end

##
bar(sizes,ssa_times;leg=false,yscale=:log10)

@show ssa_times ./ ssa_times_old