export GradientDescent, SecondMomentMatching, GramSchmidt, QRFact, SVDFact, orthonormalize!
# QR, SVD not exported due to conflict with LinearAlgebra identifiers.

struct GradientDescent
    maxiter::Int
    μ::Float64
    atol::Float64
    rtol::Float64
    function GradientDescent(;maxiter = 100, μ = 1.0, atol = 1e-8, rtol = 1e-8)
        return new(maxiter, μ, atol, rtol)
    end
end

struct QRFact end

struct SVDFact end

struct SecondMomentMatching end

struct GramSchmidt end

function orthonormalize!(U::AbstractMatrix, alg::GradientDescent)
    @unpack μ, atol, rtol, maxiter = alg
    r = size(U,2)
    K = U'*U
    A = Matrix{eltype(K)}(I, r, r)
    dA = similar(A)
    iter = 0
    ϵ = norm(A'*K*A - I)^2
    while ϵ > atol && ϵ > r*rtol && iter < maxiter
        dA .= - K*A*(A'*K*A - I)
        A .+= μ*dA
        iter += 1
        ϵ = norm(A'*K*A - I)^2
    end
    if iter == maxiter
        @warn "Gradient flow orthonormalization did not converge.
               Iterations exceeded maxiters = $maxiter.
               Primal residual: $(norm(A'*K*A - I)^2)"
    end
    U .= U*A
end

function orthonormalize!(U::AbstractMatrix, Z::AbstractMatrix, alg::GradientDescent)
    @unpack μ, atol, rtol, maxiter = alg
    K = U'*U
    r = size(U,2)
    A = Matrix{eltype(K)}(I, r, r)
    Ainv = Matrix{eltype(K)}(I, r, r)
    dA = similar(A)
    iter = 0
    ϵ = norm(A'*K*A - I)
    while ϵ > atol && ϵ > r*rtol && iter < maxiter
        dA .= - K*A*(A'*K*A - I)
        A .+= μ*dA
        Ainv .-= μ*Ainv*dA*Ainv
        iter += 1
        ϵ = norm(A'*K*A - I)
    end
    if iter == maxiter
        @warn "Gradient flow orthonormalization did not converge.
               Iterations exceeded maxiters = $maxiter.
               Primal residual: $(norm(A'*K*A - I)^2)"
    end
    U .= U*A
    Z .= Z*Ainv'
end

function orthonormalize!(U::AbstractMatrix, ::SecondMomentMatching)
    P = Symmetric(U'*U)
    Λ, V = eigen(P)
    U .= U*V*Diagonal(1 ./ sqrt.(Λ)) * V'
end

function orthonormalize!(U::AbstractMatrix, Z::AbstractMatrix, ::SecondMomentMatching)
    P = Symmetric(U'*U)
    Λ, V = eigen(P)
    U .= U*V*Diagonal(1 ./ sqrt.(Λ)) * V'
    Z .= Z*V*Diagonal(sqrt.(Λ)) * V'
end

function orthonormalize!(U::AbstractMatrix, ::QRFact)
    Q, _ = qr!(U)
    U .= Matrix(Q)
end

function orthonormalize!(U::AbstractMatrix, Z::AbstractMatrix, ::QRFact)
    Q, R = qr!(U)
    U .= Matrix(Q)
    Z .= Z*R'
end

function orthonormalize!(U::AbstractMatrix, ::SVDFact)
    Q, _, _ = svd!(U)
    U .= Q
end

function orthonormalize!(U::AbstractMatrix, Z::AbstractMatrix, ::SVDFact)
    Q, S, V = svd!(U)
    U .= Q
    Z .= Z*V*Diagonal(S)'
end

function orthonormalize!(U::AbstractMatrix, ::GramSchmidt)
    u = @view U[:,1]
    u ./= norm(u)
    for i in 2:size(U,2)
        u = @view U[:,i]
        subU = @view U[:,1:i-1]
        ips = subU'*u
        mul!(u,subU,ips,-1,1)
        u ./= norm(u)
    end
end

function orthonormalize!(U::AbstractMatrix, Z::AbstractMatrix, ::GramSchmidt)
    u = @view U[:,1]
    z = @view Z[:,1]
    r = norm(u)
    u ./= r
    z .*= r
    for i in 2:size(U,2)
        u = @view U[:,i]
        z = @view Z[:,i]
        subU = @view U[:,1:i-1]
        subZ = @view Z[:,1:i-1]
        ips = subU'*u
        mul!(u, subU, ips, -1, 1)
        mul!(subZ, z, ips', 1, 1)
        r = norm(u)
        u ./= r
        z .*= r
    end
end

# this implementation could return in the inverse of the GramSchmidt procedure in R
# function orthonormalize!(U, Z, ::GramSchmidt)
#     v = @view U[:,1]
#     R = UpperDiagonal(zeros(size(U,2), size(U,2)))
#     R[1,1] = norm(v)
#     v ./= R[1,1]
#     for i in 2:size(U,2)
#         v = @view U[:,i]
#         subU = @view U[:,1:i-1]
#         u = subU'*v
#         R[1:i-1,i] = u
#         R[i,i] = norm(v)
#         mul!(v,subU,u,-1,1)
#         v .*= R[i,i]
#     end
#     Z .= Z*R'
# end

# ToDo:
#  1. add general minimum second order moment matching orthonormalization:
#         https://link.springer.com/content/pdf/10.1007/s00211-021-01178-8.pdf


"""
Original author: Joe Ninan
Converted to Julia and optimized by Christian Gilbertson & Eric Ford
Returns a cubit interpolator for windowed sinc Filter curve.
no_of_points: number of intepolation points to use in cubic inteprolator
"""
function create_filter_curve_orig(no_of_points::Int; filter_size::Int=23, kaiserB::Int=13)
	x = range(-Int(floor(filter_size / 2)), stop=Int(floor(filter_size / 2)), length=no_of_points)
	Sinc = sinc.(x)
	Window = kaiser(length(x), kaiserB / π)
	FilterResponse = Window .* Sinc
	x_step = (x[end-1] - x[2]) / (no_of_points - 1)
	# append 0 to both ends far at the next node for preventing cubic spline
	# from extrapolating spurious values
	x = range(x[1] - x_step, stop=x[end] + x_step, length=no_of_points+2)
	#return CubicSplineInterpolation(x, multiple_append!([0.], FilterResponse, [0.]); extrapolation_bc=Line())
	return CubicSplineInterpolation(x, vcat(0., FilterResponse, 0.); extrapolation_bc=Line())
end

function create_filter_curve(no_of_points::Int; filter_size::Int=23, kaiserB::Int=13)
	ret_t = Float64
	x = range(-floor(Int,filter_size // 2), stop=floor(Int,filter_size // 2), length=no_of_points)
	# append 0 to both ends far at the next node for preventing cubic spline
	# from extrapolating spurious values
	Sinc = Array{ret_t,1}(undef,no_of_points+2)
	Window = Array{ret_t,1}(undef,no_of_points+2)
	Sinc[1] = zero(ret_t)
	Sinc[2:end-1] .= sinc.(x)
	Sinc[end] = zero(ret_t)
	Window[1] = zero(ret_t)
	Window[2:end-1] .= kaiser(length(x), kaiserB / π)
	Window[end] = zero(ret_t)
	FilterResponse = Window .* Sinc
	x_step = (x[end-1] - x[2]) / (no_of_points - 1)
	x = range(x[1] - x_step, stop=x[end] + x_step, length=no_of_points+2)
	return CubicSplineInterpolation(x, FilterResponse; extrapolation_bc=Line())
end

function next_min_arg(minarg::Integer, minvalue::Real, len::Integer)
	local sign = minvalue < 0
	local next_min_arg = minarg +sign -~sign
	if next_min_arg < 1       next_min_arg = len-next_min_arg end
	if next_min_arg > len     next_min_arg = next_min_arg-len end
	return next_min_arg
end

"""
Original author: Joe Ninan
Converted to Julia and optimized by Christian Gilbertson
Further adapted/optimized by Eric Ford
Additional optimizations possible by preallocating arrasy for minargs, Nminargs, minvalues, FilterValues and OldYCoords.
"""
function spectra_interpolate(
	newX::V1,
	oldX::V2,
	oldY::V3;
	PeriodicBoundary::Bool=false,
	filter_size::Int=23,
	kaiserB::Int=13,
	Filter::V4 = create_filter_curve(filter_size*21; filter_size=filter_size, kaiserB=kaiserB),
	 assume_sorted::Bool=false, verbose::Bool = false) where {
	 	T1<:Real, V1<:AbstractVector{T1}, T2<:Real, V2<:AbstractVector{T2}, T3<:Real, V3<:AbstractVector{T3}, T4<:Real, V4<:AbstractVector{T4}  }
	if !assume_sorted
		@assert issorted(newX)
		@assert issorted(oldX)
	end

	#Filter = create_filter_curve(filter_size*21; filter_size=filter_size, kaiserB=kaiserB)
	#pixarray = collect(-floor(Int,filter_size//2):floor(Int,filter_size//2))
	pixarray = -floor(Int,filter_size//2):floor(Int,filter_size//2)

	# oXsize = length(oldX)
    # First generate a 2D array of difference in pixel values
    # OldXminusNewX = np.array(oldX)[:,np.newaxis] - np.array(newX)
    # Find the minimum position to find nearest pixel for each each newX
	minargs = RvSpectMLBase.searchsortednearest(oldX, newX)
    # Pickout the those minumum values from 2D array
	minvalues = oldX[minargs] - newX
	Nminargs = next_min_arg.(minargs,minvalues,length(oldX))
	# EBF: used above to replace next few due to 1-based arrays, simplification and reduced temporaries
	#=
	sign = minvalues .< 0  # True means new X is infront of nearest old X
	# coordinate of the next adjacent bracketing point
	Nminargs = minargs .+sign -.~sign
	#Nminargs = Nminargs .% oXsize  # Periodic boundary  # WARN: Not right for 0-based arrays
	# In terms of pixel coordinates the shift values will be
	=#
	# EBF: consolidated to help reduce temporaries.  I'm not sure if if mattered.
	#=
	shiftvalues = minvalues./abs.(oldX[minargs]-oldX[Nminargs])
	# Coordinates to calculate the Filter values
	FilterCoords = shiftvalues .+ pixarray'
	FilterValues = Filter.(FilterCoords)
	=#
	FilterValues = Filter.(minvalues./abs.(oldX[minargs]-oldX[Nminargs]) .+ pixarray')
	# Coordinates to pick the values to be multiplied with Filter and summed
	OldYCoords = minargs .+ pixarray'
	if PeriodicBoundary
		# OldYCoords = OldYCoords .% oXsize  # Periodic boundary  # EBF: changed for 1-based arrays, TODO: WARN: need to test line below
		for i in 1:length(OldYCoords)
			if OldYCoords[i] < 1               		 OldYCoords[i] = length(oldX)-OldYCoords[i]   end
			if OldYCoords[i] > length(OldYCoords)+1  OldYCoords[i] = OldYCoords[i]-length(oldX)   end
		end

	else  # Extrapolate the last value till end..
		OldYCoords[OldYCoords .>= length(oldX) + 1] .= length(oldX)
		OldYCoords[OldYCoords .< 1] .= 1
	end
	# EBF consolidated to reduce temporaries.
	#OldYSlices = oldY[OldYCoords]
	#return vec(sum(OldYSlices .* FilterValues, dims=2))
	return vec(sum(oldY[OldYCoords] .* FilterValues, dims=2))
end


@doc raw"""
Lie-Trotter Splitting A

For a vector field $\dot{x} = f_1 (t,x) +  f_2 (t,x)$, the splitting reads
```math
\Phi_{\Delta t} = \varphi^{2}_{\Delta t} \circ \varphi^{1}_{\Delta t}
```

Reference:

    H. F. Trotter.
    On the product of semi-groups of operators.
    Proceedings of the American Mathematical Society, Volume 10, Pages 545-551, 1959.
    doi: 10.1090/S0002-9939-1959-0108732-6.

"""
function TableauLieA(::Type{T}=Float64) where {T}
    a = Array{T}([ 1 ])
    b = Array{T}([ 0 ])
    TableauSplittingNS(:LieTrotterSplittingA, 1, a, b)
end

@doc raw"""
Lie-Trotter Splitting B

For a vector field $\dot{x} = f_1 (t,x) +  f_2 (t,x)$, the splitting reads
```math
\Phi_{\Delta t} = \varphi^{1}_{\Delta t} \circ \varphi^{2}_{\Delta t}
```

Reference:

    H. F. Trotter.
    On the product of semi-groups of operators.
    Proceedings of the American Mathematical Society, Volume 10, Pages 545-551, 1959.
    doi: 10.1090/S0002-9939-1959-0108732-6.

"""
function TableauLieB(::Type{T}=Float64) where {T}
    a = Array{T}([ 0 ])
    b = Array{T}([ 1 ])
    TableauSplittingNS(:LieTrotterSplittingB, 1, a, b)
end

@doc raw"""
Strang Splitting

For a vector field $\dot{x} = f_1 (t,x) + f_2 (t,x)$, the splitting reads
```math
\Phi_{\Delta t} = \varphi^{1}_{\Delta t / 2} \circ \varphi^{2}_{\Delta t / 2} \circ \varphi^{2}_{\Delta t / 2} \circ \varphi^{1}_{\Delta t / 2}
```

For vector fields with two components, this is not the most efficient implementation.
For such cases [`TableauStrangA`](@ref) or [`TableauStrangB`](@ref) should be used instead.

References:

    Gilbert Strang.
    On the construction and comparison of difference schemes.
    SIAM Journal on Numerical Analysis, Volume 5, Pages 506-517, 1968.
    doi: 10.1137/0705041.

    Gurij Ivanovich Marchuk.
    Some applications of splitting-up methods to the solution of mathematical physics problems.
    Aplikace Matematiky, Volume 13, Pages 103-132, 1968.
    doi: 10.21136/AM.1968.103142.

"""
function TableauStrang(::Type{T}=Float64) where {T}
    a = Array{T}([ 1//2 ])
    b = Array{T}([ 1//2 ])
    TableauSplittingNS(:StrangSplitting, 2, a, b)
end

"Alias for [`TableauStrang`](@ref)"
TableauMarchuk = TableauStrang

@doc raw"""
Strang Splitting A for a vector field $\dot{x} = f_1 (t,x) + f_2 (t,x)$.

The splitting reads
```math
\Phi_{\Delta t} = \varphi^{1}_{\Delta t / 2} \circ \varphi^{2}_{\Delta t} \circ \varphi^{1}_{\Delta t / 2}
```

References:

    Gilbert Strang.
    On the construction and comparison of difference schemes.
    SIAM Journal on Numerical Analysis, Volume 5, Pages 506-517, 1968.
    doi: 10.1137/0705041.

    Gurij Ivanovich Marchuk.
    Some applications of splitting-up methods to the solution of mathematical physics problems.
    Aplikace Matematiky, Volume 13, Pages 103-132, 1968.
    doi: 10.21136/AM.1968.103142.

"""
function TableauStrangA(::Type{T}=Float64) where {T}
    a = Array{T}([ 1//2, 1//2 ])
    b = Array{T}([ 1//1, 0//1 ])
    TableauSplitting(:StrangSplittingA, 2, a, b)
end

@doc raw"""
Strang Splitting B for a vector field $\dot{x} = f_1 (t,x) + f_2 (t,x)$

The splitting reads
```math
\Phi_{\Delta t} = \varphi^{2}_{\Delta t / 2} \circ \varphi^{1}_{\Delta t} \circ \varphi^{2}_{\Delta t / 2}
```

References:

    Gilbert Strang.
    On the construction and comparison of difference schemes.
    SIAM Journal on Numerical Analysis, Volume 5, Pages 506-517, 1968.
    doi: 10.1137/0705041.

    Gurij Ivanovich Marchuk.
    Some applications of splitting-up methods to the solution of mathematical physics problems.
    Aplikace Matematiky, Volume 13, Pages 103-132, 1968.
    doi: 10.21136/AM.1968.103142.

"""
function TableauStrangB(::Type{T}=Float64) where {T}
    a = Array{T}([ 0//1, 1//1 ])
    b = Array{T}([ 1//2, 1//2 ])
    TableauSplitting(:StrangSplittingB, 2, a, b)
end

@doc raw"""
McLachlan's 2nd order symmetric, minimum error composition method

The composition reads
```math
\Phi_{\Delta t} = \varphi_{\alpha \Delta t} \circ \varphi^{*}_{(1/2 - \alpha) \Delta t} \circ \varphi_{(1/2 - \alpha) \Delta t} \circ \varphi^{*}_{\alpha \Delta t} ,
```
where the parameter $\alpha$ can be optimized, e.g., to minimize the solution error.
McLachlan arrives at $\alpha  = 0.1932$ as a generally useful value.

Reference:

    Robert I. McLachlan.
    On the Numerical Integration of Ordinary Differential Equations by Symmetric Composition Methods
    SIAM Journal on Scientific Computing, Volume 16, Pages 151-168, 1995.
    doi: 10.1137/0916010.

"""
function TableauMcLachlan2(::Type{T}=Float64; α=0.1932) where {T}
    a = Array{T}([ α, 0.5 - α ])
    b = Array{T}([ 0.5 - α, α ])
    TableauSplittingNS(:McLachlanSplitting, 2, a, b)
end

@doc raw"""
McLachlan's 4th order symmetric, minimum error composition method

The composition reads
```math
\Phi_{\Delta t} = \varphi_{\alpha_5 \Delta t} \circ \varphi^{*}_{\beta_5 \Delta t} \circ \dotsc \circ \varphi_{\alpha_2 \Delta t} \circ \varphi^{*}_{\beta_2 \Delta t} \circ \varphi_{\alpha_1 \Delta t} \circ \varphi^{*}_{\beta_1 \Delta t} ,
```
with
```math
\begin{aligned}
\beta_1 &= \alpha_5 = \frac{14 - \sqrt{19}}{108} , &
\alpha_1 &= \beta_5 = \frac{146 + 5 \sqrt{19}}{540} , & & \\
\beta_2 &= \alpha_4 = \frac{- 23 - 20 \sqrt{19}}{270} , &
\alpha_2 &= \beta_4 = \frac{-2 + 10 \sqrt{19}}{135} , &
\beta_3 &= \alpha_3 = \frac{1}{5} .
\end{aligned}
```
The coefficients are optimised to provide an integrator with minimal solution error.

Reference:

    Robert I. McLachlan.
    On the Numerical Integration of Ordinary Differential Equations by Symmetric Composition Methods
    SIAM Journal on Scientific Computing, Volume 16, Pages 151-168, 1995.
    doi: 10.1137/0916010.

"""
function TableauMcLachlan4(::Type{T}=Float64) where {T}
    a = Array{T}(@big [ (146 +  5*√19) / 540,
                        ( -2 + 10*√19) / 135,
                                     1 / 5,
                        (-23 - 20*√19) / 270,
                        ( 14 -    √19) / 108])
    TableauSplittingNS(:McLachlanSplitting, 4, a, a[end:-1:1])
end

@doc raw"""
4th order "Triple Jump" composition method.

The composition reads
```math
\Phi_{\Delta t} = \varphi_{\gamma_3 \Delta t} \circ \varphi_{\gamma_2 \Delta t} \circ \varphi_{\gamma_1 \Delta t}
```
with
```math
\gamma_1 = \gamma_3 = \frac{1}{2 - 2^{1/(p+1)}} , \qquad
\gamma_2 = - \frac{2^{1/(p+1)}}{2 - 2^{1/(p+1)}} .
```

References:

    Michael Creutz and Andreas Gocksch.
    Higher-order hybrid Monte Carlo algorithms.
    Physical Review Letters, Volume 63, Pages 9-12, 1989.
    doi: 10.1103/PhysRevLett.63.9.

    Etienne Forest.
    Canonical integrators as tracking codes (or how to integrate perturbation theory with tracking).
    AIP Conference Proceedings, Volume 184, Pages 1106-1136, 1989.
    doi: 10.1063/1.38062.

    Masuo Suzuki
    Fractal decomposition of exponential operators with applications to many-body theories and Monte Carlo simulations.
    Physics Letters A, Volume 146, Pages 319-323, 1990.
    doi: 10.1016/0375-9601(90)90962-N

    Haruo Yoshida.
    Construction of higher order symplectic integrators.
    Physics Letters A, Volume 150, Pages 262-268, 1990.
    doi: 10.1016/0375-9601(90)90092-3

"""
function TableauTripleJump(::Type{T}=Float64) where {T}
    fac = @big 2^(1/3)
    den = @big 1/(2-fac)
    a = Array{T}([ den, -fac*den ])
    TableauSplittingSS(:TripleJumpSplitting, 4, a)
end

@doc raw"""
Suzuki's 4th order "fractal" composition method

The composition reads
```math
\Phi_{\Delta t} = \varphi_{\gamma_5 \Delta t} \circ \varphi_{\gamma_4 \Delta t} \circ \varphi_{\gamma_3 \Delta t} \circ \varphi_{\gamma_2 \Delta t} \circ \varphi_{\gamma_1 \Delta t}
```
with
```math
\gamma_1 = \gamma_2 = \gamma_4 = \gamma_5 = \frac{1}{4 - 4^{1/(p+1)}} , \qquad
\gamma_3 = - \frac{4^{1/(p+1)}}{4 - 4^{1/(p+1)}} .
```

Reference:

    Masuo Suzuki
    Fractal decomposition of exponential operators with applications to many-body theories and Monte Carlo simulations.
    Physics Letters A, Volume 146, Pages 319-323, 1990.
    doi: 10.1016/0375-9601(90)90962-N

"""
function TableauSuzukiFractal(::Type{T}=Float64) where {T}
    fac = @big 4^(1/3)
    den = @big 1/(4-fac)
    a = Array{T}([ den, den, -fac*den ])
    TableauSplittingSS(:SuzukiFractalSplitting, 4, a)
end


d = 3    # number of dimensions
K = 50   # number of points along each dimension
N = 10^6  # number of points at which to interpolate
vecsize = 10

orders = Int64[K for i=1:d]
A = rand([K for i = 1:d]...)    # filtered coefficients
B = rand(N,d)                    # points at which to evaluate the splines

println("Comparison (d=",d,", K=", K, ", N=", N,")")

# interpolations.jl
using Interpolations
itp = interpolate(A, BSpline(Quadratic(Reflect())), OnGrid())

# what is the fastest way to operate on a list of points ?
LB = Vector{Float64}[copy(slice(B,i,:)) for i=1:size(B,1)]

function evaluate(itp, LB)
    s = 0.0
    for i=1:size(LB,1)
        I = LB[i]
        s += itp[I[1],I[2],I[3]]
    end
    return s
end

total_2 = evaluate(itp, LB)

println("interpolations.jl (cubic)")
@time total_2 = evaluate(itp, LB)
# exit()

using FixedSizeArrays

AA = rand(vecsize,orders...)
A0 = reinterpret(Vec{vecsize,Float64}, AA, (K,K,K))

bitp = interpolate(A0, BSpline(Cubic(Flat())), OnGrid())
function foo(itp, X)
    s = zero(eltype(itp))
    for x in X
        # s += itp[x[1]]
        s += itp[x[1],x[2],x[3]]
    end
    s
end

foo(bitp, LB)
println("interpolations.jl: (vector x",vecsize,")")

@time foo(bitp, LB)
# exit(0)
#
# orders = Int64[K for i in 1:d]
# a = rand(vecsize,orders...)
# b = reinterpret(Vec{vecsize,Float64}, a, (K,K,K))
# bitp = interpolate(b, BSpline(Cubic(Flat())), OnGrid())
# function foo(itp, X)
#     s = zero(eltype(itp))
#     for x in X
#         # s += itp[x[1]]
#         s += itp[x[1],x[2],x[3]]
#     end
#     s
# end
#
# # X = 2+150*rand(N)
# foo(bitp, XX) # warmup
# println("interpolations.jl (vector x",vecsize,")")
# @time foo(bitp, XX)

##################skyp
# with splines.jl #
###################

# splines.jl
using splines

# dummy boundaries
a = [0.0 for i=1:d]
b = [1.0 for i=1:d]

vals = rand(orders...)
coeffs = filter_coeffs(a,b,orders,vals)

# mcoeffs represent vecsize identical splines
mcoeffs = zeros(vecsize, size(coeffs)...)
for s in 1:vecsize
    mcoeffs[s,:] =  coeffs
end
# jit
# X = 150*rand(N,d)

println("splines.jl: (scalar)")
values_1 = eval_UC_spline(a,b,orders,coeffs,B);
tot = sum(values_1)
@time  values_1 = eval_UC_spline(a,b,orders,coeffs,B); tot = sum(values_1)

println("splines.jl: (vector x",vecsize,")")
values_1 = eval_UC_multi_spline(a,b,orders,mcoeffs,B);
tot = sum(values_1)
@time  values_1 = eval_UC_multi_spline(a,b,orders,mcoeffs,B); tot = sum(values_1)
@time  eval_UC_multi_spline!(a,b,orders,mcoeffs,B, values_1);


#=
Created on 03/02/2021 09:20:46
Last update: -

@author: Michiel Stock
michielfmstock@gmail.com

Example on uncertainty propagation.
=#

struct Measurement{T<:AbstractFloat} <: Number
    x::T
    σ::T
    function Measurement(x::T, σ::T) where {T<:AbstractFloat}
        if σ < zero(σ)
            error("Measurement error should be non-zero")
        end
        new{T}(x, σ)
    end
end

val(m::Measurement) = m.x;
err(m::Measurement) = m.σ;

Base.show(io::IO, measurement::Measurement) = print(io, "$(measurement.x) ± $(measurement.σ)");

±(x, σ) = Measurement(x, σ);

# scalar multiplication
Base.:*(a::Real, m::Measurement) = a * m.x ± abs(a) * m.σ;
Base.:/(m::Measurement, a::Real) = inv(a) * m;
# adding and substracting measurements
Base.:+(m1::Measurement, m2::Measurement) = (m1.x + m2.x) ± √(m1.σ^2 + m2.σ^2);
Base.:-(m1::Measurement, m2::Measurement) = (m1.x - m2.x) ± √(m1.σ^2 + m2.σ^2);
Base.:-(m::Measurement) = -m.x ± m.σ;
# adding a constant
Base.:+(m::Measurement, a::Real) = m + (a ± zero(a));
Base.:+(a::Real, m::Measurement) = m + a;
# multiplying two measurments
Base.:*(m1::Measurement, m2::Measurement) = m1.x * m2.x ± (m1.x * m2.x) * √((m1.σ / m1.x)^2 + (m2.σ / m2.x)^2);

m1 = 9.1 ± 0.3
m2 = 8.0 ± 0.2
m3 = 10.2 ± 0.3


function gencenters(side, λ)
    n = pois_rand(λ * side^2)
    return reduce(hcat, (rand(1:side, 2) for i in 1:n))
end

function gendisks(side, R, λ)
    spheres = zeros(Int8, (side + 2R + 1, side + 2R + 1))
    sphere  = zeros(Int8, (2R + 1, 2R + 1))
    centers = gencenters(side, λ)
    for i in -R:R
        for j in -R:R
            dist = i^2 + j^2
            if dist < R^2
                sphere[j+R+1, i+R+1] = 1
            end
        end
    end
    for center in (centers[:,i] for i in 1:size(centers,2))
        x = center[1]
        y = center[2]
        spheres[x:x + 2R, y:y + 2R] .|= sphere
    end
    return spheres[R+1:end-R-1, R+1:end-R-1]
end

function s2_theory(r, R, λ)
    tmp = r/(2R)
    tmp2 = (r > 2R) ? 2 : 2/π*(π + tmp*sqrt(1 - tmp^2) - acos(tmp))
    η = λ * π * R^2
    return exp(-η*abs(tmp2))
end

function l2_theory(r, R, λ)
    η = λ * π * R^2
    p = exp(-η)
    return p^(1 + 2r/(π*R))
end

function ss_theory(r, R, λ)
    if r < 2R
        A = 4R^2 - r^2
        B = acos(r/(2R))
        part = exp(-λ*(r*sqrt(A)/2 + 2π*R^2 - 2*B*R^2))
        return part * ((4A*B^2 - 8π*A*B + 4*π^2*A)*R^2*λ^2*r + 4*sqrt(A)*R^2*λ) / (A*r)
    else
        return (2π*λ*R)^2*exp(-2π*λ*R^2)
    end
end

function sv_theory(r, R, λ)
     if r < 2R
        A = 4R^2 - r^2
        B = acos(r/(2R))
        return 2R*λ*(π-B)exp(-sqrt(A)*r*λ/2 + 2B*R^2*λ - 2π*R^2*λ)
    else
        return 2π*R*λ*exp(-2π*λ*R^2)
    end
end

function pore_size_theory(r, R, λ)
    s(x) = 2π * x
    v(x) = π * x^2
    η  = v(R) * λ
    p  = exp(-η)
    s1 = s(r + R)
    v1 = v(r + R)
    return λ * s1 * exp(-λ * v1) / p
end

function chord_length_theory(r, R, λ)
    η = λ * π * R^2
    p = exp(-η)
    return 2η/(π*R)*p^(2r/(π*R))
end

mean_chord_length(R, λ) = 1/(2λ*R)

@testset "S2 on random overlapping disks generated by Poisson process" begin
    # Area = S^2, radius of a disk = R
    # Poisson parameter = λ
    S = 7000; R = 40; λ = 5e-5;

    disks = gendisks(S, R, λ)
    calc = Directional.s2(disks, 0; periodic = true) |> mean
    theory = [s2_theory(r-1, R, λ) for r in 0:length(calc) - 1]

    err = relerr.(calc, theory)
    @test maximum(err) < 0.15
end

@testset "L2 on random overlapping disks generated by Poisson process" begin
    # Area = S^2, radius of a disk = R
    # Poisson parameter = λ
    S = 7000; R = 40; λ = 5e-5; N = 700

    disks = gendisks(S, R, λ)
    calc = Directional.l2(disks, 0; len = N, periodic = true) |> mean .|> log
    theory = [(log ∘ l2_theory)(r-1, R, λ) for r in 0:N - 1]

    err = relerr.(calc, theory)
    @test maximum(err) < 0.15
end

@testset "SS on random overlapping disks generated by Poisson process" begin
    # Area = S^2, radius of a disk = R
    # Poisson parameter = λ
    S = 7000; R = 30; λ = 4e-4;

    disks = gendisks(S, R, λ)
    calc = Directional.surfsurf(disks, 0; periodic = true) |> mean
    theory = [ss_theory(r, R, λ) for r in 0:length(calc) - 1]

    err = relerr.(calc, theory)
    # Surface-surface function with r < 2R has a discontinuity and two
    # points where it goes to +Inf. Do not test it.
    @test maximum(err[2R + 10:end]) < 0.15
end

@testset "SV on random overlapping disks generated by Poisson process" begin
    # Area = S^2, radius of a disk = R
    # Poisson parameter = λ
    S = 7000; R = 30; λ = 4e-4;

    disks = gendisks(S, R, λ)
    calc = Directional.surfvoid(disks, 0; periodic = true) |> mean
    theory = [sv_theory(r, R, λ) for r in 0:length(calc) - 1]

    err = relerr.(calc, theory)
    @test maximum(err) < 0.2
end

@testset "Pore size on random overlapping disks generated by Poisson process" begin
    # Area = S^2, radius of a disk = R, Poisson parameter = λ
    S = 7000; R = 50; λ = 5e-5
    disks = gendisks(S, R, λ)

    calc = pore_size(disks; nbins = 20)
    edges = calc.edges[1]
    s = step(edges)
    theory = [integrate(x -> pore_size_theory(x, R, λ), n:0.05:n+s)
              for n in 0:s:s*(length(edges) - 2)]

    # Compare cummulative distributions instead of probability
    # densities because density is almost zero for big radii.
    calc_cdf = scan(calc.weights)
    theory_cdf = scan(theory)

    err = relerr.(calc_cdf, theory_cdf)
    @test maximum(err) < 0.1
end

@testset "Chord length on random overlapping disks generated by Poisson process" begin
    # Area = S^2, radius of a disk = R, Poisson parameter = λ
    S = 7000; R = 60; λ = 1e-4
    disks = gendisks(S, R, λ)

    calc, mc = Directional.chord_length(disks, 0; nbins = 35)
    edges = calc.edges[1]
    s = step(edges)
    theory = [integrate(x -> chord_length_theory(x, R, λ), n:0.05:n+s)
              for n in 0:s:s*(length(edges) - 2)]

    # Compare cummulative distributions instead of probability
    # densities because density is almost zero for big lengths.
    calc_cdf = scan(calc.weights)
    theory_cdf = scan(theory)

    err = relerr.(calc_cdf, theory_cdf)
    @test maximum(err) < 0.1
    @test relerr(mc, mean_chord_length(R, λ)) < 0.1
end


@testset "kernels.jl" begin
    kernel = TVExponentialKernel()

    # traits
    @test KernelFunctions.metric(kernel) === TotalVariation()

    # simple evaluation
    x, y = rand(10), rand(10)
    @test kernel(x, y) == exp(-totalvariation(x, y))

    # transformations
    @test (kernel ∘ ScaleTransform(0.1))(x, y) == exp(-0.1 * totalvariation(x, y))
    ard = rand(10)
    @test (kernel ∘ ARDTransform(ard))(x, y) == exp(-totalvariation(ard .* x, ard .* y))
end


using Sigma
using DataFrames

## Smokers and Friends: Undirected Graph
## =====================================

# noisyeq(a,b) = flip(ifelse(a == b, 1.0, 0.2))
noisyeq(a,b) = (a == b) | flip(.75)

# Given a list of smokers, and an edge graph
# Create the condition that if you have smoke friends you're more
# likely to be a smoker
function smokers_befriend_smokers(issmoker::RandArray, edges)
  condition = RandVar{Bool}[]
  for i = 1:length(edges), j = 1:length(edges)
    if edges[i,j] == 1
      push!(condition, noisyeq(issmoker[i],issmoker[j]))
    end
  end
  (&)(condition...)
end

# Make condition from observations of who is a smoker and who aint
function make_observes(issmoker, observed_smokers, observed_nonsmokers)
  (&)(vcat([issmoker[i] for i in observed_smokers],
           [!issmoker[i] for i in observed_nonsmokers])...)
end

function smokers(issmoker::RandArray{Bool})
  mrv_condition = smokers_befriend_smokers(issmoker)
  observations = make_observes(issmoker, observed_smokers, observed_nonsmokers)

end

function marginals(issmoker, observed_smokers, observed_nonsmokers)
  marginals = Dict{Int,(RandVar, RandVar)}[]
  for i = 1:nsmokers
    if (i ∉ observed_smokers) && (i ∉ observed_nonsmokers)
      marginals[i] = (issmoker[i], smoker_obs & mrv_condition)
    end
  end
  marginals
end

function test_smokers()
  datapath = joinpath(benchdir, "smokers","problem-7-friends.csv")
  edges = readtable(datapath)
  edges = edges[:,2:end] # remove "Node" column
  nsmokers = length(edges)

  # prior prob of being a smoker
  issmoker = RandArray([flip(.2) for i = 1:length(edges)])

  # observations
  observed_smokers = [1,8,9,15]
  observed_nonsmokers = [4,6]

  mrv_condition = smokers_befriend_smokers(issmoker, edges)
  observations = make_observes(issmoker, observed_smokers, observed_nonsmokers)
  issmoker, mrv_condition & observations
#   marginals()
end

issmoker, condition = test_smokers()
# @show condition.smt
# cond_sample_tlmh(issmoker, condition,1;split=rand_partial_split, solver=dreal3)
# ndims(condition)
# @show cond_sample_mh(issmoker_randvar,mrv_condition,10)

# ## Benchmarks
# # Query1: Compute the marginal probability of smoking for each of the
# # unobserved nodes given the six observed nodes:
# smoker_obs = (&)(vcat([issmoker[i] for i in observed_smokers],[issmoker[i] for i in observed_nonsmokers])...)

# fraction_true(x::Vector{Bool}) = count(identity,x)/length(x)

# marginals = Dict{Int,Float64}[]
# for i = 1:nsmokers
#   if (i ∉ observed_smokers) && (i ∉ observed_nonsmokers)
#     marginals[i] = cond_prob_mh(issmoker[i], smoker_obs & mrv_condition,10)
#   end
# end
# @show marginals


const c=299792458.

function IC(l,p,h,X,Y,Z,tilt,azimut,phase,amplitude,order)
  const Rx=vcat([1 0 0],[0 -1 0],[0 0 -1])
  const Ry=vcat([-1 0 0],[0 1 0],[0 0 -1])
  const Rz=vcat([-1 0 0],[0 -1 0],[0 0 1])
  const i0=amplitude*[sin(tilt).*cos(azimut), sin(tilt).*sin(azimut), cos(tilt)]
  POS=zeros((order+1)^3*8,7)

  count=1
  for i=0:order
    perc=round(i/order*1000)/10
    println("$perc %")
    for j=0:order
      for k=0:order
        n = int(i+j+k);
        ip = (Rx^i * Ry^j * Rz^k) * i0;
        x = (i+mod(i,2))*l+(-1)^i*X;
        y = (j+mod(j,2))*p+(-1)^j*Y;
        z = (k+mod(k,2))*h+(-1)^k*Z;
        POS[count,:]=[ip' x  y  z n ]
        POS[count+1,:]=[(Rz*ip)' x  y  -z n+1 ]
        POS[count+2,:]=[(Ry*ip)' x  -y  z n+1 ]
        POS[count+3,:]=[(Ry*Rz*ip)' x  -y  -z n+2 ]
        POS[count+4,:]=[(Rx*ip)' -x  y  z n+1 ]
        POS[count+5,:]=[(Rx*Rz*ip)' -x  y  -z n+2 ]
        POS[count+6,:]=[(Rx*Ry*ip)' -x  -y  z n+2 ]
        POS[count+7,:]=[(Rx*Ry*Rz*ip)' -x  -y  -z n+3 ]
        count+=8
      end
    end
  end
  perm=sortperm(POS[:,7])
  return POS[perm,:]
end


# Implement number base conversion

function all_your_base(digits, input_base, output_base)
    # error checking behavior
    if input_base < 2
        throw(DomainError(input_base, "input base must be >= 2"))
    end

    if output_base < 2
        throw(DomainError(output_base, "output base must be >= 2"))
    end

    # short-circuit for [0] input; doesn't matter the base
    if isempty(digits)
        return [0]
    end

    # encode the list to a single number
    total = 0
    for i in digits
        if i >= input_base || i < 0
            throw(DomainError(input_base, "all digits must satisfy 0 <= d < input base"))
        end
        total = total * input_base + i
    end

    # decode to desired base
    output = Int[]
    while total > 0
        (total, rem) = divrem(total, output_base)
        push!(output, rem)
    end

    if isempty(output)
        return [0]
    end

    return reverse!(output)
end


function svd_finite(A::Array{T,2}) where T<:ForwardDiff.Dual
    m = size(A,1)
    n = size(A,2)

    Av = map(ForwardDiff.value,A)
    num_A = length(Av)

    Uv,sv,Vv = svd(Av)

    dU = zeros(size(Uv,1),size(Uv,2),num_A)
    ds = zeros(length(sv),num_A)
    dV = zeros(size(Vv,1),size(Vv,2),num_A)

    ϵ = sqrt(eps(1.))
    for i = 1:num_A
        δ = zeros(size(Av))
        δ[i] = ϵ
        U,s,V = svd(Av + δ)
        dU[:,:,i] = (U - Uv) ./ ϵ
        ds[:,i] = (s - sv) ./ ϵ
        dV[:,:,i] = (V - Vv) ./ ϵ
    end

    Up = zeros(typeof(A[1].partials),size(Uv))
    sp = zeros(typeof(A[1].partials),size(sv))
    Vp = zeros(typeof(A[1].partials),size(Vv))
    for i = 1:size(Up,1)
        for j = 1:size(Up,2)
            Up[i,j] = sum([dU[i,j,k] * A[k].partials for k = 1:length(Av)])
        end
    end
    for i = 1:length(sp)
        sp[i] = sum([ds[i,k] * A[k].partials for k = 1:length(Av)])
    end
    for i = 1:size(Vp,1)
        for j = 1:size(Vp,2)
            Vp[i,j] = sum([dV[i,j,k] * A[k].partials for k = 1:length(Av)])
        end
    end

    U = map(T,Uv,Up)
    s = map(T,sv,sp)
    V = map(T,Vv,Vp)

    return U, s, V
end


# # Ginzburg-Landau model energy minimization

# ![landau_orig.png](landau_orig.png)

# Original

# ![landau_opt.png](landau_opt.png)

# Optimized

# In this example a basic Ginzburg-Landau model is solved.
# This example gives an idea of how the API together with ForwardDiff can be leveraged to
# performantly solve non standard problems on a FEM grid.
# A large portion of the code is there only for performance reasons,
# but since this usually really matters and is what takes the most time to optimize,
# it is included.

# The key to using a method like this for minimizing a free energy function directly,
# rather than the weak form, as is usually done with FEM, is to split up the
# gradient and Hessian calculations.
# This means that they are performed for each cell seperately instead of for the
# grid as a whole.

using ForwardDiff
import ForwardDiff: GradientConfig, HessianConfig, Chunk
using Ferrite
using Optim, LineSearches
using SparseArrays
using Tensors
using Base.Threads
# ## Energy terms
# ### 4th order Landau free energy
function Fl(P::Vec{3, T}, α::Vec{3}) where T
    P2 = Vec{3, T}((P[1]^2, P[2]^2, P[3]^2))
    return (α[1] * sum(P2) +
           α[2] * (P[1]^4 + P[2]^4 + P[3]^4)) +
           α[3] * ((P2[1] * P2[2]  + P2[2]*P2[3]) + P2[1]*P2[3])
end
# ### Ginzburg free energy
@inline Fg(∇P, G) = 0.5(∇P ⊡ G) ⊡ ∇P
# ### GL free energy
F(P, ∇P, params)  = Fl(P, params.α) + Fg(∇P, params.G)

# ### Parameters that characterize the model
struct ModelParams{V, T}
    α::V
    G::T
end

# ### ThreadCache
# This holds the values that each thread will use during the assembly.
struct ThreadCache{CV, T, DIM, F <: Function, GC <: GradientConfig, HC <: HessianConfig}
    cvP              ::CV
    element_indices  ::Vector{Int}
    element_dofs     ::Vector{T}
    element_gradient ::Vector{T}
    element_hessian  ::Matrix{T}
    element_coords   ::Vector{Vec{DIM, T}}
    element_potential::F
    gradconf         ::GC
    hessconf         ::HC
end
function ThreadCache(dpc::Int, nodespercell, cvP::CellValues{DIM, T}, modelparams, elpotential) where {DIM, T}
    element_indices  = zeros(Int, dpc)
    element_dofs     = zeros(dpc)
    element_gradient = zeros(dpc)
    element_hessian  = zeros(dpc, dpc)
    element_coords   = zeros(Vec{DIM, T}, nodespercell)
    potfunc          = x -> elpotential(x, cvP, modelparams)
    gradconf         = GradientConfig(potfunc, zeros(dpc), Chunk{12}())
    hessconf         = HessianConfig(potfunc, zeros(dpc), Chunk{12}())
    return ThreadCache(cvP, element_indices, element_dofs, element_gradient, element_hessian, element_coords, potfunc, gradconf, hessconf)
end

# ## The Model
# everything is combined into a model.
mutable struct LandauModel{T, DH <: DofHandler, CH <: ConstraintHandler, TC <: ThreadCache}
    dofs          ::Vector{T}
    dofhandler    ::DH
    boundaryconds ::CH
    threadindices ::Vector{Vector{Int}}
    threadcaches  ::Vector{TC}
end

function LandauModel(α, G, gridsize, left::Vec{DIM, T}, right::Vec{DIM, T}, elpotential) where {DIM, T}
    grid = generate_grid(Tetrahedron, gridsize, left, right)
    threadindices = Ferrite.create_coloring(grid)

    qr  = QuadratureRule{DIM, RefTetrahedron}(2)
    cvP = CellVectorValues(qr, Lagrange{DIM, RefTetrahedron, 1}())

    dofhandler = DofHandler(grid)
    push!(dofhandler, :P, 3)
    close!(dofhandler)

    dofvector = zeros(ndofs(dofhandler))
    startingconditions!(dofvector, dofhandler)
    boundaryconds = ConstraintHandler(dofhandler)
    #boundary conditions can be added but aren't necessary for optimization
    #add!(boundaryconds, Dirichlet(:P, getfaceset(grid, "left"), (x, t) -> [0.0,0.0,0.53], [1,2,3]))
    #add!(boundaryconds, Dirichlet(:P, getfaceset(grid, "right"), (x, t) -> [0.0,0.0,-0.53], [1,2,3]))
    close!(boundaryconds)
    update!(boundaryconds, 0.0)

    apply!(dofvector, boundaryconds)

    hessian = create_sparsity_pattern(dofhandler)
    dpc = ndofs_per_cell(dofhandler)
    cpc = length(grid.cells[1].nodes)
    caches = [ThreadCache(dpc, cpc, copy(cvP), ModelParams(α, G), elpotential) for t=1:nthreads()]
    return LandauModel(dofvector, dofhandler, boundaryconds, threadindices, caches)
end

# utility to quickly save a model
function Ferrite.vtk_save(path, model, dofs=model.dofs)
    vtkfile = vtk_grid(path, model.dofhandler)
    vtk_point_data(vtkfile, model.dofhandler, dofs)
    vtk_save(vtkfile)
end

# ## Assembly
# This macro defines most of the assembly step, since the structure is the same for
# the energy, gradient and Hessian calculations.
macro assemble!(innerbody)
    esc(quote
        dofhandler = model.dofhandler
        for indices in model.threadindices
            @threads for i in indices
                cache     = model.threadcaches[threadid()]
                eldofs    = cache.element_dofs
                nodeids   = dofhandler.grid.cells[i].nodes
                for j=1:length(cache.element_coords)
                    cache.element_coords[j] = dofhandler.grid.nodes[nodeids[j]].x
                end
                reinit!(cache.cvP, cache.element_coords)

                celldofs!(cache.element_indices, dofhandler, i)
                for j=1:length(cache.element_dofs)
                    eldofs[j] = dofvector[cache.element_indices[j]]
                end
                $innerbody
            end
        end
    end)
end

# This calculates the total energy calculation of the grid
function F(dofvector::Vector{T}, model) where T
    outs = fill(zero(T), nthreads())
    @assemble! begin
        outs[threadid()] += cache.element_potential(eldofs)
    end
    return sum(outs)
end

# The gradient calculation for each dof
function ∇F!(∇f::Vector{T}, dofvector::Vector{T}, model::LandauModel{T}) where T
    fill!(∇f, zero(T))
    @assemble! begin
        ForwardDiff.gradient!(cache.element_gradient, cache.element_potential, eldofs, cache.gradconf)
        @inbounds assemble!(∇f, cache.element_indices, cache.element_gradient)
    end
end

# The Hessian calculation for the whole grid
function ∇²F!(∇²f::SparseMatrixCSC, dofvector::Vector{T}, model::LandauModel{T}) where T
    assemblers = [start_assemble(∇²f) for t=1:nthreads()]
    @assemble! begin
        ForwardDiff.hessian!(cache.element_hessian, cache.element_potential, eldofs, cache.hessconf)
        @inbounds assemble!(assemblers[threadid()], cache.element_indices, cache.element_hessian)
    end
end

# We can also calculate all things in one go!
function calcall(∇²f::SparseMatrixCSC, ∇f::Vector{T}, dofvector::Vector{T}, model::LandauModel{T}) where T
    outs = fill(zero(T), nthreads())
    fill!(∇f, zero(T))
    assemblers = [start_assemble(∇²f, ∇f) for t=1:nthreads()]
    @assemble! begin
        outs[threadid()] += cache.element_potential(eldofs)
        ForwardDiff.hessian!(cache.element_hessian, cache.element_potential, eldofs, cache.hessconf)
        ForwardDiff.gradient!(cache.element_gradient, cache.element_potential, eldofs, cache.gradconf)
        @inbounds assemble!(assemblers[threadid()], cache.element_indices, cache.element_gradient, cache.element_hessian)
    end
    return sum(outs)
end

# ## Minimization
# Now everything can be combined to minimize the energy, and find the equilibrium
# configuration.
function minimize!(model; kwargs...)
    dh = model.dofhandler
    dofs = model.dofs
    ∇f = fill(0.0, length(dofs))
    ∇²f = create_sparsity_pattern(dh)
    function g!(storage, x)
        ∇F!(storage, x, model)
        apply_zero!(storage, model.boundaryconds)
    end
    function h!(storage, x)
        ∇²F!(storage, x, model)
        #apply!(storage, model.boundaryconds)
    end
    f(x) = F(x, model)

    od = TwiceDifferentiable(f, g!, h!, model.dofs, 0.0, ∇f, ∇²f)

    # this way of minimizing is only beneficial when the initial guess is completely off,
    # then a quick couple of ConjuageGradient steps brings us easily closer to the minimum.
    # res = optimize(od, model.dofs, ConjugateGradient(linesearch=BackTracking()), Optim.Options(show_trace=true, show_every=1, g_tol=1e-20, iterations=10))
    # model.dofs .= res.minimizer
    # to get the final convergence, Newton's method is more ideal since the energy landscape should be almost parabolic
    #+
    res = optimize(od, model.dofs, Newton(linesearch=BackTracking()), Optim.Options(show_trace=true, show_every=1, g_tol=1e-20))
    model.dofs .= res.minimizer
    return res
end

# ## Testing it
# This calculates the contribution of each element to the total energy,
# it is also the function that will be put through ForwardDiff for the gradient and Hessian.
function element_potential(eldofs::AbstractVector{T}, cvP, params) where T
    energy = zero(T)
    for qp=1:getnquadpoints(cvP)
        P  = function_value(cvP, qp, eldofs)
        ∇P = function_gradient(cvP, qp, eldofs)
        energy += F(P, ∇P, params) * getdetJdV(cvP, qp)
    end
    return energy
end

# now we define some starting conditions
function startingconditions!(dofvector, dofhandler)
    for cell in CellIterator(dofhandler)
        globaldofs = celldofs(cell)
        it = 1
        for i=1:3:length(globaldofs)
            dofvector[globaldofs[i]]   = -2.0
            dofvector[globaldofs[i+1]] = 2.0
            dofvector[globaldofs[i+2]] = -2.0tanh(cell.coords[it][1]/20)
            it += 1
        end
    end
end

δ(i, j) = i == j ? one(i) : zero(i)
V2T(p11, p12, p44) = Tensor{4, 3}((i,j,k,l) -> p11 * δ(i,j)*δ(k,l)*δ(i,k) + p12*δ(i,j)*δ(k,l)*(1 - δ(i,k)) + p44*δ(i,k)*δ(j,l)*(1 - δ(i,j)))

G = V2T(1.0e2, 0.0, 1.0e2)
α = Vec{3}((-1.0, 1.0, 1.0))
left = Vec{3}((-75.,-25.,-2.))
right = Vec{3}((75.,25.,2.))
model = LandauModel(α, G, (50, 50, 2), left, right, element_potential)

vtk_save("landauorig", model)
@time minimize!(model)
vtk_save("landaufinal", model)

# as we can see this runs very quickly even for relatively large gridsizes.
# The key to get high performance like this is to minimize the allocations inside the threaded loops,
# ideally to 0.


module TensorDecompositions
    using TensorOperations
    using Distributions
    using ProgressMeter
    using Base.Cartesian
    using StatsBase
    using LinearAlgebra

    export

    # types
    SparseArray,
    TensorDecompositions, PARAFAC2, CANDECOMP, CUR, Tucker,

    # TensorDecomposition methods
    rel_residue,
    compose, compose!,

    # additional tensor algebra methods
    tensorcontractmatrix, tensorcontractmatrices!, tensorcontractmatrices,

    # tensor decomposition algorithms
    hosvd,                   # High-order SVD
    candecomp,               # Canonical polyadic decomposition (CANDECOMP/PARAFAC)
    sshopm,                  # Shifted symmetric higher-order power method
    nncp,                    # Non-negative CANDECOMP
    spnntucker,              # sparse (semi) non-negative Tucker
    tensorcur3,              # Tensor-CUR for 3-mode tensors
    parafac2                 # PARAFAC2 model

    include("array_pool.jl")
    include("utils.jl")
    include("abstract_decomposition.jl")
    include("tucker.jl")
    include("hosvd.jl")
    include("candecomp.jl")
    include("sshopm.jl")
    include("nncp.jl")
    include("tensorcur.jl")
    include("parafac2.jl")
    include("spnntucker.jl")

end


@testset "Testing Ensemble Kalman Filter..." begin

using StateSpace, LinearAlgebra

# Define process model: Ricker growth model, with two populations
r1 = 0.03 							# 10% intrinsic growth rate for pop 1
r2 = 0.06 							# 14% intrinsic growth rate for pop 1
K1 = 20000 							# carrying capacity for pop 1
K2 = 10000 							# carrying capacity for pop 2

# process model
function f(x)
	[x[1] .* exp(r1 * (1 - (x[1] ./ K1))),
	 x[2] .* exp(r2 * (1 - (x[2] ./ K2)))]
end

V = Matrix(Diagonal([3000.0, 5000.0])) 		# covariance of (additive) process noise

# Define observation model
g(x) = 0.01 .* x
W = Matrix(Diagonal([100.0, 250.0]))

mod = NonlinearGaussianSSM(f, V, g, W)
x0 = MvNormal([1000.0, 500.0], [100.0 0.0; 0.0 100.0])

x1 = predict(mod, x0)
# println(mean(x1))

y1 = observe(mod, x1)
u1 = update(mod, x1, mean(y1))
# println(mean(u1))

xx, yy = simulate(mod, 200, x0)
yy[1, 50] = NaN

filt = EnKF()
fs = filter(mod, yy, x0, filt)
fs_ensemble = filter(mod, yy, x0, filt, return_ensemble=true)
@test size(fs_ensemble) == (mod.nx, filt.nparticles, size(yy, 2))
y_new = fs.observations[:, end] + randn(size(mod.V(1), 1)) / 10
update!(mod, fs, y_new)

ss = smooth(mod, fs)
@test loglikelihood(fs) < loglikelihood(ss)

end


using BSplineKit
using Test

@testset "Knots" begin
    breaks = -1:0.1:1
    for ξs ∈ (breaks, collect(breaks))
        k = 4
        ts = BSplines.AugmentedKnots{k}(ξs)
        ys = collect(ts)
        for x ∈ -1.05:0.05:1.05
            @test searchsortedlast(ts, x) == searchsortedlast(ys, x)
        end
    end
end


using Distributions, DataFrames, Plots

mutable struct DistributionalQAgent
    αp::Array{Real, 1}
    αn::Array{Real, 1}
    q::Array{Real, 1}
    function DistributionalQAgent(αp::Array{Real, 1}, αn::Array{Real, 1}, n::Real)
        q = Array{Real, 1}(zeros(n))
        new(αp, αn, q)
    end
end

function generate_lrs(min::Real, max::Real, n::Real)
    quantile(Uniform(min, max), 0.:(1 / (n - 1)):1.)
end

function update(agent::DistributionalQAgent, reward::Real)
    δ = ifelse.(reward .- agent.q .<= 0, -1., 1.)
    agent.q .+= δ .* ifelse.(δ .<= 0, agent.αn, agent.αp)
end

struct Environment
    p::Array{Real, 1}
    value::Array{Real, 1}
end

function step(env::Environment)
    i = rand(Categorical(env.p))
    env.value[i]
end

function step(env::Environment, agent::DistributionalQAgent)
    reward = step(env)
    update(agent, reward)
    DataFrame(agent.q', :auto)
end

N = 50
T = 1000
WARMUP = 100

αp = Array{Real, 1}(generate_lrs(0.01, 0.1, 50))
αn = reverse(αp)
agent = DistributionalQAgent(αp, αn, N)
env = Environment([0.5, 0.5], [1., 2.])

result = Array{Real, 2}(reduce(vcat, map(_ -> step(env, agent), 1:T)))
result_without_warmup = reshape(result[WARMUP+1:end, :], ((T - WARMUP) * N, 1))
p1 = plot(result, label="", xlabel="Trial", ylabel="Q-value")
p2 = histogram(result_without_warmup, label="", xlabel="Q-value", ylabel="counts")
plot(p1, p2)


function _sqrt(a::Bool, n)
    if a
        return n > 0 ? 2.0 : 1.0
    else
        return 1.0
    end
end

function _napier(a::Bool, n)
    if a
        return n > 0 ? Float64(n) : 2.0
    else
        return n > 1 ? n - 1.0 : 1.0
    end
end

function _pi(a::Bool, n)
    if a
        return n > 0 ? 6.0 : 3.0
    else
        return (2.0 * n - 1.0) ^ 2.0 # exponentiation operator
    end
end

function calc(f::Function, expansions::Integer)
    a, b = true, false
    r = 0.0
    for i in expansions:-1:1
        r = f(b, i) / (f(a, i) + r)
    end
    return f(a, 0) + r
end

for (v, f) in (("√2", _sqrt), ("e", _napier), ("π", _pi))
    @printf("%3s = %f\n", v, calc(f, 1000))
end


@testset "Misc/MPolyFactor" begin

  function check_factorization(num::Integer, a)
    f = factor(a)
    g = f.unit
    for i in f.fac
      num -= i[2]
      g *= i[1]^i[2]
    end
    @test num <= 0
    @test g == a
  end

  R, (x, y, z) = PolynomialRing(Hecke.QQ, ["x", "y", "z"])
  check_factorization(6, x^24-y^24*z^12)

  R, (x, y, z) = PolynomialRing(Hecke.QQ, ["x", "y", "z"])
  check_factorization(2, ((x^2+y^2+z^2+2)*(x+1)*(y+2)*(z+3)+x*y*z)*(x^2+y^2+z^2+3))

  R, (x, y, z, t) = PolynomialRing(Hecke.QQ, ["x", "y", "z", "t"])
  for n in 1:8
    count = 0
    for i in 1:2:n
        if n%i == 0
            count += 1
        end
    end
    check_factorization(count, ((1+x+y+z+t)^n+1)*((1+x+y+z+t)^n+2))
  end

  R, (x, y, z, u, v) = PolynomialRing(Hecke.QQ, ["x", "y", "z", "u", "v"])
  check_factorization(2, (y^2*v*(z+v)*x+v+y)*(v^2*(y+z)*x+z^2))

  QA, A = PolynomialRing(Hecke.QQ, "A")
  K, a = number_field(A^3 - 2, "a")
  R, (x, y, z) = PolynomialRing(K, ["x", "y", "z"])
  check_factorization(0, R(0))
  check_factorization(0, R(1))
  check_factorization(0, R(2))
  check_factorization(6, 2*x^2*y*(1+y)*(x^3-2))
  check_factorization(7, 2*x^2*y*(1+y)*(2*x^3-1)*(a*x+a^2*y+1)*(a^2*x+a*y+1)^2)
  check_factorization(2, x^6-2*y^3)
  check_factorization(2, (y*z*x^2+z-1)*(x+a*z^3*y))
  check_factorization(3, (x+y+a)*(x+a*y+1)*(x+y+2*a))
  check_factorization(3, (x^2+y^2+a)*(x^3+a*y^3+1)*(x^2+y+2*a))
  check_factorization(8, (1+x+a*y+x*y)^3*(1+x+a^2*y+x*y^2)^5)
  check_factorization(4, (x-(a+1)^4*y+z)*(a*x+y+z)*(x+a*y+z)*(x+y+a*z))

end


using PlotlyJS

function contour1()
    x = [-9, -6, -5 , -3, -1]
    y = [0, 1, 4, 5, 7]
    z = [10     10.625  12.5  15.625  20
         5.625  6.25    8.125 11.25   15.625
         2.5    3.125   5.    8.125   12.5
         0.625  1.25    3.125 6.25    10.625
         0      0.625   2.5   5.625   10]
    trace = contour(x=x, y=y, z=z)

    layout = Layout(title="Setting the X and Y Coordinates in a Contour Plot")
    plot(trace, layout)
end


using TaylorSeries
# computes all available derivatives, including value (0th order derivative)
derivatives(t::Taylor1) = derivatives!(similar(t.coeffs), t)
function derivatives!(d::AbstractVector, t::Taylor1)
    c, n = t.coeffs, t.order
    @. d = c * factorial(0:n)
end
# computes all derivatives up to mth order (including value)
function derivatives(f, x::Real, m::Int)
    d = zeros(eltype(x), m+1)
    derivatives!(d, f, x)
end
function derivatives(f, x::AbstractVector{<:Real}, m::Int)
    d = zeros(eltype(x), m+1, length(x))
    for (i, xi) in enumerate(x)
        d[:, i] .= derivatives(f, xi, m)
    end
    return d
end

function derivatives!(d::AbstractVector, f, x::Real)
    m = length(d) - 1
    @. d = 0
    d[1] = x
    d[2] = 1
    x = Taylor1(d, m) # mth order taylor series
    fx = f(x) # TODO: in place?
    derivatives!(d, fx) # overwrites coeffs vector with derivatives
end

# IDEA: could specialize for popular kernel types
# function derivatives!(d::AbstractVector, f::CovarianceFunctions.Exp, r::Real)
#     m = length(d)-1
#     @. d = f(r)
# end


#=
Various routines to discretize AR(1) processes

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date : 2014-04-10 23:55:05

References
----------

https://lectures.quantecon.org/jl/finite_markov.html
=#

import Optim
import NLopt
import Distributions: pdf, Normal, quantile

std_norm_cdf(x::T) where {T <: Real} = 0.5 * erfc(-x/sqrt(2))
std_norm_cdf(x::Array{T}) where {T <: Real} = 0.5 .* erfc(-x./sqrt(2))

@doc doc"""
Tauchen's (1996) method for approximating AR(1) process with finite markov chain

The process follows

```math
    y_t = \mu + \rho y_{t-1} + \epsilon_t
```

where ``\epsilon_t \sim N (0, \sigma^2)``

##### Arguments

- `N::Integer`: Number of points in markov process
- `ρ::Real` : Persistence parameter in AR(1) process
- `σ::Real` : Standard deviation of random component of AR(1) process
- `μ::Real(0.0)` : Mean of AR(1) process
- `n_std::Integer(3)` : The number of standard deviations to each side the process
  should span

##### Returns

- `mc::MarkovChain` : Markov chain holding the state values and transition matrix

"""
function tauchen(N::Integer, ρ::T1, σ::T2, μ=zero(promote_type(T1, T2)), n_std::Integer=3) where {T1 <: Real, T2 <: Real}
    # Get discretized space
    a_bar = n_std * sqrt(σ^2 / (1 - ρ^2))
    y = range(-a_bar, stop=a_bar, length=N)
    d = y[2] - y[1]

    # Get transition probabilities
    Π = zeros(promote_type(T1, T2), N, N)
    for row = 1:N
        # Do end points first
        Π[row, 1] = std_norm_cdf((y[1] - ρ*y[row] + d/2) / σ)
        Π[row, N] = 1 - std_norm_cdf((y[N] - ρ*y[row] - d/2) / σ)

        # fill in the middle columns
        for col = 2:N-1
            Π[row, col] = (std_norm_cdf((y[col] - ρ*y[row] + d/2) / σ) -
                           std_norm_cdf((y[col] - ρ*y[row] - d/2) / σ))
        end
    end

    # NOTE: I need to shift this vector after finding probabilities
    #       because when finding the probabilities I use a function
    #       std_norm_cdf that assumes its input argument is distributed
    #       N(0, 1). After adding the mean E[y] is no longer 0, so
    #       I would be passing elements with the wrong distribution.
    #
    #       It is ok to do after the fact because adding this constant to each
    #       term effectively shifts the entire distribution. Because the
    #       normal distribution is symmetric and we just care about relative
    #       distances between points, the probabilities will be the same.
    #
    #       I could have shifted it before, but then I would need to evaluate
    #       the cdf with a function that allows the distribution of input
    #       arguments to be [μ/(1 - ρ), 1] instead of [0, 1]

    yy = y .+ μ / (1 - ρ) # center process around its mean (wbar / (1 - rho)) in new variable

    # renormalize. In some test cases the rows sum to something that is 2e-15
    # away from 1.0, which caused problems in the MarkovChain constructor
    Π = Π./sum(Π, dims = 2)

    MarkovChain(Π, yy)
end

@doc doc"""
Rouwenhorst's method to approximate AR(1) processes.

The process follows

```math
    y_t = \mu + \rho y_{t-1} + \epsilon_t
```

where ``\epsilon_t \sim N (0, \sigma^2)``

##### Arguments
- `N::Integer` : Number of points in markov process
- `ρ::Real` : Persistence parameter in AR(1) process
- `σ::Real` : Standard deviation of random component of AR(1) process
- `μ::Real(0.0)` :  Mean of AR(1) process

##### Returns

- `mc::MarkovChain{Float64}` : Markov chain holding the state values and transition matrix

"""
function rouwenhorst(N::Integer, ρ::Real, σ::Real, μ::Real=0.0)
    σ_y = σ / sqrt(1-ρ^2)
    p  = (1+ρ)/2
    Θ = [p 1-p; 1-p p]
    ψ = sqrt(N-1) * σ_y
    m = μ / (1 - ρ)

    state_values, p = _rouwenhorst(p, p, m, ψ, N)
    MarkovChain(p, state_values)
end

function _rouwenhorst(p::Real, q::Real, m::Real, Δ::Real, n::Integer)
    if n == 2
        return [m-Δ, m+Δ],  [p 1-p; 1-q q]
    else
        _, θ_nm1 = _rouwenhorst(p, q, m, Δ, n-1)
        θN = p    *[θ_nm1 zeros(n-1, 1); zeros(1, n)] +
             (1-p)*[zeros(n-1, 1) θ_nm1; zeros(1, n)] +
             q    *[zeros(1, n); zeros(n-1, 1) θ_nm1] +
             (1-q)*[zeros(1, n); θ_nm1 zeros(n-1, 1)]

        θN[2:end-1, :] ./= 2

        return range(m-Δ, stop=m+Δ, length=n), θN
    end
end

# These are to help me order types other than vectors
@inline _emcd_lt(a::T, b::T) where {T} = isless(a, b)
# @inline _emcd_lt(a::Vector{T}, b::Vector{T}) where {T} = Base.lt(Base.Order.Lexicographic, a, b)

@doc doc"""
Accepts the simulation of a discrete state Markov chain and estimates
the transition probabilities

Let ``S = s_1, s_2, \ldots, s_N`` with ``s_1 < s_2 < \ldots < s_N`` be the discrete
states of a Markov chain. Furthermore, let ``P`` be the corresponding
stochastic transition matrix.

Given a history of observations, ``\{X\}_{t=0}^{T}`` with ``x_t \in S \forall t``,
we would like to estimate the transition probabilities in ``P`` with ``p_{ij}``
as the ith row and jth column of ``P``. For ``x_t = s_i`` and ``x_{t-1} = s_j``,
let ``P(x_t | x_{t-1})`` be defined as ``p_{i,j}`` element of the stochastic
matrix. The likelihood function is then given by

```math
  L(\{X\}^t; P) = \text{Prob}(x_1) \prod_{t=2}^{T} P(x_t | x_{t-1})
```

The maximum likelihood estimate is then just given by the number of times
a transition from ``s_i`` to ``s_j`` is observed divided by the number of times
``s_i`` was observed.

Note: Because of the estimation procedure used, only states that are observed
in the history appear in the estimated Markov chain... It can't divine whether
there are unobserved states in the original Markov chain.

For more info, refer to:

- http://www.stat.cmu.edu/~cshalizi/462/lectures/06/markov-mle.pdf
- https://stats.stackexchange.com/questions/47685/calculating-log-likelihood-for-given-mle-markov-chains

##### Arguments

- `X::Vector{T}` : Simulated history of Markov states

##### Returns

- `mc::MarkovChain{T}` : A Markov chain holding the state values and
  transition matrix

"""
function estimate_mc_discrete(X::Vector{T}, states::Vector{T}) where T
    # Get length of simulation
    capT = length(X)

    # Make sure all of the passed in states appear in X... If not
    # throw an error
    if any(!in(x, X) for x in states)
        error("One of the states does not appear in history X")
    end

    # Count states and store in dictionary
    nstates = length(states)
    d = Dict{T, Int}(zip(states, 1:nstates))

    # Counter matrix and dictionary mapping i -> states
    cm = zeros(nstates, nstates)

    # Compute conditional probabilities for each state
    state_i = d[X[1]]
    for t in 1:capT-1
        # Find next period's state
        state_j = d[X[t+1]]
        cm[state_i, state_j] += 1.0

        # Tomorrow's state is j
        state_i = state_j
    end

    # Compute probabilities using counted elements
    P = cm ./ sum(cm, dims = 2)

    return MarkovChain(P, states)
end

function estimate_mc_discrete(X::Vector{T}) where T
    # Get unique states and sort them
    states = sort!(unique(X); lt=_emcd_lt)

    return estimate_mc_discrete(X, states)
end

@doc doc"""

types specifying the method for `discrete_var`

"""
abstract type VAREstimationMethod end
struct Even <: VAREstimationMethod end
struct Quantile <: VAREstimationMethod end
struct Quadrature <: VAREstimationMethod end

@doc doc"""

Compute a finite-state Markov chain approximation to a VAR(1) process of the form

```math
    y_{t+1} = b + By_{t} + \Psi^{\frac{1}{2}}\epsilon_{t+1}
```

where ``\epsilon_{t+1}`` is an vector of independent standard normal
innovations of length `M`

```julia
P, X = discrete_var(b, B, Psi, Nm, n_moments, method, n_sigmas)
```

##### Arguments

- `b::Union{Real, AbstractVector}` : constant vector of length `M`.
                                     `M=1` corresponds scalar case
- `B::Union{Real, AbstractMatrix}` : `M x M` matrix of impact coefficients
- `Psi::Union{Real, AbstractMatrix}` : `M x M` variance-covariance matrix of
                                       the innovations
    - `discrete_var` only accepts non-singular variance-covariance matrices, `Psi`.
- `Nm::Integer > 3` : Desired number of discrete points in each dimension

##### Optional

- `n_moments::Integer` : Desired number of moments to match. The default is 2.
- `method::VAREstimationMethod` : Specify the method used to determine the grid
                                  points. Accepted inputs are `Even()`, `Quantile()`,
                                  or `Quadrature()`. Please see the paper for more details.
- `n_sigmas::Real` : If the `Even()` option is specified, `n_sigmas` is used to
                     determine the number of unconditional standard deviations
                     used to set the endpoints of the grid. The default is
                     `sqrt(Nm-1)`.

##### Returns

- `P` : `Nm^M x Nm^M` probability transition matrix. Each row
        corresponds to a discrete conditional probability
        distribution over the state M-tuples in `X`
- `X` : `M x Nm^M` matrix of states. Each column corresponds to an
        M-tuple of values which correspond to the state associated
        with each row of `P`

##### NOTES

- discrete_var only constructs tensor product grids where each dimension
  contains the same number of points. For this reason it is recommended
  that this code not be used for problems of more than about 4 or 5
  dimensions due to curse of dimensionality issues.
- Future updates will allow for singular variance-covariance matrices and
  sparse grid specifications.

##### Reference

- Farmer, L. E., & Toda, A. A. (2017).
  "Discretizing nonlinear, non‐Gaussian Markov processes with exact conditional moments,"
  Quantitative Economics, 8(2), 651-683.

"""
function discrete_var(b::Union{Real, AbstractVector},
                      B::Union{Real, AbstractMatrix},
                      Psi::Union{Real, AbstractMatrix},
                      Nm::Integer,
                      n_moments::Integer=2,
                      method::VAREstimationMethod=Even(),
                      n_sigmas::Real=sqrt(Nm-1))
    # b = zeros(2)
    # A = [0.9809 0.0028; 0.041 0.9648]
    # Sigma = [7.569e-5 0.0; 0.0 0.00068644]
    # N = 9
    # n_moments = nMoments
    # method = Quantile()
    # b, B, Psi, Nm = (zeros(2), A, Sigma, N, nMoments, Quantile())
    M, M_ = size(B, 1), size(B, 2)

    # Check size restrictions on matrices
    M == M_ || throw(ArgumentError("B must be a scalar or square matrix"))
    M == length(b) || throw(ArgumentError("b must have the same number of rows as B"))

    #% Check that Psi is a valid covariance matrix
    isposdef(Psi) || throw(ArgumentError("Psi must be a positive definite matrix"))

    # Check that Nm is a valid number of grid points
    Nm >= 3 || throw(ArgumentError("Nm must be a positive interger greater than 3"))

    # Check that n_moments is a valid number
    if n_moments < 1 || !(n_moments % 2 == 0 || n_moments == 1)
        error("n_moments must be either 1 or a positive even integer")
    end

    # warning about persistency
    warn_persistency(B, method)

    # Compute polynomial moments of standard normal distribution
    gaussian_moment = zeros(n_moments)
    c = 1
    for k=1:floor(Int, n_moments/2)
        c = (2*k-1)*c
        gaussian_moment[2*k] = c
    end
    # Compute standardized VAR(1) representation
    # (zero mean and diagonal covariance matrix)
    A, C, mu, Sigma = standardize_var(b, B, Psi, M)

    # Construct 1-D grids
    y1D, y1Dhelper = construct_1D_grid(Sigma, Nm, M, n_sigmas, method)

    # Construct all possible combinations of elements of the 1-D grids
    D = allcomb3(y1D')'

    # Construct finite-state Markov chain approximation
    # conditional mean of the VAR process at each grid point
    cond_mean = A*D
    # probability transition matrix
    P = ones(Nm^M, Nm^M)
    # normalizing constant for maximum entropy computations
    scaling_factor = y1D[:, end]
    # used to store some intermediate calculations
    temp = Matrix{Float64}(undef, Nm, M)
    # store optimized values of lambda (2 moments) to improve initial guesses
    lambda_bar = zeros(2*M, Nm^M)
    # small positive constant for numerical stability
    kappa = 1e-8

    for ii = 1:(Nm^M)

        # Construct prior guesses for maximum entropy optimizations
        q = construct_prior_guess(cond_mean[:, ii], Nm, y1D, y1Dhelper, method)

        # Make sure all elements of the prior are stricly positive
        q[q.<kappa] .= kappa

        for jj = 1:M
            # Try to use intelligent initial guesses
            if ii == 1
                lambda_guess = zeros(2)
            else
                lambda_guess = lambda_bar[(jj-1)*2+1:jj*2, ii-1]
            end

            # Maximum entropy optimization
            if n_moments == 1 # match only 1 moment
                temp[:, jj], _, _ = discrete_approximation(y1D[jj, :],
                    X -> (X'.-cond_mean[jj, ii])/scaling_factor[jj],
                    [0.0], q[jj, :], [0.0])
            else # match 2 moments first
                p, lambda, moment_error = discrete_approximation(y1D[jj, :],
                    X -> polynomial_moment(X, cond_mean[jj, ii], scaling_factor[jj], 2),
                    [0; 1]./(scaling_factor[jj].^(1:2)), q[jj, :], lambda_guess)
                if !(norm(moment_error) < 1e-5) # if 2 moments fail, just match 1 moment
                    @warn("Failed to match first 2 moments. Just matching 1.")
                    temp[:, jj], _, _ = discrete_approximation(y1D[jj, :],
                    X -> (X'.-cond_mean[jj, ii])/scaling_factor[jj],
                        [0.0], q[jj, :], [0.0])
                    lambda_bar[(jj-1)*2+1:jj*2, ii] = zeros(2,1)
                elseif n_moments == 2
                    lambda_bar[(jj-1)*2+1:jj*2, ii] = lambda
                    temp[:, jj] = p
                else # solve maximum entropy problem sequentially from low order moments
                    lambda_bar[(jj-1)*2+1:jj*2, ii] = lambda
                    for mm = 4:2:n_moments
                        lambda_guess = vcat(lambda, 0.0, 0.0) # add 0 to previous lambda
                        pnew, lambda, moment_error = discrete_approximation(y1D[jj,:],
                            X -> polynomial_moment(X, cond_mean[jj,ii],
                                                    scaling_factor[jj], mm),
                            gaussian_moment[1:mm]./(scaling_factor[jj].^(1:mm)),
                            q[jj, :], lambda_guess)
                        if !(norm(moment_error) < 1e-5)
                            @warn(
                            "Failed to match first $mm moments.  Just matching $(mm-2).")
                            break
                        else
                            p = pnew
                        end
                    end
                    temp[:, jj] = p
                end
            end
        end
        P[ii, :] .= vec(prod(allcomb3(temp), dims = 2))
    end

    X = C*D .+ mu # map grids back to original space

    M != 1 || (return MarkovChain(P, vec(X)))
    return MarkovChain(P, [X[:, i] for i in 1:Nm^M])
end

"""
check persistency when `method` is `Quadrature` and give warning if needed

##### Arguments
- `B::Union{Real, AbstractMatrix}` : impact coefficient
- `method::VAREstimationMethod` : method for grid making

##### Returns
- `nothing`

"""
function warn_persistency(B::AbstractMatrix, ::Quadrature)
    if any(eigen(B).values .> 0.9)
        @warn("The quadrature method may perform poorly for highly persistent processes.")
    end
    return nothing
end
warn_persistency(B::Real, method::Quadrature) = warn_persistency(fill(B, 1, 1), method)
warn_persistency(::Union{Real, AbstractMatrix}, ::VAREstimationMethod) = nothing

"""

return standerdized AR(1) representation

##### Arguments

- `b::Real` : constant term
- `B::Real` : impact coefficient
- `Psi::Real` : variance of innovation
- `M::Integer == 1` : must be one since the function is for AR(1)

##### Returns

- `A::Real` : impact coefficient of standardized AR(1) process
- `C::Real` : standard deviation of the innovation
- `mu::Real` : mean of the standardized AR(1) process
- `Sigma::Real` : variance of the standardized AR(1) process

"""
function standardize_var(b::Real, B::Real, Psi::Real, M::Integer)
    C = sqrt(Psi)
    A = B
    mu = [b/(1-B)] # mean of the process
    Sigma = 1/(1-B^2) #
    return A, C, mu, Sigma
end

"""

return standerdized VAR(1) representation

##### Arguments

- `b::AbstractVector` : `M x 1` constant term vector
- `B::AbstractMatrix` : `M x M` matrix of impact coefficients
- `Psi::AbstractMatrix` : `M x M` variance-covariance matrix of innovations
- `M::Intger` : number of variables of the VAR(1) model

##### Returns

- `A::Matirx` : impact coefficients of standardized VAR(1) process
- `C::AbstractMatrix` : variance-covariance matrix of standardized model innovations
- `mu::AbstractVector` : mean of the standardized VAR(1) process
- `Sigma::AbstractMatrix` : variance-covariance matrix of the standardized VAR(1) process

"""
function standardize_var(b::AbstractVector, B::AbstractMatrix,
                         Psi::AbstractMatrix, M::Integer)
    C1 = cholesky(Psi).L
    mu = ((I - B)\ I)*b
    A1 = C1\(B*C1)
    # unconditional variance
    Sigma1 = reshape(((I-kron(A1,A1))\I)*vec(Matrix(I, M, M)),M,M)
    U, _ = min_var_trace(Sigma1)
    A = U'*A1*U
    Sigma = U'*Sigma1*U
    C = C1*U
    return A, C, mu, Sigma
end

"""
construct prior guess for evenly spaced grid method

##### Arguments

- `cond_mean::AbstractVector` : conditional Mean of each variable
- `Nm::Integer` : number of grid points
- `y1D::AbstractMatrix` : grid of variable
- `::AbstractMatrix` : bounds of each grid bin
- `method::Even` : method for grid making

"""
construct_prior_guess(cond_mean::AbstractVector, Nm::Integer,
                      y1D::AbstractMatrix, ::Nothing, method::Even) =
    pdf.(Normal.(repeat(cond_mean, 1, Nm), 1), y1D)

"""
construct prior guess for quantile grid method

##### Arguments

- `cond_mean::AbstractVector` : conditional Mean of each variable
- `Nm::Integer` : number of grid points
- `::AbstractMatrix` : grid of variable
- `y1Dbounds::AbstractMatrix` : bounds of each grid bin
- `method::Quantile` : method for grid making

"""
construct_prior_guess(cond_mean::AbstractVector, Nm::Integer,
                      ::AbstractMatrix, y1Dbounds::AbstractMatrix, method::Quantile) =
    cdf.(Normal.(repeat(cond_mean, 1, Nm), 1), y1Dbounds[:, 2:end]) -
           cdf.(Normal.(repeat(cond_mean, 1, Nm), 1), y1Dbounds[:, 1:end-1])

"""
construct prior guess for quadrature grid method

##### Arguments

- `cond_mean::AbstractVector` : conditional Mean of each variable
- `Nm::Integer` : number of grid points
- `y1D::AbstractMatrix` : grid of variable
- `weights::AbstractVector` : weights of grid `y1D`
- `method::Quadrature` : method for grid making

"""
construct_prior_guess(cond_mean::AbstractVector, Nm::Integer,
                      y1D::AbstractMatrix, weights::AbstractVector, method::Quadrature) =
    (pdf.(Normal.(repeat(cond_mean, 1, Nm), 1), y1D) ./ pdf.(Ref(Normal()), y1D)).*weights'
"""

construct one-dimensional evenly spaced grid of states

##### Argument

- `Sigma::ScalarOrArray` : variance-covariance matrix of the standardized process
- `Nm::Integer` : number of grid points
- `M::Integer` : number of variables (`M=1` corresponds to AR(1))
- `n_sigmas::Real` : number of standard error determining end points of grid
- `method::Even` : method for grid making

##### Return

- `y1D` : `M x Nm` matrix of variable grid
- `nothing` : `nothing` of type `Void`

"""
function construct_1D_grid(Sigma::Union{Real, AbstractMatrix}, Nm::Integer,
                           M::Integer, n_sigmas::Real, method::Even)
    min_sigmas = sqrt(minimum(eigen(Sigma).values))
    y1Drow = collect(range(-min_sigmas*n_sigmas, stop=min_sigmas*n_sigmas, length=Nm))'
    y1D = repeat(y1Drow, M, 1)
    return y1D, nothing
end

"""

construct one-dimensional quantile grid of states

##### Argument

- `Sigma::AbstractMatrix` : variance-covariance matrix of the standardized process
- `Nm::Integer` : number of grid points
- `M::Integer` : number of variables (`M=1` corresponds to AR(1))
- `n_sigmas::Real` : number of standard error determining end points of grid
- `method::Quntile` : method for grid making

##### Return

- `y1D` : `M x Nm` matrix of variable grid
- `y1Dbounds` : bounds of each grid bin

"""
function construct_1D_grid(Sigma::AbstractMatrix, Nm::Integer,
                           M::Integer, n_sigmas::Real, method::Quantile)
    sigmas = sqrt.(diag(Sigma))
    y1D = quantile.(Normal.(0, sigmas), (2*(1:Nm)'.-1)/(2*Nm))
    y1Dbounds = hcat(fill(-Inf, M, 1),
                     quantile.(Normal.(0, sigmas), ((1:(Nm-1))')/Nm),
                     fill(Inf, M, 1))
    return y1D, y1Dbounds
end
construct_1D_grid(Sigma::Real, Nm::Integer,
                  M::Integer, n_sigmas::Real, method::Quantile) =
    construct_1D_grid(fill(Sigma, 1, 1), Nm, M, n_sigmas, method)

"""

construct one-dimensional quadrature grid of states

##### Argument

- `::ScalarOrArray` : not used
- `Nm::Integer` : number of grid points
- `M::Integer` : number of variables (`M=1` corresponds to AR(1))
- `n_sigmas::Real` : not used
- `method::Quadrature` : method for grid making

##### Return

- `y1D` : `M x Nm` matrix of variable grid
- `weights` : weights on each grid

"""
function construct_1D_grid(::ScalarOrArray, Nm::Integer,
                           M::Integer, ::Real, method::Quadrature)
    nodes, weights = qnwnorm(Nm, 0, 1)
    y1D = repeat(nodes', M, 1)
    return y1D, weights
end

"""

Return combinations of each column of matrix `A`.
It is simiplifying `allcomb2` by using `gridmake` from QuantEcon

##### Arguments

- `A::AbstractMatrix` : `N x M` Matrix

##### Returns

- `N^M x M` Matrix, combination of each row of `A`.

###### Example
```julia-repl
julia> allcomb3([1 4 7;
                 2 5 8;
                 3 6 9]) # numerical input
27×3 Array{Int64,2}:
 1  4  7
 1  4  8
 1  4  9
 1  5  7
 1  5  8
 1  5  9
 1  6  7
 1  6  8
 1  6  9
 2  4  7
 ⋮
 2  6  9
 3  4  7
 3  4  8
 3  4  9
 3  5  7
 3  5  8
 3  5  9
 3  6  7
 3  6  8
 3  6  9
```

"""
allcomb3(A::AbstractMatrix) =
    reverse(gridmake(reverse([A[:, i] for i in 1:size(A, 2)], dims = 1)...), dims = 2)

"""
Compute a discrete state approximation to a distribution with known moments,
using the maximum entropy procedure proposed in Tanaka and Toda (2013)

```julia
p, lambda_bar, moment_error = discrete_approximation(D, T, TBar, q, lambda0)
```

##### Arguments

- `D::AbstractVector` : vector of grid points of length `N`.
                        N is the number of points at which an approximation
                        is to be constructed.
- `T::Function` : A function that accepts a single `AbstractVector` of length `N`
                  and returns an `L x N` matrix of moments evaluated
                  at each grid point, where L is the number of moments to be
                  matched.
- `TBar::AbstractVector` : length `L` vector of moments of the underlying distribution
                           which should be matched

##### Optional

- `q::AbstractVector` : length `N` vector of prior weights for each point in D.
                        The default is for each point to have an equal weight.
- `lambda0::AbstractVector` : length `L` vector of initial guesses for the dual problem
                              variables. The default is a vector of zeros.

##### Returns

- `p` : (1 x N) vector of probabilties assigned to each grid point in `D`.
- `lambda_bar` : length `L` vector of dual problem variables which solve the
                maximum entropy problem
- `moment_error` : vector of errors in moments (defined by moments of
                  discretization minus actual moments) of length `L`

"""
function discrete_approximation(D::AbstractVector, T::Function, TBar::AbstractVector,
                                  q::AbstractVector=ones(N)/N, # Default prior weights
                                  lambda0::AbstractVector=zeros(Tbar))

    # Input error checking
    N = length(D)

    Tx = T(D)
    L, N2 = size(Tx)

    if N2 != N || length(TBar) != L || length(lambda0) != L || length(q) != N2
        error("Dimension mismatch")
    end

    # Compute maximum entropy discrete distribution
    options = Optim.Options(f_tol=1e-16, x_tol=1e-16)
    obj(lambda) = entropy_obj(lambda, Tx, TBar, q)
    grad!(grad, lambda) = entropy_grad!(grad, lambda, Tx, TBar, q)
    hess!(hess, lambda) = entropy_hess!(hess, lambda, Tx, TBar, q)
    res = Optim.optimize(obj, grad!, hess!, lambda0, Optim.Newton(), options)
    # Sometimes the algorithm fails to converge if the initial guess is too far
    # away from the truth. If this occurs, the program tries an initial guess
    # of all zeros.
    if !Optim.converged(res) && all(lambda0 .!= 0.0)
        @warn("Failed to find a solution from provided initial guess. Trying new initial guess.")
        res = Optim.optimize(obj, grad!, hess!, zeros(lambda0), Optim.Newton(), options)
        # check convergence
        Optim.converged(res) || error("Failed to find a solution.")
    end
    # Compute final probability weights and moment errors
    lambda_bar = Optim.minimizer(res)
    minimum_value = Optim.minimum(res)

    Tdiff = Tx .- TBar
    p = (q'.*exp.(lambda_bar'*Tdiff))/minimum_value
    grad = similar(lambda0)
    grad!(grad, Optim.minimizer(res))
    moment_error = grad/minimum_value
    return p, lambda_bar, moment_error
end
# discrete_approximation(D::AbstractVector, T::Function, TBar::Real,
#                        q::AbstractVector=ones(N)/N, lambda0::Real=0) =
#     discrete_approximation(D, T, [TBar], q, [lambda0])

"""

Compute the moment defining function used in discrete_approximation

```julia
T = polynomial_moment(X, mu, scaling_factor, mMoments)
```

##### Arguments:

- `X::AbstractVector` : length `N` vector of grid points
- `mu::Real` : location parameter (conditional mean)
- `scaling_factor::Real` : scaling factor for numerical stability.
                          (typically largest grid point)
- `n_moments::Integer` : number of polynomial moments

##### Return

- `T` : moment defining function used in `discrete_approximation`

"""
function polynomial_moment(X::AbstractVector, mu::Real,
                           scaling_factor::Real, n_moments::Integer)
    # Check that scaling factor is positive
    scaling_factor>0 || error("scaling_factor must be a positive number")
    Y = (X.-mu)/scaling_factor # standardized grid
    T = Y'.^collect(1:n_moments)
end

"""

Compute the maximum entropy objective function used in `discrete_approximation`

```julia
obj = entropy_obj(lambda, Tx, TBar, q)
```

##### Arguments

- `lambda::AbstractVector` : length `L` vector of values of the dual problem variables
- `Tx::AbstractMatrix` : `L x N` matrix of moments evaluated at the grid points
                         specified in discrete_approximation
- `TBar::AbstractVector` : length `L` vector of moments of the underlying distribution
                           which should be matched
- `q::AbstractVector` : length `N` vector of prior weights for each point in the grid.

##### Returns

- `obj` : scalar value of objective function evaluated at `lambda`

"""
function entropy_obj(lambda::AbstractVector, Tx::AbstractMatrix,
                     TBar::AbstractVector, q::AbstractVector)
    # Compute objective function
    Tdiff = Tx .- TBar
    temp = q' .* exp.(lambda'*Tdiff)
    obj = sum(temp)
    return obj
end

"""
Compute gradient of objective function

##### Returns

- `grad` : length `L` gradient vector of the objective function evaluated at `lambda`

"""
function entropy_grad!(grad::AbstractVector, lambda::AbstractVector,
                       Tx::AbstractMatrix, TBar::AbstractVector, q::AbstractVector)
    Tdiff = Tx .- TBar
    temp = q'.*exp.(lambda'*Tdiff)
    temp2 = temp.*Tdiff
    grad .= vec(sum(temp2, dims = 2))
end

"""
Compute hessian of objective function

##### Returns

- `hess` : `L x L` hessian matrix of the objective function evaluated at `lambda`

"""
function entropy_hess!(hess::AbstractMatrix, lambda::AbstractVector,
                       Tx::AbstractMatrix, TBar::AbstractVector, q::AbstractVector)
    Tdiff = Tx .- TBar
    temp = q'.*exp.(lambda'*Tdiff)
    temp2 = temp.*Tdiff
    hess .= temp2*Tdiff'
end

@doc doc"""

find a unitary matrix `U` such that the diagonal components of `U'AU` is as
close to a multiple of identity matrix as possible

##### Arguments

- `A::AbstractMatrix` : square matrix

##### Returns

- `U` : unitary matrix
- `fval` : minimum value

"""
function min_var_trace(A::AbstractMatrix)

    ==(size(A)...) || throw(ArgumentError("input matrix must be square"))

    K = size(A, 1) # size of A
    d = tr(A)/K # diagonal of U'*A*U should be closest to d
    function obj(X, grad)
        X = reshape(X, K, K)
        return (norm(diag(X'*A*X) .- d))
    end
    function unitary_constraint(res, X, grad)
        X = reshape(X, K, K)
        res .= vec(X'*X - Matrix(I, K, K))
    end

    opt = NLopt.Opt(:LN_COBYLA, K^2)
    NLopt.min_objective!(opt, obj)
    NLopt.equality_constraint!(opt, unitary_constraint, zeros(K^2))
    fval, U_vec, ret = NLopt.optimize(opt, vec(Matrix(I, K, K)))

    return reshape(U_vec, K, K), fval
end


mutable struct Measure
    mean::Float64
    error::Float64
end

Measure() = Measure(0,-1)
Measure(a::Observable) = Measure(mean(a), error(a))
Measure(a::Measure) = Measure(a.mean, a.error)

function measure_binomial(nsucc::Integer, trials::Integer)
    p = nsucc /trials
    e = sqrt(p*(1-p) / trials)
    return Measure(p, e)
end

function (&)(a::Measure, b::Measure)
	if a.error <= 0
        a.mean = b.mean
        a.error = b.error
        return
    end
    if b.error <= 0
        return
    end

    w1 = 1 / a.error^2
    w2 = 1 / b.error^2
    m = (a.mean * w1 + b.mean * w2 ) / (w1+w2)
    e = 1 / sqrt(w1+w2)
    return Measure(m, e)
end

function +(a::Measure, b::Measure)
	m = a.mean + b.mean
    e = sqrt(a.error^2 + b.error^2)
    Measure(m,e)
end

function *(a::Measure, b::Measure)
	m = a.mean*b.mean
    e = sqrt(b.mean^2*a.error^2 + a.mean^2*b.error^2)
    Measure(m,e)
end

+(a::Measure, val::Number) = a + Measure(val,0)

function *(a::Measure, val::Number)
	m = a.mean * val
    e = a.error * val
    Measure(m,e)
end

function shortshow(io::IO, a::Measure)
    print(io,a.mean, " ", a.error)
end

Base.show(io::IO, a::Measure) = shortshow(io, a)


const MultivectorArray = MVector{2^N}
const COMBINATIONS = [[], ga.combinations(1:N)...]

KVector{K}(coefs::MVector) where {K} = KVector{K,typeof(coefs)}(coefs, linear_index(K))
KVector{K,V}(coefs::V) where {K,V<:AbstractVector} = KVector{K,V}(coefs, linear_index(K))

zero(::Type{<:KVector{K}}, T) where {K} = KVector{K}(zeros(MVector{binomial(N, K),T}))
zero(::Type{<:Multivector}, T) = Multivector(zeros(MultivectorArray))

fill(v, T::Type{<:KVector{K}}) where {K} = T(fill(v, MVector{binomial(N, K)}))
fill(v, ::Type{<:Multivector}) = Multivector(fill(v, MultivectorArray))

linear_index(grade::Integer) = linear_index(COMBINATIONS, grade)
linear_index(indices) = linear_index(COMBINATIONS, indices)

indices_from_linear_index(index::Integer) = indices_from_linear_index(COMBINATIONS, index)

grade_projection(x::Blade, ::Val{K}) where {K} = grade(x) == K ? x : K > N ? scalar(zero(eltype(x))) : scalar(zero(eltype(x)))

grade(b::Blade) = length(indices_from_linear_index(b.index))

function grade(mv::Multivector)
    higher_nonzero_index = findlast(≉(0), mv.coefs)
    isnothing(higher_nonzero_index) && return N
    findfirst(≤(higher_nonzero_index), linear_index.(0:N) .+ 1) - 1
end

dual(x) = x / I

(∨)(xs...) = ∧(dual.(xs)...) ⋅ I

reverse(x::Multivector) = Multivector(map(x -> reverse_sign(length(indices_from_linear_index(x[1]))) * x[2], enumerate(x.coefs)))

"""
Get `Blade` element with index `(i, j)`.
"""
getindex(mv::Multivector, i::CartesianIndex) = getindex(mv.coefs, linear_index(SVector(i.I)))

show(io::IO, @nospecialize b::Blade) = print(io, b.coef, "v$(join(map(subscript, indices_from_linear_index(b.index))))")

@generated op_result_type(::typeof(*), T1, T2) = Multivector{MultivectorArray{promote_type(eltype(T1), eltype(T2))}}

kvector(mv::Multivector, ::Val{0}) = scalar(first(mv.coefs))

function kvector(mv::Multivector, ::Val{K}) where {K}
    start = linear_index(K)
    KVector{K}(mv.coefs[start:start + binomial(N, K) - 1])
end

kvectors(mv::Multivector) = (kvector(mv, Val(k)) for k ∈ 0:N)