using EconPDEs, Distributions

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

    r::Float64

    # utility parameters
    ρ::Float64
    γ::Float64

    amin::Float64
    amax::Float64
end

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

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

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

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

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

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


# InverseGamma distribution

export InverseGamma

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

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

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

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

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

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

# @μσ_methods InverseGamma(shape)


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

using ProjectEulerSolutions

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

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

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

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

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

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

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

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

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

    return -1
end

p061 = Problems.Problem(p061solution)

Problems.benchmark(p061, 4, 4)


## Packages
using Roots, Calculus, Plots

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

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

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

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

    return xn1
end

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

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

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

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

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

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

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

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

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

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

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

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

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

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


using OrdinaryDiffEq

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

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

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

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

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


using Plots
using BenchmarkTools

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

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

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

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

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

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

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

# Checks
vbasic == v
kprimebasic == kprime

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

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

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


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

using ProjectEulerSolutions

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

p058 = Problems.Problem(p058solution)

Problems.benchmark(p058, 0.1)


include("../helpers.jl")

using Gadfly
using Color

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

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

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

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


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

xGrid = -5:0.01:5

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

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

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

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

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


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

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

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

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

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

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

end


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

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

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

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

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

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


"""
    Berlage(; <keyword arguments>)

Create a Berlage wavelet.

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

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

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

end


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

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

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

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


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

n = 500
μ = 1e-4

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return γ
end

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return iter, state
end

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

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

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

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

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

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

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

    return γ
end

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

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

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

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

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

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

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

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

    return cost
end

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

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

    return cost
end

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

# Fields
$(FIELDS)

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

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

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

end

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

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

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

    x_i = similar(X)
    x_i .= X

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

    iters = 0
    converged = false

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

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

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

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

        conv_measure = norm(x_i .- X, 2)

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

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

        if conv_measure < abstol
            converged = true

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

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

        else
            x_i .= X
        end
    end

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


using LinearAlgebra: diagm
using SparseArrays: sparse

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

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

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

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

"""
    graph_filter(ind, order)

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

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

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

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

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

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

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


module ComplexBigMatrices

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

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

using LinearAlgebra
using GenericLinearAlgebra
using GenericSchur

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

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

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

end


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

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

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

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

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

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

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

    optimize!(mod)

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

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

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

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

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

    for t=1:maxEpochs

    mod1=Model(Mosek.Optimizer)

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

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

    U_t=value.(U)

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

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

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

    end

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

    return V_t_best*V_t_best', ofv_final
end


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

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


function chebyshev_weights(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  poly = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order[i],normalize_node(nodes[i],domain[:,i]))
  end

  weights = Array{T,N}(undef,Tuple(order.+1))

  @inbounds for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      product = one(T)
      @inbounds for j = 1:N
        product *= poly[j][s[j],i[j]]
      end

      numerator   += f[s]*product
      denominator += product^2

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  n = size(f)

  poly = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order[i],normalize_node(nodes[i],domain[:,i]))
  end

  weights = Array{T,N}(undef,Tuple(order.+1))

  @inbounds for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      num = f[s]
      den = one(T)
      @inbounds for j = 1:N
        if s[j] === 1 || s[j] === n[j]
          scale = 0.5
        else
          scale = 1.0
        end
        temp = poly[j][s[j],i[j]]
        num *= temp*scale
        den *= (temp^2)*scale
      end

      numerator   += num
      denominator += den

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  poly       = Array{Array{T,2},1}(undef,N)
  complement = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i]       = chebyshev_polynomial(order[i],normalize_node(nodes[i],domain[:,i]))
    complement[i] = pinv(poly[i])'
  end

  weights = Array{T,N}(undef,Tuple(order.+1))

  @inbounds for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      num = f[s]
      den = one(T)
      @inbounds for j = 1:N
        temp = complement[j][s[j],i[j]]
        num *= temp
        den *= temp*poly[j][s[j],i[j]]
      end

      numerator   += num
      denominator += den

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  weights = Array{T,N}(undef,Tuple(order.+1))

  @inbounds for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      product = one(T)
      @inbounds for j = 1:N
        product *= poly[j][s[j],i[j]]
      end

      numerator   += f[s]*product
      denominator += product^2

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  n = size(f)

  weights = Array{T,N}(undef,Tuple(order.+1))

  @inbounds for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      num = f[s]
      den = one(T)
      @inbounds for j = 1:N
        if s[j] == 1 || s[j] == n[j]
          scale = 0.5
        else
          scale = 1.0
        end
        temp = poly[j][s[j],i[j]]
        num *= temp*scale
        den *= (temp^2)*scale
      end

      numerator   += num
      denominator += den

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  complement = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    complement[i] = pinv(poly[i])'
  end

  weights = Array{T,N}(undef,Tuple(order.+1))

  @inbounds for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      num = f[s]
      den = one(T)
      @inbounds for j = 1:N
        temp = complement[j][s[j],i[j]]
        num *= temp
        den *= temp*poly[j][s[j],i[j]]
      end

      numerator   += num
      denominator += den

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  poly = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order,normalize_node(nodes[i],domain[:,i]))
  end

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        product = one(T)
        @inbounds for j = 1:N
          product *= poly[j][s[j],i[j]]
        end

        numerator   += f[s]*product
        denominator += product^2

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  n = size(f)

  poly = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order,normalize_node(nodes[i],domain[:,i]))
  end

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        num = f[s]
        den = one(T)
        @inbounds for j = 1:N
          if s[j] == 1 || s[j] == n[j]
            scale = 0.5
          else
            scale = 1.0
          end
          temp = poly[j][s[j],i[j]]
          num *= temp*scale
          den *= (temp^2)*scale
        end

        numerator   += num
        denominator += den

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  poly       = Array{Array{T,2},1}(undef,N)
  complement = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i]       = chebyshev_polynomial(order,normalize_node(nodes[i],domain[:,i]))
    complement[i] = pinv(poly[i])'
  end

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        num = f[s]
        den = one(T)
        @inbounds for j = 1:N
          temp = complement[j][s[j],i[j]]
          num *= temp
          den *= temp*poly[j][s[j],i[j]]
        end

        numerator   += num
        denominator += den

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::S) where {T<:AbstractFloat,N,S<:Integer}

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        product = one(T)
        @inbounds for j = 1:N
          product *= poly[j][s[j],i[j]]
        end

        numerator   += f[s]*product
        denominator += product^2

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::S) where {T<:AbstractFloat,N,S<:Integer}

  n = size(f)

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        num = f[s]
        den = one(T)
        @inbounds for j = 1:N
          if s[j] == 1 || s[j] == n[j]
            scale = 0.5
          else
            scale = 1.0
          end
          temp = poly[j][s[j],i[j]]
          num *= temp*scale
          den *= (temp^2)*scale
        end

        numerator   += num
        denominator += den

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::S) where {T<:AbstractFloat,N,S<:Integer}

  complement = Array{Array{T,2},1}(undef,N)
  @inbounds for i = 1:N
    complement[i] = pinv(poly[i])'
  end

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        num = f[s]
        den = one(T)
        @inbounds for j = 1:N
          temp = complement[j][s[j],i[j]]
          num *= temp
          den *= temp*poly[j][s[j],i[j]]
        end

        numerator   += num
        denominator += den

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_threaded(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  poly = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order[i],normalize_node(nodes[i],domain[:,i]))
  end

  weights = zeros(Tuple(order.+1))

  @inbounds @sync @qthreads for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      product = one(T)
      @inbounds for j = 1:N
        product *= poly[j][s[j],i[j]]
      end

      numerator   += f[s]*product
      denominator += product^2

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_extrema_threaded(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  n = size(f)

  poly = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order[i],normalize_node(nodes[i],domain[:,i]))
  end

  weights = zeros(Tuple(order.+1))

  @inbounds @sync @qthreads for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      num = f[s]
      den = one(T)
      @inbounds for j = 1:N
        if s[j] == 1 || s[j] == n[j]
          scale = 0.5
        else
          scale = 1.0
        end
        temp = poly[j][s[j],i[j]]
        num *= temp*scale
        den *= (temp^2)*scale
      end

      numerator   += num
      denominator += den

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_extended_threaded(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  poly       = Array{Array{T,2},1}(undef,N)
  complement = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i]       = chebyshev_polynomial(order[i],normalize_node(nodes[i],domain[:,i]))
    complement[i] = pinv(poly[i])'
  end

  weights = zeros(Tuple(order.+1))

  @inbounds @sync @qthreads for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      num = f[s]
      den = one(T)
      @inbounds for j = 1:N
        temp = complement[j][s[j],i[j]]
        num *= temp
        den *= temp*poly[j][s[j],i[j]]
      end

      numerator   += num
      denominator += den

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_threaded(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  weights = Array{T,N}(undef,Tuple(order.+1))

  @inbounds @sync @qthreads for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      product = one(T)
      @inbounds for j = 1:N
        product *= poly[j][s[j],i[j]]
      end

      numerator   += f[s]*product
      denominator += product^2

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_extrema_threaded(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  n = size(f)

  weights = Array{T,N}(undef,Tuple(order.+1))

  @inbounds @sync @qthreads for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      num = f[s]
      den = one(T)
      @inbounds for j = 1:N
        if s[j] == 1 || s[j] == n[j]
          scale = 0.5
        else
          scale = 1.0
        end
        temp = poly[j][s[j],i[j]]
        num *= temp*scale
        den *= (temp^2)*scale
      end

      numerator   += num
      denominator += den

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_extended_threaded(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  complement = Array{Array{T,2},1}(undef,N)
  @inbounds for i = 1:N
    complement[i] = pinv(poly[i])'
  end

  weights = Array{T,N}(undef,Tuple(order.+1))

  @inbounds @sync @qthreads for i in CartesianIndices(weights)

    numerator   = zero(T)
    denominator = zero(T)

    @inbounds for s in CartesianIndices(f)

      num = f[s]
      den = one(T)
      @inbounds for j = 1:N
        temp = complement[j][s[j],i[j]]
        num *= temp
        den *= temp*poly[j][s[j],i[j]]
      end

      numerator   += num
      denominator += den

    end

    weights[i] = numerator/denominator

  end

  return weights

end

function chebyshev_weights_threaded(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  poly = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order,normalize_node(nodes[i],domain[:,i]))
  end

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds @sync @qthreads for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        product = one(T)
        @inbounds for j = 1:N
          product *= poly[j][s[j],i[j]]
        end

        numerator   += f[s]*product
        denominator += product^2

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_extrema_threaded(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  n = size(f)

  poly = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i] = chebyshev_polynomial(order,normalize_node(nodes[i],domain[:,i]))
  end

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds @sync @qthreads for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        num = f[s]
        den = one(T)
        @inbounds for j = 1:N
          if s[j] == 1 || s[j] == n[j]
            scale = 0.5
          else
            scale = 1.0
          end
          temp = poly[j][s[j],i[j]]
          num *= temp*scale
          den *= (temp^2)*scale
        end

        numerator   += num
        denominator += den

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_extended_threaded(f::AbstractArray{T,N},nodes::NTuple{N,Array{T,1}},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  poly       = Array{Array{T,2},1}(undef,N)
  complement = Array{Array{T,2},1}(undef,N)

  @inbounds for i = 1:N
    poly[i]       = chebyshev_polynomial(order,normalize_node(nodes[i],domain[:,i]))
    complement[i] = pinv(poly[i])'
  end

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds @sync @qthreads for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        num = f[s]
        den = one(T)
        @inbounds for j = 1:N
          temp = complement[j][s[j],i[j]]
          num *= temp
          den *= temp*poly[j][s[j],i[j]]
        end

        numerator   += num
        denominator += den

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_threaded(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::S) where {T<:AbstractFloat,N,S<:Integer}

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds @sync @qthreads for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        product = one(T)
        @inbounds for j = 1:N
          product *= poly[j][s[j],i[j]]
        end

        numerator   += f[s]*product
        denominator += product^2

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_extrema_threaded(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::S) where {T<:AbstractFloat,N,S<:Integer}

  n = size(f)

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds @sync @qthreads for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        num = f[s]
        den = one(T)
        @inbounds for j = 1:N
          if s[j] == 1 || s[j] == n[j]
            scale = 0.5
          else
            scale = 1.0
          end
          temp = poly[j][s[j],i[j]]
          num *= temp*scale
          den *= (temp^2)*scale
        end

        numerator   += num
        denominator += den

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights_extended_threaded(f::AbstractArray{T,N},poly::NTuple{N,Array{T,2}},order::S) where {T<:AbstractFloat,N,S<:Integer}

  complement = Array{Array{T,2},1}(undef,N)
  @inbounds for i = 1:N
    complement[i] = pinv(poly[i])'
  end

  ord = (order,)
  for i = 2:N
    ord = (ord...,order)
  end

  weights = Array{T,N}(undef,ord.+1)

  @inbounds @sync @qthreads for i in CartesianIndices(weights)
    if sum(Tuple(i)) <= order+N

      numerator   = zero(T)
      denominator = zero(T)

      @inbounds for s in CartesianIndices(f)

        num = f[s]
        den = one(T)
        @inbounds for j = 1:N
          temp = complement[j][s[j],i[j]]
          num *= temp
          den *= temp*poly[j][s[j],i[j]]
        end

        numerator   += num
        denominator += den

      end

      weights[i] = numerator/denominator

    else
      weights[i] = zero(T)
    end

  end

  return weights

end

function chebyshev_weights(cheb::ChebInterpRoots)

  weights = chebyshev_weights_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  return weights

end

function chebyshev_weights(cheb::ChebInterpExtrema)

  weights = chebyshev_weights_extrema_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  return weights

end

function chebyshev_weights(cheb::ChebInterpExtended)

  weights = chebyshev_weights_extended_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  return weights

end

function chebyshev_weights(cheb::ChebInterpVertesi)

  weights = chebyshev_weights_vertesi_threaded(cheb.data,cheb.nodes,cheb.order,cheb.domain)

  return weights

end

# Functions for the one-variable case where the nodes are a vector

function chebyshev_weights(f::AbstractArray{T,1},nodes::Array{T,1},order::Union{S,Array{S,1}},domain=[one(T);-one(T)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(nodes,),order,domain)

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,1},nodes::Array{T,1},order::Union{S,Array{S,1}},domain=[one(T);-one(T)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights_extrema(f,(nodes,),order,domain)

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,1},nodes::Array{T,1},order::Union{S,Array{S,1}},domain=[one(T);-one(T)]) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights_extended(f,(nodes,),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,1},poly::Array{T,2},order::Union{S,Array{S,1}}) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights(f,(poly,),order)

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,1},poly::Array{T,2},order::Union{S,Array{S,1}}) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights_extrema(f,(poly,),order)

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,1},poly::Array{T,2},order::Union{S,Array{S,1}}) where {T<:AbstractFloat,S<:Integer}

  weights = chebyshev_weights_extended(f,(poly,),order)

  return weights

end

# Functions that allow the nodes to be in an array of arrays

function chebyshev_weights(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extrema(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extended(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extrema(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extended(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extrema(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extended(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::S) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights_extrema(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::S) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extrema(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights_extended(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::S) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extended(f,tuple(poly...),order)

  return weights

end

# Threaded functions

function chebyshev_weights_threaded(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_threaded(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_extrema_threaded(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extrema_threaded(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_extended_threaded(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::Array{S,1},domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extended_threaded(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_threaded(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_threaded(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights_extrema_threaded(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extrema_threaded(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights_extended_threaded(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::Array{S,1}) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extended_threaded(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights_threaded(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_threaded(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_extrema_threaded(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extrema_threaded(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_extended_threaded(f::AbstractArray{T,N},nodes::Array{Array{T,1},1},order::S,domain=[ones(T,1,N);-ones(T,1,N)]) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extended_threaded(f,tuple(nodes...),order,domain)

  return weights

end

function chebyshev_weights_threaded(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::S) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_threaded(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights_extrema_threaded(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::S) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extrema_threaded(f,tuple(poly...),order)

  return weights

end

function chebyshev_weights_extended_threaded(f::AbstractArray{T,N},poly::Array{Array{T,2},1},order::S) where {T<:AbstractFloat,N,S<:Integer}

  weights = chebyshev_weights_extended_threaded(f,tuple(poly...),order)

  return weights

end

const chebyshev_weights_vertesi          = chebyshev_weights_extended
const chebyshev_weights_vertesi_threaded = chebyshev_weights_extended_threaded


"""
type-II discrete cosine transform (DCT)
"""
function dct!(data::Array{T, 3}) where {T <:AbstractFloat}
	dctmat = genmat(size(data, 1))
	dctmat1 = genmat(size(data, 2))'
	@views @inbounds @fastmath for i in 1:size(data, 3)
		data[:, :, i] .= dctmat * data[:, :, i] * dctmat1
	end

	dctmat = genmat(size(data, 3))
	@views @inbounds @fastmath for j in 1:size(data, 2), i in 1:size(data, 1)
		data[i, j, :] .= dctmat * data[i, j, :]
	end
end

"""
inverse discrete cosine transform (DCT)
"""
function idct!(data::Array{T, 3}) where {T <:AbstractFloat}
	dctmat = genmat(size(data, 3))'
	@views @inbounds @fastmath for j in 1:size(data, 2), i in 1:size(data, 1)
		data[i, j, :] .= dctmat * data[i, j, :]
	end

	dctmat = genmat(size(data, 1))'
	dctmat1 = genmat(size(data, 2))
	@views @inbounds @fastmath for i in 1:size(data, 3)
		data[:, :, i] .= dctmat * data[:, :, i] * dctmat1
	end
end

function genmat(h::Int)::Array{Float64, 2}
	mat = Array{Float64, 2}(undef, h, h)
	mat[1, :] .= sqrt(1 / h)
	@inbounds @fastmath for j in 1:h
		@simd for i in 2:h
			mat[i, j] = cos(π * (i - 1) * (j - 0.5) / h) * sqrt(2 / h)
		end
	end
	mat
end


using DelimitedFiles

input = readdlm(ARGS[1], ' ', Int, '\n');

# Part 1
fuel_required(mass) = (mass ÷ 3) - 2;
part1 = reduce(+, map(fuel_required, input))
println("Part 1: ", part1)

# Part 2
function fuel_required_2(mass)
    fuel = fuel_required(mass)
    return (fuel > 0) ? fuel + fuel_required_2(fuel) : 0;
end
part2 = reduce(+, map(fuel_required_2, input))
println("Part 2: ", part2)

# Nothing too crazy here for those familiar with recursion, although reading the instructions carefully always helps.
# Notice something interesting: the part 2 function does not distribute. That is, you cannot compute
#   fuel_required_2(part1) + part1
# to get part2.
# https://twitter.com/wjholdentech/status/1201023263633219585


# Auxiliary functions for natural cubic spline interpolation

# Stores parameters for fitted spline polynomials
mutable struct Spline{T}
    x::Vector{T} # (x, y) are the original points
    y::Vector{Float64}
    params::Vector{Float64} # stores parameters for polynomials. [ a1, b1, c1, d1, a2, b2, c2, d2 ...]

    function Spline(x::Vector{T}, y::Vector{Float64}, params::Vector{Float64}) where {T}
        polynms_count = length(x) - 1
        @assert length(params) == polynms_count * 4 "params length $(length(params)) does not conform to the expected number of polynomials ($(polynms_count))" # each polynomial has 4 parameters
        return new{T}(x, y, params)
    end
end

# Given a polinomial index, returns the indexer of the last parameter before the first parameter of the referenced polinomial
_base_param_index_(poly_index::Int) = (poly_index-1)*4

# Performs natural cubic spline interpolation
function splineint(s::Spline{T}, x_out::Number) where {T<:Number}
    local poly_index::Int = 1
    local base_idx::Int

    if x_out > s.x[end]
        # Extrapolation after last point
        poly_index = length(s.x) - 1
        base_idx = _base_param_index_(poly_index)
        return (x_out - s.x[end])*(s.params[base_idx + 2] + 2*s.params[base_idx + 3]*s.x[end] + 3*s.params[base_idx + 4]*(s.x[end]^2)) + s.y[end]
    elseif x_out < s.x[1]
        # Extrapolation before first point
        base_idx = 0
        return (x_out - s.x[1])*(s.params[base_idx + 2] + 2*s.params[base_idx + 3]*s.x[1] + 3*s.params[base_idx + 4]*(s.x[1]^2)) + s.y[1]
    else
        # Interplation
        while x_out > s.x[poly_index+1]
            poly_index += 1
        end

        base_idx = _base_param_index_(poly_index)

        #P1                       P2                    P3                     ...
        #1   2    3      4        5   6    7      8     9   10   11     12
        #a + bx + cx^2 + dx^3     a + bx + cx^2 + dx^3  a + bx + cx^2 + dx^3   ...
        return s.params[base_idx + 1] + s.params[base_idx + 2]*x_out +
            s.params[base_idx + 3]*(x_out^2) + s.params[base_idx + 4]*(x_out^3)
    end
end

# Performs natural cubic spline interpolation
function splineint(s::Spline{T1}, x_out::Vector{T2}) :: Vector{Float64} where {T1, T2}
    len = length(x_out)
    y_out = Vector{Float64}(undef, len)
    for i in 1:len
        y_out[i] = splineint(s, x_out[i])
    end
    return y_out
end

# Build a Spline object by fitting 3rd order polynomials around points (x_in, y_in)
function splinefit(x_in::Vector{T}, y_in::Vector{Float64}) :: Spline{T} where {T}
    #
    # TODO: optimize. See http://www.math.ntnu.no/emner/TMA4215/2008h/cubicsplines.pdf
    #
    points_count = length(x_in)
    @assert points_count == length(y_in) "x_in and y_in doesn't conform on sizes."
    matrix_n = 4*(points_count-1) # the main matrix is a square matrix matrix_n by matrix_n

    A = zeros(matrix_n, matrix_n)
    b = zeros(matrix_n)

    # Known values for polys

    # First Point
    A[1,1] = 1.0
    A[1,2] = x_in[1]
    A[1,3] = x_in[1]^2
    A[1,4] = x_in[1]^3
    b[1] = y_in[1]

    # Last Point
    base_idx = 4*(points_count - 2)
    A[2, base_idx + 1] = 1.0
    A[2, base_idx + 2] = x_in[points_count]
    A[2, base_idx + 3] = x_in[points_count]^2
    A[2, base_idx + 4] = x_in[points_count]^3
    b[2] = y_in[points_count]

    # Inner Points
    row = 3
    for i in 2:(points_count-1)
        # Connecting to left poly
        A[row, (i-2)*4 + 1] = 1.0
        A[row, (i-2)*4 + 2] = x_in[i]
        A[row, (i-2)*4 + 3] = x_in[i]^2
        A[row, (i-2)*4 + 4] = x_in[i]^3
        b[row] = y_in[i]

        row += 1

        # Connecting to right poly
        A[row, (i-1)*4 + 1] = 1.0
        A[row, (i-1)*4 + 2] = x_in[i]
        A[row, (i-1)*4 + 3] = x_in[i]^2
        A[row, (i-1)*4 + 4] = x_in[i]^3
        b[row] = y_in[i]

        row += 1

        # Conditions on first order derivatives
        A[row, (i-2)*4 + 2] = 1.0
        A[row, (i-2)*4 + 3] = 2.0*x_in[i]
        A[row, (i-2)*4 + 4] = 3.0*(x_in[i]^2)
        A[row, (i-1)*4 + 2] = -1.0
        A[row, (i-1)*4 + 3] = -2.0*x_in[i]
        A[row, (i-1)*4 + 4] = -3.0*(x_in[i]^2)

        row += 1

        # Conditions on second order derivatives
        A[row, (i-2)*4 + 3] = 2.0
        A[row, (i-2)*4 + 4] = 6.0 * x_in[i]
        A[row, (i-1)*4 + 3] = -2.0
        A[row, (i-1)*4 + 4] = -6.0 * x_in[i]

        row += 1
    end

    # Conditions for natural cubic spline
    A[row, 3] = 2.0
    A[row, 4] = 6.0 * x_in[1]

    row += 1

    A[row, (points_count-2)*4 + 3] = 2.0
    A[row, (points_count-2)*4 + 4] = 6.0 * x_in[points_count]

    return Spline(x_in, y_in, A \ b)
end


# Copyright (c) 2021 Leonardo Banchi <leonardo.banchi@unifi.it>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

using LinearAlgebra, BlockDiagonals

" symplectic matrix "
omega(n::Int) = BlockDiagonal([[0 1.; -1. 0] for k=1:n])
omega(A::AbstractMatrix) = omega(Int(size(A,1)/2))

" arrange diagonal elements in a block form "
wdiag(D) = BlockDiagonal([ Diagonal([d,d]) for d in D])

" check if a matrix is symplectic "
issymplectic(S) = S*omega(S)*S'≈omega(S)

" generate a random symplectic matrix "
function rand_symplectic(n)
	A = randn(2n,2n)
	exp(Hermitian(A)*omega(n))
end

" generate a random covariance matrix "
function rand_covmat(n,σ=1)
	v = 1/2 .+ σ*rand(n)
	S = rand_symplectic(n)
	Symmetric(S*wdiag(v)*S')
end

function _will_init(V)
	@assert	size(V,1) == size(V,2) && iseven(size(V,1))
	Int(size(V,1)/2), omega(V)
end

"""
	D,S = williamson(V)

	find the Williamson decomposition such that
	S'*wdiag(D)*S ≈ V
"""
function williamson(V)
	n, Ω = _will_init(V)
	iv12 = inv(sqrt(Symmetric(V)))
	J = iv12*Ω*iv12
	f = schur((J-J')/2)
	T, Q = f.Schur, f.vectors
	P = BlockDiagonal([ sign(T[2k-1,2k]) > 0 ?
					   [1. 0; 0 1.] : [0 1.; 1. 0] for k=1:n])
	D = inv.(diag(P*T*P,1)[1:2:2n])
	iS = wdiag(sqrt.(D))*P*Q'*iv12
	D, Ω*iS*Ω
end

"""
	D,S = williamsonNew(V)

	find the Williamson decomposition such that
	S'*wdiag(D)*S ≈ V
	new algorithm.
"""
function williamsonNew(V)
	n, Ω = _will_init(V)
	λ = sort(real(eigvals(V*Ω*1im)))[n+1:2n]
	S = zeros(2n,2n)
    for m=1:n
	    denom = λ[m]*prod(λ[k]^2-λ[m]^2 for k=1:n if k!=m)
        Vλ = complex.(V)-(1im*λ[m])*Ω
	    minorden(k) = det(view(Vλ,setdiff(1:2n,2m),setdiff(1:2n,k)))/denom
        fam = real(minorden(2m))
	    for k=1:n
            faeven = k==m ? sqrt(fam) : minorden(2k)/sqrt(fam)
            faodd = -minorden(2k-1)/sqrt(fam)
            S[2m-1,2k] 	= -real(faodd)
            S[2m,2k] = imag(faodd)
            S[2m-1,2k-1] = real(faeven)
            S[2m,2k-1] = -imag(faeven)
        end
	end
	λ, S
end


function ksintegrateInplace(u, Lx, dt, Nt, nsave);
    u = (1+0im)*u                       # force u to be complex
    Nx = length(u)                      # number of gridpoints
    kx = vcat(0:Nx/2-1, 0:0, -Nx/2+1:-1)# integer wavenumbers: exp(2*pi*kx*x/L)
    alpha = 2*pi*kx/Lx                  # real wavenumbers:    exp(alpha*x)
    D = 1im*alpha                       # spectral D = d/dx operator
    L = alpha.^2 - alpha.^4             # spectral L = -D^2 - D^4 operator
    G = -0.5*D                          # spectral -1/2 D operator

    Nsave = div(Nt, nsave)+1            # number of saved time steps
    t = (0:Nsave)*(dt*nsave)            # t timesteps
    U = zeros(Nsave, Nx)                # matrix of u(xⱼ, tᵢ) values
    U[1,:] = u                          # assign initial condition to U
    s = 2                               # counter for saved data

    # some convenience variables
    dt2  = dt/2
    dt32 = 3*dt/2
    A_inv = (ones(Nx) - dt2*L).^(-1)
    B     =  ones(Nx) + dt2*L

    # compute in-place FFTW plans
    FFT! = plan_fft!(u, flags=FFTW.ESTIMATE)
    IFFT! = plan_ifft!(u, flags=FFTW.ESTIMATE)

    # compute nonlinear term Nn = -u u_x
    Nn  = G.*fft(u.^2);    # Nn = -1/2 d/dx u^2 = -u u_x
    Nn1 = copy(Nn);        # Nn1 = Nn at first time step
    FFT!*u;                # transform u physical -> spectral

    # timestepping loop
    for n = 1:Nt

        Nn1 .= Nn       # shift nonlinear term in time

        Nn .= u         # put u into Nn
        IFFT!*Nn;       # transform Nn = u to gridpt values, in place
        Nn .= Nn.*Nn;   # collocation calculation, set Nn = u^2
        FFT!*Nn;        # transform Nn = u^2 back to spectral coeffs
        Nn .= G.*Nn;    # apply G = -1/2d/dx  to compute Nn = -1/2 d/dx (u^2)

        # loop fusion! Julia translates this line into a single for-loop on
        #   u[i] = A_inv[i] * (B[i]*u[i] + dt32*Nn[i] - dt2*Nn1[i];
        # no temporary vectors!

        u .= A_inv .* (B .* u .+ dt32.*Nn .- dt2.*Nn1);

        if mod(n, nsave) == 0
            U[s,:] = real(ifft(u))
            s += 1
        end
    end

    t,U
end


grid_size = 25
α_vals = linspace(0.05, 0.5, grid_size)
w_bar_vals = similar(α_vals)

mcm = McCallModel()

for (i, α) in enumerate(α_vals)
    mcm.α = α
    w_bar = compute_reservation_wage(mcm)
    w_bar_vals[i] = w_bar
end

plot(α_vals,
     w_bar_vals,
     lw=2,
     α=0.7,
     xlabel="job separation rate",
     ylabel="reservation wage",
     label=L"$\bar w$ as a function of $\α$")


# we write a function that uses Monte Carlo to estimate the volume of the unit hypersphere in d dimensions using n trials
function sphereVolMC(d, n)

	# we initialize the counter variable
	count = 0

	# we perform the Monte Carlo trial n times
	for i = 1:n

		# we generate a vector with d random floats from [0,1]
		x = rand(d)

		# we calculate the squared norm of x
		norm2 = x' * x

		# if x is inside the hypersphere we increment the count
		if norm2 < 1
			count += 1
		end
	end

	# we calculate the estimated volume given by
	vol = 2^d * count / n

	# we return the estimated volume
	return vol
end


# calculate standard deviation of sensor data

module CalcStdDev
    using DataFrames, CSV, Statistics

    function main()
        data_path = joinpath(split(@__FILE__, "src")[1], "data/sensor_data_200.txt")
        df_200_mm = CSV.read(data_path, DataFrame,
                             header=["date", "time", "ir", "lidar"],
                             delim=' ')

        # calculate from definition
        zs = df_200_mm.lidar # observation
        mean_def = sum(zs) / length(zs) # mean
        diff_square = [(z - mean_def)^2 for z in zs]

        sampling_var = sum(diff_square) / length(zs) # sampling variance
        unbiased_var = sum(diff_square) / (length(zs) - 1) # unbiased variance

        sampling_std_dev = sqrt(sampling_var)
        unbiased_std_dev = sqrt(unbiased_var)

        # calculate by Statistics
        stats_std_dev = Statistics.std(zs, corrected=false)

        println("Sampling standard deviation = $sampling_std_dev")
        println("Unbiased standard deviation = $unbiased_std_dev")
        println("Standard deviation by Statistics = $stats_std_dev")
    end
end


using FileIO, Images
using Plots

pngFileName = "M:\\Matlab\\imgArrow.png"
savePath = "C:\\Users\\phiber\\Downloads\\imgArrowAnim.gif"

function main(pngFile, savePath)

    #### Reading of image to create the function in the complex plane ####

    # Load PNG file and convert from RBG to 3-dimensional matrix
    img_RGB = load(File{format"PNG"}(pngFile))
    img_Mat = permutedims(channelview(img_RGB),(2,3,1))

    # pxl is a boolean matrix where the elements are true where the drawing exists
    pxl = img_Mat[:,:,1] .!= 1
    findPxl = findall(pxl)

    # f is an nx2-dimensional matrix where each row is a coordinate for the consecutive step of the animation
    # f can thus be seen as the discrete complex function in the time domain
    f = Array{Int64}(undef,length(findPxl),2)

    #row (r) and column (c) index of an arbitrary point in the drawing, becomes the starting position of the animation
    r,c = Tuple(findPxl[1])
    f[1,:] = [r,c]
    pxl[findPxl[1]] = false #Set the "used" pixel to false to prevent it to be used again

    nearIdx = -1:1

    j = 2
    while true
        # Look at the pixels adjacent to the previous iteration's pixel and find one that is "drawn"
        near = pxl[f[j-1,1] .+ nearIdx, f[j-1,2] .+ nearIdx]
        firstNear = findfirst(near)
        if isnothing(firstNear) # if there is no more pixels, break the loop
            break
        end

        #Set the next coordinate
        rt,ct = Tuple(firstNear)
        r = r+nearIdx[rt]
        c = c+nearIdx[ct]
        f[j,:] = [r,c]
        pxl[r,c] = false #Set the "used" pixel to false to prevent it to be used again
        j += 1
    end

    f = f[begin:j-1,:]

    #Set the center of the png as the origin
    mid = ceil.(size(pxl) ./2)

    f[:,1] = f[:,1] .- mid[1]
    f[:,2] = f[:,2] .- mid[2]

    #### Complex Fourier Series  ####
    t = range(0,1,length = size(f,1)+1) # Discrete time steps
    t = t[begin:end-1]

    N = 200 #Number of non-static coefficients. Shall be even
    cp = Vector{Complex{Float64}}(undef,Int64(N/2)) #Positive coefficients
    cn = Vector{Complex{Float64}}(undef,Int64(N/2)) #Negative coefficients
    c0 = sum(f[:,1] .+ 1im.* f[:,2]) ./ size(f,1) #static coefficient

    # Calculation of Fourier coefficients
    for k=eachindex(cp)
        cp[k] = sum((f[:,1] .+ 1im.*f[:,2]) .* exp.(-k*2im*pi.*t))./size(f,1)
        cn[k] = sum((f[:,1] .+ 1im.*f[:,2]) .* exp.(k*2im*pi.*t))./size(f,1)
    end

    # Calculation of function in frequency domain
    function vecsumFunc(tt)
        #Function to calculate the sum of the phasors
        F = c0;
        for k = eachindex(cp)
            F = F + cp[k]*exp(k*2im*pi*tt) + cn[k]*exp(-k*2im*pi*tt)
        end
        return F
    end

    #### Plotting ####
    # x and y are the real and imaginary parts, respectively
    x = real(vecsumFunc.(t))
    y = imag(vecsumFunc.(t))

    println("Animating...")

    anim = @animate for i in 1:length(t)
        fourierplot(x,y,i)

        plot!([0,real(c0)],[0,imag(c0)], arrow=false,linewidth=1,color=:black, label=false)
        vecSum = c0;
        for k = eachindex(cp)
            plot!([real(vecSum),real(vecSum + cp[k]*exp(k*2im*pi*t[i]))],
             [imag(vecSum),imag(vecSum + cp[k]*exp(k*2im*pi*t[i]))],
             arrow=false,linewidth=1,color=:dodgerblue4, label=false)
            vecSum = vecSum + cp[k]*exp(k*2im*pi*t[i])

            plot!([real(vecSum),real(vecSum + cn[k]*exp(-k*2im*pi*t[i]))],
             [imag(vecSum),imag(vecSum + cn[k]*exp(-k*2im*pi*t[i]))],
             arrow=false,linewidth=1,color=:darkred, label=false)
            vecSum = vecSum + cn[k]*exp(-k*2im*pi*t[i])
        end
    end

    println("Create gif...")

    gif(anim, savePath, fps = 50)

    return
end

@userplot FourierPlot
@recipe function func(cp::FourierPlot)
    x,y,i = cp.args
    n = length(x)
    inds = circshift(1:n,1-i)
    linewidth --> [range(0,0,length = Int64(floor(n/2))); range(0,10,length = Int64(ceil(n/2)))]
    seriesalpha --> [range(0, 0, length = Int64(floor(n/2))); range(0, 1, length = Int64(ceil(n/2)))]
    aspec_ratio --> 1
    label --> false
    size --> (1080,1080)
    framestyle --> :none
    x[inds], y[inds]
end

main(pngFileName,savePath)


#---------------------------------------------------------------------#
#This function computes the Initial & Exact Solutions.
#Written by F.X. Giraldo on April 19, 2019
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
function exact_solution(coord,Npoin,case,DFloat)

    #Set some constants
    c=2*π

    #Initialize
    qe=zeros(DFloat,Npoin)
    qeₓ=zeros(DFloat,Npoin)
    fe=zeros(DFloat,Npoin)

    #Generate Grid Points
    for I=1:Npoin
        x=coord[I]
        if (case == 1) #Gaussian
            qe[I]=sin(c*x)
            qeₓ[I]=-c*cos(c*x)
            fe[I]=-c^2*sin(c*x)
        end
    end #I
    return (qe,qeₓ,fe)
end


"""
	MM1(μ, λ)

Tạo mô hình M/M/1
"""
struct MM1{T<:Any} <: AbstractMMCK
	λ::Union{T, Real}
	μ::Union{T, Real}
	ρ::Union{T, Real}
	function MM1(λ, μ)
		T = Union{typeof(μ), typeof(λ)}
		new{T}(λ, μ, λ/μ)
	end
	function MM1(λ::Real, μ::Real)
		if λ/μ >= 1
			@error "ρ phải < 1"
		end
		T = Union{typeof(μ), typeof(λ)}
		new{T}(λ, μ, λ/μ)
	end
end

pn(m::MM1, n::Int) = (1 - m.ρ) * m.ρ^n
L(m::MM1) = m.ρ / (1 - m.ρ)
Lq(m::MM1) = m.ρ ^ 2 / (1 - m.ρ)
W(m::MM1) = L(m) / m.λ
Wq(m::MM1) = Lq(m) / m.λ


# Load Julia packages (libraries)

using Pkg, DrWatson

using MonteCarloMeasurements
using StanSample
using StatisticalRethinking

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

# New col log_pop, set log() for population data
df[!, :log_pop] = map((x) -> log(x), df[!, :population]);
df[!, :society] = 1:10;

stan12_6 = "
data {
    int N;
    int N_societies;
    int total_tools[N];
    real logpop[N];
    int society[N];
}
parameters{
    real a;
    real bp;
    vector[N_societies] a_society;
    real<lower=0> sigma_society;
}
model{
    vector[N_societies] mu;
    sigma_society ~ cauchy( 0 , 1 );
    a_society ~ normal( 0 , sigma_society );
    bp ~ normal( 0 , 1 );
    a ~ normal( 0 , 10 );
    for ( i in 1:N ) {
        mu[i] = a + a_society[society[i]] + bp * logpop[i];
        mu[i] = exp(mu[i]);
    }
    total_tools ~ poisson( mu );
}
";

# Define the SampleModel.

m12_6s = SampleModel("m12.6s",  stan12_6);

# Input data for cmdstan

m12_6_data = Dict("N" => size(df, 1), "N_societies" => 10,
"total_tools" => df[!, :total_tools], "logpop" => df[!, :log_pop],
"society" => df[!, :society]);

# Sample using cmdstan's sample option

rc12_6s = stan_sample(m12_6s, data=m12_6_data);

# Describe the draws

if success(rc12_6s)
  part12_6s = read_samples(m12_6s, :particles)
  part12_6s |> display
end


## Lecture-3-Least-squares-CVX

## Here are the packages we'll use
using Convex
using Plots
plotlyjs() # need to use plotlyjs on my machine for plots.

## Step 1: Create a set of data with a linear model
m=40
n=2
A = randn(m,n)
xex = [5;1]      # "b" = 5 in ax+b and a=1
pts = -10+20*rand(m,1)
A = [ones(m,1) pts]
b = A*xex + .5*randn(m,1)

## Show the linear model
scatter(pts,b)
xlabel!("x")
ylabel!("y")

##
x = Variable(n)
problem = minimize(sumsquares(b - A*x))
solve!(problem)
xls = x.value
@show x
## Show the least squares fit
scatter(pts,b;label="data")
plot!([-11; 11], [1 -11; 1 11]*xls;label="fit")
#xaxis!([-11 11])
title!("Least-square fit")
xlabel!("x")
ylabel!("y")

## Now we add outliers
outliers = [-9.5; 9]
outvals = [20; -15]
A = [A; ones(length(outliers),1) outliers]
b = [b; outvals]
m = size(A,1)
pts = [pts;outliers]

## Show the new data
scatter(pts,b)
xlabel!("x")
ylabel!("y")

## Look at the LS fit
x = Variable(n)
problem = minimize(sumsquares(b - A*x))
solve!(problem)
xls = x.value
plot!([-11; 11], [1 -11; 1 11]*xls;label="fit")

## Solve the Huber problem and look at the fit
x = Variable(n)
problem = minimize(sum(huber(b - A*x)))
solve!(problem)
xr = x.value
plot!([-11; 11], [1 -11; 1 11]*xr;label="fit_huber")


using Plots
t = collect(-4:0.2:4)
function huber(t; λ = 1)
    flag = (abs.(t) .< λ)
    y1 = t.^2 / 2 .* flag .+ (λ * abs.(t) .- λ^2 / 2) .* (1 .- flag)
    y2 = λ * abs.(t) .- λ^2 / 2
    return y1, y2
end
y1, y2 = huber(t, λ = 2)
plot(t, y1, label = "Huber", legend = :bottomright)
plot!(t, y2, label = "Linear")


# Polynomial covariance function

@doc """
# Description
Constructor for the Polynomial kernel (covariance)

k(x,x') = σ²(xᵀx'+c)ᵈ
# Arguments:
* `lc::Float64`: Log of the constant c
* `lσ::Float64`: Log of the signal standard deviation σ
* `d::Int64`   : Degree of the Polynomial
""" ->
type Poly <: Kernel
    lc::Float64      # Log of constant
    lσ::Float64      # Log of signal std
    deg::Int64       # degree of polynomial
    Poly(lc::Float64, lσ::Float64, deg::Int64) = new(lc, lσ, deg)
end

function cov(poly::Poly, x::Vector{Float64}, y::Vector{Float64})
    c = exp(poly.lc)
    sigma2 = exp(2*poly.lσ)

    K = sigma2*(c+dot(x,y)).^poly.deg
    return K
end

get_params(poly::Poly) = Float64[poly.lc, poly.lσ]
get_param_names(poly::Poly) = [:lc, :lσ]
num_params(poly::Poly) = 2

function set_params!(poly::Poly, hyp::Vector{Float64})
    length(hyp) == 2 || throw(ArgumentError("Polynomial function has two parameters"))
    poly.lc, poly.lσ = hyp
end

function grad_kern(poly::Poly, x::Vector{Float64}, y::Vector{Float64})
    c = exp(poly.lc)
    sigma2 = exp(2*poly.lσ)

    dK_c   = c*poly.deg*sigma2*(c+dot(x,y)).^(poly.deg-1)
    dK_sigma = 2.0*sigma2*(c+dot(x,y)).^poly.deg
    dK_theta = [dK_c,dK_sigma]
    return dK_theta
end


using Images, BenchmarkTools

function calc(a,b)
    z = c = complex(a, b)
    for i in 1:3000
        z = z*z + c
        if abs(z) > 3
            return true
        end
    end
    return false
end

function buddha(points::Array{ComplexF64}, h::Int, w::Int)
    img = zeros(RGB,w,h)
    for c::ComplexF64 in points
        i::Int = 0
        z::ComplexF64 = c
        while abs(z) < 3 && i < 1000
            x::Int = trunc(Int,(real(z) + 2)*w/3); y::Int = trunc(Int,(imag(z) + 1.5)*h/3)
            0<x<w && 0<y<h && i > 20 ? img[x,y] += RGB(0,0,0.004) : nothing #log10(i)/3
            z = z*z + c
            i += 1
        end
    end
    return img
end

function main()
    h::Int = 8000
    w::Int = 8000
    points::Array{ComplexF64} = []
    for b in -1.5:3/h:1.5
        for a in -2.:3/w:1.
            calc(a,b) && append!(points, complex(a,b))
        end
    end
    println("Init done")
    return map(clamp01nan, buddha(points,2000,2000))
end

#@benchmark main()
@time save("./out/buddha.png", main())


"""
    ExpLik <: Likelihood

Exponential likelihood
```math
p(y | f) = θ\\exp(-θy),
```
where ``θ = \\exp(-f)`` and ``f`` is the latent Gaussian process.
"""
struct ExpLik <: Likelihood end

#log of probability density
function log_dens(exponential::ExpLik, f::VecF64, y::VecF64)
    #where we exponentiate for positivity f = exp(fi)
    return [-fi - exp(-fi)*yi for (fi,yi) in zip(f,y)]
end

#derivative of pdf wrt latent function
function dlog_dens_df(exponential::ExpLik, f::VecF64, y::VecF64)
    return [(yi*exp(-fi)-1) for (fi,yi) in zip(f,y)]
end

#mean and variance under likelihood
mean_lik(exponential::ExpLik, f::VecF64) = exp.(f)
var_lik(exponential::ExpLik, f::VecF64) = exp.(f).^2

get_params(exponential::ExpLik) = []
num_params(exponential::ExpLik) = 0