#  Achdou, Han, Lasry, Lions, Moll (2015)
#  Heterogeneous Agent Models in Continuous Time

##############################################################################
##
## Type
##
##############################################################################

type AiyagariProblem
    π::Float64               # productivity
    γ::Float64               # CRRA utility with parameter gamma
    α::Float64               # Production function F = K^α * L^(1-α)
    δ::Float64               # Capital depreciation
    ρ::Float64               # discount rate
    σ2::Float64
    θ::Float64
    zmean::Float64
    z::Vector{Float64}       # productivity vector
    a::Vector{Float64}       # wealth vector
    invΔ::Float64            # 1/δ in HznB algorithm
end

function AiyagariProblem(;
                         π::Float64 = 0.0,
                         γ::Float64 = 2.0,
                         α::Float64 = 0.35,
                         δ::Float64 = 0.1,
                         ρ::Float64 = 0.05,
                         σ2::Float64 = (0.10)^2,
                         θ::Float64 = 0.3,
                         zmean::Float64 = 1.0,
                         zn::Int = 20,
                         zmin::Float64 = 0.5,
                         zmax::Float64 = 1.5,
                         amin::Float64 = -1.0,
                         amax::Float64 = 50.0,
                         an::Int = 50,
                         invΔ::Float64 = 1e-3
                         )

    a = collect(linspace(amin,amax,an))
    z = collect(linspace(zmin,zmax,zn))
    AiyagariProblem(π, γ, α, δ, ρ, σ2, θ, zmean, z, a, invΔ)
end

##############################################################################
##
## Aiyagari Solution
##
##############################################################################
type AiyagariSolution{T}
    V::Vector{T}       # initial value function
    g::Vector{T}      # distribution
    K::T               # initial aggregate capital.
    r::T
    w::T
end

function AiyagariSolution{T}(ap::AiyagariProblem; K::T = 3.8)
    # initial value function
    π = ap.π ; α = ap.α ; δ = ap.δ ; z = ap.z ; a = ap.a ; γ = ap.γ ; ρ = ap.ρ
    an = length(a)
    zn = length(z)
    r = exp(π) * α * K^(α-1) - δ
    w = exp(π) * (1-α) * K^α
    V = Array(T, an*zn)
    ij = zero(Int)
    for zi in 1:zn, ai in 1:an
        ij += 1
        V[ij] = (w * z[zi] + r * a[ai])^(1-γ)/((1-γ)*ρ)
    end

    # storage arrays
    g = Array(T, an*zn)
    AiyagariSolution(V, g, K, r, w)
end

##############################################################################
##
## Aiyagari Storage
##
##############################################################################

type AiyagariArrays{T}
    C::Base.SparseMatrix.SparseMatrixCSC{T, Int}
    B::Base.SparseMatrix.SparseMatrixCSC{T, Int}
    A::Base.SparseMatrix.SparseMatrixCSC{T, Int}
    b::Vector{T}
    u::Vector{T}
end

function AiyagariArrays{T}(ap::AiyagariProblem, as::AiyagariSolution{T})
    π = ap.π ; α = ap.α ; δ = ap.δ ; z = ap.z ; a = ap.a ; γ = ap.γ ; ρ = ap.ρ; σ2 = ap.σ2 ; θ = ap.θ ; zmean = ap.zmean
    invdz = 1/(z[2] - z[1])
    invda = 1/(a[2] - a[1])
    an = length(a)
    zn = length(z)

    # create 3 sparse matrices of the right form filled with zeros
    A = spdiagm(
        (ones(T, an*zn), ones(T, an*zn-1), ones(T, an*zn-1), ones(T, (zn-1)*an), ones(T, (zn-1)*an)),
        (0, 1, -1, an, -an)
    )
    fill!(nonzeros(A), zero(T))
    B = deepcopy(A)
    C = deepcopy(A)

    # fill up matrix C
    Cvals = nonzeros(C)
    Crows = rowvals(C)
    ij = zero(Int)
    @inbounds for zi in 1:zn, ai in 1:an
        ij += 1
        krange = nzrange(C, ij)
        rows = Crows[krange]

        if zi > 1
            current =  0.5 * σ2 * invdz^2
            index = searchsortedfirst(rows, ij - an)
            Cvals[krange[index]] += current
            index = searchsortedfirst(rows, ij)
            Cvals[krange[index]] -= current
        end

        if zi < zn
            current = θ * (zmean - z[zi]) * invdz + 0.5 * σ2 * invdz^2
            index = searchsortedfirst(rows, ij + an)
            Cvals[krange[index]] += current
            index = searchsortedfirst(rows, ij)
            Cvals[krange[index]] -= current
        end
    end

    # b such that Ag = b in plank
    b = fill(zero(T), an*zn)
    b[1] = .1

    u = similar(b)

    AiyagariArrays(C, B, A, b, u)
end

##############################################################################
##
## Solve hamilton-jacobi-bellman
##
##############################################################################

    # update A and u
function  update_Au!(ap::AiyagariProblem, aa::AiyagariArrays, as::AiyagariSolution)
    γ = ap.γ  ;  a = ap.a ; z = ap.z ; invΔ = ap.invΔ ;  ρ = ap.ρ ;     z = ap.z ; a = ap.a ;
    C = aa.C  ; A = aa.A ; u = aa.u ; B = aa.B ;
    an = length(a)
    zn = length(z)
    invda = (an-1)/(a[end]-a[1])
    invγ = 1/γ
    inv1γ = 1.0/(1.0-γ)
    Cvals = nonzeros(C)
    Avals = nonzeros(A)
    copy!(Avals, Cvals)
    ij = zero(Int)
    @inbounds for zi in 1:zn, ai in 1:an
        ij += 1
        krange = nzrange(A, ij)
        rows = rowvals(A)[krange]
        if ai < an
            ∂V = (as.V[ij+1] - as.V[ij]) * invda
        else
            # state constraint boundary condition
            ∂V = (as.w * z[zi] + as.r * a[end])^(-γ)
        end
        saving = as.w * z[zi] + as.r * a[ai] - ∂V^(-invγ)
        if saving > 0 # case of positive drift
            if ai < an
                current = saving * invda
                index = searchsortedfirst(rows, ij)
                Avals[krange[index]] -= current
                index = searchsortedfirst(rows, ij + 1)
                Avals[krange[index]] += current
            end
        else
            if ai > 1
                ∂V = (as.V[ij] - as.V[ij-1]) * invda
            else
                # state constraint boundary condition
                ∂V =  (as.w * z[zi] + as.r * a[1])^(-γ)
            end
            saving = as.w * z[zi] + as.r * a[ai] - ∂V^(-invγ)
            if saving < 0 # case of negative drift
                if ai > 1
                    current = saving * invda
                    index = searchsortedfirst(rows, ij - 1)
                    Avals[krange[index]] -= current
                    index = searchsortedfirst(rows, ij)
                    Avals[krange[index]] += current
                end
            else
                ∂V = (as.w * z[zi] + as.r * a[ai])^(-γ)
            end
        end

        # update u
        u[ij] = inv1γ * ∂V^(1-invγ) + invΔ * as.V[ij]
    end
end

function update_B!(ap::AiyagariProblem, aa::AiyagariArrays, as::AiyagariSolution)
    γ = ap.γ  ;  a = ap.a ; z = ap.z ; invΔ = ap.invΔ ;  ρ = ap.ρ ;     z = ap.z ; a = ap.a ;
    C = aa.C  ; A = aa.A ; u = aa.u ; B = aa.B ;
    an = length(a)
    zn = length(z)
    Bvals = nonzeros(B)
    Brows = rowvals(B)
    Avals = nonzeros(A)
    ij = zero(Int)
    @inbounds for zi in 1:zn, ai in 1:an
        ij += 1
        for k in nzrange(B, ij)
            # loop over elements in the column ij
            row = Brows[k]
            Bvals[k] = - Avals[k]
            if row == ij
                Bvals[k] += invΔ + ρ
            end
        end
    end
end

function solve_hjb!(ap::AiyagariProblem,
                    aa::AiyagariArrays,
                    as::AiyagariSolution;
                    maxit::Int = 100,
                    crit::Float64 = 1e-6,
                    verbose::Bool = true)
    # update price vectors

    for iter in 1:maxit
        # update newV
        update_Au!(ap, aa, as)
        update_B!(ap, aa, as)
        newV = aa.B' \ aa.u
        # check convergence
        distance = chebyshev(newV, as.V)
        if distance < crit
            if verbose
                println("hjb solved : $(iter) iterations")
            end
            break
        else
            # update V using newV
            (newV, as.V) = (as.V, newV)
        end
    end
end

##############################################################################
##
## Solve fokker-planck equation
##
##############################################################################

# solve fokker-planck
function solve_fp!(ap::AiyagariProblem, aa::AiyagariArrays, as::AiyagariSolution)
    A = aa.A
    # change first row of matrix A
    Arows = rowvals(A)
    Avals = nonzeros(A)
    @inbounds for ij in 1:size(A, 2)
        for k in nzrange(A, ij)
            if Arows[k] == 1
                Avals[k] = zero(Float64)
            end
        end
    end
    Avals[1] = one(Float64)

    # solve system Ag = b
    as.g = A \ aa.b
    scale!(as.g, 1/sum(as.g))
end

##############################################################################
##
## Solve equilibrium by iterating on K
##
##############################################################################
function update_rw!(ap::AiyagariProblem, as::AiyagariSolution)
    as.r = exp(ap.π) * ap.α * as.K^(ap.α-1) - ap.δ
    as.w = exp(ap.π) * (1-ap.α) * as.K^ap.α
end

# compute agregate capital
function sum_capital{T}(a::Vector{Float64}, g::Vector{T})
    numerator = zero(Float64)
    denominator = zero(Float64)
    an = length(a)
    zn = div(length(g), length(a))
    ij = zero(Int)
    @inbounds for zi in 1:zn
        @simd for ai in 1:an
            ij += 1
            current = g[ij]
            numerator += current * a[ai]
            denominator += current
        end
    end
    return numerator / denominator
end

# solve
function solve!(ap::AiyagariProblem, aa::AiyagariArrays, as::AiyagariSolution;
                maxit::Int  = 100,      # maximum number of iterations in the HznB loop
                maxitK::Int = 100,      # maximum number of iterations in the K loop
                crit::Float64 = 1e-6,   # criterion HznB loop
                critK::Float64 = 1e-5,  # criterion K loop
                relax::Float64 = 0.99,  # relaxation parameter
                verbose = true)
    for iter in 1:maxitK
        if verbose
            println("Main loop iteration ",iter)
        end
        update_rw!(ap, as)

        # solve hamilton-jacobi-bellman equation
        solve_hjb!(ap, aa, as, maxit = maxit, crit = crit, verbose = verbose)

        # solve fokker-planck equation
        solve_fp!(ap, aa, as)

        # Update agregate capital
        newK = sum_capital(ap.a, as.g)
        if verbose
            @show newK
        end
        if abs(as.K - newK) < critK
            break
        else
            # relaxation algorithm (to ensure convergence)
            as.K = relax * as.K + (1 - relax) * newK
        end
    end
    return as
end

function solve(ap::AiyagariProblem;
    maxit::Int  = 100,      # maximum number of iterations in the HznB loop
    maxitK::Int = 100,      # maximum number of iterations in the K loop
    crit::Float64 = 1e-6,   # criterion HznB loop
    critK::Float64 = 1e-5,  # criterion K loop
    relax::Float64 = 0.99,  # relaxation parameter
    verbose = true)
    as = AiyagariSolution(ap)
    aa = AiyagariArrays(ap, as)
    solve!(ap, aa, as; maxit = maxit, maxitK = maxitK, crit = crit, critK = critK, relax = relax, verbose = verbose)
end

# as = AiyagariSolution(ap)
# HJBFiniteDifference.solve_hjb!(ap, as)


using WAVEFORM, LinearAlgebra, Test

# Test gmres
n = 100
A = randn(n,n)
A = A'*A
x = randn(100)
b = A*x

# Solve
x0 = similar(x) * 0.0
x0, rec = FGMRES(A, b, x0; tol=1e-9)

@test norm(x - x0) / norm(x0) < 1e-4

# Solve w/ preconditioner
Ai = A^-1   # perfect preconditioner
precond(x) = Ai*x

x0 = similar(x) * 0.0
x0, rec = FGMRES(A, b, x0; tol=1e-9, precond=precond)

@test norm(x - x0) / norm(x0) < 1e-4

# Test linear solver wrapper: gmres
opts = LinSolveOpts(; tol=1e-9, maxit=1000, solver=:fgmres, precond=:identity)
solver = solvesystem(A, opts)
x0 = randn(n)
x0 = solver(b, x0, true)

@test norm(x - x0) / norm(x0) < 1e-4

# Test linear solver wrapper: lufact
opts = LinSolveOpts(; tol=1e-9, maxit=1000, solver=:lufact, precond=:identity)
solver = solvesystem(A, opts)
x0 = randn(n)
x0 = solver(b, x0, true)

@test norm(x - x0) / norm(x0) < 1e-4


module MANOVABNPTest

# Imports
using DataFrames: DataFrame, filter!
using Distributions: Beta, Gamma, InverseWishart, MvNormal, pdf
using LinearAlgebra: cholesky, Cholesky, I, lowrankupdate!, lowrankdowndate!, logdet, Symmetric
using OffsetArrays: OffsetArray
using Parameters: @with_kw, @unpack
using StaticArrays: MMatrix, MVector, SVector, @MVector, @MMatrix
using StatsFuns: logmvgamma
using StatsBase: mean, cov, standardize, ZScoreTransform
using Random: randperm, randperm!, MersenneTwister

# Write your package code here.
include("01-utils.jl")
include("02-chain.jl")
include("03-model.jl")
include("04-suffstats.jl")
include("05-logl.jl")
include("06-gibbs.jl")

# Exports
export Model, SuffStats, suffstats0!, suffstats!, ChainState, log_pl, log_ml, ph0, fit, train, train_γ
export update_z!, update_γ!, update_α!
export γcode, γvector, γstr, γvec
export test_sample

end


using EffectiveWaves, Test

# large volume fraction scatterers,  small size amd on strong scatterer. This is a problamatic case.
@testset "large volume fraction and low frequency" begin
    # should pass for the list of angular frequencies below
    # ωs = [0.001, 2.0, 9.0]
    ω = 0.001

    medium = Medium(1.0,1.0+0.0im)
    species = [
        Specie(ρ=5.,r=0.004, c=1.2, volfrac=0.4),
        Specie(ρ=0.3, r=0.002, c=0.4, volfrac=0.3)
    ]

    eff_medium = effective_medium(medium, species)

    tol = 1e-7
    k_eff_low = ω/eff_medium.c
    k_effs = wavenumbers(ω, medium, species; tol = tol, num_wavenumbers=1)
    i = findmin(abs.(k_effs .- k_eff_low))[2]
    k_eff = k_effs[i]

    @test abs(k_eff - k_eff_low)/norm(k_eff_low) < 10*tol

    R = begin
        wave = EffectiveWave(ω, k_eff, medium, species)
        reflection_coefficient(ω, wave, medium, species)
    end
    R_low2 = begin
        wave = EffectiveWave(ω, k_eff_low, medium, species)
        reflection_coefficient(ω, wave, medium, species)
    end
    R_low = reflection_coefficient_halfspace(medium, eff_medium)

    @test norm(R_low - R) < tol
    @test norm(R_low2 - R) < tol
end


using FirthRegression, Base.Test, DataFrames

## Firth regression with admit data
df = readtable(Pkg.dir("GLM","data","admit.csv.gz"))
df[:rank] = pool(df[:rank])
gm = fit(FirthGLM, admit ~ gre + gpa + rank, df, Binomial(), convTol=1e-14)
@test_approx_eq_eps coef(gm) [-3.9076341278382402677, 0.0022173644223770295, 0.7875352644980999628, -0.6657797116943170446, -1.3177582553490327921, -1.5111669280800947845] 1e-8
@test_approx_eq penalized_deviance(gm.model) 431.06681398743257

## With sparse matrices
srand(1)
X = sprand(1000, 10, 0.01)
β = randn(10)
y = Bool[rand() < x for x in logistic(X * β)]

gmdense = fit(FirthGLM, full(X), y, Binomial())
gmsparse = fit(FirthGLM, X, y, Binomial())
@test_approx_eq deviance(gmsparse) deviance(gmdense)
@test_approx_eq coef(gmsparse) coef(gmdense)

## ANOVA
anv = anova(fit(FirthGLM, admit ~ gre + gpa + rank, df, Binomial(), convTol=1e-14).model,
            fit(FirthGLM, admit ~ gre + gpa, df, Binomial(), convTol=1e-14).model, convTol=1e-14)
@test_approx_eq anv.chisq 21.373030365595071
@test anv.df == 3
@test_approx_eq anv.p 8.8071001596912168e-05

## ANOVA w/o prefit models
anv = anova(fit(FirthGLM, admit ~ gre + gpa + rank, df, Binomial(), convTol=1e-14, dofit=false).model,
            fit(FirthGLM, admit ~ gre + gpa, df, Binomial(), convTol=1e-14, dofit=false).model, convTol=1e-14)
@test_approx_eq anv.chisq 21.373030365595071
@test anv.df == 3
@test_approx_eq anv.p 8.8071001596912168e-05

## ANOVA with sparse matrix
anvdense = anova(fit(FirthGLM, full(X[:, 1:5]), y, Binomial(), convTol=1e-14),
                 gmdense, convTol=1e-14)
anvsparse = anova(fit(FirthGLM, X[:, 1:5], y, Binomial(), convTol=1e-14),
                  gmsparse, convTol=1e-14)
@test_approx_eq anvdense.chisq anvsparse.chisq
@test anvdense.df == anvsparse.df
@test_approx_eq anvdense.p anvsparse.p


#Function to minimize Bregman divergences h() with respect to a given
#vector y over the submodular base polytope of a cardinality-based submodular function

# using Distributions

export card_inc_fix

function card_inc_fix(F::CardBasedAtom,
	                    w₀::AbstractArray,
										  divergence::String = "euclidean")
  # here F(S) = g(|S|) for some non-decreasing concave function g, g(0)=0.
  # currently implemented for divergences: Euclidean
  # check if g is concave, non-decreasing, g is an Array for now

  n = length(w₀)
  g = zeros(n)
  for i = 1:n
    g[i] = evaluate(F, collect(1:i))[1]
  end

	sorted_indices = sortperm(w₀, rev=true)
	w = w₀[sorted_indices]
	if(w[1] <= 0 && divergence != "euclidean")
   	error("The input is out of range for itakura-saito, logistic, entropy")
	elseif(w[1] >= 1&&divergence == "logistic")
  	error("The input is out of range for logistic")
	end

	counter = 0                             # maintains the tight set
	epsilon = zeros(n,1)

	if(divergence == "euclidean")
    x_opt = zeros(n,1)
  	# x_opt = round.(x_opt, 4)
  	feas_x_opt = (g[n]/n) * ones(n,1)
  	while(counter < n)
      for j = (counter + 1) : n
        if(counter == 0)
          epsilon[j] = (g[j] - sum(w[1:j]))/j
        else
          epsilon[j] = (g[j] - g[counter] - sum(w[counter+1:j]))/(j-counter)
        end
      end
      (e_k, I) = findmin(epsilon[counter+1:n])
      k = I + counter
    	x_opt[counter+1:k] = w[counter+1: k] + e_k
      counter = k
    	temp = counter+1
      while(temp<=n && x_opt[temp] - w[temp] <= e_k)
        x_opt[temp] = e_k+ w[temp]
        temp = temp+1
      end
    end
	end
	temp = zeros(n,1)
	for i =1:n
		temp[sorted_indices[i]] = x_opt[i]
	end
	return temp
end

card_inc_fix(p::AssocPoly, w::AbstractArray, divergence::String = "euclidean") = card_inc_fix(p.F, w, divergence)


vtype = Union{Rational{Int}, Int}
struct jp
    j::vtype
    p::Char
end

jp(v::Tuple{vtype, Char}) = jp(v[1],v[2])

import Base: length
length(jp1::jp) = 0
two_j(jp::jp) = Int(2jp.j)
#
# dealing with spin 1/2
x2(v) = @. Int(2v)
#
_over2(two) = mod(two,2)==0 ? div(two,2) : two // 2
over2(two) = _over2.(two)

⊗(p1::Char,p2::Char) = p1==p2 ? '+' : '-'
⊗(jp1::jp, jp2::jp) = [jp(j, ⊗(jp1.p, jp2.p)) for j in abs(jp1.j-jp2.j):abs(jp1.j+jp2.j)]

function str2jp(pin::String)
    p = filter(x->x!='^', pin)
    !(contains(p, '/')) && return jp(Meta.parse(p[1:end-1]), p[end])
    p[end-2:end-1] != "/2" && error("the string should be `x/2±`, while it is $(p)")
    two_j = Meta.parse(p[1:end-3])
    !(typeof(two_j) <: Int) && error("the string should be `x/2±`, while it is $(p)")
    return jp(two_j//2, p[end])
end

macro jp_str(p)
    return str2jp(p)
end

possible_ls((jp, (jp1, jp2))::Pair{jp, Tuple{jp, jp}}) = possible_ls(jp1, jp2; jp)

function possible_ls(jp1::jp, jp2::jp; jp::jp)
    ls = Vector{Tuple{Int,vtype}}(undef,0)
    for s in abs(jp1.j-jp2.j):abs(jp1.j+jp2.j)
        for l in Int(abs(jp.j-s)):Int(abs(jp.j+s))
            if jp1.p ⊗ jp2.p ⊗ jp.p  == (isodd(l) ? '-' : '+')
                push!(ls, (l,s))
            end
        end
    end
    return sort(ls, by=x->x[1])
end

function possible_lsLS(k::Int, jpR::jp, jps::T where T<:Vector{jp})
    i,j = ij_from_k(k)
    lsv = possible_ls(jps[i], jps[j]; jp = jpR)
    LSv = possible_ls(jpR, jps[k]; jp = jps[4])
    return [(ls=ls, LS=LS) for (ls, LS) in Iterators.product(lsv, LSv)]
end
function possible_lsLS(k::Int, two_s::Int, parity::Char, two_js , Ps)
    jpR = jp(two_s//2, parity)
    jps = jp.(zip(two_js .// 2, Ps))
    return possible_lsLS(k, jpR, jps)
end

jls_coupling(two_j1, two_λ1, two_j2, two_λ2, two_j, two_l, two_s) =
    sqrt((two_l+1) / (two_j+1)) *
    CG_doublearg(two_j1, two_λ1, two_j2, -two_λ2, two_s, two_λ1-two_λ2) *
    CG_doublearg(two_l,  0, two_s, two_λ1-two_λ2, two_j, two_λ1-two_λ2)


using CompScienceMeshes
using Test
using SauterSchwabQuadrature

pI = point(1,5,3)
pII = point(2,5,3)
pIII = point(7,1,0)

Sourcechart = Testchart = simplex(pI, pII, pIII)

Accuracy = 12
cf = CommonFace(SauterSchwabQuadrature._legendre(Accuracy,0.0,1.0))

function integrand(x,y)
			return(((x-pI)'*(y-pII))*exp(-im*1*norm(x-y))/(4pi*norm(x-y)))
end

function INTEGRAND(û,v̂)
	n1 = neighborhood(Testchart, û)
	n2 = neighborhood(Sourcechart, v̂)
	x = cartesian(n1)
	y = cartesian(n2)
	output = integrand(x,y)*jacobian(n1)*jacobian(n2)

return(output)
end

result = sauterschwab_parameterized(INTEGRAND, cf)-
		   verifintegral1(Sourcechart, Testchart, integrand, Accuracy)

@test norm(result) < 1.e-3

include(joinpath(dirname(@__FILE__),"numquad.jl"))

# Test the use of SauterSchwabQuadrature with the kernel generator utility
# using BEAST
using CompScienceMeshes
using StaticArrays

t1 = simplex(
    @SVector[0.180878, -0.941848, -0.283207],
    @SVector[0.0, -0.980785, -0.19509],
    @SVector[0.0, -0.92388, -0.382683])

rt = RTRefSpace{Float64}()
kernel_cf(x,y) = 1/norm(cartesian(x)-cartesian(y))
igd = generate_integrand_uv(kernel_cf, rt, rt, t1, t1)

t1 = simplex(
    @SVector[0.180878, -0.941848, -0.283207],
    @SVector[0.0, -0.980785, -0.19509],
    @SVector[0.0, -0.92388, -0.382683])

i5 = sauterschwab_parameterized(igd, CommonFace(SauterSchwabQuadrature._legendre(5,0.0,1.0)))
i10 = sauterschwab_parameterized(igd, CommonFace(SauterSchwabQuadrature._legendre(10,0.0,1.0)))

iref = numquad_cf(kernel_cf, rt, rt, t1, t1, zero(i5))

# # BEAST will arbitrate
# tqd = BEAST.quadpoints(rt, [t1], (12,))
# bqd = BEAST.quadpoints(rt, [t1], (13,))
#
# SE_strategy = BEAST.WiltonSEStrategy(
#   tqd[1,1],
#   BEAST.DoubleQuadStrategy(
# 	tqd[1,1],
# 	bqd[1,1],
#   ),
# )
#
# op = BEAST.MWSingleLayer3D(0.0, 4.0π, 0.0)
# z2 = zeros(3,3)
# BEAST.momintegrals!(op, rt, rt, t1, t1, z2, SE_strategy)
#
# @test i5  ≈ iref atol=1e-3
# @test i10 ≈ iref atol=1e-3
# @test i10 ≈ i5   atol=1e-6


using JuMP, MultilinearOpt

# using BARON
# m = Model(with_optimizer(Baron.Optimizer))

using Gurobi
m = Model(with_optimizer(Gurobi.Optimizer))

#
# Problem Topology

I = 1:5 # Input Feed Stocks
J = 1:4 # Products
L = 1:2 # Pools
K = 1:4 # Qualities

# i,l -> Bool
Tx(i,l) = (i,l) in ((1,1),(2,1),(3,1),(4,2),(5,2)) ? true : false
Ty(l,j) = true
Tz(i,j) = false

#
# Problem Parameters

# Product unit price
d = Dict(1=>16, 2=>25, 3=>15, 4=>10)

# Feed unit price
c = Dict(1=>7, 2=>3, 3=>2, 4=>10, 5=>5)

# Maximum available flow
AU(i) = Inf

# Minimum available flow
AL(i) = 0.0

# Pool size
SS(l) = Inf

# Maximum product demand
DU(j) = Dict(1=>10, 2=>25, 3=>30, 4=>10)[j]

# Minimum product demand
DL(j) = 0.0

# Feed concentration
CC = [1.0    6.0  4.0  0.5
      4.0    1.0  3.0  2.0
      4.0    5.5  3.0  0.9
      3.0    3.0  3.0  1.0
      1.0    2.7  4.0  1.6]

# Maximum Allowable Product Concentration
PU(j,k) = [3.00  3.00  3.25  0.75
           4.00  2.50  3.50  1.50
           1.50  5.50  3.90  0.80
           3.00  4.00  4.00  1.80][j,k]

# Minimum Allowable Product Concentration
PL(j,k) = 0.0

#
# Hard Bounds

@variable(m, 0 ≤ q[i in I, l in L] ≤ (Tx(i,l) ? 1 : Inf))
@variable(m, 0 ≤ y[l in L, j in J] ≤ min(SS(l), DU(j), sum(AU(i) for i in I if Tx(i,l))))
@variable(m, 0 ≤ z[i in I, j in J] ≤ min(AU(i), DU(j)))

#
# Standard Pooling Problem

@objective(m, Min, sum(c[i]*q[i,l]*y[l,j] for i in I, l in L, j in J if Tx(i,l) && Ty(l,j)) - sum(d[j]*y[l,j] for l in L, j in J if Ty(l,j)) - sum((d[j]-c[i])*z[i,j] for i in I, j in J if Tz(i,j)))

# Availability bounds
for i in I
    if AL(i) > 0
        @constraint(m, sum(q[i,l]*y[l,j] for j in J, l in L if Tx(i,l) && Ty(l,j)) + sum(z[i,j] for j in J if Tz(i,j)) ≥ AL(i))
    end
    if AU(i) < Inf
        @constraint(m, sum(q[i,l]*y[l,j] for j in J, l in L if Tx(i,l) && Ty(l,j)) + sum(z[i,j] for j in J if Tz(i,j)) ≤ AU(i))
    end
end

# Pool Capacity ========================
for l in L
    if SS(l) < Inf
        @constraint(m, sum(y[l,j] for j in J if Ty(l,j)) ≤ SS(l))
    end
end

# Product Demand =======================
for j in J
    if DL(j) > 0
        @constraint(m, sum(y[l,j] for l in L if Ty(l,j)) + sum(z[i,j] for i in I if Tz(i,j)) ≥ DL(j))
    end
    if DU(j) < Inf
        @constraint(m, sum(y[l,j] for l in L if Ty(l,j)) + sum(z[i,j] for i in I if Tz(i,j)) ≤ DU(j))
    end
end

# Product Quality ======================
for j in J, k in K
    if PL(j,k) > 0
        @constraint(m, PL(j,k)*(sum(y[l,j] for l in L if Ty(l,j)) + sum(z[i,j] for i in I if Tz(i,j))) ≤ sum(CC[i,k]*z[i,j] for i in I if Tz(i,j)) + sum(CC[i,k]*q[i,l]*y[l,j] for i in I, l in L if Tx(i,l) && Ty(l,j)))
    end
    if PU(j,k) < Inf
        @constraint(m, PU(j,k)*(sum(y[l,j] for l in L if Ty(l,j)) + sum(z[i,j] for i in I if Tz(i,j))) ≥ sum(CC[i,k]*z[i,j] for i in I if Tz(i,j)) + sum(CC[i,k]*q[i,l]*y[l,j] for i in I, l in L if Tx(i,l) && Ty(l,j)))
    end
end

# Simplex Constraint ===================
for l in L
    @constraint(m, sum(q[i,l] for i in I if Tx(i,l)) == 1)
end

# PQ Cut ===============================
for l in L, j in J
    @constraint(m, sum(q[i,l]*y[l,j] for i in I if Tx(i,l)) == y[l,j])
end

relaxbilinear!(m)
optimize!(m)


@testset "Persistence" begin

    # Intervals
    p = Interval(1.0 => 2.0)
    q = Interval(0.0 => 3.0)
    r = Interval(0.0 => 4.0)

    @test q < p
    @test q < r

    @test ComputationalHomology.birth(p) == -1.
    @test ComputationalHomology.death(p) == 3.
    @test diag(p)  == Interval(1.5 => 1.5)

    #######
    # from "Computational Topology - An Introduction" by Edelsbrunner & Harer, p. 184
    #######
    flt = Filtration(SimplicialComplex{Simplex{Int}}, Int)
    splxs=[ Simplex(1) => 1,
            Simplex(2) => 2,
            Simplex(3) => 3,
            Simplex(1,2) => 4,
            Simplex(2,3) => 5,
            Simplex(3,1) => 6,
            Simplex(1,2,3) => 7]
    for s in splxs
        push!(flt, s...)
    end
    @test size(complex(flt)) == (3,3,1)

    # Twist reduction
    ps, R = ComputationalHomology.pairs(TwistReduction, flt, reduced=true)
    @testset for (p, t) = zip(ps, [0=>1, 2=>4, 3=>5, 6=>7])
        @test p == t
    end

    # Standard reduction
    ps, R = ComputationalHomology.pairs(StandardReduction, flt)
    @testset for (p, t) = zip(ps, [2=>4, 3=>5, 6=>7, 1=>Inf])
        @test p == t
    end

    #######
    # from "Topology for Computing" by Zomorodian, pp.138-145
    #######
    flt = Filtration(SimplicialComplex{Simplex{Char}}, Int)
    splxs=[ Simplex('a') => 0,
            Simplex('b') => 0,
            Simplex('c') => 1,
            Simplex('d') => 1,
            Simplex('a', 'b') => 1,
            Simplex('b', 'c') => 1,
            Simplex('c', 'd') => 2,
            Simplex('a', 'd') => 2,
            Simplex('a', 'c') => 3,
            Simplex('a', 'b', 'c') => 4,
            Simplex('a', 'c', 'd') => 5]
    for s in splxs
        push!(flt, s...)
    end

    @testset "Intervals " for (d, itrs) in intervals(flt, length0=true)
        titrs = d == 0 ? intervals(0, 0=>Inf, 0=>1, 1=>1, 1=>2) : intervals(1, 3=>4, 2=>5)
        for itr in itrs
            @test itr ∈ titrs
        end
    end

    #######
    # Ex.1
    #######
    flt = Filtration(SimplicialComplex{Simplex{Int}}, Int)
    splxs=[ Simplex(1) => 0,
            Simplex(2) => 0,
            Simplex(3) => 0,
            Simplex(4) => 0,
            Simplex(5) => 1,
            Simplex(1,2) => 0,
            Simplex(2,3) => 0,
            Simplex(3,4) => 0,
            Simplex(4,1) => 0,
            Simplex(3,5) => 2,
            Simplex(4,5) => 3,
            Simplex(3,4,5) => 7]
    for s in splxs
        push!(flt, s...)
    end

    itr = intervals(flt, length0=true)
    @test Interval(0=>0) ∈ itr[0]
    @test Interval(1=>2) ∈ itr[0]
    @test Interval(1, 3=>7) ∈ itr[1]

    itr = intervals(flt, reduction = StandardReduction)
    @test Interval(1=>2) ∈ itr[0]
    @test Interval(1, 3=>7) ∈ itr[1]
    @test Interval(1, 0=>Inf) ∈ itr[1]

    #######
    # Ex.2
    #######
    flt = Filtration(SimplicialComplex{Simplex{Int}}, Float64)
    splxs=[ Simplex(1) => 1.0,
            Simplex(2) => 2.0,
            Simplex(3) => 3.0,
            Simplex(4) => 4.0,
            Simplex(5) => 5.0,
            Simplex(1,2) => 6.0,
            Simplex(2,3) => 7.0,
            Simplex(3,4) => 8.0,
            Simplex(4,1) => 9.0,
            Simplex(3,5) => 10.0,
            Simplex(4,5) => 11.0,
            Simplex(3,4,5) => 12.0]
    for s in splxs
        push!(flt, s...)
    end
    @test size(complex(flt)) == (5,6,1)
    @test length(flt) == 12

    # Betti numbers
    ∂ = boundary(flt)
    @test sparse(∂)[5,11] == 5
    @test length(∂[9]) > 0
    R = reduce(StandardReduction, ∂)
    @test ComputationalHomology.betti(flt, R, 0) == 1
    @test ComputationalHomology.betti(flt, R, 1) == 1
    @test ComputationalHomology.betti(flt, R, 2) == 0
    @test_throws AssertionError ComputationalHomology.betti(flt, R, 3)
    @test length(R[9]) == 0
    reduce!(StandardReduction, ∂)
    @test length(∂[9]) == 0

    # PH object
    ph = persistenthomology(flt)
    @test eltype(ph) == Tuple{Int64,Int64}
    @test length(ph) == 3
    @test group(ph, 0) == 1
    @test group(ph, 1) == 1
    @test group(ph, 2) == 0
    @test_throws AssertionError group(ph, 3)

    # PH iterator
    @testset "Method Comparison" for (g1, g2) in zip(homology(complex(flt), Int), persistenthomology(TwistReduction, flt))
        @test g1[1] == g2[1]
        @test g1[2] == g2[2]
    end
end


function Nn_series_coeff!(t::Transit_Struct{T}) where {T <: Real}
# Use series expansion to compute N_n.
# Computing leading coefficient (n=0).  Need to compute for
# n_max-3 to n_max.  Also, need just for k^2 < 1
# Nn_coeff is 2-dimensional array:
# Nn_coeff[2,jmax], where first dimension is for n_max-1 (1) to n_max (2),
# while third dimension is for the series coefficients.
# Computing leading coefficient (n=0):
coeff = zero(T)
# Loop over n_max-1 to n_max:
for j=1:2
  # m is n_max-1 to n_max (we need both since downward recursion requires two):
  m = t.n_max+j-2
  mhalf = 0.5*m
  coeff = 0.5*sqrt(pi)*exp(logabsgamma(mhalf+1.0)[1]-logabsgamma(mhalf+2.5)[1])
  # Add leading term to N_n:
  t.Nn_coeff[j,1] = coeff
  # Now, compute higher order terms until desired precision is reached:
  @inbounds for i=1:t.jmax-1
    coeff *= (4*i^2-1)/(2*i*(3+m+2*i))
    t.Nn_coeff[j,i+1] = coeff
  end
end
end


"""
VWN5 LDA correlation according to Vosko, Wilk, and Nusair, (DOI 10.1139/p80-159).
"""
function energy_per_particle(::Val{:lda_c_vwn}, ρ)
    # Ugly automatically generated code. TODO code from the equations
    x0  = sqrt(5913592351)
    x1  = 2^(2/3)*3^(1/6)
    x4  = (π*ρ)^(1/6)
    x5  = 468052052839577158149*x0/928058053502134909615000000
    x7  = 12500*x4*6^(1/3)
    x11 = 1 / (46593 * x1 * cbrt(π*ρ) + x7 + 323380 * sqrt(π*ρ))

    (  x5*atan(x0*x4/(12500*x1 + 46593*x4)) + 310907*log(x11*x7)/10000000
     + 76037485627899*log(x11*x4*(25000*x1 + 5249*x4)^2/100000)/78468213432500000)
end


using JAC: HalfInt
using JAC.AngularMomentum: clebschgordan, wigner3j, wigner6j, wigner9j

@testset "Wigner symbols" begin
    @testset "clebschgordan" begin
        @test iszero(clebschgordan(1, 0, 1, 0, 1/2, 0))
        @test iszero(clebschgordan(1, 0, 1, 0, 1, 1))
        @test clebschgordan(5/2, 3/2, 0, 0, 5//2, 3//2) ≈ 1
        @test clebschgordan(AngularJ64(5//2), AngularM64(3//2),
                            AngularJ64(0),    AngularM64(0),
                            AngularJ64(5//2), AngularM64(3//2)) ≈ 1
        @test clebschgordan(1/2,  1/2, 1/2,  1/2, 1,  1) ≈ 1
        @test clebschgordan(1/2,  1/2, 1/2, -1/2, 1,  0) ≈ √(1/2)
        @test clebschgordan(1/2, -1/2, 1/2,  1/2, 1,  0) ≈ √(1/2)
        @test clebschgordan(1/2, -1/2, 1/2, -1/2, 1, -1) ≈ 1
        @test clebschgordan(1/2,  1/2, 1/2, -1/2, 0,  0) ≈ √(1/2)
        @test clebschgordan(1/2, -1/2, 1/2,  1/2, 0,  0) ≈ -√(1/2)
        @test clebschgordan(3/2, -1/2, 2, 1, 7/2,  1/2) ≈ √(12/35)
        @test clebschgordan(3/2, -1/2, 2, 1, 5/2,  1/2) ≈ -√(5/14)
        @test clebschgordan(3/2, -1/2, 2, 1, 3/2,  1/2) ≈ 0
        @test clebschgordan(3/2, -1/2, 2, 1, 1/2,  1/2) ≈ √(3/10)
        @test clebschgordan(2,  2, 3/2, -3/2, 7/2,  1/2) ≈ √(1/35)
        @test clebschgordan(2,  2, 3/2, -3/2, 5/2,  1/2) ≈ √(6/35)
        @test clebschgordan(2,  2, 3/2, -3/2, 3/2,  1/2) ≈ √(2/5)
        @test clebschgordan(2,  2, 3/2, -3/2, 1/2,  1/2) ≈ √(2/5)
        @test clebschgordan(2,  1, 3/2, -1/2, 7/2,  1/2) ≈ √(12/35)
        @test clebschgordan(2,  1, 3/2, -1/2, 5/2,  1/2) ≈ √(5/14)
        @test clebschgordan(2,  1, 3/2, -1/2, 3/2,  1/2) ≈ 0
        @test clebschgordan(2,  1, 3/2, -1/2, 1/2,  1/2) ≈ -√(3/10)
        @test clebschgordan(2,  0, 3/2,  1/2, 7/2,  1/2) ≈ √(18/35)
        @test clebschgordan(2,  0, 3/2,  1/2, 5/2,  1/2) ≈ -√(3/35)
        @test clebschgordan(2,  0, 3/2,  1/2, 3/2,  1/2) ≈ -√(1/5)
        @test clebschgordan(2,  0, 3/2,  1/2, 1/2,  1/2) ≈ √(1/5)
        @test clebschgordan(2, -1, 3/2,  3/2, 7/2,  1/2) ≈ √(4/35)
        @test clebschgordan(2, -1, 3/2,  3/2, 5/2,  1/2) ≈ -√(27/70)
        @test clebschgordan(2, -1, 3/2,  3/2, 3/2,  1/2) ≈ √(2/5)
        @test clebschgordan(2, -1, 3/2,  3/2, 1/2,  1/2) ≈ -√(1/10)
        @test clebschgordan(5/2,  3/2, 5/2, -3/2, 5,  0) ≈ √(25/252)
        @test clebschgordan(5/2,  3/2, 5/2, -3/2, 4,  0) ≈ √(9/28)
        @test clebschgordan(5/2,  3/2, 5/2, -3/2, 3,  0) ≈ √(49/180)
        @test clebschgordan(5/2,  3/2, 5/2, -3/2, 2,  0) ≈ √(1/84)
        @test clebschgordan(5/2,  3/2, 5/2, -3/2, 1,  0) ≈ -√(9/70)
        @test clebschgordan(5/2,  3/2, 5/2, -3/2, 0,  0) ≈ -√(1/6)
        @test clebschgordan(5/2, -3/2, 5/2,  3/2, 5,  0) ≈ √(25/252)
        @test clebschgordan(5/2, -3/2, 5/2,  3/2, 4,  0) ≈ -√(9/28)
        @test clebschgordan(5/2, -3/2, 5/2,  3/2, 3,  0) ≈ √(49/180)
        @test clebschgordan(5/2, -3/2, 5/2,  3/2, 2,  0) ≈ -√(1/84)
        @test clebschgordan(5/2, -3/2, 5/2,  3/2, 1,  0) ≈ -√(9/70)
        @test clebschgordan(5/2, -3/2, 5/2,  3/2, 0,  0) ≈ √(1/6)
        # Negative numbers as arguments for J
        @info "The following 20 warnings (\"GSL Error 1\") are expected (testing error handling)"
        @test_throws DomainError clebschgordan(-5/2, 3/2, 0, 0, 5/2, 3/2)
        @test_throws DomainError clebschgordan(0, 0, -5/2, 3/2, 5/2, 3/2)
        # In the following case, GSL.jl does not print a warning since sqrt(-4.0) throws a
        # DomainError before wigner3j is called
        @test_throws DomainError clebschgordan(0, 0, 5/2, 3/2, -5/2, 3/2)
        # Arguments that are not half-integers
        @test_throws InexactError clebschgordan(5/3, 3/2, 0, 0, 5/2, 3/2)
        @test_throws InexactError clebschgordan(5/2, 4/3, 0, 0, 5/2, 3/2)
        @test_throws InexactError clebschgordan(0, 0, 5/3, 3/2, 5/2, 3/2)
        @test_throws InexactError clebschgordan(0, 0, 5/2, 4/3, 5/2, 3/2)
        @test_throws InexactError clebschgordan(5/3, 3/2, 0, 0, 5/3, 3/2)
        @test_throws InexactError clebschgordan(5/2, 4/3, 0, 0, 5/2, 4/3)
    end

    @testset "wigner3j" begin
        @test wigner3j(AngularJ64(2), AngularJ64(6), AngularJ64(4),
                       AngularM64(0), AngularM64(0), AngularM64(0)) ≈ sqrt(5/143)
        @test wigner3j(AngularJ64(2), AngularJ64(6), AngularJ64(4),
                       AngularM64(0), AngularM64(0), AngularM64(1)) == 0
        @test wigner3j(AngularJ64(2), AngularJ64(7), AngularJ64(4),
                       AngularM64(0), AngularM64(0), AngularM64(0)) == 0
        @test wigner3j(AngularJ64(3//2), AngularJ64(1), AngularJ64(5//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ -sqrt(1/10)
        @test wigner3j(AngularJ64(5//2), AngularJ64(1), AngularJ64(5//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ sqrt(1/210)
        @test wigner3j(AngularJ64(3//2), AngularJ64(0), AngularJ64(3//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ 1/2
        @test wigner3j(AngularJ64(3//2), AngularJ64(2), AngularJ64(3//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ -sqrt(1/20)
        @test wigner3j(AngularJ64(7//2), AngularJ64(2), AngularJ64(5//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ -sqrt(1/210)
        @test wigner3j(AngularJ64(7//2), AngularJ64(4), AngularJ64(5//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ sqrt(5/462)
        @test wigner3j(AngularJ64(7//2), AngularJ64(6), AngularJ64(5//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ -5*sqrt(1/858)
        @test wigner3j(AngularJ64(7//2), AngularJ64(0), AngularJ64(7//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ 1/2*sqrt(1/2)
        @test wigner3j(AngularJ64(7//2), AngularJ64(2), AngularJ64(7//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ -1/2*sqrt(5/42)
        @test wigner3j(AngularJ64(7//2), AngularJ64(4), AngularJ64(7//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ 3/2*sqrt(1/154)
        @test wigner3j(AngularJ64(7//2), AngularJ64(6), AngularJ64(7//2),
                       AngularM64(1//2), AngularM64(0), AngularM64(-1//2)) ≈ -5/2*sqrt(1/858)
        @test wigner3j(2, 6, 4, 0, 0, 0) ≈ sqrt(5/143)
        @test wigner3j(HalfInt(3/2), HalfInt(1), HalfInt(5/2), 1//2, 0, -1/2) ≈ -sqrt(1/10)
        @test wigner3j(3/2, 1, 5/2, 1/2, 0, -1/2) ≈ -sqrt(1/10)
        # Negative numbers as arguments for J
        @test_throws DomainError wigner3j(-3/2, 1, 5/2, 1//2, 0, -1/2)
        @test_throws DomainError wigner3j(3/2, -1, 5/2, 1//2, 0, -1/2)
        @test_throws DomainError wigner3j(3/2, 1, -5/2, 1//2, 0, -1/2)
        # Arguments that are not half-integers
        @test_throws InexactError wigner3j(4/3, 1, 5/2, 1//2, 0, -1/2)
        @test_throws InexactError wigner3j(3/2, 1.1, 5/2, 1//2, 0, -1/2)
        @test_throws InexactError wigner3j(3/2, 1, 5/3, 1//2, 0, -1/2)
        @test_throws InexactError wigner3j(3/2, 1, 5/2, 1//3, 0, -1/2)
        @test_throws InexactError wigner3j(3/2, 1, 5/2, 1//2, 0.1, -1/2)
        @test_throws InexactError wigner3j(3/2, 1, 5/2, 1//2, 0, -1/5)
    end

    @testset "wigner6j" begin
        @test wigner6j(AngularJ64(3), AngularJ64(3), AngularJ64(3),
                       AngularJ64(3), AngularJ64(3), AngularJ64(3)) ≈ -1/14
        @test wigner6j(AngularJ64(5), AngularJ64(5), AngularJ64(5),
                       AngularJ64(5), AngularJ64(5), AngularJ64(5)) ≈ 1/52
        @test wigner6j(AngularJ64(1), AngularJ64(3//2), AngularJ64(5//2),
                       AngularJ64(0), AngularJ64(5//2), AngularJ64(3//2)) ≈ -sqrt(1/24)
        @test wigner6j(1, 3/2, 5/2, 0, 5/2, 3/2) ≈ -sqrt(1/24)
        @test wigner6j(HalfInt(1), HalfInt(3/2), HalfInt(5/2), 0, 5//2, 3/2) ≈ -sqrt(1/24)
        # Negative numbers as arguments
        @test_throws DomainError wigner6j(-1, 3/2, 5/2, 0, 5//2, 3/2)
        @test_throws DomainError wigner6j(1, -3/2, 5/2, 0, 5//2, 3/2)
        @test_throws DomainError wigner6j(1, 3/2, -5/2, 0, 5//2, 3/2)
        @test_throws DomainError wigner6j(3/2, 5/2, 1, -5//2, 3/2, 0)
        @test_throws DomainError wigner6j(1, 3/2, 5/2, 0, -5//2, 3/2)
        @test_throws DomainError wigner6j(1, 3/2, 5/2, 0, 5//2, -3/2)
        # Arguments that are not half-integers
        @test_throws InexactError wigner6j(1.1, 3/2, 5/2, 0, 5//2, 3/2)
        @test_throws InexactError wigner6j(1, 4/3, 5/2, 0, 5//2, 3/2)
        @test_throws InexactError wigner6j(1, 3/2, 5/3, 0, 5//2, 3/2)
        @test_throws InexactError wigner6j(1, 3/2, 5/2, 0.1, 5//2, 3/2)
        @test_throws InexactError wigner6j(1, 3/2, 5/2, 0, 5//3, 3/2)
        @test_throws InexactError wigner6j(1, 3/2, 5/2, 0, 5//2, 4/3)
    end

    @testset "wigner9j" begin
        @test wigner9j(AngularJ64(1), AngularJ64(1), AngularJ64(1),
                       AngularJ64(1), AngularJ64(1), AngularJ64(1),
                       AngularJ64(1), AngularJ64(1), AngularJ64(0)) ≈ 1/18
        @test wigner9j(AngularJ64(1),    AngularJ64(3//2), AngularJ64(5//2),
                       AngularJ64(5//2), AngularJ64(0),    AngularJ64(5//2),
                       AngularJ64(3//2), AngularJ64(3//2), AngularJ64(0)) ≈ -1/24
        @test wigner9j(AngularJ64(1//2), AngularJ64(1//2), AngularJ64(1),
                       AngularJ64(3),    AngularJ64(3),    AngularJ64(2),
                       AngularJ64(5//2), AngularJ64(5//2), AngularJ64(1)) ≈ (1/21)*sqrt(2/5)
        @test wigner9j(1, 1, 1, 1, 1, 1, 1, 1, 0) ≈ 1/18
        @test wigner9j(1, AngularJ64(1), HalfInt(1), 1//1, 1.0, 1, 1, 1, 0) ≈ 1/18
        # Negative numbers as arguments
        @test_throws DomainError wigner9j(-1, 1, HalfInt(1), 1//1, 1.0, 1, 1, 1, 0)
        @test_throws DomainError wigner9j(1, -1, HalfInt(1), 1//1, 1.0, 1, 1, 1, 0)
        @test_throws DomainError wigner9j(1, 1, HalfInt(-1), 1//1, 1.0, 1, 1, 1, 0)
        @test_throws DomainError wigner9j(1, 1, HalfInt(1), -1//1, 1.0, 1, 1, 1, 0)
        @test_throws DomainError wigner9j(1, 1, HalfInt(1), 1//1, -1.0, 1, 1, 1, 0)
        @test_throws DomainError wigner9j(1, 1, HalfInt(1), 1//1, 1.0, -1, 1, 1, 0)
        @test_throws DomainError wigner9j(1, 1, HalfInt(1), 1//1, 1.0, 1, -1, 1, 0)
        @test_throws DomainError wigner9j(1, 1, HalfInt(1), 1//1, 1.0, 1, 1, -1, 0)
        @test_throws DomainError wigner9j(1, 1, HalfInt(1), 1//1, 1.0, 1, 1, 1, -1)
        # Arguments that are not half-integers
        @test_throws InexactError wigner9j(1.1, 1, HalfInt(1), 1//1, 1.0, 1, 1, 1, 0)
        @test_throws InexactError wigner9j(1, 1.1, HalfInt(1), 1//1, 1.0, 1, 1, 1, 0)
        @test_throws InexactError wigner9j(1, 1, 1.1, 1//1, 1.0, 1, 1, 1, 0)
        @test_throws InexactError wigner9j(1, 1, HalfInt(1), 1//3, 1.0, 1, 1, 1, 0)
        @test_throws InexactError wigner9j(1, 1, HalfInt(1), 1//1, 1.6, 1, 1, 1, 0)
        @test_throws InexactError wigner9j(1, 1, HalfInt(1), 1//1, 1.0, 1.1, 1, 1, 0)
        @test_throws InexactError wigner9j(1, 1, HalfInt(1), 1//1, 1.0, 1, 1.2, 1, 0)
        @test_throws InexactError wigner9j(1, 1, HalfInt(1), 1//1, 1.0, 1, 1, 1.1, 0)
        @test_throws InexactError wigner9j(1, 1, HalfInt(1), 1//1, 1.0, 1, 1, 1.3, 0.01)
    end
end


# NB: ϕ is scaled by √2π compared to standard definition
ϕ(x) = exp(-x^2/2)
ϕ(x,D) = exp((D-x)*(D+x)/2) # rescaling that can be used for improved numerical precision

# ϕ2(x) = (x*x-1)*ϕ(x) # For reference: used to find intervals where ϕ is convex or concave.

ϕ4(x) = (x2=x*x; (x2*x2 - 6x2 + 3) * ϕ(x))
# ϕ5(x) = (x2=x*x; x4=x2*x2; -x*(x4 - 10x2 + 15) * ϕ(x)) # For reference: used to find local extrema of ϕ4
ϕ6(x) = (x2=x*x; x4=x2*x2; (x4*x2 - 15x4 + 45x2 - 15) * ϕ(x))
# ϕ7(x) = (x2=x*x; x4=x2*x2; x6=x4*x2; x*(x6 - 21x4 + 105x2 - 105)*ϕ(x))                 # For reference: used to find local extrema of ϕ6
# ϕ8(x) = (x2=x*x; x4=x2*x2; x6=x4*x2; x8=x4*x4; (x8 - 28x6 + 210x4 - 420x2 + 105)*ϕ(x)) # For reference: used to find intervals where ϕ6 is convex or concave.

ϕ(::Val{4},x) = ϕ4(x)
ϕ(::Val{6},x) = ϕ6(x)

"""
	ϕbounds(a,b,x,D)

Let `x` be the mean of a set of points {xᵢ} such that xᵢ∈[`a`,`b`] ∀i.
ϕbounds computes lower and upper bounds for ∑ᵢϕ(xᵢ)/n, where n is the number of points.
`D` is a rescaling that can be used for better numerical precision. ϕ(x) := 1/√(2π) * exp(D^2/2 - x^2/2)
"""
function ϕbounds(a,b,x,D)::Tuple{Float64,Float64}
	if b-a < 1e-9 # to avoid div by zero
		(abs(x-a)<1e-9 && abs(b-x)<1e-9) || throw(ArgumentError("x=$x must be in interval [a=$a, b=$b]"))
		y = ϕ((a+b)/2,D)
		return y,y
	end
	0<=a<=x<=b || throw(ArgumentError("x=$x must be in interval [a=$a, b=$b]"))

	ya,yb = ϕ(a,D),ϕ(b,D)
	if a<1 && b>1 # not in an interval where the function is convex/concave.
		# TODO: improve fallback
		return yb,ya
	end

	# bound by line above (convex) or below (concave) ϕ in the interval
	y2 = (yb-ya)*((x-a)/(b-a)) + ya

	# lower bound (convex) or upper bound (concave) using Jensen's inqueality.
	y1 = ϕ(x,D)

	if b<=1 # concave
		y2, y1
	else # convex
		y1, y2
	end
end

"""
	ϕ4bounds(a,b,x)

Let `x` be the mean of a set of points {xᵢ} such that xᵢ∈[`a`,`b`] ∀i.
ϕ4bounds computes lower and upper bounds for ∑ᵢϕ⁽⁴⁾(xᵢ)/n, where n is the number of points.
"""
function ϕ4bounds(a,b,x)::Tuple{Float64,Float64}
	breakpoints = (0.6167065901925941, 1.8891758777537109, 3.3242574335521193)
	if b-a < 1e-9 # to avoid div by zero
		(abs(x-a)<1e-9 && abs(b-x)<1e-9) || throw(ArgumentError("x=$x must be in interval [a=$a, b=$b]"))
		y = ϕ4((a+b)/2)
		return y,y
	end
	0<=a<=x<=b || throw(ArgumentError("x=$x must be in interval [a=$a, b=$b]"))

	ai = (a>breakpoints[1]) + (a>breakpoints[2]) + (a>breakpoints[3])
	bi = (b>breakpoints[1]) + (b>breakpoints[2]) + (b>breakpoints[3])

	if ai!=bi # not in an interval where the function is convex/concave.
		# TODO: improve fallback.
		return ϕ4extrema(a,b)
	end

	ya,yb = ϕ4(a),ϕ4(b)

	# bound by line above (convex) or below (concave) ϕ4 in the interval
	y2 = (yb-ya)*((x-a)/(b-a)) + ya

	# lower bound (convex) or upper bound (concave) using Jensen's inqueality.
	y1 = ϕ4(x)

	if ai&1 == 0 # concave
		y2, y1
	else # convex
		y1, y2
	end
end

"""
	ϕ6bounds(a,b,x)

Let `x` be the mean of a set of points {xᵢ} such that xᵢ∈[`a`,`b`] ∀i.
ϕ6bounds computes lower and upper bounds for ∑ᵢϕ⁽⁶⁾(xᵢ)/n.
"""
function ϕ6bounds(a,b,x)::Tuple{Float64,Float64}
	breakpoints = (0.5390798113513751, 1.636519042435108, 2.8024858612875416, 4.144547186125894)
	if b-a < 1e-9 # to avoid div by zero
		(abs(x-a)<1e-9 && abs(b-x)<1e-9) || throw(ArgumentError("x=$x must be in interval [a=$a, b=$b]"))
		y = ϕ6((a+b)/2)
		return y,y
	end
	0<=a<=x<=b || throw(ArgumentError("x=$x must be in interval [a=$a, b=$b]"))

	ai = (a>breakpoints[1]) + (a>breakpoints[2]) + (a>breakpoints[3]) + (a>breakpoints[4])
	bi = (b>breakpoints[1]) + (b>breakpoints[2]) + (b>breakpoints[3]) + (b>breakpoints[4])

	if ai!=bi # not in an interval where the function is convex/concave.
		# TODO: improve fallback.
		return ϕ6extrema(a,b)
	end

	ya,yb = ϕ6(a),ϕ6(b)

	# bound by line above (convex) or below (concave) ϕ6 in the interval
	y2 = (yb-ya)*((x-a)/(b-a)) + ya

	# lower bound (convex) or upper bound (concave) using Jensen's inqueality.
	y1 = ϕ6(x)

	if ai&1 == 1 # concave
		y2, y1
	else # convex
		y1, y2
	end
end

ϕbounds(::Val{4},a,b,x) = ϕ4bounds(a,b,x)
ϕbounds(::Val{6},a,b,x) = ϕ6bounds(a,b,x)

function ϕ4extrema(a,b)
	# x at extrema: -√(5+√10), √(5-√10), 0, √(5-√10), √(5+√10)
	# we only need to consider x>0 since the function is even
	xExt = (1.355626179974266, 2.8569700138728056)
	yExt = (-1.8548702941445978, 0.34872676216452453)

	# the implementation assumes 0<=a<=b
	if b<=xExt[1] || a>=xExt[2]
		ϕ4(b),ϕ4(a) # decreasing function in these intervals
	elseif a>=xExt[1]
		if b<=xExt[2]
			ϕ4(a),ϕ4(b) # increasing function in this interval
		else
			min(ϕ4(a),ϕ4(b)),yExt[2] # the function covers a local maximum
		end
	elseif b<=xExt[2]
		yExt[1],max(ϕ4(a),ϕ4(b)) # the function covers the global minimum
	else
		yExt[1], max(ϕ4(a),yExt[2]) # the function covers the global minimum and a local maximum
	end
end

function ϕ6extrema(a,b)
	# ϕ6 has global minimum at 0, global maximum at x≈1.15, two more local extrema and ϕ6(x)→0 as x→∞.
	# we only need to consider x>0 since the function is even
	xExt = (1.154405394739968, 2.366759410734541, 3.750439717725742)
	yExt = (10.62963914930061, -3.513849978144481, 0.3821828509277058)
	# the implementation assumes 0<=a<=b
	if xExt[3]<=a || (xExt[1]<=a && b<=xExt[2])
		ϕ6(b),ϕ6(a) # decreasing function in these intervals
	elseif b<=xExt[1] || (xExt[2]<=a && b<=xExt[3])
		ϕ6(a),ϕ6(b) # increasing function in these intervals
	elseif b<=xExt[2]
		min(ϕ6(a),ϕ6(b)),yExt[1] # covers global maximum
	elseif xExt[1]<=a && b<=xExt[3]
		yExt[2],max(ϕ6(a),ϕ6(b)) # covers local minimum
	elseif xExt[2]<=a
		min(ϕ6(a),ϕ6(b)),yExt[3] # covers second local maximum
	elseif xExt[1]<=a
		yExt[2],max(ϕ6(a),yExt[3])  # covers the local minimum and the second local maximum
	# elseif b<=xExt[3] # same as last case
	# 	min(ϕ6(a),yExt[2]),yExt[1] # covers global maximum and local minimum
	else
		min(ϕ6(a),yExt[2]),yExt[1] # covers global maximum and local minimum
	end
end

ϕextrema(::Val{4},a,b) = ϕ4extrema(a,b)
ϕextrema(::Val{6},a,b) = ϕ6extrema(a,b)


module SSLTestGenerated
using SimpleSDMLayers
using Test

M = rand(Float64, (5, 10))
S = SimpleSDMPredictor(M, 0.0, 1.0, 0.0, 1.0)

@test maximum(S) ≤ 1.0
@test minimum(S) ≥ 0.0
@test extrema(S) == (minimum(S), maximum(S))

using Statistics

@test mean(S) ≈ mean(M)
@test std(S) ≈ std(M)
@test median(S) ≈ median(M)

M = ones(Float64, (5, 10))
S = SimpleSDMPredictor(M, 0.0, 1.0, 0.0, 1.0)
@test typeof(sqrt(S)) <: SimpleSDMLayer

S1 = SimpleSDMResponse([-1 0 1 nothing], 0.0, 1.0, 0.0, 1.0)
S2 = SimpleSDMPredictor([Float32.([-2.0 1.0 1.0])... nothing], 0.0, 1.0, 0.0, 1.0)
Sx = [S1, S2]
@test reduce(min, Sx).grid == [-2 0 1 nothing]
@test mean(Sx).grid == [Float32.([-1.5 0.5 1.0])... nothing]
@test reduce(+, Sx).grid == [-3 1 2 nothing]

end


module ManifoldsBase

import Base:
    isapprox,
    exp,
    log,
    convert,
    copyto!,
    angle,
    eltype,
    isempty,
    length,
    similar,
    show,
    +,
    -,
    *
import LinearAlgebra: dot, norm, det, cross, I, UniformScaling, Diagonal

import Markdown: @doc_str
using LinearAlgebra

"""
    Manifold{F}

A manifold type. The `Manifold` is used to dispatch to different functions on a manifold,
usually as the first argument of the function. Examples are the [`exp`](@ref)onential and
[`log`](@ref)arithmic maps as well as more general functions that are built on them like the
[`geodesic`](@ref).

The manifold is parametrized by an [`AbstractNumbers`](@ref) to distinguish for example
real (ℝ) and complex (ℂ) manifolds.

For subtypes the preferred order of parameters is: size and simple value parameters,
followed by the [`AbstractNumbers`](@ref) `field`, followed by data type parameters,
which might depend on the abstract number field type.
"""
abstract type Manifold{𝔽} end

"""
    AbstractEstimationMethod

Abstract type for defining statistical estimation methods.
"""
abstract type AbstractEstimationMethod end

"""
    AbstractInverseRetractionMethod

Abstract type for methods for inverting a retraction (see [`inverse_retract`](@ref)).
"""
abstract type AbstractInverseRetractionMethod end

"""
    AbstractRetractionMethod

Abstract type for methods for [`retract`](@ref)ing a tangent vector to a manifold.
"""
abstract type AbstractRetractionMethod end

"""
    ExponentialRetraction

Retraction using the exponential map.
"""
struct ExponentialRetraction <: AbstractRetractionMethod end

"""
    PolarRetraction <: AbstractRetractionMethod

Retractions that are based on singular value decompositions of the matrix / matrices
for point and tangent vector on a [`Manifold`](@ref)
"""
struct PolarRetraction <: AbstractRetractionMethod end

"""
    ProjectionRetraction <: AbstractRetractionMethod

Retractions that are based on projection and usually addition in the embedding.
"""
struct ProjectionRetraction <: AbstractRetractionMethod end

"""
    QRRetraction <: AbstractRetractionMethod

Retractions that are based on a QR decomposition of the
matrix / matrices for point and tangent vector on a [`Manifold`](@ref)
"""
struct QRRetraction <: AbstractRetractionMethod end

"""
    LogarithmicInverseRetraction

Inverse retraction using the [`log`](@ref)arithmic map.
"""
struct LogarithmicInverseRetraction <: AbstractInverseRetractionMethod end

"""
    PolarInverseRetraction <: AbstractInverseRetractionMethod

Inverse retractions that are based on a singular value decomposition of the
matrix / matrices for point and tangent vector on a [`Manifold`](@ref)
"""
struct PolarInverseRetraction <: AbstractInverseRetractionMethod end

"""
    ProjectionInverseRetraction <: AbstractInverseRetractionMethod

Inverse retractions that are based on a projection (or its inversion).
"""
struct ProjectionInverseRetraction <: AbstractInverseRetractionMethod end

"""
    QRInverseRetraction <: AbstractInverseRetractionMethod

Inverse retractions that are based on a QR decomposition of the
matrix / matrices for point and tangent vector on a [`Manifold`](@ref)
"""
struct QRInverseRetraction <: AbstractInverseRetractionMethod end

"""
    MPoint

Type for a point on a manifold. While a [`Manifold`](@ref) does not necessarily require this
type, for example when it is implemented for `Vector`s or `Matrix` type elements, this type
can be used for more complicated representations, semantic verification, or even dispatch
for different representations of points on a manifold.
"""
abstract type MPoint end

"""
    OutOfInjectivityRadiusError

An error thrown when a function (for example [`log`](@ref)arithmic map or
[`inverse_retract`](@ref)) is given arguments outside of its [`injectivity_radius`](@ref).
"""
struct OutOfInjectivityRadiusError <: Exception end

"""
    TVector

Type for a tangent vector of a manifold. While a [`Manifold`](@ref) does not necessarily
require this type, for example when it is implemented for `Vector`s or `Matrix` type
elements, this type can be used for more complicated representations, semantic verification,
or even dispatch for different representations of tangent vectors and their types on a
manifold.
"""
abstract type TVector end

"""
    CoTVector

Type for a cotangent vector of a manifold. While a [`Manifold`](@ref) does not necessarily
require this type, for example when it is implemented for `Vector`s or `Matrix` type
elements, this type can be used for more complicated representations, semantic verification,
or even dispatch for different representations of cotangent vectors and their types on a
manifold.
"""
abstract type CoTVector end

"""
    allocate(a)
    allocate(a, dims::Integer...)
    allocate(a, dims::Tuple)
    allocate(a, T::Type)
    allocate(a, T::Type, dims::Integer...)
    allocate(a, T::Type, dims::Tuple)

Allocate an object similar to `a`. It is similar to function `similar`, although
instead of working only on the outermost layer of a nested structure, it maps recursively
through outer layers and calls `similar` on the innermost array-like object only.
Type `T` is the new number element type [`number_eltype`](@ref), if it is not given
the element type of `a` is retained. The `dims` argument can be given for non-nested
allocation and is forwarded to the function `similar`.
"""
allocate(a, args...)
allocate(a) = similar(a)
allocate(a, dims::Integer...) = similar(a, dims...)
allocate(a, dims::Tuple) = similar(a, dims)
allocate(a, T::Type) = similar(a, T)
allocate(a, T::Type, dims::Integer...) = similar(a, T, dims...)
allocate(a, T::Type, dims::Tuple) = similar(a, T, dims)
allocate(a::AbstractArray{<:AbstractArray}) = map(allocate, a)
allocate(a::AbstractArray{<:AbstractArray}, T::Type) = map(t -> allocate(t, T), a)
allocate(a::NTuple{N,AbstractArray} where {N}) = map(allocate, a)
allocate(a::NTuple{N,AbstractArray} where {N}, T::Type) = map(t -> allocate(t, T), a)

"""
    allocate_result(M::Manifold, f, x...)

Allocate an array for the result of function `f` on [`Manifold`](@ref) `M` and arguments
`x...` for implementing the non-modifying operation using the modifying operation.

Usefulness of passing a function is demonstrated by methods that allocate results of musical
isomorphisms.
"""
function allocate_result(M::Manifold, f, x...)
    T = allocate_result_type(M, f, x)
    return allocate(x[1], T)
end

"""
    allocate_result_type(M::Manifold, f, args::NTuple{N,Any}) where N

Return type of element of the array that will represent the result of function `f` and the
[`Manifold`](@ref) `M` on given arguments `args` (passed as a tuple).
"""
function allocate_result_type(M::Manifold, f, args::NTuple{N,Any}) where {N}
    return typeof(mapreduce(eti -> one(number_eltype(eti)), +, args))
end

"""
    angle(M::Manifold, p, X, Y)

Compute the angle between tangent vectors `X` and `Y` at point `p` from the
[`Manifold`](@ref) `M` with respect to the inner product from [`inner`](@ref).
"""
angle(M::Manifold, p, X, Y) = acos(real(inner(M, p, X, Y)) / norm(M, p, X) / norm(M, p, Y))

"""
    base_manifold(M::Manifold, depth = Val(-1))

Return the internally stored [`Manifold`](@ref) for decorated manifold `M` and the base
manifold for vector bundles or power manifolds. The optional parameter `depth` can be used
to remove only the first `depth` many decorators and return the [`Manifold`](@ref) from that
level, whether its decorated or not. Any negative value deactivates this depth limit.
"""
base_manifold(M::Manifold, depth = Val(-1)) = M

"""
    check_manifold_point(M::Manifold, p; kwargs...) -> Union{Nothing,String}

Return `nothing` when `p` is a point on the [`Manifold`](@ref) `M`. Otherwise, return an
error with description why the point does not belong to manifold `M`.

By default, `check_manifold_point` returns `nothing`, i.e. if no checks are implemented, the
assumption is to be optimistic for a point not deriving from the [`MPoint`](@ref) type.
"""
check_manifold_point(M::Manifold, p; kwargs...) = nothing

"""
    check_tangent_vector(M::Manifold, p, X; kwargs...) -> Union{Nothing,String}

Check whether `X` is a valid tangent vector in the tangent space of `p` on the
[`Manifold`](@ref) `M`. An implementation should first call [`check_manifold_point(M, p;
kwargs...)`](@ref) and then validate `X`. If it is not a tangent vector, an error string
should be returned.

By default, `check_tangent_vector` returns `nothing`, i.e. if no checks are implemented, the
assumption is to be optimistic for tangent vectors not deriving from the [`TVector`](@ref)
type.
"""
check_tangent_vector(M::Manifold, p, X; kwargs...) = nothing

"""
    check_size(M::Manifold, p)
    check_size(M::Manifold, p, X)

Check whether `p` has the right [`representation_size`](@ref) for a [`Manifold`](@ref) `M`.
Additionally if a tangent vector is given, both `p` and `X` are checked to be of
corresponding correct representation sizes for points and tangent vectors on `M`.

By default, `check_size` returns `nothing`, i.e. if no checks are implemented, the
assumption is to be optimistic.
"""
function check_size(M::Manifold, p)
    n = size(p)
    m = representation_size(M)
    if length(n) != length(m)
        return DomainError(
            length(n),
            "The point $(p) can not belong to the manifold $(M), since its size $(n) is not equal to the manifolds representation size ($(m)).",
        )
    end
    if n != m
        return DomainError(
            n,
            "The point $(p) can not belong to the manifold $(M), since its size $(n) is not equal to the manifolds representation size ($(m)).",
        )
    end
end
function check_size(M::Manifold, p, X)
    mse = check_size(M, p)
    mse === nothing || return mse
    n = size(X)
    m = representation_size(M)
    if length(n) != length(m)
        return DomainError(
            length(n),
            "The tangent vector $(X) can not belong to the manifold $(M), since its size $(n) is not equal to the manifolds representation size ($(m)).",
        )
    end
    if n != m
        return DomainError(
            n,
            "The tangent vector $(X) can not belong to the manifold $(M), since its size $(n) is not equal to the manifodls representation size ($(m)).",
        )
    end
end

@doc raw"""
    copyto!(M::Manifold, q, p)

Copy the value(s) from `p` to `q`, where both are points on the [`Manifold`](@ref) `M`.
This function defaults to calling `copyto!(q, p)`, but it might be useful to overwrite the
function at the level, where also information from `M` can be accessed.
"""
copyto!(::Manifold, q, p) = copyto!(q, p)

@doc raw"""
    copyto!(M::Manifold, Y, p, X)

Copy the value(s) from `X` to `Y`, where both are tangent vectors from the tangent space at
`p` on the [`Manifold`](@ref) `M`.
This function defaults to calling `copyto!(Y, X)`, but it might be useful to overwrite the
function at the level, where also information from `p` and `M` can be accessed.
"""
copyto!(::Manifold, Y, p, X) = copyto!(Y, X)

"""
    distance(M::Manifold, p, q)

Shortest distance between the points `p` and `q` on the [`Manifold`](@ref) `M`.
"""
distance(M::Manifold, p, q) = norm(M, p, log(M, p, q))

"""
    embed(M::Manifold, p)

Embed point `p` from the [`Manifold`](@ref) `M` into the ambient space.
This method is only available for manifolds where implicitly an embedding or ambient space
is given.
Additionally, `embed` includes changing data representation, if applicable, i.e.
if the points on `M` are not represented in the same way as points on the embedding,
the representation is changed accordingly.

See also: [`EmbeddedManifold`](@ref), [`project`](@ref project(M::Manifold,p))
"""
function embed(M::Manifold, p)
    q = allocate_result(M, embed, p)
    embed!(M, q, p)
    return q
end

"""
    embed!(M::Manifold, q, p)

Embed point `p` from the [`Manifold`](@ref) `M` into an ambient space.
This method is only available for manifolds where implicitly an embedding or ambient space
is given. Not implementing this function means, there is no proper embedding for your manifold.
Additionally, `embed` might include changing data representation, if applicable, i.e.
if points on `M` are not represented in the same way as their counterparts in the embedding,
the representation is changed accordingly.

If you have more than one embedding, see [`EmbeddedManifold`](@ref) for defining a second
embedding. If your point `p` is already represented in some embedding,
see [`AbstractEmbeddedManifold`](@ref) how you can avoid reimplementing code from the embedded manifold

See also: [`EmbeddedManifold`](@ref), [`project!`](@ref project!(M::Manifold, q, p))
"""
function embed!(M::Manifold, q, p)
    return error(manifold_function_not_implemented_message(M, embed!, q, p))
end

"""
    embed(M::Manifold, p, X)

Embed a tangent vector `X` at a point `p` on the [`Manifold`](@ref) `M` into an ambient space.
This method is only available for manifolds where implicitly an embedding or ambient space
is given. Not implementing this function means, there is no proper embedding for your tangent space(s).

Additionally, `embed` might include changing data representation, if applicable, i.e.
if tangent vectors on `M` are not represented in the same way as their counterparts in the
embedding, the representation is changed accordingly.

If you have more than one embedding, see [`EmbeddedManifold`](@ref) for defining a second
embedding. If your tangent vector `X` is already represented in some embedding,
see [`AbstractEmbeddedManifold`](@ref) how you can avoid reimplementing code from the embedded manifold

See also: [`EmbeddedManifold`](@ref), [`project`](@ref project(M::Manifold, p, X))
"""
function embed(M::Manifold, p, X)
    # Note that the order is switched,
    # since the allocation by default takes the type of the first.
    Y = allocate_result(M, embed, X, p)
    embed!(M, Y, p, X)
    return Y
end

"""
    embed!(M::Manifold, Y, p, X)

Embed a tangent vector `X` at a point `p` on the [`Manifold`](@ref) `M` into the ambient
space and return the result in `Y`.
This method is only available for manifolds where implicitly an embedding or ambient space
is given.
Additionally, `embed!` includes changing data representation, if applicable, i.e.
if the tangents on `M` are not represented in the same way as tangents on the embedding,
the representation is changed accordingly. This is the case for example for Lie groups,
when tangent vectors are represented in the Lie algebra. The embedded tangents are then in
the tangent spaces of the embedded base points.

See also: [`EmbeddedManifold`](@ref), [`project!`](@ref project!(M::Manifold, Y, p, X))
"""
function embed!(M::Manifold, Y, p, X)
    return error(manifold_function_not_implemented_message(M, embed!, Y, p, X))
end

"""
    exp(M::Manifold, p, X)
    exp(M::Manifold, p, X, t::Real = 1)

Compute the exponential map of tangent vector `X`, optionally scaled by `t`,  at point `p`
from manifold the [`Manifold`](@ref) `M`.
"""
function exp(M::Manifold, p, X)
    q = allocate_result(M, exp, p, X)
    exp!(M, q, p, X)
    return q
end
exp(M::Manifold, p, X, t::Real) = exp(M, p, t * X)

"""
    exp!(M::Manifold, q, p, X)
    exp!(M::Manifold, q, p, X, t::Real = 1)

Compute the exponential map of tangent vector `X`, optionally scaled by `t`,  at point `p`
from manifold the [`Manifold`](@ref) `M`.
The result is saved to `q`.
"""
function exp!(M::Manifold, q, p, X)
    return error(manifold_function_not_implemented_message(M, exp!, q, p, X))
end
exp!(M::Manifold, q, p, X, t::Real) = exp!(M, q, p, t * X)

"""
    geodesic(M::Manifold, p, X) -> Function

Get the geodesic with initial point `p` and velocity `X` on the [`Manifold`](@ref) `M`.
 The geodesic is the curve of constant velocity that is locally distance-minimizing.
 This function returns a function of (time) `t`.

    geodesic(M::Manifold, x, v, t::Real)
    geodesic(M::Manifold, x, v, T::AbstractVector) -> AbstractVector

Return the point at time `t` or points at times `t` in `T` along the geodesic.
"""
geodesic(M::Manifold, p, X) = t -> exp(M, p, X, t)
geodesic(M::Manifold, p, X, t::Real) = exp(M, p, X, t)
geodesic(M::Manifold, p, X, T::AbstractVector) = map(t -> exp(M, p, X, t), T)

@doc doc"""
    injectivity_radius(M::Manifold, p)

Return the distance $d$ such that [`exp(M, p, X)`](@ref exp(::Manifold, ::Any, ::Any)) is
injective for all tangent vectors shorter than $d$ (i.e. has an inverse).

    injectivity_radius(M::Manifold)

Infimum of the injectivity radius of all manifold points.

    injectivity_radius(M::Manifold[, x], method::AbstractRetractionMethod)
    injectivity_radius(M::Manifold, x, method::AbstractRetractionMethod)

Distance $d$ such that
[`retract(M, p, X, method)`](@ref retract(::Manifold, ::Any, ::Any, ::AbstractRetractionMethod))
is injective for all tangent vectors shorter than $d$ (i.e. has an inverse) for point `p`
if provided or all manifold points otherwise.
"""
function injectivity_radius(M::Manifold)
    return error(manifold_function_not_implemented_message(M, injectivity_radius))
end
injectivity_radius(M::Manifold, p) = injectivity_radius(M)
function injectivity_radius(M::Manifold, p, method::AbstractRetractionMethod)
    return injectivity_radius(M, method)
end
function injectivity_radius(M::Manifold, method::AbstractRetractionMethod)
    return error(manifold_function_not_implemented_message(M, injectivity_radius, method))
end
injectivity_radius(M::Manifold, p, ::ExponentialRetraction) = injectivity_radius(M, p)
injectivity_radius(M::Manifold, ::ExponentialRetraction) = injectivity_radius(M)

"""
    inner(M::Manifold, p, X, Y)

Compute the inner product of tangent vectors `X` and `Y` at point `p` from the
[`Manifold`](@ref) `M`.

See also: [`MetricManifold`](@ref Main.Manifolds.MetricManifold)
"""
function inner(M::Manifold, p, X, Y)
    return error(manifold_function_not_implemented_message(M, inner, p, X, Y))
end

"""
    inverse_retract!(M::Manifold, X, p, q[, method::AbstractInverseRetractionMethod])

Compute the inverse retraction, a cheaper, approximate version of the
[`log`](@ref)arithmic map), of points `p` and `q` on the [`Manifold`](@ref) `M`.
Result is saved to `X`.

Inverse retraction method can be specified by the last argument, defaulting to
[`LogarithmicInverseRetraction`](@ref). See the documentation of respective manifolds for
available methods.
"""
function inverse_retract!(M::Manifold, X, p, q)
    return inverse_retract!(M, X, p, q, LogarithmicInverseRetraction())
end
function inverse_retract!(M::Manifold, X, p, q, method::LogarithmicInverseRetraction)
    return log!(M, X, p, q)
end
function inverse_retract!(M::Manifold, X, p, q, method::AbstractInverseRetractionMethod)
    return error(
        manifold_function_not_implemented_message(M, inverse_retract!, X, p, q, method),
    )
end

"""
    inverse_retract(M::Manifold, x, y)
    inverse_retract(M::Manifold, x, y, method::AbstractInverseRetractionMethod

Compute the inverse retraction, a cheaper, approximate version of the
[`log`](@ref)arithmic map), of points `p` and `q` on the [`Manifold`](@ref) `M`.

Inverse retraction method can be specified by the last argument, defaulting to
[`LogarithmicInverseRetraction`](@ref). See the documentation of respective manifolds for
available methods.
"""
function inverse_retract(M::Manifold, p, q)
    X = allocate_result(M, inverse_retract, p, q)
    inverse_retract!(M, X, p, q)
    return X
end
function inverse_retract(M::Manifold, p, q, method::AbstractInverseRetractionMethod)
    X = allocate_result(M, inverse_retract, p, q)
    inverse_retract!(M, X, p, q, method)
    return X
end

"""
    isapprox(M::Manifold, p, q; kwargs...)

Check if points `p` and `q` from [`Manifold`](@ref) `M` are approximately equal.

Keyword arguments can be used to specify tolerances.
"""
isapprox(::Manifold, x, y; kwargs...) = isapprox(x, y; kwargs...)

"""
    isapprox(M::Manifold, p, X, Y; kwargs...)

Check if vectors `X` and `Y` tangent at `p` from [`Manifold`](@ref) `M` are approximately
equal.

Keyword arguments can be used to specify tolerances.
"""
isapprox(::Manifold, p, X, Y; kwargs...) = isapprox(X, Y; kwargs...)

"""
    is_manifold_point(M::Manifold, p, throw_error = false; kwargs...)

Return whether `p` is a valid point on the [`Manifold`](@ref) `M`.

If `throw_error` is `false`, the function returns either `true` or `false`. If `throw_error`
is `true`, the function either returns `true` or throws an error. By default the function
calls [`check_manifold_point(M, p; kwargs...)`](@ref) and checks whether the returned value
is `nothing` or an error.
"""
function is_manifold_point(M::Manifold, p, throw_error = false; kwargs...)
    mpe = check_manifold_point(M, p; kwargs...)
    mpe === nothing && return true
    return throw_error ? throw(mpe) : false
end

"""
    is_tangent_vector(M::Manifold, p, X, throw_error = false; kwargs...)

Return whether `X` is a valid tangent vector at point `p` on the [`Manifold`](@ref) `M`.
Returns either `true` or `false`.

If `throw_error` is `false`, the function returns either `true` or `false`. If `throw_error`
is `true`, the function either returns `true` or throws an error. By default the function
calls [`check_tangent_vector(M, p, X; kwargs...)`](@ref) and checks whether the returned
value is `nothing` or an error.
"""
function is_tangent_vector(M::Manifold, p, X, throw_error = false; kwargs...)
    mtve = check_tangent_vector(M, p, X; kwargs...)
    mtve === nothing && return true
    return throw_error ? throw(mtve) : false
end

"""
    log(M::Manifold, p, q)

Compute the logarithmic map of point `q` at base point `p` on the [`Manifold`](@ref) `M`.
"""
function log(M::Manifold, p, q)
    X = allocate_result(M, log, p, q)
    log!(M, X, p, q)
    return X
end

"""
    log!(M::Manifold, X, p, q)

Compute the logarithmic map of point `q` at base point `p` on the [`Manifold`](@ref) `M`.
The result is saved to `X`.
"""
function log!(M::Manifold, X, p, q)
    return error(manifold_function_not_implemented_message(M, log!, X, p, q))
end

@doc doc"""
    manifold_dimension(M::Manifold)

The dimension $n=\dim_{\mathcal M}$ of real space $\mathbb R^n$ to which the neighborhood of
each point of the [`Manifold`](@ref) `M` is homeomorphic.
"""
function manifold_dimension(M::Manifold)
    return error(manifold_function_not_implemented_message(M, manifold_dimension))
end

function manifold_function_not_implemented_message(M::Manifold, f, x...)
    s = join(map(string, map(typeof, x)), ", ", " and ")
    a = length(x) > 1 ? "arguments" : "argument"
    m = length(x) > 0 ? " for $(a) $(s)." : "."
    return "$(f) not implemented on $(M)$(m)"
end

"""
    mid_point(M::Manifold, p1, p2)

Calculate the middle between the two point `p1` and `p2` from manifold `M`.
By default uses [`log`](@ref), divides the vector by 2 and uses [`exp`](@ref).
"""
function mid_point(M::Manifold, p1, p2)
    q = allocate(p1)
    return mid_point!(M, q, p1, p2)
end

"""
    mid_point!(M::Manifold, q, p1, p2)

Calculate the middle between the two point `p1` and `p2` from manifold `M`.
By default uses [`log`](@ref), divides the vector by 2 and uses [`exp!`](@ref).
Saves the result in `q`.
"""
function mid_point!(M::Manifold, q, p1, p2)
    X = log(M, p1, p2)
    return exp!(M, q, p1, X / 2)
end

@static if VERSION <= v"1.1"
    function mid_point!(
        M::Manifold,
        q::AbstractArray{T1,0},
        p1::AbstractArray{T2,0},
        p2::AbstractArray{T3,0},
    ) where {T1,T2,T3}
        X = log(M, p1, p2)
        return exp!(M, q, p1, fill(X / 2))
    end
end

"""
    norm(M::Manifold, p, X)

Compute the norm of tangent vector `X` at point `p` from a [`Manifold`](@ref) `M`.
By default this is computed using [`inner`](@ref).
"""
norm(M::Manifold, p, X) = sqrt(real(inner(M, p, X, X)))

"""
    number_eltype(x)

Numeric element type of the a nested representation of a point or a vector.
To be used in conjuntion with [`allocate`](@ref) or [`allocate_result`](@ref).
"""
number_eltype(x) = eltype(x)
function number_eltype(x::AbstractArray)
    return typeof(mapreduce(eti -> one(number_eltype(eti)), +, x))
end
function number_eltype(x::Tuple)
    return typeof(mapreduce(eti -> one(number_eltype(eti)), +, x))
end

"""
    project(M::Manifold, p)

Project point `p` from the ambient space of the [`Manifold`](@ref) `M` to `M`.
This method is only available for manifolds where implicitly an embedding or ambient space
is given. Additionally, the projection includes changing data representation, if applicable,
i.e. if the points on `M` are not represented in the same array data, the data is changed
accordingly.

See also: [`EmbeddedManifold`](@ref), [`embed`](@ref embed(M::Manifold, p))
"""
function project(M::Manifold, p)
    q = allocate_result(M, project, p)
    project!(M, q, p)
    return q
end

"""
    project!(M::Manifold, q, p)

Project point `p` from the ambient space onto the [`Manifold`](@ref) `M`. The result is
storedin `q`.
This method is only available for manifolds where implicitly an embedding or ambient space
is given. Additionally, the projection includes changing data representation, if applicable,
i.e. if the points on `M` are not represented in the same array data, the data is changed
accordingly.

See also: [`EmbeddedManifold`](@ref), [`embed!`](@ref embed!(M::Manifold, q, p))
"""
function project!(M::Manifold, q, p)
    return error(manifold_function_not_implemented_message(M, project!, q, p))
end

"""
    project(M::Manifold, p, X)

Project ambient space representation of a vector `X` to a tangent vector at point `p` on
the [`Manifold`](@ref) `M`.
This method is only available for manifolds where implicitly an embedding or ambient space
is given.
Additionally, `project` includes changing data representation, if applicable, i.e.
if the tangents on `M` are not represented in the same way as points on the embedding,
the representation is changed accordingly. This is the case for example for Lie groups,
when tangent vectors are represented in the Lie algebra. after projection the change to the
Lie algebra is perfomed, too.

See also: [`EmbeddedManifold`](@ref), [`embed`](@ref embed(M::Manifold, p, X))
"""
function project(M::Manifold, p, X)
    # Note that the order is switched,
    # since the allocation by default takes the type of the first.
    Y = allocate_result(M, project, X, p)
    project!(M, Y, p, X)
    return Y
end

"""
    project!(M::Manifold, Y, p, X)

Project ambient space representation of a vector `X` to a tangent vector at point `p` on
the [`Manifold`](@ref) `M`. The result is saved in vector `Y`.
This method is only available for manifolds where implicitly an embedding or ambient space
is given.
Additionally, `project!` includes changing data representation, if applicable, i.e.
if the tangents on `M` are not represented in the same way as points on the embedding,
the representation is changed accordingly. This is the case for example for Lie groups,
when tangent vectors are represented in the Lie algebra. after projection the change to the
Lie algebra is perfomed, too.

See also: [`EmbeddedManifold`](@ref), [`embed!`](@ref embed!(M::Manifold, Y, p, X))
"""
function project!(M::Manifold, Y, p, X)
    return error(manifold_function_not_implemented_message(M, project!, Y, p, X))
end

@doc doc"""
    representation_size(M::Manifold)

The size of an array representing a point on [`Manifold`](@ref) `M`.
Returns `nothing` by default indicating that points are not represented using an
`AbstractArray`.
"""
function representation_size(M::Manifold)
    return nothing
end

"""
    retract(M::Manifold, p, X)
    retract(M::Manifold, p, X, t::Real=1)
    retract(M::Manifold, p, X, method::AbstractRetractionMethod)
    retract(M::Manifold, p, X, t::Real=1, method::AbstractRetractionMethod)

Compute a retraction, a cheaper, approximate version of the [`exp`](@ref)onential map,
from `p` into direction `X`, scaled by `t`, on the [`Manifold`](@ref) `M`.

Retraction method can be specified by the last argument, defaulting to
[`ExponentialRetraction`](@ref). See the documentation of respective manifolds for available
methods.
"""
function retract(M::Manifold, p, X)
    q = allocate_result(M, retract, p, X)
    retract!(M, q, p, X)
    return q
end
retract(M::Manifold, p, X, t::Real) = retract(M, p, t * X)
function retract(M::Manifold, p, X, method::AbstractRetractionMethod)
    q = allocate_result(M, retract, p, X)
    retract!(M, q, p, X, method)
    return q
end
function retract(M::Manifold, p, X, t::Real, method::AbstractRetractionMethod)
    return retract(M, p, t * X, method)
end

"""
    retract!(M::Manifold, q, p, X)
    retract!(M::Manifold, q, p, X, t::Real=1)
    retract!(M::Manifold, q, p, X, method::AbstractRetractionMethod)
    retract!(M::Manifold, q, p, X, t::Real=1, method::AbstractRetractionMethod)

Compute a retraction, a cheaper, approximate version of the [`exp`](@ref)onential map,
from `p` into direction `X`, scaled by `t`, on the [`Manifold`](@ref) manifold `M`.
Result is saved to `q`.

Retraction method can be specified by the last argument, defaulting to
[`ExponentialRetraction`](@ref). See the documentation of respective manifolds for available
methods.
"""
retract!(M::Manifold, q, p, X) = retract!(M, q, p, X, ExponentialRetraction())
retract!(M::Manifold, q, p, X, t::Real) = retract!(M, q, p, t * X)
retract!(M::Manifold, q, p, X, method::ExponentialRetraction) = exp!(M, q, p, X)
function retract!(M::Manifold, q, p, X, t::Real, method::AbstractRetractionMethod)
    return retract!(M, q, p, t * X, method)
end
function retract!(M::Manifold, q, p, X, method::AbstractRetractionMethod)
    return error(manifold_function_not_implemented_message(M, retract!, q, p, method))
end

@doc doc"""
    shortest_geodesic(M::Manifold, p, q) -> Function

Get a [`geodesic`](@ref) $\gamma_{p,q}(t)$ whose length is the shortest path between the
points `p`and `q`, where $\gamma_{p,q}(0)=p$ and $\gamma_{p,q}(1)=q$. When there are
multiple shortest geodesics, there is no guarantee which will be returned.

This function returns a function of time, which may be a `Real` or an `AbstractVector`.

    shortest_geodesic(M::Manifold, p, q, t::Real)
    shortest_geodesic(M::Manifold, p, q, T::AbstractVector) -> AbstractVector

Return the point at time `t` or points at times `t` in `T` along the shortest geodesic.
"""
shortest_geodesic(M::Manifold, p, q) = geodesic(M, p, log(M, p, q))
shortest_geodesic(M::Manifold, p, q, t::Real) = geodesic(M, p, log(M, p, q), t)
shortest_geodesic(M::Manifold, p, q, T::AbstractVector) = geodesic(M, p, log(M, p, q), T)

"""
    size_to_tuple(::Type{S}) where S<:Tuple

Converts a size given by `Tuple{N, M, ...}` into a tuple `(N, M, ...)`.
"""
Base.@pure size_to_tuple(::Type{S}) where {S<:Tuple} = tuple(S.parameters...)

"""
    zero_tangent_vector!(M::Manifold, X, p)

Save to `X` a vector such that retracting `X` to the [`Manifold`](@ref) `M` at `p`
produces `p`.
"""
zero_tangent_vector!(M::Manifold, X, p) = log!(M, X, p, p)

"""
    zero_tangent_vector(M::Manifold, p)

Return the tangent vector from the tangent space at `p` on the [`Manifold`](@ref) `M`, that
represents the zero vector, i.e. such that a retraction at `p` produces `p`.
"""
function zero_tangent_vector(M::Manifold, p)
    X = allocate_result(M, zero_tangent_vector, p)
    zero_tangent_vector!(M, X, p)
    return X
end
include("errors.jl")
include("numbers.jl")
include("vector_transport.jl")
include("DecoratorManifold.jl")
include("bases.jl")
include("ValidationManifold.jl")
include("EmbeddedManifold.jl")
include("DefaultManifold.jl")
include("PowerManifold.jl")

export Manifold, MPoint, TVector, CoTVector
export AbstractDecoratorManifold
export ValidationManifold, ValidationMPoint, ValidationTVector, ValidationCoTVector
export AbstractEmbeddingType,
    TransparentIsometricEmbedding, DefaultIsometricEmbeddingType, DefaultEmbeddingType
export AbstractEmbeddedManifold, EmbeddedManifold, TransparentIsometricEmbedding
export AbstractPowerManifold, PowerManifold
export AbstractPowerRepresentation, NestedPowerRepresentation

export OutOfInjectivityRadiusError

export AbstractRetractionMethod,
    ExponentialRetraction,
    QRRetraction,
    PolarRetraction,
    ProjectionRetraction,
    PowerRetraction,
    InversePowerRetraction

export AbstractInverseRetractionMethod,
    LogarithmicInverseRetraction,
    QRInverseRetraction,
    PolarInverseRetraction,
    ProjectionInverseRetraction

export AbstractVectorTransportMethod,
    DifferentiatedRetractionVectorTransport,
    ParallelTransport,
    PoleLadderTransport,
    PowerVectorTransport,
    ProjectionTransport,
    ScaledVectorTransport,
    SchildsLadderTransport

export CachedBasis,
    DefaultBasis,
    DefaultOrthogonalBasis,
    DefaultOrthonormalBasis,
    DiagonalizingOrthonormalBasis,
    DefaultOrthonormalBasis,
    GramSchmidtOrthonormalBasis,
    ProjectedOrthonormalBasis

export CompositeManifoldError, ComponentManifoldError

export allocate,
    base_manifold,
    check_manifold_point,
    check_tangent_vector,
    check_size,
    distance,
    exp,
    exp!,
    embed,
    embed!,
    geodesic,
    get_basis,
    get_component,
    get_component!,
    get_coordinates,
    get_coordinates!,
    get_embedding,
    get_vector,
    get_vector!,
    get_vectors,
    hat,
    hat!,
    shortest_geodesic,
    injectivity_radius,
    inner,
    inverse_retract,
    inverse_retract!,
    isapprox,
    is_manifold_point,
    is_tangent_vector,
    isempty,
    length,
    log,
    log!,
    manifold_dimension,
    mid_point,
    mid_point!,
    norm,
    number_eltype,
    number_of_coordinates,
    number_system,
    power_dimensions,
    project,
    project!,
    real_dimension,
    representation_size,
    set_component!,
    show,
    retract,
    retract!,
    vector_transport_along,
    vector_transport_along!,
    vector_transport_direction,
    vector_transport_direction!,
    vector_transport_to,
    vector_transport_to!,
    vee,
    vee!,
    zero_tangent_vector,
    zero_tangent_vector!

end # module


@testset "Partition" begin
  @testset "Directional" begin
    # merge operation does not produce NaN
    dir = (0.286788, -0.496732, -0.819152)
    𝒟 = georef(CSV.File(joinpath(datadir,"nanlags.csv")), (:X,:Y,:Z))
    γ = DirectionalVariogram(dir, 𝒟, :Cu, dtol=45, maxlag=150, nlags=20)
    x, y, n = values(γ)
    @test !any(isnan.(x))
    @test !any(isnan.(y))
    @test !any(isnan.(n))

    # directional variogram and known anisotropy ratio
    img = readdlm(joinpath(datadir,"anisotropic.tsv"))
    sdata = georef((z=img,))
    γhor = DirectionalVariogram((1.,0.), sdata, :z, maxlag=50.)
    γver = DirectionalVariogram((0.,1.), sdata, :z, maxlag=50.)
    γₕ = fit(GaussianVariogram, γhor)
    γᵥ = fit(GaussianVariogram, γver)
    @test range(γₕ) / range(γᵥ) ≈ 3. atol=.1

    if visualtests
      p1 = plot(γhor, showbins=false, label="horizontal")
      plot!(γver, showbins=false, label="vertical")
      p2 = plot(γₕ, 0., 50., label="horizontal")
      plot!(γᵥ, 0., 50., label="vertical")
      plt = plot(p1, p2, layout=(2,1))
      @test_reference "data/directional.png" plt
    end
  end

  @testset "Planar" begin
    # directional equals planar rotated by 90 degrees in 2D
    img = readdlm(joinpath(datadir,"anisotropic.tsv"))
    sdata = georef((z=img,))
    γ₁ = PlanarVariogram((0.,1.), sdata, :z, maxlag=50.)
    γ₂ = DirectionalVariogram((1.,0.), sdata, :z, maxlag=50.)
    x₁, y₁, n₁ = values(γ₁)
    x₂, y₂, n₂ = values(γ₂)
    @test x₁ == x₂
    @test y₁ ≈ y₂
    @test n₁ == n₂
    γ₁ = PlanarVariogram((1.,0.), sdata, :z, maxlag=50.)
    γ₂ = DirectionalVariogram((0.,1.), sdata, :z, maxlag=50.)
    x₁, y₁, n₁ = values(γ₁)
    x₂, y₂, n₂ = values(γ₂)
    @test x₁ == x₂
    @test y₁ ≈ y₂
    @test n₁ == n₂

    # planar variogram and known anisotropy ratio
    γhor = PlanarVariogram((0.,1.), sdata, :z, maxlag=50.)
    γver = PlanarVariogram((1.,0.), sdata, :z, maxlag=50.)
    γₕ = fit(GaussianVariogram, γhor)
    γᵥ = fit(GaussianVariogram, γver)
    @test range(γₕ) / range(γᵥ) ≈ 3. atol=.1

    if visualtests
      p1 = plot(γhor, showbins=false, label="horizontal")
      plot!(γver, showbins=false, label="vertical")
      p2 = plot(γₕ, 0., 50., label="horizontal")
      plot!(γᵥ, 0., 50., label="vertical")
      plt = plot(p1, p2, layout=(2,1))
      @test_reference "data/planar.png" plt
    end
  end
end


using Statistics
using StatsFuns: xlogx, xlogy

"Cache pairwise / marginal distribution for all variables in one dataset"
mutable struct DisCache
    pairwise::Array{Float64}
    marginal::Array{Float64}
end

@inline dimension(discache::DisCache) = size(discache.marginal)[1]
DisCache(num) = DisCache(Array{Float64}(undef, num, num, 4), Array{Float64}(undef, num, 2))

#####################
# Methods for pairwise and marginal distribution
#####################
@inline get_parameters(bm::AbstractMatrix{<:Bool}, α, w=nothing) = size(bm)[2], issomething(w) ? sum(w) : size(bm)[1], @. Float64(bm), @. Float64(!bm)

function cache_distributions(bm::AbstractMatrix{<:Bool}, w::Union{Nothing, AbstractVector{<:AbstractFloat}}=nothing; α, flag=(pairwise=true, marginal=true))
    # parameters
    D, N, (m, notm) = get_parameters(bm, α, w)
    dis_cache = DisCache(D)
    base = N + 4 * α
    w = isnothing(w) ? ones(Float64, N) : w

    # pairwise distribution
    if flag.pairwise
        dis_cache.pairwise[:,:,1] = (notm' * (notm .* w) .+ α) / base   # p00
        dis_cache.pairwise[:,:,2] = (notm' * (m .* w) .+ α) / base      # p01
        dis_cache.pairwise[:,:,3] = (m' * (notm .* w) .+ α) / base      # p10
        dis_cache.pairwise[:,:,4] = (m' * (m .* w) .+ α) / base         # p11
    end
    # marginal distribution

    if flag.marginal
        dis_cache.marginal[:, 1] = (sum(notm .* w, dims=1) .+ 2 * α) / base
        dis_cache.marginal[:, 2] = (sum(m .* w, dims=1).+ 2 * α) / base
    end
    dis_cache
end

#####################
# Mutual Information
#####################

function mutual_information(dis_cache::DisCache)
    D = dimension(dis_cache)
    p0 = @view dis_cache.marginal[:, 1]
    p1 = @view dis_cache.marginal[:, 2]
    pxpy = Array{Float64}(undef, D, D, 4)
    pxpy[:,:,1] = p0 * p0'
    pxpy[:,:,2] = p0 * p1'
    pxpy[:,:,3] = p1 * p0'
    pxpy[:,:,4] = p1 * p1'
    pxy_log_pxy = @. xlogx(dis_cache.pairwise)
    pxy_log_pxpy = @. xlogy(dis_cache.pairwise, pxpy)
    dropdims(sum(pxy_log_pxy - pxy_log_pxpy,dims=3), dims=3)
end

"Calculate mutual information of given bit matrix `bm`, example weights `w`, and smoothing pseudocount `α`"
function mutual_information(bm::AbstractMatrix{<:Bool}, w::Union{Nothing, AbstractVector{<:AbstractFloat}}=nothing; α)
    dis_cache = cache_distributions(bm, w; α=α)
    mi = mutual_information(dis_cache)
    return (dis_cache, mi)
end

function mutual_information(train_x::PlainXData, w::Union{Nothing, AbstractVector{<:AbstractFloat}}=nothing; α)
    mutual_information(feature_matrix(train_x), w; α=α)
end

"Calculate set mutual information"
function set_mutual_information(mi::Matrix, sets::Vector{Vector{Var}})::Matrix
    len = length(sets)
    if len == size(mi)[1]
        return mi
    end

    pmi = zeros(len, len)
    for i in 1 : len, j in i + 1 : len
        pmi[i, j] = pmi[j, i] = mean(mi[sets[i], sets[j]])
    end
    return pmi
end

#####################
# Entropy
#####################
function entropy(dis_cache::DisCache)
    D = dimension(dis_cache)
    px_log_px = @. xlogx(dis_cache.marginal)
    - dropdims(sum(px_log_px; dims=2); dims=2)
end

function entropy(bm::AbstractMatrix{<:Bool}, w::Union{Nothing, AbstractVector{<:AbstractFloat}}=nothing; α)
    dis_cache = cache_distributions(bm, w; α=α, flag=(pairwise=true, marginal=true))
    return (dis_cache, entropy(dis_cache))
end

function sum_entropy_given_x(bm::AbstractMatrix{<:Bool}, x::Var, w::Union{Nothing, AbstractVector{<:AbstractFloat}}=nothing; α)::Float64
    @assert x <= size(bm)[2]
    vars = [1 : x-1; x+1 : size(bm)[2]]
    indexes_left = bm[:,x].== 0
    indexes_right = bm[:,x] .== 1
    w1 = sum(Float64.(indexes_left))
    w2 = sum(Float64.(indexes_right))
    w1, w2 = w1 / (w1 + w2), w2 / (w1 + w2)
    subm_left = @view bm[indexes_left, vars]
    subm_right = @view bm[indexes_right, vars]

    w_left = issomething(w) ? (@view w[indexes_left]) : nothing
    w_right = issomething(w) ? (@view w[indexes_right]) : nothing

    w1 * sum(entropy(subm_left, w_left; α=α)[2]) + w2 * sum(entropy(subm_right, w_right; α=α)[2])
end

function conditional_entropy(dis_cache::DisCache)
    D = dimension(dis_cache)
    pxy_log_pxy = @. xlogx(dis_cache.pairwise)
    pxy = dis_cache.pairwise
    log_px = log.(dis_cache.marginal)

    pxy_log_px = Array{Float64}(undef, D, D, 4)
    pxy_log_px[:, :, 1] = view(pxy, :, :, 1) .* view(log_px, :, 1)
    pxy_log_px[:, :, 2] = view(pxy, :, :, 2) .* view(log_px, :, 1)
    pxy_log_px[:, :, 3] = view(pxy, :, :, 3) .* view(log_px, :, 2)
    pxy_log_px[:, :, 4] = view(pxy, :, :, 4) .* view(log_px, :, 2)
    h_y_given_x = - dropdims(sum(pxy_log_pxy - pxy_log_px; dims=3); dims=3)
end

function conditional_entropy(bm::AbstractMatrix{<:Bool}, w::Union{Nothing, AbstractVector{<:AbstractFloat}}=nothing; α)
    dis_cache = cache_distributions(bm, w; α=α)
    return (dis_cache, conditional_entropy(dis_cache))
end

function set_conditional_entropy(ce::Matrix{<:Float64})
    dropdims(sum(ce; dims=2); dims=2)
end


#=
    Figure S4
=#
using Optim

xlim = [0.0,0.5]
xgrd = range(xlim...,length=200)

μ = 0.25
σ = [0.01,0.05,0.1]

ω = [-1.5,-0.5,0.0001,0.5,1.5]

σᵢ = σ[1]

plts = [plot() for i = 1:length(σ)]

for (i,σᵢ) in enumerate(σ)

    pltsᵢ = [plot() for i = 1:length(ω)]

    for (j,ωⱼ) in enumerate(ω)

        λ = GammaAlt(μ,σᵢ,ωⱼ)

        # Find mode
        m = optimize(x -> -pdf(λ,x),xlim...).minimizer

        if ωⱼ > 0.0
            xmap = x -> x .> m ? 2m - x : x
        elseif ωⱼ < 0.0
            xmap = x -> x .< m ? 2m - x : x
        else
            xmap = x -> x
        end

        # Plot symmetric distribution
        plot!(pltsᵢ[j],xgrd,pdf(λ,xmap.(xgrd)),lw=2.0,c=:red,α=0.5,ls=:dash)
        plot!(pltsᵢ[j],xgrd,pdf(λ,xgrd),lw=2.0,c=:black)

    end

    plts[i] = plot(pltsᵢ...,layout=grid(1,length(ω)))
    if i < length(σ)
        plot!(plts[i],xticks=(0.0:0.1:0.5,[]))
    else
        plot!(plts[i],xticks=0.0:0.1:0.5)
    end
    plot!(plts[i],subplot=1,ylabel="σ = $σᵢ")

end

figS4 = plot(plts...,layout=grid(length(σ),1),yticks=[],axis=:x,legend=:none,size=(700,400))
[plot!(figS4,subplot=i,title="ω = $ωᵢ") for (i,ωᵢ) in enumerate(ω)]
[plot!(figS4,subplot=i,xlabel="λ") for i = 11:15]
figS4

##############################################################
## Save...
##############################################################

savefig("$(@__DIR__)/FigS4.svg")


using UnivariateFunctions
using StochasticIntegrals
using DataFrames, Dates
using LinearAlgebra
using Distributions: Normal, MersenneTwister, quantile
using Statistics: var, mean, cov
using Sobol
using Distributions
using Random
const tol = 10*eps()

USD_vol = 0.0001
GBP_vol = 0.00001
GBP_FX_Vol = 0.002

today = Date(2016,12,1)
tommorow = Date(2016,12,2)
date_2020 = Date(2020,12,1)
later_date = Date(2035,12,1)
later_later_date = Date(2080,12,1)

USD_hw_a_curve = PE_Function(USD_vol)
GBP_hw_a_curve = PE_Function(GBP_vol)

brownian_corr_matrix = Hermitian([1.0 0.75 0.5 0.0;
                                  0.0 1.0 0.5 0.25;
                                  0.0 0.0 1.0 0.25;
                                  0.0 0.0 0.0 1.0])
brownian_ids = [:USD_IR, :GBP_IR, :GBP_FX, :BARC]
USD_IR_a_ito  = ItoIntegral(:USD_IR, USD_hw_a_curve)
GBP_IR_a_ito  = ItoIntegral(:GBP_IR, GBP_hw_a_curve)
GBP_FX_ito    = ItoIntegral(:GBP_FX, GBP_FX_Vol)
ito_integrals = Dict([:USD_IR_a, :GBP_IR_a, :GBP_FX] .=> [USD_IR_a_ito, GBP_IR_a_ito, GBP_FX_ito])

ito_set_ = ItoSet(brownian_corr_matrix, brownian_ids, ito_integrals)

for_covar = ForwardCovariance(ito_set_, years_from_global_base(today),
                              years_from_global_base(date_2020))
simp_covar = SimpleCovariance(ito_set_, years_from_global_base(today),
                              years_from_global_base(date_2020))

# Making ItoProcesses

ito_processes = Dict{Symbol,ItoProcess{Float64}}()
update_rates = Dict{Symbol,Exponential}()
for k in keys(ito_integrals)
    ito_processes[k] = ItoProcess( 0.0, 100.0, PE_Function(0.0), ito_integrals[k] )
    update_rates[k] = Exponential(5.0)
end

# Making the series.

for_series_sync = make_ito_process_syncronous_time_series(deepcopy(ito_processes), deepcopy(for_covar), 5.0, 1000)
abs(nrow(for_series_sync) - 3000) < 0.5
simp_series_sync = make_ito_process_syncronous_time_series(deepcopy(ito_processes), deepcopy(simp_covar), 5.0, 1000)
abs(nrow(simp_series_sync) - 3000) < 0.5

for_series_nonsync  = make_ito_process_non_syncronous_time_series(deepcopy(ito_processes), deepcopy(for_covar), update_rates , 1000)
abs(nrow(for_series_nonsync) - 1000) < 0.5
simp_series_nonsync = make_ito_process_non_syncronous_time_series(deepcopy(ito_processes), deepcopy(simp_covar), update_rates , 1000)
abs(nrow(simp_series_nonsync) - 1000) < 0.5

# Testing rolling covariance by date.
rolled_date = ForwardCovariance(for_covar, today, tommorow)
rolled_numbers = ForwardCovariance(for_covar, years_from_global_base(today), years_from_global_base(tommorow))
all(abs.(rolled_date.covariance_ .- rolled_numbers.covariance_) .< 100*eps())


using Printf
using LinearAlgebra
using SpecialFunctions

include("constants.jl")
include("alias.jl")
include("Atoms.jl")
include("utils.jl")
include("PGBF.jl")
include("CGBF.jl")
include("overlap.jl")
include("kinetic.jl")
include("nuclear.jl")
include("coulomb.jl")
include("sto3g.jl")
include("BasisSet.jl")
include("RHF.jl")

function test_main(filename)
    atoms = init_atoms_xyz(filename)
    energy, E, U = RHF(atoms)
end

#test_main("H2_v2.xyz")
#test_main("H2O.xyz")
#test_main("LiH.xyz")
test_main("I2.xyz")


using SpecialFunctions

#    _|_|_|    _|_|    _|_|_|    _|      _|
#  _|        _|    _|  _|    _|  _|      _|
#  _|        _|    _|  _|    _|    _|  _|
#    _|_|_|    _|_|    _|    _|      _|

@makepdftype myBW(x;p) = let xth = 2.961477
    AlgebraPDF.amplitudeBWsq(x+xth, p.Δm+xth, p.Γ)
end
mybw = myBW(Ext(Δm=0.1, Γ=3e-3), (0,0.22))

# convolute with average resolution
const σ_mean = 5e-3
mybw_conv = convGauss(mybw, σ_mean)

@testset "convGauss: update pars works" begin
    @test pars(updatepars(mybw_conv, (Δm=0.15,))).Δm == 0.15
end

function updategetmaxdensity(d, Δm)
    xv = range(0,0.22,length=100)
    maximum(updatepars(d,(Δm=Δm,))(xv))
end

# resolution is linear function of energy
resolution_model = FunctionWithParameters((x;p)->σ_mean+(x-0.1)*3e-2; p=∅)
mybw_conv_dep = convGauss(mybw, resolution_model)

@testset "convGauss energy-dependent gets wider" begin
    # if the peak located below 0.1 the resolution the effective resolution  < σ_mean => the peak max is higher
    @test updategetmaxdensity(mybw_conv_dep, 0.05) > updategetmaxdensity(mybw_conv, 0.05)
    # if the peak located above 0.1 the resolution the effective resolution  > σ_mean => the peak max is smaller
    @test updategetmaxdensity(mybw_conv_dep, 0.15) < updategetmaxdensity(mybw_conv, 0.15)
end