module TLWE
    using ..Key
    using ..Utils
    using LinearAlgebra
    function tlweSymEncrypt(p::Float64,α::Float64,key::Array{UInt32,1})
        a::Array{UInt32,1} = rand(0:1,length(key))
        b::UInt32 = Utils.gaussian32(p,α)[1] + a ⋅ key
        return append!(a,b)
    end
    function tlweSymDecrypt(c::Array{UInt32,1},key::Array{UInt32,1})
        return 1 - ((c[end] - c[1:end-1] ⋅ key) > UInt32(2)^31)
    end

    function bootsSymEncrypt(p::Array{Int32,1},sk::Key.SecretKey)
        return [tlweSymEncrypt(x, sk.params.α, sk.key.tlwe) for x in (2*p .- 1)*Key.μ]
    end
    function bootsSymDecrypt(c::Array{Array{UInt32,1},1},sk::Key.SecretKey)
        return [tlweSymDecrypt(x,sk.key.tlwe) for x in c]
    end
end


using UnicodePlots: heatmap
part1(input) = length(fold(input, part = 1))

function part2(input)
    coords = fold(input, part = 2)
    m = zeros(Int, 6, 39)
    for (x, y) in coords
        m[6 - y, x + 1] = 1
    end
    return heatmap(m)
end

function fold(input; part)
    coords = Tuple{Int, Int}[]
    for line in eachline(input)
        ',' in line && push!(coords, Tuple(parse.(Int, split(line, ","))))
        startswith(line, "fold along x") && fold_x!(coords, parse(Int, last(split(line, "="))))
        startswith(line, "fold along y") && fold_y!(coords, parse(Int, last(split(line, "="))))
        part == 1 && startswith(line, "fold") && break
    end
    return coords
end

fold_x!(coords, pos) = unique!(map!(c -> (mirror(c[1], pos), c[2]), coords, coords))
fold_y!(coords, pos) = unique!(map!(c -> (c[1], mirror(c[2], pos)), coords, coords))
mirror(x, pos) = x < pos ? x : 2 * pos - x


module ProximalBase

export
  # types
  DifferentiableFunction,
  ProximableFunction,
  SparseIterate, SymmetricSparseIterate, AtomIterate, numCoordinates,

  # smooth functions
  QuadraticFunction, L2Loss, LeastSquaresLoss,
  value, value_and_gradient!, gradient!,

  # proximal functions
  ProxZero,
  ProxL1,
  ProxL2, ProxL2Sq, ProxNuclear,
  ProxL1Fused,
  ProxGaussLikelihood,
  prox!, prox, cdprox!,

  # utils
  shrink, shrinkL2!,
  proxL1Fused,
  A_mul_B_row, At_mul_B_row,
  A_mul_X_mul_B, A_mul_X_mul_B_rc,
  A_mul_UUt_mul_B, A_mul_UUt_mul_B_rc,
  ind2subLowerTriangular,sub2indLowerTriangular,
  norm_diff,

  # function for SparseIterate
  nnz

using SparseArrays
using LinearAlgebra

import Base: *

include("sparseIterate.jl")

include("utils.jl")
# DifferentiableFunctions
include("differentiable_functions.jl")

# ProximableFunctions
include("proximal_functions.jl")

end


using GeoMakie, GLMakie

lons = -180:180
lats = -90:90
field = [exp(cosd(l)) + 3(y/90) for l in lons, y in lats]

# Surface example
fig = Figure()
ax = GeoAxis(fig[1,1])
surface!(ax, lons, lats, field)
display(fig)

# %% Scatter example
slons = rand(lons, 2000)
slats = rand(lats, 2000)
sfield = [exp(cosd(l)) + 3(y/90) for (l,y) in zip(slons, slats)]

fig = Figure()
ax = GeoAxis(fig[1,1])
scatter!(slons, slats; color = sfield)
display(fig)


#################################################
# Calculations required for ll and its gradient #
#################################################

function _Q(U::Matrix{Float64}, A::UpperTriangular, ρ::Vector{Float64})
    n = size(U,1)
    q = Array(Float64, n)
    Ωinv = A'Diagonal(exp(-2.0*ρ))*A
    # println("Ωinv :\n", Ωinv)
    for i in 1:n
        q[i] = (U[i,:]'Ωinv * U[i,:])[1,1]
    end
    return q
end

_L(U::Matrix{Float64}, η::Vector{Float64}) = U*η
_t(L::Vector{Float64}, Q::Vector{Float64}, ν::Float64, k::Int) = L .* sqrt((ν + k)./(Q+ν))
_sf(Q::Vector{Float64}, ν::Float64, k::Int) = sqrt((ν+k)./(ν+Q))
#_sf2(Q::Vector{Float64}, ν::Float64, k::Int) = sqrt((1+k/ν)./(1+(Q./ν))) # R uses this definition for large ν
_gQ(sf::Vector{Float64}) = (-0.5)*sf.^2 #-((ν + k)/2ν)*(1.0./(1 + q))
_∂logT₁∂t(t::Vector{Float64}, ν::Float64, k::Int) = exp(_logt₁(t, ν+k) - _logT₁(t,ν+k))
_tL(sf::Vector{Float64}) = sf
_tQ(l::Vector{Float64}, q::Vector{Float64}, sf::Vector{Float64}, ν::Float64) = (-0.5).*l.*sf./(q+ν)
function _∂log_g∂ν(q::Vector{Float64}, ν::Float64, k::Int)
    0.5 *(digamma(0.5*(ν+k)) - digamma(0.5*ν) - k/ν + ((ν + k)*q)./(ν^2 * (1+q/ν)) - log(1+q/ν))
end

function derivative(f)
    return function(x)
        # pick a small value for h
        h = x == 0 ? sqrt(eps(Float64)) : sqrt(eps(Float64)) * x

        # floating point arithmetic gymnastics
        xph = x + h
        dx = xph - x

        # evaluate f at x + h
        f1 = f(xph)

        # evaluate f at x
        f0 = f(x)

        # divide the difference by h
        return (f1 - f0) / dx
    end
end

function _∂logT₁∂ν(l::Vector{Float64}, q::Vector{Float64}, ν::Float64, k::Int)
    lg(ν1::Float64) = _logT₁(_t(l,q,ν1,k), ν1 + k)
    return derivative(lg)(ν)
end

################################################################
# Functions for translating parameters between different forms #
################################################################

function dplist2optpar(Ω::Matrix{Float64}, α::Vector{Float64}, ν::Real)
    k = length(α)
    η = α./sqrt(diag(Ω))
    Ωinv = inv(Ω)
    upper = chol(Hermitian(Ωinv))
    D = diag(upper)
    A = upper./D
    D .*= D
    ρ = -log(D)/2
    return ρ, A, η, log(ν)
end

function read_params(params::Vector{Float64}, p::Int, k::Int)
    β = reshape(params[1:(p*k)], p, k)
    ρ = params[(p*k+1):(p*k + k)]
    i0 = p*k + k
    i1 = p*k + div(k*(k+1),2)
    A = uppertri2mat(params[i0+1:i1], k)
    η = params[i1 + 1:i1 + k]
    ν = exp(params[end])
    return (β, ρ, A, η, ν)
end

function write_params(β::Matrix{Float64}, ρ::Vector{Float64}, A::Matrix{Float64}, η::Vector{Float64}, logν::Float64)
    p, k = size(β)
    params = Array(Float64, p*k + div((k+2)*(k+1),2))
    params[1:p*k] = vec(β)
    params[p*k + 1: p*k + k] = ρ
    i0 = p*k + k
    i1 = p*k + div(k*(k+1),2)
    params[i0+1:i1] = mat2uppertri(A)
    params[i1 + 1:i1 + k] = η
    params[end] = logν
    return params
end

function mat2uppertri(A::Matrix{Float64})
    k = size(A,1)
    U = Array(Float64, div(k*(k-1),2))
    for j in 1:k-1
        for i in 1:j
            U[div((j-1)*j,2) + i] = A[i,j+1]
        end
    end
    return U
end

function uppertri2mat(U::Vector{Float64}, k::Int)
    A = zeros(k, k)
    for j in 1:k-1
        for i in 1:j
            A[i,j+1] = U[div((j-1)*j,2) + i]
        end
    end
    for i in 1:k
        A[i,i] = 1.0
    end
    return A
end

function print_params(β::Matrix{Float64}, ρ::Vector{Float64}, A::Matrix{Float64}, η::Vector{Float64}, ν::Float64)
    println("β = ")
    println(β)
    println("ρ = $(ρ)")
    println("A = ")
    println(A)
    println("η = $(η)")
    println("ν = $(ν)")
end

###############################
# Log likelihood and gradient #
###############################

_U(X::Matrix{Float64}, Y::Matrix{Float64}, β::Matrix{Float64}) = Y - X*β

function nll(params::Vector{Float64}, X::Matrix{Float64}, Y::Matrix{Float64})
    n, p = size(X)
    k = size(Y,2)
    (β, ρ, A, η, ν) = read_params(params, p, k)

    U = _U(X,Y,β)
    Q = _Q(U, A, ρ)
    L = _L(U,η)
    t = _t(L, Q, ν, k)

    # logdet(D) = -2 Σᵢρᵢ
    D = Diagonal(exp(-2ρ))
    ℓ = n * ( log(2) - 0.5 * logdet(D) ) + sum( _log_g(Q, ν, k) + _logT₁(t, ν + k) )

    return -2ℓ
end

function nll_and_grad(params::Vector{Float64}, X::Matrix{Float64}, Y::Matrix{Float64})
    n, p = size(X)
    k = size(Y,2)
    (β, ρ, A, η, ν) = read_params(params, p, k)
    if any(ρ .< -30)
        return Inf, fill(Inf, length(params))
    end
    #print_params(β, ρ, A, η, ν)
    A = UpperTriangular(A)

    # println("————————————–")
    U = _U(X,Y,β)
    Q = _Q(U, A, ρ)
    L = _L(U,η)
    t = _t(L, Q, ν, k)

    D = Diagonal(exp(-2*ρ))
    Dinv = Diagonal(exp(2*ρ))
    Ωinv = A'D*A
    sf = _sf(Q,ν,k)
    ∂logT₁∂t = _∂logT₁∂t(t, ν, k)
    g_Q = _gQ(sf)
    t_Q = _tQ(L,Q,sf,ν)

    # Calculate derivatives
    ∂ℓ∂β = -2X'Diagonal(g_Q + ∂logT₁∂t.*t_Q)*U*Ωinv - X'Diagonal(∂logT₁∂t.*_tL(sf))*ones(n)*η'
    ∂ℓ∂D = Diagonal(A*U'Diagonal(g_Q + ∂logT₁∂t.*t_Q)*U*A') + 0.5 * n * Dinv
    ∂ℓ∂ρ = -2*diag(∂ℓ∂D*D)
    # println("fine...")
    # println("Q: \n", Q)
    # println("D: \n", D)
    # println("A: \n", A)
    # println("U: \n", U)
    # println("g_Q: \n", g_Q)
    # println("∂logT₁∂t: \n", ∂logT₁∂t)
    # println("t_Q: \n", t_Q)
    ∂ℓ∂A = 2 * triu(D*A*U'Diagonal(g_Q + ∂logT₁∂t.*t_Q)*U)
    # println("..and fine again...")
    ∂ℓ∂η = U'Diagonal(∂logT₁∂t.*_tL(sf))*ones(n)
    ∂ℓ∂ν = sum(_∂log_g∂ν(Q,ν,k) + _∂logT₁∂ν(L,Q,ν,k))
    ∂ℓ∂logν = ∂ℓ∂ν * ν

    #print_params(∂ℓ∂β, ∂ℓ∂ρ, ∂ℓ∂A , ∂ℓ∂η, ∂ℓ∂logν)

    # logdet(D) = -2 Σᵢρᵢ
    ℓ = n * ( log(2) + 0.5 * logdet(D) ) + sum(_log_g(Q, ν, k)) + sum(_logT₁(t, ν + k))
    #println(ℓ)
    return -2*ℓ, -2*write_params(∂ℓ∂β, ∂ℓ∂ρ, ∂ℓ∂A, ∂ℓ∂η, ∂ℓ∂logν)
end

# Fit a MvSkewTDist regression for model to data
# X = (n x p) model matrix (each row is corresponds to the predictors of one response)
# Y = (n x k) response matrix (each row correspond to one output observation)
#
# The following model is assumed:
#   yᵢ ∼ STₖ(ξᵢ, Ω, α, ν) where ξᵢ = xᵢ̱β
# for some (p x k) design matrix β
#
function fit_MvSkewTDist(X::Matrix{Float64}, Y::Matrix{Float64}, method::Optim.Optimizer=LBFGS(); kwargs...)
    # We use the following reparametrization:
    # Ω⁻¹ = Aᵀdiag(exp(-2ρ))A = AᵀDA
    # η = ω⁻¹ α
    # Now, for θ = (β, A, ρ, η, log ν)
    # ℓ(θ) = ∑ᵢ ℓᵢ(θ)
    # for observations i = 1,…,n
    # where
    # ℓᵢ(θ) = log 2 + 0.5 log |D| + log gₖ(Qᵢ; ν) + log T₁( t(Lᵢ, Qᵢ, ν); ν + k )
    # and
    # uᵢ = yᵢ - βᵀ xᵢ
    # Qᵢ = uᵢᵀ Ω⁻¹ uᵢ
    # Lᵢ = αᵀ ω⁻¹ uᵢ
    # and
    # t(L,Q,ν) = L √((ν + k)/(Q + ν))
    # T₁ denots the scalar t₁ distribution function
    # gₖ(Q; ν) = ((Γ((ν + k)/2))/((πν)^(k/2)* Γ(ν/2))) * (1 + Q/ν)^(-(ν + d)/2)
    size(X,1) == size(Y,1) || throw(ArgumentError("X and Y must have the same number of rows"))
    n,p = size(X)
    k = size(Y,2)

    obj(params::Vector{Float64}) = nll(params, X, Y)
    function grad!(params::Vector{Float64}, g::Vector{Float64})
        g[:] = nll_and_grad(params, X, Y)[2]
    end

    function obj_and_grad!(params::Vector{Float64}, g::Vector{Float64})
        nll, grad = nll_and_grad(params, X, Y)
        g[:] = grad
        return nll
    end

    func = OnceDifferentiable(obj, grad!, obj_and_grad!)
    βinit = llsq(X,Y; bias=false)
    resid = Y - X*βinit
    Ωinit = cov(resid)
    αinit = zeros(k)
    νinit = 1.0
    ρinit, Ainit, ηinit, logνinit = dplist2optpar(Ωinit, αinit, νinit)

    params = write_params(βinit, ρinit, Ainit, ηinit, logνinit)
    print_params(βinit, ρinit, Ainit, ηinit, logνinit)
    results = optimize(func, params, method, Optim.Options(;kwargs...))
    print(results)
    (β, ρ, A, η, ν) = read_params(Optim.minimizer(results), p, k)
    D = Diagonal(exp(-2*ρ))
    Ωinv = A'D*A
    Ω = inv(Ωinv)
    ω = Diagonal(sqrt(diag(Ω)))
    α = ω * η
    mat=PDMat(Symmetric(Ω))
    dist = MvSkewTDist(zeros(k), mat, α, ν)
    return β, dist
end

# Fit a multivariate Skew Normal distribution
# directly to a set of observations
#
# Y = (n x k) observation matrix (each row correspond to one output observation)
function fit_MvSkewTDist(Y::Matrix{Float64}; kwargs...)
    β, dist = fit_MvSkewTDist(ones(size(Y,1),1), Y; kwargs...)
    return MvSkewTDist(vec(β), dist.Ω, dist.α, dist.df)
end


using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

@inline function exp_ps_m5(A)
    T=promote_type(eltype(A),Float64)
    A_copy=similar(A,T); A_copy .= A;
    return exp_ps_m5!(A_copy)
end

@inline function exp_ps_m5!(A)
    T=promote_type(eltype(A),Float64) # Make it work for many 'bigger' types (matrices and scalars)
    # max_memslots=6
    n=size(A,1)
    # The first slots are precomputed nodes [:A]
    memslots2 = similar(A,T)
    memslots3 = similar(A,T)
    memslots4 = similar(A,T)
    memslots5 = similar(A,T)
    memslots6 = similar(A,T)
    # Assign precomputed nodes memslots
    memslots1=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    value_one=ValueOne()
    # Computation order: A2 A3 A4 P2 C1 P1 C0 P0
    # Computing A2 with operation: mult
    mul!(memslots2,memslots1,memslots1)
    # Computing A3 with operation: mult
    mul!(memslots3,memslots1,memslots2)
    # Computing A4 with operation: mult
    mul!(memslots4,memslots1,memslots3)
    # Computing P2 = x*I+x*A+x*A2+x*A3+x*A4
    coeff1=2.48015873015873e-5
    coeff2=2.7557319223985893e-6
    coeff3=2.755731922398589e-7
    coeff4=2.505210838544172e-8
    coeff5=2.08767569878681e-9
    memslots5 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots4
    mul!(memslots5,true,I*coeff1,true,true)
    # Computing C1 with operation: mult
    mul!(memslots6,memslots5,memslots4)
    # Deallocating P2 in slot 5
    # Computing P1 = x*I+x*A+x*A2+x*A3+x*C1
    coeff1=0.041666666666666664
    coeff2=0.008333333333333333
    coeff3=0.001388888888888889
    coeff4=0.0001984126984126984
    coeff5=1.0
    # Smart lincomb recycle C1
    memslots6 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots6
    mul!(memslots6,true,I*coeff1,true,true)
    # Computing C0 with operation: mult
    mul!(memslots5,memslots6,memslots4)
    # Deallocating P1 in slot 6
    # Deallocating A4 in slot 4
    # Computing P0 = x*I+x*A+x*A2+x*A3+x*C0
    coeff1=1.0
    coeff2=1.0
    coeff3=0.5
    coeff4=0.16666666666666666
    coeff5=1.0
    # Smart lincomb recycle A
    memslots1 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3 .+ coeff5.*memslots5
    mul!(memslots1,true,I*coeff1,true,true)
    # Deallocating A2 in slot 2
    # Deallocating A3 in slot 3
    # Deallocating C0 in slot 5
    return memslots1 # Returning P0
end


"""
    timeevolution.schroedinger(tspan, psi0, H; fout)

Integrate Schroedinger equation.

# Arguments
* `tspan`: Vector specifying the points of time for which output should be displayed.
* `psi0`: Initial state vector (can be a bra or a ket).
* `H`: Arbitrary operator specifying the Hamiltonian.
* `fout=nothing`: If given, this function `fout(t, psi)` is called every time
        an output should be displayed. ATTENTION: The state `psi` is neither
        normalized nor permanent! It is still in use by the ode solver and
        therefore must not be changed.
"""
function schroedinger(tspan, psi0::T, H::AbstractOperator{B,B};
                fout::Union{Function,Nothing}=nothing,
                kwargs...) where {B<:Basis,T<:StateVector{B}}
    tspan_ = convert(Vector{Float64}, tspan)
    dschroedinger_(t::Float64, psi::T, dpsi::T) = dschroedinger(psi, H, dpsi)
    x0 = psi0.data
    state = T(psi0.basis, psi0.data)
    dstate = T(psi0.basis, psi0.data)
    integrate(tspan_, dschroedinger_, x0, state, dstate, fout; kwargs...)
end

"""
    timeevolution.schroedinger_dynamic(tspan, psi0, f; fout)

Integrate time-dependent Schroedinger equation.

# Arguments
* `tspan`: Vector specifying the points of time for which output should be displayed.
* `psi0`: Initial state vector (can be a bra or a ket).
* `f`: Function `f(t, psi) -> H` returning the time and or state dependent Hamiltonian.
* `fout=nothing`: If given, this function `fout(t, psi)` is called every time
        an output should be displayed. ATTENTION: The state `psi` is neither
        normalized nor permanent! It is still in use by the ode solver and
        therefore must not be changed.
"""
function schroedinger_dynamic(tspan, psi0::T, f::Function;
                fout::Union{Function,Nothing}=nothing,
                kwargs...) where T<:StateVector
    tspan_ = convert(Vector{Float64}, tspan)
    dschroedinger_(t::Float64, psi::T, dpsi::T) = dschroedinger_dynamic(t, psi, f, dpsi)
    x0 = psi0.data
    state = Ket(psi0.basis, psi0.data)
    dstate = Ket(psi0.basis, psi0.data)
    integrate(tspan_, dschroedinger_, x0, state, dstate, fout; kwargs...)
end

recast!(x::D, psi::StateVector{B,D}) where {B<:Basis, D<:Vector{ComplexF64}} = (psi.data = x);
recast!(psi::StateVector{B,D}, x::D) where {B<:Basis, D<:Vector{ComplexF64}} = nothing

function dschroedinger(psi::Ket{B}, H::AbstractOperator{B,B}, dpsi::Ket{B}) where B<:Basis
    QuantumOpticsBase.gemv!(complex(0,-1.), H, psi, complex(0.), dpsi)
    return dpsi
end

function dschroedinger(psi::Bra{B}, H::AbstractOperator{B,B}, dpsi::Bra{B}) where B<:Basis
    QuantumOpticsBase.gemv!(complex(0,1.), psi, H, complex(0.), dpsi)
    return dpsi
end

function dschroedinger_dynamic(t::Float64, psi0::T, f::Function, dpsi::T) where T<:StateVector
    H = f(t, psi0)
    dschroedinger(psi0, H, dpsi)
end

function check_schroedinger(psi::Ket, H::AbstractOperator)
    check_multiplicable(H, psi)
    check_samebases(H)
end

function check_schroedinger(psi::Bra, H::AbstractOperator)
    check_multiplicable(psi, H)
    check_samebases(H)
end


#####
##### CPU horizontally periodic pressure solver
#####

function HorizontallyPeriodicPressureSolver(::CPU, grid, pressure_bcs, planner_flag=FFTW.PATIENT)
    kx², ky², kz² = generate_discrete_eigenvalues(grid, pressure_bcs)
    wavenumbers = (kx² = kx², ky² = ky², kz² = kz²)

    # Storage for RHS and Fourier coefficients is hard-coded to be Float64
    # because of precision issues with Float32.
    # See https://github.com/climate-machine/Oceananigans.jl/issues/55
    storage = zeros(Complex{Float64}, grid.Nx, grid.Ny, grid.Nz)

    @debug "Planning transforms for PressureSolver{HorizontallyPeriodic, CPU}..."
    x_bc, y_bc, z_bc = pressure_bcs.x.left, pressure_bcs.y.left, pressure_bcs.z.left
    FFTxy!  = plan_forward_transform(storage, x_bc, [1, 2], planner_flag)
    DCTz!   = plan_forward_transform(storage, z_bc, 3, planner_flag)
    IFFTxy! = plan_backward_transform(storage, x_bc, [1, 2], planner_flag)
    IDCTz!  = plan_backward_transform(storage, z_bc, 3, planner_flag)
    @debug "Planning transforms for PressureSolver{HorizontallyPeriodic, CPU} done!"

    transforms = ( FFTxy! =  FFTxy!,  DCTz! =  DCTz!,
                  IFFTxy! = IFFTxy!, IDCTz! = IDCTz!)

    return PressureSolver(HorizontallyPeriodic(), CPU(), wavenumbers, storage, transforms, nothing)
end

"""
    solve_poisson_equation!(solver::PressureSolver{HorizontallyPeriodic, CPU}, grid)

Solve Poisson equation on a uniform staggered grid (Arakawa C-grid) with
appropriate boundary conditions as specified by `solver.bcs` using planned FFTs
and DCTs. The right-hand-side RHS is stored in solver.storage which the solver
mutates to produce the solution, so it will also be stored in solver.storage.

We should describe the algorithm in detail in the documentation.
"""
function solve_poisson_equation!(solver::PressureSolver{HorizontallyPeriodic, CPU}, grid)
    Nx, Ny, Nz, _ = unpack_grid(grid)
    kx², ky², kz² = solver.wavenumbers.kx², solver.wavenumbers.ky², solver.wavenumbers.kz²

    # We can use the same storage for the RHS and the solution ϕ.
    RHS, ϕ = solver.storage, solver.storage

    solver.transforms.DCTz!  * RHS  # Calculate DCTᶻ(f) in place.
    solver.transforms.FFTxy! * RHS  # Calculate FFTˣʸ(f) in place.

    # Solve the discrete Poisson equation in spectral space. We are computing
    # the Fourier coefficients of the solution from the Fourier coefficients
    # of the RHS.
    @. ϕ = -RHS / (kx² + ky² + kz²)

    # Setting DC component of the solution (the mean) to be zero. This is also
    # necessary because the source term to the Poisson equation has zero mean
    # and so the DC component comes out to be ∞.
    ϕ[1, 1, 1] = 0

    solver.transforms.IFFTxy! * ϕ  # Calculate IFFTˣʸ(ϕ) in place.
    solver.transforms.IDCTz!  * ϕ  # Calculate IDCTᶻ(ϕ) in place.

    # Must normalize by 2N for each dimension transformed via FFTW.REDFT.
    @. ϕ = ϕ / (2Nz)

    return nothing
end

#####
##### GPU horizontally periodic pressure solver
#####

function HorizontallyPeriodicPressureSolver(::GPU, grid, pressure_bcs, no_args...)
    Nx, Ny, Nz, _ = unpack_grid(grid)

    kx², ky², kz² = generate_discrete_eigenvalues(grid, pressure_bcs) .|> CuArray
    wavenumbers = (kx² = kx², ky² = ky², kz² = kz²)

    kz⁺ = reshape(0:Nz-1,       1, 1, Nz)
    kz⁻ = reshape(0:-1:-(Nz-1), 1, 1, Nz)

    ω_4Nz⁺ = ω.(4Nz, kz⁺) |> CuArray
    ω_4Nz⁻ = ω.(4Nz, kz⁻) |> CuArray

    # The zeroth coefficient of the IDCT (DCT-III or FFTW.REDFT01) is not
    # multiplied by 2. For some reason, we only need to account for this when
    # doing a 1D IDCT (for PPN boundary conditions) but not for a 2D IDCT. It's
    # possible that the masks are effectively doing this job.
    ω_4Nz⁻[1] *= 1/2

    constants = (ω_4Nz⁺ = ω_4Nz⁺, ω_4Nz⁻ = ω_4Nz⁻)

    # Storage for RHS and Fourier coefficients is hard-coded to be Float64
    # because of precision issues with Float32.
    # See https://github.com/climate-machine/Oceananigans.jl/issues/55
    storage = zeros(Complex{Float64}, Nx, Ny, Nz) |> CuArray

    @debug "Planning transforms for PressureSolver{HorizontallyPeriodic, GPU}..."
    x_bc, y_bc, z_bc = pressure_bcs.x.left, pressure_bcs.y.left, pressure_bcs.z.left
    FFTxy!  = plan_forward_transform(storage, x_bc, [1, 2])
    FFTz!   = plan_forward_transform(storage, z_bc, 3)
    IFFTxy! = plan_backward_transform(storage, x_bc, [1, 2])
    IFFTz!  = plan_backward_transform(storage, z_bc, 3)
    @debug "Planning transforms for PressureSolver{HorizontallyPeriodic, GPU} done!"

    transforms = ( FFTxy! =  FFTxy!,  FFTz! =  FFTz!,
                  IFFTxy! = IFFTxy!, IFFTz! = IFFTz!)

    return PressureSolver(HorizontallyPeriodic(), GPU(), wavenumbers, storage, transforms, constants)
end

"""
    solve_poisson_equation!(solver::PressureSolver{HorizontallyPeriodic, GPU}, grid)

Similar to solve_poisson_equation!(solver::HPPS{CPU}, grid) except that since
the discrete cosine transform is not available through cuFFT, we perform our
own fast cosine transform (FCT) via an algorithm that utilizes the FFT.

Note that for the FCT algorithm to work, the input must have been permuted along
the dimension the FCT is to be calculated by ordering the odd elements first
followed by the even elements. For example,

    [a, b, c, d, e, f, g, h] -> [a, c, e, g, h, f, d, b]

The output will be permuted in this way and so the permutation must be undone.

We should describe the algorithm in detail in the documentation.
"""
function solve_poisson_equation!(solver::PressureSolver{HorizontallyPeriodic, GPU}, grid)
    kx², ky², kz² = solver.wavenumbers.kx², solver.wavenumbers.ky², solver.wavenumbers.kz²
    ω_4Nz⁺, ω_4Nz⁻ = solver.constants.ω_4Nz⁺, solver.constants.ω_4Nz⁻

    # We can use the same storage for the RHS and the solution ϕ.
    RHS, ϕ = solver.storage, solver.storage

    # Calculate DCTᶻ(f) in place using the FFT.
    solver.transforms.FFTz! * RHS
    @. RHS = 2 * real(ω_4Nz⁺ * RHS)

    solver.transforms.FFTxy! * RHS  # Calculate FFTˣʸ(f) in place.

    # Solve the discrete Poisson equation in spectral space. We are essentially
    # computing the Fourier coefficients of the solution from the Fourier
    # coefficients of the RHS.
    @. ϕ = -RHS / (kx² + ky² + kz²)

    # Setting DC component of the solution (the mean) to be zero. This is also
    # necessary because the source term to the Poisson equation has zero mean
    # and so the DC component comes out to be ∞.
    ϕ[1, 1, 1] = 0

    solver.transforms.IFFTxy! * ϕ  # Calculate IFFTˣʸ(ϕ̂) in place.

    # Calculate IDCTᶻ(ϕ̂) in place using the FFT.
    ϕ .*= ω_4Nz⁻
    solver.transforms.IFFTz! * ϕ

    return nothing
end


abstract type AbstractGPLayer end;

struct SVGPLayer{GPType} <: AbstractGPLayer
    dim::Int
    gps::Vector{GPType}
    function SVGPLayer(dim,μ,Σ,Z,kernel,μ₀)
        gps = [SVGP_Base(copy(μ),copy(Σ),copy(Z),deepcopy(kernel),deepcopy(μ₀)) for _ in 1:dim]
        new{SVGP_Base}(dim,gps)
    end
end

Base.length(l::SVGPLayer) = l.dim
Base.getindex(l::SVGPLayer,i::Int) = l.gps[i]

function propagate(l::AbstractGPLayer,samples::AbstractVector{<:AbstractArray})
    compute_invK.(l.gps)
    propagate.([l],samples)
end
function propagate(l::AbstractGPLayer,samples::AbstractArray{<:Real})
    ## Return the parametrization m + eps *sqrt(S)
    mu(l,samples) + randn(size(samples,1),l.dim).*sqrt.(diag_sigma(l,samples))
end

mu(l::AbstractGPLayer,x) = reduce(hcat,[mu(gp,x) for gp in l.gps])
diag_sigma(l::AbstractGPLayer,x) = reduce(hcat,[diag_sigma(gp,x) for gp in l.gps])

kl_divergence(l::AbstractGPLayer) = sum(kl_divergence(gp) for gp in l.gps)

abstract type AbstractGPBase end;

mutable struct SVGP_Base{T} <: AbstractGPBase
    dim::Int
    μ::Vector{T}
    Σ::Matrix{T}
    μ₀::AGP.PriorMean
    kernel::KernelFunctions.Kernel
    Z::Matrix
    invKmm::Matrix{T}
    function SVGP_Base(dim,Z::Matrix{T},kernel=SqExponentialKernel(),μ₀=ZeroMean()) where {T}
        new{T}(dim,zeros(T,dim),diagm(ones(T,dim)),μ₀,kernel,Z,diagm(ones(T,dim)))
    end
    function SVGP_Base(μ::Vector{T},Σ::Matrix{T},Z::Matrix,kernel::KernelFunctions.Kernel,μ₀::AGP.PriorMean) where {T}
        new{T}(length(μ),μ,Σ,μ₀,kernel,Z,diagm(ones(T,length(μ))))
    end
end

function compute_invK(gp::SVGP_Base)
    gp.invKmm = inv(KernelFunctions.kernelmatrix(gp.kernel,gp.Z,obsdim=1)+1e-5I)
end

function κ(gp::SVGP_Base,X::AbstractArray,invKmm)
    ## Return the K_XZ K
    try
        KernelFunctions.kernelmatrix(gp.kernel,X,gp.Z,obsdim=1)*invKmm
    catch e
        @show size(X)
        @show size(gp.Z)
        rethrow(e)
    end
end

_k(gp::AbstractGPBase,x,y) = KernelFunctions.kernelmatrix(gp.kernel,x,y,obsdim=1)
_diagk(gp::AbstractGPBase,x) = KernelFunctions.kerneldiagmatrix(gp.kernel,x,obsdim=1).+1e-5

mu(gp::SVGP_Base,x)= gp.μ₀ + κ(gp,x,gp.invKmm)*(gp.μ-gp.μ₀)
diag_sigma(gp::SVGP_Base,x) = diag_sigma(gp,x,κ(gp,x,gp.invKmm))
diag_sigma(gp::SVGP_Base,x,κ) = _diagk(gp,x) - AGP.opt_diag(κ*(gp.invKmm-gp.Σ),κ)

kl_divergence(gp::SVGP_Base) = 0.5*(-logdet(gp.Σ)-logdet(gp.invKmm)+AGP.opt_trace(gp.invKmm,gp.Σ)+dot(gp.μ-gp.μ₀,gp.invKmm*(gp.μ-gp.μ₀))-length(gp.μ))


function plotχ2(hval, cval, χ2, βZQ, βDQ, βDQQ, βQQ)
    χ2min = minimum(vec(χ2))
    ξHmin = hval[argmin(χ2)[1]]
    ξCmin = cval[argmin(χ2)[2]]

    clev = log10.(χ2min .+ [2.3, 6.7])
    plt = plot(framestyle=:box, legend=nothing, xguide="ξH", yguide="ξC", title="$(methyls[resi])", xlims=(-.3, .3), ylims=(-.3, .3))

    plot!(plt, hval,   hval .+ Measurements.value(sqrt(βZQ-(ξCcsa[resi]-ξHcsa[resi])^2)),
        ribbon=0hval .+ Measurements.uncertainty(sqrt(βZQ-(ξCcsa[resi]-ξHcsa[resi])^2)), c=1)
    plot!(plt, hval,   hval .- Measurements.value(sqrt(βZQ-(ξCcsa[resi]-ξHcsa[resi])^2)),
        ribbon=0hval .+ Measurements.uncertainty(sqrt(βZQ-(ξCcsa[resi]-ξHcsa[resi])^2)), c=1)
    plot!(plt, hval,  -hval .+ Measurements.value(sqrt(βDQ-(ξCcsa[resi]+ξHcsa[resi])^2)),
        ribbon=0hval .+ Measurements.uncertainty(sqrt(βDQ-(ξCcsa[resi]+ξHcsa[resi])^2)), c=2)
    plot!(plt, hval,  -hval .- Measurements.value(sqrt(βDQ-(ξCcsa[resi]+ξHcsa[resi])^2)),
        ribbon=0hval .+ Measurements.uncertainty(sqrt(βDQ-(ξCcsa[resi]+ξHcsa[resi])^2)), c=2)
    plot!(plt, hval,  3hval .+ Measurements.value(sqrt(βDQQ-(ξCcsa[resi]-3ξHcsa[resi])^2)),
        ribbon=0hval .+ Measurements.uncertainty(sqrt(βDQQ-(ξCcsa[resi]-3ξHcsa[resi])^2)), c=3)
    plot!(plt, hval,  3hval .- Measurements.value(sqrt(βDQQ-(ξCcsa[resi]-3ξHcsa[resi])^2)),
        ribbon=0hval .+ Measurements.uncertainty(sqrt(βDQQ-(ξCcsa[resi]-3ξHcsa[resi])^2)), c=3)
    plot!(plt, hval, -3hval .+ Measurements.value(sqrt(βQQ-(ξCcsa[resi]+3ξHcsa[resi])^2)),
        ribbon=0hval .+ Measurements.uncertainty(sqrt(βQQ-(ξCcsa[resi]+3ξHcsa[resi])^2)), c=4)
    plot!(plt, hval, -3hval .- Measurements.value(sqrt(βQQ-(ξCcsa[resi]+3ξHcsa[resi])^2)),
        ribbon=0hval .+ Measurements.uncertainty(sqrt(βQQ-(ξCcsa[resi]+3ξHcsa[resi])^2)), c=4)

    contour!(plt, hval, cval, log10.(χ2'), levels=clev, c=:red)

    return plt
end

function plotlinefits(ξH, ξC, resi)
    βZQ = Measurements.value((ξCcsa[resi]-ξHcsa[resi])^2 + (ξC-ξH)^2)
    βDQ = Measurements.value((ξCcsa[resi]+ξHcsa[resi])^2 + (ξC+ξH)^2)
    βDQQ = Measurements.value((ξCcsa[resi]-3ξHcsa[resi])^2 + (ξC-3ξH)^2)
    βQQ = Measurements.value((ξCcsa[resi]+3ξHcsa[resi])^2 + (ξC+3ξH)^2)

    plt = plot(frame=:box, legend=:topleft, xguide="B₀² / T²", yguide="Relaxation rate / s⁻¹")
    c0 = [10.0]

    scatter!(plt, B0², zq[resi,:],label="ZQ")
    fit = LsqFit.lmfit(c -> resid(B0², zq[resi,:], βZQ, c), c0, Float64[])
    c = fit.param[1]
    B2 = [0; B0²]
    plot!(plt, B2, line(B2, βZQ, c),primary=false)

    scatter!(plt, B0², dq[resi,:],label="DQ")
    fit = LsqFit.lmfit(c -> resid(B0², dq[resi,:], βDQ, c), c0, Float64[])
    c = fit.param[1]
    B2 = [0; B0²]
    plot!(plt, B2, line(B2, βDQ, c),primary=false)

    scatter!(plt, B0², dqq[resi,:],label="DQQ")
    fit = LsqFit.lmfit(c -> resid(B0², dqq[resi,:], βDQQ, c), c0, Float64[])
    c = fit.param[1]
    B2 = [0; B0²]
    plot!(plt, B2, line(B2, βDQQ, c),primary=false)

    scatter!(plt, B0², qq[resi,:],label="QQ")
    fit = LsqFit.lmfit(c -> resid(B0², qq[resi,:], βQQ, c), c0, Float64[])
    c = fit.param[1]
    B2 = [0; B0²]
    plot!(plt, B2, line(B2, βQQ, c),primary=false)
    ylims!(0, ylims()[2])
    return plt
end


module Shearheating

# This implements different methods to specify Shearheating of rocks
#
# If you want to add a new method here, feel free to do so.
# Remember to also export the function name in GeoParams.jl (in addition to here)

using Parameters, LaTeXStrings, Unitful
using ..Units
using GeoParams: AbstractMaterialParam
import Base.show

abstract type AbstractShearheating <: AbstractMaterialParam end

export  ComputeShearheating, ComputeShearheating!,  # calculation routines
        ConstantShearheating                        # constant

# Constant Shearheating -------------------------------------------------------
"""
    ConstantShearheating(Χ=0.0NoUnits)

Set the shear heating efficiency [0-1] parameter
```math
Χ  = cst
```
where ``\\Chi`` is the shear heating efficiency [NoUnits]

Shear heating is computed as
```math
H_s = \\Chi \\cdot \\tau_{ij}(\\dot{\\varepsilon}_{ij} - \\dot{\\varepsilon}^{el}_{ij})
```

"""
@with_kw_noshow mutable struct ConstantShearheating <: AbstractShearheating
    equation::LaTeXString   =   L"\H_s = \Chi \tau_{ij}(\dot{\varepsilon}_{ij} - \dot{\varepsilon}^{el}_{ij})"
    Χ::GeoUnit              =   0NoUnits
end

ConstantShearheating(x::Number) = ConstantShearheating(Χ = x*NoUnits)   # allow

# In-place routine
function ComputeShearheating!(H_s, τ, ε, ε_el, s::ConstantShearheating)
    @unpack Χ   = s

    if isnothing(ε_el)
        H_s = Value(Χ)*sum( τ .* ε  )
    else
        H_s = Value(Χ)*sum( τ .* (ε .- ε_el) )
    end

end

# Calculation routine
function ComputeShearheating(τ, ε, ε_el, s::ConstantShearheating)
    @unpack Χ   = s

    if isnothing(ε_el)
        H_s = Value(Χ)*sum( τ .* ε  )
    else
        H_s = Value(Χ)*sum( τ .* (ε .- ε_el) )
    end

    return H_s
end

# Print info
function show(io::IO, g::ConstantShearheating)
    print(io, "Shear heating: H_s = $(g.Χ.val) τ_ij*(ε_ij - ε^el_ij)")
end
#-------------------------------------------------------------------------

# Help info for the calculation routines
"""
    H_s = ComputeShearheating(τ, ε, ε_el,  s:<AbstractShearheating)

Computes the shear heating source term

```math
H_s = \\Chi \\cdot \\tau_{ij} ( \\dot{\\varepsilon}_{ij} - \\dot{\\varepsilon}^{el}_{ij})
```

# Parameters
- ``\\Chi`` : The efficiency of shear heating (between 0-1)
- ``\\tau_{ij}`` : The full deviatoric stress tensor [4 components in 2D; 9 in 3D]
- ``\\dot{\\varepsilon}_{ij}`` : The full deviatoric strainrate tensor
- ``\\dot{\\varepsilon}^{el}_{ij}`` : The full elastic deviatoric strainrate tensor

"""
ComputeShearheating(τ, ε, ε_el, s::AbstractShearheating)

"""
    H_s = ComputeShearheating(τ, ε,  s:<AbstractShearheating)

Computes the shear heating source term when there is no elasticity

```math
H_s = \\Chi \\cdot \\tau_{ij}  \\dot{\\varepsilon}_{ij}
```

# Parameters
- ``\\Chi`` : The efficiency of shear heating (between 0-1)
- ``\\tau_{ij}`` : The full deviatoric stress tensor [4 components in 2D; 9 in 3D]
- ``\\dot{\\varepsilon}_{ij}`` : The full deviatoric strainrate tensor
"""
ComputeShearheating(τ::Any, ε::Any, s::AbstractShearheating) = ComputeShearheating(τ, ε, nothing, s::ConstantShearheating)

"""
    ComputeShearheating!(H_s, τ, ε, ε_el,  s:<AbstractShearheating)

Computes the shear heating source term in-place

```math
H_s = \\Chi \\cdot \\tau_{ij} ( \\dot{\\varepsilon}_{ij} - \\dot{\\varepsilon}^{el}_{ij})
```

# Parameters
- ``\\Chi`` : The efficiency of shear heating (between 0-1)
- ``\\tau_{ij}`` : The full deviatoric stress tensor [4 components in 2D; 9 in 3D]
- ``\\dot{\\varepsilon}_{ij}`` : The full deviatoric strainrate tensor
- ``\\dot{\\varepsilon}^{el}_{ij}`` : The full elastic deviatoric strainrate tensor

*NOTE:*
The shear heating terms require the full deviatoric stress & strain rate tensors, i.e.:

```math
2D: \\tau_{ij} = \\left(
                \\begin{matrix}
                    \\tau_{xx} & \\tau_{xz} \\\\
                    \\tau_{zx} & \\tau_{zz}
                \\end{matrix}
            \\right)
```
Since ``\\tau_{zx}=\\tau_{xz}``, most geodynamic codes only take one of the terms into account; shear heating requires all components to be used!
"""
ComputeShearheating!(H_s, τ, ε, ε_el, s::AbstractShearheating)

"""
    ComputeShearheating!(H_s, τ, ε,  s:<AbstractShearheating)

Computes the shear heating source term `H_s` in-place when there is no elasticity

```math
H_s = \\Chi \\cdot \\tau_{ij}  \\dot{\\varepsilon}_{ij}
```

# Parameters
- ``\\Chi`` : The efficiency of shear heating (between 0-1)
- ``\\tau_{ij}`` : The full deviatoric stress tensor [4 components in 2D; 9 in 3D]
- ``\\dot{\\varepsilon}_{ij}`` : The full deviatoric strainrate tensor
"""
ComputeShearheating!(H_s::Any, τ::Any, ε::Any, s::AbstractShearheating) = ComputeShearheating!(H_s, τ, ε, nothing, s::ConstantShearheating)

end


using Perceptrons

# training a linear perceptron (solving the OR problem)
X_train        = [1.0 1.0; 0.0 1.0; 1.0 0.0; 0.0 0.0]
Y_train        = [1; 1; 1; 0.0]
X_test         = [.8 .9; .01 1; .9 0.2; 0.1 0.2]

model          = Perceptrons.fit(X_train,Y_train)
Y_pred         = Perceptrons.predict(model,X_test)

println("[Perceptron] accuracy : $(acc(Y_train,Y_pred))")

# training a voted perceptron (solving the OR problem)
model   = Perceptrons.fit(X_train,Y_train,centralize=true,mode="voted")
Y_pred  = Perceptrons.predict(model,X_test)

println("[Voted Perceptron] accuracy : $(acc(Y_train,Y_pred))")

# training a averaged perceptron (solving the OR problem)
model   = Perceptrons.fit(X_train,Y_train,centralize=true,mode="averaged")
Y_pred  = Perceptrons.predict(model,X_test)

println("[Averaged Perceptron] accuracy : $(acc(Y_train,Y_pred))")

# training a kernel perceptron (solving the XOR problem)
X_train = [1.0 1.0; 0.0 1.0; 1.0 0.0; 0.0 0.0]
Y_train = [0.0 ; 1.0; 1.0; 0.0]
X_test  = X_train .+ .03 # adding noise

model   = Perceptrons.fit(X_train,Y_train,centralize=true,mode="kernel",kernel="rbf",width=.01)
Y_pred  = Perceptrons.predict(model,X_test)

println("[Kernel Perceptron] accuracy : $(acc(Y_train,Y_pred))")

# if you want to save your model
Perceptrons.save(model,filename=joinpath(homedir(),"perceptron_model.jld"))

# if you want to load back your model
model = Perceptrons.load(filename=joinpath(homedir(),"perceptron_model.jld"))


using Distributions
using IndependentHypothesisWeighting
using Random
using MLDataPattern
import MLDataPattern: kfolds

using StatsBase
using CategoricalArrays

Random.seed!(1)
m= 10000
kf = kfolds(MLDataPattern.shuffleobs(1:m), 5)
Xs = CategoricalVector(sample(1:2, m))
Ps = rand(BetaUniformMixtureModel(0.7, 0.2), m) .* (Xs.==1) .+ rand(Uniform(), m) .* (Xs.==2)

typeof(Xs)
typeof(Ps)

α = 0.1
sum(adjust(Ps, BenjaminiHochberg()) .<= α ) # 580 discoveries

ihw_grenander = IHW(weight_learner = GrenanderLearner(), α = α)
 # 677 discoveries

ihw_grenander_fit1 = fit(ihw_grenander, Ps, Xs, kf)
sum(adjust(ihw_grenander_fit1) .<= α)

Random.seed!(1)
ihw_grenander_fit2 = fit(ihw_grenander, Ps, Xs)
sum(adjust(ihw_grenander_fit2) .<= α)


## TODO: pointmass rules
## TODO: save sigma vectors in meta to limit allocations

@rule Flow(:out, Marginalisation) (m_in::MvNormalMeanCovariance, meta::FlowMeta{M,Linearization}) where { M } = begin

    # extract parameters
    μ_in, Σ_in = mean_cov(m_in)

    # extract model
    model = getmodel(meta)

    # calculate new parameters
    (μ_out, J)   = forward_jacobian(model, μ_in)
    Σ_out   = J * Σ_in * J'

    # return distribution
    return MvNormalMeanCovariance(μ_out, Σ_out)

end

@rule Flow(:out, Marginalisation) (m_in::MvNormalMeanPrecision, meta::FlowMeta{M,Linearization}) where { M } = begin

    # extract parameters
    μ_in, Λ_in = mean_precision(m_in)

    # extract model
    model = getmodel(meta)

    # calculate new parameters
    μ_out   = forward(model, μ_in)
    Ji      = inv_jacobian(model, μ_in)
    Λ_out   = Ji' * Λ_in * Ji

    # return distribution
    return MvNormalMeanPrecision(μ_out, Λ_out)

end

@rule Flow(:out, Marginalisation) (m_in::MvNormalWeightedMeanPrecision, meta::FlowMeta{M,Linearization}) where { M } = begin

    # extract parameters
    μ_in, Λ_in = mean_precision(m_in)

    # extract model
    model = getmodel(meta)

    # calculate new parameters
    μ_out   = forward(model, μ_in)
    Ji      = inv_jacobian(model, μ_in)
    Λ_out   = Ji' * Λ_in * Ji

    # return distribution
    return MvNormalMeanPrecision(μ_out, Λ_out)

end

@rule Flow(:out, Marginalisation) (m_in::MultivariateNormalDistributionsFamily, meta::FlowMeta{M,Unscented}) where { M } = begin

    # extract parameters
    μ_in, Σ_in = mean_cov(m_in)

    # extract model
    model = getmodel(meta)
    T = eltype(model)

    # extract parameters of linearization
    approximation = getapproximation(meta)
    λ  = getλ(approximation)
    L  = getL(approximation)
    Wm = getWm(approximation)
    Wc = getWc(approximation)

    # calculate sigma points/vectors
    sqrtΣ = sqrt((L + λ)*Σ_in)
    χ = Vector{Vector{T}}(undef, 2*L + 1)
    for k = 1:length(χ)
        χ[k] = copy(μ_in)
    end
    for l = 2:L+1
        χ[l]     .+= sqrtΣ[l-1,:]
        χ[L + l] .-= sqrtΣ[l-1,:]
    end

    # transform sigma points
    Y = forward.(model, χ)

    # calculate new parameters
    μ_out = zeros(T, L)
    Σ_out = zeros(T, L, L)
    for k = 1:2*L+1
        μ_out .+= Wm[k] .* Y[k]
    end
    for k = 1:2*L+1
        Σ_out .+= Wc[k] .* ( Y[k] - μ_out ) *  ( Y[k] - μ_out )'
    end

    # return distribution
    return MvNormalMeanCovariance(μ_out, collect(Hermitian(Σ_out)))

end


using Polynomials, LinearMaps
using LinearAlgebra, SparseArrays
mutable struct CustomBOSolver{T} <: COSMO.AbstractKKTSolver
    k::Int
    σ::T
    ρ1::T
    ρ2::T
    n::Int
    m::Int
    sqrt_dim::Int
    A1::SparseMatrixCSC{T, Int}
    A::SparseMatrixCSC{T, Int}
    P::SparseMatrixCSC{T, Int}
    tmp_n::Vector{T}

    function CustomBOSolver(P::AbstractMatrix{T}, A::AbstractMatrix{T}, σ::T, ρ) where T
        n = size(P, 1)
        m = size(A, 1) - n
        ρ1 = ρ[1];
        ρ2 = ρ[end];
        k = Int(round(real(roots(Poly([-n, 0, 1/2, 1/2]))[end])))
        @assert k*(k + 1)*k/2 == n
        sqrt_dim = Int(k*(k + 1)/2)
        A1 = kron(ones(1, k), SparseMatrixCSC{T, Int}(I, sqrt_dim, sqrt_dim))
        new{T}(k, σ, ρ1, ρ2, n, m, Int(k*(k + 1)/2), A1, A, P, zeros(T, n))
    end
end

function mul_A1!(y, x, S)
    y .= 0
    idx = 1
    @inbounds for i = 1:S.k, j = 1:S.sqrt_dim
        y[j] += x[idx]
        idx += 1
    end
    #@assert norm(S.A1*x - y) <= 1e-8
end

function mul_A1t!(y, x, S)
    idx = 1
    @inbounds for i = 1:S.k, j = 1:S.sqrt_dim
        y[idx] = x[j]
        idx += 1
    end
    #@assert norm(S.A1'*x - y) <= 1e-8
end

function solve!(S::CustomBOSolver, y::AbstractVector{T}, x::AbstractVector{T}) where {T}
    # Solves the (KKT) linear system
    # | P + σI     A'  |  |y1|  =  |x1|
    # | A        -I/ρ  |  |y23|  =  |x23|
    # x1,y1: R^n, x2/y2: R^m
    # where [y1; y2] := y, [x1; x2] := x
    y1 = view(y, 1:S.n); y2 = view(y, S.n+1:S.n+S.m); y3 = view(y, S.n+S.m+1:2*S.n+S.m);
    x1 = view(x, 1:S.n); x2 = view(x, S.n+1:S.n+S.m); x3 = view(x, S.n+S.m+1:2*S.n+S.m);

    #y2 .= (S.A1*(x1 + x3*S.ρ2)/(S.ρ2 + S.σ) - x2)./(1/S.ρ1 + S.k/(S.ρ2 + S.σ))
    #y1 .= (-S.A1'*y2 + x1 + x3*S.ρ2)/(S.ρ2 + S.σ)
    #y3 .= S.ρ2* (y1 - x3)

    @. S.tmp_n = (x1 + x3*S.ρ2)/(S.ρ2 + S.σ);
    mul_A1!(y2, S.tmp_n, S)
    @. y2 = (y2 - x2)/(1/S.ρ1 + S.k/(S.ρ2 + S.σ))
    mul_A1t!(y1, y2, S)
    @. y1 = (-y1 + x1 + x3*S.ρ2)/(S.ρ2 + S.σ)
    @. y3 = S.ρ2* (y1 - x3)

    #=
    F = [S.P + S.σ*I SparseMatrixCSC(S.A'); S.A I]
    @inbounds @simd for i = (S.n + 1):size(F, 2)
        if i <= S.n + S.m
            ρ = S.ρ1
        else
            ρ = S.ρ2
        end
        F[i, i] = -1.0 / ρ
    end
    # show(stdout, "text/plain", F); println()
    @assert norm(y - F\x) <= 1e-7
    =#
end

function update_rho!(S::CustomBOSolver, ρ)
    S.ρ1 = ρ[1]
    S.ρ2 = ρ[end]
end


gen_op(CAT)

c  = @select CAT where Close;
cy = collapse(c, year, last)

# daily
Cd   = percentchange(c).values
GMMd = GMM(3, Cd, method=:kmeans, nInit=50, nIter=10)
MMd  = MixtureModel(GMMd)

#yearly
Cy   = percentchange(cy).values
GMMy = GMM(3, Cy, method=:kmeans, nInit=50, nIter=10)
MMy  = MixtureModel(GMMy)

function rday(n::Int=13089)
    rand(MMd, n)
end

function ryear(n::Int=51)
    rand(MMy, n)
end


### A Pluto.jl notebook ###
# v0.11.3

using Markdown
using InteractiveUtils

# ╔═╡ ed7c391e-d937-11ea-31e0-d5884018ca29
md"# Orthogonal Procrustes"

# ╔═╡ fd1fbc70-d936-11ea-2712-ed19c93c0c9f
begin
	using DelimitedFiles
	using Gadfly, Cairo, Colors
	using LinearAlgebra
	using Statistics

	set_default_plot_size(15cm,13cm)
	color_palette = Scale.color_discrete().f(7)[[1, 4, 6]]
	themes = [Theme(default_color=color,
			        panel_fill="white",
			        grid_color="lightgray",
                    background_color=parse(Colorant, "white")
			        ) for color in color_palette];
end

# ╔═╡ 3b0cfa3e-d937-11ea-2d53-255142465b91
B = readdlm("point_cloud.txt")

# ╔═╡ 3d667abc-d937-11ea-338d-e194f0379319
p1 = plot(x=B[1, :], y=B[2, :], Geom.point, themes[1],
         Guide.title("a point set, B"), themes[1],
         Coord.cartesian(fixed=true),
        )

# ╔═╡ 53c23896-d937-11ea-2e06-6b6331e0fe4c
draw(PNG("ptset_B.png", 5inch, 4inch, dpi=300), p1)

# ╔═╡ 7ca103a0-d937-11ea-0d5a-af0e3a9141a6
rotation_matrix2d(θ::Float64) = [cos(θ) -sin(θ); sin(θ) cos(θ)]

# ╔═╡ 83746456-d937-11ea-04a1-f15bf81a17d4
begin
	θ = π * 0.8
	R_known = rotation_matrix2d(θ)
end

# ╔═╡ 8b9acb16-d937-11ea-0e02-7f6bcec57932
begin
	ϵ = 0.01 # noise
	A = R_known * B .+ ϵ * randn(size(B)...)

	p2 = plot(x=A[1, :], y=A[2, :], Geom.point, themes[3],
	         Guide.title("point set, A"), themes[3],
	         Coord.cartesian(fixed=true),
	    )
end

# ╔═╡ a85a6086-d937-11ea-3fed-21f714cad837
draw(PNG("ptset_A.png", 5inch, 4inch, dpi=300), p2)

# ╔═╡ b04a9cfa-d937-11ea-38cd-e1a2a9d588e3
p3 = plot(layer(x=A[1, :], y=A[2, :], Geom.point, themes[3]),
         layer(x=B[1, :], y=B[2, :], Geom.point, themes[1]),
         Guide.title("orthogonal Procrustes"), themes[1],
         Coord.cartesian(fixed=true),
         Guide.manual_color_key("point set", ["A", "B"], [color_palette[3], color_palette[1]])
    )

# ╔═╡ b7f97a90-d937-11ea-1d20-a3752908fe8c
draw(PNG("before_alignment.png", 5inch, 4inch, dpi=300), p3)

# ╔═╡ bcb90346-d937-11ea-1bbb-2f481b298a95
p4 = plot(layer(x=A[1, :], y=A[2, :], Geom.point, themes[3]),
         layer(x=B[1, :], y=B[2, :], Geom.point, themes[1]),
         [layer(x=[A[1, i], B[1, i]], y=[A[2, i], B[2, i]],
            Geom.line, Theme(default_color=colorant"gray")) for i = 1:size(A)[2]]...,
         Guide.title("orthogonal Procrustes"), themes[1],
         Coord.cartesian(fixed=true),
         Guide.manual_color_key("point set", ["A", "B"], [color_palette[3], color_palette[1]])
    )

# ╔═╡ c17420c0-d937-11ea-3fad-0b62648de42f
draw(PNG("correspondence.png", 5inch, 4inch, dpi=300), p4)

# ╔═╡ c61c4986-d937-11ea-2423-51c2154ee671
F = svd(A * B')

# ╔═╡ c9fae814-d937-11ea-3e2a-8d010deb8f54
R = F.V * F.U'

# ╔═╡ ce49e546-d937-11ea-3c90-2150e15a7ce8
isapprox(R, rotation_matrix2d(-θ), rtol=0.01)

# ╔═╡ d4242594-d937-11ea-1a8e-c1b01e9563f8
A_transformed = R * A

# ╔═╡ d86f9a84-d937-11ea-2807-6de9156292c8
p5 = plot(layer(x=A_transformed[1, :], y=A_transformed[2, :], Geom.point, themes[3]),
         layer(x=B[1, :], y=B[2, :], Geom.point, themes[1]),
         Guide.title("orthogonal Procrustes"), themes[1],
         Coord.cartesian(fixed=true),
         Guide.manual_color_key("point set", ["RA", "B"], [color_palette[3], color_palette[1]])
    )

# ╔═╡ e461e8cc-d937-11ea-07f3-a1a8abd49594
draw(PNG("after_alignment.png", 5inch, 4inch, dpi=300), p5)

# ╔═╡ Cell order:
# ╟─ed7c391e-d937-11ea-31e0-d5884018ca29
# ╠═fd1fbc70-d936-11ea-2712-ed19c93c0c9f
# ╠═3b0cfa3e-d937-11ea-2d53-255142465b91
# ╠═3d667abc-d937-11ea-338d-e194f0379319
# ╠═53c23896-d937-11ea-2e06-6b6331e0fe4c
# ╠═7ca103a0-d937-11ea-0d5a-af0e3a9141a6
# ╠═83746456-d937-11ea-04a1-f15bf81a17d4
# ╠═8b9acb16-d937-11ea-0e02-7f6bcec57932
# ╠═a85a6086-d937-11ea-3fed-21f714cad837
# ╠═b04a9cfa-d937-11ea-38cd-e1a2a9d588e3
# ╠═b7f97a90-d937-11ea-1d20-a3752908fe8c
# ╠═bcb90346-d937-11ea-1bbb-2f481b298a95
# ╠═c17420c0-d937-11ea-3fad-0b62648de42f
# ╠═c61c4986-d937-11ea-2423-51c2154ee671
# ╠═c9fae814-d937-11ea-3e2a-8d010deb8f54
# ╠═ce49e546-d937-11ea-3c90-2150e15a7ce8
# ╠═d4242594-d937-11ea-1a8e-c1b01e9563f8
# ╠═d86f9a84-d937-11ea-2807-6de9156292c8
# ╠═e461e8cc-d937-11ea-07f3-a1a8abd49594


#   Unit "JoB.jl" of the Diagonalization.jl Package for Julia language
#
#   MIT License
#   Copyright (c) 2019,
#   Marco Congedo, CNRS, Grenoble, France:
#   https://sites.google.com/site/marcocongedo/home

# ? CONTENTS :
#   This unit implements a very general form of the OJoB algoithm
#   (Congedo et al., 2012)

# """
#  Orthogonal Joint Blond Source Separation (OJoB) iterative algorithm and
#  Nonorthogonal Joint Blond Source Separation (NoJoB) iterative algorithm:
#  'Orthogonal and Non-Orthogonal Joint Blind Source Separation in the
#  Least-Squares Sense'
#  by Marco Congedo, Ronald Phlypo, Jonas Chatel-Goldman, EUSIPCO 2012
#  https://hal.archives-ouvertes.fr/hal-00737835/document
#
# if `algo` is :OJoB, runs the OJoB algorithm
# if `algo` is :NoJoB, runs the NoJoB algorithm

#  These algorithms handle several use cases corresponding to several
#  combinations of the number of datasets (m) and number of observations (k):
#  1) m=1, k>2
#  2) m>1, k=1
#  3) m>1, k>1
#
#  They take as input either the data matrices or the covariance and cross-
#  covariance matrices. Argument `input` tells the algos to take as input
#  data matrices (`input=:d`) or cov/cross-cov matrices (`input=:d`)
#
#                                   INPUT
# ---------------------------------------------------------------------------
#             Data Matrices                  Cov/cross-cov matrices
# ---------------------------------------------------------------------------
#  1) a vector of k data matrices     an array (k, 1, 1) of covariance matrices
#  2) a vector of m data matrices     an array (1, m, m) of covariance matrices
#  3) k vectors of m data matrices    an array (k, m, m) of covariance matrices
#
#  If data matrices are given as input, the covariance matrices 𝒞 are computed
#  with optional keyword arguments `covEst`, `dims`, `meanX`, `trace1` and `w`.
#
#  OJoB seeks m orthogonal matrices 𝐔=U_1,...,U_m such that the sum of the
#  squares of the diagonal part of matrices U_i'𝒞(κij)U_j
#  is maximized for all i≠j=1:m and κ=1:k.
#  The covariance is maximized rotating each data matrix as: U_i' X_κi,
#  for all i=1:m and κ=1:k.
#
#  NoJoB seeks m non-singular matrices 𝐔=U_1,...,U_m with vectors
#  u_qi, for all i=1:n scaled such that the sum u_qi'(𝒞(κj)u_qju_qj^H𝒞(κj)^H)u_qi=1
#  where the sum is over κ=1:k and j=1:m, separetedly for all q=1:m
#  The covariance is maximized transforming each data matrix as: U_i' X_κi,
#  for all i=1:m and κ=1:k.
#
#  if `fullModel=true` is passed as optional keyword argument,
#  then also the sum of the squares of the diagonal part of matrices
#  U_i'𝒞(κij)U_j is maximized for all ``i=1:m`` and ``κ=1:k``.
#  This amounts to require the intra-covariance to be diagonalized as well.
#
#  if `whitening` = true is passed as optional keyword argument,
#  for OJoB:
#       the 𝐔=U_1,...,U_m are no longer constrained to be orthogonal;
#       instead they will be invertible and a scaling constraint is imposed on them
#       so as to satisfy mean_κ=1:k U_i'𝒞(κii)U_j=I for all i=1:m.
#  for NoJoB:
#       The same constraint is imposed
#
#  For both algorithms, given m whitening matrices W_i,
#  the covariance matrices will be transformed
#  as 𝒢(κij) = W_i'𝒞(κij)W_j, for all i,j=1:m and κ=1:k.
#  Then, the OJoB or NoJob algos will be run on the set 𝒢(κij)
#  and the final solutions will be given by U_iW_i, for all i=1:m
#  Using OjOB, this amounts to requiring a generalized Canonical Correlation
#  Analysis instead of a generalized Maximum Covariance Analysis.
#  By this pre-whitening the dimension of the problem can be reduced
#  requesting matrices W_i to be rectangular. This is achieved using
#  argument `eVar`. If eVar is an integer p, the matrices in 𝒢(κij) will all
#  be pxp. If it is equal to nothing the explained variance is set to
#  the default level (0.999), otherwise if eVar is a float in [0..1], this
#  will be taken as explained variance.
#  if m=1 the eigenvalues of mean_κ=1^k 𝒞(κ11) is considered, otherwise
#         the eigenvalues of mean_κ=1^k, i=1^m 𝒞(κii) is considered.
#  These eigenvalues are normalized to unit sum and the cumulated sum is
#  considered. Finally p is found on this accumulated regularized ev vector
#  using the `eVarMethod` function.
#
#  if sort=true (default) the column vectors of the 𝐔 matrices are reordered
#  so as to allow positive and sorted in descending order
#  diagonal elements of the products ``U_i'𝒞(κij)U_j`` as much as possible.
#  If the NoJoB algorithm is used the column vectors of the matrices U_i
#  are normalized to unit norm.
#  Note that if `whitening` is true the output matrices U_i will not have
#  unit norm columns as they are multiplied by the whiteners after being
#  scaled and sorted and before being returned.

#  By passing a matrix if k=1 or a vector of k matrices if k>1 as `init`,
#  you can smartly initialize `𝐔`, for example, when excluding some subjects
#  and running the OJoB again.
#
#  `tol` is the convergence to be attained.
#
#  `maxiter` is the maximum number of iterations allowed.
#
#  if `verbose`=true, the convergence attained at each iteration and other
#  information will be printed.
#
#  if `threaded`=true (default) and n>x and x>1, where x is the number
#  of threads Julia is instructed to use (the output of Threads.nthreads())
#  the algorithms run in multithreaded mode paralellising several comptations
#  over n.
#
#  Besides optionally multi-threaded, these algorithms heavely use BLAS.
#  Before running this function you may want to set the number of threades
#  Julia is instructed to use to the number of logical CPUs of your machine
#  and set `BLAS.set_num_threads(Sys.CPU_THREADS)`. See:
#  https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#Threads-1
# """
function JoB(𝐗::AbstractArray, m::Int, k::Int, input::Symbol, algo::Symbol, type;
              covEst   :: StatsBase.CovarianceEstimator=SCM,
              dims     :: Int64 = 1,
              meanX    :: Tmean = 0,
              trace1   :: Bool = false,
              w        :: Twf = ○,
          fullModel :: Bool = false,
          preWhite  :: Bool = false,
          sort      :: Bool = true,
          init      :: Union{Matrix, Nothing} = ○,
          tol       :: Real = 0.,
          maxiter   :: Int  = 1000,
          verbose   :: Bool = false,
          threaded  :: Bool = true,
      eVar     :: TeVaro = ○,
      eVarMeth :: Function = searchsortedfirst)

    k<3 && m<2 && throw(ArgumentError("Either `k` must be equal to at least 3 or `m` must be equal to at least 2"))
    #if m==1 fullModel=false end

    # if input ≠:d the input is supposed to be the cross-covariance matrices
    input==:d ? 𝒞=_crossCov(𝐗, m, k;
                            covEst=covEst, dims=dims, meanX=meanX) : 𝒞=𝐗

    𝒢=deepcopy(𝒞)
    if trace1 || w ≠ ○ _normalize!(𝒢, m, k, trace1, w) end

    if eVar isa Int && size(𝒢[1, 1, 1], 1) ≤ eVar ≤ 0
        eVar=size(𝒢[1, 1, 1], 1)
    end

    iter, conv, oldconv, 😋, conv_ = 1, 0., 1.01, false, 0.
    tol==0. ? tolerance = √eps(real(type)) : tolerance = tol

    # pre-whitening
    if preWhite
        # for the moment being computed on the average across k and m
        # if eVar is not an Int.
        # This way for all i=1:m the dimensionality reduction is fixed
        # for the case m=1 this is the classical whitening
        if !(eVar isa Int)
            G=ℍ(𝛍(𝒢[κ, i, i] for κ=1:k, i=1:m))
            p, arev = _getssd!(eVar, eigvals(G), size(𝒢[1, 1, 1], 1), eVarMeth)
            𝑾=[whitening(ℍ(𝛍(𝒢[κ, i, i] for κ=1:k)); eVar=p) for i=1:m]
        else
            𝑾=[whitening(ℍ(𝛍(𝒢[κ, i, i] for κ=1:k)); eVar=eVar) for i=1:m]
        end

        if m==1
            @inbounds for κ=1:k
                𝒢[κ, 1, 1] = 𝑾[1].F' * 𝒢[κ, 1, 1] * 𝑾[1].F
            end
        else
            @inbounds for κ=1:k, i=1:m-1, j=i+1:m
                𝒢[κ, i, j] = 𝑾[i].F' * 𝒢[κ, i, j] * 𝑾[j].F
                𝒢[κ, j, i] = 𝒢[κ, i, j]'
            end
            @inbounds for κ=1:k, i=1:m
                𝒢[κ, i, i] = 𝑾[i].F' * 𝒢[κ, i, i] * 𝑾[i].F
            end
        end
    end
    n=size(𝒢[1, 1, 1], 1)

    ⏩ = n>=Threads.nthreads() && Threads.nthreads()>1 && threaded

    # initialization of 𝐔[i], i=1:m, as the eigenvectors of sum_k,j(𝒢_k,i,j*𝒢_k,i,j')
    # see Eq. 17 of Congedo, Phlypo and Pham, 2011, or with the provided matrices
    # note: gemm supports complex matrices
    ggt(κ::Int, i::Int, j::Int) = BLAS.gemm('N', 'T', 𝒢[κ, i, j], 𝒢[κ, i, j])
    if m>1 # gmca, gcca, majd
        if init === nothing
            if algo==:NoJoB
                𝐔 = [Matrix{type}(I, n, n) for i=1:m]
            else
                if fullModel
                    𝐔 = [eigvecs(𝛍(ggt(κ, i, j) for j=1:m, κ=1:k)) for i=1:m]
                else
                    𝐔 = [eigvecs(𝛍(ggt(κ, i, j) for j=1:m, κ=1:k if j≠i)) for i=1:m]
                end
            end
        else
            𝐔 = init
        end
    else  # ajd
        if init === nothing
            if algo==:NoJoB
                𝐔 = [Matrix{type}(I, n, n)]
            else
                𝐔 = [eigvecs(𝛍(ggt(κ, 1, 1) for κ=1:k))]
            end
        else
            𝐔 = [init]
        end
    end

    function updateR!(η, i, j)  # 𝐑[η] += (𝒢[κ, i, j] * 𝐔[j][:, η]) times its transpose
        # don't use BLAS for complex data
        if ⏩ # if threaded don't share memory for Ω but use 𝛀
            if type<:Real
                @inbounds for κ=1:k
                    𝛀[η][:, κ] = BLAS.gemv('N', 𝒢[κ, i, j], 𝐔[j][:, η])
                end
                𝐑[η] += Hermitian(BLAS.gemm('N', 'T', 𝛀[η], 𝛀[η])) # (Ω * Ω')
            else
                @inbounds for κ=1:k
                    𝛀[η][:, κ] = 𝒢[κ, i, j] * 𝐔[j][:, η]
                end
                𝐑[η] += Hermitian(𝛀[η] * 𝛀[η]') # (Ω * Ω')
            end
        else
            if type<:Real
                @inbounds for κ=1:k
                    Ω[:, κ] = BLAS.gemv('N', 𝒢[κ, i, j], 𝐔[j][:, η])
                end
                𝐑[η] += Hermitian(BLAS.gemm('N', 'T', Ω, Ω)) # (Ω * Ω')
            else
                @inbounds for κ=1:k
                    Ω[:, κ] = 𝒢[κ, i, j] * 𝐔[j][:, η]
                end
                𝐑[η] += Hermitian(Ω * Ω') # (Ω * Ω')
            end
        end
    end

    function update!(m, n, i)

        function udR1!(η) # case m=1
            fill!(𝐑[η], type(0))
            updateR!(η, 1, 1)
        end

        function udRm!(η, i)  # case m>1
            fill!(𝐑[η], type(0))
            for j=1:m i≠j ? updateR!(η, i, j) : nothing end # j ≠ i
            fullModel ? updateR!(η, i, i) : nothing         # j = i
        end

        if m==1
            ⏩ ? (@threads for η=1:n udR1!(η) end) : (for η=1:n udR1!(η) end)
        else
            ⏩ ? (@threads for η=1:n udRm!(η, i) end) : (for η=1:n udRm!(η, i) end)
        end
    end

    # pre-allocate memory
    𝐑 = HermitianVector([Hermitian(zeros(type, n, n)) for η=1:n])
    ⏩ ? (𝛀 = [Matrix{type}(undef, n, k) for η=1:n]) :
         (Ω = Matrix{type}(undef, n, k))

    # # # # # here starts the algorithm

    if algo==:OJoB
        verbose && @info("Iterating OJoB algorithm...")
        while true
            conv_ =0.
            @inbounds for i=1:m # m optimizations for updating 𝐔[1]...𝐔[m]

                update!(m, n, i)

                # do 1 power iteration, not worth threading here
                for η=1:n 𝐔[i][:, η] = 𝐑[η] * 𝐔[i][:, η] end

                conv_ += PosDefManifold.ss(𝐔[i]) # square of the norms of power iteration vectors

                # Lodwin Orthogonalization and update 𝐔[i]<-UV', with svd(𝐔[i])=UwV'
                𝐔[i] = PosDefManifold.nearestOrth(𝐔[i])
            end

            conv_ =sqrt(conv_ /(n^2*m))
            iter==1 ? conv=1. : conv = abs((conv_-oldconv)/oldconv)  # relative change

            verbose && println("iteration: ", iter, "; convergence: ", conv)
            (diverging = conv < 0) && verbose && @warn("OJoB diverged at:", iter)
            (overRun = iter == maxiter) && @warn("OJoB: reached the max number of iterations before convergence:", iter)
            (😋 = 0. <= conv <= tolerance) || overRun==true ? break : nothing
            oldconv=conv_
            iter += 1
        end # while
    else
        verbose && @info("Iterating NoJoB algorithm...")

        # solve Lx=𝐑[η]*𝐔[i][:, η] for x and L'y=x for y
        # and scale the ηth column as 𝐔[i][:, η] <- y/sqrt(y'𝐑[η]t)
        function triS!(cho, η,  i)
            y=cho.U\(cho.L\(𝐑[η]*𝐔[i][:, η]))
            𝐔[i][:, η]=y*inv(sqrt(PosDefManifold.qf(y, 𝐑[η])))
        end

        # thread sum of 𝐑[η] and the solutions to triangular systems
        # only if n is at least Threads.nthreads()*2 (no worth otherwise)
        ⏩x = ⏩ && n≥Threads.nthreads()*2

        while true
            conv_ =0.
            @inbounds for e2=1:2, i=1:m # m optimizations for updating 𝐔[1]...𝐔[m]
                                        # double loop to avoid oscillating convergence
                update!(m, n, i)

                # do 1 power iteration
                cho=cholesky(⏩x ? fVec(sum, 𝐑) : sum(𝐑)) # Cholesky LL'of 𝐑[1]+...+𝐑[n]

                ⏩x ? (@threads for η=1:n triS!(cho, η,  i) end) :
                            (for η=1:n triS!(cho, η,  i) end)

                conv_+=PosDefManifold.ss(𝐔[i])
            end

            conv_=conv_/(2*n^2*m) # 2 because of the double loop
            iter==1 ? conv=1. : conv = abs((conv_-oldconv)/oldconv)  # relative change

            verbose && println("iteration: ", iter, "; convergence: ", conv)
            (diverging = conv < 0) && verbose && @warn("NoJoB diverged at:", iter)
            (overRun = iter == maxiter) && @warn("NoJoB: reached the max number of iterations before convergence:", iter)
            (😋 = 0. <= conv <= tolerance) || overRun==true ? break : nothing
            oldconv=conv_
            iter += 1
        end # while
    end

    verbose && @info("Convergence has "*(😋 ? "" : "not ")*"been attained.\n\n")

    # scale and permute the vectors of U_1,...,U_m
    if sort
        algo==:NoJoB ? for i=1:m normalizeCol!(𝐔[i], 1:size(𝐔[i], 2)) end : nothing
        λ = m==1 ? _permute!(𝐔[1], 𝒢, k, :c) :
                   _flipAndPermute!(𝐔, 𝒢, m, k, :c)
    else
        λ = m==1 ? 𝛍([𝐔[1][:, η]'*𝒢[l, 1, 1]*𝐔[1][:, η] for η=1:n] for l=1:k) :
                   𝛍([𝐔[i][:, η]'*𝒢[l, i, j]*𝐔[j][:, η] for η=1:n] for l=1:k, j=1:m, i=1:m if i≠j)
    end

    if preWhite
        algo==:OJoB ? 𝐕=[𝐔[i]'*𝑾[i].iF for i=1:m] :
                      𝐕=[pinv(𝐔[i])*𝑾[i].iF for i=1:m] # algo==:NoJoB
        for i=1:m 𝐔[i] = 𝑾[i].F * 𝐔[i] end
    else
        algo==:OJoB ? 𝐕=[Matrix(𝐔[i]') for i=1:m] :
                      𝐕=[pinv(𝐔[i]) for i=1:m] # algo==:NoJoB
    end

    return m>1 ? (𝐔, 𝐕, λ, iter, conv) : (𝐔[1], 𝐕[1], λ, iter, conv)
end


using Test, SummationByPartsOperators

BCs = (Val(:HomogeneousDirichlet), Val(:HomogeneousNeumann), Val(:NonReflecting))

for T in (Float32, Float64), left_bc in BCs, right_bc in BCs
    xmin = T(-1)
    xmax = T(1)
    N = 2^6
    u0func = sinpi
    du0func = zero
    tspan = (zero(T), one(T))

    # SBP FD
    for acc_order in 2:2:8
        D = derivative_operator(MattssonSvärdNordström2004(), 2, acc_order, xmin, xmax, N)
        semidisc = WaveEquationNonperiodicSemidiscretisation(D, left_bc, right_bc)
        ode = semidiscretise(du0func, u0func, semidisc, tspan)
        ddu = similar(ode.u0.x[1])
        semidisc(ddu, ode.u0.x..., nothing, first(tspan))
    end
end


__precompile__()
module Derivatives

using InplaceRealFFT
using LESFilter: rfftfreq, fftfreq

export dx, dy, dz, dx!, dy!, dz!

function loopdx!(fieldhat::AbstractArray{<:Complex,3}, kim::AbstractVector)
  nx,ny,nz = size(fieldhat)
  let nx = nx, ny = ny, nz = nz
  Threads.@threads for l = 1:nz
    for j = 1:ny
      @simd for i = 1:nx
      @inbounds fieldhat[i,j,l] = fieldhat[i,j,l]*kim[i]
      end
    end
  end
  end
end

function dx(field::AbstractArray{T,3},len::Real,n::Integer=1) where T<:Union{Float32,Float64}
  nx,ny,nz = size(field)

  fieldhat = rfft(field,1)
  kim = (2π .* rfftfreq(nx,len) .* im) .^ n

  loopdx!(fieldhat,kim)

  return irfft(fieldhat,nx,1)
end

function loopdy!(fieldhat::AbstractArray{<:Complex,3}, kim::AbstractVector)
  nx,ny,nz = size(fieldhat)
  let nx = nx, ny = ny, nz = nz
  Threads.@threads for l = 1:nz
    for j = 1:ny
      @simd for i = 1:nx
        @inbounds fieldhat[i,j,l] = fieldhat[i,j,l]*kim[j]
      end
    end
  end
  end
end

function dy(field::AbstractArray{T,3},len::Real,n::Integer=1) where T<:Union{Float32,Float64}
  nx,ny,nz = size(field)

  fieldhat = rfft(field,2)

  kim = (2π .* rfftfreq(ny,len) .* im) .^ n

  loopdy!(fieldhat,kim)

  return irfft(fieldhat,ny,2)
end

function loopdz!(fieldhat::AbstractArray{<:Complex,3}, kim::AbstractVector)
  nx,ny,nz = size(fieldhat)
  let nx = nx, ny = ny, nz = nz
  Threads.@threads for l = 1:nz
    for j = 1:ny
      @simd for i = 1:nx
        @inbounds fieldhat[i,j,l] = fieldhat[i,j,l]*kim[l]
      end
    end
  end
  end
end

function dz(field::AbstractArray{T,3},len::Real,n::Integer=1) where T<:Union{Float32,Float64}
  nx,ny,nz = size(field)

  fieldhat = rfft(field,3)

  kim = (2π .* rfftfreq(nz,len) .* im) .^ n

  loopdz!(fieldhat,kim)

  return irfft(fieldhat,nz,3)
end

function dx!(field::AbstractPaddedArray{T,3},len::Real,n::Integer=1) where {T<:Union{Float64,Float32}}

  nx,ny,nz = size(real(field))

  rfft!(field,1)

  kim = (2π .* rfftfreq(nx,len) .* im) .^ n

  loopdx!(InplaceRealFFT.complex_view(field),kim)

  return irfft!(field,1)
end

function dx(field::AbstractPaddedArray{T,3},len::Real,n::Integer=1) where {T<:Union{Float64,Float32}}
  fieldhat = copy(field)
  dx!(fieldhat,len,n)
  return fieldhat
end

function dy!(field::AbstractPaddedArray{T,3},len::Real,n::Integer=1) where {T<:Union{Float64,Float32}}

  nx,ny,nz = size(real(field))
  rfft!(field,1:2)

  kim = (2π .* fftfreq(ny,len) .* im) .^ n

  loopdy!(InplaceRealFFT.complex_view(field),kim)

  return irfft!(field,1:2)
end

function dy(field::AbstractPaddedArray{T,3},len::Real,n::Integer=1) where {T<:Union{Float64,Float32}}
  fieldhat = copy(field)
  dy!(fieldhat,len,n)
  return fieldhat
end

function dz!(field::AbstractPaddedArray{T,3},len::Real,n::Integer=1) where {T<:Union{Float64,Float32}}

  nx,ny,nz = size(real(field))
  rfft!(field,(1,3))

  kim = (2π .* fftfreq(nz,len) .* im) .^ n

  loopdz!(InplaceRealFFT.complex_view(field),kim)

  return irfft!(field,(1,3))
end

function dz(field::AbstractPaddedArray{T,3},len::Real,n::Integer=1) where {T<:Union{Float64,Float32}}
  fieldhat = copy(field)
  dz!(fieldhat,len,n)
  return fieldhat
end

end # module


# Various routines relating specifically to fusion reactions
#TODO: Fix Ealpha, malpha, and valpha to exact values
#TODO: Get correct values for fusion cross section and Maxwellian sigma_v
module fusion
using ..constants
using ..species

const Ealpha = 3.5e6*el	# alpha birth energy
const malpha = 4.0*mp
const valpha = sqrt(2.0*Ealpha/(malpha))

"""
fusionSourceDT()

Calculates the fusion cross section for Maxwellian D and T.
"""
function fusionSourceDT(D::SpeciesData,T::SpeciesData)

   # Get these from standard sources
   sigma_v_avg = 1.0

   return D.density*T.density*sigma_v_avg
end
function fusionSourceDT(f_D::Function,f_T::Function)

   # Calculate directly from distribution function f_D and f_T
   return stuff
end

"""
DTcrossSection(v::Float64)

Calculates the fusion cross section sigma(v)
"""
function crossSectionDT(v::Float64)

end

export Ealpha,malpha,valpha, fusionSourceDT, crossSectionDT

end


calc_I_M(reactor, R_0, B_0) = K_I(reactor) / ( B_0 * sqrt(R_0) )


function infer(M::Model, X)
  infer(M.alg, M.arc, X)
end

function update!(M::Model, y, ŷ, x, learn_rate::Float64)
    update!(M, M.alg, M.arc, y, ŷ, x, learn_rate)
end

"""
Linear Algorithm
"""
struct Linear <: Algorithm
    link
end

function Linear()::Linear
  Linear(identity)
end

function infer(alg::Linear, arc::SingleLayer, X)
  return arc.wts * X'
end

# function infer(alg::Linear, arc::MultiLayer, X)
# end

function update!(M::Model, alg::Linear, arc::SingleLayer,
                 y, ŷ, x, learn_rate::Float64)
  M.arc.wts .-= learn_rate * assessgrad(M, y, ŷ, x)
end

# function update!(M::Model, alg::Linear, arc::MultiLayer,
#                  y, ŷ, x, learn_rate::Float64)
# end


function parse_data(program_str)
    [parse(Int64, v) for v in split(program_str, ',')]
end

function get_data()
    program_str = "1,9,10,3,2,3,11,0,99,30,40,50"
    parse_data(program_str)
end

function op_add!(program, pos)
    pos_arg1, pos_arg2, pos_result = program[pos+1:pos+3]
    program[pos_result+1] = program[pos_arg1+1] + program[pos_arg2+1]
end

function op_mul!(program, pos)
    pos_arg1, pos_arg2, pos_result = program[pos+1:pos+3]
    program[pos_result+1] = program[pos_arg1+1] * program[pos_arg2+1]
end

function op_input!(program, pos)
    pos_arg = program[pos+1]
    input = parse(Int64, readline())
    program[pos_arg+1] = input
end

function op_output!(program, pos)
    pos_arg = program[pos+1]
    output = program[pos_arg+1]
    println(output)
end

function run_opcode!(program, pos)
    opcode = program[pos]
    if opcode == 1
        op_add!(program, pos)
    elseif opcode == 2
        op_mul!(program, pos)
    elseif opcode == 3
        op_input!(program, pos)
    elseif opcode == 4
        op_output!(program, pos)
    elseif opcode == 99
        return opcode
    else
        error("unknown opcode " * string(opcode) * "!")
    end
    opcode
end

function run_program!(program)
    pos = 1
    while true
        opcode = run_opcode!(program, pos)
        if opcode == 99
            return program
        end
        pos += 4
    end
end


using TensorOperations

"""

     interpolate( interp :: InterpolationType,
                  f  :: Array{Float64,1},
                  xe :: Array{Float64,1})

Computes interpolation of a 1D function via Hermite-tensor method.

Arguments:
 - interp : inteprolation type (Hermite or Radial)
 - f      : 1d array containing values at interpolation nodes (Chebyshev or Uniform)
 - xe     : vector of evaluation points

  © Anna Yurova, 2017

"""
function interpolate(interp :: InterpolationType,
                     f      :: Array{Float64,1},
		     xe     :: Array{Float64,1} )

   x = interp(xe)

   @tensoropt begin
      s[e] := x[e, c]*f[c]
   end

   s

end

"""
     interpolate( interp_x :: InterpolationType,
                  interp_y :: InterpolationType,
                  f  :: Array{Float64,2},
                  xe :: Array{Float64,1},
                  ye :: Array{Float64,1})

Computes interpolation of a 2D function via Hermite-tensor method.

Arguments:
 - interp : inteprolation type (Hermite or Radial)
 - f      : 2d array containing values at interpolation nodes (Chebyshev or Uniform)
 - xe,ye  : vector of evaluation points
"""
function interpolate(interp_x :: InterpolationType,
                     interp_y :: InterpolationType,
                     f        :: Array{Float64,2},
		     xe       :: Array{Float64,1},
		     ye       :: Array{Float64,1})

   x = interp_x(xe)
   y = interp_y(ye)

   @tensoropt begin
       s[e1,e2] := x[e1,c1]*y[e2,c2]*f[c1,c2]
   end

   s

end

"""
     interpolate( interp_x :: InterpolationType,
                  interp_y :: InterpolationType,
                  interp_z :: InterpolationType,
                  f  :: Array{Float64,3},
                  xe :: Array{Float64,1},
                  ye :: Array{Float64,1},
                  ze :: Array{Float64,1})

Computes interpolation of a 3D function via Hermite-tensor method.

Arguments:
 - interp_x : inteprolation type (Hermite or Radial)
 - interp_y : inteprolation type (Hermite or Radial)
 - interp_z : inteprolation type (Hermite or Radial)
 - f        : 3d array containing values at interpolation nodes (Chebyshev or Uniform)

Computing s = ( z ⊗ y ⊗ x) vec(f) (see Sec. 4.1 of the paper)

"""
function interpolate(interp_x :: InterpolationType,
                     interp_y :: InterpolationType,
                     interp_z :: InterpolationType,
                     f        :: Array{Float64,3},
		     xe       :: Array{Float64,1},
		     ye       :: Array{Float64,1},
		     ze       :: Array{Float64,1})

    x = interp_x(xe)
    y = interp_y(ye)
    z = interp_z(ze)

    @tensoropt begin
        s[e1,e2,e3] := x[e1,c1]*y[e2,c2]*z[e3,c3]*f[c1,c2,c3]
    end
    s
end

"""
     interpolate( interp_x :: InterpolationType,
                  interp_y :: InterpolationType,
                  interp_z :: InterpolationType,
                  interp_v :: InterpolationType,
                  f  :: Array{Float64,4},
                  xe :: Array{Float64,1},
                  ye :: Array{Float64,1},
                  ze :: Array{Float64,1},
                  ve :: Array{Float64,1})

Computes interpolation of a 4D function via Hermite-tensor method.

Arguments:
 - interp_x : inteprolation type (Hermite or Radial)
 - interp_y : inteprolation type (Hermite or Radial)
 - interp_z : inteprolation type (Hermite or Radial)
 - interp_v : inteprolation type (Hermite or Radial)
 - f        : 4d array containing values at interpolation nodes (Chebyshev or Uniform)

Computing s = (v ⊗ z ⊗ y ⊗ x) vec(f) (see Sec. 4.1 of the paper)

"""
function interpolate(interp_x :: InterpolationType,
                     interp_y :: InterpolationType,
                     interp_z :: InterpolationType,
                     interp_v :: InterpolationType,
                     f        :: Array{Float64,4},
		     xe       :: Array{Float64,1},
		     ye       :: Array{Float64,1},
		     ze       :: Array{Float64,1},
		     ve       :: Array{Float64,1})

    x = interp_x(xe)
    y = interp_y(ye)
    z = interp_z(ze)
    v = interp_v(ve)

    @tensoropt begin
        s[e1,e2,e3,e4] := x[e1,c1]*y[e2,c2]*z[e3,c3]*v[e4,c4]*f[c1,c2,c3,c4]
    end
    s
end

"""
     interpolate( interp_x :: InterpolationType,
                  interp_y :: InterpolationType,
                  interp_z :: InterpolationType,
                  interp_v :: InterpolationType,
                  interp_w :: InterpolationType,
                  f  :: Array{Float64,4},
                  xe :: Array{Float64,1},
                  ye :: Array{Float64,1},
                  ze :: Array{Float64,1},
                  ve :: Array{Float64,1},
                  we :: Array{Float64,1})

Computes interpolation of a 4D function via Hermite-tensor method.

Arguments:
 - interp_x : interpolation type (Hermite or Radial)
 - interp_y : interpolation type (Hermite or Radial)
 - interp_z : interpolation type (Hermite or Radial)
 - interp_v : interpolation type (Hermite or Radial)
 - interp_w : interpolation type (Hermite or Radial)
 - f        : 5d arrray containing values at interpolation nodes (Chebyshev or Uniform)

Extract the number of collocation and evaluation points in each dimension
Computing s = (w ⊗ v ⊗ z ⊗ y ⊗ x) vec(f) (see Sec. 4.1 of the paper)

"""
function interpolate(interp_x :: InterpolationType,
                     interp_y :: InterpolationType,
                     interp_z :: InterpolationType,
                     interp_v :: InterpolationType,
                     interp_w :: InterpolationType,
                     f        :: Array{Float64,5},
		     xe       :: Array{Float64,1},
		     ye       :: Array{Float64,1},
		     ze       :: Array{Float64,1},
		     ve       :: Array{Float64,1},
		     we       :: Array{Float64,1})

    x = interp_x(xe)
    y = interp_y(ye)
    z = interp_z(ze)
    v = interp_v(ve)
    w = interp_w(we)

    @tensoropt begin
        s[e1,e2,e3,e4,e5] := x[e1,c1]*y[e2,c2]*z[e3,c3]*v[e4,c4]*w[e5,c5]*f[c1,c2,c3,c4,c5]
    end
    s

end


using Base.Test
using StochasticProcesses

let p = PoissonProcess(5), t=linspace(0, 1, 10)
  sim(p, t, 3)
  cumsim(p, t, 3)
  distribution(p, 10)
end


using FINUFFT

using Test
using LinearAlgebra
using Random

Random.seed!(1)

nj = 10
nk = 11
ms = 12
mt = 13
mu = 14

tol = 1e-15

# nonuniform data
x = 3*pi*(1.0 .- 2*rand(nj))
y = 3*pi*(1.0 .- 2*rand(nj))
z = 3*pi*(1.0 .- 2*rand(nj))
c = rand(nj) + 1im*rand(nj)
s = rand(nk)
t = rand(nk)
u = rand(nk)
f = rand(nk) + 1im*rand(nk)

# uniform data
F1D = rand(ms) + 1im*rand(ms)
F2D = rand(ms, mt) + 1im*rand(ms,mt)
F3D = rand(ms, mt, mu) + 1im*rand(ms,mt, mu)

modevec(m) = -floor(m/2):floor((m-1)/2+1)
k1 = modevec(ms)
k2 = modevec(mt)
k3 = modevec(mu)

@testset "NUFFT" begin
    ## 1D
    @testset "1D" begin
        # 1D1
        @testset "1D1" begin
            out = complex(zeros(ms))
            ref = complex(zeros(ms))
            for j=1:nj
                for ss=1:ms
                    ref[ss] += c[j] * exp(1im*k1[ss]*x[j])
                end
            end
            # Try this one with explicit opts struct
            opts = finufft_default_opts()
            opts.spread_kerpad = 0 # This should also work
            nufft1d1!(x, c, 1, tol, out, opts)
            relerr_1d1 = norm(vec(out)-vec(ref), Inf) / norm(vec(ref), Inf)
            @test relerr_1d1 < 1e-13
            # Different caller
            out2 = nufft1d1(x, c, 1, tol, ms)
            reldiff = norm(vec(out)-vec(out2), Inf) / norm(vec(out), Inf)
            @test reldiff < 1e-14
        end

        # 1D2
        @testset "1D2" begin
            out = complex(zeros(nj))
            ref = complex(zeros(nj))
            for j=1:nj
                for ss=1:ms
                    ref[j] += F1D[ss] * exp(1im*k1[ss]*x[j])
                end
            end
            nufft1d2!(x, out, 1, tol, F1D)
            relerr_1d2 = norm(vec(out)-vec(ref), Inf) / norm(vec(ref), Inf)
            @test relerr_1d2 < 1e-13
            out2 = nufft1d2(x, 1, tol, F1D)
            reldiff = norm(vec(out)-vec(out2), Inf) / norm(vec(out), Inf)
            @test reldiff < 1e-14
        end

        # 1D3
        @testset "1D3" begin
            out = complex(zeros(nk))
            ref = complex(zeros(nk))
            for k=1:nk
                for j=1:nj
                    ref[k] += c[j] * exp(1im*s[k]*x[j])
                end
            end
            nufft1d3!(x,c,1,tol,s,out)
            relerr_1d3 = norm(vec(out)-vec(ref), Inf) / norm(vec(ref), Inf)
            @test relerr_1d3 < 1e-13
            out2 = nufft1d3(x,c,1,tol,s)
            reldiff = norm(vec(out)-vec(out2), Inf) / norm(vec(out), Inf)
            @test reldiff < 1e-14
        end
    end

    ## 2D
    @testset "2D" begin
        @testset "2D1" begin
            # 2D1
            out = complex(zeros(ms, mt))
            ref = complex(zeros(ms, mt))
            for j=1:nj
                for ss=1:ms
                    for tt=1:mt
                        ref[ss,tt] += c[j] * exp(1im*(k1[ss]*x[j]+k2[tt]*y[j]))
                    end
                end
            end
            nufft2d1!(x, y, c, 1, tol, out)
            relerr_2d1 = norm(vec(out)-vec(ref), Inf) / norm(vec(ref), Inf)
            @test relerr_2d1 < 1e-13
            out2 = nufft2d1(x, y, c, 1, tol, ms, mt)
            reldiff = norm(vec(out)-vec(out2), Inf) / norm(vec(out), Inf)
            @test reldiff < 1e-14
        end

        @testset "2D2" begin
            # 2D2
            out = complex(zeros(nj))
            ref = complex(zeros(nj))
            for j=1:nj
                for ss=1:ms
                    for tt=1:mt
                        ref[j] += F2D[ss, tt] * exp(1im*(k1[ss]*x[j]+k2[tt]*y[j]))
                    end
                end
            end
            nufft2d2!(x, y, out, 1, tol, F2D)
            relerr_2d2 = norm(vec(out)-vec(ref), Inf) / norm(vec(ref), Inf)
            @test relerr_2d2 < 1e-13
            out2 = nufft2d2(x, y, 1, tol, F2D)
            reldiff = norm(vec(out)-vec(out2), Inf) / norm(vec(out), Inf)
            @test reldiff < 1e-14
        end

        @testset "3D3" begin
            # 2D3
            out = complex(zeros(nk))
            ref = complex(zeros(nk))
            for k=1:nk
                for j=1:nj
                    ref[k] += c[j] * exp(1im*(s[k]*x[j]+t[k]*y[j]))
                end
            end
            nufft2d3!(x,y,c,1,tol,s,t,out)
            relerr_2d3 = norm(vec(out)-vec(ref), Inf) / norm(vec(ref), Inf)
            @test relerr_2d3 < 1e-13
            out2 = nufft2d3(x,y,c,1,tol,s,t)
            reldiff = norm(vec(out)-vec(out2), Inf) / norm(vec(out), Inf)
            @test reldiff < 1e-14
        end
    end

    ## 3D
    @testset "3D" begin
        @testset "3D1" begin
            # 3D1
            out = complex(zeros(ms, mt, mu))
            ref = complex(zeros(ms, mt, mu))
            for j=1:nj
                for ss=1:ms
                    for tt=1:mt
                        for uu=1:mu
                            ref[ss,tt,uu] += c[j] * exp(1im*(k1[ss]*x[j]+k2[tt]*y[j]+k3[uu]*z[j]))
                        end
                    end
                end
            end
            nufft3d1!(x, y, z, c, 1, tol, out)
            relerr_3d1 = norm(vec(out)-vec(ref), Inf) / norm(vec(ref), Inf)
            @test relerr_3d1 < 1e-13
            out2 = nufft3d1(x, y, z, c, 1, tol, ms, mt, mu)
            reldiff = norm(vec(out)-vec(out2), Inf) / norm(vec(out), Inf)
            @test reldiff < 1e-14
        end

        @testset "3D2" begin
            # 3D2
            out = complex(zeros(nj))
            ref = complex(zeros(nj))
            for j=1:nj
                for ss=1:ms
                    for tt=1:mt
                        for uu=1:mu
                            ref[j] += F3D[ss, tt, uu] * exp(1im*(k1[ss]*x[j]+k2[tt]*y[j]+k3[uu]*z[j]))
                        end
                    end
                end
            end
            nufft3d2!(x, y, z, out, 1, tol, F3D)
            relerr_3d2 = norm(vec(out)-vec(ref), Inf) / norm(vec(ref), Inf)
            @test relerr_3d2 < 1e-13
            out2 = nufft3d2(x, y, z, 1, tol, F3D)
            reldiff = norm(vec(out)-vec(out2), Inf) / norm(vec(out), Inf)
            @test reldiff < 1e-14
        end

        @testset "3D3" begin
            # 3D3
            out = complex(zeros(nk))
            ref = complex(zeros(nk))
            for k=1:nk
                for j=1:nj
                    ref[k] += c[j] * exp(1im*(s[k]*x[j]+t[k]*y[j]+u[k]*z[j]))
                end
            end
            nufft3d3!(x,y,z,c,1,tol,s,t,u,out)
            relerr_3d3 = norm(vec(out)-vec(ref), Inf) / norm(vec(ref), Inf)
            @test relerr_3d3 < 1e-13
            out2 = nufft3d3(x,y,z,c,1,tol,s,t,u)
            reldiff = norm(vec(out)-vec(out2), Inf) / norm(vec(out), Inf)
            @test reldiff < 1e-14
        end
    end
end


function density_gradient!(∇ρ, point, wfn)

    ρ    = 0.0
   ∇ρ   .= 0.0
   ϕ, ∇ϕ = eval_mo1(wfn, point)

   occ = wfn.occ

   # Run again over orbitals
   @inbounds for imo in eachindex(occ)

      @views begin
         n    = occ[imo]
         nϕ   = n   *  ϕ[imo]
          ρ  += nϕ  *  ϕ[imo]
         ∇ρ .+= nϕ .* ∇ϕ[imo,:]
      end
   end
   ∇ρ .*= 2 # + c.c.
   ∇ρ, ρ
end

function eval_mo1(wfn, point) #derivative=1)

   nmo    = wfn.nmo
   natm   = wfn.natm
   coords = wfn.coords
   ngroup = wfn.ngroup
   nuexp  = wfn.nuexp
   rcutte = wfn.rcutte
   oexp   = wfn.oexp
   nzexp  = wfn.nzexp
   ityp   = wfn.ityp
   nlm    = wfn.nlm
   coef   = wfn.coef

    ϕ   = zeros(nmo)
   ∇ϕ   = zeros(nmo,3)
    G1D = @MVector zeros(3) # Gaussian 1D functions x, y, z
   ∇G1D = @MVector zeros(3)

   @views begin
      @inbounds for ic in 1:natm

         # Atomic coordinates of this center
         xcoor = point - coords[ic,:]
         r² = sum( abs2.(xcoor) ) # ≡ dot(xcoor, xcoor)
         # Loop over different shell in this atom
         @inbounds for m in 1:ngroup[ic]

            k = nuexp[ic,m,1]
            # Skip to compute this primitive if distance is too big.
            r² > rcutte[ic,m]^2 && continue

            α = -oexp[k]
            𝟐α = 2α
            # All primitives in a shell share the same exponent.
            aexp = exp(α * r²)
            # Loop over the different primitives in this shell.
            @inbounds for jj in 1:nzexp[ic,m]

               # "i" is the original index of the primitive in the WFN.
               i = nuexp[ic,m,jj]
               itip = ityp[i]
               # Integer coefficients.
               it = nlm[itip,:]
               @inbounds for j in 1:3

                  n = it[j]
                  # r⃗ = ∑ⱼ³ rⱼ e⃗ⱼ # cartesian components
                  rⱼ = xcoor[j]
                  # G1D = r^n exp(α * r²)
                  if n == 0
                     # ∇G1D = 2αr * G1D
                     ∇G1D[j] = 𝟐α * rⱼ
                      G1D[j] = 1.0
                  elseif n == 1
                     ∇G1D[j] = 1 + 𝟐α * rⱼ * rⱼ
                      G1D[j] = rⱼ
                  elseif n == 2
                     rⱼ²   = rⱼ * rⱼ
                     ∇G1D[j] = rⱼ * ( 2.0 + 𝟐α * rⱼ² )
                      G1D[j] = rⱼ²
                  elseif n == 3
                     rⱼ²   = rⱼ  * rⱼ
                     ∇G1D[j] = rⱼ² * ( 3.0 + 𝟐α * rⱼ² )
                      G1D[j] = rⱼ  * rⱼ²
                  elseif n == 4
                     rⱼ²   = rⱼ  * rⱼ
                     ∇G1D[j] = rⱼ² * rⱼ * ( 4.0 + 𝟐α  * rⱼ² )
                      G1D[j] = rⱼ² * rⱼ²
                  elseif n == 5
                     rⱼ²   = rⱼ  * rⱼ
                     ∇G1D[j] = rⱼ² * rⱼ² * ( 5.0 + 𝟐α * rⱼ² )
                      G1D[j] = rⱼ² * rⱼ² * rⱼ
                  end
               end

               # exp(-αr₁²) * exp(-αr₂²) * exp(-αr₃²) = exp(-αr²) ≡ aexp
                 G3D =  prod(G1D) * aexp # Gxyz = Gx * Gy * Gz
               ∇xG3D = ∇G1D[1] *  G1D[2] *  G1D[3] * aexp
               ∇yG3D =  G1D[1] * ∇G1D[2] *  G1D[3] * aexp
               ∇zG3D =  G1D[1] *  G1D[2] * ∇G1D[3] * aexp

               # AO -> MO basis; run over orbitals
               @inbounds for imo in 1:nmo

                  c = coef[imo,i]
                   ϕ[imo]   += c *   G3D
                  ∇ϕ[imo,1] += c * ∇xG3D
                  ∇ϕ[imo,2] += c * ∇yG3D
                  ∇ϕ[imo,3] += c * ∇zG3D
               end
            end
         end
      end
   end
   ϕ, ∇ϕ
end


push!(LOAD_PATH, "/Users/saydjari/Dropbox/GradSchool_AKS/Doug/Projects/DHC/main")
using DHC_2DUtils
using Profile
using BenchmarkTools
using FFTW
using HDF5
using Test
using SparseArrays

filter_hash = fink_filter_hash(1, 16, nx=64, pc=2, wd=2)
filt_3d = fink_filter_bank_3dizer(filter_hash, 1, nz=128)

function S1_iso_matrix3d(fhash)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S1 coeff vector to S1iso by
    #   summing over l
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18
    # AKS 2021-Feb-22

    # Does hash contain Omega filter?
    #Omega   = haskey(fhash, "Omega_index")
    #if Omega Ω_ind = fhash["Omega_index"] end

    # unpack fhash
    Nl      = length(fhash["2d_theta_value"])
    Nj      = length(fhash["2d_j_value"])
    Nk      = length(fhash["k_value"])
    Nf      = length(fhash["filt_value"])
    ψ_ind   = fhash["psi_index"]
    #ϕ_ind   = fhash["phi_index"]

    # number of iso coefficients
    #Niso    = Omega ? Nj+2 : Nj+1
    Niso    = (Nj+1)*Nk
    Mat     = zeros(Int32, Niso, Nf)

    # first J elements of iso
    for j = 1:Nj
        for l = 1:Nl
            for k=1:Nk
                λ = ψ_ind[j,l,k]
                Mat[j+(k-1)*Nj, λ] = 1
            end
        end
    end

    j = Nj+1
    l = 1
    for k=1:Nk
        λ = ψ_ind[j,l,k]
        Mat[j+(k-1)*Nj, λ] = 1
    end

    return sparse(Mat)
end

S1_iso_matrix3d(filt_3d)

function S2_iso_matrix3d(fhash)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S2 coeff vector to S2iso by
    #   summing over l1,l2 and fixed Δl.
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18
    # AKS 2021-Feb-22

    # Does hash contain Omega filter?
    #Omega   = haskey(fhash, "Omega_index")
    #if Omega Ω_ind = fhash["Omega_index"] end

    # unpack fhash
    Nl      = length(fhash["2d_theta_value"])
    Nj      = length(fhash["2d_j_value"])
    Nk      = length(fhash["k_value"])
    Nf      = length(fhash["filt_value"])
    ψ_ind   = fhash["psi_index"]
    #ϕ_ind   = fhash["phi_index"]

    # number of iso coefficients
    #Niso    = Omega ? Nj*Nj*Nl+4*Nj+4 : Nj*Nj*Nl+2*Nj+1
    Niso    = Nj*Nj*Nk*Nk*Nl+2*Nj*Nk*Nk+Nk*Nk
    Mat     = zeros(Int32, Niso, Nf*Nf)

    #should probably think about reformatting to (Nj+1) so phi
    #phi cross terms are in the right place

    # first J*J*K*K*L elements of iso
    for j1 = 1:Nj
        for j2 = 1:Nj
            for l1 = 1:Nl
                for l2 = 1:Nl
                    for k1 = 1:Nk
                        for k2 = 1:Nk
                            DeltaL = mod(l1-l2, Nl)
                            λ1     = ψ_ind[j1,l1,k1]
                            λ2     = ψ_ind[j2,l2,k2]

                            Iiso   = k1+Nk*((k2-1)+Nk*((j1-1)+Nj*((j2-1)+Nj*DeltaL)))
                            Icoeff = λ1+Nf*(λ2-1)
                            Mat[Iiso, Icoeff] = 1
                        end
                    end
                end
            end
        end
    end

    # Next J elements are λϕ, then J elements ϕλ
    for j = 1:Nj
        for l = 1:Nl
            for k1 = 1:Nk
                for k2 =1:Nk
                    λ      = ψ_ind[j,l,k1]
                    Iiso   = Nj*Nj*Nk*Nk*Nl+k1+Nk*((k2-1)+Nk*(j-1))
                    Icoeff = λ+Nf*(ψ_ind[Nj+1,1,k2]-1)  # λϕ
                    Mat[Iiso, Icoeff] = 1

                    Iiso   = Nj*Nj*Nk*Nk*Nl+Nj*Nk*Nk+k1+Nk*((k2-1)+Nk*(j-1))
                    Icoeff = (ψ_ind[Nj+1,1,k2]-1)+Nf*(λ-1)  # ϕλ
                    Mat[Iiso, Icoeff] = 1
                end
            end
        end
    end

    # Next Nk*Nk elements are ϕ(k)ϕ(k)
    j = Nj+1
    l = 1
    for k1=1:Nk
        for k2=1:Nk
            λ1 = ψ_ind[j,l,k1]
            λ2 = ψ_ind[j,l,k2]

            Iiso = Nj*Nj*Nk*Nk*Nl+2*Nj*Nk*Nk+k1+Nk*(k2-1)
            Icoeff = λ1+Nf*(λ2-1)
            Mat[Iiso, Icoeff] = 1
        end
    end

    # # If the Omega filter exists, add more terms
    # if Omega
    #     # Next J elements are λΩ, then J elements Ωλ
    #     for j = 1:Nj
    #         for l = 1:Nl
    #             λ      = ψ_ind[j,l]
    #             Iiso   = I0+j
    #             Icoeff = λ+Nf*(Ω_ind-1)  # λΩ
    #             Mat[Iiso, Icoeff] = 1
    #
    #             Iiso   = I0+Nj+j
    #             Icoeff = Ω_ind+Nf*(λ-1)  # Ωλ
    #             Mat[Iiso, Icoeff] = 1
    #         end
    #     end
    #     # Next 3 elements are ϕΩ, Ωϕ, ΩΩ
    #     Iiso   = I0+Nj+Nj
    #     Mat[Iiso+1, ϕ_ind+Nf*(Ω_ind-1)] = 1
    #     Mat[Iiso+2, Ω_ind+Nf*(ϕ_ind-1)] = 1
    #     Mat[Iiso+3, Ω_ind+Nf*(Ω_ind-1)] = 1
    # end
    #
    return sparse(Mat)
end

S2_iso_matrix3d(filt_3d)

Nj = 4
Nk = 5
Nl = 16
Nj*Nj*Nk*Nk*Nl+2*Nj*Nk*Nk+Nk*Nk

function S1_iso_matrix(fhash)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S1 coeff vector to S1iso by
    #   summing over l
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18

    # Does hash contain Omega filter?
    Omega   = haskey(fhash, "Omega_index")
    if Omega Ω_ind = fhash["Omega_index"] end

    # unpack fhash
    Nl      = length(fhash["theta_value"])
    Nj      = length(fhash["j_value"])
    Nf      = length(fhash["filt_value"])
    ψ_ind   = fhash["psi_index"]
    ϕ_ind   = fhash["phi_index"]

    # number of iso coefficients
    Niso    = Omega ? Nj+2 : Nj+1
    Mat     = zeros(Int32, Niso, Nf)

    # first J elements of iso
    for j = 1:Nj
        for l = 1:Nl
            λ = ψ_ind[j,l]
            Mat[j, λ] = 1
        end
    end

    # Next elements are ϕ, Ω
    I0     = Nj+1
    Mat[I0, ϕ_ind] = 1
    if Omega Mat[I0+1, Ω_ind] = 1 end

    return sparse(Mat)
end

function S2_iso_matrix(fhash)
    # fhash is the filter hash output by fink_filter_hash
    # The output matrix converts an S2 coeff vector to S2iso by
    #   summing over l1,l2 and fixed Δl.
    # Matrix is stored in sparse CSC format using SparseArrays.
    # DPF 2021-Feb-18

    # Does hash contain Omega filter?
    Omega   = haskey(fhash, "Omega_index")
    if Omega Ω_ind = fhash["Omega_index"] end

    # unpack fhash
    Nl      = length(fhash["theta_value"])
    Nj      = length(fhash["j_value"])
    Nf      = length(fhash["filt_value"])
    ψ_ind   = fhash["psi_index"]
    ϕ_ind   = fhash["phi_index"]

    # number of iso coefficients
    Niso    = Omega ? Nj*Nj*Nl+4*Nj+4 : Nj*Nj*Nl+2*Nj+1
    Mat     = zeros(Int32, Niso, Nf*Nf)

    # first J*J*L elements of iso
    for j1 = 1:Nj
        for j2 = 1:Nj
            for l1 = 1:Nl
                for l2 = 1:Nl
                    DeltaL = mod(l1-l2, Nl)
                    λ1     = ψ_ind[j1,l1]
                    λ2     = ψ_ind[j2,l2]

                    Iiso   = j1+Nj*((j2-1)+Nj*DeltaL)
                    Icoeff = λ1+Nf*(λ2-1)
                    Mat[Iiso, Icoeff] = 1
                end
            end
        end
    end

    # Next J elements are λϕ, then J elements ϕλ
    for j = 1:Nj
        for l = 1:Nl
            λ      = ψ_ind[j,l]
            Iiso   = Nj*Nj*Nl+j
            Icoeff = λ+Nf*(ϕ_ind-1)  # λϕ
            Mat[Iiso, Icoeff] = 1

            Iiso   = Nj*Nj*Nl+Nj+j
            Icoeff = ϕ_ind+Nf*(λ-1)  # ϕλ
            Mat[Iiso, Icoeff] = 1
        end
    end

    # Next 1 element is ϕϕ
    I0     = Nj*Nj*Nl+Nj+Nj+1
    Icoeff = ϕ_ind+Nf*(ϕ_ind-1)
    Mat[I0, Icoeff] = 1

    # If the Omega filter exists, add more terms
    if Omega
        # Next J elements are λΩ, then J elements Ωλ
        for j = 1:Nj
            for l = 1:Nl
                λ      = ψ_ind[j,l]
                Iiso   = I0+j
                Icoeff = λ+Nf*(Ω_ind-1)  # λΩ
                Mat[Iiso, Icoeff] = 1

                Iiso   = I0+Nj+j
                Icoeff = Ω_ind+Nf*(λ-1)  # Ωλ
                Mat[Iiso, Icoeff] = 1
            end
        end
        # Next 3 elements are ϕΩ, Ωϕ, ΩΩ
        Iiso   = I0+Nj+Nj
        Mat[Iiso+1, ϕ_ind+Nf*(Ω_ind-1)] = 1
        Mat[Iiso+2, Ω_ind+Nf*(ϕ_ind-1)] = 1
        Mat[Iiso+3, Ω_ind+Nf*(Ω_ind-1)] = 1
    end

    return sparse(Mat)
end

S1_iso_matrix(filter_hash)
S2_iso_matrix(filter_hash)


using ControlSystems
using Plots

m=1;
b=1;
k=2;

sys = tf([1,0,1], [m,b,m,b,k]);

impulse(sys)

impulseplot(sys)


using DataFrames
using LinearAlgebra
using Statistics: std, var, mean, cov
using HighFrequencyCovariance
using Random

brownian_corr_matrix = Hermitian([1.0 0.75 0.5 0.0;
                                  0.0 1.0 0.5 0.25;
                                  0.0 0.0 1.0 0.25;
                                  0.0 0.0 0.0 1.0])
assets = [:BARC, :HSBC, :VODL, :RYAL]
twister = MersenneTwister(1)

ts1, true_covar, micro_noise, update_rates = generate_random_path(4, 5000; brownian_corr_matrix = brownian_corr_matrix, assets = assets, vols = [0.02,0.03,0.04,0.05], twister = deepcopy(twister))
ts2, true_covar, micro_noise, update_rates = generate_random_path(4, 50000; brownian_corr_matrix = brownian_corr_matrix, assets = assets, vols = [0.02,0.03,0.04,0.05], twister = deepcopy(twister))

# Testing constructors of SortedDataFrames
ts1 = SortedDataFrame(ts1.df, ts1.time, ts1.grouping, ts1.value,  ts1.groupingrows, ts1.time_period_per_unit)
# Subsetting to time.
ts1 = subset_to_time(ts1, 1670) # Will not delete anything
ts1 = subset_to_time(ts1, 1400) # Will delete some stuff.
# Subsetting to tick.
ts2 = subset_to_tick(ts2, 49000)

# Getting ticks per asset
tpa = ticks_per_asset(ts1)
all(values(tpa) .> 1)

iscloser(a,b) = (a.Correlation_error + a.Volatility_error < b.Correlation_error + b.Volatility_error)

# Preav Convergence
preav_estimate1 = preaveraged_covariance(ts1, assets)
preav_estimate2 = preaveraged_covariance(ts2, assets)
iscloser(calculate_mean_abs_distance(preav_estimate2, true_covar), calculate_mean_abs_distance(preav_estimate1, true_covar))
is_psd_matrix(preav_estimate1)
is_psd_matrix(preav_estimate2)
valid_correlation_matrix(preav_estimate1)
valid_correlation_matrix(preav_estimate2)

# simple Convergence
simple_estimate1 = simple_covariance(ts1, assets)
simple_estimate2 = simple_covariance(ts2, assets)
# The below fails. But no reason to  believe that the simple method would be consistant anyway.
# iscloser(calculate_mean_abs_distance(simple_estimate2, true_covar), calculate_mean_abs_distance(simple_estimate1, true_covar))
valid_correlation_matrix(simple_estimate1)
valid_correlation_matrix(simple_estimate2)

# bnhls Convergence
bnhls_estimate1 = bnhls_covariance(ts1, assets; regularisation = :eigenvalue_clean) # Need to change this regularsation because such little data means regu
bnhls_estimate2 = bnhls_covariance(ts2, assets)
iscloser(calculate_mean_abs_distance(bnhls_estimate2, true_covar), calculate_mean_abs_distance(bnhls_estimate1, true_covar))
valid_correlation_matrix(bnhls_estimate1)
valid_correlation_matrix(bnhls_estimate2)

# Preav Convergence
spectral_estimate1 = spectral_covariance(ts1, assets; num_blocks = 1) # not many observations so need to reduce the number of blocks here
spectral_estimate2 = spectral_covariance(ts2, assets)
iscloser(calculate_mean_abs_distance(spectral_estimate2, true_covar), calculate_mean_abs_distance(spectral_estimate1, true_covar))
valid_correlation_matrix(spectral_estimate1)
valid_correlation_matrix(spectral_estimate2)

# two scales Convergence
two_scales_estimate1 = two_scales_covariance(ts1, assets)
two_scales_estimate2 = two_scales_covariance(ts2, assets)
iscloser(calculate_mean_abs_distance(two_scales_estimate2, true_covar), calculate_mean_abs_distance(two_scales_estimate1, true_covar))
valid_correlation_matrix(two_scales_estimate1)
valid_correlation_matrix(two_scales_estimate2)

#############################
 # Serialisation and deserialisation

true_df = to_dataframe(true_covar)
reconstituted_df = dataframe_to_covariancematrix(true_df)
calculate_mean_abs_distance(true_covar, reconstituted_df).Correlation_error .< 10*eps()
calculate_mean_abs_distance(true_covar, reconstituted_df).Volatility_error .< 10*eps()

# Other regularisation algos:
two_scales_estimate_iden = two_scales_covariance(ts2, assets; regularisation = :identity_regularisation)
two_scales_estimate_nearest_corr = two_scales_covariance(ts2, assets; regularisation = :nearest_correlation_matrix)
two_scales_estimate_nearest_psd = two_scales_covariance(ts2, assets; regularisation = :nearest_psd_matrix)
two_scales_estimate_eigen = two_scales_covariance(ts2, assets; regularisation = :eigenvalue_clean)

# Running regularistation on a CovarianceMatrix's correlation matrix.
valid_correlation_matrix(two_scales_estimate_nearest_corr)

psd_mat = two_scales_estimate_nearest_corr

reg1 = identity_regularisation(psd_mat, ts2)
calculate_mean_abs_distance(psd_mat, reg1).Correlation_error .> 10*eps()
calculate_mean_abs_distance(psd_mat, reg1).Volatility_error .> 10*eps()

reg1_corr = identity_regularisation(psd_mat, ts2; apply_to_covariance = false)
calculate_mean_abs_distance(psd_mat, reg1_corr).Correlation_error .> 10*eps()
calculate_mean_abs_distance(psd_mat, reg1_corr).Volatility_error .< 10*eps()

reg2 = nearest_correlation_matrix(psd_mat, ts2)
calculate_mean_abs_distance(psd_mat, reg2).Correlation_error .< 10*eps() # No change as we are already psd
calculate_mean_abs_distance(psd_mat, reg2).Volatility_error .< 10*eps()

reg3 = nearest_psd_matrix(psd_mat, ts2)
calculate_mean_abs_distance(psd_mat, reg3).Correlation_error .< 10*eps() # No change as we are already psd
calculate_mean_abs_distance(psd_mat, reg3).Volatility_error .< 1000*eps()

reg3_corr = nearest_psd_matrix(psd_mat, ts2; apply_to_covariance = false)
calculate_mean_abs_distance(psd_mat, reg3_corr).Correlation_error .< 10*eps() # No change as we are already psd
calculate_mean_abs_distance(psd_mat, reg3_corr).Volatility_error .< 10*eps()

reg4 = eigenvalue_clean(psd_mat, ts2)
calculate_mean_abs_distance(psd_mat, reg4).Correlation_error .> 2*eps()
calculate_mean_abs_distance(psd_mat, reg4).Volatility_error .> 10*eps()

reg4_cov = eigenvalue_clean(psd_mat, ts2; apply_to_covariance = false)
calculate_mean_abs_distance(psd_mat, reg4_cov).Correlation_error .> 2*eps()
calculate_mean_abs_distance(psd_mat, reg4_cov).Volatility_error .< 10*eps()

# Testing blocking and regularisation.
blocking_dd = put_assets_into_blocks_by_trading_frequency(ts2, 1.3, :spectral_covariance)
block_estimate = blockwise_estimation(ts2, blocking_dd)
block_estimate = identity_regularisation(block_estimate, ts2)
iscloser(calculate_mean_abs_distance(block_estimate, true_covar), calculate_mean_abs_distance(spectral_estimate2, true_covar))

# Testing getting noise
noise = estimate_microstructure_noise(ts1, assets)
noise_from_2s = two_scales_volatility(ts1, assets)[2]
C = merge(-, noise, noise_from_2s)
maximum(collect(values(C))) < 1E-13

# Testing print functions
print(ts1)
show(ts1)
print(true_covar)
show(true_covar)
print(preav_estimate2, 3, 6)
show(preav_estimate2, 3, 6)
true # Otherwise it tends to throw from not evaluating to Bool in testing.