using Test
using DiffEqOperators, OrdinaryDiffEq, LinearAlgebra

@testset "KdV equation (Single Solition)" begin
    N = 21
    Δx = 1/(N-1)
    r = 0.5

    # x = 10:Δx:30;
    x = -10:Δx:10;
    # ϕ(x,t) = (r/2)*sech.((sqrt(r)*(x-r*t)/2)-7).^2 # solution of the single forward moving wave
    ϕ(x,t) = (1/2)*sech.((x .- t)/2).^2 # solution of the single forward moving wave

    u0 = ϕ(x,0);
    oriu = zeros(size(x));

    #const du3 = zeros(size(x));
    du = zeros(size(x));
    # const temp = zeros(size(x));

    # A = CenteredDifference(1,2,Δx,length(x),:Dirichlet0,:Dirichlet0);
    A = UpwindDifference{Float64}(1,3,Δx,length(x),-1);
    # C = CenteredDifference(3,2,Δx,length(x),:Dirichlet0,:Dirichlet0);
    #C = UpwindDifference{Float64}(3,3,Δx,length(x),-1);

    function KdV(du, u, p, t)
        # bc = DirichletBC(ϕ(-10-Δx,t),ϕ(10+Δx,t))
        bc = GeneralBC([0,1,-6*ϕ(-10,t),0,-1],[0,1,-6*ϕ(10,t),0,-1],Δx,3)
        #mul!(du3,C,bc*u)
        mul!(du,A,bc*u)
        # @. temp = -0.5*u*du - 0.25*du3
        # copyto!(du,temp)

    end

    single_solition = ODEProblem(KdV, u0, (0.,5.));
    soln = solve(single_solition,Tsit5(),abstol=1e-6,reltol=1e-6);

    #=
    for t in 0:0.5:5
        @show maximum(soln(t)-ϕ(x,t))
    end
    =#
    for t in 0:0.5:5
        @test soln(t) ≈ ϕ(x,t) atol = 0.01;
    end

    # Using Biased Upwinds with 1 offside point
    A2 = UpwindDifference{Float64}(1,3,Δx,length(x),-1,offside=1);
    function KdV(du, u, p, t)
        bc = GeneralBC([0,1,-6*ϕ(-10,t),0,-1],[0,1,-6*ϕ(10,t),0,-1],Δx,3)
        mul!(du,A2,bc*u)
    end
    single_solition = ODEProblem(KdV, u0, (0.,5.));
    soln = solve(single_solition,Tsit5(),abstol=1e-6,reltol=1e-6);
    for t in 0:0.5:5
        @test soln(t) ≈ ϕ(x,t) atol = 0.01;
    end

    A3 = UpwindDifference{Float64}(1,3,Δx*ones(length(x)+1),length(x),-1,offside=1);
    function KdV(du, u, p, t)
        bc = GeneralBC([0,1,-6*ϕ(-10,t),0,-1],[0,1,-6*ϕ(10,t),0,-1],Δx,3)
        mul!(du,A3,bc*u)
    end
    single_solition = ODEProblem(KdV, u0, (0.,5.));
    soln = solve(single_solition,Tsit5(),abstol=1e-6,reltol=1e-6);
    for t in 0:0.5:5
        @test soln(t) ≈ ϕ(x,t) atol = 0.01;
    end

    # Using Biased Upwinds with 2 offside points
    A4 = UpwindDifference{Float64}(1,4,Δx,length(x),-1,offside=2);
    function KdV(du, u, p, t)
        bc = GeneralBC([0,1,-6*ϕ(-10,t),0,-1],[0,1,-6*ϕ(10,t),0,-1],Δx,3)
        mul!(du,A4,bc*u)
    end
    single_solition = ODEProblem(KdV, u0, (0.,5.));
    soln = solve(single_solition,Tsit5(),abstol=1e-6,reltol=1e-6);
    for t in 0:0.5:5
        @test soln(t) ≈ ϕ(x,t) atol = 0.01;
    end

    A5 = UpwindDifference{Float64}(1,4,Δx,length(x),1,offside=2);
    function KdV(du, u, p, t)
        bc = GeneralBC([0,1,-6*ϕ(-10,t),0,-1],[0,1,-6*ϕ(10,t),0,-1],Δx,3)
        mul!(du,-1*A5,bc*u)
    end
    single_solition = ODEProblem(KdV, u0, (0.,5.));
    soln = solve(single_solition,Tsit5(),abstol=1e-6,reltol=1e-6);
    for t in 0:0.5:5
        @test soln(t) ≈ ϕ(x,t) atol = 0.01;
    end
end

# Conduct interesting experiments by referring to
# http://lie.math.brocku.ca/~sanco/solitons/kdv_solitons.php
#  TODO The true solution for this case seems unstable, look for a workaround
#= @testset "KdV equation (Double Solition)" begin
    N = 10
    Δx = 1/(N-1)

    x = -50:Δx:50;
    c1,c2 = 25,16
    # ϕ1(x,t) = 3*p1^2*sech(.5*(p1*(x+2+t))).^2+3*p2^2*sech(.5*(p2*(x+1+t))).^2;

    function ϕ(x,t)
        # t = t-10
        num1 = 2(c1-c2)*(c1*cosh.(√c2*(x .- c2*t)/2).^2 + c2*sinh.(√c1*(x .- c1*t)/2).^2)
        den11 = (√c1-√c2)*cosh.((√c1*(x .- c1*t) + √c2*(x .- c2*t))/2)
        den12 = (√c1+√c2)*cosh.((√c1*(x .- c1*t) - √c2*(x .- c2*t))/2)
        den1 = (den11+den12).^2
        return num1./den1
    end

    # const du3 = zeros(size(x));
    du = zeros(size(x));
    # const temp = zeros(size(x));

    A = UpwindDifference{Float64}(1,3,Δx,length(x),-1);

    # C = UpwindDifference{Float64}(3,1,Δx,length(x),-1);

    function KdV(du, u, p, t)
        bc = GeneralBC([0,1,-6*ϕ(-50,t),0,-1],[0,1,-6*ϕ(50,t),0,-1],Δx,3)
        # mul!(du3,C,bc*u)
        mul!(du,A,bc*u)
        # @. temp = -6*u*du - du3
        # copyto!(du,temp)
    end

    u0 = ϕ(x,0);
    double_solition = ODEProblem(KdV, u0, (0.,5.));
    soln = solve(double_solition,Tsit5());

    # The solution is a forward moving soliton wave with speed = 1
    for t in 0:0.1:9
        @test_skip soln(t) ≈ ϕ(x,t) atol = 0.01
    end
end =#


@testset "$(rpad("Lagrange Polynomials Tests",80))" begin

    b = LagrangeBasis([0.0, 1.0])
    l = LagrangePolynomial([0.0, 1.0], ones(2))

    @test l.b == b
    @test l.c == ones(2)

    xp = [0.0, 0.5, 1.0]
    yp = zeros(3)
    yr = [1.0, 0.5, 0.0]

    evaluate!(l, xp, yp)
    @test yp == ones(3)

    y = zeros(2)
    y[1] = 1

    evaluate!(similar(l, y), xp, yp)
    @test yp == yr

    evaluate!(b, y, xp, yp)
    @test yp == yr

    y = zeros(2)
    y[2] = 1

    evaluate!(similar(l, y), xp, yp)
    @test yp == xp

    evaluate!(b, y, xp, yp)
    @test yp == xp

end


type SDF_FD <: SDFDynamicSystem
    omg::Float64
    zta::Float64
    phi::Float64
    u0::Float64
    v0::Float64

    SDF_FD(omg::Float64, zta::Float64, phi::Float64) = new(omg, zta, phi, 0.0, 0.0)
end

function systemEquation(s::SDF_FD, p::Function)
    function dydt(y::AbstractArray{Float64,1}, t::Float64)
        [y[2]; p(t)-2.0*s.zta*s.omg*y[2]-s.phi*g*sign(y[2])-s.omg*s.omg*y[1]]
    end
    dydt
end

function response(s::SDF_FD, p::Function, t::AbstractArray{Float64,1})
    y0 = [s.u0; s.v0]

    f = systemEquation(s, p)

    y = ode4(f, y0, t)

    u = y[:,1]
    v = y[:,2]
    fdz = 2.0*s.zta*s.omg*v
    fdp = s.phi*g*sign(v)
    a = map(p,t)-fdz-fdp-s.omg*s.omg*u

    [u v a fdz fdp]
end


@testset "CDS_scaled_add_tests" begin

#test addition structured matrix
TI=Int64
TF=Float64;
n1=30;
n2=20;
N=n1*n2;
x=randn(TF,N);
y=zeros(TF,N);

#test structured matrix
#comp_grid = compgrid( (25, 25),( n1,n2 ) );
comp_grid = compgrid( (TF(25), TF(25)),( n1,n2 ) );
(TV_OP, AtA_diag, dense, TD_n, banded)=get_TD_operator(comp_grid,"TV",TF);
A=TV_OP'*TV_OP;
A=convert(SparseMatrixCSC{TF,TI},A);
(D_z, AtA_diag, dense, TD_n, banded)=get_TD_operator(comp_grid,"D_z",TF);
B=D_z'*D_z;
B=convert(SparseMatrixCSC{TF,TI},B);

C=A+B;

(R_A,offset_A)=mat2CDS(A)
(R_B,offset_B)=mat2CDS(B)
(R_C,offset_C)=mat2CDS(C)

rho=1.0;
CDS_scaled_add!(R_A,R_B,offset_A,offset_B,rho )

#test if CDS storage of the sum is the same as the sum of the CDS storage matrices
@test R_C==R_A
@test offset_C==offset_A

#test if it still works with a matrix-vector product
Cx_native = C*x;
fill!(y,0)
Cx_CDS = CDS_MVp(size(R_A,1),size(R_A,2),R_A,offset_A,x,y)
@test isapprox(Cx_native,Cx_CDS,rtol=10*eps())

#test random matrix
A=sprandn(100,100,0.01)
B=sprandn(100,100,0.01)
x=randn(100)
y=zeros(100)

A=A+B #make sure A has got all offsets of A and B

(R_A,offset_A)=mat2CDS(A)
(R_B,offset_B)=mat2CDS(B)

C=A.+B;
Cx_native = C*x;

rho=1.0;
CDS_scaled_add!(R_A,R_B,offset_A,offset_B,rho )

fill!(y,0)
Cx_CDS = CDS_MVp(size(R_A,1),size(R_A,2),R_A,offset_A,x,y)
@test isapprox(Cx_native,Cx_CDS,rtol=10*eps())

(R_C,offset_C)=mat2CDS(C)
@test R_C==R_A
@test offset_C==offset_A

end


module Schrodinger
using Compat
using Compat.LinearAlgebra, Compat.SparseArrays, Compat.Printf
using Base: tail, front, product, promote_eltype
using OrdinaryDiffEq, Optim
using OrdinaryDiffEq.DiffEqBase: AbstractParameterizedFunction

export Operator, Ket, Bra,
    data, dims, isnormalized, dimsmatch, dense, braket, isapproxhermitian,
    isunitary, isapproxunitary, isdensityop, checkdensityop,
    hermitian, symmetric,
    tensor, ptrace, expect, levelprobs, findstate,
    purity, fidelity, fidelity2, entanglement_fidelity, gate_fidelity,
    fock, basis, coherent, thermal, maxentangled, maxmixed, ket, @qb_str,
    qzero, qeye, create, destroy, numberop, projectorop,
    displacementop, squeezeop,
    rand_unitary,
    Liouvillian, SchrodingerEvo, LindbladEvo,
    Propagator, SchrodingerProp, LindbladProp,
    sesolve, mesolve, lsolve, psolve, psteady,
    grape, NormPSU, CoherentSubspaces,
    expim, gaussian, inner, scale, scale!, full,
    eigen

# VERSION-conditional definitions
if VERSION < v"2.0.0-"
    include("basepatch/v1.0.jl")
end
if VERSION < v"0.7.0-"
    include("basepatch/v0.6.jl")
end
if VERSION > v"0.7.0-"
    import Base: BitSet, print_array, sum
    import LinearAlgebra: adjoint, exp, exp!, tr, eigen
    import Arpack: eigs
    import SpecialFunctions: erf
    export eigs, eigvals, trace, normalize, normalize!
    trace(A::AbstractMatrix) = tr(A)
    sum(A::AbstractArray,i::Integer) = Base._sum(A,i)
    parseb2(s::AbstractString) = Base.tryparse_internal(Int,s,firstindex(s),lastindex(s),2,true)
    const luf = lu
end

include("quobj/types.jl")
include("quobj/basicmethods.jl")
include("math/basemath.jl")
include("math/eigen.jl")
include("math/special.jl")
include("math/superoperators.jl")
include("math/norms.jl")
include("math/ptrace.jl")
include("dynamics/liouvillian.jl")
include("dynamics/propagator.jl")
include("dynamics/constructors.jl")
include("dynamics/interface.jl")
include("control/grape.jl")
include("control/normpsu.jl")
include("control/coherentsubspaces.jl")
include("control/amplitudelimiter.jl")
include("tomography/statetomo.jl")
include("tomography/processtomo.jl")
include("misc/utils.jl")
include("misc/indexing.jl")
include("misc/tupletools.jl")
include("misc/inner.jl")
include("misc/kron.jl")
include("misc/sparsevec.jl")
include("misc/checks.jl")
include("misc/hermfact.jl")
include("library/operators.jl")
include("library/states.jl")
include("library/random.jl")
include("library/drivefuns.jl")
include("library/constants.jl")

end


include("testsystem.jl")
Random.seed!(11)
G0 = Gaussian(x0, P0)
x = rand(StateObs(G0, M.obs))
X = trace(DynamicIterators.Sampled(M), 1 => x, endtime(3))
@show x

Y = collect(t=>y for (t, (x,y)) in pairs(X))
@show Y

Xf, ll = kalmanfilter(M, 1 => G0, Y)

Xs, ll = rts_smoother(M, 1 => G0, Y)
@show Xs

QL = sqrt(Q)

P0L = sqrt(P0)

RL = sqrt(R)

μ = [ x0; Φ*x0; Φ*Φ*x0; H*x0; H*Φ*x0; H*Φ*Φ*x0; ]
L = [    P0L      0    0  0  0  0
       Φ*P0L     QL    0  0  0  0
     Φ*Φ*P0L   Φ*QL   QL  0  0  0
       H*P0L      0    0 RL  0  0
     H*Φ*P0L   H*QL    0  0 RL  0
   H*Φ*Φ*P0L H*Φ*QL H*QL  0  0 RL
]
Σ = L*L'
v = last.(Y)
@show μ
@show Σ
@show v

n = 3
F = [GaussianDistributions.conditional(Gaussian(μ, Σ), 1:i, (1:i).+n, v[1:i]) for i in 1:n]
μf = last.(mean.(F))
Σf = last.(cov.(F))
μs = mean(F[n])
Σs = diag(cov(F[n]))

o = (1:n).+n
ll2 = logpdf(Gaussian(μ[o], Σ[o, o]), v)

@test μf ≈ mean.(values(Xf))
@test Σf ≈ cov.(values(Xf))

@test μs ≈ mean.(values(Xs))
@test Σs ≈ cov.(values(Xs))

@test ll ≈ ll2

# test backward sampler

X1 = [Kalman.backward_sampler(M, 1 => G0, Y).x[1] for k in 1:1000]

@test norm(mean(X1) - mean(Xs.x[1])) < 0.005
@test norm(cov(X1) - cov(Xs.x[1])) < 0.0025


"""
    ar_sum(x)
Reduce `k = length(x)` vectors through sum to a single vector.
"""
ar_sum(x) = sum(x)

"""
    ar_prod(x)
Reduce `k = length(x)` vectors through product to a single vector.
"""
ar_prod(x) = reduce((y, z) -> y .* z, x)

"""
    arithmetic_layer()
Generate the layer of arithmetic operations of the ICN. The operations are mutually exclusive, that is only one will be selected.
"""
function arithmetic_layer()
    arithmetics = LittleDict{Symbol, Function}(
        :sum => ar_sum,
        :prod => ar_prod,
    )

    return Layer(arithmetics, true)
end


#converts gradient field into heightmap using Fourier method
function convert_gradient(scheme::Frankot, pin::AbstractArray, qin::AbstractArray)
    p = Complex{Float64}.(pin)
    q = Complex{Float64}.(qin)
    wx, wy = setup_transform_values(last(size(p)),first(size(p)))
    fft!(p)
    fft!(q)
    Z = zeros(Complex{Float64}, size(p))
    Z = (-1im .* wx .* p .+ 1im .* wy .* q) ./ (2 .* π .* (wx.^2 .+ wy.^2 .+ eps()))
    ifft!(Z)
    Z = abs.(Z)
    return Z
end

#converts gradient field into heightmap using average of two integrals
function convert_gradient(scheme::Path, p::AbstractArray, q::AbstractArray)
    R, C = size(p)
    Z₁ = zeros(Float64, size(p))
    Z₂ = zeros(Float64, size(p))

    #path 1
    for i = 2:R
        Z₁[i,1] = Z₁[i-1,1] - q[i,1]
    end
    for i = 1:R
        for j = 2:C
            Z₁[i,j] = Z₁[i,j-1] + p[i,j]
        end
    end

    #path 2
    for i = 2:C
        Z₂[1,i] = Z₂[1,i-1] + p[1,i]
    end
    for i = 2:R
        for j = 1:C
            Z₂[i,j] = Z₂[i-1,j] - q[i,j]
        end
    end
    Z = (Z₁ .+ Z₂) / 2
    return Z
end

#converts gradient field into heightmap using integral culculated from average gradient at point
function convert_gradient(scheme::SplitPath, p::AbstractArray, q::AbstractArray)
    R, C = size(p)
    Z = zeros(Float64, size(p))
    for i = 2:R
        Z[i,1]=Z[i-1,1] + q[i,1]
        Z[1,i]=Z[1,i-1] - p[1,i]
    end
    for i = 2:R
        for j = i:R
            Z[i,j] = ((Z[i,j-1] + p[i,j]) + (Z[i-1,j] - q[i,j])) / 2
            Z[j,i] = ((Z[j,i-1] + p[j,i]) + (Z[j-1,i] - q[j,i])) / 2
        end
    end
    return Z
end

#Horn and Brooks
function convert_gradient(scheme::Horn, p::AbstractArray, q::AbstractArray, iter::Real = 10000, ϵ::Real = 1)
    Z = zeros(Float64, size(p))
    Zᵏ⁺¹ = zeros(Float64, size(p))
    h = zeros(Float64, size(p))
    v = zeros(Float64, size(p))
    hv = zeros(Float64, size(p))
    R, C = size(p)
    for i = 1:R
        Z[i,1] = p[i,1]
        Z[i,R] = p[i,R]
    end
    for i = 1:C
        Z[1,i] = q[1,i]
        Z[C,i] = q[C,i]
    end
    for i = 2:(R-1)
        for j = 2:(C-1)
            h[i,j] = (p[i,j+1] - p[i,j-1]) / 2
            v[i,j] = (q[i+1,j] - q[i-1,j]) / 2
            hv[i,j] = (ϵ / 4) * (h[i,j] - v[i,j])
        end
    end
    for k = 1:iter
        copyto!(Zᵏ⁺¹, Z)
        for i = 2:(R-1)
            for j = 2:(C-1)
                Zᵏ⁺¹[i,j] = (Z[i-1,j] + Z[i+1,j] + Z[i,j-1] + Z[i,j+1]) / 4
                Zᵏ⁺¹[i,j] = Zᵏ⁺¹[i,j] - hv[i,j]
            end
        end
        Z = Zᵏ⁺¹
    end
    return Z
end

#Durou and Courteille
function convert_gradient(scheme::Durou, p::AbstractArray, q::AbstractArray, iter::Real = 1000, ϵ::Real = 1)
    Z = zeros(Float64, size(p))
    Zᵏ⁺¹ = zeros(Float64, size(p))
    h = zeros(Float64, size(p))
    v = zeros(Float64, size(p))
    hv = zeros(Float64, size(p))
    R, C = size(p)
    for i = 1:(R-1)
        for j = 1:(C-1)
            h[i,j] = (p[i,j+1] + p[i,j]) / 2
            v[i,j] = (q[i+1,j] + q[i,j]) / 2
            hv[i,j] = (ϵ / 2) * (h[i,j] - v[i,j])
        end
    end
    for k = 1:iter
        copyto!(Zᵏ⁺¹, Z)
        for i = 1:(R-1)
            for j = 1:(C-1)
                Zᵏ⁺¹[i,j] = (Z[i+1,j] + Z[i,j+1]) / 2
                Zᵏ⁺¹[i,j] = Zᵏ⁺¹[i,j] - hv[i,j]
            end
        end
        Z = Zᵏ⁺¹
    end
    return Z
end


type MEXCLPDeployment <: DeploymentModel
    m::JuMP.Model
    x::Vector{JuMP.Variable}
end
deployment(m::MEXCLPDeployment) = Int[round(Int,x) for x in JuMP.getvalue(m.x)]

function MEXCLPDeployment(p::DeploymentProblem,
                          q::Float64; # busy fraction
                          max_amb::Int = 0,
                          tol = params.δ,
                          solver = GurobiSolver(OutputFlag=0))
    (max_amb == 0) && (max_amb = p.nambulances)
    @assert max_amb > 0
    demand = vec(mean(p.demand[p.train,:],1))
    @assert length(demand) == p.nregions

    I = 1:p.nlocations
    J = 1:p.nregions
    K = 1:max_amb

    m = JuMP.Model(solver=solver)
    JuMP.@variable(m, x[1:p.nlocations] >= 0, Int)
    JuMP.@variable(m, z[1:p.nregions, 1:max_amb], Bin)

    JuMP.@objective(m, Max, sum(demand[j]*(1-q)*(q^k)*z[j,k] for j in J, k in K))

    JuMP.@constraint(m, sum(x[i] for i in I) <= p.nambulances)

    for j in J # coverage over all regions
        JuMP.@constraint(m, sum(x[i] for i in filter(i->p.coverage[j,i], I)) >= 1)
        JuMP.@constraint(m, sum(x[i] for i in filter(i->p.coverage[j,i], I)) >=
                            sum(z[j,k] for k in K))
    end

    MEXCLPDeployment(m, x)
end

solve(model::MEXCLPDeployment) = JuMP.solve(model.m)
solve(model::MEXCLPDeployment, p::DeploymentProblem) = JuMP.solve(model.m)


function Poisson(E, λ)

    (exp.(-λ).*λ.^E)/factorial.(E)

end

function cdfPoisson(x_1, x_2, λ)

    [sum(Poisson.(collect(x_1:i), λ)) for i in x_1:x_2]

end


using FFTW

export TwoDPoissonPeriodic

"""
    TwoDPoissonPeriodic

Derived type to solve the Poisson equation on 2d regular
cartesian mesh with periodic boundary conditions on both sides
- kx   : wave number in x
- ky   : wave number in y
- k2   : ``k_x^2 + k_y^2``
- nc_x : cells number in x
- nc_y : cells number in y
- dx   : x step size
- dy   : y step size
- rht  : fft(rho)
"""
struct TwoDPoissonPeriodic

    grid::TwoDGrid
    kx::Array{Float64,2}
    ky::Array{Float64,2}
    k2::Array{Float64,2}
    rht::Array{ComplexF64,2}

    function TwoDPoissonPeriodic(grid::TwoDGrid)

        nc_x = grid.nx
        nc_y = grid.ny

        rht = zeros(ComplexF64, (div(nc_x, 2) + 1, nc_y))

        kx = zeros(nc_x ÷ 2 + 1, nc_y)
        ky = zeros(nc_x ÷ 2 + 1, nc_y)
        k2 = zeros(nc_x ÷ 2 + 1, nc_y)

        kx0 = 2π / grid.dimx
        ky0 = 2π / grid.dimy

        for ik = 1:nc_x÷2+1
            kx1 = (ik - 1) * kx0
            for jk = 1:nc_y÷2
                kx[ik, jk] = kx1
                ky[ik, jk] = (jk - 1) * ky0
            end
            for jk = nc_y÷2+1:nc_y
                kx[ik, jk] = kx1
                ky[ik, jk] = (jk - 1 - nc_y) * ky0
            end
        end

        kx[1, 1] = 1.0
        k2 .= kx .* kx .+ ky .* ky
        kx .= kx ./ k2
        ky .= ky ./ k2

        new(grid, kx, ky, k2, rht)

    end

end

export solve!

"""
    solve!( poisson, phi, rho )

computes `phi` from `rho`

```math
-\\Delta phi(x,y) = rho(x,y)
```

"""
function solve!(phi, poisson::TwoDPoissonPeriodic, rho)

    poisson.rht .= rfft(rho)
    poisson.rht ./= poisson.k2
    phi .= irfft(poisson.rht, poisson.grid.nx)

end

"""
    solve!( ex, ey, poisson, rho )

solves Poisson equation to compute electric fields

```math
E(x,y) = -\\nabla \\phi(x,y) \\\\
-\\Delta \\phi(x,y) = \\rho(x,y)
```

"""
function solve!(ex, ey, poisson::TwoDPoissonPeriodic, rho)

    poisson.rht .= rfft(rho)

    poisson.rht[1, 1] = 0.0
    poisson.rht .*= -1im .* poisson.kx

    ex .= irfft(poisson.rht, poisson.grid.nx)

    poisson.rht .= rfft(rho)

    poisson.rht[1, 1] = 0.0
    poisson.rht .*= -1im .* poisson.ky

    ey .= irfft(poisson.rht, poisson.grid.nx)

end

export TwoDPoissonPIC

"""
    TwoDPoissonPIC( poisson, kernel )

- kernel : Kernel smoother taking care of charge deposition and field evaluation
- poisson : Poisson solver
- rho_dofs : Coefficients of expansion of rho (MPI global version)
- efield_dofs : Coefficients of expansion of electric field
- phi_dofs : Coefficients of expansion of potential
- rho2d : 2d version of rho_dofs to adjust to field solver format
- efield : 2d version of efield_dofs to adjust to field solver format
- phi2d : 2d version of phi_dofs to adjust to field solver format
"""
struct TwoDPoissonPIC

    ndofs::Int
    kernel::ParticleMeshCoupling2D
    poisson::TwoDPoissonPeriodic
    rho_dofs::Vector{Float64}
    efield_dofs::Vector{Vector{Float64}}
    phi_dofs::Vector{Float64}
    rho2d::Array{Float64,2}
    efield::Vector{Array{Float64,2}}
    phi2d::Array{Float64,2}

    function TwoDPoissonPIC(poisson::TwoDPoissonPeriodic, kernel::ParticleMeshCoupling2D)

        nx, ny = poisson.grid.nx, poisson.grid.ny
        ndofs = nx * ny

        rho_dofs = zeros(ndofs)
        rho_dofs_local = zeros(ndofs)
        rho_analyt_dofs = zeros(ndofs)
        efield_dofs = [zeros(ndofs), zeros(ndofs)]
        phi_dofs = zeros(ndofs)
        rho2d = zeros(nx, ny)
        efield = [zeros(nx, ny), zeros(nx, ny)]
        phi2d = zeros(nx, ny)

        new(ndofs, kernel, poisson, rho_dofs, efield_dofs, phi_dofs, rho2d, efield, phi2d)

    end

end

"""
    add_charge!(pic, position, marker_charge)

Add charge from one particle
- self : Pic Poisson solver object
- position : Position of the particle
- marker_charge : Particle weight times charge
"""
function add_charge!(pic::TwoDPoissonPIC, position, marker_charge)

    xp, yp = position

    GEMPIC.add_charge!(pic.rho_dofs, pic.kernel, xp, yp, marker_charge)

end

"""
    evaluate_rho!(pic, position)

Evaluate charge density at rho at one position
- self : Pic Poisson solver object
- position : Position of the particle
- func_value : Value of rho at given position
"""
function evaluate_rho!(pic, position)

    GEMPIC.evaluate!(kernel, position, pic.rho_dofs)

end

"""
    evaluate_phi!(pic, position)

Evaluate potential at one position
- self : Pic Poisson solver object
- position : Position of the particle
- func_value : Value of phi at given position
"""
function evaluate_rho(pic, position)

    GEMPIC.evaluate!(kernel, position, pic.phi_dofs)

end

"""
    solve!( pic )

Solve for phi and fields
- poisson : Pic Poisson solver object
"""
function solve!(pic :: TwoDPoissonPIC)

    solve_phi!(pic)
    solve_fields!(pic)

end

"""
    solve_phi!( pic )

Solve for potential
"""
function solve_phi!(pic :: TwoDPoissonPIC)

    nx, ny = pic.poisson.grid.nx, pic.poisson.grid.ny
    pic.rho2d .= reshape(pic.rho_dofs, nx, ny)
    solve!(pic.phi2d, pic.poisson, pic.rho2d)
    pic.phi_dofs .= vec(pic.phi2d)

end

"""
    solve_fields!( pic )

Solve efields from rho
"""
function solve_fields!(pic :: TwoDPoissonPIC)

    solve!(pic.efield[1], pic.efield[2], pic.poisson, pic.rho2d)

    pic.efield_dofs[1] .= vec(pic.efield[1])
    pic.efield_dofs[2] .= vec(pic.efield[2])

end

export SplittingOperator

"""
Operator splitting type for 2d2v Vlasov-Poisson
- pic :: PIC poisson solver
- pg :: Particle group
"""
struct SplittingOperator

    pic::TwoDPoissonPIC
    pg::ParticleGroup

end

"""
    operator_t(split, dt)
Push x
- split :: time splitting object
- dt   :: time step
"""
function operator_t!(split::SplittingOperator, dt)

    # x_new = x_old + dt * v

    for i_part = 1:split.pg.n_particles
        split.pg.array[1, i_part] += dt * split.pg.array[3, i_part]
        split.pg.array[2, i_part] += dt * split.pg.array[4, i_part]

    end

end

function operator_v!(split::SplittingOperator, dt)

    # v_new = v_old + dt * q/m * E

    qm = split.pg.q_over_m

    for i_part = 1:split.pg.n_particles

        # Evaluate efields at particle position

        xi = split.pg.array[1, i_part]
        yi = split.pg.array[2, i_part]

        ex, ey = GEMPIC.evaluate_multiple(split.pic.kernel, (xi, yi), split.pic.efield_dofs)

        vx_new = split.pg.array[3, i_part]
        vy_new = split.pg.array[4, i_part]

        vx_new += dt * qm * ex
        vy_new += dt * qm * ey

        split.pg.array[3, i_part] = vx_new
        split.pg.array[4, i_part] = vy_new

    end

end

export charge_deposition!

function charge_deposition!(split::SplittingOperator)

    fill!(split.pic.rho_dofs, 0.0)
    for i_part = 1:split.pg.n_particles
        xp = split.pg.array[1, i_part]
        yp = split.pg.array[2, i_part]
        wp = split.pg.array[5, i_part]
        GEMPIC.add_charge!(split.pic.rho_dofs, split.pic.kernel, xp, yp, wp)
    end
    nx, ny = split.pic.poisson.grid.nx, split.pic.poisson.grid.ny
    split.pic.rho2d .= reshape(split.pic.rho_dofs, nx, ny)

end

export solve_fields!

"""
Solve Poisson's equation for the electric field
"""
solve_fields!(split::SplittingOperator) = solve_fields!(split.pic)

export strang_splitting!

"""
Strang splitting
- split :: time splitting object
- dt   :: time step
"""
function strang_splitting!(split::SplittingOperator, dt)

    operator_t!(split, 0.5dt)
    charge_deposition!(split)
    solve_fields!(split)
    operator_v!(split, dt)
    operator_t!(split, 0.5dt)

end

export compute_field_energy

"""
    compute_field_energy

Compute the squared l2 norm of electric field
"""
compute_field_energy(pic::TwoDPoissonPIC, component::Int) =
    sum(pic.efield_dofs[component] .^ 2) * pic.poisson.grid.dx * pic.poisson.grid.dy


using Plots
plotlyjs()

include("../julia/pprocess_inhomopoisson.jl");

f = 0.5;
lambda = 20*(1+sin(2.0 * pi * f * ts));
x_inhomopoisson, rs_inhomopoisson = pprocess_inhomopoisson(lambda);

plot(ts, rs_inhomopoisson)
scatter!(x_inhomopoisson, zeros(length(x_inhomopoisson)))
xlabel!("Time (s)")
ylabel!("Firing rate (Hz)")


using ContinuumArrays, LinearAlgebra, LazyArrays, Base64, FillArrays, QuasiArrays, BandedMatrices, Test
using QuasiArrays: ApplyQuasiArray, ApplyStyle, MemoryLayout, mul, MulQuasiMatrix, Vec
import LazyArrays: MulStyle, LdivStyle, arguments
import ContinuumArrays: basis, AdjointBasisLayout, ExpansionLayout, BasisLayout, SubBasisLayout, AdjointMappedBasisLayout, MappedBasisLayout, coefficients

@testset "Splines" begin
    @testset "HeavisideSpline" begin
        H = HeavisideSpline([1,2,3])

        @test axes(H) ≡ (axes(H,1),axes(H,2)) ≡ (Inclusion(1..3), Base.OneTo(2))
        @test size(H) ≡ (size(H,1),size(H,2)) ≡ (ℵ₁, 2)

        @test_throws BoundsError H[0.1, 1]
        @test H[1.1,1] ≡ H'[1,1.1] ≡ transpose(H)[1,1.1] ≡ 1.0
        @test H[2.1,1] ≡ H'[1,2.1] ≡ transpose(H)[1,2.1] ≡ 0.0
        @test H[1.1,2] ≡ 0.0
        @test H[2.1,2] ≡ 1.0
        @test_throws BoundsError H[2.1,3]
        @test_throws BoundsError H'[3,2.1]
        @test_throws BoundsError transpose(H)[3,2.1]
        @test_throws BoundsError H[3.1,2]

        @test all(H[[1.1,2.1], 1] .=== H'[1,[1.1,2.1]] .=== transpose(H)[1,[1.1,2.1]] .=== [1.0,0.0])
        @test all(H[1.1,1:2] .=== H[1.1,:] .=== [1.0,0.0])
        @test all(H[[1.1,2.1], 1:2] .=== [1.0 0.0; 0.0 1.0])

        @test_throws BoundsError H[[0.1,2.1], 1]
        @test MemoryLayout(typeof(H)) == BasisLayout()
        @test ApplyStyle(*, typeof(H), typeof([1,2])) isa MulStyle

        f = H*[1,2]
        @test f isa ApplyQuasiArray
        @test axes(f) == (Inclusion(1.0..3.0),)
        @test f[1.1] ≈ 1
        @test f[2.1] ≈ 2

        @test @inferred(H'H) == @inferred(materialize(applied(*,H',H))) == Eye(2)
        @test summary(f) == stringmime("text/plain", f) == "HeavisideSpline([1, 2, 3]) * [1, 2]"

        @testset "sum/cumsum" begin
            H = HeavisideSpline(range(0,1;length=1000));
            x = axes(H,1)
            @test sum(H/H \ exp.(x)) ≈ ℯ-1 atol=1E-5
            @test last(cumsum(H/H \ exp.(x))) ≈ sum(H/H\exp.(x))
        end
    end

    @testset "LinearSpline" begin
        @testset "Evaluation" begin
            L = LinearSpline([1,2,3])
            @test size(L) == (ℵ₁, 3)

            @test_throws BoundsError L[0.1, 1]
            @test L[1.1,1] == L'[1,1.1] == transpose(L)[1,1.1] ≈ 0.9
            @test L[2.1,1] === L'[1,2.1] === transpose(L)[1,2.1] === 0.0
            @test L[1.1,2] ≈ 0.1
            @test L[2.1,2] ≈ 0.9
            @test L[2.1,3] == L'[3,2.1] == transpose(L)[3,2.1] ≈ 0.1
            @test_throws BoundsError L[3.1,2]

            @test L[[1.1,2.1], 1] == L'[1,[1.1,2.1]] == transpose(L)[1,[1.1,2.1]] ≈ [0.9,0.0]
            @test L[1.1,1:2] ≈ [0.9,0.1]
            @test L[[1.1,2.1], 1:2] ≈ [0.9 0.1; 0.0 0.9]

            @test_throws BoundsError L[[0.1,2.1], 1]
        end

        @testset "Expansion" begin
            L = LinearSpline([1,2,3])
            f = L*[1,2,4]

            @test basis(f) == L
            @test coefficients(f) == [1,2,4]
            @test axes(f) == (Inclusion(1.0..3.0),)
            @test f[1.1] ≈ 1.1
            @test f[2.1] ≈ 2.2

            δ = DiracDelta(1.2,1..3)
            L = LinearSpline([1,2,3])
            @test @inferred(δ'L) ≈ [0.8, 0.2, 0.0]
            @test δ'f == mul(δ',f) == dot(δ,f) == f[1.2]

            @test @inferred(L'L) == SymTridiagonal([1/3,2/3,1/3], [1/6,1/6])

            @testset "Algebra" begin
                @test 2f == f*2 == 2 .* f == f .* 2
                @test 2\f == f/2 == 2 .\ f == f ./ 2
                @test sum(f) ≈ 4.5
                @test L \ BroadcastQuasiArray(*,2,f) ≈ [2,4,8]
            end
        end

        @testset "==" begin
            L = LinearSpline([1,2,3])
            H = HeavisideSpline([1,2,3])
            @test L == L
            @test L ≠ H
            H = HeavisideSpline([1,1.5,2.5,3])
            @test_throws ArgumentError L == H
        end

        @testset "Adjoint layout" begin
            L = LinearSpline([1,2,3])
            @test MemoryLayout(L') isa AdjointBasisLayout
            @test @inferred([3,4,5]'*L') isa ApplyQuasiArray
        end

        @testset "+" begin
            L = LinearSpline([1,2,3])
            b = @inferred(L*[3,4,5] + L*[1.,2,3])
            @test ApplyStyle(\, typeof(L), typeof(b)) == LdivStyle()
            @test (L\b) == [4,6,8]
            B = BroadcastQuasiArray(+, L, L)
            @test L\B == 2Eye(3)

            @test L*[3,4,5] + L*[1.,2,3] + L*[4,5,6] == @inferred(broadcast(+, L*[3,4,5], L*[1.,2,3], L*[4,5,6])) == L*[8,11,14]

            b = L*[3,4,5] - L*[1.,2,3]
            @test (L\b) == [2,2,2]
            B = BroadcastQuasiArray(-, L, L)
            @test L\B == 0Eye(3)

            @testset "sub" begin
                v = ApplyQuasiArray(*, L[:,2:end], [1,2])
                f = L * [1,2,3]
                @test v + f == f + v == L*[1,3,5]
                @test v + v == L*[0,2,4]
            end
        end
    end

    @testset "Algebra" begin
        L = LinearSpline([1,2,3])
        f = L*[1,2,4]
        g = L*[5,6,7]

        @test MemoryLayout(f) isa ExpansionLayout
        @test MemoryLayout(2f) isa ExpansionLayout
        @test MemoryLayout(f*2) isa ExpansionLayout
        @test MemoryLayout(2\f) isa ExpansionLayout
        @test MemoryLayout(f/2) isa ExpansionLayout
        @test MemoryLayout(f+g) isa ExpansionLayout
        @test MemoryLayout(f-g) isa ExpansionLayout
        @test f[1.2] == 1.2
        @test (2f)[1.2] == (f*2)[1.2] == 2.4
        @test (2\f)[1.2] == (f/2)[1.2] == 0.6
        @test (f+g)[1.2] ≈ f[1.2] + g[1.2]
        @test (f-g)[1.2] ≈ f[1.2] - g[1.2]
    end

    @testset "Derivative" begin
        L = LinearSpline([1,2,3])
        f = L*[1,2,4]

        D = Derivative(axes(L,1))
        @test copy(D) == D

        @test D*L isa MulQuasiMatrix
        @test length((D*L).args) == 2
        @test eltype(D*L) == Float64

        M = applied(*, (D*L).args..., [1,2,4])
        @test eltype(materialize(M)) == Float64

        M = applied(*, D, L, [1,2,4])
        @test materialize(M) isa ApplyQuasiArray

        fp = D*L*[1,2,4]

        @test eltype(fp) == Float64

        @test fp isa Vec
        @test length(fp.args) == 2
        @test fp[1.1] ≈ 1
        @test fp[2.2] ≈ 2

        fp = D*f
        @test length(fp.args) == 2
        @test fp[1.1] ≈ 1
        @test fp[2.2] ≈ 2

        @testset "View derivatives" begin
            L = LinearSpline(1:5)
            H = HeavisideSpline(1:5)
            x = axes(L,1)
            D = Derivative(x)

            @test view(D, x, x) ≡ D
            @test_throws BoundsError view(D, Inclusion(0..1), x)
            jr = [1,3,4]
            @test H \ (D*L[:,jr]) == H\(L[:,jr]'D')' == (H \ (D*L))[:,jr]

            @test D*L[:,jr]*[1,2,3] == D * L * [1,0,2,3,0]
            @test L[:,jr]'D'D*L[:,jr] == (L'D'D*L)[jr,jr]

            a = affine(0..1, 1..5)
            D̃ = Derivative(axes(a,1))
            @test H[a,:] \ (D̃ * L[a,:]) == H[a,:] \ (L[a,:]'D̃')' ==  4*(H\(D*L))
            @test_throws DimensionMismatch D * L[a,:]
            @test_throws DimensionMismatch D̃ * L[:,jr]

            @test ContinuumArrays.simplifiable(*, D, L[:,jr]) isa Val{true}
            @test ContinuumArrays.simplifiable(*, L[:,jr]', D') isa Val{true}
            @test ContinuumArrays.simplifiable(*, D̃, L[a,jr]) isa Val{true}
            @test ContinuumArrays.simplifiable(*, L[a,jr]', D̃') isa Val{true}
        end

        @testset "maps broadcasted" begin
            L = LinearSpline(1:5)
            a = affine(0..1, 1..5)
            M = L[a,:]
            M̃ = L[a,1:5]
            x = axes(M,1)

            @test M == M
            @test M == M̃
            @test M̃ == M

            @test (x .* M)[0.25,:] ≈ (x .* M̃)[0.25,:] ≈ 0.25 * M[0.25,:]
            @test (exp.(x) .* M)[0.25,:] ≈ exp(0.25) * M[0.25,:]

            @test M \ [ones(x) x] ≈ [M\ones(x) M\x]
        end
    end

    @testset "Weak Laplacian" begin
        H = HeavisideSpline(0:2)
        L = LinearSpline(0:2)
        D = Derivative(axes(L,1))

        @test apply(*,L',D') isa MulQuasiMatrix
        @test MemoryLayout(typeof(L')) isa AdjointBasisLayout
        @test (L'D') isa MulQuasiMatrix

        A = (L'D') * (D*L)
        @test A isa BandedMatrix
        @test A == (D*L)'*(D*L) == [1.0 -1 0; -1.0 2.0 -1.0; 0.0 -1.0 1.0]
        @test bandwidths(A) == (1,1)
    end

    @testset "Views" begin
        L = LinearSpline(0:2)
        @test view(L,0.1,1)[1] == L[0.1,1]

        L = LinearSpline(0:2)
        B1 = view(L,:,1)
        @test B1 isa SubQuasiArray{Float64,1}
        @test size(B1) == (ℵ₁,)
        @test B1[0.1] == L[0.1,1]
        @test_throws BoundsError B1[2.2]

        B = view(L,:,1:2)
        @test B isa SubQuasiArray{Float64,2}
        @test B[0.1,:] == L[0.1,1:2]

        B = @view L[:,2:end-1]
        @test B[0.1,:] == [0.1]

        L = LinearSpline([1,2,3,4])
        @test L[:,2:3] isa SubQuasiArray
        @test axes(L[:,2:3]) ≡ (Inclusion(1..4), Base.OneTo(2))
        @test L[:,2:3][1.1,1] == L[1.1,2]
        @test_throws BoundsError L[0.1,1]
        @test_throws BoundsError L[1.1,0]

        @test MemoryLayout(typeof(L[:,2:3])) isa SubBasisLayout
        @test L\L[:,2:3] isa BandedMatrix
        @test L\L[:,2:3] == [0 0; 1 0; 0 1.0; 0 0]
        @test L[:,2:3]\L isa BandedMatrix
        @test L[:,2:3]\L == [0 1 0 0; 0 0 1 0]

        @testset "Subindex of splines" begin
            L = LinearSpline(range(0,stop=1,length=10))
            @test L[:,2:end-1] isa SubQuasiArray
            @test L[:,2:end-1][0.1,1] == L[0.1,2]
            v = randn(8)
            f = L[:,2:end-1] * v
            @test f[0.1] ≈ (L*[0; v; 0])[0.1]
        end

        @testset "sub-of-sub" begin
            L = LinearSpline([1,2,3])
            V = view(L,:,1:2)
            @test V == V
            V2 = view(V,1.1:0.1:2,:)
            @test V2 == L[1.1:0.1:2,1:2]
        end

        @testset "sub-colon" begin
            L = LinearSpline([1,2,3])
            @test L[:,1][1.1] == (L')[1,:][1.1] == L[1.1,1]
            @test L[:,1:2][1.1,:] == (L')[1:2,:][:,1.1] == L[1.1,1:2]
        end

        @testset "transform" begin
            L = LinearSpline([1,2,3])
            x = axes(L,1)
            @test (L \ x) == pinv(L)x == [1,2,3]
            @test factorize(L[:,2:end-1]) isa ContinuumArrays.ProjectionFactorization
            @test L[:,1:2] \ x == [1,2]
            @test L \ [x one(x)] ≈ [L\x L\one(x)]
            @test factorize(L) \ QuasiOnes(x, Base.OneTo(3)) ≈ L \ QuasiOnes(x, Base.OneTo(3)) ≈ ones(3,3)

            L = LinearSpline(range(0,1; length=10_000))
            x = axes(L,1)
            @test L[0.123,:]'* (L \ exp.(x)) ≈ exp(0.123) atol=1E-9
            @test L[0.123,2:end-1]'* (L[:,2:end-1] \ exp.(x)) ≈ exp(0.123) atol=1E-9

            @test L \ zeros(x) ≡ Zeros(10_000)
        end
    end

    @testset "sum" begin
        H = HeavisideSpline([1,2,3,6])
        @test sum(H; dims=1) * [1,1,1] == [5]
        x = Inclusion(0..1)
        B = H[5x .+ 1,:]
        @test sum(B; dims=1) * [1,1,1] == [1]
        @test sum(H[:,1:2]; dims=1) * [1,1] == [2]
        @test sum(H'; dims=2) == permutedims(sum(H; dims=1))

        u = H * randn(3)
        @test sum(u[5x .+ 1]) ≈ sum(view(u,5x .+ 1)) ≈ sum(u)/5

        L = LinearSpline([1,2,3,6])
        D = Derivative(axes(L,1))
        @test sum(D*L; dims=1) ≈ sum((D*L)'; dims=2)' ≈ [-1 zeros(1,2) 1]
    end

    @testset "Poisson" begin
        L = LinearSpline(range(0,stop=1,length=10))
        B = L[:,2:end-1] # Zero dirichlet by dropping first and last spline
        @test B'B == (L'L)[2:end-1,2:end-1]
        D = Derivative(axes(L,1))
        @test apply(*,D,B) isa MulQuasiMatrix
        @test D*B isa MulQuasiMatrix
        @test apply(*,D,B)[0.1,1] == (D*B)[0.1,1] == 9
        @test apply(*,D,B)[0.2,1] == (D*B)[0.2,1] == -9
        Δ = -((D*B)'*(D*B)) # Weak Laplacian
        @test Δ isa BandedMatrix

        @test B'D' isa MulQuasiMatrix
        @test length((B'D').args) == 2

        @test *(B',D',D,B) isa BandedMatrix

        @test Δ == -(*(B',D',D,B))
        @test Δ == -(B'D'D*B)
        @test Δ == -((B'D')*(D*B))
        @test_broken Δ == -B'*(D'D)*B
        @test Δ == -(B'*(D'D)*B)

        f = L*exp.(L.points) # project exp(x)
        u = B * (Δ \ (B'f))

        @test u[0.1] ≈ -0.06612902692412974
    end

    @testset "Helmholtz" begin
        L = LinearSpline(range(0,stop=1,length=10))
        B = L[:,2:end-1] # Zero dirichlet by dropping first and last spline
        D = Derivative(axes(L,1))

        A = -((B'D')*(D*B)) + 100^2*B'B # Weak Laplacian

        f = L*exp.(L.points) # project exp(x)
        u = B * (A \ (B'f))

        @test u[0.1] ≈ 0.00012678835289369413
    end

    @testset "Change-of-variables" begin
        x = Inclusion(0..1)
        y = 2x .- 1
        L = LinearSpline(range(-1,stop=1,length=10))
        @test L[y,:][0.1,:] == L[2*0.1-1,:]
        @test L[y,:]'L[y,:] isa SymTridiagonal
        @test L[y,:]'L[y,:] == 1/2*(L'L)

        D = Derivative(axes(L,1))
        H = HeavisideSpline(L.points)
        @test H\((D*L) * 2) ≈ (H\(D*L))*2 ≈ diagm(0 => fill(-9,9), 1 => fill(9,9))[1:end-1,:]

        @test MemoryLayout(L[y,:]) isa MappedBasisLayout
        @test MemoryLayout(L[y,:]') isa AdjointMappedBasisLayout
        a,b = arguments((D*L)[y,:])
        @test H[y,:]\a == Eye(9)
        @test H[y,:] \ (D*L)[y,:] isa BandedMatrix
        @test @inferred(grid(L[y,:])) ≈ (grid(L) .+ 1) ./ 2

        D = Derivative(x)
        @test (D*L[y,:])[0.1,1] ≈ -9
        @test H[y,:] \ (D*L[y,:]) isa BandedMatrix
        @test H[y,:] \ (D*L[y,:]) ≈ diagm(0 => fill(-9,9), 1 => fill(9,9))[1:end-1,:]

        B = L[y,2:end-1]
        @test MemoryLayout(typeof(B)) isa MappedBasisLayout
        @test B[0.1,1] == L[2*0.1-1,2]
        @test B\B == Eye(8)
        @test L[y,:] \ B == Eye(10)[:,2:end-1]
        @test B'B == (1/2)*(L'L)[2:end-1,2:end-1]

        @testset "algebra" begin
            f = L[y,:] * randn(10)
            g = L[y,:] * randn(10)
            @test MemoryLayout(f + g) isa ExpansionLayout
            @test (f+g)[0.1] ≈ f[0.1] + g[0.1]
        end

        @testset "vec demap" begin
            @test L[y,:] \ exp.(axes(L,1))[y] ≈ L[y,:] \ exp.(y) ≈  factorize(L[y,:]) \ exp.(y)
            @test ContinuumArrays.demap(view(axes(L,1),y)) == axes(L,1)

            @test L[y,:] \ (y .* exp.(y)) ≈ L[y,:] \ BroadcastQuasiVector(y -> y*exp(y), y)
            @test L[y,:] \ (y .* L[y,1:3]) ≈ [L[y,:]\(y .* L[y,1]) L[y,:]\(y .* L[y,2]) L[y,:]\(y .* L[y,3])]

            c = randn(size(L,2))
            @test L[y,:] \ (L[y,:] * c) ≈ c
            @test ContinuumArrays.demap(L[y,:] * c) == L*c
        end

        @testset "Mapped and BroadcastLayout{typeof(+)}" begin
            @test L[y,:] \ (y .+ y) ≈ L[y,:] \ (2y)
            @test L[y,:] \ (y .- y) ≈ zeros(10)
        end
    end

    @testset "diff" begin
        L = LinearSpline(range(-1,stop=1,length=10))
        f = L * randn(size(L,2))
        h = 0.0001;
        @test diff(f)[0.1] ≈ (f[0.1+h]-f[0.1])/h
    end

    @testset "show" begin
        x = Inclusion(0..1)
        H = HeavisideSpline([1,2,3,6])
        B = H[5x .+ 1,:]
        u = H * [1,2,3]
        @test stringmime("text/plain", B) == "HeavisideSpline([1, 2, 3, 6]) affine mapped to 0..1"
    end

    @testset "A \\ ( c .* B) == c .* (A\\B) #101" begin
        L = LinearSpline(0:5)
        @test L \ (2L) == 2(L\L)
    end

    @testset "H\\D*L" begin
        L = LinearSpline(0:5)
        H = HeavisideSpline(L)
        x = axes(L,1)
        D = Derivative(x)
        @test H\D*L == H\(D*L)
    end
end


"""
    eom_logistic3(u, p, t)

Equations of motion for a system consisting of three coupled logistic map
representing the response of two independent dynamical variables to the
forcing from a common driver. The dynamical influence goes in the directions
Z -> X and Z -> Y.

Dynamical noise may be added to each of the dynamical variables by tuning the
parameters `σz`, `σx` and `σz`. Default values for the parameter s
`r₁` and `r₂` are set such that the system exhibits chaotic behaviour.

# References
1. Runge, Jakob. Causal network reconstruction from time series: From
theoretical assumptions to practical estimation, Chaos 28, 075310 (2018);
doi: 10.1063/1.5025050
"""
function eom_logistic3(u, p, t)
    r₁, r₂, r₃, σx, σy, σz = (p...,)
    x, y, z = (u...,)

    # Independent dynamical noise for each variable.
    ηx = rand()
    ηy = rand()
    ηz = rand()

    dx = (x*(r₁ - r₁*x - z + σx*ηx)) % 1
    dy = (y*(r₂ - r₂*y - z + σy*ηy)) % 1
    dz = (z*(r₃ - r₃*z + σz*ηz)) % 1
    return SVector{3}(dx, dy, dz)
end

function logistic3(u₀, r₁, r₂, r₃, σx, σy, σz)
    p = [r₁, r₂, r₃, σx, σy, σz]
    DiscreteDynamicalSystem(eom_logistic3, u₀, p)
end

"""
    logistic3(;u₀ = rand(3), r = 4,
        σx = 0.05, σy = 0.05, σz = 0.05) -> DiscreteDynamicalSystem

Initialise a dynamical system consisting of three coupled logistic map
representing the response of two independent dynamical variables to the
forcing from a common driver. The dynamical influence goes in the directions
Z -> X and Z -> Y.

The equations of motion are

```math
dx = (x(r - r_1 x - z + σ_x η_x)) \\mod 1
dy = (y(r - r_2 y - z + σ_y η_y)) \\mod 1
dz = (z(r - r_3 z + σ_z η_z)) \\mod 1
```

Dynamical noise may be added to each of the dynamical variables by tuning the
parameters `σz`, `σx` and `σz`. Default values for the parameters
`r₁`, `r₂` and `r₃` are set such that the system exhibits chaotic behaviour,
with `r₁ = r₂ = r₃ = 4`.

# References
1. Runge, Jakob. Causal network reconstruction from time series: From theoretical assumptions to practical estimation, Chaos 28, 075310 (2018); doi: 10.1063/1.5025050
"""
logistic3(;u₀ = rand(3), r₁ = 4, r₂ = 4, r₃ = 4,
    σx = 0.05, σy = 0.05, σz = 0.05) = logistic3(u₀, r₁, r₂, r₃, σx, σy, σz)


# ---
# title: 8. String to Integer (atoi)
# id: problem8
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, String
# link: <https://leetcode.com/problems/string-to-integer-atoi/description/>
# hidden: true
# ---
#
# Implement `atoi` which converts a string to an integer.
#
# The function first discards as many whitespace characters as necessary until
# the first non-whitespace character is found. Then, starting from this
# character takes an optional initial plus or minus sign followed by as many
# numerical digits as possible, and interprets them as a numerical value.
#
# The string can contain additional characters after those that form the
# integral number, which are ignored and have no effect on the behavior of this
# function.
#
# If the first sequence of non-whitespace characters in str is not a valid
# integral number, or if no such sequence exists because either str is empty or
# it contains only whitespace characters, no conversion is performed.
#
# If no valid conversion could be performed, a zero value is returned.
#
# **Note:**
#
#   * Only the space character `' '` is considered a whitespace character.
#   * Assume we are dealing with an environment that could only store integers within the 32-bit signed integer range: `[−231,  231 − 1]`. If the numerical value is out of the range of representable values, `231 − 1` or `−231` is returned.
#
#
#
# **Example 1:**
#
#
#
#     Input: str = "42"
#     Output: 42
#
#
# **Example 2:**
#
#
#
#     Input: str = "   -42"
#     Output: -42
#     Explanation: The first non-whitespace character is '-', which is the minus sign. Then take as many numerical digits as possible, which gets 42.
#
#
# **Example 3:**
#
#
#
#     Input: str = "4193 with words"
#     Output: 4193
#     Explanation: Conversion stops at digit '3' as the next character is not a numerical digit.
#
#
# **Example 4:**
#
#
#
#     Input: str = "words and 987"
#     Output: 0
#     Explanation: The first non-whitespace character is 'w', which is not a numerical digit or a +/- sign. Therefore no valid conversion could be performed.
#
#
# **Example 5:**
#
#
#
#     Input: str = "-91283472332"
#     Output: -2147483648
#     Explanation: The number "-91283472332" is out of the range of a 32-bit signed integer. Thefore INT_MIN (−231) is returned.
#
#
#
#
# **Constraints:**
#
#   * `0 <= s.length <= 200`
#   * `s` consists of English letters (lower-case and upper-case), digits, `' '`, `'+'`, `'-'` and `'.'`.
#
#
## @lc code=start
using LeetCode

function my_atoi(s::String)::Int
    s = lstrip(s)
    if length(s) == 0
        0
    else
        try
            m = match(r"^[\-\+]?[0-9]+", s)
            if isnothing(m)
                0
            else
                parse(Int32, m.match)
            end
        catch
            if s[1] == '-'
                typemin(Int32)
            else
                typemax(Int32)
            end
        end
    end
end
## @lc code=end


struct NLADefault{T<:AbstractFloat} <: NonLinearAlgorithm end

"""
    NLADefault()
Return the array untouched, default option.
"""
NLADefault() = NLADefault{Float64}()

nla(::NLADefault{T}, x::AbstractArray{T}) where T<: AbstractFloat = x

struct NLAT1{T<:AbstractFloat} <: NonLinearAlgorithm end
"""
    NLAT1()
Apply the \$ \\text{T}_1 \$ transformation algorithm, as defined in [1] and [2]

[1] Chattopadhyay, Ashesh, et al. "Data-driven prediction of a multi-scale Lorenz 96 chaotic system using a hierarchy of deep learning methods: Reservoir computing, ANN, and RNN-LSTM." (2019).

[2] Pathak, Jaideep, et al. "Model-free prediction of large spatiotemporally chaotic systems from data: A reservoir computing approach." Physical review letters 120.2 (2018): 024102.
"""
NLAT1() = NLAT1{Float64}()
nla(::NLAT1{T}, x::AbstractArray{T}) where T<: AbstractFloat = _nlat1(x)

struct NLAT2{T<:AbstractFloat} <: NonLinearAlgorithm end
"""
    NLAT2()
Apply the \$ \\text{T}_2 \$ transformation algorithm, as defined in [1]

[1] Chattopadhyay, Ashesh, et al. "Data-driven prediction of a multi-scale Lorenz 96 chaotic system using a hierarchy of deep learning methods: Reservoir computing, ANN, and RNN-LSTM." (2019).
"""
NLAT2() = NLAT2{Float64}()
nla(::NLAT2{T}, x::AbstractArray{T}) where T<: AbstractFloat = _nlat2(x)

struct NLAT3{T<:AbstractFloat} <: NonLinearAlgorithm end
"""
    NLAT3()
Apply the \$ \\text{T}_3 \$ transformation algorithm, as defined in [1]

[1] Chattopadhyay, Ashesh, et al. "Data-driven prediction of a multi-scale Lorenz 96 chaotic system using a hierarchy of deep learning methods: Reservoir computing, ANN, and RNN-LSTM." (2019).
"""
NLAT3() = NLAT3{Float64}()
nla(::NLAT3{T}, x::AbstractArray{T}) where T<: AbstractFloat = _nlat3(x)

function _nlat1(x_old::AbstractArray{T}) where T<: AbstractFloat
    x_new = copy(x_old)
    for i =1:size(x_new, 1)
        if mod(i, 2)!=0
            x_new[i, :] = copy(x_old[i,:].*x_old[i,:])
        end
    end
    return x_new
end

function _nlat2(x_old::AbstractArray{T}) where T<: AbstractFloat
    x_new = copy(x_old)
    for i =2:size(x_new, 1)-1
        if mod(i, 2)!=0
            x_new[i, :] = copy(x_old[i-1,:].*x_old[i-2,:])
        end
    end
    return x_new
end

function _nlat3(x_old::Array{T}) where T<: AbstractFloat
    x_new = copy(x_old)
    for i =2:size(x_new, 1)-1
        if mod(i, 2)!=0
            x_new[i, :] = copy(x_old[i-1,:].*x_old[i+1,:])
        end
    end
    return x_new
end


#=
# 7. Lid-driven cavity flow
In this notebook we will simulate the flow in a square cavity with a top moving wall.
=#
using ViscousFlow
#-
using Plots

#=
## Problem specification
Take $Re=100$ for example:
=#
Re = 100

#=
## Discretization
Note that the rectangle function used for making the cavity shape requires a specified half length. The
immersed boundary projection method for internal flow requires the size of the domain to
be at least a step size greater at the boundaries (i.e. `halflength + Δx`).
=#
Δt,Δx = setstepsizes(Re,gridRe=1.0)
halflength=0.5
domain_lim=halflength+1.01*Δx
xlim, ylim = (-domain_lim,domain_lim),(-domain_lim,domain_lim)

#=
## Cavity Geometry
A square cavity can be created using the `Rectangle()` function with the half length defined above.
The `shifted=true` argument ensures that points are not placed at the corners, where
they have ill-defined normal vectors.
=#
body = Rectangle(halflength,halflength,1.5*Δx,shifted=true)
plot(body,fillrange=nothing)

#=
## Boundary Condition at the moving wall
Assign velocity to the top boundary. The `LidDrivenCavity()` function can be used
to specify the velocity value at the top wall. Note : Non-dimensional velocity = 1
=#
m = ViscousFlow.LidDrivenCavity(1.0)

#=
## Construct the system structure
The last two input `flow_side` and `static_points` must specified so the default setting in the
 `NavierStokes()` function can be overwritten.

`static_points` is set to true because the cavity wall points are not actually moving.
=#
sys = NavierStokes(Re,Δx,xlim,ylim,Δt,body,m,flow_side = InternalFlow,static_points = true)

#=
Initialize
=#
u0 = newstate(sys)

#=
Set up integrator
=#
tspan = (0.0,10.0)
integrator = init(u0,tspan,sys)

#=
## Solve
=#
step!(integrator,10)

#=
## Examine
Plot the vorticity and streamlines
=#

plot(
plot(vorticity(integrator),sys,title="Vorticity",clim=(-10,10),color=:turbo,linewidth=1.5,ylim=ylim,levels=-6:0.5:5),
plot(streamfunction(integrator),sys,title="Streamfunction",color=:black,levels=vcat(0.009:0.01:0.11,0.1145,0.11468,0.11477),ylim=ylim)
   )

#=
Make a movie:
=#
sol = integrator.sol;
@gif for (u,t) in zip(sol.u,sol.t)
    plot(vorticity(u,sys,t),sys,clim=(-10,10),levels=range(-10,10,length=30),color=:turbo)
end every 5


# Implements a simple class for propagating uncertainties using
# the formalism in ISO GUM (JCGM 100 and JCGM 102).
using Printf
using LinearAlgebra
using DataFrames
using Formatting

"""
    checkcovariance!(cov::AbstractMatrix)

Checks whether the matrix meets the criteria for a covariance matrix.
(Square, non-negative diagonal elements, symmetric cov[r,c]==cov[c,r], and the
correlation coefficient between -1.0 and 1.0).  The tests are approximate
to handle rounding errors but when a symmetric pair of values are not precisely
equal they are set equal and when the correlation coefficient is slightly outside
[-1,1], it is restricted to within these bounds. Thus the input matrix can be
modified, in ways that are probably benign, by this "check" function.
"""
function checkcovariance!(cov::AbstractMatrix{Float64}, tol = 1.0e-6)::Bool
    sz = size(cov)
    if length(sz) ≠ 2
        error("The covariance must be a matrix.")
    end
    if sz[1] ≠ sz[2]
        error("The covariance matrix must be square.")
    end
    for rc = 1:sz[1]
        if cov[rc, rc] < 0.0
            error("The diagonal elements must all be non-negative. -> ", cov[rc, rc])
        end
    end
    for r = 1:sz[1]
        for c = 1:r-1
            if !isapprox(cov[r, c], cov[c, r], atol = tol * sqrt(cov[r, r] * cov[c, c]))
                error("The variances must be symmetric. -> ", cov[r, c], " ≠ ", cov[c, r])
            end
            cov[c, r] = cov[r, c] # Now force it precisely...
        end
    end
    for r = 1:sz[1]
        for c = 1:r-1
            cc = cov[r, c] / sqrt(cov[r, r] * cov[c, c])
            if abs(cc) > 1.0 + 1.0e-12
                error(
                    "The variances must have a correlation coefficient between -1.0 and 1.0 -> ",
                    cc,
                )
            end
            if abs(cc) > 1.0
                cc = max(-1.0, min(1.0, cc))
                cov[r, c] = (cov[c, r] = cc * sqrt(cov[r, r] * cov[c, c]))
            end
        end
    end
    true
end

"""
    UncertainValues

Represents a set of related variables with the covariance matrix that represents
the uncertainty relationships between the variables.
"""
struct UncertainValues
    labels::Dict{<:Label,Int}
    values::AbstractVector{Float64}
    covariance::AbstractMatrix{Float64}
    UncertainValues(
        labels::Dict{<:Label,Int},
        values::AbstractVector{Float64},
        covar::AbstractMatrix{Float64},
    ) = checkUVS!(labels, values, covar) ? new(labels, values, covar) : error("???")
end
"""
    uvs(labels::AbstractVector{<:Label}, values::AbstractVector{Float64}, covar::AbstractMatrix{Float64})
    uvs(labels::AbstractVector{<:Label}, values::AbstractVector{Float64}, vars::AbstractVector{Float64})
    uvs(values::Pair{<:Label,UncertainValue}...)
    uvs(value::Pair{<:Label,UncertainValue})
    uvs(values::Dict{<:Label,UncertainValue})

Various methods for constructing `UncertainValues` structures from varying types of inputs.
"""
function uvs(
    labels::AbstractVector{<:Label},
    values::AbstractVector{Float64},
    covar::AbstractMatrix{Float64},
)
    @assert length(labels) == length(values)
    @assert length(labels) == size(covar, 1)
    @assert length(labels) == size(covar, 2)
    return UncertainValues(
        Dict{Label,Int}(l => i for (i, l) in enumerate(labels)),
        values,
        covar,
    )
end

function uvs(
    labels::AbstractVector{<:Label},
    values::AbstractVector{Float64},
    vars::AbstractVector{Float64},
)
    @assert length(labels) == length(values)
    @assert length(labels) == length(vars)
    return uvs(labels, values, diagm(vars))
end

uvs(values::Pair{<:Label,UncertainValue}...) = uvs(Dict(values))
uvs(value::Pair{<:Label,UncertainValue}) = uvs(Dict(value))
covariance(uvs::UncertainValues) = uvs.covariance

function uvs(values::Dict{<:Label,UncertainValue})
    labels, vals, vars = Vector{Label}(), Vector{Float64}(), Vector{Float64}()
    for (lbl, uv) in values
        push!(labels, lbl)
        push!(vals, value(uv))
        push!(vars, variance(uv))
    end
    return uvs(labels, vals, diagm(vars))
end

"""
    estimated(labels::Vector{<:Label}, samples::Matrix{Float64})

Estimate an UncertainValues based on a set of samples ('measurements').
Computes the mean values and the covariance matrix from the expectation values.
Each row `r` in samples represents a measured quantity as identified by `labels[r]`.
Each column represents a single set of measurements of all the labeled quantities.
"""
function estimated(labels::Vector{<:Label}, samples::Matrix{Float64})::UncertainValues
    @assert length(labels) == size(samples, 1) "label length must equal row count in estimated"
    nvars, nsamples = size(samples, 1), size(samples, 2)
    μ = mean.(samples[k, :] for k = 1:nvars)
    n = collect(samples[k, :] .- μ[k] for k = 1:nvars)
    Σ = Matrix{Float64}(undef, nvars, nvars)
    for j = 1:nvars, k = 1:j
        Σ[j, k] = (Σ[k, j] = dot(n[j], n[k]) / nsamples)
    end
    return uvs(labels, μ, Σ)
end

function checkUVS!(
    labels::Dict{<:Label,Int},
    values::AbstractVector{Float64},
    covar::AbstractMatrix{Float64},
)
    if length(labels) ≠ length(values)
        error("The number of labels does not match the number of values.")
    end
    if length(labels) ≠ size(covar)[1]
        error("The number of labels does not match the dimension of the covariance matrix.")
    end
    checkcovariance!(covar)
end

Base.haskey(uvs::UncertainValues, lbl::Label) = haskey(uvs.labels, lbl)

"""
    σ(lbl::Label, uvs::UncertainValues)

Returns the 1σ uncertainty associated with the specified label
"""
σ(uvs::UncertainValues, lbl::Label) = sqrt(variance(uvs, lbl))
σ(uvs::UncertainValues, lbl::Label, def) =
    haskey(uvs.labels, lbl) ? sqrt(variance(uvs, lbl)) : def

"""
    correlation(uvs::UncertainValues, a::Label, b::Label)

Returns the Pearson correlation coefficient between variables `a` and `b`.
"""
correlation(uvs::UncertainValues, a::Label, b::Label) =
    covariance(uvs, a, b) / (σ(uvs, a) * σ(uvs, b))

"""
    Base.filter(uvs::UncertainValues, labels::Vector{<:Label})::Matrix

Extract the covariance matrix associated with the variables specified in labels
into a Matrix.
"""
function Base.filter(uvs::UncertainValues, labels::Vector{<:Label})::Matrix
    idx = map(l -> uvs.labels[l], labels) # look it up once...
    m = zeros(length(idx), length(idx))
    for (r, rl) in enumerate(idx)
        for (c, cl) in enumerate(idx)
            m[c, r] = (m[r, c] = uvs.covariance[rl, cl])
        end
    end
    return m
end

Base.:*(aa::AbstractMatrix{Float64}, uvs::UncertainValues) =
    UncertainValues(uvs.labels, aa * uvs.values, aa * uvs.covariance * transpose(aa))

Base.:*(aa::Diagonal{Float64}, uvs::UncertainValues) =
    UncertainValues(uvs.labels, aa * uvs.values, aa * uvs.covariance * aa)

"""
    cat(uvss::AbstractArray{UncertainValues})::UncertainValues

Combines the disjoint UncertainValues in uvss into a single UncertainValues object.
"""
Base.cat(uvss::AbstractArray{UncertainValues})::UncertainValues = cat(uvss...)

"""
    cat(uvss::UncertainValues...)::UncertainValues

Combines the disjoint UncertainValues in uvss into a single UncertainValues object.
"""
function Base.cat(uvss::UncertainValues...)::UncertainValues
    function combinelabels(us::UncertainValues...) # Maintains order as union doesn't
        res = Label[]
        for u in us
            for lu in labels(u)
                if !(lu in res)
                    push!(res, lu)
                else
                    throw(
                        ErrorException(
                            "Unable to combine UncertainValues with duplicate labels.",
                        ),
                    )
                end
            end
        end
        return res
    end
    all = Dict{Label,Int}(lbl => i for (i, lbl) in enumerate(combinelabels(uvss...)))
    # @assert length(all) == sum(map(uvs -> length(uvs.labels), uvss)) "One or more labels were duplicated in cat(...)"
    len = length(all)
    values, covar = zeros(Float64, len), zeros(Float64, len, len)
    # Precompute the indexes for speed ( index in all, index in uvs)
    for uvs in uvss
        # Map  label indices into new label indices
        idx = collect((all[lbl], rc) for (lbl, rc) in uvs.labels)
        for (newR, oldR) in idx
            values[newR] = uvs.values[oldR]
            for (newC, oldC) in idx
                covar[newR, newC] = uvs.covariance[oldR, oldC]
            end
        end
    end
    return UncertainValues(all, values, covar)
end

Base.show(io::IO, uvs::UncertainValues) = print(
    io,
    "UVS[" *
    join(
        (
            @sprintf("%s = %-0.3g ± %-0.3g", lbl, value(uvs, lbl), σ(uvs, lbl)) for
            lbl in labels(uvs)
        ),
        ", ",
    ) *
    "]",
)

function Base.show(io::IO, ::MIME"text/plain", uvs::UncertainValues)
    trim(str, len) = str[1:min(len, length(str))] * " "^max(0, len - min(len, length(str)))
    lbls = sortedlabels(uvs)
    print(io, "Variable       Value              ")
    foreach(l -> print(io, trim(repr(l), 12)), lbls)
    for (r, rl) in enumerate(lbls)
        println(io)
        print(io, trim("$rl", 10) * @sprintf(" | %-8.3g |", value(uvs, rl)))
        print(io, r == length(uvs.labels)[1] / 2 ? "  ±  |" : "     |")
        foreach(cl -> print(io, @sprintf("   %-8.3g ", covariance(uvs, rl, cl))), lbls)
        print(io, " |")
    end
end

Base.show(io::IO, m::MIME"text/html", uvs::UncertainValues) =
# We piggyback off DataFrames here because generating a good looking table isn't easy...
    show(io, m, asa(DataFrame, uvs), summary = false, eltypes = false)

function Base.show(io::IO, ::MIME"text/markdown", uvs::UncertainValues)
    fmt(x) = @sprintf("%0.2e", x)
    cv(r, c) = r == c ? "($(fmt(σ(uvs, r))))²" : "$(fmt(covariance(uvs,r,c)))"
    esc(ss) = replace(ss, "|" => "\\|")
    ext(ss, l) = ss * repeat(" ", l - length(ss))
    lbls, rows = labels(uvs), []
    push!(rows, ["Labels", "Values", "", repr.(lbls)...])
    for lbl in lbls
        push!(
            rows,
            [repr(lbl), fmt(value(uvs, lbl)), " ", (cv(lbl, col) for col in lbls)...],
        )
    end
    cl = maximum(length(rows[j][i]) for j in eachindex(rows) for i in eachindex(rows[j]))
    insert!(
        rows,
        2,
        [
            ":$(repeat("-",cl-2))-",
            ":$(repeat("-",cl-2)):",
            ":$(repeat("-",cl-2)):",
            (":$(repeat("-",cl-2)):" for l in lbls)...,
        ],
    )
    rows[(length(rows)-3)÷2+3][3] = ext("±", cl)
    ss = join(
        (
            "| " * join((ext(rows[r][c], cl) for c in eachindex(rows[r])), " | ") * " |" for r in eachindex(rows)
        ),
        "\n",
    )
    print(io, ss)
end

"""
    sortedlabels(uvs::UncertainValues)

A alphabetically sorted list of the labels. Warning this can be slow.  Use keys(...) if you
want just a unordered set of labels.
"""
sortedlabels(uvs::UncertainValues) =
    sort([keys(uvs.labels)...], lt = (l, m) -> isless(repr(l), repr(m)))

Base.keys(uvs::UncertainValues) = keys(uvs.labels)

function Base.getindex(uvs::UncertainValues, lbl::Label)::UncertainValue
    idx = uvs.labels[lbl]
    return UncertainValue(uvs.values[idx], sqrt(uvs.covariance[idx, idx]))
end

function Base.isnan(uvs::UncertainValues, lbl::Label)::Bool
    idx = uvs.labels[lbl]
    return isnan(uvs.values[idx]) || isnan(uvs.covariance[idx,idx])
end

function Base.get(
    uvs::UncertainValues,
    lbl::Label,
    def::Union{Missing,UncertainValue},
)::Union{Missing,UncertainValue}
    idx = get(uvs.labels, lbl, -1)
    return idx ≠ -1 ? UncertainValue(uvs.values[idx], sqrt(uvs.covariance[idx, idx])) : def
end

Base.length(uvs::UncertainValues) = length(uvs.labels)

eachlabel(uvs::UncertainValues) = keys(uvs.labels)

"""
    labels(uvs::UncertainValues)::Vector{<:Label}

Returns a Vector of Label in the order in which they appear in `uvs.values` and
`uvs.covariance`.
"""
function labels(uvs::UncertainValues)::Vector{<:Label}
    res = Array{Label}(undef, length(uvs.labels))
    for (k, v) in uvs.labels
        res[v] = k
    end
    return res
end

"""
    value(uvs::UncertainValues, lbl::Label)
    value(uvs::UncertainValues, lbl::Label, defValue)

The value associate with the Label. The first implementation throws an exception if `lbl` is not present
while the second implementation returns `defValue`
"""
value(uvs::UncertainValues, lbl::Label) = uvs.values[uvs.labels[lbl]]
value(uvs::UncertainValues, lbl::Label, defValue) =
    haskey(uvs.labels, lbl) ? uvs.values[uvs.labels[lbl]] : defValue

"""
    values(uvs::UncertainValues)::Vector{Float64}

A Vector containing the Float64 value for each row in uvs.  In the same order as
labels(uvs).
"""
Base.values(uvs::UncertainValues) = uvs.values

"""
   covariance(uvs::UncertainValues, lbl1::Label, lbl2::Label)

The covariance between the two variables.
"""
covariance(uvs::UncertainValues, lbl1::Label, lbl2::Label) =
    uvs.covariance[uvs.labels[lbl1], uvs.labels[lbl2]]

covariance(uvs::UncertainValues, lbl1::Label, lbl2::Label, default) =
    haskey(uvs.labels, lbl1) && haskey(uvs.labels, lbl2) ? #
    uvs.covariance[uvs.labels[lbl1], uvs.labels[lbl2]] :
    default

"""
   variance(uvs::UncertainValues, lbl::Label)
   variance(uvs::UncertainValues, lbl::Label, default)

The variance associated with the specified Label.
"""
variance(uvs::UncertainValues, lbl::Label) =
    max(0.0, uvs.covariance[uvs.labels[lbl], uvs.labels[lbl]])
variance(uvs::UncertainValues, lbl::Label, default) =
    haskey(uvs.labels, lbl) ? uvs.covariance[uvs.labels[lbl], uvs.labels[lbl]] : default

function asa( #
    ::Type{DataFrame},
    uvss::UncertainValues,
    withCovars = true,
)::DataFrame
    lbls = labels(uvss)
    df = DataFrame(Variable = map(lbl -> "$lbl", lbls), Values = map(lbl -> value(uvss, lbl), lbls))
    if withCovars
        for (i, cl) in enumerate(lbls)
            insertcols!(
                df,
                2 + i,
                Symbol("$cl") => map(rl -> covariance(uvss, rl, cl), lbls),
            )
        end
    else
        insertcols!(df, 3, :σ => map(lbl -> σ(uvss, lbl), lbls))
    end
    return df
end

"""
    uncertainty(lbl::Label, uvs::UncertainValues, k::Float64=1.0)

The uncertainty associated with specified label (k σ where default k=1)
"""
uncertainty(uvs::UncertainValues, lbl::Label, k::Float64 = 1.0) = k * σ(uvs, lbl)

Base.indexin(uvs::UncertainValues, lbl::Label) = uvs.labels[lbl]

"""
    labeledvalues(uvs::UncertainValues)

Converts the values of an UncertainValues object into a LabeledValues object.
"""
labeledvalues(uvs::UncertainValues) = LabeledValues(labels(uvs), uvs.values)

labelsByType(ty::Type{<:Label}, uvs::UncertainValues) =
    filter(lbl -> typeof(lbl) == ty, labels(uvs))

function labelsByType(types::AbstractVector{DataType}, uvs::UncertainValues)
    lbls = labels(uvs)
    mapreduce(ty -> labelsByType(ty, lbls), append!, types)
end

"""
    extract(uvss::UncertainValues, labels::AbstractVector{<:Label})::UncertainValues
    extract(uvss::UncertainValues, labeltype::Type{T})::UncertainValues where {T<:Label}
    extract(uvss::UncertainValues, labeltypes::AbstractVector{DataType})::UncertainValues

Construct an `UncertainValues` datum containing only the specified labels/labeltype(s).
"""

function extract(uvss::UncertainValues, labels::AbstractVector{<:Label})::UncertainValues
    idx = map(l -> indexin(uvss, l), labels) # look it up once...
    vals = Float64[uvss.values[i] for i in idx]
    cov = zeros(length(idx), length(idx))
    for (r, rl) in enumerate(idx)
        for (c, cl) in enumerate(idx)
            cov[c, r] = (cov[r, c] = uvss.covariance[rl, cl])
        end
    end
    return uvs(labels, vals, cov)
end

function extract(
	uvss::UncertainValues,
    labeltype::Type{T}
)::UncertainValues where {T<:Label}
    return extract(uvss, labelsByType(labeltype, uvss))
end

function extract(
	uvss::UncertainValues,
    labeltypes::AbstractVector{DataType},
)::UncertainValues
    return extract(uvss, labelsByType(labeltypes, uvss))
end

"""

    LaTeXStrings.latexstring(uvs::UncertainValues; fmt="%0.3f", mode=:normal[|:siunutx])::LaTeXString

Formats a `UncertainValues` as a LaTeXString (`mode=:siunitx` uses `\\num{}` from the siunitx package.)
Also requires the amsmath package for `vmatrix`
"""
function LaTeXStrings.latexstring(uvs::UncertainValues; fmt="%0.3f", mode=:normal)::LaTeXString
    fmtrfunc = generate_formatter( fmt )
    lbls = labels(uvs)
    pre, post = ( mode == :siunitx ? ( raw"\num{", raw"}" ) : ( "", "" ) )
    ls = join([
            raw"\begin{vmatrix}",
            join([ repr(lbl) for lbl in lbls ], " \\\\\n"),
            raw"\end{vmatrix}"
        ],"\n")
    vals = join([
            raw"\begin{vmatrix}",
            join([ pre*fmtrfunc(value(uvs, lbl))*post for lbl in lbls ], " \\\\\n"),
            raw"\end{vmatrix}"
        ],"\n")
    covs = join([
        raw"\begin{vmatrix}",
        join( [
            join( [ pre*fmtrfunc(covariance(uvs, lbl1, lbl2))*post for lbl1 in lbls ], " & ")
                for lbl2 in lbls ], " \\\\\n"),
        raw"\end{vmatrix}" ],"\n")
    return latexstring("$ls\n=$vals\n\\pm\n$covs")
end


import SDELab2
#
T=1 # time step, final time
x0=1. # initial data
alp=-1 # parameter
bet=0
sig=0.5
quad0=SDELab2.delta_measure(x0)
###########################################
function drift_fff(u,alp,bet)
  x=u[1]; mf=u[2]
  return alp*x+bet*mf
end
##########################################
function diff_ggg(u,sig)
  x=u[1]
  return sig*x
end
#########################################
function mf_r(x)
  return [x]
end
#########################################
fcnq=SDELab2.set_fcn_mf(u->drift_fff(u,alp,bet),
          u->diff_ggg(u,sig),
          x->mf_r(x))

# test function
phi=x->[x;x.^2]
######################################
exact_mean=x0*exp((alp+bet));
exact_var=(sig^2)/(2*alp)*exp(2*alp)-sig^2/(2*alp);
exact_mom2=exact_var+exact_mean^2;
#exact_mom2=(exp(2*alp)*x0^2
#            +x0^2*exp(2*alp)*(exp(2*bet)-1)
#            +sig^2/(2*alp)*(exp(2*alp)-1))
exact=[exact_mean,exact_mom2];
println("ode soln ",exact)
##############################################
# MLMC implementation in MATLAB
mlmc_cpus= 1e3*[0.000142083385202,  0.000453974977685,   0.012572159661283, 0.051641483155560,   1.913720991516199,   8.940239873090615]
mlmc_errs=[ 0.008701679831122,   0.003964115036436,   0.000721699175333, 0.000323972021408,   0.000089514867964,   0.000056234777698]
#
dt=1e-3
solver="MFEM0"
params=SDELab2.set_opt(dt,solver)
params["eta"]=1 # determines timesteps
params["cutoff"]=10 # for diameter reduction
params["const"]=1e3 # for no_Gauss_points
params["length_scale"]=1 # for no_Gauss_points
params["skiphalf"]=true
params["periodic"]=false
params["no_partition"]=1
dtpts=[0.2,0.1,0.05,0.025,0.01,0.005,0.0025,0.001,0.0005]
errors=zeros(length(dtpts),2)
errors_no=deepcopy(errors)
errors_nd=deepcopy(errors)
cpu_times=zeros(length(dtpts))
cpu_times_no=deepcopy(cpu_times)
cpu_times_nd=deepcopy(cpu_times)

params["Solver"]="MFEM0"
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts[i]/4
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,false)
    errors_no[i,:],cpu_times_no[i]=helper(qoi, exact, cpu_time, "normal")
end
########################
params["Solver"]="MFEM0"
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,true)
    errors[i,:],cpu_times[i]=helper(qoi, exact, cpu_time, "normal")
end
################
params["Solver"]="MFEM2"
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,false)
    errors_nd[i,:],cpu_times_nd[i]=helper(qoi, exact, cpu_time, "normal")
end
#########################################
str1="--" # extrap
str2="k-." # 2nd-order
figure(1);
# first functional
i=1
time_step_err(dtpts,errors[:,i],
  dtpts,errors_no[:,i],
  dtpts,errors_nd[:,i],[1e0,1e0],utf8("mom1.pdf"))
err_cpu(errors[:,i],cpu_times[:],
  errors_no[:,i],cpu_times_no[:],
  errors_nd[:,i],cpu_times_nd[:],[1e0,1e0], utf8("mom1_.pdf"))
# second functional
i=2
time_step_err(dtpts,errors[:,i],
      dtpts,errors_no[:,i],
      dtpts,errors_nd[:,i],[1e0,1],utf8("mom2.pdf"))
err_cpu(errors[:,i], cpu_times[:],
      errors_no[:,i],cpu_times_no[:],
      errors_nd[:,i],cpu_times_nd[:],[1e0,1], utf8("mom2_.pdf"))
#########################################
println("finished preparing plots")
#


arg_dict = Dict{Any, Any}(
  # "operator_type"   => "SBPGamma",
  "operator_type"   => "SBPOmega",
  # "operator_type"   => "SBPDiagE",
  "physics"         => "Elliptic",
  "use_DG"          => true,
  "delta_t"         => 1.0e-3,
  "t_max"           => 5.0,
  "res_tol"         => 1e-12,
  "step_tol"        => 1e-12,
  "itermax"         => 1,
  "order"           => 2,
  "Relfunc_name"    => true,
  "smb_name"        => "SRCMESHES/square16x16.smb",
  "dmg_name"        => ".null",
  "numBC"           => 1,
  "BC1"             => [0, 1, 2, 3],
  "exactSolution"   =>"ExactExpTrig",
  "BC1_name"        => "DirichletTrig",
  "Diffusion"       => "poly0th",
  "SRC_name"        => "SrcExpTrigPoly0thDiffn",
  "Functional"      => "volumeAverage",
  "exactFunctional" => 1.846230857168755189823348538541e-02,
  "IC_name"         => "ICRandom",

  # "exactSolution"   =>"ExactTrig",
  # "BC1_name"        => "DirichletTrig",
  # "Diffusion"       => "poly2nd", # lambda = [x^2+1 & xy \\ xy & y^2+1]
  # "SRC_name"        => "SrcTrigPoly2ndDiffn",

  # "exactSolution"   =>"ExactTrig",
  # "BC1_name"        => "DirichletTrig",
  # "Diffusion"       => "poly6th",
  # "SRC_name"        => "SrcTrigPoly6thDiffn",

  # "exactSolution"   =>"ExactPoly2nd",
  # "Diffusion"       => "DiffnPoly0th"
  # "BC1_name"        => "DirichletPolynial2nd",
  # "SRC_name"        => "SrcPoly2nd",

  # "SRC_name"        => "SRC0",		# for energy stability test

  # "Flux_name"       => "SAT",			# SAT-BR2
  "Flux_name"       => "SAT0",		# SAT-SIPG
  # "Flux_name"       => "Shahbazi",	# Shahbazi
  "Cip"             => -3.0,
  "run_type"        => 5,
  "jac_type"        => 2,
  "jac_method"      => 2,
  "write_eigs"      => false,
  "write_jac"       => false,
  "write_conditioning" => false,
  "write_energy"    => true,
  "solve"           => true,
  "real_time"       => false
)


using SymPoly
using SymbolicUtils, Test
using DynamicPolynomials

function generate_rand_poly(x; min_deg=0, max_deg=5, sparcity=0.5)
    n = rand(min_deg:max_deg)
    p = one(x)
    r = rand(1:10) // rand(1:10)
    while true
        for i = 0:n
            if rand() > sparcity || i == n
                s = rand() < 0.5 ? +1 : -1
                p += x^i * s * r * rand(0:10)//rand(1:5)
            end
        end
        !isequal(p, 0) && return p
        return p
    end
end

function test_eq(x, f, name; n=10, min_deg=0, max_deg=10, sparcity=0.5)
    k = 0
    outcome = true
    for i = 1:n
        p = generate_rand_poly(x; max_deg, sparcity)
        q = generate_rand_poly(x; min_deg, max_deg, sparcity)

        if iszero(p) || iszero(q) continue end

        printstyled("P1 = ", p, '\n'; color=:green)
        printstyled("P2 = ", q, '\n'; color=:blue)

        try
            r = f(p, q)
            printstyled("Δ($name) = ", r, '\n'; color=:red)
            if !iszero(r)
                @warn "+/- mismatch: $r"
                outcome = false
            else
                k += 1
            end
        catch e
            println(e)
        end
    end
    outcome
end


using Revise, Test, ForwardDiff, Parameters, Setfield, Plots, LinearAlgebra
using BifurcationKit
const BK = BifurcationKit

D(f, x, p, dx) = ForwardDiff.derivative(t->f(x .+ t .* dx, p), 0.)
norminf = x -> norm(x, Inf)
####################################################################################################
function TMvf!(dz, z, p, t)
	@unpack J, α, E0, τ, τD, τF, U0 = p
	E, x, u = z
	SS0 = J * u * x * E + E0
	SS1 = α * log(1 + exp(SS0 / α))
	dz[1] = (-E + SS1) / τ
	dz[2] =	(1.0 - x) / τD - u * x * E
	dz[3] = (U0 - u) / τF +  U0 * (1.0 - u) * E
	dz
end

TMvf(z, p) = TMvf!(similar(z), z, p, 0)
dTMvf = (z,p) -> ForwardDiff.jacobian(x-> TMvf(x,p), z)

# we group the differentials together
jet  = BK.get3Jet(TMvf, dTMvf)

par_tm = (α = 1.5, τ = 0.013, J = 3.07, E0 = -2.0, τD = 0.200, U0 = 0.3, τF = 1.5, τS = 0.007) #2.87
z0 = [0.238616, 0.982747, 0.367876 ]

opts_br = ContinuationPar(pMin = -10.0, pMax = -0.9, ds = 0.04, dsmax = 0.125, nInversion = 8, detectBifurcation = 3, maxBisectionSteps = 25, nev = 3)
	opts_br = @set opts_br.newtonOptions.verbose = false
	br, = continuation(TMvf, dTMvf, z0, par_tm, (@lens _.E0), opts_br;
	printSolution = (x, p) -> (E = x[1], x = x[2], u = x[3]),
	tangentAlgo = BorderedPred(),
	plot = true, verbosity = 0, normC = norminf)

plot(br, plotfold=false, markersize=4, legend=:topleft)
####################################################################################################
hopfpt = computeNormalForm(jet..., br, 4)

# newton parameters
optn_po = NewtonPar(verbose = true, tol = 1e-8,  maxIter = 8)

# continuation parameters
opts_po_cont = ContinuationPar(dsmax = 0.1, ds= -0.0001, dsmin = 1e-4, pMax = 0., pMin=-5., maxSteps = 110, newtonOptions = (@set optn_po.tol = 1e-7), nev = 2, precisionStability = 1e-8, detectBifurcation = 3, plotEveryStep = 10, saveSolEveryStep=1)

Mt = 200 # number of sections
	br_potrap, utrap = continuation(
	jet..., br, 4, opts_po_cont,
	PeriodicOrbitTrapProblem(M = Mt);
	linearPO = :Dense,
	verbosity = 2,	plot = true,
	printSolution = (x, p) -> (xtt=reshape(x[1:end-1],3,Mt); return (max = maximum(xtt[1,:]), min = minimum(xtt[1,:]), period = x[end])),
	plotSolution = (x, p; k...) -> begin
		xtt = BK.getTrajectory(p.prob, x, p.p)
		plot!(xtt.t, xtt.u[1,:]; label = "E", k...)
		plot!(xtt.t, xtt.u[2,:]; label = "x", k...)
		plot!(xtt.t, xtt.u[3,:]; label = "u", k...)
		plot!(br,subplot=1, putbifptlegend = false)
		end,
	normC = norminf)

plot(br, br_potrap, markersize = 3)
	plot!(br_potrap.param, br_potrap.min, label = "")

####################################################################################################
# idem with Standard shooting
using DifferentialEquations

# this is the ODEProblem used with `DiffEqBase.solve`
probsh = ODEProblem(TMvf!, copy(z0), (0., 1000.), par_tm; atol = 1e-10, rtol = 1e-9)

opts_po_cont = ContinuationPar(dsmax = 0.05, ds= -0.0001, dsmin = 1e-4, pMax = 0., pMin=-5., maxSteps = 210, newtonOptions = (@set optn_po.tol = 1e-6), nev = 25, precisionStability = 1e-8, detectBifurcation = 0, plotEveryStep = 10, saveSolEveryStep=0)

br_posh, = @time continuation(
	jet..., br, 4,
	# arguments for continuation
	opts_po_cont,
	# this is where we tell that we want Standard Shooting
	ShootingProblem(15, par_tm, probsh, Rodas4(), parallel = true);
	ampfactor = 1.0, δp = 0.0005,
	usedeflation = true,
	linearPO = :autodiffDense,
	updateSectionEveryStep = 2,
	verbosity = 2,	plot = true,
	printSolution = (x, p) -> (return (max = getMaximum(p.prob, x, @set par_tm.E0 = p.p), period = getPeriod(p.prob, x, @set par_tm.E0 = p.p))),
	plotSolution = (x, p; k...) ->
		begin
			xtt = BK.getTrajectory(p.prob, x, @set par_tm.E0 = p.p)
			plot!(xtt; legend = false, k...);
			plot!(br, subplot=1, putspecialptlegend = false)
		end,
	normC = norminf)

plot(br_posh, br, markersize=3)
plot(br, br_potrap, br_posh, markersize=3)
####################################################################################################
# idem with Poincaré shooting

function TMvfExtended!(dz, z, p, t)
	# we write the first part
	TMvf!(dz, z, p, t)
	dz[4:end] .= @views ForwardDiff.derivative(t -> TMvf(z[1:3] .+ t .* z[4:end], p), 0)
end

probmono = ODEProblem(TMvfExtended!, vcat(u0, u0), (0., 1000.), par_tm; atol = 1e-10, rtol = 1e-9)

opts_po_cont = ContinuationPar(dsmax = 0.02, ds= -0.001, dsmin = 1e-5, pMax = 0., pMin=-5., maxSteps = 200, newtonOptions = NewtonPar(optn_po;tol = 1e-6, maxIter=15), nev = 25, precisionStability = 1e-8, detectBifurcation = 0, plotEveryStep = 5, saveSolEveryStep = 2)

br_popsh, = @time continuation(
	jet..., br, 4,
	# arguments for continuation
	opts_po_cont,
	# this is where we tell that we want Poincaré Shooting
	PoincareShootingProblem(9, par_tm, probsh, Rodas4P(), probmono, Rosenbrock23(), parallel = false);
	ampfactor = 1.0, δp = 0.0005,
	usedeflation = false,
	linearPO = :autodiffDense,
	updateSectionEveryStep = 2,
	verbosity = 2,	plot = true,
	printSolution = (x, p) -> (return (max = getMaximum(p.prob, x, @set par_tm.E0 = p.p), period = getPeriod(p.prob, x, @set par_tm.E0 = p.p))),
	plotSolution = (x, p; k...) ->
		begin
			xtt = BK.getTrajectory(p.prob, x, @set par_tm.E0 = p.p)
			plot!(xtt; legend = false, k...);
			plot!(br,subplot=1, putspecialptlegend = false)
			# plot!(br_potrap,subplot=1, putspecialptlegend = false)
		end,
	callbackN = (x, f, J, res, iteration, itlinear, options; kwargs...) -> (return res<1e16),
	finaliseSolution = (z, tau, step, contResult; prob=nothing, k...) ->
		begin
			isnothing(prob) && return true
			T = getPeriod(prob, z.u, @set par_tm.E0 = z.p)
			@show T
			T < 20.
		end,
	normC = norminf)

plot(br, br_popsh, markersize=3)

_po = br_popsh.sol[1].x
fl = br_popsh.functional.flow


using MLKernels

# Constructing Standard Kernels

GaussianKernel()
SquaredDistanceKernel()

ϕ = ExponentialClass(2.0, 1.0);
κ = SquaredDistanceKernel(1.0);
ψ = ϕ ∘ κ   # use \circ for '∘'

ismercer(κ)
isnegdef(κ)
ismercer(ψ)
isnegdef(ψ)

w = round(rand(5),3);
ARD(κ, w)

# Kernel Functions

x = rand(3); y = rand(3);
kernel(ψ, x[1], y[1])
kernel(ψ, x, y)

X = rand(5,3);
kernelmatrix(ψ, X)
kernelmatrix(ψ, X', true)

Y = rand(4,3);
kernelmatrix(ψ, X, Y)
kernelmatrix(ψ, X', Y', true)

kernelmatrix(ψ, [X; Y])[1:5, 6:9]

# Kernel Operations

κ1 = GaussianKernel();

2k1 + 3

κ2 = PolynomialKernel();

κ1 + κ2

κ1 * κ2

κ3 = SineSquaredKernel();

κ3^0.5

κ4 = 2ScalarProductKernel() + 3;

κ4^3

exp(κ4)

tanh(κ4)


using LinearAlgebra
using SparseArrays

"""
Construct a Heisenberg-type Hamiltonian as sparse matrix using site-dependent interaction strength, with nearest-neighbor interactions according to adjacency.
"""
function construct_hamiltonian(J::Array{<:Real,3})
    L = size(J, 3)
    @assert(size(J) == (3, L, L))

    # spin operators (Pauli matrices divided by 2)
    sigma12 = (0.5*sparse([0.  1.; 1.  0.]),
               0.5*sparse([0. -im; im  0.]),
               0.5*sparse([1.  0.; 0. -1.]))

    H = spzeros(Float64, 2^L, 2^L)

    # interaction terms
    for i in 1:L
        for j in i+1:L
            for k in 1:3
                # considering only entries in J for i < j
                if J[k, i, j] ≠ 0
                    H -= J[k, i, j] * real(kron(sparse_identity(2^(L-j)), sigma12[k], sparse_identity(2^(j-i-1)), sigma12[k], sparse_identity(2^(i-1))))
                end
            end
        end
    end

    return H
end


using Random

using Bijectors
using Turing
using Turing: Variational

using StatsFuns: softplus, invsoftplus
f = softplus
f⁻¹ = invsoftplus

# setup for plotting
using Plots, StatsPlots, LaTeXStrings
pyplot()

# used to compute closed form expression of posterior
using ConjugatePriors

# define Turing model
@model model(x) = begin
    s ~ InverseGamma(2, 3)
    m ~ Normal(0.0, sqrt(s))  # `Normal(μ, σ)` has mean μ and variance σ², i.e. parametrize with std. not variance

    for i = 1:length(x)
        x[i] ~ Normal(m, sqrt(s))
    end
end

const seeds = [125, 245, 3]
const ad_modes = [:forward_diff, :reverse_diff]

for seed ∈ seeds
    @info seed

    for ad_mode ∈ ad_modes
        @info ad_mode
        setadbackend(ad_mode)

        # set random seed
        Random.seed!(seed)

        # generate data
        x = randn(1, 2000);

        # construct model
        m = model(x)

        # ADVI
        opt = Variational.TruncatedADAGrad()   # optimizer

        advi = ADVI(10, 100)                   # <: VariationalInference
        q = Variational.meanfield(m)           # => <: VariationalPosterior

        elbo = Variational.ELBO()              # <: VariationalObjective

        μ, σs = params(q)
        θ = vcat(μ, f⁻¹.(σs))

        history = [elbo(advi, q, m, 1000)]

        # construct animation
        anim = @animate for j = 1:25
            # global q
            Variational.optimize!(elbo, advi, q, m, θ; optimizer = opt)
            μ, ω = θ[1:length(q)], θ[length(q) + 1:end]

            q = Variational.update(q, μ, f.(ω))
            samples = rand(q, 2000)

            # quick check
            println([mean(samples, dims=2), [var(x), mean(x)]])

            # plotting code assumes (samples, dim) shape so we just transpose
            samples = transpose(samples)

            # closed form computation
            # notation mapping has been verified by explicitly computing expressions
            # in "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy
            μ₀ = 0.0 # => μ
            κ₀ = 1.0 # => ν, which scales the precision of the Normal
            α₀ = 2.0 # => "shape"
            β₀ = 3.0 # => "rate", which is 1 / θ, where θ is "scale"

            # prior
            pri = NormalGamma(μ₀, κ₀, α₀, β₀)

            # posterior
            post = posterior(pri, Normal, x)

            # marginal distribution of τ = 1 / σ²
            # Eq. (90) in "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy
            # `scale(post)` = θ
            p_τ = Gamma(post.shape, scale(post))
            p_σ²_pdf = z -> pdf(p_τ, 1 / z) # τ => 1 / σ²

            # marginal of μ
            # Eq. (91) in "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy
            p_μ = TDist(2 * post.shape)

            μₙ = post.mu    # μ → μ
            κₙ = post.nu    # κ → ν
            αₙ = post.shape # α → shape
            βₙ = post.rate  # β → rate

            # numerically more stable but doesn't seem to have effect; issue is probably internal to
            # `pdf` which needs to compute ≈ Γ(1000)
            p_μ_pdf = z -> exp(logpdf(p_μ, (z - μₙ) * exp(- 0.5 * log(βₙ) + 0.5 * log(αₙ) + 0.5 * log(κₙ))))
            # p_μ_pdf1 = z -> pdf(p_μ, (z - μₙ) / √(βₙ / (αₙ * κₙ)))

            # posterior plots
            p1 = plot();
            density!(samples[:, 1], label = "s (ADVI)", color = :blue, linestyle = :dash);
            histogram!(samples[:, 1], label = "", normed = true, alpha = 0.3, color = :blue);

            # normalize using Riemann approx. because of (almost certainly) numerical issues
            Δ = 0.001
            r = 0.75:0.001:1.50
            norm_const = sum(p_σ²_pdf.(r) .* Δ)
            plot!(r, p_σ²_pdf, label = "s (posterior)", color = :red);
            vline!([var(x)], label = "s (data)", linewidth = 1.5, color = :black, alpha = 0.7);
            xlims!(0.5, 1.5);
            title!("$(j * advi.max_iters) steps");

            p2 = plot();
            density!(samples[:, 2], label = "m (ADVI)", color = :blue, linestyle = :dash);
            histogram!(samples[:, 2], label = "", normed = true, alpha = 0.3, color = :blue);

            # normalize using Riemann approx. because of (almost certainly) numerical issues
            Δ = 0.0001
            r = -0.1 + mean(x):Δ:0.1 + mean(x)
            norm_const = sum(p_μ_pdf.(r) .* Δ)
            plot!(r, z -> p_μ_pdf(z) / norm_const, label = "m (posterior)", color = :red);
            vline!([mean(x)], label = "m (data)", linewidth = 1.5, color = :black, alpha = 0.7);

            xlims!(-0.25, 0.25);

            # visualize evolution of objective wrt. optimization iterations
            obj = elbo(advi, q, m, 1000)
            @info "ELBO" obj
            push!(history, obj)
            p3 = plot();
            plot!(1:advi.max_iters:length(history) * advi.max_iters, history, label = "")
            title!("ELBO = $obj")

            # plot the latest 25 objective evaluations to visualize trend
            p4 = plot();
            plot!(history[max(1, end - 10):end], label = "")

            p = plot(p1, p2, p3, p4; layout = (4, 1))

            @info "[$j] Done"
            p
        end
        gif(anim, "advi_w_elbo_fps15_$(seed)_$(ad_mode).gif", fps = 15)
    end
end


@doc raw"""
    Sphere3

An object `s` of type [`Sphere3`](@ref) is a sphere in the three-dimensional
Euclidean space ``Ε^3``.

The sphere is oriented, i.e. its boundary has clockwise or counterclockwise
orientation. The boundary splits ``Ε^3`` into a positive and a negative
side, where the positive side is to the left of the boundary. The boundary also
splits ``Ε^3`` into a bounded and an unbounded side. Note that the sphere
can be degenerated, i.e. the squared radius may be zero.
"""
Sphere3

@doc raw"""
    Sphere3(center::Point3, r²::Real, ori::Orientation = COUNTERCLOCKWISE)

Introduces a variable `s` of type [`Sphere3`](@ref).

It is initialized to the circle with center `center`, squared radius
`squared_radius`, and orientation `ori`.

!!! info "Precondition"

    `ori` ``≠`` `COPLANAR`, and further, `squared_radius` ``≥ 0``.
"""
Sphere3(center::Point3, r²::Real, ori::Orientation = COUNTERCLOCKWISE)

@cxxdereference Sphere3(c::Point3, r²::Real, ori::Union{Orientation,Ref{Orientation}} = COUNTERCLOCKWISE) =
    Sphere3(c, convert(FT, r²), ori)

"""
    Sphere3(p::Point3, q::Point3, r::Point3, s::Point3)

Introduces a variable `s` of type [`Sphere3`](@ref).

It is initialized to the unique sphere which passes through the points `p`, `q`,
`r` and `s`. The orientation of the sphere is the orientation of the point
quadruple `p`, `q`, `r`, `s`.

!!! info "Precondition"

    `p`, `q`, `r` and `s` are not coplanar.
"""
Sphere3(p::Point3, q::Point3, r::Point3, s::Point3)

@doc raw"""
    Sphere3(p::Point3, q::Point3, r::Point3, ori::Orientation = COUNTERCLOCKWISE)

Introduces a variable `s` of type [`Sphere3`](@ref).

It is initialized to the smallest sphere which passes through the points `p`,
`q`, and `r`. The orientation of the sphere is `ori`.

!!! info "Precondition"

    `ori` ``≠`` `COPLANAR`.
"""
Sphere3(p::Point3, q::Point3, r::Point3, ori::Orientation = COUNTERCLOCKWISE)

@doc raw"""
    Sphere3(p::Point3, q::Point3, ori::Orientation = COUNTERCLOCKWISE)

Introduces a variable `s` of type [`Sphere3`](@ref).

It is initialized to the smallest sphere which passes through the points `p` and
`q`. The orientation of the sphere is `ori`.

!!! info "Precondition"

    `ori` ``≠`` `COPLANAR`.
"""
Sphere3(p::Point3, q::Point3, ori::Orientation = COUNTERCLOCKWISE)

@doc raw"""
    Sphere3(center::Point3, ori::Orientation = COUNTERCLOCKWISE)

Introduces a variable `s` of type [`Sphere3`](@ref).

It is initialized to the sphere with center `center`, squared radius zero and
orientation `ori`.

!!! info "Precondition"

    `ori` ``≠`` `COPLANAR`.

!!! info "Postcondition"

    `is_degenerate(s) == true`.
"""
Sphere3(center::Point3, ori::Orientation = COUNTERCLOCKWISE)

"""
    Sphere3(c::Circle3)

Introduces a variable `s` of type [`Sphere3`](@ref).

It is initialized to the diametral sphere of the circle.
"""
Sphere3(c::Circle3)

"""
    center(s::Sphere3)

Returns the center of `s`.
"""
center(s::Sphere3)

"""
    squared_radius(s::Sphere3)

Returns the squared radius of `s`.
"""
squared_radius(s::Sphere3)

"""
    orientation(s::Sphere3)

Returns the orientation of `s`.
"""
orientation(s::Sphere3)

"""
    ==(s₁::Sphere3, s₂::Sphere3)

Returns `true`, iff `s₁` and `s₂` are equal, i.e. if they have the same center,
same squared radius and same orientation.
"""
==(s₁::Sphere3, s₂::Sphere3)

"""
    is_degenerate(s::Sphere3)

Returns `true`, iff `s` is degenerate, i.e. if `s` has squared radius zero.
"""
is_degenerate(s::Sphere3)

"""
    oriented_side(s::Sphere3, p::Point3)

Returns either the constant [`ON_ORIENTED_BOUNDARY`](@ref),
[`ON_POSITIVE_SIDE`](@ref), or [`ON_NEGATIVE_SIDE`](@ref), iff `p` lies on the
boundary, properly on the positive side, or properly on the negative side of
`s`, resp.
"""
oriented_side(s::Sphere3, p::Point3)

"""
    bounded_side(s::Sphere3, p::Point3)

Returns [`ON_BOUNDED_SIDE`](@ref), [`ON_BOUNDARY`](@ref), or
[`ON_UNBOUNDED_SIDE`](@ref) iff p lies properly inside, on the boundary, or
properly outside of `s`, resp.
"""
bounded_side(s::Sphere3, p::Point3)

"""
    has_on_positive_side(s::Sphere3, p::Point3)

Returns `true`, iff `p` lies on the positive side of `s`.
"""
has_on_positive_side(s::Sphere3, p::Point3)

"""
    has_on_negative_side(s::Sphere3, p::Point3)

Returns `true`, iff `p` lies on the negative side of `s`.
"""
has_on_negative_side(s::Sphere3, p::Point3)

"""
    has_on_boundary(s::Sphere3, p::Point3)

Returns `true`, iff `p` lies on the boundary of `s`.
"""
has_on_boundary(s::Sphere3, p::Point3)

"""
    has_on_bounded_side(s::Sphere3, p::Point3)

Returns `true`, iff `p` lies properly inside of `s`.
"""
has_on_bounded_side(s::Sphere3, p::Point3)

"""
    has_on_unbounded_side(s::Sphere3, p::Point3)

Returns `true`, iff `p` lies properly outside of `s`.
"""
has_on_unbounded_side(s::Sphere3, p::Point3)

"""
    opposite(s::Sphere3)

Returns the sphere with the same center and squared radius as `s`, but with
opposite orientation.
"""
opposite(s::Sphere3)

"""
    orthogonal_transform(s::Sphere3, at::AffTransformation3)

Returns the circle obtained by applying `at` on `s`.

!!! info "Precondition"

    `at` is an orthogonal transformation.
"""
orthogonal_transform(s::Sphere3, at::AffTransformation3)

"""
    bbox(s::Sphere3)

Returns a bounding box containing `s`.
"""
bbox(s::Sphere3)


n = NamedArray(rand(2,4))
setnames!(n, ["one", "two"], 1)
setnames!(n, ["a", "b", "c", "d"], 2)

v = NamedArray(rand(1:100, 6), (["a", "b", "c", "d", "e", "f"],), (:index,))

m = NamedArray(rand(2,3,4,3,2,3))


using JuMP, MAT

function load_sedumi_file(filename)
   	file = matopen(filename)
   	if exists(file, "A")
      		A = read(file, "A");
   	elseif exists(file, "At")
      		A = read(file, "At")';
   	else
      		@assert false "The MATLAB file must contain either A or At as a variable"
   	end
   	b = Vector(read(file, "b")[:]);
   	c = Vector(read(file, "c")[:]);
   	K = read(file, "K")
   	close(file)

	# Cast Int entries
   	for key in ["f", "l"]
      		if haskey(K, key)
         			K[key] = Int(K[key])
      		else
         			K[key] = 0
      		end
   	end
	# Cast Vector{Int} entries
   	for key in ["q", "r", "s"]
      		if haskey(K, key)
         			@assert isa(K[key], AbstractArray) || isa(K[key], Number)

         			if isa(K[key], AbstractArray)
            				K[key] = Vector{Int}(K[key][:])
         			else # isa(K[key], Number)
            				K[key] = [Int(K[key])]
         			end
         			K[key] = K[key][K[key] .> 0]
      		else
         			K[key] = Int[]
      		end
   	end

   	return A, b, c, K
end

function solve_sedumi_jump(A, b, c, K, solver; kwargs...)
	# Solve Problems in Sedumi form, i.e.
	# minimize    c'x
	# subject to  Ax = b
	#             x \in K
	# See https://github.com/sqlp/sedumi/blob/master/sedumi.m#L49-L92

   	model = Model(with_optimizer(solver; kwargs...))

	# Fisrt handle all variables except semidefinite matrices
   	constraint_types = [MOI.Reals;
		MOI.Nonnegatives;
		fill(MOI.SecondOrderCone, length(K["q"]));
		fill(MOI.RotatedSecondOrderCone, length(K["r"]))]
   	dimensions = [K["f"]; K["l"]; K["q"]; K["r"]]

   	x = Vector{VariableRef}(undef, 0)
   	for (dimension, constraint_type) in zip(dimensions, constraint_types)
      		x_ = @variable(model, [1:dimension])
      		if constraint_type != MOI.Reals # Some solvers don't accept MOI.Reals
         			@constraint(model, x_ in constraint_type(dimension))
      		end
      		append!(x, x_)
   	end

	# Now handle semidefinite matrices
   	for dimension in K["s"]
      		x_ = @variable(model, [1:dimension, 1:dimension], PSD)
      		append!(x, reshape(x_, length(x_)))
   	end

   	@constraint(model, A * x .== b)
   	@objective(model, Min, dot(c, x))

   	JuMP.optimize!(model)
    return value.(x), JuMP.objective_value(model), JuMP.termination_status(model)
end

function solve_sedumi_jump_dual(A, b, c, K, solver; kwargs...)
	# Solve the dual of problem described in Sedumi form, i.e.
	# maximize    b'y
	# subject to  c - A'y \in K*
	# See https://github.com/sqlp/sedumi/blob/master/sedumi.m#L49-L92

   	model = Model(with_optimizer(solver; kwargs...))

	# Fisrt handle all variables except semidefinite matrices
   	constraint_types = [MOI.Zeros;
		MOI.Nonnegatives;
		fill(MOI.SecondOrderCone, length(K["q"]));
		fill(MOI.RotatedSecondOrderCone, length(K["r"]))]
   	dimensions = [K["f"]; K["l"]; K["q"]; K["r"]]

   	@variable(model, y[1:length(b)])
   	start_idx = 1
   	for (dimension, constraint_type) in zip(dimensions, constraint_types)
      		indices = start_idx:start_idx + dimension - 1
      		dimension > 0 && @constraint(model, c[indices] - A[:, indices]' * y in constraint_type(dimension))
      		start_idx += dimension
   	end
   	for dimension in K["s"]
      		indices = Vector(start_idx:start_idx + dimension^2 - 1)[upper_triangular_indices(dimension)]
      		@constraint(model, c[indices] - A[:, indices]' * y in MOI.PositiveSemidefiniteConeTriangle(dimension))
      		start_idx += dimension^2
   	end

   	@objective(model, Max, dot(b, y))

   	JuMP.optimize!(model)
    return value.(y), JuMP.objective_value(model), JuMP.termination_status(model)
end

function upper_triangular_indices(n)
   	indices = Vector{Int}(undef, Int(n * (n + 1) / 2))
   	index = 0
   	counter = 0
   	for j = 1:n, i = 1:n
      		index += 1
      		if i <= j
         			counter += 1
         			indices[counter] = index
      		end
   	end
   	return indices
end


# The function to_matrix is a fork from ControlSystems.jl
# covers the case when a vector represents the row of a matrix
to_matrix(T, A::AbstractVector, wide::Bool = false) = wide ? Matrix{T}(reshape(A, 1, length(A))) : Matrix{T}(reshape(A, length(A), 1))
to_matrix(T, A::AbstractMatrix, wide::Bool = false) = Matrix{T}(A)  # Fallback
to_matrix(T, A::Number, wide::Bool = true) = fill(T(A), 1, 1)
# Handle Adjoint Matrices
to_matrix(T, A::Adjoint{R, MT}, wide::Bool = false) where {R<:Number, MT<:AbstractMatrix} = to_matrix(T, MT(A))

#
eye(n) = Matrix{Bool}(I, n, n)
eye(m,n) = Matrix{Bool}(I, m, n)
eye(::Type{T}, n) where {T} = Matrix{T}(I, n, n)
eye(::Type{T}, m, n) where {T} = Matrix{T}(I, m, n)

function rcond(A::DenseMatrix, tola::Real = 0)
    T = eltype(A)
    T1 = T <: BlasFloat ? T : T1 = promote_type(T,Float64)
    nrmA = opnorm(A,1)
    nrmA <= tola && (return zero(real(T1)))
    istriu(A) ? (return LinearAlgebra.LAPACK.trcon!('1','U','N',copy_oftype(A,T1))) :
        (return LinearAlgebra.LAPACK.gecon!('1', LinearAlgebra.LAPACK.getrf!(copy_oftype(A,T1))[1],real(T1)(nrmA)) )
end
function rcond(A::UpperTriangular, tola::Real = 0)
    T = eltype(A)
    T1 = T <: BlasFloat ? T : T1 = promote_type(T,Float64)
    nrmA = opnorm(A,1)
    nrmA <= tola && (return zero(real(T1)))
    return LinearAlgebra.LAPACK.trcon!('1','U','N', Matrix{T1}(A))
end
function jordanblockdiag(lambda::T, ed::Vector{Int}) where T
   n = sum(ed)
   J = zeros(T,n,n)
   i = 1
   for k = 1:length(ed)
       ni = ed[k]
       J[i:i+ni-1,i:i+ni-1] = jordanblock(ni,lambda)
       i += ni
   end
   return J
end
function jordanblock( n::Integer, lambda::T) where T
   # Generate a n-by-n elementary Jordan block with eigenvalue lambda
   J = lambda*Matrix{T}(I, n, n) + diagm(1 => ones(T, n-1, 1)[:,1])
   return J
end
blockdiag(mats::AbstractMatrix...) = blockdiag(promote(mats...)...)

function blockdiag(mats::AbstractMatrix{T}...) where T
    rows = Int[size(m, 1) for m in mats]
    cols = Int[size(m, 2) for m in mats]
    res = zeros(T, sum(rows), sum(cols))
    m = 1
    n = 1
    for ind=1:length(mats)
        mat = mats[ind]
        i = rows[ind]
        j = cols[ind]
        res[m:m + i - 1, n:n + j - 1] = mat
        m += i
        n += j
    end
    return res
end
function sblockdiag(blockdims::Vector{Int},mats::Union{AbstractMatrix{T},UniformScaling}...) where T
    nmat = sum(blockdims)
    res = zeros(T, nmat, nmat)
    i = 1
    for ind = 1:length(mats)
        mat = mats[ind]
        ni = blockdims[ind]
        i1 = i:i+ni-1
        res[i1,i1] = (mat == I) ? Matrix{T}(I,ni,ni) : mat
        i += ni
    end
    return res
end