# # Nonlinear constrained optimization
#
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`ipnewton_basics.ipynb`](@__NBVIEWER_ROOT_URL__examples/generated/ipnewton_basics.ipynb)
#-
#
# The nonlinear constrained optimization interface in
# `Optim` assumes that the user can write the optimization
# problem in the following way.

# ```math
# \min_{x\in\mathbb{R}^n} f(x) \quad \text{such that}\\
# l_x \leq \phantom{c(}x\phantom{)} \leq u_x \\
# l_c \leq c(x) \leq u_c.
# ```

# For equality constraints on ``x_j`` or ``c(x)_j`` you set those
# particular entries of bounds to be equal, ``l_j=u_j``.
# Likewise, setting ``l_j=-\infty`` or ``u_j=\infty`` means that the
# constraint is unbounded from below or above respectively.

using Optim, NLSolversBase #hide
import NLSolversBase: clear! #hide

# # Constrained optimization with `IPNewton`

# We will go through examples on how to use the constraints interface
# with the interior-point Newton optimization algorithm [IPNewton](../../algo/ipnewton.md).

# Throughout these examples we work with the standard Rosenbrock function.
# The objective and its derivatives are given by

fun(x) =  (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

function fun_grad!(g, x)
g[1] = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
g[2] = 200.0 * (x[2] - x[1]^2)
end

function fun_hess!(h, x)
h[1, 1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2
h[1, 2] = -400.0 * x[1]
h[2, 1] = -400.0 * x[1]
h[2, 2] = 200.0
end;

# ## Optimization interface
# To solve a constrained optimization problem we call the `optimize`
# method
# ``` julia
# optimize(d::AbstractObjective, constraints::AbstractConstraints, initial_x::Tx, method::ConstrainedOptimizer, options::Options)
# ```

# We can create instances of `AbstractObjective` and
# `AbstractConstraints` using the types `TwiceDifferentiable` and
# `TwiceDifferentiableConstraints` from the package `NLSolversBase.jl`.

# ## Box minimization
# We want to optimize the Rosenbrock function in the box
# ``-0.5 \leq x \leq 0.5``, starting from the point ``x_0=(0,0)``.
# Box constraints are defined using, for example,
# `TwiceDifferentiableConstraints(lx, ux)`.

x0 = [0.0, 0.0]
df = TwiceDifferentiable(fun, fun_grad!, fun_hess!, x0)

lx = [-0.5, -0.5]; ux = [0.5, 0.5]
dfc = TwiceDifferentiableConstraints(lx, ux)

res = optimize(df, dfc, x0, IPNewton())
## Test the results             #src
using Test                 #src
@test Optim.converged(res)      #src
@test Optim.minimum(res) ≈ 0.25 #src

# Like the rest of Optim, you can also use `autodiff=:forward` and just pass in
# `fun`.

# If we only want to set lower bounds, use `ux = fill(Inf, 2)`

ux = fill(Inf, 2)
dfc = TwiceDifferentiableConstraints(lx, ux)

clear!(df)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                   #src
@test Optim.minimum(res) < 0.0 + sqrt(eps()) #src

# ## Defining "unconstrained" problems

# An unconstrained problem can be defined either by passing
# `Inf` bounds or empty arrays.
# **Note that we must pass the correct type information to the empty `lx` and `ux`**

lx = fill(-Inf, 2); ux = fill(Inf, 2)
dfc = TwiceDifferentiableConstraints(lx, ux)

clear!(df)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                   #src
@test Optim.minimum(res) < 0.0 + sqrt(eps()) #src

lx = Float64[]; ux = Float64[]
dfc = TwiceDifferentiableConstraints(lx, ux)

clear!(df)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                   #src
@test Optim.minimum(res) < 0.0 + sqrt(eps()) #src

# ## Generic nonlinear constraints

# We now consider the Rosenbrock problem with a constraint on
# ```math
#    c(x)_1 = x_1^2 + x_2^2.
# ```

# We pass the information about the constraints to `optimize`
# by defining a vector function `c(x)` and its Jacobian `J(x)`.

# The Hessian information is treated differently, by considering the
# Lagrangian of the corresponding slack-variable transformed
# optimization problem. This is similar to how the [CUTEst
# library](https://github.com/JuliaSmoothOptimizers/CUTEst.jl) works.
# Let ``H_j(x)`` represent the Hessian of the ``j``th component
# ``c(x)_j`` of the generic constraints.
# and ``\lambda_j`` the corresponding dual variable in the
# Lagrangian. Then we want the `constraint` object to
# add the values of ``H_j(x)`` to the Hessian of the objective,
# weighted by ``\lambda_j``.

# The Julian form for the supplied function ``c(x)`` and the derivative
# information is then added in the following way.

con_c!(c, x) = (c[1] = x[1]^2 + x[2]^2; c)
function con_jacobian!(J, x)
    J[1,1] = 2*x[1]
    J[1,2] = 2*x[2]
    J
end
function con_h!(h, x, λ)
    h[1,1] += λ[1]*2
    h[2,2] += λ[1]*2
end;

# **Note that `con_h!` adds the `λ`-weighted Hessian value of each
# element of `c(x)` to the Hessian of `fun`.**

# We can then optimize the Rosenbrock function inside the ball of radius
# ``0.5``.

lx = Float64[]; ux = Float64[]
lc = [-Inf]; uc = [0.5^2]
dfc = TwiceDifferentiableConstraints(con_c!, con_jacobian!, con_h!,
                                     lx, ux, lc, uc)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                    #src
@test Optim.minimum(res) ≈ 0.2966215688829263 #src

# We can add a lower bound on the constraint, and thus
# optimize the objective on the annulus with
# inner and outer radii ``0.1`` and ``0.5`` respectively.

lc = [0.1^2]
dfc = TwiceDifferentiableConstraints(con_c!, con_jacobian!, con_h!,
                                     lx, ux, lc, uc)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                    #src
@test Optim.minimum(res) ≈ 0.2966215688829255 #src

# **Note that the algorithm warns that the Initial guess is not an
# interior point.** `IPNewton` can often handle this, however, if the
# initial guess is such that `c(x) = u_c`, then the algorithm currently
# fails. We may fix this in the future.

# ## Multiple constraints
# The following example illustrates how to add an additional constraint.
# In particular, we add a constraint function
# ```math
#    c(x)_2 = x_2\sin(x_1)-x_1
# ```

function con2_c!(c, x)
    c[1] = x[1]^2 + x[2]^2     ## First constraint
    c[2] = x[2]*sin(x[1])-x[1] ## Second constraint
    c
end
function con2_jacobian!(J, x)
    ## First constraint
    J[1,1] = 2*x[1]
    J[1,2] = 2*x[2]
    ## Second constraint
    J[2,1] = x[2]*cos(x[1])-1.0
    J[2,2] = sin(x[1])
    J
end
function con2_h!(h, x, λ)
    ## First constraint
    h[1,1] += λ[1]*2
    h[2,2] += λ[1]*2
    ## Second constraint
    h[1,1] += λ[2]*x[2]*-sin(x[1])
    h[1,2] += λ[2]*cos(x[1])
    ## Symmetrize h
    h[2,1]  = h[1,2]
    h
end;

# We generate the constraint objects and call `IPNewton` with
# initial guess ``x_0 = (0.25,0.25)``.

x0 = [0.25, 0.25]
lc = [-Inf, 0.0]; uc = [0.5^2, 0.0]
dfc = TwiceDifferentiableConstraints(con2_c!, con2_jacobian!, con2_h!,
                                     lx, ux, lc, uc)
res = optimize(df, dfc, x0, IPNewton())
@test Optim.converged(res)                                       #src
@test Optim.minimum(res) ≈ 1.0                                   #src
@test isapprox(Optim.minimizer(res), zeros(2), atol=sqrt(eps())) #src

#md # ## [Plain Program](@id ipnewton_basics-plain-program)
#md #
#md # Below follows a version of the program without any comments.
#md # The file is also available here: [ipnewton_basics.jl](ipnewton_basics.jl)
#md #
#md # ```julia
#md # @__CODE__
#md # ```


# Inefficient: creates copies, no in place
function mergesorted(sortedarray1::Vector{T}, sortedarray2::Vector{T})::Vector{T} where T <: Real
	n = length(sortedarray1)
	m = length(sortedarray2)
	result = Vector{Int}(undef, n + m)
	i = 1
	j = 1

	while i ≤ n && j ≤ m
		if sortedarray1[i] ≤ sortedarray2[j]
			result[i + j - 1] = sortedarray1[i]
			i += 1
		else
			result[i + j - 1] = sortedarray2[j]
			j += 1
		end
	end

	if i > n
		result[i + j - 1 : end] = sortedarray2[j:end]
	else
		result[i + j - 1 : end] = sortedarray1[i:end]
	end

	return result
end

function mergesort(array::Vector{T})::Vector{T} where T <: Real
	# Split
	len = length(array)
	len == 1 && (return array)

	# Sort
	sortedleft  = mergesort(array[1:(len ÷ 2)])
	sortedright = mergesort(array[(len ÷ 2 + 1):end])

	# Combine
	return mergesorted(sortedleft, sortedright)
end

println(mergesort([6, 4, 1, 9, 8, 2, 6, 4]))


# Load Julia packages (libraries) needed  for the snippets in chapter 0

using StatisticalRethinking, CmdStan
#gr(size=(600,600));

# CmdStan uses a tmp directory to store the output of cmdstan

ProjDir = rel_path("..", "scripts", "04")
cd(ProjDir)

# ### snippet 4.7

howell1 = CSV.read(rel_path("..", "data", "Howell1.csv"), delim=';')
df = convert(DataFrame, howell1);

# ### snippet 4.8

# Use only adults

df2 = filter(row -> row[:age] >= 18, df);

# Show first 5 rows of DataFrame df

first(df2, 5)

# ### Snippet 4.14

# Generate approximate probabilities

struct Post
  mu::Float64
  sigma::Float64
  ll:: Float64
  prod::Float64
  prob::Float64
end

mu_list = repeat(range(140, 160, length=200), 200);
sigma_list = repeat(range(4, 9, length=200), inner=200);

ll = zeros(40000);
for i in 1:40000
    d1 = Normal(mu_list[i], sigma_list[i])
    ll[i] = sum(log.(pdf.(d1, df2[:height])))
end

d2 = Normal(178.0, 20.0)
d3 = Uniform(0, 50)
prod = ll + log.(pdf.(d2, mu_list)) + log.(pdf.(d3, sigma_list))
prob = exp.(prod .- maximum(prod))
post = DataFrame(mu=mu_list, sigma=sigma_list, ll=ll, prod=prod, prob=prob)
first(post, 10)

# ### Snippet 4.15

# Sample post

samples = post[sample(1:size(post, 1), Weights(post[:prob]), 10000, replace=true), :]

# ### Snippet 4.19

# Density of mu

density(samples[:mu])

# Density of sigma

density(samples[:sigma])

# ### Snippet 4.20

# Hdp mu

#hpd(samples[:mu])

# Hdp sigma

#hpd(samples[:sigma])

# End of `clip-14-20.jl`


# ------------------------------------------------------------------------------
# convenience utilities
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# compute symbol matrix
# ------------------------------------------------------------------------------

"""
```julia
computesymbolsoverrange(operator, θ; mass=nothing, θ_min=-π/2)
```

Compute the eigenvalues and eigenvectors of the symbol matrix for an operator over
  a range of θ from -π/2 to 3π/2

# Arguments:
- `operator`:       Finite element operator to compute symbol matrices for
- `numbersteps1d`:  Number of values of θ to sample in each dimension
                        Note: numbersteps1d^dimension symbol matrices will be computed

# Keyword Arguments:
- `mass`:   Mass operator to invert for comparison to analytic solution,
                default: nothing
- `θ_min`:  Bottom of range of θ, shifts range to [θ_min, θ_min + 2π],
                default: -π / 2

# Returns:
- Values of θ sampled
- Eigenvalues of symbol matrix at θ sampled
- Eigenvectors of symbol matrix at θ sampled

# Examples:
```jldoctest
using LFAToolkit;

numbersteps1d = 5

for dimension in 1:3
    # setup
    mesh = []
    if dimension == 1
        mesh = Mesh1D(1.0);
    elseif dimension == 2
        mesh = Mesh2D(1.0, 1.0);
    elseif dimension == 3
        mesh = Mesh3D(1.0, 1.0, 1.0);
    end
    diffusion = GalleryOperator("diffusion", 3, 3, mesh);

    # compute symbols
    (_, eigenvalues, _) = computesymbolsoverrange(diffusion, numbersteps1d);

    # verify
    eigenvalues = real(eigenvalues);
    if dimension == 1
        @assert minimum(eigenvalues[4, :]) ≈ 1
        @assert maximum(eigenvalues[4, :]) ≈ 4/3
    elseif dimension == 2
        @assert minimum(eigenvalues[19, :]) ≈ 2/3
        @assert maximum(eigenvalues[19, :]) ≈ 64/45
    elseif dimension == 3
        @assert minimum(eigenvalues[94, :]) ≈ 1/3
        @assert maximum(eigenvalues[94, :]) ≈ 256/225
    end
end

# output

```

```jldoctest
using LFAToolkit;

# setup
numbersteps1d = 3
mesh = Mesh2D(0.5, 0.5)
p = 6

# operators
mass = GalleryOperator("mass", p + 1, p + 1, mesh)
diffusion = GalleryOperator("diffusion", p + 1, p + 1, mesh)

# compute symbols
(_, eigenvalues, _) = computesymbolsoverrange(
    diffusion,
    numbersteps1d;
    mass = mass,
    θ_min = -π
);
eigenvalues = real(eigenvalues);

@assert minimum(eigenvalues[2, :]) ≈ π^2

# output

```
"""
function computesymbolsoverrange(
    operator::Operator,
    numbersteps1d::Int;
    mass::Union{Operator,Nothing} = nothing,
    θ_min::Float64 = -π / 2,
)
    # setup range
    dimension = operator.dimension
    numbersteps = numbersteps1d^dimension
    θ_max = θ_min + 2π
    θ_range1d = LinRange(θ_min, θ_max, numbersteps1d)
    θ_range = zeros(numbersteps, dimension)

    # setup
    numbermodes = size(operator.rowmodemap)[1]
    eigenvalues = zeros(ComplexF64, numbersteps, numbermodes)
    eigenvectors = zeros(ComplexF64, numbersteps, numbermodes, numbermodes)
    rangeiterator = []
    if dimension == 1
        rangeiterator = θ_range1d
    elseif dimension == 2
        rangeiterator = Iterators.product(θ_range1d, θ_range1d)
    elseif dimension == 3
        rangeiterator = Iterators.product(θ_range1d, θ_range1d, θ_range1d)
    end

    # compute
    for (step, θ) in enumerate(rangeiterator)
        θ = collect(θ)
        A = computesymbols(operator, θ)
        if mass != nothing
            A = computesymbols(mass, θ) \ A
        end
        currenteigen = eigen(A)
        eigenvalues[step, :] = currenteigen.values
        eigenvectors[step, :, :] = currenteigen.vectors
        θ_range[step, :] = θ
    end

    # return
    return (θ_range, eigenvalues, eigenvectors)
end

"""
```julia
computesymbolsoverrange(preconditioner, ω, θ; mass=nothing, θ_min=-π/2)
```

Compute the eigenvalues and eigenvectors of the symbol matrix for a preconditioned
  operator over a range of θ from -π/2 to 3π/2

# Arguments:
- `preconditioner`:  Preconditioner to compute symbol matries for
- `ω`:               Smoothing parameter array
- `numbersteps1d`:   Number of values of θ to sample in each dimension
                         Note: numbersteps1d^dimension symbol matrices will be computed

# Keyword Arguments:
- `mass`:   Mass operator to invert for comparison to analytic solution,
                default: nothing
- `θ_min`:  Bottom of range of θ, shifts range to [θ_min, θ_min + 2π],
                default: -π / 2

# Returns:
- Values of θ sampled
- Eigenvalues of symbol matrix at θ sampled
- Eigenvectors of symbol matrix at θ sampled

# Example:
```jldoctest
using LFAToolkit;

numbersteps1d = 5

for dimension in 1:3
    # setup
    mesh = []
    if dimension == 1
        mesh = Mesh1D(1.0);
    elseif dimension == 2
        mesh = Mesh2D(1.0, 1.0);
    elseif dimension == 3
        mesh = Mesh3D(1.0, 1.0, 1.0);
    end
    diffusion = GalleryOperator("diffusion", 3, 3, mesh);

    # preconditioner
    chebyshev = Chebyshev(diffusion);

    # compute symbols
    (_, eigenvalues, _) = computesymbolsoverrange(chebyshev, [1], numbersteps1d);

    # verify
    eigenvalues = real(eigenvalues);
    if dimension == 1
        @assert minimum(eigenvalues[4, :]) ≈ 0.15151515151515105
        @assert maximum(eigenvalues[4, :]) ≈ 0.27272727272727226
    elseif dimension == 2
        @assert minimum(eigenvalues[19, :]) ≈ -0.25495098334134725
        @assert maximum(eigenvalues[19, :]) ≈ -0.17128758445192374
    elseif dimension == 3
        @assert minimum(eigenvalues[94, :]) ≈ -0.8181818181818181
        @assert maximum(eigenvalues[94, :]) ≈ -0.357575757575757
    end
end

# output

```
"""
function computesymbolsoverrange(
    preconditioner::AbstractPreconditioner,
    ω::Array,
    numbersteps1d::Int;
    mass::Union{Operator,Nothing} = nothing,
    θ_min::Float64 = -π / 2,
)
    # setup range
    dimension = preconditioner.operator.dimension
    numbersteps = numbersteps1d^dimension
    θ_max = θ_min + 2π
    θ_range1d = LinRange(θ_min, θ_max, numbersteps1d)
    θ_range = zeros(numbersteps, dimension)

    # setup
    numbermodes = size(preconditioner.operator.rowmodemap)[1]
    eigenvalues = zeros(ComplexF64, numbersteps, numbermodes)
    eigenvectors = zeros(ComplexF64, numbersteps, numbermodes, numbermodes)
    rangeiterator = []
    if dimension == 1
        rangeiterator = θ_range1d
    elseif dimension == 2
        rangeiterator = Iterators.product(θ_range1d, θ_range1d)
    elseif dimension == 3
        rangeiterator = Iterators.product(θ_range1d, θ_range1d, θ_range1d)
    end

    # compute
    for (step, θ) in enumerate(rangeiterator)
        θ = collect(θ)
        A = computesymbols(preconditioner, ω, θ)
        if mass != nothing
            A = computesymbols(mass, θ) \ A
        end
        currenteigen = eigen(A)
        eigenvalues[step, :] = currenteigen.values
        eigenvectors[step, :, :] = currenteigen.vectors
        θ_range[step, :] = θ
    end

    # return
    return (θ_range, eigenvalues, eigenvectors)
end

"""
```julia
computesymbolsoverrange(multigrid, ω, v, θ; mass=nothing, θ_min=-π/2)
```

Compute the eigenvalues and eigenvectors of the symbol matrix for a multigrid
  preconditioned operator over a range of θ from -π/2 to 3π/2

# Arguments:
- `multigrid`:      Preconditioner to compute symbol matries for
- `p`:              Smoothing parameter array
- `v`:              Pre and post smooths iteration count array, 0 indicates no pre or post smoothing
- `numbersteps1d`:  Number of values of θ to sample in each dimension
                       Note: numbersteps1d^dimension symbol matrices will be computed

# Keyword Arguments:
- `mass`:   Mass operator to invert for comparison to analytic solution,
                default: nothing
- `θ_min`:  Bottom of range of θ, shifts range to [θ_min, θ_min + 2π],
                default: -π / 2

# Returns:
- Values of θ sampled
- Eigenvalues of symbol matrix at θ sampled
- Eigenvectors of symbol matrix at θ sampled

# Example:
```jldoctest
using LFAToolkit;

numbersteps1d = 5

using LinearAlgebra

for dimension in 1:3
    # setup
    mesh = []
    if dimension == 1
        mesh = Mesh1D(1.0);
    elseif dimension == 2
        mesh = Mesh2D(1.0, 1.0);
    elseif dimension == 3
        mesh = Mesh3D(1.0, 1.0, 1.0);
    end
    ctofbasis = TensorH1LagrangeBasis(3, 5, 1, dimension; collocatedquadrature=true);

    # operators
    finediffusion = GalleryOperator("diffusion", 5, 5, mesh);
    coarsediffusion = GalleryOperator("diffusion", 3, 5, mesh);

    # smoother
    jacobi = Jacobi(finediffusion);

    # preconditioner
    multigrid = PMultigrid(finediffusion, coarsediffusion, jacobi, [ctofbasis]);

    # compute symbols
    (_, eigenvalues, _) = computesymbolsoverrange(multigrid, [1.0], [1, 1], numbersteps1d);

    # verify
    eigenvalues = real(eigenvalues)
    if dimension == 1
       @assert maximum(eigenvalues[4, :]) ≈ 0.64
    elseif dimension == 2
       @assert maximum(eigenvalues[19, :]) ≈ 0.9082562365654528
    elseif dimension == 3
       @assert maximum(eigenvalues[94, :]) ≈ 1.4359882222222669
    end
end

# output

```
"""
function computesymbolsoverrange(
    multigrid::Multigrid,
    p::Array,
    v::Array,
    numbersteps1d::Int;
    mass::Union{Operator,Nothing} = nothing,
    θ_min::Float64 = -π / 2,
)
    # setup range
    dimension = multigrid.fineoperator.dimension
    numbersteps = numbersteps1d^dimension
    θ_max = θ_min + 2π
    θ_range1d = LinRange(θ_min, θ_max, numbersteps1d)
    θ_range = zeros(numbersteps, dimension)

    # setup
    numbermodes = size(multigrid.fineoperator.rowmodemap)[1]
    eigenvalues = zeros(ComplexF64, numbersteps, numbermodes)
    eigenvectors = zeros(ComplexF64, numbersteps, numbermodes, numbermodes)
    rangeiterator = []
    if dimension == 1
        rangeiterator = θ_range1d
    elseif dimension == 2
        rangeiterator = Iterators.product(θ_range1d, θ_range1d)
    elseif dimension == 3
        rangeiterator = Iterators.product(θ_range1d, θ_range1d, θ_range1d)
    end

    # compute
    for (step, θ) in enumerate(rangeiterator)
        θ = [abs(θ_i) > 100 * eps() ? θ_i : 100 * eps() for θ_i in θ]
        A = computesymbols(multigrid, p, v, θ)
        if mass != nothing
            A = computesymbols(mass, θ) \ A
        end
        currenteigen = eigen(A)
        eigenvalues[step, :] = currenteigen.values
        eigenvectors[step, :, :] = currenteigen.vectors
        θ_range[step, :] = θ
    end

    # return
    return (θ_range, eigenvalues, eigenvectors)
end

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


#=

Utility functions for Tapestree

Ignacio Quintero Mächler

February 6 2017

t(-_-t)

=#

"""
    rowind(x::Int64, nrow::Int64)

Get row indexing from matrix indexing.
"""
rowind(x::Int64, nrow::Int64) = mod1(x,nrow)

"""
    colind(x::Int64, nrow::Int64)

Get column indexing from matrix indexing
"""
colind(x::Int64, nrow::Int64) = cld(x, nrow)

"""
    vecind(row::Int64, col::Int64, nrow::Int64)

Get vector indexing from column and row.
"""
vecind(row::Int64, col::Int64, nrow::Int64) = row + nrow*(col - 1)

"""
    uniroot(f, approx = 1e-8, a = 0.0, b = 0.1)

Find the root of function between `0.0` and `b`.
"""
function uniroot(f; approx = 1e-8, a = 0.0, b = 0.1)
  # choose b
  while sign(f(a)::Float64)::Float64 == sign(f(b)::Float64)::Float64
    b += 0.1
  end
  m::Float64 = (a + b)/2.0::Float64

  while abs(f(m)::Float64)::Float64 > approx
    if sign(f(a)::Float64)::Float64 == sign(f(m)::Float64)::Float64
      a = m::Float64
    else
      b = m::Float64
    end
    m = (a + b)/2.0::Float64
  end
  return m::Float64
end


using Neighborhood
using DelayEmbeddings
using LinearAlgebra # for eigenvalues
using StaticArrays

export ManifoldProjection

# TODO: Would be nice to be able to generalize this to have `τ` as delay time
# instead of forcing `τ=1`. But then again, maybe the algorithm *requires* τ=1,
# similarly with how Broomhead-King needs smallest delay time.

"""
    ManifoldProjection(m, Q, k, τ=1, w=0, r=1000.0) <: Decomposition
A nonlinear noise reduction method, also known as "locally linear projections", which
works by bringing a noisy signal closer to a multi-dimensional manifold that represents
the deterministic dynamics of the signal. The method is method "IV" of [^Grassberger1993].

`m::Int` is the same as in [^Grassberger1993], the embedding dimension.
`Q` is related with the *expected* dimension `d` of the manifold of the deterministic
with approximate relation `d ≈ m-Q+1`. However, it is best t consider `Q` as an input
parameter (the dimension of the set you project locally into) and you need to vary
it to achieve the best performance.
If given a `Vector{Int}` as `Q` the algorithm will iteratively do noise reduction
to the resulting outputs (thus a vector is strongly recommended).
Duplicate entries can exist in `Q`.

`k` can be either `Int` or a `SearchType` from [Neighborhood.jl](https://julianeighbors.github.io/Neighborhood.jl/stable/#Search-types-1).
If `Int`, the `k` nearest neighbors are choosen as the neighborhood 𝓤 of each point.
The paper contains an involved process for determining optimal `k::Int`, see eq.(5.4).
`w` is just the Theiler window, while `r` is the value of the edge entries of vector R,
and probably has not much impact (the rest of the entries are 1).

In the paper too big correction vectors were rescaled to the average magnitude of corrections,
using as a criterion the distribution of their size. This is not implemented here (as it
is not clear *exactly* what it means computationally, what is "too big"?)
Contributing it is welcomed if you know how...

See also [^Schreiber1996] for an application of the same algorithm in real ECG data.

[^Schreiber1996]: Schreiber & Kaplan (1996). Nonlinear noise reduction for electrocardiograms. [Chaos, 6(1), 87–92](https://doi.org/10.1063/1.166148)

[^Grassberger1993]: Grassberger et al., (1993). On noise reduction methods for chaotic data. [Chaos 3(2), 127–141](https://doi.org/10.1063/1.165979)
"""
struct ManifoldProjection{ST <: SearchType} <: Decomposition
    m::Int
    Qs::Vector{Int}
    st::ST
    w::Int
    r::Float64
end
ManifoldProjection(m, Q::Int, args...) = ManifoldProjection(m, [Q], args...)
ManifoldProjection(m, Q, k::Int, args...) = ManifoldProjection(m, Q, NeighborNumber(k), args...)
ManifoldProjection(m, Q, st::ST, w = 0, r = 1000.0) where {ST<:SearchType} =
ManifoldProjection{ST}(m, Q, st, w, r)

function SignalDecomposition.decompose(t, s, method::ManifoldProjection)
    m = method.m
    @assert maximum(method.Qs) < m
    x = copy(s) # x is "input timeseries", iteratively corrected
    ge = GeneralizedEmbedding(tuple(0:m...))

    # Function barrier to make `m` type parameter:
    Vc = genembed(x, ge)   # corrected embedded set
    iteratively_project!(x, method, ge, Vc)
    return x, s .- x
end

function iteratively_project!(x::AbstractVector{T}, method, ge::GeneralizedEmbedding{D}, Vc) where {T, D}
    m = D-1
    # Initialize a bunch of stuff
    r = copy(x)
    R = ones(eltype(x), m+1)
    R[1] = R[end] = method.r # something "large"
    Γ = C = zeros(eltype(x), m+1, m+1)
    𝒟 = copy(C)
    ξ = zeros(eltype(x), m+1)
    θ = zeros(eltype(x), m+1)

    # Main algorithm loop, repeated for Qs
    for Q in method.Qs
        V = genembed(x, ge) # embedded set, which is "bold x" in paper
        tree = searchstructure(KDTree, V.data, Euclidean())
        𝒰 = bulkisearch(tree, V.data, method.st, Theiler(method.w); sortds = false)

        @inbounds for n in 1:length(V) # start process for every point in the embedded space
            𝒰n = 𝒰[n]
            K = length(𝒰n)
            iK = 1/K

            for i in 1:m+1; ξ[i] = iK*sum(x[k+i-1] for k in 𝒰n); end

            for j in 1:m+1, i in 1:m+1
                ξij = ξ[i]*ξ[j]
                C[i, j] = @fastmath iK*sum(x[k+i-1]*x[k+j-1] - ξij for k in 𝒰n)
                Γ[i, j] = @fastmath R[i]*C[i, j]*R[j] # make Γ
            end
            eig::Matrix{T} = eigvecs(Γ) # these are sorted by eigenvalue. Important for the algorithm.
            for j in 1:m+1, i in 1:m+1
                𝒟[i, j] = sum(eig[i, q]*eig[j, q] for q in 1:Q)
            end

            # corrections vectors:
            for i in 1:m+1
                θ[i] = (1/R[i]) * sum(𝒟[i, j] * R[j] * (ξ[j] - x[n+j-1]) for j in 1:m+1)
            end

            Vc.data[n] = V.data[n] + SVector{D, T}(θ)
        end
        # okay, now use corrected points to estimate new points, by taking the average.
        @inbounds for n in (m+1):length(V)-m
            r[n] = 1/(m+1)*sum(Vc[n-j, j+1] for j in 0:m)
        end
        # Edge cases
        @inbounds for n in 1:m
            r[n] = sum(Vc[n-j+1, j] for j in 1:n)/n
        end
        @inbounds for (z, n) in enumerate((length(V)-m):-1:(length(V)-2m))
            r[n] = sum(Vc[n-j, m+1-j] for j in 0:(z-1))/z
        end
        x .= r # set corrected vector to input vector.
    end
end


### A Pluto.jl notebook ###
# v0.14.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 71903690-5f0e-11eb-2435-a56a85d76fe0
using Plots

# ╔═╡ 7eef98bc-5f0e-11eb-0dc2-43a0e5715aa4
using PlutoUI

# ╔═╡ d5cced56-5f0c-11eb-36ce-37a624f02490
md"# Heat kernel

The heat kernel is defined by

$K(t,x,y) = \frac{1}{(4 ~ \pi ~ t)^{d/2}} ~ \exp\left(\frac{-\lvert x - y \rvert^2}{4 t}\right)$

with $d = 1,2,\cdots$ as the number of dimensions, time $t \geq 0$ and space $x \in \mathbb{R}^{d}$.
"

# ╔═╡ 4131c77a-5f0e-11eb-1aae-41015705d66e
function kernel(t, x, y; dim = 1)
    d = dim

    arg = -1 * (abs(x - y)^2)/(4*t)
    den = (4*pi*t)^(d/2)

    sol = (1/den) * exp(arg)

    return sol
end

# ╔═╡ 4118304e-5f0e-11eb-3844-5159c87d48b7
tspan = 0.001 : 0.001 : 1.0 # time

# ╔═╡ 40fa23ec-5f0e-11eb-2b02-73a75b7a6049
xspan = -3.0 : 0.1 : 3.0 # space

# ╔═╡ 401f2062-5f0e-11eb-2517-5900dd4b84ba
y = 0.0 # Offset

# ╔═╡ 679b9134-5f0e-11eb-159c-258752c5fcef
sol = zeros(length(xspan), length(tspan));

# ╔═╡ 6f5614c6-5f0e-11eb-3412-5381fbf426a0
for i = 1 : length(tspan)
    for k = 1 : length(xspan)
        sol[k,i] = kernel(tspan[i], xspan[k], y)
    end
end

# ╔═╡ ab00a9dc-5f0e-11eb-36e1-fd6e0d1acb68
@bind pos Slider(1:10:1000, default=1)

# ╔═╡ be7a3cf8-5f0e-11eb-3bb6-0d132e1b8e42
scatter(xspan, sol[:,pos], legend=false)

# ╔═╡ Cell order:
# ╟─d5cced56-5f0c-11eb-36ce-37a624f02490
# ╠═4131c77a-5f0e-11eb-1aae-41015705d66e
# ╠═4118304e-5f0e-11eb-3844-5159c87d48b7
# ╠═40fa23ec-5f0e-11eb-2b02-73a75b7a6049
# ╠═401f2062-5f0e-11eb-2517-5900dd4b84ba
# ╠═679b9134-5f0e-11eb-159c-258752c5fcef
# ╠═6f5614c6-5f0e-11eb-3412-5381fbf426a0
# ╠═71903690-5f0e-11eb-2435-a56a85d76fe0
# ╠═7eef98bc-5f0e-11eb-0dc2-43a0e5715aa4
# ╠═ab00a9dc-5f0e-11eb-36e1-fd6e0d1acb68
# ╠═be7a3cf8-5f0e-11eb-3bb6-0d132e1b8e42


using BenchmarkTools
function mygemm!(C, A, B)
# gemm General Matrix Multiplication
# syntax sugaring !
# C := C + A * B
    M, N = size(C)
    K = size(A,2)

    @assert size(C) == (M,N)
    @assert size(A) == (M,K)
    @assert size(B) == (K,N)

    @inbounds for i = 1:M
        @inbounds for j = 1:N
            @inbounds for p = 1:K
                C[i, j] = C[i, j] + A[i,p] * B[p,j]
            end
        end
    end
    return C
end
#
# C = rand(10,12)
#
# A = rand(10,15)
#
# B = rand(15,12)
#
# mygemm!(C,A,B)
#
# D = C + A*B
#
# extrema(D - C) # get minimum and maximum value

using LinearAlgebra
let # inside let are global scale variables
    M,N,K = 256, 256, 256
    C = rand(M,N)
    A = rand(M,K)
    B = rand(K,N)

    #@time D = C + A * B
    @time mul!(C,A,B)
    @time mygemm!(C,A,B)

    nsamples = 10

    t_mul! = floatmax(Float64)
    for s = 1:nsamples
        t = @elapsed mul!(C,A,B)
        t_mul! = min(t_mul!, t)
    end
    @show t_mul!

    t_mygemm! = floatmax(Float64)
    for s = 1:nsamples
        t = @elapsed mygemm!(C,A,B)
        t_mygemm! = min(t_mygemm!, t)
    end
    @show t_mygemm!

end
nothing


#=
  Euler #24 in Julia.

  Euler 24:
  """
  A permutation is an ordered arrangement of objects. For example, 3124 is one
  possible permutation of the digits 1, 2, 3 and 4. If all of the permutations are
  listed numerically or alphabetically, we call it lexicographic order. The
  lexicographic permutations of 0, 1 and 2 are:

      012   021   102   120   201   210

  What is the millionth lexicographic permutation of the digits
  0, 1, 2, 3, 4, 5, 6, 7, 8 and 9?
  """

  This Julia program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

include("Euler.jl")

#
function euler24a()
    a = Array(0:9)
    c = 1
    while c < 1000000
        a = next_permutation(a)
        c += 1
    end

    return parse(Int, join(a,""))
end

# Too slow: 15.24007616s
function euler24b()
    return parse(Int,join(all_permutations2(Array(0:9))[1000000]))
end

# Inspired by a solution on the 'net
# 0.05707951s
function euler24c()
    n = 999999
    p = 10
    eli = (1:p).|>(i->i % 10)
    answer = []
    for i in 1:p-1
        f = factorial(p-i)
        d = floor(Int,n / f)
        n %= f
        push!(answer,eli[d])
        # Note the indexOf()...
        # eli.splice(eli.indexOf(eli[d-1]),1);
        eli = filter(e->e != eli[d], eli)
    end
    return parse(Int,join(vcat(answer,eli)))
end

# all_permutations is much faster than all_permutations2
# 0.49170683s
function euler24d()
    return parse(Int,join(all_permutations(Array(0:9))[1000000]))
end

run_euler(euler24a)
# run_euler(euler24b) # too slow
# run_euler(euler24c)
# run_euler(euler24d)


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# ╔═╡ 61d79b04-7db7-11eb-1f8f-cbd858393030
md"""
Engle-Granger Two-Step Error Correction Model

$y_{t} = log(Y_{t}), x_{t} = log(X_{t})$

Order 1

$y_{t} = c + \gamma x_{t} + z_{t}$

$\delta y_{t} = y_{t} - y_{t - 1}  = \alpha z_{t - 1} + \epsilon_{t}$

Test for Cointegration
+ Ad-Fuller Test on residuals $z_{t}$
"""

# ╔═╡ 3f05c2d8-7dba-11eb-3c50-f169294281fc
md"""

https://hudsonthames.org/correct-backtest-methodology-pairs-trading/

https://docs.google.com/presentation/d/1oFpv7OUi3W9F2D30rEQRy_Ra5ejLtlCMnk9nTJnMDr4/edit#slide=id.gb70b23aaa3_1_0

https://drive.google.com/file/d/1ZeJD81OrKln8QDxm1sU63ivRgqXCcQbS/view

https://docs.google.com/presentation/d/19xUNdddCsdXg2KizcUMVvdjhTnA8Cuq293brXB8rSWM/edit#slide=id.p

https://docs.google.com/presentation/d/1Cb9NizOkvCx6-p6X3AVZyxjviub7kF2sCPy2gBiT9nk/edit#slide=id.p

https://hudsonthames.org/copula-for-pairs-trading-overview-of-common-strategies/

https://hudsonthames.org/optimal-stopping-in-pairs-trading-ornstein-uhlenbeck-model/

Tim Leung, Optimal Mean Reversion Trading with Transaction Costs and Stop-Loss Exit (2015)
https://arxiv.org/abs/1411.5062

Statistical Arbitrage Strategy in Multi-Asset Market Using Time Series Analysis

https://www.scirp.org/journal/paperinformation.aspx?paperid=100361

"""

# ╔═╡ Cell order:
# ╠═61d79b04-7db7-11eb-1f8f-cbd858393030
# ╠═3f05c2d8-7dba-11eb-3c50-f169294281fc


using BlackBoxOptim, Optim

cost_function = build_loss_objective(prob1,Tsit5(),L2Loss(t,data),
                                     maxiters=10000,verbose=false)
bound1 = Tuple{Float64, Float64}[(1, 2)]
result = bboptimize(cost_function;SearchRange = bound1, MaxSteps = 11e3)
@test result.archive_output.best_candidate[1] ≈ 1.5 atol=3e-1

cost_function = build_loss_objective(prob2,Tsit5(),L2Loss(t,data,differ_weight=nothing,data_weight=1.0),
                                     maxiters=10000,verbose=false)
bound2 = Tuple{Float64, Float64}[(1, 2),(1, 4)]
result = bboptimize(cost_function;SearchRange = bound2, MaxSteps = 11e3)
@test result.archive_output.best_candidate ≈ [1.5;3.0] atol=3e-1

cost_function = build_loss_objective(prob3,Tsit5(),L2Loss(t,data,differ_weight=10),
                                     maxiters=10000,verbose=false)
bound3 = Tuple{Float64, Float64}[(1, 2),(0, 2), (2, 4), (0, 2)]
result = bboptimize(cost_function;SearchRange = bound3, MaxSteps = 11e3)
@test result.archive_output.best_candidate ≈ [1.5;1.0;3.0;1.0] atol=5e-1

cost_function = build_loss_objective(prob3,Tsit5(),L2Loss(t,data,differ_weight=0.3,data_weight=0.7),
                                     maxiters=10000,verbose=false)
bound3 = Tuple{Float64, Float64}[(1, 2),(0, 2), (1, 4), (0, 2)]
result = bboptimize(cost_function;SearchRange = bound3, MaxSteps = 11e3)
@test result.archive_output.best_candidate ≈ [1.5;1.0;3.0;1.0] atol=5e-1

println("LSA for gradient")
cost_function = build_loss_objective(prob1,Tsit5(),L2Loss(t,data),
                                     maxiters=10000,verbose=false,flsa_gradient=true)
bound1 = Tuple{Float64, Float64}[(1, 2)]
function g_!(stor,x)
    cost_function.cost_function2(x,stor)
end
result = Optim.optimize(cost_function.cost_function,g_!,[1.0], Optim.ConjugateGradient())
@test result.minimizer[1] ≈ 1.5 atol=3e-1

println("Adjoint SA for gradient")
cost_function = build_loss_objective(prob2,Tsit5(),L2Loss(t,data),
                                     maxiters=10000,verbose=false,adjsa_gradient=true)
function g_!(stor,x)
    cost_function.cost_function2(x,stor)
end
result = Optim.optimize(cost_function.cost_function,g_!,[1.3,2.0], Optim.ConjugateGradient())
@test result.minimizer ≈ [1.5;3.0] atol=3e-1


using Distributions
using FFTW
using Plots

# Create a simple signal with two frequencies
dt = 0.001
t = 0:dt:1
n = length(t)

f = sin.(50*2pi*t) + sin.(120*2pi*t) # Sum of 2 frequencies
f += 2.5*rand(Normal(), n) # Add some noise

# Compute the Fast Fourier Transform FFT
fhat = fft(f) # Compute the fast Fourier transform
PSD = abs2.(fhat)/n # Power spectrum (power per freq)
freq = 1/(dt*n)*(0:n) # Create x-axis of frequencies in Hz

# Use the PSD to filter out noise
indices = PSD .> 100 # Find all freqs with large power
PSDclean = PSD.*indices # Zero out all others
fhat = indices.*fhat # Zero out small Fourier coeffs. in Y
ffilt = real.(ifft(fhat)) # Inverse FFT for filtered time signal

# PLOTS
p1 = plot(t, f)
plot!(t, ffilt)

L = 1:Int(floor(n/2)) # Only plot the first half of freqs
p2 = plot(freq[L], PSD[L])
plot!(freq[L], PSDclean[L])


r = 10;
A = pi*r^2; # pi has a default value!
println(A);


using Test

using Catlab.Present

# Schema
########

A, B = Ob(FreeSchema, :A), Ob(FreeSchema, :B)
f, g = Hom(:f, A, B), Hom(:g, B, A)

C = AttrType(FreeSchema.AttrType, :C)
x, y = Attr(:x, A, C), Attr(:y, B, C)

@test dom(x) == A
@test codom(x) == C
@test dom(compose(f,y)) == A
@test dom(compose(g,x)) == B
@test compose(f, compose(g,f)) == compose(compose(f,g), f)
@test compose(f, compose(g, x)) == compose(compose(f, g), x)
@test compose(id(A), f) == f
@test compose(id(A), x) == x


function step(model, scheme::ModifiedBridge, t0, x0, Δw)
  x1 = scheme.s
  t1 = scheme.t
  μ = (x1 - x0) / (t1 - t0)
  σ = sqrt((t1 - (t0 + scheme.Δt)) / (t1 - t0)) * diffusion(model, t0, x0)
  x = x0 + μ * scheme.Δt + σ * Δw
  x
end

function _normal_transition_params(model, scheme::ModifiedBridge, t0, x0, s1)
  x1 = scheme.s
  t1 = scheme.t
  μ = x0 + (x1 - x0) / (t1 - t0) * scheme.Δt
  σ = diffusion(model, t0, x0)
  Σ = (t1 - (t0 + scheme.Δt)) / (t1 - t0) * scheme.Δt * σ * σ'
  μ, Σ
end


# Load Julia packages (libraries)

using Pkg, DrWatson

using DataFrames
using StanSample
using StatisticalRethinking

df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame);

mean_ma = mean(df[!, :MedianAgeMarriage])
df[!, :MedianAgeMarriage_s] =
    convert(Vector{Float64},  (df[!, :MedianAgeMarriage]) .-
        mean_ma)/std(df[!, :MedianAgeMarriage]);
df.Divorce_SE = df[:, "Divorce SE"]

stan14_1 = "
  data {
    int N;
    vector[N] A;
    vector[N] R;
    vector[N] Dobs;
    vector[N] Dsd;
  }
  parameters {
    real a;
    real ba;
    real br;
    real<lower=0> sigma;
    vector[N] Dest;
  }
  model {
    vector[N] mu;
    // priors
    target += normal_lpdf(a | 0, 10);
    target += normal_lpdf(ba | 0, 10);
    target += normal_lpdf(br | 0, 10);
    target += cauchy_lpdf(sigma | 0, 2.5);

    // linear model
    mu = a + ba * A + br * R;

    // likelihood
    target += normal_lpdf(Dest | mu, sigma);

    // prior for estimates
    target += normal_lpdf(Dobs | Dest, Dsd);
  }
  generated quantities {
    vector[N] log_lik;
    {
      vector[N] mu;
      mu = a + ba * A + br * R;
      for(i in 1:N) log_lik[i] = normal_lpdf(Dest[i] | mu[i], sigma);
    }
  }
";

m14_1s = SampleModel("m14.1s", stan14_1)

m14_1_data = Dict(
    "N" => size(df, 1),
    "A" => df[!, :MedianAgeMarriage],
    "R" => df[!, :Marriage],
    "Dobs" => df[!, :Divorce],
    "Dsd" => df[!, :Divorce_SE]
)

rc14_1s = stan_sample(m14_1s, data=m14_1_data)

if success(rc14_1s)
    chns14_1s = read_samples(m14_1s, :dataframe)
    chns14_1s |> display
end


using MittagLefflerFunctions
using SpecialFunctions: erfcx
using PyPlot

α = 1/2
β = 1.0
r = 5

x = range(0, 10, length=201)
y = erfcx.(x)

figure(1)
m_lo = 4
style = ("-", "--", ":", "-.", "-")
for m = m_lo:2r+2-m_lo
    n = 2r + 1 - m
    E = MLPade(α, β, m, n, :svd)
    err = E.(x) - y
    plot(x, err, style[m-m_lo+1], label="m = $m, n = $n")
end
xlabel(L"$x$")
legend()
grid(true)

savefig("figure08.pdf")


using LazyInverses
using LinearAlgebra
using BenchmarkTools

n = 128
x = randn(n)
y = randn(n)
A = randn(n, n)
A = A'A
C = cholesky(A)
invC = inverse(C)

println("ternary dot-product multiplication")
@btime dot($x, $invC, $x)
@btime dot($x, $invC, $y)
@btime dot($x, $C \ $x)
@btime dot($x, $C \ $y)

k = n
X = randn(k, n)
Y = randn(n, k)
println("ternary multiplication")
@btime *($X, $invC, $X')
@btime *($X, $invC, $Y)
@btime *($X, $C \ $X')
@btime *($X, $C \ $Y)


import Base: +, -, *, /, length, iterate, zero, convert, Vector, Array
"""
This correspond to coordinate of atom.
This contain 3 fields like below.

- x <: Real
- y <: Real
- z <: Real

Some operators and functions were overloaded to make Coordinate objects broadcastable.
"""
mutable struct Coordinate{RealT <: Real}
    # This class is for make vector operation broadcastable.
    x::RealT
    y::RealT
    z::RealT

    function Coordinate(vector::VectorT
        ) where VectorT <: AbstractArray{RealT} where RealT <: Real

        if length(vector) != 3
            ArgumentError("Argument for Coordinate constructor must be array of length 3")
        end
        new{RealT}(vector[1], vector[2], vector[3])
    end

    function Coordinate(x::RealT, y::RealT, z::RealT) where RealT <: Real
        new{RealT}(x, y, z)
    end
end

function +(arg1::Coordinate{T}, arg2::Coordinate{U}) where T <: Real where U <: Real
    Coordinate(arg1.x + arg2.x, arg1.y + arg2.y, arg1.z + arg2.z)
end

function -(arg1::Coordinate{T}, arg2::Coordinate{U}) where T <: Real where U <: Real
    Coordinate(arg1.x - arg2.x, arg1.y - arg2.y, arg1.z - arg2.z)
end

function *(arg1::Coordinate{T}, arg2::U) where T <: Real where U <: Real
    Coordinate(arg1.x * arg2, arg1.y * arg2, arg1.z * arg2)
end

function *(arg1::U, arg2::Coordinate{T}) where T <: Real where U <: Real
    Coordinate(arg1 * arg2.x, arg1 * arg2.y, arg1 * arg2.z)
end

function *(arg1::Coordinate{T}, arg2::Coordinate{U}) where T <: Real where U <: Real
    arg1.x * arg2.x + arg1.y * arg2.y + arg1.z * arg2.z
end

function /(arg1::Coordinate{T}, arg2::U) where T <: Real where U <: Real
    Coordinate(arg1.x / arg2, arg1.y / arg2, arg1.z / arg2)
end

# for operator broadcast
function length(arg::Coordinate{T}) where T <: Real
    1
end

function iterate(arg::Coordinate{T}) where T <: Real
    (arg, nothing)
end

function iterate(arg::Coordinate{T}, nothing) where T <: Real
    nothing
end

function zero(arg::Coordinate{T}) where T <: Real
    Coordinate(T(0.0), T(0.0), T(0.0))
end

function convert(::Type{Coordinate{RealT}}, coord::Coordinate{<:Real}) where RealT <: Real
    return Coordinate(RealT(coord.x), RealT(coord.y), RealT(coord.z))
end

function Vector(arg::Coordinate{RealT}) where RealT <: Real
    [arg.x, arg.y, arg.z]
end

function Array(arg::Coordinate{RealT}) where RealT <: Real
    [arg.x, arg.y, arg.z]
end

function norm(arg::Coordinate{T}) where T <: Real
    sqrt(abs2(arg.x) + abs2(arg.y) + abs2(arg.z))
end

function distance(first_coord::Coordinate{T}, second_coord::Coordinate{T}) where T <: Real
    norm(first_coord - second_coord)
end


@everywhere using LowRankModels
using Test

function fit_pca(m,n,k)
	# matrix to encode
	Random.seed!(1)
	# generate a matrix with rank k
	A = randn(m,k)*randn(k,n)
	# fit a PCA model with rank k
	glrm = pca(A, k)
	glrm = share(glrm)
	p = Params()
	# just do 10 iterations
	p.max_iter = 10
	X,Y,ch = fit!(glrm)
	return A,X,Y,ch
end

@everywhere Random.seed!(1)
A,X,Y,ch = fit_pca(100,100,50)

# make sure objective went down
@test ch.objective[end] < ch.objective[1]


using ProfileView
using PolynomialBasis
using ImplicitDomainQuadrature
using BenchmarkTools
# using Revise
using CutCellDG
include("../../../test/useful_routines.jl")
include("test-bounds.jl")

interfacecenter = [1.5, 0.5]
interfaceradius = 1.0

levelsetbasis = HermiteTensorProductBasis(2)
interpgrad = InterpolatingPolynomial(2, 2, 3)
numqp = required_quadrature_order(3) + 2
dim, numpts = size(interpolation_points(levelsetbasis))
quad = tensor_product_quadrature(2, 4)
cgmesh = CutCellDG.CGMesh([0.0, 0.0], [1.0, 1.0], [1, 1], numpts)

levelset = CutCellDG.LevelSet(
    x -> -circle_distance_function(x, interfacecenter, interfaceradius)[1],
    cgmesh,
    levelsetbasis,
    quad,
)

CutCellDG.load_coefficients!(levelset, 1)
update_interpolating_gradient!(interpgrad, CutCellDG.interpolater(levelset))

function run_area_quadrature(levelset, interpgrad, numqp; numiter = 100)
    for i = 1:numiter
        pquad = area_quadrature(
            CutCellDG.interpolater(levelset),
            interpgrad,
            +1,
            [-1.0, -1.0],
            [1.0, 1.0],
            numqp,
        )
    end
end

# Run once to compile
# run_area_quadrature(levelset,interpgrad,numqp)

# @profview run_area_quadrature(levelset, interpgrad, numqp)

using IntervalArithmetic
xL,xR = [-1.,-1.],[1.,1.]
box = IntervalBox(xL,xR)
# @btime ImplicitDomainQuadrature.is_suitable(1,interpgrad,box)

curve(x) = ImplicitDomainQuadrature.curvature_measure(interpgrad(x),1,4)
# @btime sign(curve,box,tol=1e-3)
interval_arithmetic_sign_search(curve,box,1e-3,0.0,2)

@btime area_quadrature(CutCellDG.interpolater(levelset),interpgrad)


println("testing mesh")
# function test_mesh()
img = rand(100,100)
offset = [0,0]
dist = 20
mesh = create_mesh(img, offset, dist)
i1 = 100 % (dist*sin(pi/3)) / 2
@test mesh.src_nodes[1,:] ≈ [i1,0]
@test mesh.src_nodes[2,:] ≈ [i1,20]
@test mesh.src_nodes[3,:] ≈ [i1,40]
@test mesh.src_nodes[4,:] ≈ [i1,60]
@test mesh.src_nodes[5,:] ≈ [i1,80]
@test mesh.src_nodes[6,:] ≈ [i1,100]
@test mesh.src_nodes[7,:] ≈ [i1 + 20*(sin(pi/3)), 10]
@test maximum(mesh.src_nodes[:,1]) <= 100
@test maximum(mesh.src_nodes[:,2]) <= 100
@test size(mesh.src_nodes, 1) == 33
# test that src_nodes match dst_nodes, initially
function wrapper() # otherwise gives warning if a and b are already defined globally
	for (a, b) in zip(mesh.src_nodes, mesh.dst_nodes)
		@test a == b
	end
end
wrapper()
@test size(mesh.edges, 2) == 81


# AbelianSpace.jl
#
# Defines the immutable AbelianSpace for a complex vector space that is graded by
# the unitary irreps of an Abelian group.

# AbelianSpace
#--------------
immutable AbelianSpace{G<:Abelian} <: UnitaryRepresentationSpace{G}
    dims::Dict{G,Int}
    dual::Bool
end
AbelianSpace{G<:Abelian}(dims::Dict{G,Int})=AbelianSpace{G}(dims, false)

# Corresponding methods:
sectors(V::AbelianSpace) = keys(V.dims)

dim(V::AbelianSpace) = sum(values(V.dims))::Int
dim{G<:Abelian}(V::AbelianSpace{G},c::G) = get(V.dims,c,0)
dual(V::AbelianSpace) = AbelianSpace([conj(c)=>dim(V,c) for c in sectors(V)], !V.dual)
cnumber{G}(V::AbelianSpace{G}) = AbelianSpace(Dict{G,Int}(one(G)=>1),V.dual)
cnumber{G}(::Type{AbelianSpace{G}}) = AbelianSpace(Dict{G,Int}(one(G)=>1))
iscnumber{G}(V::AbelianSpace{G}) = V.dims == Dict{G,Int}(one(G)=>1)

# Comparison
==(V1::AbelianSpace,V2::AbelianSpace)=(V1.dims==V2.dims && V1.dual==V2.dual)

# Show methods
Base.show(io::IO, V::AbelianSpace) = (print(io,"AbelianSpace");showcompact(io,V.dims);print(io, V.dual ? "*" : ""))

# direct sum of AbelianSpaces
function directsum{G}(V1::AbelianSpace{G}, V2::AbelianSpace{G})
    V1.dual == V2.dual || throw(SpaceError("Direct sum of a vector space and its dual do not exist"))
    dims=Dict{G,Int}
    for c in union(sectors(V1),sectors(V2))
        dims[c]=dim(V1,c)+dim(V2,c)
    end
    return AbelianSpace(dims,V1.dual)
end

# fusing and splitting AbelianSpaces
function fuse{G}(V1::AbelianSpace{G}, V2::AbelianSpace{G}, V::AbelianSpace{G})
    dims=[c1*c2=>0 for c1 in sectors(V1), c2 in sectors(V2)]
    for c1 in sectors(V1), c2 in sectors(V2)
        dims[c1*c2]+= dim(V1,c1)*dim(V2,c2)
    end
    return V.dims == dims
end

# indexing using sectors
function to_range{G}(s::G,V::AbelianSpace{G})
    offset=0
    for c in sort!(collect(sectors(V)),rev=V.dual,by=s->s.charge)
        if c!=s
            offset+=dim(V,c)
        else
            break
        end
    end
    return offset+(1:dim(V,s))
end


## Some more extreme tests from http://ir.igsnrr.ac.cn/bitstream/311030/8840/1/%E4%BE%AF%E9%BA%9F%E7%A7%91(SCI)2.pdf
## All the methods do poorly here due to the multiplicities of the functions
## even though there should be convergnce of either |xn+1 - xn| or |f(xn)|
using Roots
using Base.Test

tests = [
         (x -> (x - sqrt(5))^4 / ((x-1)^2 + 2),    4.3,  2.236067977499790)
         (x -> (sin(x)^2 - 2x + 1)^5,              4.5,  0.71483582544138924)
         (x -> (8x*exp(-x^2) -2x - 3)^8,          -2.0, -1.7903531791589544)
         (x -> (2x*cos(x) + x^2 - 3)^10/(x^2 + 1), 3.0,  2.9806452794385368)
         (x -> (exp(-x^2 + x + 3) - x + 2)^9,      2.5,  2.4905398276083051)
         (x -> (exp(-x) + 2sin(x))^4,              3.0,  3.1627488709263654)
         (x -> (log(x^2 + 3x + 5) - 2x + 7)^8,     5.5,  5.4690123359101421)
         (x -> (sqrt(x^2 + 2x + 5) -
                           2sin(x) - x^2 + 3)^5,   2.3, 2.3319676558839640)
         (x -> (x-2)^4 / ( (x-1)^2 + 1),           1.8,  2.0000000000000000)
         (x -> abs(x - 2.5)^(15/4) * exp(x),       2.2,  2.5)
         (x -> (sqrt(x) - 1/x - 1)^7,              2.0,  2.147899035704787)
         (x -> (log(x) + sqrt(x) - 5)^3,           8.0,  8.309432694231572)
         (x -> (sin(x)*cos(x) - x^3 + 1)^9,        1.0,  1.117078770687451)
         (x -> ((x-3)*exp(x))^5,                   2.5,  3.0000000000000000)
         (x -> (log(x) + sqrt(x^4 + 1) -2)^7,      1.0,  1.222813963628973)
         ]

using Roots
ord(x) = string(floor(Integer,log10(x)))
orders = [0, 1, 2, 5, 8, 16]
out = Any[]
for (ctr, (f1, x0, xstar)) in enumerate(tests)
    m = AbstractString[]
    for order in orders
        try
            a = fzero(f1, x0, order=order, maxeval=100)
            tm = begin tic(); fzero(f1, x0, order=order, maxeval=100); toc() end
            push!(m, ord(tm))
            push!(m, ord(abs(a-xstar)))
        catch
            push!(m, "***")
            push!(m, "***")
        end
    end
    try
        newton(f1, x0)
        tm =  begin tic(); a = newton(f1, x0); toc() end
        push!(m, ord(tm))
        push!(m, ord(abs(a-xstar)))
    catch e
        push!(m, "***")
        push!(m, "***")
    end
    push!(out, m)
end

## values in pairs time/accuracy or ****** if failure
## sensitive to maxeval value
 for i in out
     for k in 1:div(length(i),2)
         @printf "%3s%3s " i[2*(k-1)+1] i[2*k]
     end
         @printf "\n"
 end


"""
$(SIGNATURES)

Compute the centroid or geometric center of the given `ConvexHull` using the
formulas given [here](https://en.wikipedia.org/wiki/Centroid#Of_a_polygon).
"""
function centroid(hull::ConvexHull)
    T = eltype(hull)
    R = arithmetic_closure(T)
    vertices = hull.vertices
    n = length(vertices)
    @inbounds begin
        if n === 0
            error()
        elseif n === 1
            return R.(vertices[1])
        elseif n === 2
            return (vertices[1] + vertices[2]) / 2
        else
            c = cross2(vertices[n], vertices[1])
            centroid = (vertices[n] + vertices[1]) * c
            double_area = c
            @simd for i in Base.OneTo(n - 1)
                c = cross2(vertices[i], vertices[i + 1])
                centroid += (vertices[i] + vertices[i + 1]) * c
                double_area += c
            end
            double_area = abs(double_area)
            centroid /= 3 * double_area
            return centroid
        end
    end
end


import SimpleDiffEq

@inline @inbounds function eom_lorenz_lorenz_bidir(u, p, t)
    c_xy, c_yx, a₁, a₂, a₃, b₁, b₂, b₃ = (p...,)
    x1, x2, x3, y1, y2, y3 = (u...,)

    dx1 = -a₁*(x1 - x2) + c_yx*(y1 - x1)
    dx2 = -x1*x3 + a₂*x1 - x2
    dx3 = x1*x2 - a₃*x3
    dy1 = -b₁*(y1 - y2) + c_xy*(x1 - y1)
    dy2 = -y1*y3 + b₂*y1 - y2
    dy3 = y1*y2 - b₃*y3

    return SVector{6}(dx1, dx2, dx3, dy1, dy2, dy3)
end

"""
    lorenz_lorenz_bidir(; u0 = rand(6),
        c_xy = 0.2, c_yx = 0.2,
        a₁ = 10, a₂ = 28, a₃ = 8/3,
        b₁ = 10, b₂ = 28, b₃ = 9/3) -> ContinuousDynamicalSystem

Initialise a bidirectionally coupled Lorenz-Lorenz system, where each
subsystem is a 3D Lorenz system [1]. Default values for the parameters
`a₁`, `a₂`, `a₃`, `b₁`, `b₂`, `b₃` are as in [1].

## Equations of motion

The dynamics is generated by the following vector field

```math
\\begin{aligned}
\\dot{x_1} &= -a_1 (x_1 - x_2) + c_{yx}(y_1 - x_1) \\\\
\\dot{x_2} &= -x_1 x_3 + a_2 x_1 - x_2 \\\\
\\dot{x_3} &= x_1 x_2 - a_3 x_3 \\\\
\\dot{y_1} &= -b_1 (y_1 - y_2) + c_{xy} (x_1 - y_1) \\\\
\\dot{y_2} &= -y_1 y_3 + b_2 y_1 - y_2 \\\\
\\dot{y_3} &= y_1 y_2 - b_3 y_3
\\end{aligned}
```

## References

1. Amigó, José M., and Yoshito Hirata. "Detecting directional couplings from
    multivariate flows by the joint distance distribution." Chaos: An
    Interdisciplinary Journal of Nonlinear Science 28.7 (2018): 075302.
"""
function lorenz_lorenz_bidir(; u0 = rand(6),
        c_xy = 0.2, c_yx = 0.2,
        a₁ = 10, a₂ = 28, a₃ = 8/3,
        b₁ = 10, b₂ = 28, b₃ = 9/3)
    ContinuousDynamicalSystem(eom_lorenz_lorenz_bidir, u0, [c_xy, c_yx, a₁, a₂, a₃, b₁, b₂, b₃])
end

function lorenzlorenz_bidir_trajectory(npts; sample_dt = 1, Ttr = 1000, dt = 0.1,
    c_xy = 0.1, c_yx = 0.1,
    u0 = rand(6),
    a₁ = 10, a₂ = 28, a₃ = 8/3,
    b₁ = 10, b₂ = 28, b₃ = 9/3)

    s = lorenz_lorenz_bidir(u0 = u0, c_xy = c_xy, c_yx = c_yx, a₁ = a₁, a₂ = a₂, a₃ = a₃, b₁ = b₁, b₂ = b₂, b₃ = b₃)

    # the system is recorded at times t0:dt:T
    T = npts*dt*sample_dt

    o = trajectory(s, T, dt = dt, Ttr = Ttr*dt, alg = SimpleDiffEq.SimpleATsit5())[1:sample_dt:end-1, :]
end

# For some initial conditions, the system wanders off and doesn't settle to an attractor. Create a function that loops until we get a good realization.
function good_lorenzlorenz_bidir_trajectory(npts;
        sample_dt = 1,
        dt = 0.1,
        c_xy = 0.1,
        c_yx = 0.1,
        Da₁ = Uniform(9.5, 10.5),
        Da₂ = Uniform(27, 29),
        Da₃ = Uniform(7.5/3, 8.5/3),
        Db₁ = Uniform(9.5, 10.5),
        Db₂ = Uniform(27, 29),
        Db₃ = Uniform(7.5/3, 8.5/3),
        a₁ = nothing,
        a₂ = nothing,
        a₃ = nothing,
        b₁ = nothing,
        b₂ = nothing,
        b₃ = nothing,
        u0 = rand(6),
        Ttr = 10000,
        n_maxtries = 300)

    n_tries = 0
    while n_tries <= n_maxtries
        a₁ == nothing ? a₁ = rand(Da₁) : nothing
        a₂ == nothing ? a₂ = rand(Da₂) : nothing
        a₃ == nothing ? a₃ = rand(Da₃) : nothing
        b₁ == nothing ? b₁ = rand(Db₁) : nothing
        b₂ == nothing ? b₂ = rand(Db₂) : nothing
        b₃ == nothing ? b₃ = rand(Db₃) : nothing

        pts = lorenzlorenz_bidir_trajectory(npts,
            sample_dt = sample_dt, dt = dt,
            c_xy = c_xy, c_yx = c_yx,
            Ttr = Ttr)

        M = Matrix(pts)

        if all(isfinite.(M)) && all(M .< 1e10) && count(M .≈ 0) < npts*0.1 && count(abs.(M) .< 1e-10) < npts*0.1 &&
            (count(abs.(M) .< 1e-12) < npts*0.1)
            return pts
        end
        println("no attractor found. trying with new initial condition and parameters")
        n_tries += 1
    end
end


function formatmatrix(inputmatrix)

  numrows = size(inputmatrix,1)
  numcols = size(inputmatrix,2)

  trans = inputmatrix'

  firstcol = trans[1,:]
  outputmatrix = map(x -> x/sum(firstcol),firstcol)

  for i=2:numcols
    row = trans[i,:]
    newrow = map(x -> x/sum(row),row)
      outputmatrix = vcat(outputmatrix,newrow)
  end
  outputmatrix'
end

mat = [1 2 3; 4 5 6]

println(formatmatrix(mat))


const PFCUT = [6:5:101]//100
const PFVAL = [10:8:26, 32:6:50, 54:4:98, 100]//100

function pricefraction{T<:FloatingPoint}(a::T)
    zero(T) <= a || error("a = ", a, ", but it must be >= 0.")
    a <= one(T) || error("a = ", a, ", but it must be <= 1.")
    convert(T, PFVAL[findfirst(a .< PFCUT)])
end

test = [0.:0.05:1., 0.51, 0.56, 0.61, rand(), rand(), rand(), rand()]

println("Testing the price fraction function")
for t in test
    println(@sprintf "    %.4f -> %.4f" t pricefraction(t))
end


"""
    snr = calculateSNR(hwfft, f, nsig=1)

Estimate the signal-to-noise ratio, given the in-band bins of a (Hann-windowed)
fft and the location of the input signal (`f` > 0). For `nsig` = 1, the input tone is
contained in `hwfft`[`f`:`f`+2]; this range is appropriate for a Hann-windowed fft.
Each increment in `nsig` adds a bin to either side. The `snr` is expressed in dB.
"""
function calculateSNR(hwfft, f, nsig=1)
    signalBins = (f-nsig+1):(f+nsig+1)
    signalBins = signalBins[signalBins .> 0]
    signalBins = signalBins[signalBins .<= length(hwfft)]
    s = norm(hwfft[signalBins])

    noiseBins = collect(1:length(hwfft))
    deleteat!(noiseBins, signalBins)
    n = norm(hwfft[noiseBins])

    if n == 0
        snr = Inf
    else
        snr = dbv(s/n)
    end

    return snr
end


using OrdinaryDiffEq, Distributions,
      DiffEqUncertainty, Test, Quadrature, Cubature, Cuba,
      FiniteDiff, Zygote, ForwardDiff, DiffEqGPU, DiffEqSensitivity

quadalgs = [HCubatureJL(), CubatureJLh(), CubatureJLp(), CubaSUAVE(), CubaDivonne(), CubaCuhre()]
batchmode = [EnsembleSerial(), EnsembleThreads(), EnsembleCPUArray()]#, EnsembleGPUArray()]

function eom!(du,u,p,t)
  @inbounds begin
    du[1] = u[2]
    du[2] = -p[1]*u[1] - p[2]*u[2]
  end
  nothing
end

u0 = [1.0;1.0]
tspan = (0.0,3.0)
p = [1.0, 2.0]
prob = ODEProblem(eom!,u0,tspan,p)

A = [0.0 1.0; -p[1] -p[2]]
u0s_dist = [Uniform(1,10), Uniform(2,6)]

@testset "Koopman Expectation, nout = 1" begin
  g(sol) = sol[1,end]
  analytical = (exp(A*tspan[end])*mean.(u0s_dist))[1]

  for alg ∈ quadalgs
    @info "$alg"
    @test expectation(g, prob, u0s_dist, p, Koopman(), Tsit5(); quadalg = alg)[1] ≈ analytical rtol=1e-2
  end
end

@testset "Koopman Expectation, nout > 1" begin
  g(sol) = sol[:,end]
  analytical = (exp(A*tspan[end])*mean.(u0s_dist))

  for alg ∈ quadalgs
    @info "$alg"
    @test expectation(g, prob, u0s_dist, p, Koopman(), Tsit5(); quadalg = alg, nout =2) ≈ analytical rtol=1e-2
  end
end

@testset "Koopman Expectation, nout = 1, batch" begin
  g(sol) = sol[1,end]
  analytical = (exp(A*tspan[end])*mean.(u0s_dist))[1]

  for bmode ∈ batchmode
    for alg ∈ quadalgs
      if alg isa HCubatureJL
        continue
      end
      @info "nout = 1, batch mode = $bmode, $alg"
      @test expectation(g, prob, u0s_dist, p, Koopman(), Tsit5(), bmode; quadalg = alg, batch = 15)[1] ≈ analytical rtol=1e-2
    end
  end
end

@testset "Koopman Expectation, nout > 1, batch" begin
  g(sol) = sol[:,end]
  analytical = (exp(A*tspan[end])*mean.(u0s_dist))

  for bmode ∈ batchmode
    for alg ∈ quadalgs
      if alg isa HCubatureJL
        continue
      end
      @info "nout = 2, batch mode = $bmode, $alg"
      res = expectation(g, prob, u0s_dist, p, Koopman(), Tsit5(), bmode;
                        quadalg = alg, batch = 15, nout=2)
      @test res ≈ analytical rtol=1e-2
    end
  end
end

############## Koopman AD ###############

@testset "Koopman Expectation AD" begin
  g(sol) = sol[1,end]
  loss(p, alg) = expectation(g, prob, u0s_dist, p, Koopman(), Tsit5(); quadalg = alg)[1]
  dp1 = FiniteDiff.finite_difference_gradient(p->loss(p, HCubatureJL()),p)
  for alg ∈ quadalgs
    @info "$alg, ForwardDiff"
    @test ForwardDiff.gradient(p->loss(p,alg),p) ≈ dp1 rtol=1e-2
    @info "$alg, Zygote"
    @test Zygote.gradient(p->loss(p,alg),p)[1] ≈ dp1 rtol=1e-2
  end
end

@testset "Koopman Expectation AD, batch" begin
  g(sol) = sol[1,end]
  loss(p, alg, bmode) = expectation(g, prob, u0s_dist, p, Koopman(), Tsit5(), bmode; quadalg = alg, batch = 10)[1]
  dp1 = FiniteDiff.finite_difference_gradient(p->loss(p, CubatureJLh(), EnsembleThreads()),p)
  for bmode ∈ batchmode
    for alg ∈ quadalgs
      if alg isa HCubatureJL #no batch support
        continue
      end
      @info "$bmode, $alg, ForwardDiff"
      @test_broken ForwardDiff.gradient(p->loss(p,alg,bmode),p) ≈ dp1 rtol=1e-2
      @info "$bmode, $alg, Zygote"
      @test_broken  Zygote.gradient(p->loss(p,alg,bmode),p)[1] ≈ dp1 rtol=1e-2
    end
  end
end

########## Central Moment ##############

function eom!(du,u,p,t)
  @inbounds begin
    du[1] = p[1]*u[1]
  end
  nothing
end

u0 = [1.0]
tspan = (0.0,3.0)
p = [-.3]
prob = ODEProblem(eom!,u0,tspan,p)

u0s_dist = [Uniform(1,10)]

@testset "Koopman Central Moment" begin
  g(sol) = sol[1,end]
  analytical = [0.0, exp(2*p[1]*tspan[end])*var(u0s_dist[1]), 0.0]

  for alg ∈ quadalgs
    if alg isa CubaDivonne || alg isa CubaCuhre  #requires 2D spatial integration
      continue
    end
    r = centralmoment(3, g, prob, u0s_dist, p, Koopman(), Tsit5();
                      ireltol = 1e-8, iabstol = 1e-8, quadalg = alg)
    if alg isa CubaSUAVE
      @test_broken r ≈ analytical rtol=1e-2
    else
      @test r ≈ analytical rtol=1e-2
    end
  end
end

@testset "Koopman Central Moment, batch" begin
  g(sol) = sol[1,end]
  analytical = [0.0, exp(2*p[1]*tspan[end])*var(u0s_dist[1]), 0.0]

  for bmode ∈ batchmode
    for alg ∈ quadalgs
      if alg isa CubaDivonne || alg isa CubaCuhre || alg isa HCubatureJL #requires 2D spatial integration
        continue
      end
      @info "batch mode = $bmode, alg = $alg"
      r = centralmoment(3, g, prob, u0s_dist, p, Koopman(), Tsit5(), bmode;
                        ireltol = 1e-8, iabstol = 1e-8, quadalg = alg, batch=15)
      if alg isa CubaSUAVE
        @test_broken r ≈ analytical rtol=1e-2
      else
        @test r ≈ analytical rtol=1e-2
      end
    end
  end
end


#################### Raftery and Lewis Diagnostic ####################

function rafterydiag(x::Vector{T}; q::Real=0.025, r::Real=0.005,
                      s::Real=0.95, eps::Real=0.001,
                      range::AbstractRange=1:1:length(x)) where {T<:Real}
  nx = length(x)
  phi = sqrt(2.0) * erfinv(s)
  nmin = ceil(Int, q * (1.0 - q) * (phi / r)^2)
  if nmin > nx
    @warn("At least $nmin samples are needed for specified q, r, and s")
    kthin = burnin = total = NaN
  else
    dichot = Int[(x .<= quantile(x, q))...]
    kthin = 0
    bic = 1.0
    local test, ntest
    while bic >= 0.0
      kthin += 1
      test = dichot[1:kthin:nx]
      ntest = length(test)
      temp = test[1:(ntest - 2)] + 2 * test[2:(ntest - 1)] + 4 * test[3:ntest]
      trantest = reshape(counts(temp, 0:7), 2, 2, 2)
      g2 = 0.0
      for i1 in 1:2, i2 in 1:2, i3 in 1:2
        tt = trantest[i1, i2, i3]
        if tt > 0
          fitted = sum(trantest[:, i2, i3]) * sum(trantest[i1, i2, :]) /
                   sum(trantest[:, i2, :])
          g2 += 2.0 * tt * log(tt / fitted)
        end
      end
      bic = g2 - 2.0 * log(ntest - 2.0)
    end
    tranfinal = counts(test[1:(ntest - 1)] + 2 * test[2:ntest], 0:3)
    alpha = tranfinal[3] / (tranfinal[1] + tranfinal[3])
    beta = tranfinal[2] / (tranfinal[2] + tranfinal[4])
    kthin *= step(range)
    m = log(eps * (alpha + beta) / max(alpha, beta)) /
        log(abs(1.0 - alpha - beta))
    burnin = kthin * ceil(m) + range.start - 1
    n = ((2.0 - alpha - beta) * alpha * beta * phi^2) /
        (r^2 * (alpha + beta)^3)
    keep = kthin * ceil(n)
    total = burnin + keep
  end
  [kthin, burnin, total, nmin, total / nmin]
end

function rafterydiag(c::AbstractChains; q::Real=0.025, r::Real=0.005,
                     s::Real=0.95, eps::Real=0.001)
  _, p, m = size(c.value)
  vals = Array{Float64}(undef, p, 5, m)
  for j in 1:p, k in 1:m
    vals[j, :, k] = rafterydiag(c.value[:, j, k], q=q, r=r, s=s, eps=eps,
                                range=c.range)
  end
  hdr = header(c) * "\nRaftery and Lewis Diagnostic:\n" *
        "Quantile (q) = $q\nAccuracy (r) = $r\nProbability (s) = $s\n"
  ChainSummary(vals, c.names, ["Thinning", "Burn-in", "Total", "Nmin",
                               "Dependence Factor"], hdr)
end


"""
# module NumerovMethod

# Examples

```jldoctest
julia>
```
"""
module NumerovMethod

using NumericalMethodsInQuantumMechanics.EigenvalueProblems: InitialCondition

export integrate

"""
    numerov_iter(y_prev, y, dx, gvec, svec)

Return the ``y[n + 1]`` step for the Numerov's method, given the current and previous step.

# Arguments
- `y_prev::Real`: the `y[n - 1]` element.
- `y::Real`: the `y[n]` element.
- `dx::Real`: the step length ``h``, need to be small.
- `gvec::AbstractArray{<:Real}`: stores `g[n - 1]`, `g[n]` and `g[n + 1]`.
- `svec::AbstractArray{<:Real}`: stores `s[n - 1]`, `s[n]` and `s[n + 1]`.
"""
function iter(yᵢ₋₁, yᵢ, dx, gvec, svec)
    gᵢ₋₁, gᵢ, gᵢ₊₁ = gvec
    sᵢ₋₁, sᵢ, sᵢ₊₁ = svec
    coeffᵢ₋₁ = -(1 + dx^2 / 12 * gᵢ₋₁)
    coeffᵢ = 2(1 - 5dx^2 / 12 * gᵢ)
    coeffᵢ₊₁ = 1 + dx^2 / 12 * gᵢ₊₁
    coeff_sᵢ = dx^2 / 12 * (sᵢ₋₁ + 10sᵢ + sᵢ₊₁)
    return (coeffᵢ * yᵢ + coeffᵢ₋₁ * yᵢ₋₁) / coeffᵢ₊₁
end
"""
    numerov_iter(y_prev, y, dx, gvec)

Same as `numerov_iter(y_prev, y, dx, gvec, svec)`, if ``s(x) ≡ 0`` on the domain.

# Arguments
- `y_prev::Real`: the `y[n - 1]` element.
- `y::Real`: the `y[n]` element.
- `dx::Real`: the step length, need to be small.
- `gvec::AbstractArray{<:Real}`: stores `g[n - 1]`, `g[n]` and `g[n + 1]`.
"""
iter(yᵢ₋₁, yᵢ, dx, gvec) = 2(12 - 5dx^2 * gvec[2]) / (12 + dx^2 * gvec[3]) * yᵢ - yᵢ₋₁

"""
    integrate(ic, r, gvec, svec)

Do the Numerov's method integration, return the solution ``y(x)``, given that ``g(x)`` and ``s(x)`` already
as vectors (already applied on ``x``).

# Arguments
- `ic::InitialCondition`: the initial condition ``y(x_0) = y_0`` and ``y'(x_0) = y_0'``, could be a guess.
- `r::AbstractRange{<:Real}`: a range, the domain ``x``.
- `gvec::AbstractArray{<:Real}`: the result of function ``g`` applied on ``x`` (range `r`).
- `svec::AbstractArray{<:Real}`: the result of function ``s`` applied on ``x`` (range `r`).
"""
function integrate(gvec, svec, ic::InitialCondition)
    if length(gvec) != length(svec)
        throw(DimensionMismatch("Integ"))
    end
    N = length(gvec)
    dx = inv(N)
    ϕ₀, ϕ′₀ = ic
    ϕ = [ϕ₀, ϕ′₀ * dx]  # ϕ₀, ϕ₁
    for i in 1:(N-2)
        ϕᵢ₊₂ = iter(ϕ[i], ϕ[i+1], dx, gvec[i:(i+2)], svec[i:(i+2)])
        push!(ϕ, ϕᵢ₊₂)
    end
    return ϕ
end
"""
    integrate(ic, r, g, s)

Same as `integrate(ic, r, gvec, svec)`, but `g` and `s` are two functions.

# Arguments
- `ic::InitialCondition`: the initial condition ``y(x_0) = y_0`` and ``y'(x_0) = y_0'``, could be a guess.
- `r::AbstractRange{<:Real}`: a range, the domain ``x``.
- `g::Function`: the function ``g``.
- `s::Function`: the function ``s``.
"""
function integrate(g::Function, s::Function, ic::InitialCondition, dx)
    vec = 0:dx:1
    gvec, svec = map(g, vec), map(s, vec)
    return integrate(gvec, svec, ic)
end

end


using UpwindDifferences
using Test

using SparseArrays
using LinearAlgebra
using DelimitedFiles
using ProgressMeter

function evaluateparameters(ps)
    θ, ρ, r, y, dt = ps[1], ps[2], ps[3], ps[4], ps[5]
    xs = LinRange(0.01, 10.0, 1000)

    reward(x, c) = c ^ (1.0 - θ) / (1.0 - θ)
    drift(x, c) = r * x + y - c
    policy(x, dv) = dv ^ (-1.0 / θ)
    zerodrift(x) = r * x + y

    R, GF, GB = zeros(1000), zeros(1000), zeros(1000)
    UR = UpwindResult(R, GF, GB)
    bs = (zeros(1000-1), zeros(1000-1))
    dv = zeros(1000-1)

    function backwardsiterate!(v0, v1)
        upwind!(UR, v1, xs, reward, drift, policy, zerodrift)
        A = sparse(Tridiagonal(-UR.GB[2:end] ./ step(xs), (UR.GB .- UR.GF) ./ step(xs), UR.GF[1:end-1] ./ step(xs)))
        v0 .= ((ρ + 1.0 / dt) * I - A) \ (UR.R .+ v1 ./ dt)
        return v0
    end

    vinit = map(x -> reward(x, zerodrift(x)) ./ ρ, xs)
    vs = (copy(vinit), copy(vinit))
    tol, maxiter = 1e-10, 10000
    err = tol
    for iter in 1:maxiter
        backwardsiterate!(vs[2], vs[1])
        err = maximum(abs.(vs[2] .- vs[1]))
        if err < tol
            return (value = vs[2], iter = iter, err = err, termination = :tolerance)
        else
            vs[1] .= vs[2]
        end
    end
    return (value = vs[1], iter = maxiter, err = err, termination = :iterations)
end

@testset "UpwindDifferences.jl" begin

    allparameters = readdlm("testdata/hjb_parameters.tsv")
    allvalues = readdlm("testdata/hjb_values.tsv")
    @showprogress for i in axes(allparameters, 1)
        ps = allparameters[i, :]
        vnew = evaluateparameters(ps).value
        vold = allvalues[i, :]
        @test maximum(abs.(vnew .- vold)) < 1e-8
    end
end


# TODO: abstract to objects (spheres) with Transmission and Reflectivity
# TODO: rasterize rays

using Plots
using LinearAlgebra
using Test

pyplot()

function Base.isapprox(n1::Nothing, n2::Nothing)
    return true
end

abstract type Object end

mutable struct Ray
    # Velocity vector
    v::Vector{Float64}

    # Position vector
    p::Vector{Float64}

end

mutable struct Lens <: Object
    # Lens position
    p::Vector{Float64}

    # Lens radius
    r::Float64

    # Index of Refraction
    ior::Float64
end

function lens_normal_at(ray, lens)
    n = normalize(ray.p .- lens.p)
    if dot(-n, ray.v) < 0
        n *= -1
    end

    return n
end

function inside_of(ray::Ray, lens)
    return inside_of(ray.p, lens)
end

function inside_of(pos, lens)
    x = lens.p[1] - pos[1]
    y = lens.p[2] - pos[2]
    if (x^2 + y^2 <= lens.r^2)
        return true
    else
        return false
    end
end

# note: light moves as a particular speed with respect to the medium it is
#       moving through, so...
#       n_2*v = n_1*l + (n_1*cos(theta_1) - n_2*cos(theta_2))*n
#       Other approximations: ior = n_1/n_2, c = -n*l
function refract!(ray, lens, ior)
    n = lens_normal_at(ray, lens)
    c = dot(-n, ray.v);
    d = 1.0 - ior^2 * (1.0 - c^2);

    if (d < 0.0)
        return zeros(2);
    end

    ray.v = ior * ray.v + (ior * c - sqrt(d)) * n;
end

abstract type Wall <: Object end;

mutable struct Mirror <: Wall
    # Normal vector
    n::Vector{Float64}

    # Position of mirror
    p::Vector{Float64}

    # Mirror size
    scale::Float64

    Mirror(in_n, in_p) = new(in_n, in_p, 2.5)
end

function is_behind(ray, mirror)
    if dot(ray.p.-mirror.p, mirror.n) >= 0
        return true
    else
        return false
    end
end

# note: for reflection, l_x -> l_x, but l_y -> -l_y
#       In this case, the y component of l = cos(theta)*n
#       so new vector: v = l + 2cos(theta)*n
function reflect!(ray, n)
    ray.v .-= 2*dot(ray.v, n).*n
end

function draw_circle(p, r, res)
    return [x .+ (p[1], p[2]) for x in Plots.partialcircle(0, 2pi, res, r)]
end

function plot_rays(positions, objects::Vector{O},
                   filename) where {O <: Object}
    plt = plot(background_color=:black, aspect_ratio=:equal)

    for i = 1:size(positions)[1]
        plot!(plt, (positions[i,:,1], positions[i,:,2]); label = "ray",
              linecolor=:white)
    end

    for object in objects
        if typeof(object) == Mirror
            dir = [-object.n[2], object.n[1]]
            extents = zeros(2,2)
            extents[1,:] .= object.p .- object.scale * dir
            extents[2,:] .= object.p .+ object.scale * dir
            plot!(plt, (extents[:,1], extents[:,2]), label = "mirror")
        elseif typeof(object) == Lens
            circle = draw_circle(object.p, object.r, 100)
            plot!(circle; label="lens", linecolor=:lightblue)
        end
    end

    savefig(plt, filename)
end

function step!(ray::Ray, dt)
    ray.p .+= .+ ray.v.*dt
end

function intersection(ray::Ray, lens::Lens)
    # threshold to avoid floating precision errors
    thresh = 0.00001

    # distance from the center of the lens
    relative_dist = lens.p - ray.p

    # distance projected onto the direction vector
    projected_relative_dist = dot(relative_dist, ray.v)

    # distance of closest approach
    # Might make more sense to do a dot product instead of a norm
    closest_approach = abs(sqrt(norm(relative_dist)^2 -
                            projected_relative_dist^2))

    # TODO: arbitrary threshold, should be fixed.
    if abs(closest_approach - lens.r) < thresh
        closest_approach = lens.r
    end

    # Note: returns nothing if tangential
    if closest_approach >= lens.r
        return nothing
    end

    # half chord length
    half_chord = sqrt(lens.r^2 - closest_approach^2)

    # checks if ray is moving away from sphere by seeing if we are
    # outside of sphere (abs(projected_relative_dist) > half_chord), and if
    # the ray is moving away from the boundary
    if (projected_relative_dist < 0 &&
        abs(projected_relative_dist) > half_chord)
        return nothing
    end

    in_lens = half_chord < projected_relative_dist
    intersect = zeros(2)
    # Adding a threshold to output so we are definitely on the other side of
    # the boundary after intersecting
    if in_lens
        intersect = (projected_relative_dist - half_chord+thresh)*ray.v
    else
        intersect = (projected_relative_dist + half_chord+thresh)*ray.v
    end

    return intersect
end

function intersection_quadratic(ray::Ray, lens::Lens)
    relative_dist = ray.p-lens.p
    a = dot(ray.v, ray.v)
    b = 2.0 * dot(relative_dist, ray.v)
    c = dot(relative_dist, relative_dist) - lens.r*lens.r
    discriminant = b*b - 4*a*c

    if discriminant < 0
        return nothing
    elseif discriminant > 0
        roots = [(-b + sqrt(discriminant)) / (2*a),
                 (-b - sqrt(discriminant)) / (2*a)]
        min = minimum(roots)
        max = maximum(roots)

        if min > 0
            return (min)*ray.v
        elseif max > 0
            return (max)*ray.v
        else
            return nothing
        end
    else
        # Returns nothing if tangential
        return nothing
        #return (-b/(2*a))*ray.v
    end
end

function intersection(ray::Ray, wall::W) where {W <: Wall}
    intersection = -dot((ray.p .- wall.p),wall.n)/dot(ray.v, wall.n)

    if isfinite(intersection) && intersection > 0 && intersection != NaN
        return intersection*ray.v
    else
        return nothing
    end
end

function intersect_test()

    lens = Lens([3.0, 1.0], 1.0, 1.5)

    @testset "intersect lens tests" begin
        rays = [Ray([1.0, 0.0], [0.0, 1.0], zeros(2,2)),
                Ray([1.0, 0.0], [3.0, 1.0], zeros(2,2)),
                Ray([sqrt(0.5), sqrt(0.5)], [2.0, 0.0], zeros(2,2)),
                Ray([-1.0, 0.0], [0.0, 1.0], zeros(2,2)),
                Ray([0.0, 1.0], [0.0, 1.0], zeros(2,2)),
                Ray([1.0, 0.0], [0.0, 0.0], zeros(2,2))]
        answers = [[2.0, 0.0], [1.0, 0.0], [1-sqrt(0.5), 1-sqrt(0.5)],
                   nothing, nothing, [3.0, 0.0]]

        for i = 1:length(rays)
            pt = intersection(rays[i], lens)
            @test isapprox(pt, answers[i])
            pt = intersection_geometric(rays[i], lens)
            @test isapprox(pt, answers[i])
        end
    end

    wall = Mirror([-1.0, 0.0], [2.0, 1.0])

    @testset "intersect wall tests" begin
        rays = [Ray([1.0, 0.0], [0.0, 1.0], zeros(2,2)),
                Ray([sqrt(0.5), sqrt(0.5)], [0.0, 0.0], zeros(2,2)),
                Ray([-sqrt(0.5), sqrt(0.5)], [0.0, 0.0], zeros(2,2)),
                Ray([0.0, 1.0], [0.0, 1.0], zeros(2,2))]

        answers = [[2.0, 0.0], [2.0, 2.0], nothing, nothing]

        for i = 1:length(rays)
            pt = intersection(rays[i], wall)
            @test isapprox(pt, answers[i])
        end
    end
end

function propagate!(rays::Vector{Ray}, objects::Vector{O},
                    num_intersections, positions) where {O <: Object}
    for i = 2:num_intersections
        for j = 1:length(rays)

            intersect_final = [Inf, Inf]
            intersected_object = nothing
            for object in objects
                intersect = intersection(rays[j], object)
                if intersect != nothing &&
                   sum(intersect[:].^2) < sum(intersect_final[:].^2)
                    intersect_final = intersect
                    intersected_object = object
                end
            end

            if intersect_final != nothing
                rays[j].p .+= intersect_final
                positions[j, i, :] .= rays[j].p
                if typeof(intersected_object) == Lens
                    ior = intersected_object.ior
                    if inside_of(rays[j], intersected_object)
                        ior = 1/intersected_object.ior
                    end

                    refract!(rays[j], intersected_object, ior)
                    if rays[j].v == zeros(2)
                        reflect!(rays[j], lens_normal_at(rays[j], lens))
                    end

                elseif typeof(intersected_object) == Mirror
                    reflect!(rays[j], intersected_object.n)
                end

                println(rays[j].v, '\t', rays[j].p)
            end
        end
        println('\n')
    end
end

function parallel_propagate(ray_num, num_intersections, box_size;
                            filename="check.png")

    rays = [Ray([1, 0],
            [0.1, float(i)]) for i = 1:ray_num]

    lenses = [Lens([10, 5], 5, 1.5)]
    mirrors = [Mirror([-1.0, 0.0], [25.0, 5.0]),
               Mirror([1.0, 0.0], [0, 5.0]),
               Mirror([0.0, 1.0], [12.5, 0.0]),
               Mirror([0.0, -1.0], [12.5, 10.0])]

    objects = vcat(lenses, mirrors)
    positions = zeros(ray_num, num_intersections, 2)

    for i = 1:length(rays)
        positions[i, 1, :] .= rays[i].p
    end

    propagate!(rays, objects, num_intersections, positions)

    plot_rays(positions, objects, filename)

    return positions

end


using LinearAlgebra
using NLopt
using Printf

mutable struct ExpFitBuffer
    p::Int
    Nt::Int
    Neval::Int

    r::Vector{Float64}
    t::Vector{Float64}
    w::Vector{Float64}
    wt::Vector{Float64}
    wrt::Vector{Float64}

    A::Vector{ComplexF64}
    B::Vector{ComplexF64}
    C::Vector{Float64}
    E::Vector{ComplexF64}
    G::Matrix{ComplexF64}
    M::Matrix{ComplexF64}
    N::Matrix{ComplexF64}
    Q::Matrix{ComplexF64}
    H::AbstractMatrix{ComplexF64}
    Bt::Matrix{ComplexF64}
    Θ::Matrix{ComplexF64}

    function ExpFitBuffer(r::Vector, t::AbstractVector, p::Integer, w::Vector)
        Nt = length(r)
        @assert Nt == length(t)

        r = copy(r)
        t = copy(t)
        wt = w.*t
        wrt = wt.*r

        A = fill(0.0im, p)
        C = fill(0.0, p)
        B = fill(0.0im, p)
        E = fill(0.0im, p)
        G = fill(0.0im, p, p)
        M = fill(0.0im, p, p)
        N = fill(0.0im, p, p)
        Q = fill(0.0im, p, p)
        if p == 1
            H = fill(0.0im, p, p)
        else
            H = Tridiagonal(M)
        end
        Bt = fill(0.0im, Nt, p)
        Θ = fill(0.0im, Nt, p)

        new(p, Nt, 0, r, t, w, wt, wrt, A, B, C, E, G, M, N, Q, H, Bt, Θ)
    end
end

function ARMA_gradient(grad::Vector, buffer::ExpFitBuffer, A::Vector{T},
    B::Vector{T}, fmodel::Vector) where {T<:Number}
    t = buffer.t
    p = buffer.p

    for i = 1:p
        buffer.Bt[:,i] = B[i].^t
        buffer.Θ[:,i] = A[i]/B[i] * t .* buffer.Bt[:,i]
    end

    buffer.M[:,:] = buffer.Bt'*(buffer.w.*buffer.Bt)
    buffer.N[:,:] = buffer.Bt'*(buffer.wt.*buffer.Bt)
    for i=1:p; buffer.N[i,:] /= B[i]; end

    for i=1:p
        buffer.E[i] = dot(buffer.wt, buffer.Bt[:,i])/B[i]
    end

    for i=1:p
        for j=1:p
            buffer.Q[i,j] = A[i] * buffer.N[i,j]
        end
    end

    buffer.G[:,:] = pinv(buffer.M) * (Diagonal(buffer.E-buffer.N*A)-buffer.Q)

    H = buffer.H
    for i=1:2:p-1
        ID = 1.0/(B[i]-B[i+1])
        H[i,i] = -B[i]*ID
        H[i+1,i] = B[i+1]*ID
        H[i,i+1] = -ID
        H[i+1,i+1] = ID
    end
    if p%2 == 1
        H[end,end] = B[end]^2
    end

    Jacobian = real((buffer.Θ + buffer.Bt*buffer.G)*buffer.H)
    grad[:] = 2*(Jacobian' * (buffer.w.*(fmodel-buffer.r)))
end

function ARMA_objective(C::Vector, grad::Vector, buffer::ExpFitBuffer)
    p = length(C)
    B = C2B(C)
    A = findA(buffer.t, buffer.r, B, w=buffer.w)
    fmodel = exponential_model(buffer.t, A, B)

    if length(grad) > 0
        ARMA_gradient(grad, buffer, A, B, fmodel)
    end
    buffer.Neval += 1
    sum(buffer.w .* (fmodel - buffer.r).^2)
end

function optimize_exponentials(data::Vector, w::Vector, guessC::Vector)
    p = length(guessC)
    if p == 0
        return sum(data[2:end].^2 .* w[2:end]), Float64[]
    end

    opt = Opt(:LD_MMA, p)
    opt = Opt(:LN_BOBYQA, p)
    xtol_abs!(opt, 1e-6)
    ftol_rel!(opt, 1e-5)
    maxeval!(opt, 100p)

    # Constraint bounds on the quadratic representation of the bases.
    # As all bases are in the unit circle, the linear coefficients are in (-2,2)
    # and the quadratic ones are in (-1,1).
    lb = fill(-2.0, size(guessC))
    ub = fill(2.0, size(guessC))
    lb[2:2:end] .= -1
    ub[2:2:end] .= +1
    if p%2 == 1
        lb[end] = -Inf
        ub[end] = +Inf
    end
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    step = 0.1*(ub-lb)
    if step[end] > 1
        step[end] = guessC[end]*0.5
    end
    initial_step!(opt, step)

    t = 0:length(data)-1
    buffer = ExpFitBuffer(data, t, p, w)
    min_objective!(opt, (x,grad)->ARMA_objective(x, grad, buffer))
    minf1 = ARMA_objective(guessC, [], buffer)
    # println("Before opt: $minf1 at $guessC")
    try
        (minf,minx,ret) = optimize(opt, guessC)
        # println("After opt: $minf at $minx after $(buffer.Neval) iterations (returned $ret)")
        # println("Moved by $(minx-guessC), improved by $(minf1-minf)\n")
        return minf, minx
    catch e
        @printf("Error in %d-order fit: %s\n", p, e)
        return minf1, guessC
    end
end


module JuliaSets
using PyPlot

export visJuliaSet, generatePlot

function visJuliaSet{T <: Real}(R::Function, x::Array{T, 1}, y::Array{T, 1},
	n_iter::Int64=1000, escape_tol::Float64=1e6)

	A = zeros(length(x), length(y))
	for i in 1:length(x)
  		for j in 1:length(y)
    		z = x[i] + im * y[j]
    		for k in 1:n_iter
      			z = R(z)
      			if abs(z) > escape_tol # diverged
        			A[i,j] = k
        			break
      			end
    		end
    		if A[i,j] == 0 # didn't diverge
      			A[i,j] = escape_tol + 1
    		end
  		end
	end
	A
end

function generatePlot(A::Array{Float64,2})
	matshow(A')
end

end # module


function henrici()
#
# M. Petkovic testing polynomials, page 123
#
    p = poly([-4.1;-3.8;-2.05;-1.85;1.95;2.15;3.9;4.05])
    z = [[-4.1;-3.8;-2.05;-1.85;1.95;2.15;3.9;4.05] ones(8)];

    p, PolyZeros(z)
end


### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# ╔═╡ 60b53bf6-dfa2-4fc5-b725-9e4239f7132c
using Colors, SpectralVis

# ╔═╡ 23835930-c0e0-11eb-21a7-49aedc5bc4b0
md"CIE1960 UCS"

# ╔═╡ 148c4957-ed02-4023-841f-9c5e82a6f796
plotly()

# ╔═╡ 144388e1-7642-4aba-b402-04a2114d9268
function cie60(c::XYZ)
	U = 2 * c.x / 3
	V = c.y
	W = 0.5 * (-c.x + 3c.y + c.z)
	u = U / (U+V+W)
	v = V / (U+V+W)
	u,v
end

# ╔═╡ f55e9217-f9b6-43e5-b29e-aa532f4adb5c
cie60(XYZ(.1, .5, .1))

# ╔═╡ 50a68b35-6e4d-4cd2-bf77-7aff52afeb83
begin
	temp = 6500
	meg = 1e7
	inv = meg / temp

	range = map(x -> 10000000/x, collect(10:500:100000))
	vals = collect(range)
	u = zeros(length(vals))
	v = zeros(length(vals))

	for k in 1:length(vals)
		t = vals[k]
		wp = blackbody_whitepoint(SPECENV,t)
		a,b = cie60(wp)
		u[k] = a
		v[k] = b
	end

	plot(xlimits=(0.15,0.25))
	scatter!(u,v,ratio = 1.0, color = :orange, markersize = 2, label = "knots")

		spl = cubicspline(u,v,:natural)
	sin = interp(spl,u[1]:0.001:u[end])
	plot!(sin[1],sin[2], color = :red, label = "spline")

	fine = map(x -> x, collect(100:10:20000))
	f = zeros(length(fine))
	g = zeros(length(fine))

	for l in 1:length(fine)
		te = fine[l]
		wpo = blackbody_whitepoint(SPECENV,te)
		c,d = cie60(wpo)
		f[l] = c
		g[l] = d
	end

	scatter!(f,g, color = :green, markersize = 1, markerstyle = :cross, label="truth")
end

# ╔═╡ ef6cd9f0-8416-472a-a4bb-43d5265d5757

# ╔═╡ acf3ed94-54cc-408f-bdae-c2eb16c9bbb7
range

# ╔═╡ b09c644b-449f-4756-aa82-c9aeb6e2265c
length(range)

# ╔═╡ Cell order:
# ╟─23835930-c0e0-11eb-21a7-49aedc5bc4b0
# ╠═60b53bf6-dfa2-4fc5-b725-9e4239f7132c
# ╠═148c4957-ed02-4023-841f-9c5e82a6f796
# ╠═144388e1-7642-4aba-b402-04a2114d9268
# ╠═f55e9217-f9b6-43e5-b29e-aa532f4adb5c
# ╠═50a68b35-6e4d-4cd2-bf77-7aff52afeb83
# ╠═ef6cd9f0-8416-472a-a4bb-43d5265d5757
# ╠═acf3ed94-54cc-408f-bdae-c2eb16c9bbb7
# ╠═b09c644b-449f-4756-aa82-c9aeb6e2265c


using Test
using GEMPIC

@testset "Maxwell 2D" begin
    function evaluate_spline_2d(nx1, nx2, degs, dofs)
        deg1, deg2 = degs
        vals = collect(reshape(dofs, nx1, nx2))
        for j in 1:nx2
            vals[:, j] .= GEMPIC.eval_uniform_periodic_spline_curve(deg1, vals[:, j])
        end
        for i in 1:nx1
            vals[i, :] .= GEMPIC.eval_uniform_periodic_spline_curve(deg2, vals[i, :])
        end
        return vec(vals)
    end

    x1min, x1max = 0.0, 2π
    nx1 = 16
    x2min, x2max = 0.0, 2π
    nx2 = 32

    mesh = TwoDGrid(x1min, x1max, nx1, x2min, x2max, nx2)

    deg = 3
    delta_t = 0.01
    nsteps = 300

    maxwell = TwoDMaxwell(mesh, deg)

    efield = [zeros(nx1 * nx2) for _ in 1:3]
    bfield = deepcopy(efield)
    efield_val = deepcopy(efield)
    bfield_val = deepcopy(efield)
    efield_ref = deepcopy(efield)
    bfield_ref = deepcopy(efield)

    x = LinRange(x1min, x1max, nx1 + 1)[1:(end - 1)] .* transpose(ones(nx2))
    y = ones(nx1) .* transpose(LinRange(x2min, x2max, nx2 + 1)[1:(end - 1)])

    w1 = sqrt(3)
    w2 = sqrt(3)

    rho = zeros(nx1 * nx2)
    rho_ref = zeros(nx1 * nx2)
    time = 0.0

    sin_k = (x, y) -> sin((x + y) - w1 * time)
    cos_k = (x, y) -> cos((x + y) - w1 * time)

    rho .= compute_rhs_from_function(maxwell, cos_k, 1, 0)

    compute_e_from_rho!(efield, maxwell, rho)

    efield_val1 = evaluate_spline_2d(nx1, nx2, (deg - 1, deg), efield[1])
    efield_val2 = evaluate_spline_2d(nx1, nx2, (deg, deg - 1), efield[2])
    efield_val3 = evaluate_spline_2d(nx1, nx2, (deg, deg), efield[3])

    efield_ref[1] .= vec(sin_k.(x, y) ./ 2)
    efield_ref[2] .= efield_ref[1]
    efield_ref[3] .= 0.0

    @test efield_val1 ≈ efield_ref[1] rtol = 1e-4
    @test efield_val2 ≈ efield_ref[2] rtol = 1e-4
    @test efield_val3 ≈ efield_ref[3]

    e1(x, y) = cos(x) * sin(y) * sin(sqrt(2) * time) / sqrt(2)
    e2(x, y) = -sin(x) * cos(y) * sin(sqrt(2) * time) / sqrt(2)
    b3(x, y) = -cos(x) * cos(y) * cos(sqrt(2) * time)

    time = -0.5 * delta_t

    bfield[1] .= l2projection(maxwell, e1, 1, 2)
    bfield[2] .= l2projection(maxwell, e2, 2, 2)
    bfield[3] .= l2projection(maxwell, b3, 3, 2)

    time = 0.0

    efield[1] .= l2projection(maxwell, e1, 1, 1)
    efield[2] .= l2projection(maxwell, e2, 2, 1)
    efield[3] .= l2projection(maxwell, b3, 3, 1)
    efield[3] .*= -1

    for istep in 1:nsteps
        compute_b_from_e!(bfield, maxwell, delta_t, efield)
        compute_e_from_b!(efield, maxwell, delta_t, bfield)
    end

    bfield_val[1] = evaluate_spline_2d(nx1, nx2, (deg, deg - 1), bfield[1])
    bfield_val[2] = evaluate_spline_2d(nx1, nx2, (deg - 1, deg), bfield[2])
    bfield_val[3] = evaluate_spline_2d(nx1, nx2, (deg - 1, deg - 1), bfield[3])

    time = (nsteps - 0.5) * delta_t
    bfield_ref[1] .= vec(e1.(x, y))
    bfield_ref[2] .= vec(e2.(x, y))
    bfield_ref[3] .= vec(b3.(x, y))

    @test bfield_ref[1] ≈ bfield_val[1] rtol = 1e-4
    @test bfield_ref[2] ≈ bfield_val[2] rtol = 1e-4
    @test bfield_ref[3] ≈ bfield_val[3] rtol = 1e-3

    efield_val[1] = evaluate_spline_2d(nx1, nx2, (deg - 1, deg), efield[1])
    efield_val[2] = evaluate_spline_2d(nx1, nx2, (deg, deg - 1), efield[2])
    efield_val[3] = evaluate_spline_2d(nx1, nx2, (deg, deg), efield[3])

    time = nsteps * delta_t

    efield_ref[1] .= vec(e1.(x, y))
    efield_ref[2] .= vec(e2.(x, y))
    efield_ref[3] .= -vec(b3.(x, y))

    @test efield_ref[1] ≈ efield_val[1] rtol = 1e-4
    @test efield_ref[2] ≈ efield_val[2] rtol = 1e-4
    @test efield_ref[3] ≈ efield_val[3] rtol = 1e-3

    efield[1] = l2projection(maxwell, cos_k, 1, 1)

    error2 = GEMPIC.inner_product(maxwell, efield[1], efield[1], 1, 1) - 2 * pi^2
    println(" Error in L2 norm squared: $error2")
    @test error2 ≈ 0 atol = 1e-5

    rho .= compute_rhs_from_function(maxwell, sin_k, 1, 1)

    compute_e_from_j!(efield[1], maxwell, rho, 1)

    efield_val1 = evaluate_spline_2d(nx1, nx2, [deg - 1, deg], efield[1])

    for i in eachindex(x, y)
        efield_ref[1][i] = cos_k(x[i], y[i]) - sin_k(x[i], y[i])
    end

    @show maximum(abs.(efield_val1 .- efield_ref[1]))
    @test efield_val1 ≈ efield_ref[1] atol = 1e-2

    time = 0.0
    rho_ref = compute_rhs_from_function(maxwell, cos_k, 1, 0)
    rho_ref .*= 2.0

    efield[1] = l2projection(maxwell, sin_k, 1, 1)
    efield[2] = l2projection(maxwell, sin_k, 2, 1)
    efield[3] = l2projection(maxwell, sin_k, 3, 1)

    compute_rho_from_e!(rho, maxwell, efield)

    @show maximum(abs.(rho .- rho_ref))

    @test rho ≈ rho_ref
end


export gaussian_kernel, sigmoid_kernel, cauchy_kernel, linear_kernel, tanh_kernel

"""
Creates a Gaussian kernel with the given distance function and `sigma` value
"""
function gaussian_kernel(dist, sigma=1.0)
    sigma2 = sigma * 2
    function fun(obj, ref)::Float64
        d = dist(obj, ref)
        (d == 0 || sigma == 0) && return 1.0
        exp(-d / sigma2)
    end

    fun
end

"""
Creates a sigmoid kernel with the given `sigma` value and distance function
"""
function sigmoid_kernel(dist, sigma=1.0)
    sqrtsigma = sqrt(sigma)
    function fun(obj, ref)::Float64
        x = dist(obj, ref)
        2 * sqrtsigma / (1 + exp(-x))
    end

    fun
end

"""
Creates a Cauchy's kernel with the given `sigma` value and distance function
"""
function cauchy_kernel(dist, sigma=1.0)
    sqsigma = sigma^2
    function fun(obj, ref)::Float64
        x = dist(obj, ref)
        (x == 0 || sqsigma == 0) && return 1.0
        1 / (1 + x^2 / sqsigma)
    end

    fun
end

"""
Creates a tanh kernel with the given `sigma` value and distance function
"""
function tanh_kernel(dist, sigma=1.0)
    function fun(obj, ref)::Float64
        x = dist(obj, ref)
        (exp(x-sigma) - exp(-x+sigma)) / (exp(x-sigma) + exp(-x+sigma))
    end

    fun
end

"""
Creates a linear kernel with the given distance function and `sigma` slope
"""
function linear_kernel(dist, sigma=1.0)
    function fun(obj, ref)::Float64
        dist(obj, ref) * sigma
    end

    fun
end


using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
using StochasticDiffEq

function roessler_dynoise(du,u, p, t)
    a, b, c, η = p
    du[1] = -u[2]-u[3]
    du[2] = u[1] + a*u[2]
    du[3] = b + u[3]*(u[1] - c)
end

# SDEProblem wants a specific function for the noise
# (This assumes diagonal noise, with the only noise acting on the
# 2nd variable)
function σ_roessler_dynoise(du, u, p, t)
    a, b, c, η = p
    du[1] = 0
    du[2] = η
    du[3] = 0
end

# Initial conditions
u₀ = [0.1, -0.2, 0.1]
# short tspan just for plot demonstrating the noisiness for different η
tspan = (0.,100.)

for η in [0.1, 2]

    p = [0.2, 0.2, 5.7, η]
    prob_sde_roessler = SDEProblem(roessler_dynoise, σ_roessler_dynoise, u₀, tspan, p)

    sol = solve(prob_sde_roessler, LambaEM())

    fig = figure()
    u = hcat(sol.u...)'
    plot3D(u[:,1], u[:,2], u[:,3])
    name = savename("roessler_dynoise", @dict(η), "png")
    wsave(plotsdir(name),fig)
end


using Plots
using LinearAlgebra

function angular_position(points)
    cpoints = points[1] + im * points[2]
    return rad2deg(angle(cpoints))
end

function perturbed_circle(point, meanradius, amplitude, frequency)
    angle = angular_position(point)
    interfaceposition = meanradius  + amplitude * sind(angle * frequency)
    distance = interfaceposition - norm(point)
    return distance
end

xrange = -1:1e-2:1
meanradius = 0.5
amplitude = 0.1
frequency = 5
contour(
    xrange,
    xrange,
    (x, y) ->
        perturbed_circle([x, y], meanradius, amplitude, frequency),
        levels=[0.0],
        aspect_ratio=:equal
)


#
#--------------------------------------#
export grad
#--------------------------------------#
"""
 Compute gradient of u∈H¹(Ω).

 Continuity isn't enforced across
 element boundaries for gradients

 [Dx] * u = [rx sx] * [Dr] * u
 [Dy]     = [ry sy]   [Ds]

"""
function grad(u  ::Array
             ,msh::Mesh)

    @unpack Dr,Ds,rx,ry,sx,sy = msh

    ux,uy = grad(u,Dr,Ds,rx,ry,sx,sy)

    return ux,uy
end
#--------------------------------------#
function grad(u,Dr,Ds,rx,ry,sx,sy)

  ur = ABu([],Dr,u)
  us = ABu(Ds,[],u)

  ux = @. rx * ur + sx * us
  uy = @. ry * ur + sy * us

  return ux,uy
end
#--------------------------------------#
export gradᵀ
#--------------------------------------#
"""
 Gradient transpose operator

 [Dx' Dy'] = [Dr' Ds'] [rx ry]
                       [sx sy]
"""
function gradᵀ(u  ::Array
              ,msh::Mesh)

    @unpack Dr,Ds,rx,ry,sx,sy = msh

    ux,uy = gradᵀ(u,Dr,Ds,rx,ry,sx,sy)

    return ux,uy
end
#--------------------------------------#
function gradᵀ(u,Dr,Ds,rx,ry,sx,sy)

  ux = ABu(Dr',[] ,rx .* u) +
       ABu([] ,Ds',sx .* u)

  uy = ABu(Dr',[] ,ry .* u) +
       ABu([] ,Ds',sy .* u)

  return ux,uy
end
#--------------------------------------#


using FastGaussQuadrature
using Serialization
#using Compat: @info

const GL_NODES, GL_WEIGHTS = gausslegendre(100)
const LGF_DIR  = joinpath(@__DIR__, "cache")
const LGF_FILE = joinpath(LGF_DIR, "lgftable.dat")

function load_lgf(N)
    if isfile(LGF_FILE)
        G = deserialize(open(LGF_FILE, "r"))
        if size(G,1) ≥ N
            return G
        end
    end
    build_lgf(N)
end

function build_lgf(N)
    @info "Building and caching LGF table"

    g = zeros(N, N)
    for y in 0:N-1, x in 0:y
        g[x+1,y+1] = lgf(x,y)
    end

    G = Symmetric(g)
    mkpath(LGF_DIR)
    #open(LGF_FILE, "w") do f
    #    serialize(f, G)
    #end
    serialize(open(LGF_FILE,"w"),G)
    G
end

quadgauss(f::Function) = dot(GL_WEIGHTS, f(GL_NODES))

function lgf(i, j)
    if i == j ==0
        return 0.0
    elseif i ≥ j
        v = quadgauss() do x
            if x == -1
                return sqrt(2)abs(i)
            else
                t = (x .+ 1)./2
                return 0.5real((1 .-
                                ( (t.-sqrt(1im))./(t.+sqrt(1im)) ).^(j.+abs(i)).*
                                ( (t.+sqrt(-1im))./(t.-sqrt(-1im)) ).^(j.-abs(i)) ))./t
            end

        end
        return 0.5v/pi
    else
        return lgf(j,i)
    end

end

const LGF_TABLE = load_lgf(1500)


using NLopt
using ForwardDiff

function a_next(c, h, states, agg_state, 𝔼V, params, hh::Renter)
  @unpack ρ, r = agg_state
  @unpack w, z = states
  a = w

  a_next = (1+r) * (a + z - c - ρ * h)
end

function objective0(c, h, states, agg_state, 𝔼V, params, hh::Renter)
  @unpack β = params

  a_next_ = a_next(c, h, states, agg_state, 𝔼V, params, hh)

  uu = c > 0 && h > 0 ? u(c,h) : 10_000 * min(c,h) - 10_000

  uu + β * 𝔼V(a_next_)

end

function constraint0(c, h, states, agg_state, 𝔼V, params, hh::Renter)
  @unpack r = agg_state

  a_next_ = a_next(c, h, states, agg_state, 𝔼V, params, hh)
  # - y_min / r <= a
  #(1+r) * m <= p * h * (1-δ) * θ
  - a_next_
end

function objective_nlopt(x::Vector, grad::Vector, args...)
  if length(grad) > 0
    ForwardDiff.gradient!(grad, XX -> objective0(XX[1], XX[2], args...), x)
  end
  objective0(x[1], x[2], args...)
end

function constraint_nlopt(x::Vector, grad::Vector, args...)
  if length(grad) > 0
    ForwardDiff.gradient!(grad, xx -> constraint0(xx[1], xx[2], args...), x)
  end
  constraint0(x[1], x[2], args...)
end

function Aiyagari.get_optimum(states, agg_state, 𝔼V, params, a_grid, hh::Renter)
  opt = Opt(:LD_MMA, 2)
  #opt = Opt(:LD_SLSQP, 2)
  lower_bounds!(opt, [eps(), eps()])
  upper_bounds!(opt, [+Inf, params.h_thres])

  xtol_rel!(opt, √eps())
  ftol_rel!(opt, eps())

  maxeval!(opt, 150)

  max_objective!(opt, (x,g) -> objective_nlopt(x, g, states, agg_state, 𝔼V, params, hh))
  inequality_constraint!(opt, (x,g) -> constraint_nlopt(x, g, states, agg_state, 𝔼V, params, hh), eps())

  guess = (states.w + states.z)/2
  (max_f, max_x, ret) = NLopt.optimize(opt, [guess, min(guess / agg_state.ρ, params.h_thres)])

  val = max_f
  c, h = max_x
  w_ = a_next(c, h, states, agg_state, 𝔼V, params, hh)

  conv = ret in [:FTOL_REACHED, :XTOL_REACHED, :SUCCESS, :LOCALLY_SOLVED]

  pol_full = (c=c, h=h, w_next=w_, ret=ret, conv=conv, count= opt.numevals)
  pol = pol_full.w_next, pol_full.h

  (pol=pol, pol_full=pol_full, val=val, conv=all(conv))
end


module MyInterpolations

export lin_int

function lin_int(grid,vals)
    function func(x)
        if x < grid[1]
            return "error"
        end
        if x > grid[end]
            return "error"
        end
        if x == grid[end]
            return vals[end]
        end
        if grid[1] <= x < length(grid)
            z = searchsortedlast(grid,x)
            slope = (vals[z+1]-vals[z])/(grid[z+1]-grid[z])
            y = slope*(x-grid[z])+vals[z]
        return y
        end
    end
    return func
end

end


export BinomialHomotopy

"""
    BinomialHomotopy

Construct the homotopy ``x^A - t γ + (1-t) b``
where each column in A denotes a monomial.
"""
struct BinomialHomotopy <: AbstractHomotopy
    A::Matrix{Int}
    b::Vector{ComplexF64}
    γ::Vector{Float64}

end

function BinomialHomotopy(A::Matrix{Int}, b::Vector{ComplexF64})
    γ = sign.(real.(b))
    BinomialHomotopy(A, b, γ)
end

cache(H::BinomialHomotopy, x, t) = HomotopyNullCache()
Base.size(H::BinomialHomotopy) = size(H.A)

function evaluate!(u, H::BinomialHomotopy, x, t, c::HomotopyNullCache)
    @unpack A, b, γ = H
    n = length(b)
    @inbounds for i = 1:n
        xAᵢ = one(x[1])
        for j = 1:n
            if A[j, i] != 0
                xAᵢ *= x[j]^A[j, i]
            end
        end
        u[i] = xAᵢ - (t * γ[i] + (1.0 - t) * b[i])
    end
    u
end

function dt!(u, H::BinomialHomotopy, x, t, c::HomotopyNullCache)
    @unpack b, γ = H
    n = length(b)
    @inbounds for i = 1:n
        u[i] = b[i] - γ[i]
    end
    u
end

function jacobian!(U, H::BinomialHomotopy, x, t, c::HomotopyNullCache)
    @unpack A, b, γ = H
    n = length(b)
    @inbounds for i = 1:n
        xAᵢ = one(x[1])
        for j = 1:n
            if A[j, i] != 0
                xAᵢ *= x[j]^A[j, i]
            end
        end
        for j = 1:n
            if A[j, i] == 0
                U[i, j] = 0
            else
                U[i, j] = A[j, i] * (@fastmath xAᵢ / x[j])
            end
        end
    end
    U
end

function evaluate_and_jacobian!(u, U, H::BinomialHomotopy, x, t, c::HomotopyNullCache)
    @unpack A, b, γ = H
    n = length(b)
    @inbounds for i = 1:n
        xAᵢ = one(x[1])
        for j = 1:n
            if A[j, i] != 0
                xAᵢ *= x[j]^A[j, i]
            end
        end
        u[i] = xAᵢ - (t * γ[i] + (1.0 - t) * b[i])

        for j = 1:n
            if A[j, i] == 0
                U[i, j] = 0
            else
                U[i, j] = A[j, i] * (@fastmath xAᵢ / x[j])
            end
        end
    end
    nothing
end

function evaluate(H::BinomialHomotopy, x, t, c::HomotopyNullCache)
    T = promote_type(eltype(x), ComplexF64)
    evaluate!(zeros(T, size(H.A, 1)), H, x, t, c)
end
(H::BinomialHomotopy)(x, t, c = cache(H, x, t)) = evaluate(H, x, t, c)

function jacobian(H::BinomialHomotopy, x, t, c::HomotopyNullCache)
    T = promote_type(eltype(x), ComplexF64)
    jacobian!(zeros(T, size(H.A)), H, x, t, c)
end
function dt(H::BinomialHomotopy, x, t, c::HomotopyNullCache)
    dt!(zeros(ComplexF64, size(H.A, 1)), H, x, t, c)
end