module SoilMoistureMod
using Parameters
using LinearAlgebra

export soil_struct, van_Genuchten,  matric_potential, root_uptake, predictor_corrector, compute_grid_settings

# --- Define soil layers

# Number of soil layers
@with_kw mutable struct soil_struct{}
   nsoi::UInt      # Number of soil layers
   z_bnds = zeros(nsoi+1)
   dz              = ones(nsoi)       # Soil layer thickness (cm)
   z_plus_onehalf  = zeros(nsoi) # Soil depth [cm] at i+1/2 interface between layers i & i+1 [negative distance from surface]
   z               = zeros(nsoi) # Soil depth [cm] at center of layer i [negative distance from surface]
   dz_plus_onehalf = zeros(nsoi) # Thickness between between z[i] & z[i+1]
   θ           = zeros(nsoi) # Soil Moisture
   ψ             = zeros(nsoi) # Matric potential
   θ_0          = 0.0  # Initial SM
   ψ_0            = 0.0  # initial matric potential
   K               = zeros(nsoi) # Hydraulic conductivity (cm H2O/s)
   cap             = zeros(nsoi) # Specific moisture capacity (/cm)
   Q_0              = 0.0 # Infiltration flux (cm H2O/s)
   QN              = 0.0 # Drainage flux (cm H2O/s)
   dθ          = 0.0 # Change in soil moisture (cm H2O)
   err             = 0.0 # Water balance error (cm H2O)
end

# Van Genuchten Function
function van_Genuchten(params, psi)

# ----------------------------------
# van Genuchten [1980] relationships
# ----------------------------------

# --- Soil parameters

theta_res = params[1];   # Residual water content
theta_sat = params[2];   # Volumetric water content at saturation
alpha = params[3];       # Inverse of the air entry potential
n = params[4];           # Pore-size distribution index
m = params[5];           # Exponent
Ksat = params[6];        # Hydraulic conductivity at saturation
ityp = params[7];        # Soil texture flag

# --- Effective saturation [Se] for specified matric potential [psi]

if (psi <= 0)
   Se = (1 + (alpha * abs(psi))^n)^-m
else()
   Se = 1
end

# --- Volumetric soil moisture [theta] for specified matric potential [psi]

theta = theta_res + (theta_sat - theta_res) * Se

# --- Hydraulic conductivity [K] for specified matric potential [psi]

if (Se <= 1)
   K = Ksat * sqrt(Se) * (1 - (1 - Se^(1/m))^m)^2

   # Special case for Haverkamp et al. (1977) sand [ityp = 1] & Yolo light clay [ityp = 2]

   if (ityp == 1)
      K = Ksat * 1.175e6 / (1.175e6 + abs(psi)^4.74)
   end
   if (ityp == 2)
      K = Ksat * 124.6/ (124.6 + abs(psi)^1.77)
   end

else()

   K = Ksat

end

# --- Specific moisture capacity [cap] for specified matric potential [psi]

if (psi <= 0)
   num = alpha * m * n * (theta_sat - theta_res) * (alpha * abs(psi))^(n-1)
   den =  (1 + (alpha * abs(psi))^n)^(m+1)
   cap = num / den
else()
   cap = 0
end

return theta, K, cap

end # function

# Van Genuchten Function
function vG(ψ, params=params)
# ----------------------------------
# van Genuchten [1980] relationships
# ----------------------------------

# --- Soil parameters
θ_res = params[1];   # Residual water content
θ_sat = params[2];   # Volumetric water content at saturation
α = params[3];       # Inverse of the air entry potential
n = params[4];           # Pore-size distribution index
m = params[5];           # Exponent
Ksat = params[6];        # Hydraulic conductivity at saturation

# --- Effective saturation [Se] for specified matric potential [psi]
if (ψ <= 0)
   Se = (1 + (α * abs(ψ))^n)^-m
else()
   Se = 1
end

# --- Volumetric soil moisture [theta] for specified matric potential [psi]
θ = θ_res + (θ_sat - θ_res) * Se

# --- Hydraulic conductivity [K] for specified matric potential [psi]
if Se <= 1
   K = Ksat * sqrt(Se) * (1 - (1 - Se^(1/m))^m)^2
else
   K = Ksat
end

# --- Specific moisture capacity [cap] for specified matric potential [psi]
if (ψ <= 0)
   num = α * m * n * (θ_sat - θ_res) * (α * abs(ψ))^(n-1)
   den =  (1 + (α * abs(ψ))^n)^(m+1)
   cap = num / den
else
   cap = 0
end

return θ, K, cap
end # function

# Other Functions
function tridiagonal_solver(a, b, c, d, n)

# Solve for U given the set of equations R * U = D; where U is a vector
# of length N; D is a vector of length N; & R is an N x N tridiagonal
# matrix defined by the vectors A, B, C each of length N. A[1] &
# C[N] are undefined and are not referenced.
#
#     |B[1] C[1] ...  ...  ...                     |
#     |A[2] B[2] C[2] ...  ...                     |
# R = |     A[3] B[3] C[3] ...                     |
#     |                    ... A[N-1] B[N-1] C[N-1]|
#     |                    ... ...    A[N]   B[N]  |
#
# The system of equations is written as:
#
#    A_i * U_i-1 + B_i * U_i + C_i * U_i+1 = D_i
#
# for i = 1 to N. The solution is found by rewriting the
# equations so that:
#
#    U_i = F_i - E_i * U_i+1

# --- Forward sweep [1 -> N] to get E & F
# Initialize E and F
e = similar(a)
f = similar(a);

e[1] = c[1] / b[1]
for i = 2: 1: n-1
   e[i] = c[i] / (b[i] - a[i] * e[i-1])
end

f[1] = d[1] / b[1]
for i = 2: 1: n
   f[i] = (d[i] - a[i] * f[i-1]) / (b[i] - a[i] * e[i-1])
end

# --- Backward substitution [N -> 1] to solve for U
u = zeros(n);
u[n] = f[n]

for i = n-1: -1: 1
   u[i] = f[i] - e[i] * u[i+1]
end

return u

end # function

function matric_potential(type, params, theta)

# --- Calculate psi for a given theta

if type == "van_Genuchten"

   theta_res = params[1];    # Residual water content
   theta_sat = params[2];    # Volumetric water content at saturation
   alpha = params[3];        # Inverse of the air entry potential
   n = params[4];            # Pore-size distribution index
   m = params[5];            # Exponent

   Se = (theta - theta_res) / (theta_sat - theta_res);
   psi = -((Se^(-1/m) - 1)^(1/n)) / alpha;

elseif type == "Campbell"

   theta_sat = params[1];    # Volumetric water content at saturation
   psi_sat = params[2];      # Matric potential at saturation
   b = params[3];            # Exponent

   psi = psi_sat * (theta / theta_sat)^-b;

end

return psi
end # function

function root_uptake(soil::soil_struct,ET)

   #Compute the soil sink terms

   densum = 0

   for i = 1:soil.nsoi
      if soil.psi[i] > soil.psidry
         soil.beta[i] = (soil.psi[i]-soil.psidry)/(soil.psiopt-soil.psidry)
      elseif soil.psi[i] <= soil.psidry
         soil.beta[i] = 0.0
      end

      densum = densum + soil.fz[i]*soil.beta[i]

   end

   # compute the sink terms

   for i = 1:soil.nsoi
      soil.sink[i] = ET*soil.fz[i]*soil.beta[i]/densum
   end

   return soil.sink

end # funnction

# Predictor Corrector Function
function predictor_corrector(soil::soil_struct, params, ET::Float64, dt)

# -------------------------------------------------------------
# Use predictor-corrector method to solve the Richards equation
# -------------------------------------------------------------

# Input
# dt                   ! Time step [s]
# soil.nsoi            ! Number of soil layers
# soil.functions       ! van Genuchten | Campbell relationships
# soil.dz_plus_onehalf ! Thickness between between z[i] & z[i+1] (cm)
# soil.dz              ! Soil layer thickness [cm]
# soil.psi0            ! Soil surface matric potential boundary condition [cm]
#
# Input/output
# soil.theta           ! Volumetric soil moisture
# soil.psi             ! Matric potential [cm]
#
# Output
# soil.K               ! Hydraulic conductivity [cm H2O/s]
# soil.cap             ! Specific moisture capacity [/cm]
# soil.Q0              ! Infiltration flux [cm H2O/s]
# soil.QN              ! Drainage flux [cm H2O/s]
# soil.dtheta          ! Change in soil moisture [cm H2O]
# soil.err             ! Water balance error (cm H2O)

# --- Save current soil moisture & matric potential for time n

theta_n = copy(soil.theta)
psi_n   = copy(soil.psi)

# --- Predictor step using implict solution for time n+1/2

# Hydraulic properties for current psi:
# theta - volumetric soil moisture
# K     - hydraulic conductivity
# cap   - specific moisture capacity

for i = 1:soil.nsoi
   if soil.functions == "van_Genuchten"
      soil.theta[i], soil.K[i], soil.cap[i] = van_Genuchten(params, soil.psi[i])
   elseif soil.functions == "Campbell"
      soil.theta[i], soil.K[i], soil.cap[i] = Campbell(params, soil.psi[i])
    end
end

# Hydraulic conductivity at i+1/2 interface between layers i & i+1 is the arithmetic mean()
K_plus_onehalf = soil.theta*0
for i = 1:soil.nsoi-1
   K_plus_onehalf[i] = 0.5 * (soil.K[i] + soil.K[i+1])
end

# Hydraulic conductivity at i=1/2 between surface (i=0) & first layer i=1

K_onehalf = soil.K[1]

# dz at i=1/2 between surface (i=0) & first layer i=1

dz_onehalf = 0.5 * soil.dz[1]

# Compute the sink terms
if soil.ssflag == 1

   soil.sink = root_uptake(soil,ET)
else

   soil.sink = soil.sink.*0
end

# Terms for tridiagonal matrix
a = zeros(soil.nsoi);
b = zeros(soil.nsoi);
c = zeros(soil.nsoi);
d = zeros(soil.nsoi);

i = 1
a[i] = 0
c[i] = -K_plus_onehalf[i] / soil.dz_plus_onehalf[i]
b[i] = soil.cap[i] * soil.dz[i] / (0.5 * dt) + K_onehalf / dz_onehalf - c[i]
d[i] = soil.cap[i] * soil.dz[i] / (0.5 * dt) * soil.psi[i] + K_onehalf / dz_onehalf * soil.psi0 + K_onehalf - K_plus_onehalf[i] - soil.sink[i]

for i = 2:soil.nsoi-1
   a[i] = -K_plus_onehalf[i-1] / soil.dz_plus_onehalf[i-1]
   c[i] = -K_plus_onehalf[i] / soil.dz_plus_onehalf[i]
   b[i] = soil.cap[i] * soil.dz[i] / (0.5 * dt) - a[i] - c[i]
   d[i] = soil.cap[i] * soil.dz[i] / (0.5 * dt) * soil.psi[i] + K_plus_onehalf[i-1] - K_plus_onehalf[i] - soil.sink[i]
end

i = soil.nsoi
a[i] = -K_plus_onehalf[i-1] / soil.dz_plus_onehalf[i-1]
c[i] = 0
b[i] = soil.cap[i] * soil.dz[i] / (0.5 * dt) - a[i] - c[i]
d[i] = soil.cap[i] * soil.dz[i] / (0.5 * dt) * soil.psi[i] + K_plus_onehalf[i-1] - soil.K[i] - soil.sink[i]

# Solve for psi at n+1/2

psi_pred = tridiagonal_solver(a, b, c, d, soil.nsoi)
# --- Corrector step using Crank-Nicolson solution for time n+1

# Hydraulic properties for psi_pred

for i = 1:soil.nsoi
   if soil.functions == "van_Genuchten"
      soil.theta[i], soil.K[i], soil.cap[i] = van_Genuchten(params, psi_pred[i])
   elseif  soil.functions == "Campbell"
      soil.theta[i], soil.K[i], soil.cap[i] = Campbell(params, psi_pred[i])
    end
end

# Hydraulic conductivity at i+1/2 interface between layers i & i+1

for i = 1:soil.nsoi-1
   K_plus_onehalf[i] = 0.5 * (soil.K[i] + soil.K[i+1])
end

# Hydraulic conductivity at i=1/2 between surface (i=0) & first layer i=1

K_onehalf = soil.K[1]

# dz at i=1/2 between surface (i=0) & first layer i=1

dz_onehalf = 0.5 * soil.dz[1]

# Compute the sink terms
if soil.ssflag == 1

   soil.sink = root_uptake(soil,ET)
else

   soil.sink = soil.sink.*0
end

# Terms for tridiagonal matrix

i = 1
a[i] = 0.0
c[i] = -K_plus_onehalf[i] / (2.0 * soil.dz_plus_onehalf[i])
b[i] = soil.cap[i] * soil.dz[i] / dt  + K_onehalf / (2.0 * dz_onehalf) - c[i]
d[i] = soil.cap[i] * soil.dz[i] / dt * soil.psi[i] + K_onehalf / (2.0 * dz_onehalf) * soil.psi0 + K_onehalf / (2.0 * dz_onehalf) * (soil.psi0 - soil.psi[i]) + c[i] * (soil.psi[i] - soil.psi[i+1]) + K_onehalf - K_plus_onehalf[i] - soil.sink[i]

for i = 2:soil.nsoi-1
   a[i] = -K_plus_onehalf[i-1] / (2.0 * soil.dz_plus_onehalf[i-1])
   c[i] = -K_plus_onehalf[i] / (2.0 * soil.dz_plus_onehalf[i])
   b[i] = soil.cap[i] * soil.dz[i] / dt - a[i] - c[i]
   d[i] = soil.cap[i] * soil.dz[i] / dt * soil.psi[i] - a[i] * (soil.psi[i-1] - soil.psi[i]) + c[i] * (soil.psi[i] - soil.psi[i+1]) + K_plus_onehalf[i-1] - K_plus_onehalf[i] - soil.sink[i]
end

i = soil.nsoi
a[i] = -K_plus_onehalf[i-1] / (2.0 * soil.dz_plus_onehalf[i-1])
c[i] = 0.0
b[i] = soil.cap[i] * soil.dz[i] / dt - a[i] - c[i]
d[i] = soil.cap[i] * soil.dz[i] / dt * soil.psi[i] - a[i] * (soil.psi[i-1] - soil.psi[i]) + K_plus_onehalf[i-1] - soil.K[i] - soil.sink[i]

# Solve for psi at n+1

soil.psi = tridiagonal_solver(a, b, c, d, soil.nsoi)
# --- Check water balance()

soil.Q0 = -K_onehalf / (2.0 * dz_onehalf) * ((soil.psi0 - psi_n[1]) + (soil.psi0 - soil.psi[1])) - K_onehalf
soil.QN = -soil.K[soil.nsoi]

soil.dtheta = 0.0
for i = 1:soil.nsoi
   soil.dtheta = soil.dtheta + (soil.theta[i] - theta_n[i]) * soil.dz[i]
end

soil.err = soil.dtheta - (soil.QN - soil.Q0) * dt

return soil

end # function

# Computational Grid
function compute_grid_settings(soil::soil_struct)
   # Set the Computational Grid for the Solver
   # Soil layer thickness (cm)
   for i = 1:soil.nsoi
      soil.dz[i] = 1.0
   end

   # Soil depth [cm] at i+1/2 interface between layers i & i+1 [negative distance from surface]

   soil.z_plus_onehalf[1] = -soil.dz[1]
   for i = 2:soil.nsoi
      soil.z_plus_onehalf[i] = soil.z_plus_onehalf[i-1] - soil.dz[i]
   end

   # Soil depth [cm] at center of layer i [negative distance from surface]

   soil.z[1]  = 0.5 * soil.z_plus_onehalf[1]
   soil.fz[1] = 1 - soil.bi^abs(soil.z[1])
   for i = 2:soil.nsoi
      soil.z[i] = 0.5 * (soil.z_plus_onehalf[i-1] + soil.z_plus_onehalf[i])

      # Assign the root fraction here
      soil.fz[i] = 1 - soil.bi^abs(soil.z[i]) - soil.fz[i-1] # if z is in cm
   end

   # Thickness between between z[i] & z[i+1]

   for i = 1:soil.nsoi-1
      soil.dz_plus_onehalf[i] = soil.z[i] - soil.z[i+1]
   end
   soil.dz_plus_onehalf[soil.nsoi] = 0.5 * soil.dz[soil.nsoi]

   return soil
end # Function

end # module


module NGPtools

import Distributions.MvNormal
import SStools.simulate

export assemble_matrices, generate, sample

"""
Observation matrix for state space model. In this case, we assume that the
Np components of the observation vector each track the state of the same
index (more specifically, its U component). All other terms are ignored.
Np is the dimension of the observation vector.
Ns is the number of underlying states.
"""
function Z(Np, Ns)
    maxind = min(3Ns, Np)
    Zmat = sparse(1:maxind, 1:maxind, 1, Np, 3Ns)
    return Zmat
end

"""
State transition matrix. Defined in (5) of Zhu and Dunson. Order of entries
is all state variables, followed by all derivatives, then all latents.
Ns is the dimension of the underlying state space.
"""
function G(Ns, δ; approx=false)
    gmat = eye(3 * Ns)
    gmat += diagm(δ * ones(2Ns), Ns)
    if approx
        gmat += diagm(δ^2/2 * ones(Ns), 2Ns)
    end

    return sparse(gmat)
end

"""
Covariance of state disturbances. Defined in (5) of Zhu and Dunson. Order of
entries is all state variables, followed by all derivatives, then all latents.
"""
function W(Ns, δ, varU::Array{Float64}, varA::Array{Float64}; approx=false)
    if approx
        wmat = δ * spdiagm([varU ; varA])
    else
        wmat = zeros(Float64, 3Ns, 3Ns)
        for idx in 1:Ns
            wmat[idx, idx] = (δ^3 / 3)varU[idx] + (δ^5 / 20)varA[idx]
            wmat[idx, Ns + idx] = (δ^2 / 2)varU[idx] + (δ^4 / 8)varA[idx]
            wmat[idx, 2Ns + idx] = (δ^3 / 6)varA[idx]
            wmat[Ns + idx, idx] = wmat[idx, Ns + idx]
            wmat[Ns + idx, Ns + idx] = δ * varU[idx] + (δ^3 / 3)varA[idx]
            wmat[Ns + idx, 2Ns + idx] = (δ^2 / 2)varA[idx]
            wmat[2Ns + idx, idx] = wmat[idx, 2Ns + idx]
            wmat[2Ns + idx, Ns + idx] = wmat[Ns + idx, 2Ns + idx]
            wmat[2Ns + idx, 2Ns + idx] = δ * varA[idx]
        end
    end

    return wmat
end

function W(Ns, δ, varU::Array{Float64}, varA::Float64; approx=false)
    return W(Ns, δ, varU, varA * ones(Ns); approx=approx)
end

function W(Ns, δ, varU::Float64, varA::Array{Float64}; approx=false)
    return W(Ns, δ, varU * ones(Ns), varA; approx=approx)
end

function W(Ns, δ, varU::Float64, varA::Float64; approx=false)
    return W(Ns, δ, varU * ones(Ns), varA * ones(Ns); approx=approx)
end

"""
Take nested GP parameters and return matrices suitable for feeding into
state space model.

Np is the dimension of the observation vector
Ns is the number of underlying states
δ is a vector of differences between observation times. If observations are
uniformly spaced, δ can be a scalar or a length-1 vector.
"""
function assemble_matrices(Np, Ns, δ, σϵ, σU, σA, σμ, σα; approx=false)
    Nt = length(δ)  # number of time points

    # calculate dimensions
    # Note: for an observation vector of size Np, we introduce 3 latent
    # variables per observation and 2 or 3 noise variables (depending on
    # approximation)
    Nm = 3 * Ns
    if approx
        Nr = 2 * Ns
    else
        Nr = 3 * Ns
    end

    # preallocate arrays:
    H = σϵ^2 * eye(Ns)
    T = Array(Float64, Nm, Nm, Nt)
    Q = Array(Float64, Nr, Nr, Nt)

    if approx
        R = sparse([2:3:Nm ; 3:3:Nm], [1:2:Nr ; 2:2:Nr], 1., Nm, Nr)
    else
        R = speye(Float64, Nm)
    end

    for t in 1:Nt
        T[:, :, t] = G(Ns, δ[t], approx=approx)
        Q[:, :, t] = W(Ns, δ[t], σU.^2, σA.^2, approx=approx)
    end

    if Nt == 1
        T = sparse(squeeze(T, 3))
        Q = squeeze(Q, 3)
        if approx
            Q = sparse(Q)
        end
    end

    a0 = zeros(Nm)
    P0 = diagm(repeat([σμ, σμ, σα], inner=[Ns]))

    return Z(Np, Ns), H, T, R, Q, a0, P0
end

"""
Generate nGP data according to the state space model. Uses matrices
constructed in assemble_matrices.

Np is the dimension of the observation vector
Ns is the number of underlying states
δ is a spacing between observations
"""
function generate(Np, Ns, Nt, δ, σϵ, σU, σA, σμ, σα; approx=false)

    Z, H, T, R, Q, a0, P0 = assemble_matrices(Np, Ns, δ, σϵ, σU, σA, σμ, σα, approx=approx)

    Nm = size(Z, 2)

    α = Array(Float64, Nm, Nt)
    y = Array(Float64, Np, Nt)

    # initialize
    α[:, 1] = rand(MvNormal(full(a0), full(P0)))

    # simulate
    for t in 1:Nt
        ϵ = rand(MvNormal(H))
        local Q_t = ndims(Q) < 3 ? Q : Q[:, :, t]
        η = rand(MvNormal(Q_t))
        y[:, t] = Z * α[:, t] + ϵ

        if t < Nt
            local T_t = ndims(T) < 3 ? T : slice(T, :, :, t)
            α[:, t + 1] = T_t * α[:, t] + R * η
        end
    end

    return y, α
end

"""
Sample Ns times from the nGP posterior, given observations y.
"""
function _sample(y, Ns, a0, P0, Z, H, T, R, Q; interleaved=true)
    Nt = size(y, 2)
    Nm = size(Z, 2)

    α_samples = Array(Float64, Nm, Nt, Ns)

    for idx in 1:Ns
        α_samples[:, :, idx] = simulate(y, a0, P0, Z, H, T, R, Q, interleaved=interleaved)
    end

    return α_samples

end

"""
Version where spacing is uniform, so number of time points must be specified
explicitly.
"""
function sample(y, Nsamples, Nstates, δ, σϵ, σU, σA, σμ, σα; approx=false)
    Np = size(y, 1)
    Z, H, T, R, Q, a0, P0 = assemble_matrices(Np, Nstates, δ, σϵ, σU, σA, σμ, σα, approx=approx)
    _sample(y, Nsamples, a0, P0, Z, H, T, R, Q)
end

"""
Version where observation model (Z and H) is specified externally.
"""
function sample(y, Nsamples, Z::Array{Float64}, H::Array{Float64}, δ::Float64,
    σU, σA, σμ, σα; approx=false, interleaved=true)
    Nstates = convert(Int64, size(Z, 2) / 3)
    Np = size(y, 1)
    _, _, T, R, Q, a0, P0 = assemble_matrices(Np, Nstates, δ, 0, σU, σA, σμ, σα, approx=approx)
    _sample(y, Nsamples, a0, P0, Z, H, T, R, Q; interleaved=interleaved)
end

end  # module


###
### logistic and logit
###

logistic(x::Real) = inv(one(x) + exp(-x))

function logistic_logjac(x::Real)
    mx = -abs(x)
    mx - 2*log1p(exp(mx))
end

logit(x::Real) = log(x / (one(x) - x))

logit_logjac(y) = -log(y) - log1p(-y)

###
### calculations
###

"""
    $SIGNATURES

Number of elements (strictly) above the diagonal in an ``n×n`` matrix.
"""
unit_triangular_dimension(n::Int) = n * (n-1) ÷ 2

###
### type calculations
###

"""
    $(SIGNATURES)

Extend element type of argument so that it is closed under the algebra used by this package.

Pessimistic default for non-real types.
"""
function extended_eltype(::Type{S}) where S
    T = eltype(S)
    T <: Real ? typeof(√(one(T))) : Any
end

extended_eltype(x::T) where T = extended_eltype(T)

####
#### random values
####

"Shared part of docstrings for keyword arguments of or passed to [`random_reals`](@ref)."
const _RANDOM_REALS_KWARGS_DOC = """
A standard multivaritate normal or Cauchy is used, depending on `cauchy`, then scaled with
`scale`. `rng` is the random number generator used.
"""

_random_reals_scale(rng::AbstractRNG, scale::Real, cauchy::Bool) =
    cauchy ? scale / abs2(randn(rng)) : scale * 1.0

"""
$(SIGNATURES)

Random real number.

$(_RANDOM_REALS_KWARGS_DOC)
"""
random_real(; scale::Real = 1, cauchy::Bool = false, rng::AbstractRNG = GLOBAL_RNG) =
    randn(rng) * _random_reals_scale(rng, scale, cauchy)

"""
$(SIGNATURES)

Random vector in ``ℝⁿ``.

$(_RANDOM_REALS_KWARGS_DOC)
"""
random_reals(n::Integer; scale::Real = 1, cauchy::Bool = false,
             rng::AbstractRNG = GLOBAL_RNG) =
                 randn(rng, n) .* _random_reals_scale(rng, scale, cauchy)


# GaussianVectors.jl

# Define and simulate normal vectors.
struct GaussianVector{R<:AbstractArray{Float64, 1}, S<:AbstractArray{Float64},
		T<:AbstractArray{Float64}}
	mean::R
	Σ::S
	var::T
	length::Int
	numberofGaussians::Int

	function GaussianVector(mean::R, Σ::S) where {R<:AbstractArray{Float64, 1}, S<:AbstractArray{Float64}}
		Base.depwarn("GaussianVector is deprecated. Use MvNormal or MvNormalCanon from the Distributions package instead.", :GaussianVector, force=true)
		length = Base.length(mean)
		size(Σ,1) == length || throw(DimensionMismatch("Dimensions of Σ and μ don't match."))
		numberofGaussians = size(Σ,2)
		var = Σ * Σ'
		new{R, S, typeof(var)}(mean, Σ, var, length, numberofGaussians)
	end
end

GaussianVector(Σ::T) where T<:AbstractArray{Float64} = GaussianVector(spzeros(size(Σ,1)), Σ)

rand(d::GaussianVector) = d.mean + d.Σ * randn((d.numberofGaussians, 1))

mean(d::GaussianVector) = d.mean

var(d::GaussianVector) = d.var

cov(d::GaussianVector) = Matrix(d.var)

invcov(d::GaussianVector) = inv(cholesky(cov(d)))

length(d::GaussianVector) = d.length

scaleGV(d::T, scale::Float64) where {T<:GaussianVector} = GaussianVector(scale * d.mean, scale * d.Σ)


## rademacher
rademacher(x) = x .* sign.(randn(nobs(x)))

const _mammen_dist = Binomial(1, (sqrt(5) + 1) / (2 * sqrt(5)))
const _mammen_val1 = -(sqrt(5) - 1) / 2
const _mammen_val2 = (sqrt(5) + 1) / 2

function mammen(x)
    r = rand(_mammen_dist, nobs(x))
    return ifelse.(r .== 1, _mammen_val1, _mammen_val2)
end

## Number of exact bootstrap runs
nrun_exact(n::Integer) = binomial(2n - 1, n)

## Jack-Knife estimate
function jack_knife_estimate(data, statistic::Function, j::Int = 1, typ::Type = typeof(statistic(data)[j]))
    n = nobs(data)
    y = zeros(typ, n)
    idx = trues(n)
    for i in 1:n
        idx[i] = false
        if i > 1
            idx[i - 1] = true
        end
        y[i] = tx(statistic(pick(data, idx)))[j]
    end
    return y
end

## Sample quantiles with Gaussian interpolation
## Davison and Hinkley, equation 5.6

function iquantile(x::AbstractVector, alpha::AbstractVector)
    x = sort(x)
    qn = Float64[iquantile(x, a, true) for a in alpha]
    return qn
end

function iquantile(x::AbstractVector, alpha::AbstractFloat, is_sorted::Bool = false)
    if !is_sorted
        x = sort(x)
    end
    n = length(x)
    k = trunc(Int, (n + 1) * alpha)
    ## infinity outside of data range
    if k == 0
        return -Inf
    elseif k > (n - 1)
        return Inf
    end
    tx = [quantile(Normal(), a) for a in [alpha, k / (n + 1), (k + 1) / (n + 1)]]
    qn = (tx[1] - tx[2]) / (tx[3] - tx[2]) * (x[k + 1] - x[k]) + x[k]
    return qn
end


module TestNullModels
  using Test
  using EcologicalNetworks

  # Generate some data

  A = [true true false; true true true; false true false]
  N = UnipartiteNetwork(A)

  n1 = null1(N)
  n2 = null2(N)
  n4 = null4(N)
  n3o = null3(N; dims=1)
  n3i = null3(N; dims=2)

  @test typeof(n1) <: UnipartiteProbabilisticNetwork
  @test typeof(n2) <: UnipartiteProbabilisticNetwork
  @test typeof(n4) <: UnipartiteProbabilisticNetwork
  @test typeof(n3o) <: UnipartiteProbabilisticNetwork
  @test typeof(n3i) <: UnipartiteProbabilisticNetwork

  B = BipartiteNetwork(A)
  @test typeof(null1(B)) <: BipartiteProbabilisticNetwork
  @test typeof(null2(B)) <: BipartiteProbabilisticNetwork
  @test typeof(null3(B; dims=2)) <: BipartiteProbabilisticNetwork
  @test typeof(null3(B; dims=1)) <: BipartiteProbabilisticNetwork
  @test typeof(null4(B)) <: BipartiteProbabilisticNetwork

  # Null model 1
  @test n1[1,1] ≈ 6.0 / 9.0
  @test n1[2,2] ≈ 6.0 / 9.0
  @test n1[3,3] ≈ 6.0 / 9.0

  # Null model 2
  @test n2[1,3] ≈ 0.5
  @test n2[2,2] ≈ 1.0
  @test n2[1,1] ≈ 2.0 / 3.0

  # Null model 2 (mult)
  @test n4[1,3] ≈ 2.0 * 1.0 / 36.0
  @test n4[2,2] ≈ 3.0 * 3.0 / 36.0
  @test n4[1,1] ≈ 2.0 * 2.0 / 36.0

  # Null model 3
  @test n3o[1,1] ≈ 2.0 / 3.0
  @test n3o[2,1] ≈ 1.0
  @test n3o[3,3] ≈ 1.0 / 3.0

  @test n3i[1,1] ≈ 2.0 / 3.0
  @test n3i[2,2] ≈ 1.0
  @test n3i[3,3] ≈ 1.0 / 3.0

  # filter

  C = copy(B)
  C[1,1] = false  # remove one interaction
  # make version of the network with the first interaction removed

  # test if the zoo of B is the same as the result of the filter of C at (1,1)
  @test linearfilter(C)[1,1] ≈ linearfilterzoo(B)[1,1]
  @test linearfilter(C)[1,1] != linearfilter(B)[1,1]

end


using NonconvexCore
using NonconvexCore: VectorOfFunctions, FunctionWrapper, CountingFunction,
                     getdim, getfunction, getfunctions, Objective, IneqConstraint,
                     EqConstraint, SDConstraint
using Test

@testset "CountingFunction" begin
    f = CountingFunction(FunctionWrapper(log, 1))
    for i in 1:3
        x = rand()
        @test log(x) == f(x)
        @test f.counter[] == i
    end
    @test getdim(f) == 1
    @test getfunction(f) isa FunctionWrapper
    @test getfunction(getfunction(f)) === log
    @test length(f) == getdim(f)
end

@testset "VectorOfFunctions" begin
    f1 = FunctionWrapper(log, 1)
    f2 = FunctionWrapper(exp, 1)
    fs = [f1, f2]
    f = VectorOfFunctions([f1, f2])
    for _ in 1:3
        x = rand()
        @test [log(x), exp(x)] == f(x)
    end
    @test getdim(f) == 2
    @test getfunctions(f) == fs
    @test getfunction(f, 1) == fs[1]
    @test getfunction(f, 2) == fs[2]
    newf = VectorOfFunctions(Any[f, FunctionWrapper(exp, 1)])
    @test getdim(newf) == 3
end

@testset "Objective" begin
    f = Objective(log)
    @test getdim(f) == 1
    for _ in 1:3
        x = rand()
        @test f(x) == log(x)
    end
    @test getfunction(f) === log
end

@testset "IneqConstraint" begin
    f = IneqConstraint(log, 1.0)
    for _ in 1:3
        x = rand()
        @test f(x) == log(x) - 1
    end
    @test getdim(f) == 1
    @test getfunction(f) === log
end

@testset "EqConstraint" begin
    f = EqConstraint(log, 1.0)
    for _ in 1:3
        x = rand()
        @test f(x) == log(x) - 1
    end
    @test getdim(f) == 1
    @test getfunction(f) === log
end

@testset "SDConstraint" begin
    f = SDConstraint(exp, 1)
    for _ in 1:3
        x = rand()
        @test f(x) == exp(x)
    end
    @test getdim(f) == 1
    @test getfunction(f) === exp

    f = SDConstraint(exp, 2)
    for _ in 1:3
        x = rand(2, 2)
        @test f(x) == exp(x)
    end
    @test getdim(f) == 2
    @test getfunction(f) === exp
end


abstract Type2FuzzySet <: AbstractFuzzySet

#####################
abstract IntervalType2FuzzySet <: Type2FuzzySet
import Base: getindex

function getindex(mf::IntervalType2FuzzySet, range)
	out = hcat([[mf[value]::Tuple{Float64,Float64}...] for value in range]...)
	(out[1,:],out[2,:])
end

getindex(mf::IntervalType2FuzzySet, values...) = getindex(mf,[values...])

export getindex

#####################
# Interval 2 Membership Functions
immutable GaussType2 <: IntervalType2FuzzySet
	upper::GaussMF
	lower::GaussMF
	GaussType2(muU,muL,sigmaU,sigmaL) = new(GaussMF(muU,sigmaU),GaussMF(muL,sigmaL))
	GaussType2(upper,lower) = new(upper,lower)

end
export GaussType2

import Base: getindex
function getindex(mf::GaussType2, value::Real)
	m1 = (mf.lower.mu <= mf.upper.mu) ? mf.lower : mf.upper
	m2 = (mf.lower.mu > mf.upper.mu) ? mf.lower : mf.upper
	if m1.mu <= value <= (m1.mu+m2.mu)/2
		mu1 = 1
		mu2 = m2[value]
	elseif (m1.mu+m2.mu)/2 < value <= m2.mu
		mu1 = 1
		mu2 = m1[value]
	else
		mu1 = m1[value]
		mu2 = m2[value]
	end
	return (max(mu1,mu2),min(mu1,mu2))
end
export getindex

immutable TriangleType2 <: IntervalType2FuzzySet
	lower::TriangleMF
	upper::TriangleMF
	# TriangleType2(muU,muL,sigmaU,sigmaL) = new(GaussMF(muU,sigmaU),GaussMF(muL,sigmaL),muU,muL)
end
export TriangleType2

import Base: getindex
function getindex(mf::TriangleType2, value::Real)
	m1 = (mf.lower.b <= mf.upper.b) ? mf.lower : mf.upper
	m2 = (mf.lower.b > mf.upper.b) ? mf.lower : mf.upper
	if m1.b <= value <= (m1.b+m2.b)/2
		mu1 = 1
		mu2 = m2[value]
	elseif (m1.b+m2.b)/2 < value <= m2.b
		mu1 = 1
		mu2 = m1[value]
	else
		mu1 = m1[value]
		mu2 = m2[value]
	end
	return (max(mu1,mu2),min(mu1,mu2))
end
export getindex

#####################


#=

    This is a port of the SPPL model disease-infection.pynb
    (The arg max value of the posterior is 330 in the SPPL model, but this is based
     on a rather crude interval. Using a less crude, but much slower, interval give 333.)

    This model:
    parameters       mean       std   naive_se      mcse         ess      rhat   ess_per_sec
        Symbol    Float64   Float64    Float64   Float64     Float64   Float64       Float64

    num_people   671.4048   61.1521     0.6115    1.3344   2135.0113    1.0012      511.6250
       num_met   335.7222   27.9166     0.2792    0.5434   2402.8246    1.0005      575.8027

    Quantiles
    parameters       2.5%      25.0%      50.0%      75.0%      97.5%
        Symbol    Float64    Float64    Float64    Float64    Float64

    num_people   560.0000   628.0000   671.0000   713.0000   793.0000
       num_met   284.0000   316.0000   335.0000   354.0000   391.0250

=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV
include("jl_utils.jl")

@model function disease_infection(num_infected=100)
    num_people ~ DiscreteUniform(500,1500) # values=range(500, 1500, 20))

    # switch (num_people) cases (n in range(500, 1500, 20)):
    num_met ~ Binomial(num_people, 0.5)
    # condition ((low[n] <= num_met) < high[n])
    # switch (num_met) cases (m in range(low[n], high[n])):
    num_infected ~ Binomial(num_met, 0.3)
end

# low  = {n : int(scipy.stats.binom.ppf(.01, n=n, p=.5)) for n in range(500, 1500)}
# high = {n : int(scipy.stats.binom.ppf(.99, n=n, p=.5)) for n in range(500, 1500)}

num_infected = 100
model = disease_infection(num_infected)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, HMC(0.1,6), 1_000)
# chns = sample(model, NUTS(), 1_000)

display(chns)


using Roots
using Statistics
using FileIO
using Parameters
include("random_matrices.jl")
include("network_sims_gpu.jl")

N = 1000
g = 10.0f0
eta = 0.0f0
s = 0.0f0

J_seeds = collect(5000:5019)
h0_seed = 6000

h0 = generate_nv(N, Float32, seed=h0_seed)

wait = 200.0f0
t_bound = 2000.0f0
dt = 0.01f0

startstep = round(Int, wait / dt) + 1

# Global array of the average temporal variances of network units
vanilla_network_avg_vars = Vector{Float32}(undef, length(J_seeds))

starttime = time()
for (i, J_seed) in enumerate(J_seeds)
    J = generate_J(N, eta, seed=J_seed)
    A = sim_activity_gpu(J, g, h0; s=s, dt=dt, t_bound=t_bound,
            minimum_norm=1e-90)
    if size(A, 2) < startstep
        avg_var = 0
    else
        avg_var = mean(A[:,startstep:end].^2)
    end
    vanilla_network_avg_vars[i] = avg_var
end
println("Ensemble integration done in ", round(time() - starttime, digits=3),
        " seconds")

target_avg_var = mean(vanilla_network_avg_vars)

out = Dict{String,Any}()
@pack! out = N, g, eta, s, J_seeds, h0_seed, wait, t_bound, dt, target_avg_var
save("vanilla_target_avg_var.jld2", out)


using CompScienceMeshes, BEAST
using Test

fn = joinpath(dirname(@__FILE__),"assets","sphere35.in")
m = readmesh(fn)
t = Maxwell3D.singlelayer(wavenumber=1.0)
X = raviartthomas(m)
numfunctions(X)

##
X1 = subset(X,1:1)
numfunctions(X1)

T1 = assemble(t,X1,X)
T2 = BEAST.assemblerow(t,X1,X)
#
# T3 = assemble(t,X,X1)
# T4 = BEAST.assemblecol(t,X,X1)
#
# @test T1 == T2

# T2 = BEAST.assembleblock(t,X,X)
@test T1 == T2

I = [3,2,7]
X1 = subset(X,I)
T2 = zeros(scalartype(t,X1,X1),numfunctions(X1),numfunctions(X1))
store(v,m,n) = (T2[m,n] += v)

test_elements, test_assembly_data,
	trial_elements, trial_assembly_data,
	quadrature_data, zlocal = BEAST.assembleblock_primer(t,X,X)

BEAST.assembleblock_body!(t,
	X, I, test_elements,  test_assembly_data,
    X, I, trial_elements, trial_assembly_data,
    quadrature_data, zlocal, store)

T1 = assemble(t,X1,X1)
@test T1 == T2

# @time BEAST.assembleblock_body!(t,
# 	X, I, test_elements,  test_assembly_data,
#     X, I, trial_elements, trial_assembly_data,
#     quadrature_data, zlocal, store)
# @time assemble(t,X1,X1)

T3 = zeros(scalartype(t,X1,X1),numfunctions(X1),numfunctions(X1))
store3(v,m,n) = (T3[m,n] += v)

blkasm = BEAST.blockassembler(t,X,X)
blkasm(I,I,store3)

T4 = assemble(t,X,X)
@test T3 == T4[I,I]

# @time blkasm(I,I)
# @time assemble(t,X1,X1)


#versioninfo()

#import Pkg; Pkg.add("Arpack")
#import Pkg; Pkg.add("ArgParse")

using LinearAlgebra
using Arpack
using SparseArrays
using ArgParse

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--N"
            help = "set Nsize (should be >=4)"
            arg_type = Int
            default = 8
    end
    return parse_args(s)
end

function make_spin()
    S0 = sparse([1,2],[1,2],[1,1])
    Sx = sparse([1,2],[2,1],[1,1])
    Sy = sparse([1,2],[2,1],[-1im,1im])
    Sz = sparse([1,2],[1,2],[1,-1])
    return S0, Sx, Sy, Sz
end

function make_interaction_list(N,Nbond)
    list_site1 = [i for i in 1:N]
    list_site2 = [mod(i,N)+1 for i in 1:N]
    list_Jxx = ones(N)
    list_Jyy = ones(N)
    list_Jzz = ones(N)
    return list_site1, list_site2, list_Jxx, list_Jyy, list_Jzz
end

function make_hamiltonian(S0,Sx,Sy,Sz,N,Nbond,
    list_site1,list_site2,list_Jxx,list_Jyy,list_Jzz)
    Ham = spzeros(2^N,2^N)
    for bond in 1:Nbond
        i1 = list_site1[bond]
        i2 = list_site2[bond]
        Jxx = list_Jxx[bond]
        Jyy = list_Jyy[bond]
        Jzz = list_Jzz[bond]
        SxSx = 1
        SySy = 1
        SzSz = 1
        for site in 1:N
            if (site==i1 || site==i2)
                SxSx = kron(SxSx,Sx)
                SySy = kron(SySy,Sy)
                SzSz = kron(SzSz,Sz)
            else
                SxSx = kron(SxSx,S0)
                SySy = kron(SySy,S0)
                SzSz = kron(SzSz,S0)
            end
#        println(bond,site,SxSx)
#        println(bond,site,SySy)
#        println(bond,site,SzSz)
        end
#        Ham += Jxx * SxSx + Jyy * SySy + Jzz * SzSz
        Ham += real(Jxx * SxSx + Jyy * SySy + Jzz * SzSz)
    end
#    println(Ham)
    return Ham
end

function main()
    parsed_args = parse_commandline()
#    println("Parsed args:")
#    for (arg,val) in parsed_args
#        println("  $arg  =>  $val")
#    end

    N = parsed_args["N"]
    Nbond = N
    println("N=",N)

    S0, Sx, Sy, Sz = make_spin()
    println("S0=",Matrix(S0))
    println("Sx=",Matrix(Sx))
    println("Sy=",Matrix(Sy))
    println("Sz=",Matrix(Sz))

    list_site1, list_site2, list_Jxx, list_Jyy, list_Jzz = make_interaction_list(N,Nbond)
    println("list_site1=",list_site1)
    println("list_site2=",list_site2)
    println("list_Jxx=",list_Jxx)
    println("list_Jyy=",list_Jyy)
    println("list_Jzz=",list_Jzz)

    @time Ham = make_hamiltonian(S0,Sx,Sy,Sz,N,Nbond,list_site1,list_site2,list_Jxx,list_Jyy,list_Jzz)

    @time ene,vec = eigs(Ham, nev=5, which=:LM)
    println("# energy:",ene[1]," ",ene[2]," ",ene[3]," ",ene[4]," ",ene[5])
end

main()


## utility functions
import OffsetArrays
import ImageFiltering
import ShiftedArrays

export cov_avg!
export boxsmooth!
export outest_bounds #

"""
    outest_bounds(cx,sx) -> px0

Helper function to find maximum padding in pixels required to accomodate all query points `cx` outside of the image size 1:`sx`.

# Arguments:
- `cx`: list of integer star centers (in either x or y)
- `sx`: image dimension along the axis indexed by `cx`

# Outputs:
- `px0`: maximum padding in pixels required to accomodate all query points
"""
function outest_bounds(cx,sx)
    px0 = 0
    sortcx = sort(cx)
    if sortcx[1] < 1
        px0 = abs(sortcx[1]-1)
    end
    if sortcx[end] > sx
        if px0 < (sortcx[end]-sx)
            px0 = (sortcx[end]-sx)
        end
    end
    return px0
end

"""
    boxsmooth!(out::AbstractArray, arr::AbstractArray, tot::Array{T,1}, widx::Int, widy::Int)

Boxcar smooths an input image (or paddedview) `arr` with window size `widx` by
`widy`. We pass the original image size `sx` and `sy` to help handle image views.

# Arguments:
- `out::AbstractArray`: preallocated output array for the boxcar smoothed image
- `arr::AbstractArray`: input array for which boxcar smoothing is computed (generally paddedview)
- `tot::Array{T,1}`: preallocated array to hold moving sums along 1 dimension
- `widx::Int`: size of boxcar smoothing window in x
- `widy::Int`: size of boxcar smoothing window in y
"""
function boxsmooth!(out::AbstractArray, arr::AbstractArray, tot::Array{T,1}, widx::Int, widy::Int) where T
    (sx, sy) = size(arr)
    Δx = (widx-1)÷2
    Δy = (widy-1)÷2
    for j=1:(sy-widy+1)
        if (j==1)
            for n = 1:widy
                @simd for m = 1:sx
                    @inbounds tot[m] += arr[m,n]
                end
            end
        else
            @simd for m = 1:sx
                @inbounds tot[m] += arr[m,j+widy-1]-arr[m,j-1]
            end
        end
        tt=zero(eltype(out))
        for i=1:(sx-widx+1)
            if (i==1)
                @simd for n=1:widx
                    @inbounds tt += tot[n]
                end
            else
                @inbounds tt += tot[i+widx-1]-tot[i-1]
            end
            @inbounds out[i,j] = tt
        end
    end
end

"""
    cov_avg!(bimage, ism, bism, in_image; Np::Int=33, widx::Int=129, widy::Int=129, ftype::Int=32)

Key function for constructing the (shifted and multiplied) versions of the input image used to quickly
estimate the local covariance matrix at a large number of locations. The main output is in the preallocated
`bism` which is used as an input to `build_cov!`.

# Arguments:
- `bimage`: preallocated output array for the boxcar smoothed unshifted image
- `ism`: preallocated intermediate array for the input image times itself shifted
- `bism`: preallocated output array to store boxcar-smoothed image products for all shifts
- `in_image`: input image the local covariance of which we want to estimate

# Keywords:
- `Np::Int`: size of local covariance matrix in pixels (default 33)
- `widx::Int`: width of boxcar window in x which determines size of region used for samples for the local covariance estimate (default 129)
- `widy::Int`: width of boxcar window in y which determines size of region used for samples for the local covariance estimate (default 129)
- `ftype::Int`: determine the Float precision, 32 is Float32, otherwise Float64
"""
function cov_avg!(bimage, ism, bism, in_image; Np::Int=33, widx::Int=129, widy::Int=129, ftype::Int=32)
    if ftype == 32
        T = Float32
    else
        T = Float64
    end

    Δx = (widx-1)÷2
    Δy = (widy-1)÷2
    halfNp = (Np-1) ÷ 2

    (sx1, sy1) = size(in_image)
    tot = zeros(T,sx1);
    boxsmooth!(bimage,in_image,tot,widx,widy)
    # loop over shifts
    for dc=0:Np-1       # column shift loop
        for dr=1-Np:Np-1   # row loop, incl negatives
            if (dr < 0) & (dc == 0)
                continue
            end
            # ism = image, shifted and multipled
            @inbounds ism .= in_image .* ShiftedArrays.circshift(in_image,(-dr, -dc))
            fill!(tot,0)
            boxsmooth!(view(bism,:,:,dr+Np,dc+1),ism,tot,widx,widy) # bism = boxcar(ism)
        end
    end
    return
end


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module Counts
using Nemo, NumberTheory

export ModuleCounts
export L000961, L002808, L005117, L013928, L025528, L065515
export L065855, L069637, L246547, L246655, L000720
export A007917, A151800, A257993
export PreviousPrime, NextPrime, PrimePiList
export takeFirst, Nth, Count, List, HilbertHotel

"""
* PreviousPrime, NextPrime, PrimePiList, takeFirst, Nth, Count, List, HilbertHotel
* L000961, L002808, L005117, L013928, L025528, L065515, L065855, L069637, L246547, L246655, L000720, A007917, A151800, A257993
"""
const ModuleCounts = ""

"""
Return a list of length len of integers ``≥ 0`` which are isA.
"""
function List(len, isA::Function)
    len ≤ 0 && return fmpz[]
    j, c = Int(0), Int(1)
    A = fill(ZZ(0), len)
    while c <= len
        if isA(j)
            A[c] = fmpz(j)
            c += 1
        end
        j += 1
    end
    A
end

"""
Iverson brackets.
"""
ι(b) = b ? 1 : 0

"""
Inverse Iverson brackets.
"""
ιι(n) = n == 0 ? true : false
 # The Unix way: A successful command returns a 0, while an unsuccessful one
 # returns a non-zero value that usually can be interpreted as an error code.

"""
Return a iterator of length n which has value 1 if isA(i) is true and otherwise 0.
"""
function Indicators(n, isA::Function)
    (ι(isA(i)) for i in 0:n - 1)
end

"""
Return a list of length len which gives the numbers of integers ≤ n which are isA. Integers start at ``n=0``.
```
julia> CountList(8, isPrime)
[0, 0, 1, 2, 2, 3, 3, 4]
```
"""
CountList(len::Int, isA::Function) = Accumulate(Indicators(len, isA))

# Consider two sequences A and invA. We say invA is the left inverse of A iff
# invA(A(n)) = n and A(n) is the least number m such that A(invA(m)) = A(n).
# A(invA(n)) = n if and only if isA(n) = true where 'isA' is the indicator
# function of the sequence A.
# We introduce 'Count' as the left inverse of 'Nth'. For example
# if Nth(n) = NthPrime(n) then Count(n) = PrimePi(n) (A000720).

"""
Return the numbers of integers in the range 0:n which are isA.
```
julia> Count(8, isPrime)
4
```
"""
Count(n::Int, isAb::Function) = Base.count((isAb(i) for i in 0:n))

"""
Return the numbers of integers in the range a:b which are isA.
```
julia> Count(3:8, isPrime)
3
```
"""
Count(r, isAb::Function) = Base.count((isAb(i) for i in r))

"""
Return a SeqArray listing the values satisfying the predicate isA for arguments ``0 ≤ x ≤ `` bound.
```
julia> FindUpTo(7, isPrime)
[2, 3, 5, 7]
```
"""
function FindUpTo(bound, isA::Function)
    bound < 0 && return fmpz[]
    filter(isA, 0:bound)
end

"""
Return the first ``n`` numbers satisfying the predicate isA.
"""
takeFirst(isA, n) = Iterators.take(Iterators.filter(isA, Iterators.countfrom(1)), n)

"""
Return a iterator listing the values satisfying the predicate isA for arguments in ``0 ≤ n ≤ bound .``
"""
function IterateUpTo(bound, isA::Function)
    (i for i in 0:bound if isA(i))
end

"""
Returns an integer which is the highest index in `b` for the value `a`. Whenever `a` is not a member of `b` it returns -1.
```
julia> L = List(10, isPrime); IndexIn(13, L)
5
```
"""
function IndexIn(a, b::AbstractArray)
    bdict = Dict(zip(b, 0:length(b)))
    get(bdict, fmpz(a), -1)
end

"""
Return the Nth integer which is isA. (For N ≤ 0 return 0.)
```
julia> Nth(7, isPrime)
17
```
"""
function Nth(N, isA::Function)
    N ≤ 0 && return 0
    n, c = Int(0), Int(0)
    while c < N
        i = isA(n)
        i && (c += 1)
        n += 1
    end
    n - 1
end

"""
Return the cumulative sum of an SeqArray.
"""
function Accumulate(A)
    R = fill(ZZ(0), length(A))
    i, acu = 1, 0
    for a in A
        acu += a
        R[i] = acu
        i += 1
    end
    R
end

"""
Return the smallest list of indicators of isA with ∑(A) = count.
```
julia> IndicatorsFind(7, isPrime)
[0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1]
```
"""
function IndicatorsFind(count, isA::Function)
    count ≤ 0 && return []
    n, c = Int(0), Int(0)
    A = Int[]
    while c < count
        i = isA(n)
        i && (c += 1)
        push!(A, ι(i))
        n += 1
    end
    A
end

"""
Return the first integer ``n ≥ 0`` such that isA(n) = true.
```
julia> First(isPrime)
2
```
"""
function First(isA::Function)
    n = 0
    while !isA(n)
        n += 1
    end
    n
end

First(A::Array{Int}) = A == [] ? nothing : first(A)

"""
Return the element at the end of the list A if A is not empty, 0 otherwise.
"""
Last(A) = A == [] ? "undef" : A[end]

"""
Trick described by David Hilbert in a 1924 lecture "Über das Unendliche".
"""
HilbertHotel(guest, hotel) = prepend!(hotel, guest)

"""
Return largest ``0 < k < n`` such that isA(k) = true or nothing if no such ``k`` exists.
```
julia> Previous(7, isPrime)
5
```
"""
function Previous(n, isA::Function)
    n == nothing && return First(isA)
    while true
        n -= 1
        isA(n) && break
        # This definition avoids throwing an exception. Alternatively:
        # n < 0 && throw(ArgumentError("Not defined for $n."))
        n < 0 && return First(isA)
    end
    n
end

"""
Return least ``k > n ≥ 0`` such that isA(k) = true. NOTE: It is assumed that such a ``k`` exists! (If not, the function will run forever.)
```
julia> Next(7, isPrime)
11
```
"""
function Next(n, isA::Function)
    ((n ≤ 0) || (n == nothing)) && return First(isA)
    while true
        n += 1
        isA(n) && break
    end
    n
end

"""
Return a list of composite numbers of length len. (Numbers which have more than one prime divisor.)
```
julia> L002808(8)
[4, 6, 8, 9, 10, 12, 14, 15]
```
"""
L002808(len) = List(len, isComposite)

"""
Return a list of the number of composite numbers ``≤ n``.
```
julia> L065855(8)
[0, 0, 0, 0, 1, 1, 2, 2]
```
"""
L065855(len) = CountList(len, isComposite)

"""
Return a list of squarefree numbers of length len. (Numbers which are not divisible by a square greater than 1.)
```
julia> L005117(8)
[1, 2, 3, 5, 6, 7, 10, 11]
```
"""
L005117(len) = List(len, isSquareFree)

"""
Return a list of the number of squarefree numbers ``< n``.
```
julia> L013928(8)
[0, 1, 2, 3, 3, 4, 5, 6]
```
"""
L013928(len) = CountList(len, isSquareFree)

"""
Return a list of powers of primes of length len. (Numbers of the form ``p^k`` where ``p`` is a prime and ``k ≥ 0``.)
```
julia> L000961(8)
[1, 2, 3, 4, 5, 7, 8, 9]
```
"""
L000961(len) = List(len, isPowerOfPrimes)

"""
Return the number of powers of primes ``≤ n``. (Powers of primes are numbers of the form ``p^k`` where ``p`` is a prime and ``k ≥ 0``.)
```
julia> L065515(8)
[0, 1, 2, 3, 4, 5, 5, 6]
```
"""
L065515(len) = CountList(len, isPowerOfPrimes)

"""
Return a list of prime powers of length len. (Numbers of the form ``p^k`` where ``p`` is a prime and ``k ≥ 1``.)
```
julia> L246655(8)
[2, 3, 4, 5, 7, 8, 9, 11]
```
"""
L246655(len) = List(len, isPrimePower)

"""
Return a list of the number of prime powers ``≤ n`` with exponents ``k ≥ 1``.
```
julia> L025528(8)
[0, 0, 1, 2, 3, 4, 4, 5]
```
"""
L025528(len) = CountList(len, isPrimePower)

"""
Return a list of perfect powers of length len. (Numbers of the form ``p^k`` where ``p`` is a prime and ``k ≥ 2``.
```
julia> L246547(8)
[4, 8, 9, 16, 25, 27, 32, 49]
```
"""
L246547(len) = List(len, isPerfectPower)

"""
Return a list of the number of prime powers ``≤ n`` with exponents ``k ≥ 2``.
```
julia> L069637(8)
[0, 0, 0, 0, 1, 1, 1, 1]
```
"""
L069637(len) = CountList(len, isPerfectPower)

# cf. also:
# A067535 Smallest squarefree number >= n.
# A070321 Largest squarefree number <= n.
# A025528 Number of prime powers <= n with exponents > 0.
# A000015 Smallest prime power >= n.
# A031218 Largest prime power <= n
# A167184 Smallest prime power >= n that is not prime.
# A081676 Largest perfect power <= n

"""
Return the largest prime in ``N`` (the semiring of natural numbers including zero) less than n for ``n ≥ 0``.
 (The `prev_prime` function of Mathematica, Maple, Magma and SageMath.)
"""
A007917(n::Int) = Previous(n, isPrime)

"""
Return the largest prime in ``Z`` (the ring of all integers) less than ``n`` for ``n ≥ 0`` (cf. A007917).
"""
PreviousPrime(n::Int) = n ∈ [0, 1, 2] ? -2 : Previous(n - 1, isPrime)

"""
Return least prime ``> n``. The next_prime function of Mathematica, Maple, Magma and SageMath (cf. A151800).
"""
NextPrime(n::Int) = Next(n, isPrime)

"""
Return least prime ``> n``. The `next_prime` function of Mathematica, Maple, Magma and SageMath.
"""
A151800(n::Int) = Next(n, isPrime)

"""
Return the list of number of primes ``≤ n`` for ``n ≥ 0``.

```
julia> PrimePiList(8)

[0, 0, 1, 2, 2, 3, 3, 4]
```
"""
PrimePiList(len::Int) = CountList(len, isPrime)

"""
Return the list of number of primes ``≤ n`` for ``n ≥ 0``.

```
julia> L000720(8)

[0, 0, 1, 2, 2, 3, 3, 4]
```
"""
L000720(len::Int) = PrimePiList(len)

"""
Return the index of the least prime not dividing ``n``.
"""
function A257993(n::Int)
    c, p = 1, 2
    while n % p == 0
        p = NextPrime(p)
        c += 1
    end
    c
end

#START-TEST-########################################################

using Test, SeqTests, SeqUtils

function test()

    indicators = [isPositive, isEven, isSquare, isPrime]
    indicatorNames = ["isPositive", "isEven", "isSquare", "isPrime"]

    len = 14
    @testset "Counts" begin

        @test Nth(96, isPrime) == 503
        @test Nth(97, isPrime) == 509
        @test Nth(98, isPrime) == 521

        # In other words: 97 is the 25-th prime.
        @test Count(96, isPrime) == 24
        @test Count(97, isPrime) == 25
        @test Count(98, isPrime) == 25

        @test List(24, isPrime)[end] == 89
        @test List(25, isPrime)[end] == 97

        for (i, isA) in enumerate(indicators)
            # This test shows that the logic behind 'Nth' and 'Count' is OK.
            for n in 1:len
               @test isA(n) == (Nth(Count(n, isA), isA) == n)
               @test     n  ==  Count(Nth(n, isA), isA)
            end
        end

        a = [A257993(n) for n in 1:10]
        b = [1, 2, 1, 2, 1, 3, 1, 2, 1, 2]
        @test all(a .== b)

        if is_oeis_installed()
            L = [
            L000961, L002808, L005117, L013928, L246547, L246655
            # L025528, L065515, L065855, L069637, L000720
            ]
            SeqTest(L, 'L')
        end
    end
end

function demo()

    indicators = [isNonnegative, isPositive, isEven, isComposite,
    isSquare, isSquareFree, isPrimePower, isPowerOfPrimes,
    isPerfectPower, isPrime
    ]
    indicatorNames = ["isNonnegative", "isPositive", "isEven", "isComposite",
    "isSquare", "isSquareFree", "isPrimePower", "isPowerOfPrimes",
    "isPerfectPower", "isPrime"
    ]

    println()
    len = 14
    for (i, isA) in enumerate(indicators)

        println("Predicate      ", indicatorNames[i])
        println("First          ", First(isA))
        println("---")

        println("Nth            ", [Nth(n, isA) for n in 0:len])
        print("List           ") ; Println(List(len, isA))
        println("FindUpTo       ", FindUpTo(len, isA))
        println("IterateUpTo    ", [k for k in IterateUpTo(len, isA)])
        println("---")

        println("IndicatorsFind ", IndicatorsFind(len, isA))
        println("Indicators     ", [k for k in Indicators(len, isA)])
        println("---")

        println("IndexIn (list) ", [IndexIn(fmpz(n), List(len, isA)) for n in 0:len])

        println("Count   (list) ", [Count(n, isA) for n in 0:len - 1])
        print("CountList      ") ; Println(CountList(len, isA))

        println("Previous(n)    ", [Previous(n, isA) for n in 0:len])
        println("Next(n)        ", [Next(n, isA)  for n in 0:len])
        println("Nth(Count(n))  ", [Nth(Count(n, isA), isA) for n in 0:len])
        println("Count(Nth(n))  ", [Count(Nth(n, isA), isA) for n in 0:len])

        println('='^66)
    end
end

"""
[A257993(n) for n in 1:10000]
    0.000635 seconds (9.20 k allocations: 221.859 KiB)
PrimePiList(10000)
    0.002290 seconds (35.82 k allocations: 637.891 KiB)
"""
function perf()
    @time [A257993(n) for n in 1:10000]
    GC.gc()
    @time PrimePiList(10000)
end

function main()
    test()
    demo()
    perf()
end

main()

end # module


using Manifolds, Manopt, Test, ManifoldsBase

@testset "Differentials (on Sphere(2))" begin
# The Adjoint Differentials test using the same variables as the differentials
# test
x = [1.,0.,0.]
y = [0.,1.,0.]
M = Sphere(2)
ξ = log(M,x,y)
# Text differentials (1) Dx of Log_xy
@test adjoint_differential_log_basepoint(M,x,x,ξ) == -ξ
@test adjoint_differential_log_argument(M,x,x,ξ) == ξ
@test adjoint_differential_exp_basepoint(M,x,zero_tangent_vector(M,x),ξ) == ξ
@test adjoint_differential_exp_argument(M,x,zero_tangent_vector(M,x),ξ) == ξ
for t in [0,0.15,0.33,0.66,0.9]
    @test adjoint_differential_geodesic_startpoint(M,x,x,t,ξ) == (1-t)*ξ
    @test norm(M,x,adjoint_differential_geodesic_endpoint(M,x,x,t,ξ) - t*ξ)  ≈ 0 atol=10.0^(-16)
end
Mp = PowerManifold(M, NestedPowerRepresentation(), 3)
xP = [x, y, x]
yP = [x, x, y]
ξP = [ξ, zero_tangent_vector(M,x), -ξ]
@test norm(
        Mp,
        xP,
        adjoint_differential_forward_logs(Mp,xP,ξP) - [-ξ, zero_tangent_vector(M,x), zero_tangent_vector(M,x)]
    ) ≈ 0 atol=4*10.0^(-16)
end


using Decimals
using Compat.Test

@testset "Equality" begin

@testset "isequal" begin
    @test isequal(Decimal(0, 2, -3), Decimal(0, 2, -3))
    @test ~isequal(Decimal(0, 2, -3), Decimal(0, 2, 3))
    @test isequal(Decimal(0, 2, -3), 0.002)
    @test isequal(Decimal(1, 2, 0), -2)
    @test ~isequal(Decimal(1, 2, 0), 2)
end

@testset "==" begin
    @test Decimal(0, 2, -3) == Decimal(0, 2, -3)
    @test Decimal(0, 2, -3) != Decimal(0, 2, 3)
    @test Decimal(0, 2, -3) == 0.002

    @test -2 == Decimal(1, 2, 0)
    @test 2 != Decimal(1, 2, 0)

    @test Decimal(1, 2, 0) == -2
    @test Decimal(1, 2, 0) != 2

    bf_pi = BigFloat(pi)
    @test Decimal(bf_pi) == bf_pi
    @test bf_pi == Decimal(bf_pi)

    bi = big"4608230166434464229556241992703"
    @test Decimal(bi) == bi
    @test bi == Decimal(bi)
end

@testset "<" begin
    @test Decimal(1, 1, 1) < Decimal(0, 1, 1)
    @test !(Decimal(0, 1, 1) < Decimal(1, 1, 1))
    @test Decimal(1, 1, 1) < Decimal(1, 0, 1)
    @test !(Decimal(1, 0, 1) < Decimal(1, 1, 1))
    @test Decimal(0, 2, -3) < Decimal(0, 2, 3)
    @test !(Decimal(0, 2, 3) < Decimal(0, 2, -3))
end

end


scipy_optimize = pyimport("scipy.optimize")

function prediction_error(
    mountaincar::MountainCar,
    params::Array{Float64},
    data::Array{MountainCarContTransition},
)
    err = []
    for i = 1:length(data)
        transition = data[i]
        predicted_state =
            step(mountaincar, transition.initial_state, transition.action, params)[1]
        # push!(err, norm(vec(predicted_state) - vec(transition.final_state)))
        push!(err, (predicted_state.position -
                    transition.final_state.position)^2 + (predicted_state.speed -
                                                          transition.final_state.speed)^2)

    end
    err
end

function get_least_squares_fit(
    mountaincar::MountainCar,
    params::MountainCarParameters,
    data::Array{MountainCarContTransition},
)
    fn(p) = prediction_error(mountaincar, p, data)
    scipy_optimize.leastsq(fn, vec(params))[1]
end

function get_least_squares_fit(
    mountaincar::MountainCar,
    params::Array{Float64},
    data::Array{MountainCarContTransition},
)
    get_least_squares_fit(
        mountaincar,
        unvec_params(params),
        data,
    )
end


"""
`AbstractPoint{N,T} <: AbstractVector{T}`

Supertype for fixedsize `N`-dimensional points with elements of type `T`.
"""
abstract type AbstractPoint{N,T} <: AbstractVector{T} end

Base.haslength(::AbstractPoint) = true
Base.length(::AbstractPoint{N}) where N = N
Base.size(::AbstractPoint{N}) where N = (N,)
Base.eltype(::AbstractPoint{N,T}) where {N,T} = T
Base.IndexStyle(::Type{<:AbstractPoint}) = Base.IndexLinear()

"""
`Point{N,T} <: AbstractPoint{N,T}`

Fixedsize `N`-dimensional point with elements of type `T`.
"""
struct Point{N,T} <: AbstractPoint{N,T}
  data :: NTuple{N,T}
end

function Point(x, y)
  values = promote(x, y)
  Point{2,eltype(values)}(values)
end

function Point(x, y, z)
  values = promote(x, y, z)
  Point{3,eltype(values)}(values)
end

using Base: unsafe_load, unsafe_convert

Point{N,T}(vec::Vector{T}) where {N,T} =
  unsafe_load(unsafe_convert(Ptr{Point{N,T}}, vec))

Base.getindex(point::Point, i) = point.data[i]

struct MetaPoint{N,T,A,M} <: AbstractPoint{N,T}
  data :: NTuple{N,T}
  meta :: NamedTuple{A,M}
end

Base.getindex(point::MetaPoint, i) = point.data[i]
Base.getindex(point::MetaPoint, sym::Symbol) = point.meta[sym]
meta(point::MetaPoint) = point.meta


using BlochSim, Test, MAT

function testB2c()

    answer = matread("matlabtestdata/testB2c.mat")

    spin = Spin(1, 600, 100, 0)
    sig = mese!(spin, 1000, 50, 1)[1]

    return sig ≈ answer["sig"]

end

function testB2d()

    answer = matread("matlabtestdata/testB2d.mat")

    spin = Spin(1, 600, 100, 0)
    sig = mese!(spin, 1000, 50, 8)

    return sig ≈ vec(answer["sig"])

end

function testB3a()

    answer = matread("matlabtestdata/testB3a.mat")

    gradz = 0.3 # G/cm
    Tg = 3 # ms
    z = π/2 / (GAMMA * gradz * Tg/1000) # cm
    spin = Spin(1, 600, 100, 0, [0,0,z])
    spgr!(spin, 10, 2, π/3, [0,0,gradz], Tg; Δθinc = 0)

    return spin.M ≈ vec(answer["M"])

end

function testB3b()

    answer = matread("matlabtestdata/testB3b.mat")

    gradz = 0.3 # G/cm
    Tg = 3 # ms
    zmax = 4π / (GAMMA * gradz * Tg/1000) # cm
    z = ((1:100)/100 .- 0.5) * zmax
    spins = map(z -> Spin(1, 600, 100, 0, [0,0,z]), z)
    map(spin -> spgr!(spin, 10, 2, π/3, [0,0,gradz], Tg; Δθinc = 0), spins)
    M = zeros(ComplexF64, 3)
    for spin in spins
        M .+= spin.M
    end
    M ./= 100

    return M ≈ vec(answer["M"])

end

function testB3c()

    answer = matread("matlabtestdata/testB3c.mat")

    spin = Spin(1, 600, 100, 0)
    spgr!(spin, 10, 2, π/3)

    return spin.M ≈ vec(answer["M"])

end

function testB5a()

    answer = matread("matlabtestdata/testB5a.mat")

    gradz = 0.3 # G/cm
    Tg = 3 # ms
    zmax = 2π / (GAMMA * gradz * Tg/1000) # cm
    z = (1:100)/100 * zmax
    spins = map(z -> Spin(1, 600, 100, 0, [0,0,z]), z)
    s = map(spin -> spgr!(spin, 10, 2, π/6, [0,0,gradz], Tg, nTR = 99), spins)
    sig = sum(s) / 100

    return sig ≈ answer["sig"]

end

function testB5b()

    answer = matread("matlabtestdata/testB5b.mat")

    α = LinRange(0, π/2, 51) # rad
    gradz = 0.3 # G/cm
    Tg = 3 # ms
    zmax = 2π / (GAMMA * gradz * Tg/1000) # cm
    z = (1:100)/100 * zmax
    srf = zeros(ComplexF64, 51)
    for i = 1:51
        spins = map(z -> Spin(1, 600, 100, 0, [0,0,z]), z)
        s = map(spin -> spgr!(spin, 10, 2, α[i], [0,0,gradz], Tg, nTR = 99), spins)
        srf[i] = sum(s) / 100
    end
    spin = Spin(1, 600, 100, 0)
    sideal = map(α -> spgr!(spin, 10, 2, α), α)

    return srf ≈ vec(answer["sig1"]) && sideal ≈ vec(answer["sig2"])

end

@testset "Sequences" begin

    @testset "MESE" begin

        @test testB2c()
        @test testB2d()

    end

    @testset "SPGR" begin

        @test testB3a()
        @test testB3b()
        @test testB3c()
        @test testB5a()
        @test testB5b()

    end

end


using DelimitedFiles
using ForwardDiff

mutable struct SplinePath
    x_coefs ::Array{Float64, 2}
    y_coefs ::Array{Float64, 2}
    breaks  ::Array{Float64, 1}
end

"""
Load spline data consisting of coefficients for the polynomials
x_i(s) = a_i(s-s_i)³ + a_i(s-s_i)² + c_i(s-s_i) + d_i and
y_i(s) = e_i(s-s_i)³ + f_i(s-s_i)² + g_i(s-s_i) + h_i and breaks.
Polynomial coefficients returned have dimensions [N, 4] and breaks is a vector of length N+1.
"""
function load_splines()
    x_coefs = readdlm("../splines/nuscenes_scene_23_x_spline_coefs.csv", ',');
    y_coefs = readdlm("../splines/nuscenes_scene_23_y_spline_coefs.csv", ',');
    breaks = readdlm("../splines/nuscenes_scene_23_breaks.csv", ',')[:];

    x_coefs, y_coefs, breaks
end

"""
Performs a bisection search over interpolating spline intervals (input as the array 'breaks')
to find interval within which the input parameter 's' resides.
"""
function find_spline_interval(s::Float64, path::SplinePath)
    breaks = path.breaks

    if s < breaks[1] || s > breaks[end]
        if s < 0 && abs(s) < 1E-2
            println("The value s = $s is slightly below 0.")
            return 1
        end
        println("The value s = $s does not lie within the valid spline domain.")
        println("Path length has been overrun.")
        return
    end

    upper_idx_lim = length(breaks)
    lower_idx_lim = 1
    interval_idx = Int(div(upper_idx_lim - lower_idx_lim, 2))

    while true
        if breaks[interval_idx] <= s <= breaks[interval_idx+1]
            return interval_idx
        elseif s > breaks[interval_idx+1]
            lower_idx_lim = max(1, interval_idx + 1)
            interval_idx = Int(div(upper_idx_lim + lower_idx_lim, 2))
        elseif s < breaks[interval_idx]
            upper_idx_lim = min(length(breaks), interval_idx)
            interval_idx = Int(div(upper_idx_lim + lower_idx_lim, 2))
        else
            println("Unexpected case with s = ", s)
            throw(UndefVarError)
        end
    end
end

"""
Evaluates the x position of an interpolated spline path at arc length parameter s.
"""
function spline_x(s, path::SplinePath, spline_idx::Integer)
    x_coefs = path.x_coefs
    breaks = path.breaks

    x = x_coefs[spline_idx, 1] * (s - breaks[spline_idx]).^3 + x_coefs[spline_idx, 2] * (s - breaks[spline_idx]).^2 +
        x_coefs[spline_idx, 3] * (s - breaks[spline_idx]) + x_coefs[spline_idx, 4]
end

"""
Evaluates the x position of an interpolated spline path at arc length parameter s.
"""
function spline_y(s, path::SplinePath, spline_idx::Integer)
    y_coefs = path.y_coefs
    breaks = path.breaks

    y = y_coefs[spline_idx, 1] * (s - breaks[spline_idx]).^3 + y_coefs[spline_idx, 2] * (s - breaks[spline_idx]).^2 +
        y_coefs[spline_idx, 3] * (s - breaks[spline_idx]) + y_coefs[spline_idx, 4]
end

"""
Reprojects the vehicle position onto the path to find the appropriate arc length
parameter s. ds is resolution for arc length search.
TODO: Fix situation where multiple possible such s may be found.
TODO: Currently this only performs a local search. Add global search for when
distance between vehicle position and closes projection is greater than some
threshold.
"""
function find_best_s(q::Array{Float64,1}, path::SplinePath; ds=0.05, enable_global_search=false, sq_dist_tol=10^2)
    x = q[1]
    y = q[2]
    s = q[5]

    if s < 0
        return 0
    end

    spline_idx = find_spline_interval(s, path)
    x_proj = spline_x(s, path, spline_idx)
    y_proj = spline_y(s, path, spline_idx)
    current_dist = (x - x_proj)^2 + (y - y_proj)^2

    ss = [max(0,s-5):ds:s; s+ds:ds:s+5;] # TODO: this indexing looks a bit weird
    spline_idces = find_spline_interval.(ss, (path,))
    xs = spline_x.(ss, (path,), spline_idces)
    ys = spline_y.(ss, (path,), spline_idces)
    sq_distances = (xs .- x).^2 + (ys .- y).^2
    best_dist, best_idx = findmin(sq_distances)

    if best_dist > sq_dist_tol && enable_global_search
        println("Local search too narrow. Performing global search.")
        ss = [path_obj.breaks[1]:ds:path_obj.breaks[end];] # TODO: check this indexing
        spline_idces = find_spline_interval.(ss, (path,))
        xs = spline_x.(ss, (path,), spline_idces)
        ys = spline_y.(ss, (path,), spline_idces)
        sq_distances = (xs .- x).^2 + (ys .- y).^2
        best_dist, best_idx = findmin(sq_distances)
        if best_dist <=  sq_dist_tol
            println("Global search successfully found a better projection.")
        end
    end

    if best_dist >= current_dist
        s
    else
        ss[best_idx]
    end
end

"""
Evaluates the derivatives dx/ds and dy/ds for path at arc length parameter s0.
"""
function dpath_ds(s0, path::SplinePath, spline_idx::Integer)
    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    ForwardDiff.derivative(s -> spline_x(s, path, spline_idx), s0),
        ForwardDiff.derivative(s -> spline_y(s, path, spline_idx), s0)
end

"""
Evaluates the angle of the tangent to the path at the reference point s0 with respect to the x-axis.
"""
function heading(s, path::SplinePath, spline_idx::Integer)
    x_coefs = path.x_coefs
    y_coefs = path.y_coefs
    breaks = path.breaks

    dx_ds, dy_ds = dpath_ds(s, path, spline_idx)
    atan(dy_ds, dx_ds)
end


# From base, but copied here to make sure we don't fail bacause base changed
function my_gcd(a::T, b::T) where T<:Union{Int8,UInt8,Int16,UInt16,Int32,UInt32,
Int64,UInt64,Int128,UInt128}
    a == 0 && return abs(b)
    b == 0 && return abs(a)
    za = trailing_zeros(a)
    zb = trailing_zeros(b)
    k = min(za, zb)
    u = unsigned(abs(a >> za))
    v = unsigned(abs(b >> zb))
    while u != v
        if u > v
            u, v = v, u
        end
        v -= u
        v >>= trailing_zeros(v)
    end
    r = u << k
    # T(r) would throw InexactError; we want OverflowError instead
    r > typemax(T) && throw(OverflowError("gcd($a, $b) overflows"))
    r % T
end


using ITensors,
      Test,
      Random

Random.seed!(1234)

@testset "BlockSparse ITensor" begin

  @testset "Constructor" begin
    i = Index([QN(0)=>1,QN(1)=>2],"i")
    j = Index([QN(0)=>3,QN(1)=>4,QN(2)=>5],"j")

    A = ITensor(QN(0),i,dag(j))

    @test flux(A) == QN(0)
    @test nnzblocks(A) == 2
  end

  @testset "Constructor (from Tuple)" begin
    i = Index([QN(0)=>1,QN(1)=>2],"i")
    j = Index([QN(0)=>3,QN(1)=>4,QN(2)=>5],"j")

    A = ITensor(QN(0), (i, dag(j)))

    @test flux(A) == QN(0)
    @test nnzblocks(A) == 2
  end

  @testset "Constructor (no flux specified)" begin
    i = Index([QN(0)=>1, QN(1)=>2], "i")
    j = Index([QN(0)=>3, QN(1)=>4, QN(2)=>5], "j")

    A = ITensor(i, dag(j))

    @test flux(A) == QN(0)
    @test nnzblocks(A) == 2
  end

  @testset "Constructor (Tuple, no flux specified)" begin
    i = Index([QN(0)=>1, QN(1)=>2], "i")
    j = Index([QN(0)=>3, QN(1)=>4, QN(2)=>5], "j")

    A = ITensor((i, dag(j)))

    @test flux(A) == QN(0)
    @test nnzblocks(A) == 2
  end

  @testset "No indices getindex" begin
    T = ITensor(QN())
    @test order(T) == 0
    @test flux(T) == nothing
    @test nnzblocks(T) == 1
    @test T[] == 0

    s = Index(QN(-1)=>1,QN(1)=>1)
    A = emptyITensor(s, dag(s'))
    B = emptyITensor(s', dag(s))
    A[1, 1] = 1
    B[2, 2] = 1
    C = A * B
    @test order(C) == 0
    @test flux(C) == nothing
    @test nnzblocks(C) == 0
    @test C[] == 0
  end

  @testset "Empty constructor" begin
    i = Index([QN(0)=>1,QN(1)=>2],"i")

    A = emptyITensor(i,dag(i'))

    @test nnzblocks(A) == 0
    @test nnz(A) == 0
    @test hasinds(A,i,i')
    @test isnothing(flux(A))

    A[i(1),i'(1)] = 1.0

    @test nnzblocks(A) == 1
    @test nnz(A) == 1
    @test flux(A) == QN(0)

    A[i(2),i'(2)] = 1.0

    @test nnzblocks(A) == 2
    @test nnz(A) == 5
    @test flux(A) == QN(0)
  end

  @testset "Check flux when setting elements" begin
    i = Index(QN(0)=>1,QN(1)=>1; tags="i")
    A = randomITensor(QN(0),i,dag(i'))
    @test_throws ErrorException A[i(1),i'(2)] = 1.0
  end

  @testset "Random constructor" begin
    i = Index([QN(0)=>1,QN(1)=>2],"i")
    j = Index([QN(0)=>3,QN(1)=>4,QN(2)=>5],"j")

    A = randomITensor(QN(1),i,dag(j))

    @test flux(A) == QN(1)
    @test nnzblocks(A) == 1

    B = randomITensor(i,dag(j))

    @test flux(B) == QN()
    @test nnzblocks(B) == 2
  end

  @testset "QN setelt" begin
    i = Index(QN(0)=>2,QN(1)=>2,tags="i")

    T = setelt(i(1))
    @test T[i(1)] ≈ 1.0
    @test T[i(2)] ≈ 0.0
    @test T[i(3)] ≈ 0.0
    @test T[i(4)] ≈ 0.0

    T = setelt(i(2))
    @test T[i(1)] ≈ 0.0
    @test T[i(2)] ≈ 1.0
    @test T[i(3)] ≈ 0.0
    @test T[i(4)] ≈ 0.0

    # Test setelt taking Pair{Index,Int}
    T = setelt(i=>3)
    @test T[i(1)] ≈ 0.0
    @test T[i(2)] ≈ 0.0
    @test T[i(3)] ≈ 1.0
    @test T[i(4)] ≈ 0.0
  end

  @testset "setindex!" begin

    @testset "Test 1" begin
      s1 = Index([QN("N",0,-1)=>1,QN("N",1,-1)=>1],"s1")
      s2 = Index([QN("N",0,-1)=>1,QN("N",1,-1)=>1],"s2")
      A = emptyITensor(s1,s2)

      @test nnzblocks(A) == 0
      @test nnz(A) == 0
      @test hasinds(A,s1,s2)
      @test isnothing(flux(A))

      A[2,1] = 1.0/sqrt(2)

      @test nnzblocks(A) == 1
      @test nnz(A) == 1
      @test A[s1(2),s2(1)] ≈ 1.0/sqrt(2)
      @test flux(A) == QN("N",1,-1)

      A[1,2] = 1.0/sqrt(2)

      @test nnzblocks(A) == 2
      @test nnz(A) == 2
      @test A[s1(2),s2(1)] ≈ 1.0/sqrt(2)
      @test A[s1(1),s2(2)] ≈ 1.0/sqrt(2)
      @test flux(A) == QN("N",1,-1)
    end

    @testset "Test 2" begin
      s1 = Index([QN("N",0,-1)=>1,QN("N",1,-1)=>1],"s1")
      s2 = Index([QN("N",0,-1)=>1,QN("N",1,-1)=>1],"s2")
      A = emptyITensor(s1,s2)

      @test nnzblocks(A) == 0
      @test nnz(A) == 0
      @test hasinds(A,s1,s2)
      @test isnothing(flux(A))

      A[1,2] = 1.0/sqrt(2)

      @test nnzblocks(A) == 1
      @test nnz(A) == 1
      @test A[s1(1),s2(2)] ≈ 1.0/sqrt(2)
      @test flux(A) == QN("N",1,-1)

      A[2,1] = 1.0/sqrt(2)

      @test nnzblocks(A) == 2
      @test nnz(A) == 2
      @test A[s1(2),s2(1)] ≈ 1.0/sqrt(2)
      @test A[s1(1),s2(2)] ≈ 1.0/sqrt(2)
      @test flux(A) == QN("N",1,-1)
    end

  end

  @testset "Multiply by scalar" begin
    i = Index([QN(0)=>1,QN(1)=>2],"i")
    j = Index([QN(0)=>3,QN(1)=>4],"j")

    A = randomITensor(QN(0),i,dag(j))

    @test flux(A) == QN(0)
    @test nnzblocks(A) == 2

    B = 2*A

    @test flux(B) == QN(0)
    @test nnzblocks(B) == 2

    for ii in dim(i), jj in dim(j)
      @test 2*A[i=>ii,j=>jj] == B[i=>ii,j=>jj]
    end
  end

  @testset "Copy" begin
    s = Index([QN(0)=>1,QN(1)=>1],"s")
    T = randomITensor(QN(0),s,s')
    cT = copy(T)
    for ss in dim(s), ssp in dim(s')
      @test T[s=>ss,s'=>ssp] == cT[s=>ss,s'=>ssp]
    end
  end

  @testset "Permute" begin
    i = Index([QN(0)=>1,QN(1)=>2],"i")
    j = Index([QN(0)=>3,QN(1)=>4],"j")

    A = randomITensor(QN(1),i,dag(j))

    @test flux(A) == QN(1)
    @test nnzblocks(A) == 1

    B = permute(A,j,i)

    @test flux(B) == QN(1)
    @test nnzblocks(B) == 1

    for ii in dim(i), jj in dim(j)
      @test A[ii,jj] == B[jj,ii]
    end
  end

  @testset "Contraction" begin
    i = Index([QN(0)=>1,QN(1)=>2],"i")
    j = Index([QN(0)=>3,QN(1)=>4],"j")

    A = randomITensor(QN(0),i,dag(j))

    @test flux(A) == QN(0)
    @test nnzblocks(A) == 2

    B = randomITensor(QN(1),j,dag(i)')

    @test flux(B) == QN(1)
    @test nnzblocks(B) == 1

    C = A*B

    @test hasinds(C,i,i')
    @test flux(C) == QN(1)
    @test nnzblocks(C) == 1
  end

  @testset "Combine and uncombine" begin

    @testset "Combine no indices" begin
      i1 = Index([QN(0,2)=>2,QN(1,2)=>2],"i1")
      A = randomITensor(QN(),i1,dag(i1'))

      C = combiner()
      c = combinedind(C)
      @test isnothing(c)
      AC = A*C
      @test nnz(AC) == nnz(A)
      @test nnzblocks(AC) == nnzblocks(A)
      @test hassameinds(AC,A)
      @test norm(AC-A*C) ≈ 0.0
      Ap = AC*dag(C)
      @test nnz(Ap) == nnz(A)
      @test nnzblocks(Ap) == nnzblocks(A)
      @test hassameinds(Ap,A)
      @test norm(A-Ap) ≈ 0.0
    end

    @testset "Combine set direction" begin
      i1 = Index([QN(0)=>2,QN(1)=>3],"i1")
      A = randomITensor(i1',dag(i1))
      @test isnothing(ITensors.checkflux(A))
      C = combiner(dag(i1); dir = ITensors.Out)
      c = combinedind(C)
      @test dir(c) == ITensors.Out
      AC = A * C
      @test nnz(AC) == nnz(A)
      @test nnzblocks(AC) == nnzblocks(A)
      @test isnothing(ITensors.checkflux(AC))
      Ap = AC*dag(C)
      @test nnz(Ap) == nnz(A)
      @test nnzblocks(Ap) == nnzblocks(A)
      @test hassameinds(Ap, A)
      @test isnothing(ITensors.checkflux(AC))
      @test A ≈ Ap
    end

    @testset "Order 2 (IndexSet constructor)" begin
      i1 = Index([QN(0,2)=>2,QN(1,2)=>2],"i1")
      A = randomITensor(QN(),i1,dag(i1'))

      iss = [i1,
             dag(i1'),
             (i1,dag(i1')),
             (dag(i1'),i1)]

      for is in iss
        C = combiner(IndexSet(is); tags="c")
        AC = A*C
        @test nnz(AC) == nnz(A)
        Ap = AC*dag(C)
        @test nnz(Ap) == nnz(A)
        @test nnzblocks(Ap) == nnzblocks(A)
        @test norm(A-Ap) ≈ 0.0
      end
    end

    @testset "Order 2" begin
      i1 = Index([QN(0,2)=>2,QN(1,2)=>2],"i1")
      A = randomITensor(QN(), (i1, dag(i1')))

      iss = [i1,
             dag(i1'),
             (i1,dag(i1')),
             (dag(i1'),i1)]

      for is in iss
        C = combiner(is; tags="c")
        AC = A*C
        @test nnz(AC) == nnz(A)
        Ap = AC*dag(C)
        @test nnz(Ap) == nnz(A)
        @test nnzblocks(Ap) == nnzblocks(A)
        @test norm(A-Ap) ≈ 0.0
      end
    end

    @testset "Order 3, Combine 2" begin
      i = Index([QN(0)=>2,QN(1)=>2],"i")

      A = randomITensor(QN(0),i,dag(i)',dag(i)'')

      C = combiner(i,dag(i)'')
      c = combinedind(C)

      AC = A*C

      @test hasinds(AC,c,i')
      @test nnz(AC) == nnz(A)

      for b in nzblocks(AC)
        @test flux(AC,b) == QN(0)
      end

      B = ITensor(QN(0),i',c)
      @test nnz(B) == nnz(AC)
      @test nnzblocks(B) == nnzblocks(AC)

      Ap = AC*dag(C)

      @test norm(A-Ap) == 0
      @test nnz(A) == nnz(Ap)
      @test nnzblocks(A) == nnzblocks(Ap)
      @test hassameinds(A,Ap)
    end

    @testset "Order 3" begin
      i1 = Index([QN(0,2)=>2,QN(1,2)=>2],"i1")
      i2 = settags(i1,"i2")
      A = randomITensor(QN(),i1,i2,dag(i1'))

      iss = [i1,
             i2,
             dag(i1'),
             (i1,i2),
             (i2,i1),
             (i1,dag(i1')),
             (dag(i1'),i1),
             (i2,dag(i1')),
             (dag(i1'),i2),
             (i1,i2,dag(i1')),
             (i1,dag(i1'),i2),
             (i2,i1,dag(i1')),
             (i2,dag(i1'),i1),
             (dag(i1'),i1,i2),
             (dag(i1'),i2,i1)]

      for is in iss
        C = combiner(is; tags="c")
        AC = A*C
        @test nnz(AC) == nnz(A)
        Ap = AC*dag(C)
        @test nnz(Ap) == nnz(A)
        @test nnzblocks(Ap) == nnzblocks(A)
        @test norm(A-AC*dag(C)) ≈ 0.0
      end
    end

    @testset "Order 4" begin
      i1 = Index([QN(0,2)=>2,QN(1,2)=>2],"i1")
      i2 = settags(i1,"i2")
      A = randomITensor(QN(),i1,i2,dag(i1'),dag(i2'))

      iss = [i1,
             i2,
             dag(i1'),
             dag(i2'),
             (i1,i2),
             (i2,i1),
             (i1,dag(i1')),
             (dag(i1'),i1),
             (i1,dag(i2')),
             (dag(i2'),i1),
             (i2,dag(i1')),
             (dag(i1'),i2),
             (i2,dag(i2')),
             (dag(i2'),i2),
             (dag(i1'),dag(i2')),
             (dag(i2'),dag(i1')),
             (i1,i2,dag(i1')),
             (i1,dag(i1'),i2),
             (i2,i1,dag(i1')),
             (i2,dag(i1'),i1),
             (dag(i1'),i1,i2),
             (dag(i1'),i2,i1),
             (i1,dag(i1'),dag(i2')),
             (i1,dag(i2'),dag(i1')),
             (dag(i1'),i1,dag(i2')),
             (dag(i1'),dag(i2'),i1),
             (dag(i2'),i1,dag(i1')),
             (dag(i2'),dag(i1'),i1),
             (i1,i2,dag(i1'),dag(i2')),
             (i1,i2,dag(i2'),dag(i1')),
             (i1,dag(i1'),i2,dag(i2')),
             (i1,dag(i1'),dag(i2'),i2),
             (i1,dag(i2'),i2,dag(i1')),
             (i1,dag(i2'),dag(i1'),i2),
             (i2,i1,dag(i1'),dag(i2')),
             (i2,i1,dag(i2'),dag(i1')),
             (i2,dag(i1'),i1,dag(i2')),
             (i2,dag(i1'),dag(i2'),i1),
             (i2,dag(i2'),i1,dag(i1')),
             (i2,dag(i2'),dag(i1'),i1),
             (dag(i1'),i2,i1,dag(i2')),
             (dag(i1'),i2,dag(i2'),i1),
             (dag(i1'),i1,i2,dag(i2')),
             (dag(i1'),i1,dag(i2'),i2),
             (dag(i1'),dag(i2'),i2,i1),
             (dag(i1'),dag(i2'),i1,i2),
             (dag(i2'),i1,dag(i1'),i2),
             (dag(i2'),i1,i2,dag(i1')),
             (dag(i2'),dag(i1'),i1,i2),
             (dag(i2'),dag(i1'),i2,i1),
             (dag(i2'),i2,i1,dag(i1')),
             (dag(i2'),i2,dag(i1'),i1)]

      for is in iss
        C = combiner(is; tags="c")
        AC = A*C
        @test nnz(AC) == nnz(A)
        Ap = AC*dag(C)
        @test nnz(Ap) == nnz(A)
        @test nnzblocks(Ap) == nnzblocks(A)
        @test norm(A-Ap) ≈ 0.0
      end
    end

    @testset "Order 4, Combine 2, Example 1" begin
      s1 = Index([QN(("Sz", 0),("Nf",0))=>1,
                  QN(("Sz",+1),("Nf",1))=>1,
                  QN(("Sz",-1),("Nf",1))=>1,
                  QN(("Sz", 0),("Nf",2))=>1],"site,n=1");
      s2 = replacetags(s1,"n=1","n=2")

      A = randomITensor(QN(),s1,s2,dag(s1)',dag(s2)')

      C = combiner(dag(s1)',dag(s2)')
      c = combinedind(C)

      AC = A*C

      @test norm(AC) ≈ norm(A)
      @test hasinds(AC,s1,s2,c)
      @test nnz(AC) == nnz(A)
      for b in nzblocks(AC)
        @test flux(AC,b) == QN()
      end

      @test nnzblocks(AC) < nnz(A)

      B = ITensor(QN(),s1,s2,c)
      @test nnz(B) == nnz(AC)
      @test nnzblocks(B) == nnzblocks(AC)

      Ap = AC*dag(C)

      @test hassameinds(A,Ap)
      @test norm(A-Ap) == 0
      @test nnz(A) == nnz(Ap)
      @test nnzblocks(A) == nnzblocks(Ap)
    end

    @testset "Order 4, Combine 2, Example 2" begin
      s1 = Index([QN(("Nf",0))=>1,
                  QN(("Nf",1))=>1],"site,n=1")
      s2 = replacetags(s1,"n=1","n=2")

      A = randomITensor(QN(),dag(s2)',s2,dag(s1)',s1)

      C = combiner(dag(s2)',dag(s1)')
      c = combinedind(C)

      AC = A*C

      @test norm(AC) ≈ norm(A)
      @test hasinds(AC,s1,s2,c)
      @test nnz(AC) == nnz(A)
      for b in nzblocks(AC)
        @test flux(AC,b) == QN()
      end

      B = ITensor(QN(),s1,s2,c)
      @test nnzblocks(B) == nnzblocks(AC)

      Ap = AC*dag(C)

      @test hassameinds(A,Ap)
      @test norm(A-Ap) == 0
      @test nnz(A) == nnz(Ap)
      @test nnzblocks(A) == nnzblocks(Ap)
    end

    @testset "Order 4, Combine 2, Example 3" begin
      s1 = Index([QN(("Nf",0))=>1,
                  QN(("Nf",1))=>1],"site,n=1")
      s2 = replacetags(s1,"n=1","n=2")

      A = randomITensor(QN(),dag(s1)',s2,dag(s2)',s1)

      C = combiner(dag(s2)',dag(s1)')
      c = combinedind(C)

      AC = A*C

      @test norm(AC) ≈ norm(A)
      @test hasinds(AC,s1,s2,c)
      @test nnz(AC) == nnz(A)
      for b in nzblocks(AC)
        @test flux(AC,b) == QN()
      end

      B = ITensor(QN(),s1,s2,c)
      @test nnzblocks(B) == nnzblocks(AC)

      Ap = AC*dag(C)

      @test hassameinds(A,Ap)
      @test norm(A-Ap) == 0
      @test nnz(A) == nnz(Ap)
      @test nnzblocks(A) == nnzblocks(Ap)
    end
  end

  @testset "Check that combiner commutes" begin
    i = Index(QN(0,2)=>2,QN(1,2)=>2; tags="i")
    j = settags(i,"j")
    A = randomITensor(QN(0,2),i,j,dag(i'),dag(j'))
    C = combiner(i,j)
    @test norm(A*dag(C')*C-A*C*dag(C')) ≈ 0.0
  end

  @testset "Combiner for block deficient ITensor" begin
    i = Index(QN(0,2)=>2,QN(1,2)=>2; tags="i")
    j = settags(i,"j")
    A = emptyITensor(i,j,dag(i'))
    A[1,1,1] = 1.0
    C = combiner(i,j; tags="c")
    AC = A*C
    Ap = AC*dag(C)
    @test norm(A-Ap) ≈ 0.0
    @test norm(Ap-A) ≈ 0.0
  end

  @testset "Combine Complex ITensor" begin
    s1 = Index(QN(("Sz",1)) => 1,QN(("Sz",-1)) => 1;tags="S=1/2,Site,n=1")
    s2 = Index(QN(("Sz",1)) => 1,QN(("Sz",-1)) => 1;tags="S=1/2,Site,n=2")

    T = randomITensor(ComplexF64,QN("Sz",0),s1,s2)

    C = combiner(s1,s2)
    CT = C*T
    @test norm(CT) ≈ norm(T)
    TT = dag(C)*CT
    @test TT ≈ T
  end

  @testset "Combiner bug #395" begin
    i1 = Index([QN(0)=>1,QN(1)=>2],"i1")
    i2 = Index([QN(0)=>1,QN(1)=>2],"i2")
    A = randomITensor(QN(),i1,i2,dag(i1)', dag(i2)')
    CL = combiner(i1,i2)
    CR = combiner(dag(i1)',dag(i2)')
    AC = A * CR * CL
    @test AC * dag(CR) * dag(CL) ≈ A
  end

  @testset "Contract to scalar" begin
    i = Index([QN(0)=>1,QN(1)=>1],"i")
    A = randomITensor(QN(0),i,dag(i'))

    c = A*dag(A)

    @test nnz(c) == 1
    @test nnzblocks(c) == 1
    @test c[] isa Float64
    @test c[] ≈ norm(A)^2
  end

  @testset "eigen" begin

    @testset "eigen hermitian" begin
      i = Index(QN(0)=>2,QN(1)=>3,QN(2)=>4; tags="i")
      j = settags(i,"j")
      k = settags(i,"k")
      l = settags(i,"l")

      A = randomITensor(QN(),i,j,dag(k),dag(l))
      A = A*prime(dag(A),(i,j))

      F = eigen(A; ishermitian=true, tags="x")

      D, U = F
      Ut = F.Vt

      @test store(U) isa NDTensors.BlockSparse
      @test store(D) isa NDTensors.DiagBlockSparse

      u = commonind(D,U)
      up = uniqueind(D,U)

      @test hastags(u,"x")
      @test plev(u) == 0
      @test hastags(up,"x")
      @test plev(up) == 1

      @test hassameinds(U, (i,j,u))
      @test hassameinds(D, (u,up))

      @test norm(A - dag(U) * D * U') ≈ 0.0 atol=1e-11
      @test norm(A - dag(U) * D * Ut) ≈ 0.0 atol=1e-11
      @test norm(A * U - U' * D) ≈ 0.0 atol=1e-11
      @test norm(A * U - Ut * D) ≈ 0.0 atol=1e-11
    end

    @testset "eigen hermitian (truncate)" begin
      i = Index(QN(0)=>2,QN(1)=>3,QN(2)=>4; tags="i")
      j = settags(i,"j")
      k = settags(i,"k")
      l = settags(i,"l")

      A = randomITensor(QN(),i,j,dag(k),dag(l))
      A = A*prime(dag(A),(i,j))
      for i = 1:4
        A = mapprime(A*A',2,1)
      end
      A = A/norm(A)

      cutoff = 1e-5
      F = eigen(A; ishermitian=true,
                   tags="x",
                   cutoff=cutoff)

      D, U, spec = F
      Ut = F.Vt

      @test store(U) isa NDTensors.BlockSparse
      @test store(D) isa NDTensors.DiagBlockSparse

      u = commonind(D, U)
      up = uniqueind(D, U)

      @test hastags(u,"x")
      @test plev(u) == 0
      @test hastags(up,"x")
      @test plev(up) == 1

      @test hassameinds(U,(i,j,u))
      @test hassameinds(D,(u,up))

      for b in nzblocks(A)
        @test flux(A,b)==QN(0)
      end
      for b in nzblocks(U)
        @test flux(U,b)==QN(0)
      end
      for b in nzblocks(D)
        @test flux(D,b)==QN(0)
      end

      Ap = dag(U) * D * U'

      @test norm(Ap-A) ≤ 1e-2
      @test norm(dag(U) * D * Ut - A) ≤ 1e-2
      @test minimum(dims(D)) == length(spec.eigs)
      @test minimum(dims(D)) < dim(i)*dim(j)

      @test spec.truncerr ≤ cutoff
      err = sqrt(1-(Ap*dag(Ap))[]/(A*dag(A))[])
      @test err ≤ cutoff
      @test err ≈ spec.truncerr rtol=1e-1
		end

    @testset "eigen non-hermitian" begin
      i = Index(QN(0)=>2,QN(1)=>3,QN(2)=>4; tags="i")
      j = settags(i,"j")

      A = randomITensor(QN(),i,j,dag(i'),dag(j'))

      F = eigen(A; tags="x")

      D, U = F
      Ut = F.Vt

      @test store(U) isa NDTensors.BlockSparse
      @test store(D) isa NDTensors.DiagBlockSparse

      u = commonind(D,U)
      up = uniqueind(D,U)

      @test hastags(u,"x")
      @test plev(u) == 0
      @test hastags(up,"x")
      @test plev(up) == 1

      @test norm(A - U' * D * dag(U)) ≉ 0.0 atol=1e-12
      @test norm(A - Ut * D * dag(U)) ≉ 0.0 atol=1e-12
      @test norm(A * U - U' * D) ≈ 0.0 atol=1e-12
      @test norm(A * U - Ut * D) ≈ 0.0 atol=1e-12
    end

    @testset "eigen non-hermitian (general inds)" begin
      i = Index(QN(0) => 2,
                QN(1) => 3,
                QN(2) => 4; tags="i")
      j = settags(i, "j")
      ĩ, j̃ = sim(i), sim(j)

      A = randomITensor(QN(), i, j, dag(ĩ),dag(j̃))

      F = eigen(A, (i, j), (ĩ, j̃); lefttags = "x", righttags = "y")

      D, U = F
      Ut = F.Vt

      @test store(U) isa NDTensors.BlockSparse
      @test store(D) isa NDTensors.DiagBlockSparse

      l = uniqueind(D, U)
      r = commonind(D, U)

      @test F.l == l
      @test F.r == r

      @test hastags(l, "x")
      @test plev(l) == 0
      @test hastags(r, "y")
      @test plev(r) == 0

      @test hassameinds(U, (ĩ, j̃, r))
      @test hassameinds(Ut, (i, j, l))

      @test norm(A * U - Ut * D) ≈ 0.0 atol=1e-12
      @test norm(A - Ut * D * dag(U)) ≉ 0.0 atol=1e-12
    end

    @testset "eigen mixed arrows" begin
      i1 = Index([QN(0)=>1,QN(1)=>2],"i1")
      i2 = Index([QN(0)=>1,QN(1)=>2],"i2")
      A = randomITensor(i1, i2, dag(i1)', dag(i2)')
      F = eigen(A, (i1, i1'), (i2', i2))
      D, U = F
      Ut = F.Vt
      @test norm(A * U - Ut * D) ≈ 0.0 atol=1e-12
    end

  end

  @testset "svd" for ElT ∈ (Float64, ComplexF64)

    @testset "svd example 1" begin
      i = Index(QN(0)=>2,QN(1)=>2; tags="i")
      j = Index(QN(0)=>2,QN(1)=>2; tags="j")
      A = randomITensor(ElT, QN(0),i,dag(j))
      for b in nzblocks(A)
        @test flux(A,b)==QN(0)
      end
      U,S,V = svd(A,i)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      for b in nzblocks(U)
        @test flux(U,b)==QN(0)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(0)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0)
      end
      @test isapprox(norm(U*S*V-A),0.0; atol=1e-14)
    end

    @testset "svd example 2" begin
      i = Index(QN(0)=>5,QN(1)=>6; tags="i")
      j = Index(QN(-1)=>2,QN(0)=>3,QN(1)=>4; tags="j")
      A = randomITensor(ElT, QN(0),i,j)
      for b in nzblocks(A)
        @test flux(A,b)==QN(0)
      end
      U,S,V = svd(A,i)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      for b in nzblocks(U)
        @test flux(U,b)==QN(0)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(0)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0)
      end
      @test isapprox(norm(U*S*V-A),0.0; atol=1e-14)
    end

    @testset "svd example 3" begin
      i = Index(QN(0)=>5,QN(1)=>6; tags="i")
      j = Index(QN(-1)=>2,QN(0)=>3,QN(1)=>4; tags="j")
      A = randomITensor(ElT, QN(0),i,dag(j))
      for b in nzblocks(A)
        @test flux(A,b)==QN(0)
      end
      U,S,V = svd(A,i)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      for b in nzblocks(U)
        @test flux(U,b)==QN(0)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(0)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0)
      end
      @test isapprox(norm(U*S*V-A),0.0; atol=1e-14)
    end

    @testset "svd example 4" begin
			i = Index(QN(0,2)=>2,QN(1,2)=>2; tags="i")
			j = settags(i,"j")

			A = randomITensor(ElT, QN(0,2),i,j,dag(i'),dag(j'))

			U,S,V = svd(A,i,j)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      for b in nzblocks(A)
        @test flux(A,b)==QN(0,2)
      end
      U,S,V = svd(A,i)
      for b in nzblocks(U)
        @test flux(U,b)==QN(0,2)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(0,2)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0,2)
      end
      @test isapprox(norm(U*S*V-A),0.0; atol=1e-14)
    end

    @testset "svd example 5" begin
			i = Index(QN(0,2)=>2,QN(1,2)=>2; tags="i")
			j = settags(i,"j")

			A = randomITensor(ElT, QN(1,2),i,j,dag(i'),dag(j'))

			U,S,V = svd(A,i,j)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      for b in nzblocks(A)
        @test flux(A,b)==QN(1,2)
      end
      U,S,V = svd(A,i)
      for b in nzblocks(U)
        @test flux(U,b)==QN(0,2)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(1,2)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0,2)
      end
      @test isapprox(norm(U*S*V-A),0.0; atol=1e-14)
    end

    @testset "svd example 6" begin
			i = Index(QN(0,2)=>2,QN(1,2)=>2; tags="i")
			j = settags(i,"j")

			A = randomITensor(ElT, QN(1,2),i,j,dag(i'),dag(j'))

			U,S,V = svd(A,i,i')

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      for b in nzblocks(A)
        @test flux(A,b)==QN(1,2)
      end
      U,S,V = svd(A,i)
      for b in nzblocks(U)
        @test flux(U,b)==QN(0,2)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(1,2)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0,2)
      end
      @test isapprox(norm(U*S*V-A),0.0; atol=1e-14)
    end

    @testset "svd truncation example 1" begin
      i = Index(QN(0)=>2,QN(1)=>3; tags="i")
      j = settags(i,"j")
      A = randomITensor(ElT, QN(0),i,j,dag(i'),dag(j'))
      for i = 1:4
        A = mapprime(A*A',2,1)
      end
      A = A/norm(A)

      cutoff = 1e-5
      U,S,V,spec = svd(A,i,j; utags="x", vtags="y", cutoff=cutoff)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      u = commonind(S,U)
      v = commonind(S,V)

      @test hastags(u,"x")
      @test hastags(v,"y")

      @test hassameinds(U,(i,j,u))
      @test hassameinds(V,(i',j',v))

      for b in nzblocks(A)
        @test flux(A,b)==QN(0)
      end
      for b in nzblocks(U)
        @test flux(U,b)==QN(0)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(0)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0)
      end

      Ap = U*S*V

      @test norm(Ap-A) ≤ 1e-2
      @test minimum(dims(S)) == length(spec.eigs)
      @test minimum(dims(S)) < dim(i)*dim(j)

      @test spec.truncerr ≤ cutoff
      err = real(1-(Ap*dag(Ap))[]/(A*dag(A))[])
      @test err ≤ cutoff
      @test isapprox(err,spec.truncerr; rtol=1e-6)
    end

    @testset "svd truncation example 2" begin
      i = Index(QN(0)=>3,QN(1)=>2; tags="i")
      j = settags(i,"j")
      A = randomITensor(ElT, QN(0),i,j,dag(i'),dag(j'))

      maxdim = 4
      U,S,V,spec = svd(A,i,j; utags="x", vtags="y", maxdim=maxdim)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      u = commonind(S,U)
      v = commonind(S,V)

      @test hastags(u,"x")
      @test hastags(v,"y")

      @test hassameinds(U,(i,j,u))
      @test hassameinds(V,(i',j',v))

      for b in nzblocks(A)
        @test flux(A,b)==QN(0)
      end
      for b in nzblocks(U)
        @test flux(U,b)==QN(0)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(0)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0)
      end

      @test minimum(dims(S)) == maxdim
      @test minimum(dims(S)) == length(spec.eigs)
      @test minimum(dims(S)) < dim(i)*dim(j)

      Ap = U*S*V
      err = 1-(Ap*dag(Ap))[]/(A*dag(A))[]
      @test isapprox(err,spec.truncerr; rtol=1e-6)
    end

    @testset "svd truncation example 3" begin
      i = Index(QN(0)=>2,QN(1)=>3,QN(2)=>4; tags="i")
      j = settags(i,"j")
      A = randomITensor(ElT, QN(1),i,j,dag(i'),dag(j'))

      maxdim = 4
      U,S,V,spec = svd(A,i,j; utags="x", vtags="y", maxdim=maxdim)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      u = commonind(S,U)
      v = commonind(S,V)

      @test hastags(u,"x")
      @test hastags(v,"y")

      @test hassameinds(U,(i,j,u))
      @test hassameinds(V,(i',j',v))

      for b in nzblocks(A)
        @test flux(A,b)==QN(1)
      end
      for b in nzblocks(U)
        @test flux(U,b)==QN(0)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(1)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0)
      end

      @test minimum(dims(S)) == maxdim
      @test minimum(dims(S)) == length(spec.eigs)
      @test minimum(dims(S)) < dim(i)*dim(j)

      Ap = U*S*V
      err = 1-(Ap*dag(Ap))[]/(A*dag(A))[]
      @test isapprox(err,spec.truncerr; rtol=1e-6)
    end

    @testset "svd truncation example 4" begin
      i = Index(QN(0,2)=>3,QN(1,2)=>4; tags="i")
      j = settags(i,"j")
      A = randomITensor(ElT, QN(1,2),i,j,dag(i'),dag(j'))

      maxdim = 4
      U,S,V,spec = svd(A,i,j; utags="x", vtags="y", maxdim=maxdim)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      u = commonind(S,U)
      v = commonind(S,V)

      @test hastags(u,"x")
      @test hastags(v,"y")

      @test hassameinds(U,(i,j,u))
      @test hassameinds(V,(i',j',v))

      for b in nzblocks(A)
        @test flux(A,b)==QN(1,2)
      end
      for b in nzblocks(U)
        @test flux(U,b)==QN(0,2)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(1,2)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0,2)
      end

      @test minimum(dims(S)) == maxdim
      @test minimum(dims(S)) == length(spec.eigs)
      @test minimum(dims(S)) < dim(i)*dim(j)

      Ap = U*S*V
      err = 1-(Ap*dag(Ap))[]/(A*dag(A))[]
      @test isapprox(err,spec.truncerr; rtol=1e-6)
    end

    @testset "svd truncation example 5" begin
      i = Index(QN(0,2)=>2,QN(1,2)=>3; tags="i")
      j = settags(i,"j")
      A = randomITensor(ElT, QN(1,2),i,j,dag(i'),dag(j'))

      maxdim = 4
      U,S,V,spec = svd(A,i,j'; utags="x", vtags="y", maxdim=maxdim)

      @test store(U) isa NDTensors.BlockSparse
      @test store(S) isa NDTensors.DiagBlockSparse
      @test store(V) isa NDTensors.BlockSparse

      u = commonind(S,U)
      v = commonind(S,V)

      @test hastags(u,"x")
      @test hastags(v,"y")

      @test hassameinds(U,(i,j',u))
      @test hassameinds(V,(i',j,v))

      for b in nzblocks(A)
        @test flux(A,b)==QN(1,2)
      end
      for b in nzblocks(U)
        @test flux(U,b)==QN(0,2)
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN(1,2)
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN(0,2)
      end

      @test minimum(dims(S)) == maxdim
      @test minimum(dims(S)) == length(spec.eigs)
      @test minimum(dims(S)) < dim(i)*dim(j)

      Ap = U*S*V
      err = 1-(Ap*dag(Ap))[]/(A*dag(A))[]
      @test isapprox(err,spec.truncerr; rtol=1e-6)
    end

    @testset "issue #231" begin
      l = Index(QN("Nf",-1,-1)=>2,
            QN("Nf", 0,-1)=>4,
            QN("Nf",+1,-1)=>2,
            tags="CMB,Link")
      s = Index(QN("Nf",0,-1)=>1,
                QN("Nf",1,-1)=>1,
                tags="Fermion,Site,n=4")
      r = Index(QN("Nf",1,-1)=>2,
                QN("Nf",0,-1)=>1,
                QN("Nf",1,-1)=>2,
                tags="Link,u")

      A = emptyITensor(ElT, l,s,dag(r))

      addblock!(A,(2,1,2))
      addblock!(A,(1,2,2))
      addblock!(A,(2,2,3))

      for b in nzblocks(A)
        @test flux(A,b)==QN()
      end

      U,S,V = svd(A,l,s)

      for b in nzblocks(U)
        @test flux(U,b)==QN()
      end
      for b in nzblocks(S)
        @test flux(S,b)==QN()
      end
      for b in nzblocks(V)
        @test flux(V,b)==QN()
      end
      @test norm(U*S*V-A) ≈ 0 atol=1e-15
    end

    @testset "SVD no truncate bug" begin
      s = Index(QN("Sz",-4) => 1,
                QN("Sz",-2) => 4,
                QN("Sz", 0) => 6,
                QN("Sz", 2) => 4,
                QN("Sz", 4) => 1)
      A = emptyITensor(ElT, s, s')
      addblock!(A, (5,2))
      addblock!(A, (4,3))
      addblock!(A, (3,4))
      addblock!(A, (2,5))
      randn!(A)
      U,S,V = svd(A,s)
      @test U*S*V ≈ A
    end

    @testset "SVD no truncate" begin
      s = Index(QN("Sz",-4) => 1,
                QN("Sz",-2) => 4,
                QN("Sz", 0) => 6,
                QN("Sz", 2) => 4,
                QN("Sz", 4) => 1)
      A = emptyITensor(ElT, s, s')
      addblock!(A, (5,1))
      addblock!(A, (4,2))
      addblock!(A, (3,3))
      addblock!(A, (2,4))
      addblock!(A, (1,5))
      U,S,V = svd(A, s)
      @test dims(S) == dims(A)
      @test U*S*V ≈ A
    end

    @testset "SVD truncate zeros" begin
      s = Index(QN("Sz",-4) => 1,
                QN("Sz",-2) => 4,
                QN("Sz", 0) => 6,
                QN("Sz", 2) => 4,
                QN("Sz", 4) => 1)
      A = emptyITensor(ElT, s, s')
      addblock!(A, (5,1))
      addblock!(A, (4,2))
      addblock!(A, (3,3))
      addblock!(A, (2,4))
      addblock!(A, (1,5))
      U,S,V = svd(A, s; cutoff=0)
      @test dims(S) == (0,0)
      @test U*S*V ≈ A
    end

  end

  @testset "Replace Index" begin
    i = Index([QN(0)=>1,QN(1)=>2],"i")
    j = Index([QN(0)=>3,QN(1)=>4,QN(2)=>5],"j")

    T1 = randomITensor(QN(1),i,j)
    T2 = copy(T1)

    k = Index([QN(0)=>1,QN(1)=>2],"k")

    replaceind!(T1,i,k)
    @test hasind(T1,k)
    @test dir(inds(T1)[1]) == dir(i)

    # Check that replaceind! keeps
    # original Arrow direction
    replaceind!(T2,i,dag(k))
    @test hasind(T2,k)
    @test dir(inds(T2)[1]) == dir(i)
    @test dir(inds(T2)[1]) != dir(dag(k))
  end

@testset "BlockSparse dag copy behavior" begin
  i = Index(QN(0)=>2,QN(1)=>2,tags="i")
  j = Index(QN(0)=>2,QN(1)=>2,tags="j")

  v1 = randomITensor(QN(1),i,j)
  orig_elt = v1[1,3]
  cv1 = dag(v1;always_copy=false)
  cv1[1,3] = 123.45
  @test v1[1,3] ≈ cv1[1,3]

  v2 = randomITensor(QN(1),i,j)
  orig_elt = v2[1,3]
  cv2 = dag(v2;always_copy=true)
  cv2[1,3] = 123.45
  @test v2[1,3] ≈ orig_elt

  v3 = randomITensor(ComplexF64,QN(1),i,j)
  orig_elt = v3[1,3]
  cv3 = dag(v3;always_copy=false)
  cv3[1,3] = 123.45
  @test v3[1,3] ≈ orig_elt

  v4 = randomITensor(ComplexF64,QN(1),i,j)
  orig_elt = v4[1,3]
  cv4 = dag(v4;always_copy=true)
  cv4[1,3] = 123.45
  @test v4[1,3] ≈ orig_elt

end

@testset "exponentiate" begin

  @testset "Simple arrows" begin
    i1 = Index([QN(0)=>1,QN(1)=>2],"i1")
    i2 = Index([QN(0)=>1,QN(1)=>2],"i2")
    A = randomITensor(QN(),i1,i2,dag(i1)', dag(i2)')
    Aexp = exp(A)
    Amat = Array(A, i1,i2, i1', i2')
    Amatexp = reshape(exp(reshape(Amat,9,9)),3,3,3,3)
    @test isapprox(norm(Array(Aexp,i1,i2,i1',i2') - Amatexp),0.0, atol=1e-14)
    @test flux(Aexp) == QN()
    @test length(setdiff(inds(Aexp),inds(A)))==0

    @test isapprox(norm(exp(A, (i1,i2), (i1',i2')) - Aexp),0.0, atol=5e-14)

    # test the case where indices are permuted
    A = randomITensor(QN(),i1,dag(i1)', dag(i2)',i2)
    Aexp = exp(A, (i1,i2), (i1',i2'))
    Amat = Array(A, i1,i2,i1', i2')
    Amatexp = reshape(exp(reshape(Amat,9,9)),3,3,3,3)
    @test isapprox(norm(Array(Aexp,i1,i2,i1',i2') - Amatexp),0.0,atol=1e-14)

    # test exponentiation in the Hermitian case
    i1 = Index([QN(0)=>2,QN(1)=>2,QN(2)=>3],"i1")
    A = randomITensor(QN(),i1,dag(i1)')
    Ad = dag(swapinds(A,IndexSet(i1),IndexSet(dag(i1)')))
    Ah = A+Ad + 1e-10*randomITensor(QN(),i1,dag(i1)')
    Amat = Array(Ah ,i1', i1)
    Aexp = exp(Ah; ishermitian=true)
    Amatexp= exp(LinearAlgebra.Hermitian(Amat))
    @test isapprox(norm(Array(Aexp,i1,i1')-Amatexp),0.0,atol=5e-14)
  end

  @testset "Mixed arrows" begin
    i1 = Index([QN(0)=>1,QN(1)=>2],"i1")
    i2 = Index([QN(0)=>1,QN(1)=>2],"i2")
    A = randomITensor(i1, i2, dag(i1)', dag(i2)')
    expA = exp(A, (i1, i1'), (i2', i2))
    @test exp(dense(A), (i1, i1'), (i2', i2)) ≈ dense(expA)
  end

end

end

nothing