#  Unit "QnLogLike.jl" of the Diagonalization.jl Package for Julia language
#
#  MIT License
#  Copyright (c) 2020-2021,
#  Marco Congedo°, Ronald Phlypo, CNRS, UGA, Grenoble-INP, France
#  Alexandre Gramfort¨, INRIA, U. Paris Saclay, France
#  ° https://sites.google.com/site/marcocongedo/
#  ¨ http://alexandre.gramfort.net/

# ? CONTENTS :
#  Quasi-Newton Approximate Joint Diagonalization (AJD) algorithm of:
#  P. Ablin, J.F. Cardoso and A. Gramfort. Beyond Pham's algorithm
#  for joint diagonalization. Proc. ESANN 2019.
#  https://hal.archives-ouvertes.fr/hal-01936887v1
#
#  Code adapted in Julia from Python code provided from the authors at:
#  https://github.com/pierreablin/qndiag/blob/master/qndiag/qndiag.py
#
#  NOTATION
#  lower-case letter: a scalar, e.g.: a
#  Upper case letter: a matrix, e.g., A
#  Bold lover case letter: a vector of matrices, e.g., 𝐀
#
#  ADVANCED QnLogLike algorithm (no PRIMITIVE version is available):
#  The algorithm takes as input a vector of k real symmetric matrices
#  𝐂 and find a non-singular matrix B such that the congruences
#  B'*𝐂_κ*B are as diagonal as possible for all κ=1:k.
#
#  `w` is an optional vector of k positive weights for each matrix in 𝐂.
#  if `w` is different from `nothing` (default), the input matrices are
#  weighted with these weights.
#  A function can be passed as the `w` argument, in which case the kth weight
#  is found as the output of the function applied to the kth matrix in 𝐂.
#  A good choice in general is the `nonD` function declared in tools.jl unit.
#
#  if `whitening` = true is passed, the Jeffrey mean of the matrices in 𝐂 is
#  computed (using the PosDefManifold.jl package) and the matrices in 𝐂
#  are pre-transformed using the whitening matrix of the mean.
#  Dimensionality reduction can be obtained at this stage using optional
#  arguments `eVar` and `eVarMeth` (see documentation of the AJD constructors).
#
#  if `sort`=true (default) the column vectors of the B matrix are normalized
#  to unit norm and permuted so as to sort in descending order the mean over
#  κ=1:k of the diagonal elements of B'*𝐂_κ*B.
#  Note that if `whitening` is true the output B will not have unit norm
#  columns, as it is multiplied by the whitener after being scaled and sorted
#  and before being returned.
#
#  if  `whitening` = false (default), a matrix can be provided with the `init`
#  argument in order to initialize B. In this case the actual AJD
#  will be given by init*B, where B is the output of the algorithm.
#
#  `tol` is the convergence to be attained. It defaults to 1e-6.
#
#  `maxiter` is the maximum number of iterations allowed. It defaults to 1000.
#
#  `𝜆min` is used to reguarize Hessian coefficients; all coefficients smaller
#  than 𝜆min  will be set to 𝜆min. It defauts tp 1e-4.
#
#  `lsmax` is the maximum number of steps in the line search. It defaults to 10.
#
#  if `verbose`=true, the convergence attained at each iteration and other
#  information will be printed.
#
#  if `threaded` is true, the algorithm runs in multi-threaded mode.
#  By default this is activated if the number of threads Julia is instructed
#  to use is higher than 1 and if the number of matrices in 𝐂 is higher than
#  twice the number of threads julia is instructed to use. For small matrices
#  and/or small number of matrices, setting threaded to false may result
#  in better performances in term of speed.
#
#  return a 5-tuple holding: B, its pseudo-inverse,
# 		the mean diagonal elements of B'*mean(𝐂)*B,
#       the number of iterations and the convergence attained.
#
#  Note on the implementation:
#  GRADIENT:
#  D_1,...,D_n are the matrices in 𝐃;
#  the jth column in D_i is divided by D_i[j, j],
#  the mean of these matrices is then taken and finally
#  the identity is subtracted
#
#  HESSIAN COEFFICIENTS
#  for each COLUMN vector dg_1,...,dg_n
#  (the diagonal part of the n matrices in 𝐃)
#  we form a m·m matrix stacking vertically m copies of these vectors
#  transposed and dividing each of them by the mth element, as
#   _                   _
#  |  (dg_i)'/dg_i[1]   |
#  |  (dg_i)'/dg_i[...] |
#  |  (dg_i)'/dg_i[m]   |
#  _                   _
#  finally we take the mean of all the n matrices created in this way.

# function to get the weights from argment `w`
function _qnlogLikeWeights!(w, 𝐂)
	if w isa Function w=[w(C) for C∈𝐂] end
	return w./mean(w)
end

function qnLogLike( 𝐂::Union{Vector{Hermitian}, Vector{Symmetric}};
				w           :: 	Twf   = ○,
				preWhite    :: 	Bool = false,
				sort        :: 	Bool = true,
				init        :: 	Union{Matrix, Nothing} = ○,
				tol         :: 	Real = 1e-6,
				maxiter     :: 	Int  = 1000,
				𝜆min        ::	 Real = 1e-4,
				lsmax       :: 	Int  = 10,
				verbose     :: 	Bool = false,
				threaded	:: 	Bool =
							begin
								thr = Threads.nthreads()
								length(𝐂) ≥ 2*thr && thr>1
							end,
			eVar     :: TeVaro = ○,
			eVarMeth :: Function = searchsortedfirst)

	# # # # # # # # #  internal functions
	# half sum of mean of log of Diagonals
	hsmld(𝐀) = 0.5*sum(mean(log, [𝔻(A) for A ∈ 𝐀]))
	hsmld(𝐀, 𝛎) = 0.5*sum(mean(log, [𝔻(A)*ν for (A, ν) ∈ zip(𝐀, 𝛎)]))

	# minus log abs det
	mlad(A) = -logabsdet(A)[1]

	# line search: update B and 𝐃
	@inline function _linesearch!(B, B₊, M, 𝐃, 𝐃₊, 𝒟, 𝐯, loss, lsmax)
	    for i ∈ 1:lsmax
	        M[:] = (1.0/i * 𝒟) + I
			B₊[:] = B * M
			### 𝐃₊ = [Hermitian(M'*D*M) for D in 𝐃]
			@threads for j ∈ eachindex(𝐃) 𝐃₊[j] = Hermitian(M'*𝐃[j]*M) end
			### iter>2 && (loss₊ = w===○ ? mlad(B₊)+hsmld(𝐃₊) : mlad(B₊)+hsmld(𝐃₊, 𝐯))
			loss₊ = w===○ ? mlad(B₊)+hsmld(𝐃₊) : mlad(B₊)+hsmld(𝐃₊, 𝐯)
	        loss₊ < loss && break
	    end
		B[:] = B₊
		𝐃[:] = 𝐃₊
	    return loss₊
	end
	# # # # # # # # #

	# pre-whiten or initialize or nothing
	W, 𝐃 = _preWhiteOrInit(𝐂, preWhite, Jeffrey, eVar, eVarMeth, init, :Hvector)
	𝐯 = w===○ ? ○ : _qnlogLikeWeights!(w, 𝐂) # if w is `nonD` function, apply it to the original input 𝐂

	# set variables
	iter, conv, 😋, sqrtn, loss₊ = 1, Inf, false, √size(𝐃[1], 1), Inf
	B, ⩫ = Matrix{eltype(𝐃[1])}(I, size(𝐃[1])), similar
	B₊, 𝒟, M, ∇, ℌ, 𝐃₊ = ⩫(B), ⩫(B), ⩫(B), ⩫(B), ⩫(B), ⩫(𝐃)
	𝕯 = [zeros(eltype(𝐃[1]), size(𝐃[1], 1)) for i = 1:length(𝐃)]
	loss = w===○ ? 	hsmld(𝐃) : hsmld(𝐃, 𝐯) #loss = Inf

	verbose && println("Iterating quasi-Newton LogLike algorithm...")
	while true
		for j ∈ eachindex(𝐃) 𝕯[j] = diag(𝐃[j]) end
	    #𝕯 = [diag(D) for D ∈ 𝐃]

	    # gradient
		w===○ ? ∇[:] = mean(D./𝔡 for (D, 𝔡) ∈ zip(𝐃, 𝕯)) - I :
				∇[:] = mean(v.*(D./𝔡) for (v, D, 𝔡) ∈ zip(𝐯, 𝐃, 𝕯)) - I
	    conv = norm(∇)/sqrtn # relative norm of ∇ with respect to the identity : ||∇-I||/||I||

	    verbose && println("iteration: ", iter, "; convergence: ", conv)
	    (overRun = iter > maxiter) && @warn("qnLogLike: reached the max number of iterations before convergence:", iter-1)
	    (😋 = conv <= tol) || overRun==true ? break : iter += 1

		# Hessian Coefficients
		w===○ ? ℌ[:] = mean(𝔡'./𝔡 for 𝔡 ∈ 𝕯) :
				ℌ[:] = mean(v.*(𝔡'./𝔡) for (v, 𝔡) ∈ zip(𝐯, 𝕯))

		# Quasi-Newton Direction 𝒟
	    𝒟[:] = -(∇' .* ℌ - ∇)./replace(x -> x<𝜆min ? 𝜆min : x, @. (ℌ'*ℌ) - 1.)

	    loss = _linesearch!(B, B₊, M, 𝐃, 𝐃₊, 𝒟, 𝐯, loss, lsmax) # Line Search
	end
	verbose && @info("Convergence has "*(😋 ? "" : "not ")*"been attained.\n\n")

	# scale and permute the vectors of B
	λ = sort ? _permute!(_scale!(B, mean(𝔻(D) for D ∈ 𝐃), size(𝐃[1], 1))...) :
	            diag(mean(𝔻(D) for D ∈ 𝐃))

	return preWhite ? (W.F*B, pinv(B)*W.iF, λ, iter, conv) :
	                  (B, pinv(B), λ, iter, conv)
end


# functions related to normal distribution

function xval(μ::Real, σ::Real, z::Number)
    if isinf(z) && iszero(σ)
        μ + one(σ) * z
    else
        μ + σ * z
    end
end
zval(μ::Real, σ::Real, x::Number) = (x - μ) / σ

# pdf
normpdf(z::Number) = exp(-abs2(z)/2) * invsqrt2π
function normpdf(μ::Real, σ::Real, x::Number)
    if iszero(σ)
        if x == μ
            z = zval(μ, one(σ), x)
        else
            z = zval(μ, σ, x)
            σ = one(σ)
        end
    else
        z = zval(μ, σ, x)
    end
    normpdf(z) / σ
end

# logpdf
normlogpdf(z::Number) = -(abs2(z) + log2π)/2
function normlogpdf(μ::Real, σ::Real, x::Number)
    if iszero(σ)
        if x == μ
            z = zval(μ, one(σ), x)
        else
            z = zval(μ, σ, x)
            σ = one(σ)
        end
    else
        z = zval(μ, σ, x)
    end
    normlogpdf(z) - log(σ)
end

# cdf
normcdf(z::Number) = erfc(-z * invsqrt2)/2
function normcdf(μ::Real, σ::Real, x::Number)
    if iszero(σ) && x == μ
        z = zval(zero(μ), σ, one(x))
    else
        z = zval(μ, σ, x)
    end
    normcdf(z)
end
# ccdf
normccdf(z::Number) = erfc(z * invsqrt2)/2
function normccdf(μ::Real, σ::Real, x::Number)
    if iszero(σ) && x == μ
        z = zval(zero(μ), σ, one(x))
    else
        z = zval(μ, σ, x)
    end
    normccdf(z)
end

# logcdf
normlogcdf(z::Number) = z < -1.0 ?
    log(erfcx(-z * invsqrt2)/2) - abs2(z)/2 :
    log1p(-erfc(z * invsqrt2)/2)
function normlogcdf(μ::Real, σ::Real, x::Number)
    if iszero(σ) && x == μ
        z = zval(zero(μ), σ, one(x))
    else
        z = zval(μ, σ, x)
    end
    normlogcdf(z)
end

# logccdf
normlogccdf(z::Number) = z > 1.0 ?
    log(erfcx(z * invsqrt2)/2) - abs2(z)/2 :
    log1p(-erfc(-z * invsqrt2)/2)
function normlogccdf(μ::Real, σ::Real, x::Number)
    if iszero(σ) && x == μ
        z = zval(zero(μ), σ, one(x))
    else
        z = zval(μ, σ, x)
    end
    normlogccdf(z)
end

norminvcdf(p::Real) = -erfcinv(2*p) * sqrt2
norminvcdf(μ::Real, σ::Real, p::Real) = xval(μ, σ, norminvcdf(p))

norminvccdf(p::Real) = erfcinv(2*p) * sqrt2
norminvccdf(μ::Real, σ::Real, p::Real) = xval(μ, σ, norminvccdf(p))

# invlogcdf. Fixme! Support more precisions than Float64
norminvlogcdf(lp::Union{Float16,Float32}) = convert(typeof(lp), _norminvlogcdf_impl(Float64(lp)))
norminvlogcdf(lp::Real) = _norminvlogcdf_impl(Float64(lp))
norminvlogcdf(μ::Real, σ::Real, lp::Real) = xval(μ, σ, norminvlogcdf(lp))

# invlogccdf. Fixme! Support more precisions than Float64
norminvlogccdf(lp::Union{Float16,Float32}) = convert(typeof(lp), -_norminvlogcdf_impl(Float64(lp)))
norminvlogccdf(lp::Real) = -_norminvlogcdf_impl(Float64(lp))
norminvlogccdf(μ::Real, σ::Real, lp::Real) = xval(μ, σ, norminvlogccdf(lp))

# norminvcdf & norminvlogcdf implementation
#
#   Rational approximations for the inverse cdf and its logarithm, from:
#
#   Wichura, M.J. (1988) Algorithm AS 241: The Percentage Points of the Normal Distribution
#   Journal of the Royal Statistical Society. Series C (Applied Statistics), Vol. 37, No. 3, pp. 477-484
#

function _norminvlogcdf_impl(lp::Float64)
    if isfinite(lp) && lp < 0.0
        q = exp(lp) - 0.5
        # qnorm_kernel(lp, q, true)
        if abs(q) <= 0.425
            _qnorm_ker1(q)
        else
            r = sqrt(q < 0 ? -lp : -log1mexp(lp))
            return copysign(_qnorm_ker2(r), q)
        end
    elseif lp >= 0.0
        lp == 0.0 ? Inf : NaN
    else # lp is -Inf or NaN
        lp
    end
end

function _qnorm_ker1(q::Float64)
    # pre-condition: abs(q) <= 0.425
    r = 0.180625 - q*q
    return q * @horner(r,
                       3.38713_28727_96366_6080e0,
                       1.33141_66789_17843_7745e2,
                       1.97159_09503_06551_4427e3,
                       1.37316_93765_50946_1125e4,
                       4.59219_53931_54987_1457e4,
                       6.72657_70927_00870_0853e4,
                       3.34305_75583_58812_8105e4,
                       2.50908_09287_30122_6727e3) /
    @horner(r,
            1.0,
            4.23133_30701_60091_1252e1,
            6.87187_00749_20579_0830e2,
            5.39419_60214_24751_1077e3,
            2.12137_94301_58659_5867e4,
            3.93078_95800_09271_0610e4,
            2.87290_85735_72194_2674e4,
            5.22649_52788_52854_5610e3)
end

function _qnorm_ker2(r::Float64)
    if r < 5.0
        r -= 1.6
        @horner(r,
                1.42343_71107_49683_57734e0,
                4.63033_78461_56545_29590e0,
                5.76949_72214_60691_40550e0,
                3.64784_83247_63204_60504e0,
                1.27045_82524_52368_38258e0,
                2.41780_72517_74506_11770e-1,
                2.27238_44989_26918_45833e-2,
                7.74545_01427_83414_07640e-4) /
        @horner(r,
                1.0,
                2.05319_16266_37758_82187e0,
                1.67638_48301_83803_84940e0,
                6.89767_33498_51000_04550e-1,
                1.48103_97642_74800_74590e-1,
                1.51986_66563_61645_71966e-2,
                5.47593_80849_95344_94600e-4,
                1.05075_00716_44416_84324e-9)
    else
        r -= 5.0
        @horner(r,
                6.65790_46435_01103_77720e0,
                5.46378_49111_64114_36990e0,
                1.78482_65399_17291_33580e0,
                2.96560_57182_85048_91230e-1,
                2.65321_89526_57612_30930e-2,
                1.24266_09473_88078_43860e-3,
                2.71155_55687_43487_57815e-5,
                2.01033_43992_92288_13265e-7) /
        @horner(r,
                1.0,
                5.99832_20655_58879_37690e-1,
                1.36929_88092_27358_05310e-1,
                1.48753_61290_85061_48525e-2,
                7.86869_13114_56132_59100e-4,
                1.84631_83175_10054_68180e-5,
                1.42151_17583_16445_88870e-7,
                2.04426_31033_89939_78564e-15)
    end
end


include("../interfaces.jl")

import PhysicalConstants.CODATA2018 as Constants # this must be moved at a high level in the package to standardize use

### Boundary conditions
function Temperature(;name,T=273.15)
    val = T
    @parameters T # how to give default parameters. how to give default to variables

    @named port = HeatPort()

    eqs =
    [
        port.T ~ T
    ]
    ODESystem(eqs, t, [], [T], systems=[port], name=name)
end

### Resistance Models
# how to have the similar effect as extend PartialResistance to avoid code rewrite
function Sphere(;name, r_inner = 0.01, r_outer=1.0, lambda=5.0)
    val0 = r_inner # don't understand why there "val*" are required for parameters with default values
    val1 = r_outer
    val2 = lambda
    @variables R(t)
    @parameters r_inner, r_outer, lambda

    @named port_a = HeatPort()
    @named port_b = HeatPort()

    eqs =
    [
        0 ~ port_a.Q_flow + port_b.Q_flow
        port_a.Q_flow ~ (port_a.T - port_b.T)/R
        R ~ 1/(4*3.14159*lambda)*(1/r_inner - 1/r_outer)
    ]
    ODESystem(eqs, t, [R], [r_inner, r_outer, lambda], systems=[port_a, port_b], name=name)
end

function Convection(;name, surfaceArea = 0.01, alpha=1000.0)
    val0 = surfaceArea
    val1 = alpha
    @variables R(t)
    @parameters surfaceArea, alpha

    @named port_a = HeatPort()
    @named port_b = HeatPort()

    eqs =
    [
        0 ~ port_a.Q_flow + port_b.Q_flow
        port_a.Q_flow ~ (port_a.T - port_b.T)/R
        R ~ 1/(alpha*max(eps(),surfaceArea))
    ]
    ODESystem(eqs, t, [R], [surfaceArea, alpha], systems=[port_a, port_b], name=name)
end

function Contact(;name, surfaceArea = 0.01, Rc_pp=1.0)
    val0 = surfaceArea
    val1 = Rc_pp
    @variables R(t)
    @parameters surfaceArea, Rc_pp

    @named port_a = HeatPort()
    @named port_b = HeatPort()

    eqs =
    [
        0 ~ port_a.Q_flow + port_b.Q_flow
        port_a.Q_flow ~ (port_a.T - port_b.T)/R
        R ~ Rc_pp/max(eps(),surfaceArea)
    ]
    ODESystem(eqs, t, [R], [surfaceArea, Rc_pp], systems=[port_a, port_b], name=name)
end

function Radiation(;name, surfaceArea = 0.01, epsilon=1.0)
    val0 = surfaceArea
    val1 = epsilon
    @variables R(t)
    @parameters surfaceArea, epsilon

    @named port_a = HeatPort()
    @named port_b = HeatPort()

    eqs =
    [
        0 ~ port_a.Q_flow + port_b.Q_flow
        port_a.Q_flow ~ (port_a.T - port_b.T)/R
        R ~ 1/(surfaceArea*5.6703744E-8*epsilon*(port_a.T^2+port_b.T^2)*(port_a.T + port_b.T))
    ]
 # exact: R ~ 1/(surfaceArea*5.6703744E-8*epsilon*(port_a.T^2+port_b.T^2)*(port_a.T + port_b.T))
 # approximate: R ~ 1/(4*surfaceArea*5.6703744E-8*epsilon*(0.5*(port_a.T + port_b.T))^3)
    # seems to be having issues :(

    # using "Constants.StefanBoltzmannConstant" causes Unitful error
    ODESystem(eqs, t, [R], [surfaceArea, epsilon], systems=[port_a, port_b], name=name)
end


"""
Import packages and files
"""
import Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

@info("Loading Packages...")
using Zygote, LinearAlgebra, Random
using Flux
using Random
using Plots

include("bayes_nn.jl")
"""
class_data.jl contains the data of the example
for Bayesian Neural Network with Turing language from Turinglang
    https://turing.ml/dev/tutorials/3-bayesnn/.
"""
include("turing_data.jl")    #
plot_data()

X = hcat(xs...)
Y = hcat(ts...)

"""
Set up a bnn for classification. The architecture and standard deviation are the
same as in the example of Turing in order to ensure correctness of the algorithm.
"""
# Create a regularization term and a Gaussain prior variance term.
alpha = 0.09
sig = sqrt(1.0 / alpha)

""" NNet with 2 input units, 3 units in 1st hidden layer,
2 units in 2nd hidden layer and 1 output unit and their correpsonding
activation function.
"""
ffnet = network([2, 3, 2, 1], [identity, tanh, tanh, sigmoid])

"""
Sample from the posterior distribution using Hamiltonian Monte Carlo.
"""
function main(net::network, X, Y, prior_sigma, maxiter, burnin)
    # matrices to store mcmc results
    sampled_weights = []
    sampled_bias = []
    lp = zeros(maxiter)

    for its in 1:maxiter
        hmc_bnn!(net, X, Y, prior_sigma, 0.05, 4)
        lp[its] = log_posterior(net, X, Y, prior_sigma)
        if its > burnin
            push!(sampled_weights, net.weights)
            push!(sampled_bias, net.biases)
        end
        if mod(its, 100) == 0
            println("HMC iterations: $its")
        end
    end

    return sampled_weights, sampled_bias, lp
end

Random.seed!(12345)
@time sampled_weights, sampled_bias, lp = main(ffnet, X, Y, sig, 5000, 1000)

# define test set
x_range = collect(range(-6,stop=6,length=25))
y_range = collect(range(-6,stop=6,length=25))

"""
Visualise the results
"""
n_end = 4000
anim = @gif for i=1:100:n_end
    plot_data();
    bnn_posterior!(ffnet, sampled_weights[i], sampled_bias[i])
    Z = [nn_forward(ffnet, [x, y]) for x=x_range, y=y_range]
    contour!(x_range, y_range, cell2array(Z), title="Iteration $i", clim = (0,1));
end every 10;

_, idx = findmax(lp)
bnn_posterior!(ffnet, sampled_weights[idx], sampled_bias[idx])
Z = [nn_forward(ffnet, [x, y]) for x=x_range, y=y_range]
plot_data()
contour!(x_range, y_range, cell2array(Z))


function InverseLUSolve(A::Array,B::Vector)
      #A = the matrix of coefficients 'A' must be squared
  #B = The Vector 'B' right hand side vector Vertical
    m,n =size(A)
    v = size(B,1)
    if m != n
        error("Matrix must be squared")
    elseif m != v
        error("Vector Length must be equal to Matrix")
    end

    N,P=InverseLU(A)

    X=N*(P*B)

    return X
end


# #TODO: finish
# function hilbert_transform(x::Array{T}; n::Int=10)
# end

function estimate_rsrange(x::Array{T})::T where T<:Real
    n = size(x,1)
    @assert n>2 "need more than two elements, have $x"
    m = sum(x)/n
    y = x .- m
    z = cumsum(y)
    r = maximum(z) - minimum(z)
    s = sqrt(sum((x.-m).^2)/n)
    return r/s
end

function npieces(x::Array{T}; minsize::Int=8)::Int where {T<:Real}
    i = 0
    n = size(x,1)
    while n > minsize
        i += 1
        n = fld(n, 2)
    end
    return i
end

function genrsdata(x::Array{T})::Matrix{T} where {T<:Real}
    depth = npieces(x)
    if depth == 0
        return [size(x,1) estimate_rsrange(x)]
    end
    rsdata = zeros(0,2)
    rsdata = [rsdata; [size(x,1) estimate_rsrange(x)]]
    a, b = divide(x)
    rsdata = [rsdata; genrsdata(a)]
    rsdata = [rsdata; genrsdata(b)]
    return rsdata
end

function divide(x::A)::Tuple{A, A} where {A<:AbstractArray}
    n = size(x,1)
    @assert n>=2
    h = n/2
    a = x[1:floor(Int,h)]
    if floor(h) == h
        b = x[(ceil(Int,h)+1):end]
    else
        b = x[ceil(Int,h):end]
    end
    return a, b
end

function estimate_hurst(x::Array{T}; intercept::Bool=false)::T where {T<:Real}
    RS = genrsdata(x)
    RS = RS[sortperm(RS[:,1]),:]
    xx = log2.(RS[:,1])
    yy = log2.(RS[:,2])
    if intercept
        _, beta = [ones(size(xx)) xx]\yy
    else
        beta = xx\yy
    end
    return beta
end

"""
```
rsrange(x::Array{T}; n::Int=100, cumulative::Bool=false, intercept::Bool=true)
```

Compute the rescaled range of a time series
"""
function rsrange(x::Array{T}; n::Int=100, cumulative::Bool=false, intercept::Bool=true) where {T<:Real}
    @assert size(x,1) >= n
    return runfun(x, estimate_rsrange; n=n, cumulative=cumulative)
end

"""
```
hurst(x::Array{T}; n::Int=100, cumulative::Bool=false, intercept::Bool=false)
```

Compute the Hurst exponent of a time series
"""
function hurst(x::Array{T}; n::Int=100, cumulative::Bool=false, intercept::Bool=false) where {T<:Real}
    @assert size(x,1) >= n
    return runfun(x, estimate_hurst; n=n, cumulative=cumulative, intercept=intercept)
end


# ==============================================================================
# Xavier Gandibleux - November 2021
#   Implemented in Julia 1.6

# ==============================================================================
# Parser of MO-MKP instances generated by E. Zitzler and M. Laumanns available at
# https://sop.tik.ee.ethz.ch/download/supplementary/testProblemSuite/?page=testProblem.php#source
#
#   NB1: Number of objectives = number of constraints
#   NB2: All coefficients P, W are integer, all variables are binary
#   NB3: All coefficients ω are floored to the integer

# ==============================================================================
# Datastructure of a multi-objective multi-dimensionnal KP with 0/1 variables
struct _MOMKP
    P  :: Matrix{Int} # profit of items for the objectives, k=1..p, j=1..n
    W  :: Matrix{Int} # weight of items for the constraints, i=1..m, j=1..n
    ω  :: Vector{Int} # capacity of knapsacks, i=1..m
end

# ==============================================================================
function readInstanceMOMKPformatZL(verbose::Bool, fname::String)

    f=open(fname)
        lines = readlines(f)
    close(f)

    # Extract the problem dimensions from the file name ------------------------
    line=(split(fname,"."))
    #line = split(lines[1]," ")
    m = parse(Int, line[3])
    p = m
    n = parse(Int, line[2])
    momkp = _MOMKP(zeros(Int,p,n),zeros(Int,m,n),zeros(Int,m))
    !verbose ? nothing : println("Instance of MKP : ",m," couples objective/constraint and ",n," variables")

    # Extract the data of the MKP from the following lines ---------------------
    i=0; j=0
    for l = 2:length(lines)
        if occursin("knapsack ", lines[l])
            i = i + 1
            j = 0
            !verbose ? nothing : println("Couple objective/constraint n°",i)
        elseif occursin("capacity", lines[l])
            line = split(lines[l],"+")
            momkp.ω[i] = floor(Int, parse(Float64, line[2]))
            !verbose ? nothing : println("ω[",i,"]=",momkp.ω[i])
        elseif occursin("item", lines[l])
            j = j + 1
        elseif occursin("weight", lines[l])
            line = split(lines[l],"+")
            momkp.W[i,j] = parse(Int, line[2])
            !verbose ? nothing : println("w[",i,",",j,"]=",momkp.W[i,j])
        elseif occursin("profit", lines[l])
            line = split(lines[l],"+")
            momkp.P[i,j] = parse(Int, line[2])
            !verbose ? nothing : println("p[",i,",",j,"]=",momkp.P[i,j])
        end
    end

    return momkp
end

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

# Example on how to use it :
# fname="knapsack.100.3"; verbose = true
# instanceZitzler = readInstanceMOMKPformatZL(verbose,fname)


__precompile__()

module MultiscaleGraphSignalTransforms

include("GraphSignal.jl")
include("GraphPartition.jl")
include("BasisSpecification.jl")
include("common.jl")
include("GHWT.jl")
include("GHWT_2d.jl")
include("GHWT_tf_1d.jl")
include("GHWT_tf_2d.jl")
include("HGLET.jl")
include("GraphSig_Plot.jl")
include("gplot.jl")
include("partition_fiedler.jl")

using Reexport
@reexport using .GraphSignal, .GraphPartition, .BasisSpecification, .GHWT, .GHWT_2d, .GHWT_tf_1d, .GHWT_tf_2d, .HGLET

export dvec2dmatrix, dmatrix2dvec, levlist2levlengths!, bsfull, bs_haar, bs_level, bs_walsh, dvec_Threshold, rs_to_region, GraphSig_Plot, gplot, gplot!, partition_fiedler
export cost_functional, dmatrix_flatten, dmatrix_ldb_flatten

## export functions of NGWP.jl
using LinearAlgebra, SparseArrays, LightGraphs, SimpleWeightedGraphs, Clustering
using JuMP, Clp, Optim, Statistics, QuadGK, Arpack
import Plots: plot, plot!, scatter, scatter!
import StatsBase: crosscor

include("dualgraph.jl")
include("eigDAG_Distance.jl")
include("eigHAD_Distance.jl")
include("eigROT_Distance.jl")
include("eigsROT_Distance.jl")
include("eigTSD_Distance.jl")
include("helpers.jl")
include("LP-HGLET.jl")
include("LP-NGWP.jl")
include("natural_distances.jl")
include("utils.jl")
include("varimax.jl")
include("NGWF.jl")
include("NGWP.jl")
include("PC-NGWP.jl")
include("SunFlowerGraph.jl")
include("VM-NGWP.jl")

export eigDAG_Distance, eigHAD_Distance, eigHAD_Affinity
export eigROT_Distance, ROT_Distance, eigEMD_Distance
export eigsROT_Distance
export eigTSD_Distance, K_functional
export natural_eigdist
export SunFlowerGraph, dualgraph
export pc_ngwp, pairclustering, mgslp
export vm_ngwp, varimax
export lp_ngwp, rising_cutoff, find_pairinds, pair_inds_shadding, lp_ngwp_analysis
export LPHGLET_Synthesis, LPHGLET_Analysis_All, HGLET_dictionary, LPHGLET_dictionary
export HGLET_DST4_Synthesis, LPHGLET_DST4_Synthesis, HGLET_DST4_dictionary, LPHGLET_DST4_dictionary
export HGLET_DST4_Analysis, LPHGLET_DST4_Analysis
export unitary_folding_operator, keep_folding!
export ngwp_analysis, ngwp_bestbasis, NGWP_jkl
export natural_eigdist
export nat_spec_filter, ngwf_all_vectors, rngwf_all_vectors, ngwf_vector, frame_approx, rngwf_lx
export scatter_gplot, scatter_gplot!, stem, stem!
export standardize_eigenvectors!, spike, characteristic, χ, sort_wavelets, transform2D
export getall_expansioncoeffs, approx_error_plot, getall_expansioncoeffs2, approx_error_plot2

end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.7.0
#     language: julia
#     name: julia-1.7
# ---

# %% [markdown]
# https://twitter.com/mkashi/status/1469153893179596801

# %%
using Plots

f(x) = x^3 - 3x^2 + 3x - 3
df(x) = 3x^2 - 6x + 3

x = 2.1
dx = @. 2.0 ^ (-60:4)

finitediff = @. (f(x + dx) - f(x)) / dx
abserr_finirediff = @. abs(finitediff - df(x))

plot(dx, abserr_finirediff; xscale=:log10, yscale=:log10, label="")
plot!(; xtick=@.(1e1^(-20:2:2)), ytick=@.(1e1^(-20:2:2)))
plot!(; xlabel="dx", ylabel="absolute error")

# %%
df(x)

# %%
using ForwardDiff
ForwardDiff.derivative(f, x)

# %%
@show finitediff;

# %%
using FiniteDifferences
@show [central_fdm(k, 1)(f, x) for k in 2:10];

# %%

0025/finite difference and autimatic differentiation.jl


# # Foundations of Bilevel Programming: Example 6
# This example is from the book Princeton Handbook of Test Problems in Local and Global Optimization
# Dempe, Chapter 9.3.2 -parg 221 [url](https://www.springer.com/gp/book/9780792358015)

# Here, only the second level is described

# Model of the problem
# First level
# ```math
# \min (x-5)^2+(2y+1)^2,\\
# \notag s.t.\\
# x \geq 0,\\
# y \geq 0,\\
# ```
# Second level
# ```math
# \min (y-1)^2-1.5xy,\\
# \notag s.t.\\
# -3x+y \leq -3,\\
# x-0.5y \leq 4,\\
# x+y \leq 7,\\
# ```

using BilevelJuMP
using Ipopt
using JuMP
using Test

model = BilevelModel(Ipopt.Optimizer, mode = BilevelJuMP.ProductMode(1e-9))

# Global variables
atol = 1e-3

# First we need to create all of the variables in the upper and lower problems:

# Upper level variables
@variable(Upper(model), x)

#Lower level variables
@variable(Lower(model), y)

# Then we can add the objective and constraints of the upper problem:

# Upper level objecive function
@objective(Upper(model), Min, (x-5)^2 + (2y+1)^2)

# Upper level constraints
@constraint(Upper(model), x >= 0)
@constraint(Upper(model), y >= 0) # only in lowrrin GAMS

# Followed by the objective and constraints of the lower problem:

# Lower objective function
@objective(Lower(model), Min, (y-1)^2 -1.5*x*y)

# Lower constraints
@constraint(Lower(model), -3x +    y <= -3)
@constraint(Lower(model),   x - 0.5y <= 4)
@constraint(Lower(model),   x +    y <= 7)

# Initial Starting conditions  #src

# Now we can solve the problem and verify the solution again that reported by
# Dempe.

optimize!(model)
primal_status(model)
termination_status(model)

@test value(x) ≈ 1 atol=atol
@test value(y) ≈ 0 atol=atol


using Symbolics
using Test
using Latexify
using ReferenceTests

@variables x y z u(x)
Dx = Differential(x)

# issue 260
@test Symbolics._toexpr(3*x/y) == :((3x) / y)
@test Symbolics._toexpr(3*x^y) == :(3x^y)

@test_reference "latexify_refs/inverse.txt" latexify(x^-1)

@test_reference "latexify_refs/frac1.txt" latexify((z + x*y^-1) / sin(z))
@test_reference "latexify_refs/frac2.txt"  latexify((3x - 7y*z^23) * (z - z^2) / x)

@test_reference "latexify_refs/minus1.txt" latexify(x - y)
@test_reference "latexify_refs/minus2.txt" latexify(x - y * z)
@test_reference "latexify_refs/minus3.txt" latexify(sin(x+y-z))

@test_reference "latexify_refs/unary_minus1.txt" latexify(-y)
@test_reference "latexify_refs/unary_minus2.txt" latexify(-y * z)

@test_reference "latexify_refs/derivative1.txt" latexify(Dx(y))
@test_reference "latexify_refs/derivative2.txt" latexify(Dx(u))
@test_reference "latexify_refs/derivative3.txt" latexify(Dx(x^2 + y^2 + z^2))

@test_reference "latexify_refs/stable_mul_ordering1.txt" latexify(x * y)
@test_reference "latexify_refs/stable_mul_ordering2.txt" latexify(y * x)

@test_reference "latexify_refs/equation1.txt" latexify(x ~ y + z)
@test_reference "latexify_refs/equation2.txt" latexify(x ~ Dx(y + z))

@test_reference "latexify_refs/equation_vec1.txt" latexify([
    x ~   y +  z
    y ~   x - 3z
    z ~ -5x + 4y
])
@test_reference "latexify_refs/equation_vec2.txt" latexify([
    Dx(u) ~   z
    Dx(y) ~   y*x
])


#### Kahan matrix

export Kahan

immutable Kahan{T<:Number,T1<:Number} <: AbstractMatrix{T}
    m::Int # dimension
    n::Int # dimension
    theta::T # angle
    pert::T1 # perturbation is pert*eps()
end # immutable

# Define its size
size(A::Kahan, r::Int) = r==1 ? A.m : A.n
size(A::Kahan) = A.m, A.n

# Index into a Kahan
function getindex(A::Kahan,i::Integer,j::Integer)
    m=minimum(size(A))
    t=tan(A.theta)
    c=1.0/sqrt(1.0+t^2)
    s=t*c
    if i>m; return 0.0
    elseif i>j; return 0.0
    elseif i==j; return s^(i-1)+A.pert*eps()*(m-i+1)
    else return -c*s^(i-1)
    end
end # getindex

# Dense version of Kahan
full(A::Kahan) =[A[i,j] for i=1:size(A,1), j=1:size(A,2)]


using Plots

# Problem parameters
α = 0.3
δ = 0.1
β = 0.9

# Capital grid
kupper = 2
klower = 0.001
n = 10
kgrid = collect(range(klower, stop = kupper, length = n))

# Iteration parameters
tolerance = 0.001
imax = 1000

# Initialize values
vnew = zeros(length(kgrid)) # initial value function guess
v = vnew .+ 2*tolerance # initialize v in a way that ensures loop will start
global cartesianindex = Array{CartesianIndex{2}, n}
i = 1

while maximum(abs.(v - vnew)) > tolerance && i<=imax
    v = vnew;
    c = zeros(n, n);
    for i in 1:n
        for j in 1:n
            c[i,j] = kgrid[i]^α + (1-δ)*kgrid[i] - kgrid[j];
            if c[i,j] <= 0
                c[i,j] = 0;
            end
        end
        (vnew, cartesianindex) = findmax(log.(c) .+ β*v', dims = 2);
    end
    i += 1
end

scatter(kgrid, vnew, title = "v(k)")

# Policy function
kprimeindex = getindex.(cartesianindex, 2)
kprime = kgrid[kprimeindex]
scatter(kgrid, kprime, title = "k'(k)")


x = randn(10)
y = randn(10)
r = x .- y

@testset "Robust Loss" begin
    δ = 0.5
    rlδ = RobustLoss(Huber(δ))
    @test rlδ isa RobustLoss{HuberRho{δ}}
    @test rlδ(r) == rlδ(x, y) == sum(ifelse(abs(rᵢ)≤δ, rᵢ^2/2, δ*(abs(rᵢ)-δ/2)) for rᵢ in r)

    rlδ = RobustLoss(Andrews(δ))
    p1(z) = -cos(π * z/δ)/(π/δ)^2
    p2(z) = p1(δ)
    @test rlδ(r) ≈ sum(ifelse(abs(rᵢ)≤δ, p1.(rᵢ), p2.(rᵢ)) for rᵢ in r)

    rlδ = RobustLoss(Bisquare(δ))
    p1(z) = δ^2/6 * (1-(1-(z/δ)^2)^3)
    p2(z) = p1(δ)
    @test rlδ(r) ≈ sum(ifelse(abs(rᵢ)≤δ, p1.(rᵢ), p2.(rᵢ)) for rᵢ in r)

    rlδ = RobustLoss(Logistic(δ))
    @test rlδ(r) ≈ sum(δ^2 * log(cosh(rᵢ/δ)) for rᵢ in r)

    rlδ = RobustLoss(FairRho(δ))
    @test rlδ(r) ≈ sum(δ^2 * (abs(rᵢ)/δ - log(1+abs(rᵢ/δ))) for rᵢ in r)

    rlδ = RobustLoss(Talwar(δ))
    p1(z) = z^2/2
    p2(z) = p1(δ)
    @test rlδ(r) ≈ sum(ifelse(abs(rᵢ)≤δ, p1.(rᵢ), p2.(rᵢ)) for rᵢ in r)
end


using EasyFit
using Plots

x = sort(rand(10)) .- 0.5
y = sort(rand(10)) .- 0.5

plot(layout=(3,2))
plot!(framestyle=:box,grid=false)

sp=1
plot!(title="Linear fit",subplot=sp)
scatter!(x,y,label="",subplot=sp)
fit = fitlinear(x,y)
plot!(fit.x,fit.y,label="",linewidth=2,subplot=sp)

sp=2
y2 = y.^2
plot!(title="Quadratic fit",subplot=sp)
scatter!(x,y2,label="",subplot=sp)
fit = fitquad(x,y2)
plot!(fit.x,fit.y,label="",linewidth=2,subplot=sp)

sp=3
y2 = y.^3
plot!(title="Cubic fit",subplot=sp)
scatter!(x,y2,label="",subplot=sp)
fit = fitcubic(x,y2)
plot!(fit.x,fit.y,label="",linewidth=2,subplot=sp)

sp=4
y2 = 0.3*exp.(5*y)
plot!(title="Mono-exponential fit",subplot=sp)
scatter!(x,y2,label="",subplot=sp)
fit = fitexp(x,y2)
plot!(fit.x,fit.y,label="",linewidth=2,subplot=sp)

sp=5
y2 = 0.3*exp.(5*y) + 0.7*exp.(3*y2)
plot!(title="Bi-exponential fit",subplot=sp)
scatter!(x,y2,label="",subplot=sp)
fit = fitexp(x,y2,n=2)
plot!(fit.x,fit.y,label="",linewidth=2,subplot=sp)

sp=6
y2 = y.^5
plot!(title="Spline fit",subplot=sp)
scatter!(x,y2,label="",subplot=sp)
fit = fitexp(x,y2,n=2)
plot!(fit.x,fit.y,label="",linewidth=2,subplot=sp)

plot!(size=(600,600))


"""
    binary_brier_score(ytrue, yscore)

Computes the binary formulation of the Brier score, defined as:

```math
\\frac{1}{N}\\sum\\limits _{t=1}^{N}(f_t-o_t)^2 \\,\\!
```

Lower values are better. Best value is 0.
"""
function binary_brier_score(
        ytrue::AbstractVector{<:Integer},
        yscore::AbstractVector{<:AbstractFloat},
        )
    if length(ytrue) != length(yscore)
        error("length(ytrue) != length(yscore)")
    end
    if length(ytrue) == 0
        error("length(ytrue) == 0")
    end
    result = mean_square_error(ytrue, yscore)
    return result
end


#for every different type of linkfunction, internally julia will find the right apply_inverse_link function

# This super type of all response distribution types
abstract type InverseLinkFunction end

##LINK FUNCTIONS##

"""inverse cauchit link."""

struct CauchitLink <: InverseLinkFunction
end

function cauchit_inverse_link(x)
  return atan(x) / pi + one(x) / 2
end

"""inverse cloglog link."""

struct CloglogLink <: InverseLinkFunction
end

function cloglog_inverse_link(x)
  return one(x) - exp(-exp(x))
end

"""inverse identity link."""

struct IdentityLink <: InverseLinkFunction
end

function identity_inverse_link(x)
  return x
end

"""inverse inverse link."""

struct InverseLink <: InverseLinkFunction
end

function inverse_inverse_link(x)
  return one(x) / x
end

"""inverse logit link."""

struct LogitLink <: InverseLinkFunction
end

function logit_inverse_link(x)
  return one(x) / (one(x) + exp(-x))
end

"""inverse log link."""

struct LogLink <: InverseLinkFunction
end

function log_inverse_link(x)
  return exp(x)
end

"""inverse probit link."""

struct ProbitLink <: InverseLinkFunction
end

function probit_inverse_link(x)
  return (one(x) + erf(x / sqrt(2 * one(x)))) / 2
end

"""inverse sqrt link."""

struct SqrtLink <: InverseLinkFunction
end

function sqrt_inverse_link(x)
  return x * x
end

##APPLY INVERSE LINK FUNCTIONS

apply_inverse_link(μ, link::LogLink) = log_inverse_link.(μ)

apply_inverse_link(μ, link::IdentityLink) = identity_inverse_link.(μ)

apply_inverse_link(μ, link::SqrtLink) = sqrt_inverse_link.(μ)

apply_inverse_link(μ, link::ProbitLink) = probit_inverse_link.(μ)

apply_inverse_link(μ, link::LogitLink) = logit_inverse_link.(μ)

apply_inverse_link(μ, link::InverseLink) = inverse_inverse_link.(μ)

apply_inverse_link(μ, link::CauchitLink) = cauchit_inverse_link.(μ)

apply_inverse_link(μ, link::CloglogLink) = cloglog_inverse_link.(μ)


using vOptGeneric
using Test, LinearAlgebra
using GLPK, GLPKMathProgInterface

function run_tests()
        m = vModel(solver = GLPKSolverMIP())

        @variable(m, x[1:6,1:6], Bin)
        @variable(m, y, Bin)
        M = 50
        p1 =   [M 4 5 M M M;
                M M 2 1 2 7;
                M M M 5 2 M;
                M M 5 M M 3;
                M M M M M 4;
                M M M M M M]
        p2 =   [M 3 1 M M M;
                M M 1 4 2 2;
                M M M 1 7 M;
                M M 1 M M 2;
                M M M M M 2;
                M M M M M M]

        @addobjective(m, Min , sum(x[i,j]*p1[i,j] for i=1:6, j=1:6))
        @addobjective(m, Min , sum(x[i,j]*p2[i,j] for i=1:6, j=1:6))
        @constraint(m, sum(x[1,j] for j = 1:6) == 1)
        @constraint(m, sum(x[i,6] for i = 1:6) == 1)
        @constraint(m, cstr[i=2:5], sum(x[i,j] for j = 1:6) - sum(x[j,i] for j =1:6) == 0)

        solve(m, method=:epsilon)

        Y_N = getY_N(m)
        f1 = map(x -> x[1], Y_N)
        f2 = map(x -> x[2], Y_N)

        @test f1 == [8.0,10.0,11.0,13.0]
        @test f2 == [9.0,7.0,5.0,4.0]
        @test getvalue(x, 1) == [0.0 1.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 1.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 1.0; 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0]
        printX_E(m)

        writeMOP(m, "test.txt")
        m = parseMOP("test.txt", solver = GLPKSolverMIP())
        solve(m, method=:dichotomy)
        Y_N = getY_N(m)
        f1 = map(x -> x[1], Y_N)
        f2 = map(x -> x[2], Y_N)
        @test f1 == [8.0, 11.0, 13.0]
        @test f2 == [9.0, 5.0, 4.0]

        m = vModel(solver = GLPKSolverMIP())
        # using CPLEX
        # m = vModel(solver = CplexSolver())

        p1 = [77,94,71,63,96,82,85,75,72,91,99,63,84,87,79,94,90,60,69,62]
        p2 = [65,90,90,77,95,84,70,94,66,92,74,97,60,60,65,97,93,60,69,74]
        w = [80,87,68,72,66,77,99,85,70,93,98,72,100,89,67,86,91,79,71,99]
        c = 900
        size = length(p1)

        @variable(m, x[1:size], Bin)
        @addobjective(m, Max, dot(x, p1))
        @addobjective(m, Max, dot(x, p2))
        @constraint(m, dot(x, w) <= c)

        solve(m, method=:chalmet)

        Y_N = getY_N(m)
        f1 = map(x -> x[1], Y_N)
        f2 = map(x -> x[2], Y_N)

        @test f1 == [918.0, 924.0, 927.0, 934.0, 935.0, 940.0, 943.0, 948.0, 949.0, 952.0, 955.0]
        @test f2 == [984.0, 975.0, 972.0, 971.0, 947.0, 943.0, 940.0, 939.0, 915.0, 909.0, 906.0]
        @test getvalue(x, 1) == [0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0]

        solve(m, method=:lex)
        Y_N = getY_N(m)
        @test Y_N == [[955.0, 906.0], [918.0, 984.0]]
end

run_tests()


#using FemAdjoint
using Test, JLD2, UnicodePlots
using LinearAlgebra, ForwardDiff, Arpack
using SparseArrays, SparseDiffTools, Symbolics
using FemAdjoint

# load mesh and data for tests
meshfile = "$(@__DIR__)/data/diskmesh.jld2"
JLD2.@load(meshfile, p, t)
np, nt = size(p, 1), size(t, 1)
ε = 1e-5
u, v = 2 * rand(np) .- 1, 2 * rand(np) .- 1
x, dp = p, 2 * rand(size(p)...) .- 1
xp, xm = p .+ ε * dp, p .- ε * dp
SK, SM = FemAdjoint.assembKM_P12D(p, t)
IK, JK = FemAdjoint.indKM_sparse(t)
IM, JM = FemAdjoint.indKM_sparse(t)
K = sparse(IK, JK, SK, np, np)
M = sparse(IM, JM, SM, np, np)

@testset "mesh data" begin
    @test size(p, 1) == 560
    @test size(t, 1) == 1045
end

@testset "first assembly" begin
    @test length(IK) == length(JK)
    @test length(IK) == length(SK)
    @test length(IM) == length(JM)
    @test length(IM) == length(SM)
end

@testset "test eval function" begin
    u, v = 2 * rand(np) .- 1, 2 * rand(np) .- 1
    @test abs(dot(K * u, v) - FemAdjoint.costKproduv(p, t, u, v)) < 1e-8
end

f = x -> FemAdjoint.costKproduv(x, t, u, v)
fp, fm = f(xp), f(xm)
@time g = ForwardDiff.gradient(f, x)
@time g = ForwardDiff.gradient(f, x)

@testset "test eval gradient dot function" begin
    @test abs(dot(g, dp) - (fp - fm) / (2 * ε)) < 1e-4
end
display(K)
print("\n")
println("directional derivative $(abs(dot(g, dp)))")

#-------------------------------------------------------------------------------
# square mesh / eigenvalue test
meshfile = "$(@__DIR__)/data/squaremesh.jld2"
JLD2.@load(meshfile, p, t, jacK, jacM, ps, ts, jacKs, jacMs)
np, nt = size(p, 1), size(t, 1)
SK, SM = FemAdjoint.assembKM_P12D(p, t)
IK, JK = FemAdjoint.indKM_sparse(t)
IM, JM = FemAdjoint.indKM_sparse(t)
K = sparse(IK, JK, SK, np, np)
M = sparse(IM, JM, SM, np, np)
Ib = unique(FemAdjoint.btri(t)[:])
K += sparse(Ib, Ib, fill(1e8, length(Ib)), np, np) # Dirichlet conditions

λ, U = eigs(K, M, nev = 10, which = :SM, tol = 0., maxiter = 30_000)
@testset "test eigenvalue" begin
    @test abs(λ[1] - pi ^ 2 / 2) < 1e-2
end
# plot
FemAdjoint.plotunicode(U[:, 4])

#-------------------------------------------------------------------------------
#p, t, jac = ps, ts, jacs
y = similar(FemAdjoint.assembKM_P12D(p, t)[1])
SK, SM = FemAdjoint.assembKM_P12D(p, t)
SKi, SMi = similar(SK), similar(SM)
FemAdjoint.assembK_P12D_inplace!(p, t, SKi)
FemAdjoint.assembM_P12D_inplace!(p, t, SMi)
@testset "test in place assemble" begin
    @test maximum(abs.(SK .- SKi)) < 1e-8
    @test maximum(abs.(SM .- SMi)) < 1e-8
end
fK(y, x) = FemAdjoint.assembK_P12D_inplace!(x, t, y)
# to speedup test use a precomputed jac matrix
#@time sparsity_pattern = Symbolics.jacobian_sparsity(fs, y, p[:])
#jacK = Float64.(sparse(sparsity_pattern))
colors = matrix_colors(jacK)
@time J = forwarddiff_color_jacobian!(jacK, fK, p[:], colorvec = colors)
@time J = forwarddiff_color_jacobian!(jacK, fK, p[:], colorvec = colors)
display(J)

#-------------------------------------------------------------------------------
# test adjoint approach for costKproduv cost function
include("test_costKproduv.jl")
#-------------------------------------------------------------------------------
# test adjoint approach for costnormU cost function
include("test_costnormU.jl")
#-------------------------------------------------------------------------------
# test direct adjoint approach for costnormU cost function (definitely simpler!!)
include("test_costnormU_direct.jl")


# https://github.com/stan-dev/blob/master/basic_estimators/normal_mixture.stan

# normal mixture, unknown proportion and means, known variance
# p(y|mu,theta) = theta * Normal(y|mu[1],1) + (1-theta) * Normal(y|mu[2],1);

const simplenormalmixturemodel = "
data {
  int<lower=0>  N;
  real y[N];
}
parameters {
  real<lower=0,upper=1> theta;
  real mu[2];
}
transformed parameters {
  real log_theta;
  real log_one_minus_theta;

  log_theta <- log(theta);
  log_one_minus_theta <- log(1.0 - theta);
}
model {
  theta ~ uniform(0,1); // equivalently, ~ beta(1,1);
  for (k in 1:2)
    mu[k] ~ normal(0,10);
  for (n in 1:N)
    increment_log_prob(
      log_sum_exp(log_theta + normal_log(y[n],mu[1],1.0),
                  log_one_minus_theta + normal_log(y[n],mu[2],1.0)));
}
"


@doc """
  Computes the error-free addition `a + b = x + y`, where `x = fl(a + b)`, for
  general floating-point numbers `a` and `b`. The algorithm requires 6 flops;
  twice as many as `err_fast_add`. However, failed branch prediction because of
  the required comparison `|a| ≥ |b|` in the general case, cause this algorithm
  to be faster on modern pipelined architectures.

  References:

    * D.E. Knuth, The Art of Computer Programming, 3rd ed., vol. 2,
      Addison-Wesley, 1998.
    * D. Møller, Quasi double-precision in floating-point addition, BIT
      Numerical Mathematics 5, pp. 37–50, 1965.
""" ->
# TODO: Extend the definition to floating-point arrays
function err_add{T<:IEEE754}(a::T, b::T)
  x = a + b
  c = x - a
  y = (a - (x-c)) + (b-c)
  x, y
end

@doc """
  Computes the error-free addition `a + b = x + y`, where `x = fl(a + b)`, for
  floating-point numbers `a` and `b` such that the exponent of `a` is larger
  than or equal to that of `b`. If `|a| ≥ |b|`, this condition is always
  satisfied. The algorithm requires 3 flops.

  References:

    * T.J. Dekker, A floating-point technique for extending the available
      precision, Numerische Mathematik 18, pp. 224–242, 1971.
  """ ->
# TODO: Extend the definition to floating-point arrays
function err_fast_add{T<:IEEE754}(a::T, b::T)
  x = a + b
  y = b - (x-a)
  x, y
end


"""
    N, E = latlon2yx(φ₀, λ₀, φ, λ)

Coordinate transformation between geographic and plane rectangular coordinates on the Gauss-Krüger Projection
This function is based on the following documents:
        https://www.gsi.go.jp/common/000061216.pdf
        https://vldb.gsi.go.jp/sokuchi/surveycalc/surveycalc/algorithm/xy2bl/xy2bl.htm
        https://vldb.gsi.go.jp/sokuchi/surveycalc/surveycalc/algorithm/bl2xy/bl2xy.htm

### input arguments
- φ₀, λ₀ : latitude of origin and central meridian
- φ , λ  : geographic latitude and longitude
### return values
- N, E   : plane rectangular coordinates northing and easting

### Examples
```julia-repl
julia> # ANSWER: y = 11543.6883, x = 22916.2436

julia> latlon2yx(36.0, 139.0+5.0/6.0, 36.103774791666666, 140.08785504166664)
(11543.688321484718, 22916.24355431881)
```
"""
function latlon2yx(φ₀, λ₀, φ, λ)

    φ₀ = deg2rad(φ₀)
    λ₀ = deg2rad(λ₀)
    φ = deg2rad(φ)
    λ = deg2rad(λ)

    S̅φ₀ = arclength(φ₀)

    ncoef1 = 2sqrt(n)/(1+n)

    t = sinh(atanh(sin(φ)) - ncoef1*atanh(ncoef1*sin(φ)))
    t̅ = sqrt(1+t^2)
    λc = cos(λ-λ₀)
    λs = sin(λ-λ₀)
    ξ′ = atan(t,λc)
    η′ = atanh(λs/t̅)

    # α terms
    αcoef = [1/2   -2/3    5/16        41/180     -127/288; # α₁
             0.0  13/48    -3/5      557/1440      281/630; # α₂
             0.0    0.0  61/240      -103/140  15061/26880; # α₃
             0.0    0.0     0.0  49561/161280     -179/168; # α₄
             0.0    0.0     0.0           0.0  34729/80460; # α₅
            ]
    αj = αcoef * [n^k for k=1:5]

    N = A̅*(ξ′ + sum([αj[j]sin(2j*ξ′)cosh(2j*η′) for j=1:5])) - S̅φ₀
    E = A̅*(η′ + sum([αj[j]cos(2j*ξ′)sinh(2j*η′) for j=1:5]))

    return [N E]
end
# -------------------------------------
"""
    E, N = lonlat2xy(λ₀, φ₀, λ, φ)

See [`lanlon2yx`](@ref)
"""
lonlat2xy(λ₀, φ₀, λ, φ) = reverse(latlon2yx(φ₀, λ₀, φ, λ), dims=2)

# -------------------------------------
function latlon2yx_ja(zone::Integer, φ, λ)
    !(0 < zone < 20) && error("zone ID must be a integer in the range of 1 to 19.")
    return latlon2yx(Origin_LatLon_Japan[zone,:]..., φ, λ)
end
# -------------------------------------
lonlat2xy_ja(zone::Integer, λ, φ)  = reverse(latlon2yx_ja(zone, φ, λ), dims=2)
# -------------------------------------


const two27 = Float64(1<<27)
const two12 = Float32(1<<12)
const two6  = Float16(1<<6)

"""
    extractscalar

Split a value relative to an integral power of 2

hi is the high order part of p, lo is the residual part
hi + lo === p
"""

for (T,P) in ((:Float64, 53), (:Float32, 24), (:Float16, 11))
  @eval begin
    @inline function extractscalar(value::$T)
        powoftwo =  ldexp(one($T), ($P- exponent(value)) >> 1) # P == precision(Float64)
        return extractscalar(powoftwo, value)
    end
  end
end

function extractscalar(powoftwo::T, value::T) where {T}
   hi = (powoftwo + value) - powoftwo
   lo = value - hi
   return hi, lo
end

"""
    two_cube(a)

Computes `s = fl(a*a*a)` and `e1 = err(a*a*a)`.
"""
function two_cube(value::T) where {T}
    hi, lo = extractscalar(value)
    hhh  = three_prod(hi, hi, hi)
    hhl  = three_prod(hi, hi, 3*lo)
    hll  = three_prod(hi, 3*lo, lo)
    lll  = three_prod(lo, lo, lo)

    himh = four_sum(hhh[1], hhl[1], hhh[2], hll[1])
    mllo = four_sum(hhl[2], hll[2], lll[1], lll[2])
    hilo = four_sum(himh[1], himh[2], mllo[1], mllo[2])

    return hilo[1], hilo[2]
end

"""
    three_cube(a)

Computes `s = fl(a*a*a)` and `e1 = err(a*a*a), e2 = err(e1)`.
"""
function three_cube(value::T) where {T}
    hi, lo = extractscalar(value)
    hhh  = three_prod(hi, hi, hi)
    hhl  = three_prod(hi, hi, 3*lo)
    hll  = three_prod(hi, 3*lo, lo)
    lll  = three_prod(lo, lo, lo)

    himh = four_sum(hhh[1], hhl[1], hhh[2], hll[1])
    mllo = four_sum(hhl[2], hll[2], lll[1], lll[2])
    hilo = four_sum(himh[1], himh[2], mllo[1], mllo[2])

    return hilo[1], hilo[2], hilo[3]
end

"""
    four_cube(a)

Computes `s = fl(a*a*a)` and `e1 = err(a*a*a), e2 = err(e1), e3 = err(e2)`.
"""
function three_cube(value::T) where {T}
    hi, lo = extractscalar(value)
    hhh  = three_prod(hi, hi, hi)
    hhl  = three_prod(hi, hi, 3*lo)
    hll  = three_prod(hi, 3*lo, lo)
    lll  = three_prod(lo, lo, lo)

    himh = four_sum(hhh[1], hhl[1], hhh[2], hll[1])
    mllo = four_sum(hhl[2], hll[2], lll[1], lll[2])
    hilo = four_sum(himh[1], himh[2], mllo[1], mllo[2])

    return hilo
end


const ymeyer = [3.478e4, 2.861e4, 2.365e4, 1.963e4, 1.637e4, 1.372e4,
                1.154e4, 9.744e3, 8.261e3, 7.030e3, 6.005e3, 5.147e3,
                4.427e3, 3.820e3, 3.307e3, 2.872e3]

const meyer = let res_init=zeros(16), jac_init=zeros(16,3), x_init=[0.02, 4.e3, 2.5e2]

    function res(x, r)
        for i = 1:16
            ti = Float64(5*i + 45)
            r[i] = x[1] * exp(x[2] / (ti + x[3])) - ymeyer[i]
        end
        return r
    end

    function jac(x, j)
        for i = 1:16
            ti = Float64(5*i + 45)
            u = ti + x[3]
            t = exp(x[2] / u)
            j[i, 1] = t
            j[i, 2] = x[1] * t/u
            j[i, 3] = -x[1] * x[2] * t / (u * u)
        end
        return j
    end

    f(;scale=1, verbose=false, print_steps=false) = testone("meyer", res, jac,
                                                   res_init, jac_init, x_init;
                                                   scale=scale, verbose=verbose,
                                                   print_steps=print_steps)
end


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, Random, StatsBase, DataFrames, Plots
	Random.seed!(1)
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.22"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	lambda, N = 1, 10^6
	xGrid = 0:6
end

# ╔═╡ 84a73e9a-2165-11eb-3620-d7ddf9c07aff
begin
	expDist = Exponential(1/lambda)
	floorData = counts(convert.(Int,floor.(rand(expDist,N))), xGrid)/N
	geomDist = Geometric(1-MathConstants.e^-lambda)
end

# ╔═╡ 84a76f3c-2165-11eb-0eeb-1bb55507bb51
begin
	plot( xGrid, floorData,
		line=:stem, marker=:circle,
		c=:blue, ms=10, msw=0, lw=4,
		label="Floor of Exponential")
	plot!( xGrid, pdf.(geomDist,xGrid),
		line=:stem, marker=:xcross,
		c=:red, ms=6, msw=0, lw=2,
		label="Geometric", ylims=(0,1),
		xlabel="x", ylabel="Probability")
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.22"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═84a73e9a-2165-11eb-3620-d7ddf9c07aff
# ╠═84a76f3c-2165-11eb-0eeb-1bb55507bb51
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3


#=
For details, please visit:
https://computationalmindset.com/en/neural-networks/experiments-with-neural-odes-in-julia.html#exp1

Parametric system to train:
x' = a1 x + b1 y + c1 e^(-d1 t)
y' = a2 x + b2 y + c2 e^(-d2 t)
x(0) = 0
y(0) = 0

to generate dataset use the following param values
a1 = 1.11
b1 = 2.43
c1 = -3.66
d1 = 1.37
a2 = 2.89
b2 = -1.97
c2 = 4.58
d2 = 2.86

Analytical solution is:
https://www.wolframalpha.com/input/?i=x%27+%3D+1.11+x+%2B+2.43+y+%2B+-3.66+e%5E%28-1.37+t%29%3B+y%27+%3D+2.89+x+-1.97+y+%2B+4.58+e%5E+%28-2.86+t%29%3B+x%280%29%3D0%3B+y%280%29%3D0%3B
=#

using Flux, DiffEqFlux, DifferentialEquations, Plots

function parametric_ode_system!(du,u,p,t)
  x, y = u
  a1, b1, c1, d1, a2, b2, c2, d2 = p
  du[1] = dx = a1*x + b1*y + c1*exp(-d1*t)
  du[2] = dy = a2*x + b2*y + c2*exp(-d2*t)
end

true_params = [1.11, 2.43, -3.66, 1.37, 2.89, -1.97, 4.58, 2.86]

an_sol_x(t) =
  -1.38778e-17 * exp(-8.99002 * t) -
  2.77556e-17 * exp(-7.50002 * t) +
  3.28757 * exp(-3.49501 * t) -
  3.18949 * exp(-2.86 * t) +
  0.258028 * exp(-1.37 * t) -
  0.356108 * exp(2.63501 * t) +
  4.44089e-16 * exp(3.27002 * t) +
  1.11022e-16 * exp(4.76002 * t)
an_sol_y(t) =
  -6.23016 * exp(-3.49501 * t) +
  5.21081 * exp(-2.86 * t) +
  1.24284 * exp(-1.37 * t) -
  0.223485 * exp(2.63501 * t) +
  2.77556e-17 * exp(4.76002 * t)

tbegin=0.0
tend=1.5
tstep=0.01
trange = tbegin:tstep:tend
u0 = [0.0,0.0]
tspan = (tbegin,tend)
p = ones(8)

prob = ODEProblem(parametric_ode_system!, u0, tspan, p)

function net()
    solve(prob, Tsit5(), p=p, saveat=trange)
end

dataset_outs = [an_sol_x.(trange), an_sol_y.(trange)]
function loss_func()
  pred = net()
  sum(abs2, dataset_outs[1] .- pred[1,:]) +
  sum(abs2, dataset_outs[2] .- pred[2,:])
end

epochs = 1000
learning_rate = 0.05
data = Iterators.repeated((), epochs)
opt = ADAM(learning_rate)

callback_func = function ()
  loss_value = loss_func()
  println("Loss: ", loss_value)
end
fparams = Flux.params(p)
Flux.train!(loss_func, fparams, data, opt, cb=callback_func)

predict_prob = ODEProblem(parametric_ode_system!, u0, tspan, p)
predict_sol = solve(prob, Tsit5(), saveat=trange)
x_predict_sol = [u[1] for u in predict_sol.u]
y_predict_sol = [u[2] for u in predict_sol.u]

println("Learned parameters:", p)

plot(trange, dataset_outs[1],
    linewidth=2, ls=:dash,
    title="Neural ODEs to fit params",
    xaxis="t",
    label="dataset x(t)",
    legend=true)
plot!(trange, dataset_outs[2],
    linewidth=2, ls=:dash,
    label="dataset y(t)")
plot!(predict_sol.t, x_predict_sol,
    linewidth=1,
    label="predicted x(t)")
plot!(predict_sol.t, y_predict_sol,
    linewidth=1,
    label="predicted y(t)")


# Linear Isotropic Covariance Function

"""
    LinIso <: Kernel

Isotropic linear kernel (covariance)
```math
k(x, x') = xᵀx'/ℓ²
```
with length scale ``ℓ``.
"""
mutable struct LinIso <: Kernel
    "Squared length scale"
    ℓ2::Float64
    "Priors for kernel parameters"
    priors::Array

    """
        LinIso(ll::Float64)

    Create `LinIso` with length scale `exp(ll)`.
    """
    LinIso(ll::Float64) = new(exp(2 * ll), [])
end

struct LinIsoData{D} <: KernelData
    XtX::D
end

function KernelData(k::LinIso, X::MatF64)
    XtX=X'*X
    LinearAlgebra.copytri!(XtX, 'U') # make sure it's symmetric
    LinIsoData(XtX)
end
kernel_data_key(k::LinIso, X::MatF64) = "LinIsoData"

_cov(lin::LinIso, xTy) = xTy ./ lin.ℓ2
function Statistics.cov(lin::LinIso, x::VecF64, y::VecF64)
    K = _cov(lin, dot(x,y))
    return K
end

@inline @inbounds function cov_ij(lin::LinIso, X::MatF64, data::LinIsoData, i::Int, j::Int, dim::Int)
    return _cov(lin, data.XtX[i, j])
end
function Statistics.cov(lin::LinIso, X::MatF64, data::LinIsoData)
    K = _cov(lin, data.XtX)
    return K
end
function cov!(cK::MatF64, lin::LinIso, X::MatF64, data::LinIsoData)
    iℓ2 = 1/lin.ℓ2
    @inbounds @simd for I in eachindex(cK,data.XtX)
        cK[I] = data.XtX[I]*iℓ2
    end
    return cK
end

get_params(lin::LinIso) = Float64[log(lin.ℓ2) / 2]
get_param_names(::LinIso) = [:ll]
num_params(lin::LinIso) = 1

function set_params!(lin::LinIso, hyp::VecF64)
    length(hyp) == 1 || throw(ArgumentError("Linear isotropic kernel only has one parameter"))
    lin.ℓ2 = exp(2 * hyp[1])
end

@inline dk_dll(lin::LinIso, xTy::Float64) = -2 * _cov(lin,xTy)
@inline function dKij_dθp(lin::LinIso, X::MatF64, i::Int, j::Int, p::Int, dim::Int)
    if p==1
        return dk_dll(lin, dotij(X,i,j,dim))
    else
        return NaN
    end
end
@inline function dKij_dθp(lin::LinIso, X::MatF64, data::LinIsoData, i::Int, j::Int, p::Int, dim::Int)
    if p==1
        return dk_dll(lin, data.XtX[i,j])
    else
        return NaN
    end
end


OPT_PROBS = Dict()

#### Himmelblau

OPT_PROBS["himmelblau"] = Dict()
OPT_PROBS["himmelblau"]["array"] = Dict()

function himmelblau!(x)
    fx = (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
    return fx
end
function himmelblau_batched_f!(X)
    F = map(himmelblau!, X)
    return F
end
function himmelblau_batched_f!(F, X)
    map!(himmelblau!, F, X)
    return F
end
function himmelblau_g!(∇f, x)
    ∇f[1] = 4.0 * x[1]^3 + 4.0 * x[1] * x[2] -
        44.0 * x[1] + 2.0 * x[1] + 2.0 * x[2]^2 - 14.0
    ∇f[2] = 2.0 * x[1]^2 + 2.0 * x[2] - 22.0 +
        4.0 * x[1] * x[2] + 4.0 * x[2]^3 - 28.0 * x[2]
    ∇f
end

function himmelblau_fg!(∇f, x)
    ∇f = himmelblau_g!(∇f, x)
    fx = himmelblau!(x)
    return fx, ∇f
end
function himmelblau_fgh!(∇f, ∇²f, x)
    ∇²f = himmelblau_h!(∇²f, x)
    fx, ∇f = himmelblau_fg!(∇f, x)
    return fx, ∇f, ∇²f
end
function himmelblau_h!(∇²f, x)
    ∇²f[1, 1] = 12.0 * x[1]^2 + 4.0 * x[2] - 44.0 + 2.0
    ∇²f[1, 2] = 4.0 * x[1] + 4.0 * x[2]
    ∇²f[2, 1] = ∇²f[1, 2]
    ∇²f[2, 2] = 2.0 + 4.0 * x[1] + 12.0 * x[2]^2 - 28.0
    return ∇²f
end
function himmelblau_hv!(hv, x, v)
    hv[1] = (12.0 * x[1]^2 + 4.0 * x[2] - 44.0 + 2.0)*v[1] + (4.0 * x[1] + 4.0 * x[2])*v[2]
    hv[2] =  (4.0 * x[1] + 4.0 * x[2])*v[1] + (2.0 + 4.0 * x[1] + 12.0 * x[2]^2 - 28.0)*v[2]
    return hv
end

OPT_PROBS["himmelblau"]["array"]["x0"] = [3.0, 1.0]
OPT_PROBS["himmelblau"]["array"]["mutating"] = ScalarObjective(himmelblau!, himmelblau_g!, himmelblau_fg!, himmelblau_fgh!, himmelblau_h!, himmelblau_hv!, himmelblau_batched_f!, nothing)

### Exponential

exponential!(x) = exp((2.0 - x[1])^2) + exp((3.0 - x[2])^2)
function exponential_g!(g, x)
    g[1] = -2.0 * (2.0 - x[1]) * exp((2.0 - x[1])^2)
    g[2] = -2.0 * (3.0 - x[2]) * exp((3.0 - x[2])^2)
    return g
end
function exponential_h!(H, x)
    H[1, 1] = 2.0 * exp((2.0 - x[1])^2) * (2.0 * x[1]^2 - 8.0 * x[1] + 9)
    H[1, 2] = 0.0
    H[2, 1] = 0.0
    H[2, 2] = 2.0 * exp((3.0 - x[2])^2) * (2.0 * x[2]^2 - 12.0 * x[2] + 19)
    return H
end
function exponential_hv!(Hv, x, v)
    Hv[1, 1] = (2.0 * exp((2.0 - x[1])^2) * (2.0 * x[1]^2 - 8.0 * x[1] + 9))*v[1]
    Hv[2, 2] = (2.0 * exp((3.0 - x[2])^2) * (2.0 * x[2]^2 - 12.0 * x[2] + 19))*v[2]
    return Hv
end
function exponential_fg!(g, x)
    fx = exponential!(x)
    g = exponential_g!(g, x)
    return fx, g
end
function exponential_fgh!(g, H, x)
    fx, g = exponential_fg!(g, x)
    H = exponential_h!(H, x)
    return fx, g, H
end

OPT_PROBS["exponential"] = Dict()
OPT_PROBS["exponential"]["array"] = Dict()
# Byttet om på x og H
OPT_PROBS["exponential"]["array"]["x0"] = [0.0, 0.0]
OPT_PROBS["exponential"]["array"]["mutating"] = ScalarObjective(exponential!, exponential_g!, exponential_fg!, exponential_fgh!, exponential_h!, exponential_hv!, nothing, nothing)

exponential_prb = Dict()
exponential_prb["twicediffed!"] = TwiceDiffed(exponential!)
exponential_prb["oncediffed!"] = OnceDiffed(exponential!)
exponential_prb["initial_x"] = [0.0, 0.0]
exponential_prb["minimizer"] = [2.0, 3.0]
exponential_prb["minimum"] = exponential!(nothing, nothing, [2.0, 3.0])
problems["unconstrained"]["exponential"] = exponential_prb


using Test
using Bigsimr
using LinearAlgebra
using Distributions

@testset "Correlation Functions" begin

    r_negdef = [
        1.00 0.82 0.56 0.44
        0.82 1.00 0.28 0.85
        0.56 0.28 1.00 0.22
        0.44 0.85 0.22 1.00
    ]

    @testset "Nearest positive definite correlation matrix" begin
        r = cor_nearPD(r_negdef)
        @test Bigsimr.iscorrelation(r)

        # Must respect input eltype
        for T in [Float64, Float32, Float64]
            @test eltype(cor_nearPD(Matrix{T}(r_negdef))) === T
        end
    end

    @testset "Nearest positive semi-definite correlation matrix" begin
        r = cor_nearPD(r_negdef, 0.0)
        λ = eigvals(r)
        @test issymmetric(r)
        @test all(λ .≥ 0)
        @test all(diag(r) .== 1.0)
        @test all(-1.0 .≤ r .≤ 1.0)
    end

    @testset "Fast near positive definite correlation matrix" begin
        r = cor_fastPD(r_negdef)
        @test Bigsimr.iscorrelation(r)

        # Must respect input eltype
        test_types = [Float64, Float32, Float16]
        for T in test_types
            @test eltype(cor_fastPD(Matrix{T}(r_negdef))) === T
        end
    end
end

@testset "Random Correlation Generation" begin

    @testset "Random postive definite correlation matrix" begin
        r = cor_randPD(100)
        @test Bigsimr.iscorrelation(r)

        # The element type must be respected
        test_types = [Float64, Float32, Float16]
        for T in test_types
            @test eltype(cor_randPD(T, 4)) === T
        end

        # Must work for numbers with integer representations
        test_types = [Float64, Float32, Float16, Rational, Int64, Int32, Int16]
        for T in test_types
            @test_nowarn cor_randPD(T(4))
            @test_nowarn cor_randPD(Float64, T(4))
            @test_nowarn cor_randPD(Float32, T(4))
            @test_nowarn cor_randPD(Float16, T(4))

            for S in test_types
                @test_nowarn cor_randPD(T(4), S(3))
                @test_nowarn cor_randPD(Float64, T(4), S(3))
                @test_nowarn cor_randPD(Float32, T(4), S(3))
                @test_nowarn cor_randPD(Float16, T(4), S(3))
            end
        end

        # `d` must not be less than 1
        @test_throws AssertionError cor_randPD(-1)
        @test_throws AssertionError cor_randPD(Float64, -1)
        @test_throws AssertionError cor_randPD(Float32, -1)
        @test_throws AssertionError cor_randPD(Float16, -1)

        # `k` must not be larger than `d`
        @test_throws AssertionError cor_randPD(4, 5)
        @test_throws AssertionError cor_randPD(Float64, 4, 5)
        @test_throws AssertionError cor_randPD(Float32, 4, 5)
        @test_throws AssertionError cor_randPD(Float16, 4, 5)

        # `k` must not be less than 1
        @test_throws AssertionError cor_randPD(4, 0)
        @test_throws AssertionError cor_randPD(Float64, 4, 0)
        @test_throws AssertionError cor_randPD(Float32, 4, 0)
        @test_throws AssertionError cor_randPD(Float16, 4, 0)
    end

    @testset "Random positive semi-definite correlation matrix" begin
        r = cor_randPSD(100)
        λ = eigvals(r)
        @test issymmetric(r)
        @test all(λ .≥ 0)
        @test all(diag(r) .== 1.0)
        @test all(-1.0 .≤ r .≤ 1.0)

        # The element type must be respected
        test_types = [Float64, Float32, Float16]
        for T in test_types
            @test eltype(cor_randPSD(T, 4)) === T
        end

        # Must work for numbers with integer representations
        test_types = [Float64, Float32, Float16, Rational, Int64, Int32, Int16]
        for T in test_types
            @test_nowarn cor_randPSD(T(4))
            @test_nowarn cor_randPSD(Float64, T(4))
            @test_nowarn cor_randPSD(Float32, T(4))
            @test_nowarn cor_randPSD(Float16, T(4))

            for S in test_types
                @test_nowarn cor_randPSD(T(4), S(3))
                @test_nowarn cor_randPSD(Float64, T(4), S(3))
                @test_nowarn cor_randPSD(Float32, T(4), S(3))
                @test_nowarn cor_randPSD(Float16, T(4), S(3))
            end
        end

        # `d` must not be less than 1
        @test_throws AssertionError cor_randPSD(-1)
        @test_throws AssertionError cor_randPSD(Float64, -1)
        @test_throws AssertionError cor_randPSD(Float32, -1)
        @test_throws AssertionError cor_randPSD(Float16, -1)

        # `k` must not be larger than `d`
        @test_throws AssertionError cor_randPSD(4, 5)
        @test_throws AssertionError cor_randPSD(Float64, 4, 5)
        @test_throws AssertionError cor_randPSD(Float32, 4, 5)
        @test_throws AssertionError cor_randPSD(Float16, 4, 5)

        # `k` must not be less than 1
        @test_throws AssertionError cor_randPSD(4, 0)
        @test_throws AssertionError cor_randPSD(Float64, 4, 0)
        @test_throws AssertionError cor_randPSD(Float32, 4, 0)
        @test_throws AssertionError cor_randPSD(Float16, 4, 0)
    end

end

@testset "Correlation Utilities" begin

    cor_types = (Pearson, Spearman, Kendall)

    @testset "Correlation calculation" begin
        # Must work for any floating point matrix or vector pair
        test_types = [Float64, Float32, Float16]
        for T in test_types
            A = rand(T, 200, 4)
            x, y = rand(T, 100), rand(T, 100)

            for C in cor_types
                @test_nowarn cor(A, C)
                @test_nowarn cor(x, y, C)

                @test_nowarn cor_fast(A, C)
            end
        end

    end

    @testset "Correlation to correlation conversion" begin

        # Converting type A -> A must result in the same matrix
        for C in cor_types
            r = cor_randPD(4)
            @test r == cor_convert(r, C, C)
        end

        # Must map (-1, 0, 1) onto itself within numerical error
        for C1 in cor_types
            for C2 in cor_types
                @test cor_convert( 0.0, C1, C2)  == 0.0
                @test cor_convert( 1.0, C1, C2)   ≤ 1.0
                @test cor_convert( 1.0, C1, C2)   ≈ 1.0
                @test cor_convert(-1.0, C1, C2)  ≥ -1.0
                @test cor_convert(-1.0, C1, C2)  ≈ -1.0
            end
        end

        # Must work for each type
        test_types = [Float64, Float32, Float16, Rational]
        for T in test_types
            for C1 in cor_types
                for C2 in cor_types
                    @test_nowarn cor_convert(T(0.5), C1,  C2)
                end
            end
        end

        # Must respect these input eltypes
        test_types = [Float64, Float32, Float16]
        for T in test_types
            for C1 in cor_types
                for C2 in cor_types
                    @test eltype(cor_convert(T(0.5), C1,  C2))  === T
                end
            end
        end
    end

    @testset "Constrain to Correlation" begin
        test_types = [Float64, Float32, Float16]
        for T in test_types
            A = rand(T, 10, 10)
            @test_nowarn cor_constrain(A)
            @test eltype(cor_constrain(A)) === T
        end
    end

    @testset "Covariance to Correlation" begin
        test_types = [Float64, Float32, Float16]
        for T in test_types
            R = cor_randPSD(T, 4, 3)
            @test_nowarn cov2cor(R)
            @test eltype(cov2cor(R)) === T
        end
    end

    @testset "Correlation Bounds" begin
        A, B = NegativeBinomial(20, 0.2), LogNormal(3, 1)
        dist_types = (A, B)

        # Must work for any univariate distribution and correlation type
        for D1 in dist_types
            for D2 in dist_types
                for C in cor_types
                    @test_nowarn cor_bounds(D1, D2, C)
                end
            end
        end

        # Must work for any number with an integer representation
        test_types = [Float64, Float32, Float16, Int64, Int32, Int16]
        for T in test_types
            @test_nowarn cor_bounds(A, B, n_samples=T(10_000))
        end

        test_types = [Float64, Float32]
        for T in test_types
            @test_throws InexactError cor_bounds(A, B, n_samples=T(10_000.5))
        end
    end

end


module entropy
export velocity_histogram, kB, entropy_2D_MB, plot_velocity_distr

using SmoothedParticles
using Plots: plot, savefig
using Ipopt
using JuMP

"""
	Histogram(xs::Vector{Float64}, ys::Vector{Float64},N,dx)

Histogram structure storing ``N`` x values ``xs`` with uniform bin width ``dx`` and ``N`` y values ``ys``.
"""
struct Histogram
	xs::Vector{Float64}
	ys::Vector{Float64}
	N::Int64
	dx::Float64
end

"""
	velocity_histogram(sys::ParticleSystem; v_max = 0, N = 10)

Building the histrogram of 2D velocities (norms) with ``v_max`` the maximum velocity in the histogram and ``N`` bins.
"""
function velocity_histogram(sys::ParticleSystem; v_max = 0.0, N = 100)::Histogram
	if v_max == 0.0 # if v_max = 0, find the maximum velocity of the particles
		for k in 1:length(sys.particles)
			v = norm(sys.particles[k].v)
			if v > v_max
				v_max = v
			end
		end
	end

	# Find the heights of the histogram bins
	dv = v_max/N # velocity increment between the bins
	vs = 0.:dv:v_max
	ns = zeros(length(vs))
	for k in 1:length(sys.particles)
		v = norm(sys.particles[k].v)
		n = Int64(round(v/dv))
		if 1 <= n <= length(ns)
			ns[n] += 1.0/(dv*length(sys.particles))
		end
	end

	return Histogram(vs,ns,100,dv)
end

"""
	kB

Boltzmann constant (in the SI units)
"""
const kB = 1.380649e-23

"""
	entropy(fMB::Histogram)::Float64

Calculate Boltzmann entropy of a 2D Maxwell-Boltzmann distribution approximated by an ``fMB`` histogram.
"""
function entropy_2D_MB(fMB::Histogram)::Float64
	@assert(fMB.xs[1] == 0) # Assuming that the histogram starts at zero velocity

	S = 0.0

	# Approximating the reduced entropy near v=0, where a numerical singularity could appear
	fMBder = (fMB.ys[2]-fMB.ys[1])/fMB.dx
	if fMBder > 0
		S = - fMB.ys[1] * (log(fMBder)*fMB.dx - fMBder*(fMB.dx^3)/6)
	end

	# Approximating the rest of entropy
	for k in 2:length(fMB.xs)
		if fMB.xs[k] != 0
			if fMB.ys[k] > 0
				S += -fMB.ys[k] * log(fMB.ys[k]/fMB.xs[k]) * fMB.dx
			end
		end
	end

	return S
end

"""
	plot_velocity_distr(sys::ParticleSystem, name::String)

Plots the distribution of velocity magnitudes among the particles of ``sys`` and saves the resulting pdf,
together with a fit of the Maxwell-Boltzmann distribution, to file ``name``. ``m`` is the mass of each particle.
Returns the temperature.
"""
function plot_velocity_distr(sys::ParticleSystem, m::Float64, name::String; v_max = 0.0)::Float64
	distr = velocity_histogram(sys, v_max=v_max, N = 100)

	# fitting the histogram to a 2D Maxwell-Boltzmann distribution
	model = Model(Ipopt.Optimizer)
	@variable(model, beta)
	@NLobjective(
			model,
			Min,
			sum((distr.ys[i] - m*beta*distr.xs[i]*exp(-0.5*m*beta*distr.xs[i]^2))^2 for i in 1:length(distr.xs)),
		)
	optimize!(model)
	beta = value(beta)

	# Plotting both the actual histogram and the fitted Maxwell-Boltzmann distribution
	ns_boltz = zeros(length(distr.xs))
	for i in 1:length(ns_boltz)
		ns_boltz[i] = m*beta*distr.xs[i]*exp(-0.5*m*beta*distr.xs[i]^2)
	end
	@show(beta)
	T = 1/(beta*kB)
	@show(T)

	p = plot(distr.xs, [distr.ys ns_boltz], label = ["data" "Maxwell-Boltzmann, T="*string(T)])
	savefig(p, name)
	return T
end

#"""
#	plot_velocity_distr(path::String, name:: String)
#
#Plot velocity distribution of particles loaded from a VTK file.
#"""
#function plot_velocity_distr(path::String, name::String)
#	@error "Not working, sorry"
#	domain = Rectangle(-box_width, -box_width, 2*box_width, 3*box_height)
#	sys = ParticleSystem(Particle, domain, h)
#	read_vtk!(sys, path, x -> Particle(x = x, type = 0.0))
#	plot_velocity_distr(sys, name)  #not yet working
#end

end


# Util functions.
module Util
using Random

export weighted_median,
       err_must_be_overriden,
       holdout

# Weighted median.
#
# @param weights Weights of values.
# @param values Values.
# @return Weighted median.
function weighted_median( weights::AbstractVector{U}, values::AbstractVector{V}) where {U,V<:Real}

  k = 1
  sorted_ind = sortperm(values)
  weight_sum = sum(weights)

  remaining_sum = weight_sum - weights[sorted_ind[k]]
  while remaining_sum > weight_sum / 2.0
    k += 1
    remaining_sum -= weights[sorted_ind[k]]
  end

  return values[sorted_ind[k]]
end

function err_must_be_overriden(s = "")
  error("Function must be overriden.")
end

# Holdout method that partitions a collection
# into two partitions.
#
# @param n Size of collection to partition.
# @param right_prop Percentage of collection placed in right partition.
# @return Two partitions of indices, left and right.
function holdout(n, right_prop)
  shuffled_indices = randperm(n)
  partition_pivot = round(Int, right_prop * n)
  right = shuffled_indices[1:partition_pivot]
  left = shuffled_indices[partition_pivot+1:end]
  return (left, right)
end

end # module


# two main type systems - static and dynamic

#Integers
a = 1;
println(typeof(a));

println(Sys.WORD_SIZE);
println(Int);
println(UInt);

# Types of Integers

println(typemin(Int8));
println(typemax(Int8));
println(typemin(UInt8));
println(typemax(UInt8));

println(typemin(Int16));
println(typemax(Int16));
println(typemin(UInt16));
println(typemax(UInt16));

println(typemin(Int32));
println(typemax(Int32));
println(typemin(UInt32));
println(typemax(UInt32));

println(typemin(Int64));
println(typemax(Int64));
println(typemin(UInt64));
println(typemax(UInt64));

println(typemin(Int128));
println(typemax(Int128));
println(typemin(UInt128));
println(typemax(UInt128));

println(typemin(Bool));
println(false == 0);

println(typemax(Bool));

# Number Bases

b = 4; #decimal
b = 0b0100; # binary
b = 0o010; # octal
b = 0x1f; # hexadecimal

# overflow handled by wrap around

a = typemax(Int64);
a = a + 1;
println(a);

# BigInt -> larger than Int128

b = BigInt(typemax(Int128) + 1);
println(b);

# floating point numbers (IEEE 754)
c = 1.2
println(typeof(c));

println(typemin(Float16));
println(typemax(Float16));

Float16; # half precision 16-bit float
Float32; # single precision 32-bit float
Float64; # double precision 64-bit float

# Simple floats

f = 1.0;
f = 1.;
f = 0.1;
f = .1;

f = 22.5f-3; # single precision
println(typeof(f));

f = 22.5e-3; # double precision
println(typeof(f));

f = 0x1fp-1; # p - base2 exponent for 64-bit hex floats

# Special floats
Inf;
-Inf;
NaN;

println(Inf / Inf);

# BigFloat for big computations

a = parse(BigFloat, "0.00000000000000000000000000000000000000001");
println(a);

# static type a variable using ::
local a::Int8;

# This throws TypeError
# local a::Int8 = 128;

# isa -> check if an object is a particular type
a = 2;
println(isa(a, Int));

a = 2.0;
println(isa(a, Int));

# supertype -> see type's supertype or parent type

println(supertype(Float64));
println(supertype(AbstractFloat));
println(supertype(Real));
println(supertype(Number));
println(supertype(Any));

# Type Conversion

x = 12;
println(typeof(x));
println(typeof(convert(UInt8, x)));

# Type Promotion

println(promote(x, 2.3));


# ---
# title: 233. Number of Digit One
# id: problem233
# author: zhwang
# date: 2022-03-12
# difficulty: Hard
# categories: Math
# link: <https://leetcode.com/problems/number-of-digit-one/description/>
# hidden: true
# ---
#
# Given an integer n, count the total number of digit 1 appearing in all non-
# negative integers less than or equal to n.
#
# **Example:**
#
#
#
#     Input: 13
#     Output: 6
#     Explanation: Digit 1 occurred in the following numbers: 1, 10, 11, 12, 13.
#
#
#
## @lc code=start
using LeetCode

function count_digit_one(n::Int)::Int
    nums = reverse!([parse(Int, i) for i in string(n + 1)])
    res, ones = 0, count(==(1), nums)
    for (i, num) in enumerate(nums)
        num == 1 && (ones -= 1)
        res += 10^(i - 1) * (num * ones + (num > 1))
        res += i > 1 && (10^(i - 2) * (i - 1) * num)
    end
    return res
end

## example: reduce 21413 to 0
## num: step => reduced num
## 2141|4|-1: 4*2*1 + 4*0 + (4>1)*10 => 21410-1
## 214|1|0-1: 1*1*10 + 1*1*1 + (1>1)*10 => 21400-1
## 21|4|00-1: 4*1*100 + 4*2*10 + (4>1)*100 => 21000-1
## 2|1|000-1: 1*0*1000 + 1*3*100 + (1>1)*1000 => 20000-1
## |2|0000-1: 2*0*10000 + 2*4*10000 + (2>1)*10000 => 0

## @lc code=end


using Random
using LinearAlgebra

using JuMP
import GLPK
using Plots

Random.seed!(314)

m = 12  # number of clients
n = 5  # number of facility locations

# Clients' locations
Xc = rand(m)
Yc = rand(m)

# Facilities' potential locations
Xf = rand(n)
Yf = rand(n)

# Fixed costs
f = ones(n);

# Distance
c = zeros(m, n)
for i in 1:m
    for j in 1:n
        c[i, j] = norm([Xc[i] - Xf[j], Yc[i] - Yf[j]], 2)
    end
end

# Display the data
scatter(Xc, Yc, label = "Clients", markershape=:circle, markercolor=:blue)
scatter!(Xf, Yf, label="Facility",
    markershape=:square, markercolor=:white, markersize=6,
    markerstrokecolor=:red, markerstrokewidth=2
)

# Create a JuMP model
ufl = Model(GLPK.Optimizer)

# Variables
@variable(ufl, y[1:n], Bin);
@variable(ufl, x[1:m, 1:n], Bin);

# Each client is served exactly once
@constraint(ufl, client_service[i in 1:m],
    sum(x[i, j] for j in 1:n) == 1
);

# A facility must be open to serve a client
@constraint(ufl, open_facility[i in 1:m, j in 1:n],
    x[i, j] <= y[j]
)

# Objective
@objective(ufl, Min, f'y + sum(c .* x));

# Solve the uncapacitated facility location problem with GLPK
optimize!(ufl)
println("Optimal value: ", objective_value(ufl))

x_ = value.(x) .> 1 - 1e-5
y_ = value.(y) .> 1 - 1e-5

# Display clients
p = scatter(Xc, Yc, markershape=:circle, markercolor=:blue, label=nothing)

# Show open facility
mc = [(y_[j] ? :red : :white) for j in 1:n]
scatter!(Xf, Yf,
    markershape=:square, markercolor=mc, markersize=6,
    markerstrokecolor=:red, markerstrokewidth=2,
    label=nothing
)

# Show client-facility assignment
for i in 1:m
    for j in 1:n
        if x_[i, j] == 1
           plot!([Xc[i], Xf[j]], [Yc[i], Yf[j]], color=:black, label=nothing)
        end
    end
end

display(p)

# Demands
a = rand(1:3, m);

# Capacities
q = rand(5:10, n);

# Display the data
scatter(Xc, Yc, label=nothing,
    markershape=:circle, markercolor=:blue, markersize= 2 .*(2 .+ a)
)

scatter!(Xf, Yf, label=nothing,
    markershape=:rect, markercolor=:white, markersize= q,
    markerstrokecolor=:red, markerstrokewidth=2
)

# Create a JuMP model
cfl = Model(GLPK.Optimizer)

# Variables
@variable(cfl, y[1:n], Bin);
@variable(cfl, x[1:m, 1:n], Bin);

# Each client is served exactly once
@constraint(cfl, client_service[i in 1:m], sum(x[i, :]) == 1)

# Capacity constraint
@constraint(cfl, capacity, x'a .<= (q .* y))

# Objective
@objective(cfl, Min, f'y + sum(c .* x));

# Solve the problem
optimize!(cfl)
println("Optimal value: ", objective_value(cfl))
@test termination_status(cfl) == MOI.OPTIMAL

x_ = value.(x) .> 1 - 1e-5
y_ = value.(y) .> 1 - 1e-5

# Display the solution
p = scatter(Xc, Yc, label=nothing,
    markershape=:circle, markercolor=:blue, markersize= 2 .*(2 .+ a)
)

mc = [(y_[j] ? :red : :white) for j in 1:n]
scatter!(Xf, Yf, label=nothing,
    markershape=:rect, markercolor=mc, markersize=q,
    markerstrokecolor=:red, markerstrokewidth=2
)

# Show client-facility assignment
for i in 1:m
    for j in 1:n
        if x_[i, j] == 1
            plot!([Xc[i], Xf[j]], [Yc[i], Yf[j]], color=:black, label=nothing)
            break
        end
    end
end

display(p)


module KrylovMethods

	import Base.BLAS

	include("cg.jl")
	include("blockCG.jl")
	include("cgls.jl")
	include("bicgstb.jl")
	include("blockBiCGSTB.jl")
	include("gmres.jl")
	include("lanczosBidiag.jl")
	include("ssor.jl")
	include("lsqr.jl")
	include("lanczosTridiag.jl")
	include("lanczos.jl")
	include("minres.jl")
	include("gs.jl")
end


# PS8 Dayuan Xie
# using Pkg
# Pkg.add("MultivariateStats")

using Optim
using HTTP
using GLM
using LinearAlgebra
using Random
using Statistics
using DataFrames
using DataFramesMeta
using Distributions
using CSV
using MultivariateStats

function allwrap()
# Question 1
url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS8-factor/nlsy.csv"
df = CSV.read(HTTP.get(url).body)

ols_lm = lm(@formula(logwage ~ black + hispanic + female + schoolt + gradHS + grad4yr), df)
println(ols_lm)

# Question 2
matrix_asvab = hcat(df.asvabAR, df.asvabCS,df.asvabMK,df.asvabNO,df.asvabPC,df.asvabWK)
cor_asvab = cor(matrix_asvab)
println(cor_asvab)

# Question 3
ols_asvab_lm = lm(@formula(logwage ~ black + hispanic + female + schoolt + gradHS + grad4yr +
                               asvabAR + asvabCS + asvabMK + asvabNO + asvabPC + asvabWK), df)
println(ols_asvab_lm)

#  It will be problematic to directly include  the six asvab variables in the regression
#  since these variables are highly correlated with each other.

# Question 4
asvabMat = (hcat(df.asvabAR, df.asvabCS,df.asvabMK,df.asvabNO,df.asvabPC,df.asvabWK))'
M = fit(PCA, asvabMat; maxoutdim=1)
asvabPCA = MultivariateStats.transform(M, asvabMat)
asvabPCA= vec(asvabPCA')
insert!(df,7, asvabPCA, :asvabPCA)
ols_PCA = lm(@formula(logwage ~ black + hispanic + female + schoolt + gradHS + grad4yr + asvabPCA), df)
println(ols_PCA)

# Question 5
asvabMat = (hcat(df.asvabAR, df.asvabCS,df.asvabMK,df.asvabNO,df.asvabPC,df.asvabWK))'
M = fit(FactorAnalysis, asvabMat; maxoutdim=1)
asvabFactorAnalysis = MultivariateStats.transform(M, asvabMat)
asvabFactorAnalysis= vec(asvabFactorAnalysis')
insert!(df,8, asvabFactorAnalysis, :asvabFactorAnalysis)
ols_FactorAnalysis = lm(@formula(logwage ~ black + hispanic + female + schoolt + gradHS + grad4yr + asvabFactorAnalysis), df)
println(ols_FactorAnalysis)

# Question 6
function lgwt(N0::Integer,a::Real=-1,b::Real=1)
    N  = N0-1
    N1 = N+1
    N2 = N+2
    xu = range(-1,stop=1,length=N1)
    y = cos.((2*(0:N) .+ 1)*pi/(2*N .+ 2))  .+  ( 0.27/N1 ) .* sin.( pi .* xu .* N/N2 )
    L  = zeros(N1,N2)
    Lp = zeros(N1,N2)
    y0 = 2
    vareps = 2e-52
    i = 0
    tracker=0
    it_max=10

    while (norm(y.-y0,Inf)>vareps && tracker<=it_max)
        d=norm(y.-y0,Inf)

        L[:,1]  .= 1
        Lp[:,1] .= 0

        L[:,2] .= y

        for k=2:N1
            L[:,k+1] = ( (2*k-1)*y .* L[:,k] .- (k-1)*L[:,k-1] )/k
        end

        Lp = (N2)*( L[:,N1] .- y .* L[:,N2] )./(1 .- y.^2)
        y0 = y
        y  = y0 - L[:,N2]./Lp
        if norm(y.-y0,Inf)==d
            tracker+=1
        end
        i+=1

    end

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

    return x,w
end

D= Normal(0,1)
ξ= rand(Normal(0,1),size(df,1),1)
M= hcat(df.asvabAR, df.asvabCS,df.asvabMK,df.asvabNO,df.asvabPC,df.asvabWK)
Xm= hcat(ones(size(df,1)), df.black, df.hispanic, df.female)
X= hcat(ones(size(df,1)), df.black, df.hispanic, df.female, df.schoolt, df.gradHS, df.grad4yr)
Y= df.logwage
L= zeros(size(df,1),1)
function likelihoodfunction(sigma,alpha,beta,gamma,delta)
    sigmaj=sigma[1:end-1]
    sigmaw=sigma[end]
    for i= 1: size(df,1)
        for j= 1:6
            part1vec=zeros(6,1)
            part1vec[j,1]=(1/sigma[1,j])*pdf.(D, (M[i,j]-Xm[i,:]*alpha-gamma*ξ[i,1])/1/sigma[1,j]))
            part1=part1vec[1,1]*part1vec[2,1]*part1vec[3,1]*part1vec[4,1]*part1vec[5,1]*part1vec[6,1]
        end
        part2=(1/sigmaw)*pdf.(D,(Y[i,1]-X[i,:]*beta-delta*ξ[i,1])/sigmaw)
        L[i,1]= part1*part2
    end
    return L
end

nodes, weights = lgwt(7,-4,4)
likelihood=zeros(size(df,1),1)
for i= 1: size(df,1)
likelihood[i,1]=log(sum(weights.* L[i,1].*pdf.(D,nodes)))
end

sigma_hat = optimize(sigma -> -sum(likelihood), rand(1,7), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000, show_trace=true))
println(sigma_hat.minimizer)

end

allwrap()


"""
    ∈(M::AbstractMatrix, A::AbstractIntervalMatrix)

Check whether a concrete matrix is an instance of an interval matrix.

### Input

- `M` -- concrete matrix
- `A` -- interval matrix

### Output

`true` iff `M` is an instance of `A`

### Algorithm

We check for each entry in `M` whether it belongs to the corresponding interval
in `A`.
"""
function ∈(M::AbstractMatrix, A::AbstractIntervalMatrix)
    @assert size(M) == size(A) "incompatible matrix sizes (M: $(size(M)), A: " *
                               "$(size(A)))"

    m, n = size(A)
    @inbounds for j in 1:n
        for i in 1:m
            if M[i, j] ∉ A[i, j]
                return false
            end
        end
    end
    return true
end

"""
    ⊆(A::AbstractIntervalMatrix, B::AbstractIntervalMatrix)

Check whether an interval matrix is contained in another interval matrix.

### Input

- `A` -- interval matrix
- `B` -- interval matrix

### Output

`true` iff `A[i, j] ⊆ B[i, j]` for all `i, j`.
"""
function ⊆(A::AbstractIntervalMatrix, B::AbstractIntervalMatrix)
    @assert size(A) == size(B) "incompatible matrix sizes $(size(A)) and " *
                               "$(size(B))"

    m, n = size(A)
    @inbounds for j in 1:n, i in 1:m
        if !(A[i, j] ⊆ B[i, j])
            return false
        end
    end
    return true
end

"""
    ∩(A::IntervalMatrix, B::IntervalMatrix)

Intersect two interval matrices.

### Input

- `A` -- interval matrix
- `B` -- interval matrix (of the same shape as `A`)

### Output

A new matrix `C` of the same shape as `A` such that
`C[i, j] = A[i, j] ∩ B[i, j]` for each `i` and `j`.
"""
function ∩(A::IntervalMatrix, B::IntervalMatrix)
    m, n = size(A)
    @assert size(A) == size(B) "incompatible matrix sizes (A: $(size(A)), B: " *
                               "$(size(B)))"

    return IntervalMatrix(map((x, y) -> x ∩ y, A, B))
end

"""
    hull(A::IntervalMatrix, B::IntervalMatrix)

Finds the interval hull of two interval matrices. This is equivalent to [`∪`](@ref).

### Input

- `A` -- interval matrix
- `B` -- interval matrix (of the same shape as `A`)

### Output

A new matrix `C` of the same shape as `A` such that
`C[i, j] = hull(A[i, j], B[i, j])` for each `i` and `j`.
"""
function hull(A::IntervalMatrix, B::IntervalMatrix)
    @assert size(A) == size(B) "incompatible matrix sizes (A: $(size(A)), B: " *
                               "$(size(B)))"

    return IntervalMatrix(map((x, y) -> hull(x, y), A, B))
end

"""
    ∪(A::IntervalMatrix, B::IntervalMatrix)

Finds the interval union (hull) of two interval matrices.
This is equivalent to [`hull`](@ref).

### Input

- `A` -- interval matrix
- `B` -- interval matrix (of the same shape as `A`)

### Output

A new matrix `C` of the same shape as `A` such that
`C[i, j] = A[i, j] ∪ B[i, j]` for each `i` and `j`.
"""
∪(A::IntervalMatrix, B::IntervalMatrix) = hull(A, B)


function ncsize!(S, present, g, i, j)
	!present[i] && return 0
	S[i] != 0 && error("the graph is not acyclic")
	S[i] = 1
	for k in neighbors(g,i)
		if k != j && present[k]
			S[i] += ncsize!(S, present, g, k, i)
		end
	end
	return S[i]
end

"""
Optimally break tree until the giant component size is less than threshold.
"""
function tree_break!(present, g, threshold)
    S = zeros(Int, nv(g))
    H = MutableBinaryMaxHeap{Tuple{Int64,Int64}}()
    for i in vertices(g)
        if present[i] && S[i] == 0
            push!(H, (ncsize!(S, present, g, i, nothing), i))
        end
    end

    attack_nodes = Int[]
    scomps = Int[]
    while !isempty(H)
        scomp, i = pop!(H)
        sender = nothing
        while true
            sizes = [(S[k],k) for k in neighbors(g,i) if k != sender && present[k]]
            isempty(sizes) && break
            M, largest = maximum(sizes)
            if M <= scomp/2
                for k in neighbors(g,i)
                    if S[k] > 1 && present[k]
                        push!(H, (S[k],k))
                    end
                end
                present[i] = false
                push!(attack_nodes, i)
                push!(scomps, scomp)
                break
            end
            S[i] = 1
            for k in neighbors(g,i)
                if k != largest && present[k]
                    S[i] += S[k]
                end
            end
            sender, i = i, largest
        end
		scomp <= threshold && break
    end
    return attack_nodes, scomps
end


using Test, PolyChaos

@test nw(EmptyQuad()) |> typeof == Array{Float64,2}
@test nw(EmptyQuad()) |> length == 0

nodes, weights = gauss(rm_legendre01(5)...)
myQuad = Quad("myQuad", length(nodes), nodes, weights)
foo(x) = 6*x^5
res = 1
@test isapprox(integrate(foo, nodes, weights), res)
@test isapprox(integrate(foo, myQuad), res)
@test_throws DomainError integrate(foo, EmptyQuad())
@test isapprox(integrate(foo, Uniform01OrthoPoly(4)), res)

foo(x) = 5*x^4
res = 1
@test isapprox(integrate(foo, Uniform_11OrthoPoly(4)), res)

d, nunc = 5, 10
op = genLaguerreOrthoPoly(d,1.23,addQuadrature = false)
opq = genLaguerreOrthoPoly(d,1.23)
mop = MultiOrthoPoly([op for i in 1:nunc],d)
mopq = MultiOrthoPoly([opq for i in 1:nunc],d)
@testset "dimenions" begin
    @test isequal(PolyChaos.dim(op), d + 1)
    @test isequal(PolyChaos.dim(opq), d + 1)
    @test isequal(PolyChaos.dim(mop), factorial(d + nunc) / (factorial(d) * factorial(nunc)))
end

@testset "degrees" begin
    @test isequal(PolyChaos.deg(op),d)
    @test isequal(PolyChaos.deg(opq),d)
    @test isequal(PolyChaos.deg(mop),d)
end

NW = [ opq.quad.nodes opq.quad.weights ]
@testset "nodes and weights" begin
    @test PolyChaos.nw(opq) == NW
    @test PolyChaos.nw([opq for i in 1:nunc])[1] == [NW[:,1] for i in 1:nunc]
    @test PolyChaos.nw([opq for i in 1:nunc])[2] == [NW[:,2] for i in 1:nunc]
    @test PolyChaos.nw(mopq)[1] == [NW[:,1] for i in 1:nunc]
    @test PolyChaos.nw(mopq)[2] == [NW[:,2] for i in 1:nunc]
end

c = [ op.α op.β ]
@testset "coefficients" begin
    @test PolyChaos.coeffs(op) == c
    @test PolyChaos.coeffs(opq) == c
    @test PolyChaos.coeffs([op for i in 1:nunc])[1] == [ c[:,1] for i in 1:nunc ]
    @test PolyChaos.coeffs([op for i in 1:nunc])[2] == [ c[:,2] for i in 1:nunc ]
    @test PolyChaos.coeffs([opq for i in 1:nunc])[1] == [ c[:,1] for i in 1:nunc ]
    @test PolyChaos.coeffs([opq for i in 1:nunc])[2] == [ c[:,2] for i in 1:nunc ]
    @test PolyChaos.coeffs(mop)[1] == [ c[:,1] for i in 1:nunc ]
    @test PolyChaos.coeffs(mop)[2] == [ c[:,2] for i in 1:nunc ]
end

@test issymmetric(op) == issymmetric(opq)

## Accounting for StaticArrays
using StaticArrays

max_degree = 4

op = Uniform01OrthoPoly(max_degree; Nrec = max_degree + 3)
nodes, weights = SVector(op.quad.nodes...), SVector(op.quad.weights...)

myfun(t) = t^2

@test integrate(myfun, op) == integrate(myfun, nodes, weights)

n_multi = 4
mop = MultiOrthoPoly([op for i in 1:n_multi], max_degree)
L = size(mop.ind)[1]

test_vector = [ rand(0:L-1) for i in 0:L ]
reference = multi2uni(test_vector, mop.ind)

@testset "multi2uni for StaticArrays" begin
    @test multi2uni(SVector(test_vector...), mop.ind) == reference
    @test multi2uni(test_vector, SMatrix{size(mop.ind)...}(mop.ind)) == reference
    @test multi2uni(SVector(test_vector...), SMatrix{size(mop.ind)...}(mop.ind)) == reference
end

tensor_dim = 2
tensor_entries = computeTensorizedSP(tensor_dim, mop)
my_getentry(row, entries) = getentry(row, tensor_entries, entries, tensor_dim)

if VERSION >= VersionNumber("1.1.1")
    @testset "getentry for StaticArrays" begin
        for row in eachrow(rand(0:L-1, L, tensor_dim))
            reference = my_getentry(Vector(row), mop.ind)
            @test my_getentry(row, mop.ind) == reference
            @test my_getentry(row, SMatrix{size(mop.ind)...}(mop.ind)) == reference
            @test my_getentry(Vector(row), SMatrix{size(mop.ind)...}(mop.ind)) == reference
        end
    end
end


using DataFrames

function lagrangianpolynomial(np::Int, x::Vector{Float64}, y::Vector{Float64}, xi::Vector{Float64})
  local yi = Vector{Float64}(undef, length(xi))
  m = sort(hcat(x, y), dims=1)
  for l in 1:length(xi)
    for i in 1:np
      term = 1.0
      for j in 1:np
        j !== i && (term *= (xi[l] - m[j, 1])/(m[i, 1] - m[j, 1]))
      end
      yi[l] += term * m[i, 2]
    end
  end
  columns=[]
  push!(columns, m[:,1])
  push!(columns, m[:,2])
  df = DataFrame(columns, [:x, :y])

  columns=[]
  push!(columns, xi)
  push!(columns, yi)
  dfxi = DataFrame(columns, [:xi, :yi])

  (df, dfxi)
end


function modalparam(A,C,dt)
	#
	# Compute modal parameters from A, C and dt
	#
	# javier.cara@upm.es, 2016-02
	#

	no,ns = size(C)

	# allocating
	wm1 = zeros(1,ns)
	zm1 = zeros(1,ns)
	vm1 = zeros(no,ns)

	# eigenvalues and eigenvectors
	D,V = eig(A)

	# C*V
	vm0 = C*V

	i = 1
	ii = 1
	fm1 = zeros(Number,no,ns) # can be real or complex numbers
	while i < ns
		d = D[i]
		if abs(imag(d)) > 1e-10 # complex number
			wm1[ii] = abs(log(d))/dt
			zm1[ii] = -real(log(d))/(wm1[ii]*dt)

			# modal shapes normalized with max component = 1
			maxval = vm0[1,i]
			for j in 2:no
				if abs(vm0[j,i]) > abs(maxval)
					maxval = vm0[j,i]
				end
			end
			vm1[:,ii] = real( vm0[:,i]/maxval )

			i = i+2
			ii = ii+1

		else
			# real eigenvalue
			i = i+1
		end
	end

	# deleting zero values
	wm = wm1[1:ii-1]
	zm = zm1[1:ii-1]
	vm = vm1[:,1:ii-1]

	# sorting frequencies
	pos=sortperm(wm)
	wm = wm[pos]
	zm = zm[pos]
	vm = vm[:,pos]

	return wm,zm,vm

end


import Base.*

export S¹
export U1

"""
    Represents a point in a one-sphere.
"""
abstract type S¹ end

"""
    Represents a point in U(1).

field: α.
"""
struct U1 <: S¹
    r::Float64
    U1(z::Complex) = begin
        @assert(isapprox(abs(z), 1), "The magnitude must be equal to 1, but it's $(abs(z)).")
        θ = angle(z) ≥ 0 ? angle(z) : π + angle(z)
        new(θ)
    end
    U1(α::Real) = begin
        #@assert(-pi ≤ α ≤ pi, "The phase angle must be in the interval [-π, π].")
        new(float(α))
    end
end

Base.angle(u::U1) = u.r
*(u1::U1, u2::U1) = U1(u1.r + u2.r)
*(u::U1, scale::Float64) = U1(u.r * scale)
*(scale::Float64, u::U1) = U1(u.r * scale)
Base.isapprox(u1::U1, u2::U1) = isapprox(u1.r, u2.r)