# this file depends on jacobi.jl from JuliaOptics/OpticsPolynomials that
# should be found in the same directory

export zernike
export zernike_series
export zernike_sum
export zernike_norm
export zernike_nm_to_fringe
export zernike_nm_to_ansi_j
export zernike_ansi_j_to_nm
export zernike_noll_to_nm
export zernike_fringe_to_nm
export zernike_zero_separation

export apply_weights

"""
    apply_weights(data, weights)

Given data of arbitrary dimensionality with multiple realizations or modes along
the final dimension, weight the final dimension by weights and return the sum.

Will have one fewer dimension than the input data.
"""
function apply_weights(data, weights)
    shape = size(data)
    prod = reshape(data, (:, shape[end])) * weights
    return reshape(prod, shape[1: end-1])
end

function kronecker(i,j)
    return i==j ? 1 : 0
end

"""
    zernike_norm(n, m)

Norm of Zernike polynomial of radial order n, azimuthal order m.

The norm is the average squared distance to zero.  By multiplying a zernike
value by the norm, the term is given unit stdev or RMS.
"""
function zernike_norm(n::Integer, m::Integer)
    return √(2 * (n+1)) / (1 + kronecker(m, 0))
end

"""
    zernike_nm_to_fringe(n, m)

Map (n,m) ANSI indices to a single fringe index.
"""
function zernike_nm_to_fringe(n::Integer, m::Integer)
    term1 = (1 + (n + abs(m))/2)^2
    term2 = 2 * abs(m)
    term3 = (1 + sign(m)) / 2
    return Int(trunc(term1 - term2 - term3)) + 1
end

"""
    zernike_nm_to_ansi_j(n, m)

Map (n,m) ANSI indices to a single ANSI j index.

See also:
    - [`zernike_ansi_j_to_nm`](@ref) (reciprocal of this function)
"""
function zernike_nm_to_ansi_j(n::Integer, m::Integer)
    return (n * (n + 2) + m) ÷ 2
end

"""
    zernike_ansi_to_ansi_j(n, m)

Map (n,m) ANSI indices to a single ANSI j index.

See also:
    - [`zernike_nm_to_ansi_j`](@ref) (reciprocal of this function)
"""
function zernike_ansi_j_to_nm(j::Integer)
    n = (-3 + √(9 + 8j))÷2
    m = 2j - n * (n + 2)
    return Int(n), Int(m)
end

"""
    zernike_noll_to_nm(j)

Map j Noll index to ANSI (n,m) indices.
"""
function zernike_noll_to_nm(j::Integer)
    n = ceil(Int, (-1 + √(1 + 8j))/2 - 1)
    if n == 0
        m = 0
    else
        nseries = (n+1) * (n+2) ÷ 2
        residual = j - nseries

        if isodd(j)
            sign = -1
        else
            sign = 1
        end

        if isodd(n)
            ms = [1,1]
        else
            ms = [0]
        end

        for i=0:(n÷2)-1
            push!(ms, ms[end]+2)
            push!(ms, ms[end])
        end

        m = ms[end+residual] * sign
    end
    return n, Int(m)
end

"""
    zernike_fringe_to_nm(j)

Map j Fringe index to ANSI (n,m) indices.
"""
function zernike_fringe_to_nm(j::Integer)
    m_n = 2 * (ceil(Int, √j) - 1)
    g_s = (m_n ÷ 2)^2 + 1
    n = m_n ÷ 2 + (j-g_s)÷2
    m = (m_n - n) * (1 - mod(j-g_s, 2) * 2)
    return n, m
end

"""
    zernike_zero_separation(n)

Minimum zero separation of Zernike polynomial of radial order n.  Useful for
computing sample count requirements.
"""
function zernike_zero_separation(n::Integer)
    return 1 / n^2
end

"""
    zernike(n, m, ρ, θ[; norm])

Zernike polynomial of radial order n and azimuthal order m, evaluated at the
point (ρ, θ).  No normalization is required of (ρ, θ), though the polynomials
are orthogonal only over the unit disk.

norm is a boolean flag indicating whether the result should be orthonormalized
(scaled to unit RMS) or not.

The zernike polynomials' radial basis is a special case of the Jacobi
polynomials under the transformation n_jacobi = (n-m)/2, α=0, β=|m|, x=2ρ^2-1.

See also: [`zernike_series`](@ref), [`zernike_sum`](@ref).
"""
function zernike(n::Integer, m::Integer, ρ, θ; norm::Bool=true)
    x = 2 .* ρ.^2 .- 1
    am = abs(m)
    n_j = (n - am) ÷ 2
    out = jacobi(n_j, 0, am, x)
    if m != 0
        if m < 0
            out .*= (ρ.^am .* sin.(m.*θ))
        else
            out .*= (ρ.^am .* cos.(m.*θ))
        end

    end
	if norm
		out .*= zernike_norm(n,m)
	end
    return out
end

"""
    zernike_series(nms, ρ, θ[; norm])

Compute a series of Zernike polynomials of orders (n,m).
Returns an array with shape (size(ρ)..., length(nms)).
That is, the _final_ dimension contains the modes and the first dimension(s)
are spatial.

See also: [`zernike`](@ref), [`zernike_sum`](@ref)
"""
function zernike_series(nms, ρ, θ; norm::Bool=true)
    # 1.  Calculate a lookup table of Jacobi polynomials for the unique elements combinations of $|m|$ and $n_j$.
    # 2.  Calculate a lookup table of $\rho^{|m|}$ for each unique $|m|$
    # 3.  Calculate a lookup table of $\sin{m\theta}$ and $\cos{m\theta}$ for each unique $m$.
    # 4.  For each input $n, m$ look up the appropriate elements and compute the Zernike polynomial.
	x = 2 .* ρ.^2 .- 1
    n_j_list = [((n - abs(m)) ÷ 2, abs(m)) for (n,m) in nms]
    abs_m_type = typeof(n_j_list[1][2])
    dtype = eltype(ρ)
    jacobi_sequences_mnj = Dict{abs_m_type,Array{typeof(n_j_list[1][1]), 1}}()
	for (nj, am) in n_j_list
		a = get!(jacobi_sequences_mnj, am) do
			Array{abs_m_type, 1}()
		end
		push!(a, nj)
	end
	for key in keys(jacobi_sequences_mnj)
		# sort may not be needed.  Inherently sorted
		# but not sure if unique guarantees not shuffling
		# input
		sort!(unique!(jacobi_sequences_mnj[key]))
	end
    # jacobi sequences_mnj now maps |m| => max n_j
    # this is step 1 at the top of the function
    jacobi_sequences = Dict{abs_m_type,Array{dtype,ndims(ρ)+1}}()
    for k in keys(jacobi_sequences_mnj)
        njs = jacobi_sequences_mnj[k]
        jacobi_sequences[k] = jacobi_series(njs, 0., dtype(k), x)
    end

    # jacobi_sequences now maps |m| => arrays with n_j along the last dimension
    # powers_of_rho is now ρ^|m| for each unique |m|
    plane_idx = ndims(ρ)+1
    sines = Dict{abs_m_type, typeof(ρ)}()
    cosines = Dict{abs_m_type, typeof(ρ)}()
    powers_of_rho = Dict{abs_m_type, typeof(ρ)}()
    out = Array{eltype(ρ), ndims(ρ)+1}(undef, size(ρ)..., length(nms))
    for i in eachindex(nms)
        n, m = nms[i]
        absm = abs(m)

        out_view = selectdim(out, plane_idx, i)
		n_j = (n-abs(m))÷2
		jac_plane = first(searchsorted(jacobi_sequences_mnj[absm], n_j))
        jacobi_piece = selectdim(jacobi_sequences[absm], plane_idx, jac_plane)
        if m == 0
            if norm
                norm_val = zernike_norm(n, m)
                out_view .= norm_val .* jacobi_piece
            else
                out_view .= jacobi_piece
            end
        else
            if m < 0
				# these evaluate the sin/cosine only as needed,
				# the do block is only called if get! would error
                azpiece = get!(sines, m) do
					sin.(m .* θ)
				end
            else
                azpiece = get!(cosines, m) do
					cos.(m .* θ)
				end
            end
            radialpiece = get!(powers_of_rho, absm) do
				ρ .^ absm
			end

            if norm
                norm_val = zernike_norm(n, m)
                out_view .= norm_val .* jacobi_piece .* radialpiece .* azpiece
            else
                out_view .= jacobi_piece .* radialpiece .* azpiece
            end
        end
    end
    return out
end

"""
    zernike_sum(nms, weights, ρ, θ[; norm])

Compute a sum of Zernike polynomials of orders (n,m) weighted by weights.

See also: [`zernike`](@ref), [`zernike_series`](@ref)
"""
function zernike_sum(nms, weights, ρ, θ, norm::Bool=true)
    data = zernike_series(nms, ρ, θ, norm=norm)
    return apply_weights(data, weights)
end


# bessel_minus_struve.jl

module bessel_minus_struve
export BesselI_minus_StruveL

using ArbNumerics

"""
Implementation of I[n, z] - L[-n, z] where I[n, z] is the modified Bessel function of the first kind, and L[-n, z] is the modified Struve function. n is an integer ≥ 0 and the order of the functions, and z∈ℜ>0 a positive real argument. Notice the order are opposite signs.

The expansion of the BesselI and StruveL functions are:

    I[n, z] = ∑_{k=0}^∞ (z/2)^{2k+n} / (Γ(k+1)⋅Γ(k+n+1))
    L[-n, z] = ∑_{k=0}^∞ (z/2)^{2k-n+1} / (Γ(k+3/2)⋅Γ(k-n+3/2))

with Γ(x) the Gamma function. The values of I[n, z] and L[-n, z] can be incredibly close, so ArbReal types are used as they allow for accurate arbitrary precision floats so that the difference, I[n, z] - L[-n, z], can be distinguished. ArbReal types are faster than BigFloat types, and the error of a float can be tracked using ball(::ArbReal)[2].

Unfortunately, as ArbReal(bits = 64) * ArbReal(bits = 128) = ArbReal(bits = 64), if the starting precision is too small to accurately produce I[n, z] - L[-n, z], then the entire summation has to be restarted at a higher precision otherwise rounding errors will be too large, which slows the algorithm sometimes. Therefore, if the precision is found to be too low, the precision is just doubled (rather than increased by some set amount) to reach a required precision quickly.

The general structure of this arbitrary precision summation algorithm is used elsewhere too since you just change the arguments and term appropriately.
"""

function BesselI_minus_StruveL(n, z, a; prec = 64) # z > 0 & n >= 0

    # Initialise precision of ArbReal to prec.
    p = prec
    setextrabits(0)
    setprecision(ArbReal, p) # ArbReal(bit = 64 + 8) has same precision as Float64 (which is why we add 8 bits).

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

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

    while abs(midpoint(term)) > err * abs(midpoint(result))
        term = ArbReal((abs(z) * a / 2)^(2 * k + n + 1) / (ArbNumerics.gamma(k + 1) * ArbNumerics.gamma(k + n + 2)) - (abs(z) * a / 2)^(2 * k - n) / (ArbNumerics.gamma(k + 3//2) * ArbNumerics.gamma(k - n + 1//2)))
        result += term
        # println("term: k = ", k, "\nterm value: ", ball(term), "\ncumulant result: ", ball(result), "\n")
        k += 1
        # Double precision if rounding error in result exceeds accuracy specified by prec.
        while radius(result) > err * abs(midpoint(result))
            p *= 2
            setprecision(ArbReal, p)
            # variables have to be re-parsed into the higher precision ArbReal type.
            # println("Not precise enough. Error = ", radius(result), " > ", ArbReal("$err"), ". Increasing precision to ", p, " bits.\n")
            n = ArbReal("$n")
            z = ArbReal("$z")
            a = ArbReal("$a")
            k = ArbReal("0")
            result = ArbReal("0.0")
            term = ArbReal("1.0")
        end
    end
    # println("z: ", ArbReal(z, bits = prec), ". Final result: ", ArbReal(result, bits = prec))
    ArbReal(result * √π * ArbNumerics.gamma(-n - 1/2) * (abs(z) / 2)^n * z / 2, bits = prec) # return to specified precision.
end
# @time bsi = BesselI_minus_StruveL(3, 55, 1)
# @show(bsi)

# β = ArbReal("2.0")
# α = ArbReal("7.0")
# v = ArbReal("5.8")
# w = ArbReal("1.6")
# R = ArbReal((v^2 - w^2) / (w^2 * v))
# a = ArbReal(sqrt(β^2 / 4 + R * β * coth(β * v / 2)))
# z = ArbReal("90.61")
# n = ArbReal("12.0")
# c = BesselI_minus_StruveL(n, z, a; prec = 24)

end # end of module


# # 110: 1D Reaction Diffusion equation with two species
# ([source code](SOURCE_URL))
#
# Solve the nonlinear coupled reaction diffusion problem
#
# ```math
# -\nabla (0.01+2u_2)\nabla u_1 + u_1u_2= 0.0001(0.01+x)
# ```
#
# ```math
# -\nabla (0.01+2u_1)\nabla u_2 - u_1u_2 = 0.0001(1.01-x)
# ```
#
#
# in $\Omega=(0,1)$ with boundary condition $u_1(0)=1$, $u_2(0)=0$ and $u_1(1)=1$, $u_2(1)=1$.
#

module Example110_ReactionDiffusion1D_TwoSpecies

using Printf
using VoronoiFVM
using ExtendableGrids
using GridVisualize

function main(;n=100,Plotter=nothing,verbose=false,unknown_storage=:sparse)
    h=1/n
    grid=VoronoiFVM.Grid(collect(0:h:1))

    eps::Vector{Float64}=[1.0,1.0]

    physics=VoronoiFVM.Physics(num_species=2,

                               reaction=function(f,u,node)
                               f[1]=u[1]*u[2]
                               f[2]=-u[1]*u[2]
                               end,

                               flux=function(f,u0,edge)
                               u=unknowns(edge,u0)

                               nspecies=2
                               f[1]=eps[1]*(u[1,1]-u[1,2])*(0.01+u[2,1]+u[2,2])
                               f[2]=eps[2]*(u[2,1]-u[2,2])*(0.01+u[1,1]+u[1,2])
                               end,

                               source=function(f,node)
                               f[1]=1.0e-4*(0.01+node[1])
                               f[2]=1.0e-4*(0.01+1.0-node[1])
                               end,

                               storage=function(f,u,node)
                               f[1]=u[1]
                               f[2]=u[2]
                               end
                               )

    sys=VoronoiFVM.System(grid,physics,unknown_storage=unknown_storage)

    enable_species!(sys,1,[1])
    enable_species!(sys,2,[1])

    boundary_dirichlet!(sys,1,1,1.0)
    boundary_dirichlet!(sys,1,2,0.0)

    boundary_dirichlet!(sys,2,1,1.0)
    boundary_dirichlet!(sys,2,2,0.0)

    inival=unknowns(sys)
    U=unknowns(sys)
    inival.=0

    control=VoronoiFVM.NewtonControl()
    control.verbose=verbose
    control.damp_initial=0.1
    u5=0
    p=GridVisualizer(Plotter=Plotter,layout=(2,1))
    for xeps in [1.0,0.5,0.25,0.1,0.05,0.025,0.01]
        eps=[xeps,xeps]
        solve!(U,inival,sys,control=control)
        inival.=U
        scalarplot!(p[1,1],grid,U[1,:],clear=true,title="U1, eps=$(xeps)")
        scalarplot!(p[2,1],grid,U[2,:],clear=true,title="U2, eps=$(xeps)",reveal=true)
        sleep(0.2)
        u5=U[5]
    end
    return u5
end

function test()
    testval=0.7117546972922056
    main(unknown_storage=:sparse) ≈ testval && main(unknown_storage=:dense) ≈ testval
end
end


# Calcul des distances a vol d'oiseau entre 2 coordonnees GPS
# X. Gandibleux - Octobre 2019

using Printf
using DelimitedFiles

# Donnees dans un fichier au format CSV séparé par des virgules avec :
# - ligne 1 : header
# - colonne 1 : nom du lieu
# - colonne 2 : position x du lieu
# - colonne 3 : position y du lieu
# - colonne 4 : latitude du lieu
# - colonne 5 : longitude du lieu
# donnees extraites de : https://data.education.gouv.fr
m = readdlm("GPSlyceesNantes.csv", ',')

# Fonction de conversion degre vers radian
function deg_rad(v)
    return pi/180*v
end

# Calcul du distancier
lieu = String[]
for i in 2:size(m,1)
    push!(lieu, m[i,1])
end

dist = zeros(size(m,1)-1, size(m,1)-1)
for i in 2:size(m,1)-1
    for j in i+1:size(m,1)

        LatitudeA  = m[i,4]
        LongitudeA = m[i,5]
        LatitudeB  = m[j,4]
        LongitudeB = m[j,5]

        # formule des sinus : http://villemin.gerard.free.fr/aGeograp/Distance.htm
        # https://www.coordonnees-gps.fr/conversion-coordonnees-gps

        d = acos(
             sin(deg_rad(LatitudeA)) * sin(deg_rad(LatitudeB))
             +
             cos(deg_rad(LatitudeA)) * cos(deg_rad(LatitudeB)) * cos(deg_rad(LongitudeB) - deg_rad(LongitudeA))
            ) * 6371

        @printf("%2d - %2d = %6.3f km | %s <-> %s \n",i-1, j-1, d, m[i,1], m[j,1])

        dist[i-1,j-1] = d
        dist[j-1,i-1] = d
    end
end

println("distancier symétrique crée : ", length(lieu), " lieux")


using Distributions
using FillArrays
using StatsPlots

using LinearAlgebra
using Random

# Set a random seed.
Random.seed!(3)

# Define Gaussian mixture model.
w = [0.5, 0.5]
μ = [-3.5, 0.5]
mixturemodel = MixtureModel([MvNormal(Fill(μₖ, 2), I) for μₖ in μ], w)

# We draw the data points.
N = 60
x = rand(mixturemodel, N);

scatter(x[1, :], x[2, :]; legend=false, title="Synthetic Dataset")

using Turing

@model function gaussian_mixture_model(x)
    # Draw the parameters for each of the K=2 clusters from a standard normal distribution.
    K = 2
    μ ~ MvNormal(Zeros(K), I)

    # Draw the weights for the K clusters from a Dirichlet distribution with parameters αₖ = 1.
    w ~ Dirichlet(K, 1.0)
    # Alternatively, one could use a fixed set of weights.
    # w = fill(1/K, K)

    # Construct categorical distribution of assignments.
    distribution_assignments = Categorical(w)

    # Construct multivariate normal distributions of each cluster.
    D, N = size(x)
    distribution_clusters = [MvNormal(Fill(μₖ, D), I) for μₖ in μ]

    # Draw assignments for each datum and generate it from the multivariate normal distribution.
    k = Vector{Int}(undef, N)
    for i in 1:N
        k[i] ~ distribution_assignments
        x[:, i] ~ distribution_clusters[k[i]]
    end

    return k
end

model = gaussian_mixture_model(x);

sampler = Gibbs(PG(100, :k), HMC(0.05, 10, :μ, :w))
nsamples = 100
nchains = 3
chains = sample(model, sampler, MCMCThreads(), nsamples, nchains);

let
    # Verify that the output of the chain is as expected.
    for i in MCMCChains.chains(chains)
        # μ[1] and μ[2] can switch places, so we sort the values first.
        chain = Array(chains[:, ["μ[1]", "μ[2]"], i])
        μ_mean = vec(mean(chain; dims=1))
        @assert isapprox(sort(μ_mean), μ; rtol=0.1)
    end
end

plot(chains[["μ[1]", "μ[2]"]]; colordim=:parameter, legend=true)

plot(chains[["w[1]", "w[2]"]]; colordim=:parameter, legend=true)

chain = chains[:, :, 1];

# Model with mean of samples as parameters.
μ_mean = [mean(chain, "μ[$i]") for i in 1:2]
w_mean = [mean(chain, "w[$i]") for i in 1:2]
mixturemodel_mean = MixtureModel([MvNormal(Fill(μₖ, 2), I) for μₖ in μ_mean], w_mean)

contour(
    range(-7.5, 3; length=1_000),
    range(-6.5, 3; length=1_000),
    (x, y) -> logpdf(mixturemodel_mean, [x, y]);
    widen=false,
)
scatter!(x[1, :], x[2, :]; legend=false, title="Synthetic Dataset")

assignments = [mean(chain, "k[$i]") for i in 1:N]
scatter(
    x[1, :],
    x[2, :];
    legend=false,
    title="Assignments on Synthetic Dataset",
    zcolor=assignments,
)

if isdefined(Main, :TuringTutorials)
    Main.TuringTutorials.tutorial_footer(WEAVE_ARGS[:folder], WEAVE_ARGS[:file])
end


using Distributions

points = [1, 14, 2, 3, 4, 7, 9, 15, 22, 4]
function GMM(points::Array; k::Int64 = 2, tol::Float64 = .001)

    prior = 1 / k
    n = length(points)
    @assert n > k "Need more data points than latent classes!"

    post = repmat([prior], n, k)
    post = hcat(post, points)

    # initial random latent gaussians
    sd = sqrt(var(points))
    gauss = [Distributions.Normal(rand(points), sd) for i in 1:k]

    est_mu(w, points) = sum(w .* points) / sum(w)
    est_var(w, mu, points) = sum(w .* (points - mu).^2) / sum(w)

    prev = post[:, 1]
    converged = false

    while !converged

        # for each point, compute the posterior of belonging to each latent gaussian
        for i in 1:n
            p = post[i, end]
            normalizer = sum([pdf(norm, p) for norm in gauss]) * prior

            for j in 1:k
                post[i, j] = (pdf(gauss[j], p) * prior) / normalizer
            end
        end

        diff = sum(abs(prev - post[:, 1]))

        if diff < tol
            converged = true
        else
            prev = post[:, 1]
        end

        mus = [est_mu(post[:, i], points) for i in 1:k]
        vars = [est_var(post[:, i], mus[i], points) for i in 1:k]
        gauss = [Distributions.Normal(mus[i], sqrt(vars[i])) for i in 1:k]
    end

    post
end


# GaussQuadrature.jl
mutable struct DNq
    n::Array{Float64, 2}
    w::Array{Float64, 1}
end

"""
    initializenode(ndims::Int, quadpts::Int)
Initialize the nodes and weights of Gauss-Hermite quadrature (Multidimentional version).
Generate `ndims` dimensional quadrature nodes and weights.

# Arguments

# Optional arguments

- `dentype`
    - `:DN` Discrete normal
    - `:GH` Gauss-Hermite quadrature

# Example
```julia
ghn = initializenode(3, 21)
```
"""
function initializenode(ndims, quadpts; θrange = (-4, 4))
    nodes = range(θrange[1], stop = θrange[2], length = quadpts)
    weights = pdf.(Normal(), nodes) ./ sum(pdf.(Normal(), nodes))
    ghn = Dict(zip(1:quadpts, nodes))
    ghw = Dict(zip(1:quadpts, weights))
    # https://stackoverflow.com/questions/54256481/converting-array-of-cartesianindex-to-2d-matrix-in-julia
    as_ints(a::AbstractArray{CartesianIndex{L}}) where L = reshape(reinterpret(Int, a), (L, size(a)...))
    idx = convert(Matrix, as_ints(vec(CartesianIndices(Array{Float64}(undef, Tuple(fill(quadpts, ndims))))))')
    n = [ghn[idx[i, j]] for i in axes(idx, 1), j in axes(idx, 2)]
    w = prod([ghw[idx[i, j]] for i in axes(idx, 1), j in axes(idx, 2)], dims = 2)[:]
    return DNq(n, w)
end


# # Logistic regression
using DataFrames
using Plots
using RDatasets
using Convex
using SCS

#-

## we'll use iris data
## predict whether the iris species is versicolor using the sepal length and width and petal length and width
iris = dataset("datasets", "iris")
## outcome variable: +1 for versicolor, -1 otherwise
Y = [species == "versicolor" ? 1.0 : -1.0 for species in iris.Species]
## create data matrix with one column for each feature (first column corresponds to offset)
X = hcat(ones(size(iris, 1)), iris.SepalLength, iris.SepalWidth, iris.PetalLength, iris.PetalWidth);

#-

## solve the logistic regression problem
n, p = size(X)
beta = Variable(p)
problem = minimize(logisticloss(-Y.*(X*beta)))

solve!(problem, SCSSolver(verbose=false))

#-

## let's see how well the model fits
using Plots
logistic(x::Real) = inv(exp(-x) + one(x))
perm = sortperm(vec(X*beta.value))
plot(1:n, (Y[perm] .+ 1)/2, st=:scatter)
plot!(1:n, logistic.(X*beta.value)[perm])


function generatemonomials(x,m,;option=0)
    mons = [];
    if option == 0
        mons = monomials(x,[1:m;], mon -> (degree(mon) == m));
    elseif 1 <= option <= n
        mons = monomials(x,[1:m;], mon -> ((degree(mon) == m)&&(degree(mon,x[option])<=1)));
    else
        throw(DomainError(option,"option must be non-negative and <=$(n)"))
    end
    return mons
end

function differentiatepow(u,var::MP.AbstractVariable;power=1)
    res = u;
    if power>1
        res = differentiate(differentiatepow(u,var,power=power-1),var);
    elseif power==1
        res = differentiate(u,var);
    elseif power < 0
        throw(DomainError(power,"power must be non-negative"));
    end
    return res;
end

function dalpha(u,alpha,vars)
    if length(alpha) != length(vars)
        throw(DomainError(alpha,"alpha must be a valid multi-index for the variables x"))
    end
    res = u;
    for i in 1:1:length(alpha)
        res = differentiatepow(res,vars[i],power=alpha[i]);
    end
    return res;
end

function polyderivative(u,poly,vars)
    pterms = terms(poly);
    pcoefs = map(coefficient,pterms);
    pmons = map(monomial,pterms);
    fixedexp = z -> exponents(prod(vars.^0)*z);
    ppows = map(fixedexp,pmons);
    Dvec = map(z->dalpha(u,z,vars),ppows);
    if length(pcoefs)!= length(Dvec)
        throw(ErrorException("poly was not well defined"))
    end
    return sum(pcoefs[i]*Dvec[i] for i in 1:1:length(Dvec));
end

function laplacian(u,vars,;power=1)
    res = u;
    if power > 1
        res = polyderivative(laplacian(u,vars,power=power-1),sum(vars.*vars),vars);
    elseif power == 1
        res = polyderivative(u,sum(vars .* vars),vars);
    elseif power < 0
        throw(DomainError(power,"power must be non-negative"));
    end
    return res;
end

function normpartialder(hvec,var::MP.AbstractVariable,vars,;power=1)
    res = hvec;
    if power>1
        der = normpartialder(hvec,var,vars,power=power-1)
        res = (der[1]-2,der[1]*var*der[2]+sum(vars.*vars)*differentiate(der[2],var));
    elseif power==1
        res = (hvec[1]-2,hvec[1]*var*hvec[2]+sum(vars.*vars)*differentiate(hvec[2],var));
    elseif power < 0
        throw(DomainError(power,"power must be non-negative"));
    end
    return res;
    return
end

function normdalpha(hvec,alpha,vars)
    if length(alpha) != length(vars)
        throw(DomainError(alpha,"alpha must be a valid multi-index for the variables x"))
    end
    res = hvec;
    for i in 1:1:length(alpha)
        res = normpartialder(res,vars[i],vars,power=alpha[i]);
    end
    return res;
end

function normpolyderivative(hvec,poly,vars)
    pterms = terms(poly);
    pcoefs = map(coefficient,pterms);
    pmons = map(monomial,pterms);
    fixedexp = z -> exponents(prod(vars.^0)*z);
    ppows = map(fixedexp,pmons);
    Dvec = map(z->normdalpha(hvec,z,x),ppows);
    if length(pcoefs)!= length(Dvec)
        throw(ErrorException("poly was not well defined"))
    end
    return Dvec
end

function generatebasissphere(m, vars)
    n = length(vars)
    genpoly = (2-n,1);
    monos = generatemonomials(vars,m,option=1);
    fixedexp = z -> exponents(prod(vars.^0)*z);
    mpows = map(fixedexp,monos)
    return map(m->normdalpha(genpoly,m,vars)[2],mpows)
end

function dcoefficient(n,m,i,j)
    d = 0;
    if i == j
        d = 1.0/(2.0^j*factorial(j)*prod([(2*m + n - 2*j - 2*l) for l in 1:j]));
    elseif i < j
        d = -dcoefficient(n,m,i,j-1)/(2*(j-i)*(2*m+n-2-2*i-2*j));
    end
    return d;
end

function harmonicdecomposition(u,vars)
    norm = sum(vars .* vars);
    n =length(vars);
    m = maxdegree(u);
    if mindegree(u) != m
        throw(DomainError(u,"polynomial must be homogeneous"));
    end
    return [ (2*i,sum([dcoefficient(n,m,i,j)*norm^(j-i)*laplacian(u,vars,power=j) for j in i:floor(Int,m/2)])) for i in 0:floor(Int,m/2) ];
end


function sieve(n :: Int)
    isprime = trues(n)
    isprime[1] = false
    for p in 2:n
        if isprime[p]
            j = p * p
            if j > n
                return findall(isprime)
            else
                for k in j:p:n
                  isprime[k] = false
                end
            end
        end
    end
end


# # Holomorphic integration on the unit disk
# In this example, we explore integration of a harmonic function:
# ```math
#   f(x,y) = \frac{x^2-y^2+1}{(x^2-y^2+1)^2+(2xy+1)^2},
# ```
# over the unit disk. In this case, we know from complex analysis that the
# integral of a holomorphic function is equal to $\pi \times f(0,0)$.
# We analyze the function on an $N\times M$ tensor product grid defined by:
# ```math
# \begin{aligned}
# r_n & = \cos\left[(n+\tfrac{1}{2})\pi/2N\right],\quad{\rm for}\quad 0\le n < N,\quad{\rm and}\\
# \theta_m & = 2\pi m/M,\quad{\rm for}\quad 0\le m < M;
# \end{aligned}
# ```
# we convert the function samples to Chebyshev×Fourier coefficients using
# `plan_disk_analysis`; and finally, we transform the Chebyshev×Fourier
# coefficients to Zernike polynomial coefficients using `plan_disk2cxf`.
#
# For the storage pattern of the arrays, please consult the
# [documentation](https://MikaelSlevinsky.github.io/FastTransforms).

using FastTransforms, LinearAlgebra, Plots
const GENFIGS = joinpath(pkgdir(FastTransforms), "docs/src/generated")
!isdir(GENFIGS) && mkdir(GENFIGS)
plotlyjs()

# Our function $f$ on the disk:
f = (x,y) -> (x^2-y^2+1)/((x^2-y^2+1)^2+(2x*y+1)^2)

# The Zernike polynomial degree:
N = 15
M = 4N-3

# The radial grid:
r = [sinpi((N-n-0.5)/(2N)) for n in 0:N-1]

# The angular grid (mod $\pi$):
θ = (0:M-1)*2/M

# On the mapped tensor product grid, our function samples are:
F = [f(r*cospi(θ), r*sinpi(θ)) for r in r, θ in θ]

# We superpose a surface plot of $f$ on top of the grid:
X = [r*cospi(θ) for r in r, θ in θ]
Y = [r*sinpi(θ) for r in r, θ in θ]
scatter3d(vec(X), vec(Y), vec(0F); markersize=0.75, markercolor=:red)
surface!(X, Y, F; legend=false, xlabel="x", ylabel="y", zlabel="f")
savefig(joinpath(GENFIGS, "zernike.html"))
###```@raw html
###<object type="text/html" data="../zernike.html" style="width:100%;height:400px;"></object>
###```

# We precompute a (generalized) Zernike--Chebyshev×Fourier plan:
α, β = 0, 0
P = plan_disk2cxf(F, α, β)

# And an FFTW Chebyshev×Fourier analysis plan on the disk:
PA = plan_disk_analysis(F)

# Its Zernike coefficients are:
U = P\(PA*F)

# The Zernike coefficients are useful for integration. The integral of $f(x,y)$
# over the disk should be $\pi/2$ by harmonicity. The coefficient of $Z_{0,0}$
# multiplied by `√π` is:
U[1, 1]*sqrt(π)

# Using an orthonormal basis, the integral of $[f(x,y)]^2$ over the disk is
# approximately the square of the 2-norm of the coefficients:
norm(U)^2, π/(2*sqrt(2))*log1p(sqrt(2))

# But there's more! Next, we repeat the experiment using the Dunkl-Xu
# orthonormal polynomials supported on the rectangularized disk.
N = 2N
M = N

# We analyze the function on an $N\times M$ mapped tensor product $xy$-grid defined by:
# ```math
# \begin{aligned}
# x_n & = \cos\left(\frac{2n+1}{2N}\pi\right) = \sin\left(\frac{N-2n-1}{2N}\pi\right),\quad {\rm for} \quad 0 \le n < N,\quad{\rm and}\\
# z_m & = \cos\left(\frac{2m+1}{2M}\pi\right) = \sin\left(\frac{M-2m-1}{2M}\pi\right),\quad {\rm for} \quad 0 \le m < M,\\
# y_{n,m} & = \sqrt{1-x_n^2}z_m.
# \end{aligned}
# ```
# Slightly more accuracy can be expected by using an auxiliary array:
# ```math
#   w_n = \sin\left(\frac{2n+1}{2N}\pi\right),\quad {\rm for} \quad 0 \le n < N,
# ```
# so that $y_{n,m} = w_nz_m$.
#
# The x grid
w = [sinpi((n+0.5)/N) for n in 0:N-1]
x = [sinpi((N-2n-1)/(2N)) for n in 0:N-1]

# The z grid
z = [sinpi((M-2m-1)/(2M)) for m in 0:M-1]

# On the mapped tensor product grid, our function samples are:
F = [f(x[n], w[n]*z) for n in 1:N, z in z]

# We superpose a surface plot of $f$ on top of the grid:
X = [x for x in x, z in z]
Y = [w*z for w in w, z in z]
scatter3d(vec(X), vec(Y), vec(0F); markersize=0.75, markercolor=:green)
surface!(X, Y, F; legend=false, xlabel="x", ylabel="y", zlabel="f")
savefig(joinpath(GENFIGS, "dunklxu.html"))
###```@raw html
###<object type="text/html" data="../dunklxu.html" style="width:100%;height:400px;"></object>
###```

# We precompute a Dunkl-Xu--Chebyshev plan:
P = plan_rectdisk2cheb(F, β)

# And an FFTW Chebyshev² analysis plan on the rectangularized disk:
PA = plan_rectdisk_analysis(F)

# Its Dunkl-Xu coefficients are:
U = P\(PA*F)

# The Dunkl-Xu coefficients are useful for integration. The integral of $f(x,y)$
# over the disk should be $\pi/2$ by harmonicity. The coefficient of $P_{0,0}$
# multiplied by `√π` is:
U[1, 1]*sqrt(π)

# Using an orthonormal basis, the integral of $[f(x,y)]^2$ over the disk is
# approximately the square of the 2-norm of the coefficients:
norm(U)^2, π/(2*sqrt(2))*log1p(sqrt(2))


# Parts of this code were taken / derived from Graphs.jl. See LICENSE for
# licensing details.

type FloydWarshallState{T}<:AbstractPathState
    dists::Matrix{T}
    parents::Matrix{Int}
end

doc"""Uses the [Floyd-Warshall algorithm](http://en.wikipedia.org/wiki/Floyd–Warshall_algorithm)
to compute shortest paths between all pairs of vertices in graph `g`. Returns a
`FloydWarshallState` with relevant traversal information, each is a
vertex-indexed vector of vectors containing the metric for each vertex in the
graph.

Note that this algorithm may return a large amount of data (it will allocate
on the order of $\mathcal{O}(nv^2)$).
"""
function floyd_warshall_shortest_paths{T}(
    g::SimpleGraph,
    distmx::AbstractArray{T, 2} = DefaultDistance()
)

    n_v = nv(g)
    dists = fill(typemax(T), (n_v,n_v))
    parents = zeros(Int, (n_v,n_v))

    # fws = FloydWarshallState(Matrix{T}(), Matrix{Int}())
    for v in 1:n_v
        dists[v,v] = zero(T)
    end
    undirected = !is_directed(g)
    for e in edges(g)
        u = src(e)
        v = dst(e)

        d = distmx[u,v]

        dists[u,v] = min(d, dists[u,v])
        parents[u,v] = u
        if undirected
            dists[v,u] = min(d, dists[v,u])
            parents[v,u] = v
        end
    end
    for w in vertices(g), u in vertices(g), v in vertices(g)
        if dists[u,w] == typemax(T) || dists[w,v] == typemax(T)
            ans = typemax(T)
        else
            ans = dists[u,w] + dists[w,v]
        end
        if dists[u,v] > ans
            dists[u,v] = dists[u,w] + dists[w,v]
            parents[u,v] = parents[w,v]
        end
    end
    fws = FloydWarshallState(dists, parents)
    # for r in 1:size(parents,1)    # row by row
    #     push!(fws.parents, vec(parents[r,:]))
    # end
    # for r in 1:size(dists,1)
    #     push!(fws.dists, vec(dists[r,:]))
    # end

    return fws
end

function enumerate_paths(s::FloydWarshallState, v::Integer)
    pathinfo = s.parents[v,:]
    paths = Vector{Int}[]
    for i in 1:length(pathinfo)
        if i == v
            push!(paths, Vector{Int}())
        else
            path = Vector{Int}()
            currpathindex = i
            while currpathindex != 0
                push!(path,currpathindex)
                currpathindex = pathinfo[currpathindex]
            end
            push!(paths, reverse(path))
        end
    end
    return paths
end

enumerate_paths(s::FloydWarshallState) = [enumerate_paths(s, v) for v in 1:size(s.parents,1)]
enumerate_paths(st::FloydWarshallState, s::Integer, d::Integer) = enumerate_paths(st, s)[d]


export modnewton

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

Newton method with Hessian shift for solving min_x f(x)

"""
function modnewton(f::Function,df::Function,H::Function,x::Vector;maxIter=20,atol=1e-8,out::Int=0,
	storeInterm::Bool=false,beta=1e-1,
		lineSearch::Function=(f,df,fk,dfk,xk,pk)->armijo(f,fk,dfk,xk,pk,maxIter=10))

    his = zeros(maxIter,4)
    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)
        tau = (minimum(diag(d2f)) > 0) ? 0 :  -minimum(diag(d2f))+beta
        nfac = 0;
        for k=1:20
            try
                nfac +=1
                LL    = cholfact(d2f + tau*speye(length(x)))
                break;
            catch err
                if isa(err,Base.LinAlg.PosDefException) && k<20
                    tau = max(2*tau,beta)
                elseif isa(err,Base.LinAlg.PosDefException)
                    flag = -2
                    his[i,4] = nfac
                    his = his[1:i,:]
                    break; break;
                else
                    throw(err)
                end
            end
        end
        his[i,4] = nfac
        pk    = - (LL\dfk)

        # 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\tLS=%d\n" i his[i,1] his[i,2] his[i,3]
        end
        # update x and H
        x  += ak*pk

        i+=1
    end
    i = min(maxIter,i)

    if out>=0
        if flag==-1
            println(@sprintf("modnewton iterated maxIter (=%d) times but reached only atol of %1.2e instead of tol=%1.2e",i,his[i,2],atol))
        elseif flag==-2
            println(@sprintf("modnewton stopped because the Hessian at iteration %d was not positive definite.",i))
        elseif flag==-3
            println(@sprintf("modnewton stopped of a line search fail at iteration %d.",i))
		elseif out>1
            println(@sprintf("modnewton 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


#!/usr/bin/env julia

#=
# data should be sorted and works by divide and conquer
# run-time complexity O(log n)
=#

function binary_search(search_for, data_array, low=1, high=-1)
  if high == -1
    high = length(data_array)
  end
  if high < low
    return -1
  end
  mid = Int(floor(low + (high - low) / 2))
  if data_array[mid] == search_for
    return mid
  end
  if data_array[mid] > search_for
    return binary_search(search_for, data_array, low, mid-1)
  else
    return binary_search(search_for, data_array, mid+1, high)
  end
end

     ## 1   2   3   4   5   6   7   8   9   10
data = [10, 11, 12, 22, 24, 27, 34, 47, 56, 67]

@time println("at index: ", binary_search(47, data))
@time println("at index: ", binary_search(12, data))
@time println("at index: ", binary_search(7, data))
@time println("at index: ", binary_search(67, data))
@time println("at index: ", binary_search(34, data))


module BifurcationLib

# a demonstration of bifurcation using the logistic map

# equation: x_n+1 = r x_n (1-x_n)

# how does equilibrium change as the parameter r changes?

using Plots
# using ArgParse

### utility functions:

# rounds x to an aribtrary number of decimal places
function decimal_round(x::AbstractFloat, num_decimal_places::Unsigned)::AbstractFloat
  rounding_factor = 10^num_decimal_places
  round(rounding_factor * x)/rounding_factor
end

### Logistic map stuff

# gives logistic map of a value x and parameter (growth rate) r
logistic_map(x, r) = r*x*(1-x)

# gives many results of logistic map, starting with nth iteration and returning
# the results of m iterations after the nth
function logistic_map_series(x, r, n, m)
  nth_x = x
  for i = 1:n
    nth_x = logistic_map(nth_x, r)
  end

  mth_x = nth_x
  series = []
  for i = 1:m
    mth_x = logistic_map(mth_x, r)
    append!(series, [mth_x])
  end
  series
end

### Plotting things

# r0: first r to use

# r_inc: amount to increment r by each time

# r_final: final r (if r_final-r0 is not divisibible by r_inc then the actual
# final r will be slightly less than r_final)

function generate_points(x0, r0, n, m, r_inc, r_final)
  x = []
  y = []

  for r = r0:r_inc:r_final
    this_x = fill(r, (m))
    this_y = logistic_map_series(x0, r, n, m)
    append!(x, this_x)
    append!(y, this_y)
  end

  x, y
end

function plot_points(x, y, n, m)
  # set backend
  gr()

  # create plot
  plot(x, y,
    seriestype = :scatter,
    title = "Values of recursive iterations of logistic map for various r",
    xlabel = "r",
    ylabel = string("next ", m, " values after first ", n, " iterations"),
    dpi = 1000,
    markersize = 1,
    markeralpha = 0.5,
    # markeralpha = 0.1,
    # markerstrokestyle = :dot,
    # markerstrokecolor = :blue,
    markerstrokewidth = 0,
    markerfillcolor = :blue,
  )
end

function generate_and_plot(x0, r0, n, m, r_inc, r_final)
  println("generating points for plot with the following parameters:")
  println("x0      = ", x0)
  println("r0      = ", r0)
  println("n       = ", n)
  println("m       = ", m)
  println("r_inc   = ", r_inc)
  println("r_final = ", r_final)
  println()

  println("total number of poins to generate: ",
    floor( (r_final-r0) * r_inc ) * (n + m)
  )
  println()

  println("generating...")
  x, y = generate_points(x0, r0, n, m, r_inc, r_final)

  println("points generated. Plotting...")
  p = plot_points(x, y, n, m)
  # display(p)
  filename = string(
      "x0=", x0, ",",
      "r0=", r0, ",",
      "n=", n, ",",
      "m=", m, ",",
      "r_inc=", r_inc, ",",
      "r_final=", r_final, ".png")
  println("saving figure to ", filename)
  savefig(p, filename)
  println("done.")
  # p
end

function example_plot()
  x0 = 0.5
  r0 = 0.0
  n = 100
  m = 100
  r_inc = 0.01
  r_final = 4.0 # seems like anything higher than this diverges quickly to -inf

  generate_and_plot(x0, r0, n, m, r_inc, r_final)
end

end # module


# ---
# title: 279. Perfect Squares
# id: problem279
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, Dynamic Programming, Breadth-first Search
# link: <https://leetcode.com/problems/perfect-squares/description/>
# hidden: true
# ---
#
# Given a positive integer _n_ , find the least number of perfect square numbers
# (for example, `1, 4, 9, 16, ...`) which sum to _n_.
#
# **Example 1:**
#
#
#
#     Input: _n_ = 12
#     Output: 3
#     Explanation:12 = 4 + 4 + 4.
#
# **Example 2:**
#
#
#
#     Input: _n_ = 13
#     Output: 2
#     Explanation:13 = 4 + 9.
#
#
## @lc code=start
using LeetCode

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


######### CmdStan batch program example  ###########

using CmdStan, Statistics, Test

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

  dyes ="
  data {
    int BATCHES;
    int SAMPLES;
    real y[BATCHES, SAMPLES];
    // vector[SAMPLES] y[BATCHES];
  }

  parameters {
    real<lower=0> tau_between;
    real<lower=0> tau_within;
    real theta;
    real mu[BATCHES];
  }

  transformed parameters {
    real sigma_between;
    real sigma_within;
    sigma_between <- 1/sqrt(tau_between);
    sigma_within <- 1/sqrt(tau_within);
  }

  model {
    theta ~ normal(0.0, 1E5);
    tau_between ~ gamma(.001, .001);
    tau_within ~ gamma(.001, .001);

    mu ~ normal(theta, sigma_between);
    for (n in 1:BATCHES)
      y[n] ~ normal(mu[n], sigma_within);
  }

  generated quantities {
    real sigmasq_between;
    real sigmasq_within;

    sigmasq_between <- 1 / tau_between;
    sigmasq_within <- 1 / tau_within;
  }
  "

  dyesdata = Dict("BATCHES" => 6,
      "SAMPLES" => 5,
      "y" => reshape([
        [1545, 1540, 1595, 1445, 1595];
        [1520, 1440, 1555, 1550, 1440];
        [1630, 1455, 1440, 1490, 1605];
        [1595, 1515, 1450, 1520, 1560];
        [1510, 1465, 1635, 1480, 1580];
        [1495, 1560, 1545, 1625, 1445]
      ], 6, 5)
    )

  global stanmodel, rc, sim, cnames

  stanmodel = Stanmodel(name="dyes", model=dyes);

  @time rc, sim, cnames = stan(stanmodel, dyesdata, ProjDir,
    CmdStanDir=CMDSTAN_HOME)

  if rc == 0
    pi = filter(p -> length(p) > 2 && p[end-1:end] == "__", cnames)
    p = filter(p -> !(p in  pi), cnames)

    chns = set_section(sim,
      Dict(
        :parameters => ["theta",
          "tau_between", "tau_within",
          "sigma_between", "sigma_within",
          "sigmasq_between", "sigmasq_within"],
        :mu => ["mu.$i" for i in 1:6],
        :internals => pi
      )
    )
    describe(chns)
    describe(chns, sections=[:mu])
    #=
    println()
    println("Test round.(mean(theta)/100.0, digits=0) ≈ 15.0")
    @test round.(mean(sim[:,10,:])/100.0, digits=0) ≈ 15.0
    =#
  end
end # cd


using Distributions

# ρ(x,y) in integrand
ρ(x,y) = (0 < x < 1) && (0 < y < 2) ? 1/2 : 0

#=
Implementation of random walk sampler Metropolis-Hastings.
Inputs:
- n: The number of states to return
- b: The standard deviation of the normal distribution used in the proposal
     distribution
Outputs:
- Returns an array of length n denoting states S_i in Metropolis-Hastings alg.
=#
function metropolis_hastings(n, b)

    # Samples
    samples = zeros(n,2)
    samples[1,:] = rand(2)

    # Conditional Q distribution
    d = Normal(0,b)

    acceptances = 0

    for i=2:n
        # Propose Y
        ξx = rand(d)
        ξy = rand(d)
        X = samples[i-1,1] + ξx
        Y = samples[i-1,2] + ξy

        # Acceptance probability
        α = min(1, ρ(X,Y)/ρ(samples[i-1,1],samples[i-1,2]))

        # Next state
        if rand() < α
            samples[i,1] = X
            samples[i,2] = Y
        else
            samples[i,:] = samples[i-1,:]
        end
    end

    return samples
end

# Obtain 10000 samples with b = 1
S = metropolis_hastings(10000,1)

# f(x) in integrand
f(x,y) = exp(sin(x*y))

# Estimate integral
M = 1000
int_estimate = 2*mean([f(x,y) for (x,y) in zip(S[M:end,1], S[M:end,2])])


# Find check if n is prime

function is_prime(n::Int64)

	if n <= 3
	    return true;
	end

	if is_divisible(n, 2)
	    return false;
	end

	return true;

	i = 3;

	while i <= sqrt(n)
		if is_divisible(n, i)
			return false;
		end
		i += 2;
	end

	return true;
end

function is_divisible(n, i)

	return n % i == 0;

end

println(is_prime(101));


"""
    calc_nchk(n::Integer,k::Integer)

Calculates binomial coefficient: n choose k
"""
function calc_nchk(n::Integer,k::Integer)
    accum::Int = 1
    for i in 1:k
        accum = accum * (n-k+i) ÷ i
    end
    return accum
end

"""
    get_nchk(n::Integer,k::Integer)

Looks up binomial coefficient from a precomputed table: n choose k
"""
@inline function get_nchk(n,k)
    return binom_coeff[n+1,k+1]
end


"""
pde_lin_test(n)

Test the linear operators.
"""
function pde_lin_test(n)
    h = 1.0 / (n + 1)
    x = collect(h:h:1.0-h)
    fdata = fishinit(n)
    z = rand(n, n)
    L2d = Lap2d(n)
    DX = Dx2d(n)
    DY = Dy2d(n)
    lapok = lap_test(n, x, z, L2d, fdata)
    discok = disc_test(n, x, L2d, DX, DY, fdata)
    lapok && discok
end

"""
disc_test()
Have I broken the discretizations?
"""
function disc_test(n, x, L2d, DX, DY, fdata)
    n2 = n * n
    ue = solexact(x)
    ux = dxexact(x)
    uy = dyexact(x)
    D2u = l2dexact(x)
    ue1 = reshape(ue, (n2,))
    uex1 = reshape(ux, (n2,))
    uey1 = reshape(uy, (n2,))
    ued21 = reshape(D2u, (n2,))
    # test DX
    dx21 = DX * ue1
    dxerr = norm(dx21 - uex1, Inf)
    # test DY
    dy21 = DY * ue1
    dyerr = norm(dy21 - uey1, Inf)
    # test Laplacian
    du21 = L2d * ue1
    d2err = norm(du21 - ued21, Inf)
    pass = (d2err < 0.75) && (dxerr < 0.1) && (dyerr < 1.e-12)
    pass || println("Discretization test fails")
    return pass
end

"""
lap_test()

Does the FFT invert the discrete Laplacian?
Does the discrete Laplacian pass a simple eigenvalue test.
"""
function lap_test(n, x, z, L2d, fdata)
    randok = rand_test(z, n, L2d, fdata)
    eigok = eig_test(n, x, fdata)
    pass = randok && eigok
    return pass
end

function rand_test(z, n, L2d, fdata)
    n2 = n * n
    z1 = reshape(z, (n2,))
    y1 = L2d * z1
    y = reshape(y1, (n, n))
    mz = fish2d(y, fdata)
    q = reshape(mz, (n2,))
    pass = (norm(q - z1, Inf) < 1.e-12)
    pass || println("rand_test fails, norm =", norm(q - z1))
    return pass
end

function eig_test(n, x, fdata)
    lambda = pi * pi * 5
    efunx = sin.(pi * x)
    efuny = sin.(2 * pi * x)
    efunu = efunx * efuny'
    vfun = fish2d(efunu, fdata)
    pass = (norm(lambda * vfun - efunu, Inf) < 1.e-2)
    pass || println("eig test fails")
    return pass
end


"""
Cross-entropy loss as defined in PyTorch:

http://pytorch.org/docs/master/nn.html#torch.nn.CrossEntropyLoss

This method works with one-hot-encoded inputs
"""
function cross_entropy_loss(ŷ::AbstractMatrix, y::AbstractMatrix)
    sublosses = -sum(y .* ŷ, 1) .+ log.(sum(exp.(ŷ), 1))
    return mean(sublosses)
end

@require CuArrays begin

    # have to split sublosses into 2 subexpressions since otherwise CUDAnative failes with
    # 'ERROR: Broadcast output type Any is not concrete'
    function cross_entropy_loss(ŷ::CuArray{T,2}, y::CuArray{T,2}) where T
        w1 = -sum(y .* ŷ, 1)
        w2 = sum(exp.(ŷ), 1)
        sublosses = w1 .+ CUDAnative.log.(w2)
        return mean(sublosses)
    end

end


# Huber loss function

"""
  HuberLoss(rho::Real=1.0, mu::Real=1.0)

Returns the function `g(x) = (mu/2)||x||² if ||x|| ⩽ rho, and rho*mu*(||x||-rho/2) otherwise`.
"""

immutable HuberLoss{R <: Real} <: ProximableFunction
  rho::R
  mu::R
  function HuberLoss(rho::R, mu::R)
    if rho <= 0.0 || mu <= 0.0
      error("parameters rho and mu must be positive")
    else
      new(rho, mu)
    end
  end
end

HuberLoss{R <: Real}(rho::R=1.0, mu::R=1.0) = HuberLoss{R}(rho, mu)

function (f::HuberLoss){T <: Union{Real, Complex}}(x::AbstractArray{T})
  normx = vecnorm(x)
  if normx <= f.rho
    return (f.mu/2)*normx^2
  else
    return f.rho*f.mu*(normx-f.rho/2)
  end
end

function prox!{T <: Union{Real, Complex}}(f::HuberLoss, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  normx = vecnorm(x)
  mugam = f.mu*gamma
  scal = (1-min(mugam/(1+mugam), mugam*f.rho/(normx)))
  for k in eachindex(y)
    y[k] = scal*x[k]
  end
  normy = scal*normx
  if normy <= f.rho
    return (f.mu/2)*normy^2
  else
    return f.rho*f.mu*(normy-f.rho/2)
  end
end

fun_name(f::HuberLoss) = "Huber loss"
fun_dom(f::HuberLoss) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::HuberLoss) = "x ↦ (μ/2)||x||² if ||x||⩽ρ, μρ(||x||-ρ/2) otherwise"
fun_params(f::HuberLoss) = string("ρ = $(f.rho), μ = $(f.mu)")

function prox_naive{T <: Union{Real, Complex}}(f::HuberLoss, x::AbstractArray{T}, gamma::Real=1.0)
  y = (1-min(f.mu*gamma/(1+f.mu*gamma), f.mu*gamma*f.rho/(vecnorm(x))))*x
  if vecnorm(y) <= f.rho
    return y, (f.mu/2)*vecnorm(y)^2
  else
    return y, f.rho*f.mu*(vecnorm(y)-f.rho/2)
  end
end


# # Introduction
#
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__notebooks/01_Introduction.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__notebooks/01_Introduction.ipynb)
#
# Let's start with a classic:

message = "Hello World!"

#-

println(message)

# Julia is a general purpose programming language designed with scientific
# computing and mathematics in mind. For this reason, some Julia expressions
# look like mathematics.
#
# For example, let's write a Julia function to calculate
# a circle’s circumference $C$ as a function of the radius $r$:
#
# $C = 2 \pi r$

r = 10

C = 2π*r

# `π` is one of the
# [mathematical constants](https://docs.julialang.org/en/v1.1/base/numbers/#General-Number-Functions-and-Constants-1)
# defined in Julia. You can write `π` by typing `\pi` and pressing `<TAB>` in
# the REPL:

π

#-

pi

# The multiplication operator `*` is not needed when a literal number
# (e.g. `2`) is placed just before a variable or constant (e.g. `π`),
# see [*Numeric Literal Coefficients*](https://docs.julialang.org/en/v1.1/manual/integers-and-floating-point-numbers/#man-numeric-literal-coefficients-1).
# That makes polynomial expressions much cleaner:

f(x) = 1.56 + 2.24x + 3.47x^2

#-

f(0.04)

# This one-line function definition, i.e. `function_name(arguments) = body`, is
# the simplest we can use in Julia.

# #### Exercise 1
#
# Define an one-line function to calculate the length of the  hypotenuse of a
# right-angled triangle with catheti (legs) $a$ and $b$ as
# $\sqrt{a^{2} + b^{2}}$ using the `sqrt()` function:

## hypotenuse(...

# If your solution is correct, the following test should pass without errors.

using Test

# The `Test` module exports the `@test` macro, useful for writing unit tests of
# the code. We are going to use `≈`, `\approx<TAB>`, (`a ≈ b` or
# `isapprox(a, b)`) test if two values can be equal taking into account
# floating point errors. :

@test hypotenuse(2, 5) ≈ hypot(2, 5)

# In real life, you can use the function `hypot` to avoid
# [underflow and overflow problems](https://www.johndcook.com/blog/2010/06/02/whats-so-hard-about-finding-a-hypotenuse/).

# ## Control Flow
#
# ### `if`

distance = hypot(1.0, 3.0, 2.5)

#-

if distance <= 2.08  # Å
    "Disulphide"
elseif distance <= 6  # Å
    "Contact"
else
    "Not interacting"
end

# Each condition should be a boolean variable (`true` or `false`) or an
# expression that evaluates to a boolean value.

# ### Short-circuit evaluation

condition = true
condition && println("It's true!")

# #### Exercise 2
#
# Can you write the equivalent expression using `if` ?

## if ...

# ### Ternary operator

condition = true
result = condition ? "👍" : "👎"

# This is equivalent to:

result = if condition
        "👍"
    else
        "👎"
end

#-

result

# ## `for` loops

dna = "ATGCAT"  # dna is a string

for base in dna
    println(base)  # base is a character
end

# `for base = dna` and `for base ∈ dna` are also a possible notations, i.e.
# using `=` or `∈` instead of `in`:

for base ∈ dna
    println(base)
end

# Julia has a nice unicode support and unicode characters can be used in
# variables names, operators, etc. You can write `∈` in the REPL by typing `\in`
# and pressing the `<TAB>` key.

# ## Functions
#
# ### One-line syntax

"Return the number of codons in the sequence."
codon_number(seq) = div(length(seq), 3) # This is a Julia comment

# A strings before the function declaration is used as docstring. You can access
# the function documentation by typing `?` in the REPL and the name of the
# function, e.g.: `?codon_number`

## Try to look at the documentation of condon_number

#-

codon_number(dna)

# ### Multi-line syntax

# Functions with longer bodies are defined using the `function` keyword:

"Returns the GC number in a given sequence."
function count_gc(string)
    count = 0
    for char in string
        if char == 'C' || char == 'G'  # || is the short-circuiting or.
            ## ' ' (instead of " ") is used to define a character.
            count += 1  # i.e. count = count + 1
        end
    end
    count  # i.e. return count
end

#-

count_gc(dna)

# By default, a function returns the result of the last evaluated expression,
# `count` in the previous example. Otherwise, the keyword `return` should be
# used:

"""
`is_dna` iterates a string and returns:
- `false` if a character different from **A, C, T or G** is found
- `true` otherwise
"""
function is_dna(string)
    for char in string
        if char != 'A' && char != 'C' && char != 'T' && char != 'G'
            return false
        end
    end
    true
end

#-

is_dna("ACHL")

# Here, the `return` keyword is used inside the loop to return `false` as soon
# as an incorrect character is found.
#
# `"""` defines multiline strings. In this example, the multiline string is the
# documentation string of the function. Markdown syntax can be used in
# docstrings.

#md ?is_dna
#nb ?is_dna

# #### Exercise 3
#
# Write a function that returns `true` if the given string (argument) has at
# least one `N` using the keywords `function`, `for`, `if` and `return`.

## has_n(...

# `@testset` is useful to aggregate tests for a particular functionality:

using Test
@testset "Exercise 3" begin
    @test has_n("ACTGN")
    @test !has_n("ACTG")
end

# #### Material
#
# You can find more complete information about
# [functions](https://docs.julialang.org/en/v1.1/manual/functions/)
# and [control flow](https://docs.julialang.org/en/v1.1/manual/control-flow/)
# statements in the Julia manual.


md"""
$$
  x[n] = A \cos(\omega n T + \phi) = A \cos(2\pi f n T + \phi)
$$
- $A$: amplitude
- $\omega$: angular frequency (`radians/sec`)
- $f = \omega / 2\pi$: frequency in Hertz (`cycles/sec`)
- $\phi$: initial phase (`radians`)
- $n$: time index
- $t = nT = n / f_s$, where $f_s$ is the sampling frequency

"""
begin
  import Pkg
  Pkg.activate(mktempdir())
  Pkg.add([
      "AudioDisplay",
      "PyPlot",
  ])
end

begin
  A = .8
  f = 1000    # i.e. 1000 cycles/sec
  ϕ = pi/2
  fs = 44100  # i.e. 44100 samples/sec
  t = range(-.002, .002, step=1/fs)
  x = A .* cos.(2πf.*t .+ ϕ)
  size(x)
end

A * cos.(2π*f.*t .+ ϕ)


function getDivGrad(n1,n2,n3)
	# D = getDivGrad(n1,n2,n3)
	# builds 3D divergence operator
	D1 = kron(speye(n3),kron(speye(n2),ddx(n1)))
	D2 = kron(speye(n3),kron(ddx(n2),speye(n1)))
	D3 = kron(ddx(n3),kron(speye(n2),speye(n1)))

	Div = [D1 D2 D3]
	return Div*Div'
end

function ddx(n)
# generate 1D finite difference on staggered grid
	return d = spdiags(ones(n)*[-1 1],[0,1],n,n+1)
end

function spdiags(B,d,m,n)
# A = spdiags(B,d,m,n)
# creates a sparse matrix from its diagonals
	d = d[:]
	p = length(d)
	len = zeros(p+1,1)
	for k = 1:p
		len[k+1] = round(Int,len[k]+length(max(1,1-d[k]):min(m,n-d[k])))
	end
	a = zeros(round(Int,len[p+1]),3)
	for k = 1:p
		# Append new d[k]-th diagonal to compact form
		i = max(1,1-d[k]):min(m,n-d[k])
		a[(round(Int,len[k])+1):round(Int,len[k+1]),:] = [i i+d[k] B[i+(m>=n)*d[k],k]]
	end

	A = sparse(round(Int,a[:,1]),round(Int,a[:,2]),a[:,3],m,n)
	return A
end


# Newey-West covariance estimator
function NeweyWest(Z,nlags=0)
#=
    Returns the Newey-West estimator of the asymptotic variance matrix
    INPUTS: Z, a nxk matrix with rows the vector zt'
            nlags, the number of lags
    OUTPUTS: omegahat, the Newey-West estimator of the covariance matrix
=#
    n,k = size(Z)
    # de-mean the variables
    Z = Z .- mean(Z,1)
    omegahat = Z'*Z/n # sample variance
    # automatic lags?
    if nlags == 0
        nlags = max(0, Int(round(n^0.25)))
    end
    # sample autocovariances
    for i = 1:nlags
       Zlag = Z[1:n-i,:]
       ZZ = Z[i+1:n,:]
       gamma = (ZZ'*Zlag)/n
       weight = 1.0 - (i/(nlags+1.0))
       omegahat += weight*(gamma + gamma')
    end
    return omegahat
end


using ApproxFun, OrdinaryDiffEq
using LinearAlgebra
using DiffEqOperators

function makesim(sts::Union{STS{:KuramotoSivashinsky}, STS{:ksiva}})
    @unpack ic, T, S, Δt, N = sts
    L = try
            L = sts.p.L
        catch e
            L = 22
        end
    # prepare initial condition
    if ic isa Union{Nothing, Integer}
        # here we integrate a pre-determined initial condition
        # for a random amount of time
        rng = MersenneTwister(isnothing(ic) ? 42 : ic)
        S += Δt*rand(rng, 10:1000)
        x = L*(0:N-1)/N
        u0 = @. cos(x) + 0.1*sin(x/8) + 0.01*cos((2π/L)*x)
    elseif ic isa AbstractVector{<:Real}
        u0 = ic
        @assert length(u0) == N
    else
        error(ICERROR)
    end
    saveat = S:Δt:(T+S)
    t, u = kuramoto_sivashinsky(u0, L, saveat)
    t = 0:Δt:T
    return @strdict u t sts
end

"""
    kuramoto_sivashinsky(u, L, saveat)

Solve the Kuramoto Sivashinsky system given by the PDE
```math
u_t = -u*u_x - u_{xx} - u_{xxxx}
```
with periodic boundary conditions and initial condition `u`.

`L` is the physical spatial length (with `length(u)` points sampled
equidistantly in it).

`saveat` dictates at which timepoints to save the system (which is solved
using `OrdinaryDiffEq`).
"""
function kuramoto_sivashinsky(u, L, saveat)
    n = length(u)                  # number of gridpoints
    F = Fourier(0..L)
    T = ApproxFun.plan_transform(F, n)
    Ti = ApproxFun.plan_itransform(F, n)

    #Linear Part
    D  = (Derivative(F) → F)[1:n,1:n]
    D2 = Derivative(F,2)[1:n,1:n]
    D4 = Derivative(F,4)[1:n,1:n]
    A = DiffEqArrayOperator(Diagonal(-D2-D4))

    #Nonlinear Part
    ks = KSFunctor(zeros(n))

    params = (D, T, Ti)
    prob = SplitODEProblem(A, ks, T*u, (0.0, saveat[end]), params)
    sol = solve(prob, ETDRK4(), dt=Float64(saveat.step), saveat=saveat,
                save_start = false, save_end = false)
    # @assert sol.t == saveat
    return sol.t, map(u -> Ti*u, sol.u)
end

struct KSFunctor{T}; tmp::Vector{T}; end
function (ks::KSFunctor)(du,u,p,t)
    D, T, Ti = p
    mul!(du,Ti,u)
    @. du = -1/2*du^2
    mul!(ks.tmp, T, du)
    mul!(du, D, ks.tmp)
end


export SG_Filter, maxDerivativeOrder, polynomialOrder, apply_SG_filter, apply_SG_filter2D

import Base: filter, length

function _Vandermonde(T::DataType=Float64;halfWidth::Int=5,degree::Int=2)::Array{T,2}

    @assert halfWidth>=0
    @assert degree>=0

    x=T[i for i in -halfWidth:halfWidth]

    n = degree+1
    m = length(x)

    V = Array{T}(undef,m, n)

    for i = 1:m
        V[i,1] = T(1)
    end
    for j = 2:n
        for i = 1:m
            V[i,j] = x[i] * V[i,j-1]
        end
    end

    return V
end

# ================================================================

struct SG_Filter{T<:AbstractFloat,N}
    _filter_set::Array{LinearFilter_DefaultCentered{T,N},1}
end

"""
    function SG_Filter(T::DataType=Float64;halfWidth::Int=5,degree::Int=2)

Creates a `SG_Filter` structure used to store Savitzky-Golay filters.

* filter length is 2*`halfWidth`+1
* polynomial degree is `degree`, which defines `maxDerivativeOrder`

You can apply these filters using the
* `apply_SG_filter`
* `apply_SG_filter2D`
functions.

Example:

```jldoctest
julia> sg = SG_Filter(halfWidth=5,degree=3);

julia> maxDerivativeOrder(sg)
3

julia> length(sg)
11

julia> filter(sg,derivativeOrder=2)
Filter(r=-5:5,c=[0.03497, 0.01399, -0.002331, -0.01399, -0.02098, -0.02331, -0.02098, -0.01399, -0.002331, 0.01399, 0.03497])

```
"""
function SG_Filter(T::DataType=Float64;halfWidth::Int=5,degree::Int=2)::SG_Filter
    @assert degree>=0
    @assert halfWidth>=1
    @assert 2*halfWidth>degree

    V=_Vandermonde(T,halfWidth=halfWidth,degree=degree)
    Q,R=qr(V)
    # breaking change in Julia V1.0,
    # see https://github.com/JuliaLang/julia/issues/27397
    #
    # before Q was a "plain" matrix, now stored in compact form
    #
    # SG=R\Q'
    #
    # must be replaced by
    #
    # Q=Q*Matrix(I, size(V))
    # SG=R\Q'
    #
    Q=Q*Matrix(I, size(V))
    SG=R\Q'

    n_filter,n_coef = size(SG)

    buffer=Array{LinearFilter_DefaultCentered{T,n_coef},1}()

    for i in 1:n_filter
        SG[i,:]*=factorial(i-1)
        push!(buffer,LinearFilter(SG[i,:]))
    end

# Returns filters set
    return SG_Filter{T,n_coef}(buffer)
end

# ================================================================

"""
    function filter(sg::SG_Filter{T,N};derivativeOrder::Int=0)

Returns the filter to be used to compute the smoothed derivatives of order `derivativeOrder`.

See: `SG_Filter`
"""
function filter(sg::SG_Filter{T,N};derivativeOrder::Int=0) where {T<:AbstractFloat,N}
    @assert 0<= derivativeOrder <= maxDerivativeOrder(sg)
    return sg._filter_set[derivativeOrder+1]
end

"""
    function length(sg::SG_Filter{T,N})

Returns filter length, this is an odd number.

See: `SG_Filter`
"""
Base.length(sg::SG_Filter{T,N}) where {T<:AbstractFloat,N} = length(filter(sg))

"""
    function maxDerivativeOrder(sg::SG_Filter{T,N})

Maximum order of the smoothed derivatives we can compute using `sg` filters.

See: `SG_Filter`
"""
maxDerivativeOrder(sg::SG_Filter{T,N}) where {T<:AbstractFloat,N} = size(sg._filter_set,1)-1

"""
    function polynomialOrder(sg::SG_Filter{T,N})

Returns the degree of the polynomial used to construct the
Savitzky-Golay filters. This is mainly a 'convenience' function, as it
is equivalent to `maxDerivativeOrder`

See: `SG_Filter`
"""
polynomialOrder(sg::SG_Filter{T,N}) where {T<:AbstractFloat,N} = maxDerivativeOrder(sg)

"""
    function apply_SG_filter(signal::Array{T,1},
                             sg::SG_Filter{T};
                             derivativeOrder::Int=0,
                             left_BE::Type{<:BoundaryExtension}=ConstantBE,
                             right_BE::Type{<:BoundaryExtension}=ConstantBE)

Applies an 1D Savitzky-Golay and returns the smoothed signal.
"""
function apply_SG_filter(signal::AbstractArray{T,1},
                         sg::SG_Filter{T};
                         derivativeOrder::Int=0,
                         left_BE::Type{<:BoundaryExtension}=ConstantBE,
                         right_BE::Type{<:BoundaryExtension}=ConstantBE) where {T<:AbstractFloat}

    return directCrossCorrelation(filter(sg,derivativeOrder=derivativeOrder),
                                  signal,
                                  left_BE,
                                  right_BE)
end

"""
    function apply_SG_filter2D(signal::Array{T,2},
                               sg_I::SG_Filter{T},
                               sg_J::SG_Filter{T};
                               derivativeOrder_I::Int=0,
                               derivativeOrder_J::Int=0,
                               min_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                               max_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                               min_J_BE::Type{<:BoundaryExtension}=ConstantBE,
                               max_J_BE::Type{<:BoundaryExtension}=ConstantBE)

Applies an 2D Savitzky-Golay and returns the smoothed signal.
"""
function apply_SG_filter2D(signal::AbstractArray{T,2},
                           sg_I::SG_Filter{T},
                           sg_J::SG_Filter{T};
                           derivativeOrder_I::Int=0,
                           derivativeOrder_J::Int=0,
                           min_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                           max_I_BE::Type{<:BoundaryExtension}=ConstantBE,
                           min_J_BE::Type{<:BoundaryExtension}=ConstantBE,
                           max_J_BE::Type{<:BoundaryExtension}=ConstantBE) where {T<:AbstractFloat}

    return directCrossCorrelation2D(filter(sg_I,derivativeOrder=derivativeOrder_I),
                                    filter(sg_J,derivativeOrder=derivativeOrder_J),
                                    signal,
                                    min_I_BE,
                                    max_I_BE,
                                    min_J_BE,
                                    max_J_BE)
end


module problem118

description = """
Pandigital prime sets
Using all of the digits 1 through 9 and concatenating them freely to form decimal integers, different sets can be formed.
Interestingly with the set {2,5,47,89,631}, all of the elements belonging to it are prime.

How many distinct sets containing each of the digits one through nine exactly once contain only prime elements?
"""

# Convert an array of digits to an integer
function digitsToInt(digits :: Array)
  n = 0
  for d in digits
    n = n*10 + d
  end
  return n
end

# calculate the number of permutations of q that are prime
# Stores results in memo, for faster lookup for the same digits in future calls
function primePermutations(q :: Vector, memo :: Dict)
  total = 0
  sort!(q)
  key = digitsToInt(q)
  if haskey(memo, key)
    return memo[key]
  end

  for p in permutations(q)
    if isprime(digitsToInt(p))
      total += 1
    end
  end
  memo[key] = total
  return total
end

# Calculate the number of sets containing permutations of digits 1-9 that are all prime
# Julia's built in partitions, permutations and isprime do all the work :)
# For a given partition of [1:9] the number of sets is the product of prime permutations of each sub-set of digits.
function run()
  total = 0
  memo = (Int=>Int)[]
  for p in partitions([1:9])
    total += prod([primePermutations(q, memo) for q in p])
  end
  return total
end

end


function randindex(c1,c2)
    # rand_index - calculates Rand Indices to compare two partitions
    # (AR, RI, MI, HI) = rand(c1,c2), where c1,c2 are vectors listing the
    # class membership, returns the "Hubert & Arabie adjusted Rand index".
    # (AR, RI, MI, HI) = rand(c1,c2) returns the adjusted Rand index,
    # the unadjusted Rand index, "Mirkin's" index and "Hubert's" index.
    #
    # See L. Hubert and P. Arabie (1985) "Comparing Partitions" Journal of
    # Classification 2:193-218

    c = counts(c1,c2,(1:maximum(c1),1:maximum(c2))) # form contingency matrix

    n = round(Int,sum(c))
    nis = sum(sum(c,2).^2)        # sum of squares of sums of rows
    njs = sum(sum(c,1).^2)        # sum of squares of sums of columns

    t1 = binomial(n,2)            # total number of pairs of entities
    t2 = sum(c.^2)                # sum over rows & columnns of nij^2
    t3 = .5*(nis+njs)

    # Expected index (for adjustment)
    nc = (n*(n^2+1)-(n+1)*nis-(n+1)*njs+2*(nis*njs)/n)/(2*(n-1))

    A = t1+t2-t3;        # no. agreements
    D = -t2+t3;          # no. disagreements

    if t1 == nc
        # avoid division by zero; if k=1, define Rand = 0
        ARI = 0
    else
        # adjusted Rand - Hubert & Arabie 1985
        ARI = (A-nc)/(t1-nc)
    end

    RI = A/t1            # Rand 1971      # Probability of agreement
    MI = D/t1            # Mirkin 1970    # p(disagreement)
    HI = (A-D)/t1        # Hubert 1977    # p(agree)-p(disagree)

    return (ARI, RI, MI, HI)
end

randindex(R::ClusteringResult, c0::AbstractVector{Int}) = randindex(assignments(R), c0)
randindex(R1::ClusteringResult, R2::ClusteringResult) = randindex(assignments(R2), assignments(R1))


#=
The following code is modified from
https://github.com/billmclean/GaussQuadrature.jl with the original license:

> The MIT License (MIT)

> Copyright (c) 2013 billmclean

> Permission is hereby granted, free of charge, to any person obtaining a copy
> of this software and associated documentation files (the "Software"), to deal
> in the Software without restriction, including without limitation the rights
> to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
> copies of the Software, and to permit persons to whom the Software is
> furnished to do so, subject to the following conditions:

> The above copyright notice and this permission notice shall be included in
> all copies or substantial portions of the Software.

> THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
> IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
> FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
> AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
> LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
> OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
> SOFTWARE.
=#

# October 2013 by Bill McLean, School of Maths and Stats,
# The University of New South Wales.
#
# Based on earlier Fortran codes
#
# gaussq.f original version 20 Jan 1975 from Stanford
# gaussq.f modified 21 Dec by Eric Grosse
# gaussquad.f95 Nov 2005 by Bill Mclean
#
# This module provides functions to compute the abscissae x[j] and
# weights w[j] for the classical Gauss quadrature rules, including
# the Radau and Lobatto variants.  Thus, the sum
#
#           n
#           ∑  w[j] f(x[j])
#          j=1
#
# approximates
#
#           hi
#          ∫  f(x) w(x) dx
#          lo
#
# where the weight function w(x) and interval lo < x < hi are as shown
# in the table below.
#
# Name                      Interval     Weight Function
#
# Legendre                 -1 < x < 1          1
# Chebyshev (first kind)   -1 < x < 1     1 / sqrt(1-x²)
# Chebyshev (second kind)  -1 < x < 1       sqrt(1-x²)
# Jacobi                   -1 < x < 1     (1-x)ᵅ (1+x)ᵝ
# Laguerre                  0 < x < ∞     xᵅ exp(-x)
# Hermite                  -∞ < x < ∞      exp(-x²)
#
# In addition to these classical rules, the module generates Gauss rules
# for logarithmic weights of the form
#
#    w(x) = x^ρ log(1/x)   for 0 < x < 1.
#
# For the Jacobi and Laguerre rules we require α > -1 and
# β > -1, so that the weight function is integrable.  Likewise, for
# log weight we require ρ > -1.
#
# Use the endpoint argument to include one or both of the end points
# of the interval of integration as an abscissa in the quadrature
# rule, as follows.
#
# endpoint = NeitherEndPoint()   Default      lo < x[j] < hi, j = 1:n.
# endpoint = LeftEndPoint()      Left Radau   lo = x[1] < x[j] < hi, j = 2:n.
# endpoint = RightEndPoint()     Right Radau  lo < x[j] < x[n] = hi, j = 1:n-1.
# endpoint = BothEndPoint()      Lobatto      lo = x[1] < x[j] < x[n] = hi, j = 2:n-1.
#
#
# The code uses the Golub and Welsch algorithm, in which the abscissae
# x[j] are the eigenvalues of a symmetric tridiagonal matrix whose
# entries depend on the coefficients in the 3-term recurrence relation
# for the othonormal polynomials generated by the weighted inner product.
#
# References:
#
#   1.  Golub, G. H., and Welsch, J. H., Calculation of Gaussian
#       quadrature rules, Mathematics of Computation 23 (April,
#       1969), pp. 221-230.
#   2.  Golub, G. H., Some modified matrix eigenvalue problems,
#       Siam Review 15 (april, 1973), pp. 318-334 (section 7).
#   3.  Stroud and Secrest, Gaussian Quadrature Formulas, Prentice-
#       Hall, Englewood Cliffs, N.J., 1966.

# Enumeration type used to specify which endpoints of the integration
# interval should be included amongst the quadrature points: neither,
# left, right, or both.

abstract type EndPoint end
struct NeitherEndPoint <: EndPoint end
struct LeftEndPoint <: EndPoint end
struct RightEndPoint <: EndPoint end
struct BothEndPoint <: EndPoint end

"""
    x, w = legendregauss(T, n, endpoint::EndPoint=Bennu.NeitherEndPoint())

Returns points `x` and weights `w` for the `n`-point Gauss-Legendre rule
for the interval `-1 < x < 1` with weight function `w(x) = 1`.

Use `endpoint=LeftEndPoint()`, `RightEndPoint() ` or `BothEndPoints()` for the
left Radau, right Radau, or Lobatto rules, respectively.
"""
function legendregauss(::Type{T}, n, endpoint=NeitherEndPoint()) where {T}
    @assert n ≥ 1
    a, b = legendrecoefficients(T, n)
    return gaussrule(-one(T), one(T), a, b, endpoint)
end

"""
    x, w = legendregauss(n, endpoint::EndPoint=Bennu.NeitherEndPoint())

Convenience function with type `T = Float64`:
"""
legendregauss(n, endpoint::EndPoint=NeitherEndPoint()) = legendregauss(Float64, n, endpoint)

"""
    x, w = legendregausslobatto(T, n)

Returns points `x` and weights `w` for the `n`-point Legendre-Gauss-Lobatto rule
for the interval `-1 ≤ x ≤ 1` with weight function `w(x) = 1`.
"""
function legendregausslobatto(::Type{T}, n) where {T}
    return legendregauss(T, n, BothEndPoint())
end

"""
    x, w = legendregausslobatto(n)

Convenience function with type `T = Float64`:
"""
legendregausslobatto(n) = legendregausslobatto(Float64, n)

function legendrecoefficients(::Type{T}, n) where {T}
    a = zeros(T, n)
    b = zeros(T, n + 1)
    b[1] = sqrt(convert(T, 2))
    for k in 2:(n + 1)
        b[k] = (k - 1) / sqrt(convert(T, (2k - 1) * (2k - 3)))
    end
    return a, b
end

"""
    x, w = gaussrule(lo, hi, a, b, endpoint, maxiterations=100)

Generates the points `x` and weights `w` for a Gauss rule with weight
function `w(x)` on the interval `lo < x < hi`.

The arrays `a` and `b` hold the coefficients (as given, for instance, by
`legendrecoefficients`) in the three-term recurrence relation for the monic
orthogonal polynomials `p(0,x)`, `p(1,x)`, `p(2,x)`, ... , that is,

    p(k, x) = (x-a[k]) p(k-1, x) - b[k]² p(k-2, x),    k ≥ 1,

where `p(0, x) = 1` and, by convention, `p(-1, x) = 0` with

              hi
    b[1]^2 = ∫  w(x) dx.
             lo

Thus, `p(k, x) = xᵏ + lower degree terms` and

     hi
    ∫  p(j, x) p(k, x) w(x) dx = 0 if j ≠ k.
    lo
"""
function gaussrule(lo, hi, a, b, endpoint::EndPoint, maxiterations=100)
    T = promote_type(typeof(lo), typeof(hi), eltype(a), eltype(b))
    n = length(a)
    @assert length(b) == n + 1
    if endpoint isa LeftEndPoint
        if n == 1
            a[1] = lo
        else
            a[n] = tridiagonalshiftsolve(n, lo, a, b) * b[n]^2 + lo
        end
    elseif endpoint isa RightEndPoint
        if n == 1
            a[1] = hi
        else
            a[n] = tridiagonalshiftsolve(n, hi, a, b) * b[n]^2 + hi
        end
    elseif endpoint isa BothEndPoint
        if n == 1
            error("Must have at least two points for both ends.")
        end
        g = tridiagonalshiftsolve(n, lo, a, b)
        t1 = (hi - lo) / (g - tridiagonalshiftsolve(n, hi, a, b))
        b[n] = sqrt(t1)
        a[n] = lo + g * t1
    end
    w = zero(a)
    tridiagonaleigenproblem!(a, b, w, maxiterations)
    for i in 1:n
        w[i] = (b[1] * w[i])^2
    end
    idx = sortperm(a)
    # Ensure end point values are exact.
    if endpoint isa LeftEndPoint || endpoint isa BothEndPoint
        a[idx[1]] = lo
    elseif endpoint isa RightEndPoint || endpoint isa BothEndPoint
        a[idx[n]] = hi
    end
    return a[idx], w[idx]
end

function tridiagonalshiftsolve(n, shift, a, b) where {T}
    #
    # Perform elimination to find the nth component s = delta[n]
    # of the solution to the nxn linear system
    #
    #     ( J_n - shift I_n ) delta = e_n,
    #
    # where J_n is the symmetric tridiagonal matrix with diagonal
    # entries a[i] and off-diagonal entries b[i], and e_n is the nth
    # standard basis vector.
    #
    t = a[1] - shift
    for i in 2:(n - 1)
        t = a[i] - shift - b[i]^2 / t
    end
    return one(t) / t
end

function tridiagonaleigenproblem!(d, e, z, maxiterations)
    #
    # Finds the eigenvalues and first components of the normalised
    # eigenvectors of a symmetric tridiagonal matrix by the implicit
    # QL method.
    #
    # d[i]   On entry, holds the ith diagonal entry of the matrix.
    #        On exit, holds the ith eigenvalue.
    #
    # e[i]   On entry, holds the [i,i-1] entry of the matrix for
    #        i = 2, 3, ..., n.  (The value of e[1] is not used.)
    #        On exit, e is overwritten.
    #
    # z[i]   On exit, holds the first component of the ith normalised
    #        eigenvector associated with d[i].
    #
    # maxiterations The maximum number of QL iterations.
    #
    # Martin and Wilkinson, Numer. Math. 12: 377-383 (1968).
    # Dubrulle, Numer. Math. 15: 450 (1970).
    # Handbook for Automatic Computation, Vol ii, Linear Algebra,
    #        pp. 241-248, 1971.
    #
    # This is a modified version of the Eispack routine imtql2.
    #
    T = promote_type(eltype(d), eltype(e), eltype(z))
    n = length(z)
    z[1] = one(T)
    z[2:n] .= zero(T)
    e[n + 1] = zero(T)

    if n == 1 # Nothing to do for a 1x1 matrix.
        return
    end
    for l in 1:n
        for j in 1:maxiterations
            # Look for small off-diagonal elements.
            m = n
            for i in l:(n - 1)
                if abs(e[i + 1]) <= eps(T) * (abs(d[i]) + abs(d[i + 1]))
                    m = i
                    break
                end
            end
            p = d[l]
            if m == l
                continue
            end
            if j == maxiterations
                msg = string("No convergence after ", j, " iterations",
                             " (try increasing maxiterations)")
                error(msg)
            end
            # Form shift
            g = (d[l + 1] - p) / (2 * e[l + 1])
            r = hypot(g, one(T))
            g = d[m] - p + e[l + 1] / (g + copysign(r, g))
            s = one(T)
            c = one(T)
            p = zero(T)
            for i in (m - 1):-1:l
                f = s * e[i + 1]
                b = c * e[i + 1]
                if abs(f) < abs(g)
                    s = f / g
                    r = hypot(s, one(T))
                    e[i + 2] = g * r
                    c = one(T) / r
                    s *= c
                else
                    c = g / f
                    r = hypot(c, one(T))
                    e[i + 2] = f * r
                    s = one(T) / r
                    c *= s
                end
                g = d[i + 1] - p
                r = (d[i] - g) * s + 2 * c * b
                p = s * r
                d[i + 1] = g + p
                g = c * r - b
                # Form first component of vector.
                f = z[i + 1]
                z[i + 1] = s * z[i] + c * f
                z[i] = c * z[i] - s * f
            end
            d[l] -= p
            e[l + 1] = g
            e[m + 1] = zero(T)
        end
    end
    return
end


using ProgressBars, Plots, StatsBase, Statistics, LaTeXStrings, JLD

function RandomWalker2D(TotalSteps)
    choices = ((1,0),(0,1),(-1,0),(0,-1))
    return reduce(.+,rand(choices,TotalSteps))
end

function RGyration(ensemble)
    MassCenter = (mean(getindex.(ensemble,1)),mean(getindex.(ensemble,2)))
    return √(sum(pos -> reduce(.+, (pos .- MassCenter).^2), ensemble)/length(ensemble))
end

TotalSteps = 40
RWNumbers = 10000000

Data = [[RandomWalker2D(TS) for i ∈ 1:RWNumbers] for TS ∈ ProgressBar(1:TotalSteps)]

RgList = [RGyration(Data[i]) for i ∈ 1:TotalSteps]
save("../../Data/Q1/Q1-Rg.jld", "RgList", RgList)

plot(1:TotalSteps,sqrt.(1:TotalSteps), label = L"Y = \sqrt{X}", c = :black, linestyle = :dash)
scatter!(RgList, label = L"Data\ Point", framestyle = :box, c = :steelblue, legend = 140)
plot!(xlabel = L"Time", ylabel = L"{R_g}_{(t)}",
    title = L"Radius\ of\ Gyration\ of\ RW\ per\ Time\ for\ 10^6 \ Particles")
savefig("../../Figs/Q1/Q1-Rg(t).pdf")


#=
Popis
  Funkcia vypocita pre kazdy vrchol hodnotu jeho closeness centrality.

Parametre
  g : datova struktura grafu :: GenericGraph
  dist_key : nazov atributu, v ktorom je ulozena vlastnost vrchola :: String
  normalize : parameter pre aplikovanie normalizacie :: Bool

Navratova hodnota
  closeness : slovnik obsahujuci mapovanie vrcholov a k nim prisluchajucich hodnot
           closeness centrality (index_vrchola => hodnota_closeness_centrality) :: Dict{Int64, Float64}()

Poznamky
  Hodnota centrality moze byt normalizovana nastavenim parametra normalize na True/False.
  Vypocet centrality je upraveny aby velkost grafu neovplyvnovala hodnotu centrality.
=#

function centrality_closeness(g::GenericGraph; dist_key::AbstractString="weight", normalize::Bool=true)

  epi = AttributeEdgePropertyInspector{Float64}(dist_key)
  closeness = Dict{Int64, Float64}()

  for v in Graphs.vertices(g)

      distances = filter(x -> x != Inf, dijkstra_shortest_paths(g, epi, [v]).dists)
      total_dist = sum(distances)

      if total_dist > 0 && num_vertices(g) > 1

        closeness[v.index] = (length(distances) - 1.0) / total_dist

        if normalize

          norm_coefficient = (length(distances)-1.0) / ( num_vertices(g) - 1 )
          closeness[v.index] *= norm_coefficient

        end

      end

  end

  return closeness

end


@doc raw"""
    SquaredEuclidean

The squared Euclidean function is defined by:

```math
f(\mathbf{x}, \mathbf{y}) = (\mathbf{x} - \mathbf{y})^{\intercal}(\mathbf{x} - \mathbf{y})
```
"""
struct SquaredEuclidean <: Metric end

@inline base_aggregate(::SquaredEuclidean, s::T, x::T, y::T) where {T} = s + (x-y)^2

@inline isstationary(::SquaredEuclidean) = true
@inline isisotropic(::SquaredEuclidean)  = true


# This is quasi derived from the Rosenbrock function
# whereby F(x,p) is almost the gradient of the Rosenbrock function
# except the Hessian matroix is made non-symmetric by slightly altering
# F(x,p)

# State function ≈ gradient of Rosenbrock.
# With p = [a, b] the parameters defining the constants
# in Rosenbrocks function, and the slight alteration to make the
# Hessian matrix non-symmetric.
# Specifically, if r(x,y) denotes the rosenbrock function,
#     F₁([x, y], p) = ∂r/∂x
#     F₂([x, y], p) = 0.5 * ∂r/∂y
F(x,p) = [
    -2 * (p[1] - x[1]) - 4 * p[2] * (x[2] - x[1]^2) * x[1]
    p[2] * (x[2] - x[1]^2)
]

# Jacobian function of F w.r.t. p
∇ₓF(x,p) = ForwardDiff.jacobian(x -> F(x,p), x)

# Define mismatch function f(x,p) and its derivative ∇ₓf(x,p)
# (Note ∇ₓF and ∇ₓf are required by the F1 method)
function state_mismatch(x)
    δ(x) = x .- 1
    return 0.5δ(x)'δ(x)
end
function parameter_mismatch(p)
    δ(p) = log.(p)
    return 0.5δ(p)'δ(p)
end
f(x,p) = state_mismatch(x) + parameter_mismatch(p)
∇ₓf(x,p) = ForwardDiff.jacobian(x -> [f(x,p)], x)

# Initialize the starting state and parameters
x₀ = rand(2)
p₀ = rand(2)
# Initialize the memory cache for storing reusable objects
mem = F1Method.initialize_mem(x₀, p₀)

# Define the functions via the F1 method
F1_objective(p) = F1Method.objective(f, F, ∇ₓF, mem, p, MyAlg())
F1_gradient(p) = F1Method.gradient(f, F, ∇ₓf, ∇ₓF, mem, p, MyAlg())
F1_Hessian(p) = F1Method.hessian(f, F, ∇ₓf, ∇ₓF, mem, p, MyAlg())

# Define the exact solution and objective for comparison
exact_solution(p) = [p[1], p[1]^2]
exact_objective(p) = f(exact_solution(p), p)

# Run the tests
@testset "objective" begin
    @test exact_objective(p₀) ≈ F1_objective(p₀)
end
@testset "gradient" begin
    @test ForwardDiff.jacobian(p -> [exact_objective(p)], p₀) ≈ F1_gradient(p₀)
end
@testset "Hessian" begin
    @test ForwardDiff.hessian(p -> exact_objective(p), p₀) ≈ F1_Hessian(p₀)
end


"""
    The Tower of Hanoi

# Brief:
    The Tower of Hanoi is a mathematical game or puzzle consisting of three rods and a number
    of disks of various diameters, which can slide onto any rod.
    The puzzle begins with the disks stacked on one rod in order of decreasing size,
    the smallest at the top, thus approximating a conical shape.
    The objective of the puzzle is to move the entire stack to the last rod,
    obeying the following rules:
        1. Only one disk may be moved at a time.
        2. Each move consists of taking the upper disk from one of the stacks and placing it on
           top of another stack or on an empty rod.
        3. No disk may be placed on top of a disk that is smaller than it.

# Complexity: O(2^n)

# Functions
    - solveUtil(tower1, tower2, tower3, n::Int, solution::Array{Pair}) - The recursive function which finds the solution and saves it to
                                                                         the array of pairs where each pair represents one move
    - solve(tower1, tower2, tower3,n::Int) - The main function for finding the set of steps to move n rings from tower 1 to tower 3, using tower 2
    - printSolution(solution::Array{Pair}) - Prints the solution given as an array of pairs by printing each pair ( step ) in a different line

# Reference:- [Wikipedia](https://en.wikipedia.org/wiki/Tower_of_Hanoi)

# Contributed by:- [Nikola Mircic](https://github.com/Nikola-Mircic)
"""

module Hanoi
    function solveUtil(tower1, tower2, tower3, n::Int, solution::Array{Pair})
        if n==1
            push!(solution, tower1 => tower3) # There is only one ring, so just move it from tower1 to tower3
        else
            solveUtil(tower1, tower3, tower2, n-1, solution) # Move n-1 rings from tower1 to tower2 to free the bottom ring
            push!(solution, tower1 => tower3) # Move the bottom ring from tower1 to tower3
            solveUtil(tower2, tower1, tower3, n-1, solution) # Now move those n-1 rings from tower2 to tower3
        end
    end

    function solve(tower1, tower2, tower3,n::Int)
        solution = Array{Pair}(undef, 0)
        solveUtil(tower1, tower2, tower3, n, solution)
        return solution
    end

    function printSolution(solution::Array{Pair})
        for step in solution
            println(step)
        end
    end
end


"""
    permutation(n, r)

    returns the number of ways to choose r items from n items
    with order and without repetition

# Arguments:
- `n`: Positive integers of items to choose from
- 'r': Positive integers of items to choose

Contributed By:- [Mayur Dahibhate](https://github.com/mayurdahibhate)
"""

function permutation(n, r)
    function factorial(n)
        fact = 1

        if n == 0 || n == 1
            return fact
        end

        for i in 1:n
            fact = fact * i
        end

        return fact
    end

    perm = factorial(n) / factorial(n - r)

    return convert(Int64, perm)
end


using QuadGK: quadgk

export fourier_coefficient, FourierSeries, eval_series, rotating_frame_series

"""
    fourier_coefficient(f::Function, frequency::Real, harmonic::Int)

Calculate the fourier coefficient `1/T ∫_0^T f(τ) e^(-2πi * frequency * harmonic * τ) dτ`
of a periodic function `f`.

# args
* `f`: the function of one variable.
* `frequency`: the frequency of periodicity of `f`.
* `harmonic`: the desired harmonic.

# returns
The fourier coefficient.
"""
function fourier_coefficient(f::Function, frequency::Real, harmonic::Int)
    T = 1/frequency
    g(t) = f(t) * exp(-2π * 1im * frequency * harmonic * t)
    integral, _ = quadgk(g, 0.0, T, atol=sqrt(eps(T)))
    return integral/T
end

"""
Representation of a Fourier series `f(t) = ∑ Aₙ e^(iωnt)`
"""
struct FourierSeries
    frequency::Real
    terms::Dict{Int, <:Number} # harmonic => coefficient
end

"""
    FourierSeries(f::Function, frequency::Real, harmonics::AbstractVector{Int})

Construct a FourierSeries from a periodic function `f` on given harmonics.

# args
* `f`: the periodic function.
* `frequency`: the frequency of periodicity of f.
* `harmonics`: a vector of harmonics to compute.

# returns
A FourierSeries.
"""
function FourierSeries(f::Function, frequency::Real, harmonics::AbstractVector{Int})
    terms = Dict(h => fourier_coefficient(f, frequency, h) for h in harmonics)
    return FourierSeries(frequency, terms)
end

"""
    eval_series(fs::FourierSeries, times::AbstractVector{<:Real})

Evaluate a FourierSeries on given times.

# args
* `fs`: a FourierSeries.
* `times`: a vector of times.

# returns
A vector of values.
"""
function eval_series(fs::FourierSeries, times::AbstractVector{<:Real})
    return reduce(+, coeff * exp.(2π * 1im * fs.frequency * harmonic * times) for (harmonic, coeff) in fs.terms)
end

"""
    rotating_frame_series(fs::FourierSeries, harmonics::AbstractVector{Int})

Compute the fourier series for `exp(i∫_0^t f(τ) dτ)` where `f` is the function given by
the FourierSeries `fs`. This function ignores a constant term in `fs` since that maps onto
an overall exponential factor which is not part of a FourierSeries.

# args
* `fs`: a FourierSeries.
* `harmonics`: the desired harmonics.

# returns
The FourierSeries.
"""
function rotating_frame_series(fs::FourierSeries, harmonics::AbstractVector{Int})
    iω = 2π * 1im * fs.frequency
    function exp_of_integral(t::Real)
        int = 0.0
        for (h, c) in fs.terms
            if h != 0
                int += c * (exp(iω * h * t) - 1)/(iω * h)
            end
        end
        return exp(1im * int)
    end
    return FourierSeries(exp_of_integral, fs.frequency, harmonics)
end


"""
Finds roots of 4th order polynomials of the form
```math
ax⁴ + bx³ + cx² + dx + e = 0
```
based on the Fast Quartic and Cubic Solver
https://github.com/NKrvavica/fqs
by Nino Krvavica.
"""
module QuarticSolver

function solvequadratic(a0, b0, c0)
    inva0 = 1/a0
    a, b = b0*inva0, c0*inva0

    a₀ = -0.5*a
    Δ = a₀^2 - b
    sqrtΔ = sqrt(Δ)

    r1 = a₀ - sqrtΔ
    r2 = a₀ + sqrtΔ

    return r1, r2
end

function solvecubic(a0, b0, c0, d0)
    inva0 = 1/a0
    a, b, c = b0*inva0, c0*inva0, d0*inva0

    third = 1/3
    thirda = a*third
    thirda² = thirda^2
    sqrt3 = sqrt(3)

    f = third*b - thirda²
    g = thirda*(2*thirda² - b) + c
    h = 0.25*g^2 + f^3  # discriminant (`Δ` or `d` in some papers)

    if f == g == h == 0  # this check seems a little strong
        r1 = -cbrt(c)
        return r1, r1, r1
    elseif h isa Complex || h <= 0   # casus irreducibilis
        j = sqrt(-f)
        k = acos(-0.5*g/j^3)
        m = cos(third*k)
        n = sqrt3*sin(third*k)
        r1 = 2*j*m - thirda
        r2 = -j*(m + n) - thirda
        r3 = -j*(m - n) - thirda
        return r1, r2, r3
    else
        sqrth = sqrt(h)
        S = cbrt(-0.5*g + sqrth)
        U = cbrt(-0.5*g - sqrth)
        SplusU = S + U
        SminusU = S - U
        tmp = SminusU*sqrt3*0.5im
        r1 = SplusU - thirda
        r2 = -0.5*SplusU - thirda + tmp
        r3 = -0.5*SplusU - thirda - tmp
        return r1, r2, r3
    end
end

function solvequartic(a0, b0, c0, d0, e0)
    inva0 = 1/a0
    a, b, c, d = b0*inva0, c0*inva0, d0*inva0, e0*inva0

    a₀ = 0.25*a
    a₀² = a₀^2

    # Subsidiary cubic equation
    p = 3*a₀² - 0.5*b
    q = a*a₀² - b*a₀ + 0.5*c
    r = 3*a₀²^2 - b*a₀² + c*a₀ - d

    # One root of cubic
    z0, _, _ = solvecubic(1, p, r, p*r-0.5*q^2)
    z0 = complex(z0)

    s = sqrt(2*p + 2*z0)
    s == 0 ? t = z0^2 + r : t = -q/s

    r0, r1 = solvequadratic(1, s, z0 + t)
    r2, r3 = solvequadratic(1, -s, z0 - t)

    return r0-a₀, r1-a₀, r2-a₀, r3-a₀
end

end # module


using DifferentialEquations
using SimpleDiffEq
using Tables
using DataFrames
using StatsPlots
using BenchmarkTools

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    @inbounds begin
        du[1] = -β*c*I/N*S
        du[2] = β*c*I/N*S - γ*I
        du[3] = γ*I
    end
    nothing
end;

δt = 0.1
tmax = 40.0
tspan = (0.0,tmax)

u0 = [990.0,10.0,0.0]; # S,I,R

p = [0.05,10.0,0.25]; # β,c,γ

prob_ode = ODEProblem(sir_ode!, u0, tspan, p);

sol_ode = solve(prob_ode, dt = δt);

df_ode = DataFrame(Tables.table(sol_ode'))
rename!(df_ode,["S","I","R"])
df_ode[!,:t] = sol_ode.t;

@df df_ode plot(:t,
    [:S :I :R],
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number")

@benchmark solve(prob_ode, dt = δt);


using LinearAlgebra, Distances
export Diffusion
export hamming_dist, jaccard_dist, diffusion_dist

"""
Compute the Hamming distance between two adjacency matrices (directed graphs)
"""
function hamming_dist(A1::Matrix{Int}, A2::Matrix{Int})
    return sum(abs(x-y) for (x,y) in zip(A1,A2))
end

(d::Distances.Hamming)(A1::Matrix{Int}, A2::Matrix{Int}) = hamming_dist(A1,A2)

(d::Distances.Hamming)(A1::Matrix{Int}, A2::Nothing) = sum(A1)
(d::Distances.Hamming)(A1::Nothing, A2::Matrix{Int}) = d(A2,A1)

"""
Compute the Jaccard distance between two adjacency matrices (directed graphs)
"""
function jaccard_dist(A1::Matrix{Int}, A2::Matrix{Int})
    return 1 - sum(min(x,y) for (x,y) in zip(A1,A2)) / sum(max(x,y) for (x,y) in zip(A1,A2))
end

(d::Distances.Jaccard)(A1::Matrix{Int}, A2::Matrix{Int}) = jaccard_dist(A1,A2)

struct Diffusion <: Metric end

"""
Compute the diffusion distance between two adjacency matrices with t the time
of diffusion.
"""
function diffusion_dist(A1::Matrix{Int}, A2::Matrix{Int}, t)
    L1 = Diagonal(vec(sum(A1, dims = 2))) - A1
    L2 = Diagonal(vec(sum(A2, dims = 2))) - A2
    function matrixExp(A, t)
        F = eigen(A)
        return F.vectors * Diagonal(exp.(F.values * t)) * transpose(F.vectors)
    end
    return sum((matrixExp(-L1, t) - matrixExp(-L2, t)) .^ 2)
end

(d::Diffusion)(A1::Matrix{Int}, A2::Matrix{Int}) = diffusion_dist(A1,A2)


module Activations

export Activation, forward, backward, parms, ∇parms, Sigmoid, Tanh, ReLU

include("./tensor_dt.jl")
using ..TensorDT: Tensor

include("./abstract_layers.jl")
import ..AbstractLayers: AbstractLayer, AL, forward, backward, parms, ∇parms

## ======================================================================
## Activation Layer
## ======================================================================
struct Activation <: AbstractLayer
  name::Symbol
  fn::Function
  der_fn::Function
  #
  _type::Symbol
  store::Dict

  function Activation(name, fn, der_fn)
    new(name, fn, der_fn, :activation, Dict())
  end
end

function forward(self::Activation, input::Tensor)::Tensor
  self.store[self.name] = self.fn.(input)
  self.store[self.name]
end

function backward(self::Activation, ∇p::Tensor)::Tensor
  self.der_fn(self.store[self.name]) .* ∇p
end

parms(self::Activation) = []
∇parms(self::Activation) = []

### Sigmoid (Activation) function
function Sigmoid()
  σ = z -> 1. ./ (1. .+ exp.(-z))
  der_σ = z -> σ(z) .* (1 .- σ(z))
  Activation(:sigmoid, σ, der_σ)
end

### Tanh (Activation) function
function Tanh()
  tanₕ = z -> (x = exp.(z); y = exp.(-z); (x .- y) ./ (x .+ y))
  der_tanₕ = z -> 1 .- tanₕ.(z) .^ 2
  Activation(:tanh, tanₕ, der_tanₕ)
end

### ReLU (Activation) function
function ReLU()
  relu_fn = z -> max(zero(eltype(z)), z)
  der_relu = z -> (z₀ =	zero(eltype(z)); z .< z₀ ? z₀ : one(eltype(z)))
  Activation(:relu, relu_fn, der_relu)
end

end  ## Module

Data Science From Scratch/src/activations.jl


@testset "Gabor wavelet" begin
# there are two methods and one struct for calculating Gabor wavelets
# calculate Gabor wavelet given ω
N=4
ω=zeros(2^N)
σ=1.0
ξ=0.0
γ = Scat.gabor1d(ω, σ, ξ)
@test length(γ)==length(ω)
@test γ==ones(length(ω))

# calculate Gabor wavelet without ω
N=4
σ=1.0/sqrt(2.0)
ξ=0.0
ω=FFTW.fftfreq(N)
γ = Scat.gabor1d(N, σ, ξ)
@test length(γ)==length(ω)
@test exp.(-ω.^2)≈γ

# test inner constructor for GaborWavelet struct
N=4
σ=1.0
ξ=0.0
j=1
γ = GaborWavelet(N, σ, ξ, j)
@test length(γ.γ)==N
@test γ.γ == Scat.gabor1d(N, σ, ξ)
end

@testset "Gaussian filter" begin
# there are two methods and one struct for calculating Gaussian lowpass filters
# calculate Gaussian lowpass filter given ω
N=4
ω=zeros(2^N)
σ=1.0
G = Scat.gauss1d(ω, σ)
@test length(G)==length(ω)
@test G==ones(length(ω))

# calculate Gaussian lowpass filter without ω
N=4
σ=1.0/sqrt(2.0)
ω=collect(0.0:N-1)./N
G = Scat.gauss1d(N, σ)
@test length(G)==length(ω)
@test exp.(-ω.^2)≈G

# test inner constructor for GaussianFilter struct
N=4
σ=1.0
ξ=0.0
j=1
G = GaussianFilter(N, σ)
@test length(G.ϕ)==N
@test G.ϕ == Scat.gauss1d(N, σ)
end

@testset "Morlet wavelet" begin
# there are two methods and one struct for calculating Morlet wavelets
# calculate Morlet wavelet given Gabor wavelet and Gaussian lowpass filter
N=4
σ=1.0/sqrt(2.0)
ξ=0.0
G=ones(N)
γ=ones(N)
ψ=Scat.morlet1d(γ, G)
@test length(ψ)==N
@test ψ==zeros(N)

# calculate Morlet wavelet without ω
N=4
σ=1.0/sqrt(2.0)
ξ=0.0
ψ=Scat.morlet1d(N, σ, ξ)
@test length(G)==N

# test both inner constructors for MorletWavelet struct
N=4
σ=1.0
ξ=0.0
j=1
# get GaborWavelet and GaussianFilter structs to build MorletWavelet
G = GaussianFilter(N, σ)
γ = GaborWavelet(N, σ, ξ, j)
ψ = MorletWavelet(γ, G)
@test length(ψ.ψ)==N
@test ψ.ψ == Scat.morlet1d(γ.γ, G.ϕ)
end