#
#--------------------------------------#
export gordonHall
#--------------------------------------#
"""
 Transfinite interpolation
"""
function gordonHall(xrm,xrp,xsm,xsp,yrm,yrp,ysm,ysp,zr,zs)

ze  = [-1,1]
Jer = interpMat(zr,ze)
Jes = interpMat(zs,ze)

xv = [xrm[1]   xrp[1]
      xrm[end] xrp[end]]

yv = [yrm[1]   yrp[1]
      yrm[end] yrp[end]]

xv = ABu(Jes,Jer,xv)
yv = ABu(Jes,Jer,yv)

#display(mesh(xv,yv,0*xv,0,90))

x = ABu([],Jer,vcat(xrm',xrp')) .+ ABu(Jes,[],hcat(xsm,xsp)) .- xv
y = ABu([],Jer,vcat(yrm',yrp')) .+ ABu(Jes,[],hcat(ysm,ysp)) .- yv

#display(mesh(x,y,0*x,0,90))

return x,y
end
#--------------------------------------#
export annulus
#--------------------------------------#
"""
 Convert (r,s) in [-1,1]^2 grid
 internal/external radii r0,r1,
 and angular extent span
"""
function annulus(r,s;r0=0.5,r1=1.0,span=2pi)

R  = @. (r1-r0)/2*(r+1) + r0
th = @. span   /2*(s+1) + 0.

x = @. R * cos(th)
y = @. R * sin(th)

return x,y
end
#--------------------------------------#
#"""
# Transforms [-1,1]^2 to annulus with
# internal/external radii r0,r1,
# and angular extent span
#"""
#function annulus(r0,r1,span,zr,zs)
#
#ze  = [-1.,1.]
#Jer = interpMat(zr,ze);
#Jes = interpMat(zs,ze);
#
#xrp = 0*zs;
#yrp = Jes*[-r0,-r1];
#
#spn = span-pi/2;
#
#xrm = Jes*[r0;r1]*cos(spn);
#yrm = Jes*[r0;r1]*sin(spn);
#
#as = Jer*[-pi/2,spn];
#xsm = r0 .* cos.(as);
#ysm = r0 .* sin.(as);
#
#xsp = r1 .* cos.(as);
#ysp = r1 .* sin.(as);
#
#pl=plot(xsm,ysm ,label="sm");
#pl=plot!(xsp,ysp,label="sp");
#pl=plot!(xrm,yrm,label="rm");
#pl=plot!(xrp,yrp,label="rp");
#pl=xlabel!("x")
#pl=ylabel!("y")
#display(pl)
#
#return gordonHall(xrm,xrp,xsm,xsp,yrm,yrp,ysm,ysp,zr,zs)
#end
#--------------------------------------#


using ODE

# Definition of parameters
const a = 1.0
const b = 0.1
const c = 1.5
const d = 0.75

function diffeq(t,u)
    du = similar(u)
    # u1: rabbits
    # u2: foxes
    du[1] = a*u[1] - b*u[1]*u[2]
    du[2] = -c*u[2] + d*b*u[1]*u[2]

    return du

end

# Simulation time
const t = collect(range(0.0,stop=18.0,length=1000))

# Initial condition
u0 = zeros(2)
u0[1] = 10
u0[2] = 5

# Solve ODE
(sol_t,sol_u) = ode45(diffeq,u0,tpoints=:specified)

#= Visualization of Results

using PyPlot
plot(sol_t,sol_u)
show()

=#


### Notation
#   small letters => vectors
# capital letters => matrices

function build_A(k, u, η, D, D2, ϕ, params)

    a = params.a
    N = params.Nϕ
    g = params.g

    k2 = k^2
    ik = 1im*k

    h = η .+ params.H
    axcos = a * cos.(ϕ)

           D = D/a
       IxCos = diagm(0=>axcos)
       IoCos = diagm(0=>1 ./ axcos)
    TwoΩxSin = diagm(0=>params.TwoΩ*sin.(ϕ));
       UoCos = diagm(0=>u./axcos)
       UxTan = diagm(0=>u./axcos.*sin.(ϕ))
       HoCos = diagm(0=>h./axcos)
       HxCos = diagm(0=>h.*axcos)
          dU = diagm(0=>D*u);
          dH = diagm(0=>D*h);

    # Form 1L-RSW Matrix
    #   [u1, v1, h1]

     U = diagm(0 => u)
     H = diagm(0 => h)
     F = diagm(0 => params.TwoΩ .+ 0*ϕ)
    dU = diagm(0 => D*u)

    A = [                UoCos             (-TwoΩxSin - UxTan + dU)[:, 2:N]    g*IoCos;
          -(TwoΩxSin + 2*UxTan)[2:N,:]/k2                     UoCos[2:N,2:N]  -g/k2*D[2:N,:];
                         HoCos                  (IxCos * D * HxCos)[:, 2:N]     UoCos];
    return A
end


using JuMP, NLopt, Test

model = Model(NLopt.Optimizer)
set_optimizer_attribute(model, "algorithm", :LD_MMA)

a1 = 2
b1 = 0
a2 = -1
b2 = 1

@variable(model, x1)
@variable(model, x2 >= 0)

@NLobjective(model, Min, sqrt(x2))
@NLconstraint(model, x2 >= (a1*x1+b1)^3)
@NLconstraint(model, x2 >= (a2*x1+b2)^3)

set_start_value(x1, 1.234)
set_start_value(x2, 5.678)

optimize!(model)

println("got ", objective_value(model), " at ", [value(x1), value(x2)])

@test_approx_eq_eps value(x1) 1/3 1e-5
@test_approx_eq_eps value(x2) 8/27 1e-5
@test_approx_eq_eps objective_value(m) sqrt(8/27) 1e-5
@test termination_status(model) == MOI.OPTIMAL


123.4 + 234.5     # This is a comment, ignored by Julia

7 / 3             # Integer division returns floating point

7 ÷ 3             # Truncated integer division (type \div and press tab)

7 % 3             # Remainder

2^7 - 4           # All integer expression which returns an integer

3 * (4.2 - 9.7)   # Use parentheses to control the order of operations

# We will later learn why the previous result is not exact

5.6e7 ^ 0.1       # Scientific notation, means 5.6⋅10⁷ (not e⁷)

1 + 2
34 / 3

println(1 + 2.1)
println("The result is ", -123 / 1000000)
34 / 3; # The output will not be displayed

typeof(1.23)

typeof(-745)

typeof(1e20)

println(typemin(Int64))
println(typemax(Int64))
println(typemin(Int128))
println(typemax(Int128))

1000^7

Int128(1000)^7


"""
```
shapeFunction(::LinearQuadrangle3D, x, y, z)
```
Computes the linear shape functions `N[1:8]` for an `(x, y, z)` point on a 3D linear quadrangle.
"""
function shapeFunction(::LinearQuadrangle3D, x, y, z)
    # N[n](x,y,z) := shape function n.
    omx = 1 - x
    omy = 1 - y
    omz = 1 - z
    opx = 1 + x
    opy = 1 + y
    opz = 1 + z

    N =
        SVector{8, typeof(x)}(
            omx * omy * omz,
            opx * omy * omz,
            opx * opy * omz,
            omx * opy * omz,
            omx * omy * opz,
            opx * omy * opz,
            opx * opy * opz,
            omx * opy * opz,
        ) / 8

    return N
end
"""
```
shapeFunction(::LinearQuadrangle3D, x::AbstractVector, y::AbstractVector, z::AbstractVector)
```
Computes the linear shape functions `N[1:8][p]` for `(x, y, z)` point `p` on a 3D linear quadrangle.
"""
function shapeFunction(
    ::LinearQuadrangle3D,
    x::AbstractVector,
    y::AbstractVector,
    z::AbstractVector,
)
    numPoints = length(x)
    @assert numPoints == length(y) == length(z)
    xType = eltype(x)
    # N[n, p](x_vec,y_vec,z_vec) := shape function n for point p.
    N = Array{SVector{8, xType}}(undef, numPoints)

    @inbounds @simd for i in eachindex(x)
        omx = 1 - x[i]
        omy = 1 - y[i]
        omz = 1 - z[i]
        opx = 1 + x[i]
        opy = 1 + y[i]
        opz = 1 + z[i]

        N[i] =
            SVector{8, xType}(
                omx * omy * omz,
                opx * omy * omz,
                opx * opy * omz,
                omx * opy * omz,
                omx * omy * opz,
                opx * omy * opz,
                opx * opy * opz,
                omx * opy * opz,
            ) / 8
    end

    return N
end

"""
```
shapeFunctionDeriv(shape<:AbstractShapeFunction, x, y, z)
```
Computes the derivatives of the linear shape functions `N[1:3, 1:8]` for an `(x, y, z)` point on a 3D linear quadrangle.

## Returns
```
dNdS[x, n](x,y,z) := x'th derivative of shape function n.
```
"""
function shapeFunctionDeriv(::LinearQuadrangle3D, x, y, z)
    # dNdS[x, n](x,y,z) := x'th derivative of shape function n.
    omx = 1 - x
    omy = 1 - y
    omz = 1 - z
    opx = 1 + x
    opy = 1 + y
    opz = 1 + z

    dNdS =
        SMatrix{3, 8, typeof(x)}(
            -omy * omz,
            -omx * omz,
            -omx * omy,
            omy * omz,
            -opx * omz,
            -opx * omy,
            opy * omz,
            opx * omz,
            -opx * opy,
            -opy * omz,
            omx * omz,
            -omx * opy,
            -omy * opz,
            -omx * opz,
            omx * omy,
            omy * opz,
            -opx * opz,
            opx * omy,
            opy * opz,
            opx * opz,
            opx * opy,
            -opy * opz,
            omx * opz,
            omx * opy,
        ) / 8

    return dNdS
end
"""
```
shapeFunctionDeriv(shape<:AbstractShapeFunction, x::AbstractVector, y::AbstractVector, z::AbstractVector)
```
Computes the derivatives of the linear shape functions `N[1:3, 1:8]` for an `(x, y, z)` point on a 3D linear quadrangle.

## Returns
```
dNdS[x, n, p](x,y,z) := x'th derivative of shape function n for point p.
```
"""
function shapeFunctionDeriv(
    ::LinearQuadrangle3D,
    x::AbstractVector,
    y::AbstractVector,
    z::AbstractVector,
)
    numPoints = length(x)
    @assert numPoints == length(y) == length(z)
    xType = eltype(x)

    # dNdS[x, n, p](x,y,z) := x'th derivative of shape function n for point p.
    dNdS = Array{SMatrix{3, 8, xType}}(undef, numPoints)

    @inbounds @simd for i in eachindex(x)
        omx = 1 - x[i]
        omy = 1 - y[i]
        omz = 1 - z[i]
        opx = 1 + x[i]
        opy = 1 + y[i]
        opz = 1 + z[i]

        dNdS[i] =
            SMatrix{3, 8, xType}(
                -omy * omz,
                -omx * omz,
                -omx * omy,
                omy * omz,
                -opx * omz,
                -opx * omy,
                opy * omz,
                opx * omz,
                -opx * opy,
                -opy * omz,
                omx * omz,
                -omx * opy,
                -omy * opz,
                -omx * opz,
                omx * omy,
                omy * opz,
                -opx * opz,
                opx * omy,
                opy * opz,
                opx * opz,
                opx * opy,
                -opy * opz,
                omx * opz,
                omx * opy,
            ) / 8
    end

    return dNdS
end


using DataDrivenDiffEq
using ModelingToolkit
using OrdinaryDiffEq
using LinearAlgebra
using Plots
gr()

# Create a
function pendulum(u, p, t)
    x = u[2]
    y = -9.81sin(u[1]) - 0.1u[2]^3 -0.2*cos(u[1])
    return [x;y]
end

u0 = [0.99π; -1.0]
tspan = (0.0, 20.0)
prob = ODEProblem(pendulum, u0, tspan)
sol = solve(prob, Tsit5(), saveat = 0.3)

plot(sol)

# Create the differential data
DX = similar(sol[:,:])
for (i, xi) in enumerate(eachcol(sol[:,:]))
    DX[:,i] = pendulum(xi, [], 0.0)
end

# Create a basis
@variables u[1:2]
@parameters w[1:2]
# Lots of polynomials
polys = Operation[1]
for i ∈ 1:5
    push!(polys, u.^i...)
    for j ∈ 1:i-1
        push!(polys, u[1]^i*u[2]^j)
    end
end

# And some other stuff
h = [w[1]*cos(u[1]); w[2]*sin(u[1]); u[1]*u[2]; u[1]*sin(u[2]); u[2]*cos(u[2]); polys...]

basis = Basis(h, u, parameters = w)
println(basis)

# Get the reduced basis via the sparse regression
# Thresholded Sequential Least Squares, works fine for more data
# than assumptions, converges fast but fails sometimes with too much noise
opt = STRRidge(1e-2)
# Enforce all 100 iterations
Ψ = SINDy(sol[:,1:25], DX[:, 1:25], basis, opt, p = [1.0; 1.0], maxiter = 100, convergence_error = 1e-5)
println(Ψ)
print_equations(Ψ)

# Lasso as ADMM, typically needs more information, more tuning
opt = ADMM(1e-2, 1.0)
Ψ = SINDy(sol[:,1:50], DX[:, 1:50], basis, opt, p = [1.0; 1.0], maxiter = 5000, convergence_error = 1e-3)
println(Ψ)
print_equations(Ψ)

# Get the associated parameters out of the result
parameters(Ψ)

# SR3, works good with lesser data and tuning
opt = SR3(1e-2, 1.0)
Ψ = SINDy(sol[:,1:end], DX[:, 1:end], basis, opt, p = [0.5; 0.5], maxiter = 5000, convergence_error = 1e-5)
println(Ψ)
print_equations(Ψ, show_parameter = true)

# Vary the sparsity threshold -> gives better results
λs = exp10.(-5:0.1:-1)
# Use SR3 with high relaxation (allows the solution to diverge from LTSQ) and high iterations
opt = SR3(1e-2, 5.0)
Ψ = SINDy(sol[:,1:10], DX[:, 1:10], basis, λs, opt, p = [1.0; 1.0], maxiter = 15000)
println(Ψ)
print_equations(Ψ)

# Transform into ODE System
sys = ODESystem(Ψ)
dudt = ODEFunction(sys)
ps = parameters(Ψ)

# Simulate
estimator = ODEProblem(dudt, u0, tspan, ps)
sol_ = solve(estimator, Tsit5(), saveat = sol.t)

# Yeah! We got it right
scatter(sol.t[1:10], sol[:,1:10]', color = :red, label = nothing)
scatter!(sol.t[11:end], sol[:,11:end]', color = :blue, label = nothing)
plot!(sol_.t, sol_[:, :]', color = :green, label = "Estimation")

plot(sol.t, abs.(sol-sol_)')
norm(sol[:,:]-sol_[:,:], 2)


# Corresponds to https://en.wikipedia.org/wiki/Logarithmic_distribution
struct Logarithmic{T<:Real} <: Distributions.DiscreteUnivariateDistribution
    p::T
    function Logarithmic(p::T) where {T <: Real}
        new{T}(p)
    end
end
function Distributions.logpdf(d::Logarithmic, x::Real)
    insupport(d, x) ? x*log(d.p) - log(x) - log(-log(1-d.p)) : log(zero(d.p))
end
function Distributions.rand(rng::Distributions.AbstractRNG, d::Logarithmic)
    # Sample a Log(p) distribution with the algorithm "LK" of Kemp (1981).
    u = rand(rng)
    if u > d.p
        return 1
    end
    q = 1 - (1-d.p)^rand(rng)
    if u < q*q
        return floor(1+log(u)/log(q))
    end
    if u < q
        return 1
    end
    return 2
end


using Diagonalizations, LinearAlgebra, PosDefManifold, Test

# Method (1) real
n, t=10, 100
X=genDataMatrix(n, t)
C=(X*X')/t
wC=whitening(Hermitian(C); simple=true)
# or, shortly
wC=whitening(ℍ(C); simple=true)

# Method (1) complex
Xc=genDataMatrix(ComplexF64, n, t)
Cc=(Xc*Xc')/t
wCc=whitening(Hermitian(Cc); simple=true)

# Method (2) real
wX=whitening(X; simple=true)
@test wC.F'*C*wC.F≈I
@test wX.F'*C*wX.F≈I
@test wX≈wC

# Method (2) complex
wXc=whitening(Xc; simple=true)
@test wCc.F'*Cc*wCc.F≈I
@test wXc.F'*Cc*wXc.F≈I
@test wXc≈wCc

# Method (3) real
k=10
Xset=[genDataMatrix(n, t) for i=1:k]

# whitening on the average covariance matrix
w=whitening(Xset)

# ... selecting subspace dimension allowing an explained variance = 0.5
w=whitening(Xset; eVar=0.5)

# ... averaging the covariance matrices using the logEuclidean metric
w=whitening(Xset; metric=logEuclidean, eVar=0.5)

# ... giving weights `w` to the covariance matrices
w=whitening(Xset; metric=logEuclidean, w=abs2.(randn(k)), eVar=0.5)

# ... subtracting the mean
w=whitening(Xset; meanX=nothing, metric=logEuclidean, w=abs2.(randn(k)), eVar=0.5)

# whitening on the average of the covariance matrices computed along dims 1
w=whitening(Xset; dims=1)

# explained variance
w.eVar

# name of the filter
w.name

using Plots
# plot regularized accumulated eigenvalues
plot(w.arev)

# plot the original covariance matrix and the whitened covariance matrix
 Cmax=maximum(abs.(C));
 h1 = heatmap(C, clim=(-Cmax, Cmax), yflip=true, c=:bluesreds, title="C");
 D=wC.F'*C*wC.F;
 h2 = heatmap(D, clim=(0, 1), yflip=true, c=:amp, title="F'*C*F");
 📈=plot(h1, h2, size=(700, 300))
# savefig(📈, homedir()*"\\Documents\\Code\\julia\\Diagonalizations\\docs\\src\\assets\\FigWhitening.png")

# Method (3) complex
k=10
Xcset=[genDataMatrix(ComplexF64, n, t) for i=1:k]

# whitening on the average covariance matrix
wc=whitening(Xcset)

# ... selecting subspace dimension allowing an explained variance = 0.5
wc=whitening(Xcset; eVar=0.5)

# ... averaging the covariance matrices using the logEuclidean metric
wc=whitening(Xcset; metric=logEuclidean, eVar=0.5)

# ... giving weights `w` to the covariance matrices
wc=whitening(Xset; metric=logEuclidean, w=abs2.(randn(k)), eVar=0.5)

# ... subtracting the mean
wc=whitening(Xcset; meanX=nothing, metric=logEuclidean, w=abs2.(randn(k)), eVar=0.5)

# whitening on the average of the covariance matrices computed along dims 1
wc=whitening(Xcset; dims=1)

# explained variance
wc.eVar

# name of the filter
wc.name


export IdentityLink, LogLink, LogitLink
export link, unlink, jacobian_link

abstract type Link end

"""
    link(args...)

The link function is a map that brings a parameter ``f`` in a subspace ``\\mathcal{F} \\subset \\mathbb{R}`` to ``\\mathbb{R}``.
"""
function link end

"""
    unlink(args...)

The unlink function is the inverse map of [`link`](@ref link). It brings ``\\tilde{f}`` in ``\\mathbb{R}`` to the subspace ``\\mathcal{F} \\subset \\mathbb{R}``.
"""
function unlink end

"""
    jacobian_link(args...)

Evaluates the derivative of the [`link`](@ref link) with respect to the parameter ``f``.
"""
function jacobian_link end

"""
    IdentityLink <: Link

Define the map ``\\tilde{f} = f`` where ``f \\in \\mathbb{R}`` and ``\\tilde{f} \\in \\mathbb{R}``
"""
struct IdentityLink <: Link end

link(::Type{IdentityLink}, param::T) where T = param
unlink(::Type{IdentityLink}, param_tilde::T) where T = param_tilde
jacobian_link(::Type{IdentityLink}, param::T) where T = one(T)

"""
    LogLink <: Link

Define the map ``\\tilde{f} = \\ln(f - a)`` where ``f \\in [a, \\infty), a \\in \\mathbb{R}`` and ``\\tilde{f} \\in \\mathbb{R}``
"""
struct LogLink <: Link end

link(::Type{LogLink}, param::T, lower_bound::T) where T = log(param - lower_bound)
unlink(::Type{LogLink}, param_tilde::T, lower_bound::T) where T = exp(param_tilde) + lower_bound
jacobian_link(::Type{LogLink}, param::T, lower_bound::T) where T = 1/(param - lower_bound)

"""
    LogitLink <: Link

Define the map ``\\tilde{f} = \\ln(\\frac{f - a}{b - f})`` where ``f \\in [a, b], a, b \\in \\mathbb{R}`` and ``\\tilde{f} \\in \\mathbb{R}``
"""
struct LogitLink <: Link end

function link(::Type{LogitLink}, param::T, lower_bound::T, upper_bound::T) where T
    return log((param - lower_bound)/(upper_bound - param))
end
function unlink(::Type{LogitLink}, param_tilde::T, lower_bound::T, upper_bound::T) where T
    return lower_bound + ((upper_bound - lower_bound)/(1 + exp(-param_tilde)))
end
function jacobian_link(::Type{LogitLink}, param::T, lower_bound::T, upper_bound::T) where T
    return (upper_bound + lower_bound)/((upper_bound - param) * (param - lower_bound))
end

const LINKS = [
    IdentityLink;
    LogLink;
    LogitLink
]


using EconPDEs, Distributions

struct AchdouHanLasryLionsMollModel
    # income process parameters
    κy::Float64
    ybar::Float64
    σy::Float64

    r::Float64

    # utility parameters
    ρ::Float64
    γ::Float64

    amin::Float64
    amax::Float64
end

function AchdouHanLasryLionsMollModel(;κy = 0.1, ybar = 1.0, σy = 0.07, r = 0.03, ρ = 0.05, γ = 2.0, amin = 0.0, amax = 500.0)
    AchdouHanLasryLionsMollModel(κy, ybar, σy, r, ρ, γ, amin, amax)
end

function (m::AchdouHanLasryLionsMollModel)(state::NamedTuple, value::NamedTuple)
    (; κy, σy, ybar, r, ρ, γ, amin, amax) = m
    (; y, a) = state
    (; v, vy_up, vy_down, va_up, va_down, vyy, vya, vaa) = value
    μy = κy * (ybar - y)
    vy = (μy >= 0) ? vy_up : vy_down

    va = va_up
    iter = 0
    @label start
    va_up = max(va_up, eps())
    c = va^(-1 / γ)
    μa = y + r * a - c
    if (iter == 0) & (μa <= 0)
        iter += 1
        va = va_down
        @goto start
    end
    if (a ≈ amin) && (μa <= 0.0)
        va = (y + r * amin)^(-γ)
        c = y + r * amin
        μa = 0.0
    end
    vt = - (c^(1 - γ) / (1 - γ) + μa * va + μy * vy + 0.5 * vyy * σy^2 - ρ * v)
    return (; vt)
end

m = AchdouHanLasryLionsMollModel()
distribution = Gamma(2 * m.κy * m.ybar / m.σy^2, m.σy^2 / (2 * m.κy))
stategrid = OrderedDict(:y => range(quantile(distribution, 0.001), quantile(distribution, 0.999), length = 10),
                        :a =>  range(m.amin, m.amax, length = 100)
                        )
yend = OrderedDict(:v => [log(y + max(a, 0.0)) for y in stategrid[:y], a in stategrid[:a]])
result = pdesolve(m, stategrid, yend)
@assert result.residual_norm <= 1e-5

# finite horizon over 20 years
yend = OrderedDict(:v => [max(a + y)^(1-m.γ)/(1-m.γ) for y in stategrid[:y], a in stategrid[:a]])
τs = range(0, stop = 100, step = 1)
result  = pdesolve(m, stategrid, yend, τs)
@assert maximum(result.residual_norm) <= 1e-5

# Check marginal value of wealth converges to 1.0 at infinity
#b = ((m.r + (m.ρ - m.r)/m.γ))^(1/(1 - 1/m.γ))
#pw = (result[:v] * (1-m.γ)).^(1/(1-m.γ)-1) .* result[:va] ./ b


using Stats
using Base.Test

@test_approx_eq autocor([1, 2, 3, 4, 5], 1) 0.4

@test iqr([1, 2, 3, 4, 5]) == [2.0, 4.0]

z = [true, true, false, false, true, false, true, true, true]
values, lengths = rle(z)
@test values == [true, false, true, false, true]
@test lengths == [2, 2, 1, 1, 3]
@test inverse_rle(values, lengths) == z

z = [true, true, false, false, true, false, true, true, true, false]
values, lengths = rle(z)
@test values == [true, false, true, false, true, false]
@test lengths == [2, 2, 1, 1, 3, 1]
@test inverse_rle(values, lengths) == z

X = [1 0; 2 1; 3 0; 4 1; 5 10]
y = [5, 3, 4, 2, 5]
@test_approx_eq cor_spearman(X, y)[1] cor_spearman(X[:,1],y)
@test_approx_eq cor_spearman(X) cor_spearman(X, X)
@test_approx_eq cor_spearman(X, y) [-0.102597835208515, -0.081110710565381]
@test_approx_eq cor_kendall(X,y) [-0.105409255338946, -0.117851130197758]

fnecdf = ecdf(randn(10000000))
@test_approx_eq_eps fnecdf([-1.96, -1.644854, -1.281552, -0.6744898, 0, 0.6744898, 1.281552, 1.644854, 1.96]) [0.025, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.975] 1e-3
@test_approx_eq_eps fnecdf(1.96) 0.975 1e-3
@test_approx_eq fnecdf([-1.96, -1.644854, -1.281552, -0.6744898, 0, 0.6744898, 1.281552, 1.644854, 1.96]) map(fnecdf, [-1.96, -1.644854, -1.281552, -0.6744898, 0, 0.6744898, 1.281552, 1.644854, 1.96])


# InverseGamma distribution

export InverseGamma

@parameterized InverseGamma(shape) ≃ Lebesgue(ℝ₊)

function logdensity(μ::InverseGamma{(:shape,)}, x)
    α = μ.shape
    xinv = 1/x

    return xlogy(α + 1, xinv) - xinv - loggamma(α)
end

Base.rand(rng::AbstractRNG, T::Type, μ::InverseGamma{(:shape,)}) = rand(rng, Dists.InverseGamma(μ.shape))

≪(::InverseGamma, ::Lebesgue{X}) where X <: Real = true

TV.as(::InverseGamma) = asℝ₊

# @μσ_methods InverseGamma(shape)


using RayTracer
using LinearAlgebra
using Test

@testset "creating_4_x_4_matrix" begin
    m = [1 2 3 4; 5.5 6.5 7.5 8.5; 9 10 11 12; 13.5 14.5 15.5 16.5]
    @test m[1,1] ≈ 1
    @test m[1,4] ≈ 4
    @test m[2,1] ≈ 5.5
    @test m[2,3] ≈ 7.5
    @test m[3,3] ≈ 11
    @test m[4,1] ≈ 13.5
    @test m[4,3] ≈ 15.5
end

@testset "equal_matrices" begin
    m = [1 2 3 4; 5.5 6.5 7.5 8.5; 9 10 11 12; 13.5 14.5 15.5 16.5]
    n = [1 2 3 4; 5.5 6.5 7.5 8.5; 9 10 11 12; 13.5 14.5 15.5 16.5]
    @test m ≈ n
end

@testset "different_matrices" begin
    m = [1 2 3 4; 5 6 7 8; 9 10 11 12; 13 14 15 16]
    n = [2 3 4 5; 6 7 8 9; 10 11 12 13; 14 15 16 1]
    @test m ≠ n
end

@testset "matrix_tuple_multiplication" begin
    a = [1 2 3 4; 2 4 4 2; 8 6 4 1; 0 0 0 1]
    b = RayTracer.tuple(1, 2, 3, 1)
    @test a*b ≈ RayTracer.tuple(18,24,33,1)
end

@testset "creating_submatix" begin
    a = [1 5 0; -3 2 7; 0 6 -3]
    b = RayTracer.submatrix(a,1,3)
    @test b ≈ [-3 2; 0 6]
end

@testset "computing minor" begin
    a = [3 5 0; 2 -1 -7; 6 -1 5]
    b = RayTracer.submatrix(a,2,1)
    @test det(b) ≈ 25
    @test RayTracer.minor(a,2,1) ≈ 25
end

@testset "computing_cofactor" begin
    a = [3 5 0; 2 -1 -7; 6 -1 5]
    @test RayTracer.minor(a,1,1) ≈ -12
    @test RayTracer.cofactor(a,1,1) ≈ -12
    @test RayTracer.minor(a,2,1) ≈ 25
    @test RayTracer.cofactor(a,2,1) ≈ -25
end

@testset "multiplying_by_translation_matrix" begin
    transform = RayTracer.translation(5,-3,2)
    p = RayTracer.point(-3,4,5)
    @test transform * p ≈ RayTracer.point(2,1,7)
end

@testset "multiplying_by_inverse_translation" begin
    transform = RayTracer.translation(5,-3,2)
    inv_transform = inv(transform)
    p = RayTracer.point(-3,4,5)
    @test inv_transform * p ≈ RayTracer.point(-8,7,3)
end

@testset "translation_does_not_affect_vectors" begin
    transform = RayTracer.translation(5,-3,2)
    v = RayTracer.vector(-3,4,5)
    @test transform*v ≈ v
end

@testset "scaling_applied_to_point" begin
    transform = RayTracer.scaling(2,3,4)
    p = RayTracer.point(-4,6,8)
    @test transform*p ≈ RayTracer.point(-8,18,32)
end

@testset "scaling_applied_to_vector" begin
    transform = RayTracer.scaling(2,3,4)
    v = RayTracer.vector(-4,6,8)
    @test transform*v ≈ RayTracer.vector(-8,18,32)
end

@testset "multiplying_inverse_scaling" begin
    transform = RayTracer.scaling(2,3,4)
    inv_transform = inv(transform)
    v = RayTracer.vector(-4,6,8)
    @test inv_transform*v ≈ RayTracer.vector(-2,2,2)
end

@testset "reflection_is_scaling_negative" begin
    transform = RayTracer.scaling(-1,1,1)
    p = RayTracer.point(2,3,4)
    @test transform*p ≈ RayTracer.point(-2,3,4)
end

@testset "rotation_around_x_axis" begin
    p = RayTracer.point(0,1,0)
    eith_turn = RayTracer.rotation_x(π/4)
    quarter_turn = RayTracer.rotation_x(π/2)
    @test eith_turn*p ≈ RayTracer.point(0,√2/2,√2/2)
    @test ≈(quarter_turn*p,RayTracer.point(0,0,1),1e-10)
end

@testset "inverse_rotation_x_axis" begin
    p = RayTracer.point(0,1,0)
    eith_turn = RayTracer.rotation_x(π/4)
    inv_transform = inv(eith_turn)
    @test inv_transform*p ≈ RayTracer.point(0,√2/2,-√2/2)
end

@testset "rotation_around_y_axis" begin
    p = RayTracer.point(0,0,1)
    eith_turn = RayTracer.rotation_y(π/4)
    quarter_turn = RayTracer.rotation_y(π/2)
    @test eith_turn*p ≈ RayTracer.point(√2/2,0,√2/2)
    @test ≈(quarter_turn*p,RayTracer.point(1,0,0),1e-10)
end

@testset "rotation_around_z_axis" begin
    p = RayTracer.point(0,1,0)
    eith_turn = RayTracer.rotation_z(π/4)
    quarter_turn = RayTracer.rotation_z(π/2)
    @test eith_turn*p ≈ RayTracer.point(-√2/2,√2/2,0)
    @test ≈(quarter_turn*p,RayTracer.point(-1,0,0),1e-10)
end

@testset "shearing_x_wrt_y" begin
    transform = RayTracer.shearing(1,0,0,0,0,0)
    p = RayTracer.point(2,3,4)
    @test transform*p ≈ RayTracer.point(5,3,4)
end

@testset "shearing_x_wrt_z" begin
    transform = RayTracer.shearing(0,1,0,0,0,0)
    p = RayTracer.point(2,3,4)
    @test transform*p ≈ RayTracer.point(6,3,4)
end

@testset "shearing_y_wrt_x" begin
    transform = RayTracer.shearing(0,0,1,0,0,0)
    p = RayTracer.point(2,3,4)
    @test transform*p ≈ RayTracer.point(2,5,4)
end

@testset "shearing_y_wrt_z" begin
    transform = RayTracer.shearing(0,0,0,1,0,0)
    p = RayTracer.point(2,3,4)
    @test transform*p ≈ RayTracer.point(2,7,4)
end

@testset "shearing_z_wrt_x" begin
    transform = RayTracer.shearing(0,0,0,0,1,0)
    p = RayTracer.point(2,3,4)
    @test transform*p ≈ RayTracer.point(2,3,6)
end

@testset "shearing_z_wrt_y" begin
    transform = RayTracer.shearing(0,0,0,0,0,1)
    p = RayTracer.point(2,3,4)
    @test transform*p ≈ RayTracer.point(2,3,7)
end

@testset "individual_transformations_sequential" begin
    p = RayTracer.point(1,0,1)
    A = RayTracer.rotation_x(π/2)
    B = RayTracer.scaling(5,5,5)
    C = RayTracer.translation(10,5,7)
    p2 = A*p
    @test ≈(p2,RayTracer.point(1,-1,0),1e-10)
    p3 = B*p2
    @test ≈(p3,RayTracer.point(5,-5,0),1e-10)
    p4 = C*p3
    @test ≈(p4,RayTracer.point(15,0,7),1e-10)
end

@testset "chained_transformations" begin
    p = RayTracer.point(1,0,1)
    A = RayTracer.rotation_x(π/2)
    B = RayTracer.scaling(5,5,5)
    C = RayTracer.translation(10,5,7)
    D = C*B*A
    @test ≈(D*p,RayTracer.point(15,0,7),1e-10)
end


# Triangle, square, pentagonal, hexagonal, heptagonal, and octagonal numbers are
# all figurate (polygonal) numbers and are generated by the following formulae:
#
# Triangle    P3,n=n(n+1)/2    1, 3, 6, 10, 15, ...
# Square      P4,n=n2          1, 4, 9, 16, 25, ...
# Pentagonal  P5,n=n(3n−1)/2   1, 5, 12, 22, 35, ...
# Hexagonal   P6,n=n(2n−1)     1, 6, 15, 28, 45, ...
# Heptagonal  P7,n=n(5n−3)/2   1, 7, 18, 34, 55, ...
# Octagonal   P8,n=n(3n−2)     1, 8, 21, 40, 65, ...
#
# The ordered set of three 4-digit numbers: 8128, 2882, 8281, has three
# interesting properties.
# 1. The set is cyclic, in that the last two digits of each number is the first
# two digits of the next number (including the last number with the first).
# 2. Each polygonal type: triangle (P3,127=8128), square (P4,91=8281), and
# pentagonal (P5,44=2882), is represented by a different number in the set.
# 4. This is the only set of 4-digit numbers with this property.
#
# Find the sum of the only ordered set of six cyclic 4-digit numbers for which
# each polygonal type: triangle, square, pentagonal, hexagonal, heptagonal, and
# octagonal, is represented by a different number in the set.

using ProjectEulerSolutions

# Solve for n given x and s: P(s,n) = x.
function polygonal_inv(s::Integer, x::Integer)::Float32
    return (sqrt(8*(s-2)*x + (s-4)^2) + (s-4))/(2*(s-2))
end

# Implements polygonal number P(s,n) = x, where s is the number of sides.
function polygonal_number(s::Integer, n::Integer)::Integer
   return div((s - 2) * n * n - n * (s - 4), 2)
end

# Helper recursive function to find the numbers that match the beginning and
# ending pairs.
function find_loop!(numbers::Array{Integer,1}, used_polys::Array{Integer,1}, end_pair::Integer, polys::Array{Array{Integer,1},1})::Bool
    d = 10^(ndigits(numbers[1])-2)
    if length(numbers) == 6 && end_pair == div(numbers[1], d)
        return true
    end
    ks = filter(x -> !(x in used_polys), 2:length(polys))
    for k in ks
        push!(used_polys, k)
        for n in polys[k]
            if div(n, d) == end_pair
                push!(numbers, n)
                found = find_loop!(numbers, used_polys, n % 100, polys)
                if found
                    return true
                end
                pop!(numbers)
            end
        end
        pop!(used_polys)
    end
    return false
end

# This is a verbose solution primarily due to manipulating the data into the
# correct data structures.  Only cycle through 4 digit polygonal numbers, then
# recursively loop through collections, remembering which polygonal numbers
# have already been used.  There is probably a graph-based solution which find
# paths through a network (6-partite graph?).
function p061solution(ndig_start::Integer=4, ndig_end::Integer=4)::Integer

    min_n = 10^(ndig_start-1)
    max_n = 10^ndig_end-1

    polys = Array{Array{Integer,1},1}()
    # Initialize data structure, making use of vectorized instructions
    for s = 3:8
        min_num = Integer(ceil(polygonal_inv(s, min_n)))
        max_num = Integer(floor(polygonal_inv(s, max_n)))

        # Calculate all 4 digit polygonal numbers for s
        poly_s = polygonal_number.(s, min_num:max_num)

        # Only keep numbers whose lower halves have two digits
        lower_s = poly_s .% 100
        inds = lower_s .> 9
        push!(polys, poly_s[inds])
    end

    # Loop through potential starting points, starting with s = 3.  The order
    # doesn't matter since all polygonal values are eventually used.
    used_polys = Array{Integer,1}([1])
    numbers = Array{Integer,1}()
    for n in polys[1]
        push!(numbers, n)
        if find_loop!(numbers, used_polys, n % 100, polys)
             return sum(numbers)
        end
        pop!(numbers)
    end

    return -1
end

p061 = Problems.Problem(p061solution)

Problems.benchmark(p061, 4, 4)


#using TensorToolbox, Test, LinearAlgebra

println("\n\n****Testing ktensor.jl")

Isz=5;J=4;K=3;
a=rand(Isz);
b=rand(J);
c=rand(K);
X=zeros(Isz,J,K);
for i=1:Isz,j=1:J,k=1:K
    X[i,j,k]=a[i]*b[j]*c[k];
end
println("\n**Test rank-1 tensor X of size: ", size(X))
println("...Testing cp_als.")
Z=cp_als(X,1)
err=norm(full(Z)-X)
println("norm(full(Xcp)-X) = ",err)
@test err ≈ 0 atol=1e-10

#Ax = rand(4,2)
#Bx = rand(3,2)
#Cx = rand(2,2)
#Xlambda=[0.5,2]
#X = ktensor(Xlambda,[Ax,Bx,Cx])
X=randktensor([4,3,2],2)
Ay = rand(4,3)
By = rand(3,3)
Cy = rand(2,3)
Ylambda=[2,3,1]
Y=ktensor(Ylambda,[Ay,By,Cy])

println("\n**Test ktensors X and Y of size: ", size(X))
println()

println("...Testing function ndims.")
println("X is of order: ",ndims(X))

println("...Testing function ttensor and full.")
T=ttensor(X)
@test norm(full(T)-full(X)) ≈ 0 atol=1e-10

println("...Testing scalar multiplication.")
err=norm(full(3*X) - 3*full(X))
println("norm(full(3*X) - 3*full(X)) = ",err)
@test norm(full(3*X) - 3*full(X)) ≈ 0 atol=1e-10

println("...Testing functions plus and minus.")
Z=X+Y
err=norm(full(Z)-(full(X)+full(Y)))
println("norm(full(Z)-(full(X)+full(Y))) = ",err)
@test err ≈ 0 atol=1e-10
Z=X-Y
err=norm(full(Z)-(full(X)-full(Y)))
println("norm(full(Z)-(full(X)-full(Y))) = ",err)
@test err ≈ 0 atol=1e-10

println("...Testing function innerprod.")
err=norm(innerprod(X,Y)-innerprod(full(X),full(Y)))
println("norm(innerprod(X,Y)-innerprod(full(X),full(Y))) = ",err)
@test err  ≈ 0 atol=1e-10
Z=rand(4,3,2)
@test norm(innerprod(X,Z)-innerprod(full(X),Z)) ≈ 0 atol=1e-10

println("...Testing function norm.")
@test abs(norm(X)-norm(full(X)))≈ 0 atol=1e-10

println("...Testing functions arrange and arrange!.")
Z=arrange(X)
err=norm(full(X)-full(Z))
println("norm(full(X)-full(Z)) = ",err)
@test err ≈ 0 atol=1e-10
Z=arrange(X,2)
@test norm(full(X)-full(Z)) ≈ 0 atol=1e-10
Z=arrange(X,[2,1])
@test norm(full(X)-full(Z)) ≈ 0 atol=1e-10
Xcopy=deepcopy(X)
arrange!(X)
@test norm(full(X)-full(Xcopy)) ≈ 0 atol=1e-10
arrange!(X,2)
@test norm(full(X)-full(Xcopy)) ≈ 0 atol=1e-10
arrange!(X,[2,1])
@test norm(full(X)-full(Xcopy)) ≈ 0 atol=1e-10

println("...Testing functions normalize and normalize!.")
Z=normalize(X)
err=norm(full(X)-full(Z))
println("norm(full(X)-full(Z)) = ",err)
@test err ≈ 0 atol=1e-10
Z=normalize(X,2)
@test norm(full(X)-full(Z)) ≈ 0 atol=1e-10
Z=normalize(X,3,normtype=1)
@test norm(full(X)-full(Z)) ≈ 0 atol=1e-10
Z=normalize(X,"sort")
@test norm(full(X)-full(Z)) ≈ 0 atol=1e-10
Z=normalize(X,factor=2)
@test norm(full(X)-full(Z)) ≈ 0 atol=1e-10
Xcopy=deepcopy(X)
normalize!(X)
@test norm(full(X)-full(Xcopy)) ≈ 0 atol=1e-10
normalize!(X,2)
@test norm(full(X)-full(Xcopy)) ≈ 0 atol=1e-10
normalize!(X,3,normtype=1)
@test norm(full(X)-full(Xcopy)) ≈ 0 atol=1e-10
normalize!(X,"sort")
@test norm(full(X)-full(Xcopy)) ≈ 0 atol=1e-10
normalize!(X,factor=2)
@test norm(full(X)-full(Xcopy)) ≈ 0 atol=1e-10

println("...Testing functions redistribute and redistribute!.")
mode=2
Z=redistribute(X,mode)
err=norm(full(X)-full(Z))
println("norm(full(X)-full(Z)) = ",err)
@test err ≈ 0 atol=1e-10
Xcopy=deepcopy(X)
redistribute!(X,mode)
@test norm(full(X)-full(Xcopy)) ≈ 0 atol=1e-10

println("...Testing functions fixsigns, fixsigns! and isequal.")
Xcopy=deepcopy(X)
X.fmat[1][:,1]=-X.fmat[1][:,1]
X.fmat[2][:,1]=-X.fmat[2][:,1]
Z=fixsigns(X)
err=norm(full(X)-full(Z))
println("norm(full(X)-full(Z)) = ",err)
@test err ≈ 0 atol=1e-10
fixsigns!(X)
@test norm(X-Xcopy) ≈ 0 atol=1e-7

println("...Testing functions ttm and ttv.")
M=MatrixCell(undef,2)
M[1]=rand(4,3)
M[2]=rand(4,2)
mode=[2,3]
Z=ttm(X,M,mode)
W=ttm(full(X),M,mode)
@test norm(W-full(Z)) ≈ 0 atol=1e-10
v=VectorCell(undef,2)
v[1]=rand(3)
v[2]=rand(2)
mode=[2,3]
Z=ttv(X,v,mode)
W=ttv(full(X),v,mode)
@test norm(W-full(Z)) ≈ 0 atol=1e-10

println("\n...Testing function tocell.")
M=tocell(X)
println("Is the output MatrixCell: ",isa(M,MatrixCell))
@test isa(M,MatrixCell)

mode=2;
println("\n...Testing function tenmat by mode $mode.")
Xn=tenmat(X,mode)
@test tenmat(full(X),mode) == Xn

println("\n...Testing function mttkrp.")
X=randktensor([5,4,3],3)
mode=1
M1=rand(5,5);
M2=rand(4,5);
M3=rand(3,5);
M=[M1,M2,M3]
println("Multiplying mode-$mode matricized tensor X by Khatri-Rao product of matrices.")
Z=mttkrp(X,M,mode)
err = norm(Z-tenmat(X,mode)*khatrirao(M3,M2))
println("Multiplication error: ",err)
@test err ≈ 0 atol=1e-10


## Packages
using Roots, Calculus, Plots

## Function for finding root
function findZero(g, init_guess)

    xn1 = init_guess;
    tol = 10^(-8);
    distance = 1;
    iter = 1;

    while distance >= tol
        xn = xn1;
        xn1 = xn - g(xn) / derivative(g, xn);
        xn1 = max(xn1, 0.0001);

        distance = abs(xn1 - xn);
        iter = iter + 1;
    end

    return xn1
end

## Function for defining utility function
function u(c)
    if c <= 0
        u = -Inf;
    else
        u = c^(1/2);
    end
    return u
end

## Parameters
β = 0.8;
πs1 = 1/2;
πs2 = 1/2;
Ys1 = 1;
Ys2 = 1/2;
M = 20; # Number of grids

## Grid for v
vmax = (πs1 * u(Ys1) + πs2 * u(Ys2)) / (1 - β);
vgrid = range(0.0, stop = vmax, length = M);

## Iterate on Value function
Pi1 = zeros(M);
cs1_grid = range(0, stop = Ys1, length = M);
ws1_grid = vgrid;
ws2_grid = vgrid;

distance = 1;
tol = 10^(-8);
iter = 1;

policy_fun = Array{CartesianIndex, 1}(undef, M);
policy_cs1 = Array{Int64, 1}(undef, M);

while distance >= tol
    println("This is iteration $iter")
    Pi = copy(Pi1);

    for vIndex in 1:M
        U = Array{Float64, 3}(undef, M, M, M)
        v = vgrid[vIndex];
        for cs1Index in 1:M, ws1Index in 1:M, ws2Index in 1:M
            cs1 = cs1_grid[cs1Index];
            ws1 = ws1_grid[ws1Index];
            ws2 = ws2_grid[ws2Index];

            # Solve for cs2
            f(cs2) = v - (πs1 * (cs1^(1/2) + β * ws1)) - (πs2 * (cs2^(1/2) + β * ws2));
            cs2 = findZero(f, 1.0);

            U_temp = (πs1 * (u(Ys1 - cs1) + β * Pi[ws1Index])) +
                     (πs2 * (u(Ys2 - cs2) + β * Pi[ws2Index]));

            if U_temp <= 0.0
                U[cs1Index, ws1Index, ws2Index] = 0.0;
            else
                U[cs1Index, ws1Index, ws2Index] = U_temp;
            end
        end
        Pi1[vIndex], policy_fun[vIndex] = findmax(U[:,:,:]);
    end

    distance = maximum(abs.(Pi1 - Pi));
    iter = iter + 1;
end

## A Check
λ = range(0.0, stop = 0.999, length = 100);
u1 = zeros(100);
u2 = zeros(100);
for i in 1:100
    A = (λ[i] / (1 - λ[i]))^2;
    cs1 = Ys1 * (A / (1 + A));
    cs2 = Ys2 * (A / (1 + A));
    u1[i] = (πs1 * u(cs1) + πs2 * u(cs2)) / (1 - β);
    u2[i] = (πs1 * u(Ys1 - cs1) + πs2 * u(Ys2 - cs2)) / (1 - β);
end

plot(vgrid, Pi1, xlabel = "v", ylabel = "P", label = "VFI", legend =:bottomleft)
plot!(u1, u2, label = "Theoretical")


using OrdinaryDiffEq

function LotkaVolterraTest_not_inplace(u,a,t)
    b,c,d = 1.0,3.0,1.0
    x,y = u[1],u[2]
    du = zeros(eltype(u),2)
    du[1] = a*x - b*x*y
    du[2] = -c*y + d*x*y
    du
end

# forward
u0 = [1.0;1.0]
tspan = (0.0,10.0)
prob = ODEProblem(LotkaVolterraTest_not_inplace,u0,tspan,1.5)

sol = solve(prob,Tsit5())
t = collect(range(0, stop=10, length=200))
using RecursiveArrayTools # for VectorOfArray
randomized = VectorOfArray([(sol(t[i]) + .01randn(2)) for i in 1:length(t)])
data = convert(Array,randomized)

# inverse
soll = solve(prob,Tsit5())

cost_function = build_loss_objective(prob,Tsit5(),L2Loss(t,data),
                                     maxiters=10000,verbose=false)
import Optim
result = Optim.optimize(cost_function, 0.0, 10.0)


#using SimpleGraphs, DrawSimpleGraphs, Plots

export KnightTourDrawing

"""
`KnightTourDrawing(r,c)` illustrates a Knight's tour of an r-by-c
chessboard returning `true` if the tour exists and `false` if not.
"""
function KnightTourDrawing(r::Int=6, c::Int=6)
    G = Knight(r,c)
    println("Searching for a Hamiltonian cycle in an $r-by-$c Knight's move graph")
    @time h = hamiltonian_cycle(G)
    println("Finished")

    if length(h)==0
        println("Sorry. This graph is not Hamiltonian")
        return false
    end

    T = vertex_type(G)
    H = SimpleGraph{T}()
    for v in G.V
        add!(H,v)
    end

    for k=1:NV(H)-1
        add!(H,h[k],h[k+1])
    end

    add!(H,h[1],h[end])

    xy = getxy(H)

    for v in H.V
        xy[v] = collect(v)
    end

    embed(H,xy)

    draw(H)

    for a=0:c
        plot!( [0.5, r+0.5], [a+0.5, a+0.5], color="black",
                linestyle=:dash)
    end

    for b=0:r
        plot!( [b+0.5, b+0.5], [0.5, c+0.5], color="black",
            linestyle=:dash)
    end

    finish()

end


using Plots
using BenchmarkTools

function vfsolve(vnew, kgrid, tolerance, imax)
    α = 0.3
    δ = 0.1
    β = 0.9

    v = vnew .+ 2*tolerance
    cartesianindex = Array{CartesianIndex{2}, length(v)}
    i = 1

    c = kgrid.^α + (1-δ)*kgrid .- kgrid'
    c[c .< 0] .= 0
    u = log.(c)
    u[kgrid.^α + (1-δ)*kgrid .- kgrid' .< 0] .= -Inf

    while maximum(abs.(v - vnew)) > tolerance && i <= imax
        v = vnew;

        (vnew, cartesianindex) = findmax(u .+ β*v', dims = 2);
        i += 1;
    end
    kprimeindex = getindex.(cartesianindex, 2)
    return (v = vnew, kprime = kgrid[kprimeindex], kprimeindex = kprimeindex)
end

kupper = 20
klower = 0.001
n = 1000
kgrid = collect(range(klower, stop = kupper, length = n))
(v, kprime, kprimeindex) = vfsolve(zeros(n), kgrid, 0.001, 1000);

scatter(kgrid, v, label = "v")
scatter(kgrid, kprime, label = "k'")

# Checks
vbasic == v
kprimebasic == kprime

@btime vfsolve(zeros(n), kgrid, 0.001, 1000);
@btime vfsolvebasic(zeros(n), kgrid, 0.001, 1000);

# Check k bounds, stepsize, tolerance, imax
findall(kprime .== maximum(kgrid))
any(kprime .== minimum(kgrid))

##
plot()
kupper = 2
klower = 0.001
tolerance = 0.001
for n in [10, 100, 500]
    kgrid = collect(range(klower, stop = kupper, length = n))
    (v, kprime, kprimeindex) = vfsolve(zeros(n), kgrid, tolerance, n);
     display(plot!(kgrid, v, label = "v, n=$n, tol=$tolerance"))
            # display(plot!(kgrid, kprime, label = "k'"))
end
##


# Starting with 1 and spiralling anticlockwise in the following way, a square
# spiral with side length 7 is formed.
#
# 37 36 35 34 33 32 31
# 38 17 16 15 14 13 30
# 39 18  5  4  3 12 29
# 40 19  6  1  2 11 28
# 41 20  7  8  9 10 27
# 42 21 22 23 24 25 26
# 43 44 45 46 47 48 49
#
# It is interesting to note that the odd squares lie along the bottom right
# diagonal, but what is more interesting is that 8 out of the 13 numbers lying
# along both diagonals are prime; that is, a ratio of 8/13 ≈ 62%.
#
# If one complete new layer is wrapped around the spiral above, a square spiral
# with side length 9 will be formed. If this process is continued, what is the
# side length of the square spiral for which the ratio of primes along both
# diagonals first falls below 10%?

using ProjectEulerSolutions

# Create the sequence of corner values and check each for primeness.  Keep a
# running total to verify that the desired threshold has been met.
function p058solution(thresh::Float64=0.44)::Integer
    nprimes = 0
    total_tested = 1
    startindex = 3
    i = 2
    while i == 2 || nprimes / total_tested > thresh
        nprimes += sum(isprime.(collect(startindex .+ (0:i:i*3))))
        startindex = startindex + 4*i + 2
        total_tested += 4
        i += 2
    end
    return i - 1
end

p058 = Problems.Problem(p058solution)

Problems.benchmark(p058, 0.1)


include("../helpers.jl")

using Gadfly
using Color

L(R) = x -> R * x * (1 - x)

k = 60
Rs = linspace(0.1, 4.0, 10) # [0.1, 0.533, 0.967, 1.4, ...]
xs = map(R -> iterate(L(R), 0.2, k), Rs)

fps = map(x -> fixedpoint(x, ϵ = 1e-3), xs) # [4, 8, 26, 10, ...]

# Values of R ∈ [3.0, 4.0] have very interesting properties.
# See Wikipedia or explore them yourself.
# Doing a Fourier Transform of the sequence will help illuminate these interesting
# properties


using Distributions, Calculus, SpecialFunctions, Plots; pyplot()

xGrid = -5:0.01:5

PhiA(x) = 0.5*(1+erf(x/sqrt(2)))
PhiB(x) = cdf(Normal(),x)

println("Maximum difference between two CDF implementations: ",
            maximum(PhiA.(xGrid) - PhiB.(xGrid)))

normalDensity(z) = pdf(Normal(),z)

d0 = normalDensity.(xGrid)
d1 = derivative.(normalDensity,xGrid)
d2 = second_derivative.(normalDensity, xGrid)

plot(xGrid, [d0 d1 d2], c=[:blue :red :green],label=[L"f(x)" L"f'(x)" L"f''(x)"])
plot!([-5,5],[0,0],  color=:black, lw=0.5, xlabel="x", xlims=(-5,5), label="")


# ---
# title: 1480. Running Sum of 1d Array
# id: problem1480
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/running-sum-of-1d-array/description/>
# hidden: true
# ---
#
# Given an array `nums`. We define a running sum of an array as `runningSum[i] =
# sum(nums[0]…nums[i])`.
#
# Return the running sum of `nums`.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,3,4]
#     Output: [1,3,6,10]
#     Explanation: Running sum is obtained as follows: [1, 1+2, 1+2+3, 1+2+3+4].
#
# **Example 2:**
#
#
#
#     Input: nums = [1,1,1,1,1]
#     Output: [1,2,3,4,5]
#     Explanation: Running sum is obtained as follows: [1, 1+1, 1+1+1, 1+1+1+1, 1+1+1+1+1].
#
# **Example 3:**
#
#
#
#     Input: nums = [3,1,2,10,1]
#     Output: [3,4,6,16,17]
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 1000`
#   * `-10^6 <= nums[i] <= 10^6`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using Sundials, Plots
pyplot(size=(700,700))

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  ## f2 routine. Compute function f(x,y).

  function f2(x, y, ydot)
      ydot[1] = y[2]
      ydot[2] = 2.0*y[1] - 3.0*y[2] + 3.0*x^2
  end

  x = 0.0:0.2:5
  res = Sundials.cvode(f2, [1.0, 0.0], collect(x))
  res |> display
  println()

  p1 = plot(x, res[:,1], color=:darkblue,
    xlabel="x", ylabel="y[1]",
    title="Example Sundials01A (using Sundials.jl)"
    )
  p2 = plot(x, res[:,2], color=:darkred,
    xlabel="x", ylabel="y[2]"
    )
  plot(p1, p2, layout=(2, 1))
  savefig("Sundials01A.png")
  gui()

end


# Conjugates for normal distribution

using Base.Test
using Distributions
using ConjugatePriors

import ConjugatePriors: NormalGamma, NormalInverseGamma
import ConjugatePriors: posterior, posterior_rand, posterior_mode, posterior_randmodel, fit_map

n = 100
w = rand(100)

# Νormal - Νormal (known sigma)

pri = Normal(1.0, 5.0)

x = rand(Normal(2.0, 3.0), n)
p = posterior((pri, 3.0), Normal, x)
@test isa(p, Normal)
@test_approx_eq mean(p)  (mean(pri) / var(pri) + sum(x) / 9.0) / (1.0 / var(pri) + n / 9.0)
@test_approx_eq var(p) inv(1.0 / var(pri) + n / 9.0)

r = posterior_mode((pri, 3.0), Normal, x)
@test_approx_eq r mode(p)

f = fit_map((pri, 3.0), Normal, x)
@test isa(f, Normal)
@test f.μ == r
@test f.σ == 3.0

p = posterior((pri, 3.0), Normal, x, w)
@test isa(p, Normal)
@test_approx_eq mean(p)  (mean(pri) / var(pri) + dot(x, w) / 9.0) / (1.0 / var(pri) + sum(w) / 9.0)
@test_approx_eq var(p) inv(1.0 / var(pri) + sum(w) / 9.0)

r = posterior_mode((pri, 3.0), Normal, x, w)
@test_approx_eq r mode(p)

f = fit_map((pri, 3.0), Normal, x, w)
@test isa(f, Normal)
@test f.μ == r
@test f.σ == 3.0

# ΙnverseGamma - Νormal (known mu)

pri = InverseGamma(1.5, 0.5)

x = rand(Normal(2.0, 3.0), n)
p = posterior((2.0, pri), Normal, x)
@test isa(p, InverseGamma)
@test_approx_eq shape(p) shape(pri) + n / 2
@test_approx_eq scale(p) scale(pri) + sum(abs2(x .- 2.0)) / 2

r = posterior_mode((2.0, pri), Normal, x)
@test_approx_eq r mode(p)

f = fit_map((2.0, pri), Normal, x)
@test isa(f, Normal)
@test f.μ == 2.0
@test_approx_eq abs2(f.σ) r

p = posterior((2.0, pri), Normal, x, w)
@test isa(p, InverseGamma)
@test_approx_eq shape(p) shape(pri) + sum(w) / 2
@test_approx_eq scale(p) scale(pri) + dot(w, abs2(x .- 2.0)) / 2

r = posterior_mode((2.0, pri), Normal, x, w)
@test_approx_eq r mode(p)

f = fit_map((2.0, pri), Normal, x, w)
@test isa(f, Normal)
@test f.μ == 2.0
@test_approx_eq abs2(f.σ) r

# Gamma - Νormal (known mu)

pri = Gamma(1.5, 2.0)

x = rand(Normal(2.0, 3.0), n)
p = posterior((2.0, pri), Normal, x)
@test isa(p, Gamma)
@test_approx_eq shape(p) shape(pri) + n / 2
@test_approx_eq scale(p) scale(pri) + sum(abs2(x .- 2.0)) / 2

r = posterior_mode((2.0, pri), Normal, x)
@test_approx_eq r mode(p)

f = fit_map((2.0, pri), Normal, x)
@test isa(f, Normal)
@test f.μ == 2.0
@test_approx_eq abs2(f.σ) inv(r)

p = posterior((2.0, pri), Normal, x, w)
@test isa(p, Gamma)
@test_approx_eq shape(p) shape(pri) + sum(w) / 2
@test_approx_eq scale(p) scale(pri) + dot(w, abs2(x .- 2.0)) / 2

r = posterior_mode((2.0, pri), Normal, x, w)
@test_approx_eq r mode(p)

f = fit_map((2.0, pri), Normal, x, w)
@test isa(f, Normal)
@test f.μ == 2.0
@test_approx_eq abs2(f.σ) inv(r)

# NormalInverseGamma - Normal

mu_true = 2.
sig2_true = 3.
x = rand(Normal(mu_true, sig2_true), n)

mu0 = 2.
v0 = 3.
shape0 = 5.
scale0 = 2.
pri = NormalInverseGamma(mu0, v0, shape0, scale0)

post = posterior(pri, Normal, x)
@test isa(post, NormalInverseGamma)

@test_approx_eq post.mu (mu0/v0 + n*mean(x))/(1./v0 + n)
@test_approx_eq post.v0 1./(1./v0 + n)
@test_approx_eq post.shape shape0 + 0.5*n
@test_approx_eq post.scale scale0 + 0.5*(n-1)*var(x) + n./v0./(n + 1./v0)*0.5*(mean(x)-mu0).^2

ps = posterior_randmodel(pri, Normal, x)

@test isa(ps, Normal)
@test insupport(ps,ps.μ) && ps.σ > zero(ps.σ)

# NormalGamma - Normal

mu_true = 2.
tau2_true = 3.
x = rand(Normal(mu_true, 1./tau2_true), n)

mu0 = 2.
nu0 = 3.
shape0 = 5.
rate0 = 2.
pri = NormalGamma(mu0, nu0, shape0, rate0)

post = posterior(pri, Normal, x)
@test isa(post, NormalGamma)

@test_approx_eq post.mu (nu0*mu0 + n*mean(x))./(nu0 + n)
@test_approx_eq post.nu nu0 + n
@test_approx_eq post.shape shape0 + 0.5*n
@test_approx_eq post.rate rate0 + 0.5*(n-1)*var(x) + n*nu0/(n + nu0)*0.5*(mean(x)-mu0).^2

ps = posterior_randmodel(pri, Normal, x)

@test isa(ps, Normal)
@test insupport(ps, ps.μ) && ps.σ > zero(ps.σ)


import StatisticalNonlocality: ou_transition_matrix, uniform_phase
import StatisticalNonlocality: discrete_laplacian

@testset "Transition Matrix Column Sum Zero" begin
    for n in 1:10
        M = ou_transition_matrix(n)
        @test all(abs.(sum(M, dims = 1)) .≤ eps(n / 2))
        if n > 1
            M = uniform_phase(n)
            @test all(abs.(sum(M, dims = 1)) .≤ eps(n / 2))
            Δ = discrete_laplacian(n)
            @test all(abs.(sum(Δ, dims = 1)) .≤ eps(n^2 / 2))
        end
    end
end

@testset "OU Transition Matrix Correctness Test for n in 1:2" begin
    n = 1
    M = ou_transition_matrix(n)
    Mh = [
        -1/2 1/2
        1/2 -1/2
    ]
    @test all(abs.(Mh - M) .≤ eps(n / 2))

    n = 2
    M = ou_transition_matrix(n)
    Mh = [
        -1 1/2 0
        1 -1 1
        0 1/2 -1
    ]
    @test all(abs.(Mh - M) .≤ eps(n / 2))
end

@testset "Uniform Phase Transition Matrix Correctness Test for n = 4" begin
    n = 4
    M = uniform_phase(n)
    Mh = [-1 1 0 0; 0 -1 1 0; 0 0 -1 1; 1 0 0 -1]
    @test all(abs.(Mh - M) .≤ eps(n / 2))

end

@testset "Eigenvalue Test for n in 1:10" begin
    for n in 1:10
        λ_exact = collect((-n):0)
        λ_approx = eigvals(ou_transition_matrix(n))
        @test all(abs.(λ_exact - λ_approx) .≤ eps(20.0 * n))
    end
end

@testset "Symmetry Test for Discrete Laplacian" begin
    for n in 4:10
        Δ = discrete_laplacian(n)
        @test norm(Δ - Δ') ≤ eps(norm(Δ))
    end
end


function expectation(g, distr::ContinuousUnivariateDistribution; kwargs...)
    return first(quadgk(x -> pdf(distr, x) * g(x), extrema(distr)...; kwargs...))
end

## Assuming that discrete distributions only take integer values.
function expectation(g, distr::DiscreteUnivariateDistribution; epsilon::Real=1e-10)
    mindist, maxdist = extrema(distr)
    # We want to avoid taking values up to infinity
    minval = isfinite(mindist) ? mindist : quantile(distr, epsilon)
    maxval = isfinite(maxdist) ? maxdist : quantile(distr, 1 - epsilon)
    return sum(x -> pdf(distr, x) * g(x), minval:maxval)
end

function expectation(g, distr::MultivariateDistribution; nsamples::Int=100, rng::AbstractRNG=GLOBAL_RNG)
    nsamples > 0 || throw(ArgumentError("number of samples should be > 0"))
    # We use a function barrier to work around type instability of `sampler(dist)`
    return mcexpectation(rng, g, sampler(distr), nsamples)
end

mcexpectation(rng, f, sampler, n) = sum(f, rand(rng, sampler) for _ in 1:n) / n

## Leave undefined until we've implemented a numerical integration procedure
# function entropy(distr::UnivariateDistribution)
#     pf = typeof(distr)<:ContinuousDistribution ? pdf : pmf
#     f = x -> pf(distr, x)
#     expectation(distr, x -> -log(f(x)))
# end

function kldivergence(P::Distribution{V}, Q::Distribution{V}; kwargs...) where {V<:VariateForm}
    return expectation(P; kwargs...) do x
        logp = logpdf(P, x)
        return (logp > oftype(logp, -Inf)) * (logp - logpdf(Q, x))
    end
end


using MCMC
using Base.Test

# Define immutable class State
immutable State
  mu::Float64
  sig2::Float64
end

@testset "Posterior Means" begin
  srand(1)
  using Distributions

  # Generate data
  const n = 1000
  const (mu,sig2) = (5,2)
  const y = rand(Normal(mu,sqrt(sig2)),n)

  # precomputes
  const ybar = mean(y)
  const (sig2_a, sig2_b) = (2.0, 1.0)
  rig(shape::Float64,rate::Float64) = 1 / rand(Gamma(shape,1/rate))

  # define an update function for State
  function update(state::State)
    const newMu = rand(Normal(ybar,sqrt(state.sig2/n)))
    const newSig2 = rig(sig2_a+n/2, sig2_b + sum((y-newMu).^2)/2)
    State(newMu,newSig2)
  end

  # run gibbs
  @time out = gibbs(State(0,1),update,10000,1000);

  ## post processing
  const postMu = map(o -> o.mu, out)
  const postSig2 = map(o -> o.sig2, out)
  @test abs(mean(postMu)-mu) < .05
  @test abs(mean(postSig2)-sig2) < .05
end

#= Personal tests:
# to test: vulture test.jl julia --color=yes test.jl

include("runtests.jl")
using RCall
R"library(rcommon) # install_github('luairthur/rcommon')"
plotPosts = R"plotPosts"
plotPosts([postMu postSig2])
=#


"""
    Berlage(; <keyword arguments>)

Create a Berlage wavelet.

# Arguments
- `dt=0.002`: sampling interval in secs.
- `f0=20.0`: central frequency in Hz.
- `m::Real=2`: exponential parameter of Berlage wavelet.
- `alpha::Real=180.0`: alpha parameter of Berlage wavelet in rad/secs.
- `phi0::Real`: phase rotation in radians.

# Example
```julia
julia> w = Berlage(); plot(w);
```
**Reference**
* Aldridge, David F., 1990, The berlage wavelet: GEOPHYSICS, 55, 1508--1511.
"""
function Berlage(; dt=0.002, f0=20.0, m::Real=2, alpha::Real=180.0,
                 phi0::Real=0.0)

    nw = floor(Int, 2.2/(f0*dt))
    t = dt*collect(0:1:nw-1)
    w = (t .^ m).* exp.(-alpha*t) .* cos.(2*pi*f0*t .+ phi0);
    w = w/maximum(w)

end


import Statistics
import Parameters.@with_kw

##### BinFixedWidth #####

"""
"""
@with_kw mutable struct BinFixedWidth
    n_bins::Integer = 10; @assert n_bins > 0
    data_min::Union{Float64,Missing} = missing
    data_span::Union{Float64,Missing} = missing
    bin_width::Union{Float64,Missing} = missing
    is_fit::Bool = false
end

"""
"""
function fit_transform!(bin::BinFixedWidth, data::AbstractArray{<:Real})
    bin.data_min = minimum(data)
    bin.data_span = maximum(data) - bin.data_min
    bin.bin_width = bin.data_span / bin.n_bins
    bin.is_fit = true
    bin
end

"""
"""
function apply_transform(bin::BinFixedWidth, data::AbstractArray{<:Real})
    if !bin.is_fit | bin.data_min === missing | \
        bin.data_span === missing | bin.bin_width === missing
        error("`bin` hasn't been fit.")
    end
    ceil.(Int,(data .- bin.data_min) ./ bin.bin_width)
end

##### BinQuantile #####

"""
"""
@with_kw mutable struct BinQuantile
    n_bins::Integer = 10; @assert n_bins > 0
    quantiles::Union{Array{Float64},Missing} = missing
    is_fit::Bool = false
end

"""
"""
function fit_transform!(bin::BinQuantile, data::AbstractArray{<:Real})
    bin.quantiles = Statistics.quantile(
        data,
        0:( 1 / (n_bins-1) ):1
    )
    bin.is_fit = true
    bin
end

"""
"""
function apply_transform(bin::BinQuantile, data::AbstractArray{<:Real})
    if !bin.is_fit | bin.quantiles === missing
        error("`bin` hasn't been fit.")
    end
    [findlast(q -> n < q, bin.quantiles) for n = data]
end


"""
    ks89RecursiveResidual(setting; indices, k)
Calculate recursive residual for the given regression setting and observation.
# Arguments
- `setting::RegressionSetting`: RegressionSetting object with a formula and dataset.
- `indices::ArrayInt,1`: Indices of observations used in the linear model.
- `k::Int`: Observation indice the recursive residual is calculated for.
# Notes
    This is a helper function for the ks89 function and it is not directly used.
# References
Kianifard, Farid, and William H. Swallow. "Using recursive residuals, calculated on
adaptively-ordered observations, to identify outliers in linear regression."
Biometrics (1989): 571-585.
"""
function ks89RecursiveResidual(setting::RegressionSetting, indices::Array{Int,1}, k::Int)
    X, y = @extractRegressionSetting setting
    return ks89RecursiveResidual(X, y, indices, k)
end

function ks89RecursiveResidual(X::Array{Float64,2}, y::Array{Float64,1}, indices::Array{Int,1}, k::Int)
    useX = X[indices, :]
    useY = y[indices]
    olsreg = ols(useX, useY)
    betas = coef(olsreg)
    XX = inv(useX'useX)
    w = (y[k] - X[k,:]' * betas) / sqrt(1 + X[k,:]' * XX * X[k,:])
    return w
end

"""
    ks89(setting; alpha = 0.05)
Perform the Kianifard & Swallow (1989) algorithm for the given regression setting.
# Arguments
- `setting::RegressionSetting`: RegressionSetting object with a formula and dataset.
- `alpha::Float64`: Optional argument of the probability of rejecting the null hypothesis.
# Examples
```julia-repl
julia> reg0001 = createRegressionSetting(@formula(stackloss ~ airflow + watertemp + acidcond), stackloss)
julia> ks89(reg0001)
2-element Array{Int64,1}:
  4
 21
```
# References
Kianifard, Farid, and William H. Swallow. "Using recursive residuals, calculated on
adaptively-ordered observations, to identify outliers in linear regression."
Biometrics (1989): 571-585.
"""
function ks89(setting::RegressionSetting; alpha=0.05)
    X = designMatrix(setting)
    y = responseVector(setting)
    return ks89(X, y, alpha=alpha)
end

function ks89(X::Array{Float64,2}, y::Array{Float64,1}; alpha=0.05)
    stdres = studentizedResiduals(X, y)
    orderingindices = sortperm(abs.(stdres))
    n, p = size(X)
    basisindices = orderingindices[1:p]
    w = zeros(Float64, n)
    s = zeros(Float64, n)
    ws = zeros(Float64, n)
    @inbounds for i in (p + 1):n
        index = orderingindices[i]
        w[index] = ks89RecursiveResidual(X, y, basisindices, index)
        s[index] = jacknifedS(X, y, index)
        ws[index] = w[index] / s[index]
        basisindices = orderingindices[1:i]
    end
    td = TDist(n - p - 1)
    q = quantile(td, alpha)
    result = filter(i -> abs.(ws[i]) > abs(q), 1:n)
    return result
end


using Random, LinearAlgebra, Krylov
using Plots
using RandomizedPreconditioners
const RP = RandomizedPreconditioners

n = 500
μ = 1e-4

# Build Spectrum (exp decay)
λs = max.(√n * exp.(-(0:n-1) ./ √n), μ)
plot(λs, lw=3, legend=false, title="Spectrum")

# Build Matrix
V = Array(qr(randn(n,n)).Q)
A = V*Diagonal(λs)*V'
A = 0.5(A + A')

# Build system
xtrue = randn(n)
b = A * xtrue

# Build preconditioner
r = 110
k = round(Int, r - 5)
Anys = NystromSketch(A, k, r)
P = RP.NystromPreconditionerInverse(Anys, μ)

# Solve system
_, stats_npc = cg(A, b; history=true)
_, stats_nys = cg(A, b; history=true, M = P)
_, stats_diag = cg(A, b; history=true, M=Diagonal(diag(A)))

nb = norm(b)
res_npc = prepend!(stats_npc.residuals, nb)
res_nys = prepend!(stats_nys.residuals, nb)
res_diag = prepend!(stats_diag.residuals, nb)

plt_cg = plot(
    res_npc,
    dpi=300,
    lw=2,
    label="No Preconditioner",
    ylabel="residual",
    xlabel="iteration",
    title="Convergence of CG",
    legend=:topright,
    yaxis=:log
)
plot!(plt_cg,
    res_nys,
    label="Nystrom Preconditioner",
    lw=2
)
plot!(plt_cg,
    res_diag,
    label="Diagonal Preconditioner",
    lw=2
)
display(plt_cg)


"""
    ot_plan(c, μ, ν; kwargs...)

Compute the optimal transport plan for the Monge-Kantorovich problem with source and target
marginals `μ` and `ν` and cost `c`.

The optimal transport plan solves
```math
\\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\int c(x, y) \\, \\mathrm{d}\\gamma(x, y)
```
where ``\\Pi(\\mu, \\nu)`` denotes the couplings of ``\\mu`` and ``\\nu``.

See also: [`ot_cost`](@ref)
"""
function ot_plan end

"""
    ot_cost(c, μ, ν; kwargs...)

Compute the optimal transport cost for the Monge-Kantorovich problem with source and target
marginals `μ` and `ν` and cost `c`.

The optimal transport cost is the scalar value
```math
\\inf_{\\gamma \\in \\Pi(\\mu, \\nu)} \\int c(x, y) \\, \\mathrm{d}\\gamma(x, y)
```
where ``\\Pi(\\mu, \\nu)`` denotes the couplings of ``\\mu`` and ``\\nu``.

See also: [`ot_plan`](@ref)
"""
function ot_cost end

#############
# Discrete OT
#############

"""
    emd(μ, ν, C, optimizer)

Compute the optimal transport plan `γ` for the Monge-Kantorovich problem with source
histogram `μ`, target histogram `ν`, and cost matrix `C` of size `(length(μ), length(ν))`
which solves
```math
\\inf_{γ ∈ Π(μ, ν)} \\langle γ, C \\rangle.
```

The corresponding linear programming problem is solved with the user-provided `optimizer`.
Possible choices are `Tulip.Optimizer()` and `Clp.Optimizer()` in the `Tulip` and `Clp`
packages, respectively.
"""
function emd(μ, ν, C, model::MOI.ModelLike)
    # check size of cost matrix
    nμ = length(μ)
    nν = length(ν)
    size(C) == (nμ, nν) || error("cost matrix `C` must be of size `(length(μ), length(ν))`")
    nC = length(C)

    # define variables
    x = MOI.add_variables(model, nC)
    xmat = reshape(x, nμ, nν)

    # define objective function
    T = float(eltype(C))
    zero_T = zero(T)
    MOI.set(
        model,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{T}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(float.(vec(C)), x), zero_T),
    )
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    # add non-negativity constraints
    for xi in x
        MOI.add_constraint(model, MOI.SingleVariable(xi), MOI.GreaterThan(zero_T))
    end

    # add constraints for source
    for (i, μi) in zip(axes(xmat, 1), μ) # eachrow(xmat) is not available on Julia 1.0
        f = MOI.ScalarAffineFunction(
            [MOI.ScalarAffineTerm(one(μi), xi) for xi in view(xmat, i, :)], zero(μi)
        )
        MOI.add_constraint(model, f, MOI.EqualTo(μi))
    end

    # add constraints for target
    for (i, νi) in zip(axes(xmat, 2), ν) # eachcol(xmat) is not available on Julia 1.0
        f = MOI.ScalarAffineFunction(
            [MOI.ScalarAffineTerm(one(νi), xi) for xi in view(xmat, :, i)], zero(νi)
        )
        MOI.add_constraint(model, f, MOI.EqualTo(νi))
    end

    # compute optimal solution
    MOI.optimize!(model)
    status = MOI.get(model, MOI.TerminationStatus())
    status === MOI.OPTIMAL || error("failed to compute optimal transport plan: ", status)
    p = MOI.get(model, MOI.VariablePrimal(), x)
    γ = reshape(p, nμ, nν)

    return γ
end

"""
    emd2(μ, ν, C, optimizer; plan=nothing)

Compute the optimal transport cost (a scalar) for the Monge-Kantorovich problem with source
histogram `μ`, target histogram `ν`, and cost matrix `C` of size `(length(μ), length(ν))`
which is given by
```math
\\inf_{γ ∈ Π(μ, ν)} \\langle γ, C \\rangle.
```

The corresponding linear programming problem is solved with the user-provided `optimizer`.
Possible choices are `Tulip.Optimizer()` and `Clp.Optimizer()` in the `Tulip` and `Clp`
packages, respectively.

A pre-computed optimal transport `plan` may be provided.
"""
function emd2(μ, ν, C, optimizer; plan=nothing)
    γ = if plan === nothing
        # compute optimal transport plan
        emd(μ, ν, C, optimizer)
    else
        # check dimensions
        size(C) == (length(μ), length(ν)) ||
            error("cost matrix `C` must be of size `(length(μ), length(ν))`")
        size(plan) == size(C) || error(
            "optimal transport plan `plan` and cost matrix `C` must be of the same size",
        )
        plan
    end
    return dot(γ, C)
end

###################################
# Semidiscrete and continuous 1D OT
###################################

"""
    ot_plan(c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution)

Compute the optimal transport plan for the Monge-Kantorovich problem with univariate
distributions `μ` and `ν` as source and target marginals and cost function `c` of
the form ``c(x, y) = h(|x - y|)`` where ``h`` is a convex function.

In this setting, the optimal transport plan is the Monge map
```math
T = F_\\nu^{-1} \\circ F_\\mu
```
where ``F_\\mu`` is the cumulative distribution function of `μ` and ``F_\\nu^{-1}`` is the
quantile function of `ν`.

See also: [`ot_cost`](@ref), [`emd`](@ref)
"""
function ot_plan(c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution)
    # Use T instead of γ to indicate that this is a Monge map.
    T(x) = quantile(ν, cdf(μ, x))
    return T
end

"""
    ot_cost(
        c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution; plan=nothing
    )

Compute the optimal transport cost for the Monge-Kantorovich problem with univariate
distributions `μ` and `ν` as source and target marginals and cost function `c` of
the form ``c(x, y) = h(|x - y|)`` where ``h`` is a convex function.

In this setting, the optimal transport cost can be computed as
```math
\\int_0^1 c(F_\\mu^{-1}(x), F_\\nu^{-1}(x)) \\mathrm{d}x
```
where ``F_\\mu^{-1}`` and ``F_\\nu^{-1}`` are the quantile functions of `μ` and `ν`,
respectively.

A pre-computed optimal transport `plan` may be provided.

See also: [`ot_plan`](@ref), [`emd2`](@ref)
"""
function ot_cost(
    c, μ::ContinuousUnivariateDistribution, ν::UnivariateDistribution; plan=nothing
)
    cost, _ = if plan === nothing
        quadgk(0, 1) do q
            return c(quantile(μ, q), quantile(ν, q))
        end
    else
        quadgk(0, 1) do q
            x = quantile(μ, q)
            return c(x, plan(x))
        end
    end
    return cost
end

################
# Discrete 1D OT
################

# internal iterator for discrete one-dimensional OT problems
# it returns tuples that consist of the indices of the source and target histograms
# and the optimal flow between the corresponding points
struct Discrete1DOTIterator{T,M,N}
    mu::M
    nu::N
end

# histograms `μ` and `ν` are expected to be iterators of the histograms where the
# corresponding support is sorted
function Discrete1DOTIterator(μ, ν)
    T = Base.promote_eltype(μ, ν)
    return Discrete1DOTIterator{T,typeof(μ),typeof(ν)}(μ, ν)
end

Base.IteratorEltype(::Type{<:Discrete1DOTIterator}) = Base.HasEltype()
Base.eltype(::Type{<:Discrete1DOTIterator{T}}) where {T} = Tuple{Int,Int,T}

Base.length(d::Discrete1DOTIterator) = length(d.mu) + length(d.nu) - 1

# we iterate through the source and target histograms
function Base.iterate(
    d::Discrete1DOTIterator{T}, (i, j, μnext, νnext)=(1, 1, iterate(d.mu), iterate(d.nu))
) where {T}
    # if we are done with iterating through the source and/or target histogram,
    # iteration is stopped
    if μnext === nothing || νnext === nothing
        return nothing
    end

    # unpack next values and states of the source and target histograms
    μiter, μstate = μnext
    νiter, νstate = νnext

    # compute next value of the iterator: indices of source and target histograms
    # and optimal flow between the corresponding points
    min_iter, max_iter = minmax(μiter, νiter)
    iter = (i, j, min_iter)

    # compute next state of the iterator
    diff = max_iter - min_iter
    state = if μiter < max_iter
        # move forward in the source histogram
        (i + 1, j, iterate(d.mu, μstate), (diff, νstate))
    else
        # move forward in the target histogram
        (i, j + 1, (diff, μstate), iterate(d.nu, νstate))
    end

    return iter, state
end

"""
    ot_plan(c, μ::DiscreteNonParametric, ν::DiscreteNonParametric)

Compute the optimal transport cost for the Monge-Kantorovich problem with univariate
discrete distributions `μ` and `ν` as source and target marginals and cost function `c`
of the form ``c(x, y) = h(|x - y|)`` where ``h`` is a convex function.

In this setting, the optimal transport plan can be computed analytically. It is returned as
a sparse matrix.

See also: [`ot_cost`](@ref), [`emd`](@ref)
"""
function ot_plan(_, μ::DiscreteNonParametric, ν::DiscreteNonParametric)
    # unpack the probabilities of the two distributions
    μprobs = probs(μ)
    νprobs = probs(ν)

    # create the iterator
    # note: support of `DiscreteNonParametric` is sorted
    iter = Discrete1DOTIterator(μprobs, νprobs)

    # create arrays for the indices of the two histograms and the optimal flow between the
    # corresponding points
    n = length(iter)
    I = Vector{Int}(undef, n)
    J = Vector{Int}(undef, n)
    W = Vector{Base.promote_eltype(μprobs, νprobs)}(undef, n)

    # compute the sparse optimal transport plan
    @inbounds for (idx, (i, j, w)) in enumerate(iter)
        I[idx] = i
        J[idx] = j
        W[idx] = w
    end
    γ = sparse(I, J, W, length(μprobs), length(νprobs))

    return γ
end

"""
    ot_cost(
        c, μ::DiscreteNonParametric, ν::DiscreteNonParametric; plan=nothing
    )

Compute the optimal transport cost for the Monge-Kantorovich problem with discrete
univariate distributions `μ` and `ν` as source and target marginals and cost function `c`
of the form ``c(x, y) = h(|x - y|)`` where ``h`` is a convex function.

In this setting, the optimal transport cost can be computed analytically.

A pre-computed optimal transport `plan` may be provided.

See also: [`ot_plan`](@ref), [`emd2`](@ref)
"""
function ot_cost(c, μ::DiscreteNonParametric, ν::DiscreteNonParametric; plan=nothing)
    return _ot_cost(c, μ, ν, plan)
end

# compute cost from scratch if no plan is provided
function _ot_cost(c, μ::DiscreteNonParametric, ν::DiscreteNonParametric, ::Nothing)
    # unpack the probabilities of the two distributions
    μprobs = probs(μ)
    νprobs = probs(ν)

    # create the iterator
    # note: support of `DiscreteNonParametric` is sorted
    iter = Discrete1DOTIterator(μprobs, νprobs)

    # compute the cost
    μsupport = support(μ)
    νsupport = support(ν)
    cost = sum(w * c(μsupport[i], νsupport[j]) for (i, j, w) in iter)

    return cost
end

# if a sparse plan is provided, we just iterate through the non-zero entries
function _ot_cost(
    c, μ::DiscreteNonParametric, ν::DiscreteNonParametric, plan::SparseMatrixCSC
)
    # extract non-zero flows
    I, J, W = findnz(plan)

    # compute the cost
    μsupport = support(μ)
    νsupport = support(ν)
    cost = sum(w * c(μsupport[i], νsupport[j]) for (i, j, w) in zip(I, J, W))

    return cost
end

# fallback: compute cost matrix (probably often faster to compute cost from scratch)
function _ot_cost(c, μ::DiscreteNonParametric, ν::DiscreteNonParametric, plan)
    return dot(plan, StatsBase.pairwise(c, support(μ), support(ν)))
end

################
# OT Gaussians
################

"""
    ot_cost(::SqEuclidean, μ::MvNormal, ν::MvNormal)

Compute the squared 2-Wasserstein distance between normal distributions `μ` and `ν` as
source and target marginals.

In this setting, the optimal transport cost can be computed as
```math
W_2^2(\\mu, \\nu) = \\|m_\\mu - m_\\nu \\|^2 + \\mathcal{B}(\\Sigma_\\mu, \\Sigma_\\nu)^2,
```
where ``\\mu = \\mathcal{N}(m_\\mu, \\Sigma_\\mu)``,
``\\nu = \\mathcal{N}(m_\\nu, \\Sigma_\\nu)``, and ``\\mathcal{B}`` is the Bures metric.

See also: [`ot_plan`](@ref), [`emd2`](@ref)
"""
function ot_cost(::SqEuclidean, μ::MvNormal, ν::MvNormal)
    return sqeuclidean(μ.μ, ν.μ) + sqbures(μ.Σ, ν.Σ)
end

"""
    ot_cost(::SqEuclidean, μ::Normal, ν::Normal)

Compute the squared 2-Wasserstein distance between univariate normal distributions `μ` and
`ν` as source and target marginals.

See also: [`ot_plan`](@ref), [`emd2`](@ref)
"""
function ot_cost(::SqEuclidean, μ::Normal, ν::Normal)
    return (μ.μ - ν.μ)^2 + (μ.σ - ν.σ)^2
end

"""
    ot_plan(::SqEuclidean, μ::MvNormal, ν::MvNormal)

Compute the optimal transport plan for the Monge-Kantorovich problem with multivariate
normal distributions `μ` and `ν` as source and target marginals and cost function
``c(x, y) = \\|x - y\\|_2^2``.

In this setting, for ``\\mu = \\mathcal{N}(m_\\mu, \\Sigma_\\mu)`` and
``\\nu = \\mathcal{N}(m_\\nu, \\Sigma_\\nu)``, the optimal transport plan is the Monge
map
```math
T \\colon x \\mapsto m_\\nu
+ \\Sigma_\\mu^{-1/2}
{\\big(\\Sigma_\\mu^{1/2} \\Sigma_\\nu \\Sigma_\\mu^{1/2}\\big)}^{1/2}\\Sigma_\\mu^{-1/2}
(x - m_\\mu).
```

See also: [`ot_cost`](@ref), [`emd`](@ref)
"""
function ot_plan(::SqEuclidean, μ::MvNormal, ν::MvNormal)
    Σμsqrt = μ.Σ^(-1 / 2)
    A = Σμsqrt * sqrt(_gaussian_ot_A(μ.Σ, ν.Σ)) * Σμsqrt
    mμ = μ.μ
    mν = ν.μ
    T(x) = mν + A * (x - mμ)
    return T
end

"""
    ot_plan(::SqEuclidean, μ::Normal, ν::Normal)

Compute the optimal transport plan for the Monge-Kantorovich problem with
normal distributions `μ` and `ν` as source and target marginals and cost function
``c(x, y) = \\|x - y\\|_2^2``.

See also: [`ot_cost`](@ref), [`emd`](@ref)
"""
function ot_plan(::SqEuclidean, μ::Normal, ν::Normal)
    mμ = μ.μ
    mν = ν.μ
    a = ν.σ / μ.σ
    T(x) = mν + a * (x - mμ)
    return T
end


# Simple ridge regression based upon the sindy-mpc
# repository, see https://arxiv.org/abs/1711.05501
# and https://github.com/eurika-kaiser/SINDY-MPC/blob/master/LICENSE

"""
$(TYPEDEF)
`STLQS` is taken from the [original paper on SINDY](https://www.pnas.org/content/113/15/3932) and implements a
sequentially thresholded least squares iteration. `λ` is the threshold of the iteration.
It is based upon [this matlab implementation](https://github.com/eurika-kaiser/SINDY-MPC/utils/sparsifyDynamics.m).
It solves the following problem
```math
\\argmin_{x} \\frac{1}{2} \\| Ax-b\\|_2 + \\lambda \\|x\\|_2
```

# Fields
$(FIELDS)

# Example
```julia
opt = STLQS()
opt = STLQS(1e-1)
opt = STLQS(Float32[1e-2; 1e-1])
```
## Note
This was formally `STRRidge` and has been renamed.
"""
mutable struct STLSQ{T} <: AbstractOptimizer{T}
    """Sparsity threshold"""
    λ::T

    function STLSQ(threshold::T = 1e-1) where T
        @assert all(threshold .> zero(eltype(threshold))) "Threshold must be positive definite"

        return new{typeof(threshold)}(threshold)
    end

end

Base.summary(::STLSQ) = "STLSQ"

function (opt::STLSQ{T})(X, A, Y, λ::U = first(opt.λ);
    maxiter = maximum(size(A)), abstol::U = eps(eltype(T)),
    progress = nothing, kwargs...) where {T,U}

    smallinds = abs.(X) .<= λ
    biginds = @. ! smallinds[:, 1]

    x_i = similar(X)
    x_i .= X

    xzero = zero(eltype(X))
    obj = xzero
    sparsity = xzero
    conv_measure = xzero

    iters = 0
    converged = false

    _progress = isa(progress, Progress)
    initial_prog = _progress ? progress.counter : 0

    @views while (iters < maxiter) && !converged
        iters += 1

        smallinds .= abs.(X) .<= λ
        X[smallinds] .= xzero

        for j in 1:size(Y, 2)
            @. biginds = ! smallinds[:, j]
            X[biginds, j] .= A[:, biginds] \ Y[:,j]
        end

        conv_measure = norm(x_i .- X, 2)

        if _progress
            obj = norm(Y - A*X, 2)
            sparsity = norm(X, 0, λ)

            ProgressMeter.next!(
            progress;
            showvalues = [
                (:Threshold, λ), (:Objective, obj), (:Sparsity, sparsity),
                (:Convergence, conv_measure)
            ]
            )
        end

        if conv_measure < abstol
            converged = true

            _progress ? (progress.counter = initial_prog + maxiter) : nothing

            #    ProgressMeter.update!(
            #    progress,
            #    initial_prog + maxiter
            #    )
            #end

        else
            x_i .= X
        end
    end

    clip_by_threshold!(X, λ; kwargs...)
    return
end


using SweepOperator, Test

# setup
n, p = 1000, 10
x = randn(n, p)
xtx = x'x

@testset "Sweep One By One" begin
    A = deepcopy(xtx)
    B = deepcopy(xtx)
    for j in 1:p
        sweep!(A, j)
        sweep!(A, j, true)
    end
    @test A ≈ B

    A = deepcopy(xtx)
    B = deepcopy(xtx)
    for j in 1:p
        sweep!(A, j, true)
    end
    for j in 1:p
        sweep!(A, j, true)
    end
    @test A ≈ B
end

@testset "Sweep All" begin
    A = deepcopy(xtx)
    B = deepcopy(xtx)
    sweep!(A, 1:p)
    sweep!(A, 1:p, true)
    @test A ≈ B
end

@testset "Linear Regression" begin
    y = x * collect(1.:p) + randn(n)
    xy = [x y]
    xytxy = xy'xy
    sweep!(xytxy, 1:p)
    @test xytxy[1:p, end] ≈ x\y
end


module TestGrad
    using Calculus2
    using Base.Test

    # multivariate test function
    function f(x::Vector)
        return sin(x[1]) + cos(x[2])
    end
    x = [0.0, 0.0]
    gr = similar(x)

    # Test in-place multivariate grad!
    grad!(gr, f, x)
    @test norm(gr - [cos(0.0), -sin(0.0)]) < 10e-8
    fill!(gr, 0.0)

    grad!(gr, f, x, method = :ad)
    @test norm(gr - [cos(0.0), -sin(0.0)]) < 10e-8
    fill!(gr, 0.0)

    ∇!(gr, f, x)
    @test norm(gr - [cos(0.0), -sin(0.0)]) < 10e-8
    fill!(gr, 0.0)

    # Test pure multivariate grad
    @test norm(grad(f, x) - [cos(0.0), sin(0.0)]) < 10e-8
    @test norm(∇(f, x) - [cos(0.0), sin(0.0)]) < 10e-8

    # Test in-place multivariate gradfunc!
    ∇f! = gradfunc!(f, x)
    ∇f!(gr, x)
    @test norm(gr - [cos(0.0), sin(0.0)]) < 10e-8

    # Test in-place multivariate gradfunc!
    ∇f! = gradfunc!(f, x, method = :ad)
    ∇f!(gr, x)
    @test norm(gr - [cos(0.0), sin(0.0)]) < 10e-8

    # Test pure multivariate gradfunc
    ∇f = gradfunc(f, x)
    @test norm(∇f(x) - [cos(0.0), sin(0.0)]) < 10e-8

    # Test pure univariate grad
    @test isapprox(grad(sin, 0.0), cos(0.0))
    @test isapprox(grad(sin, 0.0, method = :ad), cos(0.0))

    # Test pure univariate gradfunc
    @test isapprox(gradfunc(sin, 0.0)(0.0), cos(0.0))
    @test isapprox(gradfunc(sin, 0.0, method = :ad)(0.0), cos(0.0))

    # TODO: Test automatic differentation
    # TODO: Test symbolic differentiation
    # TODO: Test direction options
end


using LinearAlgebra: diagm
using SparseArrays: sparse

normalize_indicator(i::AbstractArray) = i ./ sum(i)
normalize_indicator(i::AbstractArray, dims) = i ./ sum(i, dims=dims)

to_indicator_matrix(i::AbstractVector) = diagm(i)

function union_diagonal!(A::AbstractMatrix{T}) where {T}
    for i in 1:min(size(A)...)
        A[i, i] = one(T)
    end
    return A
end

_generate_slices(n, k) = Base.ntuple(i->(i==n) ? k : Colon(), 2)

"""
    graph_filter(ind, order)

Generate a graph filter of `order`-th order moment depending on the indicator `ind`.

## Arguments
- `ind`: the indicator vector or matrix with values of 0s or 1s.
- `order`: the order of moment for a graph filter. If `order=1`, generates a first moment graph filter,
which is the mean operator. If `order=2`, generates a second moment graph filter, which is the covariance operator.
"""
function graph_filter(ind::AbstractVector, order::Integer)
    if order == 1
        return sparse(normalize_indicator(ind))
    elseif order == 2
        adj = to_indicator_matrix(ind)
        return sparse(normalize_indicator(adj))
    else
        throw(ArgumentError("order other than 1 and 2 is not supported while get $order."))
    end
end

function graph_filter(adj::AbstractMatrix, order::Integer; dims::Integer=1, self::Bool=true)
    if order == 1
        self && union_diagonal!(adj)
        return sparse(normalize_indicator(adj, dims))
    elseif order == 2
        n = dims%2 + 1
        blocks = [graph_filter(vec(adj[_generate_slices(n, k)...]), 2) for k in 1:size(adj, n)]
        return cat(blocks...,dims=(1,2))
    else
        throw(ArgumentError("order other than 1 and 2 is not supported while get $order."))
    end
end

function moment(A::AbstractMatrix, X::AbstractMatrix, order::Integer; dims::Integer=1)
    if order == 1
        return first_moment(A, X; dims=dims)
    elseif order == 2
        return second_moment(A, X; dims=dims)
    else
        throw(ArgumentError("order other than 1 and 2 is not supported while get $order."))
    end
end

function first_moment(A::AbstractMatrix, X::AbstractMatrix; dims::Integer=1)
    adj = A .> 0
    if dims == 1
        G = graph_filter(adj, 1; dims=1)
        return G*X
    elseif dims == 2
        G = graph_filter(adj, 1; dims=2)
        return X*G
    else
        throw(ArgumentError("dims other than 1 and 2 is not supported while get $dims."))
    end
end

# second_moment(A::AbstractMatrix, X::AbstractMatrix; dims::Integer=1) = second_moment(A, X, X; dims=dims)

# function second_moment(A::AbstractMatrix, X::AbstractMatrix, Y::AbstractMatrix; dims::Integer=1)
#     # x' * G_2 * y - (G_1' * x)*(G_1' * y)
#     adj = A .> 0
#     if dims == 1
#         G = graph_filter(adj, 2; dims=1)
#         return X'*G*Y - first_moment(A, X; dims=1)*first_moment(A, Y; dims=1)'
#     elseif dims == 2
#         G = graph_filter(adj, 2; dims=2)
#         return vec(X')'*G*vec(Y') - first_moment(A, X; dims=2)*first_moment(A, Y; dims=2)'
#     else
#         throw(ArgumentError("dims other than 1 and 2 is not supported while get $dims."))
#     end
# end


#Test library imports
using Test

#Arithmetic library imports
using IntervalArithmetic

#Preamble
setprecision(53)
setprecision(Interval, Float64)
setrounding(Interval, :tight)
# Set full format, and show decorations
@format full
@testset "rational_power_test" begin
        @test ^(∅, 1//3) == ∅
        @test ^(1 .. 8, 1//3) == Interval(1, 2)
        @test ^(2 .. 8, 1//3) ⊇  Interval(2^(1//3), 2)
        @test ^(1 .. 9, 1//3) ⊇  Interval(1, 9^(1//3))
        @test ^(2 .. 9, 1//3) ⊇  Interval(2^(1//3), 9^(1//3))
        @test ^(-1 .. 8, 1//3) == Interval(0, 2)
        @test ^(-2 .. 8, 1//3) ⊇  Interval(0, 2)
        @test ^(-1 .. 9, 1//3) ⊇  Interval(0, 9^(1//3))
        @test ^(-2 .. 9, 1//3) ⊇  Interval(0, 9^(1//3))
        @test ^(1 .. 8, -1//3) == Interval(0.5, 1)
        @test ^(2 .. 8, -1//3) ⊇  Interval(0.5, 2^(-1//3))
        @test ^(1 .. 9, -1//3) ⊇  Interval(9^(-1//3), 1)
        @test ^(2 .. 9, -1//3) ⊇  Interval(9^(-1//3), 2^(-1//3))
        @test ^(-1 .. 8, -1//3) == Interval(0.5, Inf)
        @test ^(-2 .. 8, -1//3) ⊇  Interval(0.5, Inf)
        @test ^(-1 .. 9, -1//3) ⊇  Interval(9^(-1//3), Inf)
        @test ^(-2 .. 9, -1//3) ⊇  Interval(9^(-1//3), Inf)
        @test ^(-2 .. 4 , 1//2) == Interval(0, 2)
        @test ^(-2 .. 8 , 1//3) == Interval(0, 2)
        @test ^(-8 .. -2 , 1//3) == ∅
        @test ^(-8 .. -2 , 1//2) == ∅
        @test ^(-8 .. -2 , -1//3) == ∅
        @test ^(-8 .. -2 , -1//2) == ∅
        @test ^(∅, 2//3) == ∅
        @test ^(1 .. 8, 2//3) == Interval(1, 4)
        @test ^(2 .. 8, 2//3) ⊇  Interval(2^(2//3), 4)
        @test ^(1 .. 9, 2//3) ⊇  Interval(1, 9^(2//3))
        @test ^(2 .. 9, 2//3) ⊇  Interval(2^(2//3), 9^(2//3))
        @test ^(-1 .. 8, 2//3) == Interval(0, 4)
        @test ^(-2 .. 8, 2//3) ⊇  Interval(0, 4)
        @test ^(-1 .. 9, 2//3) ⊇  Interval(0, 9^(2//3))
        @test ^(-2 .. 9, 2//3) ⊇  Interval(0, 9^(2//3))
        @test ^(1 .. 8, -2//3) == Interval(0.25, 1)
        @test ^(2 .. 8, -2//3) ⊇  Interval(0.25, 2^(-2//3))
        @test ^(1 .. 9, -2//3) ⊇  Interval(9^(-2//3), 1)
        @test ^(2 .. 9, -2//3) ⊇  Interval(9^(-2//3), 2^(-2//3))
        @test ^(-1 .. 8, -2//3) == Interval(0.25, Inf)
        @test ^(-2 .. 8, -2//3) ⊇  Interval(0.25, Inf)
        @test ^(-1 .. 9, -2//3) ⊇  Interval(9^(-2//3), Inf)
        @test ^(-2 .. 9, -2//3) ⊇  Interval(9^(-2//3), Inf)
        @test ^(-2 .. 4 , 3//2) == Interval(0, 8)
        @test ^(-2 .. 8 , 2//3) == Interval(0, 4)
        @test ^(-8 .. -2 , 2//3) == ∅
        @test ^(-8 .. -2 , 3//2) == ∅
        @test ^(-8 .. -2 , -2//3) == ∅
        @test ^(-8 .. -2 , -3//2) == ∅
end


module ComplexBigMatrices

"""
This module calculates exp(), log() and eigen() of a complex BigFloat Matrix.
"""

export exp, log, eigen
import .Base: exp, log

using LinearAlgebra
using GenericLinearAlgebra
using GenericSchur

"""
    exp(a::Matrix{Complex{BigFloat}})
Extends exp() to work with Matrix{Complex{BigFloat}}
"""
function exp(a::Matrix{Complex{BigFloat}})
    eig = eigen(a)
    eig.vectors*Diagonal(exp.(eig.values))*eig.vectors^(-1)
end

"""
    log(a::Matrix{Complex{BigFloat}})
Extends log() to work with Matrix{Complex{BigFloat}}
"""
function log(a::Matrix{Complex{BigFloat}})
    eig = eigen(a)
    eig.vectors*Diagonal(log.(eig.values))*eig.vectors^(-1)
end

"""
    eigen(a::Matrix{Complex{BigFloat}})
Extends eigen() to work with Matrix{Complex{BigFloat}}
"""
function eigen(a::Matrix{Complex{BigFloat}})
    LinearAlgebra.eigen(a)
end

end


using JuMP, Mosek, MosekTools, LinearAlgebra, Suppressor, StatsBase, Ipopt, Gurobi

function solveNNMFRelax(A, theK)
    mod=Model(Mosek.Optimizer)
    n=size(A,1)

    @variable(mod, X[1:n, 1:n], PSD)
    @variable(mod, Y[1:n, 1:n], Symmetric)
    @variable(mod, theta[1:n, 1:n], Symmetric)
    @constraint(mod, X.>=0.0)

    @constraint(mod, Symmetric(Matrix(1.0I, n, n).-Y) in PSDCone());

    @constraint(mod,Symmetric([theta X; X' Y]) in PSDCone());

    @constraint(mod, sum(Y[i,i] for i=1:n)<=theK)

    @objective(mod, Min, sum(theta[i,i] for i=1:n)-2.0*Compat.dot(A, X)+Compat.dot(A,A))

    optimize!(mod)

    @show ofv_dual=objective_value(mod)
    Y_0=value.(Y)
    # solve for U
    mod_U=Model(Mosek.Optimizer)
    @variable(mod_U, U[1:n,1:theK]>=0.0)
    @variable(mod_U, t)
    @constraint(mod_U, [t; 1.0; vec(U-Y_0*U)] in RotatedSecondOrderCone());
    @constraint(mod_U, Y_0*U.>=0.0)
    @suppress optimize!(mod_U)
    U_0=value.(U)

    # Solve for a SDD^+ D, where A~=U_0*D*U_0'. Note that 2x2 non-neg psd matrices are completely positive, so this is a valid inner approx.
    mod2=Model(Mosek.Optimizer)
    @variable(mod2, D[1:theK, 1:theK], Symmetric)
    @variable(mod2, M[1:theK, 1:theK, 1:2, 1:2]>=0.0);
    @constraint(mod2, defineM1[i=1:theK, j=1:theK], M[i,j,1,2]==M[i,j,2,1]);
    @constraint(mod2, defineM2[i=1:theK, j=1:theK], [M[i,j,1,1]; M[i,j,2,2]; M[i,j, 1,2]] in RotatedSecondOrderCone());
    @constraint(mod2, defineD[i=1:theK, j=(i+1):theK], D[i,j]==M[i,j,1,2])
    @constraint(mod2, defineD2[i=1:theK], D[i,i]==sum(M[i,l,2,2] for l=1:(i-1))+sum(M[i,l,1,1] for l=(i+1):theK))

    @variable(mod2, t)
    @constraint(mod2, [t; 1.0; vec(A-U_0*D*U_0')] in RotatedSecondOrderCone());
    @objective(mod2, Min, t)
    optimize!(mod2)
    X_rounded=U_0*value.(D)*U_0'

    return X_rounded, ofv_dual, objective_value(mod2), U_0*sqrt(value.(D))
end

function solveBM(A, k, V_t)
    n=size(A,1)
    ofv_prev=1e10
    ofv_best=1e10
    V_t_best=V_t
    maxEpochs=100
    rho=0.0001
    # V_t=rand(n,k)

    for t=1:maxEpochs

    mod1=Model(Mosek.Optimizer)

    @variable(mod1, U[1:n, 1:k])
    @constraint(mod1, U.>=0.0)

    @objective(mod1, Min, (sum((A[i,j]-U[i,:]'*V_t[j,:])^2 for i=1:n for j=1:n))+rho*sum((U[i,j]-V_t[i,j])^2 for i=1:n for j=1:k))
    @suppress optimize!(mod1)
    # @show objective_value(mod1)

    U_t=value.(U)

    mod2=Model(Mosek.Optimizer)
    @variable(mod2, V[1:n, 1:k])
    @constraint(mod2, V.>=0.0)
    @objective(mod2, Min, (sum((A[i,j]-U_t[i,:]'*V[j,:])^2 for i=1:n for j=1:n))+rho*sum((V[i,j]-U_t[i,j])^2 for i=1:n for j=1:k))
    @suppress optimize!(mod2)
    # @show objective_value(mod2)
    V_t=value.(V)

    @show ofv_current=norm(A.-V_t*V_t')^2 # Sometimes find a better solution near the end of the reg path but not actually at the end
    if ofv_current<ofv_best
        ofv_best=ofv_current
        V_t_best=V_t
    end
    rho=min(rho*2.0, 1e5)

    if abs(ofv_current-ofv_prev)<1e-4
        break
    end
    ofv_prev=ofv_current

    end

    ofv_final=norm(A.-V_t_best*V_t_best')^2

    return V_t_best*V_t_best', ofv_final
end


function LineLine(l1::Geometry, l2::Geometry)
    #Ips = Array{Tuple{Point2D{Float64}, Float64, Float64}, 1}()
    dx_cx = l2.Ep.x - l2.Sp.x
    dy_cy = l2.Ep.y - l2.Sp.y
    bx_ax = l1.Ep.x - l1.Sp.x
    by_ay = l1.Ep.y - l1.Sp.y
    Δ = bx_ax * dy_cy - by_ay * dx_cx
    if abs(Δ) < 1e-4 # colinear
        if l1.Ep.x == l2.Sp.x # l2 follows l1
            return (intersect=true, type="connect", Ip=Point4To2D(l1.Ep), Dist1=l1.Length, Dist2=0.0)
        elseif l2.Ep.x == l1.Sp.x # l1 follows l2
            return (intersect=true, type="connect", Ip=Point4To2D(l1.Sp), Dist1=0.0, Dist2=l2.Length)
        elseif onSegment(l1.Sp, l1.Ep, l2.Sp)
            return (intersect=true, type="overlap", Ip=Point4To2D(l2.Sp), Dist1=Distance2D(l1.Sp, l2.Sp), Dist2=0.0)
        elseif onSegment(l2.Sp, l2.Ep, l1.Sp)
            return (intersect=true, type="overlap", Ip=Point4To2D(l1.Sp), Dist1=0.0, Dist2=Distance2D(l2.Sp, l1.Sp))
        else
            return (intersect=false, type="none")
        end
    else # not colinear
        ax_cx = l1.Sp.x - l2.Sp.x
        ay_cy = l1.Sp.y - l2.Sp.y
        r = (ay_cy * dx_cx - ax_cx * dy_cy) / Δ
        s = (ay_cy * bx_ax - ax_cx * by_ay) / Δ
        if (0.0 <= r <= 1.0 &&  0.0 <= s <= 1.0)  # intersection within 2 line-segments
            Ix = l1.Sp.x + r * bx_ax
            Iy = l1.Sp.y + r * by_ay
            Ipc = Point2D(Ix, Iy)
            dist1 = Distance2D(Point4To2D(l1.Sp), Ipc)
            dist2 = Distance2D(Point4To2D(l2.Sp), Ipc)
            return (intersect=true, type="cross", Ip=Ipc, Dist1=dist1, Dist2=dist2)
        else
            return (intersect=false, type="none")
        end
    end
end

function onSegment(Sp::Point4D{Float64}, Ep::Point4D{Float64}, Pi::Point4D{Float64})
    if (Pi.x <= max(Sp.x, Ep.x) && Pi.x >= min(Sp.x, Ep.x) &&
        Pi.y <= max(Sp.y, Ep.y) && Pi.y >= min(Sp.y, Ep.y))
        return true
    else
        return false
    end
end