using ProfileView
using PolynomialBasis
using ImplicitDomainQuadrature
using BenchmarkTools
# using Revise
using CutCellDG
include("../../../test/useful_routines.jl")
# include("transformation-elasticity-solver.jl")

# TES = TransformationElasticitySolver

polyorder = 2
nelmts = 15
penaltyfactor = 1e3

################################################################################
width = 1.0              # mm
interfacecenter = [0.5, 0.5]
interfaceradius = 0.4
outerradius = 1.2
################################################################################

################################################################################
K1, K2 = 247.0, 192.0    # GPa
mu1, mu2 = 126.0, 87.0   # GPa
lambda1 = lame_lambda(K1, mu1)
lambda2 = lame_lambda(K2, mu2)
stiffness = CutCellDG.HookeStiffness(lambda1, mu1, lambda2, mu2)
theta0 = -0.067
transfstress =
    CutCellDG.plane_strain_transformation_stress(lambda1, mu1, theta0)
################################################################################

meshwidth = [width, width]
numqp = required_quadrature_order(polyorder) + 2

levelsetbasis = HermiteTensorProductBasis(2)
quad = tensor_product_quadrature(2, 4)
elasticitybasis = LagrangeTensorProductBasis(2, polyorder)

basispts = interpolation_points(elasticitybasis)
dim, numpts = size(interpolation_points(levelsetbasis))
cgmesh = CutCellDG.CGMesh([0.0, 0.0], meshwidth, [nelmts, nelmts], numpts)
dgmesh = CutCellDG.DGMesh([0.0, 0.0], meshwidth, [nelmts, nelmts], basispts)

levelset = CutCellDG.LevelSet(
    x -> -circle_distance_function(x, interfacecenter, interfaceradius)[1],
    cgmesh,
    levelsetbasis,
    quad,
)

elementsize = CutCellDG.element_size(cgmesh)
minelmtsize = minimum(elementsize)
maxelmtsize = maximum(elementsize)
penalty = penaltyfactor / minelmtsize * 0.5 * (lambda1 + mu1 + lambda2 + mu2)
tol = minelmtsize^(polyorder + 1)
boundingradius = 1.5 * maxelmtsize

cutmesh = CutCellDG.CutMesh(dgmesh, levelset)

@profview cellquads = CutCellDG.CellQuadratures(cutmesh,levelset,numqp)


module SpecialMatricesModule
export computeMatrixKinetic, computeTensorElectronRepulsionIntegrals, computeMatrixNuclearAttraction, computeMatrixOverlap, computeMatrixCoulomb, computeMatrixExchange, computeMatrixFock
using ..BaseModule
using ..AtomModule
using ..IntegralsModule
using ..BasisModule
using ..Geometry
using TensorOperations
using ..LibInt2Module
using ..ShellModule
import ..IntegralsModule.computeTensorBlockElectronRepulsionIntegrals

#HELPERS
function scatterMatrixBlocks2D(shells,totaldim::Integer,blockevaluator::Function,blocklength::Function)
  result = zeros(Float64,totaldim,totaldim)
  # column-major order (sh1 counts top to bottom, sh2 counts left to right)
  lastcol,lastrow = 0,0
  for sh1 in shells
    block_col_size = blocklength(sh1)
    lastcol += block_col_size
    for sh2 in shells
      block_row_size = blocklength(sh2)
      lastrow += block_row_size
      result[lastcol-block_col_size+1:lastcol, lastrow-block_row_size+1:lastrow] = blockevaluator(sh1,sh2)
    end
    lastrow = 0
  end
  return result
end

function contractTensorBlocksWithMatrix_CoulombLike(shells,density,totaldim::Integer,tensorblockevaluator::Function,blocklength::Function)
  result = zeros(Float64,totaldim,totaldim)
  # column-major order (sh1 counts top to bottom ("col"-direction), sh2 counts left to right ("row"-direction)
  lastcol,lastrow = 0,0
  lastcol_in_P,lastrow_in_P = 0,0
  for sh1 in shells
    block_col_size = blocklength(sh1)
    lastcol += block_col_size
    for sh2 in shells
      block_row_size = blocklength(sh2)
      lastrow += block_row_size
      for sh3 in shells
	P_block_col_size = blocklength(sh3)
	lastcol_in_P += P_block_col_size
	for sh4 in shells
	  P_block_row_size = blocklength(sh4)
	  lastrow_in_P += P_block_row_size
	  P = density[lastcol_in_P-P_block_col_size+1:lastcol_in_P, lastrow_in_P-P_block_row_size+1:lastrow_in_P]
	  tensor = tensorblockevaluator(sh1,sh2,sh3,sh4)
	  @tensor block[μ,ν] := tensor[μ,ν,λ,σ] * P[λ,σ]
	  result[lastcol-block_col_size+1:lastcol, lastrow-block_row_size+1:lastrow] += block
	end
	lastrow_in_P = 0
      end
      lastcol_in_P = 0
    end
    lastrow = 0
  end
  return result
end

function contractTensorBlocksWithMatrix_ExchangeLike(shells,density,totaldim::Integer,tensorblockevaluator::Function,blocklength::Function)
  result = zeros(Float64,totaldim,totaldim)
  # column-major order (sh1 counts top to bottom ("col"-direction), sh2 counts left to right ("row"-direction)
  lastcol,lastrow = 0,0
  lastcol_in_P,lastrow_in_P = 0,0
  for sh1 in shells
    block_col_size = blocklength(sh1)
    lastcol += block_col_size
    for sh2 in shells
      block_row_size = blocklength(sh2)
      lastrow += block_row_size
      for sh3 in shells
	P_block_col_size = blocklength(sh3)
	lastcol_in_P += P_block_col_size
	for sh4 in shells
	  P_block_row_size = blocklength(sh4)
	  lastrow_in_P += P_block_row_size
	  P = density[lastcol_in_P-P_block_col_size+1:lastcol_in_P, lastrow_in_P-P_block_row_size+1:lastrow_in_P]
	  tensor = tensorblockevaluator(sh1,sh3,sh2,sh4)
	  @tensor block[μ,ν] := tensor[μ,λ,ν,σ] * P[λ,σ]
	  result[lastcol-block_col_size+1:lastcol, lastrow-block_row_size+1:lastrow] += block
	end
	lastrow_in_P = 0
      end
      lastcol_in_P = 0
    end
    lastrow = 0
  end
  return result
end

function computeDimensions(shells::Vector{LibInt2Shell})
  totaldim = 0
  maxprims = 0
  maxlqn = 0
  for sh in shells
    totaldim += div((lqn(sh)+1)^2+(lqn(sh)+1),2)
    maxprims = maxprims < nprims(sh) ? nprims(sh) : maxprims
    maxlqn = maxlqn < lqn(sh) ? lqn(sh) : maxlqn
  end
  return (totaldim, maxprims, maxlqn)
end

#OVERLAP
function computeMatrixOverlap(basis::GaussianBasis)
  return [computeIntegralOverlap(cgb1,cgb2) for cgb1 in basis.contractedBFs, cgb2 in basis.contractedBFs]
end

function computeMatrixOverlap(shells::Vector{Shell})
  totaldim = mapreduce(ShellModule.nbf,+,0,shells)
  return scatterMatrixBlocks2D(shells,totaldim,computeMatrixBlockOverlap,ShellModule.nbf)
end

function computeMatrixOverlap(shells::Vector{LibInt2Shell})
  totaldim, maxprims, maxlqn = computeDimensions(shells)
  engine = LibInt2EngineOverlap(maxprims,LQuantumNumber(maxlqn))
  result = scatterMatrixBlocks2D(shells,totaldim,(sh1,sh2)->computeMatrixBlockOverlap(engine,sh1,sh2),sh->div((lqn(sh)+1)^2+(lqn(sh)+1),2))
  destroy!(engine)
  return result
end

#KINETIC
function computeMatrixKinetic(basis::GaussianBasis)
  return [computeIntegralKinetic(cgb1,cgb2) for cgb1 in basis.contractedBFs, cgb2 in basis.contractedBFs]
end

function computeMatrixKinetic(shells::Vector{Shell})
  totaldim = mapreduce(ShellModule.nbf,+,0,shells)
  return scatterMatrixBlocks2D(shells,totaldim,computeMatrixBlockKinetic,ShellModule.nbf)
end

function computeMatrixKinetic(shells::Vector{LibInt2Shell})
  totaldim, maxprims, maxlqn = computeDimensions(shells)
  engine = LibInt2EngineKinetic(maxprims,LQuantumNumber(maxlqn))
  result = scatterMatrixBlocks2D(shells,totaldim,(sh1,sh2)->computeMatrixBlockKinetic(engine,sh1,sh2),sh->div((lqn(sh)+1)^2+(lqn(sh)+1),2))
  destroy!(engine)
  return result
end

#NUCLEARATTRACTION
function computeMatrixNuclearAttraction(basis::GaussianBasis,geo::Geometry)
  return [sum([computeIntegralNuclearAttraction(cgb1,cgb2,atom) for atom in geo.atoms]) for cgb1 in basis.contractedBFs, cgb2 in basis.contractedBFs]
end

function computeMatrixNuclearAttraction(shells::Vector{Shell},atom::Atom)
  totaldim = mapreduce(ShellModule.nbf,+,0,shells)
  return scatterMatrixBlocks2D(shells,totaldim,(sh1,sh2)->computeMatrixBlockNuclearAttraction(sh1,sh2,atom),ShellModule.nbf)
end

function computeMatrixNuclearAttraction(shells::Vector{Shell},geo::Geometry)
  totaldim = mapreduce(ShellModule.nbf,+,0,shells)
  result = zeros(totaldim,totaldim)
  for atom in geo.atoms
    result += scatterMatrixBlocks2D(shells,totaldim,(sh1,sh2)->computeMatrixBlockNuclearAttraction(sh1,sh2,atom),ShellModule.nbf)
  end
  return result
end

function computeMatrixNuclearAttraction(shells::Vector{LibInt2Shell},atom::Atom)
  totaldim, maxprims, maxlqn = computeDimensions(shells)
  engine = LibInt2EngineNuclearAttraction(maxprims,LQuantumNumber(maxlqn),atom)
  result = scatterMatrixBlocks2D(shells,totaldim,(sh1,sh2)->computeMatrixBlockNuclearAttraction(engine,sh1,sh2),sh->div((lqn(sh)+1)^2+(lqn(sh)+1),2))
  destroy!(engine)
  return result
end

function computeMatrixNuclearAttraction(shells::Vector{LibInt2Shell},geo::Geometry)
  totaldim, maxprims, maxlqn = computeDimensions(shells)
  engine = LibInt2EngineNuclearAttraction(maxprims,LQuantumNumber(maxlqn),geo)
  result = scatterMatrixBlocks2D(shells,totaldim,(sh1,sh2)->computeMatrixBlockNuclearAttraction(engine,sh1,sh2),sh->div((lqn(sh)+1)^2+(lqn(sh)+1),2))
  destroy!(engine)
  return result
end

#ERIS
function computeTensorElectronRepulsionIntegrals(basis::GaussianBasis)
	return ERIs = [computeElectronRepulsionIntegral(μ,ν,λ,σ) for μ in basis.contractedBFs, ν in basis.contractedBFs, λ in basis.contractedBFs, σ in basis.contractedBFs]
end

function computeTensorBlockElectronRepulsionIntegrals(μ::Shell,ν::Shell,λ::Shell,σ::Shell)
  [computeElectronRepulsionIntegral(μcgbf,νcgbf,λcgbf,σcgbf) for μcgbf in expandShell(μ), νcgbf in expandShell(ν), λcgbf in expandShell(λ), σcgbf in expandShell(σ)]
end

#COULOMB
function computeMatrixCoulomb(basis::GaussianBasis, density::Matrix)
  N = length(basis.contractedBFs)

  J = zeros(N,N)
  μIndex = 0
  for (μ in basis.contractedBFs)
    μIndex+=1
    for (νIndex in 1:μIndex)
      ν = basis.contractedBFs[νIndex]
      ERImatrix = [computeElectronRepulsionIntegral(μ,ν,λ,σ) for λ in basis.contractedBFs, σ in basis.contractedBFs]
      J[μIndex,νIndex] = trace(ERImatrix*density)
      J[νIndex,μIndex] = J[μIndex,νIndex]
    end
  end
  return J
end

function computeMatrixCoulomb(electronRepulsionIntegrals::Array{Float64,4}, density::Matrix)
  @tensor J[μIndex,νIndex] := electronRepulsionIntegrals[μIndex,νIndex,λIndex,σIndex] * density[λIndex,σIndex]
end

function computeMatrixCoulomb(shells::Vector{Shell}, density)
  totaldim = mapreduce(ShellModule.nbf,+,0,shells)
  contractTensorBlocksWithMatrix_CoulombLike(shells,density,totaldim,computeTensorBlockElectronRepulsionIntegrals,ShellModule.nbf)
end

function computeMatrixCoulomb(shells::Vector{LibInt2Shell}, density::Matrix)
  totaldim, maxprims, maxlqn = computeDimensions(shells)
  engine = LibInt2EngineCoulomb(maxprims,LQuantumNumber(maxlqn))
  result = contractTensorBlocksWithMatrix_CoulombLike(shells,density,totaldim,(μ,ν,λ,σ)->computeTensorBlockElectronRepulsionIntegrals(engine,μ,ν,λ,σ),LibInt2Module.nbf)
  destroy!(engine)
  return result
end

#EXCHANGE
function computeMatrixExchange(basis::GaussianBasis, density::Matrix)
  N = length(basis.contractedBFs)

  K = zeros(N,N)
  μIndex = 0
  for (μ in basis.contractedBFs)
    μIndex+=1
    for (νIndex in 1:μIndex)
      ν = basis.contractedBFs[νIndex]
      ERImatrix = [computeElectronRepulsionIntegral(μ,λ,ν,σ) for λ in basis.contractedBFs, σ in basis.contractedBFs]
      K[μIndex,νIndex] = trace(ERImatrix*density)
      K[νIndex,μIndex] = K[μIndex,νIndex]
    end
  end
  return K
end

function computeMatrixExchange(electronRepulsionIntegrals::Array{Float64,4}, density::Matrix)
  @tensor K[μIndex,νIndex] := electronRepulsionIntegrals[μIndex,λIndex,νIndex,σIndex] * density[λIndex,σIndex]
end

function computeMatrixExchange(shells::Vector{Shell}, density)
  totaldim = mapreduce(ShellModule.nbf,+,0,shells)
  contractTensorBlocksWithMatrix_ExchangeLike(shells,density,totaldim,computeTensorBlockElectronRepulsionIntegrals,ShellModule.nbf)
end

function computeMatrixExchange(shells::Vector{LibInt2Shell}, density::Matrix)
  totaldim, maxprims, maxlqn = computeDimensions(shells)
  engine = LibInt2EngineCoulomb(maxprims,LQuantumNumber(maxlqn))
  result = contractTensorBlocksWithMatrix_ExchangeLike(shells,density,totaldim,(μ,ν,λ,σ)->computeTensorBlockElectronRepulsionIntegrals(engine,μ,ν,λ,σ),LibInt2Module.nbf)
  destroy!(engine)
  return result
end

#FOCK
function computeMatrixFock(
  basis::GaussianBasis,
  geometry::Geometry,
  density::Matrix,
	computeMatrixCoulomb=computeMatrixCoulomb,
	computeMatrixExchange=computeMatrixExchange)
  T = computeMatrixKinetic(basis)
  V = computeMatrixNuclearAttraction(basis,geometry)
  J = computeMatrixCoulomb(basis,density)
  K = computeMatrixExchange(basis,density)
  return T+V+2J-K
end

function computeMatrixFock(
  density::Matrix,
  matrixKinetic::Matrix,
  matrixNuclearAttraction::Matrix,
  electronRepulsionIntegralsColoumb::Array{Float64,4},
  electronRepulsionIntegralsExchange::Array{Float64,4}=electronRepulsionIntegralsColoumb)
  T = matrixKinetic
  V = matrixNuclearAttraction
  J = computeMatrixCoulomb(electronRepulsionIntegralsColoumb,density)
  K = computeMatrixExchange(electronRepulsionIntegralsExchange,density)
  return T+V+2J-K
end

end


using ReservoirComputing

#model parameters
res_size = 30
radius = 1.2
activation = tanh
degree = 6
sparsity = 0.5
sigma = 0.1
beta = 0.0
alpha = 1.0
nla_type = NLADefault()
in_size = 3
extended_states = false
h_steps = 2

W = init_reservoir_givendeg(res_size, radius, degree)

train_len = 50
predict_len = 12
data = ones(Float64, in_size, 100)
train = data[:, 1:1+train_len-1]
test = data[:, train_len:train_len+predict_len-1]

#test input layers functions
input_layer = [init_input_layer, init_dense_input_layer, init_sparse_input_layer, min_complex_input, irrational_sign_input]

for t in input_layer

    if t == init_sparse_input_layer
        W_in = t(res_size, in_size, sigma, sparsity)
    else
        W_in = t(res_size, in_size, sigma)
    end
    esn = ESN(W,
        train,
        W_in,
        activation = activation,
        alpha = alpha,
        nla_type = nla_type,
        extended_states = extended_states)

    @test size(W_in, 1) == res_size
    @test size(W_in, 2) == in_size
end


#= -------------------------------------------------------------------------
# @file problema2_tiempo_60.jl
#
# @date 12/07/17 18:00:56
# @author Martin Noblia
# @email mnoblia@disroot.org
#
# @brief
#
# @detail
#
#  Licence:
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or (at
# your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
---------------------------------------------------------------------------=#
using DifferentialEquations
#=------------------------------------------------------------------------------
                              constants
------------------------------------------------------------------------------=#
const  A      =  2.0
const  R      =  2.390057361376673e-5
const  cₑ     =  4.184e6
const  V      =  0.5
const  θ_ref  =  333.0
const  θ₀     =  288.0
const  θ_a    =  293.0
G = (θ_ref - θ_a) / (R * A)
tspan = (0.0, 1000.0)
θ_dot(t, θ) = ((G * A) / (cₑ * V)) + ((θ_a / (R * cₑ * V))) - ((θ / (R * cₑ * V)))
prob = ODEProblem(θ_dot, θ₀, tspan)
condition(t,θ,integrator) = (θ - 333.0)
affect!(integrator) = terminate!(integrator)
cb = ContinuousCallback(condition,affect!)
sol = solve(prob, callback=cb)


using Dice
using Dice: num_flips, num_nodes

for n = 7:7
    code = @dice begin

        function uniform(bitrange::Int, bitwidth::Int)
            ans = fill(DistBool(dicecontext(), false), bitwidth)
            for i = 1:bitrange
                ans[i] = flip(0.5)
            end
            DistInt(ans)
        end

        function gcd(a::DistInt, b::DistInt)
            for _ = 1 : 1 + max_bits(b) ÷ log2(MathConstants.golden)
                t = b
                converged = prob_equals(b, 0)
                amb, _ = (a % b)
                b = if converged
                        b
                    else
                        amb
                    end
                a = if converged
                        a
                    else
                        t
                    end
            end
            return a
        end

        a, _ = uniform(n, n+1) + 1
        b, _ = uniform(n, n+1) + 1
        g = gcd(a, b)
        prob_equals(g, 1)
    end

    @show n
    @show 2^n

    bdd = compile(code)
    @show num_flips(bdd)
    @show num_nodes(bdd)

    @time @show infer(code, :bdd)
end


"""   `calc_basis_scores_scalpels(rvs, ccfs; σ_rvs, num_basis )`
Compute the CCF basis functions and scores for a Scalpels reconstruction of CCFs
Inputs:
- rvs:  vector of estimated radial velocities
- ccfs: 2d array of CCFS of size (num_velocities, num_spectra)
Optional Inputs:
- σ_rvs:  vector of measurement uncertainties for estimated radial velocities (default: ones)
- num_basis:  number of basis vectors to use for SVD reconstruction of CCFs
- sort_by_responce:  set true to sort basis vectors by RV responce (default: true)
Output:
(scores, basis): a NamedTuple

Notes:
- Currently Scalpels weights all velocity pixels equally and uses the σ_rvs to weight each observation.
- First element of output starts with RMS with zero basis vectors (i.e., the input RVs)
"""
function calc_basis_scores_scalpels(rvs::AbstractVector{T1}, ccfs::AbstractArray{T2,2}
                ; σ_rvs::AbstractVector{T3} = ones(length(rvs)),
                num_basis::Integer = 3, assume_centered::Bool = false, sort_by_responce::Bool = true ) where { T1<:Real, T2<:Real, T3<:Real }
    @assert length(rvs) == length(σ_rvs)
    @assert length(rvs) == size(ccfs,2)
    @assert 1 <= num_basis < length(rvs)
    if assume_centered
        rvs_centered = rvs
    else
        mean_rv = mean(rvs, weights(1.0 ./ σ_rvs.^2 ))
        rvs_centered = rvs .- mean_rv
    end

    acfs = autocor(ccfs,0:size(ccfs,1)-1)
    if all(σ_rvs.==first(σ_rvs))
        acfs_minus_mean = acfs .- mean(acfs,dims=2)
    else
        acfs_minus_mean = acfs .- mean(acfs,weights(1.0 ./ σ_rvs.^2),dims=2)
    end
    #Δv_grid = convert(Float64,v_grid.step).*(0:size(acfs,1)-1)
    svd_acfs = svd(acfs_minus_mean')
    #return (basis=svd_acfs.U, scores=svd_acfs.S)  For ACF's SVD eigenvectors and eigenvalues

    alpha = svd_acfs.U'*rvs_centered
    if sort_by_responce
        idx = sortperm(abs.(alpha),rev=true)
        #println("# Ordered by alphas: ", idx[1:num_basis])
    else
        # To throw away sorting by projection onto RVs
        idx = 1:length(alpha)
    end
    U_keep = view(svd_acfs.U,:,idx[1:num_basis])
    P_keep = view(svd_acfs.Vt',:,idx[1:num_basis])
    return (scores=U_keep, basis=P_keep) # , responce=alpha[idx[1:num_basis]])

    #U_keep = view(svd_acfs.U,:,1:num_basis)
    #return (basis=U_keep, scores = view(alpha,1:num_basis) )
    #=  To clean RVs
    Δrv_shape = U_keep*U_keep'*rvs_centered
    rvs_clean = rvs .- Δrv_shape
    return rvs_clean
    =#
end

"""   `clean_rvs_scalpels(rvs, ccfs; σ_rvs, num_basis )`
Inputs:
- rvs:  vector of estimated radial velocities
- ccfs: 2d array of CCFS of size (num_velocities, num_spectra)
Optional Inputs:
- σ_rvs:  vector of measurement uncertainties for estimated radial velocities (default: ones)
- num_basis:  number of basis vectors to use for SVD reconstruction of CCFs
- sort_by_responce:  set true to sort basis vectors by RV responce (default: true)
Output:
rvs_clean: vector of estimated RVs after cleaning by scalpels

Notes:
- Currently Scalpels weights all observations equally and doesn't use the σ_rvs.
- First element of output starts with RMS with zero basis vectors (i.e., the input RVs)
"""
function clean_rvs_scalpels(rvs::AbstractVector{T1}, ccfs::AbstractArray{T2,2}
                ; σ_rvs::AbstractVector{T3} = ones(length(rvs)),
                num_basis::Integer = 3, sort_by_responce::Bool = true ) where { T1<:Real, T2<:Real, T3<:Real }
    @assert length(rvs) == length(σ_rvs)
    @assert length(rvs) == size(ccfs,2)
    if num_basis == 0   return rvs   end
    @assert 0 <= num_basis < length(rvs)
    mean_rv = mean(rvs, weights(1.0 ./ σ_rvs.^2 ))
    rvs_centered = rvs .- mean_rv

    #=
    acfs = autocor(ccfs,0:size(ccfs,1)-1)
    acfs_minus_mean = acfs .- mean(acfs,dims=2)
    #Δv_grid = convert(Float64,v_grid.step).*(0:size(acfs,1)-1)
    svd_acfs = svd(acfs_minus_mean')
    alpha = svd_acfs.U'*rvs_centered

    idx = sortperm(abs.(alpha),rev=true)
    U_keep = view(svd_acfs.U,:,idx[1:num_basis])

    #(U, scores ) = calc_basis_scores_scalpels(rvs_centered, ccfs, σ_rvs=σ_rvs, num_basis=num_basis, assume_centered=true)
    idx = sortperm(abs.(alpha),rev=true)
    alpha = U'*rvs_centered
    U_keep = view(U,:,idx[1:num_basis])
    =#

    (U_keep, ) = calc_basis_scores_scalpels(rvs_centered, ccfs, σ_rvs=σ_rvs, num_basis=num_basis, assume_centered=true, sort_by_responce=sort_by_responce)
    Δrv_shape = U_keep*U_keep'*rvs_centered
    rvs_clean = rvs .- Δrv_shape
    return rvs_clean
end

"""   `rms_clean_rvs_vs_num_basis_scalpels(rvs, ccfs; σ_rvs, max_num_basis )`
Compute RMS of estimated RVs after cleaning raw RVS with Scalpels algorithm (based on CCFs)

Inputs:
- rvs:  vector of estimated radial velocities
- ccfs: 2d array of CCFS of size (num_velocities, num_spectra)
Optional Inputs:
- σ_rvs:  vector of measurement uncertainties for estimated radial velocities (default: ones)
- max_num_basis:  maximum number of basis vectors to use for SVD reconstruction of CCFs
- sort_by_responce:  set true to sort basis vectors by RV responce (default: true)
Output:
rms_scalpels: vector of RMS estimated RVs after cleaning by scalpels as a function of the number of basis vectors

Notes:
- Currently Scalpels weights all observations equally and doesn't use the σ_rvs.
- First element of output starts with RMS with zero basis vectors (i.e., the input RVs)
"""
function rms_clean_rvs_vs_num_basis_scalpels(rvs::AbstractVector{T1}, ccfs::AbstractArray{T2,2}
                ; σ_rvs::AbstractVector{T3} = ones(length(rvs)),
                  max_num_basis::Integer = min(length(rvs)-1,default_max_num_basis),
                  sort_by_responce::Bool = true ) where { T1<:Real, T2<:Real, T3<:Real }
    rms_scalpels = map(b->std(clean_rvs_scalpels(rvs,ccfs,num_basis=b, sort_by_responce=sort_by_responce)), 0:max_num_basis)
end

"""   `calc_clean_rvs_scores_basis_scalpels(rvs, ccfs; σ_rvs, num_basis, sort_by_responce )`
Computes cleaned rvs as well as the CCF basis functions and scores for a Scalpels reconstruction of CCFs
Inputs:
- rvs:  vector of estimated radial velocities
- ccfs: 2d array of CCFS of size (num_velocities, num_spectra)
Optional Inputs:
- σ_rvs:  vector of measurement uncertainties for estimated radial velocities (default: ones)
- num_basis:  number of basis vectors to use for SVD reconstruction of CCFs
- sort_by_responce:  set true to sort basis vectors by RV responce (default: true)
Output:
(rvs, scores, basis): a NamedTuple

Notes:
- Currently Scalpels weights all velocity pixels equally and uses the σ_rvs to weight each observation.
- First element of output starts with RMS with zero basis vectors (i.e., the input RVs)
"""
function calc_clean_rvs_scores_basis_scalpels(rvs::AbstractVector{T1}, ccfs::AbstractArray{T2,2}
                ; σ_rvs::AbstractVector{T3} = ones(length(rvs)),
                num_basis::Integer = 3,
                sort_by_responce::Bool = true ) where { T1<:Real, T2<:Real, T3<:Real }
    @assert length(rvs) == length(σ_rvs)
    @assert length(rvs) == size(ccfs,2)
    if num_basis == 0   return rvs   end
    @assert 0 <= num_basis < length(rvs)
    mean_rv = mean(rvs, weights(1.0 ./ σ_rvs.^2 ))
    rvs_centered = rvs .- mean_rv

    (U_keep, basis ) = calc_basis_scores_scalpels(rvs_centered, ccfs, σ_rvs=σ_rvs, num_basis=num_basis, assume_centered=true, sort_by_responce=sort_by_responce)
    Δrv_shape = U_keep*U_keep'*rvs_centered
    rvs_clean = rvs .- Δrv_shape

    return (rvs=rvs_clean, scores=U_keep, basis=basis )
end


using DataFrames, JuMP, CPLEX, Dates
include("utils/extract_data.jl")
include("utils/display_results.jl")
include("utils/misc.jl")
include("single_depot/solve_single_depot.jl")
include("multi_depot/solve_multi_depot.jl")

function solve_problem(file, nb_vehicles, approx, verbose)

    print("\n\n")
    println("===========================================")
    println("GROOMER OPTIMIZATION PROBLEM")
    println("-----------")

    println("Extracting data from file ...")
    (data, graph) = extract_data(file)
    data[1, :nb_vehicles] = nb_vehicles

    println("-----------")

    print_info_dataset(data)

    println("-----------")

    if length(data[1, :depots])>1
        (tours, solving_time) = solve_multi_depot(data, graph, approx, verbose)
    elseif length(data[1, :depots])==1
        (tours, solving_time) = solve_single_depot(data, graph)
    else
        error("Problem with depot list. Check dataset's validity.")
    end

    println("-----------")

    display_results(tours, graph, data, solving_time)

    println("===========================================")
end


using Oceananigans.AbstractOperations: UnaryOperation, Derivative, BinaryOperation, MultiaryOperation
using Oceananigans.AbstractOperations: KernelFunctionOperation
using Oceananigans.Operators: ℑxyᶜᶠᵃ, ℑxyᶠᶜᵃ
using Oceananigans.Fields: PressureField, compute_at!
using Oceananigans.BuoyancyModels: BuoyancyField

function compute_derivative(model, ∂)
    T, S = model.tracers
    S.data.parent .= π
    @compute ∂S = ComputedField(∂(S))
    result = Array(interior(∂S))
    return all(result .≈ zero(eltype(model.grid)))
end

function compute_unary(unary, model)
    set!(model; S=π)
    T, S = model.tracers
    @compute uS = ComputedField(unary(S), data=model.pressures.pHY′.data)
    result = Array(interior(uS))
    return all(result .≈ unary(eltype(model.grid)(π)))
end

function compute_plus(model)
    set!(model; S=π, T=42)
    T, S = model.tracers
    @compute ST = ComputedField(S + T, data=model.pressures.pHY′.data)
    result = Array(interior(ST))
    return all(result .≈ eltype(model.grid)(π+42))
end

function compute_many_plus(model)
    set!(model; u=2, S=π, T=42)
    T, S = model.tracers
    u, v, w = model.velocities
    @compute uTS = ComputedField(@at((Center, Center, Center), u + T + S))
    result = Array(interior(uTS))
    return all(result .≈ eltype(model.grid)(2+π+42))
end

function compute_minus(model)
    set!(model; S=π, T=42)
    T, S = model.tracers
    @compute ST = ComputedField(S - T, data=model.pressures.pHY′.data)
    result = Array(interior(ST))
    return all(result .≈ eltype(model.grid)(π-42))
end

function compute_times(model)
    set!(model; S=π, T=42)
    T, S = model.tracers
    @compute ST = ComputedField(S * T, data=model.pressures.pHY′.data)
    result = Array(interior(ST))
    return all(result .≈ eltype(model.grid)(π*42))
end

function compute_kinetic_energy(model)
    u, v, w = model.velocities
    set!(u, 1)
    set!(v, 2)
    set!(w, 3)

    kinetic_energy_operation = @at (Center, Center, Center) (u^2 + v^2 + w^2) / 2
    @compute kinetic_energy = ComputedField(kinetic_energy_operation, data=model.pressures.pHY′.data)
    result = Array(interior(kinetic_energy))

    return all(result .≈ 7)
end

function horizontal_average_of_plus(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = sin(π * z)
    T₀(x, y, z) = 42 * z
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S + T, dims=(1, 2))

    zC = znodes(Center, model.grid)
    correct_profile = @. sin(π * zC) + 42 * zC

    result = Array(interior(ST))[:]

    return all(result .≈ correct_profile)
end

function zonal_average_of_plus(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = sin(π*z) * sin(π*y)
    T₀(x, y, z) = 42*z + y^2
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S + T, dims=1)

    yC = ynodes(Center, model.grid, reshape=true)
    zC = znodes(Center, model.grid, reshape=true)
    correct_slice = @. sin(π * zC) * sin(π * yC) + 42*zC + yC^2

    result = Array(interior(ST))

    return all(result[1, :, :] .≈ correct_slice[1, :, :])
end

function volume_average_of_times(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = 1 + sin(2π*x)
    T₀(x, y, z) = y
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S * T, dims=(1, 2, 3))

    result = Array(interior(ST))

    return all(result[1, 1, 1] .≈ 0.5)
end

function horizontal_average_of_minus(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = sin(π*z)
    T₀(x, y, z) = 42*z
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S - T, dims=(1, 2))

    zC = znodes(Center, model.grid)
    correct_profile = @. sin(π * zC) - 42 * zC

    result = Array(interior(ST))

    return all(result[1, 1, 1:Nz] .≈ correct_profile)
end

function horizontal_average_of_times(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    S₀(x, y, z) = sin(π*z)
    T₀(x, y, z) = 42*z
    set!(model; S=S₀, T=T₀)
    T, S = model.tracers

    @compute ST = AveragedField(S * T, dims=(1, 2))

    zC = znodes(Center, model.grid)
    correct_profile = @. sin(π * zC) * 42 * zC

    result = Array(interior(ST))

    return all(result[1, 1, 1:Nz] .≈ correct_profile)
end

function multiplication_and_derivative_ccf(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    w₀(x, y, z) = sin(π * z)
    T₀(x, y, z) = 42 * z
    set!(model; enforce_incompressibility=false, w=w₀, T=T₀)

    w = model.velocities.w
    T = model.tracers.T

    @compute wT = AveragedField(w * ∂z(T), dims=(1, 2))

    zF = znodes(Face, model.grid)
    correct_profile = @. 42 * sin(π * zF)

    result = Array(interior(wT))

    # Omit both halos and boundary points
    return all(result[1, 1, 2:Nz] .≈ correct_profile[2:Nz])
end

const C = Center
const F = Face

function multiplication_and_derivative_ccc(model)
    Ny, Nz = model.grid.Ny, model.grid.Nz

    w₀(x, y, z) = sin(π*z)
    T₀(x, y, z) = 42*z
    set!(model; enforce_incompressibility=false, w=w₀, T=T₀)

    w = model.velocities.w
    T = model.tracers.T

    wT_ccc = @at (C, C, C) w * ∂z(T)
    @compute wT_ccc_avg = AveragedField(wT_ccc, dims=(1, 2))

    zF = znodes(Face, model.grid)
    sinusoid = sin.(π * zF)
    interped_sin = [(sinusoid[k] + sinusoid[k+1]) / 2 for k in 1:model.grid.Nz]
    correct_profile = interped_sin .* 42

    result = Array(interior(wT_ccc_avg))

    # Omit boundary-adjacent points from comparison
    return all(result[1, 1, 2:Nz-1] .≈ correct_profile[2:Nz-1])
end

function computation_including_boundaries(arch)
    topo = (Periodic, Bounded, Bounded)
    grid = RegularRectilinearGrid(topology=topo, size=(13, 17, 19), extent=(1, 1, 1))
    model = NonhydrostaticModel(architecture=arch, grid=grid)

    u, v, w = model.velocities
    @. u.data = 1 + rand()
    @. v.data = 2 + rand()
    @. w.data = 3 + rand()

    op = @at (Center, Face, Face) u * v * w
    @compute uvw = ComputedField(op)

    return all(interior(uvw) .!= 0)
end

function operations_with_computed_field(model)
    u, v, w = model.velocities
    uv = ComputedField(u * v)
    @compute uvw = ComputedField(uv * w)
    return true
end

function operations_with_averaged_field(model)
    u, v, w = model.velocities
    UV = AveragedField(u * v, dims=(1, 2))
    wUV = ComputedField(w * UV)
    compute!(wUV)
    return true
end

function pressure_field(model)
    p = PressureField(model)
    u, v, w = model.velocities
    @compute up = ComputedField(u * p)
    return true
end

function computations_with_buoyancy_field(arch, buoyancy)
    grid = RegularRectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))
    tracers = buoyancy isa BuoyancyTracer ? :b : (:T, :S)
    model = NonhydrostaticModel(architecture=arch, grid=grid,
                                tracers=tracers, buoyancy=buoyancy)

    b = BuoyancyField(model)
    u, v, w = model.velocities

    compute!(b)

    ub = ComputedField(b * u)
    vb = ComputedField(b * v)
    wb = ComputedField(b * w)

    compute!(ub)
    compute!(vb)
    compute!(wb)

    return true # test that it doesn't error
end

function computations_with_averaged_fields(model)
    u, v, w, T, S = fields(model)

    set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3)

    # Two ways to compute turbulent kinetic energy
    U = AveragedField(u, dims=(1, 2))
    V = AveragedField(v, dims=(1, 2))

    tke_op = @at (Center, Center, Center) ((u - U)^2  + (v - V)^2 + w^2) / 2
    tke = ComputedField(tke_op)
    compute!(tke)

    return all(interior(tke)[2:3, 2:3, 2:3] .== 9/2)
end

function computations_with_averaged_field_derivative(model)

    set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3)

    u, v, w, T, S = fields(model)

    # Two ways to compute turbulent kinetic energy
    U = AveragedField(u, dims=(1, 2))
    V = AveragedField(v, dims=(1, 2))

    # This tests a vertical derivative of an AveragedField
    shear_production_op = @at (Center, Center, Center) u * w * ∂z(U)
    shear = ComputedField(shear_production_op)
    compute!(shear)

    set!(model, T = (x, y, z) -> 3 * z)

    return all(interior(shear)[2:3, 2:3, 2:3] .== interior(T)[2:3, 2:3, 2:3])
end

function computations_with_computed_fields(model)
    u, v, w, T, S = fields(model)

    set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3)

    # Two ways to compute turbulent kinetic energy
    U = AveragedField(u, dims=(1, 2))
    V = AveragedField(v, dims=(1, 2))

    u′ = ComputedField(u - U)
    v′ = ComputedField(v - V)

    tke_op = @at (Center, Center, Center) (u′^2  + v′^2 + w^2) / 2
    tke = ComputedField(tke_op)
    compute!(tke)

    return all(interior(tke)[2:3, 2:3, 2:3] .== 9/2)
end

for arch in archs
    @testset "ComputedFields [$(typeof(arch))]" begin
        @info "  Testing ComputedFields [$(typeof(arch))]..."

        grid = RegularRectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1),
                                      topology=(Periodic, Periodic, Bounded))

        buoyancy = SeawaterBuoyancy(gravitational_acceleration = 1,
                                             equation_of_state = LinearEquationOfState(α=1, β=1))

        model = NonhydrostaticModel(architecture = arch,
                                            grid = grid,
                                        buoyancy = buoyancy)

        @testset "Derivative computations [$(typeof(arch))]" begin
            @info "      Testing compute! derivatives..."
            @test compute_derivative(model, ∂x)
            @test compute_derivative(model, ∂y)
            @test compute_derivative(model, ∂z)
        end

        @testset "Unary computations [$(typeof(arch))]" begin
            @info "      Testing compute! unary operations..."
            for unary in (sqrt, sin, cos, exp, tanh)
                @test compute_unary(unary, model)
            end
        end

        @testset "Binary computations [$(typeof(arch))]" begin
            @info "      Testing compute! binary operations..."
            @test compute_plus(model)
            @test compute_minus(model)
            @test compute_times(model)

            # Basic compilation test for nested BinaryOperations...
            u, v, w = model.velocities
            @test try compute!(ComputedField(u + v - w)); true; catch; false; end
        end

        @testset "Multiary computations [$(typeof(arch))]" begin
            @info "      Testing compute! multiary operations..."
            @test compute_many_plus(model)

            @info "      Testing compute! kinetic energy..."
            @test compute_kinetic_energy(model)
        end

        @testset "Computations with KernelComputedField [$(typeof(arch))]" begin
            @test begin
                @inline trivial_kernel_function(i, j, k, grid) = 1
                op = KernelFunctionOperation{Center, Center, Center}(trivial_kernel_function, grid)
                f = ComputedField(op, arch)
                compute!(f)
                f isa ComputedField && f.operand === op
            end

            @test begin
                @inline trivial_parameterized_kernel_function(i, j, k, grid, μ) = μ
                op = KernelFunctionOperation{Center, Center, Center}(trivial_parameterized_kernel_function, grid, parameters=0.1)
                f = ComputedField(op, arch)
                compute!(f)
                f isa ComputedField && f.operand === op
            end

            @test begin
                u, v, w = model.velocities
                ζ_op = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᵃ, grid, computed_dependencies=(u, v))
                ζ = ComputedField(ζ_op) # identical to `VerticalVorticityField`
                compute!(ζ)
                ζ isa ComputedField && ζ.operand.kernel_function === ζ₃ᶠᶠᵃ
            end
        end

        @testset "Operations with ComputedField and PressureField [$(typeof(arch))]" begin
            @info "      Testing operations with ComputedField..."
            @test operations_with_computed_field(model)

            @info "      Testing PressureField..."
            @test pressure_field(model)
        end

        @testset "Horizontal averages of operations [$(typeof(arch))]" begin
            @info "      Testing horizontal averges..."
            @test horizontal_average_of_plus(model)
            @test horizontal_average_of_minus(model)
            @test horizontal_average_of_times(model)

            @test multiplication_and_derivative_ccf(model)
            @test multiplication_and_derivative_ccc(model)
        end

        @testset "Zonal averages of operations [$(typeof(arch))]" begin
            @info "      Testing zonal averges..."
            @test zonal_average_of_plus(model)
        end

        @testset "Volume averages of operations [$(typeof(arch))]" begin
            @info "      Testing volume averges..."
            @test volume_average_of_times(model)
        end

        @testset "ComputedField boundary conditions [$(typeof(arch))]" begin
            @info "      Testing boundary conditions for ComputedField..."

            set!(model; S=π, T=42)
            T, S = model.tracers

            @compute ST = ComputedField(S + T, data=model.pressures.pHY′.data)

            Nx, Ny, Nz = size(model.grid)

            @test all(ST.data[0, 1:Ny, 1:Nz]  .== ST.data[Nx+1, 1:Ny, 1:Nz])
            @test all(ST.data[1:Nx, 0, 1:Nz]  .== ST.data[1:Nx, Ny+1, 1:Nz])
            @test all(ST.data[1:Nx, 1:Ny, 0]  .== ST.data[1:Nx, 1:Ny, 1])
            @test all(ST.data[1:Nx, 1:Ny, Nz] .== ST.data[1:Nx, 1:Ny, Nz+1])

            @compute ST_face = ComputedField(@at (Center, Center, Face) S * T)

            @test all(ST_face.data[1:Nx, 1:Ny, 0] .== 0)
            @test all(ST_face.data[1:Nx, 1:Ny, Nz+2] .== 0)
        end

        @testset "Operations with AveragedField [$(typeof(arch))]" begin
            @info "      Testing operations with AveragedField..."

            T, S = model.tracers
            TS = AveragedField(T * S, dims=(1, 2))
            @test operations_with_averaged_field(model)
        end

        @testset "Compute! on faces along bounded dimensions" begin
            @info "      Testing compute! on faces along bounded dimensions..."
            @test computation_including_boundaries(arch)
        end

        EquationsOfState = (LinearEquationOfState, SeawaterPolynomials.RoquetEquationOfState,
                            SeawaterPolynomials.TEOS10EquationOfState)

        buoyancies = (BuoyancyTracer(), SeawaterBuoyancy(),
                      (SeawaterBuoyancy(equation_of_state=eos()) for eos in EquationsOfState)...)

        for buoyancy in buoyancies
            @testset "Computations with BuoyancyFields [$(typeof(arch)), $(typeof(buoyancy).name.wrapper)]" begin
                @info "      Testing computations with BuoyancyField " *
                      "[$(typeof(arch)), $(typeof(buoyancy).name.wrapper)]..."

                @test computations_with_buoyancy_field(arch, buoyancy)
            end
        end

        @testset "Computations with AveragedFields [$(typeof(arch))]" begin
            @info "      Testing computations with AveragedField [$(typeof(arch))]..."

            @test computations_with_averaged_field_derivative(model)

            u, v, w = model.velocities

            set!(model, enforce_incompressibility = false, u = (x, y, z) -> z, v = 2, w = 3)

            # Two ways to compute turbulent kinetic energy
            U = AveragedField(u, dims=(1, 2))
            V = AveragedField(v, dims=(1, 2))

            # Build up compilation tests incrementally...
            u_prime              = u - U
            u_prime_ccc          = @at (Center, Center, Center) u - U
            u_prime_squared      = (u - U)^2
            u_prime_squared_ccc  = @at (Center, Center, Center) (u - U)^2
            horizontal_twice_tke = (u - U)^2 + (v - V)^2
            horizontal_tke       = ((u - U)^2 + (v - V)^2) / 2
            horizontal_tke_ccc   = @at (Center, Center, Center) ((u - U)^2 + (v - V)^2) / 2
            twice_tke            = (u - U)^2  + (v - V)^2 + w^2
            tke                  = ((u - U)^2  + (v - V)^2 + w^2) / 2
            tke_ccc              = @at (Center, Center, Center) ((u - U)^2  + (v - V)^2 + w^2) / 2

            @test try compute!(ComputedField(u_prime             )); true; catch; false; end
            @test try compute!(ComputedField(u_prime_ccc         )); true; catch; false; end
            @test try compute!(ComputedField(u_prime_squared     )); true; catch; false; end
            @test try compute!(ComputedField(u_prime_squared_ccc )); true; catch; false; end
            @test try compute!(ComputedField(horizontal_twice_tke)); true; catch; false; end
            @test try compute!(ComputedField(horizontal_tke      )); true; catch; false; end
            @test try compute!(ComputedField(twice_tke           )); true; catch; false; end

            @test try compute!(ComputedField(horizontal_tke_ccc  )); true; catch; false; end
            @test try compute!(ComputedField(tke                 )); true; catch; false; end
            @test try compute!(ComputedField(tke_ccc             )); true; catch; false; end

            computed_tke = ComputedField(tke_ccc)
            @test (compute!(computed_tke); all(interior(computed_tke)[2:3, 2:3, 2:3] .== 9/2))
        end

        @testset "Computations with ComputedFields [$(typeof(arch))]" begin
            @info "      Testing computations with ComputedField [$(typeof(arch))]..."
            @test computations_with_computed_fields(model)
        end

        @testset "Conditional computation of ComputedField and BuoyancyField [$(typeof(arch))]" begin
            @info "      Testing conditional computation of ComputedField and BuoyancyField " *
                  "[$(typeof(arch))]..."

            set!(model, u=2, v=0, w=0, T=3, S=0)
            u, v, w, T, S = fields(model)

            uT = ComputedField(u * T)

            α = model.buoyancy.model.equation_of_state.α
            g = model.buoyancy.model.gravitational_acceleration
            b = BuoyancyField(model)

            compute_at!(uT, 1.0)
            compute_at!(b, 1.0)
            @test all(interior(uT) .== 6)
            @test all(interior(b) .== g * α * 3)

            set!(model, u=2, T=4)
            compute_at!(uT, 1.0)
            compute_at!(b, 1.0)
            @test all(interior(uT) .== 6)
            @test all(interior(b) .== g * α * 3)

            compute_at!(uT, 2.0)
            compute_at!(b, 2.0)
            @test all(interior(uT) .== 8)
            @test all(interior(b) .== g * α * 4)
        end
    end
end


using NeuralQuantum
using Test

N = 4
@testset "Hilbert spaces " begin
    hilb = HomogeneousSpin(N)
    s = state(hilb)
    vals = Int[]
    for i=1:spacedimension(hilb)
        set!(s, hilb, i)
        push!(vals, toint(s, hilb))
    end

    @test all(vals .== 1:spacedimension(hilb))

    shilb = SuperOpSpace(hilb)
    @test physical(shilb) === hilb

    s = state(shilb)
    vals = Int[]
    for i=1:spacedimension(shilb)
        set!(s, shilb, i)
        push!(vals, toint(s, shilb))
    end
    @test all(vals .== 1:spacedimension(shilb))
end

function test_hilb_iterator(hilb)
    it = states(hilb)

    all_states = collect(it)

    @test length(all_states) == spacedimension(hilb)
    @test length(it) == spacedimension(hilb)
    @test eltype(it) == typeof(state(hilb))

    vals = zeros(length(all_states))
    for (i,v) = enumerate(all_states)
        vals[i] = index(hilb, v)
    end

    @test all(vals .== 1:spacedimension(hilb))
end

@testset "Hilbert space iterator" begin
    test_hilb_iterator(HomogeneousSpin(N))
    test_hilb_iterator(HomogeneousFock(N,2))
    test_hilb_iterator(SuperOpSpace(HomogeneousSpin(N)))
end


using Test
using StaticArrays
using PolygonArea
using PolygonArea: PolarHalfPlane, Reunion, Point, vertices, complement, rotate, translate, scale, square

@testset "Polygons" begin

    @testset "creation and contained points" begin
        unit_square = square((0.0, 0.0), 1.0)
        @test Point(0.5, 0.5) in unit_square
        @test Point(1.0, 1.0) in unit_square
        @test !(Point(1.2, 1.0) in unit_square)
        @test Point(0.0, 0.0) in vertices(unit_square)
        @test Point(1.0, 0.0) in vertices(unit_square)
        @test !(Point(0.5, 0.5) in vertices(unit_square))

        @test unit_square == rectangle(bottom_left=(0.0, 0.0), top_right=(1.0, 1.0))
        @test unit_square == rectangle(bottom_left=(0.0, 0.0), top_right=(1.0, 1.0))
        @test unit_square == rectangle((0.0, 0.0), (1.0, 1.0))

        @test !(PolygonArea.ConvexPolygon{Int}([]) == square((1, 1), 1))

        rec = PolygonArea.ConvexPolygon([(0.0, 0.0), (0.0, 1.0), (1.0, 1.0), (1.0, 0.0)])
        @test (0.5, 0.5) ∈ rec

        # Should return an error...
        inv_rec = PolygonArea.ConvexPolygon([(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0)])
        @test !((0.5, 0.5) ∈ inv_rec)
    end

    @testset "properties" begin
        @test isempty(PolygonArea.ConvexPolygon{Int}([]))

        @test PolygonArea.nb_sides(rectangle((0, 0), (2, 2))) == 4

        @test area(PolygonArea.square((0.0, 0.0), 1.0)) == 1.0
        @test area(rectangle((0, 0), (4, 3))) == 12
        @test isapprox(area(circle(Point(0.0, 0.0), 1.0, 100)), π, atol=1e3)
    end

    @testset "transformations" begin
        @test Point(0.0, 1.2) in rotate(rectangle(0.0, 0.0, 1.0, 1.0), π/4)
        @test Point(0.0, 1.2) in translate(rectangle(0.0, 0.0, 1.0, 1.0), SVector(-0.3, 0.5))
        @test Point(0.0, 1.2) in translate(rectangle(0.0, 0.0, 1.0, 1.0), [-0.3, 0.5])
        @test Point(0.0, 1.2) in translate(rectangle(0.0, 0.0, 1.0, 1.0), (-0.3, 0.5))
        @test Point(1.5, 1.5) in scale(rectangle(0.0, 0.0, 1.0, 1.0), 1.5)
        @test !(Point(1.5, 1.5) in scale(rectangle(0.0, 0.0, 1.0, 1.0), 1.5, center=(1.0, 1.0)))
        λ = rand(2)
        @test area(scale(rectangle(0.0, 0.0, 1.0, 1.0), λ)) ≈ prod(λ)*area(rectangle(0.0, 0.0, 1.0, 1.0))
    end

    @testset "union and intersection" begin
        # Intersection with HalfPlane
        unit_square = square((0.0, 0.0), 1.0)
        hp = HalfPlane(1.0, 1.0, -1.5)
        c1 = unit_square ∩ hp  # Cut the top-right corner
        @test Point(0.5, 0.5) in c1
        @test !(Point(0.9, 0.9) in c1)

        c2 = unit_square ∩ complement(hp)  # Keep only the top-right corner
        @test !(Point(0.5, 0.5) in c2)
        @test Point(0.9, 0.9) in c2

        @test area(c1) + area(c2) ≈ area(unit_square)

        # Intersection with union and intersection of HalfPlanes
        left = HalfPlane(1.0, 0.0, -0.5)
        top = HalfPlane(0.0, 1.0, -0.5)
        top_left = top ∩ left
        bottom_right = complement(top) ∩ complement(left)
        @test bottom_right == complement(top ∪ left)
        @test (unit_square ∩ left) ∩ top == unit_square ∩ (left ∩ top)

        @test area(unit_square ∩ left) == 0.5
        @test area((unit_square ∩ left) ∩ top) == 0.25
        @test area(unit_square ∩ (top ∪ left)) == 0.75

        @test area(unit_square ∩ (bottom_right ∪ top_left)) ≈ 0.5
        @test area(unit_square ∩ (bottom_right ∩ top_left)) ≈ 0.0

        @test (square((0.0, 0.0), 1.0) |> complement |> typeof) == Reunion{HalfPlane{Float64}}
        @test area(square((0.0, 0.0), 3.0) \ square((1.0, 1.0), 1.0)) == 8.0

        # Intersection with ConvexPolygon
        @test area(unit_square ∩ square((0.0, 0.5), 1.0)) == 0.5
        @test area(unit_square ∩ square((0.5, 0.5), 1.0)) == 0.25

        triangle = ConvexPolygon([(0.0, 0.0), (0.5, 1.0), (1.0, 0.0)])
        @test area(unit_square ∩ triangle) == area(triangle)

        # Intersection with union of ConvexPolygon
        r1 = rectangle(0.0, 0.0, 2.0, 2.0)
        r2 = rectangle(1.0, 1.0, 2.0, 2.0)
        r3 = rectangle(1.5, 1.5, 2.5, 2.5)
        @test area((unit_square ∪ r2) ∩ (r1 ∪ r3)) == 2.0

        # Intersection with circle
        c = circle(0.0, 0.0, 1.0, 100)
        r = rectangle(0.0, -1.0, 1.0, 1.0)
        @test area(c ∩ r) ≈ area(r ∩ c) ≈ area(c)/2

        # Disjoint union
        c = circle(0.0, 0.0, 1.0, 10)
        @test PolygonArea.disjoint(c) == c

        h = rectangle(0.0, 1/3, 1.0, 2/3)
        v = rectangle(1/3, 0.0, 2/3, 1.0)
        @test area(h ∪ v) ≈ 6/9
        @test area(PolygonArea.disjoint(h ∪ v)) ≈ 5/9

        # A floating point bug: intersection of a circle with a tangeant half plane
        c = circle((0.5, 0.5), 0.4, 60) |> PolygonArea.rotate(π/5, center=(0.5, 0.5))
        r = rectangle(0.45, 0.10, 0.46, 0.11)
        @test PolygonArea.area(r ∩ c) ≈ PolygonArea.area(c ∩ r)

        #
        horizontal_rectangle = PolygonArea.rectangle((0.4, 0.2), (0.6, 0.8))
        vertical_rectangle = PolygonArea.rectangle((0.2, 0.4), (0.8, 0.6))
        cross = horizontal_rectangle ∪ vertical_rectangle
        circ = PolygonArea.circle(0.5, 0.5, 0.35, 100)
        diff = circ \ cross
        @test area(diff) ≈ area(circ) - area(PolygonArea.disjoint(cross))
    end

    @testset "generic programming" begin
        using Unitful: m
        @test area(PolygonArea.ConvexPolygon([(1m, 0m), (1m, 1m), (0m, 1m), (0m, 0m)])) == 1m^2

        # Conversion and promotion
        T = Reunion{Intersection{HalfPlane{Float64}}}
        @test convert(T, HalfPlane(1f0, 1f0, 0f0)) isa T
        @test convert(T, rectangle(1f0, 1f0, 2f0, 2f0)) isa T

        @test (HalfPlane(1, 1, 1) ∩ HalfPlane(1f0, 2f0, 3f0)) isa Intersection{HalfPlane{Float32}}
        @test (rectangle(0, 0, 2, 2) ∩ rectangle(1f0, 1f0, 3f0, 3f0)) isa ConvexPolygon{Float32}
    end
end


export get_hessian

function get_hessian(params0::AbstractArray{T}) where T
    N = length(params0)
    params = Dual.(params0, zero(T))
    hes = zeros(T, N, N)
    for i=1:N
        @inbounds i !== 1 && (params[i-1] = Dual(params0[i-1], zero(T)))
        @inbounds params[i] = Dual(params0[i], one(T))
        res = NiLang.AD.gradient(Val(1), embedding_loss, (Dual(0.0, 0.0), params))[2]
        hes[:,i] .= vec(ForwardDiff.partials.(res, 1))
    end
    hes
end


# ---
# title: 629. K Inverse Pairs Array
# id: problem629
# author: Indigo
# date: 2021-06-22
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/k-inverse-pairs-array/description/>
# hidden: true
# ---
#
# Given two integers `n` and `k`, find how many different arrays consist of
# numbers from `1` to `n` such that there are exactly `k` inverse pairs.
#
# We define an inverse pair as following: For `ith` and `jth` element in the
# array, if `i` < `j` and `a[i]` > `a[j]` then it's an inverse pair; Otherwise,
# it's not.
#
# Since the answer may be very large, the answer should be modulo 109 \+ 7.
#
# **Example 1:**
#
#
#
#     Input: n = 3, k = 0
#     Output: 1
#     Explanation:
#     Only the array [1,2,3] which consists of numbers from 1 to 3 has exactly 0 inverse pair.
#
#
#
#
# **Example 2:**
#
#
#
#     Input: n = 3, k = 1
#     Output: 2
#     Explanation:
#     The array [1,3,2] and [2,1,3] have exactly 1 inverse pair.
#
#
#
#
# **Note:**
#
#   1. The integer `n` is in the range [1, 1000] and `k` is in the range [0, 1000].
#
#
#
#
## @lc code=start
using LeetCode

function k_inverse_pairs(n::Int, k::Int)
    dp = OffsetArray(fill(0, n + 1, k + 1), -1, -1)
    dp[:, 0] .= 1
    m = 1000000007
    @inbounds for i in 1:n, j in 1:min(k, i * (i - 1) ÷ 2)
        val = dp[i - 1, j] - ((j - i) >= 0 ? dp[i - 1, j - i] : 0)
        dp[i, j] = (dp[i, j - 1] + val) % m
    end
    return (dp[n, k] + m) % m
end
## @lc code=end


using Test

m = test_page_model()
include("../src/components/CO2forcing.jl")

add_comp!(m, co2forcing)

set_param!(m, :co2forcing, :c_CO2concentration, readpagedata(m,"test/validationdata/c_co2concentration.csv"))

##running Model
run(m)

forcing=m[:co2forcing,:f_CO2forcing]
forcing_compare=readpagedata(m,"test/validationdata/f_co2forcing.csv")

@test forcing ≈ forcing_compare rtol=1e-3


# This file is a part of KCenters.jl
# License is Apache 2.0: https://www.apache.org/licenses/LICENSE-2.0.txt

using Random

export AbstractCenterSelection,
    CentroidSelection,
    RandomCenterSelection,
    MedoidSelection,
    KnnCentroidSelection,
    center

"""
    abstract type AbstractCenterSelection end

Abstract type for all center selection strategies
"""
abstract type AbstractCenterSelection end

"""
    CentroidSelection()
    center(sel::CentroidSelection, lst::AbstractVector)

Computes a center as the centroid of the `lst` set of points
"""
struct CentroidSelection <: AbstractCenterSelection end

"""
    RandomCenterSelection()
    center(sel::RandomCenterSelection, lst::AbstractVector)

Selects a random element from `lst` as representing center
"""

struct RandomCenterSelection <: AbstractCenterSelection end

"""
    MedoidSelection(dist::PreMetric, ratio::Float32)
    MedoidSelection(; dist=SqL2Distance(), ratio=0.5) = MedoidSelection(dist, convert(Float32, ratio))
    center(sel::MedoidSelection, lst::AbstractVector)

Computes the medoid of lst; if ``0 < ratio < 1`` then a sampling of `lst` (``ratio * |lst|`` elements) is used instead of the complete set
"""

struct MedoidSelection{M_<:PreMetric} <: AbstractCenterSelection
    dist::M_
    ratio::Float32
end

"""
    KnnCentroidSelection(sel1::AbstractCenterSelection, sel2::AbstractCenterSelection, dist::PreMetric, k::Int32)
    KnnCentroidSelection(; sel1=CentroidSelection(), sel2=CentroidSelection(), dist=SqL2Distance(), k=0) = KnnCentroidSelection(sel, dist, convert(Int32, k))
    center(sel::KnnCentroidSelection, lst::AbstractVector)

Computes a center using the `sel1` selection strategy, and computes the final center over the set of `k` nearest neighbors
of the initial center (from `lst`) using the `dist` distance function.
"""
struct KnnCentroidSelection{S1_<:AbstractCenterSelection, S2_<:AbstractCenterSelection, M_<:PreMetric} <: AbstractCenterSelection
    sel1::S1_
    sel2::S2_
    dist::M_
    k::Int32
end

MedoidSelection(; dist=SqL2Distance(), ratio=0.5) = MedoidSelection(dist, convert(Float32, ratio))
KnnCentroidSelection(; sel1=CentroidSelection(), sel2=CentroidSelection(), dist=SqL2Distance(), k=0) = KnnCentroidSelection(sel1, sel2, dist, convert(Int32, k))

center(::CentroidSelection, lst::AbstractVector{ObjectType}) where {ObjectType<:AbstractVector{N}} where {N<:Real} =
    mean(lst)

center(::RandomCenterSelection, lst::AbstractVector) = rand(lst)

function center(sel::MedoidSelection, lst::AbstractVector)
    if sel.ratio < 1.0
        ss = randsubseq(1:length(lst), sel.ratio)
        if length(ss) > 0
            lst = lst[ss]
        end
    end

    L = zeros(Float32, length(lst))
    n = length(lst)
    for i in 1:n
        for j in i+1:length(n)
            d = evaluate(sel.dist, lst[i], lst[j])
            L[i] += d
            L[j] += d
        end
    end

    lst[argmin(L)]
end

function center(sel::KnnCentroidSelection, lst::AbstractVector)
    c = center(sel.sel1, lst)
    seq = ExhaustiveSearch(sel.dist, lst)
    k = sel.k == 0 ? ceil(Int32, log2(length(lst))) : sel.k
    k = max(1, k)
    center(sel.sel2, lst[[id for (id, dist) in search(seq, c, k)]])
end


@testset "Particle Swarm" begin
    # TODO: Run on MultivariateProblems.UnconstrainedProblems?
    Random.seed!(100)

    function f_s(x::Vector)
        (x[1] - 5.0)^4
    end

    function rosenbrock_s(x::Vector)
        (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2
    end

    initial_x = [0.0]
    upper = [100.0]
    lower = [-100.0]
    n_particles = 4
    options = Optim.Options(iterations=100)
    res = Optim.optimize(f_s, initial_x, ParticleSwarm(lower, upper, n_particles),
                         options)
    @test norm(Optim.minimizer(res) - [5.0]) < 0.1

    initial_x = [0.0, 0.0]
    lower = [-20., -20.]
    upper = [20., 20.]
    n_particles = 5
    options = Optim.Options(iterations=300)
    res = Optim.optimize(rosenbrock_s, initial_x, ParticleSwarm(lower, upper, n_particles),
                             options)
    @test norm(Optim.minimizer(res) - [1.0, 1.0]) < 0.1
    options = Optim.Options(iterations=300, show_trace=true, extended_trace=true, store_trace=true)
    res = Optim.optimize(rosenbrock_s, initial_x, ParticleSwarm(lower, upper, n_particles), options)
    @test summary(res) == "Particle Swarm"
    res = Optim.optimize(rosenbrock_s, initial_x, ParticleSwarm(n_particles = n_particles), options)
    @test summary(res) == "Particle Swarm"
end


using Romberg, Test

@testset "Romberg.jl" begin
    # Test interfaces
    x = range(0, stop=π, length=2^8+1)
    y = sin.(x)

    @test romberg(x, y) == romberg(x, y, maxeval=9)

    # Test ispow2(length(x) - 1) == false
    x = range(0, stop=π, length=2^8)
    y = sin.(x)
    @test romberg(x, y)[1] ≈ 2   rtol=1e-12

    x = range(0, stop=π, length=2^8-1)
    y = sin.(x)
    @test romberg(x, y)[1] ≈ 2   rtol=1e-9

    # Test length(x) == 1
    x = 0:0
    y = sin.(x)
    @test romberg(x, y)[1] == 0

    # Test length(x) == 2
    x = range(0, stop=1, length=2)
    y = x.^2
    @test romberg(x,y)[1] == 0.5   # integration is inaccurate

    # Test length(x) == 3
    x = range(0, stop=1, length=3)
    y = x.^2
    @test romberg(x, y)[1] ≈ 1/3   rtol=3e-16  # should be exact up to roundoff error

    # Integrate different functions
    x = range(0, stop=π, length=2^8+1)
    y = sin.(x)
    @test romberg(x, y)[1] ≈ 2    rtol=1e-15

    x = range(0, stop=1, length=2^8+1)
    y = x.^3
    @test romberg(x, y)[1] ≈ 1/4   rtol=3e-16  # should be exact up to roundoff error

    x = range(0, stop=π/2, length=2^5+1)
    y = sin.(x).^2
    @test romberg(x, y)[1] ≈ π/4    rtol=1e-15

    m = 3
    n = 4
    x = range(0, stop=π, length=2^8+1)
    y = sin.(m*x).*cos.(n*x)
    v = romberg(x, y)
    @test v[1] ≈ 2*m/(m^2 - n^2)  rtol=1e-13
    @test v[2] < 1e-10

    # the following are singular integrands where the underlying
    # theory of Romberg integration starts to break down:

    # this one requires lots of samples...
    a = 15
    x = range(0, stop=a, length=2^16+1)
    y = sqrt.(a^2 .- x.^2)
    @test romberg(x, y)[1] ≈ π*a^2/4   rtol=1e-8
    # we can do much better if we put in the correct convergence
    # rate for this singular integrand, but this requires some
    # understanding of the theory that most people won't have…
    x = range(0, stop=a, length=2^7+1)
    y = sqrt.(a^2 .- x.^2)
    @test romberg(x, y, power=0.5)[1] ≈ π*a^2/4   rtol=1e-8

    # tricky to integrate
    x = range(1e-15, stop=1, length=2^16+1)
    y = log.(x)./(1 .+ x)
    @test romberg(x, y)[1] ≈ -π^2/12  atol=1e-4

    # make sure it works for abstractly-typed y and integer Δx
    @test romberg(1, Any[3,3,3,3,3,3,3])[1] == 6*3

    @test @inferred(romberg(1, [1//2, 1//4])) === (0.375, 0.375)
    @test @inferred(romberg(1, [[1//2,1//1], [1//4,1//1]])) == ([0.375,1.0], hypot(0.375,1.0))
end


module TestScheduler

using TapirSchedulers:
    TapirSchedulers, WORK_STEALING_SCHEDULER, DEPTH_FIRST_SCHEDULER, rollback_priority
using Test

test_simple_spawn_work_stealing() = check_simple_spawn(WORK_STEALING_SCHEDULER[])

function test_simple_spawn_depth_first()
    rollback_priority() do
        check_simple_spawn(DEPTH_FIRST_SCHEDULER[])
    end
end

function check_simple_spawn(sch)
    w = TapirSchedulers.spawn!(sch) do
        1 + 1
    end
    @test fetch(w) == 2
end

function fib(N, sch)
    if N <= 1
        return N
    end
    t = TapirSchedulers.spawn!(sch) do
        fib(N - 2, sch)
    end
    y = fib(N - 1, sch)
    return fetch(t)::Int + y
end

fib_ws(N) = fib(N, WORK_STEALING_SCHEDULER[])
function fib_df(N)
    rollback_priority() do
        fib(N, DEPTH_FIRST_SCHEDULER[])
    end
end

test_fib_work_stealing() = check_fib(fib_ws)
test_fib_depth_first() = check_fib(fib_df)

function check_fib(fib)
    @test fib(1) == 1
    @test fib(2) == 1
    @test fib(3) == 2
    @test fib(4) == 3
    @test fib(5) == 5
    @test fib(6) == 8
    @test fib(10) == 55
end

end  # module


function get_xyz_vector_from_rφz_field_vector_at_rφz(field,r,φ,z,ir,iφ,iz)::Vector
    startpoint_vector = get_xyz_vector_from_rφz_vector([r,φ,z])
    endpoint_vector = get_xyz_vector_from_rφz_vector([r,φ,z]+field[ir,iφ,iz])
    xyz_vector = endpoint_vector-startpoint_vector
    for ic in 1:size(xyz_vector,1)
        isapprox(xyz_vector[ic],0.0) ? xyz_vector[ic] = 0.0 : nothing
    end
    xyz_vector
end

function get_xy_magnitude(xyz_vector::AbstractArray)
    sqrt(xyz_vector[1]^2+xyz_vector[2]^2)
end

@recipe function f(electrical_field::Array{SVector{3,T},3}, grid::CylindricalGrid{T}; view=:Components, plane=:rz, i_fixed=3, spacing = 8, vectorscale = 0.0018, SI_factor=1/1000.) where{T <: SSDFloat}

    vectorfield = electrical_field.*SI_factor
    units = Dict(1e-3=>"mm",1e-2=>"cm",1e-1=>"dm",1=>"m",1.0=>"m")
    if view == :Components
        if plane == :rz
            layout := (2,2)
            size := (800,900)
            vectorfield_r = get_component_field(vectorfield);
            vectorfield_φ = get_component_field(vectorfield,:phi);
            vectorfield_z = get_component_field(vectorfield,:z);
            vectorfield_magn = get_magnitude_of_rφz_vector.(vectorfield);

            st := :heatmap
            # colorbar_title := "Field Strength / V m\$\^\{-1\}\$"
            @series begin
                subplot := 1
                title := "r_component"
                ylabel --> "z ["*units[SI_factor]*"]"
                grid.r ./ SI_factor, grid.z ./ SI_factor, vectorfield_r[:,i_fixed,:]'
            end
            @series begin
                subplot := 2
                colorbar_title --> "Field Strength [V / "*units[SI_factor]*"]"
                title := "φ_component"

                grid.r ./SI_factor, grid.z ./SI_factor, vectorfield_φ[:,i_fixed,:]'
            end
            @series begin
                subplot := 3
                title := "z_component"
                ylabel --> "z ["*units[SI_factor]*"]"
                xlabel --> "r ["*units[SI_factor]*"]"
                grid.r ./SI_factor, grid.z ./SI_factor, vectorfield_z[:,i_fixed,:]'
            end
            @series begin
                subplot := 4
                colorbar_title --> "Field Strength [V / "*units[SI_factor]*"]"
                xlabel --> "r ["*units[SI_factor]*"]"
                title:= "magnitude"
                grid.r ./SI_factor, grid.z ./SI_factor, vectorfield_magn[:,i_fixed,:]'
            end
        end
    elseif view == :ef
        if plane == :rφ
            vectorfield_xyz = Array{Vector{Float32}}(undef,size(vectorfield,1),size(vectorfield,2),size(vectorfield,3));
            for (iz,z) in enumerate(grid.z)
                for (iφ,φ) in enumerate(grid.φ)
                    for (ir,r) in enumerate(grid.r)
                        vectorfield_xyz[ir,iφ,iz]=get_xyz_vector_from_rφz_field_vector_at_rφz(vectorfield,r,φ,z,ir,iφ,iz)
                    end
                end
            end
        elseif plane == :xy
            vectorfield_xyz = electrical_field.field
        end
        vectorfield_xy_magn = map(x->get_xy_magnitude(x),vectorfield_xyz[:,:,i_fixed])
        max_magn = maximum(vectorfield_xy_magn)
        diff_magn = max_magn-minimum(vectorfield_xy_magn)

        size := (800,800)
        line := (:arrow,:blue,2)
        label := ""
        ylabel := "y "
        xlabel := "x "
        title := "z = $(round(grid.z[i_fixed]/SI_factor,digits=2)) / mm"
        xlims := (-1.2/SI_factor*maximum(grid.r),1.2/SI_factor*maximum(grid.r))
        ylims := (-1.2/SI_factor*maximum(grid.r),1.2/SI_factor*maximum(grid.r))
        for (ir,r) in enumerate(grid.r[1:spacing:end])
        for (iφ,φ) in enumerate(grid.φ)
                x= r*cos(φ)
            y= r*sin(φ)
            ir_actual=findfirst(x->x==r,grid.r)
            iφ_actual=findfirst(x->x==φ,grid.φ)
                xy_magn = vectorfield_xy_magn[ir_actual,iφ_actual]
                vector=vectorfield_xyz[ir_actual,iφ_actual,i_fixed]/xy_magn
                vector*=((vectorfield_xy_magn[ir_actual,iφ_actual]-0.8*minimum(vectorfield_xy_magn))/diff_magn)
                vector*=vectorscale
                @series begin
                    [x-0.5*vector[1],x+0.5*vector[1]]/SI_factor, [y-0.5*vector[2],y+0.5*vector[2]]/SI_factor
                end
            end
        end
    end
end

using Base.Math
@recipe function f( electric_field::Array{ <:StaticVector{3, T}, 3}, det::SolidStateDetector; φ=missing, z = missing, scaling_factor=1.0, spacings = [3,20]) where T
    ismissing(φ) && !ismissing(z) ? iz = searchsortednearest(det.zs,T(z)) : nothing
    !ismissing(φ) && ismissing(z) ? iφ = searchsortednearest(det.φs,T(φ)) : nothing
    guidefontsize --> 15
    tickfontsize -->15
    # labelfontsize --> 13
    if ismissing(φ)
        xy = T[0, 0]
        values = T[0, 0]
        magnitudes_xy = []
        for (iφ ,φ)  in enumerate(det.φs)
            if (iφ+spacings[1]-1)%spacings[1]==0
                for (ir, r) in enumerate(det.rs)
                    if (ir+spacings[2]-1)%spacings[2] == 0
                        xy = hcat(xy,[ r * cos(φ), r * sin(φ)])
                        values = hcat(values, electric_field[ir, iφ, iz][1:2])
                        push!(magnitudes_xy,sqrt(electric_field[ir, iφ, iz][1]^2 + electric_field[ir, iφ, iz][2]^2))
                    end
                end
            end
        end

        num = minimum([mean(diff(det.φs)*spacings[1]),mean(diff(det.rs))*spacings[2]])
        nom = maximum(magnitudes_xy)
        scaling = num/nom * scaling_factor
        for i in 1:size(xy,2)
            @series begin
                title --> "z = $(round(det.zs[iz] .* 1000,digits=2)) mm"
                xlabel --> "x / mm"
                ylabel --> "y / mm"
                aspect_ratio --> 1
                label --> ""
                color --> :red
                arrow --> true
                [xy[1,i] - 0.5*scaling*values[1,i], xy[1,i] + 0.5*scaling*values[1,i]] .* 1000, [xy[2,i] - 0.5*scaling*values[2,i], xy[2,i] + 0.5*scaling*values[2,i]] .* 1000
            end
        end

    elseif ismissing(z)
        rz = T[0, 0]
        values = T[0, 0]
        magnitudes_rz = []
        for (iz ,z)  in enumerate(det.zs)
            if (iz+spacings[1])%spacings[1]==0
                for (ir, r) in enumerate(det.rs)
                    if (ir+spacings[2]-1)%spacings[2] == 0
                        rz = hcat(rz,[ r, z])
                        values = hcat(values, electric_field[ir, iφ, iz][1:2:3])
                        push!(magnitudes_rz,sqrt(electric_field[ir, iφ, iz][1]^2 + electric_field[ir, iφ, iz][3]^2))
                    end
                end
            end
        end
        num = minimum([mean(diff(det.zs)*spacings[1]),mean(diff(det.rs))*spacings[2]])
        nom = maximum(magnitudes_rz)
        scaling = num/nom * scaling_factor
        for i in 1:size(rz,2)
            @series begin
                title --> "φ = $(round(rad2deg(det.φs[iφ]) , digits=2)) deg"
                xlabel --> "r [mm]"
                ylabel --> "z [mm]"
                aspect_ratio --> 1
                label --> ""
                color --> :red
                arrow --> true
                [rz[1,i] - 0.5*scaling*values[1,i], rz[1,i] + 0.5*scaling*values[1,i]] .* 1000, [rz[2,i] - 0.5*scaling*values[2,i], rz[2,i] + 0.5*scaling*values[2,i]] .* 1000
            end
        end
    end
end

@userplot MyQuiver
@recipe function f(gdd::MyQuiver; scaling=1.0)
    xy::Matrix = gdd.args[1]
    values = gdd.args[2]
    for i in 1:size(xy,2)
        @series begin
            arrow --> true
            [xy[1,i] - 0.5*scaling*values[1,i], xy[1,i] + 0.5*scaling*values[1,i]], [xy[2,i] - 0.5*scaling*values[2,i], xy[2,i] + 0.5*scaling*values[2,i]]
        end
    end
end

@userplot Plot_electric_field
@recipe function f(gdd::Plot_electric_field; φ = missing, r = missing, x = missing, y = missing, z = missing,
                    spacing = 4, grid_spacing=[0.0005, deg2rad(1.0), 0.0005], n_steps=3000,
                    potential=true, contours_equal_potential=true, offset = (5e-5))
    sim = gdd.args[1]
    S = get_coordinate_system(sim.electric_field.grid)
    T = Float32

    dim_array = [φ, r, x, y, z]
    dim_symbols_array = [:φ, :r, :x, :y, :z]
    if isempty(skipmissing(dim_array))
        if S == :cylindrical
            v::T = 0
            dim_number = 2
            dim_symbol = :φ
        elseif S == :cartesian
            dim_number = 1
            dim_symbol = :x
            v = sim.electric_field.grid[dim_number][ div(length(sim.electric_field.grid[dim_number]), 2) ]
        end
    elseif sum(ismissing.(dim_array)) == 4
        dim_number = findfirst(x -> !ismissing(x), dim_array)
        dim_symbol = dim_symbols_array[dim_number]
        v = dim_array[dim_number]
        if dim_symbol == :φ v = deg2rad(v) end
    else
        throw(ArgumentError("Only one keyword for a certain dimension is allowed. One of 'φ', 'r', 'x', 'y', 'z'"))
    end

    S::Symbol = get_coordinate_system(sim.detector)
    aspect_ratio --> 1
    title --> "Electric Field Lines @$(dim_symbol)=$(round(dim_symbol == :φ ? rad2deg(v) : v, digits=2))"
    xlabel --> (S == :cylindrical ? (dim_symbol == :r ? L"$\varphi$ / °" : L"$r$ / m") : (dim_symbol == :x ? L"$y$ / m" : L"$x$ / m"))
    ylabel --> L"$z$ / m"

    if potential == true
        @series begin
            # contours_equal_potential --> contours_equal_potential
            if dim_symbol == :φ
                φ --> v
            elseif dim_symbol == :z
                if S == :cylindrical error("Electric field plot not yet implemented for z in cylindrical coordinates") end
                z --> v
                if S == :cylindrical
                    proj --> :polar
                end
            elseif dim_symbol == :r
                error("Electric field plot not yet implemented for r")
                r --> v
            elseif dim_symbol == :x
                x --> v
            elseif dim_symbol == :y
                y --> v
            end
            sim.electric_potential
        end
    end
    PT = S == :cylindrical ? CylindricalPoint{T} : CartesianPoint{T}

    contacts_to_spawn_charges_for = filter!(x -> x.id !=1, SSD.Contact{T}[c for c in sim.detector.contacts])
    spawn_positions = PT[]
    grid = sim.electric_field.grid# Grid(sim.detector, init_grid_spacing = grid_spacing, full_2π = true)
    pt_offset = T[offset,0.0,offset]

    # for c in contacts_to_spawn_charges_for
    #     ongrid_positions= map(x-> CylindricalPoint{T}(grid[x...]), SSD.paint_object(sim.detector, c, grid, Val(dim_symbol), v))
    #     for position in ongrid_positions
    #         push!(spawn_positions, CylindricalPoint{T}((position + pt_offset)...))
    #         push!(spawn_positions, CylindricalPoint{T}((position - pt_offset)...))
    #     end
    # end

    ax1 = unique(range(sim.electric_field.grid[1][1], stop = sim.electric_field.grid[1][end], step = spacing * 0.001))
    ax2 = unique(range(sim.electric_field.grid[2][1], stop = sim.electric_field.grid[2][end], step = spacing * 0.001))
    if ismissing(φ) && S == :cylindrical
        ax2 = T[0]
    elseif !ismissing(φ)
        ax2 = T[φ]
    end
    ax3 = unique(range(sim.electric_field.grid[3][1], stop = sim.electric_field.grid[3][end], step = spacing * 0.001))
    for x1 in ax1
        for x2 in ax2
            for x3 in ax3
                pt::PT = PT(x1, x2, x3)
                push!(spawn_positions, pt)
            end
        end
    end
    filter!(x -> x in sim.detector && !in(x, sim.detector.contacts), spawn_positions)

    @info "$(length(spawn_positions)) drifts are now being simulated..."

    el_field_itp     = SSD.get_interpolated_drift_field(sim.electric_field.data,       sim.electric_field.grid)
    el_field_itp_inv = SSD.get_interpolated_drift_field(sim.electric_field.data .* -1, sim.electric_field.grid)
    @showprogress for (ipos, pos) in enumerate(spawn_positions)
        if ((spacing-1)+ipos)%spacing == 0

            path = CartesianPoint{T}[CartesianPoint{T}(0.0,0.0,0.0) for i in 1:n_steps]
            SSD.drift_charge!(path, sim.detector, sim.point_types, sim.electric_potential.grid, CartesianPoint(pos), T(2e-9), el_field_itp, verbose = false )
            @series begin
                c --> :white
                if dim_symbol == :z && S == :cylindrical proj --> :polar end
                label --> ""
                x, y = if dim_symbol == :φ
                    map(x -> sqrt(x[1]^2+x[2]^2), path),
                    map(x -> x[3], path)
                elseif dim_symbol == :x
                    map(x -> x[2], path),
                    map(x -> x[3], path)
                elseif dim_symbol == :y
                    map(x -> x[1], path),
                    map(x -> x[3], path)
                elseif dim_symbol == :z
                    map(x -> x[1], path),
                    map(x -> x[2], path)
                end
                x, y
            end

            path = CartesianPoint{T}[CartesianPoint{T}(0.0,0.0,0.0) for i in 1:n_steps]
            SSD.drift_charge!(path, sim.detector, sim.point_types, sim.electric_potential.grid, CartesianPoint(pos), T(2e-9), el_field_itp_inv, verbose = false )
            @series begin
                c --> :white
                if dim_symbol == :z && S == :cylindrical proj --> :polar end
                label --> ""
                x, y = if dim_symbol == :φ
                    map(x -> sqrt(x[1]^2+x[2]^2), path),
                    map(x -> x[3], path)
                elseif dim_symbol == :x
                    map(x -> x[2], path),
                    map(x -> x[3], path)
                elseif dim_symbol == :y
                    map(x -> x[1], path),
                    map(x -> x[3], path)
                elseif dim_symbol == :z
                    map(x -> x[1], path),
                    map(x -> x[2], path)
                end
                x, y
            end

        end
    end
end


function imfilter(A, kern)
    size(kern,1) == size(kern,2) || throw(ArgumentError("Only square kernels supported, you provided $(size(kern))."))
    isodd(size(kern,1)) || throw(ArgumentError("Only odd-sized kernels supported, you provided $(size(kern))."))
    A2 = LocalFilters.convolve(A, kern)
    pad = size(kern, 1) ÷ 2
    @views A2[1:pad, :] .= A[1:pad, :]
    @views A2[end-pad+1:end, :] .= A[end-pad+1:end, :]
    @views A2[:, 1:pad] .= A[:, 1:pad]
    @views A2[:, end-pad+1:end] .= A[:, end-pad+1:end]
    A2
end

function gaussian2(n)
    t = range(-2,stop=2, length=n)
    k = exp.(-t.^2)
    K = k*k'
    K ./= sum(K)
end

function gaussian(n)
    t = range(-2,stop=2, length=n)
    K = [exp(-(x-y)^2) for x in t, y in t]
    K ./= sum(K)
end


@testset "area measurement" begin

    img = [0 0 0 0 0;
           0 1 0 1 0;
           0 1 1 0 0;
           0 0 0 0 0]

    labels  = Images.label_components(img,trues(3,3))
    t = measure_components((Measurement(),), labels)
    @test length(t) == 1

    row₁ = t[1]
    @test row₁.area == 4.375

    test_image_areas = ([0], [25], [1], [1, 1, 1, 1], [17], [16], [8.5],
                        [9.25], [4, 4], [21.25], [19.25], [20], [19.875],
                        [21.5], [17], [20.75, 8.5, 12.5], [38.5], [19.625, 8.75],
                        [25.5, 11], [31.5, 11.5])

    for T in (Int, Bool, Gray{Bool}, Gray{N0f8}, Gray{N0f16}, Gray{N0f32}, Gray{Float64})
        for i = 1:20
            test_image = eval(Symbol("test_image_$(i)"))
            labels = Images.label_components(test_image,trues(3,3))
            t = measure_components((Measurement(),), labels)
            i == 1 ? (@test isempty(t)) : (@test all(select(t,:area) .== test_image_areas[i]))
        end
    end
end

@testset "perimeter" begin

    test_image_perimeters = ([0], [20], [4], [4, 4, 4, 4], [52], [48], [16],
                        [20], [10, 10], [28], [32], [34], [32],
                        [28], [30], [40, 16, 24], [76], [40, 24],
                        [72, 28], [64, 30])

    for T in (Int, Bool, Gray{Bool}, Gray{N0f8}, Gray{N0f16}, Gray{N0f32}, Gray{Float64})
        for i = 1:20
            test_image = eval(Symbol("test_image_$(i)"))
            labels = Images.label_components(test_image,trues(3,3))
            t = measure_components((Measurement(),), labels)
            i == 1 ? (@test isempty(t)) : (@test all(select(t,:perimeter₁) .== test_image_perimeters[i]))
        end
    end
end


# # Numerical Dynamic Programming

# Florian Oswald, Sciences Po, 2019

# ## Intro

# * Numerical Dynamic Programming (DP) is widely used to solve dynamic models.
# * You are familiar with the technique from your core macro course.
# * We will illustrate some ways to solve dynamic programs.
#     1. Models with one discrete or continuous choice variable
#     1. Models with several choice variables
#     1. Models with a discrete-continuous choice combination
# * We will go through:
#     1. Value Function Iteration (VFI)
#     1. Policy function iteration (PFI)
#     1. Projection Methods
#     1. Endogenous Grid Method (EGM)
#     1. Discrete Choice Endogenous Grid Method (DCEGM)

# ## Dynamic Programming Theory

# * Payoffs over time are
# 	$$U=\sum_{t=1}^{\infty}\beta^{t}u\left(s_{t},c_{t}\right) $$
# 	where $\beta<1$ is a discount factor, $s_{t}$ is the state, $c_{t}$ is the control.

# * The state (vector) evolves as $s_{t+1}=h(s_{t},c_{t})$.
# * All past decisions are contained in $s_{t}$.

# ### Assumptions

# * Let $c_{t}\in C(s_{t}),s_{t}\in S$ and assume $u$ is bounded in $(c,s)\in C\times S$.
# * Stationarity: neither payoff $u$ nor transition $h$ depend on time.
# * Write the problem as
# 	$$ v(s)=\max_{s'\in\Gamma(s)}u(s,s')+\beta v(s') $$
# * $\Gamma(s)$ is the constraint set (or feasible set) for $s'$ when the current state is $s$

# ### Existence

# **Theorem.** Assume that $u(s,s')$ is real-valued, continuous, and bounded, that $\beta\in(0,1)$, and that the constraint set $\Gamma(s)$ is nonempty, compact, and continuous. Then there exists a unique function $v(s)$ that solves the above functional equation.

# **Proof.** [@stokeylucas] <cite data-cite=stokeylucas></cite> theoreom 4.6.

# # Solution Methods

# ## Value Function Iteration (VFI)

# * Find the fix point of the functional equation by iterating on it until the distance between consecutive iterations becomes small.
# * Motivated by the Bellman Operator, and it's characterization in the Continuous Mapping Theorem.

# ## Discrete DP VFI

# * Represents and solves the functional problem in $\mathbb{R}$ on a finite set of grid points only.
# * Widely used method.
#    * Simple (+)
#    * Robust (+)
#    * Slow (-)
#    * Imprecise (-)
# * Precision depends on number of discretization points used.
# * High-dimensional problems are difficult to tackle with this method because of the curse of dimensionality.

# ### Deterministic growth model with Discrete VFI

# * We have this theoretical model:

# \begin{eqnarray}
#    V(k) &=& \max_{0<k'<f(k)} u(f(k) - k') + \beta V(k')\\
#   f(k)  &=& k^\alpha
# \end{eqnarray}

# * and we employ the followign numerical approximation:
# 	$$ V(k_i) = \max_{i'=1,2,\dots,n} u(f(k_i) - k_{i'}) + \beta V(i') $$

# * The iteration is then on successive iterates of $V$: The LHS gets updated in each iteration!

# $$
# \begin{aligned}
# 	V^{r+1}(k_i) &= \max_{i'=1,2,\dots,n} u(f(k_i) - k_{i'}) + \beta V^{r}(i') \\
# 	V^{r+2}(k_i) &= \max_{i'=1,2,\dots,n} u(f(k_i) - k_{i'}) + \beta V^{r+1}(i') \\
# 	... &
# \end{aligned}
# $$

# * And it stops at iteration $r$ if $d(V^{r},V^{r-1}) < \text{tol}$
# * You choose a measure of *distance*, $d(\cdot,\cdot)$, and a level of tolerance.
# * $V^{r}$ is usually an *array*. So $d$ will be some kind of *norm*.
# * maximal absolute distance
# * mean squared distance

# ### Exercise 1: Implement discrete VFI

# ## Checklist

# 1. Set parameter values
# 1. define a grid for state variable $k \in [0,2]$
# 1. initialize value function $V$
# 1. start iteration, repeatedly computing a new version of $V$.
# 1. stop if $d(V^{r},V^{r-1}) < \text{tol}$.
# 1. plot value and policy function
# 1. report the maximum error of both wrt to analytic solution

alpha     = 0.65
beta      = 0.95
grid_max  = 2  # upper bound of capital grid
n         = 150  # number of grid points
N_iter    = 3000  # number of iterations
kgrid     = 1e-2:(grid_max-1e-2)/(n-1):grid_max  # equispaced grid
f(x) = x^alpha  # defines the production function f(k)
tol = 1e-9


using RDatasets, StatsPlots; pyplot()

iris = dataset("datasets", "iris")
@df iris andrewsplot(:Species, cols(1:4),
	line=(fill=[:blue :red :green]), legend=:topleft)


export @ein_str, @ein, IndexSize
"""
    ein"ij,jk -> ik"(A,B)

String macro interface which understands `numpy.einsum`'s notation.
Translates strings into `EinCode`-structs that can be called to evaluate
an `einsum`.
To control evaluation order, use parentheses - instead of an `EinCode`,
a `NestedEinsumStable` is returned which evaluates the expression
according to parens.
The valid character ranges for index-labels are `a-z` and `α-ω`.

# example

```jldoctest; setup = :(using OMEinsum)
julia> a, b, c = rand(10,10), rand(10,10), rand(10,1);

julia> ein"ij,jk,kl -> il"(a,b,c) ≈ ein"(ij,jk),kl -> il"(a,b,c) ≈ a * b * c
true
```
"""
macro ein_str(s::AbstractString)
    s = replace(s, " " => "")
    m = match(r"([\(\)a-z,α-ω]*)->([a-zα-ω]*)", s)
    m == nothing && throw(ArgumentError("invalid einsum specification $s"))
    sixs, siy = m.captures
    if '(' in sixs
        return parse_nested(sixs, collect(siy))
    else
        iy  = Tuple(siy)
        ixs = Tuple(Tuple(ix) for ix in split(sixs,','))
        return EinCode(ixs, iy)
    end
end

function (code::EinCode)(xs...; size_info=nothing) where {T, N}
    size_dict = get_size_dict(getixs(code), xs)
    if !(size_info isa Nothing)
        size_dict += size_info
    end
    einsum(code, xs, size_dict)
end

einsum(code::EinCode{ixs, iy}, xs) where {ixs, iy} = einsum(code, xs, nothing)

function einsum(code::EinCode{ixs, iy}, xs, ::Nothing) where {ixs, iy}
    einsum(code, xs, get_size_dict(ixs, xs))
end

@doc raw"
    get_size_dict(ixs, xs)

return the `IndexSize` struct that is used to get the size of an index-label
in the einsum-specification with input-indices `ixs` and tensors `xs` after
consistency within `ixs` and between `ixs` and `xs` has been verified.
"
function get_size_dict(ixs::NTuple{N, NTuple{M, T} where M}, xs::NTuple{X, AbstractArray}) where {N,T,X}
    # check size of input tuples
    N != X && throw(ArgumentError("$X tensors labelled by $N indices"))

    # check tensor orders
    foreach(ixs, xs) do ix, x
        length(ix) == ndims(x) || throw(
            ArgumentError("indices $ix invalid for tensor with ndims = $(ndims(x))"))
    end
    sd = getindexsize(ixs, xs)
    check_dimensions(sd)
    return sd
end

@doc raw"
    IndexSize{N,T}(k::NTuple{N,T},v::NTuple{N,Int})

struct to hold the size of indices specified by their labels.
Note that while a dict would work, for the small sizes we usually
have, a tuple of keys and values is much faster to construct
and competitive for lookup.
"
struct IndexSize{N,T}
    indices::NTuple{N,T}
    sizes::NTuple{N,Int}
end

IndexSize(sizes::Pair...) = IndexSize(first.(sizes), last.(sizes))
IndexSize(indices::Tuple{}, sizes::Tuple{}) = IndexSize{0,Any}(indices, sizes)

Base.:+(x::IndexSize, y::IndexSize) = IndexSize((x.indices..., y.indices...), (x.sizes..., y.sizes...))

function getindexsize(ixs, xs)
    indices = TupleTools.flatten(ixs)
    sizes = TupleTools.flatten(map(size,xs))
    IndexSize(indices,sizes)
end

Base.getindex(inds::IndexSize{N,T},i::T) where {N,T} = inds.sizes[findfirst(==(i), inds.indices)]

@doc raw"
    check_dimensions(inds::IndexSize)

check whether all non-unique indexlabels point to the same
dimension - otherwise throw an error.
"
function check_dimensions(inds::IndexSize)
    for (c,i) in enumerate(inds.indices)
        j = findnext(==(i), inds.indices, c+1)
        j != nothing && inds.sizes[c] != inds.sizes[j] && return throw(
            DimensionMismatch("index $i has incosistent sizes"))
    end
    return true
end

using MacroTools
"""
    @ein A[i,k] := B[i,j] * C[j,k]     # A = B * C

Macro interface similar to that of other packages.

You may use numbers in place of letters for dummy indices, as in `@tensor`,
and need not name the output array. Thus `A = @ein [1,2] := B[1,ξ] * C[ξ,2]`
is equivalent to the above. This can also be written `A = ein"ij,jk -> ik"(B,C)`
using the numpy-style string macro.

# example

```jldoctest; setup = :(using OMEinsum)
julia> a, b = rand(2,2), rand(2,2);

julia> @ein c[i,k] := a[i,j] * b[j,k];

julia> c ≈ a * b
true

julia> c ≈ ein"ij,jk -> ik"(a,b)
true
```
"""
macro ein(exs...)
    _ein_macro(exs...)
end

primefix!(ind) = map!(i -> @capture(i, (j_)') ? Symbol(j, '′') : i, ind, ind)

function _ein_macro(ex; einsum=:einsum)
    @capture(ex, (left_ := right_)) || throw(ArgumentError("expected A[] := B[]... "))
    @capture(left, Z_[leftind__] | [leftind__] ) || throw(
        ArgumentError("can't understand LHS, expected A[i,j] etc."))
    if Z===nothing
        @gensym Z
    end
    primefix!(leftind)

    rightind, rightpairs = [], []
    @capture(right, *(factors__)) || (factors = Any[right])
    for fact in factors
        @capture(fact, A_[Aind__]) || return _nested_ein_macro(ex)
        primefix!(Aind)
        append!(rightind, Aind)
        push!(rightpairs, (A, Aind) )
    end
    unique!(rightind)
    isempty(setdiff(leftind, rightind)) || throw(
        ArgumentError("some indices appear only on the left"))

    lefttuple = Tuple(indexin(leftind, rightind))
    righttuples = [ Tuple(indexin(ind, rightind)) for (A, ind) in rightpairs ]
    rightnames = [ esc(A) for (A, ind) in rightpairs ]

    return :( $(esc(Z)) = $einsum( EinCode{($(righttuples...),), $lefttuple}(), ($(rightnames...),)) )
end


import ..bin

"""
    bin(left_bin_edges::AbstractRange, xs, ys) -> Vector{Vector{T}} where T

Distribute the elements of `ys` into `N-1` different bin vectors, based on
how the values in `xs` are distributed among the bins defined by the `N` grid
points in `left_bin_edges`. If `xs[i]` falls in the `n`-th bin interval, then `ys[i]`
is assigned to the `n`-th bin vector. If `xs[i]` lie outside the grid, then the
corresponding `ys[i]` is ignored. See also [`bin!`](@ref)

Returns `N - 1` bin vectors.

## Examples

### Getting the values in each bin:

```julia
xs = [1.2, 1.7, 2.2, 3.3, 4.5, 4.6, 7.1]
ys = [4.2, 5.1, 6.5, 4.2, 3.2, 3.1, 2.5]
left_bin_edges = 0.0:1.0:6.0
bin(left_bin_edges, xs, ys)
```

```julia
# Some example data with unevenly spaced time indices
npts = 300
time, vals = sort(rand(1:1000, npts)), rand(npts)

# See which values fall in 25 time step wide time bins ranging
# from time indices 100 to 900.
left_bin_edges = 100:25:900

bin(left_bin_edges, time, vals)
```
"""
function bin(left_bin_edges::AbstractRange{T}, xs, ys) where T
    if length(xs) != length(ys)
        msg = "`length(xs)` not equal to `length(ys)`"
        throw(ArgumentError(msg))
    end

    # Each bin is represented by (potentially differently sized) vectors of type `eltype(ys)`
    YT = eltype(ys)
    n_bins = length(left_bin_edges) - 1
    bin_vals = [Array{YT}(undef, 0) for i = 1:n_bins]

    # Assign values to bins. Can be done faster if `xs` is guaranteed to be ordered.
    g_min = minimum(left_bin_edges)
    g_step = step(left_bin_edges)

    for (x, y) in zip(xs, ys)
        arr_idx = ceil(Int, (x - g_min)/g_step)
        # Check that the value falls inside the grid; if not, ignore it
        if 0 < arr_idx <= n_bins
            push!(bin_vals[arr_idx], y)
        end
    end

    return bin_vals
end

"""
    bin(f::Function, left_bin_edges::AbstractRange, xs, ys, args...; kwargs...) -> Vector{T} where T

Distribute the elements of `ys` into `N-1` different bin vectors, based on
how the values in `xs` are distributed among the bins defined by the `N` grid
points in `left_bin_edges`. If `xs[i]` falls in the `n`-th bin interval, then `ys[i]`
is assigned to the `n`-th bin vector. If `xs[i]` lie outside the grid, then the
corresponding `ys[i]` is ignored. See also [`bin!`](@ref)

Then, apply the summary function element-wise to each of the bin vectors, with `args`
and `kwargs` as arguments and keyword arguments. Then, `N-1` summary values, one for
each bin, are returned. Empty bins are assigned `NaN` values.

Returns `N-1` bin summaries.

## Examples

### Applying a summary function to each bin

Any function that accepts a vector of values can be used in conjunction with `bin`.

```julia
xs = [1.2, 1.7, 2.2, 3.3, 4.5, 4.6, 7.1]
ys = [4.2, 5.1, 6.5, 4.2, 3.2, 3.1, 2.5]
left_bin_edges = 0.0:1.0:6.0
bin(median, left_bin_edges, xs, ys)
```

Functions with additional arguments also work (arguments and keyword
arguments must be supplied last in the function call):

```julia
xs = [1.2, 1.7, 2.2, 3.3, 4.5, 4.6, 7.1]
ys = [4.2, 5.1, 6.5, 4.2, 3.2, 3.1, 2.5]
left_bin_edges = 0.0:1.0:6.0
bin(quantile, left_bin_edges, xs, ys, [0.1])
```
"""
function bin(f::Function, left_bin_edges::AbstractRange{T}, xs, ys, args...; kwargs...) where T
    if length(xs) != length(ys)
        msg = "`length(xs)` not equal to `length(ys)`"
        throw(ArgumentError(msg))
    end

    # Each bin is represented by (potentially differently sized) vectors of type `eltype(ys)`
    YT = eltype(ys)
    n_bins = length(left_bin_edges) - 1
    bin_vals = [Array{YT}(undef, 0) for i = 1:n_bins]

    # Assign values to bins. Can be done faster if `xs` is guaranteed to be ordered.
    g_min = minimum(left_bin_edges)
    g_step = step(left_bin_edges)

    for (x, y) in zip(xs, ys)
        arr_idx = ceil(Int, (x - g_min)/g_step)
        # Check that the value falls inside the grid; if not, ignore it
        if 0 < arr_idx <= n_bins
            push!(bin_vals[arr_idx], y)
        end
    end

    # Get bin summaries by applying function to non-empty bins.
    bin_summaries = fill(NaN, n_bins)
    inds_nonempty_bins = findall(length.(bin_vals) .> 0)
    bin_summaries[inds_nonempty_bins] = f.(bin_vals[inds_nonempty_bins], args...; kwargs...)

    return bin_summaries
end

"""
    bin!(bins::Vector{AbstractVector{T}}, left_bin_edges::AbstractRange{T}, xs, ys) where T

Distribute the elements of `ys` into `N-1` different pre-allocated empty
bin vectors, based on how the values in `xs` are distributed among
the bins defined by the `N` grid points in `left_bin_edges`.
`bins` must be a vector of vector-like mutable containers.

If `xs[i]` falls in the `n`-th bin interval, then `ys[i]`
is assigned to the `n`-th bin vector. If `xs[i]` lie outside the grid,
the corresponding `ys[i]` is ignored.

See also [`bin(::AbstractRange)`](@ref).
"""
function bin!(bins::Vector{AbstractVector{T}}, left_bin_edges::AbstractRange{T}, xs, ys) where T
    if length(xs) != length(ys)
        msg = "`length(xs)` not equal to `length(ys)`"
        throw(ArgumentError(msg))
    end

    # Each bin is represented by (potentially differently sized) vectors of type `eltype(ys)`
    n_bins = length(left_bin_edges) - 1
    if length(bins) != n_bins
        msg = "`The number of pre-allocated bins ($n_bins) does not match the number of bins defined by the grid in `left_bin_edges`"
        throw(ArgumentError(msg))
    end

    ELTYPE_YS, ELTYPE_BINS  = eltype(ys), eltype(eltype(bins))

    if ELTYPE_BINS != ELTYPE_YS
        msg = "Element type of bin vectors `bins` ($ELTYPE_BINS) is not equal to the element type of `ys` ($ELTYPE_YS)"
    end

    # Assign values to bins. Can be done faster if `xs` is guaranteed to be ordered.
    g_min = minimum(left_bin_edges)
    g_step = step(left_bin_edges)

    @inbounds for (x, y) in zip(xs, ys)
        arr_idx = ceil(Int, (x - g_min)/g_step)
        # Check that the value falls inside the grid; if not, ignore it
        if 0 < arr_idx <= n_bins
            push!(bins[arr_idx], y)
        end
    end

    return bins
end

"""
    bin_mean(left_bin_edges::AbstractRange, xs, ys)

Distribute the elements of `ys` into `N - 1` different bin vectors, based on
how the values in `xs` are distributed among the bins defined by the `N` grid
points in `left_bin_edges`. Then compute the bin mean for each bin.

If `xs[i]` falls in the `n`-th bin interval, then `ys[i]` is assigned to the
`n`-th bin vector. If values fall outside the grid, they are ignored (if
`xs[i] < minimum(left_bin_edges)`, ignore `ys[i]`). After the `ys` values
have been assigned to bin vectors, apply the summary function `f` element-wise
to each of the bin vectors, with `args` and `kwargs` as arguments and keyword
arguments.

Returns `N - 1` mean values, one for each bin.

## Examples

```julia
xs = [1.2, 1.7, 2.2, 3.3, 4.5, 4.6, 7.1]
ys = [4.2, 5.1, 6.5, 4.2, 3.2, 3.2, 2.5]
left_bin_edges = 0.0:1.0:6.0
bin_mean(left_bin_edges, xs, ys)

# output
6-element Array{Float64,1}:
 NaN
   4.65
   6.5
   4.2
   3.2
 NaN
```
"""
function bin_mean(left_bin_edges::AbstractRange, xs, ys; nan_threshold = 0)
    if length(xs) != length(ys)
        msg = "`length(xs)` not equal to `length(ys)`"
        throw(ArgumentError(msg))
    end

    # Check that the value falls inside the grid; if not, ignore it
    n_bins = length(left_bin_edges) - 1
    bin_sums = fill(0.0, n_bins)
    bin_sums_n_entries = fill(0.0, n_bins)

    # Assign values to bins. Can be done faster if `xs` is guaranteed to be ordered.
    g_min = minimum(left_bin_edges)
    g_step = step(left_bin_edges)

    for (x, y) in zip(xs, ys)
        arr_idx = ceil(Int, (x - g_min)/g_step)

        if 0 < arr_idx <= n_bins
            bin_sums[arr_idx] += y
            bin_sums_n_entries[arr_idx] += 1.0
        end
    end

    # Return bin averages (entries with 0s are represented as NaNs)
    bin_avgs = bin_sums ./ bin_sums_n_entries
    bin_avgs[isapprox.(bin_avgs, 0.0)] .= NaN

    # If there aren't enough entries in a bin, set it to NaN
    # and interpolate it instead.
    bin_avgs[bin_sums_n_entries .< nan_threshold] .= NaN

    return bin_avgs
end

export bin, bin!, bin_mean


print("convert, ")

n = NamedArray(randn(5, 10))

@test convert(Array, n) == n.array
@test convert(NamedArray, n) == n

for T in [Float32, Float16]
 	for matrix in (map(T, n), convert(NamedArray{T}, n))
 		@test matrix.array == map(T, n.array)
 		@test names(matrix) == names(n)
 		@test dimnames(matrix) == dimnames(n)
 	end
end

# issue 43

x = @inferred NamedArray(Array([1 2; 3 4]), (["a","b"], [10,11]), (:rows,:cols))
@test convert(Array, x) == x.array

# Test convert don't change array type

n_sparray = NamedArray(SparseArrays.sprand(Int, 5, 5, 0.5))
@test isa(n_sparray.array, SparseMatrixCSC{Int,Int})
@test isa(convert(NamedArray{Float64}, n_sparray).array, SparseMatrixCSC{Float64,Int})


# Unit testing for element-wise expressions on scalar backend

import Devectorize
import Devectorize.@devec
import Devectorize.@inspect_devec
import Devectorize.dump_devec
import Devectorize.sqr
import Devectorize.blend

using Base.Test

a = [1., 2., 3., 4., 5., 6., 7., 8.]
b = [3., 4., 5., 6., 8., 7., 6., 5.]
c = [9., 8., 7., 6., 4., 2., 3., 1.]
cv = 12.0
cv_a = [12., 12., 12., 12., 12., 12., 12., 12.]

abc = [a b c]

#################################################
#
#   basic element-wise maps
#
#################################################

# unary

@devec r = -a
@test isequal(r, -a)

r[:] = 0.
@devec r[:] = -a
@test isequal(r, -a)

# binary

@devec r = a + b
@test isequal(r, a + b)

r[:] = 0.
@devec r[:] = a + b
@test isequal(r, a + b)

@devec r = a + 1.
@test isequal(r, a + 1.)

r[:] = 0.
@devec r[:] = a + 1.
@test isequal(r, a + 1.)

@devec r = 1. + b
@test isequal(r, 1. + b)

r[:] = 0.
@devec r[:] = 1. + b
@test isequal(r, 1. + b)

@devec r[:] = 1. + 2.
@test isequal(r, fill(3., size(a)))

# ternary

@devec r = +(a, b, c)
@test isequal(r, a + b + c)

r[:] = 0.
@devec r[:] = +(a, b, c)
@test isequal(r, a + b + c)

@devec r = +(1., b, c)
@test isequal(r, 1. + b + c)

r[:] = 0.
@devec r[:] = +(1., b, c)
@test isequal(r, 1. + b + c)

@devec r = +(a, 2., c)
@test isequal(r, a + 2. + c)

r[:] = 0.
@devec r[:] = +(a, 2., c)
@test isequal(r, a + 2. + c)

@devec r = +(a, b, 3.)
@test isequal(r, a + b + 3.)

r[:] = 0.
@devec r[:] = +(a, b, 3.)
@test isequal(r, a + b + 3.)

@devec r = +(1., 2., c)
@test isequal(r, 1. + 2. + c)

r[:] = 0.
@devec r[:] = +(1., 2., c)
@test isequal(r, 1. + 2. + c)

@devec r = +(1., b, 3.)
@test isequal(r, 1. + b + 3.)

r[:] = 0.
@devec r[:] = +(1., b, 3.)
@test isequal(r, 1. + b + 3.)

@devec r = +(a, 2., 3.)
@test isequal(r, a + 2. + 3.)

r[:] = 0.
@devec r[:] = +(a, 2., 3.)
@test isequal(r, a + 2. + 3.)

@devec r[:] = +(1., 2., 3.)
@test isequal(r, fill(6., size(a)))

# comparison and blend

@devec r = blend(a + 2 .>= b, a + 2, b)
@test isequal(r, max(a+2, b))

r[:] = 0.
@devec r[:] = blend(a + 2 .>= b, a + 2, b)
@test isequal(r, max(a+2, b))

@devec r = blend(a + 2 .<= b, a + 2, b)
@test isequal(r, min(a + 2, b))

r[:] = 0.
@devec r[:] = blend(a + 2 .<= b, a + 2, b)
@test isequal(r, min(a + 2, b))

#################################################
#
#   interaction with qualified names
#
#################################################

type S
    a::Array
    b::Array
    c::Array
end

s = S(a, b, c)

@devec r = sqr(s.a)
@test isequal(r, sqr(s.a))

@devec r = s.a + s.b .* s.c
@test isequal(r, (s.a + s.b .* s.c))

@devec r = +(s.a, s.b, s.c)
@test isequal(r, s.a + s.b + s.c)

#################################################
#
#   interaction with references
#
#################################################

@devec r = sqr(abc[:,1])
@test isequal(r, sqr(a))

@devec r = sqr(abc[:,:])
@test isequal(r, sqr(abc))

@devec r = abc[:,:] .* cv
@test isequal(r, abc .* cv)

j = 3
@devec r = abc[:,1] + b[:] .* abc[:,j]
@test isequal(r, a + b .* c)

r = zeros(8, 3)
rr = zeros(size(r))
@devec r[:,2:3] = abc[:, 1:2] + cv
rr[:,2:3] = abc[:, 1:2] + cv
@test isequal(r, rr)

#################################################
#
#   op-assignment
#
#################################################

r = a
@devec r += b
@test isequal(r, a + b)

r0 = r
@devec r[:] += c
@test r === r0
@test isequal(r, a + b + c)

r = a
@devec r .*= 2.
@test isequal(r, a * 2)

@devec r -= (c + 3)
@test isequal(r, a * 2 - (c + 3))

#################################################
#
#   blocks
#
#################################################

@devec begin
    x = a + b .* c
    y = sqr(a) + c
end
@test isequal(x, a + b .* c)
@test isequal(y, sqr(a) + c)