#=
dynamic_nelson_siegel_model.jl

    Dynamic Nelson Siegel model specification in state space formulation

@author: Quint Wiersma <q.wiersma@vu.nl>

@date: 2022/03/28
=#

"""
    DynamicNelsonSiegelModel <: StateSpaceModel

Constructor for a dynamic Nelson Siegel model instance of the state space model type
with hyperparameters `λ`, `ϕ`, and error specifications
"""
mutable struct DynamicNelsonSiegelModel{Ty, Tτ, Tλ, Tϕ, Teobs, Tefac} <: StateSpaceModel
    y::Ty               # data
    τ::Tτ               # maturities
    λ::Tλ               # decay
    ϕ::Tϕ               # autoregressive parameters
    error_obs::Teobs    # obs. eq. error specification
    error_factor::Tefac # factor eq. error specification
end
# Constructors
function DynamicNelsonSiegelModel(y::AbstractMatrix, τ::AbstractVector,
                                    error_obs::AbstractErrorModel, error_factor::AbstractErrorModel)
    # hyper paremeters
    λ= .0609
    ϕ= similar(y, 3, 3)

    return DynamicNelsonSiegelModel(y, τ, λ, ϕ, error_obs, error_factor)
end

# Methods
cov(model::DynamicNelsonSiegelModel)= cov(model.error_obs)
prec(model::DynamicNelsonSiegelModel)= prec(model.error_obs)
resid(model::DynamicNelsonSiegelModel)= resid(model.error_obs)
maturities(model::DynamicNelsonSiegelModel)= model.τ
decay(model::DynamicNelsonSiegelModel)= model.λ

function loadings!(Λ::AbstractMatrix, model::DynamicNelsonSiegelModel)
    λ= decay(model)
    τ= maturities(model)
    n= length(τ)

    Λ.= one(eltype(Λ))
    @inbounds @fastmath for i in 1:n
        Λ[i,2]= (1 - exp(-λ * τ[i])) * inv(λ * τ[i])
        Λ[i,3]= Λ[i,2] - exp(-λ * τ[i])
    end

    return nothing
end
function loadings(model::DynamicNelsonSiegelModel)
    # get dims
    n= length(maturities(model))

    Λ= Matrix{typeof(decay(model))}(undef, n, 3)
    loadings!(Λ, model)

    return Λ
end

# State space system and hyperparameters
nparams(model::DynamicNelsonSiegelModel)= 10 + nparams(model.error_obs) + nparams(model.error_factor)

function get_params!(ψ::AbstractVector, model::DynamicNelsonSiegelModel)
    idx= 0  # index counter

    # Store values
    # λ
    ψ[1]= model.λ
    idx+= 1

    # ϕ
    ψ[idx+1:idx+9].= vec(model.ϕ)
    idx+= 9

    # observation equation error model
    get_params!(view(ψ,idx+1:idx+nparams(model.error_obs)), model.error_obs)
    idx+= nparams(model.error_obs)

    # factor equation error model
    get_params!(view(ψ,idx+1:idx+nparams(model.error_fac)), model.error_factor)

    return nothing
end

function get_system!(sys::LinearTimeInvariant, model::DynamicNelsonSiegelModel, method::Symbol)
    Λ= loadings(model)  # loadings

    # Store values
    sys.T.= model.ϕ
    sys.Q.= cov(model.error_factor)
    if method === :univariate
        # Cholesky decomposition of H
        C= cholesky(cov(model))

        # Store values
        ldiv!(sys.y, C.L, model.y)
        ldiv!(sys.Z, C.L, Λ)
    elseif method === :collapsed
        # Collapsing transformation
        # Z'×H⁻¹×Z
        tmp_pn= transpose(Λ) * prec(model)
        tmp= tmp_pn * Λ
        # (Z'×H⁻¹×Z)⁻¹
        # perform pseudo inverse as Λ can contain zero columns
        pseudo= pinv(tmp)
        # pivoted Cholesky decomposition
        C= cholesky!(Hermitian(pseudo), Val(true), check=false)
        # transformation
        ip= invperm(C.p)
        U= C.U[:,ip]
        A= U * tmp_pn

        # Store values
        mul!(sys.y, A, model.y)
        sys.Z.= pinv(transpose(U))
    else
        # Store values
        sys.Z.= Λ
        sys.H.= cov(model)
    end

    return nothing
end

function store_params!(model::DynamicNelsonSiegelModel, ψ::AbstractVector)
    idx= 0  # index counter

    # Store values
    # λ
    model.λ= ψ[1]
    idx+= 1

    # ϕ
    vec(model.ϕ).= view(ψ,idx+1:idx+9)
    idx+= 9

    # observation equation error model
    store_params!(model.error_obs, view(ψ,idx+1:idx+nparams(model.error_obs)))
    idx+= nparams(model.error_obs)

    # factor equation error model
    store_params!(model.error_factor, view(ψ,idx+1:idx+nparams(model.error_factor)))

    return nothing
end

# Initialization
function init!(model::DynamicNelsonSiegelModel, init::NamedTuple, method::Symbol)
    # Model
    init_model!(model, init)

    # State space system
    create_system(model, method)
    fix_system!(sys, model, method)
    init_system!(sys, model)

    return sys
end

function init_factors(y::AbstractMatrix, Λ::AbstractMatrix)
    # Get dims
    T= size(y,2)

    # Result container
    β= similar(y, 3, T)

    # OLS
    A= transpose(Λ) * Λ
	B= A \ transpose(Λ)
    mul!(β, B, y)

    return β
end

"""
    init_ϕ!(ϕ, β)

Initialize autoregressive parameters `ϕ` using OLS with factors, storing the
result in `ϕ`.

#### Arguments
  - `β::AbstractMatrix` : factors

#### Returns
  - `ϕ::AbstractMatrix` : autoregressive parameters
"""
function init_ϕ!(ϕ::AbstractMatrix, β::AbstractMatrix)
    # Get dims
    T= size(β,2)

    # OLS
    X= view(β,:,1:T-1)
    y= view(β,:,2:T)
    ϕ.= y / X

    return nothing
end

function init_model!(model::DynamicNelsonSiegelModel, init::NamedTuple)
    # get dims
    T= size(model.y,2)

    # λ
    haskey(init, :λ) ? model.λ= init.λ : nothing
    # factors
    Λ= loadings(model)
    β= init_factors(model.y, Λ)

    # observation equation residuals
    resid(model).= model.y
    mul!(resid(model), Λ, β, -1., 1.)

    # ϕ
    haskey(init, :ϕ) ? model.ϕ.= init.ϕ : init_ϕ!(model.ϕ, β)

    # factor equation residuals
    resid(model.error_factor).= view(β,:,2:T)
    mul!(resid(model.error_factor), model.ϕ, view(β,:,1:T-1), -1., 1.)

    # observation equation error model
    init_error!(model.error_obs, init)

    # factor equation error model
    init_error!(model.error_factor, init)

    return nothing
end

function create_system(model::DynamicNelsonSiegelModel, method::Symbol)
    # get dims
    (n,T_len)= method === :collapsed ? (3, size(model.y,2)) : size(model.y)

    # types
    Te= eltype(model.y)
    Tv= Vector{Te}
    Tm= Matrix{Te}
    Td= Diagonal{Te}
    TQ= cov(model.error_factor) isa Symmetric ? Symmetric{Te} : Td
    if method === :univariate || method === :collapsed
        TH= Td
    else
        TH= cov(model) isa Symmetric ? Symmetric{Te} : Td
    end

    # State space system
    sys= LinearTimeInvariant{Tm, Tm, Tm, Tv, Tv, TH, TQ, Tv, Tm}(n, 3, T_len)

    return sys
end

function fix_system!(sys::LinearTimeInvariant, model::DynamicNelsonSiegelModel, method::Symbol)
    # infer type
    T= eltype(model.y)

    if method === :multivariate || method === :woodbury
        sys.y.= model.y
    end
    if method === :univariate || method === :collapsed
        sys.H.diag.= one(T)
    end
    sys.d.= zero(T)
    sys.c.= zero(T)
end

function init_system!(sys::StateSpaceSystem, model::DynamicNelsonSiegelModel)
    T= eltype(model.y)
    # a
    sys.a1.= zero(T)
    # P
    sys.P1.= zero(T)
    @inbounds @fastmath for i in 1:3
        sys.P1[i,i]= one(T)
    end

    return nothing
end

# Estimation
# Log-likelihood
function loglik(filter::KalmanFilter, model::DynamicNelsonSiegelModel, method::Symbol)
    # Compute log-likelihood
    ll= loglik(filter)

    # Account for transformation
    if method === :univariate
        # get dims
        T_len= size(model.y,2)

        # Cholesky decomposition of H
        C= cholesky(cov(model))

        # Add Jacobian determinant term
        d,s= logabsdet(C.L)
        ll-= T_len * (d + log(s))
    elseif method === :collapsed
        # get dims
        (n,T_len)= size(model.y)

        Λ= loadings(model)  # loadings

        # Cholesky decomposition of H
        C= cholesky(cov(model))

        # Add Jacobian determinant term
        ll-= .5 * T_len * logdet(C)

        # Z'×H⁻¹×Z
        tmp_pn= transpose(Λ) * prec(model)
        tmp= tmp_pn * Λ
        # (Z'×H⁻¹×Z)⁻¹
        # perform pseudo inverse as Λ can contain zero columns
        pseudo= pinv(tmp)
        # Projection
        M= I - Λ * pseudo * tmp_pn

        # Add projected out term
        ll-= .5 * (n - model.r) * T_len * log(2*π)
        e= similar(model.y, n)
        @inbounds @fastmath for t in 1:T_len
            mul!(e, M, view(model.y,:,t))
            ll-= .5 * dot(e, prec(model), e)
        end
    end

    return ll
end

# Forecast
function reinstantiate(model::DynamicNelsonSiegelModel, h::Integer)
    # number of time series
    n= size(model.y,1)

    # expand
    y_f= hcat(model.y, fill(NaN, n, h))

    return DynamicNelsonSiegelModel(y_f, model.τ, model.λ, model.ϕ, model.error_obs, model.error_factor)
end

function forecast(model::DynamicNelsonSiegelModel, h::Integer)
    # reinstantiate
    forecast_model= reinstantiate(model, h)

    # State space system
    sys= create_system(forecast_model, :multivariate)
    fix_system!(sys, forecast_model, :multivariate)
    init_system!(sys, forecast_model)
    get_system!(sys, forecast_model, :multivariate)

    # forecasts
    f= forecast(sys, h)

    return f
end


# Shared methods between surface meshers
# * pushback
# * isoriented
# * volume (because isoriented could be made more robust)
# * subdivision

"Calculation of precise volume. Needed for convergence check"
function volume(points,faces)

    normal0 = [0,0,1]

    s = 0

    for tri in 1:size(faces,2)
        face = faces[:,tri]
        y1 = points[:,face[1]]
        y2 = points[:,face[2]]
        y3 = points[:,face[3]]

        normaly = cross(y2-y1,y3-y1)
        normaly /= norm(normaly)

        area = norm(cross(y2-y1,y3-y1))/2
        areaproj = dot(normaly,normal0)*area
        volume = dot(y1 + y2 + y3,normal0)/3*areaproj

        s += volume
    end

    return s

    # Calculate face normal
    # Calculate projected area
    # Calculate ordinary volume
    # Calculate volume between projected and real area
    # (+) if normal is outwards
end

using ForwardDiff               # Can be made optional

function pushback(sdist::Function,x::Array{Float64,1})
    for i in 1:100
        grad = ForwardDiff.gradient(sdist,x)
        x = x - sdist(x)*grad/norm(grad)^2
        # I allow the root to be only on gradient
    end
    return x
end

### Loading mesh by filename in the /preproc directgory
### At this point it should be unimportant if it is either generated with cgal or matlab

### Creation routines

function isoriented(points::Array{T,2},triangles::Array{S,2}) where {T <: AbstractFloat, S <: Integer}

    tri = triangles[:,1]
    p1 = points[:,tri[1]]
    p2 = points[:,tri[2]]
    p3 = points[:,tri[3]]

    a1 = p2 - p1
    a2 = p3 - p1

    normal = cross(a1,a2)
    normal /= sqrt(dot(normal,normal))

    direction = p1
    direction /= sqrt(dot(p1,p1))

    angle = acos(dot(normal,direction))

    if angle<pi/2
        return true
    else
        return false
    end
end

function subdivision(faces)
    edges = []

    for ti in 1:size(faces,2)
        v1,v2,v3 = faces[:,ti]

        v2<v1 || push!(edges,(v1,v2))
        v3<v2 || push!(edges,(v2,v3))
        v1<v3 || push!(edges,(v3,v1))
    end

    rfaces = Array(Int,3,size(faces,2)*4)
    N = maximum(faces)

    for ti in 1:size(faces,2)

        v1,v2,v3 = faces[:,ti]
        e3 = findfirst(edges,v2>v1 ? (v1,v2) : (v2,v1)) + N
        e1 = findfirst(edges,v3>v2 ? (v2,v3) : (v3,v2)) + N
        e2 = findfirst(edges,v1>v3 ? (v3,v1) : (v1,v3)) + N

        rfaces[:,4*ti - 3] = [v1,e3,e2]
        rfaces[:,4*ti - 2] = [v2,e1,e3]
        rfaces[:,4*ti - 1] = [v3,e2,e1]
        rfaces[:,4*ti] = [e1,e2,e3]
    end

    return rfaces
end

function subdivision(points,faces; n=1,method=:linear)
    if method==:linear
        rpoints,rfaces = LinearSubdivision(points,faces)
    elseif method==:paraboloid
        rpoints,rfaces = ParaboloidSubdivision(points,faces)
    end

    return rpoints,rfaces
end

function LinearSubdivision(points,faces)

    rfaces = subdivision(faces)
    rpoints = Array(Float64,3,maximum(rfaces))

    for ti in 1:size(faces,2)

        v1,v2,v3 = faces[:,ti]
        e1,e2,e3 = rfaces[:,4*ti]

        rpoints[:,v1] = points[:,v1]
        rpoints[:,v2] = points[:,v2]
        rpoints[:,v3] = points[:,v3]
        rpoints[:,e1] = (points[:,v2] + points[:,v3])/2
        rpoints[:,e2] = (points[:,v1] + points[:,v3])/2
        rpoints[:,e3] = (points[:,v1] + points[:,v2])/2
    end

    return rpoints,rfaces
end

function ParaboloidSubdivision(points,faces)

    normals = Array(Float64,size(points)...)
    NormalVectors!(normals,points,faces,i->FaceVRing(i,faces))

    rfaces = subdivision(faces)
    rpoints = Array(Float64,3,maximum(rfaces))

    vproperties = []
    for xkey in 1:size(points,2)

        x = points[:,xkey]
        nx = normals[:,xkey]

        vects = Float64[]
        for i in VertexVRing(xkey,faces)
            vi = points[:,i] - x
            vects = [vects..., vi...]
        end
        vects = reshape(vects,3,div(length(vects),3))

        R = eye(3)
        C,D,E = SurfaceGeometry.ZinchenkoDifertential!(vects,nx,R)

        push!(vproperties,(C,D,E,R))
    end

    z(x,y,C,D,E) = C*x^2 + D*x*y + E*y^2
    pushon(x,C,D,E,R) = ((xp,yp,zp) = R*x; inv(R)*[xp,yp,z(xp,yp,C,D,E)])

    for ti in 1:size(faces,2)
        v1,v2,v3 = faces[:,ti]
        e1,e2,e3 = rfaces[:,4*ti]

        rpoints[:,v1] = points[:,v1]
        rpoints[:,v2] = points[:,v2]
        rpoints[:,v3] = points[:,v3]

        ### It gets overwritten once
        xe1 = (points[:,v2] + points[:,v3])/2
        rpoints[:,e1] = xe1 + (pushon(xe1-points[:,v2],vproperties[v2]...) + pushon(xe1-points[:,v3],vproperties[v3]...))/2

        xe2 = (points[:,v3] + points[:,v1])/2
        rpoints[:,e2] = xe2 + (pushon(xe2-points[:,v1],vproperties[v1]...) + pushon(xe2-points[:,v3],vproperties[v3]...))/2

        xe3 = (points[:,v1] + points[:,v2])/2
        rpoints[:,e3] = xe3 + (pushon(xe3-points[:,v1],vproperties[v1]...) + pushon(xe3-points[:,v2],vproperties[v2]...))/2
    end

    return rpoints,rfaces
end

"Subdivision of surface and pushback"
function subdivision(points,faces,sdist::Function)
    rpoints,rfaces = LinearSubdivision(points,faces)

    for i in 1:size(rpoints,2)
        rpoints[:,i] = pushback(sdist,rpoints[:,i])
    end

    return rpoints,rfaces
end


export divergence

function divergence{T}(X::DirectProductSpace{T})
    x = Space{T}[divergence(s) for s in X.factors]
    DirectProductSpace(x)
end

"""
    divergence(x)

Compute the divergence of a finite element space.
"""
function divergence(x::Space)
    ref = refspace(x)
    geo = geometry(x)
    els = elements(geo)
    fns = x.fns
    dvs = similar(fns)
    for (i,fn) in enumerate(fns)
        dvs[i] = similar(fns[i])
        for (j,sh) in enumerate(fn)
            el = els[sh.cellid]
            dvs[i][j] = divergence(ref, sh, el)
        end
    end
    divergence(x, geo, dvs)
end

"""
    curl(X)

Compute the curl of a finite element basis. The resulting set of functions might be linearly dependent because of the kernel of the curl operator.
"""
function curl(x::Space)
    ref = refspace(x)
    geo = geometry(x)
    els = elements(geo)
    crl = similar(x.fns)
    for (i,fn) in enumerate(x.fns)
        crl[i] = similar(x.fns[i], 0)
        for (j,sh) in enumerate(fn)
            el = els[sh.cellid]
            append!(crl[i], curl(ref, sh, el))
        end
    end
    curl(x, geo, crl)
end


using StatsBase

@testset "presmoothing" begin
    # df = 1
    smtab = presmoothing(sgtabX, Equate.fml₁)
    m = smtab.table.scale'smtab.table.prob
    @test m ≈ mean(smtab.raw)
    # df = 2
    smtab = presmoothing(sgtabX, Equate.fml₂)
    m = smtab.table.scale'smtab.table.prob
    s = ((smtab.table.scale .- m).^2)'smtab.table.prob
    @test s ≈ var(smtab.raw; corrected = false)
    # df = 3
    smtab = presmoothing(sgtabX, Equate.fml₃)
    m = smtab.table.scale'smtab.table.prob
    s = ((smtab.table.scale .- m).^2)'smtab.table.prob
    sk = ((smtab.table.scale .- m).^3)'smtab.table.prob / sqrt(s)^3
    # @test sk ≈ skewness(smtab.raw)
    @test round(sk; digits = 5) ≈ round(skewness(smtab.raw); digits = 5)
    # df = 4
    smtab = presmoothing(sgtabX, Equate.fml₄)
    m = smtab.table.scale'smtab.table.prob
    s = ((smtab.table.scale .- m).^2)'smtab.table.prob
    ku = ((smtab.table.scale .- m).^4)'smtab.table.prob / s^2 -3
    @test ku ≈ kurtosis(smtab.raw)
end

@testset "KernelSmoothing (Gaussian)" begin
    fml = LogLinearFormula(6)
    smtab = presmoothing(sgtabX, fml)
    bbw = EstBandwidth(smtab)
    @test bbw.ls_success
    ktab = KernelSmoothing(smtab; hX = exp(bbw.minimizer[1]))
    @test typeof(ktab) == Equate.KernelFreqTab
end

@testset "Examine smoothing" begin
    fit = presmoothing(sgtabX, 6)
    @test fit isa DataFrame
end


using OPCSPs
using OPCSPs.MVNTools

r = [0, 10.1, 5, 5, 5, 0]
cov = Float64[ 0 0 0 0 0 0; 0 0 0 0 0 0; 0 0 2 2 -2 0; 0 0 2 4 0 0; 0 0 -2 0 4 0; 0 0 0 0 0 0]

dist = MVN(r,cov)

# BAD NOTATION delta is not what it usually is
M = 1000000
rewards = Array(Float64, M)
rng = MersenneTwister(1)
for i = 1:M
    delta = rand(dist)
    if delta[3] >= 5.0
        rewards[i] = delta[3] + delta[4]
    else
        rewards[i] = delta[3] + delta[5]
    end
end

#=
using Gadfly
p = plot(rewards, Geom.histogram)
draw(SVG("not_gaussian.svg", 6inch, 3inch), p)
=#

import PyPlot
PyPlot.plt[:clf]()
PyPlot.plt[:hist](rewards, 100)

# not gaussian


# ==============================================================================
# Trop{T} types

# Fake templates to make difference between Min-Plus and Max-Plus numbers
struct Min end
struct Max end
const MM = Union{Min, Max}

# Base class for Max-Plus and Min-Plus structures
struct Trop{T <: MM} <: Number
    v::Float64
end

# Max-Plus structure
const MP = Trop{Max}

# Min-Plus structure
const MI = Trop{Min}

# Sparse matrix (classic algebra) shorter name
const SparseMatrix{T,U} = SparseMatrixCSC{T,U}

# Sparse matrix (max-plus algebra) shorter name
const SpaTrop{T,U} = SparseMatrix{Trop{T},U}
const SpaMP{U} = SparseMatrix{MP,U}
const SpaMI{U} = SparseMatrix{MI,U}

# Sparse vector (max-plus algebra) shorter name
const SpvTrop{T,U} = SparseVector{Trop{T},U}
const SpvMP{U} = SparseVector{MP,U}
const SpvMI{U} = SparseVector{MI,U}

# Dense matrix (max-plus algebra) shorter name
const ArrTrop{T,N} = Array{Trop{T},N}
const ArrMP{N} = Array{MP,N}
const ArrMI{N} = Array{MI,N}
const AbsMatMP = AbstractMatrix{MP}

# Dense vector (max-plus algebra) shorter name
const VecTrop{T} = Vector{Trop{T}}
const VecMP = Vector{MP}
const VecMI = Vector{MI}


using DataFrames, CSV, StatsPlots, MLJ, CategoricalArrays

# prepare data
data = CSV.read(joinpath(@__DIR__, "..", "data", "iris.csv"))
rename!(data,
 "sepal.length" => :sepal_len ,
 "sepal.width" => :sepal_width,
 "petal.length" => :petal_len,
 "petal.width" => :petal_width,
)
categorical!(data, :variety)
X, y_labels = data[:, Not(:variety)], data[:, :variety]
# train_size = Int(floor(size(data)[1] * 0.75))
train, test = partition(eachindex(y_labels), 0.75, shuffle=true)

# visualizing
gr(size=(600, 600))
@df data corrplot(cols(1:4), gree=false)

# use NearestNeighbors model from MLJ
model = @load KNNClassifier pkg = "NearestNeighbors" verbosity = 1
model.K = 1
knn = machine(model, X, y_labels)
fit!(knn, rows=train)
yhat = predict(knn, X[test,:])
cross_entropy(yhat, y_labels[test]) |> mean


using JuMP, LinearAlgebra, Suppressor, StatsBase, CSV, DataFrames, Compat, Test, Random#, CPLEX
include("alt_min.jl")
include("convexpenalties.jl")

results_template = DataFrame(n=Int[], r=Int[], p=Real[], err_GD_msk=Real[], err_GD_msk_am=Real[], t_gd_msk=Real[], t_gd_msk_am=Real[])
results_exact=similar(results_template, 0)

for ARG in ARGS
    array_num = parse(Int, ARG)
    for n in [50*array_num] #1-12
        p=0.25
        r=5
        gamma=20/p
        M_L=randn(n, r)
        M_R=randn(n, r)
        M=M_L*M_R'

        numEntriesSampled=Int(round(p*n^2))
        coords=[repeat(1:n, inner=size(1:n,1)) repeat(1:n, outer=size(1:n,1))];
        entries=sample(1:n^2, Int(numEntriesSampled), replace=false)

        theIndices=zeros(n,n)
        for t=1:Int(numEntriesSampled)
            theIndices[coords[entries[t],1], coords[entries[t],2]]=1.0
        end

        k=r
        t_GD_msk=@elapsed X_GD,LB, UB, Y_relax=getMatrixFrobeniusNorm_greedy(coords, entries, M, n, numEntriesSampled, k, gamma)

        @show "SDP lower bound is: " LB
        u,sigma,v,=svd(X_GD)
        @show rnk_GD=sum(sigma.>1e-2);
        @show err_GD_msk = norm(X_GD.-M)/norm(M);

        t_GD_msk_am=@elapsed X_GD=alternatingminimization_matrixcompletion_solvesdprelax(coords, entries, M, n, numEntriesSampled, k, gamma)
        @show t_GD_msk_am
        u,sigma,v,=svd(X_GD)
        @show rnk_GD=sum(sigma.>1e-2);
        @show err_GD_am = norm(X_GD.-M)/norm(M);

        push!(results_exact, [n, r, p, err_GD_msk, err_GD_am,t_GD_msk, t_GD_msk_am])

    end
end

CSV.write("results_evaluatescalability.csv", results_exact, append=true)


using NPZ
using LinearAlgebra
using Combinatorics
using ProgressBars
using Printf: @sprintf
using OnlineStats

G1 = npzread("Kernels/G1.npy")
G2 = npzread("Kernels/G2.npy")
G3 = npzread("Kernels/G3.npy")

function d_oed_exhaust(X1::Matrix{T}, X2::Matrix{T}, X3::Matrix{T}, B, ϵ, λ) where T <: Number
    X = [X3 zeros(T, size(X3));
         zeros(T, size(X3)) X3;
         X1 X2]
    # Xn = maximum(sqrt.(sum(X.^2, dims=2)))
    Xn = sqrt.(sum(X.^2, dims=2)) #normalize all rows to be 1?
    X ./= Xn

    n = size(X3, 1)
    xdas = ones(T, size(X1,1))

    function Ml(μ, X, ϵ)
        M = ϵ*I + X'*(Diagonal(μ)*X)
        Mc = cholesky(Hermitian(M))
        l = μ.^2 .*diag(X*(Mc\X'))
        return (M, Mc, l)
    end

    Ml(μ) = Ml(μ, X, ϵ)

    function f(μ, λ)
        M, Mc, l = Ml(μ)
        return -logdet(M) + λ*maximum(l)
    end

    f(μ) = f(μ, λ)

    wrap_f(x) = f(vcat(x, x, xdas))

    combs = combinations(1:n,B)

    f_vals = zeros(length(combs))
    x = zeros(n, length(combs))

    iter = ProgressBar(enumerate(combs))
    stat = Extrema()
    for (i, comb) in iter
        x[comb,i] .= 1
        f_vals[i] = wrap_f(x[:,i])
        fit!(stat, f_vals[i])
        set_description(iter, "$(values(stat))")
    end

    return (f_vals, x)
end

f_vals, x = d_oed_exhaust(G1, G2, G3, 4, 1e-1, 10.0)

npzwrite("res2.npy", x[:,argmin(f_vals)])
# ;


struct DeivVecTag end

# J(f(x))*v
function auto_jacvec!(dy, f, x, v,
                      cache1 = ForwardDiff.Dual{DeivVecTag}.(x, v),
                      cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v))
    cache1 .= Dual{DeivVecTag}.(x, v)
    f(cache2,cache1)
    dy .= partials.(cache2, 1)
end
function auto_jacvec(f, x, v)
    partials.(f(Dual{DeivVecTag}.(x, v)), 1)
end

function num_jacvec!(dy,f,x,v,cache1 = similar(v),
                     cache2 = similar(v);
                     compute_f0 = true)
    compute_f0 && (f(cache1,x))
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    @. x += ϵ*v
    f(cache2,x)
    @. x -= ϵ*v
    @. dy = (cache2 - cache1)/ϵ
end

function num_jacvec(f,x,v,f0=nothing)
    f0 === nothing ? _f0 = f(x) : _f0 = f0
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(minimum(x)))
    (f(x.+ϵ.*v) .- _f0)./ϵ
end

function num_hesvec!(dy,f,x,v,
                     cache1 = similar(v),
                     cache2 = similar(v),
                     cache3 = similar(v))
    cache = FiniteDiff.GradientCache(v[1],cache1,Val{:central})
    g = let f=f,cache=cache
        (dx,x) -> FiniteDiff.finite_difference_gradient!(dx,f,x,cache)
    end
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    @. x += ϵ*v
    g(cache2,x)
    @. x -= 2ϵ*v
    g(cache3,x)
    @. dy = (cache2 - cache3)/(2ϵ)
end

function num_hesvec(f,x,v)
    g = (x) -> FiniteDiff.finite_difference_gradient(f,x)
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    x += ϵ*v
    gxp = g(x)
    x -= 2ϵ*v
    gxm = g(x)
    (gxp - gxm)/(2ϵ)
end

function numauto_hesvec!(dy,f,x,v,
                     cache = ForwardDiff.GradientConfig(f,v),
                     cache1 = similar(v),
                     cache2 = similar(v))
    g = let f=f,x=x,cache=cache
        g = (dx,x) -> ForwardDiff.gradient!(dx,f,x,cache)
    end
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    @. x += ϵ*v
    g(cache1,x)
    @. x -= 2ϵ*v
    g(cache2,x)
    @. dy = (cache1 - cache2)/(2ϵ)
end

function numauto_hesvec(f,x,v)
    g = (x) -> ForwardDiff.gradient(f,x)
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    x += ϵ*v
    gxp = g(x)
    x -= 2ϵ*v
    gxm = g(x)
    (gxp - gxm)/(2ϵ)
end

function autonum_hesvec!(dy,f,x,v,
                     cache1 = ForwardDiff.Dual{DeivVecTag}.(x, v),
                     cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v))
    cache = FiniteDiff.GradientCache(v[1],cache1,Val{:central})
    g = (dx,x) -> FiniteDiff.finite_difference_gradient!(dx,f,x,cache)
    cache1 .= Dual{DeivVecTag}.(x, v)
    g(cache2,cache1)
    dy .= partials.(cache2, 1)
end

function autonum_hesvec(f,x,v)
    g = (x) -> FiniteDiff.finite_difference_gradient(f,x)
    partials.(g(Dual{DeivVecTag}.(x, v)), 1)
end

function num_hesvecgrad!(dy,g,x,v,
                     cache2 = similar(v),
                     cache3 = similar(v))
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    @. x += ϵ*v
    g(cache2,x)
    @. x -= 2ϵ*v
    g(cache3,x)
    @. dy = (cache2 - cache3)/(2ϵ)
end

function num_hesvecgrad(g,x,v)
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    x += ϵ*v
    gxp = g(x)
    x -= 2ϵ*v
    gxm = g(x)
    (gxp - gxm)/(2ϵ)
end

function auto_hesvecgrad!(dy,g,x,v,
                     cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v),
                     cache3 = ForwardDiff.Dual{DeivVecTag}.(x, v))
    cache2 .= Dual{DeivVecTag}.(x, v)
    g(cache3,cache2)
    dy .= partials.(cache3, 1)
end

function auto_hesvecgrad(g,x,v)
    partials.(g(Dual{DeivVecTag}.(x, v)), 1)
end

### Operator Forms

struct JacVec{F,T1,T2,xType}
    f::F
    cache1::T1
    cache2::T2
    x::xType
    autodiff::Bool
end

function JacVec(f,x::AbstractArray;autodiff=true)
    if autodiff
        cache1 = ForwardDiff.Dual{DeivVecTag}.(x, x)
        cache2 = ForwardDiff.Dual{DeivVecTag}.(x, x)
    else
        cache1 = similar(x)
        cache2 = similar(x)
    end
    JacVec(f,cache1,cache2,x,autodiff)
end

Base.size(L::JacVec) = (length(L.cache1),length(L.cache1))
Base.size(L::JacVec,i::Int) = length(L.cache1)
Base.:*(L::JacVec,v::AbstractVector) = L.autodiff ? auto_jacvec(_x->L.f(_x),L.x,v) : num_jacvec(_x->L.f(_x),L.x,v)

function LinearAlgebra.mul!(dy::AbstractVector,L::JacVec,v::AbstractVector)
    if L.autodiff
        auto_jacvec!(dy,(_y,_x)->L.f(_y,_x),L.x,v,L.cache1,L.cache2)
    else
        num_jacvec!(dy,(_y,_x)->L.f(_y,_x),L.x,v,L.cache1,L.cache2)
    end
end

struct HesVec{F,T1,T2,xType}
    f::F
    cache1::T1
    cache2::T2
    cache3::T2
    x::xType
    autodiff::Bool
end

function HesVec(f,x::AbstractArray;autodiff=true)
    if autodiff
        cache1 = ForwardDiff.GradientConfig(f,x)
        cache2 = similar(x)
        cache3 = similar(x)
    else
        cache1 = similar(x)
        cache2 = similar(x)
        cache3 = similar(x)
    end
    HesVec(f,cache1,cache2,cache3,x,autodiff)
end

Base.size(L::HesVec) = (length(L.cache2),length(L.cache2))
Base.size(L::HesVec,i::Int) = length(L.cache2)
Base.:*(L::HesVec,v::AbstractVector) = L.autodiff ? numauto_hesvec(L.f,L.x,v) : num_hesvec(L.f,L.x,v)

function LinearAlgebra.mul!(dy::AbstractVector,L::HesVec,v::AbstractVector)
    if L.autodiff
        numauto_hesvec!(dy,L.f,L.x,v,L.cache1,L.cache2,L.cache3)
    else
        num_hesvec!(dy,L.f,L.x,v,L.cache1,L.cache2,L.cache3)
    end
end

struct HesVecGrad{G,T1,T2,uType}
    g::G
    cache1::T1
    cache2::T2
    x::uType
    autodiff::Bool
end

function HesVecGrad(g,x::AbstractArray;autodiff=false)
    if autodiff
        cache1 = ForwardDiff.Dual{DeivVecTag}.(x, x)
        cache2 = ForwardDiff.Dual{DeivVecTag}.(x, x)
    else
        cache1 = similar(x)
        cache2 = similar(x)
    end
    HesVecGrad(g,cache1,cache2,x,autodiff)
end

Base.size(L::HesVecGrad) = (length(L.cache2),length(L.cache2))
Base.size(L::HesVecGrad,i::Int) = length(L.cache2)
Base.:*(L::HesVecGrad,v::AbstractVector) = L.autodiff ? auto_hesvecgrad(L.g,L.x,v) : num_hesvecgrad(L.g,L.x,v)

function LinearAlgebra.mul!(dy::AbstractVector,L::HesVecGrad,v::AbstractVector)
    if L.autodiff
        auto_hesvecgrad!(dy,L.g,L.x,v,L.cache1,L.cache2)
    else
        num_hesvecgrad!(dy,L.g,L.x,v,L.cache1,L.cache2)
    end
end


using AnimatedPlots

log_graph = StaticGraph(log, color=SFML.magenta)
exponential_graph = StaticGraph(x -> 2^x, color=SFML.cyan)

plot(log_graph)
plot(exponential_graph)

waitfor(current_window())


module TORA

using NeuralNetworkAnalysis, LaTeXStrings
using NeuralNetworkAnalysis: UniformAdditivePostprocessing
import Plots, DifferentialEquations
using Plots: plot, plot!, xlims!, ylims!, savefig, font, Measures.mm

# problem

@taylorize function TORA!(dx, x, p, t)
    x₁, x₂, x₃, x₄, u = x

    aux = 0.1 * sin(x₃)
    dx[1] = x₂
    dx[2] = -x₁ + aux
    dx[3] = x₄
    dx[4] = u
    dx[5] = zero(u)
    return dx
end

controller = read_nnet_sherlock(@modelpath "" "controller_TORA_sherlock")

X₀ = Hyperrectangle(low=[0.6, -0.7, -0.4, 0.5], high=[0.7, -0.6, -0.3, 0.6])
U = ZeroSet(1)
vars_idx = Dict(:state_vars=>1:4, :control_vars=>5)
ivp = @ivp(x' = TORA!(x), dim: 5, x(0) ∈ X₀ × U)
control_postprocessing = UniformAdditivePostprocessing(-10.0)

period = 1.0
k = 20
T = k * period
T_warmup = 2 * period

prob = ControlledPlant(ivp, controller, vars_idx, period;
                       postprocessing=control_postprocessing)

safe_states = cartesian_product(BallInf(zeros(4), 2.0), Universe(1))
predicate = X -> X ⊆ safe_states

# simulation

println("simulation")
res = @timed simulate(prob, T=T, trajectories=10, include_vertices=true)
sim = res.value
print_timed(res)

# reachability analysis

alg = TMJets(abstol=1e-10, orderT=8, orderQ=3)
alg_nn = Ai2()
splitter = BoxSplitter([4, 4, 3, 5])

function benchmark(; T=T, silent::Bool=false)
    silent || println("flowpipe construction")
    res = @timed solve(prob, T=T, alg_nn=alg_nn, alg=alg, splitter=splitter)
    sol = res.value
    silent || print_timed(res)

    silent || println("property checking")
    solh = overapproximate(sol, Hyperrectangle)
    res = @timed predicate(solh)
    silent || print_timed(res)
    if res.value
        silent || println("The property is satisfied.")
    else
        silent || println("The property may be violated.")
    end

    return solh
end

benchmark(T=T_warmup, silent=true)
res = @timed benchmark()
sol = res.value

# plotting

vars = (1, 2)
fig = plot(xlab=L"x_1", ylab=L"x_2",
           legendfontsize=15,
           tickfont=font(25, "Times"),
           guidefontsize=25,
           xguidefont=font(25, "Times"),
           yguidefont=font(25, "Times"),
           bottom_margin=0mm, left_margin=0mm, right_margin=3mm, top_margin=2mm,
           size=(500, 450))
plot!(project(X₀, vars), c=:blue, lab=L"X_0", alpha=0.6)
plot!(fig, sol, vars=vars, color=:orange, lab="")
plot!(project(X₀, vars), c=:blue, lab="", alpha=0.6)
plot_simulation!(fig, sim; vars=vars, color=:black, lab="")
xlims!(-1, 1)
ylims!(-1.5, 1)
savefig("tora_x1x2_juliareach.png")

vars = (3, 4)
fig = plot(xlab=L"x_3", ylab=L"x_4",
           legendfontsize=15,
           tickfont=font(25, "Times"),
           guidefontsize=25,
           xguidefont=font(25, "Times"),
           yguidefont=font(25, "Times"),
           bottom_margin=0mm, left_margin=0mm, right_margin=0mm, top_margin=2mm,
           size=(500, 450))
# plot!(project(X₀, vars), c=:blue, lab=L"X_0", alpha=0.6)
plot!(fig, sol, vars=vars, color=:orange, lab="")
plot!(project(X₀, vars), c=:blue, lab="", alpha=0.6)
plot_simulation!(fig, sim; vars=vars, color=:black, lab="")
xlims!(-2, 2)
ylims!(-2, 2)
savefig("tora_x3x4_juliareach.png")

end
nothing


import Base: sin, cos, tan, cot, sec, csc, atan

import Base: sinh, cosh, tanh, coth, sech, csch

import Base: abs, exp, log, zero, one

import Base: +, -, *, /, ^

struct quadcomplex
    i1
    i2
end

function mat(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    return [i1 -i2; i2 i1]
end

function mat2quadcomplex(matrix)
    out1, out2 = matrix[1, 1], matrix[2, 1]
    return quadcomplex(out1, out2)
end

## Basic operators: +, -, *, /, ^

+(quad1::quadcomplex, quad2::quadcomplex) = mat2quadcomplex(mat(quad1)+mat(quad2))
+(n, quad::quadcomplex) = mat2quadcomplex(n.+mat(quad))
+(quad::quadcomplex, n) = +(n, quad)

-(quad1::quadcomplex, quad2::quadcomplex) = mat2quadcomplex(mat(quad1)-mat(quad2))
-(n, quad::quadcomplex) = mat2quadcomplex(n.-mat(quad))
-(quad::quadcomplex) = mat2quadcomplex(-mat(quad))

*(quad1::quadcomplex, quad2::quadcomplex) = mat2quadcomplex(mat(quad1)*mat(quad2))
*(n, quad::quadcomplex) = mat2quadcomplex(n*mat(quad))
*(quad::quadcomplex, n) = *(n, quad)

/(quad1::quadcomplex, quad2::quadcomplex) = mat2quadcomplex(mat(quad1)/mat(quad2))
/(n, quad::quadcomplex) = n*inv(quad)

^(quad::quadcomplex, power) = mat2quadcomplex(mat(quad)^power)

function zero(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    return quadcomplex(zero(i1), zero(i2))
end

function one(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    return quadcomplex(one(i1), one(i2))
end

## Basic functions: abs, exp, log

function abs(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    return sqrt(i1^2+i2^2)
end

function exp(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    out1 = exp(i1)*cos(i2)
    out2 = exp(i1)*sin(i2)
    return quadcomplex(out1, out2)
end

## Basic math functions: sin, cos, tan, cot, sec, csc

function sin(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    out1 = cosh(i2)*sin(i1)
    out2 = sinh(i2)*cos(i1)
    return quadcomplex(out1, out2)
end

function cos(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    out1 = cosh(i2)*cos(i1)
    out2 = -sinh(i2)*sin(i1)
    return quadcomplex(out1, out2)
end

tan(quad::quadcomplex) = sin(quad)/cos(quad)
cot(quad::quadcomplex) = cos(quad)/sin(quad)
sec(quad::quadcomplex) = 1/cos(quad)
csc(quad::quadcomplex) = 1/sin(quad)

## Basic hyperbolic functions: sinh, cosh, tanh, coth, sech, csch

function sinh(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    out1 = cosh(i1)*cos(i2)
    out2 = sinh(i1)*sin(i2)
    return quadcomplex(out1, out2)
end

function cosh(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    out1 = sinh(i1)*cos(i2)
    out2 = cosh(i1)*sin(i2)
    return quadcomplex(out1, out2)
end

tanh(quad::quadcomplex) = sinh(quad)/cosh(quad)
coth(quad::quadcomplex) = cosh(quad)/sinh(quad)
sech(quad::quadcomplex) = 1/cosh(quad)
csch(quad::quadcomplex) = 1/sinh(quad)

function image4(quad::quadcomplex)
    i1, i2 = quad.i1, quad.i2
    return image3(i2)
end

function quadderivative(f, point, h)
    result = image4(f(quadcomplex(tricomplex(bicomplex(point+im*h, h), h), h)))/h^4
    return result
end


using Convex
Random.seed!(2)

m = 40;  n = 50
A = randn(m, n); b = randn(m, 1)
x = Variable(n)
problem = minimize(sumsquares(A * x - b), [x >= 0])

ϵ = 1e-8
opt = 12.38418747141913
solver =  LineSearchWrapper(GAP(0.5, 1.0, 1.0, eps=ϵ, verbose=1, checki=1, max_iters=10000), lsinterval=100)
solve!(problem, solver)


"""
    Module for functions shared by different parameterizations.
"""
module Common

import Thermodynamics
import CLIMAParameters
import CLIMAParameters.Planet
import CLIMAParameters.Atmos.Microphysics

const TD = Thermodynamics
const CP = CLIMAParameters
const CP_planet = CLIMAParameters.Planet
const CP_micro = CLIMAParameters.Atmos.Microphysics
const APS = CP.AbstractParameterSet

export G_func

"""
    G_func(param_set, T, Liquid())
    G_func(param_set, T, Ice())

 - `param_set` - abstract set with earth parameters
 - `T` - air temperature
 - `Liquid()`, `Ice()` - liquid or ice phase to dispatch over.

Utility function combining thermal conductivity and vapor diffusivity effects.
"""
function G_func(param_set::APS, T::FT, ::TD.Liquid) where {FT <: Real}

    _K_therm::FT = CP_micro.K_therm(param_set)
    _R_v::FT = CP_planet.R_v(param_set)
    _D_vapor::FT = CP_micro.D_vapor(param_set)

    L = TD.latent_heat_vapor(param_set, T)
    p_vs = TD.saturation_vapor_pressure(param_set, T, TD.Liquid())

    return FT(1) / (
        L / _K_therm / T * (L / _R_v / T - FT(1)) + _R_v * T / _D_vapor / p_vs
    )
end
function G_func(param_set::APS, T::FT, ::TD.Ice) where {FT <: Real}

    _K_therm::FT = CP_micro.K_therm(param_set)
    _R_v::FT = CP_planet.R_v(param_set)
    _D_vapor::FT = CP_micro.D_vapor(param_set)

    L = TD.latent_heat_sublim(param_set, T)
    p_vs = TD.saturation_vapor_pressure(param_set, T, TD.Ice())

    return FT(1) / (
        L / _K_therm / T * (L / _R_v / T - FT(1)) + _R_v * T / _D_vapor / p_vs
    )
end

end


import Theta: accola_chars_r, accola_chars_sr

@testset "Accola" begin
    @test Set(accola_chars_r([[1,1],[0,1]], [[1,0],[0,0]], [[0,1],[0,0]], [[0,0],[1,0]])) == Set([[[1,1],[0,1]], [[0,1],[0,1]], [[1,0],[0,1]], [[1,1],[1,1]], [[0,0],[0,1]], [[0,1],[1,1]], [[1,0],[1,1]], [[0,0],[1,1]]])
    @test Set(accola_chars_sr([[[1,0,0,0,0],[0,0,0,0,0]], [[0,1,0,0,0],[1,0,0,0,0]], [[0,0,1,0,0],[1,1,0,0,0]], [[0,0,0,1,0],[1,1,1,0,0]], [[0,0,0,0,1],[1,1,1,1,0]], [[1,0,0,0,0],[0,1,1,1,1]], [[0,1,0,0,0],[0,0,1,1,1]], [[0,0,1,0,0],[0,0,0,1,1]], [[0,0,0,1,0],[0,0,0,0,1]], [[1,1,1,1,0],[0,1,0,1,0]], [[1,1,1,1,0],[1,0,1,0,1]]], [[1,2,3,4], [5,6,7,8], [5,6,9,10], [5,7,9,11]])[1]) ==  Set([[[1,0,0,0,0], [0,0,0,0,0]], [[0,1,1,0,1], [1,0,1,0,1]], [[1,1,1,0,1], [1,1,0,1,0]], [[0,0,1,0,1], [0,1,1,0,1]], [[0,0,0,0,0], [0,1,1,1,1]], [[1,1,0,0,0], [1,1,0,0,0]], [[0,1,0,0,0], [1,0,1,1,1]], [[1,0,1,0,1], [0,0,0,1,0]]])
    @test Set(accola_chars()[1][1]) == Set([[[1,0,0,0,0], [0,0,0,0,0]], [[0,1,1,0,1], [1,0,1,0,1]], [[1,1,1,0,1], [1,1,0,1,0]], [[0,0,1,0,1], [0,1,1,0,1]], [[0,0,0,0,0], [0,1,1,1,1]], [[1,1,0,0,0], [1,1,0,0,0]], [[0,1,0,0,0], [1,0,1,1,1]], [[1,0,1,0,1], [0,0,0,1,0]]])
    @test begin
        τ = [0.17012929+3.36021399im 0.42793573+1.65263414im 0.55503983+2.1393021im 0.49182605+2.17345553im 0.93865305+1.81009273im; 0.42793573+1.65263414im 0.87402978+1.33203047im 0.53489381+0.68238397im 0.57601957+1.20509707im 0.50028888+1.11605622im; 0.55503983+2.1393021im 0.53489381+0.68238397im 0.21882206+2.28319153im 0.74421992+1.51352273im 0.41855281+0.87400937im; 0.49182605+2.17345553im 0.57601957+1.20509707im 0.74421992+1.51352273im 0.51070570+1.62305859im 0.52380319+1.15073893im; 0.93865305+1.81009273im 0.50028888+1.11605622im 0.41855281+0.87400937im 0.52380319+1.15073893im 0.79948032+1.38925759im];
        accola(τ) == accola(RiemannMatrix(τ))
        end
end


include("header.jl")

using Statistics
meanabs(x)=mean(abs,x)
meanabs2(x)=mean(abs2,x)

@testset "statistics" begin
    o = (:delta=>0.0001,:atol=>0.01,:rtol=>0.01)
    for T in (Float32,Float64)
        @test gradcheck(mean, randn(T,2,3); o...)
        @test gradcheck(mean, randn(T,2,3); kw=[:dims=>1], o...)
        @test gradcheck(mean, randn(T,2,3); kw=[:dims=>(1,2)], o...)
        @test gradcheck(meanabs, randn(T,2,3); o...)
        @test gradcheck(meanabs2, randn(T,2,3); o...)
        @test gradcheck(var, randn(T,2,3); o...)
        @test gradcheck(var, randn(T,2,3); kw=[:dims=>1], o...)
        @test gradcheck(var, randn(T,2,3); kw=[:dims=>(1,2)], o...)
        @test gradcheck(std, randn(T,2,3); o...)
        @test gradcheck(std, randn(T,2,3); kw=[:dims=>1], o...)
        @test gradcheck(std, randn(T,2,3); kw=[:dims=>(1,2)], o...)
    end
end


# Function for calculating Fermi surface in 2D
function getFermiSurface2D(
            unitcell:: U,
            bondInteractionMatrix:: HB,
            N_points::Int64,
            kappa::Float64;
            fermi_energy::Float64=0.0,
            enforce_hermitian::Bool=false,
            epsilon::Float64=1e-10,
            epsilon_k::Float64=1e-10,
            bounds_lower::Array{Float64,1}=-2*pi.*ones(2),
            bounds_upper::Array{Float64,1}=2*pi.*ones(2),
            refold_to_first_BZ::Bool=true
        ) where {L,NB,S,B<:AbstractBond{L,NB},U<:AbstractUnitcell{S,B},HB<:AbstractBondHamiltonian{L,NS} where {NS}}

    # list of k values that contribute to the BZ
    k_values = zeros(Float64, N_points, 2)

    LTbondInteractionMatrix = newHamiltonian(unitcell,bondInteractionMatrix)

    # funtion of energy
    function energyfunction(kvector)
        eigenvalues = eigvals(getMatrixAtK(LTbondInteractionMatrix,kvector,kappa)) .- fermi_energy
        eigenvalues = eigenvalues .* eigenvalues
        return minimum(eigenvalues)
    end

    # the current search index
    index = 1
    # search until there are enough points
    while index <= N_points

        # find a suitable starting point for the Newton algorithm
        k = Float64[rand(), rand()]
        for j in 1:length(k)
            k[j] = k[j] * bounds_lower[j] + (1-k[j]) * bounds_upper[j]
        end

        # start with the initial energy
        e0 = energyfunction(k)
        # iterate i over 100 newton steps (maximum)
        for i in 1:100
            # check if the energy is already converged
            if e0 < epsilon
                # save the k vector
                k_values[index,:] = k
                # increment the index
                index = index+1
                # break the newton loop
                break
            end
            # the current energy
            H_0 = e0
            # the gradient of the energy
            H_eps =
            [
                energyfunction(k .+ [epsilon_k, 0]),
                energyfunction(k .+ [0, epsilon_k])
            ]
            dH = (H_eps - H_0) ./ epsilon_k
            # absolute value of the gradient
            dHdH = dot(dH, dH)
            # break the Newton loop if the gradient diverges or is flattened too much
            if abs(dHdH) < 1e-20 || abs(dHdH) > 1e20
                break
            end
            # increment k
            dk = dH * (H_0 / dHdH)
            k -= dk
            # calculate a new energy
            e0 = energyfunction(k)
        end
    end

    # if refolding is enabled, refold all points
    if refold_to_first_BZ
        # get the lattice vectors
        a1 = latticeVectors(unitcell)[1]
        a2 = latticeVectors(unitcell)[2]
        # get the reciprocal lattice vectors
        b1 = [a2[2], -a2[1]]
        b2 = [a1[2], -a1[1]]
        # normalize the vectors
        b1 .*= 2*pi/sum(b1.*a1)
        b2 .*= 2*pi/sum(b2.*a2)

        # iterate over all k points
        for index in 1:N_points
            # assumend to be refolded to begin with
            refolded = true
            # iterate while still refolding possible
            while refolded
                # new iteration, not refolded
                refolded = false
                # try all possible refoldings
                for r in [b1, b2, -b1, -b2, b1+b2, -b2-b1, b1-b2, -b2+b1]
                    if sum((k_values[index,:].-r).*(k_values[index,:].-r)) < sum((k_values[index,:]).*(k_values[index,:]))
                        k_values[index,:] .-= r
                        refolded=true
                    end
                end
            end
        end
    end

    # return the list
    return k_values

end

# Function for calculating Fermi surface in 3D
function getFermiSurface3D(
            unitcell:: U,
            bondInteractionMatrix:: HB,
            N_points::Int64,
            kappa::Float64;
            fermi_energy::Float64=0.0,
            enforce_hermitian::Bool=false,
            epsilon::Float64=1e-10,
            epsilon_k::Float64=1e-10,
            bounds_lower::Array{Float64,1}= -2 .*pi .*ones(3),
            bounds_upper::Array{Float64,1}=2 .* pi .*ones(3),
            refold_to_first_BZ::Bool=true
        ) where {L,NB,S,B<:AbstractBond{L,NB},U<:AbstractUnitcell{S,B},HB<:AbstractBondHamiltonian{L,NS} where {NS}}

    # list of k values that contribute to the BZ
    k_values = zeros(Float64, N_points, 3)

    LTbondInteractionMatrix = newHamiltonian(unitcell,bondInteractionMatrix)
    # funtion of energy
    function energyfunction(kvector)
        eigenvalues = eigvals(getMatrixAtK(LTbondInteractionMatrix,kvector,kappa)) .- fermi_energy
        eigenvalues = eigenvalues .* eigenvalues
        return minimum(eigenvalues)
        #return eigmin(getMatrixAtK(LTbondInteractionMatrix,kvector))
    end

    # the current search index
    index = 1
    # search until there are enough points
    while index <= N_points

        # find a suitable starting point for the Newton algorithm
        k = Float64[rand(), rand(), rand()]
        for j in 1:length(k)
            k[j] = k[j] * bounds_lower[j] + (1-k[j]) * bounds_upper[j]
        end

        # start with the initial energy
        e0 = energyfunction(k)
        # iterate i over 100 newton steps (maximum)
        for i in 1:1000
            # check if the energy is already converged
            if e0 < epsilon
                # save the k vector
                k_values[index,:] = k
                # increment the index
                index = index+1
                # break the newton loop
                break
            end
            # the current energy
            H_0 = e0
            # the gradient of the energy
            H_eps =
            [
                energyfunction(k .+ [epsilon_k, 0, 0]),
                energyfunction(k .+ [0, epsilon_k, 0]),
                energyfunction(k .+ [0, 0, epsilon_k])
            ]
            dH = (H_eps .- H_0) ./ epsilon_k
            # absolute value of the gradient
            dHdH = dot(dH, dH)
            # break the Newton loop if the gradient diverges or is flattened too much
            if abs(dHdH) < 1e-20 || abs(dHdH) > 1e20
                break
            end
            # increment k
            dk = dH * (H_0 / dHdH)
            k -= dk
            # calculate a new energy
            e0 = energyfunction(k)
        end
    end

    # if refolding is enabled, refold all points
    if refold_to_first_BZ
        # get the lattice vectors
        a1 = latticeVectors(unitcell)[1]
        a2 = latticeVectors(unitcell)[2]
        a3 = latticeVectors(unitcell)[3]
        # get the reciprocal lattice vectors
        b1 = cross(a2, a3)
        b2 = cross(a3, a1)
        b3 = cross(a1, a2)
        # normalize the vectors
        b1 .*= 2*pi/sum(b1.*a1)
        b2 .*= 2*pi/sum(b2.*a2)
        b3 .*= 2*pi/sum(b3.*a3)
        # put together refolding
        refoldings = Array{Float64,1}[]
        for i in [-1,0,1]
        for j in [-1,0,1]
        for l in [-1,0,1]
            refolding = i*b1 + j*b2 + l*b3
            if sum(abs.(refolding))>1e-8
                push!(refoldings, refolding)
            end
        end
        end
        end
        # iterate over all k points
        for index in 1:N_points
            # assumend to be refolded to begin with
            refolded = true
            # iterate while still refolding possible
            while refolded
                # new iteration, not refolded
                refolded = false
                # try all possible refoldings
                for r in refoldings

                    if sum((k_values[index,:].-r).*(k_values[index,:].-r)) < sum((k_values[index,:]).*(k_values[index,:]))
                        k_values[index,:] .-= r
                        refolded=true
                    end
                end
            end
        end
    end

    # return the list
    return k_values

end

# obtain Fermi surface

"""
    getFermiSurface(
                unitcell::Unitcell,
                N_points::Int64
             [; fermi_energy::Float64=0.0,
                enforce_hermitian::Bool=false,
                epsilon::Float64=1e-10,
                epsilon_k::Float64=1e-10,
                bounds_lower::Array{Float64,1}=-2*pi.*ones(4),
                bounds_upper::Array{Float64,1}=2*pi.*ones(4),
                refold_to_first_BZ::Bool=true ]
            )
calculates `N_points` points which belong to the Fermi surface of the tight-binding model given by
the `Unitcell` object `unitcell`. The Fermi energy can be adjusted by using `fermi_energy=...`.
The points are located using Newton's method with random starting positions. The procedure can be
modified by using some of the optional keywords:
- `enforce_hermitian` determines if the matrix used for energy calculation is made hermitian by construction
- `epsilon` energy threshold for values to be considered close to the Fermi energy
- `epsilon_k` small k increment that is used for calculation of derivatives
- `bounds_upper` / `bounds_lower` the bounds between which the random starting location for the Newton method is searched
After the search has been completed, the values will optionally be refolded into the first brillouin zone by choosing `refold_to_first_BZ=true`.
# Examples
```julia-repl
julia> fermi_surace = getFermiSurface(unitcell, 100)    # 2D unitcell
100x2 Array{Float64}
...
julia> fermi_surace = getFermiSurface(unitcell, 100)    # 3D unitcell
100x3 Array{Float64}
...
julia> fermi_surace = getFermiSurface(unitcell, 100, fermi_energy=1.0)
100x2 Array{Float64}
...
julia> fermi_surace = getFermiSurface(unitcell, 100, refold_to_first_BZ=false)
100x2 Array{Float64}
...
```
"""

function getFermiSurface(
            unitcell:: U,
            bondInteractionMatrix:: HB,
            N_points::Int64,
            kappa::Float64;
            fermi_energy::Float64=0.0,
            enforce_hermitian::Bool=false,
            epsilon::Float64=1e-10,
            epsilon_k::Float64=1e-10,
            bounds_lower::Array{Float64,1}=-2*pi.*ones(3),
            bounds_upper::Array{Float64,1}=2*pi.*ones(3),
            refold_to_first_BZ::Bool=true
        ) where {L,NB,S,B<:AbstractBond{L,NB},U<:AbstractUnitcell{S,B},HB<:AbstractBondHamiltonian{L,NS} where {NS}}

    # check if the unitcell can be put in either method
    if length(sites(unitcell)[1].point) != length(latticeVectors(unitcell))
       println("Unitcell has not the same number of lattice vectors as dimensions")
        return zeros(1,1)
    end

    # check which function to pass to
    if length(latticeVectors(unitcell)) == 2
        # 2D case
        return getFermiSurface2D(
                unitcell,
                bondInteractionMatrix,
                N_points,
                kappa,
                fermi_energy=fermi_energy,
                enforce_hermitian=enforce_hermitian,
                epsilon=epsilon,
                epsilon_k=epsilon_k,
                bounds_lower=bounds_lower,
                bounds_upper=bounds_upper,
                refold_to_first_BZ=refold_to_first_BZ
            )

   elseif length(latticeVectors(unitcell)) == 3
        # 3D case
        return getFermiSurface3D(
                unitcell,
                bondInteractionMatrix,
                N_points,
                kappa,
                fermi_energy=fermi_energy,
                enforce_hermitian=enforce_hermitian,
                epsilon=epsilon,
                epsilon_k=epsilon_k,
                bounds_lower=bounds_lower,
                bounds_upper=bounds_upper,
                refold_to_first_BZ=refold_to_first_BZ
            )
    end
end

export getFermiSurface


using GeometryTypes, Test
function f(G, x)
    fx = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

    if !(G == nothing)
        G1 = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
        G2 = 200.0 * (x[2] - x[1]^2)
        gx = Point(G1,G2)

        return fx, gx
    else
        return fx
    end
end
f_obj = OnceDiffed(f)
minimize(f_obj, Point(3.0,3.0), GradientDescent(Inverse()), OptOptions())
minimize(f_obj, Point(3.0,3.0), BFGS(Inverse()), OptOptions())
minimize(f_obj, Point(3.0,3.0), DFP(Inverse()), OptOptions())
minimize(f_obj, Point(3.0,3.0), SR1(Inverse()), OptOptions())

minimize(f_obj, Point(3.0,3.0), GradientDescent(Direct()), OptOptions())
minimize(f_obj, Point(3.0,3.0), BFGS(Direct()), OptOptions())
minimize(f_obj, Point(3.0,3.0), DFP(Direct()), OptOptions())
minimize(f_obj, Point(3.0,3.0), SR1(Direct()), OptOptions())


if !isdefined(Main, :CustomStencil)
    include("../module/CustomStencil.jl")
    using Main.CustomStencil
end
include("misc.jl")
include("cpu_stencils.jl")

## Stencil Definition
coefs = Float64.([1; 0.5; 0.25; 0.125; 0.06])
star_stencil = @def_stencil_expression c[0]D[x,y,z] + @sum(i, 1, 4, c[i]*(
            D[x+i,y,z] + D[x,y+i,z] + D[x,y,z+i] +
            D[x-i,y,z] + D[x,y-i,z] + D[x,y,z-i]))
st_def = CreateStencilDefinition(star_stencil, coefs)
st_inst = NewStencilInstance(st_def)

## Input Data size Definition
radius = 4
nx = 8
ny = 8
nz = 8
dx = 1<<(nx)
dy = 1<<(ny)
dz = 1<<(nz)

t_steps = 4
## GPU Implementation
original_data = CreateData(dx,dy,dz)
data = PadData(radius, original_data)
dout = ApplyStencil(st_inst, data, t_steps)
gpu_out = Array(dout)


#warlpiri-properties.jl - uses warlpiri numbers to test valid properties.

@testset "warlpiri-properties" begin
    @test isempty.(wmatrix) == map(x -> (x == (WP0001 → WP0000)) || (x == (WP0000 → WP1111)), wmatrix)

    const warlpiri_allreals =
        [WP0010 → WP0001, WP0011 → WP0010, WP0100 → WP0011, WP0101 → WP0100,
         WP0110 → WP0101, WP0111 → WP0110, WP1000 → WP0111, WP1001 → WP1000,
         WP1010 → WP1001, WP1011 → WP1010, WP1100 → WP1011, WP1101 → WP1100,
         WP1110 → WP1101, WP1111 → WP1110]

    @test SigmoidNumbers.isallreals.(wmatrix) == map((x) -> x in warlpiri_allreals, wmatrix)

    #roundsinf test          ->  0000  0001  0010  0011  0100  0101  0110  0111  1000  1001  1010  1011  1100  1101  1110  1111
    const warlpiri_roundsinf = [false false false false false false false false  true  true  true  true  true  true  true false  # 0000 ->
                                false false false false false false false false  true  true  true  true  true  true  true  true  # 0001 ->
                                 true  true false false false false false false  true  true  true  true  true  true  true  true  # 0010 ->
                                 true  true  true false false false false false  true  true  true  true  true  true  true  true  # 0011 ->
                                 true  true  true  true false false false false  true  true  true  true  true  true  true  true  # 0100 ->
                                 true  true  true  true  true false false false  true  true  true  true  true  true  true  true  # 0101 ->
                                 true  true  true  true  true  true false false  true  true  true  true  true  true  true  true  # 0110 ->
                                 true  true  true  true  true  true  true false  true  true  true  true  true  true  true  true  # 0111 ->
                                 true  true  true  true  true  true  true  true  true  true  true  true  true  true  true  true  # 1000 ->
                                false false false false false false false false  true false false false false false false false  # 1001 ->
                                false false false false false false false false  true  true false false false false false false  # 1010 ->
                                false false false false false false false false  true  true  true false false false false false  # 1011 ->
                                false false false false false false false false  true  true  true  true false false false false  # 1100 ->
                                false false false false false false false false  true  true  true  true  true false false false  # 1101 ->
                                false false false false false false false false  true  true  true  true  true  true false false  # 1110 ->
                                false false false false false false false false  true  true  true  true  true  true  true false  # 1111 ->
    ]

    @test SigmoidNumbers.roundsinf.(wmatrix) == warlpiri_roundsinf

    #containszero test       ->  0000  0001  0010  0011  0100  0101  0110  0111  1000  1001  1010  1011  1100  1101  1110  1111
    const warlpiri_containszero = [ true  true  true  true  true  true  true  true  true  true  true  true  true  true  true false  # 0000 ->
                                   false false false false false false false false false false false false false false false false  # 0001 ->
                                    true  true false false false false false false false false false false false false false false  # 0010 ->
                                    true  true  true false false false false false false false false false false false false false  # 0011 ->
                                    true  true  true  true false false false false false false false false false false false false  # 0100 ->
                                    true  true  true  true  true false false false false false false false false false false false  # 0101 ->
                                    true  true  true  true  true  true false false false false false false false false false false  # 0110 ->
                                    true  true  true  true  true  true  true false false false false false false false false false  # 0111 ->
                                    true  true  true  true  true  true  true  true false false false false false false false false  # 1000 ->
                                    true  true  true  true  true  true  true  true  true false false false false false false false  # 1001 ->
                                    true  true  true  true  true  true  true  true  true  true false false false false false false  # 1010 ->
                                    true  true  true  true  true  true  true  true  true  true  true false false false false false  # 1011 ->
                                    true  true  true  true  true  true  true  true  true  true  true  true false false false false  # 1100 ->
                                    true  true  true  true  true  true  true  true  true  true  true  true  true false false false  # 1101 ->
                                    true  true  true  true  true  true  true  true  true  true  true  true  true  true false false  # 1110 ->
                                    true  true  true  true  true  true  true  true  true  true  true  true  true  true  true false  # 1111 ->
    ]

    @test SigmoidNumbers.containszero.(wmatrix) == warlpiri_containszero

    #ispositive test          ->  0000  0001  0010  0011  0100  0101  0110  0111  1000  1001  1010  1011  1100  1101  1110  1111
    const warlpiri_ispositive = [false false false false false false false false false false false false false false false false  # 0000 ->
                                 false  true  true  true  true  true  true  true false false false false false false false false  # 0001 ->
                                 false false  true  true  true  true  true  true false false false false false false false false  # 0010 ->
                                 false false false  true  true  true  true  true false false false false false false false false  # 0011 ->
                                 false false false false  true  true  true  true false false false false false false false false  # 0100 ->
                                 false false false false false  true  true  true false false false false false false false false  # 0101 ->
                                 false false false false false false  true  true false false false false false false false false  # 0110 ->
                                 false false false false false false false  true false false false false false false false false  # 0111 ->
                                 false false false false false false false false false false false false false false false false  # 1000 ->
                                 false false false false false false false false false false false false false false false false  # 1001 ->
                                 false false false false false false false false false false false false false false false false  # 1010 ->
                                 false false false false false false false false false false false false false false false false  # 1011 ->
                                 false false false false false false false false false false false false false false false false  # 1100 ->
                                 false false false false false false false false false false false false false false false false  # 1101 ->
                                 false false false false false false false false false false false false false false false false  # 1110 ->
                                 false false false false false false false false false false false false false false false false  # 1111 ->
    ]

    @test SigmoidNumbers.ispositive.(wmatrix) == warlpiri_ispositive

    #isnegative test          ->  0000  0001  0010  0011  0100  0101  0110  0111  1000  1001  1010  1011  1100  1101  1110  1111
    const warlpiri_isnegative = [false false false false false false false false false false false false false false false false  # 0000 ->
                                 false false false false false false false false false false false false false false false false  # 0001 ->
                                 false false false false false false false false false false false false false false false false  # 0010 ->
                                 false false false false false false false false false false false false false false false false  # 0011 ->
                                 false false false false false false false false false false false false false false false false  # 0100 ->
                                 false false false false false false false false false false false false false false false false  # 0101 ->
                                 false false false false false false false false false false false false false false false false  # 0110 ->
                                 false false false false false false false false false false false false false false false false  # 0111 ->
                                 false false false false false false false false false false false false false false false false  # 1000 ->
                                 false false false false false false false false false  true  true  true  true  true  true  true  # 1001 ->
                                 false false false false false false false false false false  true  true  true  true  true  true  # 1010 ->
                                 false false false false false false false false false false false  true  true  true  true  true  # 1011 ->
                                 false false false false false false false false false false false false  true  true  true  true  # 1100 ->
                                 false false false false false false false false false false false false false  true  true  true  # 1101 ->
                                 false false false false false false false false false false false false false false  true  true  # 1110 ->
                                 false false false false false false false false false false false false false false false  true  # 1111 ->
    ]

    @test SigmoidNumbers.isnegative.(wmatrix) == warlpiri_isnegative

    ################################################################################
    #rounds_positive test     ->  0000  0001  0010  0011  0100  0101  0110  0111  1000  1001  1010  1011  1100  1101  1110  1111
    const warlpiri_rounds_pos = [false false false false false false false false  true  true  true  true  true  true  true  true  # 0000 ->
                                 false false false false false false false false false false false false false false false false  # 0001 ->
                                 false false false false false false false false false false false false false false false false  # 0010 ->
                                 false false false false false false false false false false false false false false false false  # 0011 ->
                                 false false false false false false false false false false false false false false false false  # 0100 ->
                                 false false false false false false false false false false false false false false false false  # 0101 ->
                                 false false false false false false false false false false false false false false false false  # 0110 ->
                                 false false false false false false false false false false false false false false false false  # 0111 ->
                                 false false false false false false false false false false false false false false false false  # 1000 ->
                                 false false false false false false false false  true false false false false false false false  # 1001 ->
                                 false false false false false false false false  true  true false false false false false false  # 1010 ->
                                 false false false false false false false false  true  true  true false false false false false  # 1011 ->
                                 false false false false false false false false  true  true  true  true false false false false  # 1100 ->
                                 false false false false false false false false  true  true  true  true  true false false false  # 1101 ->
                                 false false false false false false false false  true  true  true  true  true  true false false  # 1110 ->
                                 false false false false false false false false  true  true  true  true  true  true  true false  # 1111 ->
                                ]

    @test SigmoidNumbers.rounds_positive.(wmatrix) == warlpiri_rounds_pos

    #test equality of all allreals.
    WP_z = zero(WP)
    WP_allreals = [next(a) → a for a in WP if a != WP_z && next(a) != WP_z]
    for a in WP_allreals, b in WP_allreals
        @test a == b
    end

    #test equality of the two null values.
    @test (next(WP_z) → WP_z) == (WP_z → prev(WP_z))
end


using CoulombIntegrals
using CompactBases
using FillArrays
using LazyArrays
using LinearAlgebra
using BandedMatrices
using Parameters
using Test

include("exact_poisson.jl")
include("hydrogenic_orbitals.jl")
include("integrals.jl")
include("test_convergence_rates.jl")

@testset "Basic test" begin
    function get_grid(grid_type, rmax, Z; ρ=0.1, ρmax=0.6, α=0.002)
        ρ /= Z

        if grid_type == :fd_uniform
            N = ceil(Int, rmax/ρ)
            FiniteDifferences(N, ρ)
        elseif grid_type == :fd_loglin
            StaggeredFiniteDifferences(0.1ρ, ρmax, α, rmax)
        elseif grid_type == :implicit_fd
            N = ceil(Int, rmax/ρ)
            ImplicitFiniteDifferences(N, ρ)
        elseif grid_type == :fedvr
            t = range(0, stop=rmax, length=41)
            FEDVR(t, Vcat(10,Fill(7,length(t)-2)))[:,2:end-1]
        elseif grid_type == :bsplines
            t = ExpKnotSet(5, -2.0, log10(rmax), 200)
            BSpline(t)[:,2:end-1]
        end
    end

    n,ℓ = 3,0
    n′,ℓ′ = 2,0
    k = abs(ℓ′-ℓ) # Multipole order

    Z = 1.0
    rmax = 40(max(n,n′)^1.05)/Z

    grid_types = [(:fd_uniform, 2e-4, 4e-8),
                  (:fd_loglin, 1e-6, 5e-10),
                  (:implicit_fd, 2e-5, 4e-8),
                  (:fedvr, 3e-7, 8e-11),
                  (:bsplines, 5e-9, 2e-9)]

    # All in the name of type-stability
    Yk_scaled(Yk_f) = (r,Z) -> Yk_f(Z*r)
    Yk_ref = exact_Yᵏs[((n,ℓ),(n′,ℓ′),k)]
    fYk = Base.Fix2(Yk_scaled(Yk_ref), Z)

    ϕf = hydredwfn(n, ℓ, Z)
    ϕ′f = hydredwfn(n′, ℓ′, Z)

    @testset "Grid = $(grid_type)" for (grid_type,tol,asymtol) in grid_types
        R = get_grid(grid_type, rmax, Z, ρ=0.1, ρmax=1.0)
        r = axes(R,1)

        Ỹᵏ = R \ fYk.(r)

        ϕ = R \ ϕf.(r)
        ϕ′ = R \ ϕ′f.(r)

        o = R \ one.(r)
        out1 = similar(o)
        out2 = similar(o)

        ρ = Density(applied(*,R,ϕ), applied(*,R,ϕ′))
        potential = CoulombRepulsionPotential(R, k)
        copyto!(potential, ρ)
        Y = potential.poisson.Y
        @test Y ≈ Ỹᵏ atol=tol

        mpot = Matrix(potential)
        @test mpot isa (R isa CompactBases.BSplineOrRestricted ? BandedMatrix : Diagonal)

        mul!(out1, potential, o)
        mul!(out2, LinearOperator(mpot, R), o)
        if R isa CompactBases.BSplineOrRestricted
            # Matricization does not yield the same behaviour for the
            # first few elements since the Dirichlet0 boundary
            # condition at r=0 does not allow the accurate
            # representation of Y/r.
            @test out1[55:end] ≈ out2[55:end] atol=tol
        else
            @test out1 ≈ out2 atol=tol
        end

        @testset "AsymptoticPoissonProblem" begin
            R̃ = R[:,1:end-5]
            apoisson = AsymptoticPoissonProblem(R, k, R̃)
            CoulombIntegrals.solve!(apoisson, ρ)
            @test Y ≈ apoisson.Y atol=asymtol
        end
    end
end

include("poisson_problem_hermiticity.jl")

@testset "Accuracy tests" begin
    rₘₐₓ = 100.0

    orbital_modes = [:symbolic]
    # # Arnoldi is not stable with very small grid spacings;
    # # shift-and-invert has to be implemented.
    # orbital_modes = [:symbolic, :arnoldi]

    coulomb_modes = [:poisson]
    # # Direct integration is simply too slow.
    # coulomb_modes = [:poisson, :direct]
    Z = 1.0

    fedvr_order = 10

    fedvr_inputs(exponents,errors) =
        ("FEDVR",
         (rₘₐₓ,N) -> begin
         t = range(0,stop=rₘₐₓ,length=N)
         amended_order = vcat(fedvr_order+2, Fill(fedvr_order,length(t)-2))
         R = FEDVR(t, amended_order)[:,2:end-1]
         R,(t[2]-t[1])/fedvr_order
         end,
         unique(ceil.(Int, 2 .^ exponents)),
         errors)

    fd_inputs(exponents,errors) =
        ("Finite-differences",
         (rₘₐₓ,N) -> begin
         R = StaggeredFiniteDifferences(rₘₐₓ,N)
         R,step(R)
         end,
         unique(ceil.(Int, 2 .^ exponents)),
         errors)

    # TODO: Add B-spline convergence tests

    @testset "Yᵏ convergence rates" begin
        @testset "$Basis" for (Basis,Rf,Ns,errors) in [
            fedvr_inputs(range(1,stop=5,length=20),
                         [(6, "local", 2e-5, 7.0)]),
            fd_inputs(range(6.5,stop=10,length=15),
                      [(6, "local", 1e-3, 2.0)])
        ]
            @test test_convergence_rates(
                Rf, rₘₐₓ, Ns, 2,
                errors,
                ["N", "ρ", "a", "b", "k", "Local δ", "Global δ", "Time"],
                orbital_modes, coulomb_modes,
                sort(collect(keys(exact_Yᵏs)))) do cache, orbital_mode, coulomb_mode, ((n,ℓ), (n′,ℓ′), k)
                    (n′ > n || n′ == n && ℓ′ > ℓ) && return nothing
                    Yᵏ_error(cache,
                             k, n, ℓ, n′, ℓ′, Z,
                             orbital_mode, coulomb_mode)
                end
        end
    end

    @testset "Fᵏ convergence rates" begin
        @testset "$Basis" for (Basis,Rf,Ns,errors) in [
            fedvr_inputs(range(2,stop=6,length=20),
                         [(11,"δF",2e-6,6.0),
                          (13,"δF′",2e-6,6.0)]),
            fd_inputs(range(7,stop=11,length=15),
                      [(11,"δF",2e-4,2.0),
                       (13,"δF′",2e-4,2.0)])
        ]
            @test test_convergence_rates(
                Rf, rₘₐₓ, Ns, 2,
                errors,
                ["N", "ρ", "a", "b", "c", "d", "k", "Exact", "Exact",
                 "F", "δ", "F′", "δ′", "F-F′", "Time"],
                orbital_modes, coulomb_modes, collect(keys(exact_Fᵏs))) do cache, orbital_mode, coulomb_mode, ((n,ℓ),(n′,ℓ′),k)
                    Fᵏ_error(cache,
                             k, n, ℓ, n′, ℓ′,
                             Z, orbital_mode, coulomb_mode)
                end
        end
    end

    @testset "Gᵏ convergence rates" begin
        @testset "$Basis" for (Basis,Rf,Ns,errors) in [
            fedvr_inputs(range(2,stop=6,length=20),
                         [(11,"δG",2e-6,6.0)]),
            fd_inputs(range(7,stop=12,length=15),
                      [(11,"δG",2e-4,2.0)])
        ]
            @test test_convergence_rates(
                Rf, rₘₐₓ, Ns, 2,
                errors,
                ["N", "ρ", "a", "b", "c", "d", "k", "Exact", "Exact",
                 "G", "δ", "Time"],
                orbital_modes, coulomb_modes, collect(keys(exact_Gᵏs))) do cache, orbital_mode, coulomb_mode, ((n,ℓ),(n′,ℓ′),k)
                    Gᵏ_error(cache,
                             k, n, ℓ, n′, ℓ′,
                             Z, orbital_mode, coulomb_mode)
                end
        end
    end
end


using LDUFacts: piv_to_perm, perm_to_piv, LDUPerm

res(A, la) = la.P' * Matrix(A) * la.P - la.L * la.D * la.U

@testset "throwing" begin
   A = [0.0 1; 1 0]
   @test_throws SingularException ldu(A, NoPivot())
   @test_throws SingularException ldu(A, DiagonalPivot())
   @test rank(ldu(A, FullPivot())) == 2
end

@testset "ldu(3,2) NoPivot" begin
    A = Matrix(Symmetric([1. 0 2; 0 1 1; 0 0 1]))
    la = ldu(A, NoPivot(); check = false)
    @test issuccess(la)
    @test rank(la) == 3
    @test la.p == [1, 2, 3]
    @test norm(res(A, la)) <= 1e-13
    @test Matrix(la.P) == I(3)
    @test propertynames(la) ⊇ (:U, :L, :D, :d, :p, :P)
end
@testset "ldu(3,2) DiagonalPivot" begin
    A = Symmetric([1. 0 2; 0 1 1; 0 0 1])
    la = ldu(A, DiagonalPivot(); check = false)
    @test issuccess(la)
    @test rank(la) == 3
    @test la.p == [1, 3, 2]
    @test norm(res(A, la)) <= 1e-13
    @test Matrix(la.P) == [1. 0 0; 0 0 1; 0 1 0]
end
@testset "ldu(3,2) FullPivot" begin
    A = Hermitian([1. 0 2; 0 1 1; 0 0 1])
    la = ldu(A, FullPivot(); check = false)
    @test issuccess(la)
    @test rank(la) == 3
    @test la.p == [1, 2, 3]
    @test norm(res(A, la)) <= 1e-13
    @test Matrix(la.P) == [1. 0 0; 0 1 0; 1 0 1]
end

@testset "ldu!(7,10), $ps" for ps in (NoPivot(), DiagonalPivot(), FullPivot())
    R = [
         0  -1   0   0   0  -1  -1
        -1   1   0   0  -1  -1   0
         1   1   1   1  -1   1   0
         1   0   1   0   0   0   1
         0   0   0  -1   0   0   0
         0  -1  -1   0  -1  -1   0
         0   0   0   1   1   0  -1
        -1   1  -1  -1   1   0  -1
         0   1  -1  -1  -1   0   0
        -1  -1   0   1   0   1  -1
    ]
    R = float.(R)
    A = R * Diagonal([1,1,1,1,-1,-1,-1]) * R'
    la = ldu(copy(A), ps, tol=-1, check=false)
    @test norm(res(A, la)) < 1e-14
    ps isa NoPivot || @test issuccess(la)
    @test rank(la) == 7
end

@testset "ldu!(rational)" begin
    A = [20 4 11 -23 20; 4 -38 -25 -46 -1; 11 -25 12 -7 27; -23 -46 -7 10 14; 20 -1 27 14 -3] .// 4
    la = ldu!(copy(A), FullPivot(), tol = 0)
    @test rank(la) == 5
    @test issuccess(la)
    @test iszero(res(A, la))
end

@testset "simple complex" begin
    A = [
        -3//1+0//1*im   3//2-2//1*im  -1//1+2//1*im
         3//2+2//1*im   2//1+0//1*im  -3//2-1//2*im
        -1//1-2//1*im  -3//2+1//2*im  -3//2+0//1*im
    ]
    la = ldu(A, DiagonalPivot(), tol = 0)
    @test issuccess(la)
    @test rank(la) == 3
    @test iszero(res(A, la))
    @test iszero(inv(A) - inv(la))
    @test det(A) == det(la)
end

@testset "advanced complex and rational" begin
    A = [
        2//1+0//1*im   2//1-3//2*im   0//1-2//1*im   3//1+3//1*im  -1//1+1//1*im
        2//1+3//2*im   1//1+0//1*im  -1//2+1//2*im  -3//2+2//1*im   1//1+0//1*im
        0//1+2//1*im  -1//2-1//2*im  -1//1+0//1*im   1//1-3//2*im  -3//2-3//2*im
        3//1-3//1*im  -3//2-2//1*im   1//1+3//2*im   3//1+0//1*im   1//1+1//2*im
       -1//1-1//1*im   1//1+0//1*im  -3//2+3//2*im   1//1-1//2*im  -1//1+0//1*im
    ]
    A = big.(A)
    A = Hermitian(A, :L)
    la = ldu(A, FullPivot(), tol = 0)
    @test issuccess(la)
    @test rank(la) == 5
    @test iszero(res(A, la))
    @test inv(la) == inv(A)
    @test logdet(la) ≈ logdet(A)
    b = big.([1+0im, 2, 3, 4, 5])
    @test la \ b == A \ b
    @test b' / la == b' / A
end

@testset "PivotLike" begin
    A = [20 4 11 -23 20; 4 -38 -25 -46 -1; 11 -25 12 -7 27; -23 -46 -7 10 14; 20 -1 27 14 -3] .// 4
    la = ldu(A, FullPivot())
    ps = PivotLike(la)
    lb = ldu(A, ps)
    @test la.piv == lb.piv
    @test la.pam == lb.pam
    @test la.ame == lb.ame
end

@testset "perm_to_piv" begin
    perm = [2, 4, 3, 9, 5, 1, 10, 7, 6, 8]
    piv = perm_to_piv(perm)
    @test piv_to_perm(piv) == perm
end

@testset "mut_by_revperm" begin
    P = LDUPerm([1, 2], [2, 0], [1+2im, 0])
    b = [10+0im, 20]
    MP = Matrix(P)
    @test MP == [1 0; 1+2im 1]
    @test P * b == MP * b
    @test P * [b b] == MP * [b b]
    @test [b'; b'] * P' == [b'; b'] * MP'
end


module Util
using ..DSP: xcorr
import Base: *
import LinearAlgebra.BLAS
using LinearAlgebra: mul!
using FFTW

export  hilbert,
        fftintype,
        fftouttype,
        fftabs2type,
        nextfastfft,
        dB,
        dBa,
        pow2db,
        amp2db,
        db2pow,
        db2amp,
        rms,
        rmsfft,
        meanfreq,
        unsafe_dot,
        shiftin!,
        finddelay,
        shiftsignal,
        shiftsignal!,
        alignsignals,
        alignsignals!

function hilbert(x::StridedVector{T}) where T<:FFTW.fftwReal
# Return the Hilbert transform of x (a real signal).
# Code inspired by Scipy's implementation, which is under BSD license.
    N = length(x)
    X = zeros(Complex{T}, N)
    p = plan_rfft(x)
    mul!(view(X, 1:(N >> 1)+1), p, x)
    for i = 2:div(N, 2)+isodd(N)
        @inbounds X[i] *= 2.0
    end
    return ifft!(X)
end
hilbert(x::AbstractVector{T}) where {T<:Real} = hilbert(convert(Vector{fftintype(T)}, x))

"""
    hilbert(x)

Computes the analytic representation of x, ``x_a = x + j
\\hat{x}``, where ``\\hat{x}`` is the Hilbert transform of x,
along the first dimension of x.
"""
function hilbert(x::AbstractArray{T}) where T<:Real
    N = size(x, 1)
    xc = Vector{fftintype(T)}(undef, N)
    X = Vector{fftouttype(T)}(undef, N)
    out = similar(x, fftouttype(T))

    p1 = plan_rfft(xc)
    Xsub = view(X, 1:(N >> 1)+1)
    p2 = plan_bfft!(X)

    normalization = 1/N
    off = 1
    for i = 1:Base.trailingsize(x, 2)
        copyto!(xc, 1, x, off, N)

        # fft
        fill!(X, 0)
        mul!(Xsub, p1, xc)

        # scale real part
        for i = 2:div(N, 2)+isodd(N)
            @inbounds X[i] *= 2.0
        end

        # ifft
        mul!(X, p2, X)

        # scale and copy to output
        @simd for j = 1:N
            @inbounds out[off+j-1] = X[j]*normalization
        end

        off += N
    end

    out
end

## FFT TYPES

# Get the input element type of FFT for a given type
fftintype(::Type{T}) where {T<:FFTW.fftwNumber} = T
fftintype(::Type{T}) where {T<:Real} = Float64
fftintype(::Type{T}) where {T<:Complex} = ComplexF64

# Get the return element type of FFT for a given type
fftouttype(::Type{T}) where {T<:FFTW.fftwComplex} = T
fftouttype(::Type{T}) where {T<:FFTW.fftwReal} = Complex{T}
fftouttype(::Type{T}) where {T<:Union{Real,Complex}} = ComplexF64

# Get the real part of the return element type of FFT for a given type
fftabs2type(::Type{Complex{T}}) where {T<:FFTW.fftwReal} = T
fftabs2type(::Type{T}) where {T<:FFTW.fftwReal} = T
fftabs2type(::Type{T}) where {T<:Union{Real,Complex}} = Float64

# Get next fast FFT size for a given signal length
const FAST_FFT_SIZES = [2, 3, 5, 7]
"""
    nextfastfft(n)

Return the closest product of 2, 3, 5, and 7 greater than or equal
to `n`. FFTW contains optimized kernels for these sizes and
computes Fourier transforms of input that is a product of these
sizes faster than for input of other sizes.
"""
nextfastfft(n) = nextprod(FAST_FFT_SIZES, n)
nextfastfft(ns::Tuple) = nextfastfft.(ns)
nextfastfft(ns...) = nextfastfft(ns)

## COMMON DSP TOOLS

struct dBconvert end
struct dBaconvert end
const dB = dBconvert()
const dBa = dBaconvert()
# for using e.g. 3dB or -3dBa
*(a::Real, ::dBconvert) = db2pow(a)
*(a::Real, ::dBaconvert) = db2amp(a)

"""
    db2pow(a)

Convert dB to a power ratio. This function call also be called
using `a*dB`, i.e. `3dB == db2pow(3)`. The inverse of `pow2db`.
"""
db2pow(a::Real) = 10^(a/10)

"""
    db2amp(a)

Convert dB to an amplitude ratio. This function call also be called
using `a*dBa`, i.e. `3dBa == db2amp(3)`. The inverse of `amp2db`.
"""
db2amp(a::Real) = 10^(a/20)

"""
    pow2db(a)

Convert a power ratio to dB (decibel), or ``10\\log_{10}(a)``.
The inverse of `db2pow`.
"""
pow2db(a::Real) = 10*log10(a)

"""
    amp2db(a)

Convert an amplitude ratio to dB (decibel), or ``20
\\log_{10}(a)=10\\log_{10}(a^2)``. The inverse of `db2amp`.
"""
amp2db(a::Real) = 20*log10(a)

"""
    rms(s)

Return the root mean square of signal `s`.
"""
rms(s::AbstractArray{T}) where {T<:Number} = sqrt(sum(abs2, s)/length(s))

"""
    rmsfft(f)

Return the root mean square of signal `s` given the FFT transform
`f = fft(s)`. Equivalent to `rms(ifft(f))`.
"""
rmsfft(f::AbstractArray{T}) where {T<:Complex} = sqrt(sum(abs2, f))/length(f)

"""
    meanfreq(x, fs)

Calculate the mean power frequency of `x` with a sampling frequency of `fs`, defined as:
```math
MPF = \\frac{\\sum_{i=1}^{F} f_i X_i^2 }{\\sum_{i=0}^{F} X_i^2 } Hz
```
where ``F`` is the Nyquist frequency, and ``X`` is the power spectral density.
"""
function meanfreq(x::AbstractVector{<:Real}, fs=2*π)
    pxx = abs2.(rfft(x))

    len = length(x)
    npoints = fld(len,2)
    freqrg = fs/len.*(0:(npoints))

    mf = sum(pxx.*freqrg)./sum(pxx)
    return mf
end

# Computes the dot product of a single column of a, specified by aColumnIdx, with the vector b.
# The number of elements used in the dot product determined by the size(A)[1].
# Note: bIdx is the last element of b used in the dot product.
function unsafe_dot(a::AbstractMatrix, aColIdx::Integer, b::AbstractVector, bLastIdx::Integer)
    aLen     = size(a, 1)
    bBaseIdx = bLastIdx - aLen
    dotprod  = a[1, aColIdx] * b[ bBaseIdx + 1]
    @simd for i in 2:aLen
        @inbounds dotprod += a[i, aColIdx] * b[bBaseIdx + i]
    end

    return dotprod
end

@inline function unsafe_dot(a::Matrix{T}, aColIdx::Integer, b::Vector{T}, bLastIdx::Integer) where T<:BLAS.BlasReal
    BLAS.dot(size(a, 1), pointer(a, size(a, 1)*(aColIdx-1) + 1), 1, pointer(b, bLastIdx - size(a, 1) + 1), 1)
end

function unsafe_dot(a::AbstractMatrix, aColIdx::Integer, b::AbstractVector{T}, c::AbstractVector{T}, cLastIdx::Integer) where T
    aLen = size(a, 1)
    bLen = length(b)
    bLen == aLen-1  || error( "length(b) must equal to length(a)[1] - 1" )
    cLastIdx < aLen || error( "cLastIdx but be < length(a)")

    dotprod = a[1, aColIdx] * b[cLastIdx]
    @simd for i in 2:aLen-cLastIdx
        @inbounds dotprod += a[i, aColIdx] * b[i+cLastIdx-1]
    end
    @simd for i in 1:cLastIdx
        @inbounds dotprod += a[aLen-cLastIdx+i, aColIdx] * c[i]
    end

    return dotprod
end

function unsafe_dot(a::T, b::AbstractArray, bLastIdx::Integer) where T
    aLen     = length(a)
    bBaseIdx = bLastIdx - aLen
    @inbounds dotprod  = a[1] * b[bBaseIdx + 1]
    @simd for i in 2:aLen
        @inbounds dotprod += a[i] * b[bBaseIdx + i]
    end

    return dotprod
end

@inline function unsafe_dot(a::Vector{T}, b::Array{T}, bLastIdx::Integer) where T<:BLAS.BlasReal
    BLAS.dot(length(a), pointer(a), 1, pointer(b, bLastIdx - length(a) + 1), 1)
end

function unsafe_dot(a::AbstractVector, b::AbstractVector{T}, c::AbstractVector{T}, cLastIdx::Integer) where T
    aLen    = length(a)
    dotprod = zero(a[1]*b[1])
    @simd for i in 1:aLen-cLastIdx
        @inbounds dotprod += a[i] * b[i+cLastIdx-1]
    end
    @simd for i in 1:cLastIdx
        @inbounds dotprod += a[aLen-cLastIdx+i] * c[i]
    end

    return dotprod
end

# Shifts b into the end a.
# julia> DSP.Util.shiftin!( [1,2,3,4], [5, 6])
# 4-element Array{Int64,1}:
#  3
#  4
#  5
#  6
function shiftin!(a::AbstractVector{T}, b::AbstractVector{T}) where T
    aLen = length(a)
    bLen = length(b)

    if bLen >= aLen
        copyto!(a, 1, b, bLen - aLen + 1, aLen)
    else

        for i in 1:aLen-bLen
            @inbounds a[i] = a[i+bLen]
        end
        bIdx = 1
        for i in aLen-bLen+1:aLen
            @inbounds a[i] = b[bIdx]
            bIdx += 1
        end
    end

    return a
end

# DELAY FINDING UTILITIES

"""
    finddelay(x, y)

Estimate the delay of x with respect to y by locating the peak of their
cross-correlation.

The output delay will be positive when x is delayed with respect y, negative if
advanced, 0 otherwise.

# Example
```jldoctest
julia> finddelay([0, 0, 1, 2, 3], [1, 2, 3])
2

julia> finddelay([1, 2, 3], [0, 0, 1, 2, 3])
-2
```
"""
function finddelay(x::AbstractVector{<: Real}, y::AbstractVector{<: Real})
    s = xcorr(y, x, padmode=:none)
    max_corr = maximum(abs, s)
    max_idxs = findall(x -> abs(x) == max_corr, s)

    center_idx = length(x)
    # Delay is position of peak cross-correlation relative to center.
    # If the maximum cross-correlation is not unique, use the position
    # closest to the center.
    d_ind = argmin(abs.(center_idx .- max_idxs))
    d = center_idx - max_idxs[d_ind]
end

"""
    shiftsignal!(x, s)

Mutating version of shiftsignals(): shift x of s samples and fill the spaces
with zeros in-place.

See also [`shiftsignal`](@ref).
"""
function shiftsignal!(x::AbstractVector, s::Integer)
    l = length(x)
    if abs(s) > l
        error("The absolute value of s must not be greater than the length of x")
    end
    if s > 0
        x[s + 1:l] = x[1:l - s]
        x[1:s] .= 0
    elseif s < 0
        x[1:l + s] = x[1 - s:l]
        x[l + s + 1:l] .= 0
    end
    x
end

"""
    shiftsignal(x, s)

Shift elements of signal x in time by a given amount s of samples and fill
the spaces with zeros. For circular shifting, use circshift.

# Example
```jldoctest
julia> shiftsignal([1, 2, 3], 2)
3-element Array{Int64,1}:
 0
 0
 1

julia> shiftsignal([1, 2, 3], -2)
3-element Array{Int64,1}:
 3
 0
 0
```

See also [`shiftsignal!`](@ref).
"""
shiftsignal(x::AbstractVector, s::Integer) = shiftsignal!(copy(x), s)

"""
    alignsignals!(x, y)

Mutating version of alignsignals(): time align x to y in-place.

See also [`alignsignals`](@ref).
"""
function alignsignals!(x, y)
    d = finddelay(x, y)
    x = shiftsignal!(x, -d)
    x, d
end

"""
    alignsignals(x, y)

Use finddelay() and shiftsignal() to time align x to y. Also return the delay
of x with respect to y.

# Example
```jldoctest
julia> alignsignals([0, 0, 1, 2, 3], [1, 2, 3])
([1, 2, 3, 0, 0], 2)

julia> alignsignals([1, 2, 3], [0, 0, 1, 2, 3])
([0, 0, 1], -2)
```

See also [`alignsignals!`](@ref).
"""
alignsignals(x, y) = alignsignals!(copy(x), y)

end # end module definition


"""
    KPath{D}

Stores an path in k-space.

#### Fields
- kpts::Vector{SVector{D, Float64}}
- arc::Vector{Float64}
- special::Vector{Float64}
- nkpts::Int

"""
struct KPath{D}
    kpts::Vector{SVector{D, Float64}}
    arc::Vector{Float64}
    special::Vector{Float64}
    nkpts::Int
end

"""
    KMesh{D}

Stores an path in k-space.

#### Fields
- kpts::Vector{SVector{D, Float64}}
- ws::Vector{Float64}
- nkpts::Int
"""
struct KMesh{D}
    kpts::Vector{SVector{D, Float64}}
    ws::Vector{Float64}
    nkpts::Int
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 59750850-b0ea-11eb-2086-87707f76e327
md"""
# [Day 24: It Hangs in the Balance](https://adventofcode.com/2015/day/24)

It's Christmas Eve, and Santa is loading up the sleigh for this year's deliveries. However, there's one small problem: he can't get the sleigh to balance. If it isn't balanced, he can't defy physics, and nobody gets presents this year.

No pressure.

Santa has provided you a list of the weights of every package he needs to fit on the sleigh. The packages need to be split into three groups of exactly the same weight, and every package has to fit. The first group goes in the passenger compartment of the sleigh, and the second and third go in containers on either side. Only when all three groups weigh exactly the same amount will the sleigh be able to fly. Defying physics has rules, you know!

Of course, that's not the only problem. The first group - the one going in the passenger compartment - needs as few packages as possible so that Santa has some legroom left over. It doesn't matter how many packages are in either of the other two groups, so long as all of the groups weigh the same.

Furthermore, Santa tells you, if there are multiple ways to arrange the packages such that the fewest possible are in the first group, you need to choose the way where the first group has the smallest quantum entanglement to reduce the chance of any "complications". The quantum entanglement of a group of packages is the [product](https://en.wikipedia.org/wiki/Product_%28mathematics%29) of their weights, that is, the value you get when you multiply their weights together. Only consider quantum entanglement if the first group has the fewest possible number of packages in it and all groups weigh the same amount.

For example, suppose you have ten packages with weights `1` through `5` and `7` through `11`. For this situation, some of the unique first groups, their quantum entanglements, and a way to divide the remaining packages are as follows:

```
Group 1;             Group 2; Group 3
11 9       (QE= 99); 10 8 2;  7 5 4 3 1
10 9 1     (QE= 90); 11 7 2;  8 5 4 3
10 8 2     (QE=160); 11 9;    7 5 4 3 1
10 7 3     (QE=210); 11 9;    8 5 4 2 1
10 5 4 1   (QE=200); 11 9;    8 7 3 2
10 5 3 2   (QE=300); 11 9;    8 7 4 1
10 4 3 2 1 (QE=240); 11 9;    8 7 5
9 8 3      (QE=216); 11 7 2;  10 5 4 1
9 7 4      (QE=252); 11 8 1;  10 5 3 2
9 5 4 2    (QE=360); 11 8 1;  10 7 3
8 7 5      (QE=280); 11 9;    10 4 3 2 1
8 5 4 3    (QE=480); 11 9;    10 7 2 1
7 5 4 3 1  (QE=420); 11 9;    10 8 2

```

Of these, although `10 9 1` has the smallest quantum entanglement (`90`), the configuration with only two packages, `11 9`, in the passenger compartment gives Santa the most legroom and wins. In this situation, the quantum entanglement for the ideal configuration is therefore `99`. Had there been two configurations with only two packages in the first group, the one with the smaller quantum entanglement would be chosen.

**What is the quantum entanglement of the first group of packages in the ideal configuration?**
"""

# ╔═╡ 45a666a1-ffaa-4517-b79b-cc6b182505b1
puzzle_input = "1
2
3
7
11
13
17
19
23
31
37
41
43
47
53
59
61
67
71
73
79
83
89
97
101
103
107
109
113"

# ╔═╡ Cell order:
# ╟─59750850-b0ea-11eb-2086-87707f76e327
# ╟─45a666a1-ffaa-4517-b79b-cc6b182505b1


using DataFrames
using GLM
using Plots
function getcolormap(name="jet", n=64)
    if name == "jet"
        return RGB{Float64}[RGB(
        clamp(min(4x - 1.5, -4x + 4.5) ,0.0,1.0),
        clamp(min(4x - 0.5, -4x + 3.5) ,0.0,1.0),
        clamp(min(4x + 0.5, -4x + 2.5) ,0.0,1.0)) for x in linspace(0.0,1.0,n)]
    end
end

const cm = getcolormap("jet", 64)

P2c(P) = cm[round(Int,(max(log10(P),-4)/4+1)*63)+1]
function Pvalues(model)
    cc = coef(model)
    se = stderr(model)
    zz = cc ./ se
    P = 2.0 * ccdf(Normal(), abs(zz))
end

"""
`scattermatrix(df::DataFrame; reglines = false)`

`scattermatrix(df::DataFrame, f::Formula; reglines = false)`

`scattermatrix(A::AbstractMatrix; reglines = false)`

`scattermatrix(m::RegressionModel)` Useful together with e.g. GLM.jl

`scattermatrix(models::Array{::RegressionModel})` Should only be used if all models have the same independent parameters, but predict different response

Plot a scatter matrix of a DataFrame, Linear regression model or a numerical matrix

Examples
```
ExperimentalAnalysis.scattermatrix(randn(10,3))
ExperimentalAnalysis.scattermatrix(DataFrame(randn(10,3)))
ExperimentalAnalysis.scattermatrix(DataFrame(randn(10,3)), reglines=true)
ExperimentalAnalysis.scattermatrix(DataFrame(randn(10,3)), f ~ x1 + x2)
ExperimentalAnalysis.scattermatrix(DataFrame(randn(10,3)), x3 ~ x1 + x2)
ExperimentalAnalysis.scattermatrix(randn(10,3),["x1", "x2", "x3"])
modela = lm(b ~ a + c, df)
modelb = lm(d ~ a + c, df)
ExperimentalAnalysis.scattermatrix([modela, modelb])
ExperimentalAnalysis.modelheatmap(["a", "b"], modela, modelb)
```
"""
function scattermatrix(df::DataFrame; reglines = false)
    N, Nparams = size(df)
    p = plot(layout=(Nparams, Nparams))
    # Plot df
    for (i,is) = enumerate(names(df))
        for (j,js) = enumerate(names(df))
            x = isa(df[Symbol(js)], Vector{Float64}) ? df[Symbol(js)] : convert(Array,df[Symbol(js)],0)
            if i != j
                y = isa(df[Symbol(js)], Vector{Float64}) ? df[Symbol(is)] : convert(Array,df[Symbol(is)],0)
                scatter!(p[i,j],x,y,legend=false,grid=true)
                if reglines
                    k = [x ones(N)]\y
                    px = [minimum(x), maximum(x)]
                    plot!(p[i,j], px, k[1].*px + k[2], c=:red)
                end
            else
                plot!(p[i,j],x,l=:histogram, legend=false)
            end
        end
    end
    names_ = string.(names(df))
    for i = 1:Nparams
        plot!(p[1,i], title=(names_[i]))
        plot!(p[i,1], ylabel=(names_[i]))
    end
    for i = 1:Nparams-1, j = 1:Nparams
        plot!(p[i,j], xticks=Float64[])
    end
    for i = 1:Nparams, j = 2:Nparams
        plot!(p[i,j], yticks=Float64[])
    end
    return p
end

function scattermatrix(df::DataFrame, f::Formula; reglines = false)
    rhs = f.rhs.args
    @assert rhs[1] == :+ "Only addition formulas supported (a ~ b + c + d + e)"
    if isa(f.lhs, Expr)
        @assert f.lhs.args[1] == :+ "Only addition formulas supported (a + b ~ c + d + e)"
        scattermatrix_someofothers(df, f, reglines=reglines)
    else
        scattermatrix(df[Symbol.(rhs[2:end])], reglines=reglines)
    end
end

function scattermatrix{T<:AbstractString}(A::AbstractMatrix, names::AbstractVector{T}; reglines = false)
    df = DataFrame(Any[A[:,i] for i = 1:size(A,2)], Symbol.(names))
    scattermatrix(df, reglines=reglines)
end

function scattermatrix{T <: Real}(A::AbstractMatrix{T}; reglines = false)
    df = DataFrame(A)
    scattermatrix(df, reglines=reglines)
end

scattermatrix(model::DataFrames.DataFrameRegressionModel) = scattermatrix([model])
function scattermatrix{T<:DataFrames.DataFrameRegressionModel}(models::AbstractVector{T})
    Nparams = size(models[1].model.pp.X,2)-1
    Nmodels = length(models)
    p = plot(layout=(Nmodels, Nparams))

    # Plot df
    for (i,model) = enumerate(models)
        df = model.mf.df
        data = model.model.pp.X[:,2:end]
        y = df[:,1]
        names_ = names(df)
        plot!(p[i,1], ylabel=(string(names_[1])))
        # P-value color coding
        P = Pvalues(model)
        yhat = predict(model)
        for j = 1:size(data,2)
            x = data[:,j]
            scatter!(p[i,j],x,y,legend=false,grid=true)
            # Regline
            minx,mini = findmin(x)
            maxx,maxi = findmax(x)
            py = yhat[[mini, maxi]]
            px = [minx, maxx]

            plot!(p[i,j], px, py, c=P2c(P[j+1]))
        end
    end
    names_ = map(x->replace(string(x),"&","\\cdot"),models[1].mf.terms.terms)
    for i = 1:Nparams
        plot!(p[1,i], title=(names_[i]))
    end
    for i = 1:Nmodels-1, j = 1:Nparams
        plot!(p[i,j], xticks=Float64[])
    end
    for i = 1:Nmodels, j = 2:Nparams
        plot!(p[i,j], yticks=Float64[])
    end
    return p

end

function Float(col)
    isa(col,DataArrays.PooledDataArray) && return Float64.(col.refs)
    isa(col,DataArrays.DataArray) && return Float64.(col.data)
    col
end

function scattermatrix_someofothers(df::DataFrame, f::Formula; reglines = false)

    Nl = length(f.lhs.args)-1
    Nr = length(f.rhs.args)-1
    p = plot(layout=(Nl, Nr))
    # Plot df
    namesl = Array(AbstractString,Nl)
    namesr = Array(AbstractString,Nr)
    for (i,is) = enumerate(f.lhs.args[2:end])
        for (j,js) = enumerate(f.rhs.args[2:end])
            namesl[i] = string(is)
            namesr[j] = string(js)
            column  = Float(df[js])
            if is != js
                scatter!(p[i,j],column,df[is],legend=false, grid=true)
                if reglines
                    A = [column ones(size(column,1))]
                    y = Float(df[is])
                    k = A\y
                    se = std(y-A*k)/√(inv(A'A)[1,1])
                    score = k[1] ./ se
                    P = 2.0 * ccdf(Normal(), abs(score))
                    column_bounds = [minimum(column), maximum(column)]
                    plot!(p[i,j], column_bounds, k[1].*column_bounds + k[2], c=P2c(P))
                end
            else
                plot!(p[i,j],df[is],l=:histogram, legend=false)
            end
        end
    end

    for i = 1:Nr
        plot!(p[1,i], title=namesr[i])
    end
    for i = 1:Nl
        plot!(p[i,1], ylabel=namesl[i])
    end
    for i = 1:Nl-1, j = 1:Nr
        plot!(p[i,j], xticks=Float64[])
    end
    for i = 1:Nl, j = 2:Nr
        plot!(p[i,j], yticks=Float64[])
    end
    return p
end

modelheatmap{T<: AbstractString}(modelnames::AbstractArray{T}, models...) = modelheatmap(modelnames, collect(models))

function modelheatmap(model::DataFrames.DataFrameRegressionModel)
    modelnames = [string(model.mf.terms.eterms[1])]
    modelheatmap(modelnames,[model])
end

function modelheatmap{T<:DataFrames.DataFrameRegressionModel}(models::AbstractArray{T})
    modelnames = AbstractString[string(models[i].mf.terms.eterms[1]) for i = eachindex(models)]
    modelheatmap(modelnames,models)
end

"""
Plot a model heatmap

example:
```
modela = lm(b ~ a + c, df)
modelb = lm(d ~ a + c, df)
ExperimentalAnalysis.modelheatmap(modela, modelb)
```

`modelheatmap(modelnames::AbstractArray, models::AbstractArray)`

`modelheatmap(models::AbstractArray)`

`modelheatmap(model::DataFrameRegressionModel)`

`modelheatmap(modelnames::AbstractArray, models...)`
"""
function modelheatmap{T<:DataFrames.DataFrameRegressionModel, D<: AbstractString}(modelnames::AbstractArray{D}, models::AbstractArray{T})
    Nmodels = length(models)
    Nparams = length(coef(models[1].model))
    P = Array(Float64, Nparams, Nmodels)
    coefnames = GLM.coeftable(models[1]).rownms
    tickvec = [replace(coefnames[i],"&","\\cdot") for i = 1:Nparams]
    for (i,mm) in enumerate(models)
        P[:,i] = Pvalues(mm.model)
    end
    logP = max(log10(P),-4)
    p = heatmap(logP', xticks=(eachindex(tickvec)-1, tickvec), yticks=(eachindex(modelnames)-1, modelnames), title="log(P)-values", colorbar=true)
end

#

# Example

function perform_example_analysis()
    A = readcsv(joinpath(dirname(@__FILE__),"results.csv"));

    df = DataFrame(Any[A[2:end,i] for i = 1:size(A,2)], Symbol.(A[1,:][:]));
    pool!(df,[:Location, :Substract, :Trial]);

    for name in names(df)
        s = Symbol(name)
        if !isa(df[s], DataArrays.PooledDataArray)
            df[s] = Float64.(df[s])
        end
    end

    df = df[14:end,:];

    model = lm(percentage ~ galacosidase +  endomannanase + Firstratio + Duration, df)

    scattermatrix(df[[:percentage, :galacosidase, :endomannanase, :Biocelulasa, :Firstratio, :Secondratio, :Duration]]);gui()
    scattermatrix(model);gui()
    modelheatmap(model);gui()

    return nothing
end