module BasesAndSamples

using Combinatorics

export basis_monomial, basis_laguerre, basis_jacobi, basis_chebyshev, basis_gegenbauer
export sample_points_simplex, sample_points_padua, sample_points_rescaled_laguerre, sample_points_chebyshev, sample_points_chebyshev_mod

"""
    basis_monomial(d::Int, x...)

Generate the monomial basis in variables x... up to degree d (inclusive).
"""
function basis_monomial(d::Int, x...)
    n = length(x)
    if n > 1
        q = Vector{typeof(prod(one.(x)))}(undef, binomial(n+d,d))
    elseif n == 1
        q = Vector{typeof(one(x...))}(undef, binomial(n+d,d))
    end
    idx = 1
    for k=0:d
        for exponent in multiexponents(n,k)
            q[idx] = prod([x[i]^exponent[i] for i=1:length(exponent)])
            idx+=1
        end
    end
    return q
end

"""
    basis_laguerre(d::Integer, alpha, x)

Generate the Laguerre polynomials with parameter alpha up to degree d (inclusive). (https://en.wikipedia.org/wiki/Laguerre_polynomials)
"""
function basis_laguerre(d::Integer, alpha, x)
    v = Vector{typeof(one(alpha) * one(x))}(undef, 1 + d)
    v[1] = one(x)
    d == 0 && return v
    v[2] = 1 + alpha - x
    d == 1 && return v
    for l = 2:d
        v[l+1] = 1 // big(l) * ((2l - 1 + alpha - x) * v[l] - (l + alpha - 1) * v[l-1])
    end
    return v
end

"""
    basis_jacobi(d::Integer, alpha, beta, x, normalized = true)

Generate the Jacobi polynomials with parameters alpha and beta up to degree d (inclusive). (https://en.wikipedia.org/wiki/Jacobi_polynomials)
"""
function basis_jacobi(d::Integer, alpha, beta, x, normalized = true)
    q = Vector{typeof(one(alpha) * one(x))}(undef, d + 1)
    q[1] = one(x)
    d == 0 && return q
    q[2] = x # normalized
    if !normalized
        q[2] *= (alpha + 1)
    end
    d == 1 && return q
    for k = 2:d
        # what if alpha+beta = -n for some integer 1 <= n < 2d ?
        q[k+1] =
            (2 * k + alpha + beta - 1) /
            BigFloat(2k * (k + alpha + beta) * (2k + alpha + beta - 2)) *
            ((2 * k + alpha + beta) * (2k + alpha + beta - 2) * x + beta^2 - alpha^2) *
            q[k] + -2 * (k + alpha - 1) * (k + beta - 1) * (2 * k + alpha + beta) * q[k-1]
    end
    return q
end

"""
    basis_chebyshev(d::Int,x)

Generate a basis of chebyshev polynomials up to degree d (inclusive). (https://en.wikipedia.org/wiki/Chebyshev_polynomials)
"""
function basis_chebyshev(d::Int,x)
    v = Vector{typeof(one(x))}(undef, 1 + d)
    v[1] = one(x)
    d == 0 && return v
    v[2] = x
    d == 1 && return v
    for l = 2:d
        v[l+1] = 2x*v[l] - v[l-1]
    end
    return v
end

"""
    basis_gegenbauer(d, n, x)

Basis for the Gegenbauer polynomials in dimension n up to degree d.
 This is the Gegenbauer polynomial with parameter lambda = n/2-1,
 or the Jacobi polynomial with alpha = beta = (n-3)/2.
 Normalized to evaluate to 1 at 1.
 Taken from arxiv/2001.00256, ancillary files, SemidefiniteProgramming.jl.
"""
function basis_gegenbauer(d, n, x)
    v = Vector{typeof(one(x))}(undef, 1 + d)
    v[1] = one(x)
    d == 0 && return v
    v[2] = x
    d == 1 && return v
    for l = 2:d
        v[l+1] = (2l + n - 4) // (l + n - 3) * x * v[l] - (l - 1) // (l + n - 3) * v[l-1]
    end
    return v
end

"""
    sample_points_simplex(n, d)

Generate the rational sample points in the unit simplex with denominator d.
"""
function sample_points_simplex(n, d)
    #rational points in the unit simplex with denominator d
    #probably not very efficient, but I dont know how to do it better for general n.
    x = [zeros(BigFloat, n) for i = 1:binomial(n + d, d)] #need n+d choose d points for a unisolvent set, if symmetry is not used.
    idx = 1
    for I in CartesianIndices(ntuple(k -> 0:d, n)) #all tuples with elements in 0:d of length n
        if sum(Tuple(I)) <= d #in unit simplex
            x[idx] = [i / BigFloat(d) for i in Tuple(I)]
            idx += 1
        end
    end
    return x
end

"""
    sample_points_padua(d)

Generate the Padua points for degree d. (https://en.wikipedia.org/wiki/Padua_points)
"""
function sample_points_padua(d)
    #padua points:
    z = [Array{BigFloat}(undef, 2) for i = 1:binomial(2 + d, d)]
    z_idx = 1
    for j = 0:d
        delta_j = j % 2 == d % 2 == 1 ? 1 : 0
        mu_j = cospi(j / d)
        for k = 1:(div(d, 2)+1+delta_j)
            eta_k = j % 2 == 1 ? cospi((2 * k - 2) / (d + 1)) : cospi((2 * k - 1) / (d + 1))
            z[z_idx] = [mu_j, eta_k]
            z_idx += 1
        end
    end
    return z
end

"""
    sample_points_rescaled_laguerre(d)

Generate 'rescaled laguerre' points, as in SDPB.
"""
function sample_points_rescaled_laguerre(d)
    #as done in sdpb: ('rescaled Laguerre')
    # x[k] = sqrt(pi)/(64*(d+1)*log( 3- 2*sqrt(2))) * (-1+4*k)^2, with k=0:d
    constant = -sqrt(BigFloat(pi)) / (64 * (d + 1) * log(3 - 2 * sqrt(BigFloat(2))))
    x = zeros(BigFloat, d + 1)
    for k = 0:d
        x[k+1] = constant * (-1 + 4 * k)^2
    end
    return [[i] for i in x]
end

"""
    sample_points_chebyshev(d, a = -1, b = 1)

Generate the d+1 chebyshev points in [a,b]. (https://en.wikipedia.org/wiki/Chebyshev_nodes)
"""
function sample_points_chebyshev(d, a = -1, b = 1)
    #roots of chebyshev polynomials of the first kind, unisolvent for polynomials up to degree d
    points = [
        (a + b) / BigFloat(2) +
        (b - a) / BigFloat(2) * cos((2k - 1) / BigFloat(2(d + 1)) * BigFloat(pi)) for
        k = 1:d+1
    ]
    return [[i] for i in points]
end

"""
    sample_points_chebyshev_mod(d, a = -1, b = 1)

Generate the d+1 modified chebyshev points in [a,b], the chebyshev points divided by cos(pi/(2(d+1))).
"""
function sample_points_chebyshev_mod(d, a = -1, b = 1)
    #roots of chebyshev polynomials of the first kind, divided by cos(pi/2(d+1)) to get a lower lebesgue constant
    points = [
        (a + b) / BigFloat(2) +
        (b - a) / BigFloat(2) * cos((2k - 1) / BigFloat(2(d + 1)) * BigFloat(pi)) /
        cos(BigFloat(pi) / 2(d + 1)) for k = 1:d+1
    ]
    return [[i] for i in points]
end

end # end module


module Transformation

export Vec, Quat, Transf, dot, cross, magnitude, norm, ×, slerp, unit

"""
    Vec - three dimensional vector.

# Examples
```julia-repl
julia> Vec(1, 2, 3)
Vec(1, 2, 3)
```
"""
struct Vec
    x::Real
    y::Real
    z::Real
    function Vec(x::Real=0, y::Real=0, z::Real=0)::Vec
        new(x, y, z)
    end
end

function Base.:+(a::Vec, b::Vec)::Vec
    Vec(
        a.x + b.x,
        a.y + b.y,
        a.z + b.z
    )
end

function Base.:-(a::Vec, b::Vec)::Vec
    Vec(
        a.x - b.x,
        a.y - b.y,
        a.z - b.z
    )
end

function Base.:-(a::Vec)::Vec
    Vec(
        -a.x,
        -a.y,
        -a.z
    )
end

"""
    dot(a::Vec, b::Vec)::Vec

Compute the dot product of two vectors.
"""
function dot(a::Vec, b::Vec)
    a.x * b.x + a.y * b.y + a.z * b.z
end

"""
    cross(a::Vec, b::Vec)::Vec

Compute the cross product of two vectors.
"""
function cross(a::Vec, b::Vec)::Vec
    Vec(
        a.y * b.z - a.z * b.y,
        a.z * b.x - a.x * b.z,
        a.x * b.y - a.y * b.x
    )
end

"""
    ×(a::Vec, b::Vec)::Vec

Compute the cross product of two vectors.
"""
function ×(a::Vec, b::Vec)::Vec
    cross(a, b)
end

function Base.:*(a::Real, b::Vec)::Vec
    Vec(
        a * b.x,
        a * b.y,
        a * b.z
    )
end

function Base.:*(a::Vec, b::Real)::Vec
    Vec(
        a.x * b,
        a.y * b,
        a.z * b
    )
end

function Base.:/(a::Vec, b::Real)::Vec
    Vec(
        a.x / b,
        a.y / b,
        a.z / b
    )
end

function Base.isequal(a::Vec, b::Vec)::Bool
    a.x == b.x && a.y == b.y && a.z == b.z
end

"""
    norm(a::Vec)::Real

Compute the Euclidean norm of the given vector.
"""
function norm(a::Vec)::Real
    √(a.x^2 + a.y^2 + a.z^2)
end

"""
    unit(a::Vec)::Vec

Compute the unit vector colinear to the given one.
"""
function unit(a::Vec)::Vec
    a / norm(a)
end

"""
    round(a::Vec)::Vec

Compute the vector with it's elements rounder to nearest value.
"""
function Base.round(a::Vec)::Vec
    Vec(
        round(a.x),
        round(a.y),
        round(a.z)
    )
end

"""
    Quat - hypercomplex number.

Unit quaternions are used for rotation representation.

# Examples
```julia-repl
julia> # 120° counterclockwise rotation around the (1, 1, 1) axis
julia> Quat(2π/3, Vec(1, 1, 1))
Quat(0.5, 0.5, 0.5, 0.5)

julia> # 45° clockwise rotation around the (1, 0, 0) axis
julia> Quat(-π/4, Vec(1, 0, 0))
Quat(0.9238795325112867, -0.3826834323650898, -0.0, -0.0)
```
"""
struct Quat
    w::Real
    x::Real
    y::Real
    z::Real
    function Quat(w::Real=1, x::Real=0, y::Real=0, z::Real=0)::Quat
        new(w, x, y, z)
    end

    function Quat(angle::Real, axis::Vec)::Quat
        w = cos(angle / 2)
        xyz = unit(axis) * sin(angle / 2)
        Quat(w, xyz.x, xyz.y, xyz.z)
    end
end

function Base.:+(a::Quat, b::Quat)::Quat
    Quat(
        a.w + b.w,
        a.x + b.x,
        a.y + b.y,
        a.z + b.z
    )
end

function Base.:-(a::Quat, b::Quat)::Quat
    Quat(
        a.w - b.w,
        a.x - b.x,
        a.y - b.y,
        a.z - b.z
    )
end

function Base.:*(a::Quat, b::Quat)::Quat
    Quat(
        a.w * b.w - a.x * b.x - a.y * b.y - a.z * b.z,
        a.w * b.x + a.x * b.w + a.y * b.z - a.z * b.y,
        a.w * b.y - a.x * b.z + a.y * b.w + a.z * b.x,
        a.w * b.z + a.x * b.y - a.y * b.x + a.z * b.w
    )
end

function Base.:*(q::Quat, v::Vec)::Vec
    r = q * Quat(0, v.x, v.y, v.z) * conj(q)
    Vec(r.x, r.y, r.z)
end

function Base.:*(a::Quat, b::Real)::Quat
    Quat(
        a.w * b,
        a.x * b,
        a.y * b,
        a.z * b
    )
end

function Base.:*(a::Real, b::Quat)::Quat
    Quat(
        a * b.w,
        a * b.x,
        a * b.y,
        a * b.z
    )
end

function dot(a::Quat, b::Quat)::Real
    a.w * b.w + a.x * b.x + a.y * b.y + a.z * b.z
end

function Base.:/(a::Quat, b::Real)::Quat
    Quat(
        a.w / b,
        a.x / b,
        a.y / b,
        a.z / b
    )
end

"""
    conj(a::Quat)::Quat

Compute the quaternion conjugate.
If the quaternion is unit, it's conjugate represents the reverse rotation.

# Examples
```julia-repl
julia> q = Quat(π, Vec(1, 1, 1))
julia> q * conj(q)
Quat(1.0, 0.0, 0.0, 0.0)
```
"""
function Base.conj(a::Quat)::Quat
    Quat(
        a.w,
        -a.x,
        -a.y,
        -a.z
    )
end

"""
    norm(a::Quat)::Real

Compute the quaternion norm.
"""
function norm(a::Quat)::Real
    √(a.w^2 + a.x^2 + a.y^2 + a.z^2)
end

"""
    unit(a::Quat)::Quat

Compute the unit quaternion colinear to the given one.
"""
function unit(a::Quat)::Quat
    a / norm(a)
end

function slerp(a::Quat, b::Quat, t)::Quat
    d = dot(a, b)
    if d < 0
        b = -b
        d = -d
    end
    if d > 0.9995
        result = a + (a - b) * t
        return norm(result)
    end
    theta0 = acos(d)
    theta = theta0 * t
    sinTheta = sin(theta)
    sinTheta0 = sin(theta0)

    s1 = sinTheta / sinTheta0
    s0 = cos(theta) - d * s1
    return a * s0 + b * s1
end

"""
    Transf - spatial transformation representation.

The transformation consists of consecutive translation on vector `v` and rotation on `q` quaternion.
"""
struct Transf
    v::Vec
    q::Quat
    function Transf(v::Vec=Vec(), q::Quat=Quat())::Transf
        new(v, q)
    end
end

function Base.:+(a::Transf, b::Transf)::Transf
    Transf(
        a.v + a.q * b.v,
        a.q * b.q
    )
end

function Base.:+(a::Transf, b::Vec)::Vec
    a.v + a.q * b
end

"""
    conj(a::Transf)::Transf

Compute such transformation, that `a + conj(a)` equals to zero transformation.

# Examples
```julia-repl
julia> t = Transf(Vec(1, 2, 3), Quat(π/3, Vec(1, 1, 0)))
julia> t + conj(t)
Transf(Vec(0.0, 0.0, 0.0), Quat(1.0, 0.0, 0.0, 0.0))
```
"""
function Base.conj(a::Transf)::Transf
    q = conj(a.q)
    Transf(
        q * (-a.v),
        q
    )
end

end


"""
    poisson_eigenvalues(N, L, dim, ::Periodic)

Return the eigenvalues satisfying the discrete form of Poisson's equation
with periodic boundary conditions along the dimension `dim` with `N` grid
points and domain extent `L`.
"""
function poisson_eigenvalues(N, L, dim, ::Periodic)
    inds = reshape(1:N, reshaped_size(N, dim)...)
    return @. (2sin((inds - 1) * π / N) / (L / N))^2
end

"""
    poisson_eigenvalues(N, L, dim, ::Bounded)

Return the eigenvalues satisfying the discrete form of Poisson's equation
with staggered Neumann boundary conditions along the dimension `dim` with
`N` grid points and domain extent `L`.
"""
function poisson_eigenvalues(N, L, dim, ::Bounded)
    inds = reshape(1:N, reshaped_size(N, dim)...)
    return @. (2sin((inds - 1) * π / 2N) / (L / N))^2
end

"""
    poisson_eigenvalues(N, L, dim, ::Flat)

Return N-element array of `0.0` reshaped to three-dimensions.
This is also the first `poisson_eigenvalue` for `Bounded` and `Periodic` directions.
"""
poisson_eigenvalues(N, L, dim, ::Flat) = reshape(zeros(N), reshaped_size(N, dim)...)


input = 347991

# Part 1
function side_length(number:: Int)
  root = ceil(sqrt(number))
  if root % 2 == 0
    return root + 1
  else
    return root
  end
end

function get_position(number:: Int)
  slength = side_length(number)
  x = floor(slength / 2)

  side_centers = slength^2 - ((slength - 1) * collect(1:4)) - floor(slength / 2)
  y = minimum(abs.(side_centers - number))

  return x, y
end

(x, y) = abs.(get_position(input))
println("Part 1: $(x + y) steps from ($x, $y)")

# Part 2
mutable struct NegativeIndexSquareArray
  size::Int
  content::Array{Int, 3}

  NegativeIndexSquareArray(size) = new(size, zeros(size, size, 4))
end

function Base.getindex(array::NegativeIndexSquareArray, x, y)
  sx = sign(x) == 0? 1 : sign(x)
  sy = sign(y) == 0? 1 : sign(y)
  (x, y) = abs.([x, y])

  if sx == 1 && sy == 1
    array.content[x + 1, y + 1, 1]
  elseif sx == 1 && sy == -1
    array.content[x + 1, y, 2]
  elseif sx == -1 && sy == 1
    array.content[x, y + 1, 3]
  else # Both negative
    array.content[x, y, 4]
  end
end

function Base.setindex!(array::NegativeIndexSquareArray, value, x, y)
  sx = sign(x) == 0? 1 : sign(x)
  sy = sign(y) == 0? 1 : sign(y)
  (x, y) = abs.([x, y])

  if sx == 1 && sy == 1
    array.content[x + 1, y + 1, 1] = value
  elseif sx == 1 && sy == -1
    array.content[x + 1, y, 2] = value
  elseif sx == -1 && sy == 1
    array.content[x, y + 1, 3] = value
  else # Both negative
    array.content[x, y, 4] = value
  end
end

spiral = NegativeIndexSquareArray(10)
spiral[0, 0] = 1
(x, y) = [1, 0]
direction = im

while true
  spiral[x, y] = spiral[x - 1, y] + spiral[x - 1, y - 1] +
    spiral[x, y - 1] + spiral[x + 1, y - 1] + spiral[x + 1, y] +
    spiral[x + 1, y + 1] + spiral[x, y + 1] + spiral[x - 1, y + 1]

  if spiral[x, y] > input
    println("Part 2: $(spiral[x, y]) at ($x, $y)")
    break
  end

  left = direction * im
  if spiral[x + real(left), y + imag(left)] == 0
    direction = left
  end
  x += real(direction)
  y += imag(direction)
end


function bisect_left(A, p, r, v)
	i = p
	if p < r
	   q = floor(Int, (p + r) / 2)
	   if v <= A[q]
		   i = bisect_left(A, p, q, v)
	   else
		   i = bisect_left(A, q + 1, r, v)
	   end
	end
	return i
end

function bisect_right(A, p, r, v)
	i = p
	if p < r
	   q = floor(Int, (p + r) / 2)
	   if v < A[q]
		   i = bisect_right(A, p, q, v)
	   else
		   i = bisect_right(A, q + 1, r, v)
	   end
	end
	return i
end

function bisect_right(A, v)
	return bisect_right(A, 1, length(A) + 1, v)
end

function bisect_left(A, v)
	return bisect_left(A, 1, length(A) + 1, v)
end


export newtoncg

"""
newtoncg(f,df,H,x)

Newton-CG method for solving min_x f(x)

"""
function newtoncg(f::Function,df::Function,H::Function,x::Vector;maxIter=20,atol=1e-8,out::Int=0,storeInterm::Bool=false,
	lineSearch::Function=(f,df,fk,dfk,xk,pk)->armijo(f,fk,dfk,xk,pk,maxIter=20),tolCG=1e-2,maxIterCG=30,P=d2f->identity)

    his = zeros(maxIter,6)
    X = (storeInterm) ? zeros(length(x),maxIter) : []
    i = 1; flag = -1; LL = []
    while i<=maxIter
        fk  = f(x)
        dfk = df(x)
        his[i,1:2] = [norm(fk) norm(dfk)]
        if storeInterm; X[:,i] = x; end;

        if(norm(dfk)<atol)
            flag = 0
            his = his[1:i,:]
            break
        end

        # get search direction
        d2f = H(x)
        PC  = P(d2f) # preconditioner
        pk,his[i,6],his[i,5],his[i,4] = cg(d2f,-dfk,out=-1,tol=tolCG,maxIter=maxIterCG,M=PC)
        if his[i,6]==-2 && norm(pk)==0
            pk = -dfk
        end
      # line search
        ak,his[i,3] = lineSearch(f,df,fk,dfk,x,pk)
        if his[i,3]==-1
            flag = -3
            his  = his[1:i,:]
            break
        end
      if out>0
            @printf "iter=%04d\t|f|=%1.2e\t|df|=%1.2e\tcgIter=%d\tLS=%1.4f\n" i his[i,1] his[i,2] his[i,4] ak
        end
        # update
        x  += ak*pk
        i+=1
    end
    i = min(maxIter,i)

    if out>=0
        if flag==-1
            println(@sprintf("newtoncg iterated maxIter (=%d) times but reached only atol of %1.2e instead of tol=%1.2e",i,his[i,2],atol))
        elseif flag==-3
            println(@sprintf("newtoncg stopped at iteration %d due to a line search fail.",i))
        elseif out>1
            println(@sprintf("newtoncg achieved desired atol of %1.2e at iteration %d.",atol,i))
        end
    end
    if storeInterm; X = X[:,1:i]; end
    return x,flag,his,X
end


# # Padua transform
# This demonstrates the Padua transform and inverse transform,
# explaining precisely the normalization and points

using FastTransforms

# We define the Padua points and extract Cartesian components:
N = 15
pts = paduapoints(N)
x = pts[:,1]
y = pts[:,2];

# We take the Padua transform of the function:
f = (x,y) -> exp(x + cos(y))
f̌ = paduatransform(f.(x , y));

# and use the coefficients to create an approximation to the function $f$:
f̃ = (x,y) -> begin
    j = 1
    ret = 0.0
    for n in 0:N, k in 0:n
        ret += f̌[j]*cos((n-k)*acos(x)) * cos(k*acos(y))
        j += 1
    end
    ret
end

# At a particular point, is the function well-approximated?
f̃(0.1,0.2) ≈ f(0.1,0.2)

# Does the inverse transform bring us back to the grid?
ipaduatransform(f̌) ≈ f̃.(x,y)


using LinearAlgebra
using Zygote: @adjoint

""" maximum eigenvalue of a symmetric square matrix """
function maxeig(symmat::Array{T,2})::T where {T}
    eig = eigen(Symmetric(symmat))
    max_idx = findmax(eig.values)[2]
    eigmax = eig.values[max_idx]
    return eigmax
end

@adjoint function maxeig(symmat::Array{T,2}) where {T}
    eig = eigen(Symmetric(symmat))
    max_idx = findmax(eig.values)[2]
    eigmax = eig.values[max_idx]
    vecmax = eig.vectors[:, max_idx]
    jac_adj = vecmax * (vecmax')
    return eigmax, c -> Tuple(jac_adj * c)
end

""" minimum eigenvalue of a symmetric square matrix """
function mineig(symmat::Array{T,2})::T where {T}
    eig = eigen(Symmetric(symmat))
    min_idx = findmin(eig.values)[2]
    eigmin = eig.values[min_idx]
    return eigmin
end

@adjoint function mineig(symmat::Array{T,2}) where {T}
    eig = eigen(Symmetric(symmat))
    min_idx = findmin(eig.values)[2]
    eigmin = eig.values[min_idx]
    vecmin = eig.vectors[:, min_idx]
    jac_adj = vecmin * (vecmin')
    return eigmax, c -> Tuple(jac_adj * c)
end

""" eigenvalue regularization """
function softplus_eigreg(symmat)
    eig = eigen(Symmetric(symmat))
    return mean(softplus.(eig.values))
end

@adjoint function softplus_eigreg(symmat)
    eig = eigen(Symmetric(symmat))
    return sum(softplus.(eig.values)),
    c -> Tuple(
        sum((
            sigmoid(eig.values[i]) * eig.vectors[:, i] * (eig.vectors[:, i]') for
            i = 1:size(symmat, 1)
        )) * c,
    )
end

function relu_eigreg(symmat)
    eig = eigen(Symmetric(symmat))
    return sum(relu.(eig.values))
end

@adjoint function relu_eigreg(symmat)
    eig = eigen(Symmetric(symmat))
    return sum(relu.(eig.values)),
    c -> Tuple(
        sum((
            0.5 * (sign(eig.values[i]) + 1) * eig.vectors[:, i] * (eig.vectors[:, i]')
            for i = 1:size(symmat, 1)
        )) * c,
    )
end

""" compute all eigenvalues """
function eigenvalues(symmat)
    eig = eigen(Symmetric(symmat))
    return eig.values
    # return eigvals(Symmetric(symmat))
end

@adjoint function eigenvalues(symmat)
    eig = eigen(Symmetric(symmat))
    return eig.values,
    c -> Tuple((eig.vectors[:, i] * (eig.vectors[:, i]')) * c for i = 1:size(symmat, 1))
    # vals = eigvals(Symmetric(symmat))
    # vecs = eigvecs(Symmetric(symmat))
    # return vals, c -> Tuple((vecs[:,i]*(vecs[:,i]'))*c for i=1:size(symmat,1))
end


include("prime.jl")
function  is_circular_prime(n::Int)
	if (!is_prime(n))
		return false
	end
	s=string(n)
	for i=1:length(s)
		if (!is_prime( parse(Int,s[i:length(s)]*s[1:i-1])))
			return false
		end

	end

	return true

end

function  compute()
	const N=1000000
	ans =sum( 1 for i =1:N if is_circular_prime(i))
	return ans

end

println(compute())


### A Pluto.jl notebook ###
# v0.17.3

using Markdown
using InteractiveUtils

# ╔═╡ 1c630e40-628c-11ec-1793-05688dabf6c9
begin
	using Pkg; Pkg.activate(@__DIR__); Pkg.instantiate()
	Pkg.precompile()

	using PlutoUI
	using Statistics
	using StatsBase
	using RDatasets
	using Plots
	using StatsPlots
	using KernelDensity
	using Distributions
	using LinearAlgebra
	using HypothesisTests
	using PyCall
	using MLBase
end;

# ╔═╡ 7242a33b-ec37-48d6-afca-8c8387a45a83
PlutoUI.TableOfContents(aside=true, indent=true, depth=3)

# ╔═╡ dd2aa8b3-9da3-418d-bf31-74574483360d
html"""

"""

# ╔═╡ abc828f1-fd14-4329-8b47-f4135300f937
md"""
# 📊 Statistics

Having a solid understanding of statistics in data science allows us to understand our data better, and allows us to create a quantifiable evaluation of any future conclusions.

[YouTube Link](https://www.youtube.com/watch?v=AAGxWEJ_eWk)
"""

# ╔═╡ a8aae3d4-198c-482f-a972-05be31be39c1
md"""
## Load dataset

In this notebook, we will use eruption data on the faithful geyser (Yellowstone, USA). The data will contain wait times between every consecutive times the geyser goes off (in minutes) and the length of the eruptions (also in minutes).

Let's get the data first using the [RDatasets.jl](https://github.com/JuliaStats/RDatasets.jl) package...

**Note:** the dataset will be loaded as a `DataFrame`.
"""

# ╔═╡ 63420679-9b77-4412-ba1c-fc471da6b1e9
data = dataset("datasets", "faithful")

# ╔═╡ 9d56b2cc-c4c6-4ce2-8d9c-33174fe5bc33
md"""
We can use the `describe` function to get some stats from our data...
"""

# ╔═╡ 67ac4a1d-eb15-473f-a548-07425fd627d8
describe(data)

# ╔═╡ 5ee886cc-2730-40cf-8497-507dcdc08c5f
md"""
Let's call each column (feature) separately...
"""

# ╔═╡ 110fda28-1b77-44d6-8485-a4ec6e012011
begin
	eruption = data[!, :Eruptions]
	# eruption = data.Eruptions

	waittime = data[!, :Waiting]
	# waittime = data.Waiting
end;

# ╔═╡ 6e093fbb-1e3a-4131-b3c2-b5891e734fa7
md"""
Now, we can plot both variables in the same scatter plot...
"""

# ╔═╡ ed9da9c2-53e4-45d6-9ec1-02c8bf54c1a5
begin
	scatter(eruption, label="Eruptions")
	scatter!(waittime, label="Wait Time", xlabel="Index", ylabel="Time (minutes)")
end

# ╔═╡ 13137e30-cfde-459a-adb4-1ddc07e626e8
md"""
## Statistics plots

As you can see, this doesn't tell us much about the data... Let's try some statistical plots.
"""

# ╔═╡ ff9d63e7-808a-4465-a68d-659a222afbae
md"""
### Boxplot
"""

# ╔═╡ c3672771-248a-47f9-a019-2de03cdd79f4
boxplot(["Eruption Length"], eruption, legend=false,
		size=(200,400), whisker_width=1, ylabel="Time (minutes)")

# ╔═╡ 46b55947-428a-483f-a00f-a232e5329fbd
md"""
Statistical plots such as a boxplot (and a violin plot), can provide a much better understanding of the data. Here, we immediately see that the median time of each eruption is about 4 minutes.
"""

# ╔═╡ e0933f5c-50b5-47cf-bae8-7a45a219fefb
md"""
### Histogram

The next plot we will see is a histogram plot. Check its arguments below...

- `x`: AbstractVector of values to be binned.
- `bins`: Integer, NTuple{2,Integer}, AbstractVector or Symbol. Default is :auto (the Freedman-Diaconis rule). For histogram-types, defines the approximate number of bins to aim for, or the auto-binning algorithm to use (`:sturges`, `:sqrt`, `:rice`, `:scott` or `:fd`). For fine-grained control pass a Vector of break values, e.g. `range(minimum(x), stop = maximum(x), length = 25)`.
- `weights`: Vector of weights for the values in x, for weighted bin counts.
- `normalize`: Bool or Symbol. Histogram normalization mode. Possible values are: `false`/`:none` (no normalization, default), `true`/`:pdf` (normalize to a discrete Probability Density Function, where the total area of the bins is 1), `:probability` (bin heights sum to 1) and `:density` (the area of each bin, rather than the height, is equal to the counts - useful for uneven bin sizes).
- `bar_position`: Symbol. Choose from `:overlay` (default), `:stack`. (warning: May not be implemented fully).
- `bar_width`: nothing or Number. Width of bars in data coordinates. When nothing, chooses based on x (or y when orientation = `:h`).
- `bar_edges`: Bool. Align bars to edges (true), or centers (the default).
- `orientation`: Symbol. Horizontal or vertical orientation for bar types. Values `:h`, `:hor`, `:horizontal` correspond to horizontal (sideways, anchored to y-axis), and `:v`, `:vert`, and `:vertical` correspond to vertical (the default).

Let's plot eruption length data...
"""

# ╔═╡ 943c8862-b909-4750-96da-7bda62c5fc8c
histogram(eruption, bins=:sqrt, label="Eruption", color=:orange,
		  xlabel="Eruption Length (minutes)", ylabel="Frequency")

# ╔═╡ ed480ae6-3302-4d0e-b83e-53a9c395d5ce
md"""
### Kernel density estimates

Next, we will see how we can fit a kernel density estimation function to our data. We will make use of the [KernelDensity.jl](https://github.com/JuliaStats/KernelDensity.jl) package.

With Kernel density estimates, we can see how things are changing...
"""

# ╔═╡ ed58e41f-3f73-40c1-92ec-5f2845c62b97
ρ = kde(eruption)

# ╔═╡ 952506dc-6cd8-4ecd-ace8-b9d914560289
md"""
If we want the histogram and the kernel density graph to be aligned we need to remember that the "density contribution" of every point added to one of these histograms is `density * nb of elements * bin width`. We do this, since KDE are densities, while the y-axis is showing the frequency.

**Note:** Read more about kernel density estimates [here](https://en.wikipedia.org/wiki/Kernel_density_estimation).
"""

# ╔═╡ 3ea44987-7ba7-4df6-a82a-c11a98ea784d
# the range of density data (x-axis)
ρ.x

# ╔═╡ 31274ae3-20da-4f52-9153-6f6b4d4472cb
# density data (y-axis)
ρ.density

# ╔═╡ 53f3e80a-2c88-405f-9131-21d07b132389
# density * nb of elements * bin width
scaled_density = ρ.density .* length(eruption) .* 0.2

# ╔═╡ ccfe5310-cc85-457c-bfbc-9671062b9088
begin
	histogram(eruption, bins=:sqrt, label="Eruption", color=:lightblue,
		  	  xlabel="Eruption Length (minutes)", ylabel="Frequency")

	plot!(ρ.x, scaled_density, linewidth=1.5, color=:black, label="KDE fit")
end

# ╔═╡ cbf471e5-523c-4550-835b-5c45b9648d8c
md"""
## Probability distributions

First, we will take a look at one probablity distribution, namely the Normal Distribution and verify that it generates a bell curve.
"""

# ╔═╡ 914412a4-69d4-4d9f-b575-03cfa03a7873
# generating random normal data
normaldata = randn(100_000)

# ╔═╡ 86c9cf10-2ade-45de-8649-a8dd5b2ae586
begin
	# generating the KDE
	ρₙ = kde(normaldata)

	# scaling density estimates
	scaled_ρₙ = ρₙ.density .* length(normaldata) * 0.1
end

# ╔═╡ e137810d-9e3c-43c8-905c-4efc2213a130
begin
	# plotting histogram
	histogram(normaldata, color=:lightblue, label="Z ~ N(0,1)",
			  xlabel="Z", ylabel="Frequency")

	# plotting KDE
	plot!(ρₙ.x, scaled_ρₙ, linewidth=3, color=:black, label="KDE fit")
end

# ╔═╡ f6430d83-a92b-4c03-a07e-8a6146cc9e1c
md"""
Another way to generate the same plot is via using the [Distributions.jl](https://github.com/JuliaStats/Distributions.jl) package and choosing the probability distribution you want, and then drawing random numbers from it.
"""

# ╔═╡ 4ffd9a7a-c62c-40a4-8ec1-4e2351a2a66a
md"""
### Normal distribution

As an example, we will use `N = Normal()` below to draw random numbers from a [Normal distribution](https://en.wikipedia.org/wiki/Normal_distribution). The default $\mu$ and $\sigma$ values are 0 and 1 (standard Normal distribution).
"""

# ╔═╡ 978e2f4c-b411-459c-8024-4f6205484f63
begin
	# drawing numbers ~N(0,1)
	N = Normal()
	gaussdata = rand(N, 100_000)

	# calculating KDE
	rhoN = kde(gaussdata)
	scaled_rhoN = rhoN.density .* length(gaussdata) .* 0.1
end

# ╔═╡ d1483a5e-0b7e-4317-adb0-8f024fad21eb
begin
	# plotting histogram
	histogram(gaussdata, color=:lightblue, label="Z ~ N(0,1)",
			  xlabel="Z", ylabel="Frequency")

	# plotting KDE
	plot!(rhoN.x, scaled_rhoN, linewidth=3, color=:black, label="KDE fit")
end

# ╔═╡ 62470113-95dc-4b92-9e1d-3118efa0bd08
md"""
### Binomial distribution

Now, let's draw random numbers from a [Binomial distribution](https://en.wikipedia.org/wiki/Binomial_distribution).
"""

# ╔═╡ 7749b6b5-cf01-4812-b1ce-6fe3204b0bcf
begin
	# drawing numbers ~B with n=40 trials
	B = Binomial(40)
	binomdata = rand(B, 100_000)

	# calculating KDE
	rhoB = kde(binomdata)
	scaled_rhoB = rhoB.density .* length(binomdata) .* 0.5
end

# ╔═╡ ff68a550-9cc6-416e-96b8-07de164e5fc4
begin
	# plotting histogram
	histogram(binomdata, color=:lightblue, label="Z ~ B(40,0.5)",
			  xlabel="Z", ylabel="Frequency")

	# plotting KDE
	plot!(rhoB.x, scaled_rhoB, linewidth=3, color=:black, label="KDE fit")
end

# ╔═╡ 693de4f8-88b1-476f-9319-64edafb049af
md"""
### Fit data to a distribution

Next, we will try to fit a given set of numbers to a parametric distribution. Thus, we are going to use the StatsBase.jl `fit` function.
"""

# ╔═╡ 7f9949d9-22e2-4d4e-92ec-6e52e01d10aa
begin
	# creating 1000 random values
	randomdata = rand(1000)
	# fitting random values to a Standard Normal distribution
	fitted_distrib = fit(Normal, randomdata)
	# drawing 1000 values from the new fitted distribution
	new_data = rand(fitted_distrib, 1000)
end

# ╔═╡ 98301aae-61ca-46ab-936a-fbb310854778
# note that `fitted_distrib` is a `Distribution`
fitted_distrib

# ╔═╡ f270c785-3a7a-4e43-9f21-b9a6e199383b
begin
	# plotting new data (Normal distribution)
	histogram(new_data, nbins=20, fillalpha=0.3, label="Normal Data")

	# plotting original data
	histogram!(randomdata, nbins=20, linecolor=:red, fillalpha=0.3,
			   label="Original Data")
end

# ╔═╡ 2d5e5740-a3c6-4448-b79f-f231a0ac9d96
md"""
We can try to fit eruptions data to a Normal distribution, although eruptions distribution does not seem to be Normal.
"""

# ╔═╡ ef651442-e846-4a5e-9f3b-0da63afeca16
begin
	erup_distrib = fit(Normal, eruption)
	normal_erup = rand(erup_distrib, 1000)
end

# ╔═╡ a5a54536-4b60-4a9c-b423-2804cce37382
erup_distrib

# ╔═╡ 25b15bda-9bb4-4ed9-90e4-c1f937051314
begin
	# plotting fitted data (Normal distribution)
	histogram(normal_erup, nbins=20, fillalpha=0.3, label="Normal Data")

	# plotting original data
	histogram!(eruption, nbins=20, linecolor=:red, fillalpha=0.3,
			   label="Original Data")
end

# ╔═╡ b82db811-d97d-4efd-81bc-a48780b23dcb
md"""
## Hypothesis testing

Next, we will perform hypothesis testing using the [HypothesisTests.jl](https://github.com/JuliaStats/HypothesisTests.jl) package.

Let's try to use the `OneSampleTTest` function to perform a one sample t-test of random sampled values.

The one-sample t-test is a statistical hypothesis test used to determine whether an unknown population mean is different from a specific value. In this case, the default value under $H_0$ is 0.
"""

# ╔═╡ af54608d-209f-4d16-9573-b4b8634d2acf
# generating 1000 random normal values
hypdata = randn(1000)

# ╔═╡ 867dab59-1c8d-44d9-9add-9f3999f1b86c
# performing one sample t-test
OneSampleTTest(hypdata)

# ╔═╡ 47c7f6df-0696-4b1e-b6c2-7eb8e782eb82
md"""
**Note:** as we can see above, we cannot reject the null hypothesis $H_0$. This was expected, since we drew sample values from a Standard Normal distribution, with $\mu=0$.

Now we can perform a one sample t-test of eruption data. We will probably be able to reject $H_0$, since the eruption length mean is close to 3.5 minutes...
"""

# ╔═╡ 6555c2a2-227f-419b-9b8b-d563767f3ea5
OneSampleTTest(eruption)

# ╔═╡ 78585d9d-5696-41dd-80f0-abab1585b008
md"""
## Correlation statistics

First of all, let's see visually if there is a correlation between `eruption` and `waittime` variables using a scatter plot...
"""

# ╔═╡ 650a1af5-b413-4794-bf51-77fa34cfb820
scatter(eruption, waittime, legend=false, color=:orange,
		xlabel="Eruption Length (minutes)", markersize=3,
		ylabel = "Time between Eruptions (minutes)")

# ╔═╡ 80db913d-faa8-43d5-9794-e17b02823cc5
md"""
These two variables may have a strong positive correlation. We can calculate both `Pearson` and `Spearman` correlation coefficients to check that.
"""

# ╔═╡ fcb84ac7-3bca-4284-8939-a55bd78a02fa
begin
	spearman_r = corspearman(eruption, waittime)
	"r(Spearman) = $spearman_r"
end

# ╔═╡ 960b8739-10d1-490d-ac49-e51527386161
begin
	pearson_r = cor(eruption, waittime)
	"r(Pearson) = $pearson_r"
end

# ╔═╡ 5bd50a69-4edb-435a-9263-e0b5d47851aa
md"""
Interesting! This means that the next time you visit Yellowstone National part ot see the faithful geysser and you have to wait for too long for it to go off, you will likely get a longer eruption!

Currently we are using the p-value of Spearman and Pearson correlation coefficients from Python. But you can follow the formula [here](https://stackoverflow.com/questions/53345724/how-to-use-julia-to-compute-the-pearson-correlation-coefficient-with-p-value) to implement your own.

Let's import the `scipy.stats` Python module.

**Note:** sometimes there are some issues getting Python and Julia to communicate as desired. We can explicity add `scipy.stats` using the [Conda.jl](https://github.com/JuliaPy/Conda.jl) package:

```julia
julia> using Conda
julia> Conda.add("scipy")
```
"""

# ╔═╡ d1ef6cd2-fda1-47f6-af12-8492ea70a539
scipy_stats = pyimport("scipy.stats");

# ╔═╡ c5f77715-02cb-437a-a45f-3f3a4dfba67b
begin
	rₛ, pvalueₛ = scipy_stats.spearmanr(eruption, waittime)
	"r(Spearman) = $(rₛ) || p-value = $(pvalueₛ)"
end

# ╔═╡ 615e7c40-bafe-475f-959d-29c5f5e845e9
begin
	rₚ, pvalueₚ = scipy_stats.pearsonr(eruption, waittime)
	"r(Pearson) = $(rₚ) || p-value = $(pvalueₚ)"
end

# ╔═╡ 32bc5a24-638e-43bc-a5c1-bd0a85db7849
md"""
## Confusion matrix & AUC

Finally, we will cover basic tools you will need such as confusion matrix and AUC scores. We use the [MLBase.jl](https://github.com/JuliaStats/MLBase.jl) package for that.

First, let's calculate a confusion matrix, given a target vector `y₁` and a prediction vector `ŷ₁`...
"""

# ╔═╡ 838a8a32-95f6-48c0-b87d-28a66ef431c5
begin
	# generating target and prediction vectors
	y₁ = [1, 1, 1, 1, 1, 1, 1, 2]
	yhat₁ = [1, 1, 2, 2, 1, 1, 1, 1]
end;

# ╔═╡ 68d7cd6c-bd76-406b-a3d7-4528ee9671fd
# compute default confusion matrix
C = confusmat(2, y₁, yhat₁)

# ╔═╡ 2bcf8647-368d-4b74-8e1a-4a4b1a577fdc
md"""
We may normalize cell values per class...
"""

# ╔═╡ 86ad3850-f631-42b2-8d19-bf20cba89c6b
C ./ sum(C, dims=2)

# ╔═╡ 38a2bb94-1d8f-4e1a-a07f-3bc1332301fd
md"""
Also, we may compute the accuracy (correct rate) from confusion matrix...

```math
accuracy = \frac{TP}{TP+TN+FP+FN}
```
"""

# ╔═╡ ea2c50c6-9404-4ba4-8fe7-4d1e1b3246ff
acc = sum(diag(C)) / length(y₁)

# ╔═╡ 70cbc1f9-171e-4664-87d1-6e73f0bc841c
md"""
Instead of calculating accuracy manually, we could use the `correctrate` function...
"""

# ╔═╡ 188c0219-8a54-4f28-95fd-9f2f525c12be
acc == correctrate(y₁, yhat₁)

# ╔═╡ 063d30ad-0ac4-4d50-b66b-9b06bf155fd8
md"""
Now, we can use the `roc` function to calculate ROC scores...
"""

# ╔═╡ f3cca2ef-17ec-44e0-996d-8a0cb2f2fcb1
begin
	# generating target and prediction vectors
	y₂ = [1, 1, 1, 1, 1, 1, 1, 0]
	yhat₂ = [1, 1, 0, 0, 1, 1, 1, 1]
end;

# ╔═╡ 9af05d0f-5daf-48f4-9949-79db70ae4bbc
ROC = MLBase.roc(y₂, yhat₂)

# ╔═╡ 92f98fe2-9154-4029-8282-45dee3f11d31
md"""
We may also use both `recall` and `precision` functions to compute these metrics...
"""

# ╔═╡ 0cfa5643-b65b-4374-a227-8f036ce32bbf
begin
	rec = recall(ROC)
	"Recall = $rec"
end

# ╔═╡ 93ea2a29-bac3-4629-bac5-5d03e8e45d80
begin
	prec = precision(ROC)
	"Precision = $prec"
end

# ╔═╡ Cell order:
# ╟─1c630e40-628c-11ec-1793-05688dabf6c9
# ╟─7242a33b-ec37-48d6-afca-8c8387a45a83
# ╟─dd2aa8b3-9da3-418d-bf31-74574483360d
# ╟─abc828f1-fd14-4329-8b47-f4135300f937
# ╟─a8aae3d4-198c-482f-a972-05be31be39c1
# ╠═63420679-9b77-4412-ba1c-fc471da6b1e9
# ╟─9d56b2cc-c4c6-4ce2-8d9c-33174fe5bc33
# ╠═67ac4a1d-eb15-473f-a548-07425fd627d8
# ╟─5ee886cc-2730-40cf-8497-507dcdc08c5f
# ╠═110fda28-1b77-44d6-8485-a4ec6e012011
# ╟─6e093fbb-1e3a-4131-b3c2-b5891e734fa7
# ╠═ed9da9c2-53e4-45d6-9ec1-02c8bf54c1a5
# ╟─13137e30-cfde-459a-adb4-1ddc07e626e8
# ╟─ff9d63e7-808a-4465-a68d-659a222afbae
# ╠═c3672771-248a-47f9-a019-2de03cdd79f4
# ╟─46b55947-428a-483f-a00f-a232e5329fbd
# ╟─e0933f5c-50b5-47cf-bae8-7a45a219fefb
# ╠═943c8862-b909-4750-96da-7bda62c5fc8c
# ╟─ed480ae6-3302-4d0e-b83e-53a9c395d5ce
# ╠═ed58e41f-3f73-40c1-92ec-5f2845c62b97
# ╟─952506dc-6cd8-4ecd-ace8-b9d914560289
# ╠═3ea44987-7ba7-4df6-a82a-c11a98ea784d
# ╠═31274ae3-20da-4f52-9153-6f6b4d4472cb
# ╠═53f3e80a-2c88-405f-9131-21d07b132389
# ╠═ccfe5310-cc85-457c-bfbc-9671062b9088
# ╟─cbf471e5-523c-4550-835b-5c45b9648d8c
# ╠═914412a4-69d4-4d9f-b575-03cfa03a7873
# ╠═86c9cf10-2ade-45de-8649-a8dd5b2ae586
# ╠═e137810d-9e3c-43c8-905c-4efc2213a130
# ╟─f6430d83-a92b-4c03-a07e-8a6146cc9e1c
# ╟─4ffd9a7a-c62c-40a4-8ec1-4e2351a2a66a
# ╠═978e2f4c-b411-459c-8024-4f6205484f63
# ╠═d1483a5e-0b7e-4317-adb0-8f024fad21eb
# ╟─62470113-95dc-4b92-9e1d-3118efa0bd08
# ╠═7749b6b5-cf01-4812-b1ce-6fe3204b0bcf
# ╠═ff68a550-9cc6-416e-96b8-07de164e5fc4
# ╟─693de4f8-88b1-476f-9319-64edafb049af
# ╠═7f9949d9-22e2-4d4e-92ec-6e52e01d10aa
# ╠═98301aae-61ca-46ab-936a-fbb310854778
# ╠═f270c785-3a7a-4e43-9f21-b9a6e199383b
# ╟─2d5e5740-a3c6-4448-b79f-f231a0ac9d96
# ╠═ef651442-e846-4a5e-9f3b-0da63afeca16
# ╠═a5a54536-4b60-4a9c-b423-2804cce37382
# ╠═25b15bda-9bb4-4ed9-90e4-c1f937051314
# ╟─b82db811-d97d-4efd-81bc-a48780b23dcb
# ╠═af54608d-209f-4d16-9573-b4b8634d2acf
# ╠═867dab59-1c8d-44d9-9add-9f3999f1b86c
# ╟─47c7f6df-0696-4b1e-b6c2-7eb8e782eb82
# ╠═6555c2a2-227f-419b-9b8b-d563767f3ea5
# ╟─78585d9d-5696-41dd-80f0-abab1585b008
# ╠═650a1af5-b413-4794-bf51-77fa34cfb820
# ╟─80db913d-faa8-43d5-9794-e17b02823cc5
# ╠═fcb84ac7-3bca-4284-8939-a55bd78a02fa
# ╠═960b8739-10d1-490d-ac49-e51527386161
# ╟─5bd50a69-4edb-435a-9263-e0b5d47851aa
# ╠═d1ef6cd2-fda1-47f6-af12-8492ea70a539
# ╠═c5f77715-02cb-437a-a45f-3f3a4dfba67b
# ╠═615e7c40-bafe-475f-959d-29c5f5e845e9
# ╟─32bc5a24-638e-43bc-a5c1-bd0a85db7849
# ╠═838a8a32-95f6-48c0-b87d-28a66ef431c5
# ╠═68d7cd6c-bd76-406b-a3d7-4528ee9671fd
# ╟─2bcf8647-368d-4b74-8e1a-4a4b1a577fdc
# ╠═86ad3850-f631-42b2-8d19-bf20cba89c6b
# ╟─38a2bb94-1d8f-4e1a-a07f-3bc1332301fd
# ╠═ea2c50c6-9404-4ba4-8fe7-4d1e1b3246ff
# ╟─70cbc1f9-171e-4664-87d1-6e73f0bc841c
# ╠═188c0219-8a54-4f28-95fd-9f2f525c12be
# ╟─063d30ad-0ac4-4d50-b66b-9b06bf155fd8
# ╠═f3cca2ef-17ec-44e0-996d-8a0cb2f2fcb1
# ╠═9af05d0f-5daf-48f4-9949-79db70ae4bbc
# ╟─92f98fe2-9154-4029-8282-45dee3f11d31
# ╠═0cfa5643-b65b-4374-a227-8f036ce32bbf
# ╠═93ea2a29-bac3-4629-bac5-5d03e8e45d80


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 17ada06e-a49a-11eb-29f8-35891d4afe6c
begin
	using Random
	using Plots
	using Base.Threads
	using Distributions
	using DataFrames
end

# ╔═╡ d79fcb85-28ca-4e32-bc4d-8c442b771cd6
md"""
# Random Election Upsets Riddler

This notebook solves the [Riddler Classic problem from April 23rd 2021](https://fivethirtyeight.com/features/can-you-cut-the-perfect-pancake/), which asks us to consider a large nation of voters who independently submit their vote between two candidates (A and B) by flipping identical fair coins. We're asked: if 80% of votes are tallied on election day (regular votes), while the remaining 20% (early votes) are tallied later, what is the probability of an upset -- i.e. what is the probability that the election day outcome doesn't match the final outcome?

As usual, I'll solve this problem in two ways. First I'll compute the result using some gratuitous math, and then I'll run a simulation of the random process to cross-check my analytical results. Since this problem isn't substantially harder if we assume that 20% of votes are submitted early, or that an arbitrary fraction of votes are submitted early, I'll solve the prompt in general and then report the solution for the prompt as a special case.
"""

# ╔═╡ 7c4713bb-7e75-4c0b-815c-e817b71e2f8b
md"""
## Analytical Solution

### Notation
First let's introduce a bit of notation. Throughout, we will use the subscripts $E$ and $R$ to indicate quantities computed respectively for the early voter and regular voter subpopulations. Similarly, we will use a subscript $T$ to refer to the total population. We will use $v$ to denote the number of votes in favor of candidate A in a population and $n$ will denote the number of voters in a population. E.g. $v_E$ is the number of voters who voted for candidate A in the early voter subpopulation, and $n_E$ is the number of voters. Finally, we will use $f$ to denote the fracction of voters who belong to a subpopulation, e.g. $f_E = n_E / n_T$.

### General Solution
Since each voter places their vote by independently flipping a fair coin, the total number of votes for candidate A is distributed as $v_\text{T} \sim \text{Binomial}(n_T,~0.5)$, and this remains true when we sub-divide into regular and early voter sub-populations. That is,

$$\begin{array}{ccl}
v_R &\sim& \text{Binomial}(n_R,~0.5). \\
v_E &\sim& \text{Binomial}(n_E,~0.5), \\
\end{array}$$

The prompt tells us that the population is "large," which is a hint that we should  think about using the Normal approximation to the Binomial. To that end, we recall that if $m \sim \text{Binomial}(n, p)$, then $m$ is approximately $\text{Normal}(\mu_m,\sigma_m)$, where $\mu_m = E[m] = pn$, and $\sigma_m = \sqrt{np(1-p)}$. Hence,

$$\begin{array}{ccl}
v_R &\approx& \text{Normal}(\mu_R,\sigma_R). \\
v_E &\approx& \text{Normal}(\mu_E,\sigma_E), \\
\end{array}$$

where $\mu_R = 0.5n_R$, $\sigma_R = 0.5 \sqrt{n_R}$ and an analogous expressions hold for $\mu_E$ and $\sigma_E$.

An election upset will occur if candidate A wins the regular vote but loses the total vote, or if candidate A loses the regular vote and wins the total vote. Since both events occur with equal probability (since candidate A and B are interchangeable), it follows that

$$P(\text{upset}) = 2 \cdot P(\text{A loses the total vote} \cap \text{A wins the regular vote}).$$

We can then re-write this as,

$$\begin{eqnarray}
P(\text{upset}) &=& 2 \cdot P(\text{A wins the regular vote} \cap \text{A loses the total vote}) \\
&=& 2 \cdot P(v_R > 0.5 \cdot n_R \text{ and }  v_T < 0.5 \cdot n_T)\\
&=& 2 \cdot P(v_R > 0.5 \cdot n_R \text{ and } v_E + v_R < 0.5 \cdot (n_E + n_R)) \\
&=& 2 \cdot P(v_R > \mu_R \text{ and } v_E + v_R < \mu_E+\mu_R) \\
&=& 2 \cdot P(v_R > \mu_R \text{ and } v_E < \mu_E + \mu_R - v_R). \\
\end{eqnarray}$$

Here we have used the fact that candidate A wins a round of votes if $v > 0.5\cdot n$, and used the expressions for $\mu_R, \mu_E$ from above.

Next, we note that $v_E$ and $v_R$ are independent so their joint density is just the product of their respective densities. Using this obsevation and the fact that both variables are normally distributed we obtain:

$$P(\text{upset}) = 2 \cdot \int_{\mu_R}^\infty \,dv_R \int_{-\infty}^{\mu_R + \mu_E - v_R}\,dv_E
\frac{e^{-\frac{1}{2}(\frac{v_R-\mu_R}{\sigma_R})^2}}{\sigma_R\sqrt{2\pi}}
\frac{e^{-\frac{1}{2}(\frac{v_E-\mu_E}{\sigma_E})^2}}{\sigma_E\sqrt{2\pi}}.$$

We will now manipulate this expression so that it matches an expression that I managed to find in a table of integrals 🙂. To that end, we make the substitution $w_R = \frac{v_R-\mu_R}{\sigma_R\sqrt{2}}$, $w_E = \frac{v_E -\mu_E}{\sigma_E\sqrt{2}}$, which (after some book-keeping) yields:

$$P(\text{upset}) = 2 \cdot \int_{0}^\infty \,dw_R \frac{e^{-w_R^2}}{\sqrt{\pi}} \int_{-\infty}^{-\frac{\sigma_R}{\sigma_E}w_R}\,dw_E \frac{e^{-w_E^2}}{\sqrt{\pi}}$$

Next, we recall that $$\text{erfc}(x) = \frac{2}{\sqrt{\pi}} \int_x^\infty e^{-t^2}\,dt$$, and that $$\text{erf}(x) = 1 - \text{erfc}(x)$$ so the above can be written:

$$\begin{eqnarray}
P(\text{upset}) &=&\int_{0}^\infty \,dw_R \frac{e^{-w_R^2}}{\sqrt{\pi}} \left(1 - \text{erf}(\frac{\sigma_R}{\sigma_E}w_R)\right)\\
&=&\int_{0}^\infty \,dw_R \frac{e^{-w_R^2}}{\sqrt{\pi}} - \frac{1}{\sqrt{\pi}} \cdot \int_{0}^\infty\,dw_R~e^{-w_R^2} \cdot \text{erf}(\frac{\sigma_R}{\sigma_E}w_R)
\end{eqnarray}$$

The first integral is equal to $\frac{1}{2}\text{erfc}(0) = \frac{1}{2}$. For the second integral, we apply formula 2 from section 4.3 of [this table of integrals involving error functions](https://nvlpubs.nist.gov/nistpubs/jres/73b/jresv73bn1p1_a1b.pdf) with $a = \sigma_R/\sigma_E$ and $b = 1$, to obtain:

$$\begin{eqnarray}
P(\text{upset})&=&\frac{1}{2} - \frac{1}{\sqrt{\pi}} \left[\frac{\sqrt{\pi}}{2} - \frac{1}{\sqrt{\pi}}\tan^{-1}\left(\frac{\sigma_E}{\sigma_R}\right)\right]\\
&=&\frac{1}{\pi}\tan^{-1}\left(\frac{\sigma_E}{\sigma_R}\right)
\end{eqnarray}$$

Finally, since $\sigma_E = \sqrt{\frac{1}{2}f_E\cdot n_\text{total}}$, $\sigma_R = \sqrt{\frac{1}{2}f_R\cdot n_\text{total}}$, and $f_E + f_R = 1$, we conclude that:

$$\boxed{P(\text{upset})=\frac{1}{\pi}\tan^{-1}\left(\sqrt{\frac{f_E}{1-f_E}}\right)}$$

### Solution for the Prompt
In particular, for the case outlined in prompt with $f_E = 0.2$, we find:

$$\boxed{P\left(\text{upset}_\text{prompt}\right) = \frac{1}{\pi} \tan^{-1}\left(\frac{1}{2}\right) \approx 0.14758}$$

"""

# ╔═╡ 70ef11be-81b9-4e4d-b83a-f7a225a9001b
function exact_prob_upset(frac_early)
	return (1/π) * atan(sqrt(frac_early / (1 - frac_early)))
end

# ╔═╡ 24e2efe2-f6f2-4648-8c28-bffedf98095b
exact_prob_upset(0.2)

# ╔═╡ 2383b70b-0693-4adf-99d4-9bcedec5d6cb
md"""
## Simulation Solution

As discussed above, the number of votes for candidate A in the early and regular vote subpopulations are both Binomial random variables. So, we  simulate the election tally by simulating appropriate Binomial random variables. From the simulated results, we next determine if A wins on election night (i.e. if more than 50% of the regular vote is for A) and determine the overall winner (i.e. if more than 50% of the total vote is for A). An election upset results if A wins in either the regular or total tally, but loses in the other tally. We simulate this process 100K times and tally up the fraction of times that an election upset occurred.

This process is repeated for 1001 evenly spaced early voter share fractions in [0,1], and the result is then compared (graphically) against the analytical solution.
"""

# ╔═╡ 562b2637-ec4f-4be2-a525-09e83c6458fa
"""
    simulate_election(frac_early, n_simulations, [n_country, p_A])

	Inputs:
	frac_early - percentage of early voters, 20% in the problem prompt
	n_simulations - number of simulations to run

	Optional:
	n_country - Has to be big-enough™. defaults to 100_001 votes
	p_A - probability that a voter selects candidate A. defaults to 0.5

	Output:
	estimated probability of an upset to the election night results
"""
function simulate_election(frac_early, n_simulations, n_country = 100_001, p_A=0.5)

	# number of early and regular voters
	n_early = floor(Int, frac_early * n_country)
	n_regular = n_country - n_early

	# distributions of votes for A for early and regular voters
	d_early = Binomial(n_early, p_A)
	d_regular = Binomial(n_regular, p_A)

	# simulate votes for A for early, regular, and total
	n_votes_A_early = rand(d_early, n_simulations)
	n_votes_A_regular = rand(d_regular, n_simulations)
	n_votes_A = n_votes_A_early + n_votes_A_regular

	# fraction of vote for A between regular and all votes
	frac_votes_A_regular = n_votes_A_regular / n_regular
	frac_votes_A = n_votes_A / n_country

	# determine if A wins on election night and if A wins overall
	a_wins_regular = frac_votes_A_regular .> 0.5
	a_wins = frac_votes_A .> 0.5

	# an upset occurs if
	#  1. A wins  regular, but loses overall (i.e. True, False)
	#  2. A loses regular, but wins overall (i.e. False, True)
	upset = a_wins_regular .!= a_wins
	prob_upset = sum(upset) / n_simulations

	return prob_upset
end

# ╔═╡ 5f82ec16-38a7-4c44-aa5a-f76d18d521dd
begin
	n_sim = 100_000
	n_pop = 10_00_001
	frac_earlys = LinRange(0, 1, 1_001)

	prob_upset_sim = map(fe->simulate_election(fe, n_sim, n_pop), frac_earlys)
	prob_upset_exact = map(exact_prob_upset, frac_earlys)
end

# ╔═╡ b2346518-43b5-4c4d-9536-244514a4d5a3
plot(
	frac_earlys,
	hcat(prob_upset_exact, prob_upset_sim),
	label=["Exact" "Simulated"],
	legend=:bottomright,
	linewidth=[5 2],
	xlabel="Fraction of Early Voters",
	ylabel="Probability of Election Upset",
	title="Probability of Election Upset \nvs Fraction of Early Voters",
	size=(800,500)
)

# ╔═╡ 06cdba4f-426d-47e5-a4c0-c35f9a41d547
md"""
The above figure shows great agreement between the exact solution discussed above and the simulation results -- so I'd wager that my exact solution is probably correct! 🙂
"""

# ╔═╡ 4a8020c4-b9f4-4c22-8919-eb0f135b1725
savefig("election_upset_probability.png")

# ╔═╡ e1bfd335-b818-4931-8e71-03b46c90c9ec
md"""
### Simulating the Prompt Solution
Finally, we'll simulate the prompt problem 10 million times to get a bit more accuracy in our simulation.
"""

# ╔═╡ e55c8934-7268-4c5f-af6c-2098a48cc094
begin
	Random.seed!(2021)
	simulate_election(0.2, 100_000_000, n_pop)
end

# ╔═╡ 537e2e5d-378b-4676-8ba4-3a50f54eb84c
md"""
Setting the seed for reproducibility (but not selecting it to cherry-pick!), the simulation result agrees with the theoretical result to around four decimal places, which is pretty good!
"""

# ╔═╡ Cell order:
# ╟─d79fcb85-28ca-4e32-bc4d-8c442b771cd6
# ╟─7c4713bb-7e75-4c0b-815c-e817b71e2f8b
# ╠═70ef11be-81b9-4e4d-b83a-f7a225a9001b
# ╠═24e2efe2-f6f2-4648-8c28-bffedf98095b
# ╟─2383b70b-0693-4adf-99d4-9bcedec5d6cb
# ╠═562b2637-ec4f-4be2-a525-09e83c6458fa
# ╠═5f82ec16-38a7-4c44-aa5a-f76d18d521dd
# ╠═b2346518-43b5-4c4d-9536-244514a4d5a3
# ╟─06cdba4f-426d-47e5-a4c0-c35f9a41d547
# ╟─4a8020c4-b9f4-4c22-8919-eb0f135b1725
# ╟─e1bfd335-b818-4931-8e71-03b46c90c9ec
# ╠═e55c8934-7268-4c5f-af6c-2098a48cc094
# ╟─537e2e5d-378b-4676-8ba4-3a50f54eb84c
# ╠═17ada06e-a49a-11eb-29f8-35891d4afe6c


export project_l1_Duchi!
function project_l1_Duchi!(v::Union{Vector{TF},Vector{Complex{TF}}}, b::TF) where {TF<:Real}
# % PROJECTONTOL1BALL Projects point onto L1 ball of specified radius.
# %
# % w = ProjectOntoL1Ball(v, b) returns the vector w which is the solution
# %   to the following constrained minimization problem:
# %
# %    min   ||w - v||_2
# %    s.t.  ||w||_1 <= b.
# %
# %   That is, performs Euclidean projection of v to the 1-norm ball of radius
# %   b.
# %
# % Author: John Duchi (jduchi@cs.berkeley.edu)
# Translated (with some modification) to Julia 1.1 by Bas Peters
if (b <= TF(0))
  error("Radius of L1 ball is negative");
end
if norm(v, 1) <= b
  return v
end
lv=length(v)
u = Vector{TF}(undef,lv)
sv= Vector{TF}(undef,lv)

#u = sort(abs(v),'descend');
#u = sort(abs.(v), rev=true) #faster than the line below
#u = sort(v, by=abs , rev=true)

#use RadixSort for Float32 (short keywords)
u=copy(v)
if TF==Float32
 u = sort!(abs.(u), rev=true,alg=RadixSort)
else
 u = sort!(abs.(u), rev=true,alg=QuickSort)
end

#sv = cumsum(u);
sv = cumsum(u)

#rho = find(u > (sv - b) ./ (1:length(u))', 1, 'last');
#tmp = (u .> ((sv.-b)./ LinSpace(1,lv,lv)))#::BitVector

rho = max(1,min(lv,findlast(u .> ((sv.-b)./ (1.0:1.0:lv)))))
#convert(TF,rho) why was this here...
#rho = findlast(tmp)::Int64

#theta = max(0, (sv(rho) - b) / rho);
theta = max.(TF(0) , (sv[rho] .- b) ./ rho)::TF

#w = sign(v) .* max(abs(v) - theta, 0);
v .= sign.(v) .* max.(abs.(v) .- theta, TF(0))

return v
end #end project_l1_Duchi
#
# function project_l1_Duchi!{TF<:Real}(v::Union{Array{TF,3},Array{Complex{TF},3}}, b::TF,mode::String)
#
# if (b < TF(0))
#   error("Radius of L1 ball is negative");
# end
# (n1,n2,n3)=size(v)
#
# ##Slice based projection
# if mode == "x_slice"
#   u = Vector{TF}(n2*n3)
#   sv= Vector{TF}(n2*n3)
#   v = reshape(v,n1,n2*n3)
#   Threads.@threads for i=1:n1
#     u = sort(abs.(v[i,:]), rev=true)
#     sv = cumsum(u)
#     rho = findlast(u .> ((sv.-b)./ (1.0:1.0:(n2*n3))))
#     theta = max.(TF(0) , (sv[rho] .- b) ./ rho)::TF
#     @inbounds v[i,:] .= sign.(v[i,:]) .* max.(abs.(v[:,i]).-theta, TF(0))
#   end
#   v = reshape(v,n1,n2,n3)
# elseif mode == "z_slice"
#   u = Vector{TF}(n1*n2)
#   sv= Vector{TF}(n1*n2)
#   v = reshape(v,n1*n2,n3)
#   Threads.@threads for i=1:n3
#     u = sort(abs.(v[:,i]), rev=true)
#     sv = cumsum(u)
#     rho = findlast(u .> ((sv.-b)./ (1.0:1.0:(n1*n2))))
#     theta = max.(TF(0) , (sv[rho] .- b) ./ rho)::TF
#     @inbounds v[:,i] .= sign.(v[:,i]) .* max.(abs.(v[:,i]).-theta, TF(0))
#   end
#   v = reshape(v,n1,n2,n3)
# elseif mode == "y_slice"
#   permutedims(v,[2;1;3]);
#   v = reshape(v,n2,n1*n3)
#   u = Vector{TF}(n1*n3)
#   sv= Vector{TF}(n1*n3)
#   Threads.@threads for i=1:size(v,1)
#     u = sort(abs.(v[i,:]), rev=true)
#     sv = cumsum(u)
#     rho = findlast(u .> ((sv.-b)./ (1.0:1.0:(n1*n3))))
#     theta = max.(TF(0) , (sv[rho] .- b) ./ rho)::TF
#     @inbounds v[i,:] .= sign.(v[i,:]) .* max.(abs.(v[:,i]).-theta, TF(0))
#   end
#   v = reshape(v,n2,n1,n3);
#   permutedims(v,[2;1;3]);
# else
#   error("provided incorred mode for l1_projection")
# end #end if
# v=vec(v)
#
# return v
# end #end project_l1_Duchi


include("../../common/numder/src/numder_module.jl")
# include("testfunctions.jl")

module opti
using ..numder
using LinearAlgebra

include("opti_1d.jl")
include("line_search.jl")
include("gradient_descent.jl")

mutable struct BFGS <: DescentMethod
    B
end
function BFGS(x::Vector{T}) where {T}
    BFGS( Matrix(one(T)*I, length(x), length(x)) )
end

function init!(m::BFGS, x)
    m.B = Matrix(one(eltype(x))*I, length(x), length(x))
end

function bfgs_matrix_update(B, g, g_prev, x, x_prev)
    dg = g - g_prev
    dx = x - x_prev

    dxdg = dot(dx, dg)

    if dxdg < 0
        println("No update happens, some accuracy errors")
        println(" dxdg = $(dxdg) ")
        return B
    end

    t1 = (dx ./ dxdg) * (dg' * B)
    t1 = t1 + t1'

    t2f = 1 + (dg' * (B * dg)) / dxdg
    t2 = (dx .* (t2f / dxdg)) * dx'

    # B = B - t1 + t2
    B += (t2 - t1)
    return B
end

function quasi_newton(
    fcn,
    x0;
    tol_options::ToleranceOptions       = ToleranceOptions(),
    search_options::LineSearchOptions   = LineSearchOptions(),
    gradient_fcn                        = numder.gradient_fw,
    line_search_fcn                     = line_search,
    log_path                            = false)

    m = BFGS(x0)

    x = x0
    f = fcn(x)
    g = gradient_fcn(fcn, x)

    step_direction  = -m.B * g
    step_size       = line_search_fcn(fcn, x, step_direction, f, search_options)
    dx              = step_size * step_direction
    g_new           = gradient_fcn(fcn, x + dx)
    dg              = g_new - g
    m.B             = m.B * dot(dg, dx) / dot(dg, dg)

    x_prev  = x
    x       = x_prev + dx
    f_prev  = f
    f       = fcn(x)
    g_prev  = g
    g       = g_new

    # out parameters
    x_sol = x
    y_sol = f
    stopping_crit = unknown::GradientDescentStoppingCrit

    iter = 0
    for outer iter = 1:tol_options.max_iter

        # update B
        # dg = g - g_prev
        # dx = x - x_prev

        # dxdg = dot(dx, dg)

        # t1 = (dx ./ dxdg) * (dg' * m.B)
        # t1 = t1 + t1'

        # t2f = 1 + (dg' * (m.B * dg)) / dxdg
        # t2 = (dx .* (t2f / dxdg)) * dx'

        # m.B = m.B - t1 + t2

        m.B = bfgs_matrix_update(m.B, g, g_prev, x, x_prev)

        # step dir
        step_direction = - m.B * g
        step_size      = line_search_fcn(fcn, x, step_direction, f, search_options)

        dx = step_direction * step_size

        # stopping criterion based on the step size
        abs_step_size = norm( dx )
        if abs_step_size <= tol_options.step_size_tol
            x_sol = x
            y_sol = f
            stopping_crit = small_step_size::GradientDescentStoppingCrit
            break
        end

        x_prev = x
        x      = x_prev + dx
        f_prev = f
        f      = fcn(x)
        g_prev = g

        # stopping criteria for function value
        if abs(f - f_prev) <= tol_options.f_abs_tol || abs(f - f_prev) / ( abs(f_prev) == 0 ? mach_eps : abs(f_prev)) < tol_options.f_rel_tol
            x_sol = x
            y_sol = f
            stopping_crit = small_function_value_delta::GradientDescentStoppingCrit
            break
        end

        # calculate the gradient at the new point
        g = gradient_fcn(fcn, x)

        # stopping criterion based on the gradient norm
        norm_g = norm(g)
        if norm_g <= tol_options.g_abs_tol
            x_sol = m.x
            y_sol = m.f
            stopping_crit = small_gradient::GradientDescentStoppingCrit
            break
        end

    end

    return (x_sol, y_sol, iter, stopping_crit)

end

function bfgs(
	fcn,
	x0;
	grad_fcn = numder.gradient_fw,
	options = ToleranceOptions() )

    n_dim = length(x0)
    H     = Matrix{Float64}(I, n_dim, n_dim)
    eye   = Matrix{Float64}(I, n_dim, n_dim)

    # first step is the steepest gradient descent
    g0        = grad_fcn(fcn, x0)
    step_size = line_search(fcn, x0, -g0)
    x1        = x0 - step_size * g0
    g1        = grad_fcn(fcn, x1)

    y            = g1 - g0
    s            = -step_size * g0
    scale_factor = dot(y, s) / dot(y, y)
    for ii = 1:n_dim
        H[ii, ii] *= scale_factor
    end

    # main loop
    iter = 0
    for outer iter = 1:options.max_iter

        # update of the matrix
        s       = x1 - x0
        y       = g1 - g0

		sdy     = dot(s, y)
		Hs 		= H * s
		sHs     = dot(s, Hs)
		theta   = 1.0
		if sdy < 0.2 * sHs
			theta = 0.8 * sHs / (sHs - sdy)
		end

		r = theta * y + (1 - theta) * Hs

        Hy      = H * y
		rho     = 1 / sdy
        rhos    = rho * s
        Hyst    = Hy * rhos'
        sytHyst = rhos * (y' * Hyst)
        # sst     = rhos * s'
        sst 	= r * r' / dot(s, r)

		Hnew    = H - (Hyst + Hyst') + sytHyst + sst
        # Hnew   = (eye - rho * s * y') * H * (eye - rho * y * s') + rho * s * s'

        # dH = norm(Htest - Hnew)
        # println(dH)

        x0        = x1
        p         = -(Hnew * g1)
        step_size = line_search(fcn, x0, p)
        dx        = p * step_size
        x1        = x0 + dx

        if norm(dx) <= options.step_size_tol
            break
        end

        # options.f_abs_tol
        # options.f_rel_tol

        g0        = g1
        g1        = grad_fcn(fcn, x1)

        if norm(g1) <= options.g_abs_tol
            break
        end

        H = Hnew
    end

    return (x1, iter)
end

end # opti

#=

function nonlinear_conjugate_gradient_descent_template(beta_method, fcn, x0, options, wolfe_condition_options)
	x_act = x0
	y_act = fcn(x_act)
	x_next = x_act
	x_sol = x_act

	# first step is a gradient descent step
	g_act = NumDiff.gradient_fw(fcn, x_act)
	direction_act = -g_act
	alpha = strong_wolfe_line_search(fcn, x_act, direction_act,
	 								wolfe_condition_options.alpha_0,
									wolfe_condition_options.beta,
									wolfe_condition_options.sigma)
	step = alpha * direction_act
	x_act = x_act + step
	y_act = fcn(x_act)
	g_prev = g_act
	direction_prev = direction_act

	# main loop
	println("x: $x_act")
	iter = 1
	for outer iter = 1:options.max_iter

		g_act = NumDiff.gradient_fw(fcn, x_act)

		## specific method is applied
		beta = beta_method(g_act, g_prev, direction_prev)

		direction_act = -g_act + beta .* direction_prev

		alpha = strong_wolfe_line_search(fcn, x_act, direction_act)
		step = alpha * direction_act
		x_next = x_act + step
		y_next = fcn(x_next)

		# convergence checks
		if norm(g_act) <= options.gradient_tolerance
			x_sol = x_next
			break
		end
		if norm(step) <= options.step_tolerance
			x_sol = x_next
			break
		end
		if abs(y_next - y_act) <= options.function_value_change_tolerance
			x_sol = x_next
			break
		end

		# updates
		y_act = y_next
		x_act = x_next
		direction_prev = direction_act
		println("x: $x_act")
	end

	if iter == options.max_iter
		x_sol = x_next
	end
	return x_sol
end

function beta_polak_ribiere(g_act, g_prev, direction_prev)
	y_hat = g_act - g_prev
	beta = dot(g_act, y_hat) / dot( g_prev, g_prev )
	beta = max(beta, 0)
	return beta
end

function beta_dai_yuan(g_act, g_prev, direction_prev)
	y_hat = g_act - g_prev
	beta = dot(g_act, g_act) / dot(y_hat, direction_prev)
	beta = max(beta, 0)
	return beta
end

function beta_hestenes_stiefel(g_act, g_prev, direction_prev)
	y_hat = g_act - g_prev
	beta = dot(y_hat, g_act) / dot(y_hat, direction_prev)
	beta = max(beta, 0)
	return beta
end

function beta_hager_zhang(g_act, g_prev, direction_prev)
	# probably implementation error
	y_hat = g_act - g_prev
	numer = dot(y_hat, direction_prev)
	t1 = y_hat - 2 .* direction_prev .* dot(y_hat, y_hat) ./ numer
	t2 = g_act ./ numer
	beta = dot(t1, t2)
	# beta = max(beta, 0)
	return beta
end

function nonlinear_conjugate_gradient_descent_polak_ribiere(
	fcn,
	x0,
	options::GradientDescentOptions = default_options(GradientDescentOptions),
	wolfe_condition_options = default_strong_wolfe_conditions_options()
	)
	x_sol = nonlinear_conjugate_gradient_descent_template(beta_polak_ribiere, fcn, x0, options, wolfe_condition_options)
end

function nonlinear_conjugate_gradient_descent_dai_yuan(
	fcn,
	x0,
	options::GradientDescentOptions = default_options(GradientDescentOptions),
	wolfe_condition_options = default_strong_wolfe_conditions_options()
	)
	x_sol = nonlinear_conjugate_gradient_descent_template(beta_dai_yuan, fcn, x0, options, wolfe_condition_options)
end

function nonlinear_conjugate_gradient_descent_hestenes_stiefel(
	fcn,
	x0,
	options::GradientDescentOptions = default_options(GradientDescentOptions),
	wolfe_condition_options = default_strong_wolfe_conditions_options()
	)
	x_sol = nonlinear_conjugate_gradient_descent_template(beta_hestenes_stiefel, fcn, x0, options, wolfe_condition_options)
end

function nonlinear_conjugate_gradient_descent_hager_zhang(
	fcn,
	x0,
	options::GradientDescentOptions = default_options(GradientDescentOptions),
	wolfe_condition_options = default_strong_wolfe_conditions_options()
	)
	x_sol = nonlinear_conjugate_gradient_descent_template(beta_hager_zhang, fcn, x0, options, wolfe_condition_options)
end

=#


using AdFem
using PyPlot
using SparseArrays

function f1func(x,y)
    18.8495559215388*pi^2*sin(pi*x)^2*sin(pi*y)*cos(pi*y) - 6.28318530717959*pi^2*sin(pi*y)*cos(pi*x)^2*cos(pi*y) + pi*sin(pi*y)*cos(pi*x)
end
function f2func(x,y)
    -18.8495559215388*pi^2*sin(pi*x)*sin(pi*y)^2*cos(pi*x) + 6.28318530717959*pi^2*sin(pi*x)*cos(pi*x)*cos(pi*y)^2 + pi*sin(pi*x)*cos(pi*y)
end

m = 50
n = 50
h = 1/n
ν = 0.5
K = ν*constant(compute_fem_laplace_matrix(m, n, h))
B = constant(compute_interaction_matrix(m, n, h))
Z = [K -B'
-B spzero(size(B,1))]

bd = bcnode("all", m, n, h)
# Due to the rank insufficiency of bilinear quadrilateral element,
# we add more pressure constraints
bd = [bd; bd .+ (m+1)*(n+1); ((1:m) .+ 2(m+1)*(n+1))]

F1 = eval_f_on_gauss_pts(f1func, m, n, h)
F2 = eval_f_on_gauss_pts(f2func, m, n, h)
F = compute_fem_source_term(F1, F2, m, n, h)
rhs = [F;zeros(m*n)]
Z, rhs = impose_Dirichlet_boundary_conditions(Z, rhs, bd, zeros(length(bd)))
sol = Z\rhs

sess = Session(); init(sess)
S = run(sess, sol)

xy = fem_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
U = @. 2*pi*sin(pi*x)*sin(pi*x)*cos(pi*y)*sin(pi*y)
figure(figsize=(12,5))
subplot(121)
visualize_scalar_on_fem_points(U, m, n, h)
title("Reference")
subplot(122)
visualize_scalar_on_fem_points(S[1:(m+1)*(n+1)], m, n, h)
title("Computed")
savefig("stokes1.png")

U = @. -2*pi*sin(pi*x)*sin(pi*y)*cos(pi*x)*sin(pi*y)
figure(figsize=(12,5))
subplot(121)
visualize_scalar_on_fem_points(U, m, n, h)
title("Reference")
subplot(122)
visualize_scalar_on_fem_points(S[(m+1)*(n+1)+1:2(m+1)*(n+1)], m, n, h)
title("Computed")
savefig("stokes2.png")

xy = fvm_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
p = @. sin(pi*x)*sin(pi*y)
figure(figsize=(12,5))
subplot(121)
visualize_scalar_on_fvm_points(p, m, n, h)
title("Reference")
subplot(122)
visualize_scalar_on_fvm_points(S[2(m+1)*(n+1)+1:end], m, n, h)
title("Computed")
savefig("stokes3.png")


module problem122

description = """
Efficient exponentiation

The most naive way of computing n15 requires fourteen multiplications:
n × n × ... × n = n15

But using a "binary" method you can compute it in six multiplications:

n × n = n^2
n^2 × n^2 = n^4
n^4 × n^4 = n^8
n^8 × n^4 = n^12
n^12 × n^2 = n^14
n^14 × n = n^15

However it is yet possible to compute it in only five multiplications:

n × n = n^2
n^2 × n = n^3
n^3 × n^3 = n^6
n^6 × n^6 = n^12
n^12 × n^3 = n^15

We shall define m(k) to be the minimum number of multiplications to compute nk; for example m(15) = 5.
For 1 ≤ k ≤ 200, find ∑ m(k).
"""

# Calculates the minimum number of multiplications to compute n^k
# given the known powers, and a maximum cutoff
function muls(k :: Int, known = [1], maxp = k-1)
  if (k % 2) == 0
    return muls(div(k,2)) + 1 # for even powers, solution is always to square the half power

  elseif known[end] == k
    return length(known)-1  # solution found

  elseif length(known) >= maxp # no solution for given known path
    return -1

  elseif (maxp - length(known)) < floor(log2(k / known[end])) # no possible solution for given knowns
    return -1
  end

  # search for solutions by summing items in the knowns path
  for i = length(known):-1:1
    for j = i:-1:1
      n = known[i] + known[j]
      if n <= k && n > known[end]
        push!(known, n)
        result = muls(k, known, maxp)
        pop!(known)
        if result == -1
          break
        elseif result == length(known)
          return result
        elseif result < maxp
          maxp = result
        end
      end
    end
  end

  return maxp
end

function solve(maxk = 200)
  total = 0
  for i = 1:maxk
    total += muls(i)
  end
  return total
end

end


function ABCDQR_kfilter(y,u,A,B,C,D,Q,R,x10,P10)
	#
	# Kalman filter for model
	#
	# x_{t+1} = A*x_{t} + B*u_{t} + w_{t}
	# y_{t}   = C*x_{t} + D*u_{t} + v_{t}
	#
	# cov(w_{t},v_{t}) = [Q 0;0 R]
	#
	# javier.cara@upm.es, 2015-10
	#

	(ny,nt) = size(y)
	nx = size(A,1)

	# allocation
	xtt = zeros(nx,nt)
	Ptt = zeros(nx,nx,nt)
	xtt1 = zeros(nx,nt+1)
	Ptt1 = zeros(nx,nx,nt+1)
	et = zeros(ny,nt)
	St = zeros(ny,ny,nt)
	Kt = zeros(nx,ny,nt)
	loglik = 0.0

	# Filter
	xtt1[:,1] = x10
	Ptt1[:,:,1] = P10
	for t in 1:nt
		#  innovations
		et[:,t] = y[:,t] - C*xtt1[:,t] - D*u[:,t]
		St[:,:,t] = C*Ptt1[:,:,t]*C' + R # et variance

		# Kalman gain
		Stinv = eye(ny,ny)/St[:,:,t] # numerically preferible to Stinv=inv(St)
		Kt[:,:,t] = Ptt1[:,:,t]*C'*Stinv # kalman gain

		# filtered values
		xtt[:,t] = xtt1[:,t] + Kt[:,:,t]*et[:,t]
		Ptt[:,:,t] = (eye(nx) - Kt[:,:,t]*C)*Ptt1[:,:,t]

		# one-step ahead prediction
		xtt1[:,t+1] = A*xtt[:,t] + B*u[:,t]
		Ptt1[:,:,t+1] = A*Ptt[:,:,t]*A' + Q

		# likelihood
		l0 = et[:,t]'*Stinv*et[:,t] # typeof(l0) = Array{Float64,1}
		loglik = loglik + log(det(St[:,:,t])) + l0[1]
	end

	loglik =  - ny*nt/2*log(2*pi) - 0.5*loglik

	return xtt,Ptt,xtt1,Ptt1,et,St,Kt,loglik

end


@doc raw"""
    PeriodicKernel([α=1 [,p=π]])

The periodic kernel is given by:

```math
\kappa(\mathbf{x},\mathbf{y}) =
\exp\left(-\alpha \sum_{i=1}^n \sin(p(x_i - y_i))^2\right)
\qquad p >0, \; \alpha > 0
```

where ``\mathbf{x}`` and ``\mathbf{y}`` are ``n`` dimensional vectors. The parameters ``p``
and ``\alpha`` are scaling parameters for the periodicity and the magnitude, respectively.
This kernel is useful when data has periodicity to it.
"""
struct PeriodicKernel{T<:AbstractFloat} <: MercerKernel{T}
    alpha::HyperParameter{T}
    PeriodicKernel{T}(α::Real) where {T<:AbstractFloat} = new{T}(
        HyperParameter(convert(T,α), interval(OpenBound(zero(T)), nothing))
    )
end
PeriodicKernel(α::T1 = 1.0) where {T1<:Real} = PeriodicKernel{floattype(T1)}(α)

@inline basefunction(::PeriodicKernel) = SineSquared()
@inline kappa(κ::PeriodicKernel{T}, z::T) where {T} = squaredexponentialkernel(z, getvalue(κ.alpha))


"""
Example modelling a linear beam element

Formulation:

K=-E*A*[dw_i!dx*du_j!dx]
"""

# local coordinates:
local_interp=InterpolationFunction(
    :(a0+a1*xi), # interpolation function expression
    [:xi], # local coordinate system degrees of freedom
    [:a0, :a1], # linear coefficients for regression
    [
        [0] # conditions to impose with regression: absolute value (order 0 derivative in the first axis)
    ];
    maximum_derivative=1
)

# absolute coordinates
abs_interp=InterpolationFunction(
    :(a0+a1*x), # interpolation function expression
    [:x], # absolute coordinate system degrees of freedom
    [:a0, :a1], # linear coefficients for regression
    [
        [0] # conditions to impose with regression: absolute value (order 0 derivative)
    ];
    maximum_derivative=1
)

# defining an integration domain:
dmn=Domain(
    local_interp, # local coordinate system
    [
        [0.0, 1.0] # boundaries for local coordinates: xi between 0 and 1
    ],
    [
        [0.0], # positions at which to impose values for interpolation:
        [1.0] # xi=0, xi=1
    ]
)

# defining a residual for within the beam:
domain_resd=Residual(
    [
        ( # variable w (weight), interpolated by InterpolationFunction abs_interp,
            :w, abs_interp, [
                (:dw!dx, [1]) # accompanied by its first order derivative
            ]
        )
        ( # variable u (displacement), interpolated by InterpolationFunction abs_interp,
            :u, abs_interp, [
                (:du!dx, [1]) # accompanied by its first order derivative
            ]
        )
    ],
    :(-E*(dw!dx*du!dx)),
    3; # Gauss-Legendre quadrature order
    extra_args=[:E, :A] # additional argument: elasticity
)

dmn_resd_func=get_volume_residual_function(domain_resd)

# defining a specific finite element:
fin=Finel(
    [
        [0.0],
        [1.0] # defined by corner points at x=0 and x=1 (corresponding to xi=0 and xi=1 in local coordinates: see Domain)
    ],
    dmn
)

# obtaining elasticity matrix:
K=zeros(Float64, 2, 2)

ws=zeros(Float64, 2)
us=zeros(Float64, 2)

E=1.0
A=1.0

for i=1:2
    ws[i]=1.0

    for j=1:2
        us[j]=1.0

        variable_vals=[ws, us]
        K[i, j]=dmn_resd_func(variable_vals, fin, E, A)

        us[j]=0.0
    end

    ws[i]=0.0
end

@assert isapprox(
    K,
    [
        -1.0 1.0;
        1.0 -1.0
    ]
)


using LinearAlgebra
using SparseArrays
using Plots

"This code solves Laplace's equation with Dirichlet boundary conditions using a
finite difference method."

m = 100 # number of points
x = LinRange(-1,1,m+2)
xint = x[2:end-1]

f(x) = 5*(Float64((x > -.5) && (x <= 0.0)) - Float64((x < .5) && (x >= 0.0)))
α,β = 1.0,pi

h = x[2]-x[1]
A = (1/h^2) * spdiagm(0=>2*ones(m),-1=>-ones(m-1),1=>-ones(m-1))
b = f.(xint)
b[1] += α/h^2
b[m] += β/h^2

# solve Au = F
uu = A\b
plot(xint,uu,linewidth=2,label="Solution")
plot!(x,uexact.(x),linestyle=:dash)
# plot!(x,.1*f.(x),linecolor=:red,ls=:dash,label="f(x)")
# plot!(leg=:topleft)

# uu - uexact.(xint)


# Headcount.jl

###### headcount #####
"""
    headcount(v, z)

Compute the Headcount Ratio of a vector `v` at a specified poverty threshold `z`.

# Examples
```julia
julia> using Inequality
julia> headcount([8, 5, 1, 3, 5, 6, 7, 6, 3], 4)
0.3333333333333333
```
"""
headcount(v::AbstractVector{<:Real}, z::Real)::Float64 = length(v[v .< z]) / length(v)

###### weighted headcount #####
"""
    headcount(v, w, z)

Compute the Headcount Ratio of a vector `v` at a specified poverty threshold `z`,
using weights given by a weight vector `w`.

Weights must not be negative, missing or NaN. The weights and data vectors must have the same length.

# Examples
```julia
julia> using Inequality
julia> headcount([8, 5, 1, 3, 5, 6, 7, 6, 3], [0.1,0.5,0.3,0.8,0.1,0.5,0.3,0.8,0.2], 4)
0.36111111111111116
```
"""
function headcount(v::AbstractVector{<:Real}, w::AbstractVector{<:Real}, z::Real)::Float64

    checks_weights(v, w)

    return sum(w[v .< z]) / sum(w)
end

function headcount(v::AbstractVector{<:Real}, w::AbstractWeights, z::Real)::Float64

    checks_weights(v, w)

    return sum(w[v .< z]) / w.sum

end

wheadcount(v::AbstractVector{<:Real}, w::AbstractVector{<:Real}, z::Real)::Float64 = headcount(v, w, z)

wheadcount(v::AbstractVector{<:Real}, w::AbstractWeights, z::Real)::Float64 = headcount(v, w, z)


function vandermonde_1d(N, r)
    V = zeros(length(r), N+1)
    for j in 1:N+1
        @. V[:, j] = jacobi(r, 0, 0, j-1)
    end
    return V
end

" Initialize the gradient of the modal basis (i) at (r) at order N. "
function ∂vandermonde_1d(N, r)
    DVr = zeros(length(r), N+1)
    for i in 0:N
        @. DVr[:, i+1] = ∂jacobi(r, 0, 0, i)
    end
    return DVr
end

" Initialize the (r) differentiation matrices on the interval, evaluated at (r) at order N. "
function Dmatrix1D(N, r, V)
    Vr = ∂vandermonde_1d(N, r)
    Dr = Vr / V
    return Dr
end


"""
    TF(r, Z)

Approximate solution of the [Thomas-Fermi equation](https://en.wikipedia.org/wiki/Thomas-Fermi_model)
for neutral atom with nuclear charge `Z`. Returns density at the distance `r`.

See Moliere, G. (1947). Theorie der streuung schneller geladener teilchen i. einzelstreuung am abgeschirmten coulomb-feld.
Zeitschrift für Naturforschung A, 2(3), 133-145.

## Example
For a density function ``\\rho(r)`` of ``N`` particles the following equation holds.
```math
N = 4\\pi\\int dr \\rho(r)r^2
```
On the grid ``r_i = \\exp(x_i)`` it is
```math
N = 4\\pi\\int dx \\rho(r(x))r^3
```
```jldoctest
julia> x = -30:0.1:20
-30.0:0.1:20.0

julia> r, n, dx = exp.(x), length(x), step(x);

julia> isapprox(4π * dx * sum(TF.(r, 18) .* r .^3), 18, atol=1e-10)
true
```
"""
function TF(r, Z)
  b = 1/2*(3π/4)^(2/3)*Z^(-1/3)
  x = r/b

  B₁, B₂, B₃ = 0.1, 0.55, 0.35
  β₁, β₂, β₃ = 6.0, 1.20, 0.30

  Z/(4π*b^3*x)*(B₁*β₁^2*exp(-β₁*x) + B₂*β₂^2*exp(-β₂*x) + B₃*β₃^2*exp(-β₃*x))
end


"""
    combinationrank(xs...)

Returns the index of the combination `xs` in the combinatorial number system. The elements
of `xs` all have to be unique, greater than zero and sorted in ascending order.
"""
function combinationrank(xs...)
    @assert length(unique(xs)) == length(xs)
    @assert sort([xs...]) == [xs...]
    @assert all(xs .> 0)
    k = length(xs)
    1 + sum(binomial.(xs .- 1, 1:k))
end

"""
    combinationunrank(k, i)

Returns the `i`-th `k`-element combination in the combinatorial number system.
"""
function combinationunrank(k, idx)
    @assert k > 0
    @assert idx > 0
    v = zeros(Int, k)
    for i = k:-1:1
        ck = _max_ck(i, idx - 1)
        v[i] = ck + 1
        idx -= binomial(ck, i)
    end
    tuple(v...)
end

function _max_ck(k, idx)
    ck = 0
    while true
        binomial(ck + 1, k) > idx && return ck
        ck += 1
    end
    error("Unreachable reached.")
end

"""
    findcombinationdiff(c1::NTuple{N,Int}, c2::NTuple{N,Int}) -> (idx1, idx2) or nothing

If the two ordered combinations of the same length have all elements except one in common,
return a pair of indices of those elements. Otherwise, return `nothing`.
"""
function findcombinationdiff(c1::NTuple{N,Int}, c2::NTuple{N,Int}) where N
    idx1::Union{Nothing,Int}, idx2::Union{Nothing,Int} = nothing, nothing
    i = 1
    while i < N
        i1 = isnothing(idx1) ? i : i + 1
        i2 = isnothing(idx2) ? i : i + 1
        if c1[i1] == c2[i2]
            i += 1
        elseif (c1[i1] < c2[i2]) && isnothing(idx1)
            idx1 = i1
        elseif (c1[i1] > c2[i2]) && isnothing(idx2)
            idx2 = i2
        else
            return nothing
        end
    end
    isnothing(idx1) && (idx1 = N)
    isnothing(idx2) && (idx2 = N)
    c1[idx1] == c2[idx2] ? nothing : (idx1, idx2)
end


export cycle_detection_floyd, cycle_detection_brent

# https://en.wikipedia.org/wiki/Cycle_detection#Algorithms

cycle_detection_floyd(f, x0) = begin
    # Main phase of algorithm: finding a repetition x_i = x_2i.
    # The hare moves twice as quickly as the tortoise and
    # the distance between them increases by 1 at each step.
    # Eventually they will both be inside the cycle and then,
    # at some point, the distance between them will be
    # divisible by the period λ.
    tortoise = f(x0) # f(x0) is the element/node next to x0.
    hare = (f ∘ f)(x0)
    while tortoise ≠ hare
        tortoise = f(tortoise)
        hare = (f ∘ f)(hare)
    end

    # At this point the tortoise position, ν, which is also equal
    # to the distance between hare and tortoise, is divisible by
    # the period λ. So hare moving in circle one step at a time,
    # and tortoise (reset to x0) moving towards the circle, will
    # intersect at the beginning of the circle. Because the
    # distance between them is constant at 2ν, a multiple of λ,
    # they will agree as soon as the tortoise reaches index μ.

    # Find the position μ of first repetition.
    μ = 0
    tortoise = x0
    while tortoise ≠ hare
        tortoise = f(tortoise)
        hare = f(hare)   # Hare and tortoise move at same speed
        μ += 1
    end

    # Find the length of the shortest cycle starting from x_μ
    # The hare moves one step at a time while tortoise is still.
    # λ is incremented until λ is found.
    λ = 1
    hare = f(tortoise)
    while tortoise ≠ hare
        hare = f(hare)
        λ += 1
    end

    λ, μ
end

cycle_detection_brent(f, x0) = begin
    # main phase: search successive powers of two
    power = λ = 1
    tortoise = x0
    hare = f(x0)  # f(x0) is the element/node next to x0.
    while tortoise ≠ hare
        if power == λ  # time to start a new power of two?
            tortoise = hare
            power *= 2
            λ = 0
        end
        hare = f(hare)
        λ += 1
    end

    # Find the position of the first repetition of length λ
    μ = 0
    tortoise = hare = x0
    for i in 0:λ-1
        hare = f(hare)
    end
    # The distance between the hare and tortoise is now λ.

    # Next, the hare and tortoise move at same speed until they agree
    while tortoise ≠ hare
        tortoise = f(tortoise)
        hare = f(hare)
        μ += 1
    end

    λ, μ
end


# ---
# title: 70. Climbing Stairs
# id: problem70
# author: zhwang
# date: 2022-01-16
# difficulty: Easy
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/climbing-stairs/description/>
# hidden: true
# ---
#
# You are climbing a staircase. It takes `n` steps to reach the top.
#
# Each time you can either climb `1` or `2` steps. In how many distinct ways can
# you climb to the top?
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 2
#     Output: 2
#     Explanation: There are two ways to climb to the top.
#     1. 1 step + 1 step
#     2. 2 steps
#
#
# **Example 2:**
#
#
#
#     Input: n = 3
#     Output: 3
#     Explanation: There are three ways to climb to the top.
#     1. 1 step + 1 step + 1 step
#     2. 1 step + 2 steps
#     3. 2 steps + 1 step
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 45`
#
#
## @lc code=start
using LeetCode

function climbing_stairs(n::Int)::Int
    n == 1 && return 1
    a, b = 1, 2
    for _ in 3:n
        a, b = b, a + b
    end
    return b
end

## @lc code=end


# Simulate a Poisson point process on a disk
# Author: H. Paul Keeler, 2019.
# Website: hpaulkeeler.com
# Repository: github.com/hpaulkeeler/posts
# For more details, see the post:
# hpaulkeeler.com/testing-the-julia-language-with-point-process-simulations/

#Note: Need the .+ for adding a scalar to an array
#Also need . for sqrt, exp, cos, sin etc and assinging scalars to arrays
#Best to use vectors instead of 1-D matrices eg x=rand(n),  NOT x=rand(n,1).

using Distributions #for random simulations
using Plots #for plotting

#Simulation window parameters
r=1; #radius of disk
xx0=0; yy0=0; #centre of disk
areaTotal=pi*r^2; #area of disk

#Point process parameters
lambda=100; #intensity (ie mean density) of the Poisson process

#Simulate Poisson point process
numbPoints=rand(Poisson(areaTotal*lambda)); #Poisson number of points
theta=2*pi*(rand(numbPoints));#angular coordinates  of Poisson points
rho=r*sqrt.(rand(numbPoints));#radial coordinates of Poisson points

#Convert polar to Cartesian coordinates
xx=rho.*cos.(theta);
yy=rho.*sin.(theta);

#Shift centre of disk to (xx0,yy0)
xx=xx.+xx0;
yy=yy.+yy0;

plot1=scatter(xx,yy,xlabel ="x",ylabel ="y", leg=false);
display(plot1);


## Maximum entropy degree distribution of a network with S species and L links

function p_k(k::Int64, S::Int64, L::Int64)
    # k: degree of a species
    # S: number of species in the food web
    # L: number of links in the food web

    # return probability of degree k from MaxEnt distribution
    # p_k = c * r^k

    # constraint of mean
    mean_degree = (2 * L) / S

    # maximum entropy distribution
    c = 1 / (mean_degree - 1)
    r = (mean_degree - 1) / mean_degree

    return c * r ^ k
end

function maxent_dd(S::Int64, L::Int64)
    # S: number of species
    # L: number of links

    # return degree distribution with maximum entropy
    # i.e. probability for each k between 1 and S
    degree_dist = zeros(S)
    for k in 1:S
        degree_dist[k] = p_k(k, S, L)
    end

    return degree_dist
end


function naiveDFT(hk::Array{T}) where T
    N = length(hk);
    Wnk = Array{Complex{T}}(undef, N, N)

    prefactor = -2 * π * im / N

    for n = 0:(N - 1), k = 0:(N - 1)
        Wnk[k + 1, n + 1] = exp(prefactor * n * k)
    end

    return Wnk * hk
end

function naiveDFTIterative(hk::Array{T}) where T
    N = length(hk);
    Hk = zeros(Complex{T}, N)

    prefactor = -2 * π * im / N

    @inbounds @simd for k = 1:N
        s = zero(Complex{T})
        km1 = k - 1
        for j = 1:N
            s += exp(prefactor * (j - 1) * km1) * hk[j]
        end
        Hk[k] = s
    end

    return Hk
end


function log_oneplusexp(E::Float64)
    return (E < 50) ? log(1+exp(E)) : E
end

function xlogy(x::Float64, y::Float64)
    if x == 0.0
        return 0.0
    else
        return x*log(y)
    end
end

function log_binom(n::Float64, k::Float64)
    @assert n >= k
    if (n == 0) || (k == 0)
       return 0
    end
    return n*log(n)-k*log(k)-(n-k)*log(n-k)
end


# fibonacci seq - LaTeX
#JL Plots

Pkg.add("LaTeXStrings")
Pkg.add("Measures")
using LaTeXStrings, Measures
gr(bg = :white)

fib = zeros(12);
for i = 1:12
    fib[i] = (((1+sqrt(5))/2)^i - ((1-sqrt(5))/2)^i)/sqrt(5);
end

Plots.plot(fib,
    marker=:circle,
    linewidth= 2,
    xlabel=L"n",
    ylabel=L"F_n",
    annotation=(5, 100, L"F_n = \frac{1}{\sqrt{5}} \left[\left( \frac{1+\sqrt{5}}{2} \right)^n - \left( \frac{1-\sqrt{5}}{2} \right)^n \right]"))


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, Plots, Measures
end;

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing6.09"

# ╔═╡ 1aa41898-3afa-11eb-2f49-570853ce6435
begin
	N = 5*10^3
	alpha = 0.05
	confLevel = 1 - alpha
	z = quantile(Normal(),1-alpha/2)
end;

# ╔═╡ 7e997ce4-3b27-11eb-167a-f74a6dff6b1c
function randCI(n,p)
    sample = rand(n) .< p
    pHat = sum(sample)/n
    serr = sqrt(pHat*(1-pHat)/n)
    (pHat - z*serr, pHat + z*serr)
end

# ╔═╡ 7e99af3e-3b27-11eb-2b94-051fc56f27cf
cover(p,ci) = ci[1] <= p && p <= ci[2]

# ╔═╡ 7e9a284c-3b27-11eb-37b2-6fcc5cbb3adb
begin
	pGrid = 0.1:0.01:0.9
	nGrid = 5:1:50
	errs = zeros(length(nGrid),length(pGrid))

	for i in 1:length(nGrid)
		for j in 1:length(pGrid)
			Random.seed!(0)
			n, p = nGrid[i], pGrid[j]
			coverageRatio = sum([cover(p,randCI(n,p)) for _ in 1:N])/N
			errs[i,j] = confLevel - coverageRatio
		end
	end
end;

# ╔═╡ 7ea98df0-3b27-11eb-292f-4d1335c99239
begin
	default(xlabel = "p", ylabel = "n",
		xticks =([1:5:length(pGrid);], minimum(pGrid):0.05:maximum(pGrid)),
		yticks =([1:5:length(nGrid);], minimum(nGrid):5:maximum(nGrid)))

	p1 = heatmap(errs, c=cgrad([:white, :black]))
	p2 = heatmap(abs.(errs) .<= 0.04, legend = false, c=cgrad([:black, :white]))
	plot(p1,p2, size = (1000,400), margin = 5mm)
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing6.09"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═1aa41898-3afa-11eb-2f49-570853ce6435
# ╠═7e997ce4-3b27-11eb-167a-f74a6dff6b1c
# ╠═7e99af3e-3b27-11eb-2b94-051fc56f27cf
# ╠═7e9a284c-3b27-11eb-37b2-6fcc5cbb3adb
# ╠═7ea98df0-3b27-11eb-292f-4d1335c99239
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


## spherical.jl : implementation of spherical covariance function

## Spherical ##
struct Spherical{T} <: IsotropicCovarianceStructure{T}
    λ::T
    σ::T
    p::T

    function Spherical{T}(λ::T, σ::T, p::T) where T
        λ > 0 || throw(DomainError(λ, "correlation length λ of spherical covariance cannot be negative or zero"))
        σ > 0 || throw(DomainError(σ, "marginal standard deviation σ of spherical covariance cannot be negative or zero"))
        p >= 1 || throw(DomainError(p, "in p-norm, p must be greater than or equal to 1"))
        isinf(p) && throw(DomainError(p, "in p-norm, p cannot be infinity"))

        new{T}(λ, σ, p)
    end
end

"""
    Spherical(λ, [σ = 1], [p = 2])

Spherical covariance structure with correlation length `λ`, (optional) marginal standard deviation `σ` and (optional) `p`-norm, defined as

``C(x, y) = \\begin{cases} σ \\left(1 - \\displaystyle\\frac{3}{2}\\frac{ρ}{λ} + \\frac{1}{2}\\left(\\frac{ρ}{λ}\\right)^3\\right) & \\text{for }ρ≤λ\\\\0 & \\text{for }ρ>λ\\end{cases}``

with ``ρ = ||x - y||_p``.

# Examples
```jldoctest
julia> Spherical(0.1)
spherical (λ=0.1, σ=1.0, p=2.0)

julia> Spherical(1.0, σ=2)
spherical (λ=1.0, σ=2.0, p=2.0)

```
See also: [`Exponential`](@ref), [`Linear`](@ref), [`Whittle`](@ref), [`Gaussian`](@ref), [`SquaredExponential`](@ref), [`Matern`](@ref)
"""
Spherical(λ::Real; σ::Real=1.0, p::Real=2) = Spherical{promote_type(typeof(λ),typeof(σ),typeof(p))}(promote(λ, σ, p)...)

# evaluate spherical covariance
apply(s::Spherical, x::Real) = max(zero(x), 1 - 3/2 * x/s.λ + 1/2 * (x/s.λ)^3)

# short name
shortname(::Spherical) = "spherical"


"""
    VaR(states,prob,alpha::Float32)

implements the Value-at-Risk at level ``\\alpha`` defined by
```math
VaR_\\alpha (Y) = \\arg \\min_x \\left( x\\in \\mathbb{R} : F_Y(x) \\geq \\alpha \\right),
```
for the random variable ``Y`` defined by `states` and `prob`.
"""
function VaR(x::Vector{Float64}, f::Vector{Float64}, α::Float32)
    ind = sortperm(x[1:length(x)])
    x = x[ind]; f = f[ind];
    i = findfirst(p -> p≥α, cumsum(f))
    if i === nothing
        return x[end]
    else
        return x[i]
    end
end

"""
    CTE(states,prob,alpha::Float32)

implements the Conditional Value-at-Risk at level ``\\alpha`` defined by
```math
CTE_\\alpha (Y) = VaR_\\alpha(Y) + \\frac{1}{1-\\alpha} \\mathbb{E} \\left( Y- VaR_\\alpha (Y) \\right)_+ ,
```
for the random variable ``Y`` defined by `states` and `prob`.
"""
function CTE(x::Vector{Float64}, f::Vector{Float64}, α::Float32)
    x = -x; α = 1-α;
    x_α = VaR(x, f, α)
    if iszero(α)
        return -x_α
    else
        tail = x .≤ x_α
        result = (sum(x[tail] .* f[tail]) - (sum(f[tail]) - α) * x_α) / α
        return -result
    end
end

"""
    EVaR2(states,prob,beta)

Solves the optimization problem associated with the primal formulation of the Entropic Value-at-Risk:

```math
EVaR_\\alpha(Y) = \\min_{x >0} \\frac{1}{x} \\left( \\beta +  \\log\\mathbb{E} e^{xY} \\right),
```
where ``Y`` is the discrete random variable defined by `states` and `prob`.
Here the optimization is done via the goldenSearch optimization routine implemented as part of this package.
"""
function EVaR2(states::Vector{Float64}, prob::Vector{Float64},beta::Float32)
	if sum(prob) == 0
		prob = ones(length(states))./ length(states)
	end
	dualZ = zeros(length(states))
	esssup = maximum(states[vec(prob.>0.0)])
	varm = dot(prob, (states .- dot(prob,states)).^2)
	if beta == 0.0
		Fvalue = dot(prob, states)
		tOpt = 0.0
		dualZ = exp.(tOpt .*states)
	else
		if varm >0.0
			tOpt = sqrt.(2 *beta/varm)
			function objective1(t)
				t=max(t,0.0)
				w = states./t
				offset = maximum(w)
    			we = exp.(w .- offset)
    			s = dot(prob, we)
    			w = log(s) + offset
				tmp = t*beta + t*w
				#tmp = t*beta + t*log(dot(prob,exp.(states./t)))
				return(tmp)
			end
			out = goldenSearch(objective1,tOpt)
			tOpt = out[1]
			Fvalue = out[2]
		else
			tOpt = 0.0
			Fvalue = 0.0
			dualZ = esssup
		end
		#if Fvalue < esssup
			dualZ = exp.(tOpt .*states)
		#else
		#	Fvalue = esssup
		#	dualZ[states .>= esssup] .= 1
		#end
		dualZ = dualZ / dot(prob,dualZ)
	end
	return(Fvalue,tOpt,dualZ) # for numerical stability
end

"""
    EVaR(states,prob,beta::Float32)

Solves the optimization problem associated with the primal formulation of the Entropic Value-at-Risk:

```math
EVaR_\\alpha(Y) = \\min_{x >0} \\frac{1}{x} \\left( \\beta +  \\log\\mathbb{E} e^{xY} \\right),
```
where ``Y`` is the discrete random variable defined by `states` and `prob`. Here, the optimization is done using JuMP and Ipopt.
"""
function EVaR(states::Vector{Float64}, prob::Vector{Float64},beta::Float32)
	if sum(prob) == 0
		prob = ones(length(states))./ length(states)
	end
	dualZ = zeros(length(states))
	esssup = maximum(states[prob.>0.0]) #error
	var = dot(prob, (states .- dot(prob,states)).^2 )
	if beta == 0.0
		EVaR = dot(prob, states)
		tOpt = 0
		dualZ = exp.(tOpt .*states)
	else
		if var >0
			tOpt = sqrt.(2 *beta/var)
			EV = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" =>0))
			function objective1(t)
				return( (beta.+ log(dot(prob,exp.(t.*states))))/t)
			end
			register(EV, :objective1, 1, objective1, autodiff=true)

	    	@variable(EV,t,start = tOpt)
			@constraint(EV, t .>=0)
	    	@NLobjective(EV, Min, objective1(t) )

			JuMP.optimize!(EV)
			EVaR = JuMP.objective_value(EV)
	    	tOpt = JuMP.value(t)
		else
			tOpt = 0
			EVaR = 0
			dualZ = esssup
		end
		if EVaR < esssup
			dualZ = exp.(tOpt .*states)
		else
			EVaR = esssup
			#dualZ[states==esssup] .= 1
		end
		#dualZ = dualZ / dot(prob,dualZ)
	end
	return(EVaR,tOpt,dualZ)
end

"""
    AVaR(states,prob,alpha::Float32)

Solves the optimization problem associated with the primal formulation of the Average Value-at-Risk:

```math
AVaR_\\alpha(Y) = \\min_{x\\in \\mathbb{R}} x + \\frac{1}{1-\\alpha} \\mathbb{E} \\left( Y - x \\right)_+,
```
where ``Y`` is the discrete random variable defined by `states` and `prob`.
"""
function AVaR(states::Vector{Float64},prob::Vector{Float64}, alpha::Float32)
    var2 = dot(prob, (states .- dot(prob,states)).^2);
     #First trivial case
    if alpha == 0.0
        return(dot(prob, states))
  #Now set up optimization problem using goldenSearch
    else
       tOpt = sqrt.(2 *alpha/var2)
       EV = Model(optimizer_with_attributes(Ipopt.Optimizer, "print_level" =>0))
		function objective1(t)
			return( t + 1/(1-alpha) * dot(prob, max.(states .- t,0)) )
		end
		register(EV, :objective1, 1, objective1, autodiff=true)
	    @variable(EV,t,start = tOpt)
		@NLobjective(EV, Min, objective1(t) )

		JuMP.optimize!(EV)
		AVaR = JuMP.objective_value(EV)
	    tOpt = JuMP.value(t)
    end
    return(AVaR, tOpt)
end


# The following reproduces the example introduced in Section 3.2.3 of Lütkepohl (2006)

using VectorAutoregressions
using CSV, DataFrames
using Dates

exdir = pkgdir(VectorAutoregressions) * "/examples";
datadf = DataFrame(CSV.File(exdir * "/lutkepohl_data.csv"));
first(datadf, 6)

# Take the subset from 1960q1 to 1978q4, and use log-difference for each series
datadf = datadf[datadf.date.<Date(1979, 3, 1), :];
Y = DataFrame(
    investment = diff(log.(datadf.invest)),
    income = diff(log.(datadf.income)),
    consumption = diff(log.(datadf.cons)),
);
first(Y, 6)

# ## Calculate lag selection criteria
ls = LagSelection(Y, 8)
# ## Estimate the VAR
v = VAR(Y, ls.selection["AIC"])
# ## Check stability of the VAR
sc = StabilityCheck(v)
# ## Check for autocorrelation in the residuals
pt = PortmanteauTest(v, 12)
bg1 = LMCorrTest(v, 4)
bg2 = LMCorrTest(v, 4, smallsample = true)


export ExtensionalSet, IntensionalSet, settuple, orderedpair, cardinality

abstract type MathematicalSet <: AbstractExpression end
istree(s::MathematicalSet) = false # TODO: refactor structure conditionals to make more sense

struct ExtensionalSet{T} <: MathematicalSet
    elements::Set{T}
end
ExtensionalSet(elements::Vector{T}) where {T} = ExtensionalSet(Set{T}(elements))
ExtensionalSet(elements::Tuple{Vararg{T}}) where {T} = ExtensionalSet(Set{T}([elements...]))
elements(es::ExtensionalSet) = es.elements
Base.length(es::ExtensionalSet) = length(elements(es))
Base.isempty(es::ExtensionalSet) = Base.isempty(elements(es))
Base.hash(es::ExtensionalSet, h::UInt) = hash(elements(es), h)
Base.:(==)(es1::ExtensionalSet, es2::ExtensionalSet) = elements(es1) == elements(es2)
cardinality(es::ExtensionalSet) = length(es)

# expression methods
variables(es::ExtensionalSet) = length(es) > 0 ? reduce(∪, variables.(elements(es))) : Set{LogicalSymbol}()
operations(es::ExtensionalSet) = length(es) > 0 ? reduce(∪, operations.(elements(es))) : Set{LogicalOperation}()

function Base.show(io::IO, es::ExtensionalSet)
    if isempty(es)
        print(io, "∅")
        return
    end

    print(io, "{")
    i = 1
    truncated_elements = truncate([elements(es)...], 5)
    for el ∈ truncated_elements
        print(io, el)
        if i < cardinality(es)
            print(io, ", ")
        end
        i += 1
    end
    if length(truncated_elements) < cardinality(es)
        missed_elements = cardinality(es) - length(truncated_elements)
        print(io, "...$(missed_elements) more")
    end
    print(io, "}")
end

orderedpair(a, b) = ExtensionalSet([ExtensionalSet([a]), ExtensionalSet([a, b])])
function settuple(a...)
    a = [a...]
    if length(a) == 0
        return ∅
    end
    x = settuple(a[1:end-1]...)
    ExtensionalSet([ExtensionalSet([x]), ExtensionalSet([x, a[end]])])
end

struct IntensionalSet <: MathematicalSet
    transform::AbstractExpression
    rule::AbstractExpression
end
IntensionalSet(transform::AbstractExpression, rules::Vector{LogicalExpression}) = IntensionalSet(transform, reduce(∧, rules))
IntensionalSet(transform::AbstractExpression, rules::Set{LogicalExpression}) = IntensionalSet(transform, [rules...])
IntensionalSet(transform::AbstractExpression, rules::Tuple{Vararg{LogicalExpression}}) = IntensionalSet(transform, [rules...])
transform(is::IntensionalSet) = is.transform
rule(is::IntensionalSet) = is.rule

# expression methods
variables(is::IntensionalSet) = variables(rule(is))
operations(is::IntensionalSet) = operations(rule(is))

function Base.show(io::IO, is::IntensionalSet)
    print(io, "{")
    print(io, transform(is))
    print(io, " | ")
    print(io, rule(is))
    print(io, "}")
end


# This file includes investigates the 2D interpolation mean square error(MSE) with respect to inteprolation point number.

using FractalTools
using Makie

# Construct interpolation data
f(x, y) = x^2 + y^2 + 1
vtx = [
    BigFloat.([0.0, 0.0]),
    BigFloat.([1.0, 0.0]),
    BigFloat.([0.5, 1.0])
    ]
freevar = 0.001
npts    = 50 : 5 : 150
ntpts   = 2 * npts[end]

# Construct test data
tpts = getdata(vtx, ntpts)

# Compute errors
mse = map(npts) do npt
    @info npt
    # Construct interpolation data
    pts = getdata(f, vtx, npt)

    # Construct interpolant
    interp = interpolate(pts, Interp2D(freevar))

    # Compute error
    fvals = map(pnt -> f(pnt...), tpts)
    ivals = map(pnt -> interp(pnt...), tpts)
    sum((fvals - ivals).^2) / length(tpts)
end

# Plot mse
fig = Figure()
ax = fig[1, 1] = Axis(fig, xlabel="Number of Points", ylabel="MSE", title="2D Interpolation MSE")
stem!(ax, npts, mse, color=:black)
save(joinpath(@__DIR__, "interp2d_mse.png"), fig)
display(fig)