using SortingAlgorithms

function main(points)
  try
    ch_points = []
    ch_points = convex_hull!(points)

    output = [[],[],[],[]]

    for ch_point in ch_points
      push!(output[1], ch_point[1])
      push!(output[2], ch_point[2])

      deleteat!(points, findfirst(points, ch_point))
    end

    for point in points
      push!(output[3], point[1])
      push!(output[4], point[2])
    end

    return output
  catch e
    if isa(e, MethodError)
      return "Please provide array of tuples where each element is Int64"
    else
      return e
    end
  end
end

function convex_hull!(points::Array{Tuple{Int64,Int64},1})
  if size(points, 1) < 3
    return points
  end

  # Sort all points
  points = sort!(union(points), alg = HeapSort)

  # Arrays for top and bottom parts of convex hull
  calculated_ch_top = []
  calculated_ch_bottom = []

  # Build top part of convex hull
  build_convex_hull_part!(calculated_ch_top, points, 1:1:length(points))

  # Build bottom part of convex hull
  build_convex_hull_part!(calculated_ch_bottom, points, length(points):-1:1)

  # Remove duplicated points:
  # End of top part == beginning of the bottom part
  pop!(calculated_ch_top)

  # End of bottom part == beginning of the top part
  pop!(calculated_ch_bottom)

  # Top + Bottom parts == convex hull
  convex_hull = [calculated_ch_top; calculated_ch_bottom]

  return convex_hull
end

function counter_clockwise(p1, p2, p3)
  return ((p2[1] - p1[1]) * (p3[2] - p1[2])) - ((p2[2] - p1[2]) * (p3[1] - p1[1]))
end

function generate_points(number_of_points, range_min, range_max)
  points = [(0,0)]
  pop!(points)
  for i = 1:number_of_points
    generated_point = (rand(range_min:range_max), rand(range_min:range_max))
    push!(points, generated_point)
  end
  return points
end

function build_convex_hull_part!(hull_points, points_to_check, iterator_range)
  for i = iterator_range
    point = points_to_check[i]

    while length(hull_points) > 1 && counter_clockwise(hull_points[end-1], hull_points[end], point) >= 0
      pop!(hull_points)
    end

    push!(hull_points, point)
  end
end


n = rand(1:10)
Z = Companion(randn(n))

#Special properties
@test Matrix(inv(Z)) ≈ inv(Matrix(Z))

#Matvec product
b = randn(n)
@test Z*b ≈ Matrix(Z)*b

m = rand(1:10)
A = randn(m, n)
@test A*Z ≈ A*Matrix(Z)

# Polynomial construction
using Polynomials
p = Polynomial([-1,0,1])
p_c = Companion(p)
v1 = [1,1]
v2 = [-1,1]
@test p_c*v1 == v1
@test p_c*v2 == -v2


using SpecialFunctions, StatsFuns

@testset "logmvgamma" begin
    @testset "type behavior" for eltya in (Float32, Float64)
        p = rand(1:50)
        a = rand(eltya)
        # add p since loggamma is only define for positive arguments
        @test typeof(logmvgamma(p, a + p)) == eltya
    end

    @testset "consistent with loggamma" for eltya in (Float32, Float64)
        #  Γ₁(a) = Γ(a), Γ₂(a) = √π Γ(a) Γ(a - 0.5), etc
        a = rand(eltya) + 1 # add one since loggamma is only define for positive arguments
        @test logmvgamma(1, a) ≈ loggamma(a)
        @test logmvgamma(2, a) ≈    eltya(0.5logπ) + loggamma(a) + loggamma(a - eltya(0.5))
        @test logmvgamma(3, a) ≈ eltya((3/2)*logπ) + loggamma(a) + loggamma(a - eltya(0.5)) + loggamma(a - one(a))
    end

    @testset "consistent with itself" for eltya in (Float32, Float64)
        #  Γᵢ(a) = (π^{i-1/2}) Γ(a) Γᵢ₋₁(a - 0.5)
        for p in 1:50
            a = rand(eltya) + p # add p since loggamma is only define for positive arguments
            @test logmvgamma(p, a) ≈ eltya((p/2-1/2)*logπ) + loggamma(a) + logmvgamma(p - 1, a - eltya(0.5))
        end
    end
end


# Isomap
# ------
# A Global Geometric Framework for Nonlinear Dimensionality Reduction,
# J. B. Tenenbaum, V. de Silva and J. C. Langford, Science 290 (5500): 2319-2323, 22 December 2000

#### Isomap type
struct Isomap{T <: Real} <: AbstractDimensionalityReduction
    k::Int
    model::KernelPCA
    component::AbstractVector{Int}

    Isomap{T}(k::Int, model::KernelPCA) where T = new(k, model)
    Isomap{T}(k::Int, model::KernelPCA, cc::AbstractVector{Int}) where T = new(k, model, cc)
end

## properties
outdim(R::Isomap) = outdim(R.model)
eigvals(R::Isomap) = principalvars(R.model)
neighbors(R::Isomap) = R.k
vertices(R::Isomap) = R.component

## show
summary(io::IO, R::Isomap) = print(io, "Isomap(outdim = $(outdim(R)), neighbors = $(neighbors(R)))")

## interface functions
function fit(::Type{Isomap}, X::AbstractMatrix{T}; k::Int=12, maxoutdim::Int=2) where {T<:Real}
    # Construct NN graph
    D, E = find_nn(X, k)
    G, C = largest_component(SimpleWeightedGraph(adjmat(D,E)))

    # Compute shortest path for every point
    n = length(C)
    DD = zeros(T, n, n)
    for i in 1:n
        dj = dijkstra_shortest_paths(G, i)
        DD[i,:] = dj.dists
    end

    broadcast!(x->-x*x/2, DD, DD)
    broadcast!((x,y)->(x+y)/2, DD, DD, DD') # remove roundoff error
    M = fit(KernelPCA, DD, kernel=nothing, maxoutdim=maxoutdim)

    return Isomap{T}(k, M, C)
end

transform(R::Isomap) = transform(R.model)


"""
# Usage

    solve(prob::FODESystem, h, NonLinearAlg())

Nonlinear algorithm for nonlinear fractional differential equations.

### References

Dingyu Xue, Northeastern University, China ISBN:9787030543981
"""
struct NonLinearAlg <: FractionalDiffEqAlgorithm end

function solve(prob::FODESystem, h, ::NonLinearAlg, L0=1e10)
    @unpack f, α, u0, tspan = prob
    t0 = tspan[1]; T = tspan[2]
    n = length(u0)
    m = round(Int, (T-t0)/h)+1
    g = genfun(1)
    g = g[:]
    u0 = u0[:]
    ha = h.^α
    z = zeros(n, m)
    x1 = copy(u0) # Here we pass the value of x0 to x1. Honestly, I kept finding this bug for almost a whole night😅

    # All of the min(m, L0+1) is to set the memory effect.
    SetMemoryEffect = Int64(min(m, L0+1))
    W = zeros(n, SetMemoryEffect) #Initializing W a n*m matrix

    @fastmath @inbounds @simd for i = 1:n
        W[i, :] = getvec(α[i], SetMemoryEffect, g)
    end

    du = zeros(n)
    @fastmath @inbounds @simd for k = 2:m
        tk = t0+(k-1)*h
        L = min(Int64(k-1), Int64(L0))
        f(du, x1, nothing, tk)

        @fastmath @inbounds @simd for i = 1:n
            x1[i] = du[i]*ha[i] - W[i, 2:L+1]'*z[i, k-1:-1:k-L] + u0[i]
        end
        z[:, k] = x1 - u0
    end

    result = (z + repeat(u0, 1, m))'

    return result
end

"""
P-th precision polynomial generate function

```math
g_p(z)=\\sum_{k=1}^p \\frac{1}{k}(1-z)^k
```
"""
function genfun(p)
    a = collect(1:p+1)
    A = Vandermonde(a)'
    return (1 .-a')*inv(A')
end

function getvec(α, n, g)
    p = length(g)-1
    b = 1 + α
    g0 = g[1]
    w = Float64[]
    push!(w, g[1]^α)

    @fastmath @inbounds @simd for m = 2:p
        M = m-1
        dA = b/M
        temp = (-(g[2:m] .*collect((1-dA):-dA:(1-b))))' *w[M:-1:1]/g0
        push!(w, temp)
    end

    @fastmath @inbounds @simd for k = p+1:n
        M = k-1
        dA = b/M
        temp = (-(g[2:(p+1)] .*collect((1-dA):-dA:(1-p*dA))))' *w[M:-1:(k-p)]/g0
        push!(w, temp)
    end
    return w
end


# logistic regression on mnist data

using MultiStochGrad

using  LinearAlgebra, Test
using Random, Distributions

include("../src/applis/LinearRegression.jl")

# the . means we use the module LinearRegressionMod in the current module, other wise refer toa package!

function test_linear_regression_scsg()
    true_coefficients = [13.37, -4.2, 3.14]
    # sample noise according to a N(0,1)
    datas = Vector{Vector{Float64}}()
    values = Vector{Float64}()
    d = Normal()
    for i in 1:100
        v = Vector{Float64}()
        push!(v,1.)
        push!(v, rand())
        push!(v, rand())
        y = dot(true_coefficients, v) + rand(d)
        push!(datas, v)
        push!(values, y)
    end
    observations= Observations(datas,values)
    nbterms = length(observations.datas)
    # define our evaluations
    linear_reg_term_value = make_linear_reg_term_value(observations)
    term_function =  TermFunction{typeof(linear_reg_term_value)}(linear_reg_term_value, nbterms, Dims{1}(3))
    #
    linear_reg_term_gradient = make_linear_reg_term_gradient(observations)
    term_gradient = TermGradient{typeof(linear_reg_term_gradient)}(linear_reg_term_gradient ,nbterms, Dims{1}(3))
    evaluator = Evaluator{typeof(linear_reg_term_value),typeof(linear_reg_term_gradient)}(term_function,term_gradient)
    #
    scsg_pb = SCSG(0.1, 0.3, 1 , 0.7)
    # solve.
    nb_iter = 50
    position = fill(1., 3)
    initial_error = compute_value(evaluator, position)
    @info "\n initial error \n\n" initial_error
    position, value = minimize(scsg_pb, evaluator, nb_iter, position)
    @info "value , position " value  position
    @info "\n\n test_linear_regression_scsg done \n \n"
    value < 0.65 ? true : false
end

function test_linear_regression_svrg()
    true_coefficients = [13.37, -4.2, 3.14]
    # sample noise according to a N(0,1)
    datas = Vector{Vector{Float64}}()
    values = Vector{Float64}()
    d = Normal()
    for i in 1:100
        v = Vector{Float64}()
        push!(v,1.)
        push!(v, rand())
        push!(v, rand())
        y = dot(true_coefficients, v) + rand(d)
        push!(datas, v)
        push!(values, y)
    end
    observations= Observations(datas,values)
    nbterms = length(observations.datas)
    # define our evaluations
    linear_reg_term_value = make_linear_reg_term_value(observations)
    term_function =  TermFunction{typeof(linear_reg_term_value)}(linear_reg_term_value, nbterms, Dims{1}(3))
    #
    linear_reg_term_gradient = make_linear_reg_term_gradient(observations)
    term_gradient = TermGradient{typeof(linear_reg_term_gradient)}(linear_reg_term_gradient ,nbterms, Dims{1}(3))
    evaluator = Evaluator{typeof(linear_reg_term_value),typeof(linear_reg_term_gradient)}(term_function,term_gradient)
    #
    svrg_pb = SVRG(100, 0.2)
    # solve.
    nb_iter = 60
    position = fill(1., 3)
    initial_error = compute_value(evaluator, position)
    @info "\n initial error \n\n" initial_error
    position, value = minimize(svrg_pb, evaluator, nb_iter, position)
    @info "value , position ", value ,  position
    @info "\n\n test_linear_regression_svrg done \n \n"
    value < 0.65 ? true : false
end


"""Julia program to implement Exponential Search algorithm.
Exponential Search Algorithm is an optimized Binary Search to search an element in sorted array.
It is specifically used when the size of array is infinite.
"""

function binary_search(arr, low, n, ele)
    high = n
    while high >= low
        mid = low + (high - low) ÷ 2
        # If the mid element is the required element, return that index
        if(arr[mid] == ele)
            return mid
        # ele is greater than mid element then ele would be present in first half of the array
        elseif (arr[mid] > ele)
            high = mid - 1
        #Else if  ele is smaller than mid element then ele would be present in last half of the array
        else
            low = mid + 1
        end
    end
    # If the element is not found return 0
    return 0
end

function exponential_search(arr, n, ele)
    if(arr[1] == ele)
        return false
    end
    i = 2
    while( i < n && arr[i] <= ele)
        i = i * 2
    end
    mini = i < (n-1) ? i : (n-1)
    return binary_search(arr, i÷2, mini, ele )
end

print("How many numbers are present in the array? ")
n = readline()
n = parse(Int, n)
if (n <= 0)
    println("Array is Empty!!!")
    exit()
end
arr = Int[]
print("Enter the numbers: ")
arr = [parse(Int, num) for num in split(readline())]
print("Which number do you want to search in the array? ")
ele = readline()
ele = parse(Int, ele)
# Sort the array in ascending order
arr = sort(arr)
res = exponential_search(arr, n, ele)
if (res == 0)
    print("The number $ele is not present in the array")
else
    print("The number $ele is present in the array.")
end

"""
Time Complexity - O(log(n)), where 'n' is the size of the array
Space Complexity - O(n)

SAMPLE INPUT AND OUTPUT

SAMPLE I
How many numbers are present in the array? 5
Enter the numbers: 1 2 3 4 5
Which number do you want to search in the array? 6
The number 6 is not present in the array

SAMPLE II
How many numbers are present in the array? 3
Enter the numbers: 3 1 2
Which number do you want to search in the array? 2
The number 2 is present in the array.
"""


"""
Package: Smoothers

    sma(x, n, center=false)

Smooth a vector of data using a simple moving average

# Arguments
- `x`: Vector of data.
- `n`: Size of the moving average.
- `center`: returns a vector of the same size of x with missing values on the tails.

# Returns
Vector of moving average values

# Examples
```julia-repl
julia> sma(1:5,3)
3-element Vector{Float64}:
 2.0
 3.0
 4.0

julia> sma(1:5,3,true)
5-element Vector{Union{Missing, Float64}}:
  missing
 2.0
 3.0
 4.0
  missing
```
"""
@inline function sma(x::AbstractVector{T}, n::Integer, center::Bool=false) where T<:Real

    n == 1 && return x
    N = length(x)
    @assert 1 <= n <= N

    P = Base.promote_op(/, T, T)
    res = Vector{P}(undef, N-n+1)

    # initial moving average value
    res[1] = ma = sum(x[1:n])/n
    for i in 1:N-n
        @inbounds res[1+i] = ma += (x[n+i] - x[i]) / P(n)
    end

    center ? vcat(repeat([missing], n÷2),res,repeat([missing], n-n÷2-1)) : res

end


"""
    rv(X, Y)
Compute the sample RV coefficient between matrices X and Y
* `X` : Matrix (n obs., p variables).
* `Y` : Matrix (n obs., q variables).

RV is bounded in [0, 1].

A dissimilarty measure between `X` and `Y` can be computed
by d = sqrt(2 * (1 - RV)).

## References

Escoufier, Y., 1973. Le Traitement des Variables Vectorielles. Biometrics 29, 751–760.
https://doi.org/10.2307/2529140

Josse, J., Pagès, J., Husson, F., 2008. Testing the significance of the RV coefficient.
Computational Statistics & Data Analysis 53, 82–91. https://doi.org/10.1016/j.csda.2008.06.012

Mayer, C.-D., Lorent, J., Horgan, G.W., 2011. Exploratory Analysis of Multiple Omics
Datasets Using the Adjusted RV Coefficient. Statistical Applications in Genetics and Molecular
Biology 10. https://doi.org/10.2202/1544-6115.1540

Smilde, A.K., Kiers, H.A.L., Bijlsma, S., Rubingh, C.M., van Erk, M.J., 2009.
Matrix correlations for high-dimensional data: the modified RV-coefficient.
Bioinformatics 25, 401–405. https://doi.org/10.1093/bioinformatics/btn634

Robert, P., Escoufier, Y., 1976. A Unifying Tool for Linear Multivariate Statistical Methods:
The RV-Coefficient. Journal of the Royal Statistical Society: Series C (Applied Statistics)
25, 257–265. https://doi.org/10.2307/2347233
"""

function rv(X, Y; centr = true)
    X = ensure_mat(X)
    Y = ensure_mat(Y)
    if centr
        X = center(X, mean(X, dims = 1))
        Y = center(Y, mean(Y, dims = 1))
    end
    YtX = Y' * X
    XtX = X' * X
    YtY = Y' * Y
    A = dot(YtX, YtX)
    B = dot(XtX, XtX)
    C = dot(YtY, YtY)
    rv = A / sqrt(B * C)
end


module day05

include("inputs.jl")
input = inputs.input_d5()

lookup = Dict('B'=>'1', 'F'=>'0', 'R'=>'1', 'L'=>'0')

function parse_bp(bp)
    row = parse(Int, join(lookup[c] for c in bp[1:7]), base=2)
    col = parse(Int, join(lookup[c] for c in bp[8:end]), base=2)
    seat_id = 8*row + col
    return seat_id
end

@assert parse_bp("BFFFBBFRRR") == 567
@assert parse_bp("FFFBBBFRRR") == 119
@assert parse_bp("BBFFBBFRLL") == 820

# Part 1
function highest_id(input)
    seat_ids = parse_bp.(input)
    return maximum(seat_ids)
end
@info "Day 5, Part 1 answer: $(highest_id(input))"

#Part 2
# function find_seat(;input=input())
function find_seat(input)
    seat_ids = parse_bp.(input)
    lo = minimum(seat_ids)
    hi = maximum(seat_ids)
    return [x for x in lo:hi if x ∉ seat_ids]
end
@info "Day 5, Part 2 answer: $(find_seat(input)[1])"

end


# This file calculates the field strength at the center of a Faraday cage. Warning, the run-time is long.
# See /Laplace.jl for a more interactive version.
using ApproxFun, SingularIntegralEquations

z_0 = 2.0
ui(x,y) = logabs(complex(x,y)-z_0)
g1(x,y) = 1/2

function electricfield(N,r)
    # Set the domains.
    cr = exp(im*2π*[0:N-1]/N)
    crl = (1-2r)cr
    crr = (1+2r)cr
    dom = ∪(Segment.(crl,crr))
    #dom = ∪(Circle.(cr,ones(length(cr))r))
    #dom = ∪(Segment.(crl[1:2:end],crr[1:2:end]) ∪ ∪(Circle.(cr[2:2:end],ones(length(cr[2:2:end]))r))

    sp = Space(dom)
    cwsp = CauchyWeight(sp⊗sp,0)
    uiΓ,⨍ = Fun(t->ui(real(t),imag(t))+0im,sp),DefiniteLineIntegral(dom)

    G = GreensFun(g1,cwsp;method=:Cholesky)

    ∂u∂n = ⨍[G]\uiΓ

    us(x,y) = -logkernel(∂u∂n,complex(x,y))/2
    ut(x,y) = ui(x,y) + us(x,y)
    h = 1e-5
    abs((ut(h,0.)-ut(-h,0.))/2h)
end

NNr1,NTr1 = [4:15],[4:15,17:2:20,30:10:80]
ENr1 = Float64[@time electricfield(N,1e-1im) for N in NNr1]
println("Done ENr1")
ETr1 = Float64[@time electricfield(N,1e-1) for N in NTr1]
println("Done ETr1")

NNr2,NTr2 = [[5:10],[15:5:20],[30:10:80]],[[5:10],[15:5:20],[30:10:100]]
ENr2 = Float64[@time electricfield(N,1e-2im) for N in NNr2]
println("Done ENr2")
ETr2 = Float64[@time electricfield(N,1e-2) for N in NTr2]
println("Done ETr2")

#Nr3 = [[5:10],[15:5:20],[30:10:90],[100:50:200]]
#@time Er3 = Float64[electricfield(N,1e-3) for N in Nr3]
#@time Er4 = Float64[electricfield(N,1e-4) for N in Nr3]


# ------ Root Finding Methods -----

"""
    brent(f, a, b; args=(), atol=2e-12, rtol=4*eps(), maxiter=100)

1D root finding using Brent's method.  Based off the brentq implementation in scipy.

**Arguments**
- `f`: scalar function, that optionally takes additional arguments
- `a`::Float, b::Float`: bracketing interval for a root - sign changes sign between: (f(a) * f(b) < 0)
- `args::Tuple`: tuple of additional arguments to pass to f
- `atol::Float`: absolute tolerance (positive) for root
- `rtol::Float`: relative tolerance for root
- `maxiter::Int`: maximum number of iterations allowed

**Returns**
- `xstar::Float`: a root of f
- `info::Tuple`: A named tuple containing:
    - `iter::Int`: number of iterations
    - 'fcalls::Int`: number of function calls
    - 'flag::String`: a convergence/error message.
"""
function brent(f, a, b; args=(), atol=2e-12, rtol=4*eps(), maxiter=100)

    xpre = a; xcur = b
    # xblk = 0.0; fblk = 0.0; spre = 0.0; scur = 0.0
    error_num = "INPROGRESS"

    fpre = f(xpre, args...)
    fcur = f(xcur, args...)
    xblk = zero(fpre); fblk = zero(fpre); spre = zero(fpre); scur = zero(fpre)
    funcalls = 2
    iterations = 0

    if fpre*fcur > 0
        error_num = "SIGNERR"
        return 0.0, (iter=iterations, fcalls=funcalls, flag=error_num)
    end
    if fpre == zero(fpre)
        error_num = "CONVERGED"
        return xpre, (iter=iterations, fcalls=funcalls, flag=error_num)
    end
    if fcur == zero(fcur)
        error_num = "CONVERGED"
        return xcur, (iter=iterations, fcalls=funcalls, flag=error_num)
    end

    for i = 1:maxiter
        iterations += 1
        if fpre*fcur < 0
            xblk = xpre
            fblk = fpre
            spre = scur = xcur - xpre
        end
        if abs(fblk) < abs(fcur)
            xpre = xcur
            xcur = xblk
            xblk = xpre

            fpre = fcur
            fcur = fblk
            fblk = fpre
        end

        delta = (atol + rtol*abs(xcur))/2
        sbis = (xblk - xcur)/2
        if fcur == zero(fcur) || abs(sbis) < delta
            error_num = "CONVERGED"
            return xcur, (iter=iterations, fcalls=funcalls, flag=error_num)
        end

        if abs(spre) > delta && abs(fcur) < abs(fpre)
            if xpre == xblk
                # interpolate
                stry = -fcur*(xcur - xpre)/(fcur - fpre)
            else
                # extrapolate
                dpre = (fpre - fcur)/(xpre - xcur)
                dblk = (fblk - fcur)/(xblk - xcur)
                stry = -fcur*(fblk*dblk - fpre*dpre)/(dblk*dpre*(fblk - fpre))
            end
            if 2*abs(stry) < min(abs(spre), 3*abs(sbis) - delta)
                # good short step
                spre = scur
                scur = stry
            else
                # bisect
                spre = sbis
                scur = sbis
            end
        else
            # bisect
            spre = sbis
            scur = sbis
        end

        xpre = xcur; fpre = fcur
        if abs(scur) > delta
            xcur += scur
        else
            xcur += (sbis > 0 ? delta : -delta)
        end

        fcur = f(xcur, args...)
        funcalls += 1
    end
    error_num = "CONVERR"
    return xcur, (iter=iterations, fcalls=funcalls, flag=error_num)
end

# TODO AN: replace w/ newer Brent method and automatic bracketing?


# !INPUT/OUTPUT PARAMETERS:
#   m : order of derivative (in,integer)
#   n : number of points (in,integer)
#   x : abscissa array (in,real(n))
#   f : function array (in,real(n))
#   g : (anti-)derivative of f (out,real(n))
# !DESCRIPTION:
#   Given function $f$ defined on a set of points $x_i$ then if $m\ge 0$ this
#   routine computes the $m$th derivative of $f$ at each point. If $m=-1$ the
#   anti-derivative of $f$ given by
#   $$ g(x_i)=\int_{x_1}^{x_i} f(x)\,dx $$
#   is calculated. Both derivatives and integrals are computed by first fitting
#   the function to a clamped cubic spline.
function fderiv!(m, n, x, f, g)

    @assert n > 0

    # automatic arrays
    cf = zeros(Float64,3,n)

    # high accuracy integration/differentiation from spline interpolation
    spline!(n, x, f, cf)
    if m == -1
        ss =0.0
        g[1] = 0.0
        for i in 1:n-1
          dx = x[i+1] - x[i]
          ss = ss + dx*( f[i] + dx*( 0.50*cf[1,i] + dx*( cf[2,i]/3 + dx*0.25*cf[3,i] ) ) )
          g[i+1] = ss
        end
    elseif m == 1
        for i in 1:n
            g[i] = cf[1,i]
        end
    elseif m == 2
        for i in 1:n
            g[i] = 2.0*cf[2,i]
        end
    elseif m == 3
        for i in 1:n
            g[i] = 6.0*cf[3,i]
        end
    else
        for i in 1:n
            g[i] = 0.0
        end
    end
    return
end


# myplotfile.jl
using Plots, Distributions

μ₁ = 10; μ₂ = 20; μ = [μ₁, μ₂]  # mean matrix
Σ = [1.0 0;0 3];                # covariance matrix (i.e. no correlation between the variables)
d = Distributions.MvNormal(μ,Σ) # multivariate normal distribution

# make a grid
ns = 3
nx = 21
ny = 31
X = range(μ₁ - ns*Σ[1,1], stop=μ₁ + ns*Σ[1,1],length=nx);
Y = range(μ₂ - ns*Σ[2,2], stop=μ₂ + ns*Σ[2,2],length=ny);

Z = permutedims(collect(Iterators.product(X,Y)))
Z = map(x->pdf(d,collect(x)), Z)

# common plot setting
plotsettings =  Dict(:xlims=>(0,20), :ylims=>(10,30),
                 :xlabel=>"x", :ylabel=>"y")

p = plot(surface(X,Y,Z, color=:blues, title="surface plot"),
         surface(X,Y,Z, color=:blues_r, title="surface plot\n(reversed colors)"),
         contourf(X,Y,Z, title="contour plot"),
         heatmap(X,Y,Z, title="heatmap"),
         layout=(2,2), size=(1200,1200); plotsettings...)

for extension in ["png","pdf"]
    savefig(p, "./img/myplot.$(extension)")
end


# libpurex:
# Exponential families paramterised by their means

include("binary_search.jl");

struct Gaussian
    σ2;
end

# convenience
Gaussian() = Gaussian(1);

struct Bernoulli
end

struct Exponential
end

struct Poisson
end

# KL divergence

rel_entr(x, y) = x==0 ? 0. : x*log(x/y);

d(expfam::Gaussian,    μ, λ) = (μ-λ)^2/(2*expfam.σ2);
d(expfam::Bernoulli,   μ, λ) = rel_entr(μ, λ) + rel_entr(1-μ, 1-λ);
d(expfam::Exponential, μ, λ) = λ == 0 ? Inf : μ/λ - log(μ/λ) - 1;
d(expfam::Poisson,     μ, λ) = rel_entr(μ, λ) - μ + λ;

# solution for λ in \min_λ d(μ, λ) - λ x
# i.e. λ satisfying (λ - μ)ϕ''(λ) == x

invh(expfam::Gaussian,    μ, x) = μ + x*expfam.σ2;
invh(expfam::Bernoulli,   μ, x) = 2μ/(1-x+sqrt((x-1)^2 + 4*x*μ));
invh(expfam::Exponential, μ, x) = x > 0 ? Inf : 2μ/(1 + sqrt(1 - 4*x*μ));
invh(expfam::Poisson,     μ, x) = μ/(1-x);

sample(rng, expfam::Gaussian,    μ) = μ + sqrt(expfam.σ2)*randn(rng);
sample(rng, expfam::Bernoulli,   μ) = rand(rng) ≤ μ;
sample(rng, expfam::Exponential, μ) = randexp(rng)*μ;

# upward and downward confidence intervals
dup(expfam::Gaussian, μ, v) = μ + sqrt(2*expfam.σ2*v);
ddn(expfam::Gaussian, μ, v) = μ - sqrt(2*expfam.σ2*v);

function dup(expfam::Bernoulli, μ, v)
    μ == 1 ? 1. : binary_search(λ -> d(expfam, μ, λ) - v, μ, 1);
end

function ddn(expfam::Bernoulli, μ, v)
    μ == 0 ? 0. : binary_search(λ -> v - d(expfam, μ, λ), 0, μ);
end

function dup(expfam::Exponential, μ, v)
    binary_search(λ -> d(expfam, μ, λ) - v, μ, μ*exp(v+1));
end

function ddn(expfam::Exponential, μ, v)
    binary_search(λ -> v - d(expfam, μ, λ), 0, μ);
end


module TestCatElements
using Test
using Catlab, Catlab.Theories, Catlab.Graphs, Catlab.CategoricalAlgebra

arr = @acset Graph begin
  V = 2
  E = 1
  src=[1]
  tgt=[2]
end
elarr = elements(arr)
@testset "Arrow Elements" begin
  @test nparts(elarr, :El)  == 3
  @test nparts(elarr, :Arr) == 2
  @test nparts(elarr, :Ob)  == 2
  @test nparts(elarr, :Hom) == 2
  @test elarr[:, :πₑ]    == [1,1,2]
  @test elarr[:, :πₐ]    == [1,2]
  @test elarr[:, :src]   == [3,3]
  @test elarr[:, :tgt]   == [1,2]
  @test elarr[:, :nameo] == [:V, :E]
  @test elarr[:, :nameh] == [:src, :tgt]
  @test elarr[:, :dom]   == [2, 2]
  @test elarr[:, :cod]   == [1, 1]
end

b = @acset Graph begin
  V = 2
  E = 1
  src = [1]
  tgt = [2]
end

ThBPG, obmap, hommap = CatElements.presentation(elements(b))
@test Symbol.(generators(ThBPG)) == [:E_1, :V_1, :V_2, :tgt_E_1, :src_E_1]

p₀ = @acset Graph begin
  V = 2
  E = 2
  src=[1,2]
  tgt=[2,1]
end

ThPetri, obmap, hommap = CatElements.presentation(elements(p₀))
@test Symbol.(generators(ThPetri)) == [:E_1, :E_2, :V_1, :V_2, :src_E_2, :tgt_E_1, :tgt_E_2, :src_E_1]

@acset_type Petri(ThPetri)

sir_eltsch = @acset Petri begin
  V_1 = 3
  V_2 = 2
  E_1 = 3
  E_2 = 3

  src_E_1 = [1,2,2]
  tgt_E_1 = [1,1,2]

  src_E_2 = [1,1,2]
  tgt_E_2 = [2,2,3]
end

@testset "DiBipartite Elements" begin
  @test length(generators(ThPetri)) == 8
  @test nparts(sir_eltsch, :V_1) == 3
  @test nparts(sir_eltsch, :V_2) == 2
  @test nparts(sir_eltsch, :E_1) == 3
  @test nparts(sir_eltsch, :E_2) == 3
  @test sir_eltsch[:, :src_E_2] == [1,1,2]
end

end


function fifth_power_digit_sum(n::Int)

	a=Int[];

	while(n>=1)
		push!(a,n%10);
		n=n-n%10;
		n/=10;
		n=Int(n)
	end
	ans=0;
	for i in a
		ans+=i^5
	end
	return ans
end

function compute()
	ans=0
	for i=2:1000000
		if i==fifth_power_digit_sum(i)
			ans+=i
		end
	end
	return ans
end

println(compute())


# # Envelopes

"""
    struct ConstantProfile

This is the trivial profile
```math
envelope(z, t) = 1
```
which gives an infinite duration for the electromagnetic field configuration (we cannot
call this a laser pusle since a pulse implicitely has a finite duration).
"""
struct ConstantProfile <: AbstractTemporalProfile end

@doc """
    GaussProfile{V,T,L}

This envelope provides a finite duration for the laser pulse and thus can provide a more
realistic description of an actual laser pulse.
```math
envelope(z, t) = \\exp\\left[-\\left(\\frac{φ}{τ}\\right)^2\\right],
```
where
```math
\\varphi = (t - t_0) - \\frac{z - z_0}{c}\\,,
```
and
- `c` is the speed of light
- `τ` is the duration of the pulse (FWHM) and has the default value 18.02fs
- `t₀` is the origin of the time axis and it is 0 by default
- `z₀` is the initial position of the intensity peak and has the default value `-4*τ*c`
"""
GaussProfile

struct GaussProfile{T,IT,L} <: AbstractTemporalProfile
    inv_τ::IT
    t₀::T
    z₀::L
end

function GaussProfile(;τ, t₀=zero(τ), z₀)
    GaussProfile(inv(τ), t₀, z₀)
end

@doc """
    Cos²Profile{V,T,L}

This envelope provides a finite duration for the laser pulse and thus can provide a more
realistic description of an actual laser pulse.
```math
envelope(z, t) =
    \\begin{cases}
    \\cos\\left[\\left(\\frac{φ}{τ}\\right)\\right]^2, & \\text{for } |t - t₀| < τ / 2\\\\
    0 \\,, & \\text{otherwise}
    \\end{cases}
```
where
```math
\\varphi = (t - t_0) - \\frac{z - z_0}{c}\\,,
```
and
- `c` is the speed of light
- `τ` is the duration of the pulse and has the default value 18.02fs
- `t₀` is the origin of the time axis and it is 0 by default
- `z₀` is the initial position of the intensity peak and is 0 by default
"""
Cos²Profile

struct Cos²Profile{T,IT,L} <: AbstractTemporalProfile
    τ::T
    inv_τ::IT
    t₀::T
    z₀::L
end

function Cos²Profile(;τ, t₀=zero(τ), z₀)
    t₀, τ = promote(t₀, τ)
    Cos²Profile(τ, inv(τ), t₀, z₀)
end

@doc """
    QuasiRectangularProfile{V,T,L}

This envelope produces a pulse with a predominant constant part of width ``Δz``
which could offer better results than the Gaussian profile in the paraxial limit
(which is considered for the spatial profiles). The shape of the envelope is given by
```math
envelope(z, t) =
    \\begin{cases}
    \\exp\\left[-\\left(\\frac{φ + Δt/2}{τ}\\right)^2\\right], & \\text{for } φ ≤ \\frac{Δt}{2}\\\\
    1\\,, & \\text{for } \\text{otherwise}\\\\
    \\exp\\left[-\\left(\\frac{φ - Δt/2}{τ}\\right)^2\\right], & \\ φ > \\frac{Δt}{2}
    \\end{cases}
```
where
```math
\\varphi = (t - t_0) - \\frac{z - z_0}{c}\\,,
```
and
- `c` is the speed of light
- `τ` is the duration of the pulse (FWHM) and has the default value 18.02fs
- `t₀` is the origin of the time axis and it is 0 by default
- `z₀` is the initial position of the intensity peak and has the default value `-4*τ*c`
- `Δt` is the duration of the flat part of the profile and the default value `10*τ`
"""
QuasiRectangularProfile

struct QuasiRectangularProfile{T,IT,L} <: AbstractTemporalProfile
    inv_τ::IT
    t₀::T
    z₀::L
    Δt::T
end

function QuasiRectangularProfile(;τ, t₀=zero(τ), z₀, Δt=10τ)
    τ, t₀, Δt = promote(τ, t₀, Δt)
    QuasiRectangularProfile(inv(τ), t₀, z₀, Δt)
end

"""
    g(z, t, par)

The time dependence of the fields defined by this package is given by
```math
g(z, t) = \\exp(i ω t) envelope(z, t),
```
where
- `z` and `t` are the position on the ``Oz`` axis and the time
- `par` are the laser parameters which pass the corresponding profile to the envelope
and
- ``ω`` is the angular frequency of the laser pulse
- ``envelope(z, t)`` is a function that can be used to control the duration of the pulse
"""
function g(z, t, laser; inv_c)
    profile = laser.profile
    ω = immutable_cache(laser, :ω)

    exp(im*ω*t) * envelope(profile, z, t; inv_c)
end

@inline envelope(::ConstantProfile, z, t; inv_c) = 1

@inline function envelope(profile::GaussProfile, z, t; inv_c)
    @unpack inv_τ, t₀, z₀ = profile
    φ = (t - t₀) - (z - z₀) * inv_c

    exp(-(φ * inv_τ)^2)
end

@inline function envelope(profile::Cos²Profile, z, t; inv_c)
    @unpack inv_τ, τ, t₀, z₀ = profile
    φ = π*((t - t₀) - (z - z₀) * inv_c)

    if abs(φ) / π < τ / 2
        cos(-(φ * inv_τ))^2
    else
        zero(t * inv_τ)
    end
end

@inline function envelope(profile::QuasiRectangularProfile, z, t; inv_c)
    @unpack inv_τ, t₀, z₀, Δt = profile
    φ = (t - t₀/2) - (z - z₀/2) * inv_c

    if φ < -Δt/2
        exp(-((φ + Δt/2) * inv_τ)^2)
    elseif φ < Δt/2
        one(φ)
    else
        exp(-((φ - Δt/2) * inv_τ)^2)
    end
end


#script for solving the acoustic duct problem
# -a0(x)*\Delta u(x,y)+a1(x)*u(x,y)=f(x,y) for (x,y) ∈ Ω:= (0,2)×(0,1)
# Here a0(x,y)=-1, a1(x,y)=k^2, f(x,y)=0
# with Neumann boundary conditions ∂u/∂x(x,y)=cos(mπy) for x = 0 and
# Robin boundary conditions ∂u/∂x(x,y) + iku = 0 for x=2

include("../src/Solver2D.jl")
include("../src/SplinePlotting.jl")

a0(x,y) = 1
a1(x,y) = -k^2
f(x,y) = 0
k = 40
m = 2
kx = sqrt(k^2 - (m*π)^2)
LHS = [im*kx -im*kx; (k-kx)*exp(-2*im*kx) (k+kx)*exp(2*im*kx)]
RHS = [1, 0]
A = LHS\RHS

exact_sol(x,y) = cos(m*π*y)*(A[1]*exp(-im*kx*x)+A[2]*exp(im*kx*x))
deriv_exact_sol(x,y) = [cos(m*π*y)*(A[1]*(-im)*kx*exp(-im*kx*x)+A[2]*im*kx*exp(im*kx*x)),
                        -sin(m*π*y)*(m*π)*(A[1]*exp(-im*kx*x)+A[2]*exp(im*kx*x))]

#Define the boundary conditions
u_bound_left(x,y) = cos(m*π*y)
u_bound_right(x,y) = 0.
bound_left = Boundary2D("Neumann", "Left", u_bound_left)
bound_right = Boundary2D("Robin", "Right", -im*k, u_bound_right)
bound_all = [bound_left, bound_right]

#Define the domain geometry
cornerLowerLeft = [0., 0.]
lengthx = 2.
lengthy = 1.
degP = [3, 3]
numSubdiv = 80
nrb = nrbsquare(cornerLowerLeft, lengthx, lengthy, degP, numSubdiv)
IGAmesh = genMesh(nrb)
#plotBasisParam(IGAmesh)

gauss_rule = [genGaussLegendre(degP[1]+1), genGaussLegendre(degP[2]+1)]
stiff = assemble_stiff2D(IGAmesh, a0, gauss_rule)
mass = assemble_mass2D(IGAmesh, a1, gauss_rule)
rhs = assemble_rhs2D(IGAmesh, f, gauss_rule)
bcdof_all, elem_all = classifyBoundary2D(IGAmesh)
lhs = mass + stiff
lhs, rhs = applyBCnurbs(IGAmesh, bound_all, lhs, rhs, bcdof_all, elem_all, gauss_rule)

println("Solving linear system")
@time sol0 = lhs\rhs
println("Plotting...")
@time plotSol(IGAmesh, real(sol0), "AcousticDuctReal")
@time plotSolError(IGAmesh, real(sol0), exact_sol, "AcousticDuctRealError")
println("Computing error")
@time relL2Err, relH1Err = compErrorNorm(IGAmesh, real(sol0), exact_sol, deriv_exact_sol, a0, gauss_rule)
println("Relative L2-norm error is $relL2Err")
println("Relative H1-seminorm error is $relH1Err")
print("Done!")


"""
Return the Inorder traversal of a complete binary tree with given `height`.
(Root node will be `0` and placed in the middle of the array
"""
function binary_inorder(height::Integer)::Vector{Int}
    n = (1 << Int(height)) - 1
    out = Vector{Int}(undef, n)

    s = 2
    while n >= 1
        t = s >> 1
        b, a = n-1, n >> 1
        for (i, j) in enumerate(a:b)
            out[i*s - t] = j
        end
        n >>= 1
        s <<= 1
    end

    return out
end


# There are exactly ten ways of selecting three from five, 12345:
#
# 123, 124, 125, 134, 135, 145, 234, 235, 245, and 345
#
# In combinatorics, we use the notation, (5 3)=10.
#
# In general, (n r)=n!/(r!(n−r)!), where r≤n, n!=n×(n−1)×...×3×2×1, and 0!=10!=1.
#
# It is not until n=23 that a value exceeds one-million: (23 10)=1144066.
#
# How many, not necessarily distinct, values of (n r) for 1≤n≤100, are greater
# than one-million?

using ProjectEulerSolutions

# We note that (n r) = (n (n-r)), and that when r = n/2, the value of (n r) is
# maximal.  Thus we calculate from (n r/2-k) for k until we get a value of more
# than 1 million.  We save more time by further constraining the range of k
# values.  This is mostly a book-keeping exercise.
function p053solution(n_max::Integer=3, thresh=1_000)::Integer
    n_min = 23
    r = 3

    total = 0
    for n in n_max:-1:n_min
        # The value of r where we surpass the threshold will only get larger
        while binomial(n, r) < thresh
            r += 1
        end
        total += (n - r * 2  + 1)
    end
    return total
end

p053 = Problems.Problem(p053solution)

Problems.benchmark(p053, 100, 1_000_000)


# Legendre-Gauss quadrature weights and nodes
# This code was originally written for Matlab
#  by Greg von Winckel - 02/25/2004
# posted on Mathworks website and licensed under BSD License
# http://fr.mathworks.com/matlabcentral/fileexchange/4540-legendre-gauss-quadrature-weights-and-nodes
#
#

function lgwt(N,a,b)

	N=N-1;
	N1=N+1; N2=N+2;
	xu=linspace(-1,1,N1)';

	y=(cos((2*(0:N)'+1)*pi/(2*N+2))+(0.27/N1)*sin(pi*xu*N/N2))';

	L=zeros(N1,N2);

	Lp=zeros(N1,N2);
	y0=2;

	while maximum(abs(y-y0))>eps()
		L[:,1]=1;
		Lp[:,1]=0;
		L[:,2]=y;
		for k=2:N1
			L[:,k+1]=( (2*k-1)*y.*L[:,k]-(k-1)*L[:,k-1] )/k;
		end
		Lp=(N2)*( L[:,N1]-y.*L[:,N2] )./(1-y.^2);
		y0=y;
		y=y0-L[:,N2]./Lp;
	end

	x=(a*(1-y)+b*(1+y))/2;

	w=(b-a)./((1-y.^2).*Lp.^2)*(N2/N1)^2;

	return x,w
end


using Test

# Thin category
###############

x,y,z = Ob(FreeThinCategory, :x, :y, :z)
f, g = Hom(:f, x, y), Hom(:g, y, z)
@test dom(compose(f,g)) == x
@test codom(compose(f,g)) == z

x,y,z,zz = Ob(FreeThinSymmetricMonoidalCategory, :x, :y, :z, :zz)
f, g, h= Hom(:f, x, y), Hom(:g, y, z), Hom(:h, otimes(y,y), z)
j = Hom(:j, y⊗z, zz)
@test dom(compose(f,g)) == x
@test codom(compose(f,g)) == z
@test dom(otimes(f,g)) == otimes(x,y)
@test codom(otimes(f,g)) == otimes(y,z)
@test dom((f⊗id(y))⋅h) == x⊗y
@test codom((f⊗id(y))⋅h) == z
@test dom((f⊗g)⋅j) == x⊗y
@test codom((f⊗g)⋅j) == zz

# Preorder
##########

x,y,z = (El(FreePreorder.El, sym) for sym in (:x, :y, :z))
f, g = Leq(:f, x, y), Leq(:g, y, z)
@test lhs(reflexive(x)) == x
@test rhs(reflexive(x)) == x
@test lhs(transitive(f,g)) == x
@test rhs(transitive(f,g)) == z


using LinearAlgebra

# coefficients from PD of π functions
π = [70 -30 0;
     -30 140 -50;
     0 -50 50]
# Force Vector
f = Float64[0 100 60]
F = transpose(f)

# Displacement vector
u = π\F
println("Nodal displacement vector (mm):")
display(u)


# hyperbolic_expansion.jl

module hypergeometric_expansion
export hypergeom_exp

include("arb_hypgeom.jl")
import .arb_hypgeom
using ArbNumerics

function arb_binomial(x, y; prec = 64)
    setextrabits(0)
    setprecision(ArbReal, prec)
    x = ArbReal("$x")
    y = ArbReal("$y")
    one = ArbReal("1")
    ArbNumerics.gamma(x + one) / (ArbNumerics.gamma(y + one) * ArbNumerics.gamma(x - y + one))
end

"""
This function implements the binomial expansion of the denominator of the integrals:

    ∫_0^1 cosh(zt) / (4a^2/β^2 - t^2)^(n+3/2) dt
    ∫_0^1 sinh(zt) / (4a^2/β^2 - t^2)^(n+3/2) dt

which is:

    ∑_{m=0}^∞ C(-n-3/2, m) (-1)^m (β/2a)^{2n+2m+3} ∫_0^1 cosh(zt) t^{2m} dt
    ∑_{m=0}^∞ C(-n-3/2, m) (-1)^m (β/2a)^{2n+2m+3} ∫_0^1 sinh(zt) t^{2m} dt.

and also substitutes the integrals for:

    ∫_0^1 cosh(zt) t^{2m} dt = ∑_{t=0}^∞ z^{2t} / ((2t+2m+1)⋅(2t)!)
    ∫_0^1 sinh(zt) t^{2m} dt = ∑_{t=0}^∞ z^{2t+1} / ((2t+2m+2)⋅(2t+1)!)

by replacing those integrals with one_F_two_fast() from arb_hypgeom.jl. h=0 gives cosh version, h=1 gives sinh version.
"""
function hypergeom_exp(z, n, β, a, h; prec = 64)
# h = 0 gives cosh, h = 1 gives sinh.

    # Initialise precision of ArbReal to prec.
    p = prec
    setextrabits(0)
    setprecision(ArbReal, p)

    z = ArbReal("$z")
    n = ArbReal("$n")
    β = ArbReal("$β")
    a = ArbReal("$a")

    m = ArbReal("0")
    result = ArbReal("0.0")
    term = ArbReal("1.0")
    err = eps(result)  # Machine accuracy of specified precision prec.

    while abs(midpoint(term)) > err * abs(midpoint(result))

        term = ArbReal(β * arb_binomial(-n - 3/2, m; prec = p) * (-1)^m * (β / (2 * a))^(2 * m) * arb_hypgeom.one_f_two_fast(z, m, h; prec = p) / a^(n + 2))
        result += term
        # println("term: m = ", m, "\nterm value: ", ball(ArbReal(term, bits = prec)), "\ncumulant result: ", ball(ArbReal(result, bits = prec)), "\n")
        m += ArbReal("1")

        # Double precision if rounding error in result exceeds accuracy specified by prec.
        if radius(result) > err * abs(midpoint(result))
            p *= 2
            setprecision(ArbReal, p)
            # println("Not precise enough. Error = ", abs(radius(result)/midpoint(result)), " > ", err, ". Increasing precision to ", p, " bits.\n")

            n = ArbReal("$n")
            z = ArbReal("$z")
            β = ArbReal("$β")
            a = ArbReal("$a")

            m = ArbReal("0")
            result = ArbReal("0.0")
            term = ArbReal("1.0")
        end
    end
    # println("z: ", ArbReal(z, bits = prec), ". Final result: ", ArbReal(result, bits = prec))
    ArbReal(result, bits = prec)
end

# β = ArbReal("2.0")
# α = ArbReal("7.0")
# v = ArbReal("5.8")
# w = ArbReal("1.6")
# R = ArbReal((v^2 - w^2) / (w^2 * v))
# a = ArbReal(sqrt(β^2 / 4 + R * β * coth(β * v / 2)))
# z = ArbReal("90.61")
# n = ArbReal("12.0")
# @time c = hypergeom_exp(z, n, β, a, 0; prec = 64)
# @show(c)
end # end module


using BenchmarkTools

# random walk using Julia for loop
function random_walk_for_loop(start, states, n)
    for _ in 1:n
        start *= states
    return start
    end
end

# random walk using matrix power
function random_walk_matrix_power(start, states, n)
    return start * states ^ n
end

# random walk until stationary state
# this assumes that there exists a stationary state
function get_stationary_state(start, states)
    π = start
    prev = 0
    n = 0
    while π != prev
      prev = π
      π *= states
      n += 1
    end
    return n, π
  end

function get_stationary_state2(start, states)
    π = start
    prev = start
    n = 1
    while true
        π *= states
        if π != prev
            prev = π
        else
            return n, π
        end
    n += 1
    end
end


"""
1-D Gradient

    gradient(a::Array, dx::Float64=1)
    :param f: Vector
    :param dx: stepsize

    :return g: gradient
"""
function gradient(f::Array{Float64, 1}, dx::Float64=1.)
    n = length(f)
    g = zeros(n)
    h = dx*(1:n)
    # Take forward differences on left and right edges
    g[1] = (f[2] - f[1])/(h[2]-h[1])
    g[n] = (f[n] - f[n-1])/(h[end]-h[end-1])

    # Take centered differences on interior points
    h = h[3:n]-h[1:(n-2)]
    g[2:(n-1)] = (f[3:n]-f[1:(n-2)])/h[1]

    return g
end

"""
2-D Gradient

    gradient2(a::Array, dx::Float64=1, dy::Float=1)
    :param a: matrix
    :param dx: stepsize
    :param dy: stepsize

    :return dxdu: derivatives along x
    :return dydv: derivatives along y
"""
function gradient2(a::Array, dx::Float64=1., dy::Float64=1.)
    m,n = size(a);
    dxdu = zeros(m,n);
    dydv = zeros(m,n);

    for i=1:m
        dxdu[i,:] = gradient(vec(a[i,:]), dx)
    end

    for i=1:n
        dydv[:,i] = gradient(a[:,i], dy)
    end

    return dxdu, dydv
end

"""
Creates Rectangular Grid in 2-D space

    meshgrid(a::LinRange,b::LinRange)
"""
function meshgrid(a::LinRange,b::LinRange)
    grid_a = [i for i in a, j in b]';
    grid_b = [j for i in a, j in b]';

    return grid_a, grid_b
end

"""
Linear interpolation

    approx(xd, yd, xi)
    :param xd: x samples
    :param yd: response samples
    :param xi: new x samples
"""
function approx(xd, yd, xi)
    nd = length(xd);
    ni = length(xi);

    yi = zeros(ni);
    for i in 1:ni
        if (xi[i] <= xd[1])
            t = (xi[i]-xd[1]) / (xd[2] - xd[1]);
            yi[i] = (1.0 - t) * yd[1] + t * yd[2];
        elseif (xd[nd] <= xi[i])
            t = (xi[i] - xd[nd-1]) / (xd[nd] - xd[nd-1]);
            yi[i] = (1.0 - t) * yd[nd-1] + t * yd[nd];
        else
            for k in 2:nd
                if (xd[k-1] <= xi[i] && xi[i] <= xd[k])
                    t = (xi[i] - xd[k-1]) / (xd[k] - xd[k-1]);
                    yi[i] = (1.0 - t) * yd[k-1] + t * yd[k];
                    break
                end
            end
        end
    end

    return yi
end

"""
Trapezoidal Integration

    trapz(x, y, dim=1)
    :param x: vector of time samples
    :param y: array of response samples
    :param dim: dimension along which to integrate
"""
function trapz(x::Array{Float64, 1}, y::Array{Float64}, dim::Integer=1)
    perm = [dim:max(ndims(y),dim); 1:dim-1];
    y = permutedims(y, perm);
    if ndims(y) == 1
        m = 1;
    else
        m = size(y,1);
    end

    if m == 1
        M = length(y);
        out = sum(diff(x).*(y[1:M-1] + y[2:M])/2.0);
    else
        out = transpose(diff(x)) * (reshape(y,Val(2))[1:m-1,:] + reshape(y,Val(2))[2:m,:])/2.0;
        siz = size(y); siz = collect(siz); siz[1] = 1;
        out = reshape(out, tuple(siz...));
        out = permutedims(out, invperm(perm));
        ind = findall(collect(size(out)).==1);
        out = dropdims(out,dims=ind[1]);
        if length(out) == 1;
            out = out[1];
        end
    end

    return out
end

"""
Cumulative Trapezoidal Integration

    cumtrapz(x, y, dim=1)
    :param x: vector describing time samples
    :param y: array describing response
    :param dim: dimension to integrate over
"""
function cumtrapz(x::Array{Float64, 1}, y::Array{Float64}, dim::Integer=1)
    perm = [dim:max(length(size(y)),dim); 1:dim-1];
    y = permutedims(y, perm);
    if ndims(y) == 1
        n = 1;
        m = length(y);
    else
        m, n = size(y);
    end

    if n == 1
        dt = diff(x)/2.0;
        z = [0; cumsum(dt.*(y[1:(m-1)] + y[2:m]))];
    else
        dt = repeat(diff(x)/2.0,1,n);
        z = [zeros(1,n); cumsum(dt.*(y[1:(m-1), :] + y[2:m, :]),dims=1)];
        z = permutedims(z, invperm(perm));
    end

    return z
end

"""
Cumulative Trapezoidal Integration using midpoint

    cumtrapzmid(x, y, c)
    :param x: time samples
    :param y: resposne samples
    :param c: midpoint
    :param mid: midpoint location
"""
function cumtrapzmid(x, y, c, mid)
    a = length(x);

    # case < mid
    fn = zeros(a);
    tmpx = x[(mid-1):-1:1];
    tmpy = y[(mid-1):-1:1];
    tmp = c .+ cumtrapz(tmpx, tmpy);
    fn[1:(mid-1)] = reverse(tmp);

    # case >= mid
    fn[mid:a] = c .+ cumtrapz(x[mid:a],y[mid:a]);

    return fn

end

"""
Multivariate Normal random number generation

    mvnrand(mu, C, n)
    :param mu: mean vector
    :param C: covariance matrix
    :param n: number of samples
"""
function mvnrand(mu, C, n)
    tmp = cholesky(C, :U, Val{true});
    R = tmp[:U];
    R = Array(R);
    R = R[:, tmp.piv];
    retval = randn(n, size(R,1)) * R;
    retval += transpose(repeat(mu, 1, n));
    return transpose(retval)
end

"""
Linear interpolation when response contains flat regions

    interp1_flat(x, y, xx)
    :param x: time samples
    :param y: response samples
    :param xx: new time samples
"""
function interp1_flat(x,y,xx)
    flat = find(diff(x).<=0);
    n = length(flat);

    if n==0
        tmp = interpolate((x,), y, Gridded(Linear()))
        yy = tmp[xx];
    else
        yy = zeros(size(xx));
        i1 = 1;
        if flat[1] == 1
            i2 = 1;
            j = xx.==x[i2];
            yy[j] = minimum(y[i2:i2+1]);
        else
            i2 = flat[1];
            j = (xx.>=x[i1]) & (xx.<=x[i2]);
            tmp = interpolate((x[i1:i2],), y[i1:i2], Gridded(Linear()))
            yy[j] = tmp[xx[j]];
            i1 = copy(i2);
        end
        for k = 2:n
            i2 = flat[k];
            if i2 > i1+1
                j = (xx.>=x[i1]) & (xx.<=x[i2]);
                yi = interpolate((x[i1+1:i2],), y[i1+1:i2], Gridded(Linear()))
                yy[j] = tmp[xx[j]];
            end
            j = xx.==x[i2];
            yy[j] = minimum(y[i2:i2+1]);
            i1 = copy(i2);
        end
        i2 = length(x);
        j = (xx.>=x[i1]) & (xx.<=x[i2]);
        if i1+1 == i2
            yy[j] = y[i2];
        else
            tmp = interpolate((x[i1+1:i2],), y[i1+1:i2], Gridded(Linear()))
            yy[j] = tmp[xx[j]];
        end
    end

    return yy
end

function basis_fourier(f_domain::Vector, numBasis::Int64, fourier_p::Int64)
    result = Array{Float64}(undef,length(f_domain),2*numBasis)
    for i in 1:2*numBasis
        j = ceil(i/2)
        if ((i % 2) == 1)
            result[:,i] = sqrt(2) * sin.(2*j*pi*f_domain/fourier_p)
        end
        if ((i % 2) == 0)
            result[:,i] = sqrt(2) * cos.(2*j*pi*f_domain/fourier_p)
        end
    end

    out = basis(f_domain,result)
    return out
end

function f_predictfunction(f::func, at; deriv=0, method="linear")
    if method=="linear"
        if (deriv==0)
            result = approx(f.x, f.y, at)
        end
        if (deriv == 1)
            fmod = approx(f.x, f.y, at)
            diffy1 = [0; diff(fmod)]
            diffy2 = [diff(fmod); 0]
            diffx1 = [0; diff(at)]
            diffx2 = [diff(at); 0]

            result = (diffy2 + diffy1) ./ (diffx2 + diffx1)
        end
    else
        error("Method not implemented")
    end

    out = func(at, result)
    return out
end

function f_basistofunction(f_domain, coef, cbasis::basis; coefconst=0)
    if (size(cbasis.matrix,2)<length(coef))
        error("coefficients exceeds basis functions")
    end

    y1 = cbasis.matrix[:,(1:length(coef))] * coef .+ coefconst
    result = func(cbasis.x, y1)
    y1 = f_predictfunction(result, f_domain)

    return y1
end


using Test, QBase

@testset "./src/constructors/brakets.jl" begin

@testset "bloch_qubit_ket()" begin
    @testset "complex bloch qubit" begin
        ψ1 = bloch_qubit_ket(π/2,π/2)

        @test ψ1 isa Ket{Complex{Float64}}
        @test ψ1 ≈ [1,im]/sqrt(2)

        ψ2 = bloch_qubit_ket(π/2,3π/2)

        @test ψ2 isa Ket{Complex{Float64}}
        @test ψ2 ≈ [1,-im]/sqrt(2)

        @test_throws DomainError bloch_qubit_ket(-0.1,π/2)
        @test_throws DomainError bloch_qubit_ket(15π/7,π/2)
        @test_throws DomainError bloch_qubit_ket(π/2,-0.1)
        @test_throws DomainError bloch_qubit_ket(π/2,15π/7)
    end

    @testset "real bloch qubit" begin
        ψ1 = bloch_qubit_ket(π/2)

        @test ψ1 isa Ket{Float64}
        @test ψ1 ≈ [1,1]/sqrt(2)

        ψ2 = bloch_qubit_ket(-π/2)

        @test ψ2 isa Ket{Float64}
        @test ψ2 ≈ [1,-1]/sqrt(2)

        @test_throws DomainError bloch_qubit_ket(-2π - 0.1)
        @test_throws DomainError bloch_qubit_ket(2π + 0.1)
    end

    @testset "equivalence between real and complex versions" for θ in 0:0.2:2π
        @test bloch_qubit_ket(θ) ≈ bloch_qubit_ket(θ,0)
        @test isapprox(bloch_qubit_ket(-θ), bloch_qubit_ket(θ,π), atol=1e-7)
    end
end

@testset "computational_basis_kets()" begin
    comp_basis_kets = computational_basis_kets(3)
    @test comp_basis_kets isa Vector{Ket{Int64}}
    @test comp_basis_kets == [[1,0,0],[0,1,0],[0,0,1]]
    @test computational_basis_kets(8) == [
        [1,0,0,0,0,0,0,0],[0,1,0,0,0,0,0,0],[0,0,1,0,0,0,0,0],[0,0,0,1,0,0,0,0],
        [0,0,0,0,1,0,0,0],[0,0,0,0,0,1,0,0],[0,0,0,0,0,0,1,0],[0,0,0,0,0,0,0,1]
    ]
end

@testset "mirror_symmetric_qubit_kets()" begin
    @test mirror_symmetric_qubit_kets(0) isa Vector{Ket{Float64}}

    @test isapprox(mirror_symmetric_qubit_kets(0), [[1,0],[1,0],[1,0]])
    @test isapprox(mirror_symmetric_qubit_kets(π/4), [[1,0],[1,1]/sqrt(2),[1,-1]/sqrt(2)])
    @test isapprox(mirror_symmetric_qubit_kets(π/3),[[1,0],[0.5,(sqrt(3)/2)],[0.5,-sqrt(3)/2]])
    @test isapprox(mirror_symmetric_qubit_kets(π/2), [[1,0],[0,1],[0,-1]], atol=1e-16)

    @test_throws DomainError mirror_symmetric_qubit_kets(-0.1)
    @test_throws DomainError mirror_symmetric_qubit_kets(3π/5)
end

@testset "planar_symmetric_qubit_kets()" begin
    @testset "orthogonal qubits" begin
        kets = planar_symmetric_qubit_kets(2)
        @test kets isa Vector{Ket{Float64}}
        @test isapprox(kets, [[1,0],[0,1]], atol=1e-7)
    end

    @testset "scanning over cases verifying right number are computed" begin
        for n in 2:100
            kets = planar_symmetric_qubit_kets(n)

            @test length(kets) == n
            @test kets isa Vector{Ket{Float64}}
        end
    end

    @test_throws(DomainError, planar_symmetric_qubit_kets(1))
end

@testset "trine_qubit_kets()" begin
    trine_kets = trine_qubit_kets()
    @test trine_kets isa Vector{Ket{Float64}}
    @test isapprox(trine_kets, mirror_symmetric_qubit_kets(π/3), atol=1e-7)
    @test isapprox(trine_kets, planar_symmetric_qubit_kets(3), atol=1e-7)
end

@testset "sic_qubit_kets()" begin
    sic_kets = sic_qubit_kets()
    @test sic_kets isa Vector{Ket{Complex{Float64}}}

    θ_sic = 2*acos(1/sqrt(3))

    @test sic_kets[1] ≈ bloch_qubit_ket(0, 0)
    @test sic_kets[2] ≈ bloch_qubit_ket(θ_sic, 0)
    @test sic_kets[3] ≈ bloch_qubit_ket(θ_sic, 2π/3)
    @test sic_kets[4] ≈ bloch_qubit_ket(θ_sic, 4π/3)
end

@testset "bb84_qubit_kets" begin
    bb84_kets = bb84_qubit_kets()
    @test bb84_kets isa Vector{Ket{Float64}}
    @test isapprox(bb84_kets, planar_symmetric_qubit_kets(4), atol=1e-7)
end

@testset "bell_kets()" begin
    kets = bell_kets()
    k0 = [1,0]
    k1 = [0,1]

    @test kets isa Vector{Ket{Float64}}
    @test length(kets) == 4
    @test kets[1] == 1/sqrt(2)*(kron(k0,k0)+kron(k1,k1))
    @test kets[2] == 1/sqrt(2)*(kron(k0,k0)-kron(k1,k1))
    @test kets[3] == 1/sqrt(2)*(kron(k0,k1)+kron(k1,k0))
    @test kets[4] == 1/sqrt(2)*(kron(k0,k1)-kron(k1,k0))
end

@testset "generalized_bell_kets()" begin
    @testset "qubit bell states" begin
        kets = generalized_bell_kets(2)

        @test kets isa Vector{Ket{Complex{Float64}}}
        @test isapprox(kets, bell_kets(), atol=1e-7)
    end

    @testset "qutrit bell states" begin
        kets = generalized_bell_kets(3)

        basis = computational_basis_kets(3)
        ψ1 = (kron(basis[1],basis[1]) + kron(basis[2],basis[2]) + kron(basis[3],basis[3]))/sqrt(3)
        ψ2 = (kron(basis[1],basis[1]) + exp(im*2*π/3)*kron(basis[2],basis[2]) + exp(im*4*π/3)*kron(basis[3],basis[3]))/sqrt(3)
        ψ3 = (kron(basis[1],basis[1]) + exp(im*4*π/3)*kron(basis[2],basis[2]) + exp(im*2*π/3)*kron(basis[3],basis[3]))/sqrt(3)

        ψ4 = (kron(basis[1],basis[2]) + kron(basis[2],basis[3]) + kron(basis[3],basis[1]))/sqrt(3)
        ψ5 = (kron(basis[1],basis[2]) + exp(im*2*π/3)*kron(basis[2],basis[3]) + exp(im*4*π/3)*kron(basis[3],basis[1]))/sqrt(3)
        ψ6 = (kron(basis[1],basis[2]) + exp(im*4*π/3)*kron(basis[2],basis[3]) + exp(im*2*π/3)*kron(basis[3],basis[1]))/sqrt(3)

        ψ7 = (kron(basis[1],basis[3]) + kron(basis[2],basis[1]) + kron(basis[3],basis[2]))/sqrt(3)
        ψ8 = (kron(basis[1],basis[3]) + exp(im*2*π/3)*kron(basis[2],basis[1]) + exp(im*4*π/3)*kron(basis[3],basis[2]))/sqrt(3)
        ψ9 = (kron(basis[1],basis[3]) + exp(im*4*π/3)*kron(basis[2],basis[1]) + exp(im*2*π/3)*kron(basis[3],basis[2]))/sqrt(3)

        ψ_match = [ψ1,ψ2,ψ3,ψ4,ψ5,ψ6,ψ7,ψ8,ψ9]

        @test kets isa Vector{Ket{Complex{Float64}}}
        @test kets ≈ ψ_match
    end
end

end


# elastic-net regularization

export ElasticNet

"""
**Elastic-net regularization**

    ElasticNet(μ=1.0, λ=1.0)

Returns the function
```math
f(x) = μ\\|x\\|_1 + (λ/2)\\|x\\|^2,
```
for nonnegative parameters `μ` and `λ`.
"""
struct ElasticNet{R <: Real} <: ProximableFunction
    mu::R
    lambda::R
    function ElasticNet{R}(mu::R, lambda::R) where {R <: Real}
        if lambda < 0 || mu < 0
            error("parameters `μ` and `λ` must be nonnegative")
        else
            new(mu, lambda)
        end
    end
end

is_separable(f::ElasticNet) = true
is_prox_accurate(f::ElasticNet) = true
is_convex(f::ElasticNet) = true

ElasticNet(mu::R=1.0, lambda::R=1.0) where {R <: Real} = ElasticNet{R}(mu, lambda)

function (f::ElasticNet)(x::AbstractArray{R}) where R <: RealOrComplex
    return f.mu*norm(x,1) + (f.lambda/2)*norm(x,2)^2
end

function prox!(y::AbstractArray{R}, f::ElasticNet{R}, x::AbstractArray{R}, gamma::R=R(1)) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    gm = gamma*f.mu
    gl = gamma*f.lambda
    for i in eachindex(x)
        y[i] = (x[i] + (x[i] <= -gm ? gm : (x[i] >= gm ? -gm : -x[i])))/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function prox!(y::AbstractArray{R}, f::ElasticNet{R}, x::AbstractArray{R}, gamma::AbstractArray{R}) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    for i in eachindex(x)
        gm = gamma[i]*f.mu
        gl = gamma[i]*f.lambda
        y[i] = (x[i] + (x[i] <= -gm ? gm : (x[i] >= gm ? -gm : -x[i])))/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function prox!(y::AbstractArray{Complex{R}}, f::ElasticNet{R}, x::AbstractArray{Complex{R}}, gamma::R=R(1)) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    gm = gamma*f.mu
    gl = gamma*f.lambda
    for i in eachindex(x)
        y[i] = sign(x[i])*max(0, abs(x[i]) - gm)/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function prox!(y::AbstractArray{Complex{R}}, f::ElasticNet{R}, x::AbstractArray{Complex{R}}, gamma::AbstractArray{R}) where R <: Real
    sqnorm2x = R(0)
    norm1x = R(0)
    for i in eachindex(x)
        gm = gamma[i]*f.mu
        gl = gamma[i]*f.lambda
        y[i] = sign(x[i])*max(0, abs(x[i]) - gm)/(1 + gl)
        sqnorm2x += abs2(y[i])
        norm1x += abs(y[i])
    end
    return f.mu*norm1x + (f.lambda/2)*sqnorm2x
end

function gradient!(y::AbstractArray{T}, f::ElasticNet{R}, x::AbstractArray{T}) where {T <: RealOrComplex, R <: Real}
    # Gradient of 1 norm
    y .= f.mu.*sign.(x)
    # Gradient of 2 norm
    y .+= f.lambda.*x
    return f.mu*norm(x,1) + (f.lambda/2)*norm(x,2)^2
end

fun_name(f::ElasticNet) = "elastic-net regularization"
fun_dom(f::ElasticNet) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::ElasticNet) = "x ↦ μ||x||_1 + (λ/2)||x||²"
fun_params(f::ElasticNet) = "μ = $(f.mu), λ = $(f.lambda)"

function prox_naive(f::ElasticNet, x::AbstractArray{R}, gamma::Real=1.0) where R <: RealOrComplex
    uz = max.(0, abs.(x) .- gamma*f.mu)/(1 + f.lambda * gamma)
    return sign.(x) .* uz, f.mu * norm(uz,1) + (f.lambda/2) * norm(uz)^2
end

function prox_naive(f::ElasticNet, x::AbstractArray{R}, gamma::AbstractArray) where R <: RealOrComplex
    uz = max.(0, abs.(x) .- gamma.*f.mu)./(1 .+ f.lambda .* gamma)
    return sign.(x) .* uz, f.mu * norm(uz,1) + (f.lambda/2) * norm(uz)^2
end


function dermatrix(N :: Integer,a :: Number,b :: Number, order = 1)
	dermatrix(N+1, N, a, b, order)
end

function dermatrix(Nx :: Integer, Nz :: Integer, a :: Number, b :: Number, order :: Integer)
	xi,wi = lgwt(Nx,a,b)
	if Nz > 1
		zi,wiz = lgwt(Nz,a,b)
	else
		zi,wiz = (a+b)/2, (b-a);
	end
	D = dermatrix(xi,zi,order)
	return D,xi,zi
end

function dermatrix(xi :: Array, zi :: Number, order = 1)
	dermatrix(xi, [zi], order)
end

function dermatrix(xi :: Array,zi :: Array, order = 1)
	D = zeros(length(zi),length(xi))
	for j = 1:length(xi)
		g = algo_diff(x-> leg_pol(x,xi[:],j), order)
		for i = 1:length(zi)
			D[i,j] = g(zi[i])
		end
	end
	D
end

function algo_diff(func :: Function, order :: Integer)
	if order == 1
		g = x-> ForwardDiff.derivative(func,x)
		return g
	elseif order == 2
		g = x->ForwardDiff.hessian(func,x)
		g(x) = g([x])
		return g
		else
			println("no more than 2nd order derivative yet")
	end
end

function dermatrix2(N)
	xi = cos(pi*(1:1:(N+2))/(N+3))
	zi = cos(pi*(1:1:(N))/(N+1))
	D = zeros(length(zi),length(xi))
	for j = 1:length(xi)
		g = algo_diff(x-> leg_pol(x,xi,j),2)
		for i = 1:length(zi)
			D[i,j] = g([zi[i]])
		end
	end
	return D,xi,zi
end


# This is a Julia adaptation of the Matlab code from O. Ledoit and M. Wolf
# solely used to test our own implementation of their estimator.
# Their code was released under the BSD-2 license with the copyright notice
# below.
# https://www.econ.uzh.ch/en/people/faculty/wolf/publications.html#9

###########################################################################
# This file is released under the BSD 2-clause license.

# Copyright (c) 2014, Olivier Ledoit and Michael Wolf
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
# IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###########################################################################

using LinearAlgebra
using Statistics

# comments marked with a `#` come from their code
# comments marked with a `#<#>` are ours

function matlab_ledoitwolf_covcor(z)
    x = copy(z)

    # de-mean returns
    t, n  = size(x)
    meanx = mean(x, dims=1)
    x -= repeat(meanx, outer=(t, 1))

    # compute sample covariance matrix
    sample = (1/t) * (x'*x)

    # compute prior
    diagsample = diag(sample)
    sqrtvar = sqrt.(diagsample)

    sqrtvar_mat = repeat(sqrtvar, outer=(1, n))

    rBar = (sum(sample./(sqrtvar_mat .* sqrtvar_mat'))-n)/(n*(n-1))
    prior = rBar * sqrtvar_mat .* sqrtvar_mat'

    #<#> replace the diagonal
    prior -= Diagonal(diag(prior))
    prior += Diagonal(diagsample)

    #<#> compute optimal shrinkage

    # what we call pi-hat
    y = x.^2;
    phiMat = y'*y/t - 2*(x'*x) .* sample/t + sample.^2
    phi = sum(phiMat)

    # what we call rho-hat
    term1    = ((x.^3)'*x)/t
    help     = x'*x/t
    helpDiag = diag(help)
    term2    = repeat(helpDiag, outer=(1, n)) .* sample
    term3    = help .* repeat(diagsample, outer=(1, n))
    term4    = repeat(diagsample, outer=(1, n)) .* sample
    thetaMat = term1 - term2 - term3 + term4
    #<#> remove the diagonal
    thetaMat -= Diagonal(diag(thetaMat))
    rho       = sum(diag(phiMat)) + rBar*sum(sqrtvar'./sqrtvar.*thetaMat)

    # what we call gamma hat
    gamma = sum((sample - prior).^2)
    kappa = (phi - rho) / gamma
    shrinkage = max(0, min(1, kappa/t))

    part_results = Dict(
        "r̄" => rBar,
        "F" => prior,
        "shrinkage" => shrinkage,
        "kappa" => kappa,
        "gamma" => gamma,
        "lwcov" => shrinkage*prior + (1-shrinkage)*sample)

    return part_results
end


function test_primes_iscoprime()
    @test iscoprime(1, 3) == true
    @test iscoprime(2, 3) == true
    @test iscoprime(11*13*17, 13*17*19) == false
    @test iscoprime(2*3, 5*7) == true
end

function test_primes_primes()
    @test primes(10) == [2, 3, 5, 7]
    @test primes(11) == [2, 3, 5, 7, 11]
    @test primes(1, 1) == primes(1:1) == Int[]
    @test primes(10, 10) == primes(10:10) == Int[]
    @test primes(10, 20) == primes(10:20) == [11, 13, 17, 19]
    @test primes(11, 19) == primes(11:19) == [11, 13, 17, 19]
    @test primes(1000, 1100) == primes(1000:1100) == [1009, 1013, 1019, 1021, 1031, 1033, 1039, 1049, 1051, 1061, 1063, 1069, 1087, 1091, 1093, 1097]
    @test eltype(primes(10)) == typeof(2)
    @test eltype(primes(10, 20)) == eltype(primes(10:20)) == typeof(11)
end

function test_primes_primesmask()
    @test primesmask(100000) == primesmask(1:100000) == [isprime(i) for i in 1:100000]
    @test primesmask(2:100000) == [isprime(i) for i in 2:100000]
    @test primesmask(1000:100000) == [isprime(i) for i in 1000:100000]
    @test primesmask(1009:100000) == [isprime(i) for i in 1009:100000]
    @test typeof(primesmask(100000)) == typeof(primesmask(1:100000)) == BitVector
end

function test_primes_primepi()
    @test primepi(10) == primepi(1:10) == primepi(2:10) == 4
    @test primepi(11) == primepi(1:11) == primepi(2:11) == 5
    @test primepi(10000) == primepi(1:10000) == primepi(2:10000) == 1229
    @test primepi(10, 20) == primepi(10:20) == 4
    @test primepi(11, 19) == primepi(11:19) == 4
end

function test_primes_nthprime()
    p = primes(10000)
    @test [nthprime(i) for i in 1:1000] == [p[i] for i in 1:1000]
    @test [nthprime(a, a + 100) for a in 1:1000] ==  [nthprime(a:(a+100)) for a in 1:1000] == [p[a:(a+100)] for a in 1:1000]
    @test typeof(nthprime(100)) == eltype(nthprime(1, 100)) == eltype(nthprime(1:100)) == typeof(100)
end

function test_primes_nextprime()
    p = primes(2000)
    @test [nextprime(i) for i in -1:1000] == [first(dropwhile(p -> p <= i, p)) for i in -1:1000]
    @test typeof(nextprime(3)) == typeof(3)
    @test typeof(nextprime(big(3))) == typeof(big(3))
end

function test_primes_prevprime()
    @test prevprime(0) == 0
    @test prevprime(2) == 0
    p = reverse(primes(2000))
    @test [prevprime(i) for i in 3:1000] == [first(dropwhile(p -> p >= i, p)) for i in 3:1000]
    @test typeof(prevprime(3)) == typeof(3)
    @test typeof(prevprime(big(3))) == typeof(big(3))
end

function test_primes_twinprimes()
    @test twinprimes(1000, 1050) == twinprimes(1000:1050) == [1019 1021; 1031 1033]
    @test length(twinprimes(20, 30)) == 0
end

function test_primes_all()
    label = "Math.NumberTheory.Primes..."
    print(rpad(label, 50, ' '))

    test_primes_iscoprime()
    test_primes_primes()
    test_primes_primesmask()
    test_primes_primepi()
    test_primes_nthprime()
    test_primes_nextprime()
    test_primes_prevprime()
    test_primes_twinprimes()

    println("PASS")
end


export Cubic, Spline, npatches, funk

# see http://mathworld.wolfram.com/CubicSpline.html

import Base: getindex, show

"""
`Cubic(a,b,c,d)` is a 3rd degree polynomial `a+bx+cx^2+dx^3`.
Use `f(x)` to evalue `f` at the value `x`.
"""
struct Cubic
    a::Number
    b::Number
    c::Number
    d::Number
end

(f::Cubic)(x::Number) = f.a + x*(f.b + x*(f.c+x*f.d))

"""
`f'` where `f` is a `Cubic` or `Spline` is the derivative
of `f`.
"""
Base.adjoint(f::Cubic) = Cubic(f.b,2f.c,3f.d,0)

struct Spline
    patches::Array{Cubic,1}
    closed::Bool
end

is_closed(S::Spline) = S.closed

function Base.adjoint(S::Spline)
    plist = adjoint.(S.patches)
    return Spline(plist,S.closed)
end

"""
`Spline(vals,kind)` returns a cubic spline based on the values in `vals`.
The resulting spline `S` will have the property that `S(1)==y[1]`,
`S(2)==y[2]`, and so on up to `S(n)==y[n]` where `n` is the length of `y`.

+ If `kind` is `:open` then the second derivatives at the end points will be zero. (This is the default.)
+ If `kind` is `:closed` then we assume that we are interpolating a periodic function where `S(n+1)==S(1)`.
"""
function Spline(y::Array{T,1},kind::Symbol=:open)::Spline where T<:Number
    n = length(y)
    if kind == :open
        @assert n>2 "Open splines must have at least three points"
        return open_spline(y)
    end
    if kind == :closed
        @assert n>3 "Closed splines must have at least four points"
        return closed_spline(y)
    end

    error("Spline type must be :open or :closed, not $kind")
end

function npatches(S::Spline)
    return length(S.patches)
end

function show(io::IO, S::Spline)
    adjective = is_closed(S) ? "Closed" : "Open"
    print(io, "$adjective spline with $(npatches(S)) patches")
end

getindex(S::Spline, idx::Int) = S.patches[idx]

"""
`funk(S)` converts the Spline `S` into a callable function
(e.g., that can be passed to `plot`).
"""
function funk(S)::Function
    return x -> S(x)
end

function (S::Spline)(x::Real)
    np = npatches(S)

    if is_closed(S)

        p = Int(floor(x))
        x = mod(x,np)
        p = Int(floor(x))

        if p==0
            p = np
            x += np
        end
        f = S[p]
        return f(x-p)
    end

    # open spline
    p = Int(floor(x))
    if x < 1
        p = 1
    end

    if p > np
        p = np
    end
    f = S[p]
    return f(x-p)
end

function open_spline(y::Array{T,1})::Spline where T<:Number
    n = length(y)
    M = zeros(n,n)
    for i=1:n
        M[i,i] = 4
    end
    M[1,1] = 2
    M[n,n] = 2

    for i=1:n-1
        M[i,i+1] = 1
        M[i+1,i] = 1
    end

    rhs = zeros(T,n)
    rhs[1] = 3*(y[2]-y[1])
    for k=2:n-1
        rhs[k] = 3*(y[k+1]-y[k-1])
    end
    rhs[n] = 3*(y[n]-y[n-1])

    D = M\rhs

    a = zeros(Number,n-1)
    b = zeros(Number,n-1)
    c = zeros(Number,n-1)
    d = zeros(Number,n-1)

    for j=1:n-1
        a[j] = y[j]
        b[j] = D[j]
        c[j] = 3*(y[j+1]-y[j])-2D[j]-D[j+1]
        d[j] = 2*(y[j]-y[j+1])+D[j]+D[j+1]
    end
    return Spline( [ Cubic(a[i],b[i],c[i],d[i]) for i=1:n-1 ] , false)
end

function closed_spline(y::Array{T,1})::Spline where T <: Number
    yy = copy(y)
    n = length(y)
    prepend!(yy,y[end-3:end])
    append!(yy,y[1:4])

    S = open_spline(yy)

    plist = S.patches[5:end-3]
    return Spline(plist,true)

end


using Random, Distributions, JLD2, CSV, DataFrames,LinearAlgebra,Statistics,FreqTables,TexTables
#TexTables wasn't in the instructions but I think it is needed for freqtable
#q1,q2,q3 should run before declaration of q4 since q3 generates "nlsw88.jld"

Random.seed!(1234)
#Question 1
function q1()
    Random.seed!(1234)

    #1.a
    #1.a.i
    A=rand(Uniform(-5,10),(10,7))

    #1.a.ii
    B=rand(Normal(-2,15),(10,7))
    #1.a.iii
    C= [A[1:5,1:5] B[1:5,6:7]]
    #1.a.iv
    D=(A.<=0).*A

    #1.b
    length(A)
    #1.c
    unique(D)
    #1.d
    E=reshape(B,70,1)
    E=vec(B)
    #1.e
    F=cat(A, B;dims=3)
    #1.f
    F=permutedims(F,(3,1,2))
    #1.g
    G=kron(B,C)
    #kron(C,F) is not possible F is not vectorized properly.
    #1.h
    JLD2.@save "matrixpractice.jld2" A B C D E F G
    #1.i
    JLD2.@save "firstmatrix.jld2" A B C D
    #1.j
    df=DataFrame(C)
    CSV.write("Cmatrix.csv",df)
    #1.k
    dfD=DataFrame(D)
    CSV.write("Dmatrix.dat",df,delim='\t')
    return A,B,C,D
end

#Question 2
function q2(a,b,c)
    #2.a
    n=length(A)
    AB=zeros(size(A))

    for i=1:n
        AB[i]=A[i]*B[i]
    end
    AB2=A.*B

    #2.b
    n=length(C)
    Cprime=[]
    for i=1:n
        if -5<=C[i]<=5
            append!(Cprime,C[i])
        end
    end

    Cprime2=[x for x in C if -5<=x<=5]

    #2.c
    X=permutedims(
    cat(ones(15169,5),
    [rand(Bernoulli(.75*(6-y)/5)) for x in 1:15169,y in 1:5],
    [rand(Normal(15+y-1,5(y-1))) for x in 1:15169,y in 1:5],
    [rand(Normal(π*(6-y)/3,1/ℯ)) for x in 1:15169,y in 1:5],
    rand(Binomial(20,0.6),15169,5),
    rand(Binomial(20,0.5),15169,5);
    dims=3),(1,3,2))

    #2.d
    β=hcat([[1+.25*(t-1), log(t), -sqrt(t), ℯ^t-ℯ^(t+1), t, t/3] for t in 1:5]...)

    #2.e
    Y=hcat([X[:,:,i]*β[:,i].+rand(Normal(0,.36)) for i in 1:5]...)
    return
end

#Question 3
function q3()
    #3.a
    #workspace() not supported on my version anymore
    df=CSV.read("nlsw88.csv"; header=true,missingstring="")
    JLD2.@save "nlsw88.jld" df

    #3.b
    describe(df)
    # the mean of neve_married is 0.104185 and married is {0,1}. so 10.4185% are never_married.

    #3.c
    tabulate(df,:race)
    #I think taubulate is in TexTables package but it wasn't listed in the required package list.

    #3.d
    summary_df=describe(df, :mean, :median, :std, :min, :max, :nunique, :q75, :q25)
    summary_df[!, :intq] = summary_df[:, :q75] - summary_df[:, :q25]
    summarystats=Matrix(summary_df[:,Not([:variable,:q25,:q75])])
    #from 3.b there are two observations missing.

    #3.e
    freqtable(df[:,:industry],df[:,:occupation])

    #3.f
    df_f=df[:,[:industry, :occupation , :wage]]
    gdf=groupby(df_f,[:industry, :occupation])
    gdf=combine(gdf, valuecols(gdf) .=> mean)
    freqtable(gdf,:wage_mean,:industry,:occupation);
    return
end

A,B,C,D=q1()
q2(A,B,C)
q3()

#Question4
function q4()
    #4.a
    JLD2.@load "firstmatrix.jld2" A B C D

    #4.b 4.c 4.e
    function matrixops(a,b)
        #calculate dot product, product and sum of elements for inputs.
        if size(a)==size(b)
            dProduct=a.*b
            product=a'*b
            sprod=sum(a+b)
            return dProduct, product,sprod
        else
            println("inputs must have the same size.")
        end
    end
    #4.d
    matrixops(A,B)
    #4.f
    matrixops(C,D)
    #4.g
    JLD2.@load "nlsw88.jld"
    matrixops(convert(Array,df.ttl_exp),convert(Array,df.wage));
end

q4()


using DataStructures

# The implementation is referencing the wikipedia page
# https://en.wikipedia.org/wiki/Matching_pursuit#The_algorithm

const default_max_iter = 20
const default_tolerance = 1e-6

function SparseArrays.sparsevec(d::DefaultDict, m::Int)
    SparseArrays.sparsevec(collect(keys(d)), collect(values(d)), m)
end

function matching_pursuit_(data::AbstractVector, dictionary::AbstractMatrix,
                           max_iter::Int, tolerance::Float64)
    n_atoms = size(dictionary, 2)

    residual = copy(data)

    xdict = DefaultDict{Int, Float64}(0.)
    for i in 1:max_iter
        if norm(residual) < tolerance
            return sparsevec(xdict, n_atoms)
        end

        # find an atom with maximum inner product
        products = dictionary' * residual
        _, maxindex = findmax(abs.(products))
        maxval = products[maxindex]
        atom = dictionary[:, maxindex]

        # c is the length of the projection of data onto atom
        a = maxval / sum(abs2, atom)  # equivalent to maxval / norm(atom)^2
        residual -= atom * a

        xdict[maxindex] += a
    end
    return sparsevec(xdict, n_atoms)
end

"""
    matching_pursuit(data::Vector, dictionary::AbstractMatrix;
                     max_iter::Int = $default_max_iter,
                     tolerance::Float64 = $default_tolerance)

Find ``x`` such that ``Dx = y`` or ``Dx ≈ y`` where y is `data` and D is `dictionary`.
```
# Arguments
* `max_iter`: Hard limit of iterations
* `tolerance`: Exit when the norm of the residual < tolerance
```
"""
function matching_pursuit(data::AbstractVector, dictionary::AbstractMatrix;
                          max_iter::Int = default_max_iter,
                          tolerance = default_tolerance)

    if tolerance <= 0
        throw(ArgumentError("`tolerance` must be > 0"))
    end

    if max_iter <= 0
        throw(ArgumentError("`max_iter` must be > 0"))
    end

    if size(data, 1) != size(dictionary, 1)
        throw(ArgumentError(
            "Dimensions must match: `size(data, 1)` and `size(dictionary, 1)`."
        ))
    end

    matching_pursuit_(data, dictionary, max_iter, tolerance)
end

"""
    matching_pursuit(data::AbstractMatrix, dictionary::AbstractMatrix;
                     max_iter::Int = $default_max_iter,
                     tolerance::Float64 = $default_tolerance)

Find ``X`` such that ``DX = Y`` or ``DX ≈ Y`` where Y is `data` and D is `dictionary`.
```
# Arguments
* `max_iter`: Hard limit of iterations
* `tolerance`: Exit when the norm of the residual < tolerance
```
"""
function matching_pursuit(data::AbstractMatrix, dictionary::AbstractMatrix;
                          max_iter::Int = default_max_iter,
                          tolerance::Float64 = default_tolerance)
    K = size(dictionary, 2)
    N = size(data, 2)

    X = spzeros(K, N)

    for i in 1:N
        X[:, i] = matching_pursuit(
            vec(data[:, i]),
            dictionary,
            max_iter = max_iter,
            tolerance = tolerance
        )
    end
    return X
end


function Newton_Spectral(nlp     :: AbstractNLPModel;
                         x       :: AbstractVector=copy(nlp.meta.x0),
                         ϵ       :: Real=√eps(eltype(x)),
                         #ϵ       :: T = 1e-6,
                         maxiter :: Int = 200,
                         Lp      :: Real = 2 # norm Lp
                         )# where T<:Real

    T = eltype(x)
    n = nlp.meta.nvar
    @info log_header([:iter, :f, :dual, :step, :slope], [Int, T, T, T, T],
                     hdr_override=Dict(:f=>"f(x)", :dual=>"‖∇f‖", :slope=>"∇fᵀd"))
    f  ::T = obj(nlp,x)
    ∇f :: Vector{T} = grad(nlp, x)

    τ₀ = 0.0005
    τ₁ = 0.999

    iter = 0
    @info log_row(Any[iter, f, norm(∇f, Lp)])

    while (norm(∇f, Lp) > ϵ) && (iter <= maxiter)
        H = Matrix(Symmetric(hess(nlp, x),:L))
        Δ, O = eigen(H)
        # Boost negative values of Δ to 1e-8
        # devise an adaptative value for γ
        γ = 1e-6
        D = abs.(Δ) + max.((γ .- abs.(Δ)), 0.0) .*ones(n)

        d = - O*diagm(1.0 ./ D)*O'*∇f

        #------------------------------------------
        # Hard coded line search
        hp0 = ∇f'*d
        t=1.0
        # Simple Wolfe forward tracking
        xt = x + t*d
        ∇ft = grad(nlp,xt)
        hp = ∇ft'*d
        ft = obj(nlp, xt)
        #  while  ~wolfe & armijo
        nbW = 0
        while (hp <= τ₁ * hp0) && (ft <= ( f + τ₀*t*hp0)) && (nbW < 10)
            t *= 5
            xt = x + t*d
            ∇ft = grad(nlp,xt)
            hp = ∇ft'*d
            ft = obj(nlp, xt)
            nbW += 1
            @debug "W", ft
        end
        tw = t

        # Simple Armijo backtracking
        nbk = 0
        while (ft > ( f + τ₀*t*hp0)) && (nbk < 20)
            t *= 0.5
            xt = x + t*d
            ft = obj(nlp, xt)
            nbk += 1
            @debug "A", ft
        end
        #------------------------------------------

        x += t*d
        f = ft
        if t!=tw   ∇ft = grad(nlp, xt) end
        ∇f = ∇ft
        iter += 1

        @info log_row(Any[iter, f, norm(∇f, Lp), t, hp0])
    end
    if iter > maxiter @warn "Maximum d'itérations"
    end

    return iter, f, norm(∇f, Lp), x
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.3
#     language: julia
#     name: julia-1.6
# ---

# %%
using LinearAlgebra
using Random
using Distributions
using StatsPlots

# %%
A = [
    2 -1
    -1 2
]

# %%
dist = MvNormal(A)

# %%
A[1,2]/√(A[1,1]*A[2,2])

# %%
n = 2^10
XY = rand(dist, n)
X, Y = XY[1,:], XY[2,:]
@show cor(X, Y)

scatter(X, Y)

# %%
function sim(A, n, L=10^5)
    dist = MvNormal(A)
    R = zeros(L)
    XY = zeros(2, n)
    for i in eachindex(R)
        rand!(dist, XY)
        R[i] = @views cor(XY[1,:], XY[2,:])
    end
    R
end

# %%
n = 2^6
R = sim(A, n)
Z = atanh.(R)
r0 = A[1,2]/√(A[1,1]*A[2,2])
z0 = atanh(r0)
histogram(Z; norm=true, alpha=0.3, bin=200)
plot!(Normal(z0, 1/√(n - 3)))
plot!(fit_mle(Normal, Z))

# %%
n = 2^8
x, y = rand(Uniform(-1, 1), n), rand(Uniform(-1, 1), n)
X, Y = x, x + y
@show cor(X, Y)
scatter(X, Y; label="")

# %%
function sim2(n, L=10^5)
    R = zeros(L)
    d = Uniform(-1, 1)
    x = zeros(n)
    y = similar(x)
    Y = similar(x)
    for i in eachindex(R)
        rand!(d, x)
        rand!(d, y)
        @. Y = x + y
        R[i] = cor(x, Y)
    end
    R
end

# %%
n = 2^6
R = sim2(n)
Z = atanh.(R)
r0 = 1/√2
z0 = atanh(r0)
histogram(Z; norm=true, alpha=0.3, bin=200)
plot!(Normal(z0, 1/√(n - 3)))
plot!(fit_mle(Normal, Z))

# %%
n = 2^8
x, y = rand(Exponential(), n), rand(Exponential(), n)
X, Y = x, x + y
@show cor(X, Y)
scatter(X, Y; label="")

# %%
function sim3(n, L=10^5)
    R = zeros(L)
    d = Exponential()
    x = zeros(n)
    y = similar(x)
    Y = similar(x)
    for i in eachindex(R)
        rand!(d, x)
        rand!(d, y)
        @. Y = x + y
        R[i] = cor(x, Y)
    end
    R
end

# %%
n = 2^6
R = sim3(n)
Z = atanh.(R)
r0 = 1/√2
z0 = atanh(r0)
histogram(Z; norm=true, alpha=0.3, bin=200)
plot!(Normal(z0, 1/√(n - 3)))
plot!(fit_mle(Normal, Z))

# %%
function simplot(dist, n, L=10^5)
    Z = zeros(L)
    x = zeros(n)
    y = similar(x)
    Y = similar(x)
    for i in eachindex(R)
        rand!(dist, x)
        rand!(dist, y)
        @. Y = x + y
        Z[i] = atanh(cor(x, Y))
    end
    r0 = 1/√2
    z0 = atanh(r0)
    histogram(Z; norm=true, alpha=0.3, bin=200)
    plot!(Normal(z0, 1/√(n - 3)))
    plot!(fit_mle(Normal, Z))
end

# %%
simplot(Normal(), 2^6)

# %%
simplot(Uniform(), 2^6)

# %%
simplot(Exponential(), 2^6)

# %%
simplot(MixtureModel([Normal(), Normal(10, 1)], [0.95, 0.05]), 2^4)

# %%
simplot(MixtureModel([Normal(), Normal(10, 1)], [0.95, 0.05]), 2^6)

# %%


# Some **really cool** maths:
#
# $$ \exp(i\pi) + 1 \quad = \quad 0 $$
#
# We can show this with some code:

x = exp(im*π) + 1

# that looks close to zero but

x ≈ 0

# however

abs(x) < eps()

# #### Conclusion
#
# The equation is proven thanks to our very rigorous proof.


using Manopt, Manifolds, Test
@testset "Alternating Gradient Descent" begin
    # Note that this is mereely an alternating gradient descent toy example
    M = Sphere(2)
    N = M × M
    data = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
    function F(N, x)
        return 1 / 2 *
               (distance(N[1], x[N, 1], data[1])^2 + distance(N[2], x[N, 2], data[2])^2)
    end
    gradF1(N, x) = -log(N[1], x[N, 1], data[1])
    gradF1!(N, Y, x) = (Y .= -log(N[1], x[N, 1], data[1]))
    gradF2(N, x) = -log(N[2], x[N, 2], data[2])
    gradF2!(N, Y, x) = (Y .= -log(N[2], x[N, 2], data[2]))
    gradF(N, x) = ProductRepr([-log(N[i], x[N, i], data[i]) for i in [1, 2]]...)
    function gradF!(N, Y, x)
        log!(N[1], Y[N, 1], x[N, 1], data[1])
        log!(N[2], Y[N, 2], x[N, 2], data[2])
        return Y .*= -1
    end
    x = ProductRepr([0.0, 0.0, 1.0], [0.0, 0.0, 1.0])

    @testset "Test gradient access" begin
        Pf = AlternatingGradientProblem(N, F, gradF)
        Pv = AlternatingGradientProblem(N, F, [gradF1, gradF2])
        Pf! = AlternatingGradientProblem(N, F, gradF!; evaluation=MutatingEvaluation())
        Pv! = AlternatingGradientProblem(
            N, F, [gradF1!, gradF2!]; evaluation=MutatingEvaluation()
        )
        for P in [Pf, Pv, Pf!, Pv!]
            Y = zero_vector(N, x)
            @test get_gradient(P, x)[N, 1] == gradF(N, x)[N, 1]
            @test get_gradient(P, x)[N, 2] == gradF(N, x)[N, 2]
            get_gradient!(P, Y, x)
            @test Y[N, 1] == gradF(N, x)[N, 1]
            @test Y[N, 2] == gradF(N, x)[N, 2]
            @test get_gradient(P, 1, x) == gradF(N, x)[N, 1]
            @test get_gradient(P, 2, x) == gradF(N, x)[N, 2]
            Y = zero_vector(N, x)
            get_gradient!(P, Y[N, 1], 1, x)
            @test Y[N, 1] == gradF(N, x)[N, 1]
            get_gradient!(P, Y[N, 2], 2, x)
            @test Y[N, 2] == gradF(N, x)[N, 2]
        end
    end
    @testset "Test high level interface" begin
        y2 = allocate(x)
        copyto!(N, y2, x)
        y3 = allocate(x)
        copyto!(N, y3, x)
        y = alternating_gradient_descent(
            N, F, [gradF1!, gradF2!], x; order_type=:Linear, evaluation=MutatingEvaluation()
        )
        alternating_gradient_descent!(
            N,
            F,
            [gradF1!, gradF2!],
            y2;
            order_type=:Linear,
            evaluation=MutatingEvaluation(),
        )
        @test isapprox(N, y, y2)
        o = alternating_gradient_descent!(
            N,
            F,
            [gradF1!, gradF2!],
            y3;
            evaluation=MutatingEvaluation(),
            order_type=:Linear,
            return_options=true,
        )
        @test isapprox(N, y, o.x; atol=10^-3)
    end
end


using SpecialFunctions

# from datetime import date
# User-provided date:
#
stosym = "wmt"
exyear = 2019
exmonth = 5
exday = 16
stopr = 92.68
strike = 94
callBid = 0.97
callAsk = 1.01
rfir = 0.0225
#
#  Initialize key variables below
#  (This is not actually required. This is your Prof's habit).
#
daystd = 1
cipd = 0.0
d1 = 0.0
durvol = 0.0
cumd1 = 0.0
cumd2 = 0.0
newcp = 0.0
tempcp = 0.0
timedecay = 0.0
tdprice = 0.0
callpr = 0.0
#
# Calculating strike value (PEG), which is somewhere between BB and BA.
# PEG coefficient is explained in lecture.
peg_coef = 0.60
spread = callAsk - callBid
callpr = callBid + (peg_coef*spread)
# Calculating days to expiry:
tnow = Dates.today()
expiry = Dates.Date(exyear, exmonth, exday)
days = Int(expiry - tnow)
#  Our method for calculating the cumulative standard normal distribution:
function csnd(dval)
	return (1.0 + erf(dval/sqrt(2.0)))/2.0
end
#
#  Below we calculate implied daily volatility using a recursive method to
#  converge to an answer.  First we make an initializing calculation, then
#  we converge inside the while loop. The convergence method shown is Prof E's
#  "Divide and Conquer", which is relatively efficient, usually using only 8
#  (+/- 2) steps. The student should recognize that the cumd1 is the delta.
#  The implied daily volatility that we are seeking is "cipd" below.
#
target = callpr
precision = 1e-4
count = 0
low = 0.0
high = 1.0
cipd = (high+low)/2
d1 = log(stopr/strike)+((rfir/365)+(cipd^2)/2)*days
durvol = cipd*sqrt(days)
cumd1 = csnd(d1/durvol)
cumd2 = csnd((d1/durvol) - durvol)
discount = exp(-rfir*days/365)
tempcp = (stopr*cumd1)-(strike*discount*cumd2)
while tempcp<=(target-precision) || tempcp>=(target+precision)
	if tempcp >= (target+precision)
		high = cipd
	else
		low = cipd
    end
	cipd = (high+low)/2
	d1 = log(stopr/strike)+((rfir/365)+(cipd^2)/2)*days
	durvol = cipd*sqrt(days)
	cumd1 = csnd(d1/durvol)
	cumd2 = csnd((d1/durvol) - durvol)
	discount = exp(-rfir*days/365)
	tempcp = (stopr*cumd1)-(strike*discount*cumd2)
	count +=1
end
#	Below we calculate one day time decay using our new value for volatility
days = days - daystd
d1 = log(stopr/strike)+((rfir/365)+(cipd^2)/2)*days
durvol = cipd*sqrt(days)
cumd1 = csnd(d1/durvol)
cumd2 = csnd((d1/durvol) - durvol)
discount = exp(-rfir*days/365)
newcp = (stopr*cumd1)-(strike*discount*cumd2)
timedecay = callpr - newcp

println( "")
println( "Days to expiry: ", days)
println( "Stock price: ", stopr)
println( "Strike price: ", strike)
println( "Call ASK: ", callAsk)
println( "Call BID: ", callBid)
println( "Call price (PEG): ", callpr)
println( "The Delta:", cumd1)
println( "One day time decay:", timedecay)
println( "The call's implied volatility: ", cipd)


#=
In England the currency is made up of pound, £, and pence, p, and there are eight coins in general circulation:

1p, 2p, 5p, 10p, 20p, 50p, £1 (100p) and £2 (200p).
It is possible to make £2 in the following way:

1×£1 + 1×50p + 2×20p + 1×5p + 1×2p + 3×1p
How many different ways can £2 be made using any number of coins?
=#
coins = [1,2,5,10,20,50,100,200]
# Number of ways of making change for n cents
# using only the first k types of coins
function coin(n,k)
  if k < 1 || n < 0
    return 0
  elseif n == 0
    return 1
  end
  coin(n, k-1)+coin(n-coins[k], k)
end
@time println(coin(200,8))


using OrdinaryDiffEq, Plots
gr()

#Half-life of Carbon-14 is 5,730 years.
C₁ = 5.730

#Setup
u₀ = 1.0
tspan = (0.0, 1.0)

#Define the problem
radioactivedecay(u,p,t) = -C₁*u

#Pass to solver
prob = ODEProblem(radioactivedecay,u₀,tspan)
sol = solve(prob,Tsit5())

#Plot
plot(sol,linewidth=2,title ="Carbon-14 half-life", xaxis = "Time in thousands of years", yaxis = "Percentage left", label = "Numerical Solution")
plot!(sol.t, t->exp(-C₁*t),lw=3,ls=:dash,label="Analytical Solution")


#Create Initially traveling gaussian
function get_gaussian(p::Param , Opp::Opps)

    #Get Variables
    Nx  = p.xnodes
    Ny  = p.ynodes

    A0  = p.A0
    σ   = p.σ
    r0  = p.r0
    ω   = p.ω
    m   = p.m

    #Get operators
    x   = Opp.xcoord
    y   = Opp.ycoord

    Dx  = Opp.Dx

    #Create initial state
    ψ_matrix = zeros( Nx , Ny )
    dψ_matrix = zeros( Nx , Ny )

    for i in 1:Nx , j in 1:Ny

        #Radial and angular coordinates
        r  = sqrt(x[i]^2 + y[j]^2)
        θ   = atan(y[j] , x[i])

        ψ_matrix[i,j] = A0*cos( m*θ ) * sin( ω*r ) * exp( - 0.5* ( (r-r0)/ σ)^2 )
        dψ_matrix[i,j] = A0*cos( m*θ ) * exp( - 0.5* ( (r-r0)/ σ)^2 ) * ( ω*cos(ω*r) - (r-r0)*sin(ω*r)/σ^2  )

    end

    ψ  = reshape(ψ_matrix , Nx*Ny)
    dψ = reshape(dψ_matrix , Nx*Ny)

    #Return initial state
    return ψ , dψ

end


"""
    cartesian_to_lat_lon(x, y, z)

Convert 3D coordinates (x, y, z) from the unit sphere to (lat, lon). Assumes "lat" is positive
with "z", equatorial plane falls at z=0  and "lon" is measured anti-clockwise (eastward)
from x-axis (y=0) about z-axis.

This is a Julia translation of MATLAB code from MITgcm [1].

[1]: http://wwwcvs.mitgcm.org/viewvc/MITgcm/MITgcm_contrib/high_res_cube/matlab-grid-generator/map_xyz2lonlat.m?view=markup
"""
cartesian_to_lat_lon(x, y, z) = cartesian_to_latitude(x, y, z), cartesian_to_longitude(x, y, z)

cartesian_to_latitude(x, y, z) = atand(z, hypot(x, y))

cartesian_to_longitude(x, y, z) = atand(y, x)


#=
The sum of the squares of the first ten natural numbers is,

12 + 22 + ... + 102 = 385
The square of the sum of the first ten natural numbers is,

(1 + 2 + ... + 10)2 = 552 = 3025
Hence the difference between the sum of the squares of the first ten natural numbers and the square of the sum is 3025 − 385 = 2640.

Find the difference between the sum of the squares of the first one hundred natural numbers and the square of the sum.
=#

function sumsquaredifferences()
  collection = 1:100
  x = sum(map((x)->x^2, collection))
  y = sum(collection)^2
  y - x
end
@time println(sumsquaredifferences())


using EconPDEs, Distributions

# Arbitrage With Holding Costs: A Utility-Based Approach (Bruce Tuckman and Jean-Luc Vila)
Base.@kwdef struct ArbitrageHoldingCosts
    c::Float64 = 0.06
    r::Float64 = 0.09
    ρ::Float64 = 5.42
    σ::Float64 = 26.72
    a::Float64 = 26.72
    T::Float64 = 100
end

function (m::ArbitrageHoldingCosts)(state::NamedTuple, y::NamedTuple, τ::Number)
    (; c, r, ρ, σ, a, T) = m
    (; z) = state
    (; F, Fz_up, Fz_down, Fzz) = y
    Fz = (z >= 0) ? Fz_up : Fz_down
    ϕ = z * (1 + z^2)^(-1/2)
    ϕz = (1 + z^2)^(-3/2)
    ϕzz = - 3 * z * (1 + z^2)^(-5/2)
    sτ = c / r * (1 - exp(- r * (T - τ)))
    μL = ρ * z - 0.5 * σ^2 * ϕzz / ϕz + c / sτ * (ϕ - 1) / ϕz
    iL = 1 / a * (μL / σ^2 + Fz)
    μS = ρ * z - 0.5 * σ^2 * ϕzz / ϕz + c / sτ * (ϕ + 1) / ϕz
    iS = 1 / a * (μS / σ^2 + Fz)
    μ, i  = 0.0, 0.0
    if iL > 0
        i = iL
        μ = μL
    elseif iS < 0
        i = iS
        μ = μS
    end
    Ft = - ((μ + σ^2 * Fz) * a * i - 0.5 * σ^2 * (a * i)^2 - ρ * z * Fz + 0.5 * σ^2 * (Fzz - Fz^2))
    return (; Ft)
end

m = ArbitrageHoldingCosts()
zn = 200
d = Normal(0, sqrt(m.σ^2 / (2 * m.ρ)))
stategrid = OrderedDict(:z => range(quantile(d, 0.00001), quantile(d, 0.99999), length = zn))
yend = OrderedDict(:F => zeros(zn))
τs = range(0, m.T, length = 100)
result = pdesolve(m, stategrid, yend, τs)
residual_norm = maximum(result.residual_norm)

### reproduce Fig 2
#d = Normal(0, sqrt(m.σ^2 / (2 * m.ρ)))
#zmin = quantile(d, 0.025)
#zmax = quantile(d, 0.975)
#idx = (state[:z] .>= zmin) .& (state[:z] .<= zmax)
#
#using Plots
#plot(result[:x][idx, 20], [result[:I_myopic][idx, 2] result[:I][idx, 2]], label = ["myopic" "all"])
#plot(result[:x][idx, 20], result[:I][idx, 2])
#plot!(result[:x][idx, 60], result[:I][idx, 60])
#plot!(result[:x][idx, 100], result[:I][idx, 100])
#


# A quadratic segment in 3D space that passes through three points: x⃗₁, x⃗₂, and x⃗₃.
# The assumed relation of the points may be seen in the diagram below.
#                 ___x⃗₃___
#            ____/        \____
#        ___/                  \
#     __/                       x⃗₂
#   _/
#  /
# x⃗₁
#
# NOTE: x⃗₃ is between x⃗₁ and x⃗₂, but not necessarily the midpoint.
# q(r) = (2r-1)(r-1)x⃗₁ + r(2r-1)x⃗₂ + 4r(1-r)x⃗₃
# See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
# Chapter 8, Advanced Data Representation, in the interpolation functions section
struct QuadraticSegment_3D{T <: AbstractFloat}
    points::NTuple{3, Point_3D{T}}
end

# Constructors
# -------------------------------------------------------------------------------------------------
QuadraticSegment_3D(p₁::Point_3D{T},
                    p₂::Point_3D{T},
                    p₃::Point_3D{T}) where {T <: AbstractFloat} = QuadraticSegment_3D((p₁, p₂, p₃))

# Methods
# -------------------------------------------------------------------------------------------------
function (q::QuadraticSegment_3D{T})(r::R) where {T <: AbstractFloat, R <: Real}
    # See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
    # Chapter 8, Advanced Data Representation, in the interpolation functions section
    rₜ = T(r)
    return (2rₜ-1)*(rₜ-1)*q.points[1] + rₜ*(2rₜ-1)*q.points[2] + 4rₜ*(1-rₜ)*q.points[3]
end

#function derivative(q::QuadraticSegment_3D{T}, r::R) where {T <: AbstractFloat, R <: Real}
#    # dq⃗/dr
#    rₜ = T(r)
#    return (4rₜ - 3)*q.points[1] + (4rₜ - 1)*q.points[2] + (4 - 8rₜ)*q.points[3]
#end
#
#function arc_length(q::QuadraticSegment_3D{T}; N::Int64=20) where {T <: AbstractFloat}
#    # Mathematica solution is pages long and can produce NaN results when the segment is
#    # straight, so numerical integration is used. (Gauss-Legengre quadrature)
#    #     1                  N
#    # L = ∫ ||q⃗'(r)||dr  ≈   ∑ wᵢ||q⃗'(rᵢ)||
#    #     0                 i=1
#    # The default number of points is N = 20, since the timing difference is very small for
#    # additional accuracy when compared with N = 15, and small N give poor accuracy.
#    w, r = gauss_legendre_quadrature(T, N)
#    return sum(norm.(w .* derivative.(q, r)))
#end
#
#function intersect(l::LineSegment_3D{T}, q::QuadraticSegment_3D{T}) where {T <: AbstractFloat}
#    # q(r) = (2r-1)(r-1)x⃗₁ + r(2r-1)x⃗₂ + 4r(1-r)x⃗₃
#    # q(r) = 2r²(x⃗₁ + x⃗₂ - 2x⃗₃) + r(-3x⃗₁ - x⃗₂ + 4x⃗₃) + x⃗₁
#    # Let D⃗ = 2(x⃗₁ + x⃗₂ - 2x⃗₃), E⃗ = (-3x⃗₁ - x⃗₂ + 4x⃗₃), F⃗ = x₁
#    # q(r) = r²D⃗ + rE⃗ + F⃗
#    # l(s) = x⃗₄ + sw⃗
#    # If D⃗ × w⃗ ≠ 0⃗
#    #   x⃗₄ + sw⃗ = r²D⃗ + rE⃗ + F⃗
#    #   sw⃗ = r²D⃗ + rE⃗ + (F⃗ - x⃗₄)
#    #   0⃗ = r²(D⃗ × w⃗) + r(E⃗ × w⃗) + (F⃗ - x⃗₄) × w⃗
#    #   Let A⃗ = (D⃗ × w⃗), B⃗ = (E⃗ × w⃗), C⃗ = (F⃗ - x⃗₄) × w⃗
#    #   0⃗ = r²A⃗ + rB⃗ + C⃗
#    #   0 = (A⃗ ⋅ A⃗)r² + (B⃗ ⋅ A⃗)r + (C⃗ ⋅ A⃗)
#    #   A = (A⃗ ⋅ A⃗), B = (B⃗ ⋅ A⃗), C = (C⃗ ⋅ A⃗)
#    #   0 = Ar² + Br + C
#    #   r = (-B - √(B²-4AC))/2A, -B + √(B²-4AC))/2A)
#    #   s = ((q(r) - p₄)⋅w⃗/(w⃗ ⋅ w⃗)
#    #   r is invalid if:
#    #     1) A = 0
#    #     2) B² < 4AC
#    #     3) r < 0 or 1 < r   (Curve intersects, segment doesn't)
#    #   s is invalid if:
#    #     1) s < 0 or 1 < s   (Line intersects, segment doesn't)
#    # If D⃗ × w⃗ = 0, we need to use line intersection instead.
#    bool = false
#    npoints = 0
#    points = [Point_3D(T, 0), Point_3D(T, 0)]
#    D⃗ = 2*(q.points[1] + q.points[2] - 2*q.points[3])
#    E⃗ = 4*q.points[3] - 3*q.points[1] - q.points[2]
#    w⃗ = l.points[2] - l.points[1]
#    A⃗ = D⃗ × w⃗
#    B⃗ = E⃗ × w⃗
#    C⃗ = (q.points[1] - l.points[1]) × w⃗
#    A = A⃗ ⋅ A⃗
#    B = B⃗ ⋅ A⃗
#    C = C⃗ ⋅ A⃗
#    if abs(A) < 1.0e-6
#        # Line intersection
#        r = (-C⃗ ⋅ B⃗)/(B⃗ ⋅ B⃗)
#        s = (q(r)- l.points[1]) ⋅ w⃗/(w⃗ ⋅ w⃗)
#        points[1] = q(r)
#        if (0 ≤ s ≤ 1) && (0 ≤ r ≤ 1)
#            bool = true
#            npoints = 1
#        end
#    elseif B^2 ≥ 4A*C
#        # Quadratic intersection
#        r₁ = (-B - √(B^2-4A*C))/2A
#        r₂ = (-B + √(B^2-4A*C))/2A
#        points[1] = q(r₁)
#        points[2] = q(r₂)
#        s₁ = (points[1] - l.points[1]) ⋅ w⃗/(w⃗ ⋅ w⃗)
#        s₂ = (points[2] - l.points[1]) ⋅ w⃗/(w⃗ ⋅ w⃗)
#
#        # Check points to see if they are valid intersections.
#        # First r,s valid?
#        if (0 ≤ r₁ ≤ 1) && (0 ≤ s₁ ≤ 1) && (points[1] ≈ l(s₁))
#            npoints += 1
#        end
#        # Second r,s valid?
#        if (0 ≤ r₂ ≤ 1) && (0 ≤ s₂ ≤ 1) && (points[2] ≈ l(s₂))
#            npoints += 1
#            # If only point 2 is valid, return it in index 1 of points
#            if npoints === 1
#                points[1] = points[2]
#            end
#        end
#        bool = npoints > 0
#    end
#    return bool, npoints, points
#end
#intersect(q::QuadraticSegment_3D, l::LineSegment_3D) = intersect(l, q)


# Sudoku Solver using backtracking

function ispossible(m::Array{Int, 4}, i::CartesianIndex{4}, n)
    return !(
        n ∈ m[  : ,   : , i[3], i[4]] || # column
        n ∈ m[i[1], i[2],   : ,   : ] || # line
        n ∈ m[  : , i[2],   : , i[4]]    # cell
    )
end

function sudokusolve!(::Array{Int, 4}, ::Nothing)
    # after `CartesianIndex(3,3,3,3)`, `iterate` returns `nothing`
    return true
end

function sudokusolve!(m::Array{Int, 4}, t::Tuple)
    return sudokusolve!(m, t[1])
end

function sudokusolve!(m::Array{Int, 4}, index::CartesianIndex{4})
    next_index = iterate(CartesianIndices(m), index)

    if m[index] != 0
        return sudokusolve!(m, next_index)
    end

    for n in 1:9
        if ispossible(m, index, n)
            m[index] = n
            if sudokusolve!(m, next_index)
                return true
            end
        end
    end

    m[index] = 0
    return false
end

function sudokusolve!(m::Array{Int, 2})
    sudokusolve!(reshape(m, 3, 3, 3, 3), CartesianIndex(1,1,1,1)) # reshape for easier manipulation
end

function string2matrix(s::String)
    return permutedims(reshape(parse.(Int, split(s, "")), 9, 9))
end

m = string2matrix("004300209005009001070060043006002087190007400050083000600000105003508690042910300")

println("\nOriginal :")
show(stdout, MIME("text/plain"), m)

sudokusolve!(m)

println("\n\nSolved :")
show(stdout, MIME("text/plain"), m)


function combined(op, kernel_function)
    function combined_kernel_function(u, v, k=nothing)
        ker = kernel_function(u, v)
        isnothing(k) ? ker : op.(k, ker)
    end
end

# Kernel functions

"""
    gaussiankernel(u, v)

Compute the Gaussian kernel, given by the formula
```
    κ(x,y) = exp(-‖x-y‖²/2).
```
`u` and `v` are matrices, where each column is a datapoint.
"""
function gaussiankernel(u, v)
    uu = sum(abs2, u, dims=1)
    vv = sum(abs2, v, dims=1)
    acc = u' * v
    @. acc = exp(acc - uu' / 2 - vv / 2)
    return acc
end

function rrule(::typeof(gaussiankernel), u, v)
    r = gaussiankernel(u, v)
    # r̄ is short for ∂l / ∂r
    function gaussiankernel_pullback(r̄)
        m = r .* r̄ # pullback the exponential
        ū = v * m'
        ū .-= sum(m', dims=1) .* u
        v̄ = u * m
        v̄ .-= sum(m, dims=1) .* v
        return NO_FIELDS, ū, v̄
    end
    return r, gaussiankernel_pullback
end

"""
    additivegaussiankernel(u, v, k=nothing)

Compute the radial basis function kernel (see [`gaussiankernel`](@ref)) on `u`, `v`.
If `k` is not `nothing`, add it to the result.
"""
const additivegaussiankernel = combined(+, gaussiankernel)

"""
    additivegaussiankernel(u, v, k=nothing)

Compute the radial basis function kernel (see [`gaussiankernel`](@ref)) on `u`, `v`.
If `k` is not `nothing`, multiply the result by it.
"""
const multiplicativegaussiankernel = combined(*, gaussiankernel)


# See https://en.wikipedia.org/wiki/Fibonacci

fibonacci(n) = ([1 1 ; 1 0] ^ n)[1, 2]

expressions = [
    :(fibonacci(0))
    :(fibonacci(1))
    :(fibonacci(6))
    :(fibonacci(91))
    :(fibonacci(big(283)))
    :(fibonacci(big(3674)))
]

for expr in expressions
    println("[] :julia-answer \"\"\"", escape_string(string(expr)), " = ", eval(expr), "\"\"\".")
end


export lambertW

#
# The Lambert W function (See reference [4])
# The following algorithm computes the Lambert-W function to machine accuracy using Halley's method for any x
# in the domain I=( -1/exp(1), ∞ ). This function can also compute the lambert-W function component-wise for
# vectors or matrices. The Lambert-W function is used in the module DESincEig to compute the mesh size h.
# Input: x:: Number, Vector or Matrix
# Output: lambertW(x)
#
function lambertW(x::Real)
    if x < -exp(-one(x))
        return throw(DomainError())
    elseif x == -exp(-one(x))
        return -one(x)
    elseif x < 0
        w0 = e*x/(1+inv(inv(sqrt(2*e*x+2))+inv(e-1)-inv(sqrt(2))))
    else
        logx = log(1+x)
        w0 = logx*(1-log(1+logx)/(2+logx))
    end
    expw0 = exp(w0)
    w1 = w0 - (w0*expw0 - x)/((w0 + 1)*expw0 -
        (w0 + 2) * (w0*expw0 - x)/(2w0 + 2))
    while abs(w1/w0 - 1) > 2eps(typeof(x))
        w0 = w1
        expw0 = exp(w0)
        w1 = w0 - (w0*expw0 - x)/((w0 + 1)*expw0 -
            (w0 + 2) * (w0*expw0 - x)/(2w0 + 2))
    end
    return w1
end
lambertW(x::Integer) = lambertW(float(x))
@vectorize_1arg Real lambertW