"""
    abstract type AbstractDivergence

An abstract type representing Csiszár φ-divergences. Subtypes should implement
[`φstar`](@ref) and [`aprox`](@ref), and optionally can implement
[`initialize_dual_potential!`](@ref) and/or [`sinkhorn_divergence!`](@ref).
"""
abstract type AbstractDivergence end

"""
    φstar(::AbstractDivergence, q::Number) -> Number

The Legendre conjugate of the function `φ` associated to the divergence.
"""
function φstar end

"""
    aprox(::AbstractDivergence, ϵ::Number, x::Number) -> Number

The anisotropic proximity operator defined in Def. 2 of [[SFVTP19](@ref)].
"""
function aprox end

# The following formulas for `φstar` and `aprox` are found in Section 2.4 of SFVTP19.

"""
    Balanced <: AbstractDivergence

Represents the divergence `Dᵩ(a|b)` which is zero if `a==b` and infinite
otherwise. Generalized by [`RG`](@ref).
"""
struct Balanced <: AbstractDivergence end
φstar(::Balanced, q) = q
aprox(::Balanced, ϵ, x) = x

"""
    KL{ρ} <: AbstractDivergence

Represents the divergence `ρ*KL(a|b)`, where `KL` is the Kullback-Leibler
divergence. The parameter `ρ` is simply a scaling.
"""
struct KL{ρ} <: AbstractDivergence end
KL(ρ::Number) = KL{ρ}()
KL() = KL(1.0)

φstar(::KL{ρ}, q) where {ρ} = ρ * (exp(q / ρ) - 1)

aprox(::KL{ρ}, ϵ, x) where {ρ} = inv(one(ρ) + ϵ / ρ) * x

"""
    RG{l,u} <: AbstractDivergence

Represents the divergence `Dᵩ(a|b)` which is zero if `l*b .<= a .<= u*b` and
infinite otherwise. Equivalent to [`Balanced`](@ref) when `l == u`.
"""
struct RG{l,u} <: AbstractDivergence end

function RG(l::Number, u::Number)
    l <= u || throw(DomainError(u - l, "Need l <= u"))
    RG{l,u}()
end

φstar(::RG{a,b}, q) where {a,b} = max(a * q, b * q)

function aprox(::RG{a,b}, ϵ, x) where {a,b}
    if (s = x - ϵ * log(a)) < 0
        s
    elseif (t = x - ϵ * log(b)) > 0
        t
    else
        0
    end
end

"""
    TV{ρ} <: AbstractDivergence

Represents the divergence `ρ*TV(u,v) = ρ*norm(u-v,1)`, where `TV` is twice the
total variation distance. The parameter `ρ` is simply a scaling.
"""
struct TV{ρ} <: AbstractDivergence end
TV(ρ::Number) = TV{ρ}()
TV() = TV(1.0)

φstar(::TV{ρ}, q) where {ρ} =
    q <= ρ ? max(-ρ, q) : throw(DomainError(q, "Must have q <= ρ"))

function aprox(::TV{ρ}, ϵ, x) where {ρ}
    if -ρ <= x <= ρ
        x
    elseif x < -ρ
        -ρ
    else
        ρ
    end
end


using Base.Test
using Permutations
p = Permutation([7:-1:1])
i = Permutation(7)
@test p*p == i
@test p == inv(p)
a =[2,3,1,6,7,8,5,4,9]
p = Permutation(a)
@test order(p) == 6
@test p^6 == Permutation(9)
@test p^-5 == p
@test array(p)==a

q = Permutation([1,5,3,9,4,8,6,7,2] )
@test p*q != q*p
@test q[2] == 5

P = matrix(p)
Q = matrix(q)
@test P*Q == matrix(q*p)
@test P' == matrix(inv(p))


# import LinearAlgebra

using LinearAlgebra: BlasInt, Cholesky

## Custom stuff

"""
    λ, C = posdefλ(Q0, Q1)

Provide an initial `λ` such that `Q0 + λ*Q1` is positive-semidefinite (or preferably,
positive-definite).  `λ` may be larger than strictly necessary. `C` is the Cholesky
decomposition of `Q0 + λ*Q1`.
"""
function posdefλ(Q0::SymTridiagonal{T}, Q1::Union{Diagonal{T},UniformScaling};
                 growauto=T(1.0001), growfail=T(1.1)) where T
    λ = zero(T)
    C = cholesky(Q0; check=false)
    if issuccess(C)
        return λ, C
    end
    dv, ev = Q0.dv, Q0.ev
    n = length(dv)
    α = 4*(cos(π/(n+1)))^2
    for i = 1:n-1
        # A symmetric tridiagonal matrix M is positive-definite if
        #     α * M[i,i+1]^2 < M[i,i]*M[i+1,i+1]   ∀ i
        # See Johnson, Neumann, and Tsatsomeros, Conditions for the positivity of
        # determinants, Linear and Multilinear Algebra 40: 241-248 (1996).
        # This criterion is sufficient but not necessary.

        # In this case M[i,i+1] = Q0[i,i+1] and M[i,i] = Q0[i,i] + λ*Q1[i,i].
        # Consequently we can turn this into a quadratic inequality in λ
        #       a*λ^2 + b*λ + c > 0
        c = dv[i]*dv[i+1] - α*ev[i]^2
        if c < 0 || dv[i] < 0
            # This term needs some positive λ
            a = T(Q1[i,i]*Q1[i+1,i+1])
            @assert a >= 0   # Q1 must be possemidef
            b = Q1[i,i]*dv[i+1] + Q1[i+1,i+1]*dv[i]
            s = sqrt(b^2 - 4*a*c)
            # Since the quadratic is convex and below 0 at λ=0, only one root has λ > 0.
            # So we always choose the (-b + s)/(2*a) solution. But write it in a
            # numerically-stable way.
            λ0 = b > 0  ?  -2*c/(s + b)  :  (-b + s)/(2*a)
            λ = max(λ, λ0)
        end
    end
    # Due to roundoff we need to check whether this actually produces a posdef matrix
    λ *= T(growauto)
    C = cholesky(Q0 + λ*Q1; check=false)
    while !issuccess(C)
        λ *= T(growfail)
        C = cholesky(Q0 + λ*Q1; check=false)
    end
    return λ, C
end

function posdefλ(Q0::AbstractMatrix{T}, Q1::Union{AbstractMatrix{T},UniformScaling};
                 growauto=T(1.0001), growfail=T(1.1)) where T
    λ = zero(T)
    C = cholesky(Q0; check=false)
    if issuccess(C)
        return λ, C
    end
    if isa(Q1, UniformScaling)
        D, _ = eigs(Q0; nev=1, which=:SR, ritzvec=false)
        λ = -T(D[1]/T(Q1.λ))
    else
        D, _ = eigs(Q0, Q1; nev=1, which=:SR, ritzvec=false)
        λ = -T(D[1])
    end
    @assert λ > 0
    λ *= T(growauto)
    C = cholesky(Q0 + λ*Q1; check=false)
    while !issuccess(C)
        λ *= T(growfail)
        C = cholesky(Q0 + λ*Q1; check=false)
    end
    return λ, C
end

asmatrix(M::UniformScaling{T}, sz) where T = M.λ * eye(T, sz...)
asmatrix(M::AbstractMatrix, sz) = Matrix(M)

## Stuff that should be in stdlib/LinearAlgebra

function cholesky(A::SymTridiagonal{T}, ::Val{false}=Val(false); check::Bool=true) where T
    n = size(A, 1)
    dv = zeros(T, n)
    ev = zeros(T, n-1)
    L = Bidiagonal{T}(dv, ev, :L)
    if A.dv[1] < 0
        check && error("A is not positive-semidefinte")
        dv[1] = A.dv[1]
        info = -1
        return Cholesky(L, 'L', convert(BlasInt, info))
    end
    dv[1] = sqrt(A.dv[1])
    info = 0
    for i = 1:n-1
        ev[i] = A.ev[i]/dv[i]
        Δ = A.dv[i+1] - ev[i]^2
        if Δ < 0
            check && error("A is not positive-semidefinte")
            info = -1
            dv[i+1] = Δ  # return the magnitude of the violation for later inspection
            break
        end
        dv[i+1] = sqrt(Δ)
    end
    return Cholesky(L, 'L', convert(BlasInt, info))
end

# cholesky(A::AbstractMatrix{T}, ::Val{false}=Val(false); check::Bool=true) where T =
#     cholfact(A)
function cholesky(A::StridedMatrix{T}, ::Val{false}=Val(false); check::Bool=true) where T
    L, info = LinearAlgebra.LAPACK.potrf!('L', copy(A))
    check && @assert info == 0
    return Cholesky(L, 'L', info)
end

# issuccess(C::Cholesky) = C.info == 0
function issuccess(C::Cholesky)
    n = size(C, 1)
    F = C.factors
    for i = 1:n
        F[i,i] >= 0 || return false
    end
    return true
end

function Base.:+(S::SymTridiagonal{T}, D::UniformScaling{T}) where T
    dv, ev = copy(S.dv), copy(S.ev)
    n = length(dv)
    for i = 1:n
        dv[i] += D[i,i]
    end
    return SymTridiagonal(dv, ev)
end

function Base.:\(C::Cholesky{T,Bidiagonal{T}}, b::AbstractVector) where T
    if C.factors.isupper #C.factors.uplo == 'U'
        R = C.factors
        return R \ (R' \ b)
    end
    L = C.factors
    return L' \ (L \ b)
end

LinearAlgebra.eigen(A::AbstractMatrix, B::UniformScaling) = eigen(A, asmatrix(B, size(A)))
LinearAlgebra.eigen(A::AbstractMatrix, B::Diagonal) = eigen(A, asmatrix(B, size(A)))


module Test_dplyap

using LinearAlgebra
using MatrixEquations
using Test

@testset "Testing positive discrete Lyapunov equation solvers" begin

n = 30
m = 5
p = 3
Ty = Float64

@testset "Positive discrete Lyapunov equations" begin

reltol = eps(float(100))
a = -.1+0.2im; b = 2im; @time u = plyapd(a,b)
@test abs(a*u*u'*a'-u*u'+b*b') < reltol

reltol = eps(float(100f0))
a = -.1f0+-.1f0*im; b = 2f0im; @time u = plyapd(a,b)
@test abs(a*u*u'*a'-u*u'+b*b')  < reltol

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ar = ar/(one(Ty) + norm(ar))
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
ac = ac/(one(Ty) + norm(ac))
br = rand(Ty,n,m)
brw = rand(Ty,n,n+m)
bc = br+im*rand(Ty,n,m)
bcw = brw+im*rand(Ty,n,n+m)
cr = rand(Ty,p,n)
crt = rand(Ty,n+p,n)
cc = cr+im*rand(Ty,p,n)
cct = crt+im*rand(Ty,n+p,n)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time u = plyapd(ar,br);
x = u*u'; @test norm(ar*x*ar'-x+br*br')/norm(x)/max(1.,norm(ar)^2) < reltol

@time u = plyapd(ar,brw);
x = u*u'; @test norm(ar*x*ar'-x+brw*brw')/norm(x)/max(1.,norm(ar)^2) < reltol

@time u = plyapd(ar,I,br);
x = u*u'; @test norm(ar*x*ar'-x+br*br')/norm(x)/max(1.,norm(ar)^2) < reltol

@time u = plyapd(ar',cr');
x = u'*u; @test norm(ar'*x*ar-x+cr'*cr)/norm(x)/max(1.,norm(ar)^2) < reltol

@time u = plyapd(ar',crt');
x = u'*u; @test norm(ar'*x*ar-x+crt'*crt)/norm(x)/max(1.,norm(ar)^2) < reltol

@time u = plyapd(ac,bc);
x = u*u'; @test norm(ac*x*ac'-x+bc*bc')/norm(x)/max(1.,norm(ac)^2) < reltol

@time u = plyapd(ac,bcw);
x = u*u'; @test norm(ac*x*ac'-x+bcw*bcw')/norm(x)/max(1.,norm(ac)^2) < reltol

@time u = plyapd(ac',cc');
x = u'*u; @test norm(ac'*x*ac-x+cc'*cc)/norm(x)/max(1.,norm(ac)^2) < reltol

@time u = plyapd(ac',cct');
x = u'*u; @test norm(ac'*x*ac-x+cct'*cct)/norm(x)/max(1.,norm(ac)^2) < reltol

@time u = plyapd(ac',cr');
x = u'*u; @test norm(ac'*x*ac-x+cr'*cr)/norm(x)/max(1.,norm(ac)^2) < reltol

@time u = plyapd(ar,bc);
x = u*u'; @test norm(ar*x*ar'-x+bc*bc')/norm(x)/max(1.,norm(ar)^2) < reltol

@time u = plyapd(ar',cc');
x = u'*u; @test norm(ar'*x*ar-x+cc'*cc)/norm(x)/max(1.,norm(ar)^2) < reltol

end
end

@testset "Positive discrete generalized Lyapunov equations" begin

reltol = eps(float(100))
a = -1+im; ee = 4; b = 2im; @time u = plyapd(a,ee,b)
@test abs(a*u*u'*a'-ee*u*u'*ee'+b*b') < reltol

reltol = eps(float(100f0))
a = -1f0+1f0*im; ee = 4f0; b = 2f0im; @time u = plyapd(a,ee,b)
@test abs(a*u*u'*a'-ee*u*u'*ee'+b*b')  < reltol

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ar = ar/(one(Ty) + norm(ar))
er = rand(Ty,n,n)
ar = er*ar
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
ac = ac/(one(Ty) + norm(ac))
ec = er+im*rand(Ty,n,n)
ac = ec*ac
br = rand(Ty,n,m)
brw = rand(Ty,n,n+m)
bc = br+im*rand(Ty,n,m)
bcw = brw+im*rand(Ty,n,n+m)
cr = rand(Ty,p,n)
crt = rand(Ty,n+p,n)
cc = cr+im*rand(Ty,p,n)
cct = crt+im*rand(Ty,n+p,n)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time u = plyapd(ar',er',cr');
x = u'*u; @test norm(ar'*x*ar-er'*x*er+cr'*cr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time u = plyapd(ar',er',crt');
x = u'*u; @test norm(ar'*x*ar-er'*x*er+crt'*crt)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time u = plyapd(ar,er,br);
x = u*u'; @test norm(ar*x*ar'-er*x*er'+br*br')/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time u = plyapd(ar,er,brw);
x = u*u'; @test norm(ar*x*ar'-er*x*er'+brw*brw')/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time u = plyapd(ac,ec,bc);
x = u*u'; @test norm(ac*x*ac'-ec*x*ec'+bc*bc')/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time u = plyapd(ac,ec,bcw);
x = u*u'; @test norm(ac*x*ac'-ec*x*ec'+bcw*bcw')/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time u = plyapd(ac',ec',cc');
x = u'*u; @test norm(ac'*x*ac-ec'*x*ec+cc'*cc)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time u = plyapd(ac',ec',cct');
x = u'*u; @test norm(ac'*x*ac-ec'*x*ec+cct'*cct)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time u = plyapd(ar,er,bc);
x = u*u'; @test norm(ar*x*ar'-er*x*er'+bc*bc')/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time u = plyapd(ar',er',cc');
x = u'*u; @test norm(ar'*x*ar-er'*x*er+cc'*cc)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time u = plyapd(ac,ec,br);
x = u*u'; @test norm(ac*x*ac'-ec*x*ec'+br*br')/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time u = plyapd(ac',ec',cr');
x = u'*u; @test norm(ac'*x*ac-ec'*x*ec+cr'*cr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol
end
end

@testset "Positive discrete 2x2 Lyapunov equations" begin

A = [-1.1 1.; -1. -1.]/10
E = [1. 1.; 0. 1.]
R = [1. 1.; 0. 1.]
reltol = eps(float(100))

U, scale, β, α = plyap2(A, R, adj = true, disc = true)
X = U'*U; @test norm(A'*X*A-X+scale^2*R'*R)/max(1,norm(X))/norm(A) < reltol &&
                norm(β*U-U*A)/max(1,norm(U))/norm(A) < reltol &&
                norm(α*U/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = plyap2(A, R, adj = false, disc = true)
X = U*U'; @test norm(A*X*A'-X+scale^2*R*R')/max(1,norm(X))/norm(A) < reltol &&
                norm(U*β-A*U)/max(1,norm(U))/norm(A) < reltol &&
                norm(U*α/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = pglyap2(A, E, R, adj = true, disc = true)
X = U'*U; @test norm(A'*X*A-E'*X*E+scale^2*R'*R)/max(1,norm(X))/norm(A) < reltol &&
                norm(β*U*E-U*A)/max(1,norm(U))/norm(A) < reltol &&
                norm(α*U*E/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = pglyap2(A, E, R, adj = false, disc = true)
X = U*U'; @test norm(A*X*A'-E*X*E'+scale^2*R*R')/max(1,norm(X))/norm(A) < reltol &&
                norm(A*U-E*U*β)/max(1,norm(U))/norm(A) < reltol &&
                norm(E*U*α/scale - scale*R)/max(1,norm(scale*R)) < reltol

A = [-1.1 1.; -1. -1.]/10; A = convert(Matrix{Float32},A)
E = [1. 1.; 0. 1.]; E = convert(Matrix{Float32},E)
R = [1. 1.; 0. 1.]; R = convert(Matrix{Float32},R)
reltol = eps(100f0)

U, scale, β, α = plyap2(A, R, adj = true, disc = true)
X = U'*U; @test norm(A'*X*A-X+scale^2*R'*R)/max(1,norm(X))/norm(A) < reltol &&
                norm(β*U-U*A)/max(1,norm(U))/norm(A) < reltol &&
                norm(α*U/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = plyap2(A, R, adj = false, disc = true)
X = U*U'; @test norm(A*X*A'-X+scale^2*R*R')/max(1,norm(X))/norm(A) < reltol &&
                norm(U*β-A*U)/max(1,norm(U))/norm(A) < reltol &&
                norm(U*α/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = pglyap2(A, E, R, adj = true, disc = true)
X = U'*U; @test norm(A'*X*A-E'*X*E+scale^2*R'*R)/max(1,norm(X))/norm(A) < reltol &&
                norm(β*U*E-U*A)/max(1,norm(U))/norm(A) < reltol &&
                norm(α*U*E/scale - scale*R)/max(1,norm(scale*R)) < reltol

U, scale, β, α = pglyap2(A, E, R, adj = false, disc = true)
X = U*U'; @test norm(A*X*A'-E*X*E'+scale^2*R*R')/max(1,norm(X))/norm(A) < reltol &&
                norm(A*U-E*U*β)/max(1,norm(U))/norm(A) < reltol &&
                norm(E*U*α/scale - scale*R)/max(1,norm(scale*R)) < reltol

end

@testset "Positive discrete Lyapunov equations - Schur form" begin

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ar = ar/(one(Ty) + norm(ar))
as,  = schur(ar)
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
ac = ac/(one(Ty) + norm(ac))
acs,  = schur(ac)
br = rand(Ty,n,m)
bc = br+im*rand(Ty,n,m)
cr = rand(Ty,p,n)
cc = cr+im*rand(Ty,p,n)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time u = plyaps(as,br,disc = true);
x = u*u'; @test norm(as*x*as'-x+br*br')/norm(x)/norm(as) < reltol

@time u = plyaps(as,ar,disc = true);
x = u*u'; @test norm(as*x*as'-x+ar*ar')/norm(x)/norm(as) < reltol

@time u = plyaps(as',cr',disc = true);
x = u'*u; @test norm(as'*x*as-x+cr'*cr)/norm(x)/norm(as) < reltol

@time u = plyaps(as',ar',disc = true);
x = u'*u; @test norm(as'*x*as-x+ar'*ar)/norm(x)/norm(as) < reltol

@time u = plyaps(acs,bc,disc = true);
x = u*u'; @test norm(acs*x*acs'-x+bc*bc')/norm(x)/norm(acs) < reltol

@time u = plyaps(acs',cc',disc = true);
x = u'*u; @test norm(acs'*x*acs-x+cc'*cc)/norm(x)/norm(as) < reltol

R = UpperTriangular(rand(Ty,n,n))
U = copy(R)
@time plyapds!(as, U, adj = false)
X = U*U'; @test norm(as*X*as'-X+R*R')/max(1,norm(X))/norm(as) < reltol

R = UpperTriangular(rand(Ty,n,n))
U = copy(R)
@time plyapds!(as,U,adj = true)
X = U'*U; @test norm(as'*X*as-X+R'*R)/max(1,norm(X))/norm(as) < reltol

R = UpperTriangular(rand(Complex{Ty},n,n))
U = copy(R)
@time plyapds!(acs,U,adj = false)
X = U*U'; @test norm(acs*X*acs'-X+R*R')/max(1,norm(X))/norm(acs) < reltol

R = UpperTriangular(complex(rand(Ty,n,n)))
U = copy(R)
@time plyapds!(acs,U,adj = true)
X = U'*U; @test norm(acs'*X*acs-X+R'*R)/max(1,norm(X))/norm(acs) < reltol

er = rand(Ty,n,n)
ar = er*ar
as, es = schur(ar,er)
ec = er+im*rand(Ty,n,n)
ac = ec*ac
acs, ecs = schur(ac,ec)

@time u = plyaps(as,es,br,disc = true);
x = u*u'; @test norm(as*x*as'-es*x*es'+br*br')/norm(x)/norm(as) < reltol

@time u = plyaps(as,es,ar,disc = true);
x = u*u'; @test norm(as*x*as'-es*x*es'+ar*ar')/norm(x)/norm(as) < reltol

@time u = plyaps(as',es',cr',disc = true);
x = u'*u; @test norm(as'*x*as-es'*x*es+cr'*cr)/norm(x)/norm(as) < reltol

@time u = plyaps(as',es',ar',disc = true);
x = u'*u; @test norm(as'*x*as-es'*x*es+ar'*ar)/norm(x)/norm(as) < reltol

@time u = plyaps(acs,ecs,bc,disc = true);
x = u*u'; @test norm(acs*x*acs'-ecs*x*ecs'+bc*bc')/norm(x)/norm(as) < reltol

@time u = plyaps(acs',ecs',cc',disc = true);
x = u'*u; @test norm(acs'*x*acs-ecs'*x*ecs+cc'*cc)/norm(x)/norm(as) < reltol

R = UpperTriangular(rand(Ty,n,n))
U = copy(R)
@time plyapds!(as,es,U,adj = false)
X = U*U'; @test norm(as*X*as'-es*X*es'+R*R')/max(1,norm(X))/norm(as) < reltol

R = UpperTriangular(rand(Ty,n,n))
U = copy(R)
@time plyapds!(as,es,U,adj = true)
X = U'*U; @test norm(as'*X*as-es'*X*es+R'*R)/max(1,norm(X))/norm(as) < reltol

R = UpperTriangular(rand(Complex{Ty},n,n)) # error
U = copy(R)
@time plyapds!(acs,ecs,U,adj = false)
X = U*U'; @test norm(acs*X*acs'-ecs*X*ecs'+R*R')/max(1,norm(X))/norm(acs) < reltol

R = UpperTriangular(rand(Complex{Ty},n,n)) # error
U = copy(R)
@time plyapds!(acs,ecs,U,adj = true)
X = U'*U; @test norm(acs'*X*acs-ecs'*X*ecs+R'*R)/max(1,norm(X))/norm(acs) < reltol
end
end
end

end


function minmeanmedianmax(vec)
    (minimum(vec),round(mean(vec),2),median(vec),maximum(vec))
end
function statcols(data)
    num_col = size(data,2)
    for ind in 1:num_col
        tmp = data[:,ind]
        print("min mean median max for $l class: ")
        println(minmeanmedianmax(tmp))
    end
end


module TestDualSimplicialSets
using Test

using LinearAlgebra: Diagonal
using SparseArrays, StaticArrays

using Catlab.CategoricalAlgebra.CSets
using CombinatorialSpaces

const Point2D = SVector{2,Float64}
const Point3D = SVector{3,Float64}

# 1D dual complex
#################

primal_s = DeltaSet1D()
add_vertices!(primal_s, 5)
add_edges!(primal_s, 1:4, repeat([5], 4))
s = DeltaDualComplex1D(primal_s)
@test nparts(s, :DualV) == nv(primal_s) + ne(primal_s)
@test nparts(s, :DualE) == 2 * ne(primal_s)

dual_v = elementary_duals(1,s,4)
@test dual_v == [edge_center(s, 4)]
@test elementary_duals(s, E(4)) == DualV(dual_v)

dual_es = elementary_duals(0,s,5)
@test length(dual_es) == 4
@test s[dual_es, :D_∂v0] == edge_center(s, 1:4)
@test elementary_duals(s, V(5)) == DualE(dual_es)

# 1D oriented dual complex
#-------------------------

primal_s = OrientedDeltaSet1D{Bool}()
add_vertices!(primal_s, 3)
add_edges!(primal_s, [1,2], [2,3], edge_orientation=[true,false])
s = OrientedDeltaDualComplex1D{Bool}(primal_s)
@test s[only(elementary_duals(0,s,1)), :D_edge_orientation] == true
@test s[only(elementary_duals(0,s,3)), :D_edge_orientation] == true

@test ∂(s, DualChain{1}([1,0,1])) isa DualChain{0}
@test d(s, DualForm{0}([1,1])) isa DualForm{1}
@test dual_boundary(1,s) == ∂(1,s)'
@test dual_derivative(0,s) == -d(0,s)'

# 1D embedded dual complex
#-------------------------

# Path graph on 3 vertices with irregular lengths.
primal_s = EmbeddedDeltaSet1D{Bool,Point2D}()
add_vertices!(primal_s, 3, point=[Point2D(1,0), Point2D(0,0), Point2D(0,2)])
add_edges!(primal_s, [1,2], [2,3], edge_orientation=true)
s = EmbeddedDeltaDualComplex1D{Bool,Float64,Point2D}(primal_s)
subdivide_duals!(s, Barycenter())
@test dual_point(s, edge_center(s, [1,2])) ≈ [Point2D(0.5,0), Point2D(0,1)]
@test volume(s, E(1:2)) ≈ [1.0, 2.0]
@test volume(s, elementary_duals(s, V(2))) ≈ [0.5, 1.0]
@test ⋆(0,s) ≈ Diagonal([0.5, 1.5, 1.0])
@test ⋆(1,s) ≈ Diagonal([1, 0.5])
@test ⋆(s, VForm([0,2,0]))::DualForm{1} ≈ DualForm{1}([0,3,0])

@test ∧(0,0,s, [1,2,3], [3,4,7]) ≈ [3,8,21]
@test ∧(s, VForm([1,2,3]), VForm([3,4,7]))::VForm ≈ VForm([3,8,21])
@test ∧(s, VForm([1,1,1]), EForm([2.5, 5.0]))::EForm ≈ EForm([2.5, 5.0])
@test ∧(s, VForm([1,1,0]), EForm([2.5, 5.0])) ≈ EForm([2.5, 2.5])
vform, eform = VForm([1.5, 2, 2.5]), EForm([13, 7])
@test ∧(s, vform, eform) ≈ ∧(s, eform, vform)

# Path graph on 5 vertices with regular lengths.
#
# Equals the graph Laplacian of the underlying graph, except at the boundary.
# Note that the DEC Laplace-Beltrami operator is *not* symmetric.
primal_s = EmbeddedDeltaSet1D{Bool,Point2D}()
add_vertices!(primal_s, 5, point=[Point2D(i,0) for i in -2:2])
add_edges!(primal_s, 1:4, 2:5, edge_orientation=true)
s = EmbeddedDeltaDualComplex1D{Bool,Float64,Point2D}(primal_s)
subdivide_duals!(s, Barycenter())
@test Δ(s, VForm([0,0,1,0,0])) ≈ VForm([0,1,-2,1,0])
@test Δ(0,s) ≈ [-2  2  0  0  0;
                 1 -2  1  0  0;
                 0  1 -2  1  0;
                 0  0  1 -2  1;
                 0  0  0  2 -2]

# 2D dual complex
#################

# Triangulated square.
primal_s = DeltaSet2D()
add_vertices!(primal_s, 4)
glue_triangle!(primal_s, 1, 2, 3)
glue_triangle!(primal_s, 1, 3, 4)
s = DeltaDualComplex2D(primal_s)
@test nparts(s, :DualV) == nv(primal_s) + ne(primal_s) + ntriangles(primal_s)
@test nparts(s, :DualE) == 2*ne(primal_s) + 6*ntriangles(primal_s)
@test nparts(s, :DualTri) == 6*ntriangles(primal_s)
@test primal_vertex(s, subsimplices(s, Tri(1)))::V == V([1,1,2,2,3,3])

dual_vs = elementary_duals(2,s,2)
@test dual_vs == [triangle_center(s,2)]
@test elementary_duals(s, Tri(2)) == DualV(dual_vs)
@test s[elementary_duals(1,s,2), :D_∂v1] == [edge_center(s,2)]
@test s[elementary_duals(1,s,3), :D_∂v1] == repeat([edge_center(s,3)], 2)
@test [length(elementary_duals(s, V(i))) for i in 1:4] == [4,2,4,2]
@test dual_triangle_vertices(s, 1) == [1,7,10]

# 2D oriented dual complex
#-------------------------

# Triangulated square with consistent orientation.
primal_s = OrientedDeltaSet2D{Bool}()
add_vertices!(primal_s, 4)
glue_triangle!(primal_s, 1, 2, 3, tri_orientation=true)
glue_triangle!(primal_s, 1, 3, 4, tri_orientation=true)
primal_s[:edge_orientation] = true
s = OrientedDeltaDualComplex2D{Bool}(primal_s)
@test sum(s[:D_tri_orientation]) == nparts(s, :DualTri) ÷ 2
@test [sum(s[elementary_duals(0,s,i), :D_tri_orientation])
       for i in 1:4] == [2,1,2,1]
@test sum(s[elementary_duals(1,s,3), :D_edge_orientation]) == 1

for k in 0:1
  @test dual_boundary(2-k,s) == (-1)^k * ∂(k+1,s)'
end
for k in 1:2
  # Desbrun, Kanso, Tong 2008, Equation 4.2.
  @test dual_derivative(2-k,s) == (-1)^k * d(k-1,s)'
end

# 2D embedded dual complex
#-------------------------

# Single triangle: numerical example from Gillette's notes on DEC, §2.13.
#
# Compared with Gillette, edges #2 and #3 are swapped in the ordering, which
# changes the discrete exterior derivative and other operators. The numerical
# values remain the same, as we verify.
primal_s = EmbeddedDeltaSet2D{Bool,Point2D}()
add_vertices!(primal_s, 3, point=[Point2D(0,0), Point2D(1,0), Point2D(0,1)])
glue_triangle!(primal_s, 1, 2, 3, tri_orientation=true)
primal_s[:edge_orientation] = true
s = EmbeddedDeltaDualComplex2D{Bool,Float64,Point2D}(primal_s)

subdivide_duals!(s, Barycenter())
@test dual_point(s, triangle_center(s, 1)) ≈ Point2D(1/3, 1/3)
@test volume(s, Tri(1)) ≈ 1/2
@test volume(s, elementary_duals(s, V(1))) ≈ [1/12, 1/12]
@test [sum(volume(s, elementary_duals(s, V(i)))) for i in 1:3] ≈ [1/6, 1/6, 1/6]
@test ⋆(0,s) ≈ Diagonal([1/6, 1/6, 1/6])
@test ⋆(1,s) ≈ Diagonal([√5/6, 1/6, √5/6])
@test ⋆(s, VForm([1,2,3]))::DualForm{2} ≈ DualForm{2}([1/6, 1/3, 1/2])
@test isapprox(δ(1,s), [ 2.236  0  2.236;
                        -2.236  1  0;
                         0     -1 -2.236], atol=1e-3)
@test δ(s, EForm([0.5,1.5,0.5])) isa VForm

@test ∧(s, VForm([2,2,2]), TriForm([2.5]))::TriForm ≈ TriForm([2.5])
vform, triform = VForm([1.5, 2, 2.5]), TriForm([7.5])
@test ∧(s, vform, triform) ≈ ∧(s, triform, vform)
eform1, eform2 = EForm([1.5, 2, 2.5]), EForm([3, 7, 10])
@test ∧(s, eform1, eform1)::TriForm ≈ TriForm([0])
@test ∧(s, eform1, eform2) ≈ -∧(s, eform2, eform1)

X♭, α = EForm([1.5, 2, 2.5]), DualForm{1}([3, 7, 10])
@test interior_product(s, X♭, α) isa DualForm{0}
@test length(interior_product_flat(1,s, X♭.data, α.data)) == 1
@test lie_derivative(s, X♭, α) isa DualForm{1}
@test length(lie_derivative_flat(1,s, X♭.data, α.data)) == 3

subdivide_duals!(s, Circumcenter())
@test dual_point(s, triangle_center(s, 1)) ≈ Point2D(1/2, 1/2)
@test ⋆(0,s) ≈ Diagonal([1/4, 1/8, 1/8])
@test ⋆(1,s) ≈ Diagonal([0.5, 0.0, 0.5])
@test δ(1,s) ≈ [ 2  0  2;
                -4  0  0;
                 0  0 -4]

subdivide_duals!(s, Incenter())
@test dual_point(s, triangle_center(s, 1)) ≈ Point2D(1/(2+√2), 1/(2+√2))
@test isapprox(⋆(0,s), Diagonal([0.146, 0.177, 0.177]), atol=1e-3)
@test isapprox(⋆(1,s), Diagonal([0.359, 0.207, 0.359]), atol=1e-3)
@test isapprox(δ(1,s), [ 2.449  0      2.449;
                        -2.029  1.172  0;
                         0     -1.172 -2.029], atol=1e-3)

end


crabs = parse.(Int, split(readline("day7/crabs.txt"), ','))

move_crabs(crabs, to) = sum(abs.(crabs .- to))
p1 = minimum(move_crabs(crabs, i) for i in 0:maximum(crabs))

@show p1

move_crabs2(crabs, to) = sum(sum(1:abs(crab - to)) for crab in crabs)
p2 = minimum(move_crabs2(crabs, i) for i in 0:maximum(crabs))

@show p2


using DiffEqFlux, GalacticOptim, OrdinaryDiffEq

function newtons_cooling(du, u, p, t)
    temp = u[1]
    k, temp_m = p
    du[1] = dT = -k*(temp-temp_m)
  end

function true_sol(du, u, p, t)
    true_p = [log(2)/8.0, 100.0]
    newtons_cooling(du, u, true_p, t)
end

function dudt_(u,p,t)
    ann(u, p).* u
end

cb = function (p,l,pred;doplot=false) #callback function to observe training
    display(l)
    # plot current prediction against data
    if doplot
      pl = scatter(t,ode_data[1,:],label="data")
      scatter!(pl,t,pred[1,:],label="prediction")
      display(plot(pl))
    end
    return false
end

u0 = Float32[200.0]
datasize = 30
tspan = (0.0f0, 1.5f0)

t = range(tspan[1], tspan[2], length=datasize)
true_prob = ODEProblem(true_sol, u0, tspan)
ode_data = Array(solve(true_prob, Tsit5(), saveat=t))

ann = FastChain(FastDense(1,8,tanh), FastDense(8,1,tanh))
pp = initial_params(ann)
prob = ODEProblem{false}(dudt_, u0, tspan, pp)

function predict_adjoint(fullp, time_batch)
    Array(solve(prob, Tsit5(), p = fullp, saveat = time_batch))
end

function loss_adjoint(fullp, batch, time_batch)
    pred = predict_adjoint(fullp,time_batch)
    sum(abs2, batch .- pred), pred
end

k = 10
train_loader = Flux.Data.DataLoader((ode_data, t), batchsize = k)

numEpochs = 300
l1 = loss_adjoint(pp, train_loader.data[1], train_loader.data[2])[1]

optfun = OptimizationFunction((θ, p, batch, time_batch) -> loss_adjoint(θ, batch, time_batch), GalacticOptim.AutoZygote())
optprob = OptimizationProblem(optfun, pp)
using IterTools: ncycle
res1 = GalacticOptim.solve(optprob, ADAM(0.05), ncycle(train_loader, numEpochs), cb = cb, maxiters = numEpochs)
@test 10res1.minimum < l1


# Templates of causal queries
module Queries

using ..Condition: cnd, |ᶜ, ConditionException, condf
using ..Interventions: intervene, Intervention, AbstractIntervention, |ᵈ

export cf, fc, cndint

"""
`cf(x, y, i::Intervention)`

Counterfactual: Given that `y` is true, what would `x` had been `i` been true

```math
(x \\mid y) \\mid \\text{do}(I)))

# Example

```
using OmegaCore, Distributions
p = 0.7
q = 0.3
order = 1 ~ Bernoulli(p)        # court orders execution wiht prob. p
Anerves = 2 ~ Bernoulli(q)      # rifleman A shoots with prob. q out of nervousness
Ashoots = order |ₚ Anerves
Bshoots = order
dead = Ashoots |ₚ Bshoots

# If the prisoner is dead, then the prisoner would be dead even if rifleman A had not shot
dead_cf = cf(dead, dead, Ashoots => false)
dead_cf = (dead |ᶜ Ashoots) |ᵈ (Ashoots => false)
```
"""
cf(x, y, i::AbstractIntervention) = (x |ᵈ i) |ᶜ y
# cf(x, y, i::AbstractIntervention) = intervene(cnd(x, y), i)
# cf(x, y, i::AbstractIntervention) =
#   let xi = x |ᵈ i
#     ω -> condf(ω, xi, y)
#   end
#   #  Bool(y(ω)) ? (x |ᵈ i)(ω) : throw(ConditionException)
cf(x, y, i...) = cf(x, y, Intervention(i...))

"""
`fc(x, y, i::Intervention)`

If in a the hypothetical world `i`, `y` were true, then what is the value of `x`

```math
x \\mid (y \\mid \\text{do}(I)))
```
# Example
```
using Distributions
order = 1 ~ Bernoulli(0.5)
Anerves = 2 ~ Bernoulli(0.5)
Ashoots = order |ₚ Anerves
Bshoots = order
dead = Ashoots |ₚ Bshoots

# If hypothetically rifleman `C` firing were to kill `B`, then `A` did not fire.
Afc = fc(Bshoots, dead, Ashoots => false)
```

"""
# fc(x, y, i::AbstractIntervention) = cnd(x, intervene(y, i))
fc(x, y, i::AbstractIntervention) =
  let yi = y |ᵈ i
    ω -> condf(ω, x, yi)
  end
fc(x, y, i...) = fc(x, y, Intervention(i...))

"""
`cndint(i::RV{<:Intervention}, y::RV) = i |ᶜ ω -> yi = y |ᵈ i(ω); yi(ω)`

Conditional intervention: conditional random variable over intervention `i`
such that `y | do(i)` is true.

# Example
```
using Distributions
x = 1 ~ Normal(0, 1)
y = 2 ~ Normal(0, 1)
e(ω) = x(ω) + y(ω) >= 11.0

# Construct distribution over interventions
x_ = 3 ~ Uniform(0, 10)
i = ω -> x => x_(ω)
icond = cndint(i, e)

# Sample from distribution over interventions such that ...
# doing that intervention will ensure that `e` is true
julia> randsample(icond, 3)
julia> randsample(icond, 3)
3-element Vector{Pair{Member{Normal{Float64}, Vector{Int64}}, Float64}}:
 [1]@Normal{Float64}(μ=0.0, σ=1.0) => 9.37047390742177
 [1]@Normal{Float64}(μ=0.0, σ=1.0) => 9.409378254490608
 [1]@Normal{Float64}(μ=0.0, σ=1.0) => 8.769913889437522

```

"""
cndint(i, y) = i |ᶜ (ω -> (y |ᵈ i(ω))(ω))

end


"""
    momentum_space_harmonic_potential(M::Integer, v_ho::Real) -> w

Set up a harmonic potential for use with momentum space Hamiltonians:
```math
\\begin{aligned}
w(k) & =  \\frac{1}{M} \\mathrm{DFT}[V_{ext}]_{k}  ,\\\\
V_\\mathrm{ext}(x) &= v_\\mathrm{ho} \\,x^2 ,
\\end{aligned}
```
where
``\\mathrm{DFT}[…]_k`` is a discrete Fourier transform performed by `fft()[k%M + 1]`.
"""
function momentum_space_harmonic_potential(M::Integer, v::Real)
    v = float(v)
    # Set up potential like in Real1DEP
    is = range(-fld(M,2); length=M) # [-M÷2, M÷2) including left boundary
    js = shift_lattice(is) # shifted such that js[1] = 0
    real_potential = [v*j^2 for j in js]
    mom_potential = fft(real_potential)
    # This should never fail for a harmonic oscillator, but it's best to check just in case.
    for x in mom_potential
        @assert iszero(real(x)) || abs(imag(x) / real(x)) < sqrt(eps(v))
    end
    # Make sure it's completely symmetric. It should be, but round-off errors can sometimes
    # make it non-symmetric.
    for i in 1:M÷2
        mom_potential[M - i + 1] = mom_potential[i + 1]
    end
    return SVector{M}(1/M .* real.(mom_potential))
end

"""
    HubbardMom1DEP(address; u=1.0, t=1.0, v_ho=1.0, dispersion=hubbard_dispersion)

Implements a one-dimensional Bose Hubbard chain in momentum space with harmonic external
potential.

```math
Ĥ = \\sum_{k} ϵ_k n_k + \\frac{u}{M}\\sum_{kpqr} a^†_{r} a^†_{q} a_p a_k δ_{r+q,p+k}
            + V̂_\\mathrm{ho} ,
```
where
```math
\\begin{aligned}
V̂_\\mathrm{ho} & = \\frac{1}{M} \\sum_{p,q}  \\mathrm{DFT}[V_{ext}]_{p-q} \\,
                    a^†_{p} a_q ,\\\\
V_\\mathrm{ext}(x) &= v_\\mathrm{ho} \\,x^2 ,
\\end{aligned}
```
is an external harmonic potential in momentum space,
``\\mathrm{DFT}[…]_k`` is a discrete Fourier transform performed by `fft()[k%M + 1]`, and
`M == num_modes(address)`.

# Arguments

* `address`: the starting address, defines number of particles and sites.
* `u`: the interaction parameter.
* `t`: the hopping strength.
* `dispersion`: defines ``ϵ_k =``` t*dispersion(k)`
    - [`hubbard_dispersion`](@ref): ``ϵ_k = -2t \\cos(k)``
    - [`continuum_dispersion`](@ref): ``ϵ_k = tk^2``
* `v_ho`: strength of the external harmonic oscillator potential ``v_\\mathrm{ho}``.

See also [`HubbardMom1D`](@ref), [`HubbardReal1DEP`](@ref),
[`Transcorrelated1D`](@ref), [`Hamiltonians`](@ref).
"""
struct HubbardMom1DEP{TT,M,AD<:AbstractFockAddress,U,T} <: AbstractHamiltonian{TT}
    add::AD # default starting address, should have M modes, U and T are model parameters
    ks::SVector{M,TT} # values for k
    kes::SVector{M,TT} # values for kinetic energy
    ep::SVector{M,TT} # external potential
end

function HubbardMom1DEP(
    add::Union{BoseFS,FermiFS2C};
    u=1.0, t=1.0, dispersion = hubbard_dispersion, v_ho=1.0,
)
    M = num_modes(add)
    U, T, V = promote(float(u), float(t), float(v_ho))
    step = 2π/M
    if isodd(M)
        start = -π*(1+1/M) + step
    else
        start = -π + step
    end
    kr = range(start; step = step, length = M)
    ks = SVector{M}(kr)
    kes = SVector{M}(T .* dispersion.(kr))

    potential = momentum_space_harmonic_potential(M, V)

    return HubbardMom1DEP{typeof(U),M,typeof(add),U,T}(add, ks, kes, potential)
end

function Base.show(io::IO, h::HubbardMom1DEP)
    print(io, "HubbardMom1DEP($(h.add); u=$(h.u), t=$(h.t))")
end

function starting_address(h::HubbardMom1DEP)
    return h.add
end
function get_offdiagonal(h::HubbardMom1DEP{<:Any,<:Any,F}, add::F, i) where {F}
    return offdiagonals(h, add)[i]
end
function num_offdiagonals(h::HubbardMom1DEP{<:Any,<:Any,F}, add::F) where {F}
    return length(offdiagonals(h, add))
end

LOStructure(::Type{<:HubbardMom1DEP{<:Real}}) = IsHermitian()

Base.getproperty(h::HubbardMom1DEP, s::Symbol) = getproperty(h, Val(s))
Base.getproperty(h::HubbardMom1DEP, ::Val{:ep}) = getfield(h, :ep)
Base.getproperty(h::HubbardMom1DEP, ::Val{:ks}) = getfield(h, :ks)
Base.getproperty(h::HubbardMom1DEP, ::Val{:kes}) = getfield(h, :kes)
Base.getproperty(h::HubbardMom1DEP, ::Val{:add}) = getfield(h, :add)
Base.getproperty(h::HubbardMom1DEP{<:Any,<:Any,<:Any,U}, ::Val{:u}) where {U} = U
Base.getproperty(h::HubbardMom1DEP{<:Any,<:Any,<:Any,<:Any,T}, ::Val{:t}) where {T} = T

ks(h::HubbardMom1DEP) = getfield(h, :ks)

@inline function momentum_transfer_diagonal(
    h::HubbardMom1DEP{<:Any,M,<:BoseFS}, map
) where {M}
    return h.u / 2M * momentum_transfer_diagonal(map)
end
@inline function momentum_transfer_diagonal(
    h::HubbardMom1DEP{<:Any,M,<:FermiFS2C}, map_a, map_b
) where {M}
    return h.u / 2M * momentum_transfer_diagonal(map_a, map_b)
end

@inline function diagonal_element(h::HubbardMom1DEP, add::BoseFS)
    map = OccupiedModeMap(add)
    return dot(h.kes, map) + momentum_transfer_diagonal(h, map) +
        momentum_external_potential_diagonal(h.ep, add, map)
end
@inline function diagonal_element(h::HubbardMom1DEP, add::FermiFS2C)
    c1, c2 = add.components
    map_a = OccupiedModeMap(c1)
    map_b = OccupiedModeMap(c2)
    return dot(h.kes, map_a) + dot(h.kes, map_b) +
        momentum_transfer_diagonal(h, map_a, map_b) +
        momentum_external_potential_diagonal(h.ep, c1, map_a) +
        momentum_external_potential_diagonal(h.ep, c2, map_b)
end

###
### offdiagonals
###
struct OffdiagonalsBoseMom1DEP{
    A<:BoseFS,T,H<:AbstractHamiltonian{T},O<:OccupiedModeMap
} <: AbstractOffdiagonals{A,T}
    hamiltonian::H
    address::A
    num_mom::Int
    num_ep::Int
    map::O
end

function offdiagonals(h::HubbardMom1DEP, a::BoseFS)
    M = num_modes(a)
    map = OccupiedModeMap(a)
    singlies = length(map)
    doublies = count(i -> i.occnum ≥ 2, map)
    num_mom = singlies * (singlies - 1) * (M - 2) + doublies * (M - 1)
    num_ep = singlies * (M - 1)
    return OffdiagonalsBoseMom1DEP(h, a, num_mom, num_ep, map)
end

function Base.getindex(s::OffdiagonalsBoseMom1DEP{A,T}, i)::Tuple{A,T} where {A,T}
    @boundscheck begin
        1 ≤ i ≤ length(s) || throw(BoundsError(s, i))
    end
    M = num_modes(s.address)
    if i ≤ s.num_mom
        new_address, onproduct = momentum_transfer_excitation(s.address, i, s.map)
        matrix_element = s.hamiltonian.u/(2*M) * onproduct
    else
        i -= s.num_mom
        new_address, matrix_element = momentum_external_potential_excitation(
            s.hamiltonian.ep, s.address, i, s.map
        )
    end
    return (new_address, matrix_element)
end

Base.size(s::OffdiagonalsBoseMom1DEP) = (s.num_mom + s.num_ep,)

struct OffdiagonalsFermiMom1D2CEP{
    A<:FermiFS2C,T,H<:AbstractHamiltonian{T},O1,O2
} <: AbstractOffdiagonals{A,T}
    hamiltonian::H
    address::A
    num_mom::Int
    num_ep_a::Int
    num_ep_b::Int
    map_a::O1
    map_b::O2
end

function offdiagonals(h::HubbardMom1DEP, a::FermiFS2C)
    M = num_modes(a)
    comp_a, comp_b = a.components
    N1 = num_particles(comp_a)
    N2 = num_particles(comp_b)
    map_a = OccupiedModeMap(comp_a)
    map_b = OccupiedModeMap(comp_b)
    num_mom = N1 * N2 * (M - 1)
    num_ep_a = N1 * (M - 1)
    num_ep_b = N2 * (M - 1)

    return OffdiagonalsFermiMom1D2CEP(h, a, num_mom, num_ep_a, num_ep_b, map_a, map_b)
end

Base.size(s::OffdiagonalsFermiMom1D2CEP) = (s.num_mom + s.num_ep_a + s.num_ep_b,)

function Base.getindex(s::OffdiagonalsFermiMom1D2CEP{A,T}, i)::Tuple{A,T} where {A,T}
    @boundscheck begin
        1 ≤ i ≤ length(s) || throw(BoundsError(s, i))
    end
    c1, c2 = s.address.components
    M = num_modes(s.address)
    if i ≤ s.num_mom
        new_c1, new_c2, onproduct = momentum_transfer_excitation(c1, c2, i, s.map_a, s.map_b)
        new_address = CompositeFS(new_c1, new_c2)
        matrix_element = s.hamiltonian.u/M * onproduct
    elseif i ≤ s.num_mom + s.num_ep_a
        i -= s.num_mom
        new_c1, matrix_element = momentum_external_potential_excitation(
            s.hamiltonian.ep, c1, i, s.map_a
        )
        new_address = CompositeFS(new_c1, c2)
    else
        i -= s.num_mom + s.num_ep_a
        new_c2, matrix_element = momentum_external_potential_excitation(
            s.hamiltonian.ep, c2, i, s.map_b
        )
        new_address = CompositeFS(c1, new_c2)
    end
    return (new_address, matrix_element)
end


#=----------------------------------------------------------------------------

Copyright (c) 2015 Peter Kovesi
pk@peterkovesi.com

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

The Software is provided "as is", without warranty of any kind.

----------------------------------------------------------------------------=#

# Definition of some utility functions that are used by the colour map
# generation code.  These functions will eventually move into an image
# processing package that I am slowing building

export bbspline, pbspline, gaussfilt1d, interp1, matprint
export normalise, normalize
export histtruncate
export sineramp, circlesineramp

#-----------------------------------------------------------------------
"""
bbspline - Basic B-spline
```
Usage:  S = bbspline(P, k, N)

Arguments:   P - [dim x Npts] array of control points
             k - order of spline (>= 2).
                 k = 2: Linear
                 k = 3: Quadratic, etc
             N - Optional number of points to evaluate along
                 spline. Defaults to 100.

Returns:     S - spline curve  [dim x N] spline points
```
See also: pbspline
"""
function bbspline(P::AbstractMatrix, k::Int, N::Int = 100)
    # PK Jan 2014

    (dim, np1) = size(P)
    n = np1-1

    @assert k >= 2   "Spline order must be 2 or greater"
    @assert np1 >= k "No of control points must be >= k"
    @assert N >= 2   "Spline must be evaluated at 2 or more points"

    # Set up open uniform knot vector from 0 - 1.
    # There are k repeated knots at each end.
    ti = collect(0:(k+n - 2*(k-1)))'
    ti = ti/ti[end]
    ti = [repeat([ti[1]], 1, k-1)  ti  repeat([ti[end]], 1, k-1)]

    nK = length(ti)

    # Generate values of t that the spline will be evaluated at
    dt = (ti[end]-ti[1])/(N-1)
    t = collect(ti[1]:dt:ti[end])'

    # Build complete array of basis functions.  We maintain two
    # arrays, one storing the basis functions at the current level of
    # recursion, and one storing the basis functions from the previous
    # level of recursion
    Blast = Array{Array{Float64}}(undef, nK-1)
    B = Array{Array{Float64}}(undef, nK-1)

    # 1st level of recursive construction
    for i = 1:nK-1
        if ti[i] < ti[i+1]
            Blast[i] = (t .>= ti[i]) .& (t .< ti[i+1])
        else
            Blast[i] = zeros(1,N)
        end
    end

    # Subsequent levels of recursive basis construction.  Note the logic to
    # handle repeated knot values where ti[i] == ti[i+1]
    for ki = 2:k
        for i = 1:nK-ki
            if (ti[i+ki-1] - ti[i]) < eps()
                V1 = zeros(1,N)
            else
                V1 = (t .- ti[i])/(ti[i+ki-1] - ti[i]) .* Blast[i]
            end

            if (ti[i+ki] - ti[i+1]) < eps()
                V2 = zeros(1,N)
            else
                V2 = (ti[i+ki] .- t)/(ti[i+ki] - ti[i+1]) .* Blast[i+1]
            end

            B[i] = V1 + V2

            # This is the ideal equation that the code above implements
            # B[i,ki] = (t - ti[i])/(ti[i+ki-1] - ti[i]) .* B[i,ki-1] + ...
            #           (ti[i+ki] - t)/(ti[i+ki] - ti[i+1]) .* B[i+1,ki-1]
        end
        # Swap B and Blast, but only if this is not the last iteration
        if ki < k
            tmp = Blast
            Blast = B
            B = tmp
        end
    end

    # Apply basis functions to the control points
    S = zeros(dim, length(t))

    for d = 1:dim, i = 1:np1
        S[d:d,:] += P[d,i]*B[i]
    end

    # Set the last point of the spline. This is not evaluated by the code above
    # because the basis functions are defined from ti[i] <= t < ti[i+1]
    S[:,end] = P[:,end]

    return S
end

#-------------------------------------------------------------------------------
"""
pbspline - Basic Periodic B-spline
```
Usage:  S = pbspline(P, k, N)

Arguments:   P - [dim x Npts] array of control points
             k - order of spline (>= 2).
                 k = 2: Linear
                 k = 3: Quadratic, etc
             N - Optional number of points to evaluate along
                 spline. Defaults to 100.

Returns:     S - spline curve  [dim x N] spline points
```
Note that the spline points are rotated so that the first point is as
close as possible to the first control point.  This is important for
the formation of cyclic paths in colour space.

See also: bbspline
"""
function pbspline(Pin::AbstractArray, k::Int, N::Int = 100)
    # PK March 2014
    # Needs a bit of tidying up and checking on domain of curve
    # Should be merged with bbspline
    # Issue: If N is too small one can get out of bounds errors at line 193

    P = copy(Pin)  # Make a copy because we will be altering P

    # For a closed spline check if 1st and last control points match.  If not
    # add another control point so that they do match
    if norm(P[:,1] - P[:,end]) > 0.01
        P = [P P[:,1]]
    end

    # Now add k - 1 control points that wrap over the first control points
    P = [P P[:,2:2+k-1]]

    (dim, np1) = size(P)
    n = np1-1

    @assert k >= 2   "Spline order must be 2 or greater"
    @assert np1 >= k "No of control points must be >= k"
    @assert N >= 2   "Spline must be evaluated at 2 or more points"

    # Form a uniform sequence. Number of knot points is m + 1 where m = n + k + 1
    ti = collect(0:(n+k+1))/(n+k+1)'
    nK = length(ti)

    # Domain of curve is [ti_k to ti_n] or [ti_(k+1) to ti_(n+1)] ???
    tstart = ti[k]
    tend = ti[n]

    dt = (tend-tstart)/(N-1)
    t = collect(tstart:dt:tend)'

    # Build complete array of basis functions.  We maintain two
    # arrays, one storing the basis functions at the current level of
    # recursion, and one storing the basis functions from the previous
    # level of recursion
    Blast = Array{Array{Float64}}(undef, nK-1)
    B = Array{Array{Float64}}(undef, nK-1)

    # 1st level of recursive construction
    for i = 1:nK-1
        if t[i] < ti[i+1]
            Blast[i] = (t .>= ti[i]) .& (t .< ti[i+1])
        else
            Blast[i] = zeros(1,N)
        end
    end

    # Subsequent levels of recursive basis construction.  Note the logic to
    # handle repeated knot values where ti[i] == ti[i+1]
    for ki = 2:k
        for i = 1:nK-ki

            if (ti[i+ki-1] - ti[i]) < eps()
                V1 = zeros(1,N)
            else
                V1 = (t .- ti[i])/(ti[i+ki-1] - ti[i]) .* Blast[i]
            end

            if (ti[i+ki] - ti[i+1]) < eps()
                V2 = zeros(1,N)
            else
                V2 = (ti[i+ki] .- t)/(ti[i+ki] - ti[i+1]) .* Blast[i+1]
            end

            B[i] = V1 + V2

            #       This is the ideal equation that the code above implements
            #       B[i,ki] = (t - ti[i])/(ti[i+ki-1] - ti[i]) .* B[i,ki-1] + ...
            #                (ti[i+ki] - t)/(ti[i+ki] - ti[i+1]) .* B[i+1,ki-1]
        end
        # Swap B and Blast, but only if this is not the last iteration
        if ki < k
            tmp = Blast
            Blast = B
            B = tmp
        end
    end

    # Apply basis functions to the control points
    S = zeros(dim, length(t))

    for d = 1:dim, i = 1:np1
        S[d:d,:] += P[d,i]*B[i]
    end

    # Finally, because of the knot arrangements, the start of the spline may not
    # be close to the first control point if the spline order is 3 or greater.
    # Normally for a closed spline this is irrelevant.  However for our purpose
    # of using closed b splines to form paths in a colour space this is important to
    # us.  The simple brute force solution used here is to search through the
    # spline points for the point that is closest to the 1st control point and
    # then rotate the spline points accordingly

    distsqrd = zeros(size(S[1,:]))
    for d = 1:dim
        distsqrd += (S[d,:] .- P[d,1]).^2
    end

    ind = argmin(distsqrd)

    return S = circshift(S, [0, -ind+1])
end

#----------------------------------------------------------------------------
# Function for applying a 1D Gaussian filter to a signal. Filtering at
# ends are done using zero padding
#
# Usage: sm = gaussfilt1d(s::Array, sigma::Real)

function gaussfilt1d(s::Array, sigma::Real)

    N = length(s)

    r = ceil(Int, 3*sigma)    # Determine filter size
    fw = 2*r + 1

    # Construct filter
    f = Float64[exp(-x.^2/(2*sigma)) for x = -r:r]
    f = f/sum(f)

    sm = zeros(size(s))

    # Filter centre section
    for i = r+1:N-r, k = 1:fw
        sm[i] += f[k] * s[i+k-r-1]
    end

    # Filter start section of array using 0 padding
    for i = 1:r, k = 1:fw
        ind = i+k-r-1
        if ind >= 1 && ind <= N
            sm[i] += f[k] * s[ind]
        end
    end

    # Filter end section of array using 0 padding
    for i = N-r+1:N, k = 1:fw
        ind = i+k-r-1
        if ind >= 1 && ind <= N
            sm[i] += f[k] * s[ind]
        end
    end

    return sm
end

#----------------------------------------------------------------------------
"""
Simple 1D linear interpolation of an array of data

```
 Usage:  yi = interp1(x, y, xi)

Arguments:  x - Array of coordinates at which y is defined.
            y - Array of values at coordinates x.
           xi - Coordinate locations at which you wish to interpolate y values.

Returns:   yi - Values linearly interpolated from y at xi.

```
Interpolates y, defined at values x, at locations xi and returns the
corresponding values as yi

x is assumed increasing but not necessarily equi-spaced.
xi values do not need to be sorted.

If any xi are outside the range of x then the corresponding value of
yi is set to the appropriate end value of y.

"""
function interp1(x, y, xi)

    N = length(xi)
    yi = zeros(size(xi))

    minx = minimum(x)
    maxx = maximum(x)

    for i = 1:N
        # Find interval in x that each xi lies within and interpolate its value

        if xi[i] <= minx
            yi[i] = y[1]

        elseif xi[i] >= maxx
            yi[i] = y[end]

        else
            left = maximum(findall(x .<= xi[i]))
            right = minimum(findall(x .> xi[i]))

            yi[i] = y[left] +  (xi[i]-x[left])/(x[right]-x[left]) * (y[right] - y[left])
        end
    end

    return yi
end

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

function matprint(arr, str="")

    @printf("\n%s\n",str)

    if ndims(arr) == 1
        (rows,) = size(arr)
        @printf("[")
        for r=1:rows
            if r < rows
                @printf("%.3f, ",arr[r])
            else
                @printf("%.3f",arr[r])
            end
        end
        @printf("]\n")

    elseif ndims(arr) == 2
        (rows,cols) = size(arr)

        @printf("[")
        for r=1:rows
            if r>1
                @printf(" ")
            end
            for c=1:cols
                @printf("%.3f ",arr[r,c])
            end
            if r<rows
                @printf("\n")
            else
                @printf("]\n")
            end
        end
    end
end

#----------------------------------------------------------------------
"""
normalise/normalize - Normalises image values to 0-1, or to desired mean and variance

```
Usage 1:      nimg = normalise(img)
```
Offsets and rescales image so that the minimum value is 0
and the maximum value is 1.

```
Usage 2:      nimg = normalise(img, reqmean, reqvar)

Arguments:  img     - A grey-level input image.
            reqmean - The required mean value of the image.
            reqvar  - The required variance of the image.
```
Offsets and rescales image so that nimg has mean reqmean and variance
reqvar.
"""
function normalise(img::Array)
    # Normalise 0 - 1
    n = img .- minimum(img)
    return n = n/maximum(n)
end

# Normalise to desired mean and variance
function normalise(img::Array, reqmean::Real, reqvar::Real)
    n = img .- mean(img)
    n = n/std(img)      # Zero mean, unit std dev
    return n = reqmean .+ n*sqrt(reqvar)
end

# For those who spell normalise with a 'z'
"""
normalize - Normalizes image values to 0-1, or to desired mean and variance
```
Usage 1:      nimg = normalize(img)
```
Offsets and rescales image so that the minimum value is 0
and the maximum value is 1.
```
Usage 2:      nimg = normalize(img, reqmean, reqvar)

Arguments:  img     - A grey-level input image.
            reqmean - The required mean value of the image.
            reqvar  - The required variance of the image.
```
Offsets and rescales image so that nimg has mean reqmean and variance
reqvar.
"""
function normalize(img::Array)
    return normalise(img)
end

function normalize(img::Array, reqmean::Real, reqvar::Real)
    return normalise(img, reqmean, reqvar)
end

#----------------------------------------------------------------------
"""
histtruncate - Truncates ends of an image histogram.

Function truncates a specified percentage of the lower and
upper ends of an image histogram.

This operation allows grey levels to be distributed across
the primary part of the histogram.  This solves the problem
when one has, say, a few very bright values in the image which
have the overall effect of darkening the rest of the image after
rescaling.

```
Usage:
1)   newimg = histtruncate(img, lHistCut, uHistCut)
2)   newimg = histtruncate(img, HistCut)

Arguments:
 Usage 1)
   img         -  Image to be processed.
   lHistCut    -  Percentage of the lower end of the histogram
                  to saturate.
   uHistCut    -  Percentage of the upper end of the histogram
                  to saturate.  If omitted or empty defaults to the value
                  for lHistCut.
 Usage 2)
   HistCut     -  Percentage of upper and lower ends of the histogram to cut.

Returns:
   newimg      -  Image with values clipped at the specified histogram
                  fraction values.  If the input image was colour the
                  lightness values are clipped and stretched to the range
                  0-1.  If the input image is greyscale no stretching is
                  applied. You may want to use normalise() to achieve this.
```
See also: normalise
"""
function  histtruncate(img::Array, lHistCut::Real, uHistCut::Real)
    # July      2001 - Original version
    # February  2012 - Added handling of NaN values in image
    # February  2014 - Code cleanup
    # September 2014 - Default for uHistCut + cleanup

    if lHistCut < 0 || lHistCut > 100 || uHistCut < 0 || uHistCut > 100
	error("Histogram truncation values must be between 0 and 100")
    end

    if ndims(img) > 2
	error("histtruncate only defined for grey scale images")
    end

    newimg = copy(img)
    sortv = sort(newimg[:])   # Generate a sorted array of pixel values.

    # Any NaN values will end up at the end of the sorted list. We
    # need to ignore these.
    N = sum(.!(isnan.(sortv)))  # Number of non NaN values.

    # Compute indicies corresponding to specified upper and lower fractions
    # of the histogram.
    lind = floor(Int, 1 + N*lHistCut/100)
    hind =  ceil(Int, N - N*uHistCut/100)

    low_val  = sortv[lind]
    high_val = sortv[hind]

    # Adjust image
    newimg[newimg .< low_val] .= low_val
    newimg[newimg .> high_val] .= high_val

    return newimg
end

function  histtruncate(img::Array, HistCut::Real)
    return histtruncate(img, HistCut, HistCut)
end

#----------------------------------------------------------------------
"""
sineramp  - Generates sine on a ramp colour map test image

The test image consists of a sine wave superimposed on a ramp function The
amplitude of the sine wave is modulated from its full value at the top of the
image to 0 at the bottom.

The image is useful for evaluating the effectiveness of different colour maps.
Ideally the sine wave pattern should be equally discernible over the full
range of the colour map.  In addition, across the bottom of the image, one
should not see any identifiable features as the underlying signal is a smooth
ramp.  In practice many colour maps have uneven perceptual contrast over their
range and often include 'flat spots' of no perceptual contrast that can hide
significant features.

```
Usage: img = sineramp(sze, amp, wavelen, p)
       img = sineramp()

Arguments:     sze - (rows, cols) specifying size of test image.
                     Defaults to (256 512)  Note the number of columns is
                     nominal and will be ajusted so that there are an
                     integer number of sine wave cycles across the image.
               amp - Amplitude of sine wave. Defaults to 12.5
           wavelen - Wavelength of sine wave in pixels. Defaults to 8.
                 p - Power to which the linear attenuation of amplitude,
                     from top to bottom, is raised.  For no attenuation use
                     p = 0.  For linear attenuation use a value of 1.  For
                     contrast sensitivity experiments use larger values of
                     p.  The default value is 2.
```
The ramp function that the sine wave is superimposed on is adjusted slightly
for each row so that each row of the image spans the full data range of 0 to
255.  Thus using a large sine wave amplitude will result in the ramp at the
top of the test image being reduced relative to the slope of the ramp at
the bottom of the image.

To start with try

```
  > img = sineramp()
```
This is equivalent to

```
  > img = sineramp((256 512), 12.5, 8, 2)
```
View it under 'gray' then try the 'jet', 'hsv', 'hot' etc colour maps.  The
results may cause you some concern!

If you are wishing to evaluate a cyclic colour map, say hsv, it is suggested
that you use the test image generated by circlesineramp().

See source code comments for more details on the default wavelength
and amplitude.

See also: circlesineramp, chirplin, chirpexp, equalisecolourmap, cmap
"""
function sineramp(sze=(256,512), amp=12.5, wavelen=8, p=2)
    #=
    The Default Wavelength:
    The default wavelength is 8 pixels.  On a computer monitor with a nominal
    pixel pitch of 0.25mm this corresponds to a wavelength of 2mm.  With a monitor
    viewing distance of 600mm this corresponds to 0.19 degrees of viewing angle or
    approximately 5.2 cycles per degree.  This falls within the range of spatial
    frequencies (3-7 cycles per degree ) at which most people have maximal
    contrast sensitivity to a sine wave grating (this varies with mean luminance).
    A wavelength of 8 pixels is also sufficient to provide a reasonable discrete
    representation of a sine wave.  The aim is to present a stimulus that is well
    matched to the performance of the human visual system so that what we are
    primarily evaluating is the colour map's perceptual contrast and not the visual
    performance of the viewer.

    The Default Amplitude:
    This is set at 12.5 so that from peak to trough we have a local feature of
    magnitude 25.  This is approximately 10% of the 256 levels in a standard
    colour map. It is not uncommon for colour maps to have perceptual flat spots
    that can hide features of this magnitude.
    =#

    # July  2013  Original version.
    # March 2014  Adjustments to make it better for evaluating cyclic colour maps.
    # June  2014  Default wavelength changed from 10 to 8.

    # ** Should I make this function return UInt8 values ?**

    # Adjust width of image so that we have an integer number of cycles of
    # the sinewave.  This is helps should one be using the test image to
    # evaluate a cyclic colour map.  However you will still see a slight
    # cyclic discontinuity at the top of the image, though this will
    # disappear at the bottom of the test image
    (rows,cols) = sze
    cycles = round(cols/wavelen)
    cols = cycles*wavelen

    # Sine wave
    x = collect(0:cols-1)'
    fx = amp*sin.( 1.0/wavelen * 2*pi*x)

    # Vertical modulating function
    A = (collect((rows-1):-1:0)'/(rows-1)).^float(p)
    img = A'*fx

    # Add ramp
    ramp = [c/(cols-1) for r = 1:rows, c = 0:(cols-1)]
    img = img + ramp*(255.0 - 2*amp)

    # Now normalise each row so that it spans the full data range from 0 to 255.
    # Again, this is important for evaluation of cyclic colour maps though a
    # small cyclic discontinuity will remain at the top of the test image.
    for r = 1:rows
        img[r,:] = normalise(img[r,:])
    end
    img = img * 255.0
    return img
end

#----------------------------------------------------------------------
"""
circlesineramp - Generates a test image for evaluating cyclic colour maps

```
Usage: (img, alpha) = circlesineramp(sze, amp, wavelen, p, hole)
       (img, alpha) = circlesineramp()

Arguments:     sze - Size of test image.  Defaults to 512x512.
               amp - Amplitude of sine wave. Defaults to pi/10
           wavelen - Wavelength of sine wave at half radius of the
                     circular test image. Defaults to 8 pixels.
                 p - Power to which the linear attenuation of amplitude,
                     from outside edge to centre, is raised.  For no
                     attenuation use p = 0.  For linear attenuation use a
                     value of 1.  The default value is 2, quadratic
                     attenuation.
              hole - Boolean flag indicating whether the test image should have
                     a 'hole' in its centre.  The default is true, to have a
                     hole, this removes the distraction of the orientation
                     singularlity at the centre.
Returns:
                im - The test image.
             alpha - Alpha mask matching the regions outside of of the
                     circular test image that are set to NaN.  Used if you
                     want to write an image with these regions transparent.
```
The test image is a circular pattern consistsing of a sine wave superimposed
on a spiral ramp function.  The spiral ramp starts at a value of 0 pointing
right, increasing anti-clockwise to a value of 2*pi as it completes the full
circle. This gives a 2*pi discontinuity on the right side of the image.  The
amplitude of the superimposed sine wave is modulated from its full value at
the outside of the circular pattern to 0 at the centre.  The default sine wave
amplitude of pi/10 means that the overall size of the sine wave from peak to
trough represents 2*(pi/10)/(2*pi) = 10% of the total spiral ramp of 2*pi.  If
you are testing your colour map over a cycle of pi you should use amp = pi/20
to obtain an equivalent ratio of sine wave to circular ramp.

The image is designed for evaluating the effectiveness of cyclic colour maps.
It is the cyclic companion to sineramp().  Ideally the sine wave pattern should
be equally discernible over all angles around the test image.  In practice
many colourmaps have uneven perceptual contrast over their range and often
include 'flat spots' of no perceptual contrast that can hide significant
features.  Try a HSV colour map.

Ideally the test image should be rendered with a cyclic colour map using
showangularim() though, in this case, rendering the image with SHOW or IMAGESC
will also be fine because all image values lie within, and use the full range
of, 0-2*pi.  However, in general, default display methods typically do not
respect data values directly and can perform inappropriate offsetting and
normalisation of the angular data before display and rendering with a colour
map.

For angular data to be rendered correctly it is important that the data values
are respected so that data values are correctly assigned to specific entries
in a cyclic colour map.  The assignment of values to colours also depends on
whether the data is cyclic over pi, or 2*pi.  SHOWANGULARIM supports this.

See also: applycycliccolourmap, sineramp, chirplin, chirpexp, equalisecolourmap, cmap
"""
function circlesineramp(sze=512, amp=pi/10, wavelen=8, p=2, hole=true)
    # September 2014  Original version.
    # October   2014  Number of cycles calculated from wave length rather than
    #                 being specified directly.

    # Set values for inner and outer radii of test pattern
    maxr = sze/2 * 0.9
    if hole
        minr = 0.15*sze
    else
        minr = 0
    end

    # Determine number of cycles to achieve desired wavelength at half radius
    meanr = (maxr + minr)/2
    circum = 2*pi*meanr
    cycles = round(circum/wavelen)

    # Angles are +ve anticlockwise and mod 2*pi
    theta = [mod(atan(-y,x), 2*pi) for y = -sze/2:sze/2, x = -sze/2:sze/2 ]
    rad = [sqrt(x^2 + y^2) for y = -sze/2:sze/2, x = -sze/2:sze/2 ]

    # Normalise radius so that it varies 0-1 over minr to maxr
    rad = (rad.-minr)/(maxr-minr)

    # Form the image
    img = amp*rad.^float(p) .* sin.(cycles*theta) + theta

    # Ensure all values are within 0-2*pi so that a simple default display
    # with a cyclic colour map will render the image correctly.
    img = mod.(img, 2*pi)

    # 'Nanify' values outside normalised radius values of 0-1
    alpha = ones(size(img))
    img[rad .> 1] .= NaN
    alpha[rad .> 1] .= 0

    if hole
        img[rad .< 0] .= NaN
        alpha[rad .< 0] .= 0
    end

   return img, alpha
end

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


using StaticArrays

struct BoundingBox{D,T}
    halflength::T
    center::SVector{D,T}
end

function BoundingBox(points::Vector{SVector{D,T}}) where {D,T}
    min_dim = Vector(points[1])
    max_dim = Vector(points[1])

    for i = 1:length(points)
        for j = 1:D
            min_dim[j] =  min_dim[j] < points[i][j] ? min_dim[j] : points[i][j]
            max_dim[j] =  max_dim[j] > points[i][j] ? max_dim[j] : points[i][j]
        end
    end

    center = @MVector zeros(D)

    length_dim = zeros(T,D)
    for j = 1:D
        length_dim[j] = max_dim[j] - min_dim[j]
        center[j] = (max_dim[j] + min_dim[j])/2.0
    end

    halflength = maximum(length_dim)/2.0

    return BoundingBox(halflength, SVector(center))
end

function getboxframe(bbox::BoundingBox{D,T}) where {D,T}
    if D == 2
        mat = zeros(5, 2)
        sign_x = [1 -1 -1 1 1]
        sign_y = [1 1 -1 -1 1]
        for i = 1:5
            mat[i,1] = bbox.center[1] + sign_x[i]*bbox.halflength
            mat[i,2] = bbox.center[2] + sign_y[i]*bbox.halflength
        end
    elseif D == 3
        sign_x = [1 -1 -1  1  1  1  1  1 NaN  1  1 -1 -1 NaN -1 -1 -1 NaN -1  1]
        sign_y = [1  1 -1 -1  1  1 -1 -1 NaN -1 -1 -1 -1 NaN -1  1  1 NaN  1  1]
        sign_z = [1  1  1  1  1 -1 -1  1 NaN -1 -1 -1  1 NaN -1 -1  1 NaN -1 -1]
        mat = zeros(length(sign_x), 3)
        for i = 1:length(sign_x)
            mat[i,1] = bbox.center[1] + sign_x[i]*bbox.halflength
            mat[i,2] = bbox.center[2] + sign_y[i]*bbox.halflength
            mat[i,3] = bbox.center[3] + sign_z[i]*bbox.halflength
        end
    else
        error("Dimension must be 2 or 3")
    end
    return mat
end

"""
    getchildbox(bbox, id)

Returns a bounding box that this is the `id`th quadrants (2D) /
octand (3D) of the bounding box `bbox`
"""
function getchildbox(bbox::BoundingBox{D,T}, id::Integer) where {D,T}

    hl = bbox.halflength / 2.0

    if D == 2
        sign_x = [1 -1 -1  1]
        sign_y = [1  1 -1 -1]

        center = bbox.center + hl .* SVector(sign_x[id], sign_y[id])
    elseif D == 3
        sign_x = [1 -1 -1  1  1 -1 -1  1]
        sign_y = [1  1 -1 -1  1  1 -1 -1]
        sign_z = [1  1  1  1 -1 -1 -1 -1]
        center = bbox.center + hl .* SVector(sign_x[id], sign_y[id],  sign_z[id])
    else
        error("We support only 2 or 3 dimensions")
    end
    return BoundingBox(hl, center)
end

function whichchildbox(bbox::BoundingBox{D,T}, point::SVector{D,T}) where {D,T}

    if point[1] >= bbox.center[1]
        if  point[2] >= bbox.center[2]
            id = 1
        else
            id = 4
        end
    else
        if point[2] >= bbox.center[2]
            id = 2
        else
            id = 3
        end
    end

    if D == 3 && point[3] < bbox.center[3]
        id += 4
    end

    return id
end


# univariate
function enclose_BranchandBound(f::Function, dom::Interval; order=10, tol=0.6)
    x0 = Interval(mid(dom))
    x = TaylorModel1(order, x0, dom)
    return branchandbound(f(x-x0).pol, dom, tol)
end

# multivariate
function enclose_BranchandBound(f::Function, dom::IntervalBox{D};
                                order=10, tol=0.6) where {D}
    x0 = mid(dom)
    set_variables(Float64, "x", order=2order, numvars=D)
    x = [TaylorModelN(i, order, IntervalBox(x0), dom) for i=1:D]
    return branchandbound(f(x...).pol, dom - x0, tol)
end

@inline _Rnext(R::Vector{<:Interval}) = Interval(minimum(inf.(R)), maximum(sup.(R)))

const Kmax = 1000

function branchandbound(p::Union{Taylor1, TaylorN},
                        dom::Interval_or_IntervalBox,
                        tol::Number)
    K = 1
    Rprev = evaluate(p, dom)
    D = bisect(dom)
    R = [evaluate(p, D[1]), evaluate(p, D[2])]
    Rnext = R[1] ∩ R[2]

    while  (Rprev.hi - Rnext.hi) >= tol * diam(Rnext) &&
           (Rprev.lo - Rnext.lo) >= tol * diam(Rnext)

        Rprev = _Rnext(R)
        R_x = sup.(R)
        R_n = inf.(R)
        max_range = maximum(R_x)
        max_index = findall(x->x == max_range, R_x)[1]
        min_range = minimum(R_n)
        min_index = findall(x->x == min_range, R_n)[1]
        l_D = length(D[1])
        K += 1
        if K == Kmax
            error("convergence not achieved in branch and bound after $Kmax steps")
        end

        D, R = divide_dom!(p, D, R, max_index)
        if max_index < min_index
            min_index += 1
            D, R = divide_dom!(p, D, R, min_index)
        end
        Rnext = _Rnext(R)
    end
    return _Rnext(R)
end

function divide_dom!(p::Union{TaylorN{T}, Taylor1{T}},
                     D::Union{Array{IntervalBox{N, W}, M}, Array{Interval{W}, M}},
                     R::Array{Interval{W}, M}, index::Number) where {N, T, M, W}
    BA = [ ((D[index][i]).hi - (D[index][i]).lo) for i = 1:length(D[1])]
    Beta1 = maximum(BA)
    β = findall(x->x == Beta1, BA)[1]
    D1, D2 = bisect(D[index], β)
    D[index] = D1
    DD = push!(D[1:index], D2)
    D = append!(DD, D[(index + 1):length(D)])
    R[index] = evaluate(p, D[index])
    RR = append!(R[1:index], evaluate(p, D[index + 1]))
    R = append!(RR, R[(index + 1):length(R)])
    return D, R
end

function enclose_binary(f, dom::Interval; kmax=3, tol=1e-3, algorithm=:IntervalArithmetic)
    y = enclose(f, dom, algorithm)
    yinf, ysup = inf(y), sup(y)
    kmax == 0 && return Interval(yinf, ysup)
    x = bisect(dom)
    fx1 = enclose(f, x[1], algorithm)
    fx2 = enclose(f, x[2], algorithm)
    ynew = hull(fx1, fx2)
    ynew_inf, ynew_sup = inf(ynew), sup(ynew)
    inf_close = abs(yinf - ynew_inf) <= tol
    sup_close = abs(ysup - ynew_sup) <= tol
    both_close = inf_close && sup_close
    inf_improves = ynew_inf > yinf
    sup_improves = ynew_sup < ysup
    both_improve = inf_improves && sup_improves
    if both_close || !both_improve
        return Interval(yinf, ysup)
    end
    yinf = max(yinf, ynew_inf)
    ysup = min(ysup, ynew_sup)
    if inf_improves
        yinf = ynew_inf
    end
    if sup_improves
        ysup = ynew_sup
    end
    e1 = enclose_binary(f, x[1], kmax=kmax-1, algorithm=algorithm)
    e2 = enclose_binary(f, x[2], kmax=kmax-1, algorithm=algorithm)
    return Interval(hull(e1, e2))
end


Sphere(x...) = HyperSphere(x...)

Circle(x...) = HyperSphere(x...)

widths(c::HyperSphere{N, T}) where {N, T} = Vec{N, T}(radius(c)*2)
radius(c::HyperSphere) = c.r
origin(c::HyperSphere) = c.center
minimum(c::HyperSphere{N, T}) where {N, T} = Vec{N, T}(origin(c)) - Vec{N, T}(radius(c))
maximum(c::HyperSphere{N, T}) where {N, T} = Vec{N, T}(origin(c)) + Vec{N, T}(radius(c))
function isinside(c::Circle, x::Real, y::Real)
    @inbounds ox, oy = origin(c)
    xD = abs(ox - x)
    yD = abs(oy - y)
    xD <= c.r && yD <= c.r
end

centered(S::Type{HyperSphere{N, T}}) where {N, T} = S(Vec{N,T}(0), T(0.5))
centered(::Type{T}) where {T <: HyperSphere} = centered(HyperSphere{ndims_or(T, 3), eltype_or(T, Float32)})


using Base:@pure

export AnalyticTransform, ExpTransform, PolyTransform,
       IdTransform, CoulombTransform, MorseTransform
import Base: convert, Dict, ==

# TODO: rename this AnalyticTransform
struct AnalyticTransform{FT, FDT, VT} <: SpaceTransform
   id::String
   f::FT
   f_d::FDT
   valid::VT
end

@pure transform(t::AnalyticTransform, r::Number) = t.f(r)
@pure transform_d(t::AnalyticTransform, r::Number) = t.f_d(r)

AnalyticTransform(id, f, f_d) = AnalyticTransform(id, f, f_d, Val(Symbol(id)))

==(T1::AnalyticTransform, T2::AnalyticTransform) = (T1.id == T2.id)

function AnalyticTransform(strans::String; fwrap = true)
   strans0 = strans
   # if @analytic is a substring then we don't do anything
   if !occursin(r"@analytic", strans)
      # but if not, then we next look for ->
      if !occursin(r"->", strans)
         # if -> is not a substring then we assume that strans is of the form
         # "(r0/r)^4" or similar i.e. explicitly uses r as the variable.
         strans = "@analytic r -> " * strans
      else
         # @analytic is not a substring but -> is a substring. e.g.
         # r -> (r0/r)^3 or s -> (r0/s)^3. we add @analytic to compile it
         strans = "@analytic " * strans
      end
   end
   ftrans = eval(Meta.parse(strans))
   if fwrap
      return AnalyticTransform(strans0, F64fun(ftrans.f), F64fun(ftrans.f_d))
   end
   return AnalyticTransform(strans0, ftrans.f, ftrans.f_d)
end

Dict(t::AnalyticTransform) = Dict( "__id__" => "NBodyIPs_AnalyticTransform",
                                "defn" => t.id )
AnalyticTransform(D::Dict) = AnalyticTransform(D["defn"])
convert(::Val{:NBodyIPs_AnalyticTransform}, D::Dict) = AnalyticTransform(D)
hash(::BASIS, t::AnalyticTransform) = hash((AnalyticTransform, t.id))

"""
Implements the space transform
```
r -> exp( - A * (r/r0 - 1) )
```

Constructor:
```
ExpTransform(A, r0)
```
"""
struct ExpTransform{TA, T} <: SpaceTransform
   A::TA
   r0::T
end

@pure transform(t::ExpTransform, r::Number) = exp( - t.A * (r/t.r0 - 1))
@pure transform_d(t::ExpTransform, r::Number) = (-t.A/t.r0) * exp( - t.A * (r/t.r0 - 1) )

# e^{-A (r/r0 - 1)) = x
# r/r0 - 1 = - (log x) / A
# r = r0 * (1 - (log x) / A)
# TODO: write a test that checks this inverse
inv_transform(t::ExpTransform, x::Number) = t.r0 * (1 - log(x) / t.A)

Dict(t::ExpTransform) =
   Dict("__id__" => "NBodyIPs_ExpTransform", "A" => t.A, "r0" => t.r0)
ExpTransform(D::Dict) = ExpTransform(D["A"], D["r0"])
convert(::Val{:NBodyIPs_ExpTransform}, D::Dict) = ExpTransform(D)
hash(::BASIS, t::ExpTransform) = hash(t)

"""
Implements the space transform
```
r -> (r0/r)^p
```

Constructor:
```
PolyTransform(p, r0)
```
"""
struct PolyTransform{TP, T} <: SpaceTransform
   p::TP
   r0::T
end

@pure transform(t::PolyTransform, r::Number) = @fastmath((t.r0/r)^t.p)
@pure transform_d(t::PolyTransform, r::Number) = @fastmath((-t.p/t.r0) * (t.r0/r)^(t.p+1))

# x = (r0/r)^p
# r x^{1/p} = r0
inv_transform(t::PolyTransform, x::Number) = t.r0 / x^(1/t.p)

Dict(t::PolyTransform) =
   Dict("__id__" => "NBodyIPs_PolyTransform",
        "p" => t.p,
        "r0" => t.r0)
PolyTransform(D::Dict) = PolyTransform(D["p"], D["r0"])
convert(::Val{:NBodyIPs_PolyTransform}, D::Dict) = PolyTransform(D)
hash(::BASIS, t::PolyTransform) = hash(t)

"""
Implements the space transform
```
r -> 1/r
```
Constructor: `CoulombTransform()`
"""
struct CoulombTransform <: SpaceTransform
end
@pure transform(t::CoulombTransform, r::Number) = @fastmath(1/r)
@pure transform_d(t::CoulombTransform, r::Number) = @fastmath(-1/(r*r))
inv_transform(t::CoulombTransform, x::Number) = 1/x

Dict(t::CoulombTransform) = Dict("__id__" => "NBodyIPs_CoulombTransform")
CoulombTransform(D::Dict) = CoulombTransform()
convert(::Val{:NBodyIPs_CoulombTransform}, D::Dict) = CoulombTransform(D)
hash(::BASIS, t::CoulombTransform) = hash(t)

"""
Implements the space transform
```
r -> exp(-r)
```
Constructor: `MorseTransform()`
"""
struct MorseTransform <: SpaceTransform
end
@pure transform(t::MorseTransform, r::Number) = @fastmath(exp(-r))
@pure transform_d(t::MorseTransform, r::Number) = @fastmath(-exp(-r))
inv_transform(t::MorseTransform, x::Number) = - log(x)

Dict(t::MorseTransform) = Dict("__id__" => "NBodyIPs_MorseTransform")
MorseTransform(D::Dict) = MorseTransform()
convert(::Val{:NBodyIPs_MorseTransform}, D::Dict) = MorseTransform(D)
hash(::BASIS, t::MorseTransform) = hash(t)

"""
Implements the space transform
```
r -> r
```
Constructor: `IdTransform()`
"""
struct IdTransform <: SpaceTransform
end

@pure transform(t::IdTransform, r::Number) = r
@pure transform_d(t::IdTransform, r::Number) = one(typeof(r))
inv_transform(t::IdTransform, x::Number) = x

Dict(t::IdTransform) = Dict("__id__" => "NBodyIPs_IdTransform")
IdTransform(D::Dict) = IdTransform()
convert(::Val{:NBodyIPs_IdTransform}, D::Dict) = IdTransform(D)
hash(::BASIS, t::IdTransform) = hash(t)


@testset "partition" begin
    for n in 1:17
        for p in 1:n
            Is = [partition(n, p, i) for i in 1:p]
            @test vcat(Is...) == collect(1:n)
            @test maximum(length, Is) <= minimum(length, Is) + 1
        end
    end
end

@testset "translate_partition" begin
    n = 17
    for p in 1:n
        for i in 1:p
            for q in 1:n
                j = translate_partition(n, p, q, i)
                @test first(partition(n, p, i)) in partition(n, q, j)
            end
        end
    end
end

@testset "align_partitions" begin
    n = 10
    p, q, i = 4, 5, 1
    @test align_partitions(n, p, q, i) == 1
    p, q, i = 4, 5, 2
    @test align_partitions(n, p, q, i) == 2
    p, q, i = 4, 5, 3
    @test align_partitions(n, p, q, i) == 2
    p, q, i = 4, 5, 4
    @test align_partitions(n, p, q, i) == 2
    p, q, i = 4, 2, 2
    @test align_partitions(n, p, q, i) == 1
    p, q, i = 4, 2, 4
    @test align_partitions(n, p, q, i) == 2

    n = 17
    for p in 1:n
        for i in 1:p
            for q in 1:n
                j = align_partitions(n, p, q, i)
                ip = translate_partition(n, q, p, j)
                @test first(partition(n, p, ip)) == first(partition(n, q, j))
            end
        end
    end
end


using Random, Plots; pyplot()
Random.seed!(808)

n = 10^6
data = tan.(rand(n)*pi .- pi/2)
averages = accumulate(+,data)./collect(1:n)

plot( 1:n, averages,
	c=:blue, legend=:none,
	xscale=:log10, xlims=(1,n), xlabel="n", ylabel="Running average")


cd(@__DIR__); include("setups/path512.jl")
pyplot(dpi = 200)

##
@time VM_NGWP = vm_ngwp(𝚽, GP_dual)
@time LP_NGWP = lp_ngwp(𝚽, Gstar_Sig.W, GP_dual; ϵ = 0.3)

dict_VM = Dict()
dict_LP = Dict()
for j in [3, 4, 5]
   for k in [1, 2, 3]
      if GP_dual.rs[k + 1, j] - GP_dual.rs[k, j] < 1
         continue
      end
      WW_VM = sort_wavelets(VM_NGWP[GP_dual.rs[k, j]:(GP_dual.rs[k + 1, j] - 1), j, :]')
      WW_LP = sort_wavelets(LP_NGWP[GP_dual.rs[k, j]:(GP_dual.rs[k + 1, j] - 1), j, :]')
      Y = [-0.5, 0.6]
      l = Int(floor(size(WW_LP, 2) / 2))
      dict_VM[(j-1, k-1, l-1)] = WW_VM[:, l]
      dict_LP[(j-1, k-1, l-1)] = WW_LP[:, l]
      plt = plot(size = (400, 350), layout = Plots.grid(2, 1), ylim = Y)
         plot!(WW_VM[:, l], c = :black, grid = false, frame = :box, lw = 0.8, xtick = false,
         lab = latexstring("\\psi^{($(j-1))}_{$(k-1), $(l-1)} (\\mathrm{VM})"), legendfontsize = 11)
         plot!(WW_LP[:, l], c = :black, grid = false, frame = :box, lw = 0.8,
         lab = latexstring("\\psi^{($(j-1))}_{$(k-1), $(l-1)} (\\mathrm{LP})"), legendfontsize = 11,
         subplot = 2)
         xticks!([1; 64:64:N], vcat(string(1), [string(k) for k in 64:64:N]),
         subplot = 2)
      savefig(plt, "../figs/Path512_VM_LP_j$(j-1)k$(k-1)l$(l-1).png")
   end
end

## Table 8.2
using PrettyTables
header = ["basis vector" "main support width" "sidelobe attenuation"]
res = Matrix{String}(undef, 0, 3)
iter = sort(collect(keys(dict_VM)))
for γ in iter
   global res
   ms = find_mainsupport(dict_VM[γ]; ϵ = 0.01)
   sa = sidelobe_attenuation(dict_VM[γ])
   j, k, l = γ
   res = vcat(res, ["ψ^{$(j)}_{$(k), $(l)} (VM)" ms[2] - ms[1] + 1 round(sa; digits = 4)])
   ms = find_mainsupport(dict_LP[γ]; ϵ = 0.01)
   sa = sidelobe_attenuation(dict_LP[γ])
   res = vcat(res, ["ψ^{$(j)}_{$(k), $(l)} (LP)" ms[2] - ms[1] + 1 round(sa; digits = 4)])
end
pretty_table(res; header = tuple(header), alignment = :c)


using Tmatrix
using Zygote: Zygote
using FiniteDifferences: FiniteDifferences

finite_difference_method = FiniteDifferences.central_fdm(5, 1; factor = 1e-6)

# inputs ====================================================
n_θ_points = 5
n_ϕ_points = 2
m, n, m_, n_ = 2, 3, 2, 5
n_max = 1
k1_r, k1_i, k2_r, k2_i = 1e5, 1e3, 2e5, 3e3
rotationally_symmetric = true
symmetric_about_plane_perpendicular_z = false
rx, rz = 1e-6, 1.2e-6
kind = "regular"
J_superscript = 11
Q_superscript = J_superscript

θ_1D_array = LinRange(1e-16, π, n_θ_points);
ϕ_1D_array = LinRange(1e-16, 2π, n_ϕ_points);
if rotationally_symmetric
    θ_array = collect(θ_1D_array)
    ϕ_array = zeros(size(θ_array))
else
    # eventually, this should be removed. I just keep it for sanity checks.
    θ_array, ϕ_array = Tmatrix.meshgrid(θ_1D_array, ϕ_1D_array)
end
r_array, n̂_array = Tmatrix.ellipsoid(rx, rz, θ_array);

epss = 1e-6

# test "J_mn_m_n__integrand_SeparateRealImag" ====================================================================================
J = Tmatrix.J_mn_m_n__integrand_SeparateRealImag(
    m,
    n,
    m_,
    n_,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
    kind,
    J_superscript,
)
J_ = Tmatrix.J_mn_m_n__integrand(
    m,
    n,
    m_,
    n_,
    complex(k1_r, k1_i) .* r_array,
    complex(k2_r, k2_i) .* r_array,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array;
    kind = kind,
    J_superscript = J_superscript,
)

println("==================================================================")
println("testing <<J_mn_m_n__integrand_SeparateRealImag>>")
println("Displaying results from _SeparateRealImag and the complex function")
println("Displaying result from _SeparateRealImag function --------------------")
display(Tmatrix.get_complex_matrix_from_concatenated_real_imag(J))
println("Displaying result from complex function ------------------------------")
display(J_)

println(
    "testing autodiff for <<J_mn_m_n__integrand_SeparateRealImag>> -----------------------------------------",
)
@time J_mn_m_n__integrand_SeparateRealImag_jacobian = Zygote.jacobian(
    Tmatrix.J_mn_m_n__integrand_SeparateRealImag,
    m,
    n,
    m_,
    n_,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
    kind,
    J_superscript,
)
function J_mn_m_n__integrand_SeparateRealImag_lean(
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
)
    Tmatrix.J_mn_m_n__integrand_SeparateRealImag(
        m,
        n,
        m_,
        n_,
        k1_r,
        k1_i,
        k2_r,
        k2_i,
        r_array,
        θ_array,
        ϕ_array,
        n̂_array,
        kind,
        J_superscript,
    )
end
@time J_mn_m_n__integrand_SeparateRealImag_lean_jacobian = Zygote.jacobian(
    J_mn_m_n__integrand_SeparateRealImag_lean,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
)
@time J_jacobian_finite_diff = FiniteDifferences.jacobian(
    finite_difference_method,
    J_mn_m_n__integrand_SeparateRealImag_lean,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
)
@time J_jacobian_manual_numerical_diff = (
    (
        J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r + epss,
            k1_i,
            k2_r,
            k2_i,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
        ) - J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r - epss,
            k1_i,
            k2_r,
            k2_i,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
        )
    ) / (2 * epss),
    (
        J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i + epss,
            k2_r,
            k2_i,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
        ) - J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i - epss,
            k2_r,
            k2_i,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
        )
    ) / (2 * epss),
    (
        J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i,
            k2_r + epss,
            k2_i,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
        ) - J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i,
            k2_r - epss,
            k2_i,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
        )
    ) / (2 * epss),
    (
        J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i,
            k2_r,
            k2_i + epss,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
        ) - J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i,
            k2_r,
            k2_i - epss,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
        )
    ) / (2 * epss),
    (
        J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i,
            k2_r,
            k2_i,
            r_array .+ epss,
            θ_array,
            ϕ_array,
            n̂_array,
        ) - J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i,
            k2_r,
            k2_i,
            r_array .- epss,
            θ_array,
            ϕ_array,
            n̂_array,
        )
    ) ./ (2 * epss),
    (
        J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i,
            k2_r,
            k2_i,
            r_array,
            θ_array .+ epss,
            ϕ_array,
            n̂_array,
        ) - J_mn_m_n__integrand_SeparateRealImag_lean(
            k1_r,
            k1_i,
            k2_r,
            k2_i,
            r_array,
            θ_array .- epss,
            ϕ_array,
            n̂_array,
        )
    ) ./ (2 * epss),
)
finite_difference_method = FiniteDifferences.central_fdm(5, 1; factor = 1e-10)
@time J_jacobian_finite_diff = FiniteDifferences.jacobian(
    finite_difference_method,
    J_mn_m_n__integrand_SeparateRealImag_lean,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
)
hcat(J_mn_m_n__integrand_SeparateRealImag_lean_jacobian[1], J_jacobian_finite_diff[1])

# test "J_mn_m_n__SeparateRealImag" ================================================================================================
J = Tmatrix.J_mn_m_n__SeparateRealImag(
    m,
    n,
    m_,
    n_,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
    kind,
    J_superscript,
    rotationally_symmetric,
)
J_ = Tmatrix.J_mn_m_n_(
    m,
    n,
    m_,
    n_,
    complex(k1_r, k1_i) .* r_array,
    complex(k2_r, k2_i) .* r_array,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array;
    kind = kind,
    J_superscript = J_superscript,
    rotationally_symmetric = rotationally_symmetric,
)

println("==================================================================")
println("Displaying results from _SeparateRealImag and the complex function")
println("Displaying result from _SeparateRealImag function --------------------")
display(Tmatrix.get_complex_matrix_from_concatenated_real_imag(J))
println("Displaying result from complex function ------------------------------")
display(J_)

# the problem that Zygote is trying to differentiate "rotationally_symmetric".
# how can I prevent Zygote from differentiating with respect to a given argument?
Zygote.jacobian(
    (m, n, m_, n_, k1_r, k1_i, k2_r, k2_i, r_array, θ_array, ϕ_array, n̂_array) ->
        Tmatrix.J_mn_m_n__SeparateRealImag(
            m,
            n,
            m_,
            n_,
            k1_r,
            k1_i,
            k2_r,
            k2_i,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
            kind,
            J_superscript,
            rotationally_symmetric,
        ),
    m,
    n,
    m_,
    n_,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
)

Zygote.jacobian(
    J_mn_m_n__integrand_SeparateRealImag_lean,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
)[5]

epss = big(1e-6);
vcat(
    (
        J_mn_m_n__integrand_SeparateRealImag_lean(
            big(k1_r),
            big(k1_i),
            big(k2_r),
            big(k2_i),
            big.(r_array .+ epss),
            big.(θ_array),
            big.(ϕ_array),
            [big.(x) for x in n̂_array],
        ) - J_mn_m_n__integrand_SeparateRealImag_lean(
            big(k1_r),
            big(k1_i),
            big(k2_r),
            big(k2_i),
            big.(r_array .- epss),
            big.(θ_array),
            big.(ϕ_array),
            [big.(x) for x in n̂_array],
        )
    ) ./ (2 * epss)...,
)

single_one_array = zeros(size(r_array));
single_one_array[1] = 1;

for epss in [big(1e-50)]
    display(
        vcat(
            (
                J_mn_m_n__integrand_SeparateRealImag_lean(
                    big(k1_r),
                    big(k1_i),
                    big(k2_r),
                    big(k2_i),
                    big.(r_array + epss * single_one_array),
                    big.(θ_array),
                    big.(ϕ_array),
                    [big.(x) for x in n̂_array],
                ) - J_mn_m_n__integrand_SeparateRealImag_lean(
                    big(k1_r),
                    big(k1_i),
                    big(k2_r),
                    big(k2_i),
                    big.(r_array - epss * single_one_array),
                    big.(θ_array),
                    big.(ϕ_array),
                    [big.(x) for x in n̂_array],
                )
            ) ./ (2 * epss)...,
        ),
    )
    println()

    display(
        vcat(
            (
                J_mn_m_n__integrand_SeparateRealImag_lean(
                    big(k1_r),
                    big(k1_i),
                    big(k2_r),
                    big(k2_i),
                    big.(r_array),
                    big.(θ_array + epss * single_one_array),
                    big.(ϕ_array),
                    [big.(x) for x in n̂_array],
                ) - J_mn_m_n__integrand_SeparateRealImag_lean(
                    big(k1_r),
                    big(k1_i),
                    big(k2_r),
                    big(k2_i),
                    big.(r_array),
                    big.(θ_array - epss * single_one_array),
                    big.(ϕ_array),
                    [big.(x) for x in n̂_array],
                )
            ) ./ (2 * epss)...,
        ),
    )
    println()
end

# test "Q_mn_m_n_SeparateRealImag"
@time Q = Tmatrix.Q_mn_m_n_SeparateRealImag(
    m,
    n,
    m_,
    n_,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
    kind,
    Q_superscript,
    rotationally_symmetric,
    symmetric_about_plane_perpendicular_z,
)
@time Q_ = Tmatrix.Q_mn_m_n_(
    m,
    n,
    m_,
    n_,
    complex(k1_r, k1_i),
    complex(k2_r, k2_i),
    complex(k1_r, k1_i) .* r_array,
    complex(k2_r, k2_i) .* r_array,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array;
    kind = kind,
    Q_superscript = Q_superscript,
    rotationally_symmetric = rotationally_symmetric,
)

println("==================================================================")
println("Displaying results from _SeparateRealImag and the complex function")
println("Displaying result from _SeparateRealImag function --------------------")
display(Tmatrix.get_complex_matrix_from_concatenated_real_imag(Q))
println("Displaying result from complex function ------------------------------")
display(Q_)

Zygote.jacobian(
    (m, n, m_, n_, k1_r, k1_i, k2_r, k2_i, r_array, θ_array, ϕ_array, n̂_array) ->
        Tmatrix.Q_mn_m_n_SeparateRealImag(
            m,
            n,
            m_,
            n_,
            k1_r,
            k1_i,
            k2_r,
            k2_i,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
            kind,
            Q_superscript,
            rotationally_symmetric,
            symmetric_about_plane_perpendicular_z,
        ),
    m,
    n,
    m_,
    n_,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
)

# test "Q_matrix_SeparateRealImag"
@time Q = Tmatrix.Q_matrix_SeparateRealImag(
    n_max,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
    kind,
    rotationally_symmetric,
    symmetric_about_plane_perpendicular_z,
)
@time Q_ = Tmatrix.Q_matrix(
    n_max,
    complex(k1_r, k1_i),
    complex(k2_r, k2_i),
    complex(k1_r, k1_i) .* r_array,
    complex(k2_r, k2_i) .* r_array,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
    kind = kind,
    rotationally_symmetric = rotationally_symmetric,
    symmetric_about_plane_perpendicular_z = false,
    verbose = false,
)

println("==================================================================")
println("Displaying results from _SeparateRealImag and the complex function")
println("Displaying result from _SeparateRealImag function --------------------")
display(Tmatrix.get_complex_matrix_from_concatenated_real_imag(Q))
println("Displaying result from complex function ------------------------------")
display(Q_)

# The jacobian calculation is so slow. For n_max = 1, it takes 52.409485 seconds. This is slow compared to 0.015419 seconds required to evaluate the function
Zygote.jacobian(
    (n_max, k1_r, k1_i, k2_r, k2_i, r_array, θ_array, ϕ_array, n̂_array) ->
        Tmatrix.Q_matrix_SeparateRealImag(
            n_max,
            k1_r,
            k1_i,
            k2_r,
            k2_i,
            r_array,
            θ_array,
            ϕ_array,
            n̂_array,
            kind,
            rotationally_symmetric,
            symmetric_about_plane_perpendicular_z,
        ),
    n_max,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
)

# test "T_matrix_SeparateRealImag"
@time T = Tmatrix.T_matrix_SeparateRealImag(
    n_max,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
    rotationally_symmetric,
    symmetric_about_plane_perpendicular_z,
)
@time T_ = Tmatrix.T_matrix(
    n_max,
    complex(k1_r, k1_i),
    complex(k2_r, k2_i),
    complex(k1_r, k1_i) .* r_array,
    complex(k2_r, k2_i) .* r_array,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array;
    rotationally_symmetric = rotationally_symmetric,
    symmetric_about_plane_perpendicular_z = false,
)

println("==================================================================")
println("Displaying results from _SeparateRealImag and the complex function")
println("Displaying result from _SeparateRealImag function --------------------")
display(Tmatrix.get_complex_matrix_from_concatenated_real_imag(T))
println("Displaying result from complex function ------------------------------")
display(T_)

@time T = Tmatrix.T_matrix_SeparateRealImag(
    n_max,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
    rotationally_symmetric,
    symmetric_about_plane_perpendicular_z,
)
@time ∂T = Zygote.jacobian(
    Tmatrix.T_matrix_SeparateRealImag,
    n_max,
    k1_r,
    k1_i,
    k2_r,
    k2_i,
    r_array,
    θ_array,
    ϕ_array,
    n̂_array,
    rotationally_symmetric,
    symmetric_about_plane_perpendicular_z,
)


####
#### Resampling schemes for particle filters
####

# Some references
#  - http://arxiv.org/pdf/1301.4019.pdf
#  - http://people.isy.liu.se/rt/schon/Publications/HolSG2006.pdf
# Code adapted from: http://uk.mathworks.com/matlabcentral/fileexchange/24968-resampling-methods-for-particle-filtering

struct ResampleWithESSThreshold{R,T<:Real}
    resampler::R
    threshold::T
end

function ResampleWithESSThreshold(resampler=resample_systematic)
    return ResampleWithESSThreshold(resampler, 0.5)
end

# More stable, faster version of rand(Categorical)
function randcat(rng::Random.AbstractRNG, p::AbstractVector{<:Real})
    T = eltype(p)
    r = rand(rng, T)
    cp = p[1]
    s = 1
    n = length(p)
    while cp <= r && s < n
        @inbounds cp += p[s += 1]
    end
    return s
end

function resample_multinomial(
    rng::Random.AbstractRNG, w::AbstractVector{<:Real}, num_particles::Integer=length(w)
)
    return rand(rng, Distributions.sampler(Distributions.Categorical(w)), num_particles)
end

function resample_residual(
    rng::Random.AbstractRNG,
    w::AbstractVector{<:Real},
    num_particles::Integer=length(weights),
)
    # Pre-allocate array for resampled particles
    indices = Vector{Int}(undef, num_particles)

    # deterministic assignment
    residuals = similar(w)
    i = 1
    @inbounds for j in 1:length(w)
        x = num_particles * w[j]
        floor_x = floor(Int, x)
        for k in 1:floor_x
            indices[i] = j
            i += 1
        end
        residuals[j] = x - floor_x
    end

    # sampling from residuals
    if i <= num_particles
        residuals ./= sum(residuals)
        rand!(rng, Distributions.Categorical(residuals), view(indices, i:num_particles))
    end

    return indices
end

"""
    resample_stratified(rng, weights, n)

Return a vector of `n` samples `x₁`, ..., `xₙ` from the numbers 1, ..., `length(weights)`,
generated by stratified resampling.

In stratified resampling `n` ordered random numbers `u₁`, ..., `uₙ` are generated, where
``uₖ \\sim U[(k - 1) / n, k / n)``. Based on these numbers the samples `x₁`, ..., `xₙ`
are selected according to the multinomial distribution defined by the normalized `weights`,
i.e., `xᵢ = j` if and only if
``uᵢ \\in [\\sum_{s=1}^{j-1} weights_{s}, \\sum_{s=1}^{j} weights_{s})``.
"""
function resample_stratified(
    rng::Random.AbstractRNG, weights::AbstractVector{<:Real}, n::Integer=length(weights)
)
    # check input
    m = length(weights)
    m > 0 || error("weight vector is empty")

    # pre-calculations
    @inbounds v = n * weights[1]

    # generate all samples
    samples = Array{Int}(undef, n)
    sample = 1
    @inbounds for i in 1:n
        # sample next `u` (scaled by `n`)
        u = oftype(v, i - 1 + rand(rng))

        # as long as we have not found the next sample
        while v < u
            # increase and check the sample
            sample += 1
            sample > m &&
                error("sample could not be selected (are the weights normalized?)")

            # update the cumulative sum of weights (scaled by `n`)
            v += n * weights[sample]
        end

        # save the next sample
        samples[i] = sample
    end

    return samples
end

"""
    resample_systematic(rng, weights, n)

Return a vector of `n` samples `x₁`, ..., `xₙ` from the numbers 1, ..., `length(weights)`,
generated by systematic resampling.

In systematic resampling a random number ``u \\sim U[0, 1)`` is used to generate `n` ordered
numbers `u₁`, ..., `uₙ` where ``uₖ = (u + k − 1) / n``. Based on these numbers the samples
`x₁`, ..., `xₙ` are selected according to the multinomial distribution defined by the
normalized `weights`, i.e., `xᵢ = j` if and only if
``uᵢ \\in [\\sum_{s=1}^{j-1} weights_{s}, \\sum_{s=1}^{j} weights_{s})``.
"""
function resample_systematic(
    rng::Random.AbstractRNG, weights::AbstractVector{<:Real}, n::Integer=length(weights)
)
    # check input
    m = length(weights)
    m > 0 || error("weight vector is empty")

    # pre-calculations
    @inbounds v = n * weights[1]
    u = oftype(v, rand(rng))

    # find all samples
    samples = Array{Int}(undef, n)
    sample = 1
    @inbounds for i in 1:n
        # as long as we have not found the next sample
        while v < u
            # increase and check the sample
            sample += 1
            sample > m &&
                error("sample could not be selected (are the weights normalized?)")

            # update the cumulative sum of weights (scaled by `n`)
            v += n * weights[sample]
        end

        # save the next sample
        samples[i] = sample

        # update `u`
        u += one(u)
    end

    return samples
end