include("../src/interface.jl")
using LinearAlgebra
using Plots
using Plots.PlotMeasures

function fctmain(n_tau, prec)
    setprecision(prec)
    u0 = BigFloat.([-34 // 100, 78 // 100, 67 // 100, -56 // 10])
    B =
        BigFloat.([
            12//100 -78//100 91//100 34//100
            -45//100 56//100 3//100 54//100
            -67//100 09//100 18//100 89//100
            -91//100 -56//100 11//100 -56//100
        ])
    alpha = BigFloat.([12 // 100, -98 // 100, 45 // 100, 26 // 100])
    beta = BigFloat.([-4 // 100, 48 // 100, 23 // 100, -87 // 100])
    println("u0=$u0")
    println("B=$B")
    println("alpha=$alpha")
    println("beta=$beta")
    fct = (u, p, t) -> B * u + t * p[1] + p[2]
    t_max = big"1.0"
    epsilon = big"0.015"
    A = [0 0 1 0; 0 0 0 0; -1 0 0 0; 0 0 0 0]
    prob = HiOscODEProblem(fct, u0, (big"0.0", t_max), (alpha, beta), A, epsilon, B)
    #    nbmaxtest=6
    nbmaxtest = 12
    ordmax = 17
    debord = 3
    pasord = 2
    y = ones(Float64, nbmaxtest, div(ordmax - debord, pasord) + 1)
    x = zeros(Float64, nbmaxtest)
    ind = 1
    for order = debord:pasord:ordmax
        nb = 100
        indc = 1
        labels = Array{String,2}(undef, 1, div(order - debord, pasord) + 1)
        ordprep = min(order + 2, 10)
        ordprep = order + 2
        println("preparation ordre $ordprep")
        par_u0 = missing
        while indc <= nbmaxtest
            println("u0=$u0")
            println("B=$B")
            println("alpha=$alpha")
            println("beta=$beta")
            @time res = solve(
                prob,
                nb_tau = n_tau,
                order = order,
                order_prep = ordprep,
                nb_t = nb,
                par_u0 = par_u0,
                dense = false,
            )
            par_u0 = res.par_u0
            sol = res[end]
            solref = getexactsol(res.par_u0.parphi, u0, t_max)
            println("solref=$solref")
            println("nb=$nb sol=$sol")
            diff = solref - sol
            x[indc] = 1.0 / nb
            println("nb=$nb dt=$(1.0/nb) normInf=$(norm(diff,Inf)) norm2=$(norm(diff))")
            resnorm = Float64(norm(diff, Inf))
            y[indc, ind] = (resnorm < 1) ? resnorm : NaN
            println("result=$y")
            println("log2(y)=$(log2.(y))")
            nb *= 2
            indc += 1
        end
        for i = debord:pasord:order
            labels[1, div((i - debord), pasord)+1] = " order=$i "
        end
        yv = y[:, 1:ind]
        p = Plots.plot(
            x,
            yv,
            xlabel = "Δt",
            xaxis = :log,
            ylabel = "error",
            yaxis = :log,
            legend = :bottomright,
            label = labels,
            marker = 2,
        )
        pp = Plots.plot(
            x,
            yv,
            xlabel = "Δt",
            xaxis = :log,
            ylabel = "error",
            yaxis = :log,
            legend = :bottomright,
            label = labels,
            marker = 2,
            bottom_margin = 30px,
            left_margin = 60px,
        )
        prec_v = precision(BigFloat)
        eps_v = convert(Float32, epsilon)
        Plots.savefig(
            p,
            "out/res9p_$(prec_v)_$(eps_v)_$(order)_$(ordprep)_$(n_tau)_exact.pdf",
        )
        Plots.savefig(
            p,
            "out/res9p_$(prec_v)_$(eps_v)_$(order)_$(ordprep)_$(n_tau)_exact.png",
        )
        Plots.savefig(
            pp,
            "out/res9pp_$(prec_v)_$(eps_v)_$(order)_$(ordprep)_$(n_tau)_exact.pdf",
        )
        Plots.savefig(
            pp,
            "out/res9pp_$(prec_v)_$(eps_v)_$(order)_$(ordprep)_$(n_tau)_exact.png",
        )
        ind += 1
    end
end

# testODESolver()

# for i=3:9
#     fctMain(2^i)
# end
fctmain(32, 512)


module Catch22
using catch22_jll
using DimensionalData
using Libdl
import Statistics.mean, Statistics.std

include("Feature.jl")
include("FeatureSet.jl")
include("FeatureArray.jl")
include("metadata.jl")
include("testdata.jl")

catch22_jll.__init__() # Initialise the c-library

zscore(𝐱::AbstractVector) = (𝐱 .- mean(𝐱))./(std(𝐱))

"""
    _catch22(𝐱::AbstractArray{Float64}, fName::Symbol)
    _catch22(fName::Symbol, 𝐱::AbstractArray{Float64})
Evaluate the feature `fName` on the time series `𝐱`. If an array is supplied, features are calculated for its columns and returned as a Vector. See `Catch22.featuredescriptions` for a summary of the 22 available time series features.

# Examples
```julia-repl
𝐱 = Catch22.testdata[:test]
Catch22._catch22(𝐱, :DN_HistogramMode_5)
```
"""
function _catch22(𝐱::AbstractVector, fName::Symbol)::Float64
    if any(isinf.(𝐱)) || any(isnan.(𝐱))
        return NaN
    end
    𝐱 = zscore(𝐱)
    𝐱 = convert(Vector{Float64}, 𝐱)
    fType = featuretypes[fName]
    if fType <: AbstractFloat
        ccall(dlsym(dlopen(ccatch22), fName), Cdouble, (Ptr{Array{Cdouble}},Cint), 𝐱, Int(size(𝐱, 1)))
    elseif fType <: Integer
        ccall(dlsym(dlopen(ccatch22), fName), Cint, (Ptr{Array{Cdouble}},Cint), 𝐱, Int(size(𝐱, 1)))
    end
end
function _catch22(X::AbstractArray{Float64, 2}, fName::Symbol)::AbstractArray{Float64, 2}
    mapslices(𝐱 -> _catch22(𝐱, fName), X, dims=[1])
end

"""
    catch22(𝐱::Vector)
    catch22(X::Array)
    catch22[featurename::Symbol](X::Array)
Evaluate all features for a time series vector `𝐱` or the columns of an array `X`.
`catch22` is a FeatureSet, which means it can be indexed by feature names (as symbols) to return a subset of the available features.
`getnames(catch22)`, `getkeywords(catch22)` and `getdescriptions(catch22)`` will also return feature names, keywords and descriptions respectively.
Features are returned in a `FeatureArray`, in which array rows are annotated by feature names. A `FeatureArray` can be converted to a regular array with `Array(F)`.

# Examples
```julia-repl
𝐱 = Catch22.testdata[:test]
𝐟 = catch22(𝐱)

X = randn(100, 10)
F = catch22(X)
F = catch22[:DN_HistogramMode_5](X)
```
"""
catch22 = FeatureSet([(x -> _catch22(x, f)) for f ∈ featurenames], featurenames, featurekeywords, featuredescriptions)
export catch22

for f = featurenames
    eval(quote
        $f = catch22[$(Meta.quot(f))]; export $f
    end)
end

"""
    DN_HistogramMode_5(x::AbstractVector{Union{Float64, Int}}) # For example
An alternative to `catch22(...)``. All features, such as `DN_HistogramMode_5`, are exported as Features and can be evaluated by calling their names.

# Examples
```julia-repl
𝐱 = Catch22.testdata[:test]
f = DN_HistogramMode_5(𝐱)
```
"""
DN_HistogramMode_5;

end


"""
    prune(x::GaussianMixture, T::Real, U::Real, J_max::Integer)

Prune the posterior Gaussian-Mixture

Arguments:
- `x::GaussianMixture` Posterior state distribution. [Gaussian Mixture]
- `T::Real` Truncation threshold.  Drop distributions with weight less than T
- `U::Real` Merging threshold.  Merge if (μ_1-μ_2)^T * Σ^-1 * (μ_1-μ_2) < U
- `J_max::Integer` Maximum number of features
"""
function prune(x::GaussianMixture, T::Real, U::Real, J_max::Integer)
	l = 0
	J_k = x.N
	n = length(x.μ[1])

	# prune indices based on truncation threshold
	I = collect(1:J_k)
	I = I[x.w.>T]

	# setup output vectors
	w_new = zeros(J_k)
	μ_new = [zeros(n) for i=1:J_k]
	Σ_new = [zeros(n,n) for i=1:J_k]

	# perform pruning
	while length(I) != 0
		l = l+1
		j = I[argmax(x.w[I])]

		# Merge close features
		L = []
		for i in I
			delta = x.μ[i]-x.μ[j]
			if dot(delta, inv(x.Σ[i])*delta) < U
				push!(L,i)
			end
		end

		# determine merged parameters
		w_tilde = sum(x.w[L])
		μ_tilde = 1/w_tilde * sum(x.w[L].*x.μ[L])
		Σ_tilde = zeros(n,n)
		for i in L
			tmp = x.Σ[i] + (μ_tilde-x.μ[i])*(μ_tilde-x.μ[i])'
			Σ_tilde = Σ_tilde + x.w[i]*tmp
		end
		Σ_tilde = 1/w_tilde * Σ_tilde

		# store values
		w_new[l] = w_tilde
		μ_new[l] = μ_tilde
		Σ_new[l] = Σ_tilde

		# prune merged features
		I = setdiff(I,L)
	end

	if l > J_max 	# only keep the J_max with highest weights
		idxs = sort!([1:length(w_new);], by=i->(w_new[i]), rev=true)
		idxs = idxs[1:J_max]
		w_new = w_new[idxs]
		μ_new = μ_new[idxs]
		Σ_new = Σ_new[idxs]
	else		# drop zeros at end
		w_new = w_new[1:l]
		μ_new = μ_new[1:l]
		Σ_new = Σ_new[1:l]
	end

	x_new = GaussianMixture(w_new,μ_new,Σ_new)
	return x_new
end


using DiffEqFlux, OrdinaryDiffEq, Test # , Plots

p = [1.5 1.0;3.0 1.0]
function lotka_volterra(du,u,p,t)
  du[1] = p[1,1]*u[1] - p[1,2]*u[1]*u[2]
  du[2] = -p[2,1]*u[2] + p[2,2]*u[1]*u[2]
end

u0 = [1.0,1.0]
tspan = (0.0,10.0)

prob = ODEProblem(lotka_volterra,u0,tspan,p)
sol = solve(prob,Tsit5())

# plot(sol)

p = [2.2 1.0;2.0 0.4] # Tweaked Initial Parameter Array
ps = Flux.params(p)

function predict_adjoint() # Our 1-layer neural network
  Array(solve(prob,Tsit5(),p=p,saveat=0.0:0.1:10.0))
end

loss_adjoint() = sum(abs2,x-1 for x in predict_adjoint())

data = Iterators.repeated((), 100)
opt = ADAM(0.1)
cb = function () #callback function to observe training
  display(loss_adjoint())
end

predict_adjoint()

# Display the ODE with the initial parameter values.
cb()
Flux.train!(loss_adjoint, ps, data, opt, cb = cb)

@test loss_adjoint() < 1


#--------------------------------------------------------------------------------------------------
# Additional functions tests
#--------------------------------------------------------------------------------------------------
using Test
using Hensel

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

@test Hensel.logistic(Hensel.logistic(Hensel.logistic(0.3, 4.0), 4.0), 4.0) ==  Hensel.logistic(0.3,4.0,3)
@test Hensel.dyadic(Hensel.dyadic(Hensel.dyadic(0.3))) ==  Hensel.dyadic(0.3,3)
@test Hensel.nyadic(Hensel.nyadic(Hensel.nyadic(0.3, 5), 5), 5) ==  Hensel.nyadic(0.3,5,3)
@test Hensel.cfshift(Hensel.cfshift(Hensel.cfshift(0.3))) ==  Hensel.cfshift(0.3,3)

@test ((y) -> Hensel.FunEnv((x) -> Hensel.dyadic(0.35,y), (0, 1), 2, 8)(y, return_type="integer")).(0:10) == [
    154,
    205,
    102,
     51,
    153,
    204,
    102,
     51,
    153,
    204,
    102
]

@test ((y) -> Hensel.FunEnv((x) -> Hensel.nyadic(0.35,2,y), (0, 1), 2, 8)(y, return_type="integer")).(0:10) == [
    154,
    205,
    102,
     51,
    153,
    204,
    102,
     51,
    153,
    204,
    102
]

@test Hensel.FunEnv(Hensel.dyadic, (0, 1), 2, 8)(1,2,return_type="integer") == 64
@test Hensel.FunEnv(Hensel.dyadic, (0, 1), 2, 8)((1,2,"integer")) == 64
@test Hensel.FunEnv(Hensel.dyadic, (0, 1), 2, 8)((1,2)) == 64
@test Hensel.FunEnv(Hensel.dyadic, (0, 1), 2, 8).([(1,1),(1,2),(1,3)]) == [128, 64, 32]
@test Hensel.FunEnv(Hensel.dyadic, (0, 1), 2, 8).([(1,1,"hensel"),(1,2,"hensel"),(1,3,"hensel")]) == [
    [0, 0, 0, 0, 0, 0, 0, 1],
    [0, 0, 0, 0, 0, 0, 1, 0],
    [0, 0, 0, 0, 0, 1, 0, 0]
]

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


# Test FMM matvec vs full matrix
push!(LOAD_PATH, string(pwd(),"/src"))
using Base.Test
using ModifiedStokesSolver
import AnalyticDomains

@testset "Matvec" begin
    # Params
    alpha = 10.0
    # Discretize
    numpanels = 20
    panelorder = 16
    curve = AnalyticDomains.starfish(n_arms = 3, amplitude=0.2)
    dcurve = discretize(curve, numpanels, panelorder)
    N = dcurve.numpoints
    # Compute
    x = rand(2*N)
    matvec = system_matvec(dcurve, alpha)
    yfmm = matvec(x)
    A = system_matrix(dcurve, alpha)
    ydir = A*x
    erel = norm(yfmm-ydir, Inf) / norm(ydir, Inf)
    @test erel < 5e-14
end


using Test
using LinearAlgebra
import sbp
using sbp.Config: Tv
using sbp.Metrics: build_jacobian, build_contravariant_basis,
                   build_covariant_metric_tensor,
                   build_contravariant_metric_tensor
using sbp.Staggered: build_all_operators_2d, build_covariant_basis_mm
using sbp.Grid: grid_xp, grid_xm, grid_2d_x, grid_2d_y

sop = sbp.OP2019
nx=10
ny=20

desc = sop.load_description()
builder = n -> sop.build_operators(desc, n)
pp, mm, pm, mp = build_all_operators_2d(builder, nx, ny)

x1, h = grid_xp(nx)
y1, h = grid_xp(ny)

x = grid_2d_x(x1, ny)
y = grid_2d_y(y1, nx)

fx = x
fy = y

a = build_covariant_basis_mm(fx, fy, mm, pm)
Ga = build_covariant_metric_tensor(a)

@testset "Jacobian" begin
        Jans = ones(mm.nx * mm.ny)
        J = build_jacobian(a)
        @test isapprox(J, Jans)
end

@testset "Covariant Metric tensor" begin
         Ga = build_covariant_metric_tensor(a)
         # Check symmetry
         @test isapprox(Ga.a12, Ga.a21)

end

@testset "Contravariant basis" begin
        J = build_jacobian(a)
        b = build_contravariant_basis(J, a)

        one = ones(mm.nx * mm.ny)
        zero = zeros(mm.nx * mm.ny)

        @test isapprox(a.x_r1 .* b.r1_x + a.y_r1 .* b.r1_y, one)
        @test isapprox(a.x_r1 .* b.r2_x + a.y_r1 .* b.r2_y, zero, atol=1.0)
        @test isapprox(a.x_r2 .* b.r2_x + a.y_r2 .* b.r2_y, one)
        @test isapprox(a.x_r2 .* b.r1_x + a.y_r2 .* b.r1_y, zero, atol=1.0)

        @testset "Contravariant Metric tensor" begin
                Gb = build_contravariant_metric_tensor(b)
                # Check symmetry
                @test isapprox(Gb.b12, Gb.b21)
                # Check orthogonality
                @test isapprox(Ga.a11 .* Gb.b11 + Ga.a12 .* Gb.b12, one)
                @test isapprox(Ga.a11 .* Gb.b12 + Ga.a12 .* Gb.b22, zero,
                               atol=1.0)
                @test isapprox(Ga.a12 .* Gb.b11 + Ga.a22 .* Gb.b12, zero,
                               atol=1.0)
                @test isapprox(Ga.a21 .* Gb.b12 + Ga.a22 .* Gb.b22, one)

        end

end


# # Model Parameters
@with_kw mutable struct CRPar
    r = 2.0
    k = 3.0
    a = 1.1
    h = 0.8
    e = 0.7
    m = 0.4
end

## CR exp resource, type I consumer
function cr_expI_model(du, u, p, t)
    @unpack r, a, e, m = p
    R, C, = u
    du[1] = r * R - a * R * C
    du[2] = e * a * R * C - m * C
    return
end

function cr_expI_model(u, p, t)
    du = similar(u)
    cr_expI_model(du, u, p, t)
    return du
end

## CR log resource, type I consumer
function cr_logI_model(du, u, p, t)
    @unpack r, k, a, e, m = p
    R, C, = u
    du[1] = r * R * (1 - R / k) - a * R * C
    du[2] = e * a * R * C - m * C
    return
end

function cr_logI_model(u, p, t)
    du = similar(u)
    cr_logI_model(du, u, p, t)
    return du
end

## CR log resource, type II consumer
function cr_logII_model(du, u, p, t)
    @unpack r, k, a, e, h, m = p
    R, C, = u
    du[1] = r * R * (1 - R / k) - a * R * C / (1 + a * h * R)
    du[2] = e * a * R * C / (1 + a * h * R) - m * C
    return
end

function cr_logII_model(u, p, t)
    du = similar(u)
    cr_logII_model(du, u, p, t)
    return du
end

## CR exp resource, type II consumer
function cr_expII_model(du, u, p, t)
    @unpack r, a, e, h, m = p
    R, C, = u
    du[1] = r * R  - a * R * C / (1 + a * h * R)
    du[2] = e * a * R * C / (1 + a * h * R) - m * C
    return
end

function cr_expII_model(u, p, t)
    du = similar(u)
    cr_expII_model(du, u, p, t)
    return du
end

#### Code that is used in multiple scripts ####

function expII_resiso(R, p)
    @unpack r, a, h = p
    return R * h * r + r / a
end

function expII_coniso(p)
    @unpack m, a, e, h = p
    return m / (a * (e - h * m))
end

# Function to calculate the jacobian at any point (with any model)
function jac(u, model, p)
    ForwardDiff.jacobian(u -> model(u, p, NaN), u)
end

# Function to calculate maximum eigenvalue (real part)
λ_stability(M) = maximum(real.(eigvals(M)))

# Function to find the interior equilibrium value for any efficiency value (for CR_logII model)
function logII_eq(p)
    @unpack r, k, a, h, m, e = p
    R = m / (a * (e - h * m))
    C = r * (-R^2 * a * h + R * a * h * k - R + k) / (a * k)
    return [R, C]
end


# epicycloid

# Makie
using CairoMakie
CairoMakie.activate!()
let
    function EpicycloidLines()
        ecycl_x(r, k, θ) = r * (k .+ 1) .* cos.(θ) .- r * cos.((k .+ 1) .* θ)
        ecycl_y(r, k, θ) = r * (k .+ 1) .* sin.(θ) .- r * sin.((k .+ 1) .* θ)
        θ = LinRange(0, 6.2π, 1000)
        tronLegacy = ["#FF410DFF", "#6EE2FFFF", "#F7C530FF", "#95CC5EFF",
            "#D0DFE6FF", "#F79D1EFF", "#748AA6FF"]

        fig = Figure(resolution = (600, 400), font = :sans)
        ax = Axis(fig[1, 1], aspect = 1, title = "Epicycloid",
            xlabel = "x(θ) = r(k+1)cos(θ) -rcos((k+1)θ)",
            ylabel = "y(θ) = r(k+1)cos(θ) -rcos((k+1)θ)")
        [CairoMakie.lines!(ecycl_x(2k, k, θ), ecycl_y(2k, k, θ), linewidth = 1.5, label = "$(k)",
            color = tronLegacy[indx]) for (indx, k) in enumerate(2:0.5:5)]
        Legend(fig[1, 2], ax, "k , r = 2k")
        colsize!(fig.layout, 1, Aspect(1, 1.0))
        fig
    end
    fig = with_theme(EpicycloidLines, theme_black())
    display(fig)
end;


# specialized methods
"""
```julia
dotdot(::Vec, ::SymmetricFourthOrderTensor, ::Vec)
```
Computes a special dot product between two vectors and a symmetric fourth order tensor
such that ``a_k C_{ikjl} b_l``.
"""
@generated function dotdot{dim}(v1::Vec{dim}, S::SymmetricTensor{4, dim}, v2::Vec{dim})
    idx(i,j,k,l) = compute_index(SymmetricTensor{4, dim}, i, j, k, l)
    exps = Expr(:tuple)
    for j in 1:dim, i in 1:dim
        ex1, ex2 = Expr[],Expr[]
        for l in 1:dim, k in 1:dim
            push!(ex1, :(get_data(v1)[$k] * get_data(S)[$(idx(i,k,j,l))]))
            push!(ex2, :(get_data(v2)[$l]))
        end
        push!(exps.args, make_muladd_exp(ex1, ex2))
    end
    return quote
        $(Expr(:meta, :inline))
        @inbounds r = $exps
        Tensor{2, dim}(r)
    end
end


mutable struct ECOSls <: ECOSWrapper
    ecos_ptr::Ptr{ECOS.Cpwork}
    A1::AbstractMatrix{Float64}
    b1::Vector{Float64}
    A2::AbstractMatrix{Float64}
    b2::Vector{Float64}
    n_ecos::Int
    n::Int
    m1::Int
    m2::Int
end

mutable struct ECOSlssplx <: ECOSWrapper
    ecos_ptr::Ptr{ECOS.Cpwork}
    A::AbstractMatrix{Float64}
    b::Vector{Float64}
    n_ecos::Int
    n::Int
    m::Int
end

"""
Solve the constrained least-squares problem

    minimize  ||A1x-b1||₂ st A2x=b2, x≥0.

Rephrases the problem as linear SOC program:

    minimize  [0]'[x]  st  [A1 0  I][x]    [b1] m1
              [1] [t]      [A2 0  0][t] =  [b2] m2
              [0] [r]               [r]
                            n  1  m1
                               x, (t,r) ∈  Rn₊ x (Qn₊)

with dimensions  len(x) = n   len(b1) = m1
                 len(r) = m1  len(b2) = m2

Solve using ECOS, which formulates linear SOC programs as

    minimize  c'x  st  A*x=b, G*x + s = h

Thus, set the ECOS variables as

  c      = [0 , 1,  0]
  A      = [A1  0   I]
            A2  0   0]
  G      = -Identity(n+1+m1)
  h      = [0   0   0]
"""
function ECOSls(A1::AbstractMatrix, b1::Vector{Float64},
                A2::AbstractMatrix, b2::Vector{Float64};
                kwargs...)
    m1, n1 = size(A1)
    m2, n2 = size(A2)
    n1 != n2 && throw(DimensionMismatch("matrix sizes not compatible"))
    n = n1

    n_ecos = n + 1 + m1  # number of vars == len(x)
    m = n_ecos           # number of inequality constraints == len(h)
    p = m1 + m2          # number of equality constraints == len(b)
    l = n                # number of vars in positive orthant
    ncones = 1           # number of SOC cones
    q = [1+m1]           # number of vars in the SOC
    e = 0                # number of exponential cones
    G = ECOS.ECOSMatrix(sparse(-1.0I,n_ecos,n_ecos))
    A = ECOS.ECOSMatrix([A1  spzeros(m1,1)  sparse(1.0I,m1,m1)
                         A2  spzeros(m2,1)      spzeros(m2,m1)])
    c = zeros(n_ecos); c[n+1] = 1.0
    h = zeros(n_ecos)
    b = vcat(b1, b2)

    ecos_ptr = ECOS.setup(n_ecos, m, p, l, ncones, q, e, G, A, c, h, b; kwargs...)
    # problem = ECOSls(ecos_ptr, A1, b1, A2, b2, n_ecos, n, m1, m2)
    ECOS.solve(ecos_ptr)
    # finalizer(cleanup!, problem)
    # return ecos_ptr
end

"""
    ls_simplex(A, b)

Setup the constrained least-squares problem

    minimize  ||Ax-b||₂  subj to  sum(x)=1, x≥0
"""
function ls_simplex(A::AbstractMatrix, b::Vector{Float64}; kwargs...)
    m, n = size(A)
    lsg = ECOSls(A, b, ones(1,n), [1.0]; kwargs...)
    return ECOSlssplx(lsg.ecos_ptr, A, b, lsg.n_ecos, n, m)
end

function solve!(p::ECOSlssplx)

    exitflag = ECOS.solve(p.ecos_ptr)

    if exitflag != ECOS.ECOS_OPTIMAL
        error("did not solve to optimality --- $exitflag")
    end

    # Extract the primal and dual solutions of the ECOS problem
    ecos_prob = unsafe_wrap(Array, p.ecos_ptr, 1)[1]
    xtr = unsafe_wrap(Array, ecos_prob.x, p.n_ecos)[:]
    du_sol_eq = unsafe_wrap(Array, ecos_prob.y, p.m)[:]
    du_sol_ineq = unsafe_wrap(Array, ecos_prob.z, p.n_ecos)[:]
    ECOS.cleanup(p.ecos_ptr)
    # Extract intended solution
    # x = xtr[1:p.n]
    # r = xtr[p.n+2:end]
    # z = du_sol_ineq[1:p.n]
    # μ = du_sol_ineq[p.n+1]

    return xtr
end


using ElectrostaticPIC1D, Random, Test; Random.seed!(0)

using ForwardDiff, QuadGK

@testset "BasisFunctions" begin

numtests = 10

function isnormalisedtest(Shape)
  outcome = true
  for _ ∈ 1:numtests
    w = rand()
    s = Shape(w)
    x0 = randn()
    integral = quadgk(x->s(x, x0), x0 - 10*w, x0 + 10*w, order=27, rtol=eps())[1]
    outcome = integral ≈ 1
    outcome || break
  end
  @test outcome
end
function supporttest(Shape)
  for _ ∈ 1:10
    w = rand()
    x0 = randn()
    b = BasisFunction(Shape(w), x0)
    outcome = true
    for _ in 1:1000
      x = 3 * randn() * w + x0
      if x ∈ b
        outcome = b(x) > 0
        outcome || @show b, x, b(x)
      else
        outcome = b(x) <= b(x0) * eps()
        outcome || @show b, x, b(x)
      end
      outcome || break
    end
    @test outcome
  end
end
function energytest(Shape)
  outcome = true
  for _ ∈ 1:numtests
    w = rand()
    x0 = randn()
    b = BasisFunction(Shape(w), x0)
    expected = quadgk(x->b(x)^2, x0 - 10*w, x0 + 10*w, order=27, rtol=eps())[1] * weight(b)^2 / 2
    result = integral(b, b, PeriodicGridBC(x0-10*w, x0+10*w)) * weight(b)^2 / 2
    outcome = expected ≈ result
    outcome || break
  end
  @test outcome
end

for S ∈ (BSpline{0}, BSpline{1}, BSpline{2}, GaussianShape)
 @testset "$S" begin
   isnormalisedtest(S)
   supporttest(S)
   energytest(S)
 end
end

@testset "Pairwise Integrals" begin
for U ∈ (DeltaFunctionShape, BSpline{0}, BSpline{1}, BSpline{2}, GaussianShape)
for V ∈ (BSpline{0}, BSpline{1}, BSpline{2}, GaussianShape)
@testset "$U $V" begin
  noverlaps = 0
  while true
    v = BasisFunction(V(5*rand()), randn())
    u = if U <: DeltaFunctionShape
      BasisFunction(U(), lower(v) + rand() * width(v))
    else
      BasisFunction(U(5*rand()), randn())
    end
    expected = if U <: DeltaFunctionShape
      v(centre(u))
    else
      quadgk(x-> u(x) * v(x), lower(u, v), upper(u, v);
             order=51, atol=0, rtol=eps())[1]
    end
    expected < eps() && continue
    try
      @inferred integral(u, v)
      @test true
    catch
      @test false
    end
    try
      @inferred integral(u, v, PeriodicGridBC(1.0))
      @test true
    catch
      @warn "integral($U, $V) cannot be inferred"
      @test false
    end
    result = integral(u, v)
    @test result ≈ expected atol=100eps()
    break
  end
end # testset U, V
end # V
end # U
end # testset integrals

end


#=
S. Boldo and J.M. Muller,
Exact and Approximated Error of the FMA,
IEEE Transactions on Computers, 60(2011), 157–164.

function fma_with_error(a::Float64, b::Float64, c::Float64)
    hi  = fma(a, b, c)
    u1, u2 = two_prod(a, b)
    v1, v2 = two_sum(c, u2)
    w1, w2 = two_sum(u1, v1)
    t = w1 - hi + w2
    md, lo = two_hilo_sum(t, v2)
    return hi, md, lo
end

  (a * b) + c == hi + md + lo
   fma(a,b,c) == hi
   |md + lo| = unit_roundoff_error(|hi|)/2
   |lo| = unit_roundoff_error(|md|)/2

function two_fma(a::Float64, b::Float64, c::Float64)
    hi  = fma(a, b, c)
    u1, u2 = two_prod(a, b)
    v1, v2 = two_sum(c, u2)
    w1, w2 = two_sum(u1, v1)
    md = (w1 - hi + w2) + v2
    return hi, md
end

=#

function Base.fma(xhi::T, xlo::T, yhi::T, ylo::T, zhi::T, zlo::T) where {T<:Float64}
   chi, c1 = two_prod(xhi, yhi)
   t0 = xlo * ylo
   t1 = fma(xhi, ylo, t0)
   c2 = fma(xlo, yhi, t1)
   c3 = c1 + c2
   dhi, dlo = two_hilo_sum(chi, c3)

   shi, slo = two_sum(zhi, dhi)
   thi, tlo = two_sum(zlo, dlo)
   c = slo + thi
   vhi, vlo = two_hilo_sum(shi, c)
   w = tlo + vlo
   hi, lo = two_hilo_sum(vhi, w)
   return Double64((hi, lo))
end

@inline function Base.fma(x::T, y::T, z::T) where {T<:Double64}
   return fma(Hi(x), Lo(x), Hi(y), Lo(y), Hi(z), Lo(z))
end

for T in (:Double64, :ComplexD64)
   @eval @inline function Base.muladd(x::$T, y::$T, z::$T)
      return x*y + z
   end
end

#=
   Quadruple-precision BLAS using Bailey's arithmetic with FMA instruction
   by S. Yamada, t. Ina, N. Sasa, Y.Idomura, M. Machida, T. Imamura
   2017 IEEE International Parallel and Distributed Processing Symposium Workshops
=#

function Base.:(*)(ahi::T, alo::T, bhi::T, blo::T) where {T<:Float64}
   p1 = ahi * bhi
   p2 = fma(ahi, bhi, -p1)
   p2 += alo*bhi + ahi*blo
   chi = p1 + p2
   clo = p2 - (chi - p1)
   return Double64((chi, clo))
end

@inline function Base.:(*)(a::T, b::T) where {T<:Double64}
   return (*)(Hi(a), Lo(a), Hi(b), Lo(b))
end


using Gadfly
using Distributions

"""
Bimodal Spike Demo
"""
function drawBimodalDemo()
    μ = [.0 2.]
    Σ = [1. .05]
    w = [.5 .5]
    xs = collect(-2:.01:μ[2]*2)

    n1 = Normal(μ[1], Σ[1])
    n2 = Normal(μ[2], Σ[2])
    p = w[1]*pdf(n1, xs) + w[2]*pdf(n2, xs)

    μ = mean(xs .* p)

    pp = plot(x=xs, y=p, Geom.line, Theme(default_color=colorant"black", line_width=3px),
              xintercept=[μ], Geom.vline(color=colorant"red", size=3px))
end


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

using Markdown
using InteractiveUtils

# ╔═╡ 59ed6c40-8025-11ec-0ace-d5b826d81255
begin
	bigbreak = HTML("<br>"^4);
	using PlutoUI
	using Test
end

# ╔═╡ c3b89eff-c3e6-47d0-acc8-661dde65305c
md"""
# Data structures
Once we start working with many pieces of data at once, it will be convenient for us to store data in structures like arrays or dictionaries (rather than just relying on variables).

Types of data structures covered:
1. Tuples and NamedTuples
2. Dictionaries
3. Arrays

As an overview, tuples and arrays are both ordered sequences of elements (so we can index into them). Dictionaries and arrays are both mutable. We'll explain this more below!

"""

# ╔═╡ a58c0b6f-c5f4-43e9-87ad-4eca62157ce4
md"""
## Tuples
We can create a tuple by enclosing an ordered collection of elements in ( ).

Syntax:
```julia
(item1, item2, ...)
```
"""

# ╔═╡ 30233243-ed9d-4bfd-a9af-fcde96677998
myfavoriteanimals = ("penguins", "cats", "sugargliders")

# ╔═╡ 15717659-eb1a-40d4-abfc-7d680f4722c8
md"We can index into this tuple,"

# ╔═╡ 9cd7ec78-75dc-4cfc-a4ef-2fd8e9a8d40c
myfavoriteanimals[1]

# ╔═╡ dae03614-aa84-483a-8e54-5f493a5a1b82
md"but since tuples are immutable, we can't update it"

# ╔═╡ 6d1bb7e6-6cff-4cee-873d-ccc82f4f0be8
myfavoriteanimals[1] = "otters"

# ╔═╡ d85c4c9a-3470-4b5a-95af-f2e9c2b04c00
md"""## Now in 1.6: NamedTuples
As you might guess, NamedTuples are just like Tuples except that each element additionally has a name! They have a special syntax using = inside a tuple:
```julia
(name1 = item1, name2 = item2, ...)
```
_Special case, for just one entry, use `(; name1 = item1)`_
"""

# ╔═╡ a8e3433a-a9d3-4081-a8fe-5fa8732c8d71
yourfavoriteanimals = (bird = "penguins", mammal = "cats", marsupial = "sugargliders")

# ╔═╡ 9bde2570-6ce5-4ad7-8b0b-dd7889f8bea1
md"
Like regular Tuples, NamedTuples are ordered, so we can retrieve their elements via indexing:"

# ╔═╡ f425304c-29c1-40eb-b606-9dc39a62e252
yourfavoriteanimals[1]

# ╔═╡ 9db7047d-d900-403c-9a10-85328d5faf44
md"They also add the special ability to access values by their name:"

# ╔═╡ 560ddf25-33e0-4075-b1f4-475f308c99a7
yourfavoriteanimals.bird

# ╔═╡ ff1774ca-7a03-474a-80c7-ff2bd9f38e61
bigbreak

# ╔═╡ e13cae8d-e8ad-4ed8-a844-e1cedb452493
md""" ## Dictionaries
If we have sets of data related to one another, we may choose to store that data in a dictionary. We can create a dictionary using the Dict() function, which we can initialize as an empty dictionary or one storing key, value pairs.

Syntax:
```julia
Dict(key1 => value1, key2 => value2, ...)
```
A good example is a contacts list, where we associate names with phone numbers.
"""

# ╔═╡ 079ece15-7872-4e24-b762-418152f232e6
myphonebook = Dict("Jenny" => "867-5309", "Ghostbusters" => "555-2368")

# ╔═╡ 37887043-b092-4d0f-bbc7-4288c4c761cd
md"""
In this example, each name and number is a "key" and "value" pair. We can grab Jenny's number (a value) using the associated key
"""

# ╔═╡ 07fb97a3-60f0-40e5-befe-5d27a0b5d597
myphonebook["Jenny"]

# ╔═╡ 5efd2a9e-7bc9-424e-812d-2fd554677b55
md"We can add another entry to this dictionary as follows"

# ╔═╡ 6647ab87-0ff0-4d61-ad10-897942b461dc
myphonebook["Kramer"] = "555-FILK"

# ╔═╡ 0b438e45-fa85-435b-9319-741379f8bf2d
myphonebook

# ╔═╡ a9ef31b2-9c59-4ed4-ba78-173a383211e2
md"We can delete Kramer from our contact list - and simultaneously grab his number - by using pop!"

# ╔═╡ d869d9fd-703b-475b-b9c2-337cbc1b967a
pop!(myphonebook, "Kramer")

# ╔═╡ 09eb599f-f454-4167-8b63-002e6b3b1ef9
myphonebook

# ╔═╡ 5c30cf8d-1295-4e42-ba93-2b49b005455c
md"Unlike tuples and arrays, dictionaries are not ordered. So, we can't index into them."

# ╔═╡ d96442d4-eaa9-4339-b55f-e48e39169e48
myphonebook[1]

# ╔═╡ 8e1fac94-4a44-499b-a49b-e4761549b166
md"
In the example above, julia thinks you're trying to access a value associated with the key 1."

# ╔═╡ 0536d875-6f77-4834-9fa4-a0a31f7fa09a
bigbreak

# ╔═╡ d54f80fa-750b-494b-be58-933d375fb8c2
md"""
## Arrays
Unlike tuples, arrays are mutable. Unlike dictionaries, arrays contain ordered collections.
We can create an array by enclosing this collection in `[ ]`.

Syntax:
```julia
[item1, item2, ...]
```
For example, we might create an array to keep track of my friends

"""

# ╔═╡ 8bd0408e-7b00-419c-930c-8e9361ad1574
myfriends = ["Ted", "Robyn", "Barney", "Lily", "Marshall"]

# ╔═╡ 3b2e4202-19be-4de5-92ae-b3911af42f7f
typeof(myfriends)

# ╔═╡ 3b3cab5a-c5e2-4587-9aae-a97cbc15f238
md"The 1 in `Array{String,1}` means this is a one dimensional vector. An `Array{String,2}` would be a 2d matrix, etc. The String is the type of each element."

# ╔═╡ c94ca702-fe13-468f-938a-be9b16a879e2
md"or to store a sequence of numbers"

# ╔═╡ b577d259-b564-47c8-b732-a0af01f30014
fibonacci = [1, 1, 2, 3, 5, 8, 13]

# ╔═╡ 24dea8be-3930-4049-9a44-b8151e164d91
typeof(fibonacci)

# ╔═╡ 7701def7-26b9-4f53-90a9-3eec7be70a16
eltype(fibonacci)

# ╔═╡ d1571780-39a7-4d8d-9ee0-b9770e235bdd
md"we can also mix types"

# ╔═╡ 98457248-d957-4376-8086-6b5866a743c9
mixture = [1, 1, 2, 3, "Ted", "Robyn"]

# ╔═╡ 0553623d-5416-4cf4-8cd8-a836dbe5c41f
typeof(mixture)

# ╔═╡ cf556b7b-1d94-43d7-99f8-0a134bd9865b
md"Once we have an array, we can grab individual pieces of data from inside that array by indexing into the array. For example, if we want the third friend listed in myfriends, we write"

# ╔═╡ 46bdba66-e0fa-4631-b0a1-9d532aa25e89
myfriends[1]

# ╔═╡ 6debf09f-ff30-4e30-93b3-8872c258908c
md"
We can use indexing to edit an existing element of an array"

# ╔═╡ 28821165-3106-4983-be9f-3155cbc4ba53
myfriends[3] = "Baby Bop"

# ╔═╡ b221fbd3-c519-4c9c-8c98-bcee6fae796d
bigbreak

# ╔═╡ 3b5b7dcf-6a56-49cb-9dd8-1b2aba81ef87
md"""
### Yes, Julia is 1-based indexing, not 0-based like Python.
_(I share your pain)._

"""

# ╔═╡ 8f398c3a-9def-4f41-a171-45d38527c83c
bigbreak

# ╔═╡ a39e4255-43ee-4ed3-8282-fd25ca0b8d53
md"""
Let's continue anyway.

We can also edit the array by using the push! and pop! functions. push! adds an element to the end of an array and pop! removes the last element of an array.

We can add another number to our fibonnaci sequence
"""

# ╔═╡ 214cde58-aa66-4499-9a5e-5a6a36c22115
push!(fibonacci, 21)

# ╔═╡ 7aa62fe2-9c5e-4640-8960-21a9084e436f
md"and then remove it"

# ╔═╡ 687f65ad-2320-4cab-b576-0b94de505c10
pop!(fibonacci)

# ╔═╡ 55e690c8-14c2-46f8-bd3c-f64e34a8d74e
fibonacci

# ╔═╡ 58fbd165-41f3-4cac-aaa3-e498e4621d1d
md"So far I've given examples of only 1D arrays of scalars, but arrays can have an arbitrary number of dimensions and can also store other arrays.

For example, the following are arrays of arrays:"

# ╔═╡ a6fc5004-5ef2-409c-9b6a-5e7e7efb25ad
favorites = [["koobideh", "chocolate", "eggs"],["penguins", "cats", "sugargliders"]]

# ╔═╡ c4d8af83-f7f9-47a0-9023-fe46e6bc1109
favorites[2][1]

# ╔═╡ 6229a275-b52d-4c81-a05b-f22d15b0cf32
numbers = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]

# ╔═╡ 53cd90d5-8a70-43a6-b72e-d61185c0b7aa
typeof(numbers)

# ╔═╡ 4fdc2f91-dc97-4c58-9cc1-486d30ad7890
bigbreak

# ╔═╡ 08924e9f-27b8-444a-b9af-9fa34cf1d61b
md"Below are examples of 2D and 3D arrays populated with random values."

# ╔═╡ daa23583-9b16-4a76-bde8-9324f3dc1908
rand(4,3)

# ╔═╡ 12e389d6-9ac2-49b5-a678-cf71787541eb
rand(4,3,2)

# ╔═╡ e791a417-a00a-479e-b433-3773dd74a48e
bigbreak

# ╔═╡ b18f9af8-130d-4bac-ba8f-230e89d229b7
X = [1 2; 2 3 ]

# ╔═╡ 26de3d3b-9229-4c34-a7d1-e0c89ef6f843
Z = [i+j for i in 1:10, j in 1:10]

# ╔═╡ 30933711-723e-4b8f-8be9-3928635e068e
Z[3:5, 2:2:6]

# ╔═╡ eb9a5d9f-9f39-401c-aa35-102414ed99de
bigbreak

# ╔═╡ 3280db88-0d68-4426-9f89-0fd3914e6ffa
md"Be careful when you want to copy arrays!"

# ╔═╡ 4fa1d5b9-36a0-491e-a034-c887ee7338d8
fibonacci

# ╔═╡ 68ce440d-53f8-4312-8911-a7348e602a1f
somenumbers = copy(fibonacci)

# ╔═╡ 1be992a6-fc03-437f-9dc1-6072f0b08327
somenumbers[1] = 404

# ╔═╡ 36086001-a2bd-4648-b2b3-3d526f0dcf71
fibonacci

# ╔═╡ 5ece33d6-4622-4385-b84c-639bd7571915
md"Editing somenumbers caused fibonacci to get updated as well!

In the above example, we didn't actually make a copy of fibonacci. **We just created a new way to access the entries in the array bound to fibonacci.**

If we'd like to make a copy of the array bound to fibonacci, we can use the copy function."

# ╔═╡ ce690520-a919-4e3d-a201-ec2baedf913d
let
	a = [[1,2],1]

	b = deepcopy(a)
	b[1][1] = 10
	b[2] = 2
	a
end

# ╔═╡ 28263f4f-37b0-43fd-9eab-6169d82b0581
bigbreak

# ╔═╡ 421ab876-efa5-4a7d-9d33-5b9eca5f528e
begin
	# let's restore the original array
	fibonacci[1] = 1
	fibonacci
end

# ╔═╡ 142b9281-196f-4d3e-ae7d-fbc1011255f6
somemorenumbers = copy(fibonacci)

# ╔═╡ 91e7e129-c7b5-4e1f-9025-bdac4d569b46
somemorenumbers[1] = 404

# ╔═╡ 0e20de93-a6f1-415c-a08e-53de67065bc4
fibonacci

# ╔═╡ be7bad99-185a-43ed-9c8d-33e5abcdbbd0

# ╔═╡ 40b0919d-64a5-4a89-95fb-5266c2c07260
md"In this last example, fibonacci was not updated. Therefore we see that the arrays bound to somemorenumbers and fibonacci are distinct."

# ╔═╡ ec0fa200-0c79-467a-af5c-97e79f37253d
let
	function gradientdecent(f, x0)

	end

	x0 = [1., 2.]
	gradientdecent( nothing, x0)
end

# ╔═╡ 55633c9d-2c1b-40c5-9475-48f18b1f8d77
bigbreak

# ╔═╡ 69e96bb3-fc25-4f44-89cd-8f2aeb7065fe
md"Variables are just associated to values!"

# ╔═╡ 9d4aeb68-6649-42f6-8a93-997a136499d4
begin
	A = [1,2,3]

	A_orig = Ref(A)  # points to the array [1,2,3]

	A .= (2 * A)

end

# ╔═╡ 2a3d5a98-751c-48c9-b102-89bc40e74251
typeof(2*A)

# ╔═╡ 457bc2c7-9aae-4b28-b216-ef9454227651
A

# ╔═╡ f35046c4-4a6e-42e5-b691-78b6ef3b0d27
md"The original array (might) still exists and we wasted memory!"

# ╔═╡ 8af669a5-4f1f-4f0e-9de7-030e3196f96b
A_orig[]

# ╔═╡ 32035322-60ec-4ed7-a340-84c6457a2464
md"""
#### Why did we not overwrite the original array?
Imagine, in the first step we did
```julia
A = [1, 2, 3]
```

$(LocalResource("images/var_1.svg", :width => 800))
After the command
```julia
A = 2*A
```
the variable is bound to a new value. **The old value remains unchanged!**

$(LocalResource("images/var_2.svg", :width => 800))

"""

# ╔═╡ 2c04e557-4afa-4acc-ace1-45f994aa3ed3
bigbreak

# ╔═╡ 83c37801-d248-442a-a360-829f740f5cd9
md"""
## Exercises
3.1
Create an array, `x` with the number 1 to 20 and add `+ 10` to all odd entries.

"""

# ╔═╡ 4170966b-cbb4-4b7b-b6e0-528f63a1c019
begin
	x = []
	# ...
end

# ╔═╡ f720b6ad-eea0-4c07-b528-822b80305000
md"""_The current solution is $( (x == [i + (i%2==1)*10 for i in 1:20]) ? "correct" : "false")._"""

# ╔═╡ 7831cafd-8ac8-41be-ad0d-a5abeeb66276
bigbreak

# ╔═╡ fc60372a-f853-4005-bc40-64e9d0615a32
md"""
3.2

Try to add `"Emergency"` as key to `ourphonebook` with the value `string(911)` with the following code
```julia
ourphonebook["Emergency"] = 112
```
Why doesn't this work?
"""

# ╔═╡ 89b74238-dd21-4063-bdb8-d64adcc2184c
ourphonebook = Dict("Jenny" => "867-5309", "Ghostbusters" => "555-2368")

# ╔═╡ b8c1cdf5-550e-49be-83b0-2a01fc583886
# enter you code here

# ╔═╡ 1b7789aa-1969-486c-a733-6e3d10d135bc

# ╔═╡ f42fa416-47d1-49b6-9d4b-365740908e03
md"How can you change the definition of `ourphonebook` such that you can add and integer value?"

# ╔═╡ 4f191e04-1972-47af-b1d9-2de4857833bb
bigbreak

# ╔═╡ 342337cc-6202-45f7-bf5d-6688b770ac75
md"""
!!! hint "Hint"
      Check out `typeof()` for these dictionaries. You can use `Dict{K,V}` instead of the generic `Dict` to create a dictionary with key type `T` and value type `V`. What should `V` be?
"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[compat]
PlutoUI = "~0.7.32"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.1"
manifest_format = "2.0"

[[deps.AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "8eaf9f1b4921132a4cff3f36a1d9ba923b14a481"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.4"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[deps.Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[deps.ColorTypes]]
deps = ["FixedPointNumbers", "Random"]
git-tree-sha1 = "024fe24d83e4a5bf5fc80501a314ce0d1aa35597"
uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
version = "0.11.0"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[deps.Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[deps.FixedPointNumbers]]
deps = ["Statistics"]
git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
version = "0.8.4"

[[deps.Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[deps.HypertextLiteral]]
git-tree-sha1 = "2b078b5a615c6c0396c77810d92ee8c6f470d238"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.3"

[[deps.IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[deps.InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[deps.JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[deps.LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[deps.Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[deps.LinearAlgebra]]
deps = ["Libdl", "libblastrampoline_jll"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[deps.Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[deps.Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "0b5cfbb704034b5b4c1869e36634438a047df065"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.2.1"

[[deps.Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[deps.PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "ColorTypes", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "ae6145ca68947569058866e443df69587acc1806"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.32"

[[deps.Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[deps.Random]]
deps = ["SHA", "Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[deps.Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[deps.Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[deps.Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[deps.Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[deps.UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[deps.Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[deps.p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─59ed6c40-8025-11ec-0ace-d5b826d81255
# ╟─c3b89eff-c3e6-47d0-acc8-661dde65305c
# ╟─a58c0b6f-c5f4-43e9-87ad-4eca62157ce4
# ╠═30233243-ed9d-4bfd-a9af-fcde96677998
# ╟─15717659-eb1a-40d4-abfc-7d680f4722c8
# ╠═9cd7ec78-75dc-4cfc-a4ef-2fd8e9a8d40c
# ╟─dae03614-aa84-483a-8e54-5f493a5a1b82
# ╠═6d1bb7e6-6cff-4cee-873d-ccc82f4f0be8
# ╟─d85c4c9a-3470-4b5a-95af-f2e9c2b04c00
# ╠═a8e3433a-a9d3-4081-a8fe-5fa8732c8d71
# ╟─9bde2570-6ce5-4ad7-8b0b-dd7889f8bea1
# ╠═f425304c-29c1-40eb-b606-9dc39a62e252
# ╟─9db7047d-d900-403c-9a10-85328d5faf44
# ╠═560ddf25-33e0-4075-b1f4-475f308c99a7
# ╟─ff1774ca-7a03-474a-80c7-ff2bd9f38e61
# ╟─e13cae8d-e8ad-4ed8-a844-e1cedb452493
# ╠═079ece15-7872-4e24-b762-418152f232e6
# ╟─37887043-b092-4d0f-bbc7-4288c4c761cd
# ╠═07fb97a3-60f0-40e5-befe-5d27a0b5d597
# ╟─5efd2a9e-7bc9-424e-812d-2fd554677b55
# ╠═6647ab87-0ff0-4d61-ad10-897942b461dc
# ╠═0b438e45-fa85-435b-9319-741379f8bf2d
# ╟─a9ef31b2-9c59-4ed4-ba78-173a383211e2
# ╠═d869d9fd-703b-475b-b9c2-337cbc1b967a
# ╠═09eb599f-f454-4167-8b63-002e6b3b1ef9
# ╟─5c30cf8d-1295-4e42-ba93-2b49b005455c
# ╠═d96442d4-eaa9-4339-b55f-e48e39169e48
# ╟─8e1fac94-4a44-499b-a49b-e4761549b166
# ╟─0536d875-6f77-4834-9fa4-a0a31f7fa09a
# ╟─d54f80fa-750b-494b-be58-933d375fb8c2
# ╠═8bd0408e-7b00-419c-930c-8e9361ad1574
# ╟─3b2e4202-19be-4de5-92ae-b3911af42f7f
# ╟─3b3cab5a-c5e2-4587-9aae-a97cbc15f238
# ╟─c94ca702-fe13-468f-938a-be9b16a879e2
# ╠═b577d259-b564-47c8-b732-a0af01f30014
# ╠═24dea8be-3930-4049-9a44-b8151e164d91
# ╠═7701def7-26b9-4f53-90a9-3eec7be70a16
# ╟─d1571780-39a7-4d8d-9ee0-b9770e235bdd
# ╠═98457248-d957-4376-8086-6b5866a743c9
# ╠═0553623d-5416-4cf4-8cd8-a836dbe5c41f
# ╟─cf556b7b-1d94-43d7-99f8-0a134bd9865b
# ╠═46bdba66-e0fa-4631-b0a1-9d532aa25e89
# ╟─6debf09f-ff30-4e30-93b3-8872c258908c
# ╠═28821165-3106-4983-be9f-3155cbc4ba53
# ╟─b221fbd3-c519-4c9c-8c98-bcee6fae796d
# ╟─3b5b7dcf-6a56-49cb-9dd8-1b2aba81ef87
# ╟─8f398c3a-9def-4f41-a171-45d38527c83c
# ╟─a39e4255-43ee-4ed3-8282-fd25ca0b8d53
# ╠═214cde58-aa66-4499-9a5e-5a6a36c22115
# ╟─7aa62fe2-9c5e-4640-8960-21a9084e436f
# ╠═687f65ad-2320-4cab-b576-0b94de505c10
# ╠═55e690c8-14c2-46f8-bd3c-f64e34a8d74e
# ╟─58fbd165-41f3-4cac-aaa3-e498e4621d1d
# ╠═a6fc5004-5ef2-409c-9b6a-5e7e7efb25ad
# ╠═c4d8af83-f7f9-47a0-9023-fe46e6bc1109
# ╠═6229a275-b52d-4c81-a05b-f22d15b0cf32
# ╠═53cd90d5-8a70-43a6-b72e-d61185c0b7aa
# ╟─4fdc2f91-dc97-4c58-9cc1-486d30ad7890
# ╟─08924e9f-27b8-444a-b9af-9fa34cf1d61b
# ╠═daa23583-9b16-4a76-bde8-9324f3dc1908
# ╠═12e389d6-9ac2-49b5-a678-cf71787541eb
# ╟─e791a417-a00a-479e-b433-3773dd74a48e
# ╠═b18f9af8-130d-4bac-ba8f-230e89d229b7
# ╠═26de3d3b-9229-4c34-a7d1-e0c89ef6f843
# ╠═30933711-723e-4b8f-8be9-3928635e068e
# ╟─eb9a5d9f-9f39-401c-aa35-102414ed99de
# ╟─3280db88-0d68-4426-9f89-0fd3914e6ffa
# ╠═4fa1d5b9-36a0-491e-a034-c887ee7338d8
# ╠═68ce440d-53f8-4312-8911-a7348e602a1f
# ╠═1be992a6-fc03-437f-9dc1-6072f0b08327
# ╠═36086001-a2bd-4648-b2b3-3d526f0dcf71
# ╟─5ece33d6-4622-4385-b84c-639bd7571915
# ╠═ce690520-a919-4e3d-a201-ec2baedf913d
# ╟─28263f4f-37b0-43fd-9eab-6169d82b0581
# ╠═421ab876-efa5-4a7d-9d33-5b9eca5f528e
# ╠═142b9281-196f-4d3e-ae7d-fbc1011255f6
# ╠═91e7e129-c7b5-4e1f-9025-bdac4d569b46
# ╠═0e20de93-a6f1-415c-a08e-53de67065bc4
# ╠═be7bad99-185a-43ed-9c8d-33e5abcdbbd0
# ╟─40b0919d-64a5-4a89-95fb-5266c2c07260
# ╠═ec0fa200-0c79-467a-af5c-97e79f37253d
# ╟─55633c9d-2c1b-40c5-9475-48f18b1f8d77
# ╟─69e96bb3-fc25-4f44-89cd-8f2aeb7065fe
# ╠═9d4aeb68-6649-42f6-8a93-997a136499d4
# ╠═2a3d5a98-751c-48c9-b102-89bc40e74251
# ╠═457bc2c7-9aae-4b28-b216-ef9454227651
# ╟─f35046c4-4a6e-42e5-b691-78b6ef3b0d27
# ╠═8af669a5-4f1f-4f0e-9de7-030e3196f96b
# ╟─32035322-60ec-4ed7-a340-84c6457a2464
# ╟─2c04e557-4afa-4acc-ace1-45f994aa3ed3
# ╟─83c37801-d248-442a-a360-829f740f5cd9
# ╠═4170966b-cbb4-4b7b-b6e0-528f63a1c019
# ╟─f720b6ad-eea0-4c07-b528-822b80305000
# ╟─7831cafd-8ac8-41be-ad0d-a5abeeb66276
# ╟─fc60372a-f853-4005-bc40-64e9d0615a32
# ╠═89b74238-dd21-4063-bdb8-d64adcc2184c
# ╠═b8c1cdf5-550e-49be-83b0-2a01fc583886
# ╟─1b7789aa-1969-486c-a733-6e3d10d135bc
# ╟─f42fa416-47d1-49b6-9d4b-365740908e03
# ╟─4f191e04-1972-47af-b1d9-2de4857833bb
# ╟─342337cc-6202-45f7-bf5d-6688b770ac75
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


using OPCSPs.MVNTools

using PyCall
using PyPlot
@pyimport matplotlib.patches as patches

function plot_confidence_ellipse(d::MVN, num_sigma::Int=3)
    assert(length(d.mean)==2)
    lambda, phi = eig(d.covariance)
    ell = patches.Ellipse(xy = d.mean,
                          width = sqrt(lambda[1])*num_sigma*2,
                          height = sqrt(lambda[2])*num_sigma*2,
                          angle = 180/pi*atan2(phi[2,1], phi[1,1]))
    ell[:set_facecolor]("none")
    gca()[:add_artist](ell)
    draw()
end

d = MVN([0,1], [1 0.08; 0.08 0.04])

n_samples = 10000
samples = Array(Float64, n_samples, 2)
rng = MersenneTwister()
for i in 1:n_samples
    samples[i,:] = rand!(rng, Array(Float64, length(d)), d)
end

figure(1)
clf()
plot_confidence_ellipse(d, 1)
plot_confidence_ellipse(d, 2)
plot_confidence_ellipse(d, 3)
xlim(-3, 3)
ylim(-3, 3)
scatter(samples[:,1], samples[:,2])


# Helper functions for nMDS based on # https://link.springer.com/content/pdf/10.1007/BF02289694.pdf
# Varibales are generally named simillaryly to thos in the original paper. However
# it felt simples to have a a vector of Is and one of Js, rather than an IJ array.
# The complete algorithm, using these functions can be found in main.jl

#### Gradient calculations (for the Euclidean case) pp.126 of Kruskal(1964)

# There are 3 terms within the sum on the right hand side of the equation
# The first term simply denotes which values affect which elements of X2 and
# whether to add or subtract them. The second and third terms can be calculated
# separately to form a length M vector and an M by L arrau respectively. Then we
# loop over L and M, using the indeces in Is and Js to add t appropriate values to X2.
# We multiply by minus S, since we are interested in the negative gradient.

# Second term
function grhs2!(d, d̂, S_star,T_star,M,res)
    @turbo for m in 1:M
        res[m] = (d[m]-d̂[m])/S_star - d[m]/T_star
    end
end
# Third term
function grhs3!(x,d, M,L,res,Is,Js)
    Imat = view(x,Is,:)
    Jmat = view(x,Js,:)
    @turbo for m in 1:M
        for l in 1:L
            res[m,l] = (Imat[m,l] - Jmat[m,l])/d[m]
        end
    end
end
# put them together
function g(grhs2,grhs3,L,M,S,Is,Js,X2,grhs)
    @inbounds for l in 1:L
       grhs .= grhs2 .* view(grhs3,:,l)
       for m in 1:M
           X2[Is[m],l] += grhs[m]
           X2[Js[m],l] -= grhs[m]
       end
   end
   return X2 *-S
end

function checkup(start,block,partition,d,M)
    println("up")

    ∑ = 0
    i = start
    while i < M
    #while partition[i] == block
        println(i,"","hi")
        if partition[i] !== block
            break
        end
        ∑+=d[i]
        i+=1
    end
    return ∑/(i -start), i-1
end
function checkdown(start,block,partition,d)
    println("down")
    ∑ = 0
    i = start
    while  i>0
        println(i,"","hip")
        if partition[i] !== block
            break
        end
        ∑+=d[i]
        i-=1

    end
    return ∑/(start-i), i+1
end

#### Monotone rgression pp.126-128 Kruscal(1964)
#The R implementation monMDS also explains this procedure https://github.com/vegandevs/vegan/blob/master/src/monoMDS.f

# The aim is to sort the distances (in DIST) into blocks, such that the mean (d̂) of each block increases
# from first to last i.e. for any given block - the block before it will have a lower d̂ and the one after it will
# have a greater d̂. In Kruskals terms, each block must be downsatisfied (block before it has lower d̂) and up
# satisfied (block after it has higher d̂).
# To achieve this we start with each element of d having it's own singleton block.
# Starting with the first block, we alternate between checking whether the block is upsatisfied and downsatisfied, whenever the block
# is not satisfied for a given direction we merge it with the adjacent block in that direction and begin
# checking that this block is up and down satified. If the current block is both up and down satisfied,
# we move on to the next block and do the same. This procedure continues until all blocks are up and down
# satisfied.
# The fitted values returned to DHAT are the d̂ values associated with the block each DIST value belongs to.

function monotone!(d,d̂,M, partition = false)
    loc = 1
    # if a starting partition is supplied e.g. because of ties, use it
    partition = partition == false ? collect(1:M) : partition
    block = view(d,partition .==loc)
    localmean = d[1]
    bounds = (1,1)
    #start at position 1, up-active i.e. check to see that the block to the right has a higher d̂
    while true#(loc < maximum(partition)) | (localmean < checkdown(bounds[1],loc-1,partition,d)[1])
    println(loc)
    println(bounds)
    println(partition[1:10])
    if loc < maximum(partition)
        upmean, upbound = checkup(bounds[2]+1,loc+1,partition,d,M)
        println(localmean,"",upmean)
        if upmean > localmean # if up satissfied
            if loc > 1 # and not first block
                downmean, downbound = checkdown(bounds[1]-1,loc-1,partition,d)#check down
                if downmean < localmean # if down satisfied also
                    loc += 1 # move to next block
                    localmean = upmean
                    bounds = (bounds[2]+1,upbound)

                else # if not down satisfied
                    partition[bounds[1]:end] .-=1
                    localmean = (localmean*(bounds[2] -bounds[1]+1)+downmean*(bounds[1]-downbound))/(bounds[2]-downbound+1)
                    bounds = (downbound,bounds[2]) # merge down
                    loc -= 1
                end
            else #if first block and up satisfied
                loc += 1 # move to next block
                localmean = upmean
                bounds = (bounds[2]+1,upbound)

            end
        else # if not upsatisfied
            partition[bounds[2]+1:end] .-=1
            localmean = (localmean*(bounds[2] -bounds[1]+1)+upmean*(upbound-bounds[2]))/(upbound-bounds[1]+1)
            bounds = (bounds[1],upbound) # merge up
            #loc += 1
            if loc > 1 # and not first block
                downmean, downbound = checkdown(bounds[1]-1,loc-1,partition,d)#check down
              #  if downmean < localmean # if down satisfied also
                 #   loc += 1 # move to next block
                if   downmean > localmean
               # else # if not down satisfied
                    partition[bounds[1]:end] .-=1
                    localmean = (localmean*(bounds[2] -bounds[1]+1)+downmean*(bounds[1]-downbound))/(bounds[2]-downbound+1)
                    bounds = (downbound,bounds[2]) # merge down
                    loc -= 1
                 end
            end
        end
    else # if last block
        println("last block")
        downmean, downbound = checkdown(bounds[1]-1,loc-1,partition,d)#check down
        if downmean < localmean # if down satisfied
            break # exit loop
        else # if not down satisfied
            partition[bounds[1]:end] .-=1
            localmean = (localmean*(bounds[2] -bounds[1]+1)+downmean*(bounds[1]-downbound))/(bounds[2]-downbound+1)
            bounds = (downbound,bounds[2]) # merge down
            loc -= 1
        end
    end

    end

    blockmax = partition[end]
    blocks = 1:blockmax

    for block in blocks
        boolmask = partition .== block
        d̂[boolmask] .= mean(d[boolmask])
    end
end

# normalize point coordinates
function normalize(x,dim)
    x .-= mean(x,dims = dim)
    return x ./std(x,dims = dim)
end

#### Step size related function pp.120-122 of Kruskal(1964)
cosθ(g,g″) = sum(g .* g″)/(sqrt(sum(g .^2))*sqrt(sum(g″ .^2)))
angle_factor(g,g″) = 4.0^(cosθ(g,g″)^3.0)
relaxation_factor(x) = 1.3/(1+(five_step_ratio(x)^5))
five_step_ratio(x) = min(1,x[end]/x[1])
good_luck_factor(x) = min(1,x[end]/x[end-1])
mag(g,K) = sqrt(sum(g .^2)/K)
x′(x,g,α,K) = x .+ g*α/mag(g,K)

# Calculate the stress (S), S⋆ an T⋆ according to Kruskals stress formula pp.115 & 125 Kruskal(1964)
function stress(d,d̂,M)
    S_star = 0
    T_star = 0
    @turbo for m in 1:M
        S_star += (d[m] - d̂[m])^2
        T_star += d[m]^2
    end
    return S_star, T_star, sqrt(S_star/T_star)
end
# calculate Euclidean distances between points
function get_dist!(X1,DIST,Is,Js,M)
    @inbounds for m in 1:M
        DIST[m] = euclidean(X1[Is[m],:],X1[Js[m],:])
    end
end


function lgd(n)
#
#  test polynomial suggested by Goedecker
#   Legendre polynomial of degree n
#
    p0 = [1.0]
    p1 = [1.0; 0]
    if n == 0
        p = p0
    elseif n == 1
        p = p1
    elseif n > 1
        for k = 1:n-1
            q1 = (2*k+1)*[p1;0]
            m = length(q1)-length(p0)
            q0 = k*[zeros(m);p0]
            p = (q1-q0)/(k+1)
            #
            p0 = p1
            p1 = p
        end
    end

    p
end


# test initilization functions

for T in (Float64, Float32)

    X = rand(T, 8, 12)

    # randinit

    W, H = NMF.randinit(X, 5)
    @test size(W) == (8, 5)
    @test size(H) == (5, 12)
    @test all(W .>= zero(T))
    @test all(H .>= zero(T))

    W, H = NMF.randinit(X, 5; zeroh=true)
    @test size(W) == (8, 5)
    @test size(H) == (5, 12)
    @test all(W .>= zero(T))
    @test all(H .== zero(T))

    W, H = NMF.randinit(X, 5; normalize=true)
    @test size(W) == (8, 5)
    @test size(H) == (5, 12)
    @test all(W .>= zero(T))
    @test all(H .>= zero(T))
    @test vec(sum(W, dims=1)) ≈ ones(5)

    ## nndsvd

    Random.seed!(5678)
    W, H = NMF.nndsvd(X, 5)
    @test size(W) == (8, 5)
    @test size(H) == (5, 12)
    @test all(W .>= zero(T))
    @test all(H .>= zero(T))

    Random.seed!(5678)
    W2, H2 = NMF.nndsvd(X, 5; zeroh=true)
    @test size(W) == (8, 5)
    @test size(H) == (5, 12)
    @test all(W2 .== W)
    @test all(H2 .== zero(T))

    W, H = NMF.nndsvd(X, 5; variant=:ar)
    @test all(W .> zero(T))
    # NMF.printf_mat(W)
end


module TestIDRs

using IterativeSolvers
using Test
using Random
using LinearAlgebra
using SparseArrays

@testset "IDR(s)" begin

n = 10
m = 6
Random.seed!(1234567)

@testset "Matrix{$T}" for T in (Float32, Float64, ComplexF32, ComplexF64)
    A = rand(T, n, n) + n * I
    b = rand(T, n)
    reltol = √eps(real(T))

    @testset "Without residual smoothing" begin
        x, history = idrs(A, b, reltol=reltol, log=true)
        @test isa(history, ConvergenceHistory)
        @test history.isconverged
        @test norm(A * x - b) / norm(b) ≤ reltol
    end

    # with smoothing
    @testset "With residual smoothing" begin
        x, history = idrs(A, b; smoothing=true, log=true)
        @test history.isconverged
        @test norm(A*x - b) / norm(b) ≤ reltol
    end
end

@testset "SparseMatrixCSC{$T, $Ti}" for T in (Float64, ComplexF64), Ti in (Int64, Int32)
    A = sprand(T, n, n, 0.5) + n * I
    b = rand(T, n)
    reltol = √eps(real(T))

    x, history = idrs(A, b, log=true)
    @test history.isconverged
    @test norm(A * x - b) / norm(b) ≤ reltol
end

@testset "Maximum number of iterations" begin
    x, history = idrs(rand(5, 5), rand(5), log=true, maxiter=2)
    @test history.iters == 2
    @test length(history[:resnorm]) == 2
end

@testset "Near solution (#222)" begin
    x = rand(5)
    A = rand(5, 5)
    b = rand(5)

    x, history = idrs!(x, A, b, log=true)
    x_new = copy(x)
    x_new, history = idrs!(x_new, A, b, log=true)

    @test x_new ≈ x
end

@testset "Termination criterion" begin
    for T in (Float32, Float64, ComplexF32, ComplexF64)
        A = T[ 2 -1  0
              -1  2 -1
               0 -1  2]
        n = size(A, 2)
        b = ones(T, n)
        x0 = A \ b
        perturbation = 10 * sqrt(eps(real(T))) * T[(-1)^i for i in 1:n]

        # If the initial residual is small and a small relative tolerance is used,
        # many iterations are necessary
        x = x0 + perturbation
        initial_residual = norm(A * x - b)
        x, ch = idrs!(x, A, b, log=true)
        @test 2 ≤ niters(ch) ≤ n

        # If the initial residual is small and a large absolute tolerance is used,
        # no iterations are necessary
        x = x0 + perturbation
        initial_residual = norm(A * x - b)
        x, ch = idrs!(x, A, b, abstol=2*initial_residual, reltol=zero(real(T)), log=true)
        @test niters(ch) == 0
    end
end

end

end # module


using StatsBase, Distances, Printf

export DTW, FixedPenaltyDTW, FixPenDTW

struct DTW <: SemiMetric
    ground_dist::Metric
end

function (d::DTW)(
    S1::Vector{T}, S2::Vector{T}
    ) where {T}
    if length(S1) < length(S2)  # This ensures first seq is longest
        d(S2, S1)
    else
        @assert (length(S1)>0) & (length(S2)>0) "both args must be either of type Nothing or of nonzero length."
        d_g = d.ground_dist
        prev_row = pushfirst!(fill(Inf, length(S2)), 0.0);
        curr_row = fill(Inf, length(S2) + 1);

        for i = 1:length(S1)
            # curr_row[1] = prev_row[1]
            for j = 1:(length(S2))
                # @show i, j, prev_row[j], d.ground_dist(S1[i], S2[j])
                cost = d_g(S1[i],S2[j])
                curr_row[j+1] = cost + min(prev_row[j], prev_row[j+1], curr_row[j])
            end
            # @show curr_row
            copy!(prev_row, curr_row)
        end
        return curr_row[end]
    end
end

function (d::DTW)(
    S1::Nothing, S2::Vector{T}
    ) where {T}

    d_g = d.ground_dist
    z = 0.0
    for x in S2
        z += d_g(x, nothing)
    end
    return z + length(S2)

end
(d::DTW)(S1::Vector{T}, S2::Nothing) where {T} = d(S2,S1)
(d::DTW)(S1::Nothing, S2::Nothing) where {T} = 0.0

function print_info(
    d::DTW,
    S1::Vector{T}, S2::Vector{T}
    ) where {T}

    d = d.ground_dist
    # First find the substitution matrix
    C = fill(Inf, length(S1)+1, length(S2)+1)
    C[1,1] = 0.0

    for j in 1:length(S2)
        for i in 1:length(S1)
            cost = d(S1[i], S2[j])
            C[i+1,j+1] = cost + min(C[i+1,j], C[i,j+1], C[i,j])
        end
    end
    # Now retrace steps to determine an optimal matching
    i, j = size(C)
    pairs = Tuple{Int,Int}[]
    pushfirst!(pairs, (i-1,j-1))
    while (i ≠ 2) | (j ≠ 2)
        i_tmp, j_tmp = argmin(view(C, (i-1):i, (j-1):j)).I
        i = i - 2 + i_tmp
        j = j - 2 + j_tmp
        pushfirst!(pairs, (i-1,j-1))
    end
    # @show outputs
    title = "Optimal Coupling"
    println(title)
    println("-"^length(title), "\n")
    # for statement in outputs
    #     println(statement)
    # end
    i_tmp,j_tmp = (0,0)
    for (i,j) in pairs
        if i == i_tmp
            println(" "^length(@sprintf("%s",S1[i])) * " ↘ $(S2[j])")
        elseif j == j_tmp
            println("$(S1[i]) ↗ " * " "^length(@sprintf("%s",S2[j])))
        else
            println("$(S1[i]) → $(S2[j])")
        end
        i_tmp, j_tmp = (i,j)
    end

end

# Penalised
# ---------

struct FixedPenaltyDTW <: Metric
    ground_dist::Metric
    ρ::Real
end

const FixPenDTW = FixedPenaltyDTW

function (d::FixPenDTW)(
    S1::Vector{T}, S2::Vector{T}
    ) where {T}
    if length(S1) < length(S2)  # This ensures first seq is longest
        d(S2, S1)
    else
        d_g = d.ground_dist
        prev_row = pushfirst!(fill(Inf, length(S2)), 0.0);
        curr_row = fill(Inf, length(S2) + 1);

        for i = 1:length(S1)
            # curr_row[1] = prev_row[1]
            for j = 1:(length(S2))
                # @show i, j, prev_row[j], d.ground_dist(S1[i], S2[j])
                cost = d_g(S1[i],S2[j])
                curr_row[j+1] = cost + min(
                    prev_row[j], # New pair (no warping)
                    prev_row[j+1] + d.ρ, # Warping
                    curr_row[j] + d.ρ)  # Warping
            end
            # @show curr_row
            copy!(prev_row, curr_row)
        end
        return curr_row[end]
    end
end

function (d::FixPenDTW)(
    S1::Nothing, S2::Vector{T}
    ) where {T}

    d_g = d.ground_dist
    z = 0.0
    for x in S2
        z += d_g(x, nothing)
    end
    return z + length(S2) * d.ρ

end
(d::FixPenDTW)(S1::Vector{T}, S2::Nothing) where {T} = d(S2,S1)
(d::FixPenDTW)(S1::Nothing, S2::Nothing) where {T} = 0.0#

function print_info(
    d::Union{FixPenDTW},
    S1::Vector{T}, S2::Vector{T}
    ) where {T}

    d_g = d.ground_dist
    # First find the substitution matrix
    C = fill(Inf, length(S1)+1, length(S2)+1)
    C[1,1] = 0.0

    for j in 1:length(S2)
        for i in 1:length(S1)
            cost = d_g(S1[i], S2[j])
            C[i+1,j+1] = cost + min(C[i+1,j] + d.ρ, C[i,j+1] + d.ρ, C[i,j])
        end
    end
    # Now retrace steps to determine an optimal matching
    i, j = size(C)
    pairs = Tuple{Int,Int}[]
    warp_shift_mat = [0 d.ρ; d.ρ 0] # Extra bit here different from DTW
    pushfirst!(pairs, (i-1,j-1))
    while (i ≠ 2) | (j ≠ 2)
        i_tmp, j_tmp = argmin(view(C, (i-1):i, (j-1):j) + warp_shift_mat).I
        i = i - 2 + i_tmp
        j = j - 2 + j_tmp
        pushfirst!(pairs, (i-1,j-1))
    end
    # @show outputs
    title = "Fixed Penalty DTW Print-out with $d_g Ground Distance and ρ=$(d.ρ)"
    println(title)
    println("-"^length(title), "\n")
    println("The cheapest way to do the tranformation...\n")
    println(S1, "---->", S2)
    println("\n...is the following series of edits...\n")
    # for statement in outputs
    #     println(statement)
    # end
    i_tmp,j_tmp = (0,0)
    for (i,j) in pairs
        if i == i_tmp
            println(" "^length(@sprintf("%s",S1[i])) * " ↘ $(S2[j])")
        elseif j == j_tmp
            println("$(S1[i]) ↗ " * " "^length(@sprintf("%s",S2[j])))
        else
            println("$(S1[i]) → $(S2[j])")
        end
        i_tmp, j_tmp = (i,j)
    end
end


using Shearlab, FFTW, Plots
function reorderCoeff(A::AbstractArray{T,2}) where T
    if size(A,1)==1
        includingZeros = reshape(sort(abs.(A), dims =
                                      length(size(A)),rev=true), (length(A,)))
    else
        includingZeros = sort(abs.(A), dims = length(size(A)),rev=true)
    end
    firstZero = findfirst(includingZeros.==0)
    if typeof(firstZero)<:Nothing
        return includingZeros
    else
        return includingZeros[1:firstZero]
    end
end
reorderCoeff(A::AbstractArray{T,1}) where T = sort(abs.(A), rev=true)

N = 32
xlocs = range(-1, 1, length=32)
rot(θ) = [cos(θ) -sin(θ); sin(θ) cos(θ)]
x=1; y=0; θ=π/3; σx = .4;σy = .1
f(xloc,yloc,σx,σy,θ) = [exp.(-1/2*[x,y]'*rot(θ)'*([σx^2 0; 0 σy^2])^(-1)*rot(θ)*[x,y]) for x in xloc, y in yloc ]

gaussianExample = f(xlocs,xlocs,.6,.3,π/6)
heatmap(gaussianExample)

shears = Shearlab.getshearletsystem2D(n, n, 4,typeBecomes=Float32)
size(shears.shearlets)
sheared = Shearlab.sheardec2D(Float32.(gaussianExample), shears)[:,:,1:(end-1)]

mutilated = copy(gaussianExample);
nZeros =floor(Int, 12/32*n); mutilated[1:nZeros,:] = zeros(nZeros,size(mutilated,2))
heatmap(mutilated)

mutsheared = Shearlab.sheardec2D(Float32.(mutilated), shears)[:,:,1:(end-1)]

reorderedShear = reorderCoeff(reshape(sheared,:))
reorderedMutShear = reorderCoeff(reshape(mutsheared,:))
pyplot()
plot(reorderedShear)
plot(log.(1:length(reorderedShear)), max.(-25, log.([reorderedShear reorderedMutShear])), labels=["original", "mutilated"])
savefig("mutilated_Gaussian_Decay_Shearlets_32_4Scales.pdf")


function smoother_kernel(Gs::T, Gf, Gpred, Φ) where {T}
    xs, Ps = meancov(Gs)
    xf, Pf = meancov(Gf)
    xpred, Ppred = meancov(Gpred)

    J = Pf*Φ'/Ppred # C/(C+w)
    xs = xf + J*(xs - xpred)
    Ps = Pf + J*(Ps - Ppred)*J'

    T(xs, Ps)
end

struct RauchTrungStriebel{T,S}
    iter::T # iterate backwards through
    M::S
end

function dyniterate(RTS::RauchTrungStriebel, ::Nothing)
    ϕ = dyniterate(RTS.iter, nothing)
    ϕ === nothing && return nothing
    (n, U), state = ϕ
    Gs, Gpred = U[1], U[2]
    n => Gs, (Gs, Gpred, state)
end

function dyniterate(RTS::RauchTrungStriebel, (Gs, Gpred, state))
    ϕ = dyniterate(RTS.iter, state)
    ϕ === nothing && return nothing
    (n, U), state = ϕ
    Gf, Gpredᵒ = U[1], U[2]
    Gs = smoother_kernel(Gs, Gf, Gpred, Phi(RTS.M.sys, n=>U))
    n => Gs, (Gs, Gpredᵒ, state)
end

function rts_smoother(M, prior, Y)
    P = filter(Y, M)

    ϕ = dyniterate(P, Start(prior))
    ϕ === nothing && error("no observations")
    (t, u), state = ϕ

    Xf = trajectory((t => (u[1], u[2]),))
    while true
        ϕ = dyniterate(P, state)
        ϕ === nothing && break
        (t, u), state = ϕ
        push!(Xf, t => (u[1], u[2])) # safe only what is needed
    end
    ll = u[3]

    n = length(Xf)
    U = Xf.x[end]
    Gf, Gpred = U[1], U[2]
    Gs = Vector{typeof(Gf)}(undef, n)
    Gs[end] = Gf
    for i in n-1:-1:1
        U = Xf.x[i]
        Gf, Gpredᵒ = U[1], U[2]
        Gs[i] = smoother_kernel(Gs[i+1], Gf, Gpred, Phi(M.sys, i => U))
        Gpred = Gpredᵒ
    end
    Trajectory(Xf.t, Gs), ll
end


##############################################################################
##
## TODO: dtype == :complex
##       See minFunc code for examples of how this could be done
##
##############################################################################

##############################################################################
##
## Graveyard of Alternative Functions
##
## In Nodedal and Wright, epsilon was calculated as
##
##   epsilon = sqrt(eps())
##
## In the examples I've tested, this strategy is best when x is small
##  and becomes very bad when x is larger
##
## In previous iteration of derivative_numer, epsilon was calculated as
##
##   epsilon = eps(max(1.0, abs(x)))^(1/3)
##
## In the examples I've tested, this strategy is bad when x is small
##  and becomes better when x is larger
##
## The current approach comes from minFunc
##
##    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))
##
## TODO: Perform a systematic comparison across several functions
##       and several ranges of inputs
##
##############################################################################

##############################################################################
##
## Derivative of f: R -> R
##
##############################################################################

function finite_difference{T <: Number}(f::Function,
                                        x::T,
                                        dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    if dtype == :forward
        return (f(x + epsilon) - f(x)) / epsilon
    elseif dtype == :central
        return (f(x + epsilon) - f(x - epsilon)) / (2 * epsilon)
    else
        error("dtype must be :forward or :central")
    end
end
finite_difference(f::Function, x::Number) = finite_difference(f, x, :central)

##############################################################################
##
## Gradient of f: R^n -> R
##
##############################################################################

function finite_difference{T <: Number}(f::Function,
                                        x::Vector{T},
                                        dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    # What is the dimension of x?
    n = length(x)

    # Storage for forward differences
    differential = Array(Float64, n)

    # Iterate over each dimension of the gradient separately.
    # Use xplusdx to store x + dx instead of creating a new vector on each pass.
    # Use xminusdx to store x - dx instead of creating a new vector on each pass.
    if dtype == :forward
        # Establish a baseline value of f(x).
        f_x = f(x)
        xplusdx = copy(x)
        for i = 1:n
            xplusdx[i] = x[i] + epsilon
            differential[i] = (f(xplusdx) - f_x) / epsilon
            xplusdx[i] = x[i]
        end
    elseif dtype == :central
        xplusdx, xminusdx = copy(x), copy(x)
        for i = 1:n
            xplusdx[i], xminusdx[i] = x[i] + epsilon, x[i] - epsilon
            differential[i] = (f(xplusdx) - f(xminusdx)) / (2 * epsilon)
            xplusdx[i], xminusdx[i] = x[i], x[i]
        end
    else
        error("dtype must be :forward or :central")
    end

    # Return the estimated gradient.
    return differential
end
function finite_difference{T <: Number}(f::Function, x::Vector{T})
    finite_difference(f, x, :central)
end

##############################################################################
##
## Jacobian derivative of f: R^n -> R^m
##
##############################################################################

function finite_difference_jacobian{T <: Number}(f::Function, x::Vector{T}, dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    # What is the dimension of x?
    n = length(x)

    # Establish a baseline for f_x
    f_x = f(x)

    # Initialize the Jacobian matrix
    J = zeros(length(f_x), n)

    # Iterate over each dimension of the gradient separately.
    if dtype == :forward
        xplusdx = copy(x)
        for i = 1:n
            xplusdx[i] = x[i] + epsilon
            J[:, i] = (f(xplusdx) - f_x) / epsilon
            xplusdx[i] = x[i]
        end
        return J
    elseif dtype == :central
        xplusdx, xminusdx = copy(x), copy(x)
        for i = 1:n
            xplusdx[i], xminusdx[i] = x[i] + epsilon, x[i] - epsilon
            J[:, i] = (f(xplusdx) - f(xminusdx)) / (2 * epsilon)
            xplusdx[i], xminusdx[i] = x[i], x[i]
        end
        return J
    else
        error("dtype must :forward or :central")
    end
end

##############################################################################
##
## Second derivative of f: R^n -> R
##
##############################################################################

function finite_difference_hessian(f::Function, g::Function, x::Number, dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    if dtype == :forward
        return (g(x + epsilon) - g(x)) / epsilon
    elseif dtype == :central
        return (g(x + epsilon) - g(x - epsilon)) / (2 * epsilon)
    else
        error("dtype must :forward or :central")
    end
end
function finite_difference_hessian(f::Function, x::Number, dtype::Symbol)
    finite_difference_hessian(f, derivative(f), x, dtype)
end
function finite_difference_hessian(f::Function, x::Number)
    finite_difference_hessian(f, derivative(f), x, :central)
end

##############################################################################
##
## Hessian of f: R^n -> R
##
##############################################################################

function finite_difference_hessian{T <: Number}(f::Function, g::Function, x::Vector{T}, dtype::Symbol)
    epsilon = 2 * sqrt(1e-12) * (1 + norm(x))

    # What is the dimension of x?
    n = length(x)

    # Initialize an empty Hessian
    H = Array(Float64, n, n)

    if dtype == :forward
        g_x = g(x)
        xplusdx = copy(x)
        for i in 1:n
            xplusdx[i] = x[i] + epsilon
            H[:, i] = (g(xplusdx) - g_x) / epsilon
            xplusdx[i] = x[i]
        end
        symmetrize!(H)
        return H
    elseif dtype == :central
        xplusdx, xminusdx = copy(x), copy(x)
        for i = 1:n
            xplusdx[i], xminusdx[i] = x[i] + epsilon, x[i] - epsilon
            H[:, i] = (g(xplusdx) - g(xminusdx)) / (2 * epsilon)
            xplusdx[i], xminusdx[i] = x[i], x[i]
        end
        symmetrize!(H)
        return H
    else
        error("dtype must :forward or :central")
    end
end
function finite_difference_hessian{T <: Number}(f::Function, x::Vector{T}, dtype::Symbol)
    finite_difference_hessian(f, gradient(f), x, dtype)
end
function finite_difference_hessian{T <: Number}(f::Function, x::Vector{T})
    finite_difference_hessian(f, gradient(f), x, :central)
end

##############################################################################
##
## Taylor Series based estimates of first and second derivatives
##
## TODO: Fill in multivariate equivalents
##
##############################################################################

# Higher precise finite difference method based on Taylor series approximation.
# h is the stepsize
function taylor_finite_difference(f::Function, x::Float64, h::Float64, dtype::Symbol)
    if dtype == :forward
        f_x = f(x)
        d = 2^3 * (2^2 * (f(x + h) - f_x) - (f(x + 2 * h) - f_x))
        d += - (2^2 * (f(x + 2 * h) - f_x) - (f(x + 4 * h) - f_x))
        d /= 3 * 2^2 * h
    elseif dtype == :central
        d = 4^5 * (2^3 * (f(x + h) - f(x - h)) - (f(x + 2 * h) - f(x - 2 * h)))
        d -= 2^3 * (f(x + 4 * h) - f(x - 4 * h)) - (f(x + 8 * h) - f(x - 8 * h))
        d /= (4^5 * (2^4 - 2^2) - (2^6 - 2^4)) * h
    else
        error("dtype must be :forward or :central")
    end
    return d
end
function taylor_finite_difference(f::Function, x::Float64, h::Float64)
    taylor_finite_difference(f, x, h, :central)
end
function taylor_finite_difference(f::Function, x::Float64, dtype::Symbol)
    taylor_finite_difference(f, x, 10e-4, dtype)
end
function taylor_finite_difference(f::Function, x::Float64)
    taylor_finite_difference(f, x, 10e-4)
end

function taylor_finite_difference_hessian(f::Function, x::Float64, h::Float64)
    f_x = f(x)
    d = 4^6 * (2^4 * (f(x + h) + f(x - h) - 2 * f_x) - (f(x + 2 * h) + f(x - 2 * h) - 2 * f_x))
    d += - (2^4 * (f(x + 4 * h) + f(x - 4 * h) - 2 * f_x) - (f(x + 8 * h) + f(x - 8 * h) - 2 * f_x))
    return d / (3 * 2^6 * (2^8 - 1) * h^2)
end

##############################################################################
##
## TODO: Implement directional_derivative()
##
##############################################################################

# The function "dirderivative" calculates directional derivatives in the direction v.
# The function supplied must have the form Vector{Float64} -> Float64
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, h::Float64, twoside::Bool)
#     derivative(t::Float64 -> f(x0 + v*t) / norm(v), 0.0, h, twoside)
# end
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, h::Float64)
#     dirderivative(f, v, x0, h, true)
# end
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, )
#     derivative(f, v, x0, 0.0001)
# end
# function dirderivative(f::Function, v::Vector{Float64})
#     x -> dirderivative(f, v, x)
# end


"""
    ProjectedPointDistribution(m::Manifold, d, proj!, x)

Generates a random point in ambient space of `m` and projects it to `m`
using function `proj!`. Generated arrays are of type `TResult`, which can be
specified by providing the `x` argument.
"""
struct ProjectedPointDistribution{TResult,TM<:Manifold,TD<:Distribution,TProj} <:
       MPointDistribution{TM}
    manifold::TM
    d::TD
    proj!::TProj
end

function ProjectedPointDistribution(
    M::Manifold,
    d::Distribution,
    proj!,
    x::TResult,
) where {TResult}
    return ProjectedPointDistribution{TResult,typeof(M),typeof(d),typeof(proj!)}(
        M,
        d,
        proj!,
    )
end

function rand(rng::AbstractRNG, d::ProjectedPointDistribution{TResult}) where {TResult}
    x = convert(TResult, rand(rng, d.d))
    return d.proj!(d.manifold, x)
end

function _rand!(rng::AbstractRNG, d::ProjectedPointDistribution, x::AbstractArray{<:Number})
    _rand!(rng, d.d, x)
    return d.proj!(d.manifold, x)
end

support(d::ProjectedPointDistribution) = MPointSupport(d.manifold)

"""
    ProjectedFVectorDistribution(type::VectorBundleFibers, x, d, project_vector!)

Generates a random vector from ambient space of manifold `type.manifold`
at point `x` and projects it to vector space of type `type` using function
`project_vector!`, see [`project_vector`](@ref) for documentation.
Generated arrays are of type `TResult`.
"""
struct ProjectedFVectorDistribution{
    TResult,
    TSpace<:VectorBundleFibers,
    TX,
    TD<:Distribution,
    TProj,
} <: FVectorDistribution{TSpace,TX}
    type::TSpace
    x::TX
    d::TD
    project_vector!::TProj
end

function ProjectedFVectorDistribution(
    type::VectorBundleFibers,
    x,
    d::Distribution,
    project_vector!,
    xt::TResult,
) where {TResult}
    return ProjectedFVectorDistribution{
        TResult,
        typeof(type),
        typeof(x),
        typeof(d),
        typeof(project_vector!),
    }(
        type,
        x,
        d,
        project_vector!,
    )
end

function rand(rng::AbstractRNG, d::ProjectedFVectorDistribution{TResult}) where {TResult}
    v = convert(TResult, reshape(rand(rng, d.d), size(d.x)))
    return d.project_vector!(d.type, v, d.x, v)
end

function _rand!(
    rng::AbstractRNG,
    d::ProjectedFVectorDistribution,
    v::AbstractArray{<:Number},
)
    # calling _rand!(rng, d.d, v) doesn't work for all arrays types
    return copyto!(v, rand(rng, d))
end

support(tvd::ProjectedFVectorDistribution) = FVectorSupport(tvd.type, tvd.x)


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
using AdFem
using Test
Random.seed!(233)

function tpfa_op(kvalue,bc,pval,m,n,h)
    tpfa_op_ = load_op_and_grad("./build/libTpfaOp","tpfa_op", multiple=true)
    kvalue,bc,pval,m_,n_,h = convert_to_tensor(Any[kvalue,bc,pval,m,n,h], [Float64,Int64,Float64,Int64,Int64,Float64])
    ii, jj, vv, rhs = tpfa_op_(kvalue,bc,pval,m_,n_,h)
    SparseTensor(ii + 1, jj + 1, vv, m*n, m*n), set_shape(rhs, m*n)
end

# TODO: specify your input parameters
m = 20
n = 10
h = 0.1
kvalue = rand(m*n)
bc = bcedge("all", m, n, h)
pval = rand(size(bc, 1))
K, rhs = compute_fvm_tpfa_matrix(kvalue,bc,pval,m,n,h)
K1, rhs1 = tpfa_op(kvalue,bc,pval,m,n,h)
sess = Session(); init(sess)
K2 = run(sess, K1)
rhs2 = run(sess, rhs1)

@test rhs≈rhs2
@test K ≈ K2

# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(x)
    # A, rhs = tpfa_op(kvalue,bc,x,m,n,h)
    # return sum(values(A)^2) + sum(rhs^2)
    A, rhs = tpfa_op(x,bc,pval,m,n,h)
    return sum(values(A)^2) + sum(rhs^2)
end

# TODO: change `m_` and `v_` to appropriate values
# m_ = constant(rand(size(bc,1)))
# v_ = rand(size(bc,1))

m_ = constant(rand(m*n))
v_ = rand(m*n)

y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


module _rMSEn
export rMSEn
using Statistics: std, mean, median, var
using DSP.Filters: filtfilt, Butterworth, Lowpass, digitalfilter
using Plots

    """
        MSx, CI = rMSEn(Sig, Mobj)

    Returns a vector of refined multiscale entropy values (`MSx`) and the complexity
    index (`CI`) of the data sequence (`Sig`) using the parameters specified by
    the multiscale object (`Mobj`) and the following default parameters:
    Scales = 3, Butterworth LPF Order = 6, Butterworth LPF cutoff frequency
    at scale (T): Fc = 0.5/T.
    If the entropy function specified by `Mobj` is SampEn or ApEn, rMSEn
    updates the threshold radius of the data sequence (Xt) at each scale
    to 0.2std(Xt) if no `r` value is provided by Mobj, or r.std(Xt) if `r`
    is specified.

        MSx, CI = rMSEn(Sig::AbstractArray{T,1} where T<:Real, Mobj::NamedTuple; Scales::Int=3,
                            F_Order::Int=6, F_Num::Float64=0.5, RadNew::Int=0, Plotx::Bool=false)

    Returns a vector of refined multiscale entropy values (`MSx`) and the complexity
    index (`CI`) of the data sequence (`Sig`) using the parameters specified by
    the multiscale object (`Mobj`) and the following 'keyword' arguments:

    # Arguments:
    `Scales`   - Number of temporal scales, an integer > 1 (default = 3) \n
    `F_Order`  - Butterworth low-pass filter order, a positive integer
                 (default: 6) \n
    `F_Num`    - Numerator of Butterworth low-pass filter cutoff frequency,
                 a scalar value in range [0 < `F_Num` < 1]. The cutoff frequency
                 at each scale (T) becomes: Fc = `F_Num/T.  (default: 0.5) \n
    `RadNew`   - Radius rescaling method, an integer in the range [1 4].
                 When the entropy specified by `Mobj` is `SampEn` or `ApEn`,
                 `RadNew` allows the radius threshold to be updated at each
                 time scale (Xt). If a radius value is specified by `Mobj` (`r`),
                 this becomes the rescaling coefficient, otherwise it is set
                 to 0.2 (default). The value of `RadNew` specifies one of the
                 following methods:\n
                 [1]    Standard Deviation          - r*std(Xt)\n
                 [2]    Variance                    - r*var(Xt) \n
                 [3]    Mean Absolute Deviation     - r*mean_ad(Xt) \n
                 [4]    Median Absolute Deviation   - r*med_ad(Xt)\n
    `Plotx`    - When `Plotx` == true, returns a plot of the entropy value at each
                 time scale (i.e. the multiscale entropy curve)  [default: false] \n

    # See also `MSobject`, `MSEn`, `cMSEn`, `hMSEn`, `SampEn`, `ApEn`, `XMSEn`

    # References:
        [1] Madalena Costa, Ary Goldberger, and C-K. Peng,
            "Multiscale entropy analysis of complex physiologic time series."
            Physical review letters
            89.6 (2002): 068102.

        [2] Vadim V. Nikulin, and Tom Brismar,
            "Comment on “Multiscale entropy analysis of complex physiologic
            time series”."
            Physical review letters
            92.8 (2004): 089803.

        [3] Madalena Costa, Ary L. Goldberger, and C-K. Peng.
            "Costa, Goldberger, and Peng reply."
            Physical Review Letters
            92.8 (2004): 089804.

        [4] José Fernando Valencia, et al.,
            "Refined multiscale entropy: Application to 24-h holter
            recordings of heart period variability in healthy and aortic
            stenosis subjects."
            IEEE Transactions on Biomedical Engineering
            56.9 (2009): 2202-2213.

        [5] Puneeta Marwaha and Ramesh Kumar Sunkaria,
            "Optimal selection of threshold value ‘r’for refined multiscale
            entropy."
            Cardiovascular engineering and technology
            6.4 (2015): 557-576.

    """
    function rMSEn(Sig::AbstractArray{T,1} where T<:Real, Mobj::NamedTuple; Scales::Int=3,
        F_Order::Int=6, F_Num::Float64=0.5, RadNew::Int=0, Plotx::Bool=false)

    (size(Sig,1)>10) ? nothing : error("Sig:   must be a numeric vector" )
    (length(Mobj) >= 1) ? nothing :  error("Mobj:    must be a multiscale entropy object created
         with the function EntropyHub.MSobject")
    (Scales>1) ? nothing : error("Scales:     must be an integer > 1")
    (F_Order>1) ? nothing :  error("F_Order:     must be an integer > 1")
    (0 < F_Num < 1) ? nothing : error("F_Num:     must be a scalar in range 0 < F_Num < 1")
    (RadNew==0 || (RadNew in 1:4 && String(Symbol(Mobj.Func)) in ("SampEn","ApEn"))) ? nothing :
    error("RadNew:  must be 0, or an integer in range [1 4] with entropy function `SampEn` or `ApEn`")

    MSx = zeros(Scales)
    Args = NamedTuple{keys(Mobj)[2:end]}(Mobj)
    (RadNew==0 && String(Symbol(Mobj.Func)) in ("SampEn","ApEn")) ? RadNew=1 : nothing
    if RadNew > 0
        if RadNew == 1
            Rnew = x -> std(x, corrected=false)
        elseif RadNew == 2
            Rnew = x -> var(x, corrected=false)
        elseif RadNew == 3
            Rnew = x -> mean(abs.(x .- mean(x)))
        elseif RadNew == 4
            Rnew = x -> median(abs.(x .- median(x)))
        end

        if haskey(Mobj,:r)
            Cx = Mobj.r
        else
            Cy = ("Standard Deviation","Variance","Mean Abs Deviation", "Median Abs Deviation")
            @warn("No radius value provided in Mobj.
                Default set to 0.2*$(Cy[RadNew]) of each new time-series.")
            Cx = .2
        end
    end

    for T = 1:Scales
        print(". ")
        Temp = refined(Sig,T,F_Order,F_Num)
        RadNew > 0 ? Args = (Args..., r=Cx*Rnew(Temp[:])) : nothing
        Tempx = Mobj.Func(Temp[:]; Args...)
        #MSx[T] = Tempx[1][end]
        typeof(Tempx)<:Tuple ? MSx[T] = Tempx[1][end] : MSx[T] = Tempx[end]
    end
    CI = sum(MSx)
    print("\n")
    if any(isnan.(MSx))
        println("Some entropy values may be undefined.")
    end

    if Plotx
        p1 = plot(1:Scales, MSx, c=RGB(8/255, 63/255, 77/255), lw=3)
        scatter!(1:Scales, MSx, markersize=6, c=RGB(1, 0, 1),
        xlabel = "Scale Factor", ylabel = "Entropy Value",
        guidefont = font(12, "arial", RGB(7/255, 54/255, 66/255)),
        tickfontsize = 10, tickfontfamily="arial", legend=false,
        title = "Refined Multiscale $(Mobj.Func)",
        plot_titlefontsize=16, plot_titlefontcolor=RGB(7/255, 54/255, 66/255))
        display(p1)
    end

    return MSx, CI
    end

    function refined(Z,sx,P1,P2)
        Yt = filtfilt(digitalfilter(Lowpass(P2/sx), Butterworth(P1)), Z)
        return Yt[1:sx:end]
    end

end

"""

Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub

"""


using Plots
gr(titlefontsize=12,size=(500,300),colorbar=false)

using FourierGPE

# Units: Oscillator
L = (15.,15.,15.)
N = (64,64,64)
sim = Sim(L,N)
@unpack_Sim sim

import FourierGPE.V
V(x,y,z,t) = 0.5*(x^2 + y^2 + z^2)

## set simulation parameters
γ = 0.5
tf = 2/γ
Nt = 200
t = LinRange(0.,tf,Nt)

## useful state functions
ψ0(x,y,z,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,y,z,0.0)/μ,0.0)+im*0.0)
healing(x,y,z,μ,g) = 1/sqrt(g*abs2(ψ0(x,y,z,μ,g)))

## make initial state
x,y,z = X
ψi = ψ0.(x,y',reshape(z,1,1,length(z)),μ,g)
ψi = randn(N)+im*randn(N)
ϕi = kspace(ψi,sim)

@pack_Sim! sim

## Evolve in k space
sol = runsim(sim)
##

# pull out the ground state
ϕg = sol[end]
ψg = xspace(ϕg,sim)
zmid = findfirst(z .≈ 0)
showpsi(x,y,ψg[:,:,zmid])

plot(x,y,abs2.(ψg[:,:,zmid]),linetype=:surface)

## Dynamics of Kohn mode
# kick it
simk = Sim(L,N)

ki = .7
ϕk = copy(ϕg)
ψk = xspace(ϕk,sim)
ψk = ψk.*exp.(im*ki*x)
ϕk = kspace(ψk,sim)
ϕi = copy(ϕk)
γ = 0.0
tf = 2*pi
t = LinRange(0.,tf,Nt)

@pack_Sim! simk

@time solk = runsim(simk)

## animate a slice
anim = @animate for i=1:Nt
    ϕ = solk[i]
    zmid = findfirst(z .≈ 0)
    ψ = xspace(ϕ,simk)[:,:,zmid]
    showpsi(x,y,ψ)
end

gif(anim,"./examples/3dKohn.gif",fps=30)
# gif(anim,"./examples/3dKohn.mov",fps=30)

# adimate slice as a surface
anim = @animate for i=1:Nt
    ϕ = solk[i]
    zmid = findfirst(z .≈ 0)
    ψ = xspace(ϕ,simk)[:,:,zmid]
    plot(x,y,abs2.(ψ),linetype=:surface)
end

gif(anim,"./examples/3dKohnSurface.mov",fps=30)

simk

## animate isosurface in Makie
using Makie, AbstractPlotting

function dense(phi)
    ψm = xspace(phi,sim)
    density = abs2.(ψm)
    pmax = maximum(density)
    return density/pmax
end

function densityfilm(sol,Nt)
    saveto="examples/3dKohn.gif"
    scene = Scene()
    tindex = Node(1)
    scene = volume(lift(i -> dense(sol[i]), tindex),
    algorithm = :iso,
    show_axis=false,
    isovalue=3f0(.15))

    R = 70
    eyepos = Vec3f0(R,R,R)
    lookat = Vec3f0(18,18,0)

    record(scene, saveto, 1:Nt) do i
        update_cam!(scene, eyepos, lookat)
        rotate_cam!(scene, 0., -0.25, 0.)
        tindex[] = i
    end
    p = scene[end];
    return
end

p = densityfilm(sol,Nt)

# p.attributes.attribute


poisson(T, n) = sparse(Tridiagonal(fill(T(-1), n-1),
                        fill(T(2), n), fill(T(-1), n-1)))
poisson(n) = poisson(Float64, n)


module HarmonicOscillatorProblem

    using GeometricIntegrators.Equations

    export harmonic_oscillator_ode, harmonic_oscillator_iode, harmonic_oscillator_pode, harmonic_oscillator_sode,
           harmonic_oscillator_dae, harmonic_oscillator_idae, harmonic_oscillator_pdae

    const Δt = 0.1
    const nt = 10

    const k = 0.5
    const ω = √k

    t₀=0.0
    q₀=[0.5, 0.0]
    z₀=[0.5, 0.0, 0.5]

    function ϑ(q)
        p = zero(q)
        p[1] = q[2]
        p[2] = 0
        return p
    end

    function hamiltonian(t, q)
        q[2]^2 / 2 + k * q[1]^2 / 2
    end

    p₀=ϑ(q₀)

    A = sqrt(q₀[2]^2 / k + q₀[1]^2)
    ϕ = asin(q₀[1] / A)

    refq = A * sin(ω * Δt * nt + ϕ)
    refp = ω * Δt * nt * A * cos(ω * Δt * nt + ϕ)
    refx = [refq, refp]

    function oscillator_ode_v(t, x, v)
        v[1] = x[2]
        v[2] = -k*x[1]
        nothing
    end

    function harmonic_oscillator_ode(x₀=q₀)
        @assert size(x₀,1) == 2
        ODE(oscillator_ode_v, x₀; h=hamiltonian)
    end

    function oscillator_pode_v(t, q, p, v)
        v[1] = p[1]
        nothing
    end

    function oscillator_pode_f(t, q, p, f)
        f[1] = -k*q[1]
        nothing
    end

    function harmonic_oscillator_pode(q₀=[q₀[1]], p₀=[p₀[1]])
        @assert size(q₀,1) == size(p₀,1) == 1
        PODE(oscillator_pode_v, oscillator_pode_f, q₀, p₀)
    end

    function oscillator_sode_v_1(t, q, v, h)
        v[1] = q[1] + h * q[2]
        v[2] = q[2]
        nothing
    end

    function oscillator_sode_v_2(t, q, v, h)
        v[1] = q[1]
        v[2] = q[2] - h * k*q[1]
        nothing
    end

    function harmonic_oscillator_sode(q₀=q₀)
        SODE((oscillator_sode_v_1, oscillator_sode_v_2), q₀)
    end

    function oscillator_iode_ϑ(t, q, p)
        p[1] = q[2]
        p[2] = 0
        nothing
    end

    function oscillator_iode_ϑ(t, q, v, p)
        oscillator_iode_ϑ(t, q, p)
    end

    function oscillator_iode_f(t, q, v, f)
        f[1] = -k*q[1]
        f[2] = v[1] - q[2]
        nothing
    end

    function oscillator_iode_g(t, q, λ, g)
        g[1] = 0
        g[2] = λ[1]
        nothing
    end

    function oscillator_iode_v(t, q, v)
        v[1] = q[2]
        v[2] = -k*q[1]
        nothing
    end

    function oscillator_iode_v(t, q, p, v)
        oscillator_iode_v(t, q, v)
    end

    function harmonic_oscillator_iode(q₀=q₀, p₀=ϑ(q₀))
        @assert size(q₀,1) == size(p₀,1) == 2
        IODE(oscillator_iode_ϑ, oscillator_iode_f,
             oscillator_iode_g, q₀, p₀;
             h=hamiltonian, v=oscillator_iode_v)
    end

    function oscillator_dae_v(t, z, v)
        v[1] = z[2]
        v[2] = -k*z[1]
        v[3] = z[2] - k*z[1]
        nothing
    end

    function oscillator_dae_u(t, z, λ, u)
        u[1] = -λ[1]
        u[2] = -λ[1]
        u[3] = +λ[1]
    end

    function oscillator_dae_ϕ(t, z, λ, ϕ)
        ϕ[1] = z[3] - z[1] - z[2]
    end

    function harmonic_oscillator_dae(z₀=z₀, λ₀=[zero(eltype(z₀))])
        @assert size(z₀,1) == 3
        @assert size(λ₀,1) == 1
        DAE(oscillator_dae_v, oscillator_dae_u, oscillator_dae_ϕ, z₀, λ₀)
    end

    function oscillator_idae_u(t, q, p, λ, u)
        u[1] = λ[1]
        u[2] = λ[2]
        nothing
    end

    function oscillator_idae_g(t, q, p, λ, g)
        g[1] = 0
        g[2] = λ[1]
        nothing
    end

    function oscillator_idae_ϕ(t, q, p, ϕ)
        ϕ[1] = p[1] - q[2]
        ϕ[2] = p[2]
        nothing
    end

    function harmonic_oscillator_idae(q₀=q₀, p₀=ϑ(q₀), λ₀=zero(q₀))
        @assert size(q₀,1) == size(p₀,1) == size(λ₀,1) == 2
        IDAE(oscillator_iode_f, oscillator_iode_ϑ,
             oscillator_idae_u, oscillator_idae_g,
             oscillator_idae_ϕ, q₀, p₀, λ₀;
             v=oscillator_iode_v)
    end

    function oscillator_pdae_v(t, q, p, v)
        v[1] = q[2]
        v[2] = -k*q[1]
        nothing
    end

    function oscillator_pdae_f(t, q, p, f)
        f[1] = -k*q[1]
        f[2] = p[1] - q[2]
        nothing
    end

    function harmonic_oscillator_pdae(q₀=q₀, p₀=ϑ(q₀), λ₀=zero(q₀))
        @assert size(q₀,1) == size(p₀,1) == 2
        PDAE(oscillator_pdae_v, oscillator_pdae_f,
             oscillator_idae_u, oscillator_idae_g,
             oscillator_idae_ϕ, q₀, p₀, λ₀)
    end

end


using DynamicHMCModels, StatsFuns

ProjDir = @__DIR__

delim = ';'
df = CSV.read(joinpath(ProjDir, "..", "..", "data", "chimpanzees.csv"), DataFrame; delim)

Base.@kwdef struct Chimpanzees_02
    "Number of actors"
    N_actors::Int
    pulled_left::Vector{Int}
    prosoc_left::Vector{Int}
    condition::Vector{Int}
    actor::Vector{Int}
end

function make_transformation(model::Chimpanzees_02)
    as((a = as(Vector, model.N_actors), bp = asℝ, bpC = asℝ))
end

model = Chimpanzees_02(; N_actors = maximum(df.actor), pulled_left = df.pulled_left,
                    prosoc_left = df.prosoc_left, condition = df.condition,
                    actor = df.actor)

function (model::Chimpanzees_02)(θ)
    @unpack pulled_left, prosoc_left, condition, actor = model
    @unpack a, bp, bpC = θ
    ℓ_likelihood = mapreduce(+, actor, condition, prosoc_left,
       pulled_left) do actor, condition, prosoc_left, pulled_left
           p = logistic(a[actor] + (bp + bpC * condition) * prosoc_left)
           logpdf(Bernoulli(p), pulled_left)
       end
    P = Normal(0, 10)
    ℓ_prior = logpdf(P, bpC) + logpdf(P, bp) + sum(a -> logpdf(P, a), a)
    ℓ_prior + ℓ_likelihood
end

P = TransformedLogDensity(make_transformation(model), model)
∇P = ADgradient(:ForwardDiff, P)
results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000)
posterior = P.transformation.(results.chain)

p = as_particles(posterior)
display(p)

# Result rethinking

rethinking = "
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
        Mean        SD       Naive SE       MCSE      ESS
a.1 -0.74503184 0.26613979 0.0042080396 0.0060183398 1000
a.2 10.77955494 5.32538998 0.0842018089 0.1269148045 1000
a.3 -1.04982353 0.28535997 0.0045119373 0.0049074219 1000
a.4 -1.04898135 0.28129307 0.0044476339 0.0056325117 1000
a.5 -0.74390933 0.26949936 0.0042611590 0.0052178124 1000
a.6  0.21599365 0.26307574 0.0041595927 0.0045153523 1000
a.7  1.81090866 0.39318577 0.0062168129 0.0071483527 1000
 bp  0.83979926 0.26284676 0.0041559722 0.0059795826 1000
bpC -0.12913322 0.29935741 0.0047332562 0.0049519863 1000
";


export rotation2D

"""
getQ(xc)

builds basis for fast computation of rigid2D
"""
function getQ(xc::Array{Float64})
	n  = round(Int,length(xc)/2)
	xc = reshape(xc,n,2)
	return  kron(speye(2), [xc ones(n)])
end

"""
rotation2D(w,x,doDerivative,Q)

returns rotated points x and derivatives.
"""
function rotation2D(w::Float64,x::Array{Float64};doDerivative::Bool=false,Q::AbstractArray{Float64,2}=getQ(x),c::Vector{Float64}=zeros(2))

R  = [ cos(w) -sin(w);sin(w) cos(w)]
g  = (eye(2)-R)*vec(c)
f  = [R[1,1];R[1,2];g[1];R[2,1];R[2,2];g[2]]
y  = Q*f

if doDerivative
	dR = [-sin(w[1]) -cos(w[1]);cos(w[1]) -sin(w[1])]
	dg = -dR*vec(c)
	df = [dR[1,1];dR[1,2];dg[1];dR[2,1];dR[2,2];dg[2]]
	dy = Q*df
	return y,dy
else
	return y
end
end


# Original example:
# http://gnuplot.sourceforge.net/demo/hidden2.html

using Gnuplot

x = LinRange(-10, 10, 25)
y = LinRange(-10, 10, 25)

@gsp    "set xyplane at 0"
@gsp :- "unset key" :-
@gsp :- "set palette rgbformulae 31,-11,32"
@gsp :- "set style fill solid 0.5"
@gsp :- "set cbrange [-1:1]"
@gsp :- title="Mixing pm3d surfaces with hidden-line plots"
@gsp :- "set hidden3d front"

f = [sin(-sqrt((x+5)^2+(y-7)^2)*0.5) for x in x, y in y]
z = [x*x-y*y for x in x, y in y]
@gsp :- x y f "w pm3d" x y z "w l lc rgb 'black'"


using Heuristics
using BenchmarkTools

# Schwefel Function
function schaffer(x)
    obj = 0.5 + (sin(x[1]^2 + x[2]^2)^2 - 0.5)/(1 + 0.001*(x[1]^2+x[2]^2))^2
    return obj
end

function waveDrop(x)
    obj = -(1 + cos(12*sqrt(x[1]^2 + x[2]^2)))/(0.5*(x[1]^2 + x[2]^2) + 2.0)
    return obj
end

# Setup Problem
d = 2
LB = -5.12*ones(d)
UB = 5.12*ones(d)

prob = Problem(waveDrop, LB, UB)
pso  = PSO(prob; numParticles = 100)

# optimize
opts = Options(;display = false, maxStallIters = 100)
@benchmark optimize!($pso, $opts)
#optimize!(pso, opts)