using StatisticalRethinking, CmdStan
#gr(size=(600,600));

ProjDir = rel_path("..", "scripts", "04")
cd(ProjDir)

howell1 = CSV.read(rel_path("..", "data", "Howell1.csv"), delim=';')
df = convert(DataFrame, howell1);

df2 = filter(row -> row[:age] >= 18, df);

first(df2, 5)

struct Post
  mu::Float64
  sigma::Float64
  ll:: Float64
  prod::Float64
  prob::Float64
end

mu_list = repeat(range(140, 160, length=200), 200);
sigma_list = repeat(range(4, 9, length=200), inner=200);

ll = zeros(40000);
for i in 1:40000
    d1 = Normal(mu_list[i], sigma_list[i])
    ll[i] = sum(log.(pdf.(d1, df2[:height])))
end

d2 = Normal(178.0, 20.0)
d3 = Uniform(0, 50)
prod = ll + log.(pdf.(d2, mu_list)) + log.(pdf.(d3, sigma_list))
prob = exp.(prod .- maximum(prod))
post = DataFrame(mu=mu_list, sigma=sigma_list, ll=ll, prod=prod, prob=prob)
first(post, 10)

samples = post[sample(1:size(post, 1), Weights(post[:prob]), 10000, replace=true), :]

density(samples[:mu])

density(samples[:sigma])

#hpd(samples[:mu])

#hpd(samples[:sigma])

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


# Linear Discriminant Analysis

#### Type to represent a linear discriminant functional

abstract type Discriminant{T} end

struct LinearDiscriminant{T<:Real} <: Discriminant{T}
    w::Vector{T}
    b::T
end

length(f::LinearDiscriminant) = length(f.w)

evaluate(f::LinearDiscriminant, x::AbstractVector) = dot(f.w, x) + f.b

function evaluate(f::LinearDiscriminant, X::AbstractMatrix)
    R = transpose(X) * f.w
    if f.b != 0
        broadcast!(+, R, R, f.b)
    end
    return R
end

predict(f::Discriminant, x::AbstractVector) = evaluate(f, x) > 0

predict(f::Discriminant, X::AbstractMatrix) = (Y = evaluate(f, X); Bool[y > 0 for y in Y])

#### function to solve linear discriminant

function ldacov(C::DenseMatrix{T},
                μp::DenseVector{T},
                μn::DenseVector{T}) where T<:Real

    w = cholesky(C) \ (μp - μn)
    ap = w ⋅ μp
    an = w ⋅ μn
    c = 2 / (ap - an)
    LinearDiscriminant(rmul!(w, c), 1 - c * ap)
end

ldacov(Cp::DenseMatrix{T},
       Cn::DenseMatrix{T},
       μp::DenseVector{T},
       μn::DenseVector{T}) where T<:Real = ldacov(Cp + Cn, μp, μn)

#### interface functions

function fit(::Type{LinearDiscriminant}, Xp::DenseMatrix{T}, Xn::DenseMatrix{T};
             covestimator::CovarianceEstimator = SimpleCovariance()) where T<:Real
    μp = vec(mean(Xp, dims=2))
    μn = vec(mean(Xn, dims=2))
    Zp = Xp .- μp
    Zn = Xn .- μn
    Cp = calcscattermat(covestimator, Zp)
    Cn = calcscattermat(covestimator, Zn)
    ldacov(Cp, Cn, μp, μn)
end

#### Multiclass LDA Stats

mutable struct MulticlassLDAStats{T<:Real, M<:AbstractMatrix{T}, N<:AbstractMatrix{T}}
    dim::Int              # sample dimensions
    nclasses::Int         # number of classes
    cweights::Vector{T}   # class weights
    tweight::T            # total sample weight
    mean::Vector{T}       # overall sample mean
    cmeans::Matrix{T}     # class-specific means
    Sw::M                 # within-class scatter matrix
    Sb::N                 # between-class scatter matrix
end

mean(S::MulticlassLDAStats) = S.mean
classweights(S::MulticlassLDAStats) = S.cweights
classmeans(S::MulticlassLDAStats) = S.cmeans

withclass_scatter(S::MulticlassLDAStats) = S.Sw
betweenclass_scatter(S::MulticlassLDAStats) = S.Sb

function MulticlassLDAStats(cweights::Vector{T},
                            mean::Vector{T},
                            cmeans::Matrix{T},
                            Sw::AbstractMatrix{T},
                            Sb::AbstractMatrix{T}) where T<:Real
    d, nc = size(cmeans)
    length(mean) == d || throw(DimensionMismatch("Incorrect length of mean"))
    length(cweights) == nc || throw(DimensionMismatch("Incorrect length of cweights"))
    tw = sum(cweights)
    size(Sw) == (d, d) || throw(DimensionMismatch("Incorrect size of Sw"))
    size(Sb) == (d, d) || throw(DimensionMismatch("Incorrect size of Sb"))
    MulticlassLDAStats(d, nc, cweights, tw, mean, cmeans, Sw, Sb)
end

function multiclass_lda_stats(nc::Int, X::AbstractMatrix{T}, y::AbstractVector{Int};
                              covestimator_within::CovarianceEstimator=SimpleCovariance(),
                              covestimator_between::CovarianceEstimator=SimpleCovariance()) where T<:Real
    # check sizes
    d = size(X, 1)
    n = size(X, 2)
    n ≥ nc || throw(ArgumentError("The number of samples is less than the number of classes"))
    length(y) == n || throw(DimensionMismatch("Inconsistent array sizes."))

    # compute class-specific weights and means
    cmeans, cweights, Z = center(X, y, nc)

    Sw = calcscattermat(covestimator_within, Z)

    # compute between-class scattering
    mean = cmeans * (cweights ./ T(n))
    U = rmul!(cmeans .- mean, Diagonal(sqrt.(cweights)))
    Sb = calcscattermat(covestimator_between, U)

    return MulticlassLDAStats(Vector{T}(cweights), mean, cmeans, Sw, Sb)
end

#### Multiclass LDA

mutable struct MulticlassLDA{T<:Real}
    proj::Matrix{T}
    pmeans::Matrix{T}
    stats::MulticlassLDAStats{T}
end

indim(M::MulticlassLDA) = size(M.proj, 1)
outdim(M::MulticlassLDA) = size(M.proj, 2)

projection(M::MulticlassLDA) = M.proj

mean(M::MulticlassLDA) = mean(M.stats)
classmeans(M::MulticlassLDA) = classmeans(M.stats)
classweights(M::MulticlassLDA) = classweights(M.stats)

withclass_scatter(M::MulticlassLDA) = withclass_scatter(M.stats)
betweenclass_scatter(M::MulticlassLDA) = betweenclass_scatter(M.stats)

transform(M::MulticlassLDA, x::AbstractVecOrMat{<:Real}) = M.proj'x

function fit(::Type{MulticlassLDA}, nc::Int, X::DenseMatrix{T}, y::AbstractVector{Int};
             method::Symbol=:gevd,
             outdim::Int=min(size(X,1), nc-1),
             regcoef::T=T(1.0e-6),
             covestimator_within::CovarianceEstimator=SimpleCovariance(),
             covestimator_between::CovarianceEstimator=SimpleCovariance()) where T<:Real

    multiclass_lda(multiclass_lda_stats(nc, X, y;
                                        covestimator_within=covestimator_within,
                                        covestimator_between=covestimator_between);
                   method=method,
                   regcoef=regcoef,
                   outdim=outdim)
end

function multiclass_lda(S::MulticlassLDAStats{T};
                        method::Symbol=:gevd,
                        outdim::Int=min(size(X,1), S.nclasses-1),
                        regcoef::T=T(1.0e-6)) where T<:Real

    P = mclda_solve(S.Sb, S.Sw, method, outdim, regcoef)
    MulticlassLDA(P, P'S.cmeans, S)
end

mclda_solve(Sb::AbstractMatrix{T}, Sw::AbstractMatrix{T}, method::Symbol, p::Int, regcoef::T) where T<:Real =
    mclda_solve!(copy(Sb), copy(Sw), method, p, regcoef)

function mclda_solve!(Sb::AbstractMatrix{T},
                      Sw::AbstractMatrix{T},
                      method::Symbol, p::Int, regcoef::T) where T<:Real

    p <= size(Sb, 1) || throw(ArgumentError("p cannot exceed sample dimension."))

    if method == :gevd
        regularize_symmat!(Sw, regcoef)
        E = eigen!(Symmetric(Sb), Symmetric(Sw))
        ord = sortperm(E.values; rev=true)
        P = E.vectors[:, ord[1:p]]

    elseif method == :whiten
        W = _lda_whitening!(Sw, regcoef)
        wSb = transpose(W) * (Sb * W)
        Eb = eigen!(Symmetric(wSb))
        ord = sortperm(Eb.values; rev=true)
        P = W * Eb.vectors[:, ord[1:p]]

    else
        throw(ArgumentError("Invalid method name $(method)"))
    end
    return P::Matrix{T}
end

function _lda_whitening!(C::AbstractMatrix{T}, regcoef::T) where T<:Real
    n = size(C,1)
    E = eigen!(Symmetric(C))
    v = E.values
    a = regcoef * maximum(v)
    for i = 1:n
        @inbounds v[i] = 1.0 / sqrt(v[i] + a)
    end
    return rmul!(E.vectors,  Diagonal(v))
end

#### SubspaceLDA

# When the dimensionality is much higher than the number of samples,
# it makes more sense to perform LDA on the space spanned by the
# within-group scatter.

struct SubspaceLDA{T<:Real}
    projw::Matrix{T}
    projLDA::Matrix{T}
    λ::Vector{T}
    cmeans::Matrix{T}
    cweights::Vector{Int}
end

indim(M::SubspaceLDA) = size(M.projw,1)
outdim(M::SubspaceLDA) = size(M.projLDA, 2)

projection(M::SubspaceLDA) = M.projw * M.projLDA

mean(M::SubspaceLDA) = vec(sum(M.cmeans * Diagonal(M.cweights / sum(M.cweights)), dims=2))
classmeans(M::SubspaceLDA) = M.cmeans
classweights(M::SubspaceLDA) = M.cweights

transform(M::SubspaceLDA, x) = M.projLDA' * (M.projw' * x)

fit(::Type{F}, X::AbstractMatrix{T}, nc::Int, label::AbstractVector{Int}) where {T<:Real, F<:SubspaceLDA} =
    fit(F, X, label, nc)

function fit(::Type{F}, X::AbstractMatrix{T},
             label::AbstractVector{Int},
             nc=maximum(label);
             normalize::Bool=false) where {T<:Real, F<:SubspaceLDA}
    d, n = size(X, 1), size(X, 2)
    n ≥ nc || throw(ArgumentError("The number of samples is less than the number of classes"))
    length(label) == n || throw(DimensionMismatch("Inconsistent array sizes."))
    # Compute centroids, class weights, and deviation from centroids
    # Note Sb = Hb*Hb', Sw = Hw*Hw'
    cmeans, cweights, Hw = center(X, label, nc)
    dmeans = cmeans .- (normalize ? mean(cmeans, dims=2) : cmeans * (cweights / T(n)))
    Hb = normalize ? dmeans : dmeans * Diagonal(convert(Vector{T}, sqrt.(cweights)))
    if normalize
        Hw /= T(sqrt(n))
    end
    # Project to the subspace spanned by the within-class scatter
    # (essentially, PCA before LDA)
    Uw, Σw, _ = svd(Hw, full=false)
    keep = Σw .> sqrt(eps(T)) * maximum(Σw)
    projw = Uw[:,keep]
    pHb = projw' * Hb
    pHw = projw' * Hw
    λ, G = lda_gsvd(pHb, pHw, cweights)
    SubspaceLDA(projw, G, λ, cmeans, cweights)
end

# Reference: Howland & Park (2006), "Generalizing discriminant analysis
# using the generalized singular value decomposition", IEEE
# Trans. Patt. Anal. & Mach. Int., 26: 995-1006.
function lda_gsvd(Hb::AbstractMatrix{T}, Hw::AbstractMatrix{T}, cweights::AbstractVector{Int}) where T<:Real
    nc = length(cweights)
    K = vcat(Hb', Hw')
    P, R, Q = svd(K, full=false)
    keep = R .> sqrt(eps(T))*maximum(R)
    R = R[keep]
    Pk = P[1:nc, keep]
    U, ΣA, W = svd(Pk)
    ncnz = sum(cweights .> 0)
    G = Q[:,keep]*(Diagonal(1 ./ R) * W[:,1:ncnz-1])
    # Normalize
    Gw = G' * Hw
    nrm = Gw * Gw'
    G = G ./ reshape(sqrt.(diag(nrm)), 1, ncnz-1)
    # Also get the eigenvalues
    Gw = G' * Hw
    Gb = G' * Hb
    λ = diag(Gb * Gb')./diag(Gw * Gw')
    λ, G
end

function center(X::AbstractMatrix{T}, label::AbstractVector{Int}, nc=maximum(label)) where T<:Real
    d, n = size(X,1), size(X,2)
    # Calculate the class weights and means
    cmeans = zeros(T, d, nc)
    cweights = zeros(Int, nc)
    for j = 1:n
        k = label[j]
        for i = 1:d
            cmeans[i,k] += X[i,j]
        end
        cweights[k] += 1
    end
    for j = 1:nc
        cw = cweights[j]
        cw == 0 && continue
        for i = 1:d
            cmeans[i,j] /= cw
        end
    end
    # Compute differences from the means
    dX = Matrix{T}(undef, d, n)
    for j = 1:n
        k = label[j]
        for i = 1:d
            dX[i,j] = X[i,j] - cmeans[i,k]
        end
    end
    cmeans, cweights, dX
end


function Laplacian{T<:AbstractFloat}(nz::Int,nx::Int,dz::T,dx::T)

# This function computes the 2D Laplacian operator/matrix that acts on a
# vectorized 2D field. The vectorization is assumed to be done using column
# major order.
#
# INPUTS:     nz        - Number of grid points in the z-direction
#             nx        - Number of grid points in the x-direction
#             dz        - Grid spacing in z-direction
#             dx        - Grid spacing in x-direction
#
# OUTPUTS:    L         - 2D Laplace operator

    Dzz = (1/dz^2)*spdiagm((ones(nz-1),-2*ones(nz),ones(nz-1)),(-1,0,1))
    Dxx = (1/dx^2)*spdiagm((ones(nx-1),-2*ones(nx),ones(nx-1)),(-1,0,1))

    Mzz = kron(speye(nx),Dzz)
    Mxx = kron(Dxx,speye(nz))

    L = Mzz + Mxx

    return L

end


function crra_utility(;
    β = 0.9,
    σ = 2.0,
    γ = 2.0,
    Π = 0.5 * ones(2, 2),
    G = [0.1, 0.2],
    Θ = ones(Float64, 2),
    transfers = false
    )
    function U(c, n)
        if σ == 1.0
            U = log(c)
        else
            U = (c.^(1.0 - σ) - 1.0) / (1.0 - σ)
        end
        return U - n.^(1 + γ) / (1 + γ)
    end
    # Derivatives of utility function
    Uc(c,n) =  c.^(-σ)
    Ucc(c,n) = -σ * c.^(-σ - 1.0)
    Un(c,n) = -n.^γ
    Unn(c,n) = -γ * n.^(γ - 1.0)
    n_less_than_one = false
    return Model(β, Π, G, Θ, transfers,
                 U, Uc, Ucc, Un, Unn, n_less_than_one)
end


function modalparam_test()
	# State space model for testing modalparam
	# and modalparami
	#
	# javier.cara@upm.es

	M = [35.0 0.0;0.0 17.5] # mass matrix
	K = [12250.0 -3500.0;-3500.0 3500.0] # stiffness matrix
	g = [0.02;0.02] # damping ratios

	# teorethical eigenvalues and eigenvectors
	W2,V = eig(K,M)
	W = sqrt(W2)

	# damping matrix
	Mm = V'*M*V
	Gm = 2*Mm*diagm(W)*diagm(g)
	G = inv(V')*Gm*inv(V) # damping matrix

	# state-space matrices
	Minv = [ 1/M[1,1] 0.0;0.0 1/M[2,2] ]
	Ac = [zeros(2,2) eye(2,2);-Minv*K -Minv*G] # continuous A matrix
	dt = 0.02 # time step
	A = expm(Ac*dt) # discrete A matrix
	C = [-Minv*K -Minv*G] # discrete C matrix
	Bc = [zeros(2,2);Minv]
	B = (A - eye(4))*inv(Ac)*Bc

	# eigenvectors with max. component = 1
	V1 = zeros(2,2)
	for j in 1:2
		maxval = V[1,j]
		if abs(V[2,j]) > abs(maxval)
			maxval = V[2,j]
		end
		V1[:,j] = V[:,j]/maxval
	end
	# updated mass masses
	Mm1 = V1'*M*V1
	mm1 = diag(Mm1)

	ssm = Dict("M"=>M,"K"=>K,"G"=>G,"wm"=>W,"zm"=>g,"Vm"=>V1,"mm"=>mm1,
							"dt"=>dt,"A"=>A,"B"=>B,"C"=>C)

	return ssm

end


#=
Example 3:
  Fit the following analytical model with 3 parameters:
  (p1  +  p2 * x  +  p3 * x^2)  *  cos(x)

A `FuncWrap` component is used to wrap a common Julia function
=#
using Random, DataFitting

# Define the analytic model function
f(x, p1, p2, p3) = @. (p1  +  p2 * x  +  p3 * x^2)  *  cos(x)

# "True" parameter values:
params = [1, 2, 3];

# "True" physical quantities
x = 0:0.1:15
y = f(x, params...);

# Add some noise to simulate a measurement process
noise = 0.1 .* y
y .+= noise .* randn(length(x));

# Wrap empirical measures and uncertainties in a `Measure` object
data = Measures(y, noise)

# Prepare the model with a `FuncWrap` component.  Provide the guess
# values with the `params` vector.
model = Model(Domain(x), :comp1 => DataFitting.FuncWrap(f, params...))

# Fit the model to empirical data
result = fit!(model, data)

# Print best fit parameters
for i in 1:length(result[:comp1].p)
    println("p[$i] = ", result[:comp1].p[i].val, " ± ", result[:comp1].p[i].unc)
end

# Plot data and best fit model
using Gnuplot
@gp    x data.val data.unc "w yerr" :-
@gp :- x model[1] "w line"


# p32.jl - solve u_xx = exp(4x), u(-1)=0, u(1)=1 (compare p13.jl)

N = 16;
(D,x) = cheb(N);
D2 = D^2;
D2 = D2[2:N,2:N];                   # boundary conditions
f = exp(4*x[2:N]);
u = D2\f;                           # Poisson eq. solved here
u = [0;u;0] + (x+1)/2;
clf();
plot(x,u,".",markersize=10);
xx = -1:.01:1;
uu = polyval(polyfit(x,u),xx);
plot(xx,uu); grid(true);
exact = @. (exp(4*xx) - sinh(4)*xx - cosh(4))/16 + (xx+1)/2;
title("max err = $(signif(norm(uu-exact,Inf),4))",fontsize=12);


################################################################################
# 3-D model

"""
    dett(::Type{DeforModelRed3D},  C::FMat{T}) where {T}

Compute the determinant of a general square matrix.
"""
function dett(::Type{DeforModelRed3D},  C::FMat{T}) where {T}
	return (C[1,1] * C[2,2] * C[3,3] +
		C[1,2] * C[2,3] * C[3,1] +
		C[1,3] * C[2,1] * C[3,2] -
		C[1,3] * C[2,2] * C[3,1] -
		C[1,2] * C[2,1] * C[3,3] -
		C[1,1] * C[2,3] * C[3,2])
end

"""
    strainvdet(::Type{DeforModelRed3D},  Cv::FVec{T}) where {T}

Compute the determinant of a symmetric strain-like square matrix represented
as a vector. Remember that the shear strain components are twice the entries
of the matrix representation.
"""
function strainvdet(::Type{DeforModelRed3D},  Cv::FVec{T}) where {T}
	return (Cv[1] * Cv[2] * Cv[3] +
		Cv[4]/2 * Cv[6]/2 * Cv[5]/2 +
		Cv[5]/2 * Cv[4]/2 * Cv[6]/2 -
		Cv[5]/2 * Cv[2] * Cv[5]/2 -
		Cv[4]/2 * Cv[4]/2 * Cv[3] -
		Cv[1] * Cv[6]/2 * Cv[6]/2)
end

"""
    strainvtr(::Type{DeforModelRed3D},  Cv::FVec{T}) where {T}

Compute the trace of a symmetric strain-like square matrix represented as a
vector.
"""
function strainvtr(::Type{DeforModelRed3D},  Cv::FVec{T}) where {T}
	return (Cv[1] + Cv[2] + Cv[3])
end

"""
    strainttov!(::Type{DeforModelRed3D}, v::FVec{T}, t::FMat{T}) where {T}

Convert a symmetric matrix of 3x3 strain components  into a 6-component vector.
"""
function strainttov!(::Type{DeforModelRed3D}, v::FVec{T}, t::FMat{T}) where {T}
	v[1] = t[1,1];
	v[2] = t[2,2];
	v[3] = t[3,3];
	v[4] = t[1,2] + t[2,1];
	v[5] = t[1,3] + t[3,1];
	v[6] = t[3,2] + t[2,3];
	return v
end

"""
    strainvtot!(::Type{DeforModelRed3D}, t::FMat{T}, v::FVec{T}) where {T}

Convert a strain 3-vector to a  matrix of 2x2 strain components (symmetric tensor).
"""
function strainvtot!(::Type{DeforModelRed3D}, t::FMat{T}, v::FVec{T}) where {T}
	t[1,1] = v[1];
	t[2,2] = v[2];
	t[3,3] = v[3];
	t[1,2] = v[4]/2.;
	t[2,1] = v[4]/2.;
	t[1,3] = v[5]/2.;
	t[3,1] = v[5]/2.;
	t[3,2] = v[6]/2.;
	t[2,3] = v[6]/2.;
    return t
end

"""
    stressttov!(::Type{DeforModelRed3D}, v::FVec{T}, t::FMat{T}) where {T}

Convert a symmetric matrix of 3x3 stress components to a 6-component vector.
"""
function stressttov!(::Type{DeforModelRed3D}, v::FVec{T}, t::FMat{T}) where {T}
	v[1] = t[1,1];
	v[2] = t[2,2];
	v[3] = t[3,3];
	v[4] = 1/2.0*(t[1,2] + t[2,1]);
	v[5] = 1/2.0*(t[1,3] + t[3,1]);
	v[6] = 1/2.0*(t[3,2] + t[2,3]);
	return v
end

"""
    stressvtot!(::Type{DeforModelRed3D}, t::FMat{T}, v::FVec{T}) where {T}

Convert a 6-vector to a  matrix of 3x3 stress components (symmetric tensor).
"""
function stressvtot!(::Type{DeforModelRed3D}, t::FMat{T}, v::FVec{T}) where {T}
	t[1,1] = v[1];
	t[2,2] = v[2];
	t[3,3] = v[3];
	t[1,2] = v[4];
	t[2,1] = v[4];
	t[1,3] = v[5];
	t[3,1] = v[5];
	t[3,2] = v[6];
	t[2,3] = v[6];
	return t
end

################################################################################
# 2-D plane strain model

"""
    dett(::Type{DeforModelRed2DStrain},  C::FMat{T}) where {T}

Compute the determinant of a general square matrix.
"""
function dett(::Type{DeforModelRed2DStrain},  C::FMat{T}) where {T}
	return (C[1,1] * C[2,2] - C[1,2] * C[2,1])
end

"""
    strainvdet(::Type{DeforModelRed2DStrain},  Cv::FVec{T}) where {T}

Compute the determinant of a symmetric strain-like square matrix represented
as a vector. Remember that the shear strain components are twice the entries
of the matrix representation.
"""
function strainvdet(::Type{DeforModelRed2DStrain},  Cv::FVec{T}) where {T}
	return (Cv[1] * Cv[2] - Cv[3]/2 * Cv[3]/2)
end

"""
    strainvtr(::Type{DeforModelRed2DStrain},  Cv::FVec{T}) where {T}

Compute the trace of a symmetric strain-like square matrix represented as a
vector.
"""
function strainvtr(::Type{DeforModelRed2DStrain},  Cv::FVec{T}) where {T}
	return (Cv[1] + Cv[2])
end

"""
    strainttov!(::Type{DeforModelRed2DStrain}, v::FVec{T}, t::FMat{T}) where {T}

Convert a symmetric matrix of 2x2 strain components  into a 3-component vector.
"""
function strainttov!(::Type{DeforModelRed2DStrain}, v::FVec{T}, t::FMat{T}) where {T}
    v[1] = t[1,1];
    v[2] = t[2,2];
    v[3] = t[1,2] + t[2,1];
    return v
end

"""
    strainvtot!(::Type{DeforModelRed2DStrain}, t::FMat{T}, v::FVec{T}) where {T}

Convert a strain 3-vector to a  matrix of 2x2 strain components (symmetric tensor).
"""
function strainvtot!(::Type{DeforModelRed2DStrain}, t::FMat{T}, v::FVec{T}) where {T}
    t[1,1] = v[1];
    t[2,2] = v[2];
    t[1,2] = v[3]/2.;
    t[2,1] = v[3]/2.;
    return t
end

"""
    stressttov!(::Type{DeforModelRed2DStrain}, v::FVec{T}, t::FMat{T}) where {T}

Convert a symmetric matrix of 2x2 stress components to a 3-component vector.
"""
function stressttov!(::Type{DeforModelRed2DStrain}, v::FVec{T}, t::FMat{T}) where {T}
	v[1] = t[1,1];
	v[2] = t[2,2];
	v[3] = 0.5*(t[1,2] + t[2,1]);
	return v
end

"""
    stressvtot!(::Type{DeforModelRed2DStrain}, t::FMat{T}, v::FVec{T}) where {T}

Convert a vector to a  matrix of 2x2 stress components (symmetric tensor).

If `v` has 4 entries, also the `t[3,3]` matrix entry is set.

The stress vector components need to be ordered as:
    sigmax, sigmay, tauxy, sigmaz,
which is the ordering used for the plane-strain model reduction.
"""
function stressvtot!(::Type{DeforModelRed2DStrain}, t::FMat{T}, v::FVec{T}) where {T}
    t[1,1] = v[1];
    t[2,2] = v[2];
    t[1,2] = v[3];
    t[2,1] = v[3];
    (length(v) == 4)  && (t[3,3] = v[4]);
    return t
end

################################################################################
# 2-D plane stress model

"""
    stressttov!(::Type{DeforModelRed2DStress}, v::FVec{T}, t::FMat{T}) where {T}

Convert a symmetric matrix of 2x2 stress components to a 3-component vector.
"""
function stressttov!(::Type{DeforModelRed2DStress}, v::FVec{T}, t::FMat{T}) where {T}
	v[1] = t[1,1];
	v[2] = t[2,2];
	v[3] = 0.5*(t[1,2] + t[2,1]);
	return v
end

"""
    stressvtot!(::Type{DeforModelRed2DStress}, t::FMat{T}, v::FVec{T}) where {T}

Convert a 3-vector to a  matrix of 2x2 stress components (symmetric tensor).
"""
function stressvtot!(::Type{DeforModelRed2DStress}, t::FMat{T}, v::FVec{T}) where {T}
    t[1,1] = v[1];
    t[2,2] = v[2];
    t[1,2] = v[3];
    t[2,1] = v[3];
    return t
end

################################################################################
# 2-D axially symmetric stress model

"""
    stressvtot!(::Type{DeforModelRed2DAxisymm}, t::FMat{T}, v::FVec{T}) where {T}

Convert a 4-vector to a  matrix of 3x3 stress components (tensor).

Convert a 4-vector to a *symmetric* matrix of 3x3 stress components (tensor).

The stress vector components need to be ordered as:
    sigmax, sigmay, sigmaz, tauxy.
"""
function stressvtot!(::Type{DeforModelRed2DAxisymm}, t::FMat{T}, v::FVec{T}) where {T}
    t[1,1] = v[1];
    t[2,2] = v[2];
    t[1,2] = v[4];
    t[2,1] = v[4];
    t[3,3] = v[3];
    return t
end


# Used "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy as
# a reference.  Note that there were some typos in that document so the code
# here may not correspond exactly.

immutable NormalWishart <: Distribution
    dim::Int
    zeromean::Bool
    mu::Vector{Float64}
    kappa::Float64
    Tchol::Cholesky{Float64}  # Precision matrix (well, sqrt of one)
    nu::Float64

    function NormalWishart(mu::Vector{Float64}, kappa::Real,
                                  Tchol::Cholesky{Float64}, nu::Real)
        # Probably should put some error checking in here
        d = length(mu)
        zmean::Bool = true
        for i = 1:d
            if mu[i] != 0.
                zmean = false
                break
            end
        end
        new(d, zmean, mu, float64(kappa), Tchol, float64(nu))
    end
end

function NormalWishart(mu::Vector{Float64}, kappa::Real,
                       T::Matrix{Float64}, nu::Real)
    NormalWishart(mu, kappa, cholfact(T), nu)

end

function insupport(::Type{NormalWishart}, x::Vector{Float64}, Lam::Matrix{Float64})
    return (all(isfinite(x)) &&
           size(Lam, 1) == size(Lam, 2) &&
           isApproxSymmmetric(Lam) &&
           size(Lam, 1) == length(x) &&
           hasCholesky(Lam))
end

pdf(nw::NormalWishart, x::Vector{Float64}, Lam::Matrix{Float64}) =
        exp(logpdf(nw, x, Lam))

function logpdf(nw::NormalWishart, x::Vector{Float64}, Lam::Matrix{Float64})
    if !insupport(NormalWishart, x, Lam)
        return -Inf
    else
        p = length(x)

        nu = nw.nu
        kappa = nw.kappa
        mu = nw.mu
        Tchol = nw.Tchol
        hnu = 0.5 * nu
        hp = 0.5 * p

        # Normalization
        logp::Float64 = hp*(log(kappa) - float64(log2π))
        logp -= hnu * logdet(Tchol)
        logp -= hnu * p * log(2.)
        logp -= lpgamma(p, hnu)

        # Wishart (MvNormal contributes 0.5 as well)
        logp += (hnu - hp) * logdet(Lam)
        logp -= 0.5 * trace(Tchol \ Lam)

        # Normal
        z = nw.zeromean ? x : x - mu
        logp -= 0.5 * kappa * dot(z, Lam * z)

        return logp

    end
end

function rand(nw::NormalWishart)
    Lam = rand(Wishart(nw.nu, nw.Tchol))
    mu = rand(MvNormal(nw.mu, inv(Lam) ./ nw.kappa))
    return (mu, Lam)
end


# # POVM simulation
# This notebook shows how we can check how much depolarizing noise a qubit positive operator-valued measure (POVM) can take before it becomes simulable by projective measurements. The general method is described in [arXiv:1609.06139](https://arxiv.org/abs/1609.06139). The question of simulability by projective measurements boils down to an SDP problem. Eq. (8) from the paper defines the noisy POVM that we obtain subjecting a POVM $\mathbf{M}$ to a depolarizing channel $\Phi_t$:
#
# $$
# \left[\Phi_t\left(\mathbf{M}\right)\right]_i := t M_i + (1-t)\frac{\mathrm{tr}(M_i)}{d} \mathbb{1}.
# $$
#
# If this visibility $t\in[0,1]$ is one, the POVM $\mathbf{M}$ is simulable.
#
# We will use Convex.jl to solve the SDP problem.

using Convex, SCS, LinearAlgebra
if VERSION < v"1.2.0-DEV.0"
    (I::UniformScaling)(n::Integer) = Diagonal(fill(I.λ, n))
     LinearAlgebra.diagm(v::AbstractVector) = diagm(0 => v)
end

# For the qubit case, a four outcome qubit POVM $\mathbf{M} \in\mathcal{P}(2,4)$ is simulable if and only if
#
# $M_{1}=N_{12}^{+}+N_{13}^{+}+N_{14}^{+},$
#
# $M_{2}=N_{12}^{-}+N_{23}^{+}+N_{24}^{+},$
#
# $M_{3}=N_{13}^{-}+N_{23}^{-}+N_{34}^{+},$
#
# $M_{4}=N_{14}^{-}+N_{24}^{-}+N_{34}^{-},$
#
# where Hermitian operators $N_{ij}^{\pm}$ satisfy $N_{ij}^{\pm}\geq0$ and $N_{ij}^{+}+N_{ij}^{-}=p_{ij}\mathbb{1}$, where $i<j$ , $i,j=1,2,3,4$ and $p_{ij}\geq0$ as well as $\sum_{i<j}p_{ij}=1$, that is, the $p_{ij}$ values form a probability vector. This forms an SDP feasibility problem, which we can rephrase as an optimization problem by adding depolarizing noise to the left-hand side of the above equations and maximizing the visibility $t$:
#
# $\max_{t\in[0,1]} t$
#
# such that
#
# $t\,M_{1}+(1-t)\,\mathrm{tr}(M_{1})\frac{\mathbb{1}}{2}=N_{12}^{+}+N_{13}^{+}+N_{14}^{+},$
#
# $t\,M_{2}+(1-t)\,\mathrm{tr}(M_{2})\frac{\mathbb{1}}{2}=N_{12}^{-}+N_{23}^{+}+N_{24}^{+},$
#
# $t\,M_{3}+(1-t)\,\mathrm{tr}(M_{3})\frac{\mathbb{1}}{2}=N_{13}^{-}+N_{23}^{-}+N_{34}^{+},$
#
# $t\,M_{4}+(1-t)\,\mathrm{tr}(M_{4})\frac{\mathbb{1}}{2}=N_{14}^{-}+N_{24}^{-}+N_{34}^{-}$.
#
# We organize these constraints in a function that takes a four-output qubit POVM as its argument:

function get_visibility(K)
    noise = real([tr(K[i])*I(2)/2 for i=1:size(K, 1)])
    P = [[ComplexVariable(2, 2) for i=1:2] for j=1:6]
    q = Variable(6, Positive())
    t = Variable(1, Positive())
    constraints = [P[i][j] in :SDP for i=1:6 for j=1:2]
    constraints += sum(q)==1
    constraints += t<=1
    constraints += [P[i][1]+P[i][2] == q[i]*I(2) for i=1:6]
    constraints += t*K[1] + (1-t)*noise[1] == P[1][1] + P[2][1] + P[3][1]
    constraints += t*K[2] + (1-t)*noise[2] == P[1][2] + P[4][1] + P[5][1]
    constraints += t*K[3] + (1-t)*noise[3] == P[2][2] + P[4][2] + P[6][1]
    constraints += t*K[4] + (1-t)*noise[4] == P[3][2] + P[5][2] + P[6][2]
    p = maximize(t, constraints)
    solve!(p, SCS.Optimizer(verbose=0))
    return p.optval
end

# We check this function using the tetrahedron measurement (see Appendix B in [arXiv:quant-ph/0702021](https://arxiv.org/abs/quant-ph/0702021)). This measurement is non-simulable, so we expect a value below one.

function dp(v)
    I(2) + v[1]*[0 1; 1 0] + v[2]*[0 -im; im 0] + v[3]*[1 0; 0 -1]
end
b = [ 1  1  1;
     -1 -1  1;
     -1  1 -1;
      1 -1 -1]/sqrt(3)
M = [dp(b[i, :]) for i=1:size(b,1)]/4;
get_visibility(M)

# This value matches the one [we obtained](https://github.com/peterwittek/ipython-notebooks/blob/master/Simulating_POVMs.ipynb) using [PICOS](http://picos.zib.de/).


"""
    gpdfitnew(x)

Estimate the paramaters for the Generalized Pareto Distribution (GPD). Returns empirical Bayes estimate for the parameters of the two-parameter generalized Parato distribution given the data.

# Arguments
* `x::AbstractArray`: One dimensional data array.

# Returns
* `k::Float64, sigma::Float64`: Estimated parameter values.

# Notes
* This function returns a negative of Zhang and Stephens's k, because it is more common parameterization.
"""
function gpdfitnew(x::AbstractArray)
    if ndims(x) != 1 || length(x) <= 1
        throw(DimensionMismatch())
    end
    sort!(x)

    n = length(x)
    m = 80 + trunc(Int, sqrt(n))
    prior = 3

    bs = collect(1:m) .- 0.5
    bs .= (1 .- sqrt.(m ./ bs)) ./ prior ./ x[trunc(Int64, n/4 + 0.5)] .+ 1 / x[end]

    ks = mean(log1p.(-bs .* x'), dims=2)[:,1]
    L = n * (log.(bs ./ -ks) .- ks .- 1)
    w = similar(L)
    w .= (1 ./ sum(exp.(L .- L'), dims=1)')[:,1]

    # remove negligible weights
    dii = w .>= (10 * eps(Float64))
    if ~all(dii)
        w = w[dii]
        bs = bs[dii]
    end
    # normalise w
    w ./= sum(w)
    # posterior mean for b
    b = sum(bs .* w)
    # Estimate for k, note that we return a negative of Zhang and
    # Stephens's k, because it is more common parameterisation.
    k = mean(log1p.(-b .* x))
    # estimate for sigma
    sigma = -k / b

    return k, sigma

end


abstract type ClusterModel end

"""
    TotalClusterSS( Clustered::ClusterModel )

Returns a scalar of the total sum of squares for a ClusterModel object.
"""
function TotalClusterSS( Clustered::ClusterModel )
    GrandMean = StatsBase.mean( Clustered.X , dims = 1 )
    return sum((Clustered.X .- GrandMean).^ 2)
end

"""
    WithinClusterSS( Clustered::ClusterModel )

Returns a scalar of the within cluter sum of squares for a ClusterModel object.
"""
function WithinClusterSS( Clustered::ClusterModel )
    Clusters = unique( Clustered.Assignments )
    WithinSS = zeros( size(Clustered.X) )
    for cluster in Clusters
        ElementsInCluster = vec(Clustered.Assignments .== cluster)
        samples = Clustered.X[ ElementsInCluster, : ]
        clusterMean = StatsBase.mean( samples , dims = 1 )
        WithinSS[ElementsInCluster,:] = (samples .- clusterMean) .^ 2
    end
    return sum( WithinSS )
end

"""
    BetweenClusterSS( Clustered::ClusterModel )

Returns a scalar of the between cluster sum of squares for a ClusterModel object.
"""
function BetweenClusterSS( Clustered::ClusterModel )
    Clusters = unique( Clustered.Assignments )
    GrandMean = StatsBase.mean( Clustered.X , dims = 1 )
    BetweenSS = 0.0
    for cluster in Clusters
        ElementsInCluster = vec(Clustered.Assignments .== cluster)
        samples = Clustered.X[ ElementsInCluster, : ]
        clusterMean = StatsBase.mean( samples , dims = 1 )
        BetweenSS += size(samples)[1] * sum( ( clusterMean .- GrandMean ) .^ 2)
    end
    return BetweenSS
end

struct KMeansClustering <: ClusterModel
    X
    Centroids
    Assignments
end

"""
    KMeans( X, Clusters; tolerance = 1e-8, maxiters = 200 )

Returns a ClusterModel object after finding clusterings for data in `X` via MacQueens K-Means algorithm. `Clusters` is the K parameter, or the # of clusters.

MacQueen, J. B. (1967). Some Methods for classification and Analysis of Multivariate Observations. Proceedings of 5th Berkeley Symposium on Mathematical Statistics and Probability. 1. University of California Press. pp. 281–297.
"""
function KMeans( X, Clusters; tolerance = 1e-8, maxiters = 200 )
    (Xrows, Xcols) = size( X )
    ResultVector = zeros( Xrows )
    #1) Initialize K-Means centroids: Randomly select C samples
    rndsmpl = collect( StatsBase.sample( collect(1 : Xrows), Clusters, replace = false ) )
    Centroids = X[ rndsmpl, :]
    newCentroid = zeros( Clusters, Xcols )
    #2) finding which many points are nearest to our centroid
    #3) Calculate the mean value of each dimension to make a new centroid and repeat.
    for iter in 1 : maxiters
        DistsToCentroids = SquareEuclideanDistance( X, Centroids )
        ResultVector = last.(Tuple.(argmin( DistsToCentroids, dims = 2 )))

        for cluster in 1:Clusters
            Assigned = (ResultVector .== cluster)[:,1]
            newCentroid[ cluster, : ] = StatsBase.mean( X[ Assigned, : ], dims = 1 )
        end
        ToleranceCheck = ( 1.0 / Clusters ) * sum( ( Centroids .- newCentroid ) .^ 2)
        if ToleranceCheck <= tolerance ; break; end
        Centroids = newCentroid
    end
    return KMeansClustering( X, Centroids, ResultVector )
end


"""
    angleofrefraction(n1, n2, β)

Compute the angle of refraction.

# Arguments
* `n1`: refractive index of the medium in which the wave is propagating towards the interface
* `n2`: refractive index of the other medium
* `β`:  incident or reflected angle
"""
function angleofrefraction(n1, n2, β)
    asin(n1 / n2 * sin(β))
end

"""
    fresnel_x(n1, n2, β)

Compute the diagonal elements \$L_{xx}\$ of the fresnel factor.

# Arguments
* `n1`: refractive index of the medium in which the wave is propagating towards the interface
* `n2`: refractive index of the other medium
* `β`:  incident or reflected angle

https://doi.org/10.1103/PhysRevB.59.1263
"""
function fresnel_x(n1, n2, β)
    γ = angleofrefraction(n1, n2, β)
    2n1 * cos(γ) / (n1 * cos(γ) + n2 * cos(β))
end

"""
    fresnel_y(n1, n2, β)

Compute the diagonal elements \$L_{yy}\$ of the fresnel factor.

# Arguments
* `n1`: refractive index of the medium in which the wave is propagating towards the interface
* `n2`: refractive index of the other medium
* `β`:  incident or reflected angle

https://doi.org/10.1103/PhysRevB.59.1263
"""
function fresnel_y(n1, n2, β)
    γ = angleofrefraction(n1, n2, β)
    2n1 * cos(β) / (n1 * cos(β) + n2 * cos(γ))
end

"""
    fresnel_z(n1, n2, n′, β)

Compute the diagonal elements \$L_{yy}\$ of the fresnel factor.

# Arguments
* `n1`: refractive index of the medium in which the wave is propagating towards the interface
* `n2`: refractive index of the other medium
* `n′`: refractive index of the interface
* `β`:  incident or reflected angle

https://doi.org/10.1103/PhysRevB.59.1263
"""
function fresnel_z(n1, n2, n′, β)
    γ = angleofrefraction(n1, n2, β)
    2n2 * cos(β) / (n1 * cos(γ) + n2 * cos(β)) * (n1 / n′)^2
end

"""
    ninterface(n2)

Calculate the index of refraction of the interfacial layer according to

``\\left(\\frac{1}{n^\\prime}\\right)^2 = \\frac{4n_2^2 + 2}{n_2^2(n_2 ^2+5)}``.

# Arguments
* `n2`: refractive index of the medium in which the wave is **not** propagating towards the interface

https://doi.org/10.1103/PhysRevB.59.1263
"""
function ninterface(n2)
    1 / sqrt((4 * n2^2 + 2) / (n2^2 * (n2^2 + 5)))
end


function gradexpr(ex::SymbolicVariable, wrt::SymbolicVariable)
    return ex == wrt ? 1 : 0
end

gradexpr(ex::Number, wrt::SymbolicVariable) = 0

function gradexpr(ex::Expr, wrt::Any)
    if ex.head != :call
        error("Unrecognized expression $ex")
    end
    simplify(
        gradexpr(
            SymbolParameter(ex.args[1]),
            ex.args[2:end],
            wrt
        )
    )
end

function gradexpr{T}(x::SymbolParameter{T}, args::Any, wrt::Any)
    error(
        @sprintf(
            "Derivative of function %s not supported",
            string(T)
        )
    )
end

# The Power Rule
function gradexpr(::SymbolParameter{:^}, args::Any, wrt::Any)
    x, y = args[1], args[2]
    xp, yp = gradexpr(x, wrt), gradexpr(y, wrt)
    if xp == 0 && yp == 0
        return 0
    elseif yp == 0
        return :($y * $xp * ($x^($y - 1)))
    else
        return :($x^$y * ($xp * $y / $x + $yp * log($x)))
    end
end

function gradexpr(::SymbolParameter{:+}, args::Any, wrt::Any)
    termdiffs = {:+}
    for y in args
        x = gradexpr(y, wrt)
        if x != 0
            push!(termdiffs, x)
        end
    end
    n = length(termdiffs)
    if n == 1
        return 0
    elseif n == 2
        return termdiffs[2]
    else
        return Expr(:call, termdiffs...)
    end
end

function gradexpr(::SymbolParameter{:-}, args::Any, wrt::Any)
    termdiffs = {:-}
    # first term is special, can't be dropped
    term1 = gradexpr(args[1], wrt)
    push!(termdiffs, term1)
    for y in args[2:end]
        x = gradexpr(y, wrt)
        if x != 0
            push!(termdiffs, x)
        end
    end
    n = length(termdiffs)
    if term1 != 0 && n == 2 && length(args) >= 2
        # if all of the terms but the first disappeared, we return the first
        return term1
    elseif term1 == 0 && n == 2
        return 0
    else
        return Expr(:call, termdiffs...)
    end
end

# The Product Rule
# d/dx (f * g) = (d/dx f) * g + f * (d/dx g)
# d/dx (f * g * h) = (d/dx f) * g * h + f * (d/dx g) * h + ...
function gradexpr(::SymbolParameter{:*}, args::Any, wrt::Any)
    n = length(args)
    resargs = Array(Any, n)
    for i in 1:n
       newargs = Array(Any, n)
       for j in 1:n
           if j == i
               newargs[j] = gradexpr(args[j], wrt)
           else
               newargs[j] = args[j]
           end
       end
       resargs[i] = Expr(:call, :*, newargs...)
    end
    return Expr(:call, :+, resargs...)
end

# The Quotient Rule
# d/dx (f / g) = ((d/dx f) * g - f * (d/dx g)) / g^2
function gradexpr(::SymbolParameter{:/}, args::Any, wrt::Any)
    x, y = args[1], args[2]
    xp, yp = gradexpr(x, wrt), gradexpr(y, wrt)
    if xp == 0 && yp == 0
        return 0
    elseif xp == 0
        return :(-$yp * $x / $y^2)
    elseif yp == 0
        return :($xp / $y)
    else
        return :(($xp * $y - $x * $yp) / $y^2)
    end
end

# This table is used in other packages, and if someone changes it they should
# notify the following packages:
# * https://github.com/scidom/DualNumbers.jl
derivative_rules = [
    ( :sqrt,        :(  xp / 2 / sqrt(x)                         ))
    ( :cbrt,        :(  xp / 3 / cbrt(x)^2                       ))
    ( :square,      :(  xp * 2 * x                               ))
    ( :log,         :(  xp / x                                   ))
    ( :log10,       :(  xp / x / log(10)                         ))
    ( :log2,        :(  xp / x / log(2)                          ))
    ( :log1p,       :(  xp / (x + 1)                             ))
    ( :exp,         :(  xp * exp(x)                              ))
    ( :exp2,        :(  xp * log(2) * exp2(x)                    ))
    ( :expm1,       :(  xp * exp(x)                              ))
    ( :sin,         :(  xp * cos(x)                              ))
    ( :cos,         :( -xp * sin(x)                              ))
    ( :tan,         :(  xp * (1 + tan(x)^2)                      ))
    ( :sec,         :(  xp * sec(x) * tan(x)                     ))
    ( :csc,         :( -xp * csc(x) * cot(x)                     ))
    ( :cot,         :( -xp * (1 + cot(x)^2)                      ))
    ( :sind,        :(  xp * pi / 180 * cosd(x)                  ))
    ( :cosd,        :( -xp * pi / 180 * sind(x)                  ))
    ( :tand,        :(  xp * pi / 180 * (1 + tand(x)^2)          ))
    ( :secd,        :(  xp * pi / 180 * secd(x) * tand(x)        ))
    ( :cscd,        :( -xp * pi / 180 * cscd(x) * cotd(x)        ))
    ( :cotd,        :( -xp * pi / 180 * (1 + cotd(x)^2)          ))
    ( :asin,        :(  xp / sqrt(1 - x^2)                       ))
    ( :acos,        :( -xp / sqrt(1 - x^2)                       ))
    ( :atan,        :(  xp / (1 + x^2)                           ))
    ( :asec,        :(  xp / abs(x) / sqrt(x^2 - 1)              ))
    ( :acsc,        :( -xp / abs(x) / sqrt(x^2 - 1)              ))
    ( :acot,        :( -xp / (1 + x^2)                           ))
    ( :asind,       :(  xp * 180 / pi / sqrt(1 - x^2)            ))
    ( :acosd,       :( -xp * 180 / pi / sqrt(1 - x^2)            ))
    ( :atand,       :(  xp * 180 / pi / (1 + x^2)                ))
    ( :asecd,       :(  xp * 180 / pi / abs(x) / sqrt(x^2 - 1)   ))
    ( :acscd,       :( -xp * 180 / pi / abs(x) / sqrt(x^2 - 1)   ))
    ( :acotd,       :( -xp * 180 / pi / (1 + x^2)                ))
    ( :sinh,        :(  xp * cosh(x)                             ))
    ( :cosh,        :(  xp * sinh(x)                             ))
    ( :tanh,        :(  xp * sech(x)^2                           ))
    ( :sech,        :( -xp * tanh(x) * sech(x)                   ))
    ( :csch,        :( -xp * coth(x) * csch(x)                   ))
    ( :coth,        :( -xp * csch(x)^2                           ))
    ( :asinh,       :(  xp / sqrt(x^2 + 1)                       ))
    ( :acosh,       :(  xp / sqrt(x^2 - 1)                       ))
    ( :atanh,       :(  xp / (1 - x^2)                           ))
    ( :asech,       :( -xp / x / sqrt(1 - x^2)                   ))
    ( :acsch,       :( -xp / abs(x) / sqrt(1 + x^2)              ))
    ( :acoth,       :(  xp / (1 - x^2)                           ))
    ( :erf,         :(  xp * 2 * exp(-square(x)) / sqrt(pi)      ))
    ( :erfc,        :( -xp * 2 * exp(-square(x)) / sqrt(pi)      ))
    ( :erfi,        :(  xp * 2 * exp(square(x)) / sqrt(pi)       ))
    ( :gamma,       :(  xp * digamma(x) * gamma(x)               ))
    ( :lgamma,      :(  xp * digamma(x)                          ))
    # note: only covers the 1-arg version of :airy
    ( :airy,        :(  xp * airyprime(x)                        ))
    ( :airyprime,   :(  xp * airy(2, x)                          ))
    ( :airyai,      :(  xp * airyaiprime(x)                      ))
    ( :airybi,      :(  xp * airybiprime(x)                      ))
    ( :airyaiprime, :(  xp * x * airyai(x)                       ))
    ( :airybiprime, :(  xp * x * airybi(x)                       ))
    ( :besselj0,    :( -xp * besselj1(x)                         ))
    ( :besselj1,    :(  xp * (besselj0(x) - besselj(2, x)) / 2   ))
    ( :bessely0,    :( -xp * bessely1(x)                         ))
    ( :bessely1,    :(  xp * (bessely0(x) - bessely(2, x)) / 2   ))
    ## ( :erfcx,   :(  xp * (2 * x * erfcx(x) - 2 / sqrt(pi))   ))  # uncertain
    ## ( :dawson,  :(  xp * (1 - 2x * dawson(x))                ))  # uncertain
]

for (funsym, exp) in derivative_rules
    @eval function gradexpr(
        ::SymbolParameter{$(Meta.quot(funsym))},
        args::Any,
        wrt::Any
    )
        x = args[1]
        xp = gradexpr(x, wrt)
        if xp != 0
            return @sexpr($exp)
        else
            return 0
        end
    end
end

derivative_rules_bessel = [
    ( :besselj,    :(  xp * (besselj(nu - 1, x) - besselj(nu + 1, x)) / 2   ))
    ( :besseli,    :(  xp * (besseli(nu - 1, x) + besseli(nu + 1, x)) / 2   ))
    ( :bessely,    :(  xp * (bessely(nu - 1, x) - bessely(nu + 1, x)) / 2   ))
    ( :besselk,    :( -xp * (besselk(nu - 1, x) + besselk(nu + 1, x)) / 2   ))
    ( :hankelh1,   :(  xp * (hankelh1(nu - 1, x) - hankelh1(nu + 1, x)) / 2 ))
    ( :hankelh2,   :(  xp * (hankelh2(nu - 1, x) - hankelh2(nu + 1, x)) / 2 ))
]

# 2-argument bessel functions
for (funsym, exp) in derivative_rules_bessel
    @eval function gradexpr(
        ::SymbolParameter{$(Meta.quot(funsym))},
        args::Any,
        wrt::Any
    )
        nu, x = args[1], args[2]
        xp = gradexpr(x, wrt)
        if xp != 0
            return @sexpr($exp)
        else
            return 0
        end
    end
end

### Other functions from julia/base/math.jl we might want to define
### derivatives for. Some have two arguments.

## atan2
## hypot
## beta, lbeta, eta, zeta, digamma

function gradexpr(ex::Expr, targets::Vector{Symbol})
    n = length(targets)
    exprs = Array(Any, n)
    for i in 1:n
        exprs[i] = gradexpr(ex, targets[i])
    end
    return exprs
end


using Distributions, Plots, LaTeXStrings; pyplot()

alphas = [0.5, 1.5, 1]
lam = 2

lambda(dist::Weibull) = shape(dist)*scale(dist)^(-shape(dist))
theta(lam,alpha) = (alpha/lam)^(1/alpha)

dists = [Weibull.(a,theta(lam,a)) for a in alphas]

hA(dist,x) = pdf(dist,x)/ccdf(dist,x)
hB(dist,x) = lambda(dist)*x^(shape(dist)-1)

xGrid = 0.01:0.01:10
hazardsA = [hA.(d,xGrid) for d in dists]
hazardsB = [hB.(d,xGrid) for d in dists]

println("Maximum difference between two implementations of hazard: ",
    maximum(maximum.(hazardsA-hazardsB)))

Cl = [:blue :red :green]
Lb = [L"\lambda=" * string(lambda(d)) * ",   " * L"\alpha =" * string(shape(d))
        for d in dists]

plot(xGrid, hazardsA, c=Cl, label=reshape(Lb, 1,:), xlabel="x",
	ylabel="Instantaneous failure rate", xlims=(0,10), ylims=(0,10))


using SchumakerSpline
tol = 10*eps()

x = [1,2,3,4,5,6,7,8,9,10,11,12]
y = log.(x) + sqrt.(x)

typeof(x) == Vector{Int}

spline = Schumaker(x,y)
for i in 1:length(x)
    abs(evaluate(spline, x[i]) - y[2]) < tol
end
# Evaluation with a Float64.
evaluate(spline, 11.5)

# Testing second derivatives
xArray = range(1, stop=6, length=1000)
second_derivatives = evaluate.(spline, xArray,2)
maximum(second_derivatives) < tol

# Testing Integrals
analytic_integral(lhs,rhs) = rhs*log(rhs) - rhs + (2/3) * rhs^(3/2) - ( lhs*log(lhs) - lhs + (2/3) * lhs^(3/2) )
lhs = 2.0
rhs = 2.5
numerical_integral = evaluate_integral(spline, lhs,rhs)
abs(analytic_integral(lhs,rhs) - numerical_integral) < 0.01

lhs = 2.1
rhs = 2.11
numerical_integral = evaluate_integral(spline, lhs,rhs)
abs(analytic_integral(lhs,rhs) - numerical_integral) < 0.01

lhs = 1
rhs = 4
numerical_integral = evaluate_integral(spline, lhs,rhs)
abs(analytic_integral(lhs,rhs) - numerical_integral) < 0.1


# Minimally adapted from https://github.com/scipy/scipy/blob/b5d8bab88af61d61de09641243848df63380a67f/scipy/special/_spherical_bessel.pxd#L74
function sphericalbesselj(n::Integer, x::T) where {T<:AbstractFloat}
    isnan(x) && return x
    n < 0 && throw(DomainError(n, "n must be non-negative"))
    isinf(x) && return zero(T)
    iszero(x) && return iszero(n) ? one(T) : zero(T)

    if n > 0 && n ≥ x
        # We don't use SpecialFunctions.sphericalbesselj directly because it errors out on large x
        return T(SpecialFunctions.sphericalbesselj(n, x))
    else
        return _sphericalbesselj(n, x)
    end
end

function _sphericalbesselj(n::Integer, x::T) where {T<:AbstractFloat}
    invx = inv(x)
    s0 = sin(x) * invx
    iszero(n) && return s0

    s1 = (s0 - cos(x)) * invx
    isone(n) && return s1

    sn = zero(T)
    for idx in 2:n
        sn = (2idx - 1) * invx * s1 - s0
        s0 = s1
        s1 = sn
    end
    return sn
end

# Adapted from https://github.com/numpy/numpy/blob/4adc87dff15a247e417d50f10cc4def8e1c17a03/numpy/polynomial/legendre.py#L832-L914
function legval(x, c::AbstractVector)
    nd = length(c)
    nd ≥ 2 || return last(c)

    c0, c1 = @inbounds c[nd - 1], c[nd]
    @inbounds for j in (nd - 2):-1:1
        k = j / (j + 1)
        c0, c1 = c[j] - c1 * k, c0 + c1 * x * (k + 1)
    end
    return c0 + c1 * x
end

# Adapted from https://github.com/numpy/numpy/blob/4adc87dff15a247e417d50f10cc4def8e1c17a03/numpy/polynomial/legendre.py#L1126-L1176
"""
legvander(x, deg)

Pseudo-Vandermonde matrix of degree `deg`.
"""
function legvander(x::AbstractVector{T}, deg::Integer) where {T}
    deg ≥ 0 || throw(DomainError(deg, "degree needs to be non-negative"))

    vsize = (length(x), deg + 1)
    v = Matrix{T}(undef, vsize...)

    # Use forward recursion to generate the entries. This is not as accurate
    # as reverse recursion in this application but it is more efficient.
    @inbounds begin
        v[:, 1] .= one(T)
        if deg > 0
            v[:, 2] .= x
            for i in 2:deg
                invi = inv(i)
                @views @. v[:, i + 1] = v[:, i] * x * (2 - invi) - v[:, i - 1] * (1 - invi)
            end
        end
    end

    return v
end

# Adapted from https://github.com/numpy/numpy/blob/4adc87dff15a247e417d50f10cc4def8e1c17a03/numpy/polynomial/legendre.py#L612-L701
"""
    legder
"""
function legder(c::AbstractMatrix{T}) where {T}
    c = copy(c)
    n, m = size(c)
    n -= 1
    der = Matrix{T}(undef, n, m)
    @views @inbounds for j in n:-1:2
        @. der[j, :] = (2j - 1) * c[j + 1, :]
        @. c[j - 1, :] += c[j + 1, :]
    end
    @views @inbounds @. der[1, :] = c[2, :]
    return der
end


"""
AXLeastSquares(y, X)

An Ordinary Least Squares implementation.
"""

struct AXLeastSquares <: AXEstimator
    β::Array{Float64} # Coefficient
    y::Matrix{Float64} # Response
    X::Matrix{Float64} # Covariates

    # Constructor Function
    function AXLeastSquares(y::Matrix{Float64}, X::Matrix{Float64})
        β = inv(X' * X) * (X' * y)
        new(β, y, X)
    end
end #AXLeastSquares

# Methods ======================================================================
"""
infer(fit::AXLeastSquares; heteroskedastic, print_df)

A method to calculate standard errors of an AXLeastSquares object.
"""
function infer(fit::AXLeastSquares; heteroskedastic::Bool=false, print_df::Bool=true)
    # Retrieve necessary parameters
    N = length(fit.y)
    K = size(fit.X, 2) # number of covariates

    # Calculate the variance-covariance matrix
    resid  = fit.y - predict(fit)
    XX     = fit.X' * fit.X
    residX = fit.X .* resid
    if heteroskedastic
        vcv = inv(XX) * (residX' * residX) * inv(XX) .* (N / (N-K))
    else
        # homoskedasticity
        vcv = inv(XX) .* (resid' * resid) ./ (N-K)
    end

    se = sqrt.(diag(vcv))
    t_stat = fit.β ./ se
    p_val = 2 * cdf.(Normal(), -abs.(t_stat))
    r2 = 1 - sum(resid.^2)/sum((fit.y.-mean(fit.y)).^2)

    # Print estimates
    if print_df
        out_df = DataFrame(hcat(fit.β, se, t_stat, p_val, r2), :auto)
        rename!(out_df, ["coef", "se", "t-stat", "p-val", "R-sqaure"])
        display(out_df)
    end

    # Organize and return output
    output = (β = fit.β, se = se, t = t_stat, p = p_val, r = r2)
    return output
end #infer.AXLeastSquares


"""
    idct2!(x, fft_type::FFTType = Radix2Type())

Inverse discrete cosine transform of type II.

X = IDCT2(Y) inverts the DCT2 transform, returning the
original vector if Y was obtained using Y = DCT2(X). It is based on
the staggered-grid definition
                          N
   X(j) = 2/N*( Y(1)/2 +  ∑ Y(k) cos (π*(k-1)*(j-1/2)/N) )
                         k=2

X = IDCT2(Y,N) pads or truncates the vector Y to length N
before transforming.

If Y is a matrix, the IDCT2 operation is applied to
each column.

                          N
   X(j) = 2/N*( Y(1)/2 +  ∑ Y(k) cos (π*(k-1)*(j-1/2)/N) )
                         k=2
"""
function idct2!(x, fft_type=Radix2Type())
    x .= idct2(x, fft_type)
end
function idct2(a, fft_type::FFTType=Radix2Type())
    FT = eltype(a)
    if min(size(a)...)==1
        if size(a,2)>1
            do_trans = true;
        else
            do_trans = false;
        end
        a = a[:];
    else
        do_trans = false;
    end
    n = size(a,1);
    m = size(a,2);
    # Pad or truncate a if necessary
    if size(a,1)<n
      aa = zeros(n,m);
      aa[1:size(a,1),:] .= a;
    else
      aa = a[1:n,:];
    end
    y = Array{Complex{FT},2}(undef, 2*n, m)
    fill!(y, 0)
    y[1:n,:] .= aa;
    i = Complex(0,1);
    e = 0.5 * exp.(-i*FT(0.5) * π * (0:n-1) ./ n);
    for l = 1:m
        y[1:n,l] .= y[1:n,l] .* e;
    end
    y[n+2:2*n,:] .= reverse(conj.(y[2:n,:]); dims=1);
    yy = fft(y, 1);
    b = (2/n) * yy[1:n,:];
    if isreal(a)
        b = real.(b)
    end
    if do_trans
        b = b'
    end
    return b
end


#=
   From Allen Downey
   https://twitter.com/AllenDowney/status/1063460117029535746
   """
   Today's Bayesian problem of the week: Suppose we visit a wild animal preserve where we
   know that the only animals are lions and tigers and bears, but we don't know how
   many of each there are.

   During the tour, we see 3 lions, 2 tigers, and 1 bear. Assuming that every animal had an equal
   chance to appear in our sample, estimate the prevalence of each species.

   What is the probability that the next animal we see is a bear?
   """

   Also see: https://towardsdatascience.com/estimating-probabilities-with-bayesian-modeling-in-python-7144be007815
    From https://towardsdatascience.com/estimating-probabilities-with-bayesian-modeling-in-python-7144be007815
    """
    Species: lions    Prevalence: 44.44%.
    Species: tigers   Prevalence: 33.33%.
    Species: bears    Prevalence: 22.22%.
    """

   Distributions of variable animal7
   lion       =>    4526  (0.452600)
   tiger      =>    3251  (0.325100)
   bear       =>    2223  (0.222300)

   Cf ~/blog/lions_tigers_and_bears.blog
      ~/webppl/lions_tigers_and_bears.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function lions_tigers_and_bears()
    lion = 1
    tiger = 2
    bear = 3

    # Prior
    # probLion  ~ Beta(1,1)
    # probTiger ~ Beta(1,1)
    # probBear  ~ Beta(1,1)

    # The Dirichlet distribution ensures that the sum of probabilities is 1
    # i.e. we don't have to ensure this via some specific constraint.
    v = [3, 2, 1] # We saw three lions, two tigers, and one bear
    # v = [2,2,2]
    x ~ Dirichlet(v)

    # The probabilities to calculate ("aliased" for simplicity)
    probLion  ~ Dirac(x[1])
    probTiger ~ Dirac(x[2])
    probBear  ~ Dirac(x[3])

    n = 7
    o = tzeros(n)
    for i in 1:n
        o[i] ~ Categorical([probLion,probTiger,probBear])
    end

    # It shouldn't matter in what order we see the different animals.
    true ~ Dirac(o[1] == lion)
    true ~ Dirac(o[2] == lion)
    true ~ Dirac(o[3] == lion)
    true ~ Dirac(o[4] == tiger)
    true ~ Dirac(o[5] == tiger)
    true ~ Dirac(o[6] == bear)

    animal7 ~ Dirac(o[7])

    # probbear7 ~ Dirac(o[7] == bear)
    # problion7 ~ Dirac(o[7] == lion)
    # probtiger7 ~ Dirac(o[7] == tiger)

end

model = lions_tigers_and_bears()
num_chns = 4

# chns = sample(model, Prior(), 10_000)

# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(15), 10_000)
# chns = sample(model, SMC(1000), 10_000)
# chns = sample(model, IS(), 10_000)

#
display(chns)

show_var_dist_pct(chns,:animal7,["lion","tiger","bear"])


123 == 123

123 == 100 + 23

123 == -123

typeof(true)

typeof(false)

typeof(45 == 45)

23 ≤ -5

4 ≠ 44

1.0 == 1

1.000000000001 == 1

3 < 4 && -15 > -10

3 < 4 || -15 > -10

!(-4 ≥ 4) && (1 > 2 || 1 < 2)

1 < 2 <= 2 < 3 == 3 > 2 >= 1 == 1 < 3 != 5

x = -2.3
-3 ≤ x ≤ 4

# Examples of if-else statements

function check_x(x)
    if 0 ≤ x ≤ 10
        println("x is between 0 and 10")
    elseif !(x < 0)
        println("x is not between 0 and 10, and x is not negative (meaning x > 10)")
    else
        println("x is not between 0 and 10, and x is not not negative (meaning x is negative)")
    end
end

check_x(π)

check_x(10.000001)

check_x(2^63)        # trick question - causes Int64 overflow

x = -3
if x ≥ 0 && sqrt(x) < 10
    # do something
else
    println("Sorry, need a non-negative number with square root strictly less than 10")
end

function fibonacci(n)
    # First special cases
    if n == 0
        return 0
    elseif n == 1
        return 1
    else
        # General case n ≥ 2 - initialize and iterate
        Fkold = 0
        Fk = 1
        for k = 2:n
            Fknew = Fk + Fkold
            # Swap previous values
            Fkold = Fk
            Fk = Fknew
        end
        return Fk
    end
end

fibonacci(10)

fibonacci(1000)   # Again be careful with overflow of Int64, but this one is OK

function PE1(n)
    s = 0
    for i = 1:n-1
        if i % 3 == 0 || i % 5 == 0
            s += i
        end
    end
    s
end

PE1(10)

PE1(1000)


using Plots
using LinearAlgebra

"
This code shows how errors for different finite difference methods converge as h->0.
"

# define forward, centered differences, D3 = some other difference approx
D1(u,x,h) = (u(x+h) - u(x))/h
D2(u,x,h) = (u(x+h) - u(x-h))/(2*h)
D3(u,x,h) = (u(x-2*h) - 6*u(x-h) + 3*u(x) + 2*u(x+h))/(6*h)

include("fdcoeffV.jl")
x = -2:2
a = fdcoeffV(1,x[3],x) # 1 = order of derivative, x̄, stencil points x
D4(u,x,h,c) = dot(c,(u(x-2*h),u(x-h),u(x),u(x+h),u(x+2*h)))/h

# function, exact derivative
u(x) = x^2 + sin(1+pi*x)
dudx(x) = 2*x + pi*cos(1+pi*x)
x = 0.0 # point at which we approximate the derivative

# measure how error changes as h decreases
hvec = 2.0 .^(-(0:8))
err = (zeros(size(hvec)),zeros(size(hvec)),zeros(size(hvec)),zeros(size(hvec)))
for (i,h) in enumerate(hvec)
    err[1][i] = abs(dudx(x) - D1(u,x,h))
    err[2][i] = abs(dudx(x) - D2(u,x,h))
    err[3][i] = abs(dudx(x) - D3(u,x,h))
    err[4][i] = abs(dudx(x) - D4(u,x,h,a))
end

plot(xaxis=:log,yaxis=:log,legend=:bottomright)
labels = ("Forward diff","Centered diff","Some other 4-point diff","5-point centered diff")
for i = 1:4
    plot!(hvec,err[i],marker=:dot,label=labels[i])
    plot!(hvec,hvec.^i,linestyle=:dash,label="Order $(i)")
end
display(plot!())


using Plots
gr(titlefontsize=12,size=(500,300),transpose=true,colorbar=false)

using FourierGPE
# Units
# this example works in oscillator units

## Initialize simulation

L = (40.0,40.0)
N = (256,256)
sim = Sim(L,N)
@unpack_Sim sim

## simulation parameters
μ = 15.0

# potential
import FourierGPE.V
V(x,y,t) = 0.5*(x^2 + y^2)

# TF state
ψtf(x,y,μ,g) = sqrt(μ/g)*sqrt(max(1.0-V(x,y,0.0)/μ,0.0)+im*0.0)

# make initial state
x,y = X
ψi = ψtf.(x,y',μ,g)
ϕi = kspace(ψi,sim)

@pack_Sim! sim

## evolve
sol = runsim(sim)

## pull out the ground state
ϕg = sol[end]
ψg = xspace(ϕg,sim)
showpsi(x,y,ψg)

R(w) = sqrt(2*μ/w^2)
R(1)
plot(x, abs2.(ψtf.(x,0.,μ,g)))
plot!(x, abs2.(ψg[:,128]))

## free expansion
V(x,y,t) = 0.0

# sim parameters
γ = 0.0
tf = 10.0
t = LinRange(ti,tf,Nt)
ϕi = ϕg

sime = Sim(L,N)
@pack! sime = tf,t,γ,ϕi

## evolve
sole = runsim(sime)

ϕf = sole[end]
ψf = xspace(ϕf,sime)
showpsi(x,y,ψf)

anim = @animate for i=1:Nt-6
    ψ = xspace(sole[i],sime)
    showpsi(x,y,ψ)
end

gif(anim,"./examples/expand.gif",fps=30)


"""
    first_po(n::Int64, k::Int64) -> po_idx::Vector{Int64}

Notation-Generator: Calculate the first generator matrix in PO notation (padding ones indices) by n and k.
"""
function first_po(n::Int64, k::Int64)::Vector{Int64}
    return collect(1:(n-k))
end

"""
    next_po(current_po::Vector{Int64}, n::Int64, i=0) -> next::Vector{Int64}

Notation-Generator: Calculates the next generator matrix in PO notation (padding ones indices) in relation to a given generator matrix in PO notation.
"""
function next_po(current_po::Vector{Int64}, n::Int64, i=0)::Vector{Int64}
    if i > size(current_po,1)
        println("Error: i=$i")
        return nothing
    end
    next = copy(current_po)
    if current_po[end-i] + 1 > n - i
        return next_po(current_po, n, i+1)
    else
        next[end-i] += 1
        for j in (size(next,1) - i + 1):size(next,1)
           next[j] = next[j-1] +1
        end
        return next
    end
end

"""
    construct_generators(n::Int64, k::Int64, prnt = false) -> Ω::Int64, all_pos, all_generating_matrices

Matrix- and Notation-Generator: Calculate all generating matrices (including PO notations) and the number of generating matrices Ω.
"""
function construct_generators(n::Int64, k::Int64, prnt = false)
    # calculate number of possible generating matrices
    Ω = binomial(n,k)
    println("There are Ω = ($n choose $k) = $Ω  possible basic permutations.")

    # Generate all possible padding ones positions first
    # number of padding ones
    num_pos = n - k
    # save first PO
    all_pos = []
    push!(all_pos, first_po(n,k))
    # save first PO as matrix
    all_generating_matrices = SparseMatrixCSC{ℤ₂, Int64}[]
    push!(all_generating_matrices, po2matrix(all_pos[1],n))

    if prnt
        println(
            "1:\t ", all_pos[1],
            "\t↔ AI:\t ", po2ai(all_pos[1],n),
            "\tA² Test: ", multiplyZ2(all_generating_matrices[1],all_generating_matrices[1]) == id(n),
            ": ", all_generating_matrices[1], "\n\n")
    end

    # construct all other generating matrices
    for i in 2:Ω
        next = next_po(all_pos[end],n)
        push!(all_pos, next)
        push!(all_generating_matrices, po2matrix(all_pos[end],n))
        if prnt
            println(
                "$i:\t", next,
                "\t↔ AI:\t ", po2ai(all_pos[end],n),
                "\tA² Test: ", multiplyZ2(all_generating_matrices[end],all_generating_matrices[end]) == id(n),
                ": ", all_generating_matrices[end], "\n\n")
        end
    end

    # return all necessary variables
    return Ω, all_pos, all_generating_matrices
end


# # Lasso, Ridge and Elastic Net Regressions
#
# This notebook presents a simple implementation of Lasso and elastic net regressions.

#-

# ## Load Packages and Extra Functions

using DelimitedFiles, LinearAlgebra, Statistics, Plots, Convex, SCS

import MathOptInterface
const MOI = MathOptInterface

# # Loading Data
#
# We use the diabetes data from Efron et al, downloaded from https://web.stanford.edu/~hastie/StatLearnSparsity_files/DATA/diabetes.html and then converted from a tab to a comma delimited file.
#
# All data series are standardised (see below) to have zero means and unit standard deviation, which improves the numerical stability. (Efron et al do not standardise the scale of the response variable.)

(x,header) = readdlm("aux_files/diabetes.csv",',',header=true)
#display(header)
#display(x)

x = (x .- mean(x,dims=1))./std(x,dims=1)          #standardise

(Y,X) = (x[:,end],x[:,1:end-1]);                  #to get traditional names
xNames = header[1:end-1];

# # Lasso, Ridge and Elastic Net Regressions
#
# (a)  The regression is $Y = Xb + u$,
# where $Y$ and $u$ are $T \times 1$, $X$ is $T \times K$, and $b$ is the $K$-vector of regression coefficients.
#
# (b) We want to minimize $(Y-Xb)'(Y-Xb) + \gamma \sum |b_i| + \lambda \sum b_i^2$.
#
# (c) We can equally well minimise $b'Qb - 2c'b + \gamma \sum |b_i| + \lambda \sum b_i^2$,
# where $Q = X'X$ and $c=X'Y$
#
# (d) Lasso: $\gamma>0,\lambda=0$; Ridge: $\gamma=0,\lambda>0$; elastic net: $\gamma>0,\lambda>0$.

"""
    LassoEN(Y,X,γ,λ)

Do Lasso (set γ>0,λ=0), ridge (set γ=0,λ>0) or elastic net regression (set γ>0,λ>0).

## Input
- `Y::Vector`:     T-vector with the response (dependent) variable
- `Y::VecOrMat`:   TxK matrix of covariates (regressors)
- `γ::Number`:     penalty on sum(abs.(b))
- `λ::Number`:     penalty on sum(b.^2)

"""
function LassoEN(Y,X,γ,λ=0.0)

    K = size(X,2)

    b_ls = X\Y                    #LS estimate of weights, no restrictions

    Q  = X'X
    c  = X'Y                      #c'b = Y'X*b

    b  = Variable(K)              #define variables to optimize over
    L1 = quadform(b,Q)            #b'Q*b
    L2 = dot(c,b)                 #c'b
    L3 = norm(b,1)                #sum(|b|)
    L4 = sumsquares(b)            #sum(b^2)

    Sol = minimize(L1-2*L2+γ*L3+λ*L4)      #u'u + γ*sum(|b|) + λsum(b^2), where u = Y-Xb
    solve!(Sol,()->SCS.Optimizer(verbose = false))
    Sol.status == MOI.OPTIMAL ? b_i = vec(evaluate(b)) : b_i = NaN

    return b_i, b_ls

end

# The next cell makes a Lasso regression for a single value of γ.

K = size(X,2)
γ = 100

(b,b_ls) = LassoEN(Y,X,γ)

println("OLS and Lasso coeffs (with γ=$γ)")
display([["" "OLS" "Lasso"];xNames b_ls b])

# # Redo the Lasso Regression with Different Gamma Values
#
#
# We now loop over $\gamma$ values.
#
# Remark: it would be quicker to put this loop inside the `LassoEN()` function so as to not recreate `L1`-`L4`.

nγ = 101
γM = range(0; stop=600, length=nγ)             #different γ values

bLasso   = fill(NaN,size(X,2),nγ)       #results for γM[i] are in bLasso[:,i]
for i = 1:nγ
    bLasso[:,i], = LassoEN(Y,X,γM[i])
end

#-

plot(log.(γM),bLasso',
     title = "Lasso regression coefficients",
     xlabel = "log(γ)",
     label = permutedims(xNames),
     size = (600,400))

# # Ridge Regression
#
# We use the same function to do a ridge regression. Alternatively, do `b = inv(X'X + λ*I)*X'Y`.

nλ = 101
λM = range(0; stop=3000, length=nλ)

bRidge  = fill(NaN,size(X,2),nλ)
for i = 1:nλ
    bRidge[:,i], = LassoEN(Y,X,0,λM[i])
end

#-

plot(log.(λM),bRidge',
     title = "Ridge regression coefficients",
     xlabel = "log(λ)",
     label = permutedims(xNames),
     size = (600,400))

# # Elastic Net Regression

λ = 200
println("redo the Lasso regression, but with λ=$λ: an elastic net regression")

bEN   = fill(NaN,size(X,2),nγ)
for i = 1:nγ
    bEN[:,i], = LassoEN(Y,X,γM[i],λ)
end

#-

plot(log.(γM),bEN',
     title = "Elastic Net regression coefficients",
     xlabel = "log(γ)",
     label = permutedims(xNames),
     size = (600,400))
#-


#=
  Port of WebPPL code from https://github.com/probmods/ppaml2016/blob/gh-pages/chapters/5-election.md
  """
  Basic model
  Learning a state-wide preference from poll data
  """

  Clinton: 1 Trump: 2
  Distributions of variable (num:0)
  2.00000 =>    8750  (0.875000)
  1.00000 =>    1250  (0.125000)

  According to the WebPPL model, the probability that Clinton wins (won) is 0.12207
=#

using Turing
include("jl_utils.jl")

@model function sampleElection()
    function trueStatePref()
        pref ~ Beta(1,1)
        counts = Dict(:trump=>304, :clinton=>276)
        total = sum(values(counts))
        c ~ Binomial(total, pref)
        true ~ Dirac(c == counts[:clinton])
        return pref
    end
    clinton = 0
    trump = 1
    pref = trueStatePref()
    turnout = 2400000
    clintonVotes ~ Binomial(turnout,pref)
    trumpVotes ~ Dirac(turnout - clintonVotes)
    winner ~ Dirac(clintonVotes > trumpVotes ? clinton : trump)

    clintonWins ~ Dirac(winner == clinton)
    trumpWins ~ Dirac(winner == trump)
    return winner
end

model = sampleElection()

# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(15), 10_000)
# chns = sample(model, IS(), 10_000)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, SMC(), MCMCThreads(), 10_000, 4)

display(chns)

chns_params = Turing.MCMCChains.get_sections(chns, :parameters)
genq = generated_quantities(model, chns_params)
println("Clinton: 0 Trump: 1")
show_var_dist_pct(genq)


# # Hovercraft Path Planning
# In this case study, we seek to determine an optimal control policy for the
# trajectory of a hovercraft that travels to a set of dynamic waypoints while
# trying to minimize the thrust input.

# ## Formulation

# The corresponding dynamic optimization problem is expressed:
# ```math
# \begin{aligned}
# 	&&\underset{x(t), v(t), u(t)}{\text{min}} &&& \int_{t \in T} |u(t)|_2^2 dt  \\
# 	&&\text{s.t.} &&& v(0) = v0\\
# 	&&&&& \frac{dx}{dt} = v(t), && t \in T\\
#     &&&&& \frac{dv}{dt} = u(t), && t \in T\\
#     &&&&& x(t_i) = xw_i, && i \in I
# \end{aligned}
# ```
# where ``x(t)`` is the Cartesian position, ``v(t)`` is the velocity, ``u(t)`` is
# the thrust input, ``xw_i, \ i \in I,`` are the waypoints, and ``T`` is the time
# horizon.

# ## Model Definition

# Let's implement this in `InfiniteOpt` and first import the packages we need:
using InfiniteOpt, Ipopt

# Next we'll specify our waypoint data:
xw = [1 4 6 1; 1 3 0 1] # positions
tw = [0, 25, 50, 60];    # times

# We initialize the infinite model and opt to use the Ipopt solver:
m = InfiniteModel(optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 0));

# Let's specify our infinite parameter which is time ``t \in [0, 60]``:
@infinite_parameter(m, t in [0, 60], num_supports = 61)

# Now let's specify the decision variables:
@variables(m, begin
    ## state variables
    x[1:2], Infinite(t)
    v[1:2], Infinite(t)
    ## control variables
    u[1:2], Infinite(t), (start = 0)
end)

# Specify the objective:
@objective(m, Min, ∫(u[1]^2 + u[2]^2, t))

# Set the initial conditions with respect to the velocity:
@constraint(m, [i = 1:2], v[i] == 0, DomainRestrictions(t => 0))

# Define the point physics ODEs which serve as our system model:
@constraint(m, [i = 1:2], ∂(x[i], t) == v[i])
@constraint(m, [i = 1:2], ∂(v[i], t) == u[i])

# Ensure we hit all the waypoints:
@constraint(m, [i = 1:2, j = eachindex(tw)], x[i] == xw[i, j], DomainRestrictions(t => tw[j]))

# ## Problem Solution

# Optimize the model:
optimize!(m)

# Extract the results:
x_opt = value.(x);

# Plot the results:
using Plots
scatter(xw[1,:], xw[2,:], label = "Waypoints", background_color = :transparent)
plot!(x_opt[1], x_opt[2], label = "Trajectory")
xlabel!("x_1")
ylabel!("x_2")

# That's it, now we have our optimal trajectory!

# ### Maintenance Tests
# These are here to ensure this example stays up to date.
using Test
@test termination_status(m) == MOI.LOCALLY_SOLVED
@test has_values(m)
@test x_opt isa Vector{<:Vector{<:Real}}


# Functions

# Point
#rand(rng::Random.AbstractRNG, ::Random.SamplerType{Point{T, N}}) where N where T = Point{T, N}(NTuple{N, T}(rand(rng, T, N)))
distance(p1::Point{T, N}, p2::Point{T, N}) where N where T = sqrt(sum((p1.coordinates .- p2.coordinates).^2))

function area(arr::Array{Point{T, 2}, 1}) where T
    length(arr) <= 2 && error("Area for two points cannot be computed!")
    area = 0.0
    x = 2
    y = 3
    flag = false
    while(x <= length(arr) && y <= length(arr))
        a = distance(arr[1], arr[x])
        b = distance(arr[x], arr[y])
        c = distance(arr[1], arr[y])
        s = (a + b + c)/2.0
        area += sqrt(s*(s - a)*(s - b)*(s - c))
        if flag == false
            flag = true
            x = y + 1
        else
            flag = false
            y = x + 1
        end
    end
    return area
end

function convexHull(points::Array{Point{T,2},1}) where T
    if length(points) < 3
        error("Minimum three points required")
    end
    n = length(points)
    function orientation(p::Point{T, 2}, q::Point{T, 2}, r::Point{T, 2}) where T <: Real
        v = (q.coordinates[2] - p.coordinates[2])*(r.coordinates[1] - q.coordinates[1])- (q.coordinates[1] - p.coordinates[1])*(r.coordinates[2] - q.coordinates[2])
        if v == 0 return 0 end
        return v > 0 ? 1 : 2
    end
    hull = Array{Point{T, 2}, 1}(undef, 0)
    l = 1
    for i in 2:n
        if points[i].coordinates[1] < points[l].coordinates[1]
            l = i
        end
    end
    p = l - 1
    q = -1
    while true
        push!(hull, points[p + 1])
        q = (p + 1) % n
        for i in 0:(n - 1)
            if(orientation(points[p + 1], points[i + 1], points[q+1]) == 2)
                q = i
            end
        end
        p = q
        if p == l - 1
            break
        end
    end
    return hull
end
# Line
function distance(p::Point{T, N}, l::Line{T, N}) where N where T
    t = dot(p.coordinates .- l.p.coordinates, l.v)/dot(l.v, l.v)
    return norm(p.coordinates .- l.p.coordinates .- t*l.v)
end
distance(l::Line{T, N}, p::Point{T, N}) where N where T = distance(p::Point{T, N}, l::Line{T, N})

function distance(l1::Line{T, N}, l2::Line{T, N}; err = 1e-10) where N where T
    if norm(l1.v .- l2.v) ≦ err
        return distance(l1.p, l2)
    else
        return zero(T)
    end
end

# Plane

function distance(p::Point{T, N}, q::Plane{T, N}) where N where T
    x = distance(p, q.p)
    y = distance(p, q.n)
    return sqrt(x^2 - y^2)
end

function boundingbox(a::Array{Point{T, N}, 1}) where N where T
    mn = Array{T,1}(undef, N)
    mn .= typemax(T)
    mx = Array{T,1}(undef, N)
    mx .= typemin(T)
    for i in 1:N, j in 1:length(a)
        mn[i] = min(mn[i], a[j].coordinates[i])
        mx[i] = max(mx[i], a[j].coordinates[i])
    end
    return (Point{T, N}(mn), Point{T, N}(mx))
end


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

A = Float64[
  1/2 1/3 1/4 1/5 1/6;
  1/3 1/4 1/5 1/6 1/7;
  1/4 1/5 1/6 1/7 1/8;
  1/5 1/6 1/7 1/8 1/9;
  1/6 1/7 1/8 1/9 1/10;
]
b = ones(5)

(Al, Au, Ap) = lu(A)

@show y = Al \ b[Ap]
println()

@show x = Au \ y
println()

@show As = sparse(A)
println()

@show As\b
println()

@show F = lu(As)
println()

@show F.L * F.U == F.Rs .* A[F.p, :]
println()

@show F.L*F.U
println()

@show sparse(F.Rs.* A[F.p, F.q])
println()

@show lu(A)
println()

end


"""
Polynomial smoothing with the Savitzky Golay filters. Adapted for julia >= 1.0.
Sources: https://github.com/BBN-Q/Qlab.jl/blob/master/src/SavitskyGolay.jl
         https://gist.github.com/lnacquaroli/c97fbc9a15488607e236b3472bcdf097
Requires LinearAlgebra and DSP modules loaded.
"""

# TODO: could probably improve performance with some @inbounds sprinkling
function savitzkyGolay(x::AbstractVector, windowSize::Int, polyOrder::Int; deriv::Int=0)

  isodd(windowSize) || throw("Window size must be an odd integer.")
  polyOrder < windowSize || throw("Polynomial order must me less than window size.")

  halfWindow = Int( ceil((windowSize-1)/2) )

  # Setup the S matrix of basis vectors
  S = zeros.(windowSize, polyOrder+1)
  for ct = 0:polyOrder
    S[:,ct+1] = (-halfWindow:halfWindow).^(ct)
  end

  ## Compute the filter coefficients for all orders

  # From the scipy code it seems pinv(S) and taking rows should be enough
  G = S * pinv(S' * S)

  # Slice out the derivative order we want
  filterCoeffs = G[:, deriv+1] * factorial(deriv)

  # Pad the signal with the endpoints and convolve with filter
  paddedX = vcat(x[1]*ones(halfWindow), x, x[end]*ones(halfWindow))
  y = conv(filterCoeffs[end:-1:1], paddedX)

  # Return the valid midsection
  return y[2*halfWindow+1:end-2*halfWindow]

end


using Random, Plots, StatsPlots, Distributions, DataFrames, Optim
# Bayesian Inference for Bernoulli parameter (inference success probability)
# Define HPD (highest posterior density) interval
# HPD are difficult to calculate and not be used nowadays.
# Maybe this? https://github.com/tpapp/HighestDensityRegions.jl
"""
HPD interval for beta dist
Optim.jl
http://julianlsolvers.github.io/Optim.jl/v0.9.3/user/tipsandtricks/#dealing-with-constant-parameters

    Inputs
    ------
    ci0:  initial value for calculating HPD interval numerically
    α:    parameter 1 of Beta dist
    β:    parameter 2 of Beta dist
    prob: probabilty of HPD interval

    Outputs
    -------
    HPD interval

"""
function beta_hpdi(ci0, α, β, prob)
    """
    Continuous equation for HPD interval

        Inputs
        ------
        v: HPD interval (2-element Vector)
        a: parameter 1 of Beta dist
        b: parameter 2 of Beta dist
        p: probabilty of HPD interval

        Outputs
        -------
        HPD interval condition equations (2-element Vector)

    """
    function hpdi_conditions(v, a, b, p)
        # HPD interval probability = p
        eq1 = cdf(Beta(a, b), v[2]) - cdf(Beta(a, b), v[1]) - p
        # p(lower|D) = p(upper|D)
        eq2 = pdf(Beta(a, b), v[2]) - pdf(Beta(a, b), v[1])
        return [eq1, eq2]
    end
    return nothing
end

function hpdi_conditions(v, a, b, p)
    # HPD interval probability = p
    eq1 = cdf(Beta(a, b), v[2]) - cdf(Beta(a, b), v[1]) - p
    # p(lower|D) = p(upper|D)
    eq2 = pdf(Beta(a, b), v[2]) - pdf(Beta(a, b), v[1])
    return [eq1, eq2]
end

function hpdi_conditions(v)
    # HPD interval probability = p
    eq1 = cdf(Beta(14, 38), v[2]) - cdf(Beta(14, 38), v[1]) - 0.95
    # p(lower|D) = p(upper|D)
    eq2 = pdf(Beta(14, 38), v[2]) - pdf(Beta(14, 38), v[1])
    return [eq1, eq2]
end

# Optimization
hpdi_conditions([0.17, 0.41])
Optim.optimize(hpdi_conditions, [0.17, 0.41])
hpdi_conditions([0, 0])
Optim.optimize(hpdi_conditions, [0, 0])
NLsolve.nlsolve(hpdi_conditions)

# Define calculation of posterior for Bernoulli parameter
"""
Posterior Distribution (Beta dist) for Bernoulli parameter π

    Inputs
    ------
    data: data (0 or 1)
    a0: parameter 1 of prior (Beta dist)
    b0: parameter 2 of prior (Beta dist)
    prob: credible interval probability (0 < prob < 1)

    Outputs
    -------
    results: posterior summary stats table (DataFrames)
    a: parameter 1 of posterior (Beta dist)
    b: parameter 2 of posterior (Beta dist)

"""
function bernoulli_stats(data, a0, b0, prob)
    # posterior
    n = size(data)[1]
    sum_data = sum(data)
    a = sum_data + a0
    b = n - sum_data + b0
    d_posterior = Beta(a, b)
    # mean
    mean_π = mean(d_posterior)
    # median
    median_π = median(d_posterior)
    # mode
    mode_π = mode(d_posterior)
    # Analytic Mode for Beta Dist
    # (a - 1.0) / (a + b - 2.0)
    # std
    sd_π = std(d_posterior)
    # credible interval
    # ci upper boud
    ci_π_upper = quantile(d_posterior, prob - (1-prob)/2)
    # ci lower boud
    ci_π_lower = quantile(d_posterior, (1-prob)/2)
    ci_π = [ci_π_lower, ci_π_upper]

    posterior_stats = reshape([mean_π, median_π, mode_π, sd_π, ci_π], 1, 5)
    stats_string = ["Mean", "Median", "Mode", "SD", "CI"]
    results = DataFrame(posterior_stats, stats_string)
    return results, a, b
end

# Generate moc data from Bernoulli
p = 0.25
n = 50
Random.seed!(99)
data = rand(Bernoulli(p), n)

# Set prior
a0 = 1.0
b0 = 1.0

# Calculate posterior statistics
prob = 0.95
results, a, b = bernoulli_stats(data, a0, b0, prob)
print(results)

# Visualize Posterior
plot(Beta(a0, b0), xlims=(0, 1), ylims=(0, 7), xlabel="Success Probability q", ylabel="Probability Density", label="prior: Beta(1, 1)", linestyle=:dash)
plot!(Beta(a, b), label="posterior: Beta($a, $b)", linestyle=:solid)

# Plot posterior by different prior
Random.seed!(99)
data = rand(Bernoulli(p), 250) # more data
value_size = [10, 50, 250]
value_a0 = [1.0, 6.0]
value_b0 = [1.0, 4.0]
styles = [:dot, :dashdot, :dash, :solid]

l = @layout [grid(1, 2)]
plots = []
for index = 1:2
    style_index = 1
    a0_i = value_a0[index]
    b0_i = value_b0[index]
    pl = plot(Beta(a0_i, b0_i), linestyle=styles[style_index], label="prior: Beta($a0_i, $b0_i)", color="cornflower blue",
             xlims=(0, 1), ylims=(0, 15.5), xlabel="Success Probability q", ylabel="Probability Density")
    for n_j in value_size
        style_index += 1
        sum_data = sum(data[1:n_j])
        a_j = sum_data + a0_i
        b_j = n_j - sum_data + b0_i
        plot!(Beta(a_j, b_j), linestyle=styles[style_index], label="posterior: Beta($a_j, $b_j)", color="cornflower blue")
    end
    push!(plots, pl)
end
plot(plots..., layout=l)

# Animation
# https://stackoverflow.com/questions/55794068/animating-subplots-using-plots-jl-efficiently
Random.seed!(99)
data = rand(Bernoulli(p), 250) # more data
value_size = Vector(1:1:100)
value_a0 = [1.0, 6.0]
value_b0 = [1.0, 4.0]

gr(fmt = :png)
anim = @animate for n_j in value_size
    sum_data = sum(data[1:n_j])
    a_j_1 = sum_data + value_a0[1]
    b_j_1 = n_j - sum_data + value_b0[1]
    # all_p[1][1][:z] = Beta(a_j, b_j)
    p1 = plot(Beta(a_j_1, b_j_1), color="cornflower blue", label="posterior: Beta($a_j_1, $b_j_1)", xlims=(0, 1), ylims=(0, 15.5), xlabel="Success Probability q", ylabel="Probability Density", title="n=$n_j")
    a0_1 = value_a0[1]
    b0_1 = value_b0[1]
    plot!(Beta(a0_1, b0_1), linestyle=:dot, label="prior: Beta($a0_1, $b0_1)")

    a_j_2 = sum_data + value_a0[2]
    b_j_2 = n_j - sum_data + value_b0[2]
    # all_p[2][1][:z] = Beta(a_j, b_j)
    p2 = plot(Beta(a_j_2, b_j_2), color="cornflower blue", label="posterior: Beta($a_j_2, $b_j_2)", xlims=(0, 1), ylims=(0, 15.5), xlabel="Success Probability q", ylabel="Probability Density", title="n=$n_j")
    a0_2 = value_a0[2]
    b0_2 = value_b0[2]
    plot!(Beta(a0_2, b0_2), linestyle=:dot, label="prior: Beta($a0_2, $b0_2)")

    layout = @layout [grid(1, 2)]
    plot(p1, p2; layout)
end
gif(anim, "posterior.gif", fps=10)


"""
    AMHCopula{d,T}

Fields:
  - θ::Real - parameter

Constructor

    AMHCopula(d, θ)

The [AMH](https://en.wikipedia.org/wiki/Copula_(probability_theory)#Most_important_Archimedean_copulas) copula in dimension ``d`` is parameterized by ``\\theta \\in [0,1)``. It is an Archimedean copula with generator :

```math
\\phi(t) = 1 - \\frac{1-\\theta}{e^{-t}-\\theta}
```
"""
struct AMHCopula{d,T} <: ArchimedeanCopula{d}
    θ::T
end
AMHCopula(d,θ) = (0 <= θ <= 1) ? AMHCopula{d,typeof(θ)}(θ) : @error "Theta must be between 0 and 1"
ϕ(  C::AMHCopula,t) = (1-C.θ)/(exp(t)-C.θ)
ϕ⁻¹(  C::AMHCopula,t) = log(C.θ + (1-C.θ)/t)

τ(C::AMHCopula) = 1 - 2(C.θ+(1-C.θ)^2*log(1-C.θ))/(3C.θ^2) # no closed form inverse...
function τ⁻¹(::Type{AMHCopula},τ)
    if τ == zero(τ)
        return τ
    end
    if τ > 1/3
        @warn "AMHCopula cannot handle kendall tau's greater than 1/3. We capped it to 1/3."
        return 1
    end
    return Roots.fzero(θ -> 1 - 2(θ+(1-θ)^2*log(1-θ))/(3θ^2) - τ,0.5)
end

radial_dist(C::AMHCopula) = 1 + Distributions.Geometric(1-C.θ)


(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module SequencesDemo

using IntegerSequences

println("\nProduct 1*2*3")
p = ∏([1, 2, 3])
println(p)

println("\nList of first 10000 Kolakoski numbers (timing):")
@time KolakoskiList(10000)

println("\nList of first 10000 partition numbers (timing):")
@time PartitionNumberList(10000)

println("\nList of first 10000 Clausen numbers (timing):")
@time ClausenNumberList(10000)

println("\nList of first 10000 Ramanujan tau numbers (timing):")
@time RamanujanTauList(10000)

println("\nRamanujan tau of n = 10000:")
r = RamanujanTau(10000)
println(r)

println("\nAll partitions of 1, 2, 3, 4 and 5.")
for n in 1:5 Partition(n); println() end

println("\nNumber of points in square lattice on the
circle of radius √n.")
SeqShow(L004018(12))

println("\nThe product of the prime numbers dividing n.")
println(typeof([Radical(n) for n in 0:9]))
SeqShow([Radical(n) for n in 0:9])

println("\nOr:")
A = ZArray(9, Radical)
SeqShow(A)

println("\n=============")

println("\nThe exponential transform of Pascal's triangle.")
T = T055883(8)
ShowAsΔ(T)

println("\nThe row sums of the triangle above.")
Println(RowSums(T))

println("\nThe factorial of 100.")
f = F!(100)
println(f)

println("\nThe rising factorial ↑(20, 80).") # (a.k.a. Pochhammer)
rf = 20 ↑ 80
println(rf)

println("\nThe falling factorial ↓(80, 60).")
ff = 80 ↓ 60
println(ff)

println("\nThe number of divisors of n = 25920.")
t = τ(25920)
println(t)

println("\nThe sum of divisors of n = 25920.")
s = σ(25920)
println(s)

println("\nNumber of acyclic orientations of the
Turán graph T(2n, n) for n = 20.")
g = V033815(20)
println(g)

println("\nThe binomial function defined for all fmpz.
The same way as Maple and Mathematica compute the binomial.
First the familiar Pascal case:")
for n in 0:8
    for k in 0:n print(lpad(Binomial(n, k), 4)) end
    println()
end

println("\nAn extended region on the 2-dim lattice:")
for n in -5:5
    for k in -5:5
        print(lpad(Binomial(n, k), 5))
    end
    println()
end

println("\nKnuth, Graham and Patashnik write in
Concrete Mathematics: Hear us, O mathematicians
of the world! Let us not wait any longer! We can
make many formulas clearer by defining a new notation
now! Let us agree to write m ⊥ n, and to say m is
prime to n, if m and n are relatively prime.\n")
for n in 1:6
    for m in 1:6 print(⊥(m, n), " " ) end
    println()
end

println("\nUsing conversion to integer this can also be written as:\n")
for n in 1:6
    for m in 1:6 print(Int(⊥(m, n)), " " ) end
    println()
end

# Generating the OLMS logo.
# ⍊(n, k) means k is strong prime to n.
# n is strong prime to k iff n is prime to k and n does not divide k-1.
# Insert the line sum(L) == n - 5 && println(n) if you want to see why this
# triangle is called the Save-Prime-Triangle.

println("\nThe Save-Prime-Triangle, the logo of the OLMS
(open library of mathematical sequences).\n")
for n in 5:23
    T = [k for k in 1:n if ⍊(n, k)]
    L = [Int(k ∈ T) for k in 3:(n - 2)]
    println(L)
end

println("\n... have fun with Sequences!")

end # module


function HPDE_∂F∂x(func_F,u,Δx,dim)
    Uₓᴸ, Uₓᴿ = GeneralFluxLimiter(u,Δx)
    n = size(u,dim)
    ∂f∂x = copy(u)

    for i=1:n
        A⁺ , A⁻ = LRJacobian(func_F,u[:,i])
        ∂f∂x[:,i] = (A⁻ * Uₓᴸ[:,i] .+ A⁺ * Uₓᴿ[:,i])
    end
    return ∂f∂x
end

function HPDE_∂u∂t(prob::HPDEProblem,u)

    ∂F∂x = HPDE_∂F∂x(prob.func_F,u,prob.Grid.Δx,2)

    # future extentions
    #∂F∂x = HPDE_∂F∂x(prob.func_F,u,prob.Grid.Δx,1)
    #∂G∂y = HPDE_∂F∂x(prob.func_G,u,prob.Grid.Δy,2)
    #∂H∂z = HPDE_∂F∂x(prob.func_H,u,prob.Grid.Δz,3)
    #S = prob.func_S(u)

    ∂G∂y = 0
    ∂H∂z = 0
    #S = func_S(u)
    S = 0

    ∂u∂t = S .- (∂F∂x .+ ∂G∂y .+ ∂H∂z)

    return ∂u∂t
end

function LRJacobian(func_F,u)
    A = ForwardDiff.jacobian(func_F, u)
    #A = (abs.(A) .> 1e-10) .* A
    #println(A)
    Λ = diagm(eigvals(A))
    Λ⁺ = 0.5 .* (Λ .+ abs.(Λ))
    Λ⁻ = 0.5 .* (Λ .- abs.(Λ))
    P = eigvecs(A)
    P⁻¹ = inv(P)
    A⁺ = P * Λ⁺ * P⁻¹
    A⁻ = P * Λ⁻ * P⁻¹
    return A⁺, A⁻
end

function GeneralFluxLimiter(u,Δx)
    ϵ = 1.0e-8
    ∂u∂x(U⁺,U⁻) = (U⁺ .- U⁻) ./ Δx
    func_Ψ(r) = 0
    #func_Ψ(r) = κ_Scheme_new(r,1)
    func_Ψ(r) = SuperBee(r)
    #func_Ψ(r) = Koren(r)
    #func_Ψ(r) = 0
    func_r(Uₓ₁,Uₓ₂) = (Uₓ₁ .+ ϵ) ./ (Uₓ₂ .+ ϵ)
    func_U(Δxᵢ,r,uₓ) = (Δxᵢ / 2) .* func_Ψ(r) .* (uₓ .+ ϵ)
    Uᵢ₋₂ = [u[:,1] u[:,1] u[:,1:end-2]]
    Uᵢ₋₁ = [u[:,1] u[:,1:end-1]]
    Uᵢ   = copy(u)
    Uᵢ₊₁ = [u[:,2:end] u[:,end]]
    Uᵢ₊₂ = [u[:,3:end] u[:,end] u[:,end]]

    Uₓi₊¾ = ∂u∂x(Uᵢ₊₂,Uᵢ₊₁)
    Uₓi₊½ = ∂u∂x(Uᵢ₊₁,Uᵢ)
    Uₓi₋½ = ∂u∂x(Uᵢ,Uᵢ₋₁)
    Uₓi₋¾ = ∂u∂x(Uᵢ₋₁,Uᵢ₋₂)

    Uᴸi₊½ = Uᵢ₊₁ .- func_U(Δx,func_r(Uₓi₊½,Uₓi₊¾),Uₓi₊¾)
    Uᴸi₋½ = Uᵢ   .- func_U(Δx,func_r(Uₓi₋½,Uₓi₊½),Uₓi₊½)
    Uᴿi₊½ = Uᵢ   .+ func_U(Δx,func_r(Uₓi₊½,Uₓi₋½),Uₓi₋½)
    Uᴿi₋½ = Uᵢ₋₁ .+ func_U(Δx,func_r(Uₓi₋½,Uₓi₋¾),Uₓi₋¾)

    Uₓᴸ = (Uᴸi₊½ .- Uᴸi₋½) ./ Δx
    Uₓᴿ = (Uᴿi₊½ .- Uᴿi₋½) ./ Δx
    #return Uᴿi₊½, Uᴿi₋½, Uᴸi₊½ , Uᴸi₋½
    #return Uₓi₊¾, Uₓi₊½, Uₓi₋½, Uₓi₋¾
    return Uₓᴸ, Uₓᴿ
end

################## METHODS ####################

function Koren(r)
    return max.(0,min.(2 .* r,min.(1/3 .+ 2 /3 .* r,2)))
end

function UpWind_FO(r)
    return 0
end

function UpWind_SO(r)
    return κ_Scheme_new(r,-1)
end

function SCD(r)
    return κ_Scheme_new(r,1)
end

function UpWind_Q(r)
    return κ_Scheme_new(r,0.5)
end

function UpWind_C(r)
    return κ_Scheme_new(r,1/3.0)
end

function SuperBee(r)
    return max.(0,min.(r,1))
end


# Converts an FFT into a gammatonegram. Based on the gammatonegram/fft2gammatonemx MATLAB
# functions by Dan Ellis.
# D. P. W. Ellis (2009). "Gammatone-like spectrograms", web resource.
# http://www.ee.columbia.edu/~dpwe/resources/matlab/gammatonegram/

type Gammatonegram{T, F<:Real} <: DSP.Periodograms.TFR{T}
	amplitude::Matrix{T}
	frequencies::Vector{F}
	time::FloatRange{Float64}
end

function gammatonegram(x,sr::Integer,twin::Real,thop::Real,N::Integer,fmin,fmax,width)
	nfft = int(2^(ceil(log(2*twin*sr)/log(2))))
    nhop = iround(thop*sr)
    nwin = iround(twin*sr)
    (W,F) = fft2gammatonemx(nfft, sr, N, width, fmin, fmax)
    # perform FFT and weighting in amplitude domain
	S = stft(x, nwin, nwin-nhop; nfft=nfft, fs=sr, window=hanning)
	Y = 1/nfft*W*abs(S)
	Gammatonegram(Y, vec(F), ((0:size(Y,2)-1)*(nhop)+nwin/2)/sr)
end

function fft2gammatonemx(nfft::Integer, sr::Integer, N::Integer, width, fmin, fmax)
	# Allocating matrix to store the filterbank weights
	W = zeros(N, int(nfft/2+1));
	# Gammatone filterbank constants
	EarQ = 9.26449
	minBW = 24.7
	order = 1

	cfreqs = -(EarQ*minBW) + exp((1:N)'*(-log(fmax + EarQ*minBW) + log(fmin + EarQ*minBW))/N) * (fmax + EarQ*minBW);
	cfreqs = fliplr(cfreqs)
	GTord = 4
	ucirc = exp(im*2*pi*[0:(nfft/2)]/nfft)
	for k=1:N
	    cf = cfreqs[k];
	    ERB = width*((cf/EarQ).^order + minBW^order).^(1/order);
	    B = 1.019*2*pi*ERB;
	    r = exp(-B/sr);
	    theta = 2*pi*cf/sr;
	    pole = r*exp(im*theta);
	    T = 1/sr;

		A11 = -(2*T*cos(2*cf*pi*T)./exp(B*T) + 2*sqrt(3+2^1.5)*T*sin(2*
	                                                      cf*pi*T)./exp(B*T))/2;
	    A12 = -(2*T*cos(2*cf*pi*T)./exp(B*T) - 2*sqrt(3+2^1.5)*T*sin(2*
	                                                      cf*pi*T)./exp(B*T))/2;
	    A13 = -(2*T*cos(2*cf*pi*T)./exp(B*T) + 2*sqrt(3-2^1.5)*T*sin(2*
	                                                      cf*pi*T)./exp(B*T))/2;
	    A14 = -(2*T*cos(2*cf*pi*T)./exp(B*T) - 2*sqrt(3-2^1.5)*T*sin(2*
	                                                      cf*pi*T)./exp(B*T))/2;
	    zros = -[A11 A12 A13 A14]/T;

	    gain =  abs((-2*exp(4*im*cf*pi*T)*T +
	                2*exp(-(B*T) + 2*im*cf*pi*T).*T.*
	                (cos(2*cf*pi*T) - sqrt(3 - 2^(3/2))*
	                 sin(2*cf*pi*T))) .*
	               (-2*exp(4*im*cf*pi*T)*T +
	                2*exp(-(B*T) + 2*im*cf*pi*T).*T.*
	                (cos(2*cf*pi*T) + sqrt(3 - 2^(3/2)) *
	                 sin(2*cf*pi*T))).*
	               (-2*exp(4*im*cf*pi*T)*T +
	                2*exp(-(B*T) + 2*im*cf*pi*T).*T.*
	                (cos(2*cf*pi*T) -
	                 sqrt(3 + 2^(3/2))*sin(2*cf*pi*T))) .*
	               (-2*exp(4*im*cf*pi*T)*T + 2*exp(-(B*T) + 2*im*cf*pi*T).*T.*
	                (cos(2*cf*pi*T) + sqrt(3 + 2^(3/2))*sin(2*cf*pi*T))) ./
	               (-2 ./ exp(2*B*T) - 2*exp(4*im*cf*pi*T) +
	                2*(1 + exp(4*im*cf*pi*T))./exp(B*T)).^4);
	    W[k,:] = ((T^4)/gain) *
	        abs(ucirc-zros[1]).*abs(ucirc-zros[2]) .*
	        abs(ucirc-zros[3]).*abs(ucirc-zros[4]) .*
	        (abs((pole-ucirc).*(pole'-ucirc)).^-GTord);
	end
	W, cfreqs
end


#!/usr/bin/julia

# Trizen
# 28 August 2020
# https://github.com/trizen

# Generate all the positive divisors of n.

using Primes

function divisors(n)

    d = Int64[1]

    for (p,e) in factor(n)
        t = Int64[]
        r = 1

        for i in 1:e
            r *= p
            for u in d
                push!(t, u*r)
            end
        end

        append!(d, t)
    end

    return sort(d)
end

println(divisors(5040))


function kernel(x::REALMAT, x₀::REALMAT, kernel_struct::Kernel)
    @assert size(x, 2) == size(x₀, 2) "The number of columns of x and x₀ do not match."
    # Dispatch to kernel
    K = kernel_calculation(x, x₀, kernel_struct::Kernel)
    return K
end
kernel(x::REALMAT, kernel_struct::Kernel) = kernel(x, x, kernel_struct)

# Linear
function kernel_calculation(x, x₀, k::Linear)
    K = [dot(x[i1, :], x[i2, :]) for i1 in 1:size(x, 1), i2 in 1:size(x₀, 1)]
    return K * k.σ
end

# Radial Basis Function (Squared Exponential)
function kernel_calculation(x, x₀, k::RBF)
    K = zeros(size(x, 1), size(x₀, 1))
    for i1 in 1:size(x, 1)
        for i2 in 1:size(x₀, 1)
            K[i1, i2] = k.σ * exp(-sum((x[i1, :] .- x₀[i2, :]).^2)/(2*k.l^2))
        end
    end
    return K
end

# Rational Quadratic
function kernel_calculation(x, x₀, k::RationalQuadratic)
    K = zeros(size(x, 1), size(x₀, 1))
    for i1 in 1:size(x, 1)
        for i2 in 1:size(x₀, 1)
            K[i1, i2] = k.σ * (1.0 + sum((x[i1, :] .- x₀[i2, :])^2)/(2 * k.a * k.l^2))^-k.a
        end
    end
    return K
end

# Periodic
function kernel_calculation(x, x₀, k::Periodic)
    K = zeros(size(x, 1), size(x₀, 1))
    for i1 in 1:size(x, 1)
        for i2 in 1:size(x₀, 1)
            K[i1, i2] = k.σ * exp(-2 * sin(sum(abs.(x[i1, :] .- x₀[i2, :])))^2/k.l^2)
        end
    end
    return K
end

# Locally Periodic
function kernel_calculation(x, x₀, k::LocallyPeriodic)
    K = zeros(size(x, 1), size(x₀, 1))
    for i1 in 1:size(x, 1)
        for i2 in 1:size(x₀, 1)
            K[i1, i2] = k.σ * exp(-2 * sin(sum(abs.(x[i1, :] .- x₀[i2, :])))^2/k.l^2) * exp(-sum((x[i1, :] .- x[i2, :].^2)/(2*k.l^2)))
        end
    end
    return K
end


function evalPsiFromRectangleCoil(r0,z0,dr,dz,Ic,quadRule_r,quadRule_z, xr,xz; derivative=0)

	surf=dr*dz
	jc=Ic/surf

	#define quadrature points and weigths in the rectangle
	xgr=r0 .+ quadRule_r.x *0.5 *dr
	xgz=z0 .+ quadRule_z.x *0.5 *dz

	xgr_rep=repeat(xgr,size(xgz,1))
	xgz_rep=reshape(repeat(xgz,1,size(xgr,1))',size(xgr,1)*size(xgz,1))

	wgr=quadRule_r.w *0.5 *dr
	wgz=quadRule_z.w *0.5 *dz

	wgr_rep=repeat(wgr,size(wgz,1))
	wgz_rep=reshape(repeat(wgz,1,size(wgr,1))',size(wgr,1)*size(wgz,1))

	#
	G = @.Green(xr,xz,xgr_rep,xgz_rep, derivative=derivative)
	G_PXPG = getindex.(G,:G_xy)

	val=wgr_rep.*wgz_rep.*G_PXPG*jc
	val=sum(val)

	if(derivative==0)
		return (psi=val,)
	end
	if(derivative>0)
		dxrG = getindex.(G,:dxrG_xy)
		dxzG = getindex.(G,:dxzG_xy)

		dxrval=wgr_rep.*wgz_rep.*dxrG*jc
		dxzval=wgr_rep.*wgz_rep.*dxzG*jc

		dxrval=sum(dxrval)
		dxzval=sum(dxzval)

		return (psi=val, dxrpsi=dxrval, dxzpsi=dxzval)

	end

end

function assemblyRhsRectangleCoil(r0,z0,dr,dz,Ic,quadRule_r, quadRule_z, BEMesh, quadRule)

	# dimension
	n=size(BEMesh.edges,1)

	#
	rhs=zeros(n)

	#quadrature rule on BEMesh
	xGauss=quadRule.x
	wGauss=quadRule.w
	nQuad=size(quadRule.w,1)

	#in order to remap the quadrature we need the midpoints and length of edges
	e=collect(1:n)
	edgeToNode=BEMesh.edges
	h=BEMesh.length

    P=0.5*(BEMesh.coordinates[edgeToNode[:,1],:]+BEMesh.coordinates[edgeToNode[:,2],:])
	#@show P

	#quadrature points on all edges
	PG=zeros(size(P,1),2,nQuad)
	for i in 1:nQuad
		PG[:,:,i]=P + xGauss[i]*0.5*(BEMesh.coordinates[edgeToNode[:,2],:]-BEMesh.coordinates[edgeToNode[:,1],:])

		evalC=@.evalPsiFromRectangleCoil((r0,),(z0,),(dr,),(dz,),(Ic,),(quadRule_r,),(quadRule_z,), PG[:,1,i], PG[:,2,i])

		psiC = getindex.(evalC,:psi)

		rhs=rhs+0.5*wGauss[i]*h.*psiC
	end

	return rhs

end


"""
    randBall(nDims, nPts)

Draw nPts random points uniformely distributed in an nDims-dimensional ball.

# Implementation
Samples uniformely in a unit cube, then throws away unnecessary points.
Efficient only for low number of dimensions (nDims=1,2,3).

# Examples
```julia-repl
julia> using FPRS; using Plots;
julia> x = randBall(2,50);
julia> figure(); scatter(x[1,:], x[2,:]);
```
"""
function randBall(nDims::Integer, nPts::Integer)

    res = zeros(nDims, nPts)
    numInBall = 0

    while numInBall < nPts
        newPoint = 2.0 * rand(nDims) .- 1.0
        if norm(newPoint) <= 1
            numInBall += 1
            res[1:end, numInBall] = newPoint
        end
    end

    return res
end

"""
    randRotationMatrix(deflection=1.0; randSeed)

Return a random rotation matrix in 3d.

Sources:
-http://blog.lostinmyterminal.com/python/2015/05/12/random-rotation-matrix.html
-http://www.realtimerendering.com/resources/GraphicsGems/gemsiii/rand_rotation.c

# Arguments
- `deflection`: number between 0 and 1 controls the size of the perturbation.
                For 0, no rotation; for 1, competely random rotation.
                Small deflection => small perturbation.

# Examples
```julia-repl
julia> mx = randRotationMatrix();
```
"""
function randRotationMatrix(deflection=1.0)
    # Random.seed!(seedVal) # do not reseed rng

    theta = 2.0*deflection* pi * rand() # Rotation about the pole (Z)
    phi = 2.0* pi * rand() # For direction of pole deflection
    z = 2.0*deflection * rand() # For magnitude of pole deflection

    # Compute a vector V used for distributing points over the sphere
    # via the reflection I - V Transpose(V).  This formulation of V
    # will guarantee that if x[1] and x[2] are uniformly distributed,
    # the reflected points will be uniform on the sphere.  Note that V
    # has length sqrt(2) to eliminate the 2 in the Householder matrix.

    r =  sqrt(z)
    Vx = sin(phi) * r
    Vy = cos(phi) * r
    Vz = sqrt(2.0 - z)

    # Compute the row vector S = Transpose(V) * R, where R is a simple
    # rotation by theta about the z-axis.  No need to compute Sz since
    # it's just Vz.

    st = sin(theta)
    ct = cos(theta)
    Sx = Vx * ct - Vy * st
    Sy = Vx * st + Vy * ct

    # Construct the rotation matrix  ( V Transpose(V) - I ) R, which
    # is equivalent to V S - R.

    M = zeros(3,3)
    M[1,1] = Vx * Sx - ct;
    M[1,2] = Vx * Sy - st;
    M[1,3] = Vx * Vz;

    M[2,1] = Vy * Sx + st;
    M[2,2] = Vy * Sy - ct;
    M[2,3] = Vy * Vz;

    M[3,1] = Vz * Sx;
    M[3,2] = Vz * Sy;
    M[3,3] = 1.0 - z;   # This equals Vz * Vz - 1.0

    return M
end


# ffr, 19 April 2016

# for the moment, I decided to not include L, because it can
# be calculated from A and B
struct LF1dGrid
    N::Int64
    A::Float64
    B::Float64
    h::Float64
    grid::Array{Float64,1}
    D1jl::Array{Float64,2}
    D2jl::Array{Float64,2}
end

# Periodic LF
function init_LF1d_p( N::Int64, A::Float64, B::Float64, verbose=false )
    # Check argument
    if N % 2 == 0
        error("N must be an odd number\n")
    end
    L = (B - A)
    h = L/N
    grid = zeros(Float64,N)
    for i = 1:N
        grid[i] = A + 0.5*(B-A)*(2*i-1)/N
    end
    D1jl = zeros(Float64,N,N)
    D2jl = zeros(Float64,N,N)
    #
    # Diagonal elements
    #
    Nprimed = (N-1)/2
    for j = 1:N
        # the diagonal elements of D1jl are already zero
        D2jl[j,j] = -(2.0*pi/L)^2 * Nprimed * (Nprimed+1)/3.
    end
    #
    # Off diagonal elements
    #
    for j = 1 : N
        for l = j+1 : N
            #
            nn = j - l
            #
            tt1 = pi/L * (-1.0)^nn
            tt2 = sin(pi*nn/N)
            #
            tt3 = (2.0*pi/L)^2 * (-1.0)^nn * cos(pi*nn/N)
            tt4 = 2.0*sin(pi*nn/N)^2
            #
            D1jl[j,l] =  tt1/tt2
            D1jl[l,j] = -tt1/tt2
            #
            D2jl[j,l] = -tt3/tt4
            D2jl[l,j] = -tt3/tt4
        end
    end
    #
    LF = LF1dGrid( N, A, B, h, grid, D1jl, D2jl )
    if verbose
        @printf("Allocated: periodic 1d LBF grid: N, h = %d , %f\n", N, h)
    end
    #
    return LF
end # function init_LF1d_p

# cluster Lagrange function
function init_LF1d_c( N::Int64, A::Float64, B::Float64, verbose=false )
    L = B - A
    h = L/(N + 1.)
    #
    grid = zeros(Float64,N)
    for i = 1:N
        grid[i] = A + i*(B-A)/(N+1)
    end
    #
    D1jl = zeros(Float64,N,N)  # XXX this is not yet used
    D2jl = zeros(Float64,N,N)
    #
    # Diagonal part
    #
    pre = -pi^2/(2.0*L^2)
    for i = 1:N
        t1 = ( 2.0*(N+1)^2 + 1 )/3.0
        t2 = sin( i*pi/(N+1) )^2
        D2jl[i,i] = pre*( t1 - 1.0/t2 )
    end
    #
    # Off-diagonal
    #
    for l = 1 : N
        for j = l+1 : N
            nnm = l - j
            nnp = l + j
            pre = -pi^2 / (2*L^2) * (-1.0)^nnm
            t1 = sin( pi*nnm/2.0/(N+1) )^2
            t2 = sin( pi*nnp/2.0/(N+1) )^2
            #
            D2jl[l,j] = pre*( 1.0/t1 - 1.0/t2 )
            D2jl[j,l] = pre*( 1.0/t1 - 1.0/t2 )  # XXX is it faster to just recalculate?
        end
    end
    LF = LF1dGrid( N, A, B, h, grid, D1jl, D2jl )
    if verbose
        @printf("Allocated: cluster 1d LBF grid: N, h = %5d , %10.5f\n", N, h)
    end
    return LF
end

# Lagrange-sinc function
function init_LF1d_sinc( N::Int64, h::Float64, verbose=false )
    # Choice for A and B
    A = -(N-1)*h/2.0
    B =  (N-1)*h/2.0
    #
    grid = zeros(Float64,N)
    for i = 1:N
        grid[i] = A + (i-1)*h
    end
    #
    D1jl = zeros(Float64,N,N)  # XXX this is not yet used
    D2jl = zeros(Float64,N,N)
    #
    # Diagonal part
    #
    for i = 1:N
        D2jl[i,i] = -pi^2 / 3.0 / h^2
    end
    #
    # Off-diagonal
    #
    for j = 1 : N
        for i = j+1 : N
            D2jl[i,j] = -2.0*(-1.0)^(i-j)/( grid[i] - grid[j] )^2
            D2jl[j,i] = D2jl[i,j]  # XXX is it faster to just recalculate?
        end
    end
    LF = LF1dGrid( N, A, B, h, grid, D1jl, D2jl )
    if verbose
        @printf("Allocated: 1d Lagrange-sinc grid: N, h = %5d , %10.5f\n", N,  h)
    end
    return LF
end