using Revise
using NLSolvers
using LinearAlgebra
using SparseDiffTools
using SparseArrays
using IterativeSolvers
using DoubleFloats
using ForwardDiff
using Test

@testset "nlequations interface" begin
#
# Have Anderosn use FixedPointProblem
#
#

function theta(x)
   if x[1] > 0
       return atan(x[2] / x[1]) / (2.0 * pi)
   else
       return (pi + atan(x[2] / x[1])) / (2.0 * pi)
   end
end

function F_fletcher_powell!(Fx, x)
    if ( x[1]^2 + x[2]^2 == 0 )
        dtdx1 = 0;
        dtdx2 = 0;
    else
        dtdx1 = - x[2] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
        dtdx2 =   x[1] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
    end
    Fx[1] = -2000.0*(x[3]-10.0*theta(x))*dtdx1 +
    200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 );
    Fx[2] = -2000.0*(x[3]-10.0*theta(x))*dtdx2 +
    200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 );
    Fx[3] =  200.0*(x[3]-10.0*theta(x)) + 2.0*x[3];
    Fx
end

function F_jacobian_fletcher_powell!(Fx, Jx, x)
    ForwardDiff.jacobian!(Jx, F_fletcher_powell!, Fx, x)
    Fx, Jx
end

jv = JacVec(F_fletcher_powell!, rand(3); autodiff=false)
function jvop(x)
    jv.u .= x
    jv
end
prob_obj = NLSolvers.NEqObjective(F_fletcher_powell!, nothing, F_jacobian_fletcher_powell!, jvop)

prob = NEqProblem(prob_obj)

x0 = [-1.0, 0.0, 0.0]
res = solve(prob, x0, LineSearch(Newton(), Backtracking()))
@test norm(res.info.best_residual, Inf) < 1e-12

#x0 = [-1.0, 0.0, 0.0]
#res = solve(prob, x0, Anderson(), NEqOptions())
#@test norm(res.info.best_residual, Inf) < 1e-12

x0 = [-1.0, 0.0, 0.0]
res = solve(prob, x0, TrustRegion(Newton()), NEqOptions())
@test norm(res.info.best_residual, Inf) < 1e-12

#x0 = [-1.0, 0.0, 0.0]
#res = solve(prob, x0, TrustRegion(DBFGS()), NEqOptions())
#@test norm(res.info.best_residual, Inf) < 1e-12

#x0 = [-1.0, 0.0, 0.0]
#res = solve(prob, x0, TrustRegion(BFGS()), NEqOptions())
#@test norm(res.info.best_residual, Inf) < 1e-12

x0 = [-1.0, 0.0, 0.0]
res = solve(prob, x0, TrustRegion(SR1()), NEqOptions())
@test norm(res.info.best_residual, Inf) < 1e-9

x0 = [-1.0, 0.0, 0.0]
state = (z=copy(x0), d=copy(x0), Fx=copy(x0), Jx=zeros(3,3))
res = solve(prob, x0, LineSearch(Newton(), Backtracking()), NEqOptions(), state)
@test norm(res.info.best_residual, Inf) < 1e-12

#x0 = [-1.0, 0.0, 0.0]
#res = solve(prob, x0, InexactNewton(FixedForceTerm(1e-3), 1e-3, 300), NEqOptions())
#@test norm(res.info.best_residual, Inf) < 1e-12

#x0 = [-1.0, 0.0, 0.0]
#@show solve(prob, x0, InexactNewton(EisenstatWalkerA(), 1e-8, 300), NEqOptions())
#@test norm(res.info.best_residual, Inf) < 1e-12

#x0 = [-1.0, 0.0, 0.0]
#@show solve(prob, x0, InexactNewton(EisenstatWalkerB(), 1e-8, 300), NEqOptions())
#@test norm(res.info.best_residual, Inf) < 1e-12

x0 = [-1.0, 0.0, 0.0]
res = solve(prob, x0, NLSolvers.DFSANE(), NLSolvers.NEqOptions(maxiter=1000))
@test norm(res.info.best_residual, Inf) < 1e-0

function dfsane_exponential(Fx, x)
    Fx[1] = exp(x[1]-1)-1
    for i = 2:length(Fx)
  	    Fx[i] = i*(exp(x[i]-1)-x[i])
    end
    Fx
end

function FJ_dfsane_exponential!(Fx, Jx, x)
  ForwardDiff.jacobian!(Jx, dfsane_exponential, Fx, x)
  Fx, Jx
end
prob_obj = NLSolvers.NEqObjective(dfsane_exponential, nothing, FJ_dfsane_exponential!, nothing)

n = 5000
x0 = fill(n/(n-1), n)
res = solve(NEqProblem(prob_obj), x0, NLSolvers.DFSANE(), NLSolvers.NEqOptions(maxiter=1000))
@test norm(res.info.best_residual, Inf)<1e-5

n = 1000
x0 = fill(n/(n-1), n)
res = solve(NEqProblem(prob_obj), x0, LineSearch(NLSolvers.Newton()), NLSolvers.NEqOptions(f_abstol=1e-6)) #converges well but doing this for time
@test norm(res.info.best_residual, Inf)<1e-6

n = 5000
x0 = fill(n/(n-1), n)
res = solve(NEqProblem(prob_obj), x0, TrustRegion(NLSolvers.Newton()), NLSolvers.NEqOptions(f_abstol=1e-6)) #converges well but doing this for time
@test norm(res.info.best_residual, Inf)<1e-6

#n = 1000
#x0 = fill(n/(n-1), n)
#@show solve(NEqProblem(prob_obj), x0, TrustRegion(NLSolvers.BFGS()), NLSolvers.NEqOptions())

function dfsane_exponential2(Fx, x)
    Fx[1] = exp(x[1])-1
    for i = 2:length(Fx)
        Fx[i] = i/10*(exp(x[i])+x[i-1]-1)
    end
    Fx
end
dfsane_prob2 = NLSolvers.NEqObjective(dfsane_exponential2, nothing, nothing,nothing)
n = 500
x0 = fill(1/(n^2), n)
res = solve(NEqProblem(dfsane_prob2), x0, NLSolvers.DFSANE(), NLSolvers.NEqOptions())
@test norm(res.info.best_residual, Inf) < 1e-5
res = solve(NEqProblem(dfsane_prob2), Double64.(x0), NLSolvers.DFSANE(), NLSolvers.NEqOptions())
@test norm(res.info.best_residual, Inf) < 1e-5
n = 2000
x0 = fill(1/(n^2), n)
res = solve(NEqProblem(dfsane_prob2), x0, NLSolvers.DFSANE(), NLSolvers.NEqOptions())
@test norm(res.info.best_residual, Inf) < 1e-5

# Take from problems
prob = NEqProblem(NLE_PROBS["rosenbrock"]["array"]["mutating"])
x0 = copy(NLE_PROBS["rosenbrock"]["array"]["x0"])
X0 = [[-0.8, 1.0], [-0.7, 1.0], [-0.7, 1.2], [-0.5, 0.7], [0.5, 0.4], [-1.2, 1.0]]
for x0 in X0
    res = solve(prob, x0, TrustRegion(Newton(), Dogleg()), NEqOptions())
    @test norm(res.info.best_residual, Inf) < 1e-12
    res = solve(prob, x0, TrustRegion(Newton(), NWI()), NEqOptions())
    @test norm(res.info.best_residual, Inf) < 1e-12
    res = solve(prob, x0, TrustRegion(Newton(), NTR()), NEqOptions())
    @test norm(res.info.best_residual, Inf) < 1e-12
    res = solve(prob, x0, LineSearch(Newton(), Static(0.6)), NEqOptions())
    @test norm(res.info.best_residual, Inf) < 1e-8
    #res = solve(prob, copy(initial), DFSANE(), NEqOptions())
    #res = solve(prob, copy(initial), InexactNewton(), NEqOptions())
end

@testset "fixedpoints" begin
# function G(x, Gx)
#     V1 = [1.0, 0.0] .+ 0.99*[0.1 0.9; 0.5 0.5]*x
#     V2 = [0.0, 2.0] .+ 0.99*[0.5 0.5; 1.0 0.0]*x
#     K = max(maximum(V1), maximum(V2))
#     Gx .= K .+ log.(exp.(V1 .- K) .+ exp.(V2 .- K))
#   end
#   fp1 = NLSolvers.fixedpoint!(G, zeros(2), Anderson(10000000,1, nothing,nothing))
#   fp2 = NLSolvers.fixedpoint!(G, zeros(2), Anderson(2, 2, 0.3, 1e2))
#   fp3 = NLSolvers.fixedpoint!(G, zeros(2), Anderson(2, 2, 0.01, 1e2))
#   @test all(fp1.x .≈ fp2.x)
#   @test all(fp1.x .≈ fp3.x)

#   fp1 = NLSolvers.solve(NEqProblem((x, F)->G(x, F).-x), zeros(2), Anderson(10000000,1, nothing,nothing))
#   fp2 = NLSolvers.solve(NEqProblem((x, F)->G(x, F).-x), zeros(2), Anderson(2, 2, 0.3, 1e2))
#   fp3 = NLSolvers.solve(NEqProblem((x, F)->G(x, F).-x), zeros(2), Anderson(2, 2, 0.01, 1e2))
#   @test all(fp1.x .≈ fp2.x)
#   @test all(fp1.x .≈ fp3.x)
end

@testset "complementarity" begin
# using NLsolve

# M = [0  0 -1 -1 ;
#      0  0  1 -2 ;
#      1 -1  2 -2 ;
#      1  2 -2  4 ]

# q = [2; 2; -2; -6]

# function f!(x, fvec)
#     fvec = M * x + q
# end

# r = mcpsolve(f!, [0., 0., 0., 0.], [Inf, Inf, Inf, Inf],
#              [1.25, 0., 0., 0.5], reformulation = :smooth, autodiff = true)

# x = r.zero  # [1.25, 0.0, 0.0, 0.5]
# @show dot( M*x + q, x )  # 0.5

# sol = [2.8, 0.0, 0.8, 1.2]
# @show dot( M*sol + q, sol )  # 0.0

end
end

@testset "lsqfit" begin
# boxbod_f(x, b) = b[1]*(1-exp(-b[2]*x))

# ydata = [109, 149, 149, 191, 213, 224]
# xdata = [1, 2, 3, 5, 7, 10]

# start1 = [1.0, 1.0]
# start2 = [100.0, 0.75]

# using Plots
# using NLSolvers

# nd = NonDiffed(t->sum(abs2, boxbod_f.(xdata, Ref(t)).-ydata))
# bounds = (fill(0.0, 2), fill(500.0, 2))
# problem = MinProblem(; obj=nd, bounds=bounds)
# @show minimize!(problem, zeros(2), APSO(), OptimizationOptions())
# @show minimize!(nd, [200.0, 1], NelderMead(), OptimizationOptions())

# @show minimize!(nd, minimize!(nd, [200.0, 10.0], NelderMead(), OptimizationOptions()).info.minimizer, NelderMead(), OptimizationOptions())
# @show minimize!(nd, minimize!(nd, [200.0, 5.0], NelderMead(), OptimizationOptions()).info.minimizer, NelderMead(), OptimizationOptions())

# function F(b, F, J=nothing)
#   @. F = b[1]*(1 - exp(-b[2]*xdata)) - ydata
#   if !isa(J, Nothing)
#   	@. J[:, 1] = 1 - exp(-b[2]*xdata)
#   	@. @views J[:, 2] = xdata*b[1]*(J[:, 1]-1)
#     return F, J
#   end
#   F
# end

# Fc = zeros(6)
# minimize!(lsqwrap, [100.0, 1.0], NelderMead(), OptimizationOptions())

# lsqwrap = NLSolvers.LsqWrapper(F, zeros(6), zeros(6,2))
# minimize!(MinProblem(;obj=lsqwrap,bounds=([0.0,0.0],[250.0,2.0])), [100.0, 1.0], APSO(), OptimizationOptions())

# function F(F, b)
#   @. F = b[1]*(1 - exp(-b[2]*xdata)) - ydata

#   F
# end
# function F(J, F, b)
#   @. F = b[1]*(1 - exp(-b[2]*xdata)) - ydata
#   if !isa(J, Nothing)
#   	@. J[:, 1] = 1 - exp(-b[2]*xdata)
#   	@. @views J[:, 2] = xdata*b[1]*(J[:, 1]-1)
#     return F, J
#   end
#   F
# end
# OnceDiffed(F)(rand(2), rand(6), rand(6,2)

# Fc = zeros(6)

# lsqwrap = NLSolvers.LsqWrapper(OnceDiffed(F), zeros(6), zeros(6,2))
# minimize!(lsqwrap, [100.0, 1.0], LineSearch(LBFGS()), OptimizationOptions())

# #using Plots
# #theme(:ggplot2)
# #gr(size=(500,500))
# #X = range(000.0, 350.0; length=420)
# #Y = range(0.0, 3.00; length=420)
# #contour(X, Y, (x, y)->nd([x, y]);
# #       fill=true,
# #       c=:turbid,levels=200, ls=:dash,
# #       xlims=(minimum(X), maximum(X)),
# #       ylims=(minimum(Y), maximum(Y)),
# #       colorbar=true)

end


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ b0f0ff46-903b-11eb-3a2f-31e7f1f1ab7a
using ImageDraw, ColorTypes, BenchmarkTools,ColorVectorSpace,ImageMorphology

# ╔═╡ 837217b6-903c-11eb-180f-f14a7e934ff7
begin
	using HistogramThresholding
using TestImages # For the moonsurface image.
using ImageContrastAdjustment # For the build_histogram() function.

img = testimage("moonsurface")
edges, counts = build_histogram(img,256)
t = find_threshold(UnimodalRosin(), counts[1:end], edges)
img_final = zeros(Bool, axes(img))
for i in CartesianIndices(img)
        img_final[i] = img[i] < t ? 0 : 1
end
img_final
end

# ╔═╡ 5bbb0dd4-903e-11eb-0098-9b83c4f34379
using FileIO

# ╔═╡ 5f8fe6da-9040-11eb-086f-53680b92c269
using Noise,ImageShow

# ╔═╡ bb2be408-903b-11eb-3232-99614dda4abc
 begin
vert=CartesianIndex{2}[]
	push!(vert, CartesianIndex(2,2))
	push!(vert, CartesianIndex(3,2))
    push!(vert, CartesianIndex(4,2))
	push!(vert, CartesianIndex(5,2))
    push!(vert, CartesianIndex(6,3))
	push!(vert, CartesianIndex(7,4))
	push!(vert, CartesianIndex(8,4))
	push!(vert, CartesianIndex(9,3))
	push!(vert, CartesianIndex(10,2))
	push!(vert, CartesianIndex(11,3))
	push!(vert, CartesianIndex(12,4))
	push!(vert, CartesianIndex(13,5))
	push!(vert, CartesianIndex(12,6))
	push!(vert, CartesianIndex(11,7))
	push!(vert, CartesianIndex(10,6))#u
	push!(vert, CartesianIndex(9,7))
	push!(vert, CartesianIndex(8,8))
	push!(vert, CartesianIndex(8,9))
	push!(vert, CartesianIndex(8,10))
	push!(vert, CartesianIndex(8,11))
	push!(vert, CartesianIndex(9,11))
	push!(vert, CartesianIndex(10,11))
	push!(vert, CartesianIndex(11,11))
	push!(vert, CartesianIndex(12,11))
	push!(vert, CartesianIndex(13,11))
	push!(vert, CartesianIndex(13,12))
	push!(vert, CartesianIndex(13,13))
	push!(vert, CartesianIndex(12,13))
	push!(vert, CartesianIndex(11,13))
	push!(vert, CartesianIndex(10,13))
	push!(vert, CartesianIndex(9,13))
	push!(vert, CartesianIndex(8,13))
	push!(vert, CartesianIndex(7,13))
	push!(vert, CartesianIndex(6,13))
	push!(vert, CartesianIndex(6,12))
	push!(vert, CartesianIndex(6,11))
	push!(vert, CartesianIndex(5,11))
	push!(vert, CartesianIndex(4,11))
	push!(vert, CartesianIndex(3,11))
	push!(vert, CartesianIndex(2,11))
	push!(vert, CartesianIndex(3,10))
	push!(vert, CartesianIndex(4,9))
	push!(vert, CartesianIndex(5,8))
	push!(vert, CartesianIndex(4,7))
	push!(vert, CartesianIndex(4,6))
	push!(vert, CartesianIndex(5,5))
	push!(vert, CartesianIndex(4,4))
    push!(vert, CartesianIndex(3,3))
	push!(vert, CartesianIndex(2,2))
end

# ╔═╡ a6e92884-903b-11eb-3830-0191de19b831
# begin
# img = zeros(Bool,14,14)
# img[1,1] = 1
# img[13,13] =1
# img[1,13] =1
# img[13,1] =1
# end

# ╔═╡ 8554a5be-903d-11eb-069c-11673c9c9356
img

# ╔═╡ e67457e4-903b-11eb-3b1b-89259950a651
convex_hull_test = convexhull(img_final)

# ╔═╡ 926e9572-903d-11eb-2ec2-e15a01a9524a
# begin
# 	img_final1 = zeros(Gray{Bool}, axes(img))
# for i in CartesianIndices(convex_hull_test)
#         img_final1[i] = convex_hull_test[i][1] < 1 ? 0 : 1
# end
# img_final1
# Gray.(img_final1)
# end

# ╔═╡ 47fe6612-903e-11eb-280d-7f745e6c0612
begin
img_coins = load("../1.jpg")
end

# ╔═╡ a60870b6-903e-11eb-263c-05a406696259
begin
edges1, counts1 = build_histogram(Gray.(img_coins),256)
t1 = find_threshold(UnimodalRosin(), counts1[1:end], edges1)
img_final_coins = zeros(Bool, axes(Gray.(img_coins)))
for i in CartesianIndices(img_coins)
        img_final_coins[i] = Gray.(img_coins[i]) < t1 ? 0 : 1
end
img_final_coins
end

# ╔═╡ 9cd44928-9040-11eb-1f50-8fdfbdb00ab8
Gray.(img_final_coins)

# ╔═╡ 6b951ff6-9040-11eb-287e-6dec9dc1d7fa
img_final_coins1 = salt_pepper(img_final_coins, 0.5)

# ╔═╡ d903f85e-903f-11eb-36ef-d5435d6ed4c6
Gray.(img_final_coins1)

# ╔═╡ 865d4368-903e-11eb-36fc-3745cdadf810
begin
fill_image=imfill(img_final_coins, (1,2))
Gray.(fill_image)
end

# ╔═╡ 24ebf4fa-9041-11eb-245a-31e91bd0f2ee
img_thinning=thinning(img_final_coins,algo=GuoAlgo())

# ╔═╡ 34e54dc0-9041-11eb-325e-43f60e27a486
Gray.(img_thinning)

# ╔═╡ 7236322a-9041-11eb-13f2-1153e9fe7571
Gray.([img_final_coins img_thinning])

# ╔═╡ 007b900c-9042-11eb-27da-611a1b40a7a8
cleared_img = clearborder(img_final_coins, 20)

# ╔═╡ 0b500786-9042-11eb-2b74-79c766338419
Gray.([img_final_coins cleared_img])

# ╔═╡ Cell order:
# ╠═b0f0ff46-903b-11eb-3a2f-31e7f1f1ab7a
# ╠═bb2be408-903b-11eb-3232-99614dda4abc
# ╠═a6e92884-903b-11eb-3830-0191de19b831
# ╠═8554a5be-903d-11eb-069c-11673c9c9356
# ╠═837217b6-903c-11eb-180f-f14a7e934ff7
# ╠═e67457e4-903b-11eb-3b1b-89259950a651
# ╠═926e9572-903d-11eb-2ec2-e15a01a9524a
# ╠═5bbb0dd4-903e-11eb-0098-9b83c4f34379
# ╠═5f8fe6da-9040-11eb-086f-53680b92c269
# ╠═47fe6612-903e-11eb-280d-7f745e6c0612
# ╠═a60870b6-903e-11eb-263c-05a406696259
# ╠═9cd44928-9040-11eb-1f50-8fdfbdb00ab8
# ╠═6b951ff6-9040-11eb-287e-6dec9dc1d7fa
# ╠═d903f85e-903f-11eb-36ef-d5435d6ed4c6
# ╠═865d4368-903e-11eb-36fc-3745cdadf810
# ╠═24ebf4fa-9041-11eb-245a-31e91bd0f2ee
# ╠═34e54dc0-9041-11eb-325e-43f60e27a486
# ╠═7236322a-9041-11eb-13f2-1153e9fe7571
# ╠═007b900c-9042-11eb-27da-611a1b40a7a8
# ╠═0b500786-9042-11eb-2b74-79c766338419


function matrix_cone(S::Type{<:MOI.AbstractSymmetricMatrixSetTriangle},
                     side_dimension)
    if iszero(side_dimension)
        # Some solvers such as Mosek does not support 0-dimensional PSD cone
        return EmptyCone()
    elseif isone(side_dimension)
        # PSD constraints on 1x1 matrices are equivalent to the nonnegativity
        # of the only entry.
        return MOI.Nonnegatives(1)
    elseif side_dimension == 2
        # PSD constraints on 2x2 matrices are SOC representable.
        return PositiveSemidefinite2x2ConeTriangle()
    else
        # PSD constraints on nxn matrices with n ≥ 3 is not SOC representable,
        # see [F18].
        #
        # [F18] Fawzi, Hamza
        # On representing the positive semidefinite cone using the second-order
        # cone.
        # Mathematical Programming (2018): 1-10.
        return S(side_dimension)
    end
end

"""
    struct DiagonallyDominantConeTriangle <: MOI.AbstractSymmetricMatrixSetTriangle
        side_dimension::Int
    end

See Definition 4 of [AM17] for a precise definition of the last two items.

[AM17] Ahmadi, A. A. & Majumdar, A.
*DSOS and SDSOS Optimization: More Tractable Alternatives to Sum of Squares and Semidefinite Optimization*
ArXiv e-prints, **2017**.
"""
struct DiagonallyDominantConeTriangle <: MOI.AbstractSymmetricMatrixSetTriangle
    side_dimension::Int
end

function matrix_cone(S::Type{DiagonallyDominantConeTriangle},
                     side_dimension)
    if iszero(side_dimension)
        return EmptyCone()
    elseif isone(side_dimension)
        return MOI.Nonnegatives(1)
    else
        # With `side_dimension` = 2, we want to avoid using SOC and only use LP
        return S(side_dimension)
    end
end

"""
    struct ScaledDiagonallyDominantConeTriangle <: MOI.AbstractSymmetricMatrixSetTriangle
        side_dimension::Int
    end

See Definition 4 of [AM17] for a precise definition of the last two items.

[AM17] Ahmadi, A. A. & Majumdar, A.
*DSOS and SDSOS Optimization: More Tractable Alternatives to Sum of Squares and Semidefinite Optimization*
ArXiv e-prints, **2017**.
"""
struct ScaledDiagonallyDominantConeTriangle <: MOI.AbstractSymmetricMatrixSetTriangle
    side_dimension::Int
end

function matrix_cone(S::Type{ScaledDiagonallyDominantConeTriangle},
                     side_dimension)
    if iszero(side_dimension)
        return EmptyCone()
    elseif isone(side_dimension)
        return MOI.Nonnegatives(side_dimension)
    elseif side_dimension == 2
        return PositiveSemidefinite2x2ConeTriangle()
    else
        return S(side_dimension)
    end
end

# isbits types, nothing to copy
function Base.copy(set::Union{DiagonallyDominantConeTriangle,
                              ScaledDiagonallyDominantConeTriangle})
    return set
end


"""
  Clustering.jl file

Clustering and collaborative filtering (via clustering) algorithms

- [Importable source code (most up-to-date version)](https://github.com/sylvaticus/BetaML.jl/blob/master/src/Clustering.jl) - [Julia Package](https://github.com/sylvaticus/BetaML.jl)
- [Demonstrative static notebook](https://github.com/sylvaticus/BetaML.jl/blob/master/notebooks/Clustering.ipynb)
- [Demonstrative live notebook](https://mybinder.org/v2/gh/sylvaticus/BetaML.jl/master?filepath=notebooks%2FClustering.ipynb) (temporary personal online computational environment on myBinder) - it can takes minutes to start with!
- Theory based on [MITx 6.86x - Machine Learning with Python: from Linear Models to Deep Learning](https://github.com/sylvaticus/MITx_6.86x) ([Unit 4](https://github.com/sylvaticus/MITx_6.86x/blob/master/Unit%2004%20-%20Unsupervised%20Learning/Unit%2004%20-%20Unsupervised%20Learning.md))
- New to Julia? [A concise Julia tutorial](https://github.com/sylvaticus/juliatutorial) - [Julia Quick Syntax Reference book](https://julia-book.com)
"""

"""
    Clustering module (WIP)

Provide clustering methods and missing values imputation / collaborative filtering / reccomendation systems using clustering methods as backend.

The module provides the following functions. Use `?[function]` to access their full signature and detailed documentation:

- [`initRepresentatives(X,K;initStrategy,Z₀)`](@ref initRepresentatives): Initialisation strategies for Kmean and Kmedoids
- [`kmeans(X,K;dist,initStrategy,Z₀)`](@ref kmeans): Classical KMean algorithm
- [`kmedoids(X,K;dist,initStrategy,Z₀)`](@ref kmedoids): Kmedoids algorithm
- [`gmm(X,K;p₀,mixtures,tol,verbosity,minVariance,minCovariance,initStrategy)`](@ref gmm): gmm algorithm over GMM
- [`predictMissing(X,K;p₀,mixtures,tol,verbosity,minVariance,minCovariance)`](@ref predictMissing): Impute mixing values ("matrix completion") using gmm as backbone. Note that this can be used for collaborative filtering / reccomendation systems often with better results than traditional algorithms as k-nearest neighbors (KNN)

{Spherical|Diagonal|Full}Gaussian mixtures for `gmm` / `predictMissing` are already provided. User defined mixtures can be used defining a struct as subtype of `Mixture` and implementing for that mixture the following functions:
- `initMixtures!(mixtures, X; minVariance, minCovariance, initStrategy)`
- `lpdf(m,x,mask)` (for the e-step)
- `updateParameters!(mixtures, X, pₙₖ; minVariance, minCovariance)` (the m-step)

"""
module Clustering

using LinearAlgebra, Random, Statistics, Reexport, CategoricalArrays
import Distributions

using  ForceImport
@force using ..Api
@force using ..Utils

export initRepresentatives, kmeans, kmedoids, gmm, predictMissing

abstract type AbstractMixture end
include("Mixtures.jl")

## Helper functions

"""
  initRepresentatives(X,K;initStrategy,Z₀)

Initialisate the representatives for a K-Mean or K-Medoids algorithm

# Parameters:
* `X`: a (N x D) data to clusterise
* `K`: Number of cluster wonted
* `initStrategy`: Whether to select the initial representative vectors:
  * `random`: randomly in the X space
  * `grid`: using a grid approach [default]
  * `shuffle`: selecting randomly within the available points
  * `given`: using a provided set of initial representatives provided in the `Z₀` parameter
 * `Z₀`: Provided (K x D) matrix of initial representatives (used only together with the `given` initStrategy) [default: `nothing`]
 * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Returns:
* A (K x D) matrix of initial representatives

# Example:
```julia
julia> Z₀ = initRepresentatives([1 10.5;1.5 10.8; 1.8 8; 1.7 15; 3.2 40; 3.6 32; 3.6 38],2,initStrategy="given",Z₀=[1.7 15; 3.6 40])
```
"""
function initRepresentatives(X,K;initStrategy="grid",Z₀=nothing,rng = Random.GLOBAL_RNG)
    X  = makeMatrix(X)
    (N,D) = size(X)
    # Random choice of initial representative vectors (any point, not just in X!)
    minX = minimum(X,dims=1)
    maxX = maximum(X,dims=1)
    Z = zeros(K,D)
    if initStrategy == "random"
        for i in 1:K
            for j in 1:D
                Z[i,j] = rand(rng,Distributions.Uniform(minX[j],maxX[j]))
            end
        end
    elseif initStrategy == "grid"
        for d in 1:D
                # same "space" for each class on each dimension
                Z[:,d] = collect(range(minX[d] + (maxX[d]-minX[d])/(K*2) , stop=maxX[d] - (maxX[d]-minX[d])/(K*2)  , length=K))
                #ex: collect(range(minX[d], stop=maxX[d], length=K))
                #collect(range(s+(e-s)/(K*2), stop=e-(e-s)/(K*2), length=K))
        end
    elseif initStrategy == "given"
        if isnothing(Z₀) error("With the `given` strategy you need to provide the initial set of representatives in the Z₀ parameter.") end
        Z₀ = makeMatrix(Z₀)
        Z = Z₀
    elseif initStrategy == "shuffle"
        zIdx = shuffle(rng,1:size(X)[1])[1:K]
        Z = X[zIdx, :]
    else
        error("initStrategy \"$initStrategy\" not implemented")
    end
    return Z
end

## Basic K-Means Algorithm (Lecture/segment 13.7 of https://www.edx.org/course/machine-learning-with-python-from-linear-models-to)

"""
  kmeans(X,K;dist,initStrategy,Z₀)

Compute K-Mean algorithm to identify K clusters of X using Euclidean distance

# Parameters:
* `X`: a (N x D) data to clusterise
* `K`: Number of cluster wonted
* `dist`: Function to employ as distance (see notes). Default to Euclidean distance.
* `initStrategy`: Whether to select the initial representative vectors:
  * `random`: randomly in the X space
  * `grid`: using a grid approach [default]
  * `shuffle`: selecting randomly within the available points
  * `given`: using a provided set of initial representatives provided in the `Z₀` parameter
* `Z₀`: Provided (K x D) matrix of initial representatives (used only together with the `given` initStrategy) [default: `nothing`]
* `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Returns:
* A tuple of two items, the first one being a vector of size N of ids of the clusters associated to each point and the second one the (K x D) matrix of representatives

# Notes:
* Some returned clusters could be empty
* The `dist` parameter can be:
  * Any user defined function accepting two vectors and returning a scalar
  * An anonymous function with the same characteristics (e.g. `dist = (x,y) -> norm(x-y)^2`)
  * One of the above predefined distances: `l1_distance`, `l2_distance`, `l2²_distance`, `cosine_distance`

# Example:
```julia
julia> (clIdx,Z) = kmeans([1 10.5;1.5 10.8; 1.8 8; 1.7 15; 3.2 40; 3.6 32; 3.3 38; 5.1 -2.3; 5.2 -2.4],3)
```
"""
function kmeans(X,K;dist=(x,y) -> norm(x-y),initStrategy="grid",Z₀=nothing,rng = Random.GLOBAL_RNG)
    X  = makeMatrix(X)
    (N,D) = size(X)
    # Random choice of initial representative vectors (any point, not just in X!)
    minX = minimum(X,dims=1)
    maxX = maximum(X,dims=1)
    Z₀ = initRepresentatives(X,K,initStrategy=initStrategy,Z₀=Z₀,rng=rng)
    Z  = Z₀
    cIdx_prev = zeros(Int64,N)

    # Looping
    while true
        # Determining the constituency of each cluster
        cIdx      = zeros(Int64,N)
        for (i,x) in enumerate(eachrow(X))
            cost = Inf
            for (j,z) in enumerate(eachrow(Z))
               if (dist(x,z)  < cost)
                   cost    =  dist(x,z)
                   cIdx[i] = j
               end
            end
        end

        # Determining the new representative by each cluster
        # for (j,z) in enumerate(eachrow(Z))
        for j in  1:K
            Cⱼ = X[cIdx .== j,:] # Selecting the constituency by boolean selection
            #debug  =  sum(Cⱼ,dims=1)
            #debug2 = size(Cⱼ)[1]
            if size(Cⱼ)[1] > 0
                Z[j,:] = sum(Cⱼ,dims=1) ./ size(Cⱼ)[1]
            else
                # move toward the center if no costituency
                xAvg = mean(X,dims=1)'
                Z[j,:] = Z[j,:] .+ ((xAvg - Z[j,:]) .* 0.01)
                #debug  = sum(Cⱼ,dims=1)
                # debug2 = size(Cⱼ)[1]
            end
            #Z[j,:] = median(Cⱼ,dims=1) # for l1 distance
        end

        # Checking termination condition: clusters didn't move any more
        if cIdx == cIdx_prev
            return (cIdx,Z)
        else
            cIdx_prev = cIdx
        end

    end
end

## Basic K-Medoids Algorithm (Lecture/segment 14.3 of https://www.edx.org/course/machine-learning-with-python-from-linear-models-to)
"""
  kmedoids(X,K;dist,initStrategy,Z₀)

Compute K-Medoids algorithm to identify K clusters of X using distance definition `dist`

# Parameters:
* `X`: a (n x d) data to clusterise
* `K`: Number of cluster wonted
* `dist`: Function to employ as distance (see notes). Default to Euclidean distance.
* `initStrategy`: Whether to select the initial representative vectors:
  * `random`: randomly in the X space
  * `grid`: using a grid approach
  * `shuffle`: selecting randomly within the available points [default]
  * `given`: using a provided set of initial representatives provided in the `Z₀` parameter
 * `Z₀`: Provided (K x D) matrix of initial representatives (used only together with the `given` initStrategy) [default: `nothing`]
 * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Returns:
* A tuple of two items, the first one being a vector of size N of ids of the clusters associated to each point and the second one the (K x D) matrix of representatives

# Notes:
* Some returned clusters could be empty
* The `dist` parameter can be:
  * Any user defined function accepting two vectors and returning a scalar
  * An anonymous function with the same characteristics (e.g. `dist = (x,y) -> norm(x-y)^2`)
  * One of the above predefined distances: `l1_distance`, `l2_distance`, `l2²_distance`, `cosine_distance`

# Example:
```julia
julia> (clIdx,Z) = kmedoids([1 10.5;1.5 10.8; 1.8 8; 1.7 15; 3.2 40; 3.6 32; 3.3 38; 5.1 -2.3; 5.2 -2.4],3,initStrategy="grid")
```
"""
function kmedoids(X,K;dist=(x,y) -> norm(x-y),initStrategy="grid",Z₀=nothing,rng = Random.GLOBAL_RNG)
    X  = makeMatrix(X)
    (n,d) = size(X)
    # Random choice of initial representative vectors
    Z₀ = initRepresentatives(X,K,initStrategy=initStrategy,Z₀=Z₀,rng=rng)
    Z = Z₀
    cIdx_prev = zeros(Int64,n)

    # Looping
    while true
        # Determining the constituency of each cluster
        cIdx      = zeros(Int64,n)
        for (i,x) in enumerate(eachrow(X))
            cost = Inf
            for (j,z) in enumerate(eachrow(Z))
               if (dist(x,z) < cost)
                   cost =  dist(x,z)
                   cIdx[i] = j
               end
            end
        end

        # Determining the new representative by each cluster (within the points member)
        #for (j,z) in enumerate(eachrow(Z))
        for j in  1:K
            Cⱼ = X[cIdx .== j,:] # Selecting the constituency by boolean selection
            nⱼ = size(Cⱼ)[1]     # Size of the cluster
            if nⱼ == 0 continue end # empty continuency. Let's not do anything. Stil in the next batch other representatives could move away and points could enter this cluster
            bestCost = Inf
            bestCIdx = 0
            for cIdx in 1:nⱼ      # candidate index
                 candidateCost = 0.0
                 for tIdx in 1:nⱼ # target index
                     candidateCost += dist(Cⱼ[cIdx,:],Cⱼ[tIdx,:])
                 end
                 if candidateCost < bestCost
                     bestCost = candidateCost
                     bestCIdx = cIdx
                 end
            end
            Z[j,:] = reshape(Cⱼ[bestCIdx,:],1,d)
        end

        # Checking termination condition: clusters didn't move any more
        if cIdx == cIdx_prev
            return (cIdx,Z)
        else
            cIdx_prev = cIdx
        end
    end

end

## The gmm algorithm (Lecture/segment 16.5 of https://www.edx.org/course/machine-learning-with-python-from-linear-models-to)

# no longer true with the numerical trick implemented
# - For mixtures with full covariance matrix (i.e. `FullGaussian(μ,σ²)`) the minCovariance should NOT be set equal to the minVariance, or if the covariance matrix goes too low, it will become singular and not invertible.
"""
  gmm(X,K;p₀,mixtures,tol,verbosity,minVariance,minCovariance,initStrategy)

Compute Expectation-Maximisation algorithm to identify K clusters of X data, i.e. employ a Generative Mixture Model as the underlying probabilistic model.

X can contain missing values in some or all of its dimensions. In such case the learning is done only with the available data.
Implemented in the log-domain for better numerical accuracy with many dimensions.

# Parameters:
* `X`  :           A (n x d) data to clusterise
* `K`  :           Number of cluster wanted
* `p₀` :           Initial probabilities of the categorical distribution (K x 1) [default: `nothing`]
* `mixtures`:      An array (of length K) of the mixture to employ (see notes) [def: `[DiagonalGaussian() for i in 1:K]`]
* `tol`:           Tolerance to stop the algorithm [default: 10^(-6)]
* `verbosity`:     A verbosity parameter regulating the information messages frequency [def: `STD`]
* `minVariance`:   Minimum variance for the mixtures [default: 0.05]
* `minCovariance`: Minimum covariance for the mixtures with full covariance matrix [default: 0]. This should be set different than minVariance (see notes).
* `initStrategy`:  Mixture initialisation algorithm [def: `kmeans`]
* `maxIter`:       Maximum number of iterations [def: `-1`, i.e. ∞]
* `rng`:           Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Returns:
* A named touple of:
  * `pₙₖ`:      Matrix of size (N x K) of the probabilities of each point i to belong to cluster j
  * `pₖ`:       Probabilities of the categorical distribution (K x 1)
  * `mixtures`: Vector (K x 1) of the estimated underlying distributions
  * `ϵ`:        Vector of the discrepancy (matrix norm) between pⱼₓ and the lagged pⱼₓ at each iteration
  * `lL`:       The log-likelihood (without considering the last mixture optimisation)
  * `BIC`:      The Bayesian Information Criterion (lower is better)
  * `AIC`:      The Akaike Information Criterion (lower is better)

 # Notes:
 - The mixtures currently implemented are `SphericalGaussian(μ,σ²)`,`DiagonalGaussian(μ,σ²)` and `FullGaussian(μ,σ²)`
 - Reasonable choices for the minVariance/Covariance depends on the mixture. For example 0.25 seems a reasonable value for the SphericalGaussian, 0.05 seems better for the DiagonalGaussian, and FullGaussian seems to prefer either very low values of variance/covariance (e.g. `(0.05,0.05)` ) or very big but similar ones (e.g. `(100,100)` ).
 - For `initStrategy`, look at the documentation of `initMixtures!` for the mixture you want. The provided gaussian mixtures support `grid`, `kmeans` or `given`. `grid` is faster (expecially if X contains missing values), but `kmeans` often provides better results.

 # Resources:
 - [Paper describing gmm with missing values](https://doi.org/10.1016/j.csda.2006.10.002)
 - [Class notes from MITx 6.86x (Sec 15.9)](https://stackedit.io/viewer#!url=https://github.com/sylvaticus/MITx_6.86x/raw/master/Unit 04 - Unsupervised Learning/Unit 04 - Unsupervised Learning.md)
 - [Limitations of gmm](https://www.r-craft.org/r-news/when-not-to-use-gaussian-mixture-model-gmm-clustering/)

# Example:
```julia
julia> clusters = gmm([1 10.5;1.5 0; 1.8 8; 1.7 15; 3.2 40; 0 0; 3.3 38; 0 -2.3; 5.2 -2.4],3,verbosity=HIGH)
```
"""
function gmm(X,K;p₀=nothing,mixtures=[DiagonalGaussian() for i in 1:K],tol=10^(-6),verbosity=STD,minVariance=0.05,minCovariance=0.0,initStrategy="kmeans",maxIter=-1,rng = Random.GLOBAL_RNG)
    if verbosity > STD
        @codeLocation
    end
    # debug:
    #X = [1 10.5;1.5 missing; 1.8 8; 1.7 15; 3.2 40; missing missing; 3.3 38; missing -2.3; 5.2 -2.4]
    #K = 3
    #p₀=nothing; tol=0.0001; msgStep=1; minVariance=0.25; initStrategy="grid"
    #mixtures = [SphericalGaussian() for i in 1:K]
    # ---------
    X     = makeMatrix(X)
    (N,D) = size(X)
    pₖ    = isnothing(p₀) ? fill(1/K,K) : p₀

    # no longer true with the numerical trick implemented
    #if (minVariance == minCovariance)
    #    @warn("Setting the minVariance equal to the minCovariance may lead to singularity problems for mixtures with full covariance matrix.")
    #end

    msgStepMap = Dict(NONE => 0, LOW=>100, STD=>20, HIGH=>5, FULL=>1)
    msgStep    = msgStepMap[verbosity]

    # Initialisation of the parameters of the mixtures
    initMixtures!(mixtures,X,minVariance=minVariance,minCovariance=minCovariance,initStrategy=initStrategy,rng=rng)

    pₙₖ = zeros(Float64,N,K) # The posteriors, i.e. the prob that item n belong to cluster k
    ϵ = Float64[]

    # Checking dimensions only once (but adding then inbounds doesn't change anything. Still good
    # to provide a nice informative message)
    if size(pₖ,1) != K || length(mixtures) != K
        error("Error in the dimensions of the inputs. Please check them.")
    end

    # finding empty/non_empty values
    Xmask     =  .! ismissing.(X)
    #XdimCount = sum(Xmask, dims=2)

    lL = -Inf
    iter = 1
    while(true)
        oldlL = lL
        # E Step: assigning the posterior prob p(j|xi) and computing the log-Likelihood of the parameters given the set of data
        # (this last one for informative purposes and terminating the algorithm)
        pₙₖlagged = copy(pₙₖ)
        logpₙₖ = log.(pₙₖ)
        lL = 0
        for n in 1:N
            if any(Xmask[n,:]) # if at least one true
                Xu = X[n,Xmask[n,:]]
                #=if (length(ϵ) == 2)
                    println("here I am")
                    for m in mixtures[3:end]
                        println(m.μ)
                        println(m.σ²)
                        println(Xu)
                        println(Xmask[n,:])
                        lpdf(m,Xu,Xmask[n,:])
                        println("here I am partially")
                    end
                    println("here I am dead")
                end=#
                logpx = lse([log(pₖ[k] + 1e-16) + lpdf(mixtures[k],Xu,Xmask[n,:]) for k in 1:K])
                lL += logpx
                #px = sum([pⱼ[k]*normalFixedSd(Xu,μ[k,XMask[n,:]],σ²[k]) for k in 1:K])
                #println(n)
                for k in 1:K
                    logpₙₖ[n,k] = log(pₖ[k] + 1e-16)+lpdf(mixtures[k],Xu,Xmask[n,:])-logpx
                end
            else
                logpₙₖ[n,:] = log.(pₖ)
            end
        end
        pₙₖ = exp.(logpₙₖ)

        push!(ϵ,norm(pₙₖlagged - pₙₖ))

        # M step: find parameters that maximise the likelihood
        # Updating the probabilities of the different mixtures
        nₖ = sum(pₙₖ,dims=1)'
        n  = sum(nₖ)
        pₖ = nₖ ./ n
        updateParameters!(mixtures, X, pₙₖ; minVariance=minVariance,minCovariance=minCovariance)

        # Information. Note the likelihood is whitout accounting for the new mu, sigma
        if msgStep != 0 && (length(ϵ) % msgStep == 0 || length(ϵ) == 1)
            println("Iter. $(length(ϵ)):\tVar. of the post  $(ϵ[end]) \t  Log-likelihood $(lL)")
        end

        # Closing conditions. Note that the logLikelihood is those without considering the new mu,sigma
        if ((lL - oldlL) <= (tol * abs(lL))) || (maxIter > 0 && iter == maxIter)
            npars = npar(mixtures) + (K-1)
            #BIC  = lL - (1/2) * npars * log(N)
            BICv = bic(lL,npars,N)
            AICv = aic(lL,npars)
        #if (ϵ[end] < tol)
           return (pₙₖ=pₙₖ,pₖ=pₖ,mixtures=mixtures,ϵ=ϵ,lL=lL,BIC=BICv,AIC=AICv)
       else
            iter += 1
       end
    end # end while loop
end # end function

#using BenchmarkTools
#@benchmark clusters = emGMM([1 10.5;1.5 10.8; 1.8 8; 1.7 15; 3.2 40; 3.6 32; 3.3 38; 5.1 -2.3; 5.2 -2.4],3,msgStep=0)
#@benchmark clusters = emGMM([1 10.5;1.5 0; 1.8 8; 1.7 15; 3.2 40; 0 0; 3.3 38; 0 -2.3; 5.2 -2.4],3,msgStep=0,missingValue=0)
#@benchmark clusters = emGMM([1 10.5;1.5 0; 1.8 8; 1.7 15; 3.2 40; 0 0; 3.3 38; 0 -2.3; 5.2 -2.4],3,msgStep=0,missingValue=0)
#@benchmark clusters = emGMM([1 10.5;1.5 missing; 1.8 8; 1.7 15; 3.2 40; missing missing; 3.3 38; missing -2.3; 5.2 -2.4],3,msgStep=0)
#@code_warntype gmm([1 10.5;1.5 0; 1.8 8; 1.7 15; 3.2 40; 0 0; 3.3 38; 0 -2.3; 5.2 -2.4],3,msgStep=0,missingValue=0)
#using Profile
#Juno.@profiler (for i = 1:1000 gmm([1 10.5;1.5 0; 1.8 8; 1.7 15; 3.2 40; 0 0; 3.3 38; 0 -2.3; 5.2 -2.4],3,msgStep=0,missingValue=0) end)
#Profile.clear()
#Profile.print()

#  - For mixtures with full covariance matrix (i.e. `FullGaussian(μ,σ²)`) the minCovariance should NOT be set equal to the minVariance, or if the covariance matrix goes too low, it will become singular and not invertible.
"""
  predictMissing(X,K;p₀,mixtures,tol,verbosity,minVariance,minCovariance)

Fill missing entries in a sparse matrix (i.e. perform a "matrix completion") assuming an underlying Gaussian Mixture probabilistic Model (GMM) and implementing
an Expectation-Maximisation algorithm.

While the name of the function is `predictMissing`, the function can be also used for system reccomendation / collaborative filtering and GMM-based regressions. The advantage over traditional algorithms as k-nearest neighbors (KNN) is that GMM can "detect" the hidden structure of the observed data, where some observation can be similar to a certain pool of other observvations for a certain characteristic, but similar to an other pool of observations for other characteristics.

Implemented in the log-domain for better numerical accuracy with many dimensions.

# Parameters:
* `X`  :           A (N x D) sparse matrix of data to fill according to a GMM model
* `K`  :           Number of mixtures (latent classes) to consider [def: 3]
* `p₀` :           Initial probabilities of the categorical distribution (K x 1) [default: `nothing`]
* `mixtures`:      An array (of length K) of the mixture to employ (see notes) [def: `[DiagonalGaussian() for i in 1:K]`]
* `tol`:           Tolerance to stop the algorithm [default: 10^(-6)]
* `verbosity`:     A verbosity parameter regulating the information messages frequency [def: `STD`]
* `minVariance`:   Minimum variance for the mixtures [default: 0.05]
* `minCovariance`: Minimum covariance for the mixtures with full covariance matrix [default: 0]. This should be set different than minVariance (see notes).
* `initStrategy`:  Mixture initialisation algorithm [def: `grid`]
* `maxIter`:       Maximum number of iterations [def: `-1`, i.e. ∞]
* `rng`:           Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

# Returns:
* A named touple of:
  * `̂X̂`    : The Filled Matrix of size (N x D)
  * `nFill`: The number of items filled
  * `lL`   : The log-likelihood (without considering the last mixture optimisation)
  * `BIC` :  The Bayesian Information Criterion (lower is better)
  * `AIC` :  The Akaike Information Criterion (lower is better)

  # Notes:
  - The mixtures currently implemented are `SphericalGaussian(μ,σ²)`,`DiagonalGaussian(μ,σ²)` and `FullGaussian(μ,σ²)`
  - For `initStrategy`, look at the documentation of `initMixtures!` for the mixture you want. The provided gaussian mixtures support `grid`, `kmeans` or `given`. `grid` is faster, but `kmeans` often provides better results.
  - The algorithm requires to specify a number of "latent classes" (mlixtures) to divide the dataset into. If there isn't any prior domain specific knowledge on this point one can test sevaral `k` and verify which one minimise the `BIC` or `AIC` criteria.

# Example:
```julia
julia>  cFOut = predictMissing([1 10.5;1.5 missing; 1.8 8; 1.7 15; 3.2 40; missing missing; 3.3 38; missing -2.3; 5.2 -2.4],3)
```
"""
function predictMissing(X,K=3;p₀=nothing,mixtures=[DiagonalGaussian() for i in 1:K],tol=10^(-6),verbosity=STD,minVariance=0.05,minCovariance=0.0,initStrategy="kmeans",maxIter=-1,rng = Random.GLOBAL_RNG)
    if verbosity > STD
        @codeLocation
    end
    emOut = gmm(X,K;p₀=p₀,mixtures=mixtures,tol=tol,verbosity=verbosity,minVariance=minVariance,minCovariance=minCovariance,initStrategy=initStrategy,maxIter=maxIter,rng=rng)
    (N,D) = size(X)
    nDim  = ndims(X)
    nmT   = nonmissingtype(eltype(X))
    #K = size(emOut.μ)[1]
    XMask = .! ismissing.(X)
    nFill = (N * D) - sum(XMask)
    X̂ = copy(X)
    for n in 1:N
        for d in 1:D
            if !XMask[n,d]
                 X̂[n,d] = sum([emOut.mixtures[k].μ[d] * emOut.pₙₖ[n,k] for k in 1:K])
            end
        end
    end
    X̂ = convert(Array{nmT,nDim},X̂)
    return (X̂=X̂,nFill=nFill,lL=emOut.lL,BIC=emOut.BIC,AIC=emOut.AIC)
end

# MLJ interface
include("Clustering_MLJ.jl")

end


"""
    yyyymmdd2yyyy(d)

Takes an Int64-formatted date and returns the year in Int64 format.
```@example
dt = 20100811
@show yyyymmdd2yyyy(dt)
```
"""
yyyymmdd2yyyy(d) = div(d,10000)

"""
    yyyymmdd2yyyymm(d)

Takes an Int64-formatted date and returns the year and month in Int64 format.
```@example
dt = 20110909
@show yyyymmdd2yyyymm(dt)
```
"""
yyyymmdd2yyyymm(d) = div(d,100)

"""
    yyyymmdd2date(yyyymmdd::Number)

Takes in an Int-64 formatted date and returns a Date object in y-m-d format.
```@example
dt = 19980430
@show yyyymmdd2date(dt)
```
"""
function yyyymmdd2date(yyyymmdd::Number)
    yyyy=div(yyyymmdd,10000)
    mm=div(mod(yyyymmdd,10000),100)
    dd=mod(yyyymmdd,100)
    Date(yyyy,mm,dd)
end

"""
    yyyymm2date(yyyymm::Number)

Takes in a Int64-formatted year and date and returns a Date object,
with the last day of the month attached.

# Example
```@example
dt = 199804
@show yyyymm2date(dt)
```
"""
function yyyymm2date(yyyymm::Number)
    yyyy=div(yyyymm,100)
    mm=mod(yyyymm,100)
    Dates.lastdayofmonth(Date(yyyy,mm,1))
end

const dduuuyyyydateformat = Dates.DateFormat("dduuuyyyy")
"""
    dduuuyyyy2date(dduuuyyyy)

Takes in a string of the form "31Jan1997" and returns a Date
object in the format y-m-d.
```@example
d = "31Jan1997"
@show dduuuyyyy2date(d)
```
"""
function dduuuyyyy2date(dduuuyyyy)
    Date(dduuuyyyy,dduuuyyyydateformat)
end

const uuuddyyyydateformat = Dates.DateFormat("uuu dd, yyyy")
"""
    uuuddyyyy2date(uuuddyyyy)

Takes in a string of the form "Month Day, Year" and returns a Date in the
format y-m-d
```@example
u = "Jan 31, 1891"
@show uuuddyyyy2date(u)
```
"""
function uuuddyyyy2date(uuuddyyyy)
    Date(uuuddyyyy,uuuddyyyydateformat)
end

"""
    datetimestring2datetime(str,datetimeformat)

Takes in a string of the form "" and converts it to the user-specified
DateTime format.
```@example
using Dates #hide
str = "1989-07-07"
dt_fmt = DateFormat("y-m-d")
@show datetimestring2datetime(str, dt_fmt)
```
"""
function datetimestring2datetime(str,datetimeformat)
  try
    dt=DateTime(str,datetimeformat)
    y=Dates.year(dt)
    if 70<=y<=99
      y+=1900
    elseif 0<=y<70
      y+=2000
    end
    dt=Dates.DateTime(y,Dates.month(dt),Dates.day(dt),Dates.hour(dt),Dates.minute(dt),Dates.second(dt))
  catch
    missing
  end
end

"""
    date2yyyymmdd(date)

Takes in a Date object and converts it back to Int64.
```@example
using Dates #hide
dt = Date(2014,1,29)
@show date2yyyymmdd(dt)
```
"""
function date2yyyymmdd(date)
    try
        Dates.year(date)*10000+Dates.month(date)*100+Dates.day(date)
    catch
        missing
    end
end

"""
    date2yyyymm(date)

Takes in a Date object and transforms it back to Int64.
```@example
using Dates #hide
dt = Date(2010,8,11)
@show date2yyyymm(dt)
```
"""
function date2yyyymm(date)
    try
        Dates.year(date)*100+Dates.month(date)
    catch
        missing
    end
end

"""
    date2yyyy(date)

Takes in a Date object and returns the year as an Int64 type.
```@example
using Dates #hide
dt = Date(2011,7,24)
@show date2yyyy(dt)
```
"""
function date2yyyy(date)
    try
        Dates.year(date)
    catch
        missing
    end
end

"""
    datetimestring2yyyymmdd(str,datetimeformat)

Takes in a string and a DateTime format, converts the date back into Int64 type.
```@example
using Dates #hide
str = "1999-07-24"
dt_fmt = DateFormat("y-m-d")
@show datetimestring2yyyymmdd(str,dt_fmt)
```
"""
function datetimestring2yyyymmdd(str,datetimeformat)
    d = datetimestring2datetime(str,datetimeformat)
    if ismissing(d)
        missing
    else
        date2yyyymmdd(d)
    end
end

"""
    date2yyyyqq(d)

Takes in a Date object, formatted as y-m-d, and returns the year and quarter in Int64 type.
```@example
using Dates #hide
dt = Date(2009,6,22)
@show date2yyyyqq(dt)
```
"""
function date2yyyyqq(d)
  try
    date2yyyy(d)*100 + Dates.quarterofyear(d)
  catch
    missing
  end
end

"""
    date2yyyyww(d)

Takes in a Date object, formatted as y-m-d, and returns the year and
week in Int64 format.
```@example
using Dates #hide
dt = Date(2011,9,8)
@show dt
@show date2yyyyww(dt)
```
"""
function date2yyyyww(d)
  try
    date2yyyy(d)*100 + Dates.week(d)
  catch
    missing
  end
end

"""
    yearqtr2date(s, dlm="-Q")

convert something like 2018-Q3 into end of quarter date
`dlm` defines the string delimiter characters passed to split()
"""
function yearqtr2date(s, dlm="-Q")
  sp = split(s, dlm)
  y = parse(Int,sp[1])
  q = parse(Int,sp[2])
  Dates.lastdayofquarter(Date(y) + Dates.Month(3*(q-1)))
end

"""
    yearmonth2date(s, dlm="-")

convert something like 2018-11 into year month
`dlm` defines the string delimiter characters passed to split()
"""
function yearmonth2date(s, dlm="-")
  sp = split(s, dlm)
  y = parse(Int,sp[1])
  m = parse(Int,sp[2])
  Dates.lastdayofmonth(Date(y,m,1))
end


export group_by_type
export rectangular_lattice
export triangular_lattice
export wedge
export line
export remove_overlaps!
export remove_shape!
export nearest_neighbor_pairs

"""
    group_by_type(particles; ptype)

Returns a group of particles with type(s) `ptype`.
"""
function group_by_type(particles::Array{Particle, 1}; ptype::Union{Symbol, Array{Symbol, 1}})
    if !(ptype isa Array)
        ptype = [ptype]
    end

    pgroup = Array{Particle, 1}()
    for particle in particles
        if particle.ptype in ptype
            push!(pgroup, particle)
        end
    end
    return pgroup
end

"""
    rectangular_lattice(; s_x, s_y, M_x, M_y)

Generates a rectangular lattice with lattice constants `s_x, s_y` in the x, y
directions.  Each cell is duplicated `M_x, M_y` times in the x, y directions.
"""
function rectangular_lattice(; s_x::Float64, s_y::Float64, M_x::Int64, M_y::Int64)
    positions = Array{Array{Float64, 1}, 1}()
    for i = 0 : M_x - 1, j = 0 : M_y - 1
        push!(positions, [i * s_x, j * s_y])
    end
    return positions, M_x * s_x, M_y * s_y
end

"""
    triangular_lattice(; s, M_x, M_y)

Generates a triangular lattice with lattice constant `s`.  Each cell consists of
points `(0, 0)` and `(s / 2, s * sqrt(3) / 2)`.  The cells are duplicated
`M_x, M_y` times in the x, y directions.  Note that there are other
differently oriented triangular lattices.
"""
function triangular_lattice(; s::Float64, M_x::Int64, M_y::Int64)
    cell = [[0.0, 0.0], [s / 2, s * sqrt(3) / 2]]
    positions = Array{Array{Float64, 1}, 1}()
    for i = 0 : M_x - 1, j = 0 : M_y - 1
        push!(positions, cell[1] + [i * s, j * s * sqrt(3)])
        push!(positions, cell[2] + [i * s, j * s * sqrt(3)])
    end
    return positions, M_x * s, M_y * s * sqrt(3)
end

"""
    line(; from, to, spacing)

Generates a line of points from point `from` to point `to` with approximate
spacing `spacing`.  Note that the spacing is automatically adjusted slightly so
that points are equally spaced and include the start and end points.
"""
function line(; from::Array{Float64, 1}, to::Array{Float64, 1}, spacing::Float64)
    edge_length = sqrt((to[1] - from[1])^2 + (to[2] - from[2])^2)
    M = ceil(Int64, edge_length / spacing)

    edge = [from]
    for i = 1 : M
        push!(edge, from + i / M * (to - from))
    end
    return edge
end

"""
    wedge(; angle, edge_length, spacing, orientation, tip)

Generates a V-shaped wedge.  The angle of the wedge is `angle` with each arm a
length `edge_length`.  The points are separated by approximately `spacing` so
that the start and end points are included.  `orientation` is the angle relative
to the x-axis and tip is the x, y location of the tip of the wedge.
"""
function wedge(; angle::Float64, edge_length::Float64, spacing::Float64, orientation::Float64 = 0.0, tip::Array{Float64, 1})
    M = ceil(Int64, edge_length / spacing)
    approx_spacing = edge_length / M

    funnel = [tip]
    for i = 1 : M
        push!(funnel, tip + i * approx_spacing * [cos(angle / 2 + orientation), sin(angle / 2 + orientation)])
        push!(funnel, tip + i * approx_spacing * [cos(-angle / 2 + orientation), sin(-angle / 2 + orientation)])
    end
    return funnel
end

"""
    remove_overlaps!(positions; fixed_positions, minimum_distance, period_x, period_y)

Removes points from `positions` that are within a `minimum_distance` of points
in `fixed_positions`.  This is used to avoid large forces due to overlaps.
"""
function remove_overlaps!(positions::Array{Array{Float64, 1}, 1}; fixed_positions::Array{Array{Float64, 1}, 1}, minimum_distance::Float64, period_x::Float64 = -1.0, period_y::Float64 = -1.0)
    indices_to_remove = Array{Int64, 1}()
    for (n, (x, y)) in enumerate(positions)
        for (xx, yy) in fixed_positions
            Δx = wrap_displacement(x - xx; period = period_x)
            Δy = wrap_displacement(y - yy; period = period_y)

            if Δx^2 + Δy^2 < minimum_distance^2
                push!(indices_to_remove, n)
                break
            end
        end
    end
    deleteat!(positions, indices_to_remove)
end

"""
    remove_shape!(positions; shape, in_or_out)

Remove positions that are either inside (:in) or outside (:out) of a
parametric shape.
"""
function remove_shape!(positions::Array{Array{Float64, 1}, 1}; shape::Array{Array{Float64, 1}, 1}, in_or_out::Symbol)
    indices_to_remove = Array{Int64, 1}()
    N_shape = length(shape)
    for (n, (x, y)) in enumerate(positions)
        θ = 0.0
        for p = 1 : N_shape
            next_p = p % N_shape + 1
            Δx_1, Δy_1 = shape[p][1] - x, shape[p][2] - y
            Δx_2, Δy_2 = shape[next_p][1] - x, shape[next_p][2] - y
            cross = (Δx_1 * Δy_2 - Δx_2 * Δy_1) / sqrt((Δx_1^2 + Δy_1^2) * (Δx_2^2 + Δy_2^2))

            θ += asin(cross)
        end

        winding_number = round(θ / (2 * pi))
        if (winding_number == 0.0 && in_or_out == :out) || (winding_number != 0.0 && in_or_out == :in)
            push!(indices_to_remove, n)
        end
    end
    deleteat!(positions, indices_to_remove)
end

"""
    nearest_neighbor_pairs(particles; maximum_distance)

Returns pairs of particles that are within `maximum_distance` of eachother.
"""
function nearest_neighbor_pairs(particles::Array{Particle, 1}; maximum_distance::Float64 = 0.0, period_x::Float64 = -1.0, period_y::Float64 = -1.0)
    pairs = Array{NTuple{2, Particle}, 1}()

    N = length(particles)
    for i = 1 : N, j = i + 1 : N
        Δx = wrap_displacement(particles[i].x - particles[j].x; period = period_x)
        Δy = wrap_displacement(particles[i].y - particles[j].y; period = period_y)

        if Δx^2 + Δy^2 < maximum_distance^2
            push!(pairs, (particles[i], particles[j]))
        end
    end

    return pairs
end


# this function makes sure that the number of examples in each class are the same
function equalize_inputs(train_X::Array{Float32,2}, train_Y::Array{Int,2}, num_classes)
    train_X_len = size(train_X)[2]
    train_Y_len = size(train_Y)[2]
    @assert train_X_len == train_Y_len
    len = train_X_len

    class_counter = zeros(Int, num_classes)
    for i in 1:len
        class_counter[train_Y[1, i]+1] += 1
    end
    println("Total entries for each class: ", class_counter)
    max_per_class = minimum(class_counter)

    new_train_X = Array{Float32, 2}(undef, size(train_X)[1], min(train_X_len, max_per_class*num_classes))
    new_train_Y = Array{Float32, 2}(undef, size(train_Y)[1], min(train_Y_len, max_per_class*num_classes))

    counters = zeros(Int, num_classes)
    iter = 1
    new_counter = 1
    while (iter < len) && any(counters .< max_per_class) # Exits when all of the counters are more than the max limit
#         println(counters, new_counter)
        if counters[train_Y[1, iter]+1] < max_per_class
            counters[train_Y[1, iter]+1] += 1
            new_train_X[:, new_counter] = train_X[:, iter]
            new_train_Y[:, new_counter] = train_Y[:, iter]
            new_counter += 1
        end
        iter += 1
    end
    println("Counters: ", counters)
    println("Total entries added: ", new_counter)
    println("Total entries checked: ", iter)
    return new_train_X, new_train_Y
end

# Example usage:
# new_train_X, new_train_Y = equalize_inputs(train_X, train_Y, 2)


export chain_complex, isexact, homology, free_resolution

######################################################################
#
# Lift of homomorphisms
#
######################################################################
#=
  G
  | phi
  V
  F <- H
    psi
 and Im(phi) subset Im(psi), then G -> H can be constructed
=#

@doc Markdown.doc"""
    lift(phi::Map, psi::Map) -> Map
Given $\phi: G\to F$ and $\psi:H \to F$ s.th. $\Im(\phi) \subseteq \Im(\psi)$
return the map $G\to H$ to make the diagram commute.
"""
function lift(phi::Map, psi::Map)
  x = [haspreimage(psi, image(phi, g)) for g = gens(domain(phi))]
  @assert all(t -> t[1], x)
  return hom(domain(phi), domain(psi), [t[2] for t = x])
end

@doc Markdown.doc"""
    zero_map(G::GrpAbFinGen) -> Map
Create the map $G \to \{0\}$.
"""
function zero_map(G::GrpAbFinGen)
  Z = abelian_group([1])
  set_name!(Z, "Zero")
  return hom(G, Z, [Z[0] for i=1:ngens(G)])
end

######################################################################
#
# complex/ free resolution
#
######################################################################
function iszero(h::T) where {T <: Map{<:GrpAbFinGen, <:GrpAbFinGen}}
  return all(x -> iszero(h(x)), gens(domain(h)))
end

mutable struct ChainComplex{T}
  @declare_other
  maps::Array{<:Map, 1}
  direction::Symbol
  exact::Array{Bool, 1}
  function ChainComplex(A::S; check::Bool = true, direction:: Symbol = :left) where {S <:Array{<:Map{<:T, <:T}, 1}} where {T}
    if check
      @assert all(i-> iszero(A[i]*A[i+1]), 1:length(A)-1)
    end
    r = new{T}()
    r.maps = A
    r.direction = direction
    return r
  end
  function ChainComplex(X::Type, A::S; check::Bool = true, direction:: Symbol = :left) where {S <:Array{<:Map, 1}}
    if check
      @assert all(i-> iszero(A[i]*A[i+1]), 1:length(A)-1)
    end
    r = new{X}()
    r.maps = A
    r.direction = direction
    return r
  end

end

length(C::ChainComplex) = length(C.maps)
Base.map(C::ChainComplex, i::Int) = C.maps[i]
obj(C::ChainComplex, i::Int) = (i==0 ? domain(C.maps[1]) : codomain(C.maps[i]))

obj_type(C::ChainComplex{T}) where {T} = T
map_type(C::ChainComplex) = eltype(C.maps)

function show(io::IO, C::ChainComplex)
  @show_name(io, C)
  @show_special(io, C)

  Cn = get_special(C, :name)
  if Cn === nothing
    Cn = "C"
  end
  name_mod = String[]
  name_map = String[]
  mis_map = Tuple{Int, <:Map}[]
  mis_mod = Tuple{Int, <:Any}[]

  if C.direction == :left
    rng = 0:length(C)
    arr = ("--", "-->")
    dir = 1
  else
    rng = length(C)+1:-1:1
    arr = ("<--", "--")
    dir = 0
  end

  for i=1:length(C)
    phi = map(C, i)
    if get_special(phi, :name) !== nothing
      push!(name_map, get_special(phi, :name))
    else
      push!(name_map, "")
      push!(mis_map, (i, phi))
    end
  end
  for i=0:length(C)
    M = obj(C, i)
    if get_special(M, :name) !== nothing
      push!(name_mod, get_special(M, :name))
    else
      push!(name_mod, "$(Cn)_$i")
      push!(mis_mod, (i, M))
    end
  end

  io = IOContext(io, :compact => true)
  for i=rng
    if i == first(rng)
      print(io, name_mod[i+dir])
      continue
    end
    if name_map[i] != ""
      print(io, " ", arr[1], " ", name_map[i], " ", arr[2], " ", name_mod[i+dir])
    else
      print(io, " ", arr[1], arr[2], " ", name_mod[i+dir])
    end
  end
  if length(mis_mod) > 0 # || length(mis_map) > 0
    print(io, "\nwhere:\n")
    for (i, M) = mis_mod
      print(io, "\t$(Cn)_$i = ", M, "\n")
    end
#    for (i, phi) = mis_map
#      print(io, "\tphi_$i = ", phi, "\n")
#    end
  end
end

@doc Markdown.doc"""
    chain_complex(A::Map{GrpAbFinGen, GrpAbFinGen, <:Any, <:Any}...) -> ChainComplex{GrpAbFinGen}
Given maps $A_i$ s.th. $\Im(A_i) \subseteq \Kern(A_{i+1})$, this creates
the chain complex.
"""
function chain_complex(A::Map{GrpAbFinGen, GrpAbFinGen, <:Any, <:Any}...)
  return ChainComplex(collect(A))
end

function chain_complex(A::Array{<:Map{GrpAbFinGen, GrpAbFinGen, <:Any, <:Any}, 1})
  return ChainComplex(A)
end

Base.lastindex(C::ChainComplex) = length(C)
getindex(C::ChainComplex{T}, u::UnitRange) where {T} = ChainComplex(T, C.maps[u], check = false)

@doc Markdown.doc"""
    isexact(C::ChainComplex) -> Bool
Tests is the complex $A_i: G_i \to G_{i+1}$
is exact, ie. if $\Im(A_i) = \Kern(A_{i+1})$.
"""
function isexact(C::ChainComplex)
  return all(i->iseq(image(C.maps[i])[1], kernel(C.maps[i+1])[1]), 1:length(C)-1)
end

@doc Markdown.doc"""
    free_resolution(G::GrpAbFinGen) -> ChainComplex{GrpAbFinGen}
A free resultion for $G$, ie. a chain complex terminating in
$G \to \{0\}$ that is exact.
"""
function free_resolution(G::GrpAbFinGen)
  A = free_abelian_group(ngens(G))
  R = rels(G)
  B = free_abelian_group(nrows(R))
  h_A_G = hom(A, G, gens(G))
  h_B_A = hom(B, A, [GrpAbFinGenElem(A, R[i, :]) for i=1:ngens(B)])
  Z = abelian_group(Int[1])
  set_name!(Z, "Zero")
  return chain_complex(hom(Z, B, [B[0]]), h_B_A, h_A_G, hom(G, Z, [Z[0] for i = 1:ngens(G)]))
end

mutable struct ChainComplexMap{T} <: Map{ChainComplex{T}, ChainComplex{T}, HeckeMap, ChainComplexMap}
  header::MapHeader{ChainComplex{T}, ChainComplex{T}}
  maps::Array{<:Map{<:T, <:T}, 1}
  function ChainComplexMap(C::ChainComplex{T}, D::ChainComplex{T}, A::S; check::Bool = !true) where {S <: Array{<:Map{<:T, <:T}, 1}} where {T}
    r = new{T}()
    r.header = MapHeader(C, D)
    r.maps = A
    return r
  end
end

@doc Markdown.doc"""
    hom(C::ChainComplex{T}, D::ChainComplex{T}, phi::Map{<:T, <:T}) where {T} -> ChainComplexMap
Given chain complexes $C_i: G_i \to G_{i+1}$ and $D_i: H_i \to H_{i+1}$
as well as a map $\phi = \phi_n: G_n \to H_n$, lift $\phi$ to
the entire complex: $\phi_i: G_i \to H_i$ s.th. all squares commute.
"""
function hom(C::ChainComplex{T}, D::ChainComplex{T}, phi::Map{<:T, <:T}) where {T}
  @assert length(C) == length(D)
  @assert domain(C.maps[end]) == domain(phi)
  @assert domain(D.maps[end]) == codomain(phi)

  h = [phi]
  for i=length(C)-1:-1:1
    push!(h, lift(C.maps[i]*h[end], D.maps[i]))
  end
  return ChainComplexMap(C, D, reverse(h))
end

@doc Markdown.doc"""
    hom(C::ChainComplex{T}, G::T) -> ChainComplex{T}
Given a complex $A_i: G_i \to G_{i+1}$ and a module $G$,
compute the derived complex $\hom(G_i, G)$.
"""
function hom(C::ChainComplex{T}, G::T) where {T}
  A = map_type(C)[]
  H = [hom(domain(C.maps[1]), G)]
  H = vcat(H, [hom(codomain(f), G) for f = C.maps])

  R = map_type(C)[]
  for i=1:length(C)
    A = H[i+1][1] # hom(C_i+1, G)
    B = H[i][1]   # hom(C_i  , G)
    #need map from A -> B
    #   C.maps[i] : E -> D
    D = codomain(C.maps[i])
    E = domain(C.maps[i])
    #  H[2][i+1]: A -> Hom(D, G)
    #  H[2][i]  : B -> hom(E, G)
    g = elem_type(obj_type(C))[]
    for h = gens(A)
      phi = H[i+1][2](h) # D -> G
      psi = C.maps[i] * phi
      push!(g, preimage(H[i][2], psi))
    end
    push!(R, hom(A, B, g))
  end
  return ChainComplex(reverse(R))
end

@doc Markdown.doc"""
    hom(C::ChainComplex{T}, G::T) -> ChainComplex{T}
Given a complex $A_i: G_i \to G_{i+1}$ and a module $G$,
compute the derived complex $\hom(G, G_i)$.
"""
function hom(G::T, C::ChainComplex{T}) where {T}

  A = map_type(C)[]
  H = [hom(G, domain(C.maps[1]))]
  H = vcat(H, [hom(G, codomain(f)) for f = C.maps])

  R = map_type(C)[]
  for i=1:length(C)
    A = H[i+1][1] # hom(G, C_i+1)
    B = H[i][1]   # hom(G, C_i)
    #need map from A -> B
    #   C.maps[i] : E -> D
    D = codomain(C.maps[i])
    E = domain(C.maps[i])
    #  H[2][i+1]: A -> Hom(G, D)
    #  H[2][i]  : B -> hom(G, E)
    g = elem_type(obj_type(C))[]
    for h = gens(B)
      phi = H[i][2](h) # G -> E
      psi = phi * C.maps[i]
      push!(g, preimage(H[i+1][2], psi))
    end
    push!(R, hom(B, A, g))
  end
  return ChainComplex(R)
end

@doc Markdown.doc"""
    homology(C::ChainComplex{GrpAbFinGen}) -> Array{GrpAbFinGen, 1}
Given a complex $A_i: G_i \to G_{i+1}$,
compute the homology, ie. the modules $H_i = \Kern A_{i+1}/\Im A_i$
"""
function homology(C::ChainComplex)
  H = obj_type(C)[]
  for i=1:length(C)-1
    push!(H, quo(kernel(C.maps[i+1])[1], image(C.maps[i])[1])[1])
  end
  return H
end

function snake_lemma(C::ChainComplex{T}, D::ChainComplex{T}, A::Array{<:Map{T, T}, 1}) where {T}
  @assert length(C) == length(D) == 3
  @assert length(A) == 3
  @assert domain(A[1]) == obj(C,0) && codomain(A[1]) == obj(D, 1)
  @assert domain(A[2]) == obj(C,1) && codomain(A[2]) == obj(D, 2)
  @assert domain(A[3]) == obj(C,2) && codomain(A[3]) == obj(D, 3)

  ka, mka = kernel(A[1])
  kb, mkb = kernel(A[2])
  kc, mkc = kernel(A[3])
  ca, mca = cokernel(A[1])
  cb, mcb = cokernel(A[2])
  cc, mcc = cokernel(A[3])

  MT = map_type(C)
  res = MT(C)[]
  push!(res, MT(mka * map(C, 1) * inv(mkb)))
  push!(res, MT(mkb * map(C, 2) * inv(mkc)))
  #now the snake
  push!(res, MT(mkc * inv(map(C, 2)) * A[2] * inv(map(D, 2)) * mca))
  #and the boring rest
  push!(res, MT(inv(mca) * map(D, 2) * mcb))
  push!(res, MT(inv(mcb) * map(D, 3) * mcc))
  return chain_complex(res...)
end

@doc Markdown.doc"""
    tensor_product(C::ChainComplex{T}, G::T) -> ChainComplex{T}
Given a complex $A_i: G_i \to G_{i+1}$ and a module $G$,
compute the derived complex $G_i \otimes G$.
"""
function tensor_product(C::ChainComplex{T}, G::T) where {T}
  A = map_type(C)[]
  H = [tensor_product(domain(C.maps[1]), G, task  = :none)]
  H = vcat(H, [tensor_product(codomain(f), G, task = :none) for f = C.maps])

  R = map_type(C)[]
  I = identity_map(G)
  for i = 1:length(C)
    push!(R, hom(H[i], H[i+1], [C.maps[i], I]))
  end
  return chain_complex(R)
end


mutable struct LassoSCOP <: LassoRoot
    ## It's just a collection of data.

    A::Matrix # Feature matrix
    y::Matrix # label vector
    μ::Matrix # feature mean
    u::Number # label mean
    Z::Matrix # Starndardized Matrix
    l::Matrix # Zero mean label

    OptModel::Model  # The JuMP model for getting it right.
    λ::Float64       # the regularization parameter.

    LassoPath::Matrix # going along a fixed row is the fixing
    # the feature while varying the lambda quantity.
    λs::Vector # the lambda values.

    # Paramters for the iterator
    Tol::Float64
    λMin::Float64

    function LassoSCOP(A::Matrix, y::Union{Matrix, Vector}, λ::Float64=0.0)
        A = copy(A)
        m, _ = size(A)
        @assert size(y, 1) == m ""*
        "The rows of X should match of the columns of y, but the size of"*
        string("X, Y is: ", size(A), " ", size(y))

        y = copy(y)
        if ndims(y) == 1
            y = reshape(y, (length(y), 1))
        end
        μ = mean(A, dims=1)::Matrix
        u = mean(y)
        Z = A .- μ
        l = y .- u
        OptModel = MakeLassoOptimizationProblem(Z, l, λ)
        this = new(A, y, μ, u, Z, l, OptModel, λ)
        this.Tol = 1e-4
        this.λMin = 1e-8
        return this
    end
end

# ==============================================================================
# Type specific functions
# ==============================================================================

function Changeλ!(this::LassoSCOP, λ)
    """
        Change the Lasso regularizer of the current model, this mutates
        the state of the object.
    """
    model = this.OptModel
    x = model[:x]  # objects can be indexed with symbols!
    η = model[:η]
    y = this.l
    A = this.Z
    @objective(model, Min, λ*sum(η) + sum((A*x - y).^2))

end

function SolveLasso(
        this::LassoSCOP;
        x0::T=nothing
    )::Vector{Float64} where {T <: Union{Vector{Float64}, Nothing}}
    """
        Solve for the weights of the current model,
        given the current configuration of the model.
    """
    if !(x0 === nothing)
        setvalue.(this.OptModel[:x], x0) # warm start!
    end
    optimize!(this.OptModel)
    TernimationStatus = termination_status(this.OptModel)
    # @assert TernimationStatus == MOI.OPTIMAL "Terminated with non-optimal value when solving for x. "*
    # string("The status is: ", TernimationStatus)*"\n this is the results \n $(OptResults)"

    if !(TernimationStatus == MOI.OPTIMAL)
        Warn("\nWarning: convergence status for solver: $(TernimationStatus)")
        Warn("Current Value λ: $(this.λ)")
        PrintTitle("Here is the summary for the solution: ")
        display(solution_summary(this.OptModel))
    end
    return value.(this.OptModel[:x])
end

function Base.iterate(this::LassoSCOP)
    """
        Iterate through the parameters for the Lasso program:
            * λ
            * solutions
    """
    u = this.u
    A = this.Z
    y = this.l

    function λMax(A, y)
        ToMax = A'*(y .- u)
        ToMax *= 2
        return maximum(abs.(ToMax))
    end
    λ = λMax(A, y)
    Changeλ!(this, λ)
    x = SolveLasso(this)

    return (x, λ), (x, λ)
end

function Base.iterate(this::LassoSCOP, state::Tuple{Vector{Float64}, Float64})
    x, λ = state
    λ /= 2
    setvalue.(this.OptModel[:x], x)
    Changeλ!(this, λ)
    y = SolveLasso(this)

    if norm(x - y, Inf) <= this.Tol || λ <= this.λMin
        return nothing
    end

    return (y, λ), (y, λ)

end


using DiffEqBase, OrdinaryDiffEq, ApproxFun, DiffEqApproxFun
using Base.Test

S=Fourier()
u0=Fun(θ->cos(cos(θ-0.1))-cos(cos(0-0.1)),S)
c=Fun(cos,S)

prob = ODEProblem((u,p,t)->u''+(c+1)*u',u0,(0.,1.))

@time sol=solve(prob,Euler(),dt=1/1000)

tstops =  [0.0,0.00630546,0.0196241,0.0365809,0.0580536,0.0840865,0.1151,
           0.151529,0.195094,0.24658,0.292971,0.316851,0.338594,0.36075,
           0.384077,0.408548,0.433682,0.458922,0.483926,0.508619,0.533097,
           0.557498,0.581922,0.606414,0.630967,0.655551,0.68014,0.704717,
           0.729278,0.753828,0.778375,0.802923,0.827475,0.85203,0.876586,
           0.901142,0.925698,0.950253,0.974808,0.999362,1.0]

@time sol=solve(prob,Tsit5(),dt=1/40, adaptive=false, tstops=tstops)

#@time sol=solve(prob,Tsit5())
#@time sol=solve(prob,Rosenbrock23(),dt=1/40, adaptive=false, tstops=tstops)


# identity matrix function
identm(x) = [row == col ? true : false for row in 1:x, col in 1:x]


using Test
using ModelingToolkit, Graphs, DiffEqJump

import ModelingToolkit: value

#################################
#  testing for Jumps / all dgs
#################################
@parameters k1 k2
@variables t S(t) I(t) R(t)
j₁ = MassActionJump(k1, [0 => 1], [S => 1])
j₂ = MassActionJump(k1, [S => 1], [S => -1])
j₃ = MassActionJump(k2, [S => 1, I => 1], [S => -1, I => 1])
j₄ = MassActionJump(k2, [S => 2, R => 1], [R => -1])
j₅ = ConstantRateJump(k1 * I, [R ~ R + 1])
j₆ = VariableRateJump(k1 * k2 / (1 + t) * S, [S ~ S - 1, R ~ R + 1])
eqs = [j₁, j₂, j₃, j₄, j₅, j₆]
@named js = JumpSystem(eqs, t, [S, I, R], [k1, k2])
S = value(S);
I = value(I);
R = value(R);
k1 = value(k1);
k2 = value(k2);
# eq to vars they depend on
eq_sdeps = [Variable[], [S], [S, I], [S, R], [I], [S]]
eq_sidepsf = [Int[], [1], [1, 2], [1, 3], [2], [1]]
eq_sidepsb = [[2, 3, 4, 6], [3, 5], [4]]
deps = equation_dependencies(js)
@test all(i -> isequal(Set(eq_sdeps[i]), Set(deps[i])), 1:length(eqs))
depsbg = asgraph(js)
@test depsbg.fadjlist == eq_sidepsf
@test depsbg.badjlist == eq_sidepsb

# eq to params they depend on
eq_pdeps = [[k1], [k1], [k2], [k2], [k1], [k1, k2]]
eq_pidepsf = [[1], [1], [2], [2], [1], [1, 2]]
eq_pidepsb = [[1, 2, 5, 6], [3, 4, 6]]
deps = equation_dependencies(js, variables = parameters(js))
@test all(i -> isequal(Set(eq_pdeps[i]), Set(deps[i])), 1:length(eqs))
depsbg2 = asgraph(js, variables = parameters(js))
@test depsbg2.fadjlist == eq_pidepsf
@test depsbg2.badjlist == eq_pidepsb

# var to eqs that modify them
s_eqdepsf = [[1, 2, 3, 6], [3], [4, 5, 6]]
s_eqdepsb = [[1], [1], [1, 2], [3], [3], [1, 3]]
ne = 8
bg = BipartiteGraph(ne, s_eqdepsf, s_eqdepsb)
deps2 = variable_dependencies(js)
@test isequal(bg, deps2)

# eq to eqs that depend on them
eq_eqdeps = [[2, 3, 4, 6], [2, 3, 4, 6], [2, 3, 4, 5, 6], [4], [4], [2, 3, 4, 6]]
dg = SimpleDiGraph(6)
for (eqidx, eqdeps) in enumerate(eq_eqdeps)
    for eqdepidx in eqdeps
        add_edge!(dg, eqidx, eqdepidx)
    end
end
dg3 = eqeq_dependencies(depsbg, deps2)
@test dg == dg3

# var to vars that depend on them
var_vardeps = [[1, 2, 3], [1, 2, 3], [3]]
ne = 7
dg = SimpleDiGraph(3)
for (vidx, vdeps) in enumerate(var_vardeps)
    for vdepidx in vdeps
        add_edge!(dg, vidx, vdepidx)
    end
end
dg4 = varvar_dependencies(depsbg, deps2)
@test dg == dg4

#####################################
#       testing for ODE/SDEs
#####################################
@parameters k1 k2
@variables t S(t) I(t) R(t)
D = Differential(t)
eqs = [D(S) ~ k1 - k1 * S - k2 * S * I - k1 * k2 / (1 + t) * S,
    D(I) ~ k2 * S * I,
    D(R) ~ -k2 * S^2 * R / 2 + k1 * I + k1 * k2 * S / (1 + t)]
noiseeqs = [S, I, R]
@named os = ODESystem(eqs, t, [S, I, R], [k1, k2])
deps = equation_dependencies(os)
S = value(S);
I = value(I);
R = value(R);
k1 = value(k1);
k2 = value(k2);
eq_sdeps = [[S, I], [S, I], [S, I, R]]
@test all(i -> isequal(Set(eq_sdeps[i]), Set(deps[i])), 1:length(deps))

@parameters k1 k2
@variables t S(t) I(t) R(t)
@named sdes = SDESystem(eqs, noiseeqs, t, [S, I, R], [k1, k2])
deps = equation_dependencies(sdes)
@test all(i -> isequal(Set(eq_sdeps[i]), Set(deps[i])), 1:length(deps))

deps = variable_dependencies(os)
s_eqdeps = [[1], [2], [3]]
@test deps.fadjlist == s_eqdeps

#####################################
#       testing for nonlin sys
#####################################
@variables x y z
@parameters σ ρ β

eqs = [0 ~ σ * (y - x),
    0 ~ ρ - y,
    0 ~ y - β * z]
@named ns = NonlinearSystem(eqs, [x, y, z], [σ, ρ, β])
deps = equation_dependencies(ns)
eq_sdeps = [[x, y], [y], [y, z]]
@test all(i -> isequal(Set(deps[i]), Set(value.(eq_sdeps[i]))), 1:length(deps))


using Zeros
if VERSION < v"0.7-"
    using Base.Test
else
    using Test
end

# Real
Z = Zero()
@test Zero(0) === Z
@test Z === Zero()
@test Z == Z
@test Z == 0
@test sizeof(Z) == 0
@test float(Z) === 0.0
@test -Z === Z
@test +Z === Z
@test 2*Z === Z
@test 2.0*Z === Z
@test Z*3 === Z
@test Z/2 === Z
@test Z-Z === Z
@test Z+Z === Z
@test Z*Z === Z
@test 1-Z === 1
@test 1.0-Z === 1.0
@test Z-1 == -1
@test Z+1 == 1
@test 2+Z == 2
@test (Z < Z) == false
@test (Z > Z) == false
@test Z <= Z
@test Z >= Z
@test Z < 3
@test Z > -2.0
@test ldexp(Z, 3) === Z
@test copysign(Z, 3) === Z
@test copysign(Z, -1) === Z
@test flipsign(Z, -1) === Z
@test sign(Z) === Z
@test round(Z) === Z
@test floor(Z) === Z
@test ceil(Z) === Z
@test trunc(Z) === Z
@test significand(Z) === Z
@test !isodd(Z)
@test iseven(Z)
@test string(Z) == "0̸"
@test fma(Z,1,Z) === Z
@test muladd(Z,1,Z) === Z
@test fma(Z,1,3) === 3
@test muladd(Z,1,3) === 3
@test fma(Z,Z,Z) === Z
@test muladd(Z,Z,Z) === Z
@test mod(Z, 3) === Z
@test rem(Z, 3) === Z
@test modf(Z) === (Z, Z)

#Complex
@test Z*im === Z
@test im*Z === Z
@test real(Z) === Z
@test imag(Z) === Z
@test Z+(2+3im) === 2+3im
@test (2+3im)+Z === 2+3im
@test Z-(2+3im) === -2-3im
@test (2+3im)-Z === 2+3im
@test Z*(2+3im) === Z
@test (2.0+3.0im)*Z === Z
@test Z/(2+3im) === Z

@test Complex(Z,Z) == Z
@test Complex(1,Z) == 1
@test Complex(1.0,Z) == 1.0
@test Complex(true,Z) == true

# testzero()
@test testzero(3) === 3
@test testzero(3+3im) === 3+3im
@test testzero(0) === Z
@test testzero(0+0im) === Z

# Array functions
for cf in [(T)->T, (T)->Complex{T}]
    for T in [UInt8, UInt16, UInt32, UInt64, UInt128, Int8, Int16, Int32, Int64, Int128, BigInt, Float16, Float32, Float64, BigFloat]
        A = ones(cf(T),10)
        @test zero!(A) === A
        @test all(A .== zero(cf(T)))
    end
end

# Test error handling
@test_throws InexactError Zero(1)
@test_throws InexactError convert(Zero, 0.1)
@test_throws DivideError 1.0/Z
@test_throws DivideError (1.0+2.0im)/Z
@test_throws DivideError Z/Z

# Test `MyComplex` example type
include("mycomplex_example.jl")
@test MyImaginary(2)*MyImaginary(3) === MyReal(-6)


#this file tests whether using a for loop to apply a function of two variables
# f(x, y) to all pairs of entries of two vectors u, v is faster or if
#if using syntactic sugar f.(X, Y) is faster

g = (x, y)->sin(sqrt(x*y+10000))

x = collect(1:1:400)
y = collect(1:1:400)
yadj = y'

println("for loop time")
@time begin
    z = zeros(400, 400)
    for i = 1:400
        for j = 1:400
            z[i, j] = g(x[i], y[j])
        end
    end
end
println("syntactic sugar time")
@time w = g.(x, yadj)
10
println(norm(w-z))

#why is broadcasting slower T_T


# # Chapter 4
using BenchmarkTools

##Globals

p = 2

function pow_array(x::Vector{Float64})
  s = 0.0
    for y in x
      s = s + y^p
    end
    return s
end
t=rand(100000);

@btime pow_array($t)

@code_warntype pow_array(t)

### const

const p2 = 2
function pow_array2(x::Vector{Float64})
    s = 0.0
    for y in x
        s = s + y^p2
    end
    return s
end

 @btime pow_array2($t)

 @code_warntype pow_array2(t)

 function pow_array3(x::Vector{Float64})
    return pow_array_inner(x, p)
end

function pow_array_inner(x, pow)
   s = 0.0
   for y in x
      s = s + y^pow
   end
   return s
end

@btime pow_array3($t)

## Inlining

function f(x)
    a=x*5
    b=a+3
end

g(x) = f(2*x)

@code_typed g(3)

@code_llvm g(3)

function f(x)
    a=x*5
    b=a+3
    c=a-4
    if c < 0
        throw(DomainError())
    elseif c < 2
        d=c^3
    else
        d=c^2
    end
end

@inline function f_in(x)
    a=x*5
    b=a+3
    c=a-4
    if c < 0
        throw(DomainError())
    elseif c < 2
        d=c^3
    else
        d=c^2
    end
end

g_in(x) = f_in(2*x)

@noinline function f_ni(x)
    a=x*5
    b=a+3
end

g_ni(x) = f_ni(2*x)

@code_typed g_ni(3)

#Constant Propagation

sqr(x) = x * x

sqr2() = sqr(2)

@code_typed sqr2()

#Macros

function poly_naive(x, a...)
  p=zero(x)
  for i = 1:length(a)
    p = p + a[i] *  x^(i-1)
  end
  return p
end

f_naive(x) = poly_naive(x, 1,2,3,4,5)

x=3.5

@btime f_naive(3.5)

function poly_horner(x, a...)
  b=zero(x)
    for i = length(a):-1:1
      b = a[i] + b * x
    end
    return b
end

f_horner(x) = poly_horner(x, 1,2,3,4,5)

@btime f_horner($x)

macro horner(x, p...)
    ex = esc(p[end])
    for i = length(p)-1:-1:1
        ex = :(muladd(t, $ex, $(esc(p[i]))))
    end
    Expr(:block, :(t = $(esc(x))), ex)
end

f_horner_macro(x) = @horner(x, 1,2,3,4,5,6,7,8,9,10)

@btime f_horner_macro($x)

# Generated Functions

function prod_dim(x::Array{T, N}) where {T, N}
    s = 1
    for i = 1:N
        s = s * size(x, i)
    end
    return s
end

prod_dim(rand(10, 5, 5))

@generated function prod_dim_gen(x::Array{T, N}) where {T, N}
     ex = :(1)
     for i = 1:N
          ex = :(size(x, $i) * $ex)
     end
     return ex
end

prod_dim_gen(rand(10, 5, 5))

function prod_dim_gen_impl(x::Array{T, N}) where {T,N}
    ex = :(1)
    for i = 1:N
       ex = :(size(x, $i) * $ex)
    end
    return ex
end

#Keyword arguments
named_param(x; y=1, z=1) = x^y + x^z

pos_param(x,y,z) = x^y + x^z

@btime named_param(4; y=2, z=3)

@btime pos_param(4, 2, 3)


# ---
# title: 81. Search in Rotated Sorted Array II
# id: problem81
# author: Qling
# date: 2021-04-09
# difficulty: Medium
# categories: Array, Binary Search
# link: <https://leetcode.com/problems/search-in-rotated-sorted-array-ii/description/>
# hidden: true
# ---
#
# Suppose an array sorted in ascending order is rotated at some pivot unknown to
# you beforehand.
#
# (i.e., `[0,0,1,2,2,5,6]` might become `[2,5,6,0,0,1,2]`).
#
# You are given a target value to search. If found in the array return `true`,
# otherwise return `false`.
#
# **Example 1:**
#
#
#
#     Input: nums = [2,5,6,0,0,1,2], target = 0
#     Output: true
#
#
# **Example 2:**
#
#
#
#     Input: nums = [2,5,6,0,0,1,2], target = 3
#     Output: false
#
# **Follow up:**
#
#   * This is a follow up problem to [Search in Rotated Sorted Array](https://leetcode.com/problems/search-in-rotated-sorted-array/description/), where `nums` may contain duplicates.
#   * Would this affect the run-time complexity? How and why?
#
#
## @lc code=start
using LeetCode

function search_rotate(nums::Vector{Int}, target::Int)::Bool
    left, right = 1, length(nums)
    while left <= right
        mid = (left + right) >> 1
        (nums[mid] == target) && return true
        if nums[mid] == nums[left]
            left += 1
        elseif nums[mid] <= nums[right]
            if nums[mid] < target <= nums[right]
                left = mid + 1
            else
                right = mid - 1
            end
        else
            if nums[left] <= target < nums[mid]
                right = mid - 1
            else
                left = mid + 1
            end
        end
    end

    return false
end
## @lc code=end


using Primes;
using DelimitedFiles;
using SHA;

function primeHash(value::BigInt)
  _value = unsafe_wrap(Vector{UInt8}, Ptr{UInt8}(value.d), (value.size*8,))
  res = sha256(_value)
  bn = BigInt(0)
  ccall((:__gmpz_import, :libgmp), Cvoid, (Ref{BigInt}, Csize_t, Cint, Csize_t, Cint, Csize_t, Ptr{Cvoid}), bn, 4, 1, 8, 0, 0, res)
  while !isprime(bn)
    bn-=1
  end
  return bn
end

function primeN(bn::BigInt)
  while !isprime(bn)
    bn-=1
  end
  return bn
end

let
  randstate = zeros(UInt8,32)
  global bigrand
  ccall((:__gmp_randinit_default,:libgmp),Cvoid,(Ptr{Cvoid},),randstate)
  function rand(min::BigInt, max::BigInt)
    delta = max - min
    x = BigInt()
    ccall((:__gmpz_urandomm,:libgmp),Cvoid,(Ref{BigInt},Ptr{Cvoid},Ref{BigInt}),x,randstate,delta)
    return x+min
  end

  bigrand(max::BigInt) = bigrand(BigInt(0), max)
end

function normalize(rem::BigInt, divisor::BigInt)
  return powermod(rem, -1, divisor)
end

function test()
  mantice = rand(UInt8, 2^20) .% 2
  primes = map(x -> primeN(x<<10), BigInt(1):BigInt(2^20))
  println(length(unique(primes))-2^20)
end

test()


using InfiniteOpt, Ipopt, Plots
using KNITRO

# Parmeters
# Time span
t0 = 0
tf = 5
# Initial values
x10 = 0.7

x20 = 0.4

# Model
model = InfiniteModel(KNITRO.Optimizer)

# infinite_parameter
@infinite_parameter(model, t ∈ [t0, tf], num_supports = 101)

# variable
@variable(model, x1 ≥ 0, Infinite(t))
@variable(model, x2 ≥ 0, Infinite(t))

@variable(model, 0 ≤ u ≤ 1, Infinite(t), start = 0)

# objective
@objective(model, Max, ∫(log(x2), t))

# constraint
@constraint(model, x1(0) == x10)
@constraint(model, x2(0) == x20)
@constraint(model, x1_constr, ∂(x1, t) == u * x1)
@constraint(model, x2_constr, ∂(x2, t) == (1 - u) * x2)
@constraint(model, x1_min, x1 >= 0)
@constraint(model, x2_min, x2 >= 0)

# Optimization

print(model)

optimize!(model)

# Visulization
ts = value(t)
u = value(u)
x1 = value(x1)
x2 = value(x2)

plot(ts, u)
plot(ts, x1)
plot!(ts, x2)


using FractionalDiffEq
h=0.005; tf=(0, 30)
alpha = [0.99, 0.99, 0.99]
x0 = [1, 0, 1]

function f(du, u)
    a, b, c = 10, 28, 8/3
    du[1] = a*(u[2]-u[1])
    du[2] = u[1]*(b-u[3])-u[2]
    du[3] = u[1]*u[2]-c*u[3]
end
prob = FODESystem(f, alpha, x0, tf)
result = solve(prob, h, PIEX())

using Plots
plot(result[:, 1], result[:, 2])


function VannaVolga(putStrike, ATMStrike, callStrike, kStrike, kFlag, T, r, b, putVol, ATMVol, callVol)

# Reversal and BF

RRVol = callVol - putVol
BFVol = (callVol + putVol) / 2 - ATMVol

# Put

vega1 = Vega("c", ATMStrike, putStrike, T, r, b, ATMVol)
volga1 = DVegaDVol("c", ATMStrike, putStrike, T, r, b, ATMVol)
vanna1 = DDeltaDVol("c", ATMStrike, putStrike, T, r, b, ATMVol)

# ATM

vega2 = Vega("c", ATMStrike, ATMStrike, T, r, b, ATMVol)
volga2 = DVegaDVol("c", ATMStrike, ATMStrike, T, r, b, ATMVol)
vanna2 = DDeltaDVol("c", ATMStrike, ATMStrike, T, r, b, ATMVol)

# Call

vega3 = Vega("c", ATMStrike, callStrike, T, r, b, ATMVol)
volga3 = DVegaDVol("c", ATMStrike, callStrike, T, r, b, ATMVol)
vanna3 = DDeltaDVol("c", ATMStrike, callStrike, T, r, b, ATMVol)

# Any Strike K

vegaK = Vega("c", ATMStrike, kStrike, T, r, b, ATMVol)
volgaK = DVegaDVol("c", ATMStrike, kStrike, T, r, b, ATMVol)
vannaK = DDeltaDVol("c", ATMStrike, kStrike, T, r, b, ATMVol)

weight1 = (vegaK / vega1) * ((log(ATMStrike / kStrike) * log(callStrike / kStrike)) / (log(ATMStrike / putStrike) * log(callStrike / putStrike)))
weight2 = (vegaK / vega2) * ((log(kStrike / putStrike) * log(callStrike / kStrike)) / (log(ATMStrike / putStrike) * log(callStrike / ATMStrike)))
weight3 = (vegaK / vega3) * ((log(kStrike / putStrike) * log(kStrike / ATMStrike)) / (log(callStrike / putStrike) * log(callStrike / ATMStrike)))

o1 = weight1 * (GBlackScholes("c", ATMStrike, putStrike, T, r, b, putVol) - GBlackScholes("c", ATMStrike, putStrike, T, r, b, ATMVol))
o2 = weight2 * (GBlackScholes("c", ATMStrike, ATMStrike, T, r, b, ATMVol) - GBlackScholes("c", ATMStrike, ATMStrike, T, r, b, ATMVol))
o3 = weight3 * (GBlackScholes("c", ATMStrike, callStrike, T, r, b, callVol) - GBlackScholes("c", ATMStrike, callStrike, T, r, b, ATMVol))

premiumK = GBlackScholes(kFlag, ATMStrike, kStrike, T, r, b, ATMVol) + sum([o1, o2, o3])

impliedK = putVol * ((log(ATMStrike / kStrike) * log(callStrike / kStrike)) / (log(ATMStrike / putStrike) * log(callStrike / putStrike))) +
                ATMVol * ((log(kStrike / putStrike) * log(callStrike / kStrike)) / (log(ATMStrike / putStrike) * log(callStrike / ATMStrike))) +
                callVol * ((log(kStrike / putStrike) * log(kStrike / ATMStrike)) / (log(callStrike / putStrike) * log(callStrike / ATMStrike)))

[premiumK, impliedK]

end


#
#
#
export gradTV, gradTV2, gradIntrICTV12, forwardLogs
export gradDistance
@doc doc"""
    gradDistance(M,y,x[, p=2])

compute the (sub)gradient of the distance (squared)

```math
f(x) = \frac{1}{2} d^p_{\mathcal M}(x,y)
```

to a fixed [`MPoint`](@ref)` y` on the [`Manifold`](@ref) `M` and `p` is an
integer. The gradient reads

```math
  \nabla f(x) = -d_{\mathcal M}^{p-2}(x,y)\log_xy
```

for $p\neq 1$ or $x\neq  y$. Note that for the remaining case $p=1$,
$x=y$ the function is not differentiable. This function returns then the
[`zeroTVector`](@ref)`(M,x)`, since this is an element of the subdifferential.

# Optional

* `p` – (`2`) the exponent of the distance,  i.e. the default is the squared
  distance
"""
gradDistance(M,y,x,p::Int=2) = (p==2) ? -log(M,x,y) : -distance(M,x,y)^(p-2)*log(M,x,y)

@doc doc"""
    ∇u,∇v = gradIntrICTV12(M,f,u,v,α,β)

compute (sub)gradient of the intrinsic infimal convolution model using the mid point
model of second order differences, see [`costTV2`](@ref), i.e. for some $f\in\mathcal M$
on a [`Power`](@ref) manifold $\mathcal M$ this function computes the (sub)gradient of

```math
E(u,v) =
\frac{1}{2}\sum_{i\in\mathcal G} d_{\mathcal M}(g(\frac{1}{2},v_i,w_i),f_i)
+ \alpha
\bigl(
\beta\mathrm{TV}(v) + (1-\beta)\mathrm{TV}_2(w)
\bigr),
```
where both total variations refer to the intrinsic ones, [`gradTV`](@ref) and
[`gradTV2`](@ref), respectively.
"""
function gradIntrICTV12(M::mT,f::P,u::P,v::P,α,β) where {mT <: Manifold, P <: MPoint}
  c = midPoint(M,u,v,f)
  iL = log(M,c,f)
  return AdjDxGeo(M,u,v,1/2,iL) + α*β*gradTV(M,u), AdjDyGeo(M,u,v,1/2,iL) + α * (1-β) * gradTV2(M,v)
end
@doc doc"""
    gradTV(M,(x,y),[p=1])

compute the (sub) gradient of $\frac{1}{p}d^p_{\mathcal M}(x,y)$ with respect
to both $x$ and $y$.
"""
function gradTV(M::mT where {mT <: Manifold}, xT::Tuple{P,P} where {P <: MPoint}, p::Number=1)
  x = xT[1];
  y = xT[2];
  if p==2
      return (-log(M,x,y), -log(M,y,x))
  else
    d = distance(M,x,y);
    if d==0 # subdifferential containing zero
      return (zeroTVector(M,x),zeroTVector(M,y))
    else
      return (-log(M,x,y)/(d^(2-p)), -log(M,y,x)/(d^(2-p)))
    end
  end
end
@doc doc"""
    ξ = gradTV(M,λ,x,[p])
Compute the (sub)gradient $\partial F$ of all forward differences orrucirng,
in the power manifold array, i.e. of the function

$F(x) = \sum_{i}\sum_{j\in\mathcal I_i} d^p(x_i,x_j)$

where $i$ runs over all indices of the [`Power`](@ref) manifold `M`
and $\mathcal I_i$ denotes the forward neighbors of $i$.

# Input
* `M` – a [`Power`](@ref) manifold
* `x` – a [`PowPoint`](@ref).

# Ouput
* ξ – resulting tangent vector in $T_x\mathcal M$.
"""
function gradTV(M::Power,x::PowPoint,p::Int=1)::PowTVector
  R = CartesianIndices(M.powerSize)
  d = length(M.powerSize)
  maxInd = last(R)
  ξ = zeroTVector(M,x)
  c = costTV(M,x,p,0)
  for i in R # iterate over all pixel
    di = 0.
    for k in 1:d # for all direction combinations
      ek = CartesianIndex(ntuple(i  ->  (i==k) ? 1 : 0, d) ) #k th unit vector
      j = i+ek # compute neighbor
      if all( map(<=, j.I, maxInd.I)) # is this neighbor in range?
        if p != 1
          g = (c[i]==0 ? 1 : 1/c[i]) .* gradTV(M.manifold,(x[i],x[j]),p) # Compute TV on these
        else
          g = gradTV(M.manifold,(x[i],x[j]),p) # Compute TV on these
        end
        ξ[i] += g[1]
        ξ[j] += g[2]
      end
    end # directions
  end # i in R
  return ξ
end

@doc doc"""
    ξ = forwardLogs(M,x)

compute the forward logs $F$ (generalizing forward differences) orrucirng,
in the power manifold array, the function

```math
$F_i(x) = \sum_{j\in\mathcal I_i} \log_{x_i} x_j,\quad i \in \mathcal G,
```

where $\mathcal G$ is the set of indices of the [`Power`](@ref) manifold `M`
and $\mathcal I_i$ denotes the forward neighbors of $i$.

# Input
* `M` – a [`Power`](@ref) manifold
* `x` – a [`PowPoint`](@ref).

# Ouput
* `ξ` – resulting tangent vector in $T_x\mathcal M$ representing the logs, where
  $\mathcal N$ is thw power manifold with the number of dimensions added to `size(x)`.
"""
function forwardLogs(M::Power, x::PowPoint{P,Nt}) where {P <: MPoint, Nt}
  sX = size(x)
  R = CartesianIndices(sX)
  d = length(sX)
  maxInd = [last(R).I...] # maxInd as Array
  if d > 1
    d2 = fill(1,d+1)
    d2[d+1] = d
  else
    d2 = 1
  end
  N = Power(M.manifold, (prod(sX)*d,) )
  xT = PowPoint(repeat(getValue(x),inner=d2))
  ξ = zeroTVector(N,xT)
  for i in R # iterate over all pixel
    for k in 1:d # for all direction combinations
      I = [i.I...] # array of index
      J = I .+ 1 .* (1:d .== k) #i + e_k is j
      if all( J .<= maxInd ) # is this neighbor in range?
        j = CartesianIndex{d}(J...) # neigbbor index as Cartesian Index
        ξ[i,k] = log(M.manifold,x[i],x[j]) # Compute log and store in kth entry
      end
    end # directions
  end # i in R
  return ξ
end

@doc doc"""
    gradTV2(M,(x,y,z),p)

computes the (sub) gradient of $\frac{1}{p}d_2^p(x,y,z)$ with respect
to $x$, $y$, and $z$, where $d_2$ denotes the second order absolute difference
using the mid point model, i.e. let
```math
  \mathcal C = \bigl\{ c\in \mathcal M \ |\ g(\tfrac{1}{2};x_1,x_3) \text{ for some geodesic }g\bigr\}
```
denote the mid points between $x$ and $z$ on the manifold $\mathcal M$.
Then the absolute second order difference is defined as

```math
d_2(x,y,z) = \min_{c\in\mathcal C_{x,z}} d(c,y).
```

While the (sub)gradient with respect to $y$ is easy, the other two require
the evaluation of an [`adjointJacobiField`](@ref).
See [Illustration of the Gradient of a Second Order Difference](@ref secondOrderDifferenceGrad)
for its derivation.
"""
function gradTV2(M::mT where {mT <: Manifold}, xT::Tuple{P,P,P} where {P <: MPoint}, p::Number=1)
  x = xT[1];
  y = xT[2];
  z = xT[3];
  c = midPoint(M,x,z,y) # nearest mid point of x and z to y
  d = distance(M,y,c)
  innerLog = -log(M,c,y)
  if p==2
      return ( AdjDxGeo(M,x,z,1/2,innerLog), -log(M,y,c), AdjDyGeo(M,x,z,1/2,innerLog))
  else
    if d==0 # subdifferential containing zero
      return (zeroTVector(M,x),zeroTVector(M,y),zeroTVector(M,z))
    else
      return ( AdjDxGeo(M,x,z,1/2,innerLog/(d^(2-p))), -log(M,y,c)/(d^(2-p)), AdjDyGeo(M,x,z,1/2,innerLog/(d^(2-p))) )
    end
  end
end
@doc doc"""
    gradTV2(M,x [,p=1])

computes the (sub) gradient of $\frac{1}{p}d_2^p(x_1,x_2,x_3)$
with respect to all $x_1,x_2,x_3$ occuring along any array dimension in the
[`PowPoint`](@ref) `x`, where `M` is the corresponding [`Power`](@ref) manifold.
"""
function gradTV2(M::Power,x::PowPoint,p::Int=1)::PowTVector
  R = CartesianIndices(M.powerSize)
  d = length(M.powerSize)
  minInd, maxInd = first(R), last(R)
  ξ = zeroTVector(M,x)
  c = costTV2(M,x,p,false)
  for i in R # iterate over all pixel
    di = 0.
    for k in 1:d # for all direction combinations (TODO)
      ek = CartesianIndex(ntuple(i  ->  (i==k) ? 1 : 0, d) ) #k th unit vector
      jF = i+ek # compute forward neighbor
      jB = i-ek # compute backward neighbor
      if all( map(<=, jF.I, maxInd.I) ) && all( map(>=, jB.I, minInd.I)) # are neighbors in range?
        if p != 1
          g = (c[i] == 0 ? 1 : 1/c[i]) .* gradTV2(M.manifold,(x[jB],x[i],x[jF]),p) # Compute TV2 on these
        else
          g = gradTV2(M.manifold,(x[jB],x[i],x[jF]),p) # Compute TV2 on these
        end
        ξ[jB] += g[1]
        ξ[i] += g[2]
        ξ[jF] += g[3]
      end
    end # directions
  end # i in R
  return ξ
end


using Combinatorics

function masked(mask, lis)
    combos = []
    idx = 1
    for step in mask
        if(step < 1)
            push!(combos, Array{Int,1}[])
        else
            push!(combos, sort(lis[idx:idx+step-1]))
            idx += step
        end
    end
    Array{Array{Int, 1}, 1}(combos)
end

function orderedpartitions(mask)
    tostring(masklis) = replace("$masklis", r"Array{Int\d?\d?,1}|Int\d?\d?", "")
    join([tostring(lis) for lis in unique([masked(mask, p)
                        for p in permutations(1:sum(mask))])], "\n")
end

println(orderedpartitions([2, 0, 2]))
println(orderedpartitions([1, 1, 1]))


# TODO: better names for the functions

function updatecoeffs!{T<:Real}(coeff0, coeff1, coeff2, rhs, model, v, t::T, x,
                                a::T, Δτ::T, Δx::T)
    # TODO: should we just remove this dispatch and pass in ones(n-2)*a instead?
    # Updates
    # coeffn = values in linear system
    # rhs    = value function at previous timestep + f at current timestep

    # Input
    # model  = HJBOneDim object
    # v      = value function at previous timestep
    # t      = value of (forward) time
    # x      = vector of x-values
    # a      = policy-values on interior
    # Δτ     = time-step size
    # Δx     = spacial step length
    taux = Δτ/Δx
    htaux2 = 0.5*Δτ/Δx^2

    n = length(coeff0)
    for j = 2:n-1
        @inbounds begin
            bval = model.b(t,x[j],a)
            sval2 = model.σ(t,x[j],a)^2
            coeff1[j] = -(sval2*htaux2 + max(bval,0.)*taux)
            coeff2[j-1] = -(sval2*htaux2 - min(bval,0.)*taux)
            coeff0[j] = 1.-coeff1[j]-coeff2[j-1]
            rhs[j] = v[j] + Δτ*model.f(t,x[j],a)
        end
    end

end

function policytimestep(model::HJBOneDim,
                        v, avals, x, Δx, Δτ, ti::Int)
    taux = Δτ/Δx
    htaux2 = 0.5*Δτ/Δx^2
    t = (ti-1)*Δτ
    n = length(v)

    # TODO: redo this thing
    newind = ones(Int, n)
    @inbounds vnew = fill(-maxintfloat(typeof(v[1])), n)

    ind12 = zeros(Bool, length(vnew))

    coeff0 = ones(v)   # v_i
    coeff1 = zeros(n-1) # v_{i+1} # TODO: type stability
    coeff2 = zeros(n-1) # v_{i-1} # TODO: type stability
    rhs = zeros(v)
    # Dirichlet conditions
    if model.bcond[1] == true
        @inbounds rhs[1] = model.Dfun[1](t)
    end
    if model.bcond[2] == true
        @inbounds rhs[end] = model.Dfun[2](t)
    end

    for i = 1:length(avals)
        @inbounds begin
            a = avals[i]
            vold = vnew

            updatecoeffs!(coeff0, coeff1, coeff2, rhs, model, v, t, x, a, Δτ, Δx)
            Mat = spdiagm((coeff2, coeff0, coeff1), -1:1, n, n)
            # Move to sparse(I,J,K) instead?
            if model.bcond[1] == false
                bval = taux*model.b(t,x[1],a)
                sval2 = htaux2*model.σ(t,x[1],a)^2
                Mat[1,1:3] = [1.-(sval2-bval), -(bval-2*sval2), -sval2]
                rhs[1] = v[1] + Δτ*model.f(t,x[1],a)
            end

            if model.bcond[2] == false
                bval = taux*model.b(t,x[end],a)
                sval2 = htaux2*model.σ(t,x[end],a)^2
                Mat[end,end-2:end] = [-sval2, bval+2*sval2, 1.-(sval2+bval)]
                rhs[end] = v[end] + Δτ*model.f(t,x[end],a)
            end

            # TODO: Use Krylov solver for high-dimensional PDEs
            vnew = Mat\rhs

            ind12[:] = vold .> vnew
            vnew[ind12] = vold[ind12]
            newind[!ind12] = i
        end
    end

    @inbounds pol = avals[newind]

    return vnew, pol
end

function timeloopconstant(model::HJBOneDim, N::Int,
                          Δτ, vinit, avals, x, Δx)
    # Pass v and pol by reference?
    K = length(x)
    v = zeros(K, N+1)
    pol = zeros(K, N)

    @inbounds v[:,end] = vinit # We use forward time t instead of backward time τ

    for j = N:-1:1
        @inbounds begin
            # t = (j-1)*Δτ
            # TODO: pass v-column, pol-column by reference?
            v[:,j], pol[:,j] = policytimestep(model, v[:,j+1],
                                              avals, x, Δx, Δτ, j)
        end
    end

    return v, pol
end


# ---
# title: 553. Optimal Division
# id: problem553
# author: Indigo
# date: 2021-06-23
# difficulty: Medium
# categories: Math, String
# link: <https://leetcode.com/problems/optimal-division/description/>
# hidden: true
# ---
#
# Given a list of **positive integers** , the adjacent integers will perform the
# float division. For example, [2,3,4] -> 2 / 3 / 4.
#
# However, you can add any number of parenthesis at any position to change the
# priority of operations. You should find out how to add parenthesis to get the
# **maximum** result, and return the corresponding expression in string format.
# **Your expression should NOT contain redundant parenthesis.**
#
# **Example:**
#
#
#
#     Input: [1000,100,10,2]
#     Output: "1000/(100/10/2)"
#     Explanation:
#     1000/(100/10/2) = 1000/((100/10)/2) = 200
#     However, the bold parenthesis in "1000/( **(** 100/10 **)** /2)" are redundant,
#     since they don't influence the operation priority. So you should return "1000/(100/10/2)".
#
#     Other cases:
#     1000/(100/10)/2 = 50
#     1000/(100/(10/2)) = 50
#     1000/100/10/2 = 0.5
#     1000/100/(10/2) = 2
#
#
# **Note:**
#
#   1. The length of the input array is [1, 10].
#   2. Elements in the given array will be in range [2, 1000].
#   3. There is only one optimal division for each test case.
#
#
## @lc code=start
using LeetCode

function optimal_division(nums::Vector{Int})
    (length(nums) == 1) && return string(nums[1])
    (length(nums) == 2) && return string(nums[1]) * "/" * string(nums[1])
    res = "$(nums[1])/($(nums[2])"
    for i in 3:length(nums)
        res *= "/$(nums[i])"
    end
    return res * ")"
end
## @lc code=end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 2ad11ff8-b8c8-11eb-1c46-f3babe900e58
using Luxor

# ╔═╡ de5764b5-7203-41d5-94e1-438bcd8c75fb
md"Based on [inconvergent/differential](https://github.com/inconvergent/differential-line/tree/d951ed459a99fd87713b8b5a2aee06685119c007)"

# ╔═╡ 328c272f-4d85-4bb5-928a-8d95883498ba
function draw_points(points)
	for i in 1:length(points)
		next_i = i < length(points) ? i + 1 : 1
		line(points[i], points[next_i], :stroke)
		# circle(points[i], 1, :fill)
	end
end

# ╔═╡ 4b35feec-702d-402f-8015-5d77ea9a3b67
function attraction_step(points, farl, nearl, step)
	n_points = length(points)
	d_points = [Point(0, 0) for _ in 1:n_points]
	# for every point, attract and reject
	for i in 1:n_points
		next_i = i < n_points ? i + 1 : 1
		prev_i = i > 1 ? i - 1 : n_points

		v = points[i]

		# for every other point, attract or reject
		for j in 1:n_points
			w = points[j]
			diff = v - w
			norm = sqrt(v.x^2 + v.y^2)

			if j == next_i || j == prev_i
				# it is one of the neighbors
				# don't want to be too close
				if norm >= nearl
					# attract
					d_points[i] -= diff/norm*step
				end
			else
				# reject if not too far
				# closer points reject stronger
				if norm < farl
					d_points[i] += diff*(farl/norm-1)*step
				end
			end
		end
	end

	for i in 1:n_points
		points[i] += d_points[i]
	end
end

# ╔═╡ 4e291e01-6601-4102-8038-6f082e0ffedd
@draw begin
	n_init_points = 40
	radius = 100
	n_steps = 400
	draw_every = 10
	step = 10
	farl = 40
	nearl = 2
	setline(1)

	angles = sort(rand(n_init_points)*2pi)

	points = [Point(radius*cos(a), radius*sin(a)) for a in angles]

	draw_points(points)

	for i in 1:n_steps
		attraction_step(points, farl, nearl, step)
		if i % draw_every == 0
			draw_points(points)
		end
	end
end

# ╔═╡ Cell order:
# ╠═de5764b5-7203-41d5-94e1-438bcd8c75fb
# ╠═2ad11ff8-b8c8-11eb-1c46-f3babe900e58
# ╠═328c272f-4d85-4bb5-928a-8d95883498ba
# ╠═4b35feec-702d-402f-8015-5d77ea9a3b67
# ╠═4e291e01-6601-4102-8038-6f082e0ffedd