import .Tracker
using .Tracker: Tracker,
                TrackedReal,
                TrackedVector,
                TrackedMatrix,
                TrackedArray,
                TrackedVecOrMat,
                @grad,
                track,
                data,
                param
using LinearAlgebra

# Different from Tracker.istracked
_istracked(x::AbstractArray{<:TrackedReal}) = true
_istracked(::TrackedArray) = false

Tracker.dual(x::Bool, p) = x
Base.prevfloat(r::TrackedReal) = track(prevfloat, r)
@grad function prevfloat(r::Real)
    prevfloat(data(r)), Δ -> Δ
end
Base.nextfloat(r::TrackedReal) = track(nextfloat, r)
@grad function nextfloat(r::Real)
    nextfloat(data(r)), Δ -> Δ
end
for i = 0:2, c = Tracker.combinations([:AbstractArray, :TrackedArray, :TrackedReal, :Number], i), f = [:hcat, :vcat]
    if :TrackedReal in c
        cnames = map(_ -> gensym(), c)
        @eval Base.$f($([:($x::$c) for (x, c) in zip(cnames, c)]...), x::Union{TrackedArray,TrackedReal}, xs::Union{AbstractArray,Number}...) =
            track($f, $(cnames...), x, xs...)
    end
end
@grad function vcat(x::Real)
    vcat(data(x)), (Δ) -> (Δ[1],)
end
@grad function vcat(x1::Real, x2::Real)
    vcat(data(x1), data(x2)), (Δ) -> (Δ[1], Δ[2])
end
@grad function vcat(x1::AbstractVector, x2::Real)
    vcat(data(x1), data(x2)), (Δ) -> (Δ[1:length(x1)], Δ[length(x1)+1])
end

function Base.copy(
    A::TrackedArray{T, 2, <:Adjoint{T, <:AbstractTriangular{T, <:AbstractMatrix{T}}}},
) where {T <: Real}
    return track(copy, A)
end
@grad function Base.copy(
    A::TrackedArray{T, 2, <:Adjoint{T, <:AbstractTriangular{T, <:AbstractMatrix{T}}}},
) where {T <: Real}
    return copy(data(A)), ∇ -> (copy(∇),)
end

Base.:*(A::TrackedMatrix, B::AbstractTriangular) = track(*, A, B)
Base.:*(A::AbstractTriangular{T}, B::TrackedVector) where {T} = track(*, A, B)
Base.:*(A::AbstractTriangular{T}, B::TrackedMatrix) where {T} = track(*, A, B)
Base.:*(A::Adjoint{T, <:AbstractTriangular{T}}, B::TrackedMatrix) where {T} = track(*, A, B)
Base.:*(A::Adjoint{T, <:AbstractTriangular{T}}, B::TrackedVector) where {T} = track(*, A, B)

_eps(::Type{<:TrackedReal{T}}) where {T} = _eps(T)

# AD implementations
function jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::Real
)
    return data(Tracker.gradient(b, x)[1])
end
function jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::AbstractVector{<:Real}
)
    # We extract `data` so that we don't return a `Tracked` type
    return data(Tracker.jacobian(b, x))
end

# implementations for Shift bijector
function _logabsdetjac_shift(a::TrackedReal, x::Real, ::Val{0})
    return param(_logabsdetjac_shift(data(a), data(x), Val(0)))
end
function _logabsdetjac_shift(a::TrackedReal, x::AbstractVector{<:Real}, ::Val{0})
    return param(_logabsdetjac_shift(data(a), data(x), Val(0)))
end
function _logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractVector{<:Real},
    ::Val{1}
)
    return param(_logabsdetjac_shift(data(a), data(x), Val(1)))
end
function _logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractMatrix{<:Real},
    ::Val{1}
)
    return param(_logabsdetjac_shift(data(a), data(x), Val(1)))
end

# Log bijector

@grad function logabsdetjac(b::Log{1}, x::AbstractVector)
    return -sum(log, data(x)), Δ -> (nothing, -Δ ./ data(x))
end
@grad function logabsdetjac(b::Log{1}, x::AbstractMatrix)
    return -vec(sum(log, data(x); dims = 1)), Δ -> (nothing, .- Δ' ./ data(x))
end

# implementations for Scale bijector
# Adjoints for 0-dim and 1-dim `Scale` using `Real`
function _logabsdetjac_scale(a::TrackedReal, x::Real, ::Val{0})
    return track(_logabsdetjac_scale, a, data(x), Val(0))
end
@grad function _logabsdetjac_scale(a::Real, x::Real, ::Val{0})
    return _logabsdetjac_scale(data(a), data(x), Val(0)), Δ -> (inv(data(a)) .* Δ, nothing, nothing)
end
# Need to treat `AbstractVector` and `AbstractMatrix` separately due to ambiguity errors
function _logabsdetjac_scale(a::TrackedReal, x::AbstractVector, ::Val{0})
    return track(_logabsdetjac_scale, a, data(x), Val(0))
end
@grad function _logabsdetjac_scale(a::Real, x::AbstractVector, ::Val{0})
    da = data(a)
    J = fill(inv.(da), length(x))
    return _logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
function _logabsdetjac_scale(a::TrackedReal, x::AbstractMatrix, ::Val{0})
    return track(_logabsdetjac_scale, a, data(x), Val(0))
end
@grad function _logabsdetjac_scale(a::Real, x::AbstractMatrix, ::Val{0})
    da = data(a)
    J = fill(size(x, 1) / da, size(x, 2))
    return _logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
# adjoints for 1-dim and 2-dim `Scale` using `AbstractVector`
function _logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    return track(_logabsdetjac_scale, a, data(x), Val(1))
end
@grad function _logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    # ∂ᵢ (∑ⱼ log|aⱼ|) = ∑ⱼ δᵢⱼ ∂ᵢ log|aⱼ|
    #                 = ∂ᵢ log |aᵢ|
    #                 = (1 / aᵢ) ∂ᵢ aᵢ
    #                 = (1 / aᵢ)
    da = data(a)
    J = inv.(da)
    return _logabsdetjac_scale(da, data(x), Val(1)), Δ -> (J .* Δ, nothing, nothing)
end
function _logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    return track(_logabsdetjac_scale, a, data(x), Val(1))
end
@grad function _logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    da = data(a)
    Jᵀ = repeat(inv.(da), 1, size(x, 2))
    return _logabsdetjac_scale(da, data(x), Val(1)), Δ -> (Jᵀ * Δ, nothing, nothing)
end
# TODO: implement analytical gradient for scaling a vector using a matrix
# function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     track(_logabsdetjac_scale, a, data(x), Val{1})
# end
# @grad function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     throw
# end

# implementations for Stacked bijector
function logabsdetjac(b::Stacked, x::TrackedMatrix{<:Real})
    return mapvcat(eachcol(x)) do c
        logabsdetjac(b, c)
    end
end
# TODO: implement custom adjoint since we can exploit block-diagonal nature of `Stacked`
function (sb::Stacked)(x::TrackedMatrix{<:Real})
    return eachcolmaphcat(sb, x)
end

# Simplex adjoints

function _simplex_bijector(X::TrackedVecOrMat, b::SimplexBijector)
    return track(_simplex_bijector, X, b)
end
function _simplex_inv_bijector(Y::TrackedVecOrMat, b::SimplexBijector)
    return track(_simplex_inv_bijector, Y, b)
end

@grad function _simplex_bijector(X::AbstractVector, b::SimplexBijector)
    Xd = data(X)
    return _simplex_bijector(Xd, b), Δ -> (simplex_link_jacobian(Xd)' * Δ, nothing)
end
@grad function _simplex_inv_bijector(Y::AbstractVector, b::SimplexBijector)
    Yd = data(Y)
    return _simplex_inv_bijector(Yd, b), Δ -> (simplex_invlink_jacobian(Yd)' * Δ, nothing)
end

@grad function _simplex_bijector(X::AbstractMatrix, b::SimplexBijector)
    Xd = data(X)
    return _simplex_bijector(Xd, b), Δ -> begin
        maphcat(eachcol(Xd), eachcol(Δ)) do c1, c2
            simplex_link_jacobian(c1)' * c2
        end, nothing
    end
end
@grad function _simplex_inv_bijector(Y::AbstractMatrix, b::SimplexBijector)
    Yd = data(Y)
    return _simplex_inv_bijector(Yd, b), Δ -> begin
        maphcat(eachcol(Yd), eachcol(Δ)) do c1, c2
            simplex_invlink_jacobian(c1)' * c2
        end, nothing
    end
end

replace_diag(::typeof(log), X::TrackedMatrix) = track(replace_diag, log, X)
@grad function replace_diag(::typeof(log), X)
    Xd = data(X)
    f(i, j) = i == j ? log(Xd[i, j]) : Xd[i, j]
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = i == j ? ∇[i, j]/Xd[i, j] : ∇[i, j]
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

replace_diag(::typeof(exp), X::TrackedMatrix) = track(replace_diag, exp, X)
@grad function replace_diag(::typeof(exp), X)
    Xd = data(X)
    f(i, j) = ifelse(i == j, exp(Xd[i, j]), Xd[i, j])
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = ifelse(i == j, ∇[i, j]*exp(Xd[i, j]), ∇[i, j])
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

logabsdetjac(b::SimplexBijector, x::TrackedVecOrMat) = track(logabsdetjac, b, x)
@grad function logabsdetjac(b::SimplexBijector, x::AbstractVector)
    xd = data(x)
    return logabsdetjac(b, xd), Δ -> begin
        (nothing, simplex_logabsdetjac_gradient(xd) * Δ)
    end
end
@grad function logabsdetjac(b::SimplexBijector, x::AbstractMatrix)
    xd = data(x)
    return logabsdetjac(b, xd), Δ -> begin
        (nothing, maphcat(eachcol(xd), Δ) do c, g
            simplex_logabsdetjac_gradient(c) * g
        end)
    end
end

for header in [
    (:(u::TrackedArray), :w),
    (:u, :(w::TrackedArray)),
    (:(u::TrackedArray), :(w::TrackedArray)),
]
    @eval begin
        function get_u_hat($(header...))
            if u isa TrackedArray
                T = typeof(u)
            else
                T = typeof(w)
            end
            x = w' * u
            return (u .+ (planar_flow_m(x) - x) .* w ./ sum(abs2, w))::T
        end
    end
end

for header in [
    (:(z::TrackedArray), :w, :b),
    (:z, :(w::TrackedArray), :b),
    (:z, :w, :(b::TrackedReal)),
    (:(z::TrackedArray), :(w::TrackedArray), :b),
    (:(z::TrackedArray), :w, :(b::TrackedReal)),
    (:z, :(w::TrackedArray), :(b::TrackedReal)),
    (:(z::TrackedArray), :(w::TrackedArray), :(b::TrackedReal)),
]
    @eval begin
        function ψ($(header...))
            if z isa AbstractMatrix
                if z isa TrackedMatrix
                    T = typeof(z)
                elseif w isa TrackedVector
                    T = matrixof(typeof(w))
                else
                    T = matrixof(typeof(b))
                end
            else
                if z isa TrackedVector
                    T = typeof(z)
                elseif w isa TrackedVector
                    T = typeof(w)
                else
                    T = vectorof(typeof(b))
                end
            end
            return ((1 .- tanh.(w' * z .+ b).^2) .* w)::T    # for planar flow from eq(11)
        end
    end
end

for header in [
    (:(u::TrackedArray), :w, :b, :(z::AbstractVecOrMat)),
    (:u, :(w::TrackedArray), :b, :(z::AbstractVecOrMat)),
    (:u, :w, :(b::TrackedReal), :(z::AbstractVecOrMat)),
    (:u, :w, :b, :(z::TrackedVecOrMat)),
    (:(u::TrackedArray), :(w::TrackedArray), :b, :(z::AbstractVecOrMat)),
    (:(u::TrackedArray), :w, :(b::TrackedReal), :(z::AbstractVecOrMat)),
    (:(u::TrackedArray), :w, :b, :(z::TrackedVecOrMat)),
    (:u, :(w::TrackedArray), :(b::TrackedReal), :(z::AbstractVecOrMat)),
    (:u, :(w::TrackedArray), :b, :(z::TrackedVecOrMat)),
    (:u, :w, :(b::TrackedArray), :(z::TrackedVecOrMat)),
    (:(u::TrackedArray), :(w::TrackedArray), :(b::TrackedReal), :(z::AbstractVecOrMat)),
    (:(u::TrackedArray), :(w::TrackedArray), :b, :(z::TrackedVecOrMat)),
    (:(u::TrackedArray), :w, :(b::TrackedReal), :(z::TrackedVecOrMat)),
    (:u, :(w::TrackedArray), :(b::TrackedReal), :(z::TrackedVecOrMat)),
    (:(u::TrackedArray), :(w::TrackedArray), :(b::TrackedReal), :(z::TrackedVecOrMat)),
]
    @eval begin
        function _planar_transform($(header...))
            u_hat = get_u_hat(u, w)
            if z isa AbstractVector
                temp = w' * z + b + zero(eltype(u_hat))
                if z isa TrackedVector
                    T = typeof(z)
                elseif u_hat isa TrackedVector
                    T = typeof(u_hat)
                else
                    T = vectorof(typeof(temp))
                end
            else
                temp = w' * z .+ (b + zero(eltype(u_hat)))
                if z isa TrackedMatrix
                    T = typeof(z)
                elseif u_hat isa TrackedVector
                    T = matrixof(typeof(u_hat))
                else
                    T = matrixof(typeof(temp'))
                end
            end
            transformed::T = z .+ u_hat .* tanh.(temp) # from eq(10)
            return (transformed = transformed, u_hat = u_hat)
        end
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractVector)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:α_, :β, :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedVecOrMat)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedVector)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
]
    @eval begin
        function _radial_transform($(header...))
            α = softplus(α_)            # from A.2
            β_hat = -α + softplus(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = vectorof(typeof(β_hat))
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = typeof(z_0)
            else
                T = TV = typeof(z)
            end
            Tr = promote_type(eltype(z), eltype(z_0))
            r::Tr = norm((z .- z_0)::TV)
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractMatrix)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:α_, :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
]
    @eval begin
        function _radial_transform($(header...))
            α = softplus(α_)            # from A.2
            β_hat = -α + softplus(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = matrixof(TV)
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = matrixof(TV)
            else
                T = typeof(z)
                TV = vectorof(T)
            end
            r::TV = mapvcat(norm, eachcol(z .- z_0))
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end

function matrixof(::Type{TrackedArray{T, 1, Vector{T}}}) where {T <: Real}
    return TrackedArray{T, 2, Matrix{T}}
end
function matrixof(::Type{TrackedReal{T}}) where {T <: Real}
    return TrackedArray{T, 2, Matrix{T}}
end
function vectorof(::Type{TrackedArray{T, 2, Matrix{T}}}) where {T <: Real}
    return TrackedArray{T, 1, Vector{T}}
end
function vectorof(::Type{TrackedReal{T}}) where {T <: Real}
    return TrackedArray{T, 1, Vector{T}}
end

(b::Exp)(x::TrackedVector) = exp.(x)::vectorof(float(eltype(x)))
(b::Log)(x::TrackedVector) = log.(x)::vectorof(float(eltype(x)))
logabsdetjac(b::Log{0}, x::TrackedVector) = .-log.(x)::vectorof(float(eltype(x)))
logabsdetjac(b::Log{1}, x::TrackedMatrix) = - vec(sum(log.(x); dims = 1))

_logabsdetjac_shift(a::Real, x::TrackedVector{T}, ::Val{0}) where {T<:Real} = zeros(T, length(x))
_logabsdetjac_shift(a::T1, x::TrackedVector{T2}, ::Val{1}) where {T1<:Union{Real, TrackedVector}, T2<:Real} = zero(T2)
_logabsdetjac_shift(a::T1, x::TrackedMatrix{T2}, ::Val{1}) where {T1<:Union{Real, TrackedVector}, T2<:Real} = zeros(T2, size(x, 2))


"""
spin18.jl - sampling robustness
"""

WDIR = joinpath(@__DIR__, "../../")
include(joinpath(WDIR, "src", "spin", "spin.jl"))

using Altro
using HDF5
using LinearAlgebra
using RobotDynamics
using StaticArrays
using TrajectoryOptimization
const RD = RobotDynamics
const TO = TrajectoryOptimization

# paths
const EXPERIMENT_META = "spin"
const EXPERIMENT_NAME = "spin18"
const SAVE_PATH = joinpath(WDIR, "out", EXPERIMENT_META, EXPERIMENT_NAME)

# problem
const CONTROL_COUNT = 1
const STATE_COUNT = 2
const SAMPLE_COUNT = 8
const ASTATE_SIZE_BASE = STATE_COUNT * HDIM_ISO + 3 * CONTROL_COUNT
const ASTATE_SIZE = ASTATE_SIZE_BASE + SAMPLE_COUNT * HDIM_ISO
const ACONTROL_SIZE = CONTROL_COUNT
const INITIAL_STATE1 = [1., 0, 0, 0]
const INITIAL_STATE2 = [0., 1, 0, 0]
const INITIAL_STATE3 = [1., 0, 0, 1] ./ sqrt(2)
const INITIAL_STATE4 = [1., -1, 0, 0] ./ sqrt(2)
# state indices
const STATE1_IDX = 1:HDIM_ISO
const STATE2_IDX = STATE1_IDX[end] + 1:STATE1_IDX[end] + HDIM_ISO
const INTCONTROLS_IDX = STATE2_IDX[end] + 1:STATE2_IDX[end] + CONTROL_COUNT
const CONTROLS_IDX = INTCONTROLS_IDX[end] + 1:INTCONTROLS_IDX[end] + CONTROL_COUNT
const DCONTROLS_IDX = CONTROLS_IDX[end] + 1:CONTROLS_IDX[end] + CONTROL_COUNT
const S1_IDX = DCONTROLS_IDX[end] + 1:DCONTROLS_IDX[end] + HDIM_ISO
const S2_IDX = S1_IDX[end] + 1:S1_IDX[end] + HDIM_ISO
const S3_IDX = S2_IDX[end] + 1:S2_IDX[end] + HDIM_ISO
const S4_IDX = S3_IDX[end] + 1:S3_IDX[end] + HDIM_ISO
const S5_IDX = S4_IDX[end] + 1:S4_IDX[end] + HDIM_ISO
const S6_IDX = S5_IDX[end] + 1:S5_IDX[end] + HDIM_ISO
const S7_IDX = S6_IDX[end] + 1:S6_IDX[end] + HDIM_ISO
const S8_IDX = S7_IDX[end] + 1:S7_IDX[end] + HDIM_ISO
# control indices
const D2CONTROLS_IDX = 1:CONTROL_COUNT

# model
struct Model <: AbstractModel
    namp::Float64
end
@inline RD.state_dim(::Model) = ASTATE_SIZE
@inline RD.control_dim(::Model) = ACONTROL_SIZE

# This cost puts a gate error cost on
# the sample states and a LQR cost on the other terms.
# The hessian w.r.t the state and controls is constant.
struct Cost{N,M,T} <: TO.CostFunction
    Q::Diagonal{T, SVector{N,T}}
    R::Diagonal{T, SVector{M,T}}
    q::SVector{N, T}
    c::T
    hess_astate::Symmetric{T, SMatrix{N,N,T}}
    target_states::Array{SVector{HDIM_ISO, T}, 1}
    q_ss1::T
    q_ss2::T
    q_ss3::T
    q_ss4::T
end

function Cost(Q::Diagonal{T,SVector{N,T}}, R::Diagonal{T,SVector{M,T}},
              xf::SVector{N,T}, target_states::Array{SVector{HDIM_ISO}, 1},
              q_ss1::T, q_ss2::T, q_ss3::T, q_ss4::T) where {N,M,T}
    q = -Q * xf
    c = 0.5 * xf' * Q * xf
    hess_astate = zeros(N, N)
    # For reasons unknown to the author, throwing a -1 in front
    # of the gate error Hessian makes the cost function work.
    # This is strange, because the gate error Hessian has been
    # checked against autodiff.
    hess_state1 = -1 * q_ss1 * hessian_gate_error_iso2(target_states[1])
    hess_state2 = -1 * q_ss2 * hessian_gate_error_iso2(target_states[2])
    hess_state3 = -1 * q_ss3 * hessian_gate_error_iso2(target_states[3])
    hess_state4 = -1 * q_ss4 * hessian_gate_error_iso2(target_states[4])
    hess_astate[S1_IDX, S1_IDX] = hess_state1
    hess_astate[S2_IDX, S2_IDX] = hess_state2
    hess_astate[S3_IDX, S3_IDX] = hess_state3
    hess_astate[S4_IDX, S4_IDX] = hess_state4
    hess_astate[S5_IDX, S5_IDX] = hess_state1
    hess_astate[S6_IDX, S6_IDX] = hess_state2
    hess_astate[S7_IDX, S7_IDX] = hess_state3
    hess_astate[S8_IDX, S8_IDX] = hess_state4
    hess_astate += Q
    hess_astate = Symmetric(SMatrix{N, N}(hess_astate))
    return Cost{N,M,T}(Q, R, q, c, hess_astate, target_states, q_ss1, q_ss2, q_ss3, q_ss4)
end

@inline TO.state_dim(cost::Cost{N,M,T}) where {N,M,T} = N
@inline TO.control_dim(cost::Cost{N,M,T}) where {N,M,T} = M
@inline Base.copy(cost::Cost{N,M,T}) where {N,M,T} = Cost{N,M,T}(
    cost.Q, cost.R, cost.q, cost.c, cost.hess_astate,
    cost.target_states, cost.q_ss1, cost.q_ss2, cost.q_ss3, cost.q_ss4
)

@inline TO.stage_cost(cost::Cost{N,M,T}, astate::SVector{N}) where {N,M,T} = (
    0.5 * astate' * cost.Q * astate + cost.q'astate + cost.c
    + cost.q_ss1 * gate_error_iso2(astate, cost.target_states[1], S1_IDX[1] - 1)
    + cost.q_ss2 * gate_error_iso2(astate, cost.target_states[2], S2_IDX[1] - 1)
    + cost.q_ss3 * gate_error_iso2(astate, cost.target_states[3], S3_IDX[1] - 1)
    + cost.q_ss4 * gate_error_iso2(astate, cost.target_states[4], S4_IDX[1] - 1)
    + cost.q_ss1 * gate_error_iso2(astate, cost.target_states[1], S5_IDX[1] - 1)
    + cost.q_ss2 * gate_error_iso2(astate, cost.target_states[2], S6_IDX[1] - 1)
    + cost.q_ss3 * gate_error_iso2(astate, cost.target_states[3], S7_IDX[1] - 1)
    + cost.q_ss4 * gate_error_iso2(astate, cost.target_states[4], S8_IDX[1] - 1)
)

@inline TO.stage_cost(cost::Cost{N,M,T}, astate::SVector{N}, acontrol::SVector{M}) where {N,M,T} = (
    TO.stage_cost(cost, astate) + 0.5 * acontrol' * cost.R * acontrol
)

function TO.gradient!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T}) where {N,M,T}
    E.q = (cost.Q * astate + cost.q + [
        @SVector zeros(ASTATE_SIZE_BASE);
        cost.q_ss1 * jacobian_gate_error_iso2(astate, cost.target_states[1], S1_IDX[1] - 1);
        cost.q_ss2 * jacobian_gate_error_iso2(astate, cost.target_states[2], S2_IDX[1] - 1);
        cost.q_ss3 * jacobian_gate_error_iso2(astate, cost.target_states[3], S3_IDX[1] - 1);
        cost.q_ss4 * jacobian_gate_error_iso2(astate, cost.target_states[4], S4_IDX[1] - 1);
        cost.q_ss1 * jacobian_gate_error_iso2(astate, cost.target_states[1], S5_IDX[1] - 1);
        cost.q_ss2 * jacobian_gate_error_iso2(astate, cost.target_states[2], S6_IDX[1] - 1);
        cost.q_ss3 * jacobian_gate_error_iso2(astate, cost.target_states[3], S7_IDX[1] - 1);
        cost.q_ss4 * jacobian_gate_error_iso2(astate, cost.target_states[4], S8_IDX[1] - 1);
    ])
    return false
end

function TO.gradient!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T},
                      acontrol::SVector{M,T}) where {N,M,T}
    TO.gradient!(E, cost, astate)
    E.r = cost.R * acontrol
    E.c = 0
    return false
end

function TO.hessian!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T}) where {N,M,T}
    E.Q = cost.hess_astate
    return true
end

function TO.hessian!(E::TO.QuadraticCostFunction, cost::Cost{N,M,T}, astate::SVector{N,T},
                     acontrol::SVector{M,T}) where {N,M,T}
    TO.hessian!(E, cost, astate)
    E.R = cost.R
    E.H .= 0
    return true
end

# dynamics
function RD.discrete_dynamics(::Type{RD.RK3}, model::Model, astate::StaticVector,
                              acontrols::StaticVector, time::Real, dt::Real) where {SC}
    camp = astate[CONTROLS_IDX[1]]
    negi_hc = camp * NEGI_H1_ISO
    h_prop = exp((FQ_NEGI_H0_ISO + negi_hc) * dt)
    state1 = h_prop * astate[STATE1_IDX]
    state2 = h_prop * astate[STATE2_IDX]
    intcontrols = astate[INTCONTROLS_IDX[1]] + dt * astate[CONTROLS_IDX[1]]
    controls = astate[CONTROLS_IDX[1]] + dt * astate[DCONTROLS_IDX[1]]
    dcontrols = astate[DCONTROLS_IDX[1]] + dt * acontrols[D2CONTROLS_IDX[1]]

    hp_prop = exp((FQ_NEGI_H0_ISO + (camp + model.namp) * NEGI_H1_ISO) * dt)
    hn_prop = exp((FQ_NEGI_H0_ISO + (camp - model.namp) * NEGI_H1_ISO) * dt)
    s1 = hp_prop * astate[S1_IDX]
    s2 = hp_prop * astate[S2_IDX]
    s3 = hp_prop * astate[S3_IDX]
    s4 = hp_prop * astate[S4_IDX]
    s5 = hn_prop * astate[S5_IDX]
    s6 = hn_prop * astate[S6_IDX]
    s7 = hn_prop * astate[S7_IDX]
    s8 = hn_prop * astate[S8_IDX]

    astate_ = [
        state1; state2; intcontrols; controls; dcontrols;
        s1; s2; s3; s4; s5; s6; s7; s8;
    ]

    return astate_
end

# main
function run_traj(;gate_type=xpiby2, evolution_time=56.8, solver_type=altro,
                  sqrtbp=false, integrator_type=rk3, qs=[1e0, 1e0, 1e0, 1e-1, 1e0, 1e0, 1e0, 1e0, 1e-1],
                  dt_inv=Int64(1e1), smoke_test=false, constraint_tol=1e-8, al_tol=1e-4,
                  pn_steps=2, max_penalty=1e11, verbose=true, save=true, max_iterations=Int64(2e5),
                  namp=NAMP_PREFACTOR)
    # model configuration
    model = Model(namp)
    n = state_dim(model)
    m = control_dim(model)
    t0 = 0.
    tf = evolution_time

    # initial state
    x0 = SVector{n}([
        INITIAL_STATE1;
        INITIAL_STATE2;
        zeros(3 * CONTROL_COUNT);
        repeat([INITIAL_STATE1; INITIAL_STATE2; INITIAL_STATE3; INITIAL_STATE4], 2);
    ])
    # target state
    gate_unitary = GT_GATE[gate_type]
    target_states = Array{SVector{HDIM_ISO}, 1}(undef, 4)
    target_states[1] = gate_unitary * INITIAL_STATE1
    target_states[2] = gate_unitary * INITIAL_STATE2
    target_states[3] = gate_unitary * INITIAL_STATE3
    target_states[4] = gate_unitary * INITIAL_STATE4
    xf = SVector{n}([
        target_states[1];
        target_states[2];
        zeros(3 * CONTROL_COUNT);
        repeat([target_states[1]; target_states[2];
                target_states[3]; target_states[4]], 2);
    ])

    # control amplitude constraint
    x_max = SVector{n}([
        fill(Inf, STATE_COUNT * HDIM_ISO);
        fill(Inf, CONTROL_COUNT);
        fill(MAX_CONTROL_NORM_0, 1); # control
        fill(Inf, CONTROL_COUNT);
        fill(Inf, SAMPLE_COUNT * HDIM_ISO);
    ])
    x_min = SVector{n}([
        fill(-Inf, STATE_COUNT * HDIM_ISO);
        fill(-Inf, CONTROL_COUNT);
        fill(-MAX_CONTROL_NORM_0, 1); # control
        fill(-Inf, CONTROL_COUNT);
        fill(-Inf, SAMPLE_COUNT * HDIM_ISO);
    ])
    # control amplitude constraint at boundary
    x_max_boundary = SVector{n}([
        fill(Inf, STATE_COUNT * HDIM_ISO);
        fill(Inf, CONTROL_COUNT);
        fill(0, 1); # control
        fill(Inf, CONTROL_COUNT);
        fill(Inf, SAMPLE_COUNT * HDIM_ISO);
    ])
    x_min_boundary = SVector{n}([
        fill(-Inf, STATE_COUNT * HDIM_ISO);
        fill(-Inf, CONTROL_COUNT);
        fill(0, 1); # control
        fill(-Inf, CONTROL_COUNT);
        fill(-Inf, SAMPLE_COUNT * HDIM_ISO);
    ])

    # initial trajectory
    dt = dt_inv^(-1)
    N = Int(floor(evolution_time * dt_inv)) + 1
    U0 = [SVector{m}([
        fill(1e-4, CONTROL_COUNT);
    ]) for k = 1:N-1]
    X0 = [SVector{n}([
        fill(NaN, n);
    ]) for k = 1:N]
    Z = Traj(X0, U0, dt * ones(N))

    # cost function
    Q = Diagonal(SVector{n}([
        fill(qs[1], STATE_COUNT * HDIM_ISO); # ψ1, ψ2
        fill(qs[2], 1); # ∫a
        fill(qs[3], 1); # a
        fill(qs[4], 1); # ∂a
        fill(0, SAMPLE_COUNT * HDIM_ISO);
    ]))
    Qf = Q * N
    R = Diagonal(SVector{m}([
        fill(qs[9], CONTROL_COUNT);
    ]))
    cost_k = Cost(Q, R, xf, target_states, qs[5], qs[6], qs[7], qs[8])
    cost_f = Cost(Qf, R, xf, target_states, N * qs[5], N * qs[6], N * qs[7], N * qs[8])
    objective = TO.Objective(cost_k, cost_f, N)

    # must satisfy control amplitude bound
    control_bnd = BoundConstraint(n, m, x_max=x_max, x_min=x_min)
    # must statisfy conrols start and end at 0
    control_bnd_boundary = BoundConstraint(n, m, x_max=x_max_boundary, x_min=x_min_boundary)
    # must reach target state, must have zero net flux
    target_astate_constraint = GoalConstraint(xf, [STATE1_IDX; STATE2_IDX; INTCONTROLS_IDX])
    # must obey unit norm
    norm_constraints = [NormConstraint(n, m, 1, TO.Equality(), idxs) for idxs in (
        STATE1_IDX, STATE2_IDX, S1_IDX, S2_IDX, S3_IDX, S4_IDX, S5_IDX, S6_IDX, S7_IDX, S8_IDX,
    )]
    constraints = ConstraintList(n, m, N)
    add_constraint!(constraints, control_bnd, 2:N-2)
    add_constraint!(constraints, control_bnd_boundary, N-1:N-1)
    add_constraint!(constraints, target_astate_constraint, N:N);
    for norm_constraint in norm_constraints
        add_constraint!(constraints, norm_constraint, 2:N-1)
    end

    # solve problem
    prob = Problem{IT_RDI[integrator_type]}(model, objective, constraints, x0, xf, Z, N, t0, evolution_time)
    solver = ALTROSolver(prob)
    verbose_pn = verbose ? true : false
    verbose_ = verbose ? 2 : 0
    projected_newton = solver_type == altro ? true : false
    constraint_tolerance = solver_type == altro ? constraint_tol : al_tol
    iterations_inner = smoke_test ? 1 : 300
    iterations_outer = smoke_test ? 1 : 30
    n_steps = smoke_test ? 1 : pn_steps
    set_options!(solver, square_root=sqrtbp, constraint_tolerance=constraint_tolerance,
                 projected_newton_tolerance=al_tol, n_steps=n_steps,
                 penalty_max=max_penalty, verbose_pn=verbose_pn, verbose=verbose_,
                 projected_newton=projected_newton, iterations_inner=iterations_inner,
                 iterations_outer=iterations_outer, iterations=max_iterations)
    Altro.solve!(solver)

    # post-process
    acontrols_raw = TO.controls(solver)
    acontrols_arr = permutedims(reduce(hcat, map(Array, acontrols_raw)), [2, 1])
    astates_raw = TO.states(solver)
    astates_arr = permutedims(reduce(hcat, map(Array, astates_raw)), [2, 1])
    Q_raw = Array(Q)
    Q_arr = [Q_raw[i, i] for i in 1:size(Q_raw)[1]]
    Qf_raw = Array(Qf)
    Qf_arr = [Qf_raw[i, i] for i in 1:size(Qf_raw)[1]]
    R_raw = Array(R)
    R_arr = [R_raw[i, i] for i in 1:size(R_raw)[1]]
    cidx_arr = Array(CONTROLS_IDX)
    d2cidx_arr = Array(D2CONTROLS_IDX)
    cmax = TO.max_violation(solver)
    cmax_info = TO.findmax_violation(TO.get_constraints(solver))
    iterations_ = Altro.iterations(solver)

    result = Dict(
        "acontrols" => acontrols_arr,
        "controls_idx" => cidx_arr,
        "d2controls_dt2_idx" => d2cidx_arr,
        "evolution_time" => evolution_time,
        "astates" => astates_arr,
        "Q" => Q_arr,
        "Qf" => Qf_arr,
        "R" => R_arr,
        "cmax" => cmax,
        "cmax_info" => cmax_info,
        "dt" => dt,
        "sample_count" => SAMPLE_COUNT,
        "solver_type" => Integer(solver_type),
        "sqrtbp" => Integer(sqrtbp),
        "max_penalty" => max_penalty,
        "constraint_tol" => constraint_tol,
        "al_tol" => al_tol,
        "gate_type" => Integer(gate_type),
        "save_type" => Integer(jl),
        "integrator_type" => Integer(integrator_type),
        "iterations" => iterations_,
        "max_iterations" => max_iterations,
    )

    # save
    if save
        save_file_path = generate_file_path("h5", EXPERIMENT_NAME, SAVE_PATH)
        println("Saving this optimization to $(save_file_path)")
        h5open(save_file_path, "cw") do save_file
            for key in keys(result)
                write(save_file, key, result[key])
            end
        end
        result["save_file_path"] = save_file_path
    end

    return result
end


struct GeneralUnstructuredMesh_2D <: UnstructuredMesh_2D
    name::String
    points::Vector{Point_2D}
    edges::Vector{<:SVector{L, UInt32} where {L}}
    materialized_edges::Vector{<:Edge_2D}
    faces::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}}
    materialized_faces::Vector{<:Face_2D}
    edge_face_connectivity::Vector{SVector{2, UInt32}}
    face_edge_connectivity::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}}
    boundary_edges::Vector{Vector{UInt32}}
    face_sets::Dict{String, Set{UInt32}}
end

function GeneralUnstructuredMesh_2D(;
        name::String = "DefaultMeshName",
        points::Vector{Point_2D} = Point_2D[],
        edges::Vector{<:SVector{L, UInt32} where {L}} = SVector{2, UInt32}[],
        materialized_edges::Vector{<:Edge_2D} = LineSegment_2D[],
        faces::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}} = SVector{4, UInt32}[],
        materialized_faces::Vector{<:Face_2D} = Triangle_2D[],
        edge_face_connectivity::Vector{SVector{2, UInt32}} = SVector{2, UInt32}[],
        face_edge_connectivity ::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}} = SVector{3, UInt32}[],
        boundary_edges::Vector{Vector{UInt32}} = Vector{UInt32}[],
        face_sets::Dict{String, Set{UInt32}} = Dict{String, Set{UInt32}}()
    ) where {F <: AbstractFloat, U <: Unsigned}
        return GeneralUnstructuredMesh_2D(name,
                                          points,
                                          edges,
                                          materialized_edges,
                                          faces,
                                          materialized_faces,
                                          edge_face_connectivity,
                                          face_edge_connectivity,
                                          boundary_edges,
                                          face_sets,
                                         )
end

Base.broadcastable(mesh::GeneralUnstructuredMesh_2D) = Ref(mesh)

# Return a mesh with boundary edges
function add_boundary_edges(mesh::M; boundary_shape="Unknown"
    ) where {M <: UnstructuredMesh_2D}
    if 0 === length(mesh.edge_face_connectivity)
        mesh = add_connectivity(mesh)
    end
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = boundary_edges(mesh, boundary_shape),
             face_sets = mesh.face_sets
            )
end

# Return a mesh with face/edge connectivity and edge/face connectivity
function add_connectivity(mesh::UnstructuredMesh_2D)
    return add_edge_face_connectivity(mesh)
end

# Return a mesh with its edges
function add_edges(mesh::M) where {M <: UnstructuredMesh_2D}
    return M(name = mesh.name,
             points = mesh.points,
             edges = edges(mesh),
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a mesh with edge/face connectivity
function add_edge_face_connectivity(mesh::M) where {M <: UnstructuredMesh_2D}
    if 0 === length(mesh.face_edge_connectivity)
        mesh = add_face_edge_connectivity(mesh)
    end
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = edge_face_connectivity(mesh),
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a mesh with every field created
function add_everything(mesh::UnstructuredMesh_2D)
    return add_materialized_faces(
             add_materialized_edges(
               add_boundary_edges(mesh, boundary_shape = "Rectangle")))
end

# Return a mesh with face/edge connectivity
function add_face_edge_connectivity(mesh::M) where {M <: UnstructuredMesh_2D}
    if 0 === length(mesh.edges)
        mesh = add_edges(mesh)
    end
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = face_edge_connectivity(mesh),
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a mesh with materialized edges
function add_materialized_edges(mesh::M) where {M <: UnstructuredMesh_2D}
    if 0 === length(mesh.edges)
        mesh = add_edges(mesh)
    end
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = materialize_edges(mesh),
             faces = mesh.faces,
             materialized_faces = mesh.materialized_faces,
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a mesh with materialized faces
function add_materialized_faces(mesh::M) where {M <: UnstructuredMesh_2D}
    return M(name = mesh.name,
             points = mesh.points,
             edges = mesh.edges,
             materialized_edges = mesh.materialized_edges,
             faces = mesh.faces,
             materialized_faces = materialize_faces(mesh),
             edge_face_connectivity = mesh.edge_face_connectivity,
             face_edge_connectivity = mesh.face_edge_connectivity,
             boundary_edges = mesh.boundary_edges,
             face_sets = mesh.face_sets
            )
end

# Return a vector of the faces adjacent to the face of ID face
function adjacent_faces(face::UInt32, mesh::UnstructuredMesh_2D)
    edges = mesh.face_edge_connectivity[face]
    the_adjacent_faces = UInt32[]
    for edge in edges
        faces = mesh.edge_face_connectivity[edge]
        for face_id in faces
            if face_id != face && face_id != 0
                push!(the_adjacent_faces, face_id)
            end
        end
    end
    return the_adjacent_faces
end

# Area of face
function area(face::SVector{N, UInt32}, points::Vector{Point_2D}) where {N}
    return area(materialize_face(face, points))
end

# Return the area of a face set
function area(mesh::UnstructuredMesh_2D, face_set::Set{UInt32})
    if 0 < length(mesh.materialized_faces)
        return mapreduce(x->area(mesh.materialized_faces[x]), +, face_set)
    else
        return mapreduce(x->area(mesh.faces[x], mesh.points), +, face_set)
    end
end

# Return the area of a face set by name
function area(mesh::UnstructuredMesh_2D, set_name::String)
    return area(mesh, mesh.face_sets[set_name])
end

# Bounding box of a vector of points
function boundingbox(points::Vector{Point_2D})
    x = getindex.(points, 1)
    y = getindex.(points, 2)
    return Rectangle_2D(minimum(x), minimum(y), maximum(x), maximum(y))
end

# Bounding box of a vector of points
function boundingbox(points::SVector{L, Point_2D}) where {L}
    x = getindex.(points, 1)
    y = getindex.(points, 2)
    return Rectangle_2D(minimum(x), minimum(y), maximum(x), maximum(y))
end

# Return a vector containing vectors of the edges in each side of the mesh's bounding shape, e.g.
# For a rectangular bounding shape the sides are North, East, South, West. Then the output would
# be [ [e1, e2, e3, ...], [e17, e18, e18, ...], ..., [e100, e101, ...]]
function boundary_edges(mesh::UnstructuredMesh_2D, boundary_shape::String)
    # edges which have face 0 in their edge_face connectivity are boundary edges
    the_boundary_edges = UInt32.(findall(x->x[1] === 0x00000000, mesh.edge_face_connectivity))
    if boundary_shape == "Rectangle"
        # Sort edges into NESW
        bb = boundingbox(mesh.points)
        y_north = bb.ymax
        x_east  = bb.xmax
        y_south = bb.ymin
        x_west  = bb.xmin
        p_NW = Point_2D(x_west, y_north)
        p_NE = bb.tr
        p_SE = Point_2D(x_east, y_south)
        p_SW = bb.bl
        edges_north = UInt32[]
        edges_east = UInt32[]
        edges_south = UInt32[]
        edges_west = UInt32[]
        # Insert edges so that indices move from NW -> NE -> SE -> SW -> NW
        for edge ∈  the_boundary_edges
            epoints = edgepoints(mesh.edges[edge], mesh.points)
            if all(x->abs(x[2] - y_north) < 1e-4, epoints)
                insert_boundary_edge!(edge, p_NW, edges_north, mesh)
            elseif all(x->abs(x[1] - x_east) < 1e-4, epoints)
                insert_boundary_edge!(edge, p_NE, edges_east, mesh)
            elseif all(x->abs(x[2] - y_south) < 1e-4, epoints)
                insert_boundary_edge!(edge, p_SE, edges_south, mesh)
            elseif all(x->abs(x[1] - x_west) < 1e-4, epoints)
                insert_boundary_edge!(edge, p_SW, edges_west, mesh)
            else
                @error "Edge $edge could not be classified as NSEW"
            end
        end
        return [ edges_north, edges_east, edges_south, edges_west ]
    else
        return [ convert(Vector{UInt32}, the_boundary_edges) ]
    end
end

# SVector of MVectors of point IDs representing the 3 edges of a triangle
function edges(face::SVector{3, UInt32})
    edges = SVector( MVector{2, UInt32}(face[1], face[2]),
                     MVector{2, UInt32}(face[2], face[3]),
                     MVector{2, UInt32}(face[3], face[1]) )
    # Order the linear edge vertices by ID
    for edge in edges
        if edge[2] < edge[1]
            e1 = edge[1]
            edge[1] = edge[2]
            edge[2] = e1
        end
    end
    return edges
end

# SVector of MVectors of point IDs representing the 4 edges of a quadrilateral
function edges(face::SVector{4, UInt32})
    edges = SVector( MVector{2, UInt32}(face[1], face[2]),
                     MVector{2, UInt32}(face[2], face[3]),
                     MVector{2, UInt32}(face[3], face[4]),
                     MVector{2, UInt32}(face[4], face[1]) )
    # Order the linear edge vertices by ID
    for edge in edges
        if edge[2] < edge[1]
            e1 = edge[1]
            edge[1] = edge[2]
            edge[2] = e1
        end
    end
    return edges
end

# SVector of MVectors of point IDs representing the 3 edges of a quadratic triangle
function edges(face::SVector{6, UInt32})
    edges = SVector( MVector{3, UInt32}(face[1], face[2], face[4]),
                     MVector{3, UInt32}(face[2], face[3], face[5]),
                     MVector{3, UInt32}(face[3], face[1], face[6]) )
    # Order the linear edge vertices by ID
    for edge in edges
        if edge[2] < edge[1]
            e1 = edge[1]
            edge[1] = edge[2]
            edge[2] = e1
        end
    end
    return edges
end

# SVector of MVectors of point IDs representing the 4 edges of a quadratic quadrilateral
function edges(face::SVector{8, UInt32})
    edges = SVector( MVector{3, UInt32}(face[1], face[2], face[5]),
                     MVector{3, UInt32}(face[2], face[3], face[6]),
                     MVector{3, UInt32}(face[3], face[4], face[7]),
                     MVector{3, UInt32}(face[4], face[1], face[8]) )
    # Order th linear edge vertices by ID
    for edge in edges
        if edge[2] < edge[1]
            e1 = edge[1]
            edge[1] = edge[2]
            edge[2] = e1
        end
    end
    return edges
end

# The unique edges from a vector of triangles or quadrilaterals represented by point IDs
function edges(mesh::UnstructuredMesh_2D)
    edges_filtered = sort(unique(reduce(vcat, edges.(mesh.faces))))
    return [ SVector(e.data) for e in edges_filtered ]
end

# A vector of length 2 SVectors, denoting the face ID each edge is connected to. If the edge
# is a boundary edge, face ID 0 is returned
function edge_face_connectivity(mesh::UnstructuredMesh_2D)
    # Each edge should only border 2 faces if it is an interior edge, and 1 face if it is
    # a boundary edge.
    # Loop through each face in the face_edge_connectivity vector and mark each edge with
    # the faces that it borders.
    if length(mesh.edges) === 0
        @error "Does not have edges!"
    end
    if length(mesh.face_edge_connectivity) === 0
        @error "Does not have face/edge connectivity!"
    end
    edge_face = [MVector{2, UInt32}(0, 0) for _ in eachindex(mesh.edges)]
    for (iface, edges) in enumerate(mesh.face_edge_connectivity)
        for iedge in edges
            # Add the face id in the first non-zero position of the edge_face conn. vec.
            if edge_face[iedge][1] == 0
                edge_face[iedge][1] = iface
            elseif edge_face[iedge][2] == 0
                edge_face[iedge][2] = iface
            else
                @error "Edge $iedge seems to have 3 faces associated with it!"
            end
        end
    end
    return [SVector(sort(two_faces).data) for two_faces in edge_face]
end

# Return an SVector of the points in the edge (Linear)
function edgepoints(edge::SVector{2, UInt32}, points::Vector{Point_2D})
    return SVector(points[edge[1]], points[edge[2]])
end

# Return an SVector of the points in the edge (Quadratic)
function edgepoints(edge::SVector{3, UInt32}, points::Vector{Point_2D})
    return SVector(points[edge[1]], points[edge[2]], points[edge[3]])
end

# Return an SVector of the points in the edge
function edgepoints(edge_id::UInt32, mesh::UnstructuredMesh_2D)
    return edgepoints(mesh.edges[edge_id], mesh.points)
end

# Return an SVector of the points in the face (Triangle)
function facepoints(face::SVector{3, UInt32}, points::Vector{Point_2D})
    return SVector(points[face[1]], points[face[2]], points[face[3]])
end

# Return an SVector of the points in the face (Quadrilateral)
function facepoints(face::SVector{4, UInt32}, points::Vector{Point_2D})
    return SVector(points[face[1]], points[face[2]], points[face[3]], points[face[4]])
end

# Return an SVector of the points in the face (Triangle6)
function facepoints(face::SVector{6, UInt32}, points::Vector{Point_2D})
    return SVector(points[face[1]], points[face[2]], points[face[3]],
                   points[face[4]], points[face[5]], points[face[6]])
end

# Return an SVector of the points in the face (Quadrilateral8)
function facepoints(face::SVector{8, UInt32}, points::Vector{Point_2D})
    return SVector(points[face[1]], points[face[2]], points[face[3]], points[face[4]],
                   points[face[5]], points[face[6]], points[face[7]], points[face[8]])
end

# Return an SVector of the points in the face
function facepoints(edge_id::UInt32, mesh::UnstructuredMesh_2D)
    return facepoints(mesh.faces[edge_id], mesh.points)
end

# Find the faces which share the vertex of ID v.
function faces_sharing_vertex(v::Integer, mesh::UnstructuredMesh_2D)
    shared_faces = UInt32[]
    for i ∈ 1:length(mesh.faces)
        if v ∈  mesh.faces[i]
            push!(shared_faces, UInt32(i))
        end
    end
    return shared_faces
end

# Return the face containing point p.
function findface(p::Point_2D, mesh::UnstructuredMesh_2D)
    if 0 < length(mesh.materialized_faces)
        return UInt32(findface_explicit(p, mesh.materialized_faces))
    else
        return UInt32(findface_implicit(p, mesh.faces, mesh.points))
    end
end

# Find the face containing the point p, with explicitly represented faces
function findface_explicit(p::Point_2D, faces::Vector{<:Face_2D})
    for i ∈ 1:length(faces)
        if p ∈  faces[i]
            return i
        end
    end
    return 0
end

# Return the face containing the point p, with implicitly represented faces
function findface_implicit(p::Point_2D,
                            faces::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}},
                            points::Vector{Point_2D})
    for i ∈ 1:length(faces)
        bool = p ∈  materialize_face(faces[i], points)
        if bool
            return i
        end
    end
    return 0
end

# Return the intersection algorithm that will be used for l ∩ mesh
function get_intersection_algorithm(mesh::UnstructuredMesh_2D)
    if length(mesh.materialized_edges) !== 0
        return "Edges - Explicit"
    elseif length(mesh.edges) !== 0
        return "Edges - Implicit"
    elseif length(mesh.materialized_faces) !== 0
        return "Faces - Explicit"
    else
        return "Faces - Implicit"
    end
end

# Insert the boundary edge into the correct place in the vector of edge indices, based on
# the distance from some reference point
function insert_boundary_edge!(edge_index::UInt32, p_ref::Point_2D, edge_indices::Vector{UInt32},
                               mesh::UnstructuredMesh_2D)
    # Compute the minimum distance from the edge to be inserted to the reference point
    insertion_distance = minimum(distance.(Ref(p_ref),
                                           edgepoints(mesh.edges[edge_index], mesh.points)))
    # Loop through the edge indices until an edge with greater distance from the reference point
    # is found, then insert
    nindices = length(edge_indices)
    for i ∈ 1:nindices
        epoints = edgepoints(mesh.edges[edge_indices[i]], mesh.points)
        edge_distance = minimum(distance.(Ref(p_ref), epoints))
        if insertion_distance < edge_distance
            insert!(edge_indices, i, edge_index)
            return nothing
        end
    end
    insert!(edge_indices, nindices+1, edge_index)
    return nothing
end

# Intersect a line with the mesh. Returns a vector of intersection points, sorted based
# upon distance from the line's start point
function intersect(l::LineSegment_2D, mesh::UnstructuredMesh_2D)
    # Edges are faster, so they are the default
    if length(mesh.edges) !== 0
        if 0 < length(mesh.materialized_edges)
            return intersect_edges_explicit(l, mesh.materialized_edges)
        else
            return intersect_edges_implicit(l, mesh.edges, mesh.points)
        end
    else
        if 0 < length(mesh.materialized_faces)
            return intersect_faces_explicit(l, mesh.materialized_faces)
        else
            return intersect_faces_implicit(l, mesh.faces, mesh.points)
        end
    end
end

# Intersect a line with an implicitly defined edge
function intersect_edge_implicit(l::LineSegment_2D,
                                 edge::SVector{L, UInt32} where {L},
                                 points::Vector{Point_2D})
    return l ∩ materialize_edge(edge, points)
end

# Intersect a line with materialized edges
function intersect_edges_explicit(l::LineSegment_2D, edges::Vector{LineSegment_2D})
    # A vector to hold all of the intersection points
    intersection_points = Point_2D[]
    for edge in edges
        npoints, point = l ∩ edge
        # If the intersections yields 1 or more points, push those points to the array of points
        if 0 < npoints
            push!(intersection_points, point)
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with implicitly defined edges
function intersect_edges_explicit(l::LineSegment_2D, edges::Vector{QuadraticSegment_2D})
    # A vector to hold all of the intersection points
    intersection_points = Point_2D[]
    for edge in edges
        npoints, points = l ∩ edge
        # If the intersections yields 1 or more points, push those points to the array of points
        if 0 < npoints
            append!(intersection_points, points[1:npoints])
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with a vector of implicitly defined linear edges
function intersect_edges_implicit(l::LineSegment_2D,
                                  edges::Vector{SVector{2, UInt32}},
                                  points::Vector{Point_2D})
    # A vector to hold all of the intersection points
    intersection_points = Point_2D[]
    # Intersect the line with each of the faces
    for edge in edges
        npoints, point = intersect_edge_implicit(l, edge, points)
        if 0 < npoints
            push!(intersection_points, point)
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with a vector of implicitly defined quadratic edges
function intersect_edges_implicit(l::LineSegment_2D,
                                  edges::Vector{SVector{3, UInt32}},
                                  points::Vector{Point_2D})
    # An array to hold all of the intersection points
    intersection_points = Point_2D[]
    # Intersect the line with each of the faces
    for edge in edges
        npoints, ipoints = intersect_edge_implicit(l, edge, points)
        if 0 < npoints
            append!(intersection_points, ipoints[1:npoints])
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with an implicitly defined face
function intersect_face_implicit(l::LineSegment_2D,
                                 face::SVector{L, UInt32} where {L},
                                 points::Vector{Point_2D})
    return l ∩ materialize_face(face, points)
end

# Intersect a line with explicitly defined linear faces
function intersect_faces_explicit(l::LineSegment_2D, faces::Vector{<:Face_2D} )
    # An array to hold all of the intersection points
    intersection_points = Point_2D[]
    for face in faces
        npoints, points = l ∩ face
        # If the intersections yields 1 or more points, push those points to the array of points
        if 0 < npoints
            append!(intersection_points, points[1:npoints])
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# Intersect a line with implicitly defined faces
function intersect_faces_implicit(l::LineSegment_2D,
                                  faces::Vector{<:SArray{S, UInt32, 1, L} where {S<:Tuple, L}},
                                  points::Vector{Point_2D})
    # An array to hold all of the intersection points
    intersection_points = Point_2D[]
    # Intersect the line with each of the faces
    for face in faces
        npoints, ipoints = intersect_face_implicit(l, face, points)
        # If the intersections yields 1 or more points, push those points to the array of points
        if 0 < npoints
            append!(intersection_points, ipoints[1:npoints])
        end
    end
    sort_intersection_points!(l[1], intersection_points)
    return intersection_points
end

# If a point is a vertex
function isvertex(p::Point_2D, mesh::UnstructuredMesh_2D)
    for point in mesh.points
        if p ≈ point
            return true
        end
    end
    return false
end

# Return a LineSegment_2D from the point IDs in an edge
function materialize_edge(edge::SVector{2, UInt32}, points::Vector{Point_2D})
    return LineSegment_2D(edgepoints(edge, points))
end

# Return a QuadraticSegment_2D from the point IDs in an edge
function materialize_edge(edge::SVector{3, UInt32}, points::Vector{Point_2D})
    return QuadraticSegment_2D(edgepoints(edge, points))
end

# Return a LineSegment_2D or QuadraticSegment_2D
function materialize_edge(edge_id::UInt32, mesh::UnstructuredMesh_2D)
    return materialize_edge(mesh.edges[edge_id], mesh.points)
end

# Return a materialized edge for each edge in the mesh
function materialize_edges(mesh::UnstructuredMesh_2D)
    return materialize_edge.(mesh.edges, Ref(mesh.points))
end

# Return a Triangle_2D from the point IDs in a face
function materialize_face(face::SVector{3, UInt32}, points::Vector{Point_2D})
    return Triangle_2D(facepoints(face, points))
end

# Return a Quadrilateral_2D from the point IDs in a face
function materialize_face(face::SVector{4, UInt32}, points::Vector{Point_2D})
    return Quadrilateral_2D(facepoints(face, points))
end

# Return a Triangle6_2D from the point IDs in a face
function materialize_face(face::SVector{6, UInt32}, points::Vector{Point_2D})
    return Triangle6_2D(facepoints(face, points))
end

# Return a Quadrilateral8_2D from the point IDs in a face
function materialize_face(face::SVector{8, UInt32}, points::Vector{Point_2D})
    return Quadrilateral8_2D(facepoints(face, points))
end

# Return an SVector of the points in the edge
function materialize_face(face_id::UInt32, mesh::UnstructuredMesh_2D)
    return materialize_face(mesh.faces[face_id], mesh.points)
end

# Return a materialized face for each face in the mesh
function materialize_faces(mesh::UnstructuredMesh_2D)
    return materialize_face.(mesh.faces, Ref(mesh.points))
end

# Return the number of edges in a face
function num_edges(face::SVector{L, UInt32}) where {L}
    if L % 3 === 0
        return 0x00000003
    elseif L % 4 === 0
        return 0x00000004
    else
        # Error
        return 0x00000000
    end
end

function remap_points_to_hilbert(points::Vector{Point_2D})
    bb = boundingbox(points)
    npoints = length(points)
    # Generate a Hilbert curve
    hilbert_points = hilbert_curve(bb, npoints)
    nhilbert_points = length(hilbert_points)
    # For each point, get the index of the hilbert points that is closest
    point_indices = Vector{Int64}(undef, npoints)
    for i ∈ 1:npoints
        min_distance = 1.0e10
        for j ∈ 1:nhilbert_points
            pdistance = distance²(points[i], hilbert_points[j])
            if pdistance < min_distance
                min_distance = pdistance
                point_indices[i] = j
            end
        end
    end
    return sortperm(point_indices)
end

function reorder_points_to_hilbert!(mesh::UnstructuredMesh_2D)
    # Points
    # point_map     maps  new_points[i] == mesh.points[point_map[i]]
    # point_map_inv maps mesh.points[i] == new_points[point_map_inv[i]]
    point_map  = remap_points_to_hilbert(mesh.points)
    point_map_inv = UInt32.(sortperm(point_map))
    # reordered to resemble a hilbert curve
    permute!(mesh.points, point_map)

    # Adjust face indices
    # Point IDs have changed, so we need to change the point IDs referenced by the faces
    new_faces_vec = [ point_map_inv[face] for face in mesh.faces]
    for i in 1:length(mesh.faces)
        mesh.faces[i] = SVector(point_map_inv[mesh.faces[i]])
    end
    return nothing
end

function reorder_faces_to_hilbert!(mesh::UnstructuredMesh_2D)
    face_map  = UInt32.(remap_points_to_hilbert(centroid.(materialize_faces(mesh))))
    permute!(mesh.faces, face_map)
    return nothing
end

function reorder_to_hilbert!(mesh::UnstructuredMesh_2D)
    reorder_points_to_hilbert!(mesh)
    reorder_faces_to_hilbert!(mesh)
    return nothing
end

# Return the ID of the edge shared by two adjacent faces
function shared_edge(face1::UInt32, face2::UInt32, mesh::UnstructuredMesh_2D)
    for edge1 in mesh.face_edge_connectivity[face1]
        for edge2 in mesh.face_edge_connectivity[face2]
            if edge1 == edge2
                return edge1
            end
        end
    end
    return 0x00000000
end

# How to display a mesh in REPL
function Base.show(io::IO, mesh::UnstructuredMesh_2D)
    mesh_type = typeof(mesh)
    println(io, mesh_type)
    println(io, "  ├─ Name      : $(mesh.name)")
    size_MB = Base.summarysize(mesh)/1E6
    if size_MB < 1
        size_KB = size_MB*1000
        println(io, "  ├─ Size (KB) : $size_KB")
    else
        println(io, "  ├─ Size (MB) : $size_MB")
    end
    println(io, "  ├─ Points    : $(length(mesh.points))")
    nedges = length(mesh.edges)
    println(io, "  ├─ Edges     : $nedges")
    println(io, "  │  ├─ Linear         : $(count(x->x isa SVector{2, UInt32},  mesh.edges))")
    println(io, "  │  ├─ Quadratic      : $(count(x->x isa SVector{3, UInt32},  mesh.edges))")
    println(io, "  │  └─ Materialized?  : $(length(mesh.materialized_edges) !== 0)")
    println(io, "  ├─ Faces     : $(length(mesh.faces))")
    println(io, "  │  ├─ Triangle       : $(count(x->x isa SVector{3, UInt32},  mesh.faces))")
    println(io, "  │  ├─ Quadrilateral  : $(count(x->x isa SVector{4, UInt32},  mesh.faces))")
    println(io, "  │  ├─ Triangle6      : $(count(x->x isa SVector{6, UInt32},  mesh.faces))")
    println(io, "  │  ├─ Quadrilateral8 : $(count(x->x isa SVector{8, UInt32},  mesh.faces))")
    println(io, "  │  └─ Materialized?  : $(length(mesh.materialized_faces) !== 0)")
    println(io, "  ├─ Connectivity")
    println(io, "  │  ├─ Edge/Face : $(0 < length(mesh.edge_face_connectivity))")
    println(io, "  │  └─ Face/Edge : $(0 < length(mesh.face_edge_connectivity))")
    println(io, "  ├─ Boundary edges")
    nsides = length(mesh.boundary_edges)
    if nsides !== 0
        println(io, "  │  ├─ Edges : $(mapreduce(x->length(x), +, mesh.boundary_edges))")
    else
        println(io, "  │  ├─ Edges : 0")
    end
    println(io, "  │  └─ Sides : $nsides")
    println(io, "  └─ Face sets : $(length(keys(mesh.face_sets)))")
end

# Sort intersection points, deleting points that are less than minimum_segment_length apart
function sort_intersection_points!(p::Point_2D, points::Vector{Point_2D})
    if 2 <= length(points)
        sortpoints!(p, points)
        # Eliminate any points for which the distance between consecutive points
        # is less than the minimum segment length
        delete_ids = Int64[]
        id_start = 1
        for id_stop ∈ 2:length(points)
            if distance²(points[id_start], points[id_stop]) < minimum_segment_length^2
                push!(delete_ids, id_stop)
            else
                id_start = id_stop
            end
        end
        deleteat!(points, delete_ids)
    else
        return points
    end
end

# Return a mesh with name name, composed of the faces in the set face_ids
function submesh(name::String, mesh::M) where {M <: UnstructuredMesh_2D}
    # Setup faces and get all vertex ids
    face_ids = mesh.face_sets[name]
    submesh_faces = Vector{Vector{UInt32}}(undef, length(face_ids))
    vertex_ids = Set{UInt32}()
    for (i, face_id) in enumerate(face_ids)
        face_vec = collect(mesh.faces[face_id].data)
        submesh_faces[i] = face_vec
        union!(vertex_ids, Set(face_vec))
    end
    # Need to remap vertex ids in faces to new ids
    vertex_ids_sorted = sort(collect(vertex_ids))
    vertex_map = Dict{UInt32, UInt32}()
    for (i,v) in enumerate(vertex_ids_sorted)
        vertex_map[v] = i
    end
    submesh_points = Vector{Point_2D}(undef, length(vertex_ids_sorted))
    for (i, v) in enumerate(vertex_ids_sorted)
        submesh_points[i] = mesh.points[v]
    end
    # remap vertex ids in faces
    for face in submesh_faces
        for (i, v) in enumerate(face)
            face[i] = vertex_map[v]
        end
    end
    # At this point we have points, faces, & name.
    # Just need to get the face sets
    submesh_face_sets = Dict{String, Set{UInt32}}()
    for face_set_name in keys(mesh.face_sets)
        set_intersection = mesh.face_sets[face_set_name] ∩ face_ids
        if length(set_intersection) !== 0
            submesh_face_sets[face_set_name] = set_intersection
        end
    end
    # Need to remap face ids in face sets
    face_map = Dict{UInt32, UInt32}()
    for (i,f) in enumerate(face_ids)
        face_map[f] = i
    end
    for face_set_name in keys(submesh_face_sets)
        new_set = Set{UInt32}()
        for fid in submesh_face_sets[face_set_name]
            union!(new_set, face_map[fid])
        end
        submesh_face_sets[face_set_name] = new_set
    end
    return M(name = name,
             points = submesh_points,
             faces = [SVector{length(f), UInt32}(f) for f in submesh_faces],
             face_sets = submesh_face_sets
            )
end

# Plot
# -------------------------------------------------------------------------------------------------
if enable_visualization
    function convert_arguments(LS::Type{<:LineSegments}, mesh::UnstructuredMesh_2D)
        if 0 < length(mesh.materialized_edges)
            return convert_arguments(LS, mesh.materialized_edges)
        elseif 0 < length(mesh.edges)
            return convert_arguments(LS, materialize_edges(mesh))
        else
            return convert_arguments(LS, materialize_faces(mesh))
        end
    end

    function convert_arguments(P::Type{<:Mesh}, mesh::UnstructuredMesh_2D)
        if 0 < length(mesh.materialized_faces)
            return convert_arguments(P, mesh.materialized_faces)
        else
            return convert_arguments(P, materialize_faces(mesh))
        end
    end
end


"""
    struct CSF

Represents a single _configuration state function_ (CSF) adapted to angular momentum
symmetries. Depending on the type parameters, it can represent both non-relativistic CSFs in
[``LS``-coupling](@ref) and relativistic CSFs in [``jj``-coupling](@ref).

A CSF is defined by the following information:

* The configuration, i.e. an ordered list of subshells together with their occupations.
* A list of intermediate coupling terms (including the seniority quantum number to label
  states in degenerate subspaces) for each subshell in the configuration.
* A list of coupling terms for the coupling tree. The coupling is assumed to be done by,
  first, coupling two orbitals together, then coupling that to the next orbital, and so on.

An instance of a [`CSF`](@ref) object does not specify the ``J_z`` or ``L_z``/``S_z``
quantum number(s).

# Constructors

    CSF(configuration::Configuration, subshell_terms::Vector, terms::Vector)

Constructs an instance of a [`CSF`](@ref) from the information provided. The arguments have
different requirements depending on whether it is `configuration` is based on relativistic
or non-relativistic orbitals.

* If the configuration is based on [`Orbital`](@ref)s, `subshell_terms` must be a list of
  [`IntermediateTerm`](@ref)s and `terms` a list of [`Term`](@ref)s.
* If it is a configuration of [`RelativisticOrbital`](@ref)s, both `subshell_terms` and
  `terms` should both be a list of half-integer values.
"""
struct CSF{O<:AbstractOrbital, T<:Union{Term,HalfInteger}, S}
    config::Configuration{<:O}
    subshell_terms::Vector{IntermediateTerm{T,S}}
    terms::Vector{T}

    function CSF(config::Configuration{O},
                 subshell_terms::Vector{IntermediateTerm{T,S}},
                 terms::Vector{T}) where {O<:Union{<:Orbital,<:RelativisticOrbital},
                                          T<:Union{Term,HalfInt}, S}
        length(subshell_terms) == length(peel(config)) ||
            throw(ArgumentError("Need to provide $(length(peel(config))) subshell terms for $(config)"))
        length(terms) == length(peel(config)) ||
            throw(ArgumentError("Need to provide $(length(peel(config))) terms for $(config)"))

        for (i,(orb,occ,_)) in enumerate(peel(config))
            st = subshell_terms[i]
            assert_unique_classification(orb, occ, st) ||
                throw(ArgumentError("$(st) is not a unique classification of $(orb)$(to_superscript(occ))"))
        end

        new{O,T,S}(config, subshell_terms, terms)
    end

    CSF(config, subshell_terms::Vector{<:IntermediateTerm}, terms::Vector{<:Real}) =
        CSF(config, subshell_terms, convert.(HalfInt, terms))
end

"""
    const NonRelativisticCSF = CSF{<:Orbital,Term}
"""
const NonRelativisticCSF = CSF{<:Orbital,Term}

"""
    const RelativisticCSF = CSF{<:RelativisticOrbital,HalfInt}
"""
const RelativisticCSF = CSF{<:RelativisticOrbital,HalfInt}

Base.:(==)(a::CSF{O,T}, b::CSF{O,T}) where {O,T} =
    (a.config == b.config) && (a.subshell_terms == b.subshell_terms) && (a.terms == b.terms)

# We possibly want to sort on configuration as well
Base.isless(a::CSF, b::CSF) = last(a.terms) < last(b.terms)

num_electrons(csf::CSF) = num_electrons(csf.config)

"""
    orbitals(csf::CSF{O}) -> Vector

Access the underlying list of orbitals

```jldoctest
julia> csf = first(csfs(rc"1s2 2p-2"))
1s²   2p-²
     0     0
      0     0+

julia> orbitals(csf)
2-element Vector{RelativisticOrbital{Int64}}:
 1s
 2p-
```
"""
orbitals(csf::CSF) = orbitals(csf.config)

"""
    csfs(::Configuration) -> Vector{CSF}
    csfs(::Vector{Configuration}) -> Vector{CSF}

Generate all [`CSF`](@ref)s corresponding to a particular configuration or a set of
configurations.
"""
function csfs end

function csfs(config::Configuration)
    map(allchoices(intermediate_terms(peel(config)))) do subshell_terms
        map(intermediate_couplings(subshell_terms)) do coupled_terms
            CSF(config, subshell_terms, coupled_terms[2:end])
        end
    end |> c -> vcat(c...) |> sort
end

csfs(configs::Vector{Configuration{O}}) where O = vcat(map(csfs, configs)...)

Base.length(csf::CSF) = length(peel(csf.config))
Base.getindex(csf::CSF, i::Integer) =
    (peel(csf.config)[i],csf.subshell_terms[i],csf.terms[i])

Base.iterate(csf::CSF, (el, i)=(length(csf)>0 ? csf[1] : nothing,1)) =
    i > length(csf) ? nothing : (el, (csf[i==length(csf) ? i : i+1],i+1))

function Base.show(io::IO, csf::CSF)
    c = core(csf.config)
    p = peel(csf.config)
    nc = length(c)
    if nc > 0
        show(io, c)
        write(io, " ")
    end

    for (i,((orb,occ,state),st,t)) in enumerate(csf)
        show(io, orb)
        occ > 1 && write(io, to_superscript(occ))
        write(io, "($(st)|$(t))")
        i != lastindex(p) && write(io, " ")
    end
    print(io, iseven(parity(csf.config)) ? "+" : "-")
end

function Base.show(io::IO, ::MIME"text/plain", csf::RelativisticCSF)
    c = core(csf.config)
    nc = length(c)
    cw = length(string(c))

    w = 0
    p = peel(csf.config)
    for (orb,occ,state) in p
        w = max(w, length(string(orb))+length(digits(occ)))
    end

    w += 1

    cfmt = FormatExpr("{1:$(cw)s} ")
    ofmt = FormatExpr("{1:<$(w+1)s}")
    ifmt = FormatExpr("{1:>$(w+1)d}")
    rfmt = FormatExpr("{1:>$(w-1)d}/2")

    nc > 0 && printfmt(io, cfmt, c)
    for (orb,occ,state) in p
        printfmt(io, ofmt, join([string(orb), occ > 1 ? to_superscript(occ) : ""], ""))
    end
    println(io)

    nc > 0 && printfmt(io, cfmt, "")
    for it in csf.subshell_terms
        j = it.term
        if denominator(j) == 1
            printfmt(io, ifmt, numerator(j))
        else
            printfmt(io, rfmt, numerator(j))
        end
    end
    println(io)

    nc > 0 && printfmt(io, cfmt, "")
    print(io, " ")
    for j in csf.terms
        if denominator(j) == 1
            printfmt(io, ifmt, numerator(j))
        else
            printfmt(io, rfmt, numerator(j))
        end
    end
    print(io, iseven(parity(csf.config)) ? "+" : "-")
end

export CSF, NonRelativisticCSF, RelativisticCSF, csfs


const EXPONENTS = [
    1e0,   1e1,   1e2,   1e3,   1e4,   1e5,   1e6,   1e7,   1e8,    1e9,
    1e10,  1e11,  1e12,  1e13,  1e14,  1e15,  1e16,  1e17,  1e18,  1e19,
    1e20,  1e21,  1e22,  1e23,  1e24,  1e25,  1e26,  1e27,  1e28,  1e29,
    1e30,  1e31,  1e32,  1e33,  1e34,  1e35,  1e36,  1e37,  1e38,  1e39,
    1e40,  1e41,  1e42,  1e43,  1e44,  1e45,  1e46,  1e47,  1e48,  1e49,
    1e50,  1e51,  1e52,  1e53,  1e54,  1e55,  1e56,  1e57,  1e58,  1e59,
    1e60,  1e61,  1e62,  1e63,  1e64,  1e65,  1e66,  1e67,  1e68,  1e69,
    1e70,  1e71,  1e72,  1e73,  1e74,  1e75,  1e76,  1e77,  1e78,  1e79,
    1e80,  1e81,  1e82,  1e83,  1e84,  1e85,  1e86,  1e87,  1e88,  1e89,
    1e90,  1e91,  1e92,  1e93,  1e94,  1e95,  1e96,  1e97,  1e98,  1e99,
    1e100, 1e101, 1e102, 1e103, 1e104, 1e105, 1e106, 1e107, 1e108, 1e109,
    1e110, 1e111, 1e112, 1e113, 1e114, 1e115, 1e116, 1e117, 1e118, 1e119,
    1e120, 1e121, 1e122, 1e123, 1e124, 1e125, 1e126, 1e127, 1e128, 1e129,
    1e130, 1e131, 1e132, 1e133, 1e134, 1e135, 1e136, 1e137, 1e138, 1e139,
    1e140, 1e141, 1e142, 1e143, 1e144, 1e145, 1e146, 1e147, 1e148, 1e149,
    1e150, 1e151, 1e152, 1e153, 1e154, 1e155, 1e156, 1e157, 1e158, 1e159,
    1e160, 1e161, 1e162, 1e163, 1e164, 1e165, 1e166, 1e167, 1e168, 1e169,
    1e170, 1e171, 1e172, 1e173, 1e174, 1e175, 1e176, 1e177, 1e178, 1e179,
    1e180, 1e181, 1e182, 1e183, 1e184, 1e185, 1e186, 1e187, 1e188, 1e189,
    1e190, 1e191, 1e192, 1e193, 1e194, 1e195, 1e196, 1e197, 1e198, 1e199,
    1e200, 1e201, 1e202, 1e203, 1e204, 1e205, 1e206, 1e207, 1e208, 1e209,
    1e210, 1e211, 1e212, 1e213, 1e214, 1e215, 1e216, 1e217, 1e218, 1e219,
    1e220, 1e221, 1e222, 1e223, 1e224, 1e225, 1e226, 1e227, 1e228, 1e229,
    1e230, 1e231, 1e232, 1e233, 1e234, 1e235, 1e236, 1e237, 1e238, 1e239,
    1e240, 1e241, 1e242, 1e243, 1e244, 1e245, 1e246, 1e247, 1e248, 1e249,
    1e250, 1e251, 1e252, 1e253, 1e254, 1e255, 1e256, 1e257, 1e258, 1e259,
    1e260, 1e261, 1e262, 1e263, 1e264, 1e265, 1e266, 1e267, 1e268, 1e269,
    1e270, 1e271, 1e272, 1e273, 1e274, 1e275, 1e276, 1e277, 1e278, 1e279,
    1e280, 1e281, 1e282, 1e283, 1e284, 1e285, 1e286, 1e287, 1e288, 1e289,
    1e290, 1e291, 1e292, 1e293, 1e294, 1e295, 1e296, 1e297, 1e298, 1e299,
    1e300, 1e301, 1e302, 1e303, 1e304, 1e305, 1e306, 1e307, 1e308,
]

pow10(exp) = (@inbounds v = EXPONENTS[exp+1]; return v)

maxexponent(::Type{Int16}) = 4
maxexponent(::Type{Int32}) = 38
maxexponent(::Type{Int64}) = 308

minexponent(::Type{Int16}) = -5
minexponent(::Type{Int32}) = -38
minexponent(::Type{Int64}) = -308

inttype(::Type{Float16}) = Int16
inttype(::Type{Float32}) = Int32
inttype(::Type{Float64}) = Int64

const BIGN = UInt8('N')
const LITTLEN = UInt8('n')
const BIGA = UInt8('A')
const LITTLEA = UInt8('a')
const BIGI = UInt8('I')
const LITTLEI = UInt8('i')
const BIGF = UInt8('F')
const LITTLEF = UInt8('f')
const BIGT = UInt8('T')
const LITTLET = UInt8('t')
const BIGY = UInt8('Y')
const LITTLEY = UInt8('y')
const BIGE = UInt8('E')
const LITTLEE = UInt8('e')

ParsingException(::Type{<:AbstractFloat}, exp::Signed, row, col) = ParsingException("error parsing a `$T` value on column $col, row $row; exponent out of range: $exp")

function scale(exp, v::T, frac, row, col) where T
    if exp >= 0
        max_exp = maxexponent(T)
        exp > max_exp && throw(ParsingException(T, exp, row, col))
        if exp > 15
            return Float64(Base.TwicePrecision{Float64}(v) * Base.TwicePrecision{Float64}(pow10(exp)))
        else
            return v * pow10(exp)
        end
    else
        min_exp = minexponent(T)
        if exp < min_exp
            -exp + min_exp > -min_exp && throw(ParsingException(T, exp, row, col))
            return Float64(Base.TwicePrecision{Float64}(v) / Base.TwicePrecision{Float64}(pow10(-exp + min_exp)))
        else
            if exp > 15
                return Float64(Base.TwicePrecision{Float64}(v) / Base.TwicePrecision{Float64}(pow10(-exp)))
            else
                return v / pow10(-exp)
            end
        end
    end
end

function parsefield(io::IO, ::Type{T}, opt::CSV.Options, row, col, state, ifmissing::Function) where {T <: Union{Float16, Float32, Float64}}
    mark(io)
    @checkmissingstart()
    minussign = plussign = false
    if b == MINUS # check for leading '-' or '+'
        minussign = true
        c = peekbyte(io)
        if (NEG_ONE < c < TEN) || c == opt.decimal
            b = readbyte(io)
        end
    elseif b == PLUS
        plussign = true
        c = peekbyte(io)
        if (NEG_ONE < c < TEN) || c == opt.decimal
            b = readbyte(io)
        end
    end
    # float digit parsing
    iT = inttype(T)
    v = zero(iT)
    parseddigits = false
    while NEG_ONE < b < TEN
        parseddigits = true
        # process digits
        v *= iT(10)
        v += iT(b - ZERO)
        eof(io) && (state[] = EOF; result = T(v); @goto done)
        b = readbyte(io)
    end
    # if we didn't get any digits and character isn't leading dot, check for NaN/Inf
    if !parseddigits && b != opt.decimal
        if minussign || plussign # skip sign character, if any
            eof(io) && @goto checkmissingend
            b = readbyte(io)
        end
        if b == LITTLEN || b == BIGN
            eof(io) && @goto checkmissingend
            b = readbyte(io)
            (!(b == LITTLEA || b == BIGA) || eof(io)) && (reset(io); b = readbyte(io); @goto checkmissingend)
            b = readbyte(io)
            !(b == LITTLEN || b == BIGN) && (reset(io); b = readbyte(io); @goto checkmissingend)
            result = T(NaN)
            eof(io) && (state[] = EOF; @goto done)
            b = readbyte(io)
            @goto checkdone
        elseif b == LITTLEI || b == BIGI
            eof(io) && @goto checkmissingend
            b = readbyte(io)
            (!(b == LITTLEN || b == BIGN) || eof(io)) && (reset(io); b = readbyte(io); @goto checkmissingend)
            b = readbyte(io)
            !(b == LITTLEF || b == BIGF) && (reset(io); b = readbyte(io); @goto checkmissingend)
            result = T(Inf)
            eof(io) && (state[] = EOF; @goto done)
            b = readbyte(io)
            if b == LITTLEI || b == BIGI
                # read the rest of INFINITY
                eof(io) && (state[] = EOF; @goto done)
                b = readbyte(io)
                b == LITTLEN || b == BIGN || @goto checkdone
                eof(io) && (state[] = EOF; @goto done)
                b = readbyte(io)
                b == LITTLEI || b == BIGI || @goto checkdone
                eof(io) && (state[] = EOF; @goto done)
                b = readbyte(io)
                b == LITTLET || b == BIGT || @goto checkdone
                eof(io) && (state[] = EOF; @goto done)
                b = readbyte(io)
                b == LITTLEY || b == BIGY || @goto checkdone
                eof(io) && (state[] = EOF; @goto done)
                b = readbyte(io)
            end
            @goto checkdone
        else
            @goto checkmissingend
        end
    end
    # parse fractional part
    frac = 0
    result = T(v)
    if b == opt.decimal
        eof(io) && (state[] = EOF; parseddigits ? @goto(done) : @goto(error))
        b = readbyte(io)
    elseif b == LITTLEE || b == BIGE
        @goto parseexp
    else
        @goto checkdone
    end

    while NEG_ONE < b < TEN
        frac += 1
        # process digits
        v *= iT(10)
        v += iT(b - ZERO)
        eof(io) && (state[] = EOF; result = scale(-frac, v, 0, row, col); @goto done)
        b = readbyte(io)
    end
    # parse potential exp
    if b == LITTLEE || b == BIGE
        @label parseexp
        eof(io) && (state[] = EOF; result = scale(-frac, v, 0, row, col); @goto done)
        b = readbyte(io)
        exp = zero(iT)
        negativeexp = false
        if b == MINUS
            negativeexp = true
            b = readbyte(io)
        elseif b == PLUS
            b = readbyte(io)
        end
        parseddigits = false
        while NEG_ONE < b < TEN
            parseddigits = true
            # process digits
            exp *= iT(10)
            exp += iT(b - ZERO)
            eof(io) && (state[] = EOF; result = scale(ifelse(negativeexp, -exp, exp) - frac, v, frac, row, col); @goto done)
            b = readbyte(io)
        end
        result = parseddigits ? scale(ifelse(negativeexp, -exp, exp) - frac, v, frac, row, col) : scale(-frac, v, 0, row, col)
    else
        result = scale(-frac, v, 0, row, col)
    end

    @label checkdone
    @checkdone(done)
    @goto checkmissingend

    @label checkmissingend
    @checkmissingend()
    @goto error

    @label done
    return T(ifelse(minussign, -result, result))

    @label missing
    return ifmissing(row, col)

    @label error
    throw(ParsingException(T, b, row, col))
end

# (-73.99378204345703, -73.99378204345702)
#  (-73.95227813720703, -73.95227813720702)
#  (-73.98616027832031, -73.98616027832033)
#  (-74.00163269042969, -74.00163269042967)
#  (-73.96940612792969, -73.9694061279297)
#  (-73.96797943115234, -73.96797943115236)
#  (-73.95426940917969, -73.95426940917967)
#  (-73.97286224365234, -73.97286224365236)
#  (-73.99149322509766, -73.99149322509767)
#  (-73.97639465332031, -73.97639465332033)
#  (-73.97297668457031, -73.97297668457033)
#  (-73.98991394042969, -73.98991394042967)
#  (-73.98424530029297, -73.98424530029298)
#  (-73.97872161865234, -73.97872161865236)
#  (-73.99348449707031, -73.99348449707033)
#  (-73.96598815917969, -73.96598815917967)
#  (-73.98743438720703, -73.98743438720702)


# Default implmenentation for CuArrays
using CUDA: @workspace, @argout

memsize(a::CuArray) = sizeof(a)

scalar(C::CuArray) = ndims(C)==0 ? collect(C)[] : throw(DimensionMismatch())

function add!(α, A::CuArray{<:Any, N}, CA::Symbol,
                β, C::CuArray{<:Any, N}, indCinA) where {N}

    T = eltype(C)
    N == length(indCinA) || throw(IndexError("Invalid permutation of length $N: $indCinA"))
    CA == :N || CA == :C ||
        throw(ArgumentError("Value of conjA should be :N or :C instead of $CA"))
    opA = (T <: Real || CA == :N) ? CUTENSOR_OP_IDENTITY : CUTENSOR_OP_CONJ
    opAC = CUTENSOR_OP_ADD

    descA = CuTensorDescriptor(A; op = opA)
    descC = CuTensorDescriptor(C; op = CUTENSOR_OP_IDENTITY)
    descD = descC
    typeCompute = convert(cudaDataType, T)
    modeA = collect(Cint, 1:N)
    modeC = collect(Cint, indCinA)
    stream = CuDefaultStream()
    if β == zero(β)
        cutensorPermutation(handle(), T[α], A, descA, modeA, C, descC, modeC,
                            typeCompute, stream)
    else
        cutensorElementwiseBinary(handle(), T[α], A, descA, modeA, T[β], C, descC,
                                  modeC, C, descC, modeC, opAC, typeCompute, stream)
    end

    return C
end

# TODO: the following seems broken due to bug in CUTENSOR?#
# function add!(α, A::CuArray{T, N}, CA::Symbol,
#                 β, C::CuArray{Complex{T}, N}, indCinA) where {T<:CublasReal,N}
#
#     N == length(indCinA) || throw(IndexError("Invalid permutation of length $N: $indCinA"))
#     CA == :N || CA == :C ||
#         throw(ArgumentError("Value of conjA should be :N or :C instead of $CA"))
#     opid = CUTENSOR_OP_IDENTITY
#     opAC = CUTENSOR_OP_ADD
#
#     if β == zero(β)
#         fill!(C, zero(Complex{T}))
#     elseif β != one(β)
#         rmul!(C, β)
#     end
#     sizeC = size(C)
#     stridesC = 2 .* strides(C)
#     Cr = reinterpret(T, C, (2*length(C),))
#
#     @show stridesC, sizeC
#
#     descA = CuTensorDescriptor(A; op = opid)
#     descCr = CuTensorDescriptor(Cr; size = sizeC, strides = stridesC, op = opid)
#     typeCompute = cudaDataType(T)
#     modeA = collect(Cint, 1:N)
#     modeC = collect(Cint, indCinA)
#     stream = CuDefaultStream()
#     cutensorElementwiseBinary(handle(), T[real(α)], A, descA, modeA, T[1], Cr, descCr,
#                               modeC, Cr, descCr, modeC, opAC, typeCompute, stream)
#     if imag(α) != 0
#         Ci = view(Cr, 2:length(Cr))
#         descCi = CuTensorDescriptor(Ci; size = sizeC, strides = stridesC, op = opid)
#         cutensorElementwiseBinary(handle(), T[imag(α)], A, descA, modeA, T[1], Ci, descCi,
#                                   modeC, Ci, descCi, modeC, opAC, typeCompute, stream)
#     end
#     return C
# end

function trace!(α, A::CuArray, CA::Symbol, β, C::CuArray,
                indCinA, cindA1, cindA2)

    T = eltype(C)
    NA, NC = ndims(A), ndims(C)
    NC == length(indCinA) ||
        throw(IndexError("Invalid selection of $NC out of $NA: $indCinA"))
    NA-NC == 2*length(cindA1) == 2*length(cindA2) ||
        throw(IndexError("invalid number of trace dimension"))

    opA = (T <: Real || CA == :N) ? CUTENSOR_OP_IDENTITY : CUTENSOR_OP_CONJ
    opReduce = CUTENSOR_OP_ADD

    sizeA = i->size(A, i)
    strideA = i->stride(A, i)
    tracesize = sizeA.(cindA1)
    tracesize == sizeA.(cindA2) || throw(DimensionMismatch("non-matching trace sizes"))
    size(C) == sizeA.(indCinA) || throw(DimensionMismatch("non-matching sizes"))

    newstrides = (strideA.(indCinA)..., (strideA.(cindA1) .+ strideA.(cindA2))...)
    newsize = (size(C)..., tracesize...)
    descA = CuTensorDescriptor(A; op = opA, size = newsize, strides = newstrides)
    descC = CuTensorDescriptor(C; op = CUTENSOR_OP_IDENTITY)
    descD = descC
    typeCompute = cutensorComputeType(T)
    modeA = collect(Cint, 1:NA)
    modeC = collect(Cint, 1:NC)
    stream = CuDefaultStream()
    @workspace fallback=1<<13 size=@argout(
            cutensorReductionGetWorkspace(handle(),
                A, descA, modeA,
                C, descC, modeC,
                C, descC, modeC,
                opReduce, typeCompute,
                out(Ref{UInt64}(C_NULL)))
        )[] workspace->begin
            cutensorReduction(handle(),
                T[α], A, descA, modeA,
                T[β], C, descC, modeC,
                        C, descC, modeC,
                opReduce, typeCompute,
                workspace, sizeof(workspace), stream)
        end
    return C
end

function contract!(α, A::CuArray, CA::Symbol,
                    B::CuArray, CB::Symbol,
                    β, C::CuArray,
                    oindA::IndexTuple, cindA::IndexTuple,
                    oindB::IndexTuple, cindB::IndexTuple,
                    indCinoAB::IndexTuple, syms::Union{Nothing, NTuple{3,Symbol}} = nothing)

    pA = (oindA...,cindA...)
    (length(pA) == ndims(A) && isperm(pA)) ||
        throw(IndexError("invalid permutation of length $(ndims(A)): $pA"))
    pB = (oindB...,cindB...)
    (length(pB) == ndims(B) && isperm(pB)) ||
        throw(IndexError("invalid permutation of length $(ndims(B)): $pB"))
    (length(oindA) + length(oindB) == ndims(C)) ||
        throw(IndexError("non-matching output indices in contraction"))
    (ndims(C) == length(indCinoAB) && isperm(indCinoAB)) ||
        throw(IndexError("invalid permutation of length $(ndims(C)): $indCinoAB"))

    sizeA = i->size(A, i)
    sizeB = i->size(B, i)
    sizeC = i->size(C, i)

    csizeA = sizeA.(cindA)
    csizeB = sizeB.(cindB)
    osizeA = sizeA.(oindA)
    osizeB = sizeB.(oindB)

    csizeA == csizeB ||
        throw(DimensionMismatch("non-matching sizes in contracted dimensions"))
    sizeAB = let osize = (osizeA..., osizeB...)
        i->osize[i]
    end
    sizeAB.(indCinoAB) == size(C) ||
        throw(DimensionMismatch("non-matching sizes in uncontracted dimensions"))

    TC = eltype(C)

    CA == :N || CA == :C ||
        throw(ArgumentError("Value of conjA should be :N or :C instead of $CA"))
    CB == :N || CB == :C ||
        throw(ArgumentError("Value of conjB should be :N or :C instead of $CB"))

    opA = (TC <: Real || CA == :N) ? CUTENSOR_OP_IDENTITY : CUTENSOR_OP_CONJ
    opB = (TC <: Real || CB == :N) ? CUTENSOR_OP_IDENTITY : CUTENSOR_OP_CONJ
    opC = CUTENSOR_OP_IDENTITY

    strideA = i->stride(A, i)
    strideB = i->stride(B, i)

    cstrideA = strideA.(cindA)
    cstrideB = strideB.(cindB)
    ostrideA = strideA.(oindA)
    ostrideB = strideB.(oindB)

    descA = CuTensorDescriptor(A; op = opA, size = (osizeA..., csizeA...),
                                    strides = (ostrideA..., cstrideA...))
    descB = CuTensorDescriptor(B; op = opB, size = (osizeB..., csizeB...),
                                    strides = (ostrideB..., cstrideB...))
    descC = CuTensorDescriptor(C)
    T = eltype(C)
    typeCompute = cutensorComputeType(T)
    opOut = CUTENSOR_OP_IDENTITY

    NoA = length(osizeA)
    NoB = length(osizeB)
    Nc = length(csizeA)
    modeoA = ntuple(n->n, NoA)
    modeoB = ntuple(n->NoA+n, NoB)
    modec = ntuple(n->NoA+NoB+n, Nc)

    modeA = collect(Cint, (modeoA..., modec...))
    modeB = collect(Cint, (modeoB..., modec...))
    modeC = collect(Cint, indCinoAB)

    algo = CUTENSOR_ALGO_DEFAULT
    stream = CuDefaultStream()
    pref = CUTENSOR_WORKSPACE_RECOMMENDED

    alignmentRequirementA = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), A, descA, alignmentRequirementA)
    alignmentRequirementB = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), B, descB, alignmentRequirementB)
    alignmentRequirementC = Ref{UInt32}(C_NULL)
    cutensorGetAlignmentRequirement(handle(), C, descC, alignmentRequirementC)
    desc = Ref{cutensorContractionDescriptor_t}()
    cutensorInitContractionDescriptor(handle(),
                                      desc,
                   descA, modeA, alignmentRequirementA[],
                   descB, modeB, alignmentRequirementB[],
                   descC, modeC, alignmentRequirementC[],
                   descC, modeC, alignmentRequirementC[],
                   typeCompute)

    find = Ref{cutensorContractionFind_t}()
    cutensorInitContractionFind(handle(), find, algo)

    @workspace fallback=1<<27 size=@argout(
            cutensorContractionGetWorkspace(handle(), desc, find, pref,
                                            out(Ref{UInt64}(C_NULL)))
        )[] workspace->begin
            plan_ref = Ref{cutensorContractionPlan_t}()
            cutensorInitContractionPlan(handle(), plan_ref, desc, find, sizeof(workspace))

            cutensorContraction(handle(), plan_ref, T[α], A, B, T[β], C, C,
                                workspace, sizeof(workspace), stream)
        end

    return C
end

# overwrite similar_from_indices to return zero initialized arrays
function similar_from_indices(T::Type, ind::IndexTuple, A::CuArray, CA::Symbol)
    sz = similarstructure_from_indices(T, ind, A, CA)
    return fill!(similar(A, T, sz), zero(T))
end
function similar_from_indices(T::Type, poA::IndexTuple, poB::IndexTuple,
                                p1::IndexTuple, p2::IndexTuple,
                                A::CuArray, B::CuArray, CA::Symbol, CB::Symbol)
    sz = similarstructure_from_indices(T, poA, poB, p1, p2, A, B, CA, CB)
    return fill!(similar(A, T, sz), zero(T))
end

# overwrite cached_similar_from_indices in order not to use cache for CuArray objects
function cached_similar_from_indices(sym::Symbol, T::Type, p1::IndexTuple, p2::IndexTuple, A::CuArray, CA::Symbol)
    # also zero fill to avoid problems with cutensorElementwiseBinary
    return similar_from_indices(T, p1, p2, A, CA)
end

function cached_similar_from_indices(sym::Symbol, T::Type, poA::IndexTuple, poB::IndexTuple,
    p1::IndexTuple, p2::IndexTuple, A::CuArray, B::CuArray, CA::Symbol, CB::Symbol)
    # also zero fill to avoid problems with cutensorElementwiseBinary
    return similar_from_indices(T, poA, poB, p1, p2, A, B, CA, CB)
end