# simple example of method of simulated moments estimated by
# NUTS MCMC

# load packages
using DelimitedFiles, TransformVariables, LogDensityProblems, DynamicHMC, MCMCDiagnostics,
    Parameters, Statistics, Distributions, ForwardDiff, LinearAlgebra

include("SVlib.jl")

# Define a structure for the problem
# Should hold the data and  the parameters of prior distributions.
struct MSM_Problem{Tm <: Vector{Float64}, Tn <: Int, Tshocks_u <: Array{Float64,2}, Tshocks_e <: Array{Float64,2}}
    "statistic"
    m::Tm
    "sample size"
    n::Tn
    "shocks"
    shocks_u::Tshocks_u
    shocks_e::Tshocks_e

end

# Make the type callable with the parameters *as a single argument*.
function (problem::MSM_Problem)(θ)
    @unpack m, n, shocks_u, shocks_e = problem   # extract the data
    @unpack σe, ρ, σu = θ         # extract parameters (only one here)
    S = size(shocks_u,2)
    k = size(m,1)
    ms = zeros(eltype(SVmodel(σe, ρ, σu, n, shocks_u[:,1], shocks_e[:,1])), S, k)
    Threads.@threads for s = 1:S
        ms[s,:] = SVmodel(σe, ρ, σu, n, shocks_u[:,s], shocks_e[:,s])
    end
    mbar = mean(ms,dims=1)[:]
    Σ = cov(ms)
    x = (m .- mbar)
    logL = try
        logL = -0.5*log(det(Σ)) - 0.5*x'*inv(Σ)*x
    catch
        logL = -Inf
    end
end
function main()
# generate data
σe = exp(-0.736/2.0)
ρ = 0.9
σu = 0.363
n = 1000 # sample size
burnin = 100
S = 100 # number of simulations
shocks_u = randn(n+burnin,1)
shocks_e = randn(n+burnin,1)
# load the example data
y = readdlm("../svdata.txt")
m =  sqrt(n)*aux_stat(y) # statistic
m = m[:]
# or use a new draw
#m = SVmodel(σe, ρ, σu, n, shocks_u, shocks_e)
shocks_u = randn(n+burnin,S) # fixed shocks for simulations
shocks_e = randn(n+burnin,S) # fixed shocks for simulations
# original problem, without transformation of parameters
p = MSM_Problem(m, n, shocks_u, shocks_e)
# define the transformation of parameters (in this case, priors are uniform on segments of real line)
# σu ~ U(0,3), ρ ~U(0,1), σe ~ U(0,3)
problem_transformation(p::MSM_Problem) =
        as((σe=as(Real, 0.0, 3.0), ρ=as(Real, 0.0, 1.0) ,σu=as(Real,0.0,3.0)))
# Wrap the problem with the transformation
t = problem_transformation(p)
P = TransformedLogDensity(t, p)
# use AD for the gradient
∇P = ADgradient(:ForwardDiff, P)
# Sample from the posterior. `chain` holds the chain (positions and
# diagnostic information), while the second returned value is the tuned sampler
# which would allow continuation of sampling.
n = dimension(problem_transformation(p))
chain, NUTS_tuned = NUTS_init_tune_mcmc(∇P, 1000)
# We use the transformation to obtain the posterior from the chain.
posterior = transform.(Ref(t), get_position.(chain));
return posterior, chain
end


export qmr

"""
 Quasi Minimal Residual (QMR) method applied to the complex symmetric system of
 linear equations Ax=b. The algorithm employed here was modified from Algorithm 8.1
 in the paper:
    Freund R W and Nachtigal N M. An implementation of the QMR method based on
      coupled two-term recurrences. SIAM J Sci., Comput., 1994, 15 (2): 313-337.

 - by Bo Han, May 2016, hanbo1735@163.com, China University of Geosciences.

Input:
    A       :: Function   - computing A*x.
    b       :: Vector     - right hand side vector.
    tol     :: Real       - error tolerance.
    maxIter :: Int        - maximum number of iterations.
    M       :: Function   - preconditioners, either matrices or functions computing M\\x.
    x0      :: Vector     - starting guess.
    out     :: Int        - flag for output (-1: no output, 0: only errors,
                                              1: final status, 2: relres at each iteration).

Output:
    x       :: Vector     - solution.
    flag    :: Int        - exit flag ( 0: desired tolerance achieved
                                       -1: maxIter reached without converging
                                       -2: delta equal to zero
                                       -3: epsilon equal to zero
                                       -9: right hand side was zero ).
    err     :: Real       - residual, i.e., norm(A*x-b)/norm(b).
    iter    :: Int        - number of iterations.
    resvec  :: Vector     - residual at each iteration.

"""
function qmr(A::SparseMatrixCSC{T,Int},b::Array{T,1}; kwargs...) where {T}
	Ax = zeros(T,size(A,2))                  # pre-allocate
	return qmr(x -> mul!(Ax,A,x,1.0,0.0),b;kwargs...)
end

qmr(A,b::Vector;kwargs...) = qmr(x -> A*x ,b;kwargs...)

function qmr(A::Function, b::Vector; tol::Real=1e-6, maxIter::Int=100, M=x->copy(x),
             x::Vector=[],out::Int=0)
	println("WARNING: QMR: this (new) funtion does not seem to work very well at the moment. Need to check algorithm correctness. ")
    n  = length(b)
    if norm(b)==0; return zeros(eltype(b),n),-9,0,0,zeros(1); end
    Mf = isa(M,Function) ? M : x -> M\x

    # allocate vector: p,d,s
    p = zeros(eltype(b),n)
    d = zeros(eltype(b),n)
    s = zeros(eltype(b),n)

    # specify scalar
    c_n      =  one(eltype(b))
    eps      =  one(eltype(b))
    theta_n  =  zero(eltype(b))
    eta      =  -one(eltype(b))

    if isempty(x)
        x = zeros(eltype(b),n)
        r  = copy(b)
    else
		r  = b - A(x)
    end

	if eltype(b) <: Complex
       x = complex(x)
    end

    resvec = zeros(maxIter+1)
    bnrm2 = norm( b )
    err   = norm( r ) / bnrm2; resvec[1] = err

    rho_np = norm( r )
    v = r / rho_np
    # v = BLAS.scal(n,one(eltype(r))/rho_np,r,1)

    if out==2
        println("=== qmr ===")
        println(@sprintf("%4s\t%7s","iter","relres"))
    end

    iter = 1
    flag = -1

    # start the iteration
    while iter <= maxIter
        c_nm     = c_n
        rho_n    = rho_np
        theta_nm = theta_n

        mv = Mf(v)
        #delta = At_mul_B(v, mv)    # delta = v.' * mv
        delta = transpose(v) * mv
        abs(delta) < 1e-40 && (flag=-2; break)

        # Update p:   p = M \ v - p * (rho*delta/eps)
        sctmp = -rho_n * delta / eps
        BLAS.scal!(n,sctmp,p,1)                 #  - p * (rho*delta/eps)
        BLAS.axpy!(n,one(eltype(p)),mv,1,p,1)   # p = p + mv

        ap   = A(p)
        # eps  = At_mul_B(p, ap)    # eps  = p.' * ap
        eps  = transpose(p) * ap
        abs(eps) < 1e-40 && (flag=-3; break)
        beta = eps / delta

        # Update v:   v = A * p - beta * v,   v=v/norm(v)
        BLAS.scal!(n,-beta,v,1)                  #  -beta * v
        BLAS.axpy!(n,one(eltype(v)),ap,1,v,1)    # v = v + ap
        rho_np = norm( v )
        BLAS.scal!(n,one(eltype(v))/rho_np,v,1)  # v = v / rho_np

        theta_n = rho_np^2 / c_nm / (real(beta)^2+imag(beta)^2)    # theta_n^2
        c_n     = 1.0 / (1.0 + theta_n)                            # c_n^2
        eta     = -eta * rho_n * c_n / (beta * c_nm)
        k2      = theta_nm * c_n

        # Update d:   d = eta * p + k2 * d
        BLAS.scal!(n,one(eltype(b))*k2,d,1)     # k2 * d
        BLAS.axpy!(n,eta,p,1,d,1)               # d = d + eta * p

        # Update x:  x = x + d
        BLAS.axpy!(n,one(eltype(d)),d,1,x,1)

        # Update s:  s = eta * ap + k2 * s
        BLAS.scal!(n,one(eltype(s))*k2,s,1)     # k2 * s
        BLAS.axpy!(n,eta,ap,1,s,1)              # s = s + eta * ap

        # Update r:  r = r - s
        BLAS.axpy!(n,-one(eltype(r)),s,1,r,1)

        err = norm( r ) / bnrm2
        resvec[iter+1] = err

        out==2 && println(@sprintf("%3d\t%1.2e",iter,resvec[iter+1]))
        iter += 1

        err <= tol && (flag=0; break)
    end  # iter = 1:maxIter
    iter = min(iter, maxIter)

    if out>=0
      if flag==-1
        println(@sprintf("qmr iterated maxIter (=%d) times but reached only residual norm %1.2e instead of tol=%1.2e.",
						 maxIter,resvec[iter],tol))
      elseif flag==-2
        println(@sprintf("qmr: delta equal to zero at iteration %d. Returned residual has norm %1.2e.", iter,resvec[iter]))
      elseif flag==-3
        println(@sprintf("qmr: epsilon equal to zero at iteration %d. Returned residual has norm %1.2e.", iter,resvec[iter]))
      elseif out>=1
	    println(@sprintf("qmr achieved desired tolerance at iteration %d. Residual norm is %1.2e.",iter,resvec[iter+1]))
	  end
    end

	return x, flag,resvec[iter+1],iter,resvec[1:iter+1]

end


import DifferentialEquations;	const DE = DifferentialEquations
"""
	dOmega!(dOmega, Omega, args, t)
Differential equation for the first terms of the Magnus expansion.
#Arguments
- dOmega: 2-dimensional array with the time derivative of the Magnus expansion
terms.
- Omega: 2-dimensional array with the Magnus expansion terms. The nth column is
the nth term of the Magnus expansion: Omega[:,n] = Omega_n.
- args: Array with the many different arguments of the function.
args[1] = order: The number of terms of the Magnus expansion that will be computed.\n
args[2] = A: the matrix of function that defines the linear ODE (dx/dt = A*x).\n
args[3] = S: array of arrays used to store the S operators. ``S[n] = S_n``, where
S[n] is a 2-dimensional array. ``(S[n])[:,j] = S_n^{(j)}``.\n
args[4] = B: array where B[n] = b_n/n!. b_n is the n'th Bernoulli number.\n
args[5] = Snj: array to temporarily store S_n^{(j)}\n
args[6] = temp1 and args[7] = temp2: temporary arrays to store data.\n
- t: time variable.\n
"""
function dOmega!(dOmega, Omega, args, t)
	order = args[1]::Integer
	A = args[2]
	comm! = args[3]
	S = args[4]
	B = args[5]
	Snj = args[6]
	temp = args[7]
	temp2 = args[8]

	dOmega[:, 1] .= A(t)

	S2 = S[2]
	Omega1 = view(Omega, :, 1)
	dOmega1 = view(dOmega, :, 1)
	comm!(Snj, Omega1, dOmega1)
	S2[:, 1] .= Snj
	dOmega[:, 2] .= B[1].*Snj
	for n = 3: order
		Sn = S[n]
		update_S!(Snj, S, n, 1, comm!, Omega, Omega1, dOmega1, temp2)
		temp .= B[1].*Snj
		for j = 2: n - 1
			update_S!(Snj, S, n, j, comm!, Omega, Omega1, dOmega1, temp2)
			if j%2 == 0
				temp .= temp .+ B[j].*Snj
			end
		end
		dOmega[:, n] .= temp
	end
	return nothing
end
"""
	update_S!(Snj, S, n, j, Omega, Omega1, dOmega1, temp2)
Calculate ``S_n^{(j)}`` recursivelly and store the result in S and Snj. The result
is stored in Snj just for efficiency purposes.
#Arguments
See dOmega! documentation.
"""
function update_S!(Snj, S, n, j, comm!, Omega, Omega1, dOmega1, temp2)
	Sn = S[n]
	if j == 1
		Omegam = view(Omega, :, n - 1)
		comm!(Snj, Omegam, dOmega1)
	else
		comm!(Snj, Omega1, view(S[n - 1], :, j - 1))
		for m = 2: n - j
			Omegam = view(Omega, :, m)
			comm!(temp2, Omegam, view(S[n - m], :, j - 1))
			Snj .= Snj .+ temp2
		end
	end
	Sn[:, j] .= Snj
	return nothing
end
"""
	bernoulli(n)
Calculate the n'th Bernoulli number. This function was taken from Rosetta Code
(https://rosettacode.org/wiki/Bernoulli_numbers#Julia).
#Arguments:
- n: the index of the Bernoulli number.
"""
function bernoulli(n)
    A = Vector{Rational{Int64}}(undef, n + 1)
    for m = 0: n
        A[m + 1] = 1//(m + 1)
        for j = m: -1: 1
            A[j] = j*(A[j] - A[j + 1])
        end
    end
    return A[1]
end
"""
	solveOmega(tf, order, A, numdims)
Calculate the Magnus expansion.
#Arguments
- tf: the time at which the Magnus expansion is calculated. It is assumed that
the time evolution goes from 0 to tf.
- order: the number of calculated terms in the Magnus expansion.
- A: the matrix of function that defines the linear ODE (dx/dt = A*x).
- numdims: the number of dimensions of the ODE.
"""
function solveOmega(tf, order, A, comm!, numdims)
	order_aux = order
	order = max(2, order)

	numtype = typeof(tf)
	S = createS(order, numdims, numtype)
	B = zeros(Rational{Int64}, order)
	for j = 1: order
		B[j] = bernoulli(j)//factorial(j)
	end
	B = convert(Array{numtype}, B)
	B = Tuple(B)

	Omega0 = zeros(numtype, numdims, order)
	tspan = (zero(tf), tf)
	temp1 = zeros(numtype, numdims)
	temp2 = zeros(numtype, numdims)
	temp3 = zeros(numtype, numdims)
	args = (order, A, comm!, S, B, temp1, temp2, temp3)
	prob = DE.ODEProblem(dOmega!, Omega0, tspan, args)
	sol = DE.solve(prob, DE.Vern9(), saveat=[tf], reltol=1e-11, abstol=1e-14,
				   maxiters = 10^6)

	if sol.retcode == :Success
		return sol(tf)[:, 1:order_aux]
	else
		error("Calculation of the Magnus series was unsuccessful.")
	end
end
"""
	createS(order, numdims, numtype)
Create the S array, which is an argument of the dOmega! function.
#Arguments:
-numtype: The type of the numbers used in problem.
See solveOmega documentation.
"""
function createS(order, numdims, numtype)
	S = [zeros(numtype, numdims, 1)]
	for n=2:order
		push!(S, zeros(numtype, numdims, n-1))
	end
	return S
end

### TEST ###


println("\n\n\nRunning tests ...")

using RandomFerns
using Base.Test

using FunctionalData

D = 2
N = 100
C = 10
S = 10
data = @p map (1:C) (x->randn(D,N).+5x.+randn(D,1)) | unstack | flatten
labels = @p map (1:C) (x->x*ones(1,N)) | unstack | flatten

a = RandFerns(data,labels);
# testdata = @p meshgrid linspace(-15,15,200) linspace(-15,15,200)
r = predict(a, data)
@test sum(vec(r).!=vec(labels)) < 100

a = RandFerns(data,labels; regression = true);
r = predict(a, data)
@test mean(abs(vec(r)-vec(labels))) < 0.5

println("   done running tests!")


type VcovWhite <: AbstractVcovMethod end

type VcovWhiteData <: AbstractVcovMethodData end

VcovMethodData(::VcovWhite, ::AbstractDataFrame) = VcovWhiteData()

function vcov!(v::VcovWhiteData, x::VcovData)
    S = shat!(v, x)
    return sandwich(x.crossmatrix, S)
end

function shat!{T}(::VcovWhiteData, x::VcovData{T, 1})
    X = x.regressors
    res = x.residuals
    Xu = scale!(res, X)
    S = At_mul_B(Xu, Xu)
    scale!(S, size(X, 1) / x.df_residual)
    return S
end

# S_{(l-1) * K + k, (l'-1)*K + k'} = \sum_i X[i, k] res[i, l] X[i, k'] res[i, l']
function shat!{T}(::VcovWhiteData, x::VcovData{T, 2})
    X = x.regressors
    res = x.residuals
    nobs = size(X, 1)
    dim = size(X, 2) * size(res, 2)
    S = fill(zero(Float64), (dim, dim))
    temp = fill(zero(Float64), nobs, dim)
    index = zero(Int)
    @inbounds for k in 1:size(X, 2), l in 1:size(res, 2)
        index += 1
        @simd for i in 1:nobs
            temp[i, index] = X[i, k]* res[i, l]
        end
    end
    S = At_mul_B(temp, temp)
    scale!(S, nobs / x.df_residual)
    return S
end


module PLaplacianTests
using SparseMatricesCSR
using Gridap
using Gridap.Algebra
using GridapDistributed
using PartitionedArrays
using Test
using SparseArrays

function main(parts)
  main(parts,FullyAssembledRows(),SparseMatrixCSR{0,Float64,Int})
  main(parts,SubAssembledRows(),SparseMatrixCSC{Float64,Int})
end

function main(parts,strategy,local_matrix_type)

  output = mkpath(joinpath(@__DIR__,"output"))

  domain = (0,4,0,4)
  cells = (4,4)
  model = CartesianDiscreteModel(parts,domain,cells)

  k = 1
  u((x,y)) = (x+y)^k
  σ(∇u) =(1.0+∇u⋅∇u)*∇u
  dσ(∇du,∇u) = 2*∇u⋅∇du + (1.0+∇u⋅∇u)*∇du
  f(x) = -divergence(y->σ(∇(u,y)),x)

  Ω = Triangulation(strategy,model)
  dΩ = Measure(Ω,2*k)
  r(u,v) = ∫( ∇(v)⋅(σ∘∇(u)) - v*f )dΩ
  j(u,du,v) = ∫( ∇(v)⋅(dσ∘(∇(du),∇(u))) )dΩ

  reffe = ReferenceFE(lagrangian,Float64,k)
  V = TestFESpace(model,reffe,dirichlet_tags="boundary")
  U = TrialFESpace(u,V)

  assem=SparseMatrixAssembler(local_matrix_type,Vector{Float64},U,V,strategy)
  op = FEOperator(r,j,U,V,assem)

  uh = zero(U)
  b,A = residual_and_jacobian(op,uh)
  _A = copy(A)
  _b = copy(b)
  residual_and_jacobian!(_b,_A,op,uh)
  @test (norm(_b-b)+1) ≈ 1
  x = similar(b,Float64,axes(A,2))
  fill!(x,1)
  @test (norm(A*x-_A*x)+1) ≈ 1

  # This leads to a dead lock since the printing of the trace seems to lead to collective operations
  # show_trace = i_am_main(parts)
  show_trace = true # The output in MPI will be ugly, but fixing this would require to edit NLSolvers package.
  nls = NLSolver(show_trace=show_trace, method=:newton)
  solver = FESolver(nls)
  uh = solve(solver,op)

  Ωo = Triangulation(model)
  dΩo = Measure(Ωo,2*k)
  eh = u - uh
  @test sqrt(sum(∫( abs2(eh) )dΩo)) < 1.0e-9

end

end # module


immutable MPoly{T}
    terms::OrderedDict{Vector{Int},T}
    vars::Vector{Symbol}
end

terms(p::MPoly) = p.terms

eltype{T}(::MPoly{T}) = T

vars(p::MPoly) = p.vars

nvars(p::MPoly) = length(vars(p))

zero{T}(::Type{MPoly{T}}; vars::Vector{Symbol}=Symbol[]) = MPoly{T}(OrderedDict{Vector{Int},T}(), vars)
zero{T}(p::MPoly{T}) = zero(MPoly{T}, vars=vars(p))

one{T}(::Type{MPoly{T}}; vars::Vector{Symbol}=Symbol[]) = MPoly{T}(OrderedDict(zeros(Int, length(vars)) => one(T)), vars)
one{T}(p::MPoly{T}) = one(MPoly{T}, vars=vars(p))

function generators{T}(::Type{MPoly{T}}, vars::Symbol...)
     m = eye(Int, length(vars))
     [MPoly{T}(OrderedDict(m[:,i] => one(T)), [vars...]) for i = 1:length(vars)]
end

generator{T}(::Type{MPoly{T}}, var::Symbol) =
     MPoly{T}(OrderedDict([1] => one(T)), [var])

call{T}(::Type{MPoly{T}}, var::Symbol) =
     MPoly{T}(OrderedDict([1] => one(T)), [var])

call(::Type{MPoly}, var::Symbol) =
     MPoly{Float64}(var)

promote_rule{T,U}(::Type{MPoly{T}}, ::Type{MPoly{U}}) = MPoly{promote_type(T, U)}
promote_rule{T,U}(::Type{MPoly{T}}, ::Type{U}) = MPoly{promote_type(T, U)}

function convert{T}(P::Type{MPoly{T}}, p::MPoly{T})
    p
end

function convert{T}(P::Type{MPoly{T}}, p::MPoly)
    r = zero(P, vars=vars(p))
    for (m, c) in p
        r[m] = convert(T, c)
    end
    r
end

convert{T}(::Type{MPoly{T}}, c::T) =
    MPoly{T}(OrderedDict(Int[] => c))

monomials(p::MPoly) = keys(terms(p))

getindex{T}(p::MPoly{T}, m::Vector{Int}) =
    get(terms(p), m, zero(T))

getindex(p::MPoly, m::Int...) = p[[m...]]

function setindex!{T}(p::MPoly{T}, v, m::Int...)
    v = convert(T, v)
    if method_exists(isapprox, (T,T)) && isapprox(v, zero(T))
        delete!(terms(p), [m...])
    else
        terms(p)[[m...]] = v
    end
end

setindex!{T}(p::MPoly{T}, v, m::Vector{Int}) = p[m...] = v

start(p::MPoly) = start(terms(p))
next(p::MPoly, state) = next(terms(p), state)
done(p::MPoly, state) = done(terms(p), state)

copy{T}(p::MPoly{T}) = MPoly{T}(copy(terms(p)), copy(vars(p)))

function deg(p::MPoly)
    d = 0
    for m in monomials(p)
        s = sum(m)
        if s > d
            d = s
        end
    end
    d
end

function ==(p::MPoly, q::MPoly)
    for (m, c) in p
        if !isapprox(q[m], c)
            return false
        end
    end
    for (m, c) in q
        if !isapprox(p[m], c)
            return false
        end
    end
    true
end


"""
Tableau for the explicit 4-stage RS1 method due to Andreas Rossler
  Method cited in Table 5.2 in
  Andreas Rossler, "Second order Runge-Kutta methods for Stratonovich stochastic differential equations",
  BIT Numerical Mathematics (2007) 47
  According to the paper, the method has weak order 2.0.
"""
function getTableauRosslerRS1()

    A0 = [[0.0   0.0   0.0   0.0]
          [0.0   0.0   0.0   0.0]
          [1.    0.0   0.0   0.0]
          [0.0   0.0   0.0   0.0]]

    A1 = [[0.0   0.0   0.0   0.0]
          [0.0   0.0   0.0   0.0]
          [1.    0.0   0.0   0.0]
          [1.    0.0   0.0   0.0]]

    A2 =  zeros(Float64,4,4)

    B0 = [[0.0     0.0    0.0   0.0]
          [0.0     0.0    0.0   0.0]
          [1. / 4.   3. / 4.  0.0   0.0]
          [0.0     0.0    0.0   0.0]]

    B1 = [[0.0     0.0    0.0   0.0]
          [2. / 3.   0.0    0.0   0.0]
          [1. / 12.  1. / 4.  0.0   0.0]
          [-5. / 4.  1. / 4.  2.    0.0]]

    B2 = [[0.0     0.0    0.0   0.0]
          [1.      0.0    0.0   0.0]
          [-1.     0.0    0.0   0.0]
          [0.0     0.0    0.0   0.0]]

    B3 = [[0.0     0.0    0.0   0.0]
          [0.0     0.0    0.0   0.0]
          [1. / 4.   3. / 4.  0.0   0.0]
          [1. / 4.   3. / 4.  0.0   0.0]]

    α  =  [0.0, 0.0, 0.5, 0.5]
    β1 =  [1. / 8., 3. / 8., 3. / 8., 1. / 8.]
    β2 =  [0.0, -0.25, 0.25, 0.0]

    c0 =  [0.0, 0.0, 1., 0.0]
    c1 =  [0.0, 0.0, 1., 1.]
    c2 =  zeros(Float64, 4)

    TableauWERK(:RosslerRS1_explicit_method, A0, A1, A2, B0, B1, B2, B3, α, β1, β2, c0, c1, c2)
end

"""
Tableau for the explicit 4-stage RS2 method due to Andreas Rossler
  Method cited in Table 5.3 in
  Andreas Rossler, "Second order Runge-Kutta methods for Stratonovich stochastic differential equations",
  BIT Numerical Mathematics (2007) 47
  According to the paper, the method has weak order 2.0.
"""
function getTableauRosslerRS2()

    A0 = [[0.0      0.0      0.0   0.0]
          [2. / 3.  0.0      0.0   0.0]
          [1. / 6.  1. / 2.  0.0   0.0]
          [0.0      0.0      0.0   0.0]]

    A1 = [[0.0   0.0   0.0   0.0]
          [0.0   0.0   0.0   0.0]
          [1.    0.0   0.0   0.0]
          [1.    0.0   0.0   0.0]]

    A2 =  zeros(Float64,4,4)

    B0 = [[0.0       0.0      0.0   0.0]
          [0.0       0.0      0.0   0.0]
          [1. / 4.   3. / 4.  0.0   0.0]
          [0.0       0.0      0.0   0.0]]

    B1 = [[0.0       0.0      0.0   0.0]
          [2. / 3.   0.0      0.0   0.0]
          [1. / 12.  1. / 4.  0.0   0.0]
          [-5. / 4.  1. / 4.  2.    0.0]]

    B2 = [[0.0     0.0    0.0   0.0]
          [1.      0.0    0.0   0.0]
          [-1.     0.0    0.0   0.0]
          [0.0     0.0    0.0   0.0]]

    B3 = [[0.0       0.0      0.0   0.0]
          [0.0       0.0      0.0   0.0]
          [1. / 4.   3. / 4.  0.0   0.0]
          [1. / 4.   3. / 4.  0.0   0.0]]

    α  =  [0.25, 0.25, 0.5, 0.0]
    β1 =  [1. / 8., 3. / 8., 3. / 8., 1. / 8.]
    β2 =  [0.0, -0.25, 0.25, 0.0]

    c0 =  [0.0, 2. / 3., 2. / 3., 0.0]
    c1 =  [0.0, 0.0, 1., 1.]
    c2 =  zeros(Float64, 4)

    TableauWERK(:RosslerRS2_explicit_method, A0, A1, A2, B0, B1, B2, B3, α, β1, β2, c0, c1, c2)
end


"""
    test_bqg_rossbywave(stepper, dt, nsteps, dev; kwargs...)

Evolves a Rossby wave and compares with the analytic solution.
"""
function test_bqg_rossbywave(stepper, dt, nsteps, dev::Device=CPU())
    nx = 64
    Lx = 2π
     β = 2.0
     μ = 0.0
     ν = 0.0
     T = Float64

  # the following if statement is called so that all the cases of
  # Problem() fuction are tested
  if stepper=="ForwardEuler"
    eta = zeros(dev, T, (nx, nx))
  else
    eta(x, y) = 0*x
  end

  prob = BarotropicQG.Problem(dev; nx=nx, Lx=Lx, eta=eta, β=β, μ=μ, ν=ν, stepper=stepper, dt=dt)
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  x, y = gridpoints(g)

  # the Rossby wave initial condition
   ampl = 1e-2
  kwave = 3 * 2π/g.Lx
  lwave = 2 * 2π/g.Ly
      ω = -p.β*kwave/(kwave^2 + lwave^2)
     ζ0 = @. ampl*cos(kwave*x)*cos(lwave*y)
    ζ0h = rfft(ζ0)

  BarotropicQG.set_zeta!(prob, ζ0)

  stepforward!(prob, nsteps)
  dealias!(sol, g)
  BarotropicQG.updatevars!(prob)

  ζ_theory = @. ampl*cos(kwave*(x - ω/kwave*cl.t)) * cos(lwave*y)

  return isapprox(ζ_theory, v.zeta, rtol=g.nx*g.ny*nsteps*1e-12)
end

"""
    test_bqg_stochasticforcingbudgets(dev; kwargs...)

Tests if the energy budget is closed for BarotropicQG problem with stochastic forcing.
"""
function test_bqg_stochasticforcingbudgets(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, T=Float64)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

    gr = TwoDGrid(dev, n, L)
  x, y = gridpoints(gr)

  # Forcing
  kf, dkf = 12.0, 2.0
  ε = 0.1

  CUDA.@allowscalar Kr = ArrayType(dev)([gr.kr[i] for i=1:gr.nkr, j=1:gr.nl])

  forcingcovariancespectrum = zeros(dev, T, (gr.nkr, gr.nl))
  @. forcingcovariancespectrum = exp( -(sqrt(gr.Krsq) - kf)^2 / (2 * dkf^2) )
  CUDA.@allowscalar @. forcingcovariancespectrum[gr.Krsq .< 2^2] = 0
  CUDA.@allowscalar @. forcingcovariancespectrum[gr.Krsq .> 20^2] = 0
  CUDA.@allowscalar @. forcingcovariancespectrum[Kr .< 2π/L] = 0
  ε0 = parsevalsum(forcingcovariancespectrum .* gr.invKrsq / 2, gr) / (gr.Lx * gr.Ly)
  forcingcovariancespectrum .= ε / ε0 * forcingcovariancespectrum

  Random.seed!(1234)

  function calcFq!(Fqh, sol, t, clock, vars, params, grid)
    eta = ArrayType(dev)(exp.(2π * im * rand(T, size(sol))) / sqrt(clock.dt))
    CUDA.@allowscalar eta[1, 1] = 0
    @. Fqh = eta * sqrt(forcingcovariancespectrum)
    return nothing
  end

  prob = BarotropicQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcFq=calcFq!, stochastic=true)

  BarotropicQG.set_zeta!(prob, 0*x)

  E = Diagnostic(BarotropicQG.energy,             prob, nsteps=nt)
  D = Diagnostic(BarotropicQG.energy_dissipation, prob, nsteps=nt)
  R = Diagnostic(BarotropicQG.energy_drag,        prob, nsteps=nt)
  W = Diagnostic(BarotropicQG.energy_work,        prob, nsteps=nt)
  diags = [E, D, W, R]

  stepforward!(prob, diags, nt)

  BarotropicQG.updatevars!(prob)

  dEdt_numerical = (E[2:E.i] - E[1:E.i-1]) / prob.clock.dt

  # If the Ito interpretation was used for the work
  # then we need to add the drift term
  # dEdt_computed = W[2:E.i] + ε - D[1:E.i-1] - R[1:E.i-1]      # Ito
  dEdt_computed = W[2:E.i] - D[1:E.i-1] - R[1:E.i-1]        # Stratonovich

  return isapprox(dEdt_numerical, dEdt_computed, rtol=1e-3)
end

"""
    test_bqg_deterministicforcingbudgets(dev; kwargs...)

Tests if the energy budget is closed for BarotropicQG problem with deterministic forcing.
"""
function test_bqg_deterministicforcingbudgets(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ = 1e-1
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  gr = TwoDGrid(dev, n, L)
  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly

  # Forcing = 0.01cos(4x)cos(5y)cos(2t)
  f = @. 0.01 * cos(4k₀*x) * cos(5l₀*y)
  fh = rfft(f)

  function calcFq!(Fqh, sol, t, clock, vars, params, grid)
    @. Fqh = fh * cos(2t)
    return nothing
  end

  prob = BarotropicQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt,
   stepper="RK4", calcFq=calcFq!, stochastic=false)

  BarotropicQG.set_zeta!(prob, 0*x)

  E = Diagnostic(BarotropicQG.energy,             prob, nsteps=nt)
  D = Diagnostic(BarotropicQG.energy_dissipation, prob, nsteps=nt)
  R = Diagnostic(BarotropicQG.energy_drag,        prob, nsteps=nt)
  W = Diagnostic(BarotropicQG.energy_work,        prob, nsteps=nt)
  diags = [E, D, W, R]

  stepforward!(prob, diags, round(Int, nt))

  BarotropicQG.updatevars!(prob)

  dEdt_numerical = (E[3:E.i] - E[1:E.i-2]) / (2 * prob.clock.dt)
  dEdt_computed  = W[2:E.i-1] - D[2:E.i-1] - R[2:E.i-1]

  residual = dEdt_numerical - dEdt_computed

  return isapprox(dEdt_numerical, dEdt_computed, atol=1e-10)
end

"""
    test_bqg_nonlinearadvection(dt, stepper, dev; kwargs...)

Tests the advection term in the BarotropicQG module by timestepping a
test problem with timestep dt and timestepper identified by the string stepper.
The test problem is derived by picking a solution ζf (with associated
streamfunction ψf) for which the advection term J(ψf, ζf) is non-zero. Next, a
forcing Ff is derived according to Ff = ∂ζf/∂t + J(ψf, ζf) - ν∇²ζf. One solution
to the vorticity equation forced by this Ff is then ζf. (This solution may not
be realized, at least at long times, if it is unstable.)
"""
function test_bqg_advection(dt, stepper, dev::Device=CPU(); n=128, L=2π, ν=1e-2, nν=1, μ=0.0)
  n, L  = 128, 2π
  ν, nν = 1e-2, 1
   μ = 0.0
  tf = 1.0
  nt = round(Int, tf/dt)

  gr = TwoDGrid(dev, n, L)
  x, y = gridpoints(gr)

  ψf = @. sin(2x) * cos(2y) + 2sin(x) * cos(3y)
  qf = @. -8sin(2x) * cos(2y) - 20sin(x) * cos(3y)

  Ff = @. -(
    ν*( 64sin(2x) * cos(2y) + 200sin(x) * cos(3y) )
    + 8 * ( cos(x) * cos(3y) * sin(2x) * sin(2y) - 3cos(2x) * cos(2y) * sin(x) * sin(3y) )
    )

  Ffh = rfft(Ff)

  # Forcing
  function calcFq!(Fqh, sol, t, cl, v, p, g)
    Fqh .= Ffh
    return nothing
  end

  prob = BarotropicQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt, stepper=stepper, calcFq=calcFq!)

  BarotropicQG.set_zeta!(prob, qf)

  stepforward!(prob, round(Int, nt))

  BarotropicQG.updatevars!(prob)

  return isapprox(prob.vars.q, qf, rtol=rtol_barotropicQG)
end

"""
    test_bqg_formstress(dt, stepper, dev; kwargs...)

Tests the form stress term that forces the domain-averaged zonal flow U(t).
"""
function test_bqg_formstress(dt, stepper, dev::Device=CPU(); n=128, L=2π, ν=0.0, nν=1, μ=0.0)
  n, L  = 128, 2π
  ν, nν = 1e-2, 1
  μ = 0.0
  tf = 1
  nt = 1

  gr  = TwoDGrid(dev, n, L)
  x, y = gridpoints(gr)

  ζ₀ = @. -20 * sin(10x) * cos(10y)
  topoPV(x, y) = @. cos(10x) * cos(10y)
  F(t) = 0 #no forcing

  answer = 0.25 # this is what <v*eta> should be

  prob = BarotropicQG.Problem(dev; nx=n, Lx=L, ν=ν, nν=nν, μ=μ, dt=dt, stepper=stepper, eta=topoPV, calcFU=F)
  BarotropicQG.set_zeta!(prob, ζ₀)
  BarotropicQG.updatevars!(prob)

  # Step forward
  stepforward!(prob, nt)
  return CUDA.@allowscalar isapprox(prob.timestepper.N[1, 1], answer, rtol=rtol_barotropicQG)
end

"""
    test_bqg_energyenstrophy(dev)

Tests the energy and enstrophy function for a BarotropicQG problem.
"""
function test_bqg_energyenstrophy(dev::Device=CPU())
  nx, Lx  = 64, 2π
  ny, Ly  = 64, 3π
  g  = TwoDGrid(dev, nx, Lx, ny, Ly)
  k₀, l₀ = 2π/Lx, 2π/Ly # fundamental wavenumbers
  x, y = gridpoints(g)

  energy_calc = 29/9
  enstrophy_calc = 2701/162

    eta = @. cos(10k₀*x) * cos(10l₀*y)
   psi0 = @. sin(2k₀*x) * cos(2l₀*y) + 2sin(k₀*x) * cos(3l₀*y)
  zeta0 = @. - ((2k₀)^2+(2l₀)^2) * sin(2k₀*x) * cos(2l₀*y) - (k₀^2+(3l₀)^2) * 2sin(k₀*x) * cos(3l₀*y)

  prob = BarotropicQG.Problem(dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, eta = eta, stepper="ForwardEuler")
  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid
  BarotropicQG.set_zeta!(prob, zeta0)
  BarotropicQG.updatevars!(prob)

  energyzeta0 = BarotropicQG.energy(prob)
  enstrophyzeta0 = BarotropicQG.enstrophy(prob)

  return isapprox(energyzeta0, energy_calc, rtol=rtol_barotropicQG) && isapprox(enstrophyzeta0, enstrophy_calc, rtol=rtol_barotropicQG) &&
  BarotropicQG.addforcing!(prob.timestepper.N, sol, cl.t, cl, v, p, g)==nothing
end

"""
    test_bqg_energyenstrophy(dev)

Tests the mean flow U(t) energy and enstrophy function for a BarotropicQG problem.
"""
function test_bqg_meanenergyenstrophy(dev::Device=CPU())
  nx, Lx  = 64, 2π
  ny, Ly  = 96, 3π
  g = TwoDGrid(dev, nx, Lx, ny, Ly)
  k₀, l₀ = 2π/Lx, 2π/Ly # fundamental wavenumbers
  x, y = gridpoints(g)

  calcFU(t) = 0.0
  eta(x, y) = @. cos(10x) * cos(10y)
  psi0 = @. sin(2k₀*x) * cos(2l₀*y) + 2sin(k₀*x) * cos(3l₀*y)
 zeta0 = @. - ((2k₀)^2+(2l₀)^2) * sin(2k₀*x) * cos(2l₀*y) - (k₀^2+(3l₀)^2) * 2sin(k₀*x) * cos(3l₀*y)
  β = 10.0
  U = 1.2

  energy_calc = 29/9
  enstrophy_calc = 2701/162

  prob = BarotropicQG.Problem(dev; nx=nx, Lx=Lx, ny=ny, Ly=Ly, β=β, eta=eta, calcFU = calcFU,
                                    stepper="ForwardEuler")

  BarotropicQG.set_zeta!(prob, zeta0)
  BarotropicQG.set_U!(prob, U)
  BarotropicQG.updatevars!(prob)

  energyU = BarotropicQG.meanenergy(prob)
  enstrophyU = BarotropicQG.meanenstrophy(prob)

  energyzeta0 = BarotropicQG.energy(prob)
  enstrophyzeta0 = BarotropicQG.enstrophy(prob)

  return (isapprox(energyU, 0.5*U^2, rtol=rtol_barotropicQG) &&
    isapprox(enstrophyU, β*U, rtol=rtol_barotropicQG) &&
    isapprox(energyzeta0, energy_calc, rtol=rtol_barotropicQG) &&
    isapprox(enstrophyzeta0, enstrophy_calc, rtol=rtol_barotropicQG)
  )
end

"""
    test_bqg_problemtype(dev, T)

Tests the BarotropicQG problem constructor for different DataType `T`.
"""
function test_bqg_problemtype(dev, T)
  prob = BarotropicQG.Problem(dev; T=T)

  A = ArrayType(dev)

  return (typeof(prob.sol)<:A{Complex{T}, 2} && typeof(prob.grid.Lx)==T && eltype(prob.grid.x)==T && typeof(prob.vars.u)<:A{T, 2})
end


to_int(s) = parse(Int, s)

using DataStructures

function pass_day(fish)
    zero_count = length(filter(==(0), fish))
    fish = map(grow, fish)
    for _ = 1:zero_count
        append!(fish, 8)
    end
    fish
end

function solve(filename, num_days)
    fish = parse_file(filename)
    for day = 1:num_days
        fish = pass_day(fish)
    end
    length(fish)
end

solve("6eg.txt", 80) == 5934
solve("6.txt", 80)

grow(n) = n == 0 ? 6 : n - 1
parse_file(filename) = parse.(Int, split(read(filename, String), ","))

function solve2(filename, num_days)
    fishes = parse_file(filename)
    counts = counter(fishes)
    for _ in 1:num_days
        new_counts = counter(Int)
        for (fish, count) in counts
            new_counts[grow(fish)] = count
        end
        new_counts[8] += new_counts[0]
        counts = new_counts
    end
    sumcounts
end
solve2("6eg.txt", 256)
solve2("6eg.txt", 256) == 26984457539
solve2("6.txt", 80) == 5934


# Function for Gibbs sampling multivariate density via slice sampling
function mv_slice_sampler(x::Array{Float64}, f::Function, opts::Options)
  P = length(x)
  gxp = 0
  for p = 1:length(x)
    function fp(w)  # want to fix all but the p'th value of this array
      x[p] = w
      f(x)
    end
    # TODO: initialize at last point
    # TODO: Pass on options to univariate slice sampler
    x[p], gxp = slice_sampler(x[p], fp)
  end
  x, gxp
end
mv_slice_sampler(x::Array{Float64}, d::Union(Density,DifferentiableDensity), opts::Options) = mv_slice_sampler(x, d.f, opts)


module TestGradient
    using Base.Test
    import FiniteDiff

    f1(x) = sin(x[1]) + cos(x[2])
    function g1!(out, x)
        out[1] = cos(x[1])
        out[2] = -sin(x[2])
        return
    end

    f2(x) = 2.3 * sin(x[1]) + 3.1 * cos(x[2])
    function g2!(out, x)
        out[1] = 2.3 * cos(x[1])
        out[2] = -3.1 * sin(x[2])
        return
    end

    f3(x) = 0.1 * (x[1] - 1.0)^2 - 0.3 * (x[2] - 0.7)^2
    function g3!(out, x)
        out[1] = 0.1 * 2.0 * (x[1] - 1.0)
        out[2] = -0.3 * 2.0 * (x[2] - 0.7)
        return
    end

    function check_error(y_true, y_approximate, err = 1 // 1_000)
        n = length(y_true)
        @test length(y_approximate) == n
        for i in 1:n
            if y_true[i] > eps(y_true[i])
                @test abs(y_true[i] - y_approximate[i]) / abs(y_true[i]) < err
            end
        end
    end

    function run_tests(n::Integer)
        funcs = (
            (f1, g1!),
            (f2, g2!),
            (f3, g3!),
        )

        x = Array(Float64, 2)
        output = Array(Float64, 2)
        buffer = Array(Float64, 2)
        buffer_complex = Array(Complex{Float64}, 2)
        gr = Array(Float64, 2)

        n_dims = length(x)

        modes = (
            FiniteDiff.ForwardMode(),
            FiniteDiff.BackwardMode(),
            FiniteDiff.CentralMode(),
            FiniteDiff.ComplexMode(),
        )

        for (f, f′!) in funcs
            for itr in 1:n
                for i in 1:n_dims
                    x[i] = itr
                end

                # Evaluate the true gradient and store it into gr.
                f′!(gr, x)

                for mode in modes
                    if mode != FiniteDiff.ComplexMode()
                        FiniteDiff.gradient!(output, f, x, buffer, mode)
                    else
                        FiniteDiff.gradient!(output, f, x, buffer_complex, mode)
                    end
                    check_error(gr, output)

                    output2 = FiniteDiff.gradient(f, x, mode)
                    check_error(gr, output2)

                    tmp! = FiniteDiff.gradient(f, mode, mutates = true)
                    if mode != FiniteDiff.ComplexMode()
                        tmp!(output, x, buffer)
                    else
                        tmp!(output, x, buffer_complex)
                    end
                    check_error(gr, output)

                    tmp = FiniteDiff.gradient(f, mode, mutates = false)
                    output2 = tmp(x)
                    check_error(gr, output2)
                end
            end
        end
    end

    @testset "gradient tests" begin
        run_tests(100)
    end
end


using TSPDrone
using Test, Statistics
# include("../src/main.jl")

include("test_instances.jl")

@testset verbose = true "TSPDrone.jl" begin

    @testset verbose = true "Cost Matrix" begin
        for t = 1:10
            n = 10
            x = rand(n)
            y = rand(n)
            dist_mtx = zeros(n, n)
            for i in 1:n, j in 1:n
                dist_mtx[i, j] = sqrt( (x[i] - x[j])^2 +(y[i] - y[j])^2 )
            end
            result1 = solve_tspd(x, y, 1.0, 0.5)
            result2 = solve_tspd(dist_mtx, dist_mtx .* 0.5)

            @show result1.total_cost, result1.truck_route, result1.drone_route
            @show result2.total_cost, result2.truck_route, result2.drone_route

            @test result1.total_cost ≈ result2.total_cost
            # @test tr1 == tr2
            # @test dr1 == dr2
        end
    end

    @testset verbose = true "Test Instances" begin
        checkTestInstances()
    end

    @testset verbose = true "print_summary" begin
        n = 10
        x = rand(n); y = rand(n);
        truck_cost_factor = 1.0
        drone_cost_factor = 0.5
        result = solve_tspd(x, y, truck_cost_factor, drone_cost_factor)
        print_summary(result)
        report = generate_summary(result)
        @test length(split(report, "\n")) > 10
    end

end


using Test, LinearMaps, LinearAlgebra

@testset "block maps" begin
    @testset "hcat" begin
        for elty in (Float32, Float64, ComplexF64), n2 = (0, 20)
            A11 = rand(elty, 10, 10)
            A12 = rand(elty, 10, n2)
            L = @inferred hcat(LinearMap(A11), LinearMap(A12))
            @test @inferred(LinearMaps.MulStyle(L)) === matrixstyle
            @test L isa LinearMaps.BlockMap{elty}
            A = [A11 A12]
            x = rand(10+n2)
            @test size(L) == size(A)
            @test Matrix(L) ≈ A
            @test L * x ≈ A * x
            L = @inferred hcat(LinearMap(A11), LinearMap(A12), LinearMap(A11))
            A = [A11 A12 A11]
            @test Matrix(L) ≈ A
            A = [I I I A11 A11 A11]
            L = @inferred hcat(I, I, I, LinearMap(A11), LinearMap(A11), LinearMap(A11))
            @test L == [I I I LinearMap(A11) LinearMap(A11) LinearMap(A11)]
            x = rand(elty, 60)
            @test L isa LinearMaps.BlockMap{elty}
            @test L * x ≈ A * x
            A11 = rand(elty, 11, 10)
            A12 = rand(elty, 10, n2)
            @test_throws DimensionMismatch hcat(LinearMap(A11), LinearMap(A12))
        end
    end

    @testset "vcat" begin
        for elty in (Float32, Float64, ComplexF64)
            A11 = rand(elty, 10, 10)
            L = @inferred vcat(LinearMap(A11))
            @test L == [LinearMap(A11);]
            @test Matrix(L) ≈ A11
            A21 = rand(elty, 20, 10)
            L = @inferred vcat(LinearMap(A11), LinearMap(A21))
            @test L isa LinearMaps.BlockMap{elty}
            @test @inferred(LinearMaps.MulStyle(L)) === matrixstyle
            A = [A11; A21]
            x = rand(10)
            @test size(L) == size(A)
            @test Matrix(L) ≈ A
            @test L * x ≈ A * x
            A = [I; I; I; A11; A11; A11]
            L = @inferred vcat(I, I, I, LinearMap(A11), LinearMap(A11), LinearMap(A11))
            @test L == [I; I; I; LinearMap(A11); LinearMap(A11); LinearMap(A11)]
            x = rand(elty, 10)
            @test L isa LinearMaps.BlockMap{elty}
            @test L * x ≈ A * x
            A11 = rand(elty, 10, 11)
            A21 = rand(elty, 20, 10)
            @test_throws DimensionMismatch vcat(LinearMap(A11), LinearMap(A21))
        end
    end

    @testset "hvcat" begin
        for elty in (Float32, Float64, ComplexF64)
            A11 = rand(elty, 10, 10)
            A12 = rand(elty, 10, 20)
            A21 = rand(elty, 20, 10)
            A22 = rand(elty, 20, 20)
            A = [A11 A12; A21 A22]
            @inferred hvcat((2,2), LinearMap(A11), LinearMap(A12), LinearMap(A21), LinearMap(A22))
            L = [LinearMap(A11) LinearMap(A12); LinearMap(A21) LinearMap(A22)]
            @test @inferred(LinearMaps.MulStyle(L)) === matrixstyle
            @test @inferred !issymmetric(L)
            @test @inferred !ishermitian(L)
            x = rand(30)
            @test L isa LinearMaps.BlockMap{elty}
            @test size(L) == size(A)
            @test L * x ≈ A * x
            @test Matrix(L) ≈ A
            A = [I A12; A21 I]
            @inferred hvcat((2,2), I, LinearMap(A12), LinearMap(A21), I)
            L = @inferred hvcat((2,2), I, LinearMap(A12), LinearMap(A21), I)
            @test L isa LinearMaps.BlockMap{elty}
            @test size(L) == (30, 30)
            @test Matrix(L) ≈ A
            @test L * x ≈ A * x
            A = rand(elty, 10,10); LA = LinearMap(A)
            B = rand(elty, 20,30); LB = LinearMap(B)
            @test [LA LA LA; LB] isa LinearMaps.BlockMap{elty}
            @test Matrix([LA LA LA; LB]) ≈ [A A A; B]
            @test [LB; LA LA LA] isa LinearMaps.BlockMap{elty}
            @test Matrix([LB; LA LA LA]) ≈ [B; A A A]
            @test [I; LA LA LA] isa LinearMaps.BlockMap{elty}
            @test Matrix([I; LA LA LA]) ≈ [I; A A A]
            A12 = LinearMap(rand(elty, 10, 21))
            A21 = LinearMap(rand(elty, 20, 10))
            @test_throws DimensionMismatch A = [I A12; A21 I]
            @test_throws DimensionMismatch A = [I A21; A12 I]
            @test_throws DimensionMismatch A = [A12 A12; A21 A21]
            @test_throws DimensionMismatch A = [A12 A21; A12 A21]

            # basic test of "misaligned" blocks
            M = ones(elty, 3, 2) # non-square
            A = LinearMap(M)
            B = [I A; A I]
            C = [I M; M I]
            @test B isa LinearMaps.BlockMap{elty}
            @test Matrix(B) == C
            @test Matrix(transpose(B)) == transpose(C)
            @test Matrix(adjoint(B)) == C'
        end
    end

    @testset "adjoint/transpose" begin
        for elty in (Float32, Float64, ComplexF64), transform in (transpose, adjoint)
            A12 = rand(elty, 10, 10)
            A = [I A12; transform(A12) I]
            L = [I LinearMap(A12); transform(LinearMap(A12)) I]
            @test @inferred(LinearMaps.MulStyle(L)) === matrixstyle
            if elty <: Complex
                if transform == transpose
                    @test @inferred issymmetric(L)
                else
                    @test @inferred ishermitian(L)
                end
            end
            if elty <: Real
                @test @inferred ishermitian(L)
                @test @inferred issymmetric(L)
            end
            x = rand(elty, 20)
            @test L isa LinearMaps.LinearMap{elty}
            @test size(L) == size(A)
            @test L * x ≈ A * x
            @test Matrix(L) ≈ A
            Lt = @inferred transform(L)
            @test Lt isa LinearMaps.LinearMap{elty}
            @test Lt * x ≈ transform(A) * x
            Lt = @inferred transform(LinearMap(L))
            @test Lt * x ≈ transform(A) * x
            @test Matrix(Lt) ≈ Matrix(transform(A))
            A21 = rand(elty, 10, 10)
            A = [I A12; A21 I]
            L = [I LinearMap(A12); LinearMap(A21) I]
            Lt = @inferred transform(L)
            @test Lt isa LinearMaps.LinearMap{elty}
            @test Lt * x ≈ transform(A) * x
            @test Matrix(Lt) ≈ Matrix(transform(A))
        end
    end
end


function calc_rho( Focc, v::Array{Float64,2} )
    Ncols   = size(v)[2]
    Npoints = size(v)[1]
    rho = zeros(Float64,Npoints)
    for ic = 1:Ncols
        rho[:] = rho[:] + v[:,ic] .* v[:,ic] * Focc[ic]
    end
    return rho
end


using Test
using InvariantSets
using LinearAlgebra: I
using Polyhedra # needed for prerset_vrep, preset_hrep
using CDDLib # needed for preset_hrep

@testset "Preset computation" begin
    n = 2
    m = 1
    A = rand(n,n)
    B = rand(n,m)
    X = BallInf(zeros(n), 10.)
    U = BallInf(zeros(m), .1)
    K = [0.0890246  0.484408]
    F, g = tosimplehrep(X)

    P = preset(A, X)
    @test P ⊆ HPolytope(F*A,g) && HPolytope(F*A,g) ⊆ P
    # @test isequivalent(preset(A, X), HPolytope(F*A,g))

    D = state_constraints(X, U, K)
    D2 = InvariantSets.state_constraints_ops(X, U, K)
    @test D ⊆ D2 && D2 ⊆ D

    P_2 = InvariantSets._preset_vrep(A, B, U, D)
    P_3 = InvariantSets._preset_hrep(A, B, U, D)
    @test P_2 ⊆ P_3 && P_3 ⊆ P_2
end

@testset "Maximum invariant set" begin
    A2 = [1.0 1.0; 0.0 1.0]
    B2 = [1.0; 0.5]
    Q2 = I
    R2 = 90*I
    # K2 = dlqr(A2,B2,Q2,R2)
    K2 = [0.0890246  0.484408]
    X2 = BallInf(zeros(2), 2.)
    U2 = BallInf(zeros(1), 2.)
    S2 = A2-B2*K2

    max_PI_1 =  maximum_invariant_set(S2, X2)
    max_PI_2 =  maximum_invariant_set(S2, X2, 5)
    max_PI_3 =  maximum_invariant_set(S2, X2, 4)

    @test max_PI_1 ⊆ max_PI_2 && max_PI_2 ⊆ max_PI_1
    @test !(max_PI_1 ⊆ max_PI_3 && max_PI_3 ⊆ max_PI_1)

    max_PI_4 = maximum_invariant_set(A2, B2, X2, U2, K2)
    @test max_PI_1 ⊆ max_PI_4 && max_PI_4 ⊆ max_PI_1
end

@testset "Maximum control invariant set" begin
    A2 = [1.0 1.0; 0.0 1.0]
    B2 = [1.0; 0.5]
    X2 = BallInf(zeros(2), 5.)
    U2 = BallInf(zeros(1), 1.)
    max_CPI_1 =  maximum_control_invariant_set(A2, B2, X2, U2; preset=InvariantSets._preset_vrep)
    max_CPI_2 =  maximum_control_invariant_set(A2, B2, X2, U2; preset=InvariantSets._preset_hrep)
    max_CPI_3 =  maximum_control_invariant_set(A2, B2, X2, U2, 5)
    @test max_CPI_1 ⊆ max_CPI_2 && max_CPI_2 ⊆ max_CPI_1
    @test !(max_CPI_1 ⊆ max_CPI_3 && max_CPI_3 ⊆ max_CPI_1)
end


include("euler/euler.jl")
using .Calculus: fibonacci_index, fibonacci_numbers
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(n::Int)::Int
    index = fibonacci_index(n)
    fibonacci = fibonacci_numbers(index + 1, Int)
    last_sum, new_sum = 0, 0
    result = sum(fibonacci[1:4]) - 1
    for i ∈ 5:index
        last_sum, new_sum = new_sum, new_sum + last_sum + fibonacci[i - 3]
        result += new_sum + fibonacci[i]
    end
    return result
end

compute(13)

compute(23_416_728_348_467_685)

@benchmark compute(23_416_728_348_467_685)


using LinearAlgebra
using Distributions
using Random

struct SimWorld{A2<:Array{Float64,2}, F<:Float64, I<:Int64, A1<:Array{Float64,1}, AA<:Array{Array{Float64,2},1}}
    A::A2
    C::A2
    muInit::A1
    emitVar::A1
    endT::F
    dt::F
    endI::I
    allT::A1
    allAs::AA
end

function getStates(SW::SimWorld)
    Z = (SW.muInit .+ sqrt.(SW.emitVar).*randn(size(SW.muInit)))
    Zs = [A*Z for A in SW.allAs]
    Ys = [SW.C*z .+ rand!(MvNormal(zeros(size(SW.C*SW.C',1)), SW.C*diagm(SW.emitVar)*SW.C'), similar(SW.C*SW.muInit)) for z in Zs];
    return Zs::Array{Array{Float64,1},1}, Ys::Array{Array{Float64,1},1}
end

SimWorld(A::Array{Float64,2}, C::Array{Float64,2}, muPrior::Array{Float64,1}, emitVar::Array{Float64,1}, endT::Float64=1000., dt::Float64=.1) =
    SimWorld(A, C, muPrior, emitVar, endT, dt, Integer(ceil(endT/dt)))
SimWorld(A::Array{Float64,2}, C::Array{Float64,2}, muPrior::Array{Float64,1}, emitVar::Array{Float64,1}, endT::Float64, dt::Float64, endI::Int) =
    SimWorld(A, C, muPrior, emitVar, endT, dt, endI,
                Array{Float64,1}(0:dt:endI*dt),
                [A^t for t in 0:endI])


#=
 Example 1: Track the pathways of a user-defined water mass.
 Steps: (a) define the water mass 1). by a rectangular surface patch
            dyed with passive tracer concentration of 1,
            or 2. load a pre-defined surface patch in d_all.mat.
        (b) propagate the dye with the matrix A, with the result
            being the fraction of water originating from the
            surface region.
 See Section 2b of Gebbie & Huybers 2010, esp. eqs. (15)-(17).
=#
using Revise
using TMI, PyPlot, PyCall

TMIversion = "modern_90x45x33_GH10_GH12"
A, Alu, γ, TMIfile, L, B = config_from_nc(TMIversion)

#- define the surface patch by the bounding latitude and longitude.
latbox = [50,60]; # 50 N -> 60 N, for example.

# mutable due to wraparound: don't use an immutable tuple
lonbox = [-50,0]; # 50 W -> prime meridian

# do numerical analysis
c = trackpathways(Alu,latbox,lonbox,γ);

# do plotting (could be a function)
plotextent(latbox, lonbox)

# plot a section at 330 east longitude (i.e., 30 west)
lon_section = 330; # only works if exact
lims = 0:5:100
tlabel = "water-mass fraction [%]"
sectionplot(100c,lon_section,lims,titlelabel = tlabel)


@testset "frohlich_α" begin

    # NaCl ∼ 5 (Feynman 1955)
    @test frohlich_α(2.3, 5.6, 4.9e13 / (2 * π), 1.0) ≈ 5.0 atol = 0.3

    # CdTe ∼ 0.39 (Stone), 0.29 (Devreese)
    @test frohlich_α(7.1, 10.4, 5.08e12, 0.095) ≈ 0.3 atol = 0.1

    # GaAs ∼ 0.068 (Devreese)
    @test frohlich_α(10.89, 12.9, 8.46e12, 0.063) ≈ 0.068 atol = 0.01

end


using Flux, Statistics
using Flux: onecold, crossentropy

include("myplots.jl")
include("model-commons.jl")

epochs = 300
eta = 0.001

# Train the ANN model
function train_ann(d::DataSet; hidden=750, epochs=epochs, eta=eta, plotloss=true)
  @info("Constructing ANN model...")
  m = Chain(
    Dense(size(d.X_train[:,1],1), hidden, relu),
    Dense(hidden, size(d.Y_train[:,1],1)),
    softmax)

  loss(x, y) = crossentropy(m(x), y)
  accuracy(x, y) = mean(onecold(m(x)) .== onecold(y))

  dataset = [(d.X_train, d.Y_train)]
  evalcb = () -> @show(loss(d.X_train, d.Y_train))
  opt = ADAM(eta)

  # Check for NaN values so we know if we have bad data before training
  # @info("Checking for NaN values")
  # for i in 1:size(X_train,2)
  #   if (isnan(loss(X_train[:,i], Y_train[:,i])))
  #     @info("$i is NaN")
  #   end
  # end
  # @info("Done.")

  @info("Training model...")

  # Train for a specified number of epochs
  if plotloss
    plt = create_plot(0, loss(d.X_train, d.Y_train).data)
  end
  for step in 1:epochs
    @info("Epoch $step")
    Flux.train!(loss, params(m), dataset, opt)
    error = loss(d.X_train, d.Y_train).data
    @show(error)
    @show(accuracy(d.X_test, d.Y_test))
    if plotloss
      update_plot(plt, step, error)
    end
  end

  @info("Done")

  return m
end

# 250 epochs, image size 150x100, eta 0.001, model = 15000 relu | 750 sigmoid | 10 softmax => 77.6% accuracy


struct Monomials{dim,order, T}
    indices::Union{Vector{NTuple{dim,Int}},Vector{Int}}
    centroid::Tensors.Vec{dim,T}
    diameter::T
end
function Monomials(dim,degree, centroid, diameter)
    if dim == 1
        return Monomials1D(degree, centroid, diameter)
    elseif dim == 2
        return Monomials2D(degree, centroid, diameter)
    else
        throw("Monomials of dim $dim are not implemented!")
    end
end
function Monomials1D(degree::Int, centroid::Tensors.Vec{1,T}, diameter::T) where {T}
    indices = 0:degree
    Monomials{1,degree, T}(collect(indices), centroid, diameter)
end
function Monomials2D(degree::Int, centroid::Tensors.Vec{2,T}, diameter::T) where {T}
    indices = get_monomial2DIndices(degree)
    Monomials{2,degree, T}(indices, centroid, diameter)
end

function get_monomial2DIndices(degree)
    if degree == 0
        return [(0,0)]
    else
        tmp = Vector{NTuple{2,Int}}()
        a = (degree,0)
        for i in 1:(binomial(degree+1,1))
            push!(tmp,a)
            a = (a[1]-1,a[2]+1)
        end
        return [get_monomial2DIndices(degree-1)...,tmp...]
    end
end

getnbasefunctions(::Monomials{2,order}) where {order} = Int((order+1)*(order+2)/2)
getnbasefunctions(::Monomials{1,order}) where {order} = order+1

# Evaluate monomials
function value(ip::Monomials{dim,order}, k::Int, ξ::Tensors.Vec{dim,T}) where {order, T, dim}
    if k > getnbasefunctions(ip);throw(ArgumentError("no shape function $k for interpolation $ip"));end
    prod((1/ip.diameter*(ξ-ip.centroid)).^ip.indices[k])
end

function gradient_value(ip::Monomials{dim,order}, k::Int, ξ::Tensors.Vec{dim,T}) where {dim,order,T}
    if k >getnbasefunctions(ip);throw(ArgumentError("no shape function $k for interpolation $ip"));end
    Tensors.gradient(ξ -> value(ip, k, ξ), ξ)
end

function laplace_value(ip::Monomials{dim,order}, k::Int, ξ::Tensors.Vec{dim,T}) where {dim,order,T}
    if k >getnbasefunctions(ip);throw(ArgumentError("no shape function $k for interpolation $ip"));end
    Tensors.laplace(ξ -> value(ip, k, ξ), ξ)
end

struct RotMonomials{dim,order, T}
    indices::Union{Vector{NTuple{dim,Int}},Vector{Int}}
    centroid::Tensors.Vec{dim,T}
    diameter::T
end
function RotMonomials(dim,order, centroid, diameter)
    if dim == 2
        return RotMonomials2D(order, centroid, diameter)
    else
        throw("Monomials of dim $dim are not implemented!")
    end
end
function RotMonomials2D(order::Int, centroid::Tensors.Vec{2,T}, diameter::T) where {T}
    indices = get_monomial2DIndices(order-1)
    RotMonomials{2,order, T}(indices, centroid, diameter)
end

getnbasefunctions(::RotMonomials{2,order}) where {order} = Int((order)*(order+1)/2)

# Evaluate monomials

rot(ξ::Tensors.Vec{2}) = Tensors.Vec((-ξ[2],ξ[1]))

function value(ip::RotMonomials{dim,order}, k::Int, ξ::Tensors.Vec{dim,T}) where {order, T, dim}
    if k > getnbasefunctions(ip);throw(ArgumentError("no shape function $k for interpolation $ip"));end
    rot(ξ)*prod((1/ip.diameter*(ξ-ip.centroid)).^ip.indices[k])
end

function gradient_value(ip::RotMonomials{dim,order}, k::Int, ξ::Tensors.Vec{dim,T}) where {dim,order,T}
    if k >getnbasefunctions(ip);throw(ArgumentError("no shape function $k for interpolation $ip"));end
    Tensors.gradient(ξ -> value(ip, k, ξ), ξ)
end

function laplace_value(ip::RotMonomials{dim,order}, k::Int, ξ::Tensors.Vec{dim,T}) where {dim,order,T}
    if k >getnbasefunctions(ip);throw(ArgumentError("no shape function $k for interpolation $ip"));end
    Tensors.laplace(ξ -> value(ip, k, ξ), ξ)
end


using Yao, YaoExtensions
using Test, YaoArrayRegister, LinearAlgebra

@testset "test reflect gate" begin
    reg0 = rand_state(3)
    mirror = randn(1<<3)*im; mirror[:] /= norm(mirror)
    rf = YaoExtensions.reflect(mirror)
    reg = copy(reg0)
    apply!(reg, rf)

    v0, v1 = vec(reg.state), vec(reg0.state)
    @test mat(rf) * statevec(reg0) ≈ statevec(apply!(copy(reg0), rf))
    @test statevec(rf.psi)' * v0 ≈ statevec(rf.psi)' * v1
    @test v0 - statevec(rf.psi)'*v0* statevec(rf.psi) ≈
        -(v1- statevec(rf.psi)' * v1 * statevec(rf.psi))

    @test mat(ComplexF32, rf) == ComplexF32.(mat(rf))

    @test rf == YaoExtensions.reflect(rf.psi)
    # copy do not occur on register
    @test copy(rf).psi === rf.psi

    @test ishermitian(rf)
    @test isreflexive(rf)
    @test isunitary(rf)
end


using Random
using SumProductNetworks

println("make dataset A and B")

N = 1000
D = 4

A = hcat([randn(N) .+ (2 + 0.5*i) for i in 1:D]...)
B = hcat([randn(N) .+ (4 + 0.5*i) for i in 1:D]...)

println("generate an SPN by LearnSPN with dataset A")
spn = generate_spn(A, :learnspn)
updatescope!(spn)

println("mean log-likelihood for...")
println("   dataset A: ", mean(logpdf(spn, A)))
println("   dataset B: ", mean(logpdf(spn, B)))


##################################################################################
# This file is part of StateSpaceEcon.jl
# BSD 3-Clause License
# Copyright (c) 2020, Bank of Canada
# All rights reserved.
##################################################################################

"""
    SteadyStateSolver

A module that is part of StateSpaceEcon package.
Contains methods for finding a steady state of a model.
"""
module SteadyStateSolver

using SparseArrays
using LinearAlgebra
using ForwardDiff
using DiffResults
using Printf

using ModelBaseEcon

include("steadystate/1dsolvers.jl")
include("steadystate/presolve.jl")
include("steadystate/initial.jl")
include("steadystate/global.jl")
include("steadystate/solverdata.jl")
include("steadystate/nr.jl")
include("steadystate/lm.jl")
include("steadystate/diagnose.jl")
include("steadystate/sssolve.jl")

end # module

using .SteadyStateSolver
export clear_sstate!
export initial_sstate!
export check_sstate
export diagnose_sstate
export sssolve!


using BenchmarkTools
using Printf
using Transducers

benchmark = false

exampleFilename = "src/advent/day13/example.txt"
example1Filename = "src/advent/day13/example1.txt"
example2Filename = "src/advent/day13/example2.txt"
example3Filename = "src/advent/day13/example3.txt"
example4Filename = "src/advent/day13/example4.txt"
example5Filename = "src/advent/day13/example5.txt"
inputFilename = "src/advent/day13/input.txt"

struct Notes
    t::Int64
    buses::Array{Union{Missing, Int64}}
end

function readData(filename)
    f = open(filename)
    t = parse(Int64, readline(f))
    buses = split(readline(f), ",") |>
            Map(s -> s == "x" ? Missing() : parse(Int64, s)) |>
            collect
    close(f)
    Notes(t, buses)
end

exampleNotes = readData(exampleFilename)
example1Notes = readData(example1Filename)
example2Notes = readData(example2Filename)
example3Notes = readData(example3Filename)
example4Notes = readData(example4Filename)
example5Notes = readData(example5Filename)
inputNotes = readData(inputFilename)

#### part1

function part1(notes)
    result = notes.buses |>
            Filter(bus -> !ismissing(bus)) |>
            Map(bus -> (bus = bus, t = notes.t + (bus - (notes.t % bus)) % bus)) |>
            foldxl((a, b) -> a.t < b.t ? a : b)
    result.bus * (result.t - notes.t)
end

@assert part1(exampleNotes) == 295

@show part1(inputNotes)
benchmark && @btime part1(inputNotes);

#### part2

function part2(notes)
    buses = NamedTuple{(:number, :index), Tuple{Int64, Int64}}[]
    for i = 1:length(notes.buses)
        if !ismissing(notes.buses[i])
            push!(buses, (number = notes.buses[i], index = i))
        end
    end
    sort!(buses; rev = true, by = bus -> bus.number)
    referenceBus = buses[1]
    otherBuses = @view buses[2:length(buses)]

    t = referenceBus.number
    while true
        matched = true
        for bus in otherBuses
            t2 = t + bus.index - referenceBus.index
            if t2 % bus.number != 0
                matched = false
                break
            end
        end
        matched && break
        t += referenceBus.number
    end
    t - (referenceBus.index - 1)
end

@assert part2(exampleNotes) == 1068781
@assert part2(example1Notes) == 3417
@assert part2(example2Notes) == 754018
@assert part2(example3Notes) == 779210
@assert part2(example4Notes) == 1261476
@assert part2(example5Notes) == 1202161486

#@show part2(inputNotes)
#benchmark && @btime part2(inputNotes);


export sources

"""
`sources(f::DiscreteFunction)` returns a list of all inputs to `f`
that are not the output of `f`. That is, this is a list of the vertices
with out-degree 0 in the graph of `f`.
"""
function sources(f::DiscreteFunction)
    n = length(f)
    S = setdiff( Set(1:n), image(f))
    return sort(collect(S))
end


# This file was generated, do not modify it.

using MLJ, PrettyPrinting
X, y = @load_iris
@load DecisionTreeClassifier

dtc = DecisionTreeClassifier()
r   = range(dtc, :max_depth, lower=1, upper=5)

tm = TunedModel(model=dtc, ranges=[r, ], measure=cross_entropy)

m = machine(tm, X, y)
fit!(m)

fitted_params(m).best_model.max_depth

tm = TunedModel(model=dtc, ranges=r, operation=predict_mode,
                measure=misclassification_rate)
m = machine(tm, X, y)
fit!(m)
fitted_params(m).best_model.max_depth

r = report(m)
r.best_measurement

using PyPlot
figure(figsize=(8,6))
plot(r.parameter_values, r.measurements)

xticks(1:5, fontsize=12)
yticks(fontsize=12)
xlabel("Maximum depth", fontsize=14)
ylabel("Misclassification rate", fontsize=14)
ylim([0, 1])

savefig("assets/literate/A-model-tuning-hpt.svg") # hide

X = (x1=rand(100), x2=rand(100), x3=rand(100))
y = 2X.x1 - X.x2 + 0.05 * randn(100);

dtr = @load DecisionTreeRegressor
forest = EnsembleModel(atom=dtr)

params(forest) |> pprint

r1 = range(forest, :(atom.n_subfeatures), lower=1, upper=3)
r2 = range(forest, :bagging_fraction, lower=0.4, upper=1.0)
tm = TunedModel(model=forest, tuning=Grid(resolution=12),
                resampling=CV(nfolds=6), ranges=[r1, r2],
                measure=rms)
m = machine(tm, X, y)
fit!(m);

r = report(m)
r.best_measurement

figure(figsize=(8,6))

vals_sf = r.parameter_values[:, 1]
vals_bf = r.parameter_values[:, 2]

tricontourf(vals_sf, vals_bf, r.measurements)
xlabel("Number of sub-features", fontsize=14)
ylabel("Bagging fraction", fontsize=14)
xticks([1, 2, 3], fontsize=12)
yticks(fontsize=12)

savefig("assets/literate/A-model-tuning-hm.svg") # hide


import FireflyModule.Firefly

function arc_distance(f1::Firefly, f2::Firefly)::Int
    """Calculates the arc distance between two firefly paths.
    Example:
        Firefly i -> 1 2 4 6 3 5 7
        Firefly j -> 1 2 4 3 5 7 6
        1-2, 2-4, 3-5, 5-7 arcs in firefly i are also presented
        in firefly j. But 4-6, 6-3 arcs are missing in
        firefly j. so the arc distance between firefly i
        and firefly j is 2 here.
    """
    missmatched_pair_counts = 0
    f1_pairs = zip(f1.path[1:end-1], f1.path[2:end])
    f2_pairs = zip(f2.path[1:end-1], f2.path[2:end])

    for f1_pair in f1_pairs
        if f1_pair ∉ f2_pairs missmatched_pair_counts += 1 end
    end

    return missmatched_pair_counts
end


module DA_NlayerkernelTests

using Test
using Parameters, SpecialFunctions

abstract type DiffusionParameters end

include(joinpath(dirname(@__FILE__), "..", "..","src/forwardmodels/Diffusion Approximation/DAcylinder_layered.jl"))

## Test the derivations with original

@inline function α_coeff!(α, μa, D, sn)
    @inbounds for ind in 1:length(μa)
        α[ind] = sqrt(μa[ind] / D[ind] + sn^2)
    end
    return α
end

# greens function for 2 layer cylinder as presented by Liemert
function green_2layer(α, sn, μa, D, z, z0, zb, l, n)
    α = α_coeff!(α, μa, D, sn)

    β3 = sinh(α[2]*(l[2] + zb[2]))
    γ3 = cosh(α[2]*(l[2] + zb[2]))

    g1 = sinh(α[1]*(z0 + zb[1]))*sinh(α[1]*(z + zb[1])) / (D[1]*α[1]*exp(α[1]*(l[1] + zb[1])))
    g1 *= (D[1]*α[1]*n[1]^2*β3 - D[2]*α[2]*n[2]^2*γ3) / (D[1]*α[1]*n[1]^2*β3*cosh(α[1]*(l[1] + zb[1])) + (D[2]*α[2]*n[2]^2*γ3*sinh(α[1]*(l[1] + zb[1]))))

    g = (exp(-α[1]*abs(z - z0)) - exp(-α[1]*(z + z0 + 2zb[1]))) / (2*D[1]*α[1])

    return g + g1
end
# greens function for 3 layer cylinder as presented by Liemert
function green_3layer(α, sn, μa, D, z, z0, zb, l, n)
    α = α_coeff!(α, μa, D, sn)

    β3 = D[2]*α[2]*n[2]^2*cosh(α[2]*l[2])*sinh(α[3]*(l[3] + zb[2])) + D[3]*α[3]*n[3]^2*sinh(α[2]*l[2])*cosh(α[3]*(l[3] + zb[2]))
    γ3 = D[2]*α[2]*n[2]^2*sinh(α[2]*l[2])*sinh(α[3]*(l[3] + zb[2])) + D[3]*α[3]*n[3]^2*cosh(α[2]*l[2])*cosh(α[3]*(l[3] + zb[2]))

    g1 = sinh(α[1]*(z0 + zb[1]))*sinh(α[1]*(z + zb[1])) / (D[1]*α[1]*exp(α[1]*(l[1] + zb[1])))
    g1 *= (D[1]*α[1]*n[1]^2*β3 - D[2]*α[2]*n[2]^2*γ3) / (D[1]*α[1]*n[1]^2*β3*cosh(α[1]*(l[1] + zb[1])) + (D[2]*α[2]*n[2]^2*γ3*sinh(α[1]*(l[1] + zb[1]))))

    g = (exp(-α[1]*abs(z - z0)) - exp(-α[1]*(z + z0 + 2zb[1]))) / (2*D[1]*α[1])

    return g + g1
end

### test the 2 layer case
α = [1.2, 1.5]
sn = 2.2
D = [0.2, 0.5]
μa = [0.1, 0.14]
z0 = 0.2
zb = [0.2, 0.3]
l = [1.0, 2.0]
n = [1.0, 1.0]
n2 = @. D * n^2
z = 0.0

@test _green_Nlaycylin_top(sn, μa, D, z, z0, zb, l, n2, 2) ≈ green_2layer(α, sn, μa, D, z, z0, zb, l, n)

### test the 3 layer case
α = [1.2, 1.5, 1.3]
sn = 5.8
D = [0.2, 0.5, 0.7]
μa = [0.1, 0.14, 0.2]
z0 = 0.2
zb = [0.2, 0.3, 0.3]
l = [1.0, 2.0, 1.2]
n = [1.0, 1.0, 1.1]
n2 = @. D * n^2
z = 0.0

@test _green_Nlaycylin_top(sn, μa, D, z, z0, zb, l, n2, 3) ≈ green_3layer(α, sn, μa, D, z, z0, zb, l, n)

end # module


@generated function gen_harmonic_mean(x...)
    sum = :(0)
    for i = 1:length(x)
        sum = :(1/x[$i] + $sum)
    end
    mean = :(length(x) / $sum)
    return mean
end


"""
    f, p = logsmooth(X, inBin, nbin=8, n=3)

Smooth the fft, `X`, and convert it to dB. Use `nbin` bins from 0 to 3×`inBin`,
thereafter increase bin sizes by a factor of 1.1, staying less than 2¹⁰.
For the `n` sets of bins `inBin`+[0:2], 2×`inBin`+[0:2], ... , `n`×`inBin`+[0:2],
don't do averaging. This way, the noise BW and the scaling of the tone and its
harmonics are unchanged. Unfortunately, harmonics above the nth appear smaller
than they really are because their energy is averaged over many bins.
"""
function logsmooth(X, inBin, nbin=8, n=3)
    N = length(X)
    N2 = N÷2
    f1 = rem(inBin-1, nbin) + 1

    startbin = [f1:nbin:inBin-1; inBin:inBin+2]
    for i = 1:n
        startbin = [startbin; (startbin[end]+1):nbin:(i*inBin-1); i*inBin .+ (0:2)]
    end
    m = startbin[length(startbin)] + nbin
    while m < N2
        startbin = [startbin; m]
        nbin = min(nbin*1.1, 2^10)
        m = round(Int, m .+ nbin)
    end
    stopbin = [startbin[2:length(startbin)] .- 1; N2]

    f = ((startbin+stopbin)/2 .- 1)/N
    p = zeros(size(f))
    for i = 1:length(f)
        p[i] = dbp(norm(X[startbin[i]:stopbin[i]])^2 / (stopbin[i] - startbin[i] + 1))
    end

    return f, p
end


@doc raw"
    random_point(M::Manifold)

generate a random point on a manifold. By default it uses `random_point(M,:Gaussian)`.
"
random_point(M::Manifold) = random_point(M, Val(:Gaussian))
@doc raw"
    random_point(M::Manifold, s::Symbol, options...)

generate a random point using a noise model given by `s` with its additional `options`
just passed on.
"
random_point(M::Manifold, s::Symbol, options...) = random_point(M, Val(s), options...)

@doc raw"""
    random_point(M::AbstractPowerManifold, options...)

generate a random point on the `AbstractPowerManfold` `M` given `options` that are
passed on.
"""
function random_point(
    M::AbstractPowerManifold{𝔽,Mt,NestedPowerRepresentation}, options...
) where {𝔽,Mt}
    return [random_point(M.manifold, options...) for i in get_iterator(M)]
end
function random_point(
    M::AbstractPowerManifold{𝔽,Mt,ArrayPowerRepresentation}, options...
) where {𝔽,Mt}
    return cat(
        [random_point(M.manifold, options...) for i in get_iterator(M)]...;
        dims=length(representation_size(M.manifold)) + 1,
    )
end

@doc raw"""
    random_point(M::AbstractGroupManifold, options...)

On an abstract group manifold, the random point is taken from the internally stored `M.manifold`.
"""
random_point(M::AbstractGroupManifold, kwargs...) = random_point(M.manifold, kwargs...)
function random_point(M::AbstractGroupManifold, s::Symbol, options...)
    return random_point(M, Val(s), options...)
end

@doc raw"""
    random_point(M::Circle, :Uniform)

return a random point on the [`Circle`](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold) $\mathbb S^1$ by
picking a random element from $[-\pi,\pi)$ uniformly.
"""
random_point(::Circle, ::Val{:Uniform}) = sym_rem(rand() * 2 * π)
random_point(M::Circle) = random_point(M, Val(:Uniform)) # introduce different default

@doc raw"""
    random_point(M::Euclidean[,:Gaussian, σ::Float64=1.0])

generate a random point on the `Euclidean` manifold `M`, where the
optional parameter determines the type of the entries of the
resulting point on the Euclidean space d.
"""
random_point(M::Euclidean) = randn(representation_size(M))
random_point(M::Euclidean, ::Val{:Gaussian}, σ=1.0) = σ * randn(manifold_dimension(M))

@doc raw"""
    random_point(M::FixedRankMatrices, options...)

return a random point on the FixedRankMatrices manifold.
The orthogonal matrices are sampled from the Stiefel manifold
and the singular values are sampled uniformly at random.
"""
function random_point(::FixedRankMatrices{m,n,k}, options...) where {m,n,k}
    U = random_point(Stiefel(m, k), options...)
    S = sort(rand(k); rev=true)
    V = random_point(Stiefel(n, k), options...)
    return SVDMPoint(U, S, V')
end

@doc raw"""
    random_point(M::Grassmannian, :Gaussian [, σ=1.0])

return a random point `x` on `Grassmannian` manifold `M` by
generating a random (Gaussian) matrix with standard deviation `σ` in matching
size, which is orthonormal.
"""
function random_point(::Grassmann{n,k,𝔽}, ::Val{:Gaussian}, σ::Float64=1.0) where {n,k,𝔽}
    V = σ * randn(𝔽 === ℝ ? Float64 : ComplexF64, (n, k))
    A = qr(V).Q[:, 1:k]
    return A
end

@doc raw"""
    random_point(M::ProductManifold, options...)

return a random point `x` on `Grassmannian` manifold `M` by
generating a random (Gaussian) matrix with standard deviation `σ` in matching
size, which is orthonormal.
"""
function random_point(M::ProductManifold, o...)
    return ProductRepr([random_point(N, o...) for N in M.manifolds]...)
end

@doc raw"""
    random_point(M::Rotations, :Gaussian [, σ=1.0])

return a random point `p` on the manifold `Rotations`
by generating a (Gaussian) random orthogonal matrix with determinant $+1$. Let $QR = A$
be the QR decomposition of a random matrix $A$, then the formula reads $p = QD$
where $D$ is a diagonal matrix with the signs of the diagonal entries of $R$, i.e.
````math
D_{ij}=\begin{cases}
\operatorname{sgn}(R_{ij}) & \text{if} \; i=j \\
0 & \, \text{otherwise.}
\end{cases}
````
It can happen that the matrix gets -1 as a determinant. In this case, the first
and second columns are swapped.
"""
function random_point(M::Rotations, ::Val{:Gaussian}, σ::Real=1.0)
    d = Int(0.5 + sqrt(0.25 + 2 * manifold_dimension(M)))
    if d == 1
        return ones(1, 1)
    else
        A = randn(Float64, d, d)
        s = diag(sign.(qr(A).R))
        D = Diagonal(s)
        C = qr(A).Q * D
        if det(C) < 0
            C[:, [1, 2]] = C[:, [2, 1]]
        end
        return C
    end
end

@doc raw"""
    random_point(M::SymmetricPositiveDefinite, :Gaussian[, σ=1.0])

gerenate a random symmetric positive definite matrix on the
`SymmetricPositiveDefinite` manifold `M`.
"""
function random_point(
    ::SymmetricPositiveDefinite{N}, ::Val{:Gaussian}, σ::Float64=1.0
) where {N}
    D = Diagonal(1 .+ rand(N)) # random diagonal matrix
    s = qr(σ * randn(N, N)) # random q
    return Matrix(Symmetric(s.Q * D * transpose(s.Q)))
end

@doc raw"""
    random_point(M::Stiefel, :Gaussian[, σ=1.0])

return a random (Gaussian) point `x` on the `Stiefel` manifold `M` by generating a (Gaussian)
matrix with standard deviation `σ` and return the orthogonalized version, i.e. return the Q
component of the QR decomposition of the random matrix of size $n×k$.
"""
function random_point(::Stiefel{n,k,𝔽}, ::Val{:Gaussian}, σ::Float64=1.0) where {n,k,𝔽}
    A = σ * randn(𝔽 === ℝ ? Float64 : ComplexF64, n, k)
    return Matrix(qr(A).Q)
end

@doc raw"""
    random_point(M::Sphere, :Gaussian[, σ=1.0])
return a random point on the Sphere by projecting a normal distirbuted vector
from within the embedding to the sphere.
"""
function random_point(M::Sphere, ::Val{:Gaussian}, σ::Float64=1.0)
    return project(M, σ * randn(manifold_dimension(M) + 1))
end

@doc raw"""
    random_point(M::TangentBundle, options...)

generate a random point on the tangent bundle by calling a [`random_point`](@ref) and a
[`random_tangent`](@ref) with the given `options...`
"""
function random_point(M::TangentBundle, options...)
    p = random_point(M.manifold, options...)
    X = random_tangent(M.manifold, p, options...)
    return ProductRepr(p, X)
end

@doc raw"""
    random_tangent(M::AbstractGroupManifold, p, options...)

On an abstract group manifold, the random tangent is taken from the internally stored `M.manifold`s tangent space at `p`.
"""
function random_tangent(M::AbstractGroupManifold, p, kwargs...)
    return random_tangent(M.manifold, p, kwargs...)
end
function random_tangent(M::AbstractGroupManifold, p, s::Symbol, options...)
    return random_tangent(M, p, Val(s), options...)
end

@doc raw"""
    random_tangent(M, p, options...)

generate a random tangent vector in the tangent space of `p` on `M`. By default
this is a `:Gaussian` distribution.
"""
random_tangent(M::Manifold, p, options...) = random_tangent(M, p, :Gaussian, options...)
function random_tangent(M::Manifold, p, s::Symbol, options...)
    return random_tangent(M, p, Val(s), options...)
end

@doc raw"""
    random_tangent(M::Circle, p [, :Gaussian, σ=1.0])

return a random tangent vector from the tangent space of the point `p` on the
[Circle](https://juliamanifolds.github.io/Manifolds.jl/stable/interface.html#ManifoldsBase.Manifold) $\mathbb S^1$ by using a normal distribution with
mean 0 and standard deviation 1.
"""
random_tangent(::Circle, p, ::Val{:Gaussian}, σ::Real=1.0) = σ * randn()

function random_tangent(M::Euclidean, p, ::Val{:Gaussian}, σ::Float64=1.0)
    return σ * randn(manifold_dimension(M))
end

@doc raw"""
    random_tangent(M::Grassmann, p[,type=:Gaussian, σ=1.0])

return a (Gaussian) random vector from the tangent space $T_p\mathrm{Gr}(n,k)$ with mean
zero and standard deviation `σ` by projecting a random Matrix onto the  `p`.
"""
function random_tangent(M::Grassmann, p, ::Val{:Gaussian}, σ::Float64=1.0)
    Z = σ * randn(eltype(p), size(p))
    X = project(M, p, Z)
    X = X ./ norm(X)
    return X
end

@doc raw"""
    random_tangent(M::FixedRankMatrices, p, options...)

generate a random tangent vector in the tangent space of the point `p` on the
`FixedRankMatrices` manifold `M`.
"""
function random_tangent(::FixedRankMatrices{m,n,k}, p, options...) where {m,n,k}
    Up = randn(m, k, options...)
    Vp = randn(n, k, options...)
    A = randn(k, k, options...)
    return UMVTVector(Up - p.U * p.U' * Up, A, Vp' - Vp' * p.Vt' * p.Vt)
end

@doc raw"""
    random_tangent(M::Hyperbolic, p, :Gaussian [, σ=1.0])

generate a random point on the Hyperbolic manifold by projecting a point from the embedding
with respect to the Minkowsky metric.
"""
function random_tangent(M::Hyperbolic, p, ::Val{:Gaussian}, σ=1.0)
    Y = σ * randn(eltype(p), size(p))
    X = project(M, p, Y)
    return X
end

function random_tangent(M::PowerManifold, p, options...)
    rep_size = representation_size(M.manifold)
    X = zero_tangent_vector(M, p)
    for i in get_iterator(M)
        X[M, i] = random_tangent(M.manifold, p[M, i], options...)
    end
    return X
end

@doc raw"""
    random_tangent(M::ProductManifold, p)

generate a random tangent vector in the tangent space of the point `p` on the
`ProductManifold` `M`.
"""
function random_tangent(M::ProductManifold, p, options...)
    X = map(
        (m, p) -> random_tangent(m, p, options...),
        M.manifolds,
        submanifold_components(M, p),
    )
    return ProductRepr(X...)
end

@doc raw"""
    random_tangent(M::Rotations, p[, type=:Gaussian, σ=1.0])

return a random tangent vector in the tangent space
$T_x\mathrm{SO}(n)$ of the point `x` on the `Rotations` manifold `M` by generating
a random skew-symmetric matrix. The function takes the real upper triangular matrix of a
(Gaussian) random matrix $A$ with dimension $n\times n$ and subtracts its transposed matrix.
Finally, the matrix is normalized.
"""
function random_tangent(M::Rotations, p, ::Val{:Gaussian}, σ::Real=1.0)
    d = manifold_dimension(M)
    if d == 0
        return zeros(1, 1)
    else
        A = σ * randn(Float64, d, d)
        A = triu(A, 1) - transpose(triu(A, 1))
        A = (1 / norm(A)) * A
        return A
    end
end

@doc raw"""
    random_tangent(M::Sphere, p[, :Gaussian, σ=1.0])

return a random tangent vector in the tangent space of `p` on the `Sphere` `M`.
"""
function random_tangent(M::Sphere, p, ::Val{:Gaussian}, σ::Float64=1.0)
    n = σ * randn(size(p)) # Gaussian in embedding
    return project(M, p, n) #project to TpM (keeps Gaussianness)
end

@doc raw"""
    random_tangent(M::Stiefel, p[,type=:Gaussian, σ=1.0])

return a (Gaussian) random vector from the tangent space $T_p\mathrm{St}(n,k)$ with mean
zero and standard deviation `σ` by projecting a random Matrix onto the  `p`.
"""
function random_tangent(M::Stiefel, p, ::Val{:Gaussian}, σ::Float64=1.0)
    Z = σ * randn(eltype(p), size(p))
    X = project(M, p, Z)
    X = X ./ norm(X)
    return X
end

@doc raw"""
    random_tangent(M::SymmetricPositiveDefinite, p[, :Gaussian, σ = 1.0])

generate a random tangent vector in the tangent space of the point `p` on the
`SymmetricPositiveDefinite` manifold `M` by using a Gaussian distribution
with standard deviation `σ` on an ONB of the tangent space.
"""
function random_tangent(M::SymmetricPositiveDefinite, p, ::Val{:Gaussian}, σ::Float64=0.01)
    # generate ONB in TxM
    I = one(p)
    B = get_basis(M, p, DiagonalizingOrthonormalBasis(I))
    Ξ = get_vectors(M, p, B)
    Ξx = vector_transport_to.(Ref(M), Ref(I), Ξ, Ref(p), Ref(ParallelTransport()))
    return sum(σ * randn(length(Ξx)) .* Ξx)
end

@doc raw"""
    random_tangent(M::SymmetricPositiveDefinite, p, :Rician [,σ = 0.01])

generate a random tangent vector in the tangent space of `p` on
the `SymmetricPositiveDefinite` manifold `M` by using a Rician distribution
with standard deviation `σ`.
"""
function random_tangent(::SymmetricPositiveDefinite, p, ::Val{:Rician}, σ::Real=0.01)
    # Rician
    C = cholesky(Hermitian(p))
    R = C.L + sqrt(σ) * triu(randn(size(p, 1), size(p, 2)), 0)
    return R * R'
end

@doc raw"""
    random_tangent(M::TangentBundle, p, options...)

generate a random tangent vector at `p` on the tangent bundle by calling
[`random_tangent`](@ref) with the given `options...` twice.
"""
function random_tangent(M::TangentBundle, p, options...)
    X = random_tangent(M.manifold, p[M, :point], options...)
    Y = random_tangent(M.manifold, p[M, :point], options...)
    return ProductRepr(X, Y)
end


"""
    spar = splineparams(T, D1, Nseg, pcof)

Constructor for struct splineparams, which sets up the parameters for a regular B-spline function
(without carrier waves).

# Arguments
- `T:: Float64`: Duration of spline function
- `D1:: Int64`: Number of basis functions in each spline
- `Nseg:: Int64`:  Number of splines (real, imaginary, different ctrl func)
- `pcof:: Array{Float64, 1}`: Coefficient vector. Must have D1*Nseg elements

# External links
* [Spline Wavelet](https://en.wikipedia.org/wiki/Spline_wavelet#Quadratic_B-spline) on Wikipedia.
"""
struct splineparams
    T::Float64
    D1::Int64 # Number of coefficients per spline
    tcenter::Array{Float64,1}
    dtknot::Float64
    pcof::Array{Float64,1} # pcof should have D1*Nseg elements
    Nseg::Int64 # Number of segments (real, imaginary, different ctrl func)
    Ncoeff:: Int64 # Total number of coefficients

# new, simplified constructor
    function splineparams(T, D1, Nseg, pcof)
        dtknot = T/(D1 -2)
        tcenter = dtknot.*(collect(1:D1) .- 1.5)
        new(T, D1, tcenter, dtknot, pcof, Nseg, Nseg*D1)
    end

end

# bspline2: Evaluate quadratic bspline function
"""
    f = bspline2(t, splineparam, splinefunc)

Evaluate a B-spline function. See also the `splineparams` constructor.

# Arguments
- `t::Float64`: Evaluate spline at parameter t ∈ [0, param.T]
- `param::splineparams`: Parameters for the spline
- `splinefunc::Int64`: Spline function index ∈ [0, param.Nseg-1]
"""
@inline function bspline2(t::Float64, param::splineparams, splinefunc::Int64)
  f = 0.0

  dtknot = param.dtknot
  width = 3*dtknot

  offset = splinefunc*param.D1

  k = max.(3, ceil.(Int64,t./dtknot + 2)) # Unsure if this line does what it is supposed to
  k = min.(k, param.D1)

  # 1st segment of nurb k
  tc = param.tcenter[k]
  tau = (t .- tc)./width
  f = f + param.pcof[offset+k] * (9/8 .+ 4.5*tau + 4.5*tau^2) # test to remove square for extra speed

  # 2nd segment of nurb k-1
  tc = param.tcenter[k-1]
  tau = (t - tc)./width
  f = f .+ param.pcof[offset+k.-1] .* (0.75 - 9 *tau^2)

  # 3rd segment of nurb k-2
  tc = param.tcenter[k.-2]
  tau = (t .- tc)./width
  f = f + param.pcof[offset+k.-2] * (9/8 - 4.5*tau + 4.5*tau.^2)
end

"""
    g = gradbspline2(t, param, splinefunc)

Evaluate the gradient of a spline function with respect to all coefficient.
NOTE: the index of the spline functions are 0-based. For a set of
coupled controls, mod(`splinefunc`,2)=0 corresponds to ∇ p_j(t) and mod(`splinefunc`,2) = 1
corresponds to ∇ q_j(t), where j = div(splinefunc,2).

# Arguments
- `t::Float64`: Evaluate spline at parameter t ∈ [0, param.T]
- `param::splineparams`: Spline parameter object
- `splinefunc::Int64`: Spline function index ∈ [0, param.Nseg-1]
"""
@inline function gradbspline2(t::Float64,param::splineparams, splinefunc::Int64)

# NOTE: param.Nseg used to be '2'
  g = zeros(param.Nseg*param.D1) # real and imag parts for both f and g

  dtknot = param.dtknot
  width = 3*dtknot

  offset = splinefunc*param.D1

  k = max.(3, ceil.(Int64,t./dtknot .+ 2)) # t_knot(k-1) < t <= t_knot(k), but t=0 needs to give k=3
  k = min.(k, param.D1) # protect agains roundoff that sometimes makes t/dt > N_nurbs-2

  #1st segment of nurb k
  tc = param.tcenter[k]
  tau = (t .- tc)./width
  g[offset .+ k] = (9/8 .+ 4.5.*tau .+ 4.5.*tau.^2);

  #2nd segment of nurb k-1
  tc = param.tcenter[k.-1]
  tau = (t .- tc)./width
  g[offset .+ k.-1] = (0.75 .- 9 .*tau.^2)

  # 3rd segment og nurb k-2
  tc = param.tcenter[k.-2]
  tau = (t .- tc)./width
  g[offset .+ k.-2] = (9/8 .- 4.5.*tau .+ 4.5.*tau.^2);
  return g
end

"""
    bcpar = bcparams(T, D1, Ncoupled, Nunc, omega, pcof)

General constructor of `struct bcparams` for setting up B-splines with carrier waves.

    bcpar = bcparams(T, D1, omega, pcof)

Simplified constructor for the case when there are no uncoupled controls and `Ncoupled = size(omega,1)`.

# Arguments
- `T:: Float64`: Duration of spline function
- `D1:: Int64`: Number of basis functions in each segment
- `Ncoupled::Int64`: Number of coupled controls in the simulation
- `Nunc::Int64`: Number of uncoupled controls in the simulation
- `omega::Array{Float64,2}`: Carrier wave frequencies
- `pcof:: Array{Float64, 1}`: Coefficient vector. Must have D1*Nseg elements

# First dimensions of the `omega` array:
- Without uncoupled controls, `Nunc=0` and `size(omega,1) = Ncoupled`.
- With uncoupled controls, `Nunc > 0` and `size(omega,1) = Ncoupled + Nunc`.

# Second dimension of the `omega` array:
- `size(omega, 2) = Nfreq`

# Ordering of the `pcof` array:
First consider the case without uncoupled control functions, `Nunc = 0`:
Then the `pcof` array then has `2*Ncoupled*Nfreq*D1` elements.
Each `ctrl ∈ [1,Ncoupled]` and `freq ∈ [1,Nfreq]` corresponds to `D1` elements in
the `pcof` vector. For the case `Ncoupled = 2` and `Nfreq = 2`, the elements are ordered according to

| ctrl    | freq    | α_1  | α_2 |
| ------- | -------- | -------- | -------|
| 1      | 1        | 1:D1             | D1+1:2 D1 |
| 1      | 2        | 2 D1+1: 3 D1 | 3 D1+1:4 D1 |
| 2      | 1        | 4 D1+1: 5 D1 | 5 D1+1:6 D1 |
| 2      | 2        | 6 D1+1: 7 D1 | 7 D1+1: 8D1 |

If there are uncoupled controls, `Nunc > 0`, the `pcof` array should have `(2*Ncoupled + Nunc)*Nfreq*D1` elements.
The last `Nunc*Nfreq*D1` elements correspond to the uncoupled control functions and are ordered in a corresponding way.

# External links
* [Spline Wavelet](https://en.wikipedia.org/wiki/Spline_wavelet#Quadratic_B-spline) on Wikipedia.
"""
struct bcparams
    T ::Float64
    D1::Int64 # number of B-spline coefficients per control function
    om::Array{Float64,2} #Carrier wave frequencies [rad/s], size Nfreq
    tcenter::Array{Float64,1}
    dtknot::Float64
    pcof::Array{Float64,1} # coefficients for all 2*Ncoupled splines, size Ncoupled*D1*Nfreq*2 (*2 because of sin/cos)
    Nfreq::Int64 # Number of frequencies
    Ncoeff:: Int64 # Total number of coefficients
    Ncoupled::Int64 # Number of B-splines functions for the coupled ctrl Hamiltonians
    Nunc::Int64 # Number of B-spline functions  for the UNcoupled ctrl Hamiltonians

    # New constructor to allow defining number of symmetric Hamiltonian terms
    function bcparams(T::Float64, D1::Int64, Ncoupled::Int64, Nunc::Int64, omega::Array{Float64,2}, pcof::Array{Float64,1})
        dtknot = T/(D1 -2)
        tcenter = dtknot.*(collect(1:D1) .- 1.5)
        Nfreq = size(omega,2)
        nCoeff = Nfreq*D1*2*(Ncoupled + Nunc)
        if nCoeff != length(pcof)
            println("nCoeff = ", nCoeff, " Nfreq = ", Nfreq, " D1 = ", D1, " Ncoupled = ", Ncoupled, " Nunc = ", Nunc, " len(pcof) = ", length(pcof))
            throw(DimensionMismatch("Inconsistent number of coefficients and size of parameter vector (nCoeff ≠ length(pcof)."))
        end
        new(T, D1, omega, tcenter, dtknot, pcof, Nfreq, nCoeff, Ncoupled, Nunc)
    end

end

# simplified constructor (assumes no uncoupled terms)
function bcparams(T::Float64, D1::Int64, omega::Array{Float64,2}, pcof::Array{Float64,1})
  dtknot = T/(D1 -2)
  tcenter = dtknot.*(collect(1:D1) .- 1.5)
  Ncoupled = size(omega,1) # should check that Ncoupled >=1
  Nfreq = size(omega,2)
  Nunc = 0
  nCoeff = Nfreq*D1*2*Ncoupled
  if nCoeff != length(pcof)
    throw(DimensionMismatch("Inconsistent number of coefficients and size of parameter vector (nCoeff ≠ length(pcof)."))
  end
  bcparams(T, D1, Ncoupled, Nunc, omega, pcof)
end

"""
    f = bcarrier2(t, params, func)

Evaluate a B-spline function with carrier waves. See also the `bcparams` constructor.

# Arguments
- `t::Float64`: Evaluate spline at parameter t ∈ [0, param.T]
- `param::params`: Parameters for the spline
- `func::Int64`: Spline function index ∈ [0, param.Nseg-1]
"""
@inline function bcarrier2(t::Float64, params::bcparams, func::Int64)
    # for a single oscillator, func=0 corresponds to p(t) and func=1 to q(t)
    # in general, 0 <= func < 2*Ncoupled + Nunc

    # compute basic offset: func 0 and 1 use the same spline coefficients, but combined in a different way
    osc = div(func, 2) # osc is base 0; 0<= osc < Ncoupled
    q_func = func % 2 # q_func = 0 for p and q_func=1 for q

    f = 0.0 # initialize

    dtknot = params.dtknot
    width = 3*dtknot

    k = max.(3, ceil.(Int64,t./dtknot + 2)) # pick out the index of the last basis function corresponding to t
    k = min.(k, params.D1) #  Make sure we don't access outside the array

    if func < 2*(params.Ncoupled + params.Nunc)
        # Coupled and uncoupled controls
        @fastmath @inbounds @simd for freq in 1:params.Nfreq
            fbs1 = 0.0 # initialize
            fbs2 = 0.0 # initialize
            # offset in parameter array (osc = 0,1,2,...
            # Vary freq first, then osc
            offset1 = 2*osc*params.Nfreq*params.D1 + (freq-1)*2*params.D1
            offset2 = 2*osc*params.Nfreq*params.D1 + (freq-1)*2*params.D1 + params.D1

            # 1st segment of nurb k
            tc = params.tcenter[k]
            tau = (t .- tc)./width
            fbs1 += params.pcof[offset1+k] * (9/8 .+ 4.5*tau + 4.5*tau^2)
            fbs2 += params.pcof[offset2+k] * (9/8 .+ 4.5*tau + 4.5*tau^2)

            # 2nd segment of nurb k-1
            tc = params.tcenter[k-1]
            tau = (t - tc)./width
            fbs1 += params.pcof[offset1+k.-1] .* (0.75 - 9 *tau^2)
            fbs2 += params.pcof[offset2+k.-1] .* (0.75 - 9 *tau^2)

            # 3rd segment of nurb k-2
            tc = params.tcenter[k.-2]
            tau = (t .- tc)./width
            fbs1 += params.pcof[offset1+k-2] * (9/8 - 4.5*tau + 4.5*tau.^2)
            fbs2 += params.pcof[offset2+k-2] * (9/8 - 4.5*tau + 4.5*tau.^2)

            #    end # for carrier phase
            # p(t)
            if q_func==1
                f += fbs1 * sin(params.om[osc+1,freq]*t) + fbs2 * cos(params.om[osc+1,freq]*t) # q-func
            else
                f += fbs1 * cos(params.om[osc+1,freq]*t) - fbs2 * sin(params.om[osc+1,freq]*t) # p-func
            end
        end # for freq
    end # if
    # else

    #   # uncoupled control
    #   @fastmath @inbounds @simd for freq in 1:params.Nfreq
    #     fbs = 0.0 # initialize

    #     # offset to uncoupled parameters
    #     # offset = (2*params.Ncoupled)*params.D1*params.Nfreq + (fun - 2*params.Ncoupled)*params.D1*params.Nfreq  +  (freq-1)*params.D1
    #     offset = func*params.D1*params.Nfreq  +  (freq-1)*params.D1

    #     # 1st segment of nurb k
    #     if k <= params.D1  # could do the if statement outside the for-loop instead
    #       tc = params.tcenter[k]
    #       tau = (t .- tc)./width
    #       fbs += params.pcof[offset+k] * (9/8 .+ 4.5*tau + 4.5*tau^2)
    #     end

    #     # 2nd segment of nurb k-1
    #     if k >= 2 && k <= params.D1 + 1
    #       tc = params.tcenter[k-1]
    #       tau = (t - tc)./width
    #       fbs += params.pcof[offset+k.-1] .* (0.75 - 9 *tau^2)
    #     end

    #     # 3rd segment of nurb k-2
    #     if k >= 3 && k <= params.D1 + 2
    #       tc = params.tcenter[k.-2]
    #       tau = (t .- tc)./width
    #       fbs += params.pcof[offset+k-2] * (9/8 - 4.5*tau + 4.5*tau.^2)
    #     end

    #     # FMG: For now this would assume that every oscillator would have
    #     # an uncoupled term if uncoupled terms exist in the problem. Perhaps
    #     # update parameter struct to track terms if they are only present on some
    #     # oscillators?
    #     ind = func - 2*params.Ncoupled
    #     f += fbs * cos(params.om[ind+1,freq]*t) # spl is 0-based
    #   end # for

    return f
end

"""
    gradbcarrier2!(t, params, func, g) -> g
Evaluate the gradient of a control function with respect to all coefficient.

NOTE: the index of the control functions is 0-based. For a set of
coupled controls, mod(`func`,2)=0 corresponds to ∇ p_j(t) and mod(`func`,2) = 1
corresponds to ∇ q_j(t), where j = div(func,2).

# Arguments
- `t::Float64`: Evaluate spline at parameter t ∈ [0, param.T]
- `params::bcparams`: Parameters for the spline
- `func::Int64`: Control function index ∈ [0, param.Nseg-1]
- `g::Array{Float64,1}`: Preallocated array to store calculated gradient
"""
function gradbcarrier2!(t::Float64, params::bcparams, func::Int64, g::Array{Float64,1})

    # compute basic offset: func 0 and 1 use the same spline coefficients, but combined in a different way
    osc = div(func, 2)
    q_func = func % 2 # q_func = 0 for p and q_func=1 for q

    # allocate array for returning the results
    # g = zeros(length(params.pcof)) # cos and sin parts

    g .= 0.0

    dtknot = params.dtknot
    width = 3*dtknot

    k = max.(3, ceil.(Int64,t./dtknot .+ 2)) # t_knot(k-1) < t <= t_knot(k), but t=0 needs to give k=3
    k = min.(k, params.D1) # protect agains roundoff that sometimes makes t/dt > N_nurbs-2

    if func < 2*(params.Ncoupled + params.Nunc)
        @fastmath @inbounds @simd for freq in 1:params.Nfreq

            # offset in parameter array (osc = 0,1,2,...
            # Vary freq first, then osc
            offset1 = 2*osc*params.Nfreq*params.D1 + (freq-1)*2*params.D1
            offset2 = 2*osc*params.Nfreq*params.D1 + (freq-1)*2*params.D1 + params.D1

            #1st segment of nurb k
            tc = params.tcenter[k]
            tau = (t .- tc)./width
            bk = (9/8 .+ 4.5.*tau .+ 4.5.*tau.^2)
            if q_func==1
                g[offset1 .+ k] = bk * sin(params.om[osc+1,freq]*t)
                g[offset2 .+ k] = bk * cos(params.om[osc+1,freq]*t)
            else # p-func
                g[offset1 .+ k] = bk * cos(params.om[osc+1,freq]*t)
                g[offset2 .+ k] = -bk * sin(params.om[osc+1,freq]*t)
            end

            #2nd segment of nurb k-1
            tc = params.tcenter[k.-1]
            tau = (t .- tc)./width
            bk = (0.75 .- 9 .*tau.^2)
            if q_func==1
                g[offset1 .+ (k-1)] = bk * sin(params.om[osc+1,freq]*t)
                g[offset2 .+ (k-1)] = bk * cos(params.om[osc+1,freq]*t)
            else # p-func
                g[offset1 .+ (k-1)] = bk * cos(params.om[osc+1,freq]*t)
                g[offset2 .+ (k-1)] = -bk * sin(params.om[osc+1,freq]*t)
            end

            # 3rd segment og nurb k-2
            tc = params.tcenter[k.-2]
            tau = (t .- tc)./width
            bk = (9/8 .- 4.5.*tau .+ 4.5.*tau.^2)
            if q_func==1
                g[offset1 .+ (k-2)] = bk * sin(params.om[osc+1,freq]*t)
                g[offset2 .+ (k-2)] = bk * cos(params.om[osc+1,freq]*t)
            else # p-func
                g[offset1 .+ (k-2)] = bk * cos(params.om[osc+1,freq]*t)
                g[offset2 .+ (k-2)] = -bk * sin(params.om[osc+1,freq]*t)
            end

        end #for freq
    # else
    #     # uncoupled control case
    #   @fastmath @inbounds @simd for freq in 1:params.Nfreq

    #     # offset
    #     offset = func*params.D1*params.Nfreq  +  (freq-1)*params.D1
    #     ind = func - 2*params.Ncoupled

    #     #1st segment of nurb k
    #     if k <= params.D1  # could do the if statement outside the for-loop instead
    #       tc = params.tcenter[k]
    #       tau = (t .- tc)./width
    #       g[offset .+ k] = (9/8 .+ 4.5.*tau .+ 4.5.*tau.^2)*cos(params.om[ind+1,freq]*t)
    #     end

    #     #2nd segment of nurb k-1
    #     if k >= 2 && k <= params.D1 + 1
    #       tc = params.tcenter[k.-1]
    #       tau = (t .- tc)./width
    #       g[offset .+ k.-1] = (0.75 .- 9 .*tau.^2)*cos(params.om[ind+1,freq]*t)
    #     end

    #     # 3rd segment og nurb k-2
    #     if k >= 3 && k <= params.D1 + 2
    #       tc = params.tcenter[k.-2]
    #       tau = (t .- tc)./width
    #       g[offset .+ k.-2] = (9/8 .- 4.5.*tau .+ 4.5.*tau.^2)*cos(params.om[ind+1,freq]*t)
    #     end

    #   end #for freq

       end #if
 end


# Copyright (c) 2018: Matthew Wilhelm & Matthew Stuber.
# This code is licensed under MIT license (see LICENSE.md for full details)
#############################################################################
# ReverseMcCormick.jl
# A Reverse McCormick operator library in Julia
# See https://github.com/PSORLab/ReverseMcCormick.jl
#############################################################################
# src/reverse_operators/inverse_hyperbolic.jl
# Contains definitions of reverse McCormick operators of asinh, acosh, atanh,
# asech, acsch, acoth
#############################################################################

"""
$(SIGNATURES)

Reverse McCormick operator for `asinh`.
"""
function asinh_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    x = x ∩ sinh(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `acosh`.
"""
function acosh_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    y = y ∩ Interval{Float64}(0.0, Inf)
    x = x ∩ cosh(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `atanh`.
"""
function atanh_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    x = x ∩ tanh(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `asech`.
"""
function asech_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    y = y ∩ Interval{Float64}(0.0, Inf)
    x = x ∩ sech(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `acsch`.
"""
function acsch_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    x = x ∩ csch(y)
    y, x
end

"""
$(SIGNATURES)

Reverse McCormick operator for `acoth`.
"""
function acoth_rev(y::MC, x::MC)
    isempty(y) && (return y, y)
    x = x ∩ coth(y)
    y, x
end


# Power iteration algorithm for computing the spectral norm
function power_iteration!(W, u, n_iterations)
    v = nothing
    for i=1:n_iterations
        Wu = W' * u
        v = Wu ./ (norm(Wu) + eps(Float32))
        Wv = W * v
        u .= Wv ./ (norm(Wv) + eps(Float32))
    end
    u, v
end

# Compute the maximum singular value
msv(u, v, W) = u' * W * v

## Dense Layer with Spectral Normalization
struct DenseSN{F,S<:AbstractMatrix, B, I<:Int, VV<:AbstractArray}
  weight::S
  bias::B
  σ::F
  n_iterations::I # Number of power iterations for computing max singular value
  u::VV # Left vector for power iteration
  function DenseSN(W::M, bias=true, σ::F=identity, n_iterations::I=1, u::VV=randn(Float32, size(W,1), 1)) where {M<:AbstractMatrix, F, I<:Int, VV<:AbstractArray}
    b = Flux.create_bias(W, bias, size(W,1))
    new{F,M,typeof(b), I, VV}(W, b, σ, n_iterations, u)
  end
end

function DenseSN(in::Integer, out::Integer, σ=identity; init=Flux.glorot_uniform, bias=true, n_iterations=1, u=randn(Float32, out, 1))
  DenseSN(init(out, in), bias, σ, n_iterations, u)
end

Flux.@functor DenseSN

Flux.trainable(a::DenseSN) = (a.weight, a.bias)

function (a::DenseSN)(x::AbstractVecOrMat)
  W, b, σ = a.weight, a.bias, a.σ
  u, v = Zygote.ignore(() -> power_iteration!(W, a.u, a.n_iterations))
  σ.((W ./ msv(u, v, W))*x .+ b)
end

(a::DenseSN)(x::AbstractArray) = reshape(a(reshape(x, size(x,1), :)), :, size(x)[2:end]...)

function Base.show(io::IO, l::DenseSN)
  print(io, "DenseSN(", size(l.weight, 2), ", ", size(l.weight, 1))
  l.σ == identity || print(io, ", ", l.σ)
  print(io, ")")
end

## Convluational layer with Spectral Normalization
struct ConvSN{N,M,F,A,V, I<:Int, VV<:AbstractArray}
  σ::F
  weight::A
  bias::V
  stride::NTuple{N,Int}
  pad::NTuple{M,Int}
  dilation::NTuple{N,Int}
  n_iterations::I # Number of power iterations for computing max singular value
  u::VV # Left vector for power iteration
end

function ConvSN(w::AbstractArray{T,N}, b::Union{Flux.Zeros, AbstractVector{T}}, σ = identity;
              stride = 1, pad = 0, dilation = 1, n_iterations = 1) where {T,N}
  stride = Flux.expand(Val(N-2), stride)
  dilation = Flux.expand(Val(N-2), dilation)
  pad = Flux.calc_padding(Conv, pad, size(w)[1:N-2], dilation, stride)
  u = randn(Float32, size(to2D(w), 1), 1)
  return ConvSN(σ, w, b, stride, pad, dilation, n_iterations, u)
end

function ConvSN(k::NTuple{N,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
            init = Flux.glorot_uniform,  stride = 1, pad = 0, dilation = 1,
            weight = Flux.convfilter(k, ch, init = init), bias = Flux.zeros(ch[2]), n_iterations = 1) where N
  ConvSN(weight, bias, σ, stride = stride, pad = pad, dilation = dilation, n_iterations = n_iterations)
end

Flux.@functor ConvSN

Flux.trainable(a::ConvSN) = (a.weight, a.bias)

function (c::ConvSN)(x::AbstractArray)
  σ, b = c.σ, reshape(c.bias, ntuple(_->1, length(c.stride))..., :, 1)
  cdims = DenseConvDims(x, c.weight; stride=c.stride, padding=c.pad, dilation=c.dilation)
  u, v = Zygote.ignore(() -> power_iteration!(to2D(c.weight), c.u, c.n_iterations))
  σ.(conv(x, c.weight ./ msv(u, v, to2D(c.weight)), cdims) .+ b)
end

function Base.show(io::IO, l::ConvSN)
  print(io, "ConvSN(", size(l.weight)[1:ndims(l.weight)-2])
  print(io, ", ", size(l.weight, ndims(l.weight)-1), "=>", size(l.weight, ndims(l.weight)))
  l.σ == identity || print(io, ", ", l.σ)
  print(io, ")")
end