import Distributions, Random, StatsFuns, DataFrames

"""
    simulate_data_logistic(coefs; n::Int=100, noise::Number=0.0, groupnames=:random)

Generate a DataFrame of variables related a binary dependent variable by specified regression `coefs`.

!!! warning

    This function, adapted from [this](https://stats.stackexchange.com/questions/46523/how-to-simulate-artificial-data-for-logistic-regression) thread, doesn't work as expected (See [#27](https://github.com/neuropsychology/Psycho.jl/issues/27)).

# Multiple Variables / Groups
- If `coefs` is a vector (*e.g., `[0.1, 0.2]`), the DataFrame will contain `length(coefs)` variables (`Var1, Var2, ...`). Altough uncorrelated between them, they are  correlated to the outcome (`y`) by the specified coefs.
- If `coefs` is a vector of vectors (*e.g., `[[0.1], [0.2]]`), it will create `length(coefs)` groups, *i.e.*, stacked DataFrames with a correlation between the variables and the outcome varying between groups. It is possible to specify the `groupnames`.

# Arguments
- `coefs`: Regression coefficients. Can be a number, a vector of numbers or a vector of vectors.
- `n::Int`: Number of observations.
- `noise::Number`: The SD of the random gaussian noise.
- `groupnames::Vector`: Vector of group names (default to `:random`).
- `kwargs...`: Arguments to pass to other functions.

!!! note

    **Ideas / help required:**
    - Different group sizes (See [#9](https://github.com/neuropsychology/Psycho.jl/issues/9))

# Examples
```julia
simulate_data_logistic(0.2)
```
"""
function simulate_data_logistic end

function simulate_data_logistic(coefs::Vector{<:Number}; n::Int=100, noise::Number=0.0)
    # https://stats.stackexchange.com/questions/46523/how-to-simulate-artificial-data-for-logistic-regression
    n_var = length(coefs)
    X = randn(n, n_var)
    Z = X * coefs .+ 1

    p = StatsFuns.logistic.(Z)
    y = Float64.(rand(length(p)) .< p)

    # Add noise
    if noise != 0
      X = X .+ Random.rand(Distributions.Normal(0, noise), n, n_var)
    end

    data = DataFrames.DataFrame(hcat(y, X),
      Symbol.(vcat(["y"], "Var" .* string.(1:n_var))))
    return data

end

function simulate_data_logistic(coefs::Number; n::Int=100, noise::Number=0.0)
  simulate_data_logistic([coefs], n=n, noise=noise)
end

function simulate_data_logistic(coefs::Vector{<:Vector}; n::Int=100, noise::Number=0.0, groupnames=:random, kwargs...)
  data = simulate_data_groups(simulate_data_logistic, coefs=coefs, n=n, noise=noise, groupnames=groupnames, kwargs...)
  return data
end


export AbstractSimpleShape

abstract type AbstractSimpleShape{K} <: AbstractShape{K} end

# Below, typing x::SFloat{K} generates an error in gradient().
ndir(x::SReal{K}, s::AbstractSimpleShape{K}) where {K} = normalize(gradient(x -> level(x,s), x))

function project(x::SReal{K}, s::AbstractSimpleShape{K}) where {K}
    n̂ = ndir(x, s)
    f(α) = level(x + α*n̂, s)

    α = newton(f, 0.0).sol
    pt = x + α*n̂

    return (pt=pt, ndir=n̂)
end

# Return the minimum axis-aligned bounding box of the shape.
function bounds(s::AbstractSimpleShape{K}) where {K}
    c = center(s)
    δc = τᵣ*abs.(c) .+ τₐ  # SFloat{K}; all-positive entries

    # Calculate the negative-end corner of the axis-aligned bounding box.
    bₙ = SVec(ntuple(k->
        begin
            c′ = c - δc[k] * SVec(ntuple(k′->(k′==k), Val(K)))  # SFloat{K}; perturb c towards negative direction of k-axis
            bₖ = lagrange(x->x[k], x->level(x,s), c′, rtol=τₐ).sol[k]  # Float
        end, Val(K)))

    # Calculate the positive-end corner of the axis-aligned bounding box.
    bₚ = SVec(ntuple(k->
        begin
            c′ = c + δc[k] * SVec(ntuple(k′->(k′==k), Val(K)))  # SFloat{K}; perturb c towards positive direction of k-axis
            bₖ = lagrange(x->x[k], x->level(x,s), c′, rtol=τₐ).sol[k]  # Float
        end, Val(K)))

    # return c + τᵣ*(bₙ-c), c + τᵣ*(bₚ-c)  # make box slightly larger
    return bₙ, bₚ
end

include("ball.jl")
include("ellipsoid.jl")
include("halfspace.jl")


include("decrement.jl")

const ixVector = Vector{Float64}

mutable struct InterestRate <: Decrement
    ix::Vector{Float64}
    ifx
    # constructor with a predefined vector
    function InterestRate(v::ixVector)
        new(v,(x -> v[x]))
    end

    # constructor with a real number argument converts to a function that will produce an interest rate
    function InterestRate(i::Real)
        new(Float64[],(x -> i))
    end

    # constructor with (any) argument assumes a function that will produce an interest rate
    function InterestRate(f)
        new(Float64[],f)
    end
    # constructor with (any) argument assumes a function that will produce an interest rate
    # and provided rates for the first n... time periods
    function InterestRate(f, x...)
        new(collect(x),f)
    end
end

# the interst during time x
function i(iv::InterestRate,x)
    if x < 0 # looking for nth item from end of vector
        return iv.ix[end + Int(x) + 1]
    elseif length(iv.ix) < x
        append!(iv.ix, iv.ifx(length(iv.ix)+1))
        return i(iv,x)
    else
        return iv.ix[Int(x)]
    end
end

# the discount rate from time x to x+t
function tvx(iv::InterestRate,t,x, v=1.0)
    if t > 0
        return tvx(iv,t-1,x+1,v / (1.0 + i(iv,x)))
    else
        return v
    end
end

# ω (omega), the ultimate end of the given table
function ω(i::InterestRate)
    if length(i.ix) == 0
        return Inf
    else
        return length(i.ix)
    end
end
w(iv::InterestRate) = ω(iv)

### Convienence functions

# the discount rate at time x
vx(iv::InterestRate,x) = tvx(iv,1.0,x,1.0)

# the discount rate at time 1
v(iv::InterestRate) = vx(iv,1.0)


### A Pluto.jl notebook ###
# v0.16.0

using Markdown
using InteractiveUtils

# ╔═╡ a2a09c0b-d3ec-4b46-9bca-7c69b794f8ce
md"# Solutions to exercises"

# ╔═╡ 40956165-26d0-4d61-b795-033f6f2a0674
md"## Setup"

# ╔═╡ 197fd00e-9068-46ea-935f-ddf6a6bfbbdd
md"Inspect Julia version:"

# ╔═╡ f6d4f8c4-e441-45c4-af2a-25235e544a31
VERSION

# ╔═╡ 45740c4d-b789-45dc-8af5-148d95ea2900
md"The following instantiates a package environment."

# ╔═╡ 42b0f1e1-16c9-4238-a6bf-47194d7e8e12
md"""
The package environment has been created using **Julia 1.6** and may not
instantiate properly for other Julia versions.
"""

# ╔═╡ 754d60ac-4cb0-443a-828a-4cc485149963
begin
  using Pkg
  Pkg.activate(@__DIR__)
  Pkg.instantiate()
end

# ╔═╡ 19dae74b-24bd-428f-9e54-3f7b3ca87ecb
md"## Exercise 1"

# ╔═╡ 3af8467f-4988-44dc-ba1f-6363f43ec697
md"Write a function named `total` that adds the elements of its vector input."

# ╔═╡ cd6e9fce-e54a-48c3-95ea-2955abd45275
md"### Solution"

# ╔═╡ 992f71c3-12d6-4ab6-b1b3-dd47e367ba54
function total(v)
    sum = 0.0 # better is zero(v)
    for i in 1:length(v)
        sum = sum + v[i]
    end
    return sum
end

# ╔═╡ 61304232-d701-4218-8d7e-8e449afd1c48
total(1:10)

# ╔═╡ 18960c40-199a-48fa-a949-7f3bd292fe31
md"The built-in function is called `sum`."

# ╔═╡ b04e2fdc-8f1a-4964-8514-99938a2932db
md"## Exercise 2"

# ╔═╡ 860bb453-c2f1-446d-a0de-1721c1bc2df2
md"""
Generate a 1000 random samples from the standard normal
distribution. Create a second such sample, and add the two samples
point-wise.  Compute the (sample) mean and variance of the combined
samples. In the same plot, show a frequency-normalized histogram of
the combined samples and a plot of the pdf for normal distribution
with zero mean and variance `2`.
"""

# ╔═╡ cd6e9fce-e54a-48c3-bca8-7eec7950fd82
md"### Solution"

# ╔═╡ fe0bd76b-d2bb-4f68-9873-1b1430e635cc
begin
  using Distributions, Plots, Statistics

  samples1 = randn(1000); # or rand(Normal(), 1000)
  samples2 = randn(1000);

  samples = samples1 .+ samples2;

  mu = mean(samples)
  var = std(samples)^2

  @show mu var
end

# ╔═╡ 37a8bacf-d991-42ac-b43e-d0ebe87da176
begin
  d = Normal(0, sqrt(2))
  f(x) = pdf(d, x)

  xs = -5:(0.1):5
  ys = f.(xs);

  using Plots

  plt = histogram(samples, normalize=true)
  plot!(xs, ys)
  plt
end

# ╔═╡ 4f939d4a-e802-4b7d-9009-4cfb2012998f
md"## Exercise 3"

# ╔═╡ 63f410ea-c37a-4761-abd3-cb77d7a79683
md"The following shows that named tuples share some behaviour with dictionaries:"

# ╔═╡ cb2c42bf-21b5-4e04-879e-dc128f3db7b3
begin
  t = (x = 1, y = "cat", z = 4.5)
  keys(t)
end

# ╔═╡ b8996fa3-c046-4fb1-a367-9aa3c6cf34d0
t[:y]

# ╔═╡ 6584c100-ffd6-4fa1-bf32-94657e65284e
md"""
Write a function called `dict` that converts a named tuple to an
actual dictionary. You can create an empty dictionary using `Dict()`.
"""

# ╔═╡ cd6e9fce-e54a-48c3-9afd-65a5b5facac9
md"### Solution"

# ╔═╡ 834576cc-c6dc-4d3e-b6c8-54f5ed90a964
function dict(t)
    d = Dict()
    for k in keys(t)
        d[k] = t[k]
    end
    return d
end

# ╔═╡ dfdc3d18-4274-491d-9292-fe822525fc77
dict(t)

# ╔═╡ a8888dc5-ef44-4ee7-ae5e-4b83dcbc9675
md"Or alternatively:"

# ╔═╡ 5b64ff6d-923b-4eab-8a29-32f11452654a
dict(t) = Dict(k => t[k] for k in keys(t))

# ╔═╡ 135dac9b-0bd9-4e1d-a5f1-151fcbe229a2
md"""
---

*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*
"""

# ╔═╡ Cell order:
# ╟─a2a09c0b-d3ec-4b46-9bca-7c69b794f8ce
# ╟─40956165-26d0-4d61-b795-033f6f2a0674
# ╟─197fd00e-9068-46ea-935f-ddf6a6bfbbdd
# ╠═f6d4f8c4-e441-45c4-af2a-25235e544a31
# ╟─45740c4d-b789-45dc-8af5-148d95ea2900
# ╟─42b0f1e1-16c9-4238-a6bf-47194d7e8e12
# ╠═754d60ac-4cb0-443a-828a-4cc485149963
# ╟─19dae74b-24bd-428f-9e54-3f7b3ca87ecb
# ╟─3af8467f-4988-44dc-ba1f-6363f43ec697
# ╟─cd6e9fce-e54a-48c3-95ea-2955abd45275
# ╠═992f71c3-12d6-4ab6-b1b3-dd47e367ba54
# ╠═61304232-d701-4218-8d7e-8e449afd1c48
# ╟─18960c40-199a-48fa-a949-7f3bd292fe31
# ╟─b04e2fdc-8f1a-4964-8514-99938a2932db
# ╟─860bb453-c2f1-446d-a0de-1721c1bc2df2
# ╟─cd6e9fce-e54a-48c3-bca8-7eec7950fd82
# ╠═fe0bd76b-d2bb-4f68-9873-1b1430e635cc
# ╠═37a8bacf-d991-42ac-b43e-d0ebe87da176
# ╟─4f939d4a-e802-4b7d-9009-4cfb2012998f
# ╟─63f410ea-c37a-4761-abd3-cb77d7a79683
# ╠═cb2c42bf-21b5-4e04-879e-dc128f3db7b3
# ╠═b8996fa3-c046-4fb1-a367-9aa3c6cf34d0
# ╟─6584c100-ffd6-4fa1-bf32-94657e65284e
# ╟─cd6e9fce-e54a-48c3-9afd-65a5b5facac9
# ╠═834576cc-c6dc-4d3e-b6c8-54f5ed90a964
# ╠═dfdc3d18-4274-491d-9292-fe822525fc77
# ╟─a8888dc5-ef44-4ee7-ae5e-4b83dcbc9675
# ╠═5b64ff6d-923b-4eab-8a29-32f11452654a
# ╟─135dac9b-0bd9-4e1d-a5f1-151fcbe229a2


# make detection more robust
# 2. sub-grid detection based on fitting, rather than accurate zero finding
# rationale: core location isn't a good proxy for vortex location.
# vortex is a bulk flow excitation, whereas any small noise will shift zero - without shifting bulk

#---
using Plots, Revise, VortexDistributions

# unitary?
v0 = [.2 .4 1]
w0 = PointVortex(v0)
@test rawData(w0) == v0

# @test foundNear(1)

# no vortices
Nx = 400; Ny = 400
Lx = 200;
Ly = 200;
x = LinRange(-Lx / 2, Ly / 2, Nx)
y = LinRange(-Ly / 2, Ly / 2, Ny)

psi0 = one.(x * y') |> complex;
psi = Torus(psi0, x, y);
vortfound = findvortices(psi)
@test vortfound == nothing

#--- debugging functions
import VortexDistributions: findvortices, findvortices_grid, findvortices_interp, findvortices_jumps

function findvortices(psi::Field,interp=true)
    if interp
        return findvortices_interp(psi)
    else
        return findvortices_grid(psi)
    end
end

function findvortices_grid(psi::Torus;shift=true)
    vort = findvortices_jumps(psi,shift=shift)
    rawvort = rawData(vort)
    vort = sortslices(rawvort,dims=1)
    return PointVortex(vort)
end

function findvortices_jumps(psi::Field;shift=true)
    @unpack x,y,ψ = psi
    phase = angle.(ψ)

    # Compute all nearest neighbour jumps
    diffx = phasejumps(phase,1); diffy = phasejumps(phase,2)

    # Compute all plaquette loop integrals with in-place memory recycling
    circshift!(phase,diffx,(0,1))
    diffx .-= phase; diffx .-= diffy
    circshift!(phase,diffy,(1,0))
    diffx .+= phase

    # Find all windings
    ixp,iyp,vp = findwhere(diffx .> 0.0)
    xp = x[ixp]; yp = y[iyp]; np = length(vp)
    ixn,iyn,vn = findwhere(diffx .< 0.0)
    xn = x[ixn]; yn = y[iyn]; nn = length(vn)

    if shift
        dx = x[2]-x[1]; dy = y[2] - y[1]
        xp .-= dx/2; yp .-= dy/2; xn .-= dx/2; yn .-= dy/2
    end

    vortices = [xn yn -vn; xp yp vp]

    return PointVortex(vortices)
end

function findvortices_interp(psi::Field)
    vort = findvortices_grid(psi,shift=true)

        return vort
    end
end

# function findvortices_interp(psi::Field)
#     vort = findvortices_grid(psi,shift=true)
#     vort = removeedgevortices(vort,psi)
#     #TODO: allow for interp with periodic data (here edges are stripped)
#
#     for (j,vortex) in enumerate(vort)
#         try
#         vortz,psiz = corezoom(vortex,psi)
#         vortz,psiz = corezoom(vortz,psiz)
#         vortz,psiz = corezoom(vortz,psiz)
#         vort[j] = vortz
#         catch nothing
#         end
#         return vort
#     end
# end

#--- one vortex
Nx = 400; Ny = 400
Lx = 200;
Ly = 200;
x = LinRange(-Lx / 2, Ly / 2, Nx)
y = LinRange(-Ly / 2, Ly / 2, Ny)

psi0 = one.(x * y') |> complex;
psi = Torus(psi0, x, y);
vort = randVortex(1, psi)
vortex!(psi, vort)

vort = PointVortex(vort)

vortfound = findvortices(psi)

vfdata = rawData(vortfound)

vdata = rawData(vort)

heatmap(x,y,abs2.(psi.ψ),aspect_ratio=1)
heatmap(x,y,angle.(psi.ψ),aspect_ratio=1)

# Let's try fitting a Gaussian to the density and compare with root finding
# Do this with and without local grid noise

#--- local field (within 2 healing lengths)
x1 = x[@. abs(x -vort[1].xv)<2]
y1 = y[@. abs(x -vort[1].yv)<2]

psi1 = psi.ψ[(@. abs(x -vort[1].xv)<2) , @. abs(x -vort[1].yv)<2]

heatmap(x1,y1,abs2.(psi1),aspect_ratio=1)
heatmap(x1,y1,angle.(psi1),aspect_ratio=1)

m = abs2.(psi1)
m = maximum(m) .- m
heatmap(x1,y1,m,aspect_ratio=1)

using DataFitting

gaussian(x, y, p1, p2, p3,p4) = @.  exp(-((x-p1)^2  +  (y-p2)^2)/p3^2)*p4

dom = CartesianDomain(x1,y1)

data = Measures(m) #Measures(d + randn(rng, size(d)), 1.)

model1 = Model(:comp1 => FuncWrap(gaussian, vort[1].xv, vort[1].yv, 1.,1.))
prepare!(model1, dom, :comp1)
@time result = fit!(model1, data)

#--- add local noise
phi = Torus(copy(psi.ψ),x,y)
phi.ψ .+= (randn()+im*randn())*0.2

vort = findvortices_interp(phi)
x1 = x[@. abs(x -vort[1].xv)<2]
y1 = y[@. abs(x -vort[1].yv)<2]

phi1 = phi.ψ[(@. abs(x -vort[1].xv)<2) , @. abs(x -vort[1].yv)<2]

heatmap(x1,y1,abs2.(phi1),aspect_ratio=1)
heatmap(x1,y1,angle.(phi1),aspect_ratio=1)
heatmap(x,y,angle.(phi.ψ),aspect_ratio=1)

m = abs2.(phi1)
m = maximum(m) .- m
heatmap(x1,y1,m,aspect_ratio=1)

using DataFitting

gaussian(x, y, p1, p2, p3,p4) = @.  exp(-((x-p1)^2  +  (y-p2)^2)/p3^2)*p4

dom = CartesianDomain(x1,y1)

data = Measures(m) #Measures(d + randn(rng, size(d)), 1.)

model1 = Model(:comp1 => FuncWrap(gaussian, vort[1].xv, vort[1].yv, 1.,1.))
prepare!(model1, dom, :comp1)
@time result = fit!(model1, data)

vfoundnoise = findvortices(phi) |> rawData


#==
This is the modified ZDT example taken from
Esmaile Khorram, K Khaledian, and M Khaledyan.
A numerical method for
constructing the pareto front of multi-objective optimization problems.
Journal of Computational and Applied Mathematics, 261:158–171, 2014.

This is an example of a disconnected Pareto front.
The standard equality-based NBI method fails miserably here, and
will hit an infeasibility error somewhere along the way.
It will be interesting to see how it works with
the inequality-extended NBI method
==#

using JuMP, MultiJuMP, Ipopt

m = multi_model(solver=IpoptSolver())
n = 30

l = -ones(n); l[1] = 0
u = ones(n)
@variable(m, l[i] <= x[i=1:n] <= u[i])
@NLexpression(m, f1, x[1])
@NLexpression(m, g, 1 + 9*sum{x[j]^2, j=2:n}/(n-1))
@NLexpression(m, h, 1 - sqrt((f1/g)) - (f1/g)*sin(10*π*f1))
@NLexpression(m, f2, g*h)

iv1 = zeros(n)
obj1 = SingleObjective(f1, sense = :Min,
                       iv = Dict{Symbol, Any}(:x => iv1))

# As the problem is nonconvex, we have to supply the
# initial value to get the global minimum of f2
iv2 = zeros(n)
iv2[1] = 0.85
obj2 = SingleObjective(f2, sense = :Min,
                       iv = Dict{Symbol, Any}(:x => iv2))

multim = get_multidata(m)
multim.objectives = [obj1, obj2]
multim.pointsperdim = 60
# NB: solve fails with an infeasibility error 65% through the algorithm,
# You can still plot the results by pasting in the plot_comparison() function below
solve(m, method = NBI(true))

function pltfun(x1)
    setvalue(x[1], x1)
    setvalue(x[2:n], zeros(n-1))
    getvalue(f2)
end

using Plots
plt = plot(multim, label="NBI points",
           legend = :bottomleft)
xvals = linspace(0,1,200)
plot!(plt, xvals, pltfun, markercolor="orange",
      label="Image boundary")


###########################################################################
# Cross-correlation
#
# Function carries out cross-correlation to determine the bearing
# angles is the vector of angles corresponding to the samples
#
# Returns vector containing shifts at every degree
#	ith index corresponds to (i-1) shift
#
# The idea is that the reference is of length 360,
#  but the sample need not be.
# Therefore, a vector of sampled angles is provided.
# TODO: give this a thorough scrubbing. It is important to me that it works
function bearing_cc(angles, gains, ref_angles, ref_gains)
	cval = 0.0
	max_cval = -Inf
	max_i = 0.0

	norm_gains = copy(gains)
	normalize!(norm_gains)

	for i = 0:359
		cval = cross_correlate(angles, norm_gains, ref_angles, ref_gains, i)
		if (cval > max_cval)
			max_cval = cval
			max_i = i
		end
	end
	return 360.0 - max_i
end
# returns max correlation coefficient
function bearing_ccn(angles, gains, ref_angles, ref_gains)
	cval = 0.0
	max_cval = -Inf
	max_i = 0.0

	norm_gains = copy(gains)
	normalize!(norm_gains)

	for i = 0:359
		cval = cross_correlate(angles, norm_gains, ref_angles, ref_gains, i)
		if (cval > max_cval)
			max_cval = cval
			max_i = i
		end
	end
	return max_cval
end

function bearing_cc(angles, gains, ref_file::AbstractString)
	ref_gp = GainPattern(ref_file)
	return bearing_cc(angles, gains, ref_gp.angles, ref_gp.meangains)
end

# Returns array of bearing values
# Go through all samples in test_file
# Assumes all samples are the same length in test_gp
function bearing_cc(test_file::AbstractString, ref_file::AbstractString)
	test_gp = GainPattern(test_file)
	ref_gp = GainPattern(ref_file)

	num_samples = length(test_gp.angles)
	num_exp = length(test_gp.samples[1])

	aoa_array = zeros(num_exp)
	angles = zeros(num_samples)
	gains = zeros(num_samples)
	for i = 1:num_exp
		for j = 1:num_samples
			# Loop through all samples in this experiment...
			angles[j] = test_gp.angles[j]
			gains[j] = test_gp.samples[j][i]
		end
		aoa_array[i] = bearing_cc(angles, gains, ref_gp.angles, ref_gp.meangains)
	end
	return aoa_array
end

function bearing_ccn(test_file::AbstractString, ref_file::AbstractString)
	test_gp = GainPattern(test_file)
	ref_gp = GainPattern(ref_file)

	num_samples = length(test_gp.angles)
	num_exp = length(test_gp.samples[1])

	aoa_array = zeros(num_exp)
	angles = zeros(num_samples)
	gains = zeros(num_samples)
	for i = 1:num_exp
		for j = 1:num_samples
			# Loop through all samples in this experiment...
			angles[j] = test_gp.angles[j]
			gains[j] = test_gp.samples[j][i]
		end
		aoa_array[i] = bearing_ccn(angles, gains, ref_gp.angles, ref_gp.meangains)
	end
	return aoa_array
end

# Helper function to perform cross-correlation
# This is the cross-correlation function used in Graefenstein 2009
# Handle null observations
function cross_correlate(angles, gains, ref_angles, ref_gains, shift)
	c = 0.0
	len = length(angles)
	for i = 1:len
		if validgain(gains[i])
			c += gains[i] * interp_linear(angles[i], shift, ref_angles, ref_gains)
		end
	end
	return c
end


using PyPlot
using DelimitedFiles

R = readdlm("result.txt")
n = [1, 4, 9, 16, 25]

close("all")
loglog(n[end:-1:1], R[1:5, 3], "or-", label = "1 Core")
loglog(n[4:-1:1], R[6:9, 3], "xg-", label = "2 Cores")
loglog(n[3:-1:1], R[10:12, 3], "^b-", label = "3 Cores")
loglog(n[end:-1:1], R[1:5, 5], "or--")
loglog(n[4:-1:1], R[6:9, 5], "xg--")
loglog(n[3:-1:1], R[10:12, 5], "^b--")
loglog(n[end-2:-1:1], 8 ./n[end-2:-1:1], "k--", label = "\$\\mathcal{O}(1/n_{{proc}})\$")
legend()
title("Strong Scaling for Forward Computation")
xlabel("MPI Processors")
ylabel("Time (sec)")
savefig("forward_strong.png")

close("all")
loglog(n[end:-1:1], R[1:5, 4], "or-", label = "1 Core")
loglog(n[4:-1:1], R[6:9, 4], "xg-", label = "2 Cores")
loglog(n[3:-1:1], R[10:12, 4], "^b-", label = "3 Cores")
loglog(n[end:-1:1], R[1:5, 6], "or--")
loglog(n[4:-1:1], R[6:9, 6], "xg--")
loglog(n[3:-1:1], R[10:12, 6], "^b--")
loglog(n[end-2:-1:1], 12 ./n[end-2:-1:1], "k--", label = "\$\\mathcal{O}(1/n_{{proc}})\$")
legend()
title("Strong Scaling for Gradient Backpropagation")
xlabel("MPI Processors")
ylabel("Time (sec)")
savefig("backward_strong.png")


using .Curves: ECPoint, gx, gy
using .Specs: PRG

abstract type Group end

Base.broadcasted(f::Function, x::Group, y::AbstractVector{<:Integer}) = f.((x for i in 1:length(y)), y)

Base.broadcasted(::typeof(*), x::G, y::Vector{G}) where G <: Group = (x for i in 1:length(y)) .* y
Base.broadcasted(::typeof(*), x::Vector{G}, y::G) where G <: Group =  x .* (y for i in 1:length(x))

order(x::G) where G <: Group = order(G)

import Base./
/(x::G, y::G) where G <: Group = x * inv(y)

name(x::G) where G <: Group = name(G)

Base.convert(::Type{Vector{G}}, x::Vector) where G <: Group = G[ G <| i for i in x]

Base.rand(prg::PRG, ::Type{G}, N::Integer; nr::Integer = 0) where G <: Group = Vector{G} <| rand(prg, spec(G), N; nr)

struct ECGroup{P<:ECPoint} <: Group
    x::P
end

Base.convert(::Type{ECGroup{P}}, x) where P <: ECPoint = ECGroup{P}(P <| x)
Base.convert(::Type{G}, x::G) where G <: ECGroup = x

Base.:*(x::G, y::G) where G <: ECGroup = G(x.x + y.x)

function Base.:^(x::G, n::Integer) where G <: ECGroup
    n_mod = mod(n, order(G))

    @assert n_mod != 0 "A bad exponent"

    return G(n_mod * x.x)
end

Base.inv(g::G) where G <: ECGroup = g^(order(G) - 1)

order(::Type{ECGroup{P}}) where P = order(P)

Base.isvalid(x::ECGroup) = isvalid(x.x)

Base.:(==)(x::G, y::G) where G <: ECGroup = x.x == y.x

modulus(::Type{ECGroup{P}}) where P <: ECPoint = modulus(P)

name(::Type{ECGroup{P}}) where P <: ECPoint = name(P)

Base.isless(x::G, y::G) where G <: ECGroup = isless(x.x, y.x)

Curves.gx(g::ECGroup) = gx(g.x)
Curves.gy(g::ECGroup) = gy(g.x)

function Base.show(io::IO, g::G) where G <: ECGroup
    show(io, G)
    print(io, " <| (")
    show(io, gx(g))
    print(io, ", ")
    show(io, gy(g))
    print(io, ")")
end

struct PGroup{S} <: Group
    g::BigInt

    PGroup(p::Integer, q::Integer; name=nothing) = PGroup{static(; p, q, name)}

    function PGroup{S}(x::BigInt) where S
        #@assert 1 < x < S.p "Not in range"
        # Relaxing a little so that group products could happen propeerly
        # But is a presence of 1 in multiplicative products a vulnerability?
        @assert 0 < x < S.p "Not in range"
        new{S}(x)
    end

    PGroup{S}(x::Integer) where S = PGroup{S}(BigInt(x))

end

modulus(::Type{PGroup{S}}) where S = BigInt(S.p)
modulus(::G) where G <: PGroup = modulus(G)

order(::Type{PGroup{S}}) where S = BigInt(S.q)

name(::Type{PGroup}) = nothing

name(::Type{PGroup{S}}) where S =  isnothing(S.name) ? nothing : convert(Symbol, S.name)

Base.show(io::IO, g::PGroup) = print(io, value(g))

Base.show(io::IO, ::Type{PGroup}) = print(io, "PGroup")

#dublicate at Fields.
function trimnumber(x::String)
    if length(x) < 30
        return x
    else
        return x[1:10] * "..." * x[end-10:end]
    end
end

trimnumber(x::Integer) = trimnumber(string(x))

groupstr(m) = "𝐙/($(trimnumber(m)))"

function Base.show(io::IO, ::Type{G}) where G <: PGroup

    if name(G) == nothing
        print(io, groupstr(modulus(G)))
    else
        print(io, name(G))
    end
end

function Base.display(::Type{G}) where G <: PGroup
    show(G)
    print(" (order = $(order(G)))")
end

value(g::PGroup) = g.g

Base.convert(::Type{P}, x::Integer) where P <: PGroup = P(BigInt(x))

Base.isvalid(g::G) where G <: PGroup = value(g) != 1 && powermod(value(g), order(G), modulus(G)) == 1

import Base.*
*(x::G, y::G) where G <: PGroup = G(mod(value(x) * value(y), modulus(G)))

import Base.^
function ^(x::G, n::Integer) where G <: PGroup

    n_mod = mod(n, order(G))
    @assert n_mod != 0 "A bad exponent"
    @assert value(x) != 1 "A value 1 is not part of prime group"

    return G(powermod(value(x), n_mod, modulus(G)))
end

import Base.inv
inv(x::G) where G <: PGroup = G(modinv(value(x), modulus(G)))

Base.:(==)(x::G, y::G) where G <: PGroup = x.g == y.g

Base.isless(x::G, y::G) where G <: PGroup = value(x) < value(y)

#Base.convert(::Type{G}, x::Integer) where G <: PGroup = G(x)

# function Base.prod(x::Vector{G}) where G <: PGroup

#     p = modulus(G)1
#     s = value(x[1])

#     for i in x[2:end]
#         s = mod(s * value(i), p)
#     end

#     return G(s)
# end


# run: ~/.julia/bin/mpiexecjl -n 4 julia --project scripts/heat_1D_mpi.jl
using Plots, Printf, MAT
import MPI

do_save = true

# MPI functions
@views function update_halo(A, neighbors_x, comm)
    if neighbors_x[1] >= 0 # MPI_PROC_NULL?
        sendbuf = A[2]
        recvbuf = zeros(1)
        MPI.Send(sendbuf,  neighbors_x[1], 0, comm)
        MPI.Recv!(recvbuf, neighbors_x[1], 1, comm)
        A[1] = recvbuf[1]
    end
    if neighbors_x[2] >= 0 # MPI_PROC_NULL?
        sendbuf = A[end-1]
        recvbuf = zeros(1)
        MPI.Send(sendbuf,  neighbors_x[2], 1, comm)
        MPI.Recv!(recvbuf, neighbors_x[2], 0, comm)
        A[end] = recvbuf[1]
    end
    return
end

@views function heat_1D_mpi()
    # MPI
    MPI.Init()
    dims        = [0]
    comm        = MPI.COMM_WORLD
    nprocs      = MPI.Comm_size(comm)
    MPI.Dims_create!(nprocs, dims)
    comm_cart   = MPI.Cart_create(comm, dims, [0], 1)
    me          = MPI.Comm_rank(comm_cart)
    coords      = MPI.Cart_coords(comm_cart)
    neighbors_x = MPI.Cart_shift(comm_cart, 0, 1)
    if (me==0) println("nprocs=$(nprocs), dims[1]=$(dims[1])") end
    # Physics
    lx   = 10.0
    λ    = 1.0
    ρCp  = 1.0
    nt   = 100
    # Numerics
    nx   = 32   # local
    nx_g = dims[1]*(nx-2) + 2 # global
    # Derived numerics
    dx   = lx/nx_g            # global
    dt   = dx^2/ρCp/λ/2.1
    # Array allocation
    qx   = zeros(nx-1)
    # Initial condition
    x0   = coords[1]*(nx-2)*dx
    xc   = [x0 + ix*dx - dx/2 - 0.5*lx  for ix=1:nx]
    T    = exp.(.-xc.^2)
    # Time loop
    for it = 1:nt
        if (it==11) global t0 = Base.time() end
        qx         .= .-λ.*diff(T)./dx
        T[2:end-1] .= T[2:end-1] .- dt./ρCp.*diff(qx)./dx
        update_halo(T, neighbors_x, comm_cart)
     end
    time_s = (Base.time()-t0)
    if (me==0) @printf("Time = %1.4e s, T_eff = %1.2f GB/s \n", time_s, round((2/1e9*nx*sizeof(lx))/(time_s/(nt-10)), sigdigits=2)) end
    if do_save file = matopen("$(@__DIR__)/T_$(me).mat", "w"); write(file, "T", Array(T)); close(file) end
    MPI.Finalize()
    return
end

heat_1D_mpi()


# testset for using them in loops
@testset "Loop Induction Variables" begin
	@show @__LINE__
    f(x) = cos(x) * log(x)
    function avxmax(v)
        max_x = -Inf
        @avx for i ∈ eachindex(v)
            x = f(i)
            max_x = max(max_x, x)
        end
        max_x
    end
    function avxextrema(v)
        max_x = -Inf
        min_x = Inf
        @avx for i ∈ eachindex(v)
            x = f(i)
            max_x = max(max_x, x)
            min_x = min(min_x, x)
        end
        min_x, max_x
    end

    v = 1:19
    minref, maxref = extrema(f, v)
    @test maxref ≈ avxmax(v)
    minavx, maxavx = avxextrema(v);
    @test minref ≈ minavx
    @test maxref ≈ maxavx
end


using ApproxFun, Base.Test

import ApproxFun: expα, expβ, expγ

#
# This tests the numerical stability of the formulæ appearing in Eq. (2.5) of:
#
# A.-K. Kassam and L. N. Trefethen, Fourth-order time-stepping for stiff PDEs, SIAM J. Sci. Comput., 26:1214--1233, 2005.
#

x = logspace(-15,2,18)

@test norm(expα(x)./expα(big(x))-1,Inf) < eps()
@test norm(expβ(x)./expβ(big(x))-1,Inf) < eps()
@test norm(expγ(x)./expγ(big(x))-1,Inf) < eps()


using Plexifications
using Base.Test: @test, @test_throws

@test_throws ErrorException begin
    inv(DualComplex(0, 0.0))
end

@test begin
    !iszerodivisor(DualComplex(1,2,3,4))
end

@test begin
    iszerodivisor(DualComplex(0,0,3,4))
end

@test begin
    a = Plexifications.randomBigFloat()
    isreal(DualComplex(a))
end

@test begin
    a = 1
    !isreal(DualComplex(a, a))
end

@test begin
    z = random(DualComplex{Int})
    z == +(z)
end

@test begin
    a = Plexifications.randomBigFloat()
    real(DualComplex(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, DualComplex(1, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1, H: 2, A: 3, HA: 4]"
    l == r
end

@test begin
    x = random(DualComplex{BigFloat})
    y = random(DualComplex{BigFloat})
    x != y
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(DualComplex{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(DualComplex{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(DualComplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(DualComplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Cmplex{BigInt})
    y = random(DualComplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    z = random(DualComplex{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    z = random(DualComplex{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    barA(barA(x)) == x
end

@test begin
    x = random(DualComplex{BigInt})
    barH(barH(x)) == x
end

@test begin
    x = random(DualComplex{BigInt})
    barH(barA(x)) == barA(barH(x))
end

@test begin
    x = random(DualComplex{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = barA(x * y)
    r = barA(x) * barA(y)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = barA(x + y)
    r = barA(x) + barA(y)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = barA(x - y)
    r = barA(x) - barA(y)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = barH(x * y)
    r = barH(x) * barH(y)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = barH(x + y)
    r = barH(x) + barH(y)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = barH(x - y)
    r = barH(x) - barH(y)
    l == r
end

@test begin
    x = random(DualComplex{Rational{BigInt}})
    y = random(DualComplex{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(DualComplex{Rational{BigInt}})
    y = random(DualComplex{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(DualComplex{Rational{BigInt}})
    y = random(DualComplex{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    z = random(DualComplex{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(DualComplex{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    z = random(DualComplex{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    z = random(DualComplex{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    z = random(DualComplex{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    n = norm(x)
    n > big(0) || n == big(0)
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = norm(x * y)
    r = norm(x) * norm(y)
    l == r
end

@test_throws ErrorException begin
    inv(DualComplex(0, 0.0))
end

@test_throws ErrorException begin
    random(DualComplex{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(DualComplex{Int})
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(DualComplex{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(DualComplex{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(DualComplex{BigInt})
    z == z + zero(DualComplex{BigInt})
end

@test begin
    z = random(DualComplex{BigInt})
    z == z * one(DualComplex{BigInt})
end

@test begin
    z = random(DualComplex{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(DualComplex{BigInt})
    z == z * one(z)
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = commutator(x, y)
    r = zero(DualComplex{BigInt})
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = alternatorL(x, y)
    r = zero(DualComplex{BigInt})
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = alternatorR(x, y)
    r = zero(DualComplex{BigInt})
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    l = flexator(x, y)
    r = zero(DualComplex{BigInt})
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    z = random(DualComplex{BigInt})
    l = associator(x, y, z)
    r = zero(DualComplex{BigInt})
    l == r
end

@test begin
    x = random(DualComplex{BigInt})
    y = random(DualComplex{BigInt})
    z = random(DualComplex{BigInt})
    l = jacobiator(x, y, z)
    r = 3 * x * y * z
    l == r
end

@test begin
    l = DualComplex(Cmplex(1,2))
    r = DualComplex(1,2,0)
    l == r
end

@test begin
    l = DualComplex(1, 2, 3.0)
    r = DualComplex(1.0, 2.0, 3.0)
    l == r
end

@test begin
    l = DualComplex(1, 2, 3, 4.0)
    r = DualComplex(1.0, 2.0, 3.0, 4.0)
    l == r
end


using Manopt, Manifolds, Test, Random
using LinearAlgebra: I, eigvecs, tr
Random.seed!(42)

@testset "Riemannian quasi-Newton Methods" begin
    @testset "Mean of 3 Matrices" begin
        # Mean of 3 matrices
        A = [18.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0]
        B = [0.0 0.0 0.0 0.009; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0]
        C = [0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0; -5.0 0.0 0.0 0.0]
        ABC = [A, B, C]
        x_solution = mean(ABC)
        F(::Euclidean, x) = 0.5 * norm(A - x)^2 + 0.5 * norm(B - x)^2 + 0.5 * norm(C - x)^2
        gradF(::Euclidean, x) = -A - B - C + 3 * x
        M = Euclidean(4, 4)
        x = zeros(Float64, 4, 4)
        x_lrbfgs = quasi_Newton(
            M, F, gradF, x; stopping_criterion=StopWhenGradientNormLess(10^(-6))
        )
        @test norm(x_lrbfgs - x_solution) ≈ 0 atol = 10.0^(-14)
        # with Options
        lrbfgs_o = quasi_Newton(
            M,
            F,
            gradF,
            x;
            stopping_criterion=StopWhenGradientNormLess(10^(-6)),
            return_options=true,
        )
        @test lrbfgs_o.x == x_lrbfgs
        # with Cached Basis
        x_lrbfgs_cached = quasi_Newton(
            M,
            F,
            gradF,
            x;
            stopping_criterion=StopWhenGradientNormLess(10^(-6)),
            basis=get_basis(M, x, DefaultOrthonormalBasis()),
        )
        @test x_lrbfgs_cached == x_lrbfgs

        x_lrbfgs_cached_2 = quasi_Newton(
            M,
            F,
            gradF,
            x;
            stopping_criterion=StopWhenGradientNormLess(10^(-6)),
            basis=get_basis(M, x, DefaultOrthonormalBasis()),
            memory_size=-1,
        )
        @test isapprox(M, x_lrbfgs_cached_2, x_lrbfgs; atol=1e-5)

        x_clrbfgs = quasi_Newton(
            M,
            F,
            gradF,
            x;
            cautious_update=true,
            stopping_criterion=StopWhenGradientNormLess(10^(-6)),
        )
        @test norm(x_clrbfgs - x_solution) ≈ 0 atol = 10.0^(-14)

        x_rbfgs_Huang = quasi_Newton(
            M,
            F,
            gradF,
            x;
            memory_size=-1,
            stepsize=WolfePowellBinaryLinesearch(
                ExponentialRetraction(), ParallelTransport()
            ),
            stopping_criterion=StopWhenGradientNormLess(10^(-6)),
        )
        @test norm(x_rbfgs_Huang - x_solution) ≈ 0 atol = 10.0^(-14)

        for T in [InverseBFGS(), BFGS(), InverseDFP(), DFP(), InverseSR1(), SR1()]
            for c in [true, false]
                x_direction = quasi_Newton(
                    M,
                    F,
                    gradF,
                    x;
                    direction_update=T,
                    cautious_update=c,
                    memory_size=-1,
                    stopping_criterion=StopWhenGradientNormLess(10^(-12)),
                )
                @test norm(x_direction - x_solution) ≈ 0 atol = 10.0^(-14)
            end
        end
    end
    @testset "Rayleigh Quotient Minimzation" begin
        n = 4
        rayleigh_atol = 1e-8
        A = [2.0 1.0 0.0 3.0; 1.0 3.0 4.0 5.0; 0.0 4.0 3.0 2.0; 3.0 5.0 2.0 6.0]
        A = (A + A') / 2
        M = Sphere(n - 1)
        F(::Sphere, X) = X' * A * X
        gradF(::Sphere, X) = 2 * (A * X - X * (X' * A * X))
        x_solution = abs.(eigvecs(A)[:, 1])

        x = Matrix{Float64}(I, n, n)[n, :]
        x_lrbfgs = quasi_Newton(
            M,
            F,
            gradF,
            x;
            basis=get_basis(M, x, DefaultOrthonormalBasis()),
            memory_size=-1,
            stopping_criterion=StopWhenGradientNormLess(1e-9),
        )
        @test norm(abs.(x_lrbfgs) - x_solution) ≈ 0 atol = rayleigh_atol

        x_clrbfgs = quasi_Newton(
            M,
            F,
            gradF,
            x;
            cautious_update=true,
            stopping_criterion=StopWhenGradientNormLess(1e-9),
        )

        x_cached_lrbfgs = quasi_Newton(
            M,
            F,
            gradF,
            x;
            basis=get_basis(M, x, DefaultOrthonormalBasis()),
            memory_size=-1,
            stopping_criterion=StopWhenGradientNormLess(1e-9),
        )
        @test norm(abs.(x_cached_lrbfgs) - x_solution) ≈ 0 atol = rayleigh_atol

        for T in [
                InverseDFP(),
                DFP(),
                Broyden(0.5),
                InverseBroyden(0.5),
                Broyden(0.5, :Davidon),
                Broyden(0.5, :InverseDavidon),
                InverseBFGS(),
                BFGS(),
            ],
            c in [true, false]

            x_direction = quasi_Newton(
                M,
                F,
                gradF,
                x;
                direction_update=T,
                cautious_update=c,
                memory_size=-1,
                stopping_criterion=StopWhenGradientNormLess(5 * 1e-8),
            )
            @test norm(abs.(x_direction) - x_solution) ≈ 0 atol = rayleigh_atol
        end
    end
    @testset "Brockett" begin
        struct GradF
            A::Matrix{Float64}
            N::Diagonal{Float64,Vector{Float64}}
        end
        function (gradF::GradF)(::Stiefel, X::Array{Float64,2})
            AX = gradF.A * X
            XpAX = X' * AX
            return 2 .* AX * gradF.N .- X * XpAX * gradF.N .- X * gradF.N * XpAX
        end

        n = 4
        k = 2
        M = Stiefel(n, k)
        A = [2.0 1.0 0.0 3.0; 1.0 3.0 4.0 5.0; 0.0 4.0 3.0 2.0; 3.0 5.0 2.0 6.0]
        F(::Stiefel, X) = tr((X' * A * X) * Diagonal(k:-1:1))
        gradF = GradF(A, Diagonal(Float64.(collect(k:-1:1))))

        x = Matrix{Float64}(I, n, n)[:, 2:(k + 1)]
        x_inverseBFGSCautious = quasi_Newton(
            M,
            F,
            gradF,
            x;
            memory_size=8,
            vector_transport_method=ProjectionTransport(),
            retraction_method=QRRetraction(),
            cautious_update=true,
            stopping_criterion=StopWhenGradientNormLess(1e-6),
        )

        x_inverseBFGSHuang = quasi_Newton(
            M,
            F,
            gradF,
            x;
            memory_size=8,
            stepsize=WolfePowellBinaryLinesearch(QRRetraction(), ProjectionTransport()),
            vector_transport_method=ProjectionTransport(),
            retraction_method=QRRetraction(),
            cautious_update=true,
            stopping_criterion=StopWhenGradientNormLess(1e-6),
        )
        @test isapprox(M, x_inverseBFGSCautious, x_inverseBFGSHuang; atol=2e-5)
    end
end


#
# inline-issue-minimal.jl -
#
# Demonstrate a reduction of performances with Julia-1.6.0 compared to 1.5.4.
# This version is similar to `inline-issue.jl` but does not use
# `MayOptimize.jl`.
#
module InlineIssue

using BenchmarkTools

# Catmull-Rom spline:
function spline(x::T) where {T<:AbstractFloat}
    abs_x = abs(x)
    if abs_x ≥ 2
        zero(T)
    elseif abs_x ≤ 1 # (((3/2)*x - (5/2))*x^2 + 1)
        ((T(3)/2)*abs_x - (T(5)/2))*abs_x^2 + 1
    else # (((5/2) - (1/2)*x)*x - 4)*x + 2 = (2 - x)^2 (1 - x)/2
        (2 - abs_x)^2*(1 - abs_x)/2
    end
end

# Compute all weights (10 operations).
function compute_weights(::typeof(spline),
                         t::T) where {T<:AbstractFloat}
    u = 1 - t
    v = (T(-1)/2)*t*u
    w1 = v*u
    w4 = v*t
    w = w4 - w1
    w2 = u - w1 + w
    w3 = t - w4 - w
    return (w1, w2, w3, w4)
end

function inbounds_map!(f::Function,
                       dst::Array{T,N},
                       src::Array{T,N}) where {T<:AbstractFloat,N}
    @inbounds for i in eachindex(dst, src)
        dst[i] = f(src[i])
    end
    return dst
end

function simd_map!(f::Function,
                   dst::Array{T,N},
                   src::Array{T,N}) where {T<:AbstractFloat,N}
    @inbounds @simd for i in eachindex(dst, src)
        dst[i] = f(src[i])
    end
    return dst
end

function compute_weights!(f::Function,
                          dst::Array{T,2},
                          src::Array{T,1}) where {T<:AbstractFloat}
    @assert size(dst) == (4,length(src))
    @inbounds @simd for i in eachindex(src)
        w1, w2, w3, w4 = compute_weights(f, src[i])
        dst[1,i] = w1
        dst[2,i] = w2
        dst[3,i] = w3
        dst[4,i] = w4
    end
    return dst
end

@inline function inlined_compute_weights!(f::Function,
                                          dst::Array{T,2},
                                          src::Array{T,1}) where {T<:AbstractFloat}
    @assert size(dst) == (4,length(src))
    @inbounds @simd for i in eachindex(src)
        w1, w2, w3, w4 = compute_weights(f, src[i])
        dst[1,i] = w1
        dst[2,i] = w2
        dst[3,i] = w3
        dst[4,i] = w4
    end
    return dst
end

function runtests(; T::Type{<:AbstractFloat}=Float32, n::Int=1_000)
    x = rand(T, n)
    y = Array{T}(undef, n)
    t = rand(T, n) .- one(T)/2
    z = Array{T}(undef, 4, n)
    print("Tests with Julia-", VERSION, ", T=", T, ", n=", n, "\n")
    print(" ├─ Call spline ($n times):\n");
    print(" │   ├─ map! ───────────────────────"); @btime $(map!)($spline, $y, $x)
    print(" │   ├─ inbounds_map! ──────────────"); @btime $(inbounds_map!)($spline, $y, $x)
    print(" │   └─ simd_map! ──────────────────"); @btime $(simd_map!)($spline, $y, $x)
    print(" └─ Computation of weights with spline ($n times):\n");
    print("     ├─ compute_weights! ───────────"); @btime $(compute_weights!)($spline, $z, $t)
    print("     └─ inlined_compute_weights! ───"); @btime $(inlined_compute_weights!)($spline, $z, $t)

end

runtests(T=Float32)
#println()
#runtests(T=Float64)

end # module


"""
# DlogGF

A library containing algorithms for computing discrete logarithm in finite
field.

# References

1. R. Barbulescu, P. Gaudry, A. Joux and E. Thomé: A quasi-polynomial algorithm
  for discrete logarithm in finite fields of small characteristic.
2. Stephen C. Pohlig and Martin E. Hellman: An Improved Algorithm for Computing
  Logarithms over GF(p) and Its Cryptographic Signifiance, IEEE transactions on
  information theory, vol. it-24, no. 1, January 1978.
3. Jincheng Zhuang and Qi Cheng: On Generating Coset Representatives of
  PGL_2(F-q) in PGL_2(F_q²)
4. R. Granger, T. Kleinjung and J. Zumbrägel: On the Powers of 2
"""
module DlogGF
using Nemo, Primes

function __init__()

    # Welcome message
    println("")
    print("Welcome to \n")
    print_with_color(:red, "
         8I  ,dPYb,
         8I  IP'`Yb
         8I  I8  8I                          .oooooo.    oooooooooooo
         8I  I8  8'                         d8P'  `Y8b   `888'     `8
   ,gggg,8I  I8 dP    ,ggggg,     ,gggg,gg 888            888
  dP'  'Y8I  I8dP    dP'  'Y8ggg dP'  'Y8I 888            888oooo8
 i8'    ,8I  I8P    i8'    ,8I  i8'    ,8I 888     ooooo  888    '
 d8,   ,d8b,,d8b,_ ,d8,   ,d8'  d8,   ,d8I `88.    .88'   888
 'Y8888P'`Y88P''Y88P'Y8888P'    'Y8888P'888  `Y8bood8P'   o888o
                                      ,d8I'
                                    ,dP'8I
                                   ,8'  8I
                                   I8   8I
                                   `8, ,8I
                                    `Y8P'
")
end

# Basic functions

include("basics.jl")

# Random suite

include("random.jl")

# Composite types

include("compositeTypes.jl")

# Linear algebra

include("linalg.jl")

# Generating the cosets

include("pglCosets.jl")

# Pohlig Hellman suite

include("pohligHellman.jl")

# Linear factors

include("linfactors.jl")

# BGJT algorithm

include("bgjt.jl")

# GKZ algorithm

include("gkz.jl")

# Factor base

include("factorbase.jl")

# Debug

include("debug.jl")

# end of module
end


# This file includes an example of memory operation.
# In this example `mem` delays its input for one step size time.

using Causal
using Plots

# Construct a model
ti, dt, tf = 0., 1., 100.
numtaps = 5  # Number of buffer taps in Memory.
delay = dt   # One step size delay.
model = Model(clock=Clock(ti, dt, tf))

@defmodel model begin
    @nodes begin
        gen = RampGenerator()
        mem = Memory(delay=dt, numtaps=numtaps)
        writer = Writer(input=Inport(2))
    end
    @branches begin
        gen     =>  mem
        mem[1]  =>  writer[1]
        gen[1]  =>  writer[2]
    end
end

# Simulate model
simulate!(model, ti, dt, tf)

# Read simulation data
t, x = read(getnode(model, :writer).component)
u = getnode(model, :gen).component.readout.(t .- delay)
err = u - x[:, 2]

# Plots simulation data
n1, n2 = 1, 5
marker = (:circle, 2)
p1 = plot(t[n1:n2], x[n1:n2, 1], label=:mem, marker=marker)
plot!(t[n1:n2], x[n1:n2, 2], label=:gen, marker=marker)
plot!(t[n1:n2], u[n1:n2], label=:real, marker=marker)
p2 = plot(t[n1:n2], err[n1:n2], label=:error, marker=marker)
display(plot(p1, p2, layout=(2,1)))


using Diagonalizations, LinearAlgebra, PosDefManifold, Test

##  Create data for testing the case k>1, m>1 ##
# `t` is the number of samples,
# `m` is the number of datasets,
# `k` is the number of observations,
# `n` is the number of variables,
# `noise` must be smaller than 1.0. The smaller the noise, the more data are correlated
# Output k vectors of m data data matrices
function getData(t, m, k, n, noise)
    # create m identical data matrices and rotate them by different
    # random orthogonal matrices V_1,...,V_m
    𝐕=[randU(n) for i=1:m] # random orthogonal matrices
    # variables common to all subjects with unique variance profile across k
    X=[(abs2.(randn(n))).*randn(n, t) for s=1:k]
    # each subject has this common part plus a random part
    𝐗=[[𝐕[i]*((1-noise)*X[s] + noise*randn(n, t)) for i=1:m] for s=1:k]
    return 𝐗, 𝐕
end

function getData(::Type{Complex{T}}, t, m, k, n, noise) where {T<:AbstractFloat}
    # create m identical data matrices and rotate them by different
    # random orthogonal matrices V_1,...,V_m
    𝐕=[randU(ComplexF64, n) for i=1:m] # random orthogonal matrices
    # variables common to all subjects with unique variance profile across k
    X=[(abs2.(randn(n))).*randn(ComplexF64, n, t) for s=1:k]
    # each subject has this common part plus a random part
    𝐗=[[𝐕[i]*((1-noise)*X[s] + noise*randn(ComplexF64, n, t)) for i=1:m] for s=1:k]
    return 𝐗, 𝐕
end

# REAL data
# do joint blind source separation of non-stationary data
t, m, n, k, noise = 200, 5, 4, 6, 0.1
Xset, Vset=getData(t, m, k, n, noise)
𝒞=Array{Matrix}(undef, k, m, m)
for s=1:k, i=1:m, j=1:m 𝒞[s, i, j]=(Xset[s][i]*Xset[s][j]')/t end

aX=majd(Xset; fullModel=true, algorithm=:OJoB)
# the spForm index of the estimated demixing matrices times the true
# mixing matrix must be low
@test mean(spForm(aX.F[i]'*Vset[i]) for i=1:m)<0.1

# test the same using NoJoB algorithm
aX=majd(Xset; fullModel=true, algorithm=:NoJoB)
@test mean(spForm(aX.F[i]'*Vset[i]) for i=1:m)<0.1

# plot the original cross-covariance matrices and the rotated
# cross-covariance matrices

# Get all products 𝐔[i]' * 𝒞[l, i, j] * 𝐔[j]
function _rotate_crossCov(𝐔, 𝒞, m, k)
    𝒮=Array{Matrix}(undef, k, m, m)
    @inbounds for l=1:k, i=1:m, j=1:m 𝒮[l, i, j]=𝐔[i]'*𝒞[l, i, j]*𝐔[j] end
    return 𝒮
end

# Put all `k` cross-covariances in a single matrix
# of dimension m*n x m*n for visualization
function 𝒞2Mat(𝒞::AbstractArray, m, k)
    n=size(𝒞[1, 1, 1], 1)
    C=Matrix{Float64}(undef, m*n, m*n)
    for i=1:m, j=1:m, x=1:n, y=1:n C[i*n-n+x, j*n-n+y]=𝒞[k, i, j][x, y] end
    return C
end

using Plots

Cset=[𝒞2Mat(𝒞, m, s) for s=1:k]
 Cmax=maximum(maximum(abs.(C)) for C ∈ Cset)
 h1 = heatmap(Cset[1], clim=(-Cmax, Cmax), yflip=true, c=:bluesreds, title="all cross-cov, k=1")
 h2 = heatmap(Cset[2], clim=(-Cmax, Cmax), yflip=true, c=:bluesreds, title="all cross-cov, k=2")
 h3 = heatmap(Cset[2], clim=(-Cmax, Cmax), yflip=true, c=:bluesreds, title="all cross-cov, k=3")
 h4 = heatmap(Cset[2], clim=(-Cmax, Cmax), yflip=true, c=:bluesreds, title="all cross-cov, k=4")
 h5 = heatmap(Cset[2], clim=(-Cmax, Cmax), yflip=true, c=:bluesreds, title="all cross-cov, k=5")
 h6 = heatmap(Cset[2], clim=(-Cmax, Cmax), yflip=true, c=:bluesreds, title="all cross-cov, k=6")
 📈=plot(h1, h2, h3, h4, h5, h6, size=(1200,550))
# savefig(📈, homedir()*"\\Documents\\Code\\julia\\Diagonalizations\\docs\\src\\assets\\FigmAJD1.png")

𝒮=_rotate_crossCov(aX.F, 𝒞, m, k)
 Sset=[𝒞2Mat(𝒮, m, s) for s=1:k]
 Smax=maximum(maximum(abs.(S)) for S ∈ Sset)
 h11 = heatmap(Sset[1], clim=(-Smax, Smax), yflip=true, c=:bluesreds, title="rotated cross-cov, k=1")
 h12 = heatmap(Sset[2], clim=(-Smax, Smax), yflip=true, c=:bluesreds, title="rotated cross-cov, k=2")
 h13 = heatmap(Sset[2], clim=(-Smax, Smax), yflip=true, c=:bluesreds, title="rotated cross-cov, k=3")
 h14 = heatmap(Sset[2], clim=(-Smax, Smax), yflip=true, c=:bluesreds, title="rotated cross-cov, k=4")
 h15 = heatmap(Sset[2], clim=(-Smax, Smax), yflip=true, c=:bluesreds, title="rotated cross-cov, k=5")
 h16 = heatmap(Sset[2], clim=(-Smax, Smax), yflip=true, c=:bluesreds, title="rotated cross-cov, k=6")
 📉=plot(h11, h12, h13, h14, h15, h16, size=(1200,550))
# savefig(📉, homedir()*"\\Documents\\Code\\julia\\Diagonalizations\\docs\\src\\assets\\FigmAJD2.png")

# COMPLEX data
# do joint blind source separation of non-stationary data
t, m, n, k, noise = 200, 5, 4, 6, 0.1
Xcset, Vcset=getData(ComplexF64, t, m, k, n, noise)
𝒞=Array{Matrix}(undef, k, m, m)
for s=1:k, i=1:m, j=1:m 𝒞[s, i, j]=(Xcset[s][i]*Xcset[s][j]')/t end

aXc=majd(Xcset; fullModel=true, algorithm=:OJoB)
# the spForm index of the estimated demixing matrices times the true
# mixing matrix must be low
@test mean(spForm(aXc.F[i]'*Vcset[i]) for i=1:m)<0.1

# test the same using NoJoB algorithm
aXc=majd(Xcset; fullModel=true, algorithm=:NoJoB)
@test mean(spForm(aXc.F[i]'*Vcset[i]) for i=1:m)<0.1


using Match

struct Instruction
    action::Char
    value::Int
end

mutable struct Ship
    direction::Complex{Int}
    position::Complex{Int}
end

function move!(s::Ship, i::Instruction)
    @match i.action begin
        'N' => (s.position += i.value * im)
        'S' => (s.position -= i.value * im)
        'E' => (s.position += i.value)
        'W' => (s.position -= i.value)
        'L' => (s.direction *= (im)^(i.value ÷ 90))
        'R' => (s.direction *= (-im)^(i.value ÷ 90))
        'F' => (s.position += s.direction * i.value)
        _ => error(i)
    end
end

function part1(lines)
    ship = Ship(1, 0)
    for line in lines
        inst = Instruction(line[1], parse(Int, line[2:end]))
        move!(ship, inst)
    end
    return abs(real(ship.position)) + abs(imag(ship.position))
end

const DATA = readlines("day12.txt")

println("Part 1: $(part1(DATA))")

function move2!(s::Ship, i::Instruction, waypoint::Complex{Int})
    @match i.action begin
        'N' => (waypoint += i.value * im)
        'S' => (waypoint -= i.value * im)
        'E' => (waypoint += i.value)
        'W' => (waypoint -= i.value)
        'L' => (waypoint *= (im)^(i.value ÷ 90))
        'R' => (waypoint *= (-im)^(i.value ÷ 90))
        'F' => (s.position += waypoint * i.value)
        _ => error(i)
    end
    return waypoint
end

function part2(lines)
    ship = Ship(1, 0)
    waypoint = complex(10, 1)
    for line in lines
        inst = Instruction(line[1], parse(Int, line[2:end]))
        waypoint = move2!(ship, inst, waypoint)
    end
    return abs(real(ship.position)) + abs(imag(ship.position))
end

println("Part 2: $(part2(DATA))")


input = read(stdin, String)

constraints, myticket, nearby = split(input, "\n\n")
myticket = [parse(Int, s) for s in split(split(myticket, '\n')[2], ',', keepempty=false)]

# Constraints : (min1, max1, min2, max2)
constr = []
for c in split(constraints, '\n', keepempty=false)
    colon = findfirst(':', c)
    part1, _, part2 = split(c[colon + 1: end], keepempty=false)
    part1 = strip(part1)
    part2 = strip(part2)
    (min1, max1), (min2, max2) = split(part1, '-', keepempty=false), split(part2, '-', keepempty=false)

    push!(constr, (parse(Int, min1), parse(Int, max1), parse(Int, min2), parse(Int, max2)))
end

errorrate = 0

# poss[index in tickets][index in constraint] = Possible that ticket matches constraint index
poss = ones(Bool, 20, 20)
for ticket in split(nearby, '\n', keepempty=false)[2:end]
    global errorrate, part2prod

    nb = [parse(Int, s) for s in split(ticket, ',', keepempty=false)]

    for (ni, n) in enumerate(nb)
        valid = false

        for (ci, (min1, max1, min2, max2)) in enumerate(constr)
            if min1 <= n <= max1 || min2 <= n <= max2
                valid = true

                break
            end
        end

        if !valid
            errorrate += n
        else
            # Part 2
            for (ci, (min1, max1, min2, max2)) in enumerate(constr)
                if !(min1 <= n <= max1 || min2 <= n <= max2)
                    global poss[ni, ci] = false
                end
            end
        end
    end
end

# Part 2
class2idx = Dict()

# Sums (all values in range 1 to 20) to start from min sum to max sum
sums = [(i, sum(poss[i, begin : end])) for i in 1:20]
sort!(sums, lt=((_, a), (_, b)) -> a < b)

for i in 1:20
    ci = sums[i][1]

    # There is only one 1
    class = findfirst(poss[ci, begin : end])

    # Remove all 1s of this class (already taken)
    for j in 1:20
        poss[j, class] = false
    end

    # Set class
    class2idx[class] = ci
end

part2prod = reduce(*, [myticket[class2idx[i]] for i in 1:6])

println("Part 1 : $errorrate")
println("Part 2 : $part2prod")


struct PinnedFoot <: Constraints
    n
    ineq
    q_pinned
end

function pinned_foot_constraint(model, q_pinned, T)
    n = 2 * (T + 1)
    ineq = (1:0)
    PinnedFoot(n, ineq, q_pinned)
end

function constraints!(c, Z, con::PinnedFoot, model, idx, h, T)
    k(y) = kinematics_2(model, y, body = :leg_1, mode = :ee)

    for t = 1:T
        if t == 1
            q⁻ = view(Z, idx.x[t][1:model.nq])
            c[1:2] = k(q⁻) - k(con.q_pinned)
        end
        q = view(Z, idx.x[t][model.nq .+ (1:model.nq)])
        c[2 + (t - 1) * 2 .+ (1:2)] = k(q) - k(con.q_pinned)
    end
    nothing
end

function constraints_jacobian!(∇c, Z, con::PinnedFoot, model, idx, h, T)
    k(y) = kinematics_2(model, y, body = :leg_1, mode = :ee)
    shift = 0
    for t = 1:T
        if t == 1
            q⁻ = view(Z, idx.x[t][1:model.nq])
            r_idx = (1:2)
            c_idx = idx.x[t][1:model.nq]
            len = length(r_idx) * length(c_idx)
            ∇c[shift .+ (1:len)] = vec(ForwardDiff.jacobian(k, q⁻))
            shift += len
        end
        q = view(Z, idx.x[t][model.nq .+ (1:model.nq)])
        r_idx = 2 + (t - 1) * 2 .+ (1:2)
        c_idx = idx.x[t][model.nq .+ (1:model.nq)]
        len = length(r_idx) * length(c_idx)
        ∇c[shift .+ (1:len)] = vec(ForwardDiff.jacobian(k, q))
        shift += len
    end

    return nothing
end

function constraints_sparsity(con::PinnedFoot, model, idx, T;
    shift_row = 0, shift_col = 0)
    row = []
    col = []

    for t = 1:T
        if t == 1
            r_idx = shift_row .+ (1:2)
            c_idx = shift_col .+ idx.x[t][1:model.nq]
            row_col!(row, col, r_idx, c_idx)
        end
        r_idx = shift_row + 2 + (t - 1) * 2 .+ (1:2)
        c_idx = shift_col .+ idx.x[t][model.nq .+ (1:model.nq)]
        row_col!(row, col, r_idx, c_idx)
    end

    return collect(zip(row, col))
end


"""
    PMC(; epsilon, K, prior, ess_min = 0.5, proposal = defaultproposal(prior))

Population Monte Carlo ABC structure, with `epsilon` schedule, `K` particles,
`prior`, resampling threshold `ess_min` and `proposal` distribution.
The `epsilon` schedule should be an iterable, e.g. an array of numbers,
[`EpsilonExponentialDecay`](@ref) or [`EpsilonLinearDecay`](@ref).
One dimensional problems should be defined in vectorized form (
see [`run!`](@ref) for an example).

# Example
```julia-repl
pmc = PMC(epsilon = EpsilonExponentialDecay(1, .01, .9), K = 20,
                 prior = MultivariateUniform([-1, -1], [1, 1]))
model(theta) = .1 * randn(2) .+ theta
data = [.3, -.2]
run!(pmc, model, data)

mean(pmc)
particles(pmc)
weights(pmc)
```

[Beaumont, M. A., Cornuet, J., Marin, J., and Robert, C. P. (2009)
Adaptive approximate Bayesian computation. Biometrika,96, 983-990.
](http://dx.doi.org/10.1093/biomet/asp052)
"""
struct PMC{E,P,D,D1,T} <: AbstractABC
    epsilon::E
    log_ess_min::Float64
    particles::Vector{P}
    log_weights::Vector{Float64}
    prior::D
    proposal::D1
    distance::T
end
function PMC(; epsilon, prior::Distribution{T}, K = 10^3, ess_min = 0.5,
             proposal = defaultproposal(prior), distance = euclidean) where T
    particles = init_particles(prior, K)
    PMC(epsilon, log(ess_min), particles, fill(-log(K), K),
        prior, proposal(particles), distance)
end
function Base.show(io::IO, mime::MIME"text/plain", pmc::PMC)
    println(io, "PMC(K = $(length(pmc.particles)), ess_min = $(exp(pmc.log_ess_min)))")
    print(io, "  epsilon schedule: ")
    show(io, mime, pmc.epsilon)
    print(io, "  prior: ")
    show(io, mime, pmc.prior)
    print(io, "  proposal: ")
    show(io, mime, pmc.proposal)
end
weights(pmc::PMC) = weights(pmc.log_weights)
particles(pmc::PMC) = pmc.particles

sample_particles!(particles, distribution, eps, model, data, distance) =
    sample_particles!(Random.GLOBAL_RNG, particles, distribution, eps, model, data, distance)
function sample_particles!(rng::Random.AbstractRNG,
                           particles, distribution, eps, model, data, distance)
    n_sims = 0
    Threads.@threads for particle in particles
        while true
            rand!(rng, distribution, particle)
            sim_data = model(particle)
            n_sims += 1
            distance(sim_data, data) < eps && break
        end
    end
    n_sims
end
resample(particles, weights; verbose = false, log_weights = false) = resample!(Random.GLOBAL_RNG, copy(particles), weights, verbose = verbose, log_weights = log_weights)[1]
function resample!(rng::Random.AbstractRNG, particles, weights;
                   verbose = true,
                   log_weights = true)
    verbose && println("Resampling...")
    idxs = wsample(1:length(weights), log_weights ? exp.(weights) : weights, length(weights))
    old_particles = deepcopy(particles) # only particles[idxs] would be needed
    for (i, j) in enumerate(idxs)
        copyto!(particles[i], old_particles[j])
    end
    weights .= log_weights ? -log(length(weights)) : 1/length(weights)
    particles, weights
end
function update_log_weights!(log_weights, particles, prior, proposal)
    for i in eachindex(log_weights)
        log_weights[i] = logpdf(prior, particles[i]) - logpdf(proposal, particles[i])
    end
    log_weights .-= logsumexp(log_weights)
    log_weights
end

"""
    run!([rng], method, model, data;
         callback = () -> nothing, maxfevals = Inf, verbose = true)

Run `method` on `model` and `data`.
The function `callback` gets evaluated after every iteration.
The method stops after the first iteration that reaches more than `maxfevals` calls to the model.
# Example
```julia-repl
using Distributions, LinearAlgebra
pmc = PMC(epsilon = [1, .5, .2, .1, .01, .001], K = 10^3,
                 prior = TruncatedMultivariateNormal([1.], 2I,
                                                     lower = [0], upper = [3]))
model(theta) = theta * randn() .+ 1.2
data = [.5]
callback() = @show mean(pmc)
run!(pmc, model, data, callback = callback, maxfevals = 10^5)

using StatsBase, StatsPlots
h = fit(Histogram, vcat(particles(pmc)...), nbins = 50)
StatsPlots.plot(h)
```
"""
run!(pmc, model, data; kwargs...) = run!(Random.GLOBAL_RNG, pmc, model, data; kwargs...)
# Beaumont et al. 2009 (see also Lintusaari et al. 2016
function run!(rng::AbstractRNG, pmc::PMC, model, data;
              callback = () -> nothing, maxfevals = Inf, verbose = true)
    start_time = now()
    n_sims = Int[]
    K = length(pmc.particles)
    push!(n_sims, sample_particles!(rng, pmc.particles, pmc.prior,
                                    first(pmc.epsilon), model, data, pmc.distance))
    update!(pmc.proposal, pmc.particles, pmc.log_weights)
    callback()
    verbose && @printf "%10.s %13.s %8.s %8.s %8.s\n" "iteration" "elapsed" "fevals" "epsilon" "ess"
    verbose && @printf "%10.s %13.s %8.s %8.e %8.s\n" 0 round(now() - start_time, Second) n_sims[1] first(pmc.epsilon) 1
    for (i, eps) in enumerate(Iterators.drop(pmc.epsilon, 1))
        push!(n_sims, sample_particles!(rng, pmc.particles, pmc.proposal, eps, model, data, pmc.distance))
        update_log_weights!(pmc.log_weights, pmc.particles, pmc.prior, pmc.proposal)
        l_ess = log_ess(pmc.log_weights) - log(K)
        if l_ess < pmc.log_ess_min
            resample!(rng, pmc.particles, pmc.log_weights, verbose = verbose)
        end
        update!(pmc.proposal, pmc.particles, pmc.log_weights)
        callback()
        verbose && @printf "%10.s %13.s %8.s %8.e %8.s\n" i round(now() - start_time, Second) sum(n_sims) eps exp(l_ess)
        sum(n_sims) > maxfevals && break
    end
    (weights = exp.(pmc.log_weights),
     particles = pmc.particles,
     n_sims = n_sims)
end

struct Particle{T}
    θ::Vector{Float64}
    x::T
    distances::Vector{Float64}
end
function Base.copyto!(dest::Particle, bc::Particle)
    dest.θ .= bc.θ
    for k in eachindex(dest.x)
        dest.x[k] .= bc.x[k]
    end
    dest.distances .= bc.distances
    dest
end
theta(p::Particle) = p.θ

"""
    AdaptiveSMC(; alpha = 0.9, epsilon = .1, ess_min = 0.5, M = 1,
                  prior, K = 10^3, proposal = defaultproposal(prior))

Adaptive Sequential Monte Carlo structure with `K` particles,
`M` calls of the model per parameter value, final `epsilon`,
decrease parameter `alpha` and resampling threshold `ess_min`.
See also [`PMC`](@ref).

# Example
```
using Distributions, LinearAlgebra
asmc = AdaptiveSMC(prior = MultivariateNormal(zeros(2), 4I))
model(theta) = vcat([.3*randn(2) .+ theta for _ in 1:2]...) # 2 i.i.d samples
data = [.4, -.3, .5, -.2]
run!(asmc, model, data)

mean(asmc)
using StatsPlots
p = hcat(particles(asmc)...); StatsPlots.scatter(p[1, :], p[2, :])
```

[Del Moral, P., Doucet, A., and Jasra, A. (2012) An adaptive
sequential Monte Carlo method for approximate Bayesian
computation. Statistics and Computing, 22, 1009-1020
](http://dx.doi.org/10.1007/s11222-011-9271-y)
"""
struct AdaptiveSMC{P,P1,D,T} <: AbstractABC
    α::Float64
    ϵ::Float64
    Nt::Float64
    prior::P
    proposal::P1
    particles::Vector{Particle{T}}
    weights::Vector{Float64}
    distance::D
end
function AdaptiveSMC(; alpha = 0.9, epsilon = .1, ess_min = 0.5, M = 1,
                     prior, K = 10^3,
                     proposal = defaultproposal(prior),
                     distance = euclidean,
                     datatype = distance == euclidean ? Vector{Float64} :
                                                        Vector{Vector{Float64}}
                    )
    particles = [Particle(randn(eltype(prior), length(prior)),
                          [datatype(undef, 0) for _ in 1:M],
                          zeros(M)) for _ in 1:K]
    AdaptiveSMC(alpha, float(epsilon), float(ess_min * K), prior,
                proposal(theta.(particles)),
                particles, zeros(K), distance)
end
Distributions.mean(x::Union{<:PMC, <:AdaptiveSMC}) = sum(particles(x) .* weights(x))
weights(asmc::AdaptiveSMC) = asmc.weights
particles(asmc::AdaptiveSMC) = theta.(asmc.particles)
function Base.show(io::IO, mime::MIME"text/plain", asmc::AdaptiveSMC)
    println(io, "AdaptiveSMC(α = $(asmc.α), ϵ = $(asmc.ϵ), K = $(length(asmc.particles)), ess_min = $(asmc.Nt/length(asmc.particles)), M = $(length(asmc.particles[1].distances)))")
    print(io, "  prior: ")
    show(io, mime, asmc.prior)
    print(io, "  proposal: ")
    show(io, mime, asmc.proposal)
end
function ess!(weights, particles, factor, e)
    for i in eachindex(weights)
        weights[i] = factor[i] * sum(particles[i].distances .< e)
    end
    maximum(weights) == 0 && return 0
    weights ./= sum(weights)
    ess(weights)
end
function getfactor(weights, particles, e_old)
    factor = zeros(length(weights))
    smallest_dist = Inf
    for i in eachindex(factor)
        weights[i] == 0 && continue
        tmp = minimum(particles[i].distances)
        tmp < smallest_dist && (smallest_dist = tmp)
        denom = sum(particles[i].distances .< e_old)
        factor[i] = weights[i] / denom
    end
    smallest_dist, factor
end
function new_epsilon!(weights, e_old, particles, alpha, e_min)
    smallest_dist, factor = getfactor(weights, particles, e_old)
    target = alpha * ess(weights)
    e = find_zero(e -> target - ess!(weights, particles, factor, e),
                  (smallest_dist, e_old), Bisection())
    e = max(e, e_min)
    e, ess!(weights, particles, factor, e)
end
function sample_particles!(rng::Random.AbstractRNG,
                           particles::Vector{<:Particle},
                           distribution, model, data, weights, distance;
                           indexed = false)
    n_sims = 0
    for (i, (w, particle)) in enumerate(zip(weights, particles))
        w == 0 && continue
        if indexed
            Distributions._rand!(rng, distribution, particle.θ, i)
        else
            rand!(rng, distribution, particle.θ)
        end
        for k in eachindex(particle.x)
            particle.x[k] = model(particle.θ)
            n_sims += 1
            particle.distances[k] = distance(particle.x[k], data)
        end
    end
    n_sims
end
function mh_accept!(particles, old_particles, prior, ϵ, weights)
    for i in eachindex(particles)
        weights[i] == 0 && continue
        logw = logpdf(prior, particles[i].θ) - logpdf(prior, old_particles[i].θ) +
               log(mean(particles[i].distances .< ϵ)/mean(old_particles[i].distances .< ϵ))
        if logw >= 0 || rand() < exp(logw)
            continue
        end
        copyto!(particles[i], old_particles[i]) # proposal rejected
    end
    particles
end

function run!(rng::Random.AbstractRNG, asmc::AdaptiveSMC, model, data;
              callback = () -> nothing,
              maxfevals = Inf,
              verbose = true)
    start_time = now()
    n_sims = Int[]
    epsilons = Float64[]
    i = 0
    eps = Inf
    K = length(asmc.particles)
    asmc.weights .= fill(1/K, K)
    push!(n_sims, sample_particles!(rng, asmc.particles, asmc.prior, model, data, asmc.weights, asmc.distance))
    update!(asmc.proposal, theta.(asmc.particles), asmc.weights, log_weights = false)
    callback()
    verbose && @printf "%10.s %13.s %8.s %8.s %8.s\n" "iteration" "elapsed" "fevals" "epsilon" "ess"
    verbose && @printf "%10.s %13.s %8.s %8.e %8.s\n" 0 round(now() - start_time, Second) n_sims[1] eps 1
    while true
        i += 1
        eps, ess = new_epsilon!(asmc.weights, eps, asmc.particles, asmc.α, asmc.ϵ)
        push!(epsilons, eps)
        ess < asmc.Nt && resample!(rng, asmc.particles, asmc.weights,
                                  log_weights = false, verbose = verbose)
        old_particles = deepcopy(asmc.particles)
        update!(asmc.proposal, theta.(asmc.particles), asmc.weights, log_weights = false)
        push!(n_sims, sample_particles!(rng, asmc.particles, asmc.proposal, model, data, asmc.weights, asmc.distance, indexed = true))
        mh_accept!(asmc.particles, old_particles, asmc.prior, eps, asmc.weights)
        ess = ess!(asmc.weights, asmc.particles, ones(K), eps)
        callback()
        verbose && @printf "%10.s %13.s %8.s %8.e %8.s\n" i round(now() - start_time, Second) sum(n_sims) eps ess/K
        (eps ≈ asmc.ϵ || sum(n_sims) > maxfevals) && break
    end
    (weights = asmc.weights,
     particles = asmc.particles,
     n_sims = n_sims,
     epsilons = epsilons)
end


using RipQP, LinearOperators, LLSModels, QPSReader, QuadraticModels, Quadmath
using DelimitedFiles, LinearAlgebra, MatrixMarket, SparseArrays, Test

qps1 = readqps("QAFIRO.SIF") #lower bounds
qps2 = readqps("HS21.SIF") # low/upp bounds
qps3 = readqps("HS52.SIF") # free bounds
qps4 = readqps("AFIRO.SIF") # LP
qps5 = readqps("HS35MOD.SIF") # fixed variables

include("solvers/test_augmented.jl")
include("solvers/test_newton.jl")
include("solvers/test_normal.jl")

include("test_multi.jl")
include("test_alternative_methods.jl")
include("test_alternative_problems.jl")


module BestBasis
export
    # cost functions
    CostFunction,
    LSDBCost,
    JBBCost,
    BBCost,
    LoglpCost,
    NormCost,
    DifferentialEntropyCost,
    ShannonEntropyCost,
    LogEnergyEntropyCost,
    coefcost,
    # tree cost
    tree_costs,
    # best tree selection
    bestbasis_treeselection,
    # best basis types
    BestBasisType,
    LSDB,
    JBB,
    BB,
    SIBB,
    # best basis tree
    bestbasis_tree

using
    Wavelets,
    LinearAlgebra,
    Statistics,
    AverageShiftedHistograms,
    Parameters

using
    ..Utils,
    ..DWT,
    ..SIWPD

## COST COMPUTATION
"""
Cost function abstract type.

**See also:** [`LSDBCost`](@ref), [`JBBCost`](@ref), [`BBCost`](@ref)
"""
abstract type CostFunction end

"""
    LSDBCost <: CostFunction

Cost function abstract type specifically for LSDB.

**See also:** [`CostFunction`](@ref), [`JBBCost`](@ref), [`BBCost`](@ref)
"""
abstract type LSDBCost <: CostFunction end

"""
    JBBCost <: CostFunction

Cost function abstract type specifically for JBB.

**See also:** [`CostFunction`](@ref), [`LSDBCost`](@ref), [`BBCost`](@ref)
"""
abstract type JBBCost <: CostFunction end

"""
    BBCost <: CostFunction

Cost function abstract type specifically for BB.

**See also:** [`CostFunction`](@ref), [`LSDBCost`](@ref), [`JBBCost`](@ref)
"""
abstract type BBCost <: CostFunction end

@doc raw"""
    LoglpCost <: JBBCost

``\log \ell^p`` information cost used for JBB. Typically, we set `p=2` as in
Wickerhauser's original algorithm.

**See also:** [`CostFunction`](@ref), [`JBBCost`](@ref), [`NormCost`](@ref)
"""
@with_kw struct LoglpCost <: JBBCost
    p::Number = 2
end

@doc raw"""
    NormCost <: JBBCost

``p``-norm information cost used for JBB.

**See also:** [`CostFunction`](@ref), [`JBBCost`](@ref), [`LoglpCost`](@ref)
"""
@with_kw struct NormCost <: JBBCost
    p::Number = 1
end

"""
    DifferentialEntropyCost <: LSDBCost

Differential entropy cost used for LSDB.

**See also:** [`CostFunction`](@ref), [`LSDBCost`](@ref)
"""
struct DifferentialEntropyCost <: LSDBCost end

"""
    ShannonEntropyCost <: LSDBCost

Shannon entropy cost used for BB.

**See also:** [`CostFunction`](@ref), [`BBCost`](@ref),
    [`LogEnergyEntropyCost`](@ref)
"""
struct ShannonEntropyCost <: BBCost end

"""
    LogEnergyEntropyCost <: LSDBCost

Log energy entropy cost used for BB.

**See also:** [`CostFunction`](@ref), [`BBCost`](@ref),
    [`ShannonEntropyCost`](@ref)
"""
struct LogEnergyEntropyCost <: BBCost end

# cost functions for individual best basis algorithm
function coefcost(x::T, et::ShannonEntropyCost, nrm::T) where T<:AbstractFloat
    s = (x/nrm)^2
    if s == 0.0
        return -zero(T)
    else
        return -s*log(s)
    end
end

function coefcost(x::T, et::LogEnergyEntropyCost, nrm::T) where T<:AbstractFloat
    s = (x/nrm)^2
    if s == 0.0
        return -zero(T)
    else
        return -log(s)
    end
end

function coefcost(x::AbstractArray{T}, et::BBCost, nrm::T=norm(x)) where
        T<:AbstractFloat

    @assert nrm >= 0
    sum = zero(T)
    nrm == sum && return sum
    for i in eachindex(x)
        @inbounds sum += coefcost(x[i], et, nrm)
    end
    return sum
end

# cost functions for joint best basis (JBB)
function coefcost(x::AbstractArray{T}, et::LoglpCost) where T<:Number
    xᵖ = abs.(x) .^ et.p
    return sum(log.(xᵖ))
end

function coefcost(x::AbstractArray{T}, et::NormCost) where T<:Number
    return norm(x, et.p)
end

# cost functions for least statistically dependent basis (LSDB)
function coefcost(x::AbstractVector{T}, et::DifferentialEntropyCost) where
        T<:AbstractFloat

    N = length(x)                           # length of vector x
    M = 50                                  # arbitrary large number M

    nbins = ceil(Int64, (30 * N)^(1/5))     # number of bins per histogram
    mbins = ceil(Int64, M/nbins)            # number of histograms calculated

    σ = std(x)                              # standard deviation of x
    # setup range for ASH
    s = 0.5
    δ = (maximum(x) - minimum(x) + σ)/((nbins+1)*mbins-1)
    rng = (minimum(x) - s*σ):δ:(maximum(x) + s*σ)
    # compute ASH
    epdf = ash(x, rng=rng, m=mbins, kernel=Kernels.triangular)
    ent = 0
    for k in 1:N
        ent -= (1/N) * log(AverageShiftedHistograms.pdf(epdf, x[k]))
    end
    return ent
end

function coefcost(x::AbstractArray{T,2}, et::DifferentialEntropyCost) where
        T<:Number

    cost = 0
    for i in axes(x,1)
        @inbounds cost += coefcost(x[i,:], DifferentialEntropyCost())
    end
    return cost
end

## BEST BASIS TYPES
"""
Abstract type for best basis. Current available types are:
- [`LSDB`](@ref)
- [`JBB`](@ref)
- [`BB`](@ref)
- [`SIBB`](@ref)
"""
abstract type BestBasisType end

"""
    LSDB([; cost=DifferentialEntropyCost(), redundant=false])

Least Statistically Dependent Basis (LSDB). Set `redundant=true` when running
LSDB with redundant wavelet transforms such as SWT or ACWT.

**See also:** [`BestBasisType`](@ref), [`JBB`](@ref), [`BB`](@ref),
    [`SIBB`](@ref)
"""
@with_kw struct LSDB <: BestBasisType
    cost::LSDBCost = DifferentialEntropyCost()
    redundant::Bool = false
end

"""
    JBB([; cost=LoglpCost(2), redundant=false])

Joint Best Basis (JBB). Set `redundant=true` when running JBB with redundant
wavelet transforms such as SWT or ACWT.

**See also:** [`BestBasisType`](@ref), [`LSDB`](@ref), [`BB`](@ref),
    [`SIBB`](@ref)
"""
@with_kw struct JBB <: BestBasisType    # Joint Best Basis
    cost::JBBCost = LoglpCost(2)
    redundant::Bool = false
end

"""
    BB([; cost=LoglpCost(2), redundant=false])

Best Basis (BB). Set `redundant=true` when running BB with redundant wavelet
transforms such as SWT or ACWT.

**See also:** [`BestBasisType`](@ref), [`LSDB`](@ref), [`JBB`](@ref),
    [`SIBB`](@ref)
"""
@with_kw struct BB <: BestBasisType     # Individual Best Basis
    cost::BBCost = ShannonEntropyCost()
    redundant::Bool = false
end

"""
    SIBB([; cost=ShannonEntropyCost()])

Shift Invariant Best Basis (SIBB).

**See also:** [`BestBasisType`](@ref), [`LSDB`](@ref), [`JBB`](@ref),
    [`BB`](@ref)
"""
@with_kw struct SIBB <: BestBasisType   # Shift invariant best basis
    cost::BBCost = ShannonEntropyCost()
end

## TREE COST
"""
    tree_costs(X, method)

Returns the cost of each node in a binary tree in order to find the best basis.

**See also:** [`bestbasistree`](@ref), [`bestbasis_treeselection`](@ref)
"""
function tree_costs(X::AbstractArray{T,3}, method::LSDB) where T<:AbstractFloat
    L = size(X, 2)
    n = size(X, 1)

    if method.redundant
        costs = Vector{T}(undef, L)
        for i in eachindex(costs)
            j = floor(Integer, log2(i))
            costs[i] = coefcost(X[:,i,:], method.cost) / (1<<j)
        end
    else
        costs = Vector{T}(undef, 2^L - 1)
        i = 1
        for lvl in 0:(L-1)
            n₀ = nodelength(n, lvl)
            for node in 0:(2^lvl-1)
                rng = (node * n₀ + 1):((node + 1) * n₀)
                costs[i] = coefcost(X[rng, lvl+1, :], method.cost)
                i += 1
            end
        end
    end
    return costs
end

function tree_costs(X::AbstractArray{T,3}, method::JBB) where T<:AbstractFloat
    # For each level, compute mean, sum of squares, and variance
    (n, L, N) = size(X)                     # L = levels if wpd, nodes if swpd
    EX = sum(X, dims = 3) ./ N              # calculate E(X)
    EX² = sum(X .^ 2, dims = 3) ./ N        # calculate E(X²)
    VarX = EX² - (EX) .^ 2                  # calculate Var(X)
    VarX = reshape(VarX, (n, L))
    σ = VarX .^ 0.5                         # calculate σ = √Var(X)
    @assert all(σ .>= 0)

    if method.redundant
        costs = Vector{T}(undef, L)
        for i in eachindex(costs)
            j = floor(Integer, log2(i))
            costs[i] = coefcost(σ[:, i], method.cost) / (1<<j)
        end
    else
        costs = Vector{T}(undef, 2^L - 1)
        i = 1   # iterates over the nodes for the costs variable
        for lvl in 0:(L-1)
            n₀ = nodelength(n, lvl)
            for node in 0:(2^lvl-1)
                rng = (node * n₀ + 1):((node + 1) * n₀)
                coef = σ[rng, lvl+1]
                costs[i] = coefcost(coef, method.cost)
                i += 1
            end
        end
    end
    return costs
end

function tree_costs(X::AbstractArray{T,2}, method::BB) where T<:AbstractFloat
    nrm = norm(X[:,1])
    L = size(X, 2)                      # count of levels if wpd, nodes if swpd
    n = size(X,1)

    if method.redundant                # swpd
        costs = Vector{T}(undef, L)
        for i in eachindex(costs)
            j = floor(Integer, log2(i))
            costs[i] = coefcost(X[:, i], method.cost, nrm) / (1<<j)
        end
    else
        costs = Vector{T}(undef, 2^L - 1)
        i = 1
        for lvl in 0:(L-1)
            n₀ = nodelength(n, lvl)
            for node in 0:(2^lvl-1)
                rng = (node * n₀ + 1):((node + 1) * n₀)
                costs[i] = coefcost(X[rng, lvl+1, :], method.cost, nrm)
                i += 1
            end
        end
    end
    return costs
end

function tree_costs(y::AbstractArray{T,2}, tree::AbstractVector{BitVector},
        method::SIBB) where T<:Number

    nn = length(tree)
    ns = size(y,1)
    @assert size(y,2) == nn
    tree_costs = Vector{Vector{Union{T,Nothing}}}(undef, nn)
    nrm = norm(y[:,1])

    for i in eachindex(tree)
        level = floor(Int, log2(i))
        len = nodelength(ns, level)
        # number of nodes corresponding to Ω(i,j)
        costs = Vector{Union{AbstractFloat,Nothing}}(nothing, length(tree[i]))
        for j in eachindex(tree[i])
            if tree[i][j]
                shift = j-1                     # current shift
                nstart = shift*len + 1
                nend = (shift+1) * len
                costs[j] = coefcost(y[nstart:nend,i], method.cost, nrm)
            end
        end
        tree_costs[i] = costs
    end
    return tree_costs
end

## BEST TREE SELECTION
"""
    bestbasis_treeselection(costs, n[, type=:min])

Computes the best tree based on the given cost vector.

**See also:** [`bestbasistree`](@ref), [`tree_costs`](@ref)
"""
function bestbasis_treeselection(costs::AbstractVector{T}, n::Integer,
        type::Symbol=:min) where T<:AbstractFloat

    @assert length(costs) <= 2*n - 1
    L = floor(Int, log2(length(costs)))
    bt = maketree(n, L, :full)
    if type == :min
        @inbounds begin
            for i in reverse(1:(1<<L-1))
                childcost = costs[getchildindex(i,:left)] + costs[getchildindex(i,:right)]
                if childcost < costs[i]     # child cost < parent cost
                    costs[i] = childcost
                else
                    delete_subtree!(bt, i)
                end
            end
        end
    elseif type == :max
        @inbounds begin
            for i in reverse(1:(1<<L-1))
                childcost = costs[getchildindex(i,:left)] + costs[getchildindex(i,:right)]
                if childcost > costs[i]     # child cost < parent cost
                    costs[i] = childcost
                else
                    delete_subtree!(bt, i)
                end
            end
        end
    else
        throw(ArgumentError("Accepted types are :min and :max only"))
    end
    return bt
end

function bestbasis_treeselection(costs::AbstractVector{Tc},
        tree::AbstractVector{Tt}) where
        {Tc<:AbstractVector{<:Union{Number,Nothing}}, Tt<:BitVector}

    @assert length(costs) == length(tree)
    bt = deepcopy(tree)
    bc = deepcopy(costs)
    nn = length(tree)
    for i in reverse(eachindex(bt))
        if getchildindex(i,:left) > nn                        # current node is at bottom level
            continue
        end
        level = floor(Int, log2(i))
        for j in eachindex(bt[i])       # iterate through all available shifts
            if !bt[i][j]                # node of current shift does not exist
                continue
            end
            @assert bt[getchildindex(i,:left)][j] ==
                    bt[getchildindex(i,:right)][j] ==
                    bt[getchildindex(i,:left)][j+1<<level] ==
                    bt[getchildindex(i,:right)][j+1<<level] == true || continue
            # child cost of current shift of current node
            cc = bc[getchildindex(i,:left)][j] + bc[getchildindex(i,:right)][j]
            # child cost of circshift of current node
            scc = bc[getchildindex(i,:left)][j+1<<level] +
                  bc[getchildindex(i,:right)][j+1<<level]
            mincost = min(cc, scc, bc[i][j])
            # mincost=parent, delete subtrees of both cc & scc
            if mincost == bc[i][j]
                delete_subtree!(bt, getchildindex(i,:left), j)
                delete_subtree!(bt, getchildindex(i,:right), j)
                delete_subtree!(bt, getchildindex(i,:left), j+1<<level)
                delete_subtree!(bt, getchildindex(i,:right), j+1<<level)
            # mincost=cc, delete subtrees of scc
            elseif mincost == cc
                bc[i][j] = mincost
                delete_subtree!(bt, getchildindex(i,:left), j+1<<level)
                delete_subtree!(bt, getchildindex(i,:right), j+1<<level)
            # mincost=scc, delete subtrees of cc
            else
                bc[i][j] = mincost
                delete_subtree!(bt, getchildindex(i,:left), j)
                delete_subtree!(bt, getchildindex(i,:right), j)
            end
        end
    end
    # ensure only 1 node selected from each Ω(i,j)
    @assert all(map(node -> sum(node), bt) .<= 1)
    return bt
end

# deletes subtree due to inferior cost
function delete_subtree!(bt::BitVector, i::Integer)
    @assert 1 <= i <= length(bt)
    bt[i] = false
    if (getchildindex(i,:left)) < length(bt)
        if bt[getchildindex(i,:left)]
            delete_subtree!(bt, getchildindex(i,:left))
        end
        if bt[getchildindex(i,:right)]
            delete_subtree!(bt, getchildindex(i,:right))
        end
    end
    return nothing
end

function delete_subtree!(bt::AbstractVector{BitVector}, i::Integer, j::Integer)
    @assert 1 <= i <= length(bt)
    level = floor(Int, log2(i))
    bt[i][j] = false
    if (getchildindex(i,:left)) < length(bt)          # current node can contain subtrees
        if bt[getchildindex(i,:left)][j]              # left child of current shift
            delete_subtree!(bt, getchildindex(i,:left), j)
        end
        if bt[getchildindex(i,:right)][j]            # right child of current shift
            delete_subtree!(bt, getchildindex(i,:right), j)
        end
        if bt[getchildindex(i,:left)][j+1<<level]     # left child of added shift
            delete_subtree!(bt, getchildindex(i,:left), j+1<<level)
        end
        if bt[getchildindex(i,:right)][j+1<<level]   # right child of added shift
            delete_subtree!(bt, getchildindex(i,:right), j+1<<level)
        end
    end
    return nothing
end

## BEST BASIS TREES
"""
    bestbasistree(X[, method])

Extension to the best basis tree function from Wavelets.jl. Given a set of
decomposed signals, returns different types of best basis trees based on the
methods specified. Available methods are the joint best basis ([`JBB`](@ref)),
least statistically dependent basis ([`LSDB`](@ref)), individual regular
best basis ([`BB`](@ref)), and shift-invariant best basis ([`SIBB`](@ref)).

# Examples
```julia
bestbasistree(X, JBB())

bestbasistree(X, SIBB())
```

**See also:** [`getbasiscoef`](@ref), [`getbasiscoefall`](@ref)
"""
function Wavelets.Threshold.bestbasistree(X::AbstractArray{T,3},
        method::LSDB) where T<:AbstractFloat

    costs = tree_costs(X, method)
    besttree = bestbasis_treeselection(costs, size(X,1))
    return besttree
end

function Wavelets.Threshold.bestbasistree(X::AbstractArray{T,3},
        method::JBB) where T<:AbstractFloat

    costs = tree_costs(X, method)
    besttree = bestbasis_treeselection(costs, size(X,1))
    return besttree
end

function Wavelets.Threshold.bestbasistree(X::AbstractArray{T,3},
        method::BB) where T<:AbstractFloat

    n = size(X,1)
    besttree = falses(n-1, size(X,3))
    @inbounds begin
        for i in axes(besttree,2)
            costs = tree_costs(X[:,:,i], method)
            besttree[:,i] = bestbasis_treeselection(costs, n)
        end
    end
    return besttree
end

function Wavelets.Threshold.bestbasistree(X::AbstractArray{T,2},
        method::BB) where T<:AbstractFloat

    costs = tree_costs(X, method)
    besttree = bestbasis_treeselection(costs, size(X,1))
    return besttree
end

function Wavelets.Threshold.bestbasistree(y::AbstractArray{T,2}, d::Integer,
        method::SIBB) where T<:Number                                           # TODO: find a way to compute bestbasis_tree without input d

    nn = size(y,2)
    L = maxtransformlevels((nn+1)÷2)
    ns = size(y,1)
    tree = makesiwpdtree(ns, L, d)
    costs = tree_costs(y, tree, method)
    besttree = bestbasis_treeselection(costs, tree)
    return besttree
end

function Wavelets.Threshold.bestbasistree(X::AbstractArray{T,3};
        method::BestBasisType=JBB()) where T<:AbstractFloat
    return bestbasistree(X, method)
end

end # end of module


using CSV, DataFrames, Plots, LinearAlgebra, Interact, StatsBase, Statistics

f(x) = typeof(x) <: Number ? Float32(x) : parse(Float32, x)

#Load data, clean data, recast data
kaggleData = CSV.read("/Users/javier/Desktop/SynData/synthetic-derivative/Dataset/kag_risk_factors_cervical_cancer.csv", DataFrame)
data = Array(kaggleData)
data[data .== "?"] .= -1
data = f.(data)

#Load data, clean data, recast Syn data
synData = CSV.read("/Users/javier/Desktop/SynData/synthetic-derivative/Dataset/synData-3.csv", DataFrame)
synData = CSV.read("/Users/javier/Desktop/SynData/synthetic-derivative/Dataset/synData-4.csv", DataFrame)
dataSyn = Array{Int}(synData)
# data[data .== "?"] .= -1
dataSyn = Array{Int}(f.(dataSyn[1:5000,:]))

#PC
function doPCA(arr)
    l = size(arr,1)
    # mean2=[arr[j,i]-sum(arr[:,i])/(l) for j=1:l,i=1:size(arr,2)]
    mean2=[arr[j,i]-mean(arr[:,i]) for j=1:l,i=1:size(arr,2)]

    #Cov matrix
    m=transpose(mean2)*mean2/l
    # plot(m, seriestype=:heatmap)

    #Eigenvalues and Eigenvectors
    d,v = eigen(m)
    dNorm=d/maximum(d)

    m, dNorm, v
end

mReal, dReal, vReal = doPCA(data)
mSyn, dSyn, vSyn = doPCA(dataSyn)
mean(abs2, mReal .- mSyn)
#Plot eigenvalues
plot(dReal, dSyn, ms=10, markershapes = [:circle], markerstrokewidth=0, lw=1, size=(500,500))
plot!(0:1,x->x, color=:red)

#Plot eigenvectors
@manipulate for i in 1:size(vReal,2)
    plot(abs.(vReal[:,i]), abs.(vSyn[:,i]), frame=:box, xrotation=0, ms=5, markershapes = [:circle], markerstrokewidth=0, lw=0, size=(500,500), tickfont = font(12, "Helvetica"), legendfont = font(12, "Helvetica"), label="$(dReal[i])", legend=(0.15,0.9), xlim=(-0.1,1), ylim=(-0.1,1), rightmargin = 5Plots.mm)
    plot!(0:1, x->x)
end

#Plot histogram
@manipulate for i in 1:size(dataSyn,2)
    lab = names(kaggleData)[i]
    plot(Int.(dataSyn[:,i]), seriestype=:barhist, frame=:box, width=0, opacity=0.5, normalize=true, title="$lab", label="Synth", tickfont = font(12, "Helvetica"), xlabel="Values", ylabel="Freq")
    plot!(data[:,i], seriestype=:barhist, width=0, opacity=0.5, normalize=true, label="Real")
end
typeof(dataSyn)
#=Plot eigenvector inner product as heatmap
    A = ∑|λᵢ⟩⟨cᵢ| and B = ∑|β⟩⟨c|
    Then, C = Aᵀ⋅B = ∑⟨λᵢ|βⱼ⟩|cᵢ⟩⟨cᵢ|
    Notice that ⟨λᵢ|βⱼ⟩ is the correlation between the eigenvector of the synthetic
    data covariance matrix (SDCM) and the real data covariance matrix (RDCM). Hence, a heatmap
    of C shows the correlation between all eigenvectors from the SDCM with the RDCM
    one-on-one.
=#
plot(vSyn' * vSyn, seriestype=:heatmap, xticks=(0:35,reverse(names(kaggleData))), yticks=(0:35,reverse(names(kaggleData))), xrotation=80, bottommargin = 30Plots.mm,  size=(700,700))
plot(vReal' * vSyn, seriestype=:heatmap, xticks=(0:35,reverse(names(kaggleData))), yticks=(0:35,reverse(names(kaggleData))), xrotation=80, bottommargin = 30Plots.mm,  size=(700,700))
plot(abs.(vReal' * vSyn), seriestype=:heatmap, xticks=(0:35,reverse(names(kaggleData))), yticks=(0:35,reverse(names(kaggleData))), xrotation=80, bottommargin = 30Plots.mm,  size=(700,700))

#=
    From the previous, we now compute the eigenvectors of C. Notice that if C corresponded to a completely correlated system, i.e., ⟨λᵢ|βⱼ⟩ = δᵢⱼ, then C = Identity. In that case, the spectrum is degenerate and
    the eigenvalues all equal 1. In our case, notice the eigenvalues are complex, yet the absolute value
    are all close to 1.
=#
dCorr, vCorr = eigen(vReal' * vSyn)
dCorr, vCorr = eigen(vReal' * vReal .- 0.001*randn(36,36))
dCorr, vCorr = eigen(vReal' * randn(36,36))
dCorr, vCorr = eigen(vReal' * vSyn[sample(collect(1:36),36, replace=false),:])
dCorr, vCorr = eigen(zeros(36,36))

plot(abs.(dCorr))

@manipulate for i in 1:size(vReal,2)
    plot(abs.(vReal[:,i]), frame=:box, xrotation=80, ms=10, markershapes = [:circle], markerstrokewidth=0, lw=1, size=(500,500), tickfont = font(12, "Helvetica"), legendfont = font(12, "Helvetica"), label="$(dReal[i])", legend=(0.75,0.9), xticks=(0:35,names(kaggleData)), bottommargin = 40Plots.mm)
end

# using HTTP, JSON, BSON, Dates
#
# a = HTTP.request("GET", "https://data-live.flightradar24.com/zones/fcgi/feed.js")
#
# a.body
#
# aa = JSON.parse(String(a.body))

#TO EXPORT
PATH = "/Users/javier/Dropbox/Aplicaciones/Overleaf/SynthData/Figs/"

begin
    fig = plot([abs.(vReal[:,i]) for i in 31:36], [abs.(vSyn[:,i]) for i in 31:36], frame=:box, xrotation=0, ms=5, markershapes = [:circle], markerstrokewidth=0, lw=0, size=(500,500), tickfont = font(12, "Helvetica"), legendfont = font(12, "Helvetica"), label="", legend=(0.15,0.9), xlim=(-0.1,1), ylim=(-0.1,1), rightmargin = 5Plots.mm, xlabel="Real Data Covariance Matrix Eigenvectors", ylabel="Synth Data Covariance Matrix Eigenvectors", title="5 largest eigenvalues")
    fig = plot!(0:1, x->x, legend=:none, lw=2, c=:black)
    savefig(fig, PATH * "SR-DCM-EigenVectors.png")
end

begin
    fig = plot(dReal, dSyn, ms=10, markershapes = [:circle], markerstrokewidth=0, lw=1, size=(500,500), tickfont = font(12, "Helvetica"), frame=:box, legend=:topleft, label="Eigenvalues", legendfont = font(12, "Helvetica"), xlabel="Real Data Covariance Matrix", ylabel="Synth Data Covariance Matrix")
    fig = plot!(0:1,x->x, color=:black, lw=2, label="y=x")
    savefig(fig, PATH * "SR-DCM-EigenValues.png")
end

s(x) = split(x, " ")[1]
begin
    for i in 1:size(dataSyn,2)
        lab = names(kaggleData)[i]
        name = "$i-" * s.(names(kaggleData))[i]
        fig = plot(Int.(dataSyn[:,i]), seriestype=:barhist, frame=:box, width=0, opacity=0.5, normalize=true, title="$lab", label="Synth", tickfont = font(12, "Helvetica"), xlabel="Values", ylabel="Freq")
        fig = plot!(data[:,i], seriestype=:barhist, width=0, opacity=0.5, normalize=true, label="Real")
        savefig(fig, PATH * "Hist-$name.png")
    end
end

begin
    fig = plot(abs.(vReal' * vSyn), seriestype=:heatmap, xticks=(0:35,reverse(names(kaggleData))), yticks=(0:35,reverse(names(kaggleData))), xrotation=80, bottommargin = 30Plots.mm,  size=(700,700))
    savefig(fig, PATH * "SandR-DCMInnerProduct.png")
    fig = plot(vSyn' * vSyn, seriestype=:heatmap, xticks=(0:35,reverse(names(kaggleData))), yticks=(0:35,reverse(names(kaggleData))), xrotation=80, bottommargin = 30Plots.mm,  size=(700,700))
    savefig(fig, PATH * "SandS-DCMInnerProduct.png")
    fig = plot(abs.(vReal' * (vReal .- 0.4*randn(36,36))), seriestype=:heatmap, xticks=(0:35,reverse(names(kaggleData))), yticks=(0:35,reverse(names(kaggleData))), xrotation=80, bottommargin = 30Plots.mm,  size=(700,700))
    savefig(fig, PATH * "R-DCMInnerProductwNoise.png")
end

begin
    dCorr, vCorr = eigen(vReal' * vSyn)
    fig = plot(sort(abs.(dCorr), rev=true), markershapes = [:circle], markerstrokewidth=0, lw=1,
    frame=:box, legend=:none, xlabel="Eigenvalues", ylabel="Amplitude")
    savefig(fig, PATH * "eigenvaluesForSRDCMInnerproduct.png")
    dCorr, vCorr = eigen(vReal' * vReal .- 0.001*randn(36,36))
    fig = plot(sort(abs.(dCorr), rev=true), markershapes = [:circle], markerstrokewidth=0, lw=1,
    frame=:box, legend=:none, xlabel="Eigenvalues", ylabel="Amplitude")
    savefig(fig, PATH * "eigenvaluesForRDCMInnerproductwNoise.png")
    dCorr, vCorr = eigen(vReal' * vReal .- 0.000*randn(36,36))
    fig = plot(sort(abs.(dCorr), rev=true), markershapes = [:circle], markerstrokewidth=0, lw=1,
    frame=:box, legend=:none, xlabel="Eigenvalues", ylabel="Amplitude")
    savefig(fig, PATH * "eigenvaluesForRDCMInnerproduct.png")
    dCorr, vCorr = eigen(ones(36,36))
    fig = plot(sort(abs.(dCorr), rev=true), markershapes = [:circle], markerstrokewidth=0, lw=1,
    frame=:box, legend=:none, xlabel="Eigenvalues", ylabel="Amplitude")
    savefig(fig, PATH * "eigenvaluesForOnes.png")
    # dCorr, vCorr = eigen(vReal' * randn(36,36))
    # dCorr, vCorr = eigen(vReal' * vSyn[sample(collect(1:36),36, replace=false),:])

end
#################################
#PRIVACY
##################################
sample(1:10, 11, replace=false)
function replaceRow(ar1, row, idx)
    ar3 = copy(ar1)
    ar3[idx,:] = row
    return ar3
end

function replaceRows(ar1,ar2, nrows)
    samp1 = sample(1:size(ar1,1),nrows, replace=false)
    samp2 = sample(1:size(ar2,1),nrows, replace=false)
    # err = Vector{Float32}(undef, size(ar1,1))
    temp = copy(ar1)
    for i in 1:nrows
        temp = replaceRow(temp, ar2[samp2[i],:],samp1[i])
        # err[i] = mean(abs, ar1 .- temp)
    end
    temp
end

data1 = replaceRow(data, dataSyn[1,:],1)

err = [sum(abs, data .- replaceRows(data,dataSyn, j)) for i in 1:500, j in 0:100]
reshape(mean(err, dims=1),:)
plot(reshape(mean(err, dims=1),:), ribbon=reshape(std(err, dims=1),:))

mReal, dReal, vReal = doPCA(data)
mSyn, dSyn, vSyn = doPCA(replaceRows(data,dataSyn, 100))

@manipulate for nrows in 1:858
    mSyn, dSyn, vSyn = doPCA(replaceRows(data,dataSyn, nrows))
    plot(dReal, dSyn, ms=10, markershapes = [:circle], markerstrokewidth=0, lw=1, size=(500,500))
    plot!(0:1,x->x, color=:red)
end