# Alternative Krylov phiv methods using error estimates of Saad to automatically
# terminate Arnoldi/Lanczos iterations.
# Currently only expv for Lanczos is implemented.

########################################
# Cache types
abstract type SubspaceCache{T} end
abstract type HermitianSubspaceCache{T} <: SubspaceCache{T} end

mutable struct StegrCache{T,R<:Real} <: HermitianSubspaceCache{T}
    v::Vector{T} # Subspace-propagated vector
    w::Vector{T}
    sw::Stegr.StegrWork{R}
    StegrCache(::Type{T}, n::Integer) where T = new{T,real(T)}(
        Vector{T}(undef, n), Vector{T}(undef, n),
        Stegr.StegrWork(real(T), n))
end

"""
    expT!(α, β, t, cache)

Calculate the subspace exponential `exp(t*T)` for a tridiagonal
subspace matrix `T` with `α` on the diagonal and `β` on the
super-/subdiagonal, diagonalizing via `stegr!`.
"""
function expT!(α::AbstractVector{R}, β::AbstractVector{R}, t::Number,
               cache::StegrCache{T,R}) where {T,R<:Real}
    LAPACK.stegr!(α, β, cache.sw)
    sel = 1:length(α)
    @inbounds for i = sel
        cache.w[i] = exp(t*cache.sw.w[i])*cache.sw.Z[1,i]
    end
    mul!(@view(cache.v[sel]), @view(cache.sw.Z[sel,sel]), @view(cache.w[sel]))
end

get_subspace_cache(Ks::KrylovSubspace{T,U}) where {T,U<:Complex} =
    error("Subspace exponential caches not yet available for non-Hermitian matrices.")
get_subspace_cache(Ks::KrylovSubspace{T,U}) where {T,U<:Real} =
    StegrCache(T, Ks.maxiter)

########################################
# Phiv with error estimate as termination condition
"""
    expv!(w, t, A, b, Ks, cache)

Alternative interface for calculating the action of `exp(t*A)` on the
vector `b`, storing the result in `w`. The Krylov iteration is
terminated when an error estimate for the matrix exponential in the
generated subspace is below the requested tolerance. `Ks` is a
`KrylovSubspace` and `typeof(cache)<:HermitianSubspaceCache`, the
exact type decides which algorithm is used to compute the subspace
exponential.
"""
function expv!(w::AbstractVector{T}, t::Number, A, b::AbstractVector{T},
               Ks::KrylovSubspace{T, B, B}, cache::HSC;
               atol::B=1.0e-8, rtol::B=1.0e-4,
               m=min(Ks.maxiter, size(A,1)),
               verbose::Bool=false) where {B, T <: Number, HSC <: HermitianSubspaceCache}
    if m > Ks.maxiter
        resize!(Ks, m)
    else
        Ks.m = m # might change if error estimate is below requested tolerance
    end

    V, H = getV(Ks), getH(Ks)
    Ks.beta = norm(b)
    @. V[:, 1] = b / Ks.beta

    ε = atol + rtol * Ks.beta
    verbose && @printf("Initial norm: β₀ %e, stopping threshold: %e\n", Ks.beta, ε)

    α = @diagview(H)
    β = @diagview(H,-1)
    n = size(V, 1)

    for j ∈ 1:m
        lanczos_step!(j, A, V, α, β)
        expT!(@view(α[1:j]), @view(β[1:j]), t, cache)

        # This is practical error estimate Er₂ from
        #
        #   Saad, Y. (1992). Analysis of some Krylov subspace
        #   approximations. SIAM Journal on Numerical Analysis.
        σ = β[j]*Ks.beta*abs(cache.v[j])
        verbose && @printf("iter %d, α[%d] %e, β[%d] %e, σ %e\n",j, j, α[j], j, β[j], σ)
        if σ < ε
            Ks.m = j
            break
        end
    end
    verbose && println("Krylov subspace size: ", Ks.m)

    lmul!(Ks.beta, mul!(w, @view(Ks.V[:,1:Ks.m]), @view(cache.v[1:Ks.m])))
end


#--------------------------------------------------------------------------------------------------
# LinMap tests
#--------------------------------------------------------------------------------------------------
using Test
using Hensel

#--------------------------------------------------------------------------------------------------

@test Hensel.LinMap((1.0,2.0), 2, 10)(1) == 1.50048828125
@test Hensel.LinMap((1.0,2.0), 2, 10)([1,0,0,0,0,0,0,0,0,0]) == 1.50048828125
@test Hensel.LinMap((1.0,2.0), 2, 10)([1,1,1,1,1,1,1,1,1,1]) == 1.99951171875

@test Hensel.LinMap((1.0,2.0), 2, 10)(1.50048828125) == 1
@test Hensel.LinMap((1.0,2.0), 2, 10)(1.50048828125, return_type="hensel") == [1,0,0,0,0,0,0,0,0,0]
@test Hensel.LinMap((1.0,2.0), 2, 10)(1.99951171875, return_type="hensel") == [1,1,1,1,1,1,1,1,1,1]

#--------------------------------------------------------------------------------------------------


using Images, ShapeFromShading
using Test

@testset "ShapeFromShading" begin
    include("shah.jl")
    include("pentland.jl")
    include("discreteshape.jl")
    include("discreteshapebound.jl")
    include("photometric.jl")
    include("deterministic_custom.jl")
    include("deterministic.jl")
end

@testset "Integration" begin
    include("frankot.jl")
    include("path.jl")
    include("splitpath.jl")
    include("horn.jl")
    include("durou.jl")
end

@testset "SyntheticSurface" begin
    include("synthsphere.jl")
    include("synthvase.jl")
    include("tent.jl")
    include("dem.jl")
    include("cake.jl")
    include("cake2.jl")
    include("ripple.jl")
    include("ripple2.jl")
    include("synthgaussian.jl")
    include("supergaussian.jl")
    include("estimatealbedo.jl")
end


using Base.Test
using LinearControl

# Test System
A = [1 2 3; 2 1 3; 3 2 1]
B = [1; 0; 0]
C = [0 0 1]
D = [0]

num = [1]
den = [1 2 10]

# State Space
myss = SS(A,B,C,D)

# Transfer Function
mytf = TF(num,den)

# Linear Model
mylm1 = LinearModel(myss)
mylm2 = LinearModel(mytf)
@assert is(mylm1,mylm2)

@assert is(mylm1.tf, mytf)
@assert is(mylm2.ss, myss)


"""
    tempcorr(w::CTASensor, E, temp, T0)

Corrects the anemometer output due to differences in flow temperature.

 * `w::CTASensor` object
 * `E` Anemometer output
 * `temp` Temperature during measurement
 * `T0` Reference calibration temperature. If not provided, use ref. temp of Rw

This correction assumes that only the temperature changed and not fluid properties. Another function that takes that into account will be developped later.

The correction works by using the definition of convective heat transfer:

``Q = h A (Tw - Too)``

where ``Tw`` is the operating temperature and ``Too`` is the flow temperature. Imagine that for the *same* flow speed the flow temperature changes to ``Too'``. Now, the assuming the the convection coefficient (``h``) remains the same, the heat output will change:

``Q' = h A (Tw - Too')``

Since the anemometer output is proportial to the voltage accross the sensor element, dividing the above equations and setting

``Q = α²E²/Rw``

we get the following equation for output correction:

``E' = E × ( (Tw - Too') / (Tw - Too) )^(1/2)``

If between the resistance element and the flow is an insulator (very common in thermistors and hot film probes), the correction equation remains the same event though this insulator should be taken into account:

``Q = h A (Tw - Too - Q⋅β)``

where ``Q⋅β`` corrects the sensor temperature (``Tw``) to the outer surface mean temperature.

"""
function tempcorr(w::CTASensor, E, temp, T0)
    Tw = optemperature(w)
    return E * sqrt( (Tw - T0) / (Tw - temp) )
end

tempcorr(w::CTASensor, E, temp) = tempcorr(w, E, temp, temperature(w.R))

"""
    tempcorr(w::CCASensor, E, temp, T0)

Corrects the anemometer output due to differences in flow temperature.

 * `w::CCASensor` Constant current sensor
 * `E` Anemometer output
 * `temp` Temperature during measurement
 * `T0` Reference calibration temperature. If not provided, use ref. temp of R.

This correction assumes that only the temperature changed and not fluid properties. Another function that takes that into account will be developped later.

The correction works by using the definition of convective heat transfer:

``Q = h A (Tw - Too)``

where ``Tw`` is the operating temperature and ``Too`` is the flow temperature. Imagine that for the *same* flow speed the flow temperature changes to ``Too'``. Now, the assuming the the convection coefficient (``h``) remains the same, the heat output will change:

``Q' = h A (Tw - Too')``

Since the anemometer output is proportial to the voltage accross the sensor element, dividing the above equations and setting

``Q = α²E²/Rw``

we get the following equation for output correction:

``E' = E × ( (Tw - Too') / (Tw - Too) )^(1/2)``

If between the resistance element and the flow is an insulator (very common in thermistors and hot film probes), the correction equation remains the same event though this insulator should be taken into account:

``Q = h A (Tw - Too - Q⋅β)``

where ``Q⋅β`` corrects the sensor temperature (``Tw``) to the outer surface mean temperature.

"""
function tempcorr(w::CCASensor{Resistor}, E, temp, T0)

    α = w.R.α
    return E / (1.0 + α*(temp - T0))
end

tempcorr(w::CCASensor{Resistor}, E, temp) = tempcorr(w, E, temp, temperature(w.R))

function funroot(f, x1, x2, eps=1e-7, maxiter=100)
    y1 = f(x1)
    y2 = f(x2)
    if abs(y1) < eps return x1 end
    if abs(y2) < eps return x2 end

    err = 0.0

    for i in 1:maxiter
        a = (y2-y1) / (x2 - x1)
        b = y1 - a*x1
        xn = -b/a
        err = abs(xn - x2)
        if err < eps
            return xn
        end
        x1 = x2
        x2 = xn
        y1 = y2
        y2 = f(x2)
        if abs(y2) < eps
            return x2
        end

    end

    error("funroot failed to converge after $maxiter iterations and with $err error!")
    return x2
end

function tempcorr(w::CCASensor, E, temp, T0; eps=1e-7, maxiter=100)

    R = w.R
    g = gain(w)
    I = current(w)
    Tw = temperature(R, E/(g*I))
    To = T0
    ΔT = Tw - temp

    f = Eo -> E*(temperature(R, Eo/(g*I)) - To) - Eo * ΔT
    return funroot(f, E, 1.01*E, eps, maxiter)
    #error("tempcorr(w::CCASensor, $E, $temp failed to converge in $maxiter iterations with $err error!")
    return Eo
end

tempcorr(w::CCASensor, E, temp; eps=1e-7, maxiter=100) =
    tempcorr(w, E, temp, temperature(w.R), eps=eps, maxiter=maxiter)


using Test, NeuroAnalysis, BenchmarkTools, DataFrames, Plots, FileIO, FFTW

@testset "NeuroAnalysis" begin

    include("spiketest.jl")
    include("imagetest.jl")
    include("conditiontest.jl")
    include("regressiontest.jl")
    include("visualizationtest.jl")

## 1D grating
# @manipulate for f in 0:0.1:1, p in 0:0.001:1
#     plot([x->gratingf(x,f=f,phase=p),x->cas(x,f=f,phase=sin2cas(p))],-2,2,labels=["sin" "cas"])
# end
## 2D grating
# @manipulate for θ in 0:0.01:π, f in 0:0.05:1, sinp in 0:0.001:1, kx in -1:0.05:1,ky in -1:0.05:1, casp in 0:0.001:1
#     x=y=-2:0.05:2
#     sg = [gratingf(i,j,θ=θ,f=f,phase=sinp) for j in reverse(y),i in x]
#     scg = [cas(i,j;sin2cas(θ,f,sinp)...) for j in reverse(y),i in x]
#     cg = [cas(i,j,kx=kx,ky=ky,phase=casp) for j in reverse(y),i in x]
#     csg = [gratingf(i,j;cas2sin(kx,ky,casp)...) for j in reverse(y),i in x]
#
#     p = plot(layout=(2,3),leg=false,size=(600,400),clims=(-1,1),frame=:none,aspect_ratio=:equal,yflip=true)
#     heatmap!(p,subplot=1,sg,color=:grays,title="sin")
#     heatmap!(p,subplot=2,scg,color=:grays,title="sin2cas")
#     heatmap!(p,subplot=3,sg.-scg,color=:grays,title="sin - sin2cas")
#     heatmap!(p,subplot=4,cg,color=:grays,title="cas")
#     heatmap!(p,subplot=5,csg,color=:grays,title="cas2sin")
#     heatmap!(p,subplot=6,cg.-csg,color=:grays,title="cas - cas2sin")
#     p
# end

## functions
plot(vmf,-π,π)
plot(gvmf,-π,π)
plot(gratingf,-3,3)
plot(gaussianf,-3,3)
plot(gaborf,-3,3)
plot(dogf,-3,3)

funnames = ["grating","gaussian","gabor","dog"]
plot([gratingf,gaussianf,gaborf,dogf],-3,3,labels=permutedims(funnames),lw=[2 2 3 3],xlabel="y′")

x=y=-3:0.05:3;z=[]
push!(z,[gratingf(i,j,θ=0.25π,f=0.5) for j in y,i in x])
push!(z,[gaussianf(i,j,σ₁=0.8,σ₂=0.5,θ=0.25π) for j in y,i in x])
push!(z,[gaborf(i,j,σ₁=0.8,σ₂=0.5,θ=0.25π,f=0.5) for j in y,i in x])
push!(z,[dogf(i,j,σₑ₁=0.8,σₑ₂=0.5,σᵢ₁=1,σᵢ₂=0.7,θₑ=0.25π,θᵢ=0.25π) for j in y,i in x])

p = plot(layout=(2,2),legend=false,size=(600,600))
foreach(i->heatmap!(p,z[i],subplot=i,aspect_ratio=:equal,frame=:none,color=:coolwarm,clims=(-1,1),title=funnames[i]),1:4)
p

fs = 20
f = 2
t = 0:1/fs:10
x = cos.(2π*(f*t))
plot(t,x,xlabel="Second")
Fs = dft(x,fs,0,f)
FF = rfft(x)
plot(abs.(FF))

i = round(Int,f*length(x)/fs) + 1
@test FF[i] ≈ Fs[2]
@test FF[1] ≈ Fs[1]

end


# Tseng, "On Accelerated Proximal Gradient Methods for Convex-Concave
# Optimization" (2008).
#
# Beck, Teboulle, "A Fast Iterative Shrinkage-Thresholding Algorithm
# for Linear Inverse Problems", SIAM Journal on Imaging Sciences, vol. 2,
# no. 1, pp. 183-202 (2009).

using Base.Iterators
using ProximalAlgorithms.IterationTools
using ProximalOperators: Zero
using LinearAlgebra
using Printf

"""
    ForwardBackwardIteration(; <keyword-arguments>)

Instantiate the forward-backward splitting algorithm (see [1, 2]) for solving
optimization problems of the form

    minimize f(Ax) + g(x),

where `f` is smooth and `A` is a linear mapping (for example, a matrix).

# Arguments
- `x0`: initial point.
- `f=Zero()`: smooth objective term.
- `A=I`: linear operator (e.g. a matrix).
- `g=Zero()`: proximable objective term.
- `Lf=nothing`: Lipschitz constant of the gradient of x ↦ f(Ax).
- `gamma=nothing`: stepsize to use, defaults to `1/Lf` if not set (but `Lf` is).
- `adaptive=false`: forces the method stepsize to be adaptively adjusted.
- `minimum_gamma=1e-7`: lower bound to `gamma` in case `adaptive == true`.
- `fast=false`: enables Nesterov acceleration.

# References
- [1] Tseng, "On Accelerated Proximal Gradient Methods for Convex-Concave
Optimization" (2008).
- [2] Beck, Teboulle, "A Fast Iterative Shrinkage-Thresholding Algorithm
for Linear Inverse Problems", SIAM Journal on Imaging Sciences, vol. 2, no. 1,
pp. 183-202 (2009).
"""

@Base.kwdef struct ForwardBackwardIteration{R,C<:Union{R,Complex{R}},Tx<:AbstractArray{C},Tf,TA,Tg}
    f::Tf = Zero()
    A::TA = I
    g::Tg = Zero()
    x0::Tx
    Lf::Maybe{R} = nothing
    gamma::Maybe{R} = Lf === nothing ? nothing : (1 / Lf)
    adaptive::Bool = false
    minimum_gamma::R = real(eltype(x0))(1e-7)
    fast::Bool = false
end

Base.IteratorSize(::Type{<:ForwardBackwardIteration}) = Base.IsInfinite()

Base.@kwdef mutable struct ForwardBackwardState{R,Tx,TAx}
    x::Tx             # iterate
    Ax::TAx           # A times x
    f_Ax::R           # value of smooth term
    grad_f_Ax::TAx    # gradient of f at Ax
    At_grad_f_Ax::Tx  # gradient of smooth term
    gamma::R          # stepsize parameter of forward and backward steps
    y::Tx             # forward point
    z::Tx             # forward-backward point
    g_z::R            # value of nonsmooth term (at z)
    res::Tx           # fixed-point residual at iterate (= z - x)
    theta::R = one(real(eltype(x)))
    z_prev::Tx = copy(x)
end

f_model(state::ForwardBackwardState) = f_model(state.f_Ax, state.At_grad_f_Ax, state.res, state.gamma)

function Base.iterate(iter::ForwardBackwardIteration{R}) where {R}
    x = copy(iter.x0)
    Ax = iter.A * x
    grad_f_Ax, f_Ax = gradient(iter.f, Ax)

    gamma = iter.gamma

    if gamma === nothing
        # compute lower bound to Lipschitz constant of the gradient of x ↦ f(Ax)
        xeps = x .+ R(1)
        grad_f_Axeps, f_Axeps = gradient(iter.f, iter.A * xeps)
        L = norm(iter.A' * (grad_f_Axeps - grad_f_Ax)) / R(sqrt(length(x)))
        gamma = R(1) / L
    end

    # compute initial forward-backward step
    At_grad_f_Ax = iter.A' * grad_f_Ax
    y = x - gamma .* At_grad_f_Ax
    z, g_z = prox(iter.g, y, gamma)

    # compute initial fixed-point residual
    res = x - z

    state = ForwardBackwardState(
        x=x, Ax=Ax, f_Ax=f_Ax, grad_f_Ax=grad_f_Ax, At_grad_f_Ax=At_grad_f_Ax,
        gamma=gamma, y=y, z=z, g_z=g_z, res=res,
    )

    return state, state
end

function Base.iterate(iter::ForwardBackwardIteration{R}, state::ForwardBackwardState{R,Tx,TAx}) where {R,Tx,TAx}
    Az, f_Az, grad_f_Az, At_grad_f_Az = nothing, nothing, nothing, nothing
    a, b, c = nothing, nothing, nothing

    # backtrack gamma (warn and halt if gamma gets too small)
    while iter.gamma === nothing || iter.adaptive == true
        if state.gamma < iter.minimum_gamma
            @warn "parameter `gamma` became too small ($(state.gamma)), stopping the iterations"
            return nothing
        end
        f_Az_upp = f_model(state)
        Az = iter.A * state.z
        grad_f_Az, f_Az = gradient(iter.f, Az)
        tol = 10 * eps(R) * (1 + abs(f_Az))
        if f_Az <= f_Az_upp + tol
            break
        end
        state.gamma *= 0.5
        state.y .= state.x .- state.gamma .* state.At_grad_f_Ax
        state.g_z = prox!(state.z, iter.g, state.y, state.gamma)
        state.res .= state.x .- state.z
    end

    if iter.fast == true
        theta1 = (R(1) + sqrt(R(1) + 4 * state.theta^2)) / R(2)
        extr = (state.theta - R(1)) / theta1
        state.theta = theta1
        state.x .= state.z .+ extr .* (state.z .- state.z_prev)
        state.z_prev, state.z = state.z, state.z_prev
    else
        state.x, state.z = state.z, state.x
    end

    # TODO: if iter.fast == true, in the adaptive case we should still be able
    # to save some computation by extrapolating Ax and (if f is quadratic)
    # f_Ax, grad_f_Ax, At_grad_f_Ax.
    if iter.fast == false && (iter.gamma === nothing || iter.adaptive == true)
        state.Ax = Az
        state.f_Ax = f_Az
        state.grad_f_Ax = grad_f_Az
    else
        mul!(state.Ax, iter.A, state.x)
        state.f_Ax = gradient!(state.grad_f_Ax, iter.f, state.Ax)
    end

    mul!(state.At_grad_f_Ax, adjoint(iter.A), state.grad_f_Ax)
    state.y .= state.x .- state.gamma .* state.At_grad_f_Ax
    state.g_z = prox!(state.z, iter.g, state.y, state.gamma)

    state.res .= state.x .- state.z

    return state, state
end

# Solver

struct ForwardBackward{R, K}
    maxit::Int
    tol::R
    verbose::Bool
    freq::Int
    kwargs::K
end

function (solver::ForwardBackward)(x0; kwargs...)
    stop(state::ForwardBackwardState) = norm(state.res, Inf) / state.gamma <= solver.tol
    disp((it, state)) =
        @printf("%5d | %.3e | %.3e\n", it, state.gamma, norm(state.res, Inf) / state.gamma)
    iter = ForwardBackwardIteration(; x0=x0, solver.kwargs..., kwargs...)
    iter = take(halt(iter, stop), solver.maxit)
    iter = enumerate(iter)
    if solver.verbose
        iter = tee(sample(iter, solver.freq), disp)
    end
    num_iters, state_final = loop(iter)
    return state_final.z, num_iters
end

ForwardBackward(; maxit=10_000, tol=1e-8, verbose=false, freq=100, kwargs...) =
    ForwardBackward(maxit, tol, verbose, freq, kwargs)


#=
Linear Regression using Flux.jl

pkg> add Flux#master
pkg> add Plots

Reference:
https://qiita.com/cometscome_phys/items/f58174c0dad7ecb811ed

=#

using Random
using Flux
using Plots
#using CuArrays

function vandermonde_matrix(xs,maxdim)
    φ(xs,k) = xs.^k
    ret = zeros(Float32,maxdim,length(xs))
    for k in 0:maxdim-1
        ret[k+1,:]=φ(xs,k)
    end
    ret
end

function target_func(x)
    a₀,a₁,b₀ = 3.0,2.0,1.0
    return @. a₀*x + a₁*x^2 + b₀+ 3*cos(20*x)
end

function train(epochs,batchsize)
    k=4
    xs = range(-2,length=100,stop=2)
    X = vandermonde_matrix(xs,k)
    ys = target_func(xs)

    model = Dense(k,1)
    lossfn(x,y)=Flux.mse(model(x),y)
    optimizer = ADAM(params(model))
    for e in 1:epochs
        shuffle_indices=shuffle(1:length(ys))
        loss=0
        for i in 1:floor(Int,length(ys)/batchsize)
            batch_indice=shuffle_indices[(i-1)*batchsize+1:i*batchsize]
            batch_X=X[:,batch_indice]
            batch_ys=ys[batch_indice]
            data=[(batch_X[:,i], batch_ys[i]) for i in 1:batchsize]
            Flux.train!(lossfn, data, optimizer)
            loss=sum(lossfn(batch_X[:,i],batch_ys[i]) for i in 1:batchsize)/batchsize
        end
        if e%100==0
            @show e
            @show loss
        end
    end

    y_predict = [model(X[:,i]).data[1] for i=1:length(ys)]
    plt = plot(xs, [ys,y_predict], marker=:circle,label=["Data","Estimation"])

end

function main()
    train(1000,16)
end

main()


# k-space trajectory utilities

export KSpaceCartesianSampling
export kspace_sampling, kspace_Cartesian_sampling, downscale, downscale_phase_encode_index

## k-space trajectory

coord(K::AbstractKSpaceSampling) = K.K

struct KSpaceFixedSizeSampling{T}<:AbstractKSpaceFixedSizeSampling{T}
    K::AbstractArray{T,3} # Array (nt,nk,3)
end

Base.getindex(K::AbstractKSpaceFixedSizeSampling, t::Integer) = K.K[t,:,:]
Base.size(K::AbstractKSpaceFixedSizeSampling) = size(K.K)[1:2]

kspace_sampling(K::AbstractArray{T,3}) where {T<:Real} = KSpaceFixedSizeSampling{T}(K)

struct KSpaceCartesianSampling{T}<:AbstractKSpaceCartesianSampling{T}
    X::RegularCartesianSpatialSampling{T}
    K::AbstractArray{T,3} # Array (nt,nk,3)
    phase_encoding::NTuple{2,Integer}
    readout::Integer
    subsampling::AbstractVector{<:Integer}
end

function kspace_Cartesian_sampling(X::RegularCartesianSpatialSampling{T}; phase_encoding::NTuple{2,Integer}=(1,2), subsampling::Union{Nothing,AbstractVector{<:Integer}}=nothing) where {T<:Real}

    # Set phase-encode/readout dimensions
    nx, ny, nz = X.n
    hx, hy, hz = X.h
    ((phase_encoding == (1,2)) || (phase_encoding == (2,1))) && (readout = 3)
    ((phase_encoding == (1,3)) || (phase_encoding == (3,1))) && (readout = 2)
    ((phase_encoding == (2,3)) || (phase_encoding == (3,2))) && (readout = 1)
    nt = (X.n[phase_encoding[1]], X.n[phase_encoding[2]]); nk = X.n[readout]
    perm = (phase_encoding..., readout)

    # Full k-space grid
    kx = T(pi)/hx*collect(coord_norm(nx))
    ky = T(pi)/hy*collect(coord_norm(ny))
    kz = T(pi)/hz*collect(coord_norm(nz))
    Kx = reshape(permutedims(repeat(reshape(kx,:,1,1); outer=(1,ny,nz)), perm), prod(nt), nk)
    Ky = reshape(permutedims(repeat(reshape(ky,1,:,1); outer=(nx,1,nz)), perm), prod(nt), nk)
    Kz = reshape(permutedims(repeat(reshape(kz,1,1,:); outer=(nx,ny,1)), perm), prod(nt), nk)
    K = cat(Kx, Ky, Kz; dims=3)

    # Subsampling and reordering
    isnothing(subsampling) && (subsampling = 1:prod(nt))
    K = K[subsampling, :, :]

    return KSpaceCartesianSampling{T}(X, K, phase_encoding, readout, subsampling)

end

function coord_norm(n::Integer)
    (mod(n,2) == 0) ? (c = -div(n, 2):div(n, 2)-1) : (c = -div(n-1, 2):div(n-1, 2))
    return c/norm(c, Inf)
end

function downscale(K::KSpaceCartesianSampling{T}; fact::Integer=1) where {T<:Real}
    (fact == 0) && (return K)
    X_h = downscale(K.X; fact=fact)
    phase_encode_idx_h, readout_idx_h = downscale_phase_encode_index(K; fact=fact, readout=true)
    K_h = K.K[phase_encode_idx_h, readout_idx_h, :]
    return KSpaceCartesianSampling{T}(X_h, K_h, K.phase_encoding, K.readout, K.subsampling)
end

function downscale_phase_encode_index(K::KSpaceCartesianSampling{T}; fact::Integer=1, readout::Bool=false) where {T<:Real}
    k_max = T(pi)./K.X.h
    pe = K.phase_encoding
    r = K.readout
    phase_encode_idx_h = findall((K.K[:,1,pe[1]] .< k_max[pe[1]]/2^fact) .&& (K.K[:,1,pe[1]] .>= -k_max[pe[1]]/2^fact) .&& (K.K[:,1,pe[2]] .< k_max[pe[2]]/2^fact) .&& (K.K[:,1,pe[2]] .>= -k_max[pe[2]]/2^fact))
    readout_idx_h = findall((K.K[1,:,r] .< k_max[r]/2^fact) .&& (K.K[1,:,r] .>= -k_max[r]/2^fact))
    readout ? (return (phase_encode_idx_h, readout_idx_h)) : (return phase_encode_idx_h)
end

function downscale(d::AbstractArray{CT,2}, K::KSpaceCartesianSampling{T}; fact::Integer=1, flat::Bool=false, coeff::Real=1) where {T<:Real,CT<:RealOrComplex{T}}
    (fact == 0) && (return d)
    i_pe, i_r = downscale_phase_encode_index(K; fact=fact, readout=true)
    return anti_aliasing_filter(K; fact=fact, flat=flat, coeff=coeff).*d[i_pe, i_r]/T(sqrt(2.0^(3*fact)))
end


@testset "Cylinder" begin
    N = 1000
    pos = rand(SVector{3,Float32}, N)
    cylinder = Cylinder(0.5, 0.5, 0.2, 0.8, 0.1)

    search_correct = true
    for i = 1:length(pos)
        v = pos[i]
        dv = v - SVector(0.5,0.5,0.5)
        expected = dv[1]^2 + dv[2]^2 < 0.01 && abs(dv[3]) < 0.3
        actual = v ∈ cylinder
        search_correct &= actual == expected
    end
    @test search_correct

    @test boundingbox(cylinder) == BoundingBox(0.4, 0.4, 0.2, 0.6, 0.6, 0.8)

    @test volume(cylinder) ≈ 0.01884955592153876
end


function miyak00()
#
#  test polynomial suggested by Goedecker
#  square of Fibocacci polynomial
#
    p = poly([(1.1+1.1im)*ones(4);(3.2+2.3im)*[1;1];2.1+1.5im]);
    z = [[(1.1+1.1im);(3.2+2.3*im);2.1+1.5*im] [4;2;1]];
    p, PolyZeros(z)
end


##################################################
# Computation of the NR diffusion coefficients
##################################################

##################################################
# Preliminary useful functions
##################################################

"""
    logCoulomb(iBath)

Coulomb logarithm of the family indexed by iBath in the source file.
Defined as the logarithm of the ratio of the mass of the SMBH over the mass of an element of the family iBath.

# Remark
Using the convention where log is the Neperian logarithm (in base e).

"""
function logCoulomb(iBath::Int64)
    m = tabmBath[iBath]
    Q = mBH/m
    return log(Q)
end

"""
    prefac_J(E)

Numerical prefactor J(E) used to convert the PDF f(E,R) to the number density N(E,R), where R = j^2.
Depends on the binding energy E, which is the opposite of the gravitational energy.

Explicit formula is N(E,R) = J(E)*f(E,R).
"""
function prefac_J(E::Float64)
    return (2.0)^(1/2)*pi^3*(G*mBH)^3*E^(-5/2)
end

"""
    E_bind(a)

Binding energy of the system, which is the opposite of the gravitational energy.
Defined as `E(a) = G*mBH / (2a)`, where `a` is the semi-major axis of the orbit, G is the gravitational constant and mBH the mass of the SMBH.
"""
function E_bind(a::Float64)
    return (G*mBH)/(2.0*a)
end

"""
    j_to_R(j)

Converts reduced angular momentum j = J/Jc(a) into R = j^2, where J is the angular momentum and Jc(a) the angular momentum of a circular orbit.
"""
function j_to_R(j::Float64)
    return j^2
end

"""
    DF_E(E,iBath)

R-averaged distribution function in binding enery E, defined in E space, for the family indexed by iBath in the source file.
R is the squared reduced angular momentum j^2.
"""
function DF_E(E::Float64,iBath::Int64)
    E0    = E_bind(taba0Bath[iBath])
    NBar  = tabNbarBath[iBath]
    gamma = tabgammaBath[iBath]
    return NBar/prefac_J(E) * (G*mBH)/(2.0*E^2) * (E0/E)^(2.0-gamma)
end

"""
    gamma_prefactor(iBath)

Gamma prefactor which appears in Fi integrals, for the family indexed by iBath in the source file.
Used in the computation of the non-resonant diffusion coefficients.
"""
function gamma_prefactor(iBath::Int64)
    m = tabmBath[iBath]
    return 4.0*pi*(G*m)^2 * logCoulomb(iBath)
end

"""
    xpm(R)

xplus and xminus quantities which appear in the bounds of the integrals involved in the computation of the non-resonant diffusion coefficients.

# Remark
x = E(r)/gravitational potential(r)
"""
function xpm(R::Float64)
    return (1.0+sqrt(1.0-R))/2.0,(1.0-sqrt(1.0-R))/2.0
end

"""
    inverse_reduced_E_bind(Ep,E)

Reduced binding energy used in the computation of integrals Fi.
Defined as Ep over E.
"""
function inverse_reduced_E_bind(Ep::Float64,E::Float64)
    return Ep/E
end

##################################################
# Bounds in A_k, B_k and C_k (xp>1/s) integrals
##################################################

"""
    _beta(R,s)

Bound in Ck integral when x_plus > 1/s, where s is the reduced binding energy and R is the squared reduced angular momentum.
"""
function _beta(R::Float64,s::Float64)
    xp,xm = xpm(R)
    res = 1.0/(xm*s)
    res = sqrt(res)
    return acosh(res)
end

"""
    _eta(R)

Bound in Ak integral when xp <= 1/s, where s is the reduced binding energy and R is the squared reduced angular momentum.
"""
function _eta(R::Float64)
    xp,xm = xpm(R)
    res = 1.0/(2.0*xm)
    res = sqrt(res)
    return acosh(res)
end

"""
    _delta(R)

Bound in Bk integral when xp <= 1/s, where s is the reduced binding energy and R is the squared reduced angular momentum.
"""
function _delta(R::Float64)
    xp,xm = xpm(R)
    res = 1.0/(2.0*xp)
    res = sqrt(res)
    return acos(res)
end

##################################################
# Computation of the NR diffusion coefficients.
##################################################

"""
    _fB(theta,s,R)

Computes all seven integrands within Bk integrals for 1<=k<=7.
Returns a 7-uple, ordered from k=1 to k=7.
See `Conrad et Kulsrud (1978)`.

# Remarks
- All the integrands are computed at the same time.
- Avoids multiple evalution of cos(theta), xplus and xminus.
- Useful when xplus <= 1/s.

# Arguments
- `theta::Float64`: Integration variable.
- `s    ::Float64`: Reduced binding energy.
- `R    ::Float64`: Squared reduced angular momentum j^2.
"""
function _fB(theta::Float64,s::Float64,R::Float64)
    xp,xm = xpm(R)
    ym    = xm/xp
    csTh  = cos(theta)

    fB1 = csTh^3/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(1)/sqrt(1.0-xp*csTh^2)^(1)
    fB2 = csTh^5/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(1)/sqrt(1.0-xp*csTh^2)^(3)
    fB3 = csTh^7/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(1)/sqrt(1.0-xp*csTh^2)^(1)
    fB4 = csTh/sqrt(csTh^2-ym)   * sqrt(1.0-s*xp*csTh^2)^(3)/sqrt(1.0-xp*csTh^2)^(1)
    fB5 = csTh^3/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(3)/sqrt(1.0-xp*csTh^2)^(3)
    fB6 = csTh^5/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(3)/sqrt(1.0-xp*csTh^2)^(5)
    fB7 = csTh^7/sqrt(csTh^2-ym) * sqrt(1.0-s*xp*csTh^2)^(3)/sqrt(1.0-xp*csTh^2)^(3)
    return fB1,fB2,fB3,fB4,fB5,fB6,fB7
end

"""
    _fAC(theta,s,R)

Computes all 14 integrands within Ak and Ck integrals for 1<=k<=7.
Returns a 7-uple, ordered from k=1 to k=7.
See `Conrad et Kulsrud (1978)`.

# Remarks
- All the integrands are computed at the same time.
- Avoids multiple evalution of cos(theta), xplus and xminus.
- Useful when xplus <= 1/s for Ak.
- Useful when xplus > 1/s for Ck.

# Arguments
- `theta::Float64`: Integration variable.
- `s    ::Float64`: Reduced binding energy.
- `R    ::Float64`: Squared reduced angular momentum j^2.
"""
function _fAC(theta::Float64,s::Float64,R::Float64)
    xp,xm = xpm(R)
    ym    = xp/xm
    cshTh = cosh(theta)

    fAC1 = cshTh^3/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(1)/sqrt(1.0-xm*cshTh^2)^(1)
    fAC2 = cshTh^5/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(1)/sqrt(1.0-xm*cshTh^2)^(3)
    fAC3 = cshTh^7/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(1)/sqrt(1.0-xm*cshTh^2)^(1)
    fAC4 = cshTh  /sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(3)/sqrt(1.0-xm*cshTh^2)^(1)
    fAC5 = cshTh^3/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(3)/sqrt(1.0-xm*cshTh^2)^(3)
    fAC6 = cshTh^5/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(3)/sqrt(1.0-xm*cshTh^2)^(5)
    fAC7 = cshTh^7/sqrt(ym-cshTh^2) * sqrt(1.0-s*xm*cshTh^2)^(3)/sqrt(1.0-xm*cshTh^2)^(3)
    return fAC1,fAC2,fAC3,fAC4,fAC5,fAC6,fAC7
end

"""
    _C_NR(s,R,nbC=500)

Computes all seven Ck integrals for 1<=k<=7 used for the computation of the non-resonance diffusion coefficients.
Midpoint sampling rule is used for integration.
Returns a 7-uple, ordered from k=1 to k=7.
See `Conrad et Kulsrud (1978)`.

# Remarks
- All the integrands are computed at the same time.
- Avoids multiple evalution of cos(theta), xplus and xminus, etc.

# Arguments
- `s  ::Float64`: Reduced binding energy.
- `R  ::Float64`: Squared reduced angular momentum j^2.
- `nbC::Float64`: Sampling number of the computed integrals. Set as 500 by default.
"""
function _C_NR(s::Float64,R::Float64,nbC::Int64=500)
    xp,xm = xpm(R)
    C1,C2,C3,C4,C5,C6,C7 = 0.0,0.0,0.0,0.0,0.0,0.0,0.0

    if xp>1/s
        supC = _beta(R,s)
        stepC = supC/nbC
        prefactorC = 4.0*stepC/pi
    else
        supA = _eta(R)
        supB = _delta(R)
        stepA = supA/nbC
        stepB = supB/nbC
        prefactorA = 4.0*stepA/pi
        prefactorB = 4.0*stepB/pi
        A1,A2,A3,A4,A5,A6,A7 = 0.0,0.0,0.0,0.0,0.0,0.0,0.0
        B1,B2,B3,B4,B5,B6,B7 = 0.0,0.0,0.0,0.0,0.0,0.0,0.0
    end
    for jTheta=1:nbC
        if xp>1/s
            thetaC = (jTheta-0.5)*stepC
            fAC1,fAC2,fAC3,fAC4,fAC5,fAC6,fAC7 = _fAC(thetaC,s,R)
            C1 += fAC1
            C2 += fAC2
            C3 += fAC3
            C4 += fAC4
            C5 += fAC5
            C6 += fAC6
            C7 += fAC7
        else
            thetaA = (jTheta-0.5)*stepA
            thetaB = (jTheta-0.5)*stepB
            fAC1,fAC2,fAC3,fAC4,fAC5,fAC6,fAC7 = _fAC(thetaA,s,R)
            fB1,fB2,fB3,fB4,fB5,fB6,fB7 = _fB(thetaB,s,R)
            A1 += fAC1
            A2 += fAC2
            A3 += fAC3
            A4 += fAC4
            A5 += fAC5
            A6 += fAC6
            A7 += fAC7
            B1 += fB1
            B2 += fB2
            B3 += fB3
            B4 += fB4
            B5 += fB5
            B6 += fB6
            B7 += fB7
        end
    end
    if xp>1/s
        C1 *= prefactorC*xm
        C2 *= prefactorC*xm^2
        C3 *= prefactorC*xm^3
        C4 *= prefactorC
        C5 *= prefactorC*xm
        C6 *= prefactorC*xm^2
        C7 *= prefactorC*xm^3
    else
        A1 *= prefactorA*xm
        A2 *= prefactorA*xm^2
        A3 *= prefactorA*xm^3
        A4 *= prefactorA
        A5 *= prefactorA*xm
        A6 *= prefactorA*xm^2
        A7 *= prefactorA*xm^3
        B1 *= prefactorB*xp
        B2 *= prefactorB*xp^2
        B3 *= prefactorB*xp^3
        B4 *= prefactorB
        B5 *= prefactorB*xp
        B6 *= prefactorB*xp^2
        B7 *= prefactorB*xp^3
        C1 = A1 + B1
        C2 = A2 + B2
        C3 = A3 + B3
        C4 = A4 + B4
        C5 = A5 + B5
        C6 = A6 + B6
        C7 = A7 + B7
    end
    return C1,C2,C3,C4,C5,C6,C7
end

"""
    _F_NR_0(E,iBath)

Analytic evaluation of the F0 integral, at binding energy E and for the family indexed by iBath in the source file.
See `Conrad et Kulsrud (1978)`.
"""
function _F_NR_0(E::Float64,iBath::Int64)
    a0 = taba0Bath[iBath]
    E0 = E_bind(a0)
    NBar  = tabNbarBath[iBath]
    gamma = tabgammaBath[iBath]
    gamma_pref = gamma_prefactor(iBath)
    return 4.0*pi*gamma_pref* a0 * 1.0/prefac_J(E)*NBar/(gamma-0.5) * (E0/E)^(3.0-gamma)
end

"""
    _F_NR(E,R,iBath,nbC=500)

Computes all seven Fk integrals for 1<=k<=7 used for the computation of the non-resonance diffusion coefficients.
Midpoint sampling rule is used for integration.
Returns a 7-uple, ordered from k=1 to k=7.
See `Conrad et Kulsrud (1978)`.

# Remarks
- All the integrands are computed at the same time.
- Avoids multiple evalution of cos(theta), xplus and xminus, etc.

# Arguments
- `E    ::Float64`: Binding energy.
- `R    ::Float64`: Squared reduced angular momentum j^2.
- `iBath::Float64`: Index in the source file of the considered family.
- `nbC  ::Float64`: Sampling number of the computed integrals. Set as 500 by default.
"""
function _F_NR(E::Float64,R::Float64,iBath::Int64,nbF::Int64=500)
    xp,xm = xpm(R)
    inf = E
    sup = E/xm
    step = E/nbF * (1.0/xm - 1.0)
    F1,F2,F3,F4,F5,F6,F7 = 0.0,0.0,0.0,0.0,0.0,0.0,0.0
    for kEnergy=1:nbF
        Ep = inf+(kEnergy-0.5)*step
        s  = inverse_reduced_E_bind(Ep,E)
        C1,C2,C3,C4,C5,C6,C7 = _C_NR(s,R)
        DF = DF_E(Ep,iBath)
        F1 += DF*C1
        F2 += DF*C2
        F3 += DF*C3
        F4 += DF*C4
        F5 += DF*C5
        F6 += DF*C6
        F7 += DF*C7
    end
    prefactor = 4.0*pi*gamma_prefactor(iBath)*step
    F1 *= prefactor
    F2 *= prefactor
    F3 *= prefactor
    F4 *= prefactor
    F5 *= prefactor
    F6 *= prefactor
    F7 *= prefactor
    return F1,F2,F3,F4,F5,F6,F7
end

##################################################
# Diffusion coefficients
##################################################

"""
    DNR_EJ(E,j,iBath,m_test=0.0)

Non-resonant diffusion coefficients in (E,R) coordinates and Bar-Or&Alexander (2016) convention.
Returns the 5-uple `(DE, DEE/E, DJ/Jc, DJJ/Jc^2, DEJ/Jc)` for the family indexed by iBath.

# Arguments
- `E     ::Float64`: Binding energy.
- `j     ::Float64`: Reduced angular momentum `j=J/Jc`.
- `iBath ::Float64`: Index in the source file of the considered family.
- `m_test::Float64`: Mass of a test particle.
"""
function DNR_EJ(E::Float64,j::Float64,iBath::Int64,m_test::Float64=0.0)
    R = j^2
    F1,F2,F3,F4,F5,F6,F7 = _F_NR(E,R,iBath)
    F0 = _F_NR_0(E,iBath)
    m = tabmBath[iBath]
    D_E = -F0 + m_test/m * F1
    D_EE_over_E = 4/3*(F0 + F4)
    D_J_over_Jc = 1.0/(j*E) * ( (5.0 - 3.0*j^2)/(12.0)*F0 - j^2*(m_test + m)/(2.0*m)*F2 + F3 - 1/3*F7 )
    D_JJ_over_JcSQ = 1.0/E* ( (5.0 - 3.0*j^2)/6.0*F0 + 1/2*j^2*F6 - 1/2*j^2*F2 + 2.0*F3 - 2/3*F7 )
    D_EJ_over_Jc = -2/3*j*(F0 + F5)
    return D_E,D_EE_over_E,D_J_over_Jc,D_JJ_over_JcSQ,D_EJ_over_Jc
end

"""
    DNR_J_red(a,j)

Total non-resonant reduced diffusion coefficient DJ/Jc in (E,R) coordinates and Bar-Or&Alexander (2016) convention.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_J_red(a::Float64,j::Float64)
    res = 0.0
    E = E_bind(a)
    for iBath=1:nbBath
        res += DNR_EJ(E,j,iBath)[3]
    end
    return res
end

"""
    DNR_JJ_red(a,j)

Total non-resonant reduced diffusion coefficient DJJ/Jc^2 in (E,R) coordinates and Bar-Or&Alexander (2016) convention.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_JJ_red(a::Float64,j::Float64)
    res = 0.0
    E = E_bind(a)
    for iBath=1:nbBath
        res += DNR_EJ(E,j,iBath)[4]
    end
    return res
end

"""
    DNR_red(a,j,iBath)

Non-resonant reduced diffusion coefficients in (a,j) coordinates and Bar-Or&Alexander (2016) convention.
Returns the 5-uple `(Da, Daa, Dj, Djj, Daj)` for the family indexed by iBath.

# Arguments
- `a     ::Float64`: Semi-major axis.
- `j     ::Float64`: Reduced angular momentum `j=J/Jc`.
- `iBath ::Float64`: Index in the source file of the considered family.
"""
function DNR_red(a::Float64,j::Float64,iBath::Int64)
    E = E_bind(a)
    cst = 2.0/(G*mBH)
    D_E,D_EE_over_E,D_J_over_Jc,D_JJ_over_JcSQ,D_EJ_over_Jc = DNR_EJ(E,j,iBath)
    D_a = cst*a^2*(-D_E + D_EE_over_E)
    D_aa = cst*a^3*D_EE_over_E
    D_j = D_J_over_Jc + j/(2.0*E)*D_E + 1.0/(2.0*E)*D_EJ_over_Jc - j/(8.0*E)*D_EE_over_E
    D_jj = D_JJ_over_JcSQ + j/E*D_EJ_over_Jc + j^2/(4.0*E)*D_EE_over_E
    D_aj = -cst*a^2*(D_EJ_over_Jc + j/2*D_EE_over_E)
    return D_a,D_aa,D_j,D_jj,D_aj
end

"""
    DNR_aj(a,j)

Total non-resonant reduced diffusion coefficients in (a,j) coordinates and Bar-Or&Alexander (2016) convention.
Returns the 5-uple `(Da, Daa, Dj, Djj, Daj)`.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_aj(a::Float64,j::Float64) #Wrapped function to give total D_NR_jj
    D_a,D_aa,D_j,D_jj,D_aj = 0.0,0.0,0.0,0.0,0.0
    for iBath=1:nbBath
        Di_a,Di_aa,Di_j,Di_jj,Di_aj = DNR_red(a,j,iBath)
        D_a  += Di_a
        D_aa += Di_aa
        D_j  += Di_j
        D_jj += Di_jj
        D_aj += Di_aj
    end
    return D_a,D_aa,D_j,D_jj,D_aj
end

"""
    DNR_jj(a,j)

Total non-resonant reduced diffusion coefficient Djj in (a,j) coordinates and Bar-Or&Alexander (2016) convention.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_jj(a::Float64,j::Float64)
    return DNR_aj(a,j)[4]
end

"""
    DNR_j(a,j)

Total non-resonant reduced diffusion coefficient Dj in (a,j) coordinates and Bar-Or&Alexander (2016) convention.

# Arguments
- `a::Float64`: Semi-major axis.
- `j::Float64`: Reduced angular momentum `j=J/Jc`.
"""
function DNR_j(a::Float64,j::Float64)
    return DNR_aj(a,j)[3]
end


# Author: Vincent Arel-Bundock
# Contact: varel@umich.edu
# License: BSD-3
# Original code from the python statsmodels project
# https://github.com/statsmodels/statsmodels

function bandwidth_hall_sheather(n::Integer, q::Real, alpha::Real)
    z = quantile(Normal(), q)
    num = 1.5 * pdf(Normal(), z)^2
    den = 2.0 * z^2 + 1
    h = n^(-1/3) * quantile(Normal(), (1 - alpha / 2))^(2/3) * (num / den)^(1/3)
    return h
end

function kernel_epanechnikov!(u::Vector, x::Vector)
    for i in 1:length(x)
        if abs(x[i]) <= 1
            @inbounds u[i] = 0.75 * (1 - x[i]^2)
        else
            @inbounds u[i] = 0
        end
    end
    return
end

function kernel_epanechnikov(x::Vector)
u = similar(x)
kernel_epanechnikov!(u, x)
return u
end

function qreg_vcov(y::Vector, X::Matrix, beta::Vector, q::Real)
    resid = y - X * beta
    n = length(resid)
    iqre = quantile(resid, .75) - quantile(resid, .25)
    h = bandwidth_hall_sheather(n, q, .05)
    h1 = min(std(y), iqre / 1.34)
    h2 = quantile(Normal(), q + h) - quantile(Normal(), q - h)
    h = h1 * h2
    # TODO: This line could be optimized, but is a hassle
    fhat0 = 1 / (n * h) * sum(kernel_epanechnikov(resid / h))
    d = copy(resid)
    for i in 1:n
        if d[i] > 0
            @inbounds d[i] = (q / fhat0)^2
        else
            @inbounds d[i] = ((1 - q) / fhat0)^2
        end
    end

    xtx = X' * X
    xtdx = broadcast(*, X, d)' * X
    vcov = (xtx \ xtdx) / xtx

    return vcov
end


"""
    hadi1994(multivariateData)

Perform Hadi (1994) algorithm for a given multivariate data.

# Arguments
- `multivariateData::Array{Float64, 2}`: Multivariate data.

# Description
Algorithm starts with an initial subset and enlarges the subset to
obtain robust covariance matrix and location estimates. This algorithm
is an extension of `hadi1992`.

# Output
- `["outliers"]`: Array of indices of outliers
- `["critical.chi.squared"]`: Threshold value for determining being an outlier
- `["rth.robust.distance"]`: rth robust distance, where (r+1)th robust distance is the first one that exceeds the threshold.

# Examples
```julia-repl
julia> multidata = hcat(hbk.x1, hbk.x2, hbk.x3);

julia> hadi1994(multidata)
Dict{Any,Any} with 3 entries:
  "outliers"              => [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
  "critical.chi.squared" => 7.81473
  "rth.robust.distance"   => 5.04541
```

# Reference
Hadi, Ali S. "A modification of a method for the dedection of outliers in multivariate samples"
Journal of the Royal Statistical Society: Series B (Methodological) 56.2 (1994): 393-396.
"""
function hadi1994(multivariateData::Array{Float64,2}; alpha=0.05)
    n, p = size(multivariateData)
    h = Int(round((n + p + 1.0) / 2.0))
    cnp = (1 + ((2) / (n - 1 - 3p)) + ((p + 1) / (n - p)))^2
    chi_50_quantile = quantile(Chisq(p), 0.50)
    critical_quantile = quantile(Chisq(p), 1 - alpha)
    allindices = collect(1:n)

    meds = coordinatwisemedians(multivariateData)
    Sm = (1.0 / (n - 1.0)) * (multivariateData .- meds')' * (multivariateData .- meds')
    mah0 = diag(mahalanobisSquaredMatrix(multivariateData, meanvector=meds, covmatrix=Sm))
    ordering_indices_mah0 = sortperm(mah0)
    best_indices_mah0 = ordering_indices_mah0[1:h]
    starting_data = multivariateData[best_indices_mah0, :]

    Cv = coordinatwisemedians(starting_data)
    Sv = (1.0 / (h - 1.0)) * (starting_data .- Cv')' * (starting_data .- Cv')
    mah1 = diag(mahalanobisSquaredMatrix(multivariateData, meanvector=Cv, covmatrix=Sv))
    ordering_indices_mah1 = sortperm(mah1)

    r = p + 1
    basic_subset_indices = []
    basic_subset = []
    sorted_mah1 = []
    Cb = []
    Sb = []
    cfactor = 0
    isFullRank = false

    while r < n

        isFullRank = false

        while !isFullRank
            basic_subset_indices = ordering_indices_mah1[1:r]
            basic_subset = multivariateData[basic_subset_indices, :]
            Cb = applyColumns(mean, basic_subset)
            Sb = cov(basic_subset)
            cfactor = cnp * sqrt(sort(mah1)[h]) / chi_50_quantile
            r += 1

            if (det(cfactor * Sb) == 0 && r < n)
                isFullRank = false
            else
                isFullRank = true
            end
        end

        mah1 = diag(mahalanobisSquaredMatrix(multivariateData, meanvector=Cb, covmatrix=(cfactor * Sb)))

        ordering_indices_mah1 = sortperm(mah1)
        basic_subset_indices = ordering_indices_mah1[1:r]

        sorted_mah1 = sort(mah1)
        if sorted_mah1[r] >= critical_quantile
            break
        end
    end

    outlierset = setdiff(allindices, basic_subset_indices)

    result = Dict()
    result["outliers"] = sort(outlierset)
    result["critical.chi.squared"] = critical_quantile
    result["rth.robust.distance"] = sorted_mah1[r - 1]
    return result
end


#
# Extended Kalman Filter Types and Methods
#

"""
    ExtendedKalmanFilter(sys::NonlinearSystem, obs::NonlinearObserver,
                         estimate::AbstractState[, consider_states::Vector{UInt}])

Extended Kalman Filter type. The internal filter model takes on a specific form
depending on the type of the initial estimate. For a `DiscreteState`,

\$x_k = F(k, x_{k-1}) + w_k\$
\$y_k = H(k, x_k) + v_k\$

where \$w_k \\sim N(0, Q)\$ and \$v_k \\sim N(0, R)\$. For a `ContinuousState`,

\$\\dot{x}(t_k) = F(t_k, x(t_k)) + w(t_k)\$
\$y(t_k) = H(t_k, x(t_k)) + v(t_k)\$

Construction of an ExtendedKalmanFilter requires an initial estimate.
Internally, the initial estimate is an uncertain state type. The constructor
automatically converts absolute initial estimates to uncertain states with zero
covariance.

The `consider_states` input contains a list of indices of state elements to be
considered in the filtering process and not updated.
"""
immutable ExtendedKalmanFilter{T,S<:AbstractUncertainState{T}} <:
                                                       AbstractKalmanFilter{T,S}
    sys::NonlinearSystem{T}
    obs::NonlinearObserver{T}
    estimate::S
    consider_states::Vector{UInt16}

    function ExtendedKalmanFilter(sys, obs, estimate::S, consider_states::Vector
                                  ) where {T, S<:AbstractUncertainState{T}}
        if !allunique(consider_states)
            throw(ArgumentError("Consider state indices must be unique"))
        end
        for idx = 1:length(consider_states)
            if consider_states[idx] > length(estimate.x)
                throw(DimensionMismatch(
                    "Consider indices extend beyond length of initial state."))
            end
        end
        new{T,S}(NonlinearSystem(sys), NonlinearObserver(obs), estimate,
                 consider_states)
    end
end
ExtendedKalmanFilter(sys, obs, estimate) =
    ExtendedKalmanFilter(sys, obs, estimate, [])
ExtendedKalmanFilter(sys, obs, estimate::AbstractAbsoluteState, consider_states) =
    ExtendedKalmanFilter(sys, obs, make_uncertain(estimate), consider_states)

const DiscreteEKF{T} =
    ExtendedKalmanFilter{T,UncertainDiscreteState{T}} where T
const ContinuousEKF{T} =
    ExtendedKalmanFilter{T,UncertainContinuousState{T}} where T

"""
    kalman_predict(kf::KalmanFilter, t)
"""
function kalman_predict(ekf::ExtendedKalmanFilter, t)
    xk = predict(ekf.estimate, ekf.sys, t)
    yk = predict(xk, ekf.obs)
    Pxy = xk.P * ekf.obs.dH_dx(t, ekf.estimate.x)'
    return xk, yk, Pxy
end

"""
    extended_kalman_filter!(state, system, observer, measurement[, archive])

Extended Kalman Filter convenience function. Performs a single Kalman Filter
iteration, updating the provided state in-place. This function is not as
efficient as constructing creating a static `ExtendedKalmanFilter` and using the
`process!` method, but allows for increased flexibility.

Use `extended_kalman_filter` for not-in-place state updates.
"""
function extended_kalman_filter!(state, system, observer, measurement)
    ekf = ExtendedKalmanFilter(system, observer, state)
    process!(ekf, measurement)
end
function extended_kalman_filter!(state, system, observer, measurement, archive)
    ekf = ExtendedKalmanFilter(system, observer, state)
    process!(ekf, measurement, archive)
end

"""
extended_kalman_filter(state, system, observer, measurement[, archive])

Extended Kalman Filter convenience function. Performs a single Kalman Filter
iteration. This function is not as efficient as constructing creating a static
`ExtendedKalmanFilter` and using the `process!` method, but allows for increased
flexibility.

Use `extended_kalman_filter!` for in-place state updates.
"""
function extended_kalman_filter(state, system, observer, measurement)
    out_state = deepcopy(state)
    extended_kalman_filter!(out_state, system, observer, measurement)
    return out_state
end
function extended_kalman_filter(state, system, observer, measurement, archive)
    out_state = deepcopy(state)
    extended_kalman_filter!(out_state, system, observer, measurement, archive)
    return out_state
end


struct MatrixBingham{N,T,n,k} <: ParameterizedMeasure{N}
    par::NamedTuple{N,T}
    manifold::Stiefel{n,k,ℝ}
end
function MatrixBingham(n, k; kwargs...)
    par = NamedTuple(kwargs)
    return MatrixBingham(par, Stiefel(n, k))
end

Manifolds.base_manifold(d::MatrixBingham) = getfield(d, :manifold)

function MeasureTheory.basemeasure(μ::MatrixBingham{(:B,)})
    return normalize(Hausdorff(base_manifold(μ), μ.B))
end

function MeasureTheory.logdensity(d::MatrixBingham{(:B,)}, x::AbstractMatrix)
    B = d.B
    return dot(x, B, x) - log(pFq((k//2,), (n//2,), B))
end


using UncertaintyQuantification

x = RandomVariable.(Uniform(-π, π), [:x1, :x2, :x3])
a = Parameter(7, :a)
b = Parameter(0.05, :b)

inputs = [x; a; b]

ishigami = Model(
    df -> sin.(df.x1) .+ df.a .* sin.(df.x2) .^ 2 .+ df.b .* (df.x3 .^ 4) .* sin.(df.x1), :y
)

data = sample(inputs, SobolSampling(500))
evaluate!(ishigami, data)

phs = PolyharmonicSpline(data, 2, :y)

si, _ = sobolindices([ishigami], inputs, :y, MonteCarlo(10000))
si_phs, _ = sobolindices([phs], inputs, :y, SobolSampling(10000))

println("First order sobol indices of the ishigami function: $si")
println("First order sobol indices of the polyharmonic spline meta mode: $si_phs")


module Issue614

using Gridap
using Gridap.CellData

domain = (0,2,0,1)
partition = (2,2)
model = CartesianDiscreteModel(domain,partition)
const k = 2
reffeᵤ = ReferenceFE(lagrangian,VectorValue{2,Float64},k)
reffe_p = ReferenceFE(lagrangian,Float64,k-1)
V₀ = TestFESpace(model, reffeᵤ, conformity=:H1)
U = TrialFESpace(V₀)
Q = TestFESpace(model, reffe_p, conformity=:C0)
P = TrialFESpace(Q)
Y = MultiFieldFESpace([V₀,Q])
X = MultiFieldFESpace([U,P])

Ω = Triangulation(model)

degree = 1
xh = interpolate_everywhere([VectorValue(0.0,1.0),x->x[1]],X)

Γc = BoundaryTriangulation(model)
dΓc = Measure(Γc, degree)
nΓc = get_normal_vector(Γc)

for i in 0:1000
    uh, ph = xh
    α = ph*nΓc
    CD, CL = ∑( ∫( α )dΓc )
end

end # module


#Reference
#岑松, 龙志飞, 对转角场和剪应变场进行合理插值的厚板元[J]. 工程力学, 1998,15(3): 1―14.
function K_TMT(elm::Tria)::Matrix{Float64}
    E₀,ν₀=elm.material.E,elm.material.ν
    G₀=elm.material.G
    center=elm.center
    t=elm.t
    T=elm.T[1:3,1:3]
    x₁,y₁,z₁=T*(elm.node1.loc-center)
    x₂,y₂,z₂=T*(elm.node2.loc-center)
    x₃,y₃,z₃=T*(elm.node3.loc-center)
    K=Matrix{Float64}(undef,12,12)
    A=elm.A
    a₁,a₂,a₃=[xᵢ[j]*yᵢ[m]-xᵢ[m]*yᵢ[j] for (j,m) in zip([2,3,1],[3,1,2])]
    b₁,b₂,b₃=[yᵢ[j]-yᵢ[m] for (j,m) in zip([2,3,1],[3,1,2])]
    c₁,c₂,c₃=[-xᵢ[j]+xᵢ[m] for (j,m) in zip([2,3,1],[3,1,2])]
    l₁=norm(elm.node3.loc-elm.node2.loc)
    l₂=norm(elm.node1.loc-elm.node3.loc)
    l₃=norm(elm.node1.loc-elm.node2.loc)

    δ₁=(t/d₁)^2/(5/6*(1-ν₀)+2*(t/d₁)^2)
    δ₂=(t/d₂)^2/(5/6*(1-ν₀)+2*(t/d₂)^2)
    δ₃=(t/d₃)^2/(5/6*(1-ν₀)+2*(t/d₃)^2)

    F=-3/2/A*[c₁/l₁^2*(b₂*L₃+b₃*L₂) c₂/l₂^2*(b₃*L₁+b₁*L₃) c₃/l₃^2*(b₁*L₂+b₂*L₁);
             -b₁/l₁^2*(c₂*L₃+c₃*L₂) -b₂/l₂^2*(c₃*L₁+c₁*L₃) -b₃/l₃^2*(c₁*L₂+c₂*L₁);
             M₁ M₂ M₃]

    Δ=[1-2δ₁ 0 0;
       0 1-2δ₂ 0;
       0 0 1-2δ₃]

    G=[0 0 0 -2 -c₁ b₁ 2 -c₁ b₁;
       2 -c₂ b₂ 0 0 0 -2 -c₂ b₂;
       -2 -c₃ b₃ 2 -c₃ b₃ 0 0 0]

    Bb⁰=-1/2/A*[0 b₁ 0 0 b₂ 0 0 b₃ 0;
               0 0 c₁ 0 0 c₂ 0 0 c₃;
               0 c₁ b₁ 0 c₂ b₂ 0 c₃ b₃]

    FDG=F*Δ*G

    Bb=Bb⁰+FDG

    Δʼ=[δ₁ 0 0;
        0 δ₂ 0;
        0 0 δ₃]

    D₀=E₀*t^3/12/(1-ν₀^2)
    D=D₀*[1 ν₀ 0;
          ν₀ 1 0;
          0  0 (1-ν₀)/2]
    k=5/6

    function BtDB(x)
        x,y=x[1],x[2]
        L₁=1/2/A*(a₁+b₁*x+c₁*y)
        L₂=1/2/A*(a₂+b₂*x+c₂*y)
        L₃=1/2/A*(a₃+b₃*x+c₃*y)

        M₁=1/l₁^2*((c₁*c₂-b₁*b₂)*L₃+(c₃*c₁-b₃*b₁)*L₂)
        M₂=1/l₂^2*((c₂*c₃-b₂*b₃)*L₁+(c₁*c₂-b₁*b₂)*L₃)
        M₁=1/l₃^2*((c₃*c₁-b₃*b₁)*L₂+(c₂*c₃-b₂*b₃)*L₁)

        H=1/2/A*[b₃*L₂-b₂*L₃ b₁*L₃-b₃*L₁ b₂*L₁-b₁*L₂;
                 c₃*L₂-c₂*L₃ c₁*L₃-c₃*L₁ c₂*L₁-c₁*L₂]

        Bs=H*Δʼ*G

        C=k*G₀*t*[1 0;0 1]

        Kb=transpose(Bb)*D*Bb
        Ks=transpose(Bs)*C*Bs
        K=Kb+Ks
    end
    K=hcubature(BtDB,[-1,-1],[1,1])[1]
    # #left-hand system to right-hand system
    # I=1:9
    # J=[1,3,2,4,6,5,7,9,8,10,12,11]
    # L=sparse(I,J,1.,12,12)
    # for i in [1,3,4,6,7,9,10,12]
    #     for j in [2,5,8,11]
    #         K[i,j]=K[j,i]=-K[i,j]
    #     end
    # end
    # K=L'*sparse(K)*L
    #9x9 to 18x18
    I=1:9
    J=[3,4,5,9,10,11,15,16,17]
    L=sparse(I,J,1.,9,18)
    return L'*K*L
end


function rbf(γ::Array, X₁::Matrix, X₂::Matrix)::Matrix
      F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
      K = Array{eltype(γ)}(undef, N₁, N₂)
      for j in 1:N₂
            for i in 1:N₁
                  sqd = zero(eltype(γ))
                  for k in 1:F
                        sqd += γ[k]*(X₁[k,i] - X₂[k,j])^2
                  end
                  K[i,j] = exp(-0.5*sqd)
            end
      end
      return K
end

function ∇ₓrbf(γ::Array, X₁::Matrix, X₂::Matrix)::Array  # derivative on X1
    F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
    ∇ₓK = Array{eltype(γ)}(undef, F, N₁, N₂)
    for j in 1:N₂
          for i in 1:N₁
                sqd = zero(eltype(γ))
                for m in 1:F
                      sqd += γ[m]*(X₁[m,i]-X₂[m,j])^2
                end
                for l in 1:F
                      ∇ₓK[l,i,j] = -γ[l]*(X₁[l,i]-X₂[l,j])*exp(-0.5*sqd)
                end
          end
    end
    return ∇ₓK
end

function linear(γ::Array, X₁::Matrix, X₂::Matrix)::Array
      F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
      K = Array{eltype(γ)}(undef, N₁, N₂)
      for j in 1:N₂
            for i in 1:N₁
                  dot_prod = zero(eltype(γ))
                  for m in 1:F
                        dot_prod += γ[m]*X₁[m,i]*X₂[m,j]
                  end
                  K[i,j] = dot_prod + γ[F+1]
            end
      end
      return K
end

function ∇ₓlinear(γ::Array, X₁::Matrix, X₂::Matrix)::Array
      F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
      ∇ₓK = Array{eltype(γ)}(undef, F, N₁, N₂)
      for j in 1:N₂
            for i in 1:N₁
                  for l in 1:F
                        ∇ₓK[l,i,j] = γ[l]*X₂[l,j]
                  end
            end
      end
      return ∇ₓK
end

function mixed_effect_kernel(θ::Array, X₁::Matrix, X₂::Matrix)::Matrix
      F, N₁, N₂ = size(X₁, 1), size(X₁, 2), size(X₂, 2)
      K = Array{eltype(θ)}(undef, N₁, N₂)
      M = ones(F, F)
      D = diagm(ones(F))
      for j in 1:N₂
            for i in 1:N₁
                  sqd = zero(eltype(θ))
                  for k in 1:F
                        sqd += θ[k]*(X₁[k,i] - X₂[k,j])^2
                  end
                  K[i,j] = exp(-0.5*sqd)
            end
      end
      𝐾 = kron(K, θ[F+1]*M+(1-θ[F+1])*D)
      return 𝐾
end


@testset "Multivectors" begin
    @test is_homogeneous(mv_1)
    @test grades(mv_1) == [1]
    @test grades(mv_2) == [1,2,3]
    @test all(nonzero_blades(mv_1) .== Any[1v1, 1v2])
    @test all(nonzero_blades(mv_2) .== Any[1v1, 1v12, 1v123])
    @test mv_1[2] == 1
    @test mv_1[3] == 1
    @test mv_2[CartesianIndex(1,2)] == 1
    @test mv_2[CartesianIndex(1,2,3)] == 1
    @test sum(nonzero_blades(mv_1)) == sum(blades(mv_1)) == mv_1
    @test sum(nonzero_blades(mv_2)) == sum(blades(mv_2)) == mv_2

    @test 1v2 + 1v12 == Multivector([0, 0, 1, 0, 1, 0, 0, 0])
    @test 1v1 + mv_1 == Multivector([0, 2, 1, 0, 0, 0, 0, 0])
    @test v1 + v2 == mv_1
end


### A Pluto.jl notebook ###
# v0.19.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el)
        el
    end
end

# ╔═╡ 5323b5ac-0c51-4952-8f27-4f489c1f2d1a
begin
    import Pkg
    Pkg.activate() # disable Pluto's package management
end

# ╔═╡ d8aeb9f4-990c-11eb-0dbb-078cda4928d5
begin
    using GSL # provides the special functions
    using GR
    using PlutoUI
end

# ╔═╡ 6127fb6a-22fc-4f35-8228-cc06ffebf9c0
begin
    a0 = 5.291772109217e-11

    # The unitless radial coordinate
    ρ(r, n) = 2r/(n*a0)
end

# ╔═╡ 7aa443bd-b706-41e2-bdc8-0d8c10c99b02
# The θ and ϕ dependence
function Yml(m::Int, l::Int, θ::Real, ϕ::Real)
    (-1.0)^m*sf_legendre_Plm(l, abs(m), cos(θ))*ℯ^(im*m*ϕ)
end

# ╔═╡ 778c4fa2-68c0-494f-ae22-477189a27f2b
# The Radial dependence
function R(n::Int, l::Int, ρ::Real)
    sf_laguerre_n(n-l-1, 2*l+1, ρ)*ℯ^(-ρ/2)*ρ^l
end

# ╔═╡ cf4d7a50-afea-4f10-b08d-0f0150b7da7a
# A normalization: This is dependent on the choice of polynomial representation
function norm(n::Int, l::Int)
    sqrt((2/n)^3 * factorial(n-l-1)/(2n*factorial(n+l)))
end

# ╔═╡ d3ff17f0-3dd6-4ffe-a045-5b573aec02a5
# Generates an Orbital Funtion of (r, θ, ϕ) for a specificied n, l, m.
function Orbital(n::Int, l::Int, m::Int)
    # we make sure l and m are within proper bounds
    if l > n || abs(m) > l
        throw(DomainError())
    end
    Ψ(ρ, θ, ϕ) = norm(n, l) * R(n, l, ρ) * abs(Yml(m, l, θ, ϕ))
    Ψ
end

# ╔═╡ 94158386-447f-40cd-a321-8b333053d558
function CarttoSph(x::Array, y::Array, z::Array)
    r = sqrt.(x.^2+y.^2+z.^2)
    θ = acos.(z./r)
    ϕ = atan.(y./x)
    r, θ, ϕ
end

# ╔═╡ ad5ffb61-b2cc-499a-87a4-ad71833de766
function calculate_electronic_density(n, l, m)
    N = 50
    r = 1e-9
    s = LinRange(-r, r, N)
    x, y, z = meshgrid(s, s, s)

    r, θ, ϕ = CarttoSph(x, y, z)
    Ψ = Orbital(n, l, m)

    Ψv = zeros(Float32, N, N, N)
    for i in 1:N
        for j in 1:N
            for k in 1:N
                Ψv[i, j, k] = abs(Ψ(ρ(r[i, j, k], n), θ[i, j, k], ϕ[i, j, k]))
            end
        end
    end

    (Ψv.-minimum(Ψv))./(maximum(Ψv).-minimum(Ψv))
end

# ╔═╡ a4c51114-4bf9-4d40-b123-1629a3be54e1
@bind Ψ Select(["2,0,0", "3,0,0", "2,1,0", "3,1,0", "3,1,1", "2,1,1", "3,2,0", "3,2,1", "3,2,2", "4,0,0", "4, 1,0", "4,2,0", "4,2,1", "4,2,2", "4,3,0", "4,3,1", "4,3,2", "4,3,3"], default="3,2,0")

# ╔═╡ 11cfffdf-91d6-4909-bb59-e8d12a95e5e8
@bind ϕ Slider(0:360, default=30)

# ╔═╡ 9899f86c-2fba-4d72-abab-acbef1dc6988
@bind iso Slider(0.1:0.05:0.8, default=0.2)

# ╔═╡ 16697f21-f879-48d1-ad41-f39e1dde5830
begin
    n, m, l = parse.(Int, split(Ψ, ','))
    Ψv = calculate_electronic_density(n, m, l)
    isosurface(Ψv, isovalue=iso, rotation=ϕ)
end

# ╔═╡ Cell order:
# ╠═5323b5ac-0c51-4952-8f27-4f489c1f2d1a
# ╠═d8aeb9f4-990c-11eb-0dbb-078cda4928d5
# ╠═6127fb6a-22fc-4f35-8228-cc06ffebf9c0
# ╠═7aa443bd-b706-41e2-bdc8-0d8c10c99b02
# ╠═778c4fa2-68c0-494f-ae22-477189a27f2b
# ╠═cf4d7a50-afea-4f10-b08d-0f0150b7da7a
# ╠═d3ff17f0-3dd6-4ffe-a045-5b573aec02a5
# ╠═94158386-447f-40cd-a321-8b333053d558
# ╠═ad5ffb61-b2cc-499a-87a4-ad71833de766
# ╠═a4c51114-4bf9-4d40-b123-1629a3be54e1
# ╠═11cfffdf-91d6-4909-bb59-e8d12a95e5e8
# ╠═9899f86c-2fba-4d72-abab-acbef1dc6988
# ╠═16697f21-f879-48d1-ad41-f39e1dde5830


function build_gram_matrix(q::Vector,
                           monos::AbstractVector{<:MP.AbstractMonomial})
    return GramMatrix(MultivariateMoments.SymMatrix(q, length(monos)),
                         monos)
end

function build_moment_matrix(q::Vector,
                             monos::AbstractVector{<:MP.AbstractMonomial})
    return MomentMatrix(MultivariateMoments.SymMatrix(q, length(monos)),
                        monos)
end

abstract type SOSLikeCone <: PolyJuMP.PolynomialSet end
Base.broadcastable(cone::SOSLikeCone) = Ref(cone)

struct SOSPolynomialSet{DT <: AbstractSemialgebraicSet,
                        CT <: SOSLikeCone,
                        BT <: PolyJuMP.AbstractPolynomialBasis,
                        MT <: MP.AbstractMonomial,
                        MVT <: AbstractVector{MT},
                        NPT <: Tuple} <: MOI.AbstractVectorSet
    domain::DT
    cone::CT
    basis::Type{BT}
    monomials::MVT
    newton_polytope::NPT
    mindegree::Int
    maxdegree::Int
end


using JuMP
using CPLEX

include("parser.jl")

start = time()

function Mymodel(MyFileName::String)
  n, P, alpha, Survie, indice_danger, cout = read_instance(MyFileName)

  # Create the model
  m = Model(CPLEX.Optimizer)

  ## Variables
  @variable(m, x[1:n,1:n], Bin)
  @variable(m, y[1:n,1:n], Bin)
  @variable(m, z[1:P,1:n,1:n] <= 0)

  ## Constraints
  @constraint(m, [k in 1:P, i in 1:n, j in 1:n ; indice_danger[k] == 0], z[k,i,j] >= x[i,j]*log(1-Survie[k,i,j]))
  @constraint(m, [k in 1:P, i in 1:n, j in 1:n ; indice_danger[k] == 1], z[k,i,j] >= y[i,j]*log(1-Survie[k,i,j]))
  @constraint(m, [k in 1:P], sum(z[k,i,j] for i in 1:n, j in 1:n) <= log(1-alpha[k]))
  @constraint(m, [i in 1:n, j in 1:n ; i == 1 || j == 1 || i == n || j == n], y[i] <= 0)    # bordure => y = 0
  @constraint(m, [i in 2:n-1, j in 2:n-1], y[i,j] <= x[i,j])
  @constraint(m, [i in 2:n-1, j in 2:n-1], y[i,j] <= x[i-1,j])
  @constraint(m, [i in 2:n-1, j in 2:n-1], y[i,j] <= x[i+1,j])
  @constraint(m, [i in 2:n-1, j in 2:n-1], y[i,j] <= x[i,j-1])
  @constraint(m, [i in 2:n-1, j in 2:n-1], y[i,j] <= x[i,j+1])
  @constraint(m, [i in 2:n-1, j in 2:n-1], y[i,j] <= x[i-1,j-1])
  @constraint(m, [i in 2:n-1, j in 2:n-1], y[i,j] <= x[i-1,j+1])
  @constraint(m, [i in 2:n-1, j in 2:n-1], y[i,j] <= x[i+1,j-1])
  @constraint(m, [i in 2:n-1, j in 2:n-1], y[i,j] <= x[i+1,j+1])

  ## Objective
  @objective(m, Min, sum(x[i,j]*cout[i,j] for i in 1:n, j in 1:n))

  #resolution
  optimize!(m)

  # solution_summary(m, verbose=true)
  z_vals = JuMP.getvalue.( m[:z] )
  # println(z_vals)

  for k in 1:P
    println(-(exp(sum(z_vals[k,i,j] for i in 1:n, j in 1:n))-1))
  end

  # println(-(exp(sum(z[k,i,j] for i in 1:n, j in 1:n))-1))

end

Mymodel("outputn10P6.txt")


module MDP

    export Model,
           ProbModel,
           DynamicModel,
           valueIteration,
           finiteHorizon

    using Printf

    abstract type Model end

    struct ProbModel <: Model
        bellmanUpdate! :: Function
        objective  :: Function
        contractionFactor :: Float64
        stateSize  :: Int
        actionSize :: Int

        P :: Array{ Matrix{Float64}, 1}
        cost   :: Matrix{Float64}
        reward :: Matrix{Float64}

        function ProbModel(c, P; objective=:Max)
            (n, m) = size(c)

            if length(P) != m
                error("Number of transition matrices does not match number of actions.")
            end

            P_concatenated = vcat(P...)
            if size(P_concatenated) != (n*m, n)
                error("Size of transition and reward matrices are inconsistent")
            end

            is_square(Pi)          = size(Pi) == (n,n)
            is_row_stochastic(Pi)  = isapprox(sum(Pi; dims=2), ones(n); atol = 100*eps(Float64))
            is_stopping_action(Pi) = Pi == zero(Pi)
            for Pi in P
                if !is_square(Pi)
                    error("Transition matrix is not a square matrix")
                elseif !(is_row_stochastic(Pi) || is_stopping_action(Pi))
                    error("Transition matrix is not row stochastic.")
                end
            end

            if objective != :Max && objective != :Min
                error("Model objective must be :Max or :Min")
            else
                obj, compare = (objective == :Max) ? (maximum, >) : (minimum, <)

                function bellman!(vUpdated, gUpdated, v; discount=1)
                    Q = c + discount * reshape(P_concatenated * v, n, m);

                    # vUpdated = m.objective(Q, (), 2)
                    withIndex!(vUpdated, gUpdated, compare, Q)
                end

                # See Puterman Thm 6.6.6
                contractionFactor = 1 - sum(minimum(P_concatenated; dims=1))

                new(bellman!, obj, contractionFactor, n, m, P, c, c)
            end
        end
    end

    struct DynamicModel <: Model
        bellmanUpdate! :: Function # (valueFunction; discount=1.0) -> (valueFunction, policy)
        objective  :: Function
        contractionFactor :: Float64

        function DynamicModel(bellmanUpdate!; contractionFactor=1, objective=:Max)
            if objective != :Max && objective != :Min
                error("Model objective must be :Max or :Min")
            else
                obj = (objective == :Max) ? maximum : minimum
                new(bellmanUpdate!, obj, contractionFactor)
            end
        end

    end

    function valueIteration(m::Model, initial_v;
                    discount   = 0.95,
                    iterations = 1_000,
                    tolerance  = 1e-4)

        scaledDiscount  = (discount < 1) ? (1-discount)/discount : 1
        scaledTolerance = scaledDiscount * tolerance / 2

        update!(vUpdated, gUpdated, v) = m.bellmanUpdate!(vUpdated, gUpdated, v; discount=discount)

        v_previous = copy(initial_v)
        v          = copy(initial_v)
        g          = zeros(Int, size(v))

        update!(v, g, initial_v)
        copy!(v_previous, v)

        v_precision = spanNorm(v, initial_v)

        if abs( 1 - m.contractionFactor * discount ) < 4*eps(Float64)
            warn("Contraction factor too small to guarantee convergece. Value iteration may not converge.")
        else
            # See Puterman Prop 6.6.5
            # We compare with zero to allow overflow errors when v_precision is 0.
            if abs(v_precision) < 4*eps(Float64)
                iteration_bound = 1
            else
                iteration_bound = floor(Int, log( scaledTolerance/v_precision ) / log( m.contractionFactor*discount ))
            end

            @info "value iteration will converge in at most $iteration_bound iterations"
            if (iterations <= iteration_bound)
                @warn "Value iteration may not converge. Iterations $iterations less than estimated bound $iteration_bound"
            end
        end

        iterationCount  = 1;
        while (v_precision > scaledTolerance && iterationCount < iterations)
            iterationCount += 1

            copy!(v_previous, v)
            update!(v, g, v_previous)

            v_precision = spanNorm(v, v_previous)
        end

        if (v_precision > scaledTolerance)
            warn(@sprintf("Value iteration did not converge.
                 Reached precision %e at iteration %d", 2*v_precision/scaledDiscount, iterationCount))
        else
            @info @sprintf("Reached precision %e at iteration %d", 2*v_precision/scaledDiscount, iterationCount)
        end

        # Renormalize v -- See Puterman 6.6.12 for details
        v .+= m.objective(v - v_previous)/scaledDiscount

        return (v, g)
    end

    # For the probability model, it is possible to automatically initialize
    # the initial vector
    function valueIteration(m::ProbModel;
                    initial_v  = zeros(m.stateSize),
                    discount   = 0.95,
                    iterations = 1_000,
                    tolerance  = 1e-4)

        return valueIteration(m, initial_v;
                    discount   = discount,
                    iterations = iterations,
                    tolerance  = tolerance)
    end

    function finiteHorizon(m::Model, final_v;
                           discount = 1.0,
                           horizon :: Int = 10)

        update!(vUpdated, gUpdated, v) = m.bellmanUpdate!(vUpdated, gUpdated, v; discount=discount)

        v = [ zero(final_v) for stage = 1 : horizon ]
        g = [ zeros(Int,   size(final_v)) for stage = 1 : horizon ]

        update!(v[horizon], g[horizon], final_v)
        for stage in horizon-1: -1 : 1
          update!(v[stage], g[stage], v[stage+1])
        end
        return (v,g)
    end

    function spanNorm(x, y)
        # z = x - y
        # return max(z) - min(z)
        # Optimized code. This could have been done using @devec, but
        # I don't want to add a dependency just for one function.
        z = vec(x - y)
        max_z, min_z = -Inf, Inf
        for elem in z
            if max_z < elem
                max_z = elem
            end
            if min_z > elem
                min_z = elem
            end
        end
        return max_z - min_z
    end

    # Julia does not have an in-built function that returns the minimum and the
    # arg min.
    # function withIndex{T}(f::Function, x::Matrix{T})
    #     idx = vec(mapslices(f,x,2))
    #     val :: Vector{T} = [x[i,idx[i]] for i = 1:size(x,1) ]

    #     return val,idx
    # end

    # A more direct implementation
    function withIndex!(val::AbstractArray{T,1}, idx::AbstractArray{Int,1}, compare::Function, x::AbstractArray{T,2}) where T
        (n, m) = size(x)

        for i=1:n
            idx[i], val[i] = 1, x[i,1]
            for j=2:m
                if compare(x[i,j], val[i])
                    idx[i], val[i] = j, x[i,j]
                end
            end
        end
    end
end


@symbol_func function B_rho(cur_reactor::AbstractReactor, cur_rho::AbstractSymbol, cur_theta::AbstractSymbol)
  cur_angle = cur_theta * ( cur_reactor.pi / 180.0 )

  cur_B = cur_reactor.B_0

  cur_B /= 1 + cur_reactor.epsilon * cur_rho * cos(cur_angle)

  cur_B
end


"""
    BernoulliCodeDistribution

Distribution where bits are each independent Bernoulli trials (coin flips)
"""
mutable struct BernoulliCodeDistribution <: AbstractBinaryVectorDistribution
    p::Vector{Float64}
    # I::Vector{Int} # moving this to metadata
    metadata::Dict{Any,Any}
    cache::Dict{Any,Any}

    function BernoulliCodeDistribution(P::Vector{Float64}; kwargs...)
        new(P, Dict(kwargs), Dict(:pdf=>spzeros(2^length(P))))
    end
end

function show(io::IO, P::BernoulliCodeDistribution)
    println(io,"Bernoulli Code Distribution")
    println(io,"N_neurons: $(n_bits(P))")
    println(io,"p_i's:     $(P.p)")
    show_metadata(io, P)
end

==(P1::BernoulliCodeDistribution, P2::BernoulliCodeDistribution) = P1.p == P2.p

n_bits(B::BernoulliCodeDistribution) = length(B.p)

function pdf(Pr::BernoulliCodeDistribution, x)
    if length(x) != n_bits(Pr)
        error("BernoulliCodeDistribution pdf: out of domain error")
    else
        # minor inefficiency: this won't differentiate between a 0.0 value that
        # is stored vs. not stored in the sparse vector. But, this would only
        # come up if one of the p_i = 0.
        # also, might be worth splitting this back into two cases (Vector{Bool}
        # vs. BitVector)
        # idx = 1 + dot([2^i for i = 0:(n_bits(Pr) - 1)], x)
        idx = 1 + _binary_to_int(x)
        if Pr.cache[:pdf][idx] == 0.0
            Pr.cache[:pdf][idx] = prod([x[i] ? Pr.p[i] : (1-Pr.p[i]) for i = 1:length(Pr.p)])
        end
        return Pr.cache[:pdf][idx]
        # if Pr.cache[:pdf][idx] > 0.0
        #     return Pr.cache[:pdf][idx]
        # else
        #     Pr.cache[:pdf][idx] = prod([x[i] ? Pr.p[i] : (1-Pr.p[i]) for i = 1:length(Pr.p)])
        #     return Pr.cache[:pdf][idx]
        # end
    end
    # length(x) == length(Pr.p) ? prod([x[i] ? Pr.p[i] : (1-Pr.p[i]) for i = 1:length(Pr.p)]) : error("BernoulliCodeDistribution: out of domain error")
end

function expectation_matrix(B::BernoulliCodeDistribution)
    em = B.p * B.p' # bits are independent coin flips
    em[1:(n_bits(B)+1):end] = B.p # set the diagonal to the probabilities instead of squares
    return em
end

# I'm going to be lazy here and allow this value to be computed each time. The
# parentheses around the second assignment expression are just for human
# readability.
entropy(P::BernoulliCodeDistribution) = (P.metadata[:entropy] = -sum([(p_i * log(p_i) + (1-p_i) * log(1-p_i)) for p_i in filter(x -> 0 < x < 1, P.p)]))
entropy2(P::BernoulliCodeDistribution) = (P.metadata[:entropy2] = -sum([(p_i * log2(p_i) + (1-p_i) * log2(1-p_i)) for p_i in filter(x -> 0 < x < 1, P.p)]))

random(B::BernoulliCodeDistribution, n_samples=1) = rand(n_bits(B), n_samples) .< B.p


# Random access lists (Okasaki, 1995)
# (une structure utilisant de la récursion polymorphe)

struct Pair fst; snd end
struct Zero    s0    end
struct One  e; s1    end

function length(l::Nothing) 0                    end
function length(l::Zero   ) 2 * length(l.s0)     end
function length(l::One    ) 2 * length(l.s1) + 1 end

function get(l::Zero, i)
    x = get(l.s0, div(i, 2))
    if i % 2 == 0 x.fst else x.snd end
end
function get(l::One, i)
    if i == 0 return l.e end
    x = get(l.s1, div(i - 1, 2))
    if i % 2 == 1 x.fst else x.snd end
end

function cons(x, l::Nothing) One(x, nothing)                end
function cons(x, l::Zero   ) One(x, l.s0)                   end
function cons(x, l::One    ) Zero(cons(Pair(x, l.e), l.s1)) end

function sequence(i, j)
    if j < i
        nothing
    else
        cons(i, sequence(i + 1, j))
    end
end

s = sequence(34, 55)
for i = 0: 21
    println(get(s, i))
end


using Revise
module BayesMNIST

using Statistics, LinearAlgebra, StatsBase, Flux, Flux.Data.MNIST, DistributionsAD
using Flux: softmax, onehotbatch#, params
using DynamicHMC, LogDensityProblems, TransformVariables, DynamicHMC.Diagnostics, Parameters
using CuArrays, Random, Distributions, Zygote, BenchmarkTools, InteractiveUtils#, Plots
using CUDAapi
if has_cuda()
    @info "CUDA is on"
    import CuArrays
    CuArrays.allowscalar(false)
end
import ForwardDiff

using Revise

#load the data
#the purpose is to explore flux in a Bayesian context

Random.seed!(11)

const D=784
const K=10

#Lasso Logistic Categorical regression
struct LCRegression{TY, TX, Tατ, Tθτ}
  Y::TY
  X::TX

  ατβ::Tατ
  θτβ::Tθτ

end

function (problem::LCRegression)(θ)
    #unpack parameters and data
    @unpack Y, X, ατβ, θτβ = problem
    @unpack βs, τβs = θ
    Ncovariates, Npoints = size(X,1),size(X,2)
    Nclasses = size(βs,2)

    @assert size(βs,1) == Ncovariates
    @assert size(βs,2) == length(τβs) == Nclasses

    #p(τβᵢ)
    logτβpri = sum([logpdf(Gamma(ατβ, θτβ), τβs[i]) for i ∈ 1:Nclasses])

    #p(βᵢ|τβᵢ)
    logβpri = sum([logpdf(MvNormal(zeros(Ncovariates), 1/τβs[i]), βs[:,i]) for i ∈ 1:Nclasses])

    μs = softmax(X' * βs, dims=2)
    #μs = (xᵢ->softmax(βs'*xᵢ)).(X)
    #  display(μs[1])
    #p(Y|β,τβ)
    #loglik = sum([logpdf(Categorical(μs[n, :]), Y[n]) for n ∈ 1:Npoints])
    #println(size(μs[1]))
    #println(size(Y[1]))
    #display(logpdf(Multinomial(1, μs[1]), Y[1]))
    loglik = sum([logpdf(Multinomial(1,μs[n,:]), Y[:,n]) for n ∈ 1:Npoints])

    return logτβpri + logβpri + loglik
end

#hyperpriors
const Πmwe = Dict{Symbol, Float64}(
  :ατβ=>1.0,
  :θτβ=>1.0
)

function LCtransform(p::LCRegression)
    return as((βs=as(Array, D+1, K),
      τβs=as(Array, asℝ₊, K)))
end

function dynamicmnistmwe(Npoints::Int = 1000,
    ::Type{V} = AbstractVector, ::Type{T}=Float32) where {V,T}
  imgs = MNIST.images()
  labels = MNIST.labels()

  #make each image into a feature vector
  #concatenate each feature vector into a matrix
  preX =hcat(float.(reshape.(imgs,:))...)

  #add an intercept
  preX = vcat(preX, ones(1,size(preX,2)))[:,1:Npoints]
  #X = (Vector).(eachcol(preX)) #|>gpu
  X = preX
  preY = onehotbatch(labels[1:Npoints], 0:9)
  Y = Matrix(preY)

  #Y=labels[1:5_000]

  #make the problem
  p = LCRegression(Y, X, Πmwe[:ατβ], Πmwe[:θτβ])

  #transform the support
  t = LCtransform(p)
  P = TransformedLogDensity(t, p)

  #differentiate via AD
  ∇P =  ADgradient(:Zygote, P)

  @info "beginning MCMC"

  testβs = rand(D+1,K)# |>gpu
  testτβs = rand(K)

  #InteractiveUtils.@code_warntype p((βs=testβs, τβs=testτβs))
  #@btime $p((βs=$testβs, τβs=$testτβs))
  results = mcmc_with_warmup(Random.GLOBAL_RNG, ∇P, 1000, reporter = LogProgressReport())
  serialize("results.jls", results)

  posterior = transform.(t, results.chain)

  β_posterior = first.(transform.(t, results.chain))
  display(mean(β_posterior))

end

#dynamicmnistmwe()

end # module


include("workspace.jl")

#make circuit
qpe = Circuit(4)
X!(qpe,4)
H!(qpe,1:3)
λ = π/4

reps = 1
for counting_qubit ∈ 1:3
    for i ∈ 1:reps
        CU!(qpe, λ, counting_qubit, 4)
    end
    reps *= 2
end

function qft_dagger!(C::Circuit, n::Int)
    #Don't forget the swaps!
    for qubit in 1:div(n,2)-1
        Swap!(C, qubit, n-qubit)
    end
    for j in 1:n-1
        for m in 1:j-1
            CU!(C,-π/(2^(j-m)), m, j)
        end
        H!(C,j)
    end
end

qft_dagger!(qpe,4)
data = shots(qpe,10_000)
graph!(data)

##qpe2
qpe2 = Circuit(4)

X!(qpe2,4)
H!(qpe2,1:3)

λ2 = 2*π/3
reps = 1
for counting_qubit ∈ 1:3
    for i ∈ 1:reps
        CU!(qpe2, λ2, counting_qubit, 4)
    end
    reps *= 2
end

qft_dagger!(qpe2,4)
data2 = shots(qpe2,10_000)
graph!(data2)

##qpe3
qpe3 = Circuit(6)

H!(qpe3,1:5)
X!(qpe3,6)

reps = 1
for counting_qubit ∈ 1:5
    for i ∈ 1:reps
        CU!(qpe3, λ2, counting_qubit, 6)
    end
    reps *= 2
end

qft_dagger!(qpe3, 6)
data3 = shots(qpe3,10_000)
graph!(data3)


## Makes the model associated with a fully adapted auxiliary particle filter for
## a linear Gaussian state space model. This was introduced in:
## M. K. Pitt and N. Shephard. Filtering via simulation: Auxiliary particle
## filters. J. Amer. Statist. Assoc., 94(446):590--599, 1999.

function makeLGAPFModel(theta::LGTheta, ys::Vector{Float64})
  n::Int64 = length(ys)

  RC2v0::Float64 = theta.R + theta.C * theta.C * theta.v0
  v0RoverRC2v0::Float64 = theta.v0*theta.R/RC2v0
  sqrtv0RoverRC2v0::Float64 = sqrt(v0RoverRC2v0)

  mu1::Float64 = v0RoverRC2v0*(theta.x0/theta.v0 + theta.C*ys[1]/theta.R)
  tmp1::Float64 = theta.C * theta.x0 - ys[1]
  lG1::Float64 = -0.5 * log(2 * π * RC2v0) - tmp1 * 0.5/RC2v0 * tmp1

  RC2Q::Float64 = theta.R + theta.C * theta.C * theta.Q
  invRC2Qover2::Float64 = 0.5/RC2Q
  QRoverRC2Q::Float64 = theta.Q * theta.R / RC2Q
  sqrtQRoverRC2Q::Float64 = sqrt(QRoverRC2Q)

  invRover2::Float64 = 0.5/theta.R
  logncG::Float64 = -0.5 * log(2 * π * RC2Q)

  @inline function lG(p::Int64, particle::Float64Particle, ::Nothing)
    if p == n
      return p == 1 ? lG1 : 0.0
    end
    v::Float64 = theta.C * theta.A * particle.x - ys[p+1]
    v = logncG - v * invRC2Qover2 * v
    p == 1 && (v += lG1)
    return v
  end
  @inline function M!(newParticle::Float64Particle, rng::RNG, p::Int64,
    particle::Float64Particle, ::Nothing)
    if p == 1
      newParticle.x = mu1 + sqrtv0RoverRC2v0*randn(rng)
    else
      mu::Float64 = QRoverRC2Q*(theta.A*particle.x/theta.Q + theta.C*ys[p]/theta.R)
      newParticle.x = mu + sqrtQRoverRC2Q*randn(rng)
    end
  end
  return SMCModel(M!, lG, n, Float64Particle, Nothing)
end


using .LazySets

function (opl::LinearOettliPrager)(A, b)
    n = length(b)
    Ac = mid.(A)
    bc = mid.(b)
    Ar = IntervalArithmetic.radius.(A)
    br = IntervalArithmetic.radius.(b)

    polytopes = Vector{HPolytope}(undef, 2^n)
    orthants = DiagDirections(n)

    @inbounds for (i, d) in enumerate(orthants)
        D = Diagonal(d)
        Ard = -Ar*D

        A1 = [Ac + Ard; -Ac + Ard; -D]
        b1 = [br + bc; br - bc; zeros(n)]

        polytopes[i] =  HPolytope(A1, b1)
    end
    return identity.(filter!(!isempty, polytopes))
end

_default_precondition(_, ::LinearOettliPrager) = NoPrecondition()


# Load Julia packages (libraries) needed.

using Pkg, DrWatson

#@quickactivate "StstisticalRethinkingStan"
using StanSample
using StatisticalRethinking

using NamedTupleTools

function quap(sm::SampleModel)
  s = read_samples(sm, :dataframe)
  ntnames = (:coef, :vcov, :converged, :distr, :params)
  n = Symbol.(names(s))
  coefnames = tuple(n...,)
  p = quap(s)
  c = [mean(p[k]) for k in n]
  cvals = reshape(c, 1, length(n))
  coefvalues = tuple(cvals...,)
  v = Statistics.covm(Array(s), cvals)

  ntvalues = tuple(
    namedtuple(coefnames, coefvalues),
    v, true, MvNormal(c, v), n
  )

  namedtuple(ntnames, ntvalues)
end

q = quap(m5_5s)
r = rand(q.distr, 4000)'

mean(r, dims=1) |> display


module vanderpol
using ReachabilityAnalysis
@taylorize function vanderpol!(dx, x, params, t)
    local μ = 1.0
    dx[1] = x[2]
    dx[2] = (μ * x[2]) * (1 - x[1]^2) - x[1]
    return dx
end

function model(X0)
    S = @system(x' = vanderpol!(x), dim:2)
    return IVP(S, X0)
end
end # module


"""
    f∂f(f,x)

returns f and ∂f/∂x evaluated in `x`, using `ForwardDiff.jl`, `DiffResults.jl` and `StaticArrays.jl` to calculate everything in one pass.
"""
function f∂f(f,x::T) where T
    _f(z) = f(only(z))
    x_vec =   SVector(x)
    ∂result = DiffResults.GradientResult(x_vec)
    _∂f =  ForwardDiff.gradient!(∂result, _f,x_vec)
    fx =  DiffResults.value(_∂f)
    ∂f∂x = only(DiffResults.gradient(_∂f))
    return fx,∂f∂x
end

"""
    f∂f∂2f(f,x)

returns f,∂f/∂x,and ∂²f/∂²x and evaluated in `x`, using `ForwardDiff.jl`, `DiffResults.jl` and `StaticArrays.jl` to calculate everything in one pass.
"""
function f∂f∂2f(f,x::T) where T
    _f(z) = f(only(z))
    x_vec =   SVector(x)
    ∂result = DiffResults.HessianResult(x_vec)
    _∂f =  ForwardDiff.hessian!(∂result, _f,x_vec)
    fx =  DiffResults.value(_∂f)
    ∂f∂x = only(DiffResults.gradient(_∂f))
    ∂²f∂²x =  only(DiffResults.hessian(_∂f))
    return fx,∂f∂x,∂²f∂²x
end

"""
    ad_newton(f,x0;kwargs...)

    Performs newton root finding (via Roots.jl), but the derivatives are calculated with ForwardDiff.jl
    kwargs are passed to `Roots.find_zero`
"""
function ad_newton(f,x0;kwargs...)
    function f0(x)
        fx,∂f∂x = f∂f(f,x)
        return fx,fx/∂f∂x
    end
    return Roots.find_zero(f0,x0,Roots.Newton();kwargs...)
end


using SpectralClustering
using Distances
using Test
using LinearAlgebra
using Statistics
using Clustering
using Images
using Random
using Distributions

import LightGraphs.LinAlg: adjacency_matrix
number_of_vertices = 5
Random.seed!(0)

function two_gaussians(N::Integer = 500; std_1=5, std_2 = 5, center_1=[15,5], center_2=[-15, 5])
    d1 = (randn(2, N) * std_1) .+ center_1
    d2 = (randn(2, N) * std_2) .+ center_2
    labels = round.(Integer, vcat(zeros(N), ones(N)))
    return (hcat(d1, d2), labels)
end

function three_gaussians(N::Integer = 250; )
    d1 = (randn(2, N) * 1.5) .+ [5, 0]
    d2 = (randn(2, N) * 1) .+ [0, 0]
    d3 = (randn(2, N) * 1.5) .+ [-5, 0]
    labels = round.(Integer, vcat(zeros(N), ones(N), ones(N)* 2))
    return (hcat(d1, d2, d3), labels)
end

@testset "Graph Creation" begin
    @testset "KNNNeighborhood" begin
        function weight(i::Integer, neigh, v, m)
            return Distances.colwise(SqEuclidean(), m, v)
        end
        data = convert.(Float64, [2 2; -2 -2; 2 -2; -2 2;0  0]')
        knnconfig = KNNNeighborhood(data, 3);
        graph = create(knnconfig, weight, data);
        v1 = graph.vertices[1]
        @test v1.number_of_edges == 3
        for edge in v1
            v2 = target_vertex(edge, v1)
            @test norm(data[:, v2.id] - [-2, -2]) > 0.001
        end
        @test graph.vertices[2].connections == Set([4,3,5])
        @test graph.vertices[end].number_of_edges == 4
    end
    @testset "RandomNeighborhood" begin
        data, labels = two_gaussians()
        randomconfig = RandomNeighborhood(5)
        graph = create(randomconfig, ones, data);
        @test all([v.number_of_edges for v in graph.vertices] .>= 5)
    end
    @testset "CliqueNeighborhood" begin
        data, labels = two_gaussians(15)
        clique = CliqueNeighborhood()
        graph = create(clique, ones, data);
        @test all([v.number_of_edges for v in graph.vertices] .== 29)
        @test all([v.degree for v in graph.vertices] .== 29)
    end
    @testset "PixelNeighborhood" begin
        img = fill(RGB(1, 0, 0), 20, 20)
        nconfig = PixelNeighborhood(1)
        graph = create(nconfig, ones, img)
        @test length(graph.vertices) == size(img, 1) * size(img, 2)
        @test graph.vertices[1].number_of_edges == 3
        @test graph.vertices[148].number_of_edges == 8
    end
    @testset "Local Sale Knn" begin
        function weight(i::Integer, neigh, v, m)
            return Distances.colwise(Euclidean(), m, v)
        end
        X = [5.0 5; 5 4; 4 4; 6 6; -10 -10; -9 -9; -8 -8; -11 -11]'
        knnconfig = KNNNeighborhood(X, 3)
        scale = local_scale(knnconfig, weight, X, k = 3)
        @test isapprox(scale[1], sqrt(2))
        @test isapprox(scale[end], sqrt(18))
    end
    @testset "Local Sale Image" begin
        function weight(i::Integer, neigh, v, m)
            col_dist = Distances.colwise(Euclidean(), m[3:end, :], v[3:end])
            xy_dist = Distances.colwise(Euclidean(), m[1:2, :], v[1:2])
            return hcat(col_dist, xy_dist)
        end
        X = rand(RGB, 50, 50)
        knnconfig = PixelNeighborhood(3)
        scale = local_scale(knnconfig, weight, X, k = 9)
        @test size(scale) == (2, 50 * 50)
    end
end;

@testset "Embedding" begin
    @testset "NgLaplacian" begin
        function weight(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 15)
        end
        (data, labels) = two_gaussians()

        knnconfig = KNNNeighborhood(data, 7)
        graph = create(knnconfig, weight, data)
        emb = embedding(NgLaplacian(1), graph)
        pred_clustering = convert(Array{Int64}, (emb .<= mean(emb)))
        @test randindex(pred_clustering, labels)[4] > 0.9
    end
    @testset "ShiMalikLaplacian" begin
        function weight(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 15)
        end
        (data, labels) = two_gaussians()
        knnconfig = KNNNeighborhood(data, 15)
        graph = create(knnconfig, weight, data)
        emb = embedding(ShiMalikLaplacian(1), graph)
        pred_clustering = convert(Array{Int64}, (emb .<= mean(emb)))
        @test randindex(pred_clustering, labels)[4] > 0.9
    end
    @testset "PartialGroupingConstraints" begin
        function weight(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 0.7)
        end
        N = 150
        (d, labels) = three_gaussians(N)
        knnconfig = KNNNeighborhood(d, 100)
        graph = create(knnconfig, weight, d)

        indices_clus_1 = [1, 2, 3, 4, 5]
        indices_clus_2 = [N+1, N+2, N+3, N+4]
        indices_clus_3 = [2*N+1, 2*N+2, 2*N+3, 2*N+4]

        constraints = Vector{Integer}[ vcat(indices_clus_1, indices_clus_2) ] ;
        emb_1 = embedding(PartialGroupingConstraints(1, smooth=true),  graph, constraints)
        labels_1 = vcat(zeros(Integer, N*2), ones(Integer, N))

        constraints = Vector{Integer}[ vcat(indices_clus_2, indices_clus_3) ]
        emb_2 = embedding(PartialGroupingConstraints(1, smooth=true),  graph, constraints)
        labels_2 = vcat(zeros(Integer, N), ones(Integer, N*2))

        pred_clustering = convert(Array{Int64}, (emb_1 .<= mean(emb_1)))
        @test randindex(pred_clustering, labels_1)[4] > 0.85
        @test randindex(pred_clustering, labels_2)[4] < 0.5

        pred_clustering = convert(Array{Int64}, (emb_2 .<= mean(emb_2)))
        @test randindex(pred_clustering, labels_2)[4] > 0.85
        @test randindex(pred_clustering, labels_1)[4] < 0.5
    end
    @testset "YuShiPopout" begin
        function weight(i::Integer, ineigh, vi, vneigh)
            intensity_dist = Distances.colwise(Euclidean(), vi[3:end], vneigh[3:end, :])
            xy_dist = Distances.colwise(Euclidean(), vi[1:2], vneigh[1:2, :])
            a = 5
            b = 0.05
            return (pdf.(Normal(0, a*15), xy_dist) - pdf.(Normal(0, a), xy_dist)) .*
                   (pdf.(Normal(0, b*100), intensity_dist) - pdf.(Normal(0, b), intensity_dist))
        end
        function attraction(i::Integer, ineigh, vi, vneigh)
            diff = weight(i, ineigh, vi, vneigh)
            diff[diff.<0] .= 0
            return diff
        end
        function repulsion(i::Integer, ineigh, vi, vneigh)
            diff = weight(i, ineigh, vi, vneigh)
            diff[diff.>0] .= 0
            return abs.(diff)
        end
        img = zeros(31,31)
        img[8:25, 3:12] .= 0.9
        img[3:12, 5:28] .= 0.2
        img[8:25, 25:30] .= 0.6
        img = Gray.(img + randn(31, 31)*0.03)

        labels = zeros(Integer, 31, 31)
        labels[8:25, 3:12] .= 1
        labels[3:12, 5:28] .= 2
        labels[8:25, 25:30] .= 3

        nconfig = PixelNeighborhood(4)
        graph_attraction = create(nconfig, attraction, img);
        graph_repulsion = create(nconfig, repulsion, img);

        emb_config = YuShiPopout(3,  false)
        cluster_result = clusterize(emb_config, KMeansClusterizer(4), graph_attraction, graph_repulsion)
        @test randindex(cluster_result.assignments[:], labels[:])[4] > 0.9
    end
end
@testset "Clustering" begin
    @testset "KMeans Clustering" begin
        function weight(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 15)
        end
        (data, labels) = two_gaussians()
        knnconfig = KNNNeighborhood(data, 7)
        graph = create(knnconfig, weight, data)
        (data, labels) = two_gaussians()
        pred_clustering = clusterize(NgLaplacian(2), KMeansClusterizer(2), graph)
        @test randindex(pred_clustering.assignments, labels)[4] > 0.9
    end
    @testset "YuEigenvectorRotation" begin
        function weight(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 15)
        end
        (data, labels) = two_gaussians()
        knnconfig = KNNNeighborhood(data, 7)
        graph = create(knnconfig, weight, data)
        (data, labels) = two_gaussians()
        pred_clustering = clusterize(NgLaplacian(2), YuEigenvectorRotation(500), graph)
        @test randindex(pred_clustering.assignments, labels)[4] > 0.9
    end

end
@testset "Landmark Selection" begin
    @testset "RandomLandmarkSelection" begin
        r = RandomLandmarkSelection()
        data = rand(2, 25)
        s = select_landmarks(r, 15, data)
        @test length(s) == 15
        @test length(unique(s)) == length(s)
        @test minimum(s)>=1 && maximum(s)<= 25
    end
    @testset "EvenlySpacedLandmarkSelection" begin
        e = EvenlySpacedLandmarkSelection()
        data = rand(2, 25)
        s = select_landmarks(e, 5, data)
        @test length(s) == 5
        @test all(diff(s) .== 5)
        @test minimum(s)>=1 && maximum(s)<= 25
    end
    @testset "BresenhamLandmarkSelection" begin
        e = BresenhamLandmarkSelection()
        data = rand(2, 25)
        s = select_landmarks(e, 5, data)
        @test length(s) == 5
        @test all(diff(s) .> 0)
        @test minimum(s)>=1 && maximum(s)<= 25
    end
end
@testset "Approximate Embedding" begin
    @testset "Nystrom" begin
        @testset "Data embedding" begin
            function weight(i::Integer, neigh, v, m)
                return exp.(-Distances.colwise(SqEuclidean(), m, v) / 15)
            end
            (data, labels) = two_gaussians(6000)
            embedding_config = NystromMethod(EvenlySpacedLandmarkSelection(), 1000, weight, 1)
            emb = embedding(embedding_config, data)
            pred_clustering = convert(Array{Int64}, (emb .<= mean(emb)))
            @test randindex(pred_clustering, labels)[4] > 0.9
        end
        @testset "Image embedding" begin
            function weight(i::Integer,j::Vector{<:Integer},pixel_i, neighbors_data)
                data_diff = pixel_i[3:5] .- neighbors_data[3:5,:]
                a = exp.(-abs.(data_diff)./(2*0.1^2))
                a = prod(a, dims=1)
                return vec(a)
            end

            img = fill(RGB(0,0,0), 50, 50)

            cluster_1 = CartesianIndices(size(img))[5:20, 5:20]
            cluster_2 = CartesianIndices(size(img))[35:42, 35:47]
            img[cluster_1] .= RGB(1, 0, 0)
            img[cluster_2] .= RGB(0, 0, 1)

            labels = zeros(Integer, 50*50)
            labels[LinearIndices(size(img))[cluster_1][:]] .= 1
            labels[LinearIndices(size(img))[cluster_2][:]] .= 2

            embedding_config = NystromMethod(EvenlySpacedLandmarkSelection(),
                                            500, weight, 1)
            emb = embedding(embedding_config, img)
            emb = vec(round.(emb, digits=2))
            clusters = zeros(Integer, 50*50)
            for (i, val) in enumerate(unique(emb))
                clusters[findall(emb.==val)] .= i -1
            end
            @test randindex(clusters, labels)[4] > 0.95
        end
    end
    @testset "DNCuts" begin
        function weight(i::Integer,j::Vector{<:Integer},pixel_i, neighbors_data)
            data_diff = pixel_i[3:5] .- neighbors_data[3:5,:]
            a = exp.(-((data_diff).^2)./(0.1))
            a = prod(a, dims=1)
            return vec(a)
        end
        img = fill(RGB(0,0,0), 50, 50)

        cluster_1 = CartesianIndices(size(img))[5:20, 5:20]
        cluster_2 = CartesianIndices(size(img))[35:42, 35:47]
        img[cluster_1] .= RGB(1, 0, 0)
        img[cluster_2] .= RGB(0, 0, 1)

        labels = ones(Integer, 50*50)
        labels[LinearIndices(size(img))[cluster_1][:]] .= 2
        labels[LinearIndices(size(img))[cluster_2][:]] .= 3

        nconfig = PixelNeighborhood(4)
        graph = create(nconfig, weight, img);

        dncuts = DNCuts(2, 2, size(img))
        emb = embedding(dncuts, graph)
        pred_clustering = clusterize(dncuts, KMeansClusterizer(3), graph)
        @test randindex(pred_clustering.assignments, labels)[4] > 0.9
    end
    @testset "LandmarkBasedRepresentation" begin
        function weight(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 15)
        end
        (data, labels) = two_gaussians(6000)
        cfg = LandmarkBasedRepresentation(
                BresenhamLandmarkSelection(),
                500,
                25,
                2,
                weight,
                true)
        emb = embedding(cfg, data)
        pred_clustering = convert(Array{Int64}, (emb[:, 1] .<= mean(emb[:, 1])))
        @test randindex(pred_clustering, labels)[4] > 0.9
    end
end
@testset "MultiView" begin
    @testset "CoRegularizedMultiView" begin
        function weight_1(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 15)
        end
        function weight_2(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 45)
        end
        (data, labels) = two_gaussians(500, center_1=[-15, -15], center_2=[9, 9])
        knnconfig = KNNNeighborhood(data, 7)
        graph_1 = create(knnconfig, weight_1, data);
        graph_2 = create(knnconfig, weight_2, data);
        coreg = CoRegularizedMultiView([View(1, 0.001),
                                        View(1, 0.001)])
        emb = embedding(coreg, [graph_1, graph_2])
        pred_clustering = convert(Array{Int64}, (emb[:, 1] .<= mean(emb[:, 1])))
        @test randindex(pred_clustering, labels)[4] > 0.9
    end
    @testset "KernelProduct" begin
        function weight_1(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 15)
        end
        function weight_2(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 45)
        end
        (data, labels) = two_gaussians(500, center_1=[-15, -15], center_2=[9, 9])
        knnconfig = KNNNeighborhood(data, 7)
        graph_1 = create(knnconfig, weight_1, data);
        graph_2 = create(knnconfig, weight_2, data);
        kernel_addition = KernelProduct(NgLaplacian(2))
        emb = embedding(kernel_addition, [graph_1, graph_2])
        pred_clustering = convert(Array{Int64}, (emb[:, 1] .<= mean(emb[:, 1])))
        @test randindex(pred_clustering, labels)[4] > 0.9
    end
    @testset "KernelAddition" begin
        function weight_1(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 15)
        end
        function weight_2(i::Integer, neigh, v, m)
            return exp.(-Distances.colwise(SqEuclidean(), m, v) / 45)
        end
        (data, labels) = two_gaussians(500, center_1=[-15, -15], center_2=[9, 9])
        knnconfig = KNNNeighborhood(data, 7)
        graph_1 = create(knnconfig, weight_1, data);
        graph_2 = create(knnconfig, weight_2, data);
        kernel_addition = KernelAddition(NgLaplacian(2))
        emb = embedding(kernel_addition, [graph_1, graph_2])
        pred_clustering = convert(Array{Int64}, (emb[:, 1] .<= mean(emb[:, 1])))
        @test randindex(pred_clustering, labels)[4] > 0.9
    end
end

@testset "Utils" begin
    @testset "Normalization" begin
        a = rand(50, 3)
        @test isapprox(norm(SpectralClustering.normalize_rows(a)[1, :]), 1)
        @test isapprox(norm(SpectralClustering.normalize_cols(a)[:, 1]), 1)
        a = rand(50, 1)
        @test isapprox(norm(SpectralClustering.normalize_rows(a)[1, :]), a[1])
    end
end
@testset "DataAcces" begin
    a = rand(50, 3)
    @test number_of_patterns(a) == 3
    @test get_element(a, 1) == a[:, 1]
    a = a'
    @test number_of_patterns(a) == 50
    @test get_element(a, 1) == a[:, 1]

    a = rand(RGB, 50 ,50)
    @test number_of_patterns(a) == 50*50
end


export QFT, qft
using BitBasis
using LinearAlgebra
using FFTW

"""
    QFT{N} <: PrimitiveBlock{N}

Quantum Fourier Transform node. See also [`qft_circuit`](@ref).
"""
struct QFT{N} <: PrimitiveBlock{N} end

"""
    qft(n)

Create a Quantum Fourier Transform block. See also [`qft_circuit`](@ref).
"""
qft(n) = QFT{n}()

Yao.mat(::Type{T}, q::QFT{N}) where {T, N} = T.(applymatrix(q))

function _apply!(r::ArrayReg, x::QFT)
    α = sqrt(1<<nactive(r))
    st = state(invorder!(r))
    ifft!(st, 1)
    lmul!(α, st)
    return r
end

function _apply!(r::ArrayReg, ::Daggered{<:QFT})
    # (reg.state = invorder_firstdim(fft!(reg|>state, 1)/sqrt(1<<nactive(reg))); reg)
    α = sqrt(1 << nactive(r))
    st = state(r)
    fft!(st, 1)
    st ./= α
    invorder!(r)
    return r
end

# traits
Yao.ishermitian(q::QFT{N}) where N = N==1
Yao.isreflexive(q::QFT{N}) where N = N==1
Yao.isunitary(q::QFT{N}) where N = true

function Yao.print_block(io::IO, pb::QFT{N}) where N
    printstyled(io, "qft(1-$N)"; bold=true, color=:blue)
end

function Yao.print_block(io::IO, pb::Daggered{N,<:QFT}) where {N, T}
    printstyled(io, "iqft(1-$N)"; bold=true, color=:blue)
end