include("../psmsdst.jl")

"""
    solvegreedy(problem)

Solves the given PSMSDST problem using a greedy heuristic: for every pair (j,m)
of available jobs and machines, the pair that produces the lowest increment
in the objective function is programmed, until a complete solution is reached.
"""
function solvegreedy(problem::PSMSDSTProblem)::Vector{Vector{Int}}

    machines = 1:problem.nummachines

    # Create resulting vector for the solution
    solution = [ Vector{Int}() for machineid in machines ]

    # Create a set for the jobs that have not been programmed yet
    availablejobs = Set{Int}(1:problem.numjobs)

    # Create a vector for the finish times of each machine
    machinefinishtimes = [ 0 for machineid in machines ]

    # Variable that holds the current objective function
    currentmaxfinishtime = maximum(machinefinishtimes)

    # Matrix that contains how the machine finish time would increase when
    # a job is programmed on it.
    delays = [ problem.durations[job] for job in 1:problem.numjobs,
                machine in machines ]

    # Matrix that indicates if programming a job on a machine would invalidate
    # the pre-calculated delay.
    createsdelayupdate = [ true for job in 1:problem.numjobs,
        machine in machines ]

    # The id of the machine that will be updated because the cached delays
    # where invalidated.
    updatedelays = 0

    # Program jobs until a complete solution is reached
    while !isempty(availablejobs)

        # Cartesian product, (j,m) pairs of jobs and machines.
        options = Iterators.product(availablejobs,machines)

        if updatedelays > 0
            # The delays of a machine need to be recalculated.
            machine = updatedelays
            for job in availablejobs

                # The initial delay is just the duration of the job
                delay = problem.durations[job]

                if length(solution[machine]) > 0
                    lastjob = solution[machine][end]
                    lastfamily = problem.families[lastjob]
                    thisfamily = problem.families[job]

                    # If the family of the next job is different to the
                    # family of this job, this means that all the delays are
                    # invalidated because the setup times will change.
                    createsdelayupdate[job,machine] = thisfamily != lastfamily

                    # Add the setup time to the delay
                    delay += problem.setuptimes[lastfamily,thisfamily]
                end

                delays[job,machine] = delay
            end
        end

        # Look for the best option
        bestoption = nothing
        bestcost = typemax(Int64)

        for (job,machine) in options
            # Finish time of the machine if the job was programmed on it.
            virtualmachinefinish = (machinefinishtimes[machine]
                +delays[job,machine])

            # Objective function if the job was programmed on it.
            cost = max(currentmaxfinishtime,virtualmachinefinish)

            if cost < bestcost
                bestoption = (job,machine,virtualmachinefinish,
                    createsdelayupdate[job,machine])
                bestcost = cost
            end
        end

        job,machine,finishtime,createsupdate = bestoption

        # If the job creates an update in the delays, they have to be
        # re-calculated in the next iteration.
        if createsupdate
            updatedelays = machine
        else
            updatedelays = 0
        end

        machinefinishtimes[machine] = finishtime

        currentmaxfinishtime = maximum(machinefinishtimes)
        push!(solution[machine],job)
        delete!(availablejobs,job)
    end

    return solution
end

#
#    getsolver()
#
#Returns a pair containing the name of the solver, and a pre-parameterized
#function that only takes the problem as argument.
#
#Internally it also solves a small problem using the solver, ensuring that
#the code for the solver is compiled ahead of time and there is no reduced
#performance during evaluation.
#
## Examples
#```julia-repl
#julia> include("solvers/solver.jl")(param1=42.0,param2=64.0)
#"Solver param1=42 param2=64"=>#1 (generic function with 1 method)
#```

function getsolver()::Pair{String,Function}
    solver = problem -> solvegreedy(problem)

    presolve(solver)

    return "Greedy"=>solver
end


abstract type FluidFlowProblem end
abstract type SteadyStateProblem <: FluidFlowProblem end
abstract type TimeDependantProblem <: FluidFlowProblem end

abstract type InitialValueProblem <: FluidFlowProblem end
abstract type DoubleDistributionProblem <: FluidFlowProblem end

has_external_force(::FluidFlowProblem) = false
function velocity_gradient(
    ::FluidFlowProblem,
    x::T,
    y::T,
    timestep::Real = 0.0,
) where {T <: Real}
    return zeros(T, 2, 2)
end

function range(problem::FluidFlowProblem)
    Δx = problem.domain_size[1] / problem.NX
    Δy = problem.domain_size[2] / problem.NY

    x_range = range(Δx / 2, stop = problem.domain_size[1] - Δx / 2, length = problem.NX)
    y_range = range(Δy / 2, stop = problem.domain_size[1] - Δy / 2, length = problem.NY)

    return x_range, y_range
end

boundary_conditions(problem::FluidFlowProblem) = BoundaryCondition[]

function deviatoric_tensor(
    q::Quadrature,
    problem::FluidFlowProblem,
    x::T,
    y::T,
    time::Real = 0.0,
) where {T <: Real}
    a = velocity_gradient(problem, x, y, time)
    ν = viscosity(problem)

    σ = -ν * [
        2 * a[1, 1] a[1, 2] + a[2, 1]
        a[1, 2] + a[2, 1] 2 * a[2, 2]
    ]

    return σ
end

function pressure_tensor(
    q::Quadrature,
    problem::FluidFlowProblem,
    x::T,
    y::T,
    time::Real = 0.0,
) where {T <: Real}
    A = problem.A
    B = problem.B

    p = pressure(q, problem, x, y, time) * I
    return p - deviatoric_tensor(q, problem, x, y, time)
end

function force(
    problem::FluidFlowProblem,
    x_idx::Int,
    y_idx::Int,
    time::Real = 0.0,
) where {Int <: Integer}
    x_range, y_range = range(problem)

    x = x_range[x_idx]
    y = y_range[y_idx]

    return force(problem, x, y, time)
end

is_steady_state(problem::FluidFlowProblem) = problem.static
is_time_dependant(problem::FluidFlowProblem) = !problem.static

# Dimensionless
viscosity(problem::FluidFlowProblem) = problem.ν * delta_x(problem)^2 / delta_t(problem)
heat_diffusion(problem::FluidFlowProblem) =
    problem.κ * delta_x(problem)^2 / delta_t(problem)
reynolds(problem::FluidFlowProblem) = problem.NY * problem.u_max / problem.ν

function delta_t(problem::FluidFlowProblem)
    return delta_x(problem) * problem.u_max
end

function delta_x(problem::FluidFlowProblem)
    if (problem.NX > problem.NY)
        return problem.domain_size[1] * (1 / problem.NX)
    end

    return problem.domain_size[2] * (1 / problem.NY)
end

lattice_viscosity(problem) = problem.ν #::FluidFlowProblem)
lattice_density(q, problem::FluidFlowProblem, x, y, t = 0.0) = density(q, problem, x, y, t)
lattice_velocity(q, problem::FluidFlowProblem, x, y, t = 0.0) =
    problem.u_max * velocity(problem, x, y, t)
lattice_pressure(q, problem::FluidFlowProblem, x, y, t = 0.0) =
    problem.u_max^2 * pressure(q, problem, x, y, t)
lattice_force(problem::FluidFlowProblem, x, y, t = 0.0) =
    problem.u_max * delta_t(problem) * force(problem, x, y, t)
lattice_temperature(q, problem::FluidFlowProblem, x, y, t = 0.0) =
    pressure(q, problem, x, y) / density(q, problem, x, y)

dimensionless_viscosity(problem) = problem.ν * delta_x(problem)^2 / delta_t(problem)
dimensionless_density(problem::FluidFlowProblem, ρ) = ρ
dimensionless_velocity(problem::FluidFlowProblem, u) = u / problem.u_max
dimensionless_pressure(q, problem::FluidFlowProblem, p) = p# (p - 1.0) / (q.speed_of_sound_squared * problem.u_max^2 )
dimensionless_temperature(q, problem::FluidFlowProblem, T) = T #* q.speed_of_sound_squared
dimensionless_force(problem::FluidFlowProblem, F) = F / (problem.u_max * delta_t(problem))
dimensionless_stress(problem::FluidFlowProblem, σ) = begin
    # Rescale to dimensionless number (TODO check why problem.u_max)
    factor = 1 / problem.u_max^2

    return σ * factor
end

function equilibrium(q::Quadrature, problem::FluidFlowProblem, x::T, y::T) where {T <: Real}
    equilibrium(
        q,
        lattice_density(q, problem, x, y),
        lattice_velocity(q, problem, x, y),
        lattice_temperature(q, problem, x, y),
    )
end

include("taylor_green_vortex.jl")
include("decaying_shear_flow.jl")
include("poiseuille.jl")
include("couette_flow.jl")
include("lid_driven_cavity.jl")
include("linear_hydrodynamics_modes.jl")
include("second_order_convergence.jl")

# error(::Val{:density}, node, solution) = density(node) - density(solution)


# types.jl

# Physical Constants

"Planck's constant, (kgm²s⁻¹)."
const hbar = const ħ = 1.054571817e-34
"Electron charge, (kgm²s⁻²)."
const eV = const q = const ElectronVolt = 1.602176634e-19
"Electron mass, (kg)."
const me = MassElectron = 9.1093837015e-31
"Boltzmann's constant, (kgm²K⁻¹)."
const Boltzmann = const kB = 1.380649e-23
"Permittivity of free space, (C²N⁻¹m⁻²)."
const ε_0 = 8.85418682e-12
"Speed of light, (ms⁻¹)."
const c = 299792458

# Structures

"Polaron. Structure to store data of polaron solution and other parameters, for each temperature or frequency."
struct Polaron
    "T, temperature (K)."
    T
    "Kμ, Kadanoff mobility (cm²V⁻¹s⁻¹)."
    Kμ
    "Hμ, Hellwarth mobility (cm²V⁻¹s⁻¹)."
    Hμ
    "FHIPμ, FHIP mobility (cm²V⁻¹s⁻¹)."
    FHIPμ
    "k, spring constant."
    k
    "M, renormalised (phonon-drag) mass (mₑ)."
    M
    "Osaka free energy components (A,B,C) and total (F) (unitless). See Hellwarth et al. 1999 PRB Part IV."
    A; B; C; F
    "Tau, relaxation time from Kadanoff Boltzmann transport equation (s)."
    Tau
    "v and w, raw variational parameters (unitless)."
    v; w
    "βred, reduced thermodynamic beta (unitless)."
    βred
    "rfsi, Feynman polaron radius (Schultz) (m)."
    rfsi
    "rfsmallalpha, small-alpha asymptotic approximation (unitless)."
    rfsmallalpha

    # Setup of simulation. These parameters are sent to the function.

    "α, Fröhlich alpha coupling parameter (unitless)."
    α
    "mb, Band effective mass (mₑ)."
    mb
    "ω, effective dielectric frequency (2π THz)."
    ω
end

# structure initialisation
Polaron()=Polaron([],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[],[])

"NewPolaron. Structure to store data of polaron solution and other parameters, for each temperature or frequency."
struct NewPolaron
    "α, Fröhlich alpha coupling parameter (unitless)"
    α
    "T, temperature (K)."
    T
    "β, reduced Thermodynamic beta (unitless)."
    β
    "ω, phonon frequency (2π THz)."
    ω
    "v, variational parameters (unitless)."
    v
    "w, variational parameters (unitless)."
    w
    "κ, fictitious spring constant (multiples of mₑ) (kgs⁻²)."
    κ
    "M, fictitious particle (multiples of mₑ) (kg)."
    M
    "F, free energy (meV)."
    F
    "Ω, electric field frequencies (multiples of phonon frequency ω) (2π THz)."
    Ω
    "Z, complex impedence (V/A)."
    Z
    "σ, complex conductivity (A/V)."
    σ
    "μ, FHIP mobility (cm²V⁻¹s⁻¹)."
    μ
end

# Broadcast Polaron data.
function Base.show(io::IO, x::NewPolaron)
    flush(stdout)
    print(io, "-------------------------------------------------\n Polaron Information: \n-------------------------------------------------\n", "Fröhlich coupling | α = ", round.(x.α, digits = 3), " | sum(α) = ", round(sum(x.α), digits = 3),"\nTemperatures | T = ", round.(x.T, digits = 3), " K \nReduced thermodynamic | β = ", round.(x.β, digits = 3), "\nPhonon frequencies | ω = ", round.(x.ω, digits = 3), " 2π THz\nVariational parameters | v = ", round.(x.v, digits = 3), " ω | w = ", round.(x.w, digits = 3), " ω\nFictitious spring constant | κ = ", round.(x.κ, digits = 3), " kg/s²\nFictitious mass | M = ", round.(x.M, digits = 3), " kg\nFree energy | F = ", round.(x.F, digits = 3), " meV\nElectric field frequency | Ω = ", round.(Float64.(x.Ω), digits = 3),  " 2π THz\nComplex impedance | Z = ", x.Z .|> y -> round.(ComplexF64.(y), digits = 3), " V/A\nComplex conductivity | σ = ", x.σ .|> y -> round.(ComplexF64.(y), digits = 3), " A/V\nMobility | μ = ", round.(Float64.(x.μ), digits = 3), " cm²/Vs")
end


@testset "Local quantities" begin

    # Test Dynamics
    T = Float64
    dt = 0.01
    model = DoubleIntegratorGame(p=3, d=2)
    x = rand(SVector{model.n,T})
    u = rand(SVector{model.m,T})
    x_dot = RobotDynamics.dynamics(model, x, u)
    x1 = x + dt*x_dot

    z = KnotPoint(x,u,dt)
    @test typeof(z) <: AbstractKnotPoint
    @test norm(dynamics_residual(model, z, x1), 1) < 1e-3
    @test (@ballocated dynamics_residual($model, $z, $x1)) == 0

    N = 10
    k = 5
    probsize = ProblemSize(N,model)
    pdtraj = PrimalDualTraj(probsize, dt)
    @test (@ballocated dynamics_residual($model, $pdtraj, $k)) == 0

    # Test dynamics Jacobian
    function ∇dynamics_easy(∇f, model, z)
        dt = z.dt
        x = state(z)
        u = control(z)
        function local_dyn_x(x)
            z = KnotPoint(x, u, dt)
            return discrete_dynamics(RK2, model, z)
        end
        function local_dyn_u(u)
            z = KnotPoint(x, u, dt)
            return discrete_dynamics(RK2, model, z)
        end
        ∇f[:,1:n] = ForwardDiff.jacobian(local_dyn_x, x)
        ∇f[:,n .+ (1:m)] = ForwardDiff.jacobian(local_dyn_u, u)
        return ∇f
    end

    T = Float64
    dt = 0.2
    model = UnicycleGame(p=3)
    x = rand(SVector{model.n,T})
    u = rand(SVector{model.m,T})
    z = KnotPoint(x,u,dt)

    n = model.n
    m = model.m
    ∇f1 = zeros(MMatrix{n,n+m,T,n*(n+m)})
    ∇f2 = zeros(MMatrix{n,n+m,T,n*(n+m)})
    ∇f3 = zeros(MMatrix{n,n+m,T,n*(n+m)})
    RobotDynamics.discrete_jacobian!(RK2,∇f1,model,z)
    ∇dynamics_easy(∇f2, model, z)
    ∇dynamics!(∇f3, model, z)
    @test norm(∇f1 - ∇f2, 1) < 1e-10
    @test norm(∇f1 - ∇f3, 1) < 1e-10
    @test (@ballocated ∇dynamics!($∇f3, $model, $z)) == 0
    @test (@ballocated ∇dynamics!($∇f3, $model, $pdtraj, $k)) == 0

end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    supportfun(geometry, direction)

Support function of `geometry` for given `direction`.

## References

* Gilbert, E., Johnson, D., Keerthi, S. 1988. [A fast
  Procedure for Computing the Distance Between Complex
  Objects in Three-Dimensional Space]
  (https://ieeexplore.ieee.org/document/2083)
"""
function supportfun(g::Geometry, d::Vec)
  v = vertices(g)
  c = centroid(g)
  i = argmax([(vᵢ - c) ⋅ d for vᵢ in v])
  v[i]
end

function supportfun(b::Ball, d::Vec)
  r = (radius(b) / norm(d)) * d
  center(b) + r
end


using Test
using LinearAlgebra
using SparseArrays
using Distributed

using jInv.Utils
using jInv.Mesh
using jInv.InverseSolve

# build domain and true image
domain = [0.0 1.0 0.0 1.0]
n      = [32,32]
Minv   = getRegularMesh(domain,n)
xc     = getCellCenteredGrid(Minv)
xtrue = sin.(2*pi*xc[:,1]).*sin.(pi*xc[:,2])

# get noisy data
A     = sparse(1.0I,Minv.nc,Minv.nc)
ids   = sort(rand(1:Minv.nc,round(Int64,Minv.nc*.8)))
A     = A[ids,:]
btrue = A*xtrue
bdata = btrue + .1*randn(size(btrue))/norm(btrue,Inf)
Wd    = ones(length(btrue))

# generate misfit param
i1     = (1:round(Int64,size(A,1)/4));
i2     = (round(Int64,size(A,1)/4)+1:round(Int64,2*size(A,1)/4));
i3 	   = (round(Int64,2*size(A,1)/4)+1 : round(Int64,3*size(A,1)/4));
i4 	   = (round(Int64,3*size(A,1)/4)+1 : size(A,1));
pFor1  = LSparam(A[i1,:],[])
pFor2  = LSparam(A[i2,:],[])

sigmaBack = zeros(length(xtrue))
Iact         = sparse(1.0I,Minv.nc,Minv.nc);
gl1          = getGlobalToLocal(1.0,sigmaBack)
gl2          = getGlobalToLocal(Iact,sigmaBack)

bd1          = bdata[i1]
bd2          = bdata[i2]
bd3          = bdata[i3]
bd4          = bdata[i4]
Wd1          = Wd[i1]
Wd2          = Wd[i2]
Wd3          = Wd[i3]
Wd4          = Wd[i4]

pMis =  getMisfitParam(pFor1,Wd1,bd1,SSDFun,fMod,gl1)

pMisRefs    = Array{RemoteChannel}(undef, 2)
pMisRefs[1] = initRemoteChannel(getMisfitParam,workers()[1],pFor1,Wd1,bd1,SSDFun,fMod,gl1)
pMisRefs[2] = initRemoteChannel(getMisfitParam,workers()[min(2,nworkers())],pFor2,Wd2,bd2,SSDFun,fMod,gl2)

pForRefs = Array{RemoteChannel}(undef, 4)
pForRefs[1] = initRemoteChannel(LSparam,workers()[1],A[i3,:],[])
pForRefs[2] = initRemoteChannel(LSparam,workers()[min(2,nworkers())],A[i4,:],[])
pForRefs[3] = initRemoteChannel(LSparam,workers()[min(3,nworkers())],A[i3,:],[])
pForRefs[4] = initRemoteChannel(LSparam,workers()[min(4,nworkers())],A[i4,:],[])

alpha        = 1.
x0           = (sum(xtrue)/length(xtrue))*ones(Minv.nc)

boundsLow    = minimum(xtrue)*ones(Minv.nc)
boundsHigh   = maximum(xtrue)*ones(Minv.nc)
pInv         = getInverseParam(Minv,fMod,diffusionReg,alpha,x0,boundsLow,boundsHigh)
display(pInv)

# Test with local misfit
# evaluate distance
sig,dsig = pInv.modelfun(x0)
Dc,F,dF,d2F,pMis,tMis = computeMisfit(sig,pMis,true)
# get Hessian
Hm = getHessian(sig,pMis,d2F)
# multiply with random direction and compare results
x = randn(size(Hm,2))
t1 = HessMatVec(x,pMis,sig,d2F)
t2 = Hm*x
@test norm(t1-t2,Inf)/norm(t1,Inf) < 1e-12

# Test with single remote ref misfit param
sig,dsig = pInv.modelfun(x0)
Dc,F,dF,d2F,pMisR,tMis = computeMisfit(sig,pMisRefs[1:1],true)
# get Hessian
Hm = getHessian(sig,pMisR[1],d2F[1])
# multiply with random direction and compare results
x = randn(size(Hm,2))
t1 = HessMatVec(x,pMisR,sig,d2F)
t2 = Hm*x
@test norm(t1-t2,Inf)/norm(t1,Inf) < 1e-12

# Test array of remote ref misfit param
# get Hessian
Hm = getHessian(sig,pMisR,d2F)
# multiply with random direction and compare results
x = randn(size(Hm,2))
t1 = HessMatVec(x,pMisR,sig,d2F)
t2 = Hm*x
@test norm(t1-t2,Inf)/norm(t1,Inf) < 1e-12

# Test array of misfit params
pMisA = [pMis;pMis;pMis]
sig,dsig = pInv.modelfun(x0)
Dc,F,dF,d2F,pMisA,tMis = computeMisfit(sig,pMisA,true)

# get Hessian
Hm = getHessian(sig,pMisA,d2F)
# multiply with random direction and compare results
x = randn(size(Hm,2))
t1 = HessMatVec(x,pMisA,sig,d2F)
t2 = Hm*x
@test norm(t1-t2,Inf)/norm(t1,Inf) < 1e-12


"""
Transverse electric (TE) or transverse magnetic (TM) polarisation.
"""
@enum Polarisation TE TM

"""
    Solver(basis::PlaneWaveBasis, epc::Matrix{ComplexF64}, muc::Matrix{ComplexF64})

A plane-wave expansion method solver, where `epc` and `muc` are the convolution
matrices of the permittivity and permeability, respectively, for the given
`basis` of plane waves.
"""
struct Solver{T}
    basis::PlaneWaveBasis
    epc::T
    muc::T
    Kx::T
    Ky::T
end

"""
    Solver(geometry::Geometry, basis::PlaneWaveBasis; GPU::Bool=false)

Approximate the geometry using the given `basis` of plane waves.
"""
function Solver(geometry::Geometry, basis::PlaneWaveBasis; GPU=false)
    epc = convmat(geometry.ep, basis)
    muc = convmat(geometry.mu, basis)
    Kx = diagm(Complex.(basis.kxs))
    Ky = diagm(Complex.(basis.kys))
    if GPU
        epc = CuArray(epc)
        muc = CuArray(muc)
        Kx = CuArray(Kx)
        Ky = CuArray(Ky)
    end
    return Solver(basis, epc, muc, Kx, Ky)
end

"""
    Solver(geometry::Geometry, cutoff::Int)

Approximate the geometry using a basis of plane waves truncated in a circle.

The circle has a diameter of `cutoff` Brillouin zones. Increasing the `cutoff`
will increase the number of plane waves leading to a more accurate solution.
It is assumed that `norm(b1) == norm(b2)`.
"""
function Solver(geometry::Geometry, cutoff::Int; GPU=false)
    basis = PlaneWaveBasis(geometry, cutoff)
    return Solver(geometry, basis, GPU=GPU)
end

"""
    Solver(geometry::Geometry, cutoff_b1::Int, cutoff_b2::Int)

Approximate the geometry using a basis of plane waves truncated in a rhombus.

The rhombus has lengths `cutoff_b1` and `cutoff_b2` in the `b1` and `b2`
directions, respectively.
"""
function Solver(geometry::Geometry, cutoff_b1::Int, cutoff_b2::Int; GPU=false)
    basis = PlaneWaveBasis(geometry, cutoff_b1, cutoff_b2)
    return Solver(geometry, basis, GPU=GPU)
end

"""
    convmat(mat::AbstractMatrix, basis::PlaneWaveBasis)

Generate convolution matrices, see Raymond Rumpf's CEM Lecture #18,
"Maxwell's Equations in Fourier Space", for further reading.
"""
function convmat(mat::AbstractMatrix, basis::PlaneWaveBasis)
    @assert length(basis.ps) == length(basis.qs)
    M = length(basis.ps)
    # Modified from Rumpf CEM Lecture #18
    # "Maxwell's Equations in Fourier Space"
    mat = fftshift(fft(fftshift(mat))) / length(mat)
    p0 = 1 + size(mat,1)÷2
    q0 = 1 + size(mat,2)÷2
    convmat = zeros(ComplexF64, M, M)
    for irow in 1:M, icol in 1:M
        prow, qrow = basis.ps[irow], basis.qs[irow]
        pcol, qcol = basis.ps[icol], basis.qs[icol]
        pfft = prow - pcol
        qfft = qrow - qcol
        convmat[irow,icol] = mat[p0+pfft,q0+qfft]
    end
    return convmat
end

"""
    solve(solver::Solver, k::AbstractVector{<:Real}, polarisation::Polarisation; bands=:)

Calculate the eigenmodes of a photonic crystal at position `k` in reciprocal space.
"""
function solve(solver::Solver, k::AbstractVector{<:Real}, polarisation::Polarisation; bands=:)

    basis = solver.basis
    epc, muc = solver.epc, solver.muc
    Kx = solver.Kx + k[1]*I
    Ky = solver.Ky + k[2]*I

    # Get left and right hand sides of the generalised eigenvalue problem
    if polarisation == TE
        LHS = Kx/epc*Kx + Ky/epc*Ky
        RHS = muc
        label = L"H_z"
    elseif polarisation == TM
        LHS = Kx/muc*Kx + Ky/muc*Ky
        RHS = epc
        label = L"E_z"
    end

    # LHS and RHS may be CUDA (GPU) arrays
    # Convert back to standard Julia arrays to perform eigen on the CPU
    LHS = Array(LHS)
    RHS = Array(RHS)

    # Sometimes the generalised eigenvalue problem solver
    # fails near Γ when the crystals are symmetric.
    # In these cases, rewrite as a regular eigenvalue problem
    freqs_squared, modes_data = try
        eigen(LHS, RHS)
    catch
        eigen(RHS \ LHS)
    end

    freqs = sqrt.(freqs_squared)
    # Sort by increasing frequency
    idx = sortperm(freqs, by=real)
    freqs = freqs[idx][bands]
    modes_data = modes_data[:,idx][:,bands]
    # Eigenmodes are weighted by the RHS of the generalised eigenvalue problem
    weighting = RHS
    modes = Mode[]
    for i in 1:length(freqs)
        mode = Mode(k, freqs[i], modes_data[:,i], weighting, basis, label)
        push!(modes, mode)
    end
    return modes
end

"""
    solve(solver::Solver, x::BrillouinZoneCoordinate, polarisation::Polarisation; bands=:)

Calculate the eigenmodes of a photonic crystal at position `k=x.p*b1 + x.q*b2` in reciprocal space.
"""
function solve(solver::Solver, x::BrillouinZoneCoordinate, polarisation::Polarisation; bands=:)
    k = get_k(x, solver.basis)
    return solve(solver, k, polarisation, bands=bands)
end


using StatsBase, Graphs, SimpleWeightedGraphs
include("sampling.jl")

"""
"""
function random_walks(G, k::Integer, num_walks::Integer, nodes, node_probabilities,
        bias_nodes::Bool=false, weighted::Bool=false)
    walks = Vector{Tuple}(undef, num_walks)
    Threads.@threads for i in range(1, num_walks)
        # Initializ walk to garbage value (always enters loop)
        walk = [0]
	while length(walk) < k+1
            start_node = sample_start_node(nodes, node_probabilities, bias_nodes)
            # Start node must have positive out degree
            while sum(G.weights[start_node,:]) == 0
                start_node = sample_start_node(nodes, node_probabilities, bias_nodes)
	    end
            if !weighted
                walk = Tuple(randomwalk(G, start_node, k+1))
            else
                walk = weighted_rw(G, start_node, k+1)
            end
        end
        walks[i] = walk
    end
    return walks
end

"""
Compute out weight basd edge probabilitis.
"""
function edge_probabilities(G, node, nbrs)
    edge_probs = convert(Vector{Float64}, G.weights[node,nbrs])
    edge_probs ./= sum(edge_probs)
    return ProbabilityWeights(edge_probs)
end

"""
Computed a weighted random walk of length
len through G from start_node.

WARNING: Assume start_node has at least 1 len-edge
walk originating from it
"""
function weighted_rw(G, start_node, len)
    walk = Vector{Int64}()
    push!(walk, start_node)
    curr_node = start_node
    max_attempts = 5
    attempts = 0
    while length(walk) < len
        nbrs = findall(>(0), G.weights[curr_node,:])
        edge_probs = edge_probabilities(G, curr_node, nbrs)
        nxt_node = sample(nbrs, edge_probs, 1)[1]
        tries = 0
        mtries = length(nbrs)*2
        while sum(G.weights[nxt_node,:]) == 0 && tries < mtries
            nxt_node = sample(nbrs, edge_probs, 1)[1]
            tries += 1
        end

        if tries < mtries
            push!(walk, nxt_node)
            curr_node = nxt_node
        else
            # start over if we get stuck
            walk = Vector{Int64}()
            push!(walk, start_node)
            curr_node = start_node
            attempts += 1
            if attempts > max_attempts
                break
            end
        end
    end
    return Tuple(walk)
end


using ITensors,
  ITensors.NDTensors,
  ITensorGPU,
  LinearAlgebra, # For tr()
  Combinatorics, # For permutations()
  CUDA,
  Test

@testset "cuITensor $T Contractions" for T in (Float64, ComplexF64)
  mi, mj, mk, ml, ma = 2, 3, 4, 5, 6, 7
  i = Index(mi, "i")
  j = Index(mj, "j")
  k = Index(mk, "k")
  l = Index(ml, "l")
  a = Index(ma, "a")
  @testset "Test contract cuITensors" begin
    Aij = cuITensor(randomITensor(T, i, j))
    Aji = cuITensor(randomITensor(T, j, i))
    Bij = cuITensor(randomITensor(T, i, j))
    Aik = cuITensor(randomITensor(T, i, k))
    Ajk = cuITensor(randomITensor(T, j, k))
    Ajl = cuITensor(randomITensor(T, j, l))
    Akl = cuITensor(randomITensor(T, k, l))
    Dv = rand(T, mi)
    D = itensor(ITensors.tensor(NDTensors.Diag(CuVector(Dv)), IndexSet(i, i')))
    Ev = rand(T, mi)
    E = itensor(ITensors.tensor(NDTensors.Diag(CuVector(Ev)), IndexSet(i, i'')))
    @testset "Test contract cuITensors (Matrix*Diag -> Matrix)" begin
      C = Aij * D
      @test collect(CuArray(C)) ≈ collect(CuMatrix(Aij, j, i)) * diagm(0 => Dv)
    end
    @testset "Test contract cuDiagITensors (Diag*Diag -> Diag)" begin
      C = E * D
      cC = CuArray(C)
      @test collect(cC) ≈ diagm(0 => Ev) * diagm(0 => Dv)
    end
    @testset "Test contract cuDiagITensors (UniformDiag*Diag -> Diag)" begin
      scal = itensor(ITensors.tensor(NDTensors.Diag(2.0), IndexSet(i, i'')))
      C = scal * D
      @test collect(CuArray(C)) ≈ 2.0 .* diagm(0 => Dv)
      C = D * scal
      @test collect(CuArray(C)) ≈ 2.0 .* diagm(0 => Dv)
    end
    @testset "Test contract cuITensors (Matrix*UniformDiag -> Matrix)" begin
      scal = itensor(ITensors.tensor(NDTensors.Diag(T(2.0)), IndexSet(i, i')))
      C = scal * Aij
      @test cpu(C) ≈ 2.0 * cpu(replaceind(Aij, i, i')) atol = 1e-4
      C = Aij * scal
      @test_broken cpu(C) ≈ 2.0 * cpu(replaceind(permute(Aij, j, i), i, i')) atol = 1e-4
    end
  end # End contraction testset
end

@testset "cuITensor $T1, $T2 Contractions" for T1 in (Float64, ComplexF64),
  T2 in (Float64, ComplexF64)

  mi, mj, mk, ml, ma = 2, 3, 4, 5, 6, 7
  i = Index(mi, "i")
  j = Index(mj, "j")
  k = Index(mk, "k")
  l = Index(ml, "l")
  a = Index(ma, "a")
  @testset "Test contract cuITensors" begin
    Aij = cuITensor(randomITensor(T1, i, j))
    Aji = cuITensor(randomITensor(T1, j, i))
    Bij = cuITensor(randomITensor(T1, i, j))
    Dv = rand(T2, mi)
    D = itensor(ITensors.tensor(NDTensors.Diag(CuVector(Dv)), IndexSet(i, i')))
    Ev = rand(T2, mi)
    E = itensor(ITensors.tensor(NDTensors.Diag(CuVector(Ev)), IndexSet(i, i'')))
    @testset "Test contract cuITensors (Matrix*Diag -> Matrix)" begin
      C = Aij * D
      cC = CuArray(C)
      @test collect(cC) ≈ collect(CuMatrix(Aij, j, i)) * diagm(0 => Dv)
    end
    @testset "Test contract cuDiagITensors (Diag*Diag -> Diag)" begin
      C = E * D
      cC = CuArray(C)
      @test collect(cC) ≈ diagm(0 => Ev) * diagm(0 => Dv)
    end
    @testset "Test contract cuDiagITensors (UniformDiag*Diag -> Diag)" begin
      scal = itensor(ITensors.tensor(NDTensors.Diag(T2(2.0)), IndexSet(i, i'')))
      C = scal * D
      cC = CuArray(C)
      @test collect(cC) ≈ 2.0 .* diagm(0 => Dv)
      C = D * scal
      cC = CuArray(C)
      @test collect(cC) ≈ 2.0 .* diagm(0 => Dv)
    end
    @testset "Test contract cuITensors (Matrix*UniformDiag -> Matrix)" begin
      scal = itensor(ITensors.tensor(NDTensors.Diag(T2(2.0)), IndexSet(i, i')))
      C = scal * Aij
      cC = CuArray(C)
      @test collect(cC) ≈ array(2.0 * cpu(replaceind(Aij, i, i'))) atol = 1e-4
      C = Aij * scal
      cC = CuArray(C)
      @test_broken collect(cC) ≈ array(2.0 * cpu(replaceind(permute(Aij, j, i), i, i'))) atol =
        1e-4
    end
  end # End contraction testset
end


using NetworkEpidemics
using LightGraphs
using Plots

N = 1000
g = complete_graph(N)
g = erdos_renyi(N, 0.3)
g = random_regular_graph(N,15)
M = 10
h = complete_graph(M)

β = 0.1
D = 0.1

# SI

mp = Metapopulation(h, D, SI(β))
mpx = Metaplex(g, h, D, SI(β))

x0_mp = Array{Int,2}(undef, M,2)
x0_mp[2:M,:] .= [100 0]
x0_mp[1, :] .= [90, 10]

x0_i = fill(1, N)
x0_i[1:10] .= 2
x0_μ = [div(i-1,div(N,M))+1 for i in 1:N]

tmax = 100.0
nmax = 12000

@time ts_mp, output_mp = gillespie(mp, x0_mp, nmax=nmax)
@time ts_av_mp, u_av_mp = average(mp, x0_mp, nmax=nmax, tmax=tmax, progressbar=true)
@time ts_mf_mp, u_mf_mp  = meanfield(mp, x0_mp, tmax=tmax, saveat=ts_av_mp)

@time ts_mpx, output_mpx = gillespie(mpx, [x0_i,x0_μ], nmax=nmax)
@time ts_av_mpx, u_av_mpx = average(mpx, [x0_i,x0_μ], nmax=nmax, tmax=tmax, progressbar=true)
@time ts_mf_mpx, u_mf_mpx = meanfield(mpx, [x0_i, x0_μ], tmax=tmax, saveat=ts_av_mpx)

plot(ts_mp, output_mp[2],
    line=:steppre,
    label="",
    title="Metapopulation")
plot(ts_mpx, output_mpx[2],
    line=:steppre,
    label="",
    title="Metaplex")

plot(ts_av_mp, u_av_mp[2],
    label="",
    linestyle=:dash,
    title="Metapopulation")
plot!(ts_mf_mp, u_mf_mp[2], label="")

plot(ts_av_mpx, u_av_mpx[2],
    label="",
    linestyle=:dash,
    title="Metaplex")
plot!(ts_mf_mpx, u_mf_mpx[2], label="")

plot(ts_av_mp, sum(u_av_mp[2], dims=2),
    label="Average (Metapopulation)",
    linestyle=:dash,
    legend=:bottomright,
    title="Metapopulation vs Metaplex (SI) - RRG(k=15)"
    )
plot!(ts_mf_mp, sum(u_mf_mp[2], dims=2),
    label="Meanfield (Metapopulation)")
plot!(ts_av_mpx, sum(u_av_mpx[2], dims=2),
    label="Average (Metaplex)",
    linestyle=:dash,
    )
plot!(ts_mf_mpx, sum(u_mf_mpx[2], dims=2),
    label="Meanfield (Metaplex)")

# SIS
β = 0.1
γ = 1.0
D = 0.1

mp = Metapopulation(h, D, SIS(β,γ))
mpx = Metaplex(g, h, D, SIS(β,γ))
c_p = ContactProcess(g, SIS(β,γ))

x0_mp = Array{Int,2}(undef, M,2)
x0_mp[2:M,:] .= [100 0]
x0_mp[1, :] .= [90, 10]

x0_i = fill(1, N)
x0_i[1:10] .= 2
x0_μ = [div(i-1,div(N,M))+1 for i in 1:N]

x0_cp = x0_i .== 2

tmax = 10.0
nmax = 8000

@time ts_mp, output_mp = gillespie(mp, x0_mp, nmax=nmax)
@time ts_av_mp, u_av_mp = average(mp, x0_mp, nmax=nmax, tmax=tmax, progressbar=true)
@time ts_mf_mp, u_mf_mp  = meanfield(mp, x0_mp, tmax=tmax, saveat=ts_av_mp)

@time ts_mpx, output_mpx = gillespie(mpx, [x0_i,x0_μ], tmax=tmax, nmax=nmax)
@time ts_av_mpx, u_av_mpx = average(mpx, [x0_i,x0_μ], nmax=nmax, tmax=tmax, progressbar=true)
@time ts_mf_mpx, u_mf_mpx = meanfield(mpx, [x0_i, x0_μ], tmax=tmax, saveat=ts_av_mpx)

@time ts_cp, output_cp = gillespie(c_p, x0_cp, nmax=nmax)
@time ts_av_cp, u_av_cp = average(c_p, x0_cp, nmax=nmax, tmax=tmax, progressbar=true)
@time ts_mf_cp, u_mf_cp = meanfield(c_p, x0_cp, tmax=tmax, saveat=ts_av_cp)

plot(ts_mp, output_mp[2],
    line=:steppre,
    label="",
    title="Metapopulation")
plot(ts_mpx, output_mpx[2],
    line=:steppre,
    label="",
    title="Metaplex")
plot(ts_cp, output_cp,
    line=:steppre,
    label="",
    title="Contact Process")

plot(ts_av_mp, u_av_mp[2],
    label="",
    linestyle=:dash,
    title="Metapopulation")
plot!(ts_mf_mp, u_mf_mp[2], label="")

plot(ts_av_mpx, u_av_mpx[2],
    label="",
    linestyle=:dash,
    title="Metaplex")
plot!(ts_mf_mpx, u_mf_mpx[2], label="")

plot(ts_av_mp, sum(u_av_mp[2], dims=2)/N,
    label="Average (Metapopulation)",
    linestyle=:dash,
    legend=:bottomright,
    title="Metapopulation vs Metaplex (SIS) - RRG(k=15)",
    xlabel="Time",
    ylabel="Fraction of infected individuals"
    )
plot!(ts_mf_mp, sum(u_mf_mp[2], dims=2)/N,
    label="Meanfield (Metapopulation)")
plot!(ts_av_mpx, sum(u_av_mpx[2], dims=2)/N,
    label="Average (Metaplex)",
    linestyle=:dash,
    )
plot!(ts_mf_mpx, sum(u_mf_mpx[2], dims=2)/N,
    label="Meanfield (Metaplex)")
plot!(ts_mf_cp, sum(u_mf_cp, dims=2)/N,
    label="Meanfield (Contact Process)")

# SIR
β = 0.1
γ = 1.0
D = 0.1

mp = Metapopulation(h, D, SIR(β,γ))
mpx = Metaplex(g, h, D, SIR(β,γ))

x0_mp = Array{Int,2}(undef, M, 3)
x0_mp[2:M,:] .= [100 0 0]
x0_mp[1, :] .= [90, 10, 0]

x0_i = fill(1, N)
x0_i[1:10] .= 2
x0_μ = [div(i-1,div(N,M))+1 for i in 1:N]

tmax = 18.0
nmax = 5000

@time ts_mp, output_mp = gillespie(mp, x0_mp, nmax=nmax)
@time ts_av_mp, u_av_mp = average(mp, x0_mp, nmax=nmax, tmax=tmax, progressbar=true)
@time ts_mf_mp, u_mf_mp  = meanfield(mp, x0_mp, tmax=tmax, saveat=ts_av_mp)

@time ts_mpx, output_mpx = gillespie(mpx, [x0_i,x0_μ], tmax=tmax, nmax=nmax)
@time ts_av_mpx, u_av_mpx = average(mpx, [x0_i,x0_μ], nmax=nmax, tmax=tmax, progressbar=true)
@time ts_mf_mpx, u_mf_mpx = meanfield(mpx, [x0_i, x0_μ], tmax=tmax, saveat=ts_av_mpx)

plot(ts_mp, output_mp[2],
    line=:steppre,
    label="",
    title="Metapopulation")
plot(ts_mpx, output_mpx[2],
    line=:steppre,
    label="",
    title="Metaplex")

plot(ts_av_mp, u_av_mp[2],
    label="",
    linestyle=:dash,
    title="Metapopulation")
plot!(ts_mf_mp, u_mf_mp[2], label="")

plot(ts_av_mpx, u_av_mpx[2],
    label="",
    linestyle=:dash,
    title="Metaplex")
plot!(ts_mf_mpx, u_mf_mpx[2], label="")

plot(ts_av_mp, sum(u_av_mp[2], dims=2)/N,
    label="Average (Metapopulation)",
    linestyle=:dash,
    legend=:topright,
    title="Metapopulation vs Metaplex (SIR) - RRG(k=15)",
    xlabel="Time",
    ylabel="Fraction of infected individuals"
    )
plot!(ts_mf_mp, sum(u_mf_mp[2], dims=2)/N,
    label="Meanfield (Metapopulation)")
plot!(ts_av_mpx, sum(u_av_mpx[2], dims=2)/N,
    label="Average (Metaplex)",
    linestyle=:dash,
    )
plot!(ts_mf_mpx, sum(u_mf_mpx[2], dims=2)/N,
    label="Meanfield (Metaplex)")

plot(ts_av_mp, sum(u_av_mp[3], dims=2)/N,
    label="Average (Metapopulation)",
    linestyle=:dash,
    legend=:bottomright,
    title="Metapopulation vs Metaplex (SIR) - RRG(k=15)",
    xlabel="Time",
    ylabel="Fraction of Recovered individuals"
    )
plot!(ts_mf_mp, sum(u_mf_mp[3], dims=2)/N,
    label="Meanfield (Metapopulation)")
plot!(ts_av_mpx, sum(u_av_mpx[3], dims=2)/N,
    label="Average (Metaplex)",
    linestyle=:dash,
    )
plot!(ts_mf_mpx, sum(u_mf_mpx[3], dims=2)/N,
    label="Meanfield (Metaplex)")


module Day08

using LinearAlgebra: dot

export day08

const TESTINPUT = """be cfbegad cbdgef fgaecd cgeb fdcge agebfd fecdb fabcd edb | fdgacbe cefdb cefbgd gcbe
edbfga begcd cbg gc gcadebf fbgde acbgfd abcde gfcbed gfec | fcgedb cgb dgebacf gc
fgaebd cg bdaec gdafb agbcfd gdcbef bgcad gfac gcb cdgabef | cg cg fdcagb cbg
fbegcd cbd adcefb dageb afcb bc aefdc ecdab fgdeca fcdbega | efabcd cedba gadfec cb
aecbfdg fbg gf bafeg dbefa fcge gcbea fcaegb dgceab fcbdga | gecf egdcabf bgf bfgea
fgeab ca afcebg bdacfeg cfaedg gcfdb baec bfadeg bafgc acf | gebdcfa ecba ca fadegcb
dbcfg fgd bdegcaf fgec aegbdf ecdfab fbedc dacgb gdcebf gf | cefg dcbef fcge gbcadfe
bdfegc cbegaf gecbf dfcage bdacg ed bedf ced adcbefg gebcd | ed bcgafe cdgba cbgef
egadfb cdbfeg cegd fecab cgb gbdefca cg fgcdab egfdb bfceg | gbdfcae bgc cg cgb
gcafb gcf dcaebfg ecagb gf abcdeg gaef cafbge fdbac fegbdc | fgae cfgab fg bagce"""

function parseLine(line)
    combinations, output = split.(split(line, "|"))
    keys = _solve(Set.(combinations))
    return getindex.(Ref(keys),Set.(output))
end

function _solve(combinations)
    find(pred, list) = list[findfirst(pred, list)]
    haslength(L) = x -> length(x) == L
    subsetof(s1) = s2 -> issubset(s2, s1)
    contains(s1) = s2 -> issubset(s1, s2)
    unequalto(s1) = s2 -> s1 !== s2 # can check identity in this case
    and(preds...) = x -> all(pred(x) for pred in preds)

    one = find(haslength(2), combinations)
    seven = find(haslength(3), combinations)
    four = find(haslength(4), combinations)
    eight = find(haslength(7), combinations)
    three = find(and(haslength(5), contains(one)), combinations)
    nine = find(and(haslength(6), contains(three)), combinations)
    five = find(and(haslength(5), subsetof(nine), unequalto(three)), combinations)
    two = find(and(haslength(5), unequalto(five), unequalto(three)), combinations)
    six = find(and(haslength(6), contains(five), unequalto(nine)), combinations)
    zero = find(and(haslength(6), unequalto(six), unequalto(nine)), combinations)
    return Dict([zero=>0, one=>1, two=>2, three=>3, four=>4, five=>5, six=>6, seven=>7, eight=>8, nine=>9])
end

function readInput(inputfile="inputs/input08.txt")
    linedata = [parseLine(line) for line in eachline(inputfile)]
    return linedata # transpose
end

function day08(inputdata = readInput("inputs/input08.txt"))
    return (day08Part1(inputdata), day08Part2(inputdata))
end

function day08Part1(outputs)
    # count 1,4,7,8
    count = 0
    for line in outputs
        for out in line
            if out==1 || out==4 || out==7 || out==8
                count += 1
            end
        end
    end
    return count
end

function day08Part2(outputs)
    magnitudes = collect(10 .^ (3:-1:0))
    count = 0
    for line in outputs
        count += dot(line, magnitudes)
    end
    return count
end

end #module


# f(x)=sf*tanh(x)

sf=2.5

Fstanh(x::Array{Float64,2})=(sf*tanh(x),[]);
Fstanh_inplace(value,aux,x::Array{Float64,2})=copy!(value,sf*tanh(x))

Dstanh(derivativeIDX,f_c,faux_c,grad_c,grad_n,x)=axpy!(sf,grad_c.*(1.-(f_c/sf).^2),grad_n)

if PROC=="GPU" ## TODO
    function Fstanh(x::CudaArray)
        tmp=copy(x)
        stanh!(sf,x,tmp)
        return (tmp,[]) # problem here is that I don't know how to free tmp -- this will cause memory leak I think.
    end
    Fstanh_inplace(value::CudaArray,aux,x::CudaArray)=stanh!(sf,x,value)

    Dstanh(derivativeIDX,f_c,faux_c,grad_c::CudaArray,grad_n::CudaArray,x::CudaArray)=Dstanh!(sf,grad_c,f_c,grad_n)
end

Derivative[Fstanh]=Dstanh
Inplace[Fstanh]=Fstanh_inplace

stanh(A::ADnode)=ADnode(Fstanh,A)
export stanh


module TestPkg

# __precompile__(false)

using TaylorIntegration
# using BenchmarkTools
using InteractiveUtils

order = 25
N = 200
x0 = 10randn(2N)
μ = 1e-7rand(N)
x0[N] = - sum(x0[1:N-1].*μ[1:N-1])/μ[N]
x0[2N] = - sum(x0[N+1:2N-1].*μ[1:N-1])/μ[N]

@taylorize xdot1(x, p, t) = b1-x^2

@taylorize function f1!(dq, q, params, t)
    local N = Int(length(q)/2)
    local _eltype_q_ = eltype(q)
    local μ = params
    X = Array{_eltype_q_}(undef, N, N)
    accX = Array{_eltype_q_}(undef, N) #acceleration
    for j in 1:N
        accX[j] = zero(q[1])
        dq[j] = q[N+j]
    end
    #compute accelerations
    for j in 1:N
        for i in 1:N
            if i == j
            else
                X[i,j] = q[i]-q[j]
                temp_001 = accX[j] + (μ[i]*X[i,j])
                accX[j] = temp_001
            end #if i != j
        end #for, i
    end #for, j
    for i in 1:N
        dq[N+i] = accX[i]
    end
    nothing
end

function f2!(dq, q, params, t)
    local N = Int(length(q)/2)
    local _eltype_q_ = eltype(q)
    local μ = params
    X = Array{_eltype_q_}(undef, N, N)
    accX = Array{_eltype_q_}(undef, N) #acceleration
    for j in 1:N
        accX[j] = zero(q[1])
        dq[j] = q[N+j]
    end
    #compute accelerations
    for j in 1:N
        for i in 1:N
            if i == j
            else
                X[i,j] = q[i]-q[j]
                temp_001 = accX[j] + (μ[i]*X[i,j])
                accX[j] = temp_001
            end #if i != j
        end #for, i
    end #for, j
    for i in 1:N
        dq[N+i] = accX[i]
    end
    nothing
end

function TaylorIntegration.jetcoeffs!(::Val{f2!}, t::Taylor1{_T}, q::AbstractVector{Taylor1{_S}}, dq::AbstractVector{Taylor1{_S}}, params) where {_T <: Real, _S <: Number}
    order = t.order
    local N = Int(length(q) / 2)
    local _eltype_q_ = eltype(q)
    local μ = params
    X = Array{_eltype_q_}(undef, N, N)
    accX = Array{_eltype_q_}(undef, N)
    for j = 1:N
        accX[j] = Taylor1(zero(constant_term(q[1])), order)
        dq[j] = Taylor1(identity(constant_term(q[N + j])), order)
    end
    tmp337 = Array{Taylor1{_S}}(undef, size(X))
    tmp337 .= Taylor1(zero(_S), order)
    temp_001 = Array{Taylor1{_S}}(undef, size(tmp337))
    temp_001 .= Taylor1(zero(_S), order)
    for j = 1:N
        for i = 1:N
            if i == j
            else
                X[i, j] = Taylor1(constant_term(q[i]) - constant_term(q[j]), order)
                tmp337[i, j] = Taylor1(constant_term(μ[i]) * constant_term(X[i, j]), order)
                temp_001[i, j] = Taylor1(constant_term(accX[j]) + constant_term(tmp337[i, j]), order)
                accX[j] = Taylor1(identity(constant_term(temp_001[i, j])), order)
            end
        end
    end
    for i = 1:N
        dq[N + i] = Taylor1(identity(constant_term(accX[i])), order)
    end
    for __idx = eachindex(q)
        (q[__idx]).coeffs[2] = (dq[__idx]).coeffs[1]
    end
    for ord = 1:order - 1
        ordnext = ord + 1
        for j = 1:N
            TaylorSeries.zero!(accX[j], q[1], ord)
            TaylorSeries.identity!(dq[j], q[N + j], ord)
        end
        for j = 1:N
            for i = 1:N
                if i == j
                else
                    TaylorSeries.subst!(X[i, j], q[i], q[j], ord)
                    TaylorSeries.mul!(tmp337[i, j], μ[i], X[i, j], ord)
                    TaylorSeries.add!(temp_001[i, j], accX[j], tmp337[i, j], ord)
                    TaylorSeries.identity!(accX[j], temp_001[i, j], ord)
                end
            end
        end
        for i = 1:N
            TaylorSeries.identity!(dq[N + i], accX[i], ord)
        end
        for __idx = eachindex(q)
            (q[__idx]).coeffs[ordnext + 1] = (dq[__idx]).coeffs[ordnext] / ordnext
        end
    end
    return nothing
end

ex1 = :(function f3!(dq, q, params, t)
    local N = Int(length(q)/2)
    local _eltype_q_ = eltype(q)
    local μ = params
    X = Array{_eltype_q_}(undef, N, N)
    accX = Array{_eltype_q_}(undef, N) #acceleration
    for j in 1:N
        accX[j] = zero(q[1])
        dq[j] = q[N+j]
    end
    #compute accelerations
    for j in 1:N
        for i in 1:N
            if i == j
            else
                X[i,j] = q[i]-q[j]
                temp_001 = accX[j] + (μ[i]*X[i,j])
                accX[j] = temp_001
            end #if i != j
        end #for, i
    end #for, j
    for i in 1:N
        dq[N+i] = accX[i]
    end
    nothing
end)

ex2 = :(xdot2(x, p, t) = b1-x^2)

ex3 = :(function harm_osc!(dx, x, p, t)
    local ω = p[1]
    local ω2 = ω^2
    dx[1] = x[2]
    dx[2] = - (ω2 * x[1])
    return nothing
end)

nex1, narr1 = TaylorIntegration._make_parsed_jetcoeffs(ex1)
nex2, narr2 = TaylorIntegration._make_parsed_jetcoeffs(ex2)
nex3, narr3 = TaylorIntegration._make_parsed_jetcoeffs(ex3)

greet(f, parse_eqs) = begin
    t = Taylor1(order)
    x = Taylor1.(x0, t.order)
    dx = similar(x)

    @show @which TaylorIntegration.__jetcoeffs!(Val(parse_eqs), f, t, x, dx, similar(x), μ)
    @show @which TaylorIntegration.jetcoeffs!(Val(f), t, x, dx, μ)
    @show methods(TaylorIntegration.jetcoeffs!)
    # @btime TaylorIntegration.__jetcoeffs!(Val($parse_eqs), $f, $t, $x, $dx, similar($x), μ)
end

end # module


save_results = false

@testset "Quadrotor" begin
# load expected results
res = load(joinpath(@__DIR__,"quadrotor.jld2"))

# Set up problem and solver
prob, opts = Problems.Quadrotor()
times = TO.get_times(prob)
u0 = zeros(prob.model)[2]
U = [u0 + 0.1*SA[sin(t), sin(t), -sin(t), -sin(t)] for t in times]
initial_controls!(prob, U)
rollout!(prob)
states(prob)[end]

solver = ALTROSolver(prob, opts, save_S=true)
n,m,N = size(solver)
ilqr = Altro.get_ilqr(solver)
Altro.initialize!(ilqr)
X = states(ilqr)
U = controls(ilqr)
@test X ≈ res["X"] atol=1e-6
@test U ≈ res["U"] atol=1e-6

# State diff Jacobian
TO.state_diff_jacobian!(ilqr.G, ilqr.model, ilqr.Z)
G = ilqr.G
@test G[1] != zeros(n,n)
@test G[N] != zeros(n,n)
@test G ≈ res["G"] atol=1e-6

# Dynamics Jacobians
TO.dynamics_expansion!(TO.integration(ilqr), ilqr.D, ilqr.model, ilqr.Z)
TO.error_expansion!(ilqr.D, ilqr.model, ilqr.G)
A_ = [Matrix(D.A_) for D in ilqr.D]
B_ = [Matrix(D.B_) for D in ilqr.D]
A = [Matrix(D.A) for D in ilqr.D]
B = [Matrix(D.B) for D in ilqr.D]
@test A_ ≈ res["A_"] atol=1e-6
@test B_ ≈ res["B_"] atol=1e-6
@test A ≈ res["A"] atol=1e-6
@test B ≈ res["B"] atol=1e-6

# Unconstrained Cost Expansion
TO.cost_expansion!(ilqr.quad_obj, ilqr.obj.obj, ilqr.Z, init=true, rezero=true)
TO.error_expansion!(ilqr.E, ilqr.quad_obj, ilqr.model, ilqr.Z, ilqr.G)
hess0 = [[E.Q E.H'; E.H E.R] for E in ilqr.E]
grad0 = [Vector([E.q; E.r]) for E in ilqr.E]
@test hess0 ≈ res["hess0"] atol=1e-6
@test grad0 ≈ res["grad0"] atol=1e-6

# AL Cost Expansion
TO.cost_expansion!(ilqr.quad_obj, ilqr.obj, ilqr.Z, init=true, rezero=true)
TO.error_expansion!(ilqr.E, ilqr.quad_obj, ilqr.model, ilqr.Z, ilqr.G)
hess = [[E.Q E.H'; E.H E.R] for E in ilqr.E]
grad = [Vector([E.q; E.r]) for E in ilqr.E]
@test hess ≈ res["hess"] atol=1e-6
@test grad ≈ res["grad"] atol=1e-6

# Backward Pass
ilqr.ρ[1] = ilqr.opts.bp_reg_min*0
ΔV = Altro.static_backwardpass!(ilqr)
K = Matrix.(ilqr.K)
d = Vector.(ilqr.d)
S = [Matrix(S.Q) for S in ilqr.S]
s = [Vector(S.q) for S in ilqr.S]
Qzz = [[E.Q E.H'; E.H E.R] for E in ilqr.Q]
Qz = [Vector([E.q; E.r]) for E in ilqr.Q]
@test norm(K - res["K"]) < 1e-6
@test norm(d - res["d"]) < 1e-5
@test S ≈ res["S"] atol=1e-6
@test s ≈ res["s"] atol=1e-6
@test Qzz ≈ res["Qzz"] atol=1e-6
@test Qz  ≈ res["Qz"] atol=1e-6
@test ΔV ≈ res["ΔV"] atol=1e-6

if save_results
    @save joinpath(@__DIR__,"quadrotor.jld2") X U G A B A_ B_ hess0 grad0 hess grad K d S s Qzz Qz ΔV
end

end


using ComplexValues
const Spherical = ComplexValues.Spherical
const Polar = ComplexValues.Polar

using Test

@testset "Constructions of Polar" begin
	@test Polar(-5+1im) isa Number
	@test Polar(1.0im) isa Number
	@test Polar(-2) isa Number
	@test Polar(10.0f0) isa Number
	@test Polar(2.0f0,pi) isa Number
	@test Polar(pi,pi) isa Number
end

@testset "Constructions of Spherical" begin
	@test Spherical(-5+1im) isa Number
	@test Spherical(1.0im) isa Number
	@test Spherical(-2) isa Number
	@test Spherical(10.0f0) isa Number
	@test Spherical(2.0f0,pi) isa Number
	@test Spherical(pi/2,pi) isa Number
	@test S2coord(Spherical(1)) ≈ [1,0,0]
end

@testset "Conversions between Polar, Spherical" begin
	for z in [-5+1im,1.0im,-2,10f0,Polar(Inf,-pi/5),Spherical(1,.5)]
		@test Polar(Spherical(z)) ≈ z
	end
	for z in [-5+1im,1.0im,-2,10f0]
		@test Polar{Float64}(Spherical{Float64}(z)) ≈ z
	end
end

@testset "Conversions in and out of Complex" begin
	for z in [-5+1im,1.0im,-2,10f0,Polar(Inf,-pi/5),Spherical(1,.5)]
		@test Polar(Complex(z)) ≈ Polar(z)
		@test Complex(Polar(z)) ≈ Complex(z)
		@test Spherical(Complex(z)) ≈ Spherical(z)
		@test Complex(Spherical(z)) ≈ Complex(z)
	end
end

@testset "Zero and infinity for $T" for T in [Polar,Spherical]
    z = T(Inf)
    @test isinf(z)
    @test isinf(abs(z))
    @test iszero(1/z)
    @test iszero(abs(inv(z)))
    @test iszero(Complex(4/z))
    z = T(0)
    @test iszero(z)
    @test isinf(3im/z)
    z = zero(T)
    @test iszero(z)
    @test isinf(3im/z)
end

@testset "Binary functions on $S,$T" for S in [Polar,Spherical], T in [Polar,Spherical]
	u = S(3.0+5.0im)
	v = T(-5+1im)
	@test Complex(u+v) ≈ Complex(u)+Complex(v)
	@test Complex(u*v) ≈ Complex(u)*Complex(v)
	@test Complex(u/v) ≈ Complex(u)/Complex(v)
	@test Complex(u-v) ≈ Complex(u)-Complex(v)
	@test Complex(u^v) ≈ Complex(u)^Complex(v)
end

@testset "Unary functions on $T" for T in [Polar,Spherical]
    u = T(3.0+5.0im)
    @test Complex(u-4) ≈ Complex(u)-4
    @test Complex(4-u) ≈ 4-Complex(u)
    @test Complex(cos(exp(u))) ≈ cos(exp(Complex(u)))
    @test Complex(sqrt(inv(u))) ≈ sqrt(inv(Complex(u)))
end


"""
`InitialGuessODE`: Initial guess for ordinary differential equations

### Fields

* `int`: interpolation structure
* `v`:   vector  field
* `Δt`:  time step
* `s`:   number of extrapolation stages (for initialisation)
"""
mutable struct InitialGuessODE{DT, TT, VT, IT <: Interpolator}
    int::IT
    v::VT
    Δt::TT
    s::Int

    function InitialGuessODE{DT,TT,VT,IT}(interp, v, Δt) where {DT,TT,VT,IT}
        new(interp, v, Δt, get_config(:ig_extrapolation_stages))
    end
end

function InitialGuessODE(interp, equation::ODE{DT,TT,VT}, Δt::TT) where {DT,TT,VT}
    int = interp(zero(DT), one(DT), Δt, ndims(equation))
    InitialGuessODE{DT, TT, VT, interp}(int, equation.v, Δt)
end

function InitialGuessODE(interp, equation::IODE{DT,TT,ΘT,FT,GT,HT,VT}, Δt::TT) where {DT,TT,ΘT,FT,GT,HT,VT}
    int = interp(zero(DT), one(DT), Δt, ndims(equation))
    InitialGuessODE{DT, TT, VT, interp}(int, equation.v, Δt)
end

function InitialGuessODE(interp, equation::VODE{DT,TT,AT,FT,GT,VT}, Δt::TT) where {DT,TT,AT,FT,GT,VT}
    int = interp(zero(DT), one(DT), Δt, ndims(equation))
    InitialGuessODE{DT, TT, VT, interp}(int, equation.v, Δt)
end

"Initialise initial guess, i.e., given t₀, t₁, q₁ compute q₀, v₀, v₁."
function initialize!(ig::InitialGuessODE{DT,TT,VT,IT},
                t₁::TT,
                q₁::SolutionVector{DT},
                v₁::SolutionVector{DT},
                t₀::TT,
                q₀::SolutionVector{DT},
                v₀::SolutionVector{DT}) where {DT,TT,VT,IT}
    midpoint_extrapolation(ig.v, t₁, t₀, q₁, q₀, ig.s)
    ig.v(t₀, q₀, v₀)
    ig.v(t₁, q₁, v₁)
end

"compute vector field of new solution"
function update!(ig::InitialGuessODE{DT,TT}, t₁::TT, q₁::SolutionVector{DT}, v₁::Vector{DT}) where {DT,TT}
    ig.v(t₁, q₁, v₁)
end

function CommonFunctions.evaluate!(ig::InitialGuessODE{DT,TT},
                q₀::SolutionVector{DT},
                v₀::SolutionVector{DT},
                q₁::SolutionVector{DT},
                v₁::SolutionVector{DT},
                guess::SolutionVector{DT},
                c::TT=one(TT)) where {DT,TT}

    if q₀ == q₁
        @warn "q₀ and q₁ in initial guess are identical! Setting q=q₁."
        guess .= q₁
    else
        evaluate!(ig.int, q₀, q₁, v₀, v₁, one(TT)+c, guess)
    end
end

function CommonFunctions.evaluate!(ig::InitialGuessODE{DT,TT},
                q₀::SolutionVector{DT},
                v₀::SolutionVector{DT},
                q₁::SolutionVector{DT},
                v₁::SolutionVector{DT},
                guess_q::SolutionVector{DT},
                guess_v::SolutionVector{DT},
                c::TT=one(TT)) where {DT,TT}
                @assert length(guess_q) == length(guess_v)

    if q₀ == q₁
        @warn "q₀ and q₁ in initial guess are identical! Setting q=q₁ and v=0."
        guess_q .= q₁
        guess_v .= 0
    else
        evaluate!(ig.int, q₀, q₁, v₀, v₁, one(TT)+c, guess_q, guess_v)
    end
end


using Test, Mill, Flux
using Mill: p_map, inv_p_map, r_map, inv_r_map
import Mill: bagnorm

@testset "basic aggregation functionality" begin
    W = [1, 1/2, 1/2, 1/8, 1/3, 13/24]
    X = Matrix{Float64}(reshape(1:12, 2, 6))
    bags = BAGS[1]
    @test SegmentedMean(2)(X, bags) ≈ [1.0 4.0 9.0; 2.0 5.0 10.0]
    @test SegmentedSum(2)(X, bags) ≈ length.(bags)' .* SegmentedMean(2)(X, bags)
    @test SegmentedMax(2)(X, bags) ≈ [1.0 5.0 11.0; 2.0 6.0 12.0]
    @test SegmentedMeanMax(2)(X, bags) ≈ cat(SegmentedMean(2)(X, bags), SegmentedMax(2)(X, bags), dims=1)
    @test SegmentedMean(2)(X, bags, W) ≈ [1.0 4.0 236/24; 2.0 5.0 260/24]
    @test SegmentedSum(2)(X, bags, W) ≈ [bagnorm(W, b) for b in bags]' .* SegmentedMean(2)(X, bags, W)
    @test SegmentedMax(2)(X, bags, W) ≈ SegmentedMax(2)(X, bags)
end

@testset "matrix weights" begin
    W = abs.(randn(6))
    W_mat = vcat(W, 2*W)
    X = Matrix{Float64}(reshape(1:12, 2, 6))
    for bags in BAGS
        @test SegmentedSum(2)(X, bags, W_mat) ≈ SegmentedSum(2)(X, bags, W)
        @test SegmentedMean(2)(X, bags, W_mat) ≈ SegmentedMean(2)(X, bags, W)
        @test SegmentedMax(2)(X, bags, W_mat) ≈ SegmentedMax(2)(X, bags, W)
    end
end

@testset "pnorm functionality" begin
    dummy = randn(2)
    for t = 1:10
        a, b, c, d, ρ1, ρ2, c1, c2, w1, w2 = randn(10)
        p1, p2 = p_map(ρ1), p_map(ρ2)
        w1 = abs(w1)
        w2 = abs(w2)
        bags = ScatteredBags([[1,2]])
        agg = SegmentedPNorm([ρ1, ρ2], [c1, c2], dummy)
        @test agg([a b; c d], bags) ≈ [
                                     (1/2*(abs(a-c1)^p1 + abs(b-c1)^p1))^(1/p1);
                                     (1/2*(abs(c-c2)^p2 + abs(d-c2)^p2))^(1/p2)
                                    ]
        @test agg([a b; c d], bags) ≈ [
                                     (1/2*(abs(a-c1)^p1 + abs(b-c1)^p1))^(1/p1);
                                     (1/2*(abs(c-c2)^p2 + abs(d-c2)^p2))^(1/p2)
                                    ]
        @test agg([a b; c d], bags, [w1, w2]) ≈ [
                                               (1/(w1+w2)*(w1*abs(a-c1)^p1 + w2*abs(b-c1)^p1))^(1/p1);
                                               (1/(w1+w2)*(w1*abs(c-c2)^p2 + w2*abs(d-c2)^p2))^(1/p2)
                                              ]
        for bags in BAGS
            X = randn(2, 6)
            agg = SegmentedPNorm(inv_p_map.([1+1e-16, 1+1e-16]), [0, 0], dummy)
            @test agg(X, bags) ≈ SegmentedMean(2)(abs.(X), bags)
            agg = SegmentedPNorm(inv_p_map.([2, 2]), [0, 0], dummy)
            @test agg(X, bags) ≈ hcat([sqrt.(sum(X[:, b] .^ 2, dims=2) ./ length(b)) for b in bags]...)
        end
    end
end

@testset "lse functionality" begin
    dummy = randn(2)
    for t = 1:10
        a, b, c, d, ρ1, ρ2 = randn(6)
        r1, r2 = r_map(ρ1), r_map(ρ2)
        @test SegmentedLSE([ρ1, ρ2], dummy)([a b; c d], ScatteredBags([[1,2]])) ≈ [
                                                                               1/r1*log(1/2*(exp(a*r1)+exp(b*r1)));
                                                                               1/r2*log(1/2*(exp(c*r2)+exp(d*r2)))
                                                                              ]
        for bags in BAGS
            X = randn(2, 6)
            W = abs.(randn(6))
            r1, r2 = randn(2)
            # doesn't use weights
            @test all(SegmentedLSE([ρ1, ρ2], dummy)(X, bags) .== SegmentedLSE([ρ1, ρ2], dummy)(X, bags, W))
            # the bigger value of r, the closer we are to the real maximum
            @test isapprox(SegmentedLSE([100, 100], dummy)(X, bags), SegmentedMax(2)(X, bags), atol=0.1)
        end
    end
end

@testset "pnorm numerical stability" begin
    k, d = 10, 5
    dummy, c = randn(d), zeros(d)
    b = AlignedBags([1:k])
    ρ1 = inv_p_map.(ones(d))
    ρ2 = randn(d)
    p2 = p_map(ρ2)
    Z = 1e5 .+ 1e3 .* randn(d, k)
    W = abs.(randn(k)) .+ 1e-2
    for X in [Z, -Z, randn(d, k)]
        @test SegmentedPNorm(ρ1, c, dummy)(X, b) ≈ sum(abs.(X); dims=2) ./ k
        @test SegmentedPNorm(ρ1, c, dummy)(X, b, W) ≈ sum(W' .* abs.(X); dims=2) / sum(W)
        @test SegmentedPNorm(ρ1, c, dummy)(X, b, W) ≈ sum(W' .* abs.(X); dims=2) / sum(W)
        for i in 1:k
            @test SegmentedPNorm(ρ1, c, dummy)(repeat(X[:, i], 1, k), b) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ2, c, dummy)(repeat(X[:, i], 1, k), b) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ1, c, dummy)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ1, c, dummy)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ2, c, dummy)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
            @test SegmentedPNorm(ρ2, c, dummy)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
        end
    end
end

@testset "lse numerical stability" begin
    k, d = 10, 5
    dummy = randn(d)
    b = AlignedBags([1:k])
    ρ1 = inv_r_map.(1e15 .+ 1e5 .* randn(d))
    ρ2 = inv_r_map.(abs.(1e5 .* randn(d)))
    Z = 1e5 .+ 1e3 .* randn(d, k)
    W = abs.(randn(k)) .+ 1e-2
    for X in [Z, -Z, randn(d, k)]
        @test SegmentedLSE(ρ1, dummy)(X, b) ≈ maximum(X; dims=2)
        # doesn't use weights
        @test SegmentedLSE(ρ1, dummy)(X, b, W) ≈ maximum(X; dims=2)

        # implementation immune to underflow not available yet
        @test_skip @test SegmentedLSE(ρ2, dummy)(X, b) ≈ sum(X; dims=2) ./ k
        # doesn't use weights
        @test_skip @test SegmentedLSE(ρ2, dummy)(X, b, W) ≈ sum(X; dims=2) ./ k

        for i in 1:k
            @test SegmentedLSE(randn(d), dummy)(repeat(X[:, i], 1, k), b) ≈ X[:, i]
            # doesn't use weights
            @test SegmentedLSE(randn(d), dummy)(repeat(X[:, i], 1, k), b, W) ≈ X[:, i]
        end
    end
end

@testset "missing values" begin
    dummy = randn(2)
    C = randn(2)
    for bags in [AlignedBags([0:-1]), AlignedBags([0:-1, 0:-1, 0:-1])]
        @test SegmentedMean(C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedSum(C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedMax(C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedLSE(dummy, C)(missing, bags) == repeat(C, 1, length(bags))
        @test SegmentedPNorm(dummy, dummy, C)(missing, bags) == repeat(C, 1, length(bags))
    end

    # default values C are indeed filled in
    for bags in vcat(BAGS2)
        idcs = isempty.(bags.bags)
        l = maximum(maximum.(bags.bags[.!idcs]))
        X = randn(2, l)
        @test SegmentedMean(C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
        @test SegmentedSum(C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
        @test SegmentedMax(C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
        @test SegmentedLSE(dummy, C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
        @test SegmentedPNorm(dummy, dummy, C)(X, bags)[:, idcs] == repeat(C, 1, sum(idcs))
    end
end


using Bolt
using Test
using DelimitedFiles
using LinearAlgebra

@testset "FFTLog" begin
    N = 64
    μ = 0
    q = 0.0
    r₀ = 1.0
    L = 8.0
    Nhalf = N ÷ 2
    n = range(-Nhalf,Nhalf,length=N)
    r = r₀ .* 10 .^ (n .* L ./ N )
    pl = Bolt.plan_fftlog(r, μ, q, 1.0; kropt=true)
    aₙ = r .^ (μ + 1) .* exp.(-r.^2 / 2)
    y = similar(r, ComplexF64)
    fftdata = readdlm("data/fftlog_example.txt", ' ', Float64, '\n')

    # test forward
    mul!(y, pl, aₙ)
    f_ref = fftdata[:,2]
    @test all(abs.(y .- f_ref) .< 1e-15)
    @test isapprox(y, f_ref)

    # test backward
    y2 = similar(r, ComplexF64)
    ldiv!(y2, pl, y)
    @test all(abs.(y2 .- aₙ) .< 1e-15)

end

# test U with more naive result
# U_μ(μ, x) = 2^x * gamma((μ + 1 + x)/2) / gamma((μ + 1 - x)/2)

## hand-written ℋ derivative for testing
# function ℋ′(x, par::AbstractCosmoParams)
#     a = x2a(x)
#     return -H₀(par) * (2par.Ω_r + (par.Ω_b + par.Ω_m) * a - 2Ω_Λ(par) * a^4) /
#         (2 * a * √(par.Ω_r + (par.Ω_b + par.Ω_m) * a + Ω_Λ(par) * a^4))
# end


export Vandermonde

immutable Vandermonde{T} <: AbstractArray{T, 2}
	c :: Vector{T}
end

#XXX Inefficient but works
getindex(V::Vandermonde, i, j) = getindex(full(V), i, j)
isassigned(V::Vandermonde, i, j) = isassigned(full(V), i, j)

size(V::Vandermonde, r::Int) = (r==1 || r==2) ? length(V.c) :
    throw(ArgumentError("Invalid dimension $r"))
size(V::Vandermonde) = length(V.c), length(V.c)

function full{T}(V::Vandermonde{T})
	n=size(V, 1)
	M=Array(T, n, n)
	for i=1:n
		M[:,i] = V.c.^(i-1)
	end
	M
end


"""
does the same thing as `x[..., :-1]` in pytorch.
"""
function head(x::T) where {T <: AbstractArray}
  return x[[axes(x, i) for i in 1:ndims(x)-1]..., 1:end-1]
end

"""
returns: x[[axes(x, i) for i in 1:ndims(x)-1]..., l]
"""
function slice(x::T, l::Integer) where {T <: AbstractArray}
  return x[[axes(x, i) for i in 1:ndims(x)-1]..., l]
end

"""
returns: x[[axes(x, i) for i in 1:ndims(x)-1]..., end]
"""
function tail(x::T) where {T <: AbstractArray}
  return x[[axes(x, i) for i in 1:ndims(x)-1]..., end]
end

"""
logsumexp
"""
function logsumexp(logx::T; dims::Integer=1) where {T <: AbstractArray}
  d = dims < 0 ? ndims(logx) + 1 + dims : dims
  mx = maximum(logx, dims=d)
  return mx .+ log.(sum(exp.(logx .- mx), dims=d))
end

"""
return: dropdims(logsumexp(logx, dims=dims), dims=dims)
"""
function logsumexpdd(logx::T; dims::Integer=1) where {T <: AbstractArray}
  d = dims < 0 ? ndims(logx) + 1 + dims : dims
  return dropdims(logsumexp(logx, dims=d), dims=d)
end

"""
return: dropdims(sum(x, dims=dims), dims=dims)
"""
function sumdd(x::T; dims::Integer=1) where {T <: AbstractArray}
  d = dims < 0 ? ndims(x) + 1 + dims : dims
  return dropdims(sum(x, dims=d), dims=d)
end

"""
return: dropdims(prod(x, dims=dims), dims=dims)
"""
function proddd(x::T; dims::Integer=1) where {T <: AbstractArray}
  d = dims < 0 ? ndims(x) + 1 + dims : dims
  return dropdims(prod(x, dims=d), dims=d)
end

function lpdf_normal(x::X, m::M, s::S) where {X <: Real, M <: Real, S<:Real}
  z = (x - m) / s
  # return oftype(z, -0.5 * log(2*pi) - z^2 * 0.5 - log(s))
  return -0.5 * log(2*pi) - z^2 * 0.5 - log(s)
end

"""
log pdf of gaussian mixture model
"""
function lpdf_gmm(x::TX, m::TM, s::TS, w::TW; dims::Integer, dropdim::Bool=false) where {TX, TW, TM, TS}
  if dropdim
    return logsumexpdd(log.(w) .+ lpdf_normal.(x, m, s), dims=dims)
  else
    return logsumexp(log.(w) .+ lpdf_normal.(x, m, s), dims=dims)
  end
end

"""
cat(a, b, dims=ndims(A) + 1)
"""
function stack(a::A, b::B) where {A, B}
  @assert size(a) == size(b)
  return cat(a, b, dims=ndims(a) + 1)
end

function lpdf_dirichlet(p::AbstractArray, a::AbstractArray)
  if size(a) != size(p)
    a = reshape(a, (one.(size(p)[1:end-1])..., length(a)))
  end
  K = ndims(p)
  return Tracker.lgamma.(sumdd(a, dims=K)) .- sumdd(Tracker.lgamma.(a), dims=K) .+ sumdd((a .- 1) .* log.(p), dims=K)
end

function lpdf_beta(p::P, a::A, b::B) where {P <: Real, A <: Real, B <: Real}
  if 0 < p < 1
    return Tracker.lgamma(a + b) - Tracker.lgamma(a) - Tracker.lgamma(b) + (a-1)*log(p) + (b-1)*log1p(-p)
  else
    return -Inf
  end
end

function lpdf_gamma(x::X, shape::A, scale::B) where {X <: Real, A <: Real, B <: Real}
  if x <= 0
    return -Inf
  else
    return -(Tracker.lgamma(shape) + shape*log(scale)) + (shape - 1)*log(x) - x / scale
  end
end

function lpdf_lognormal(x::X, m::A, s::B) where {X <: Real, A <: Real, B <: Real}
  if x <= 0
    return -Inf
  else
    lx = log(x)
    return lpdf_normal(lx, m, s) - lx
  end
end

function lpdf_uniform(x::X, a::A, b::B) where {X <: Real, A <: Real, B <: Real}
  @assert b > a
  if a <= x <= b
    return -log(b - a)
  else
    return -Inf
  end
end

compute_lpdf(dist::Dirichlet, x::X) where X = lpdf_dirichlet(x, dist.alpha)
compute_lpdf(dist::Normal, x::X) where X = lpdf_normal.(x, dist.μ, dist.σ)
compute_lpdf(dist::Beta, x::X) where X = lpdf_beta.(x, dist.α, dist.β)
compute_lpdf(dist::Gamma, x::X) where X = lpdf_gamma.(x, dist.α, dist.θ)
compute_lpdf(dist::LogNormal, x::X) where X = lpdf_lognormal.(x, dist.μ, dist.σ)
compute_lpdf(dist::Uniform, x::X) where X = lpdf_uniform.(x, dist.a, dist.b)

function sigmoid_safe(x::T; eps::Float64=1e-8) where {T <: Real}
  return T(clamp(sigmoid(x), eps, 1 - eps))
end

function logit_safe(p::T; eps::Float64=1e-8) where {T <: Real}
  p = T(clamp(p, eps , 1 - eps))
  return log(p) - log1p(-p)
end


@testset "set operations" begin
    @testset "union" begin
        d1 = UnitDisk()
        d2 = (-.9..0.9)^2
        d3 = ProductDomain(-.5 .. -.1, ChebyshevInterval())
        d4 = (0.0..1.5)
        d5 = [1.0,3.0]

        @test convert(Domain{SVector{2,Float64}}, d3) isa Domain{SVector{2,Float64}}

        u1 = UnitDisk() ∪ (-.9..0.9)^2
        u2 = u1 ∪ d3
        @test dimension(u1) == 2
        @test dimension(u2) == 2
        @test boundingbox(u1) == ChebyshevInterval()^2

        u3 = d3 ∪ u1
        u4 = u1 ∪ u2
        @test SVector(0.,.15) ∈ u3
        @test SVector(0.,.15) ∈ u4
        @test SVector(1.1,.75) ∉ u3
        @test SVector(1.1,.75) ∉ u4

        ũ1 = UnionDomain(d1,d2)
        @test u1 == ũ1
        @test UnionDomain{SVector{2,Float64}}(d1,d2) == ũ1
        ũ1 = UnionDomain((d1,d2))
        @test u1 == ũ1
        ũ2 = UnionDomain([d1,d2])
        @test ũ2 == ũ2
        @test u1 == ũ2
        @test UnionDomain{SVector{2,Float64}}(d1) isa UnionDomain

        # Don't create a union with two identical elements
        @test UnitDisk() ∪ UnitDisk() isa UnitDisk

        # union with non-Domain type that implements domain interface
        u45 = (0.0..1.5) ∪ [1.0,3.0]
        @test u45 isa Domain{Float64}
        @test u45 isa UnionDomain
        @test eltype(component(u45,1)) == Float64
        @test eltype(component(u45,2)) == Float64
        @test 0.2 ∈ u45
        @test 1.2 ∈ u45
        @test 3 ∈ u45
        @test -1.2 ∉ u45
        @test convert(Domain{BigFloat}, u45) isa Domain{BigFloat}

        u45b = (0.0..1.5) ∪ [1,3]
        @test u45b isa Domain{Float64}
        @test component(u45b,2) isa AbstractArray{Float64}
        @test [1,3] ∪ (0.0..1.5) isa Domain{Float64}

        @test issubset([0,1], 0..1)
        @test !issubset([0,1,2], 0..1)
        @test issubset(Set([0,1]), 0..1)
        @test !issubset(Set([0,2]), 0..1)

        @test uniondomain() == EmptySpace{Any}()
        @test uniondomain(0..1) == 0..1
        @test uniondomain(0..1, [0,1]) == 0..1
        @test uniondomain([0,1], 0..1) == 0..1

        @test uniondomain([0,1], [0.0,1.0]) == [0,1]
        @test uniondomain([0,2], [0.0,1.0]) isa UnionDomain

        # larger union expressions
        @test uniondomain(0..1, 1..3, Point(0.4), 2..5, FullSpace(), Point(-0.2)) isa FullSpace
        @test uniondomain(0..1, 1..3, Point(0.4)) == 0..3
        @test uniondomain(0..1, Point(0.4), 1..3) == 0..3
        @test uniondomain(Point(0.4), 0..1, 1..3) == 0..3

        # ordering doesn't matter
        @test UnionDomain(d1,d2) == UnionDomain(d2,d1)

        @test UnionDomain((d1,d2)) == UnionDomain(d1,d2)
        @test UnionDomain(d1) isa UnionDomain
        @test UnionDomain(UnionDomain(d1,d2),d3) == UnionDomain(d3,UnionDomain(d1,d2))

        @test convert(Domain, [0..1,2..3,3..4]) isa UnionDomain{Int}
        @test convert(Domain{Float64}, [0..1,2..3,3..4]) isa UnionDomain{Float64}
        @test convert(Domain, Set([0..1,2..3,3..4])) isa UnionDomain{Int}
        @test convert(Domain{Float64}, Set([0..1,2..3,3..4])) isa UnionDomain{Float64}
        @test convert(Domain, Set([1,2,3])) isa UnionDomain{Int}
        @test 2 ∈ convert(Domain, Set([1,2,3]))
        @test 2 ∈ convert(Domain{Float64}, Set([1,2,3]))
        @test 4 ∉ convert(Domain, Set([1,2,3]))

        @test interior(uniondomain(0..1, 2..3)) == uniondomain(OpenInterval(0,1),OpenInterval(2,3))
        @test closure(uniondomain(OpenInterval(0,1),OpenInterval(2,3))) == uniondomain(0..1, 2..3)

        @test uniondomain(0..1, 2..3) \ (0.5..2.5) == uniondomain(0..0.5, 2.5..3.0)
        @test uniondomain(0..1, 2..3) \ uniondomain(0.5..2.5, -2..(-1)) == uniondomain(0..0.5, 2.5..3.0)

        @test !isempty(u1)
        show(io, textmime, u1)
        @test String(take!(io)) == "UnitDisk() ∪ ((-0.9..0.9) × (-0.9..0.9))"

        # repeated union
        @test ncomponents(uniondomain(UnitBall{Float64}(), UnitInterval(), UnitInterval())) == 2
        @test ncomponents(uniondomain(UnitInterval(), UnitBall{Float64}(), UnitInterval())) == 2
        @test ncomponents(uniondomain(UnitInterval(), UnitInterval(), UnitBall{Float64}())) == 2
    end

    @testset "intersect" begin
        @test IntersectDomain(0..1) == IntersectDomain((0..1))
        @test IntersectDomain{Float64}(0..1) == IntersectDomain(0.0..1.0)
        @test IntersectDomain{Float64}(0..1, 1..2) == IntersectDomain((0..1, 1..2))
        @test intersectdomain(0..1, 1..2) == Point(1)

        @test intersectdomain(0..1, 0.5..1.5) == (0..1) & (0.5..1.5)

        # intersection of productdomains
        i1 = intersectdomain((-.4..0.4)^2, (-.5 .. 0.5) × (-.1.. 0.1))
        @test i1 == productdomain(-0.4..0.4, -0.1..0.1)
        show(io,i1)
        @test String(take!(io)) == "(-0.4..0.4) × (-0.1..0.1)"
        @test intersectdomain(productdomain(UnitDisk(),-1..1), productdomain(-1..1, UnitDisk())) isa IntersectDomain
        i2 = UnitDisk() & (-.4..0.4)^2
        show(io, textmime, i2)
        @test String(take!(io)) == "UnitDisk() ∩ ((-0.4..0.4) × (-0.4..0.4))"
        @test dimension(i1) == 2
        @test dimension(i2) == 2

        i3 = ((-.5 .. 0.5) × (-.1.. 0.1)) & i2
        i4 = i2 & ((-.5 .. 0.5) × (-.1.. 0.1))
        i5 = i3 & i2
        @test SVector(0.,.05) ∈ i3
        @test SVector(0.,.05) ∈ i4
        @test SVector(0.,.75) ∉ i3
        @test SVector(0.,.75) ∉ i4

        d45 = IntersectDomain(0.0..1.5, [1.0,3.0])
        @test d45 isa Domain{Float64}
        @test 1.0 ∈ d45
        @test 1.1 ∉ d45
        @test convert(Domain{BigFloat}, d45) isa Domain{BigFloat}

        @test intersectdomain() == EmptySpace{Any}()
        @test intersectdomain(UnitDisk()) == UnitDisk()

        @test (0..1) ∩ [1.5] isa IntersectDomain{Float64}
        @test [0.5] ∩ (1..2) isa IntersectDomain{Float64}

        @test intersectdomain(0..1, [0,1]) == [0,1]
        @test intersectdomain([0,1], 0..1) == [0,1]
        @test intersectdomain([0,1], [0.0,1.0]) == [0,1]
        @test intersectdomain([0,2], [0.0,1.0]) isa IntersectDomain

        @test IntersectDomain(UnitDisk(), UnitSquare()) == IntersectDomain(UnitSquare(), UnitDisk())

        @test intersectdomain(uniondomain(0..1, 2..3), uniondomain(0..0.5, 3.5..4.5)) == 0..0.5
        @test intersectdomain(uniondomain(0..1, 2..3), 2..4) == 2..3
        @test intersectdomain(2..4, uniondomain(0..1, 2..3)) == 2..3

        # repeated intersection
        @test ncomponents(intersectdomain(UnitBall{Float64}(), UnitInterval(), UnitInterval())) == 2
        @test ncomponents(intersectdomain(UnitInterval(), UnitBall{Float64}(), UnitInterval())) == 2
        @test ncomponents(intersectdomain(UnitInterval(), UnitInterval(), UnitBall{Float64}())) == 2
        # larger intersection expressions
        @test intersectdomain(0..1, 1..3, Point(0.4), 2..5, FullSpace(), Point(-0.2)) isa EmptySpace
        @test intersectdomain(0..1, 1..3, Point(1.0)) == Point(1.0)
        @test intersectdomain(0..1, Point(1.0), 1..3) == Point(1.0)
        @test intersectdomain(Point(1.0), 0..1, 1..3) == Point(1.0)

        @test boundingbox(UnitSphere() ∩ 2UnitBall()) == (-1..1)^3
    end

    @testset "setdiff" begin
        d1 = UnitDisk() \ ProductDomain(-.5..0.5, -.1..0.1)
        @test dimension(d1) == 2
        @test SVector(0.,.74) ∈ d1
        @test SVector(0.,.25) ∈ d1
        @test SVector(0.,-.05) ∉ d1
        @test SVector(0.5, 0.1) ∉ d1
        @test SVector(1.01, 0.1) ∉ d1
        @test approx_in(SVector(1.01, 0.1), d1, 0.1)
        show(io, textmime, d1)
        @test String(take!(io)) == "UnitDisk() \\ ((-0.5..0.5) × (-0.1..0.1))"
        @test setdiff(d1, ProductDomain(-0.5..0.5, -.1..0.1)) == d1

        d2 = SetdiffDomain(0.0..3.0, [1.0, 2.5])
        @test d2 isa Domain{Float64}
        @test d2 isa SetdiffDomain
        @test 0.99 ∈ d2
        @test_throws MethodError approx_in(3.01, d2, 0.1)
        @test 1.0 ∉ d2
        @test convert(Domain{BigFloat}, d2) isa Domain{BigFloat}

        @test (0..1) \ [0.5] isa SetdiffDomain{Float64}
        d3 = [0,5] \ (0..3)
        @test d3 isa SetdiffDomain{Int}
        @test 0 ∉ d3
        @test 5 ∈ d3

        @test setdiff(0..1, 2..3) == setdiffdomain(0..1, 2..3)
        @test setdiff(0..1, 0.5) == setdiffdomain(0..1, 0.5)
        @test setdiff(0.5, 0..1) == setdiffdomain(0.5, 0..1)

        @test setdiff(0..1, EmptySpace()) == 0..1
        @test setdiff(0..1, 0.0..1.0) == EmptySpace()

        @test (0..1)^2 \ UnitCircle() == UnitInterval()^2 \ UnitCircle()
    end

    @testset "arithmetic" begin
        d1 = (0..1)
        d2 = (2..3)
        d = UnionDomain(d1) ∪ UnionDomain(d2)

        @test d .+ 1 == UnionDomain(d1 .+ 1) ∪ (d2 .+ 1)
        @test d .- 1 == UnionDomain(d1 .- 1) ∪ (d2 .- 1)
        @test 2 * d  == UnionDomain(2 * d1)  ∪ (2 * d2)
        @test d * 2 == UnionDomain(d1 * 2) ∪ (d2 * 2)
        @test d / 2 == UnionDomain(d1 / 2) ∪ (d2 / 2)
        @test 2 \ d == UnionDomain(2 \ d1) ∪ (2 \ d2)

        @test infimum(d) == minimum(d) == 0
        @test supremum(d) == maximum(d) == 3
    end

    @testset "different types" begin
        d̃1 = (0..1)
        d1 = (0f0.. 1f0)
        d2 = (2..3)

        @test UnionDomain(d1) ∪ d2 == UnionDomain(d̃1) ∪ d2
    end

    @testset "disk × interval" begin
        d = (0.0..1) × UnitDisk()
        @test SVector(0.1, 0.2, 0.3) ∈ d

        d = UnitDisk() × (0.0..1)
        @test SVector(0.1, 0.2, 0.3) ∈ d
    end

    @testset "Wrapped domain" begin
        d = 0..1.0
        w = DomainSets.WrappedDomain(d)
        @test w isa Domain
        @test 0.5 ∈ w

        d1 = DomainSets.WrappedDomain{Any}([0,5])
        @test 0 ∈ d1
        @test 1 ∉ d1

        @test convert(Domain{Float64}, [0.4,0.5]) isa Domain{Float64}
    end

    @testset "more set operations" begin
        D = UnitInterval()^3
        S = 2 * UnitBall()
        @testset "joint domain" begin
            DS = D ∪ S
            @test SA[0.0, 0.6, 0.0] ∈ DS
            @test SA[0.9, 0.6,-2.5] ∉ DS
        end

        @testset "domain intersection" begin
            DS = D ∩ S
            @test SA[0.1, 0.1, 0.1] ∈ DS
            @test SA[0.1, -0.1, 0.1] ∉ DS
        end
        @testset "domain difference" begin
            DS = D\S
            @test SA[0.1, 0.1, 0.1] ∉ DS

            D1 = 2 * D
            D2 = D * 2
            D3 = D / 2

            @test SA[2., 2., 2.] ∈ D1
            @test SA[0.9, 0.6,-2.5] ∉ D1
            @test SA[2., 2., 2.] ∈ D2
            @test SA[0.9, 0.6,-2.5] ∉ D2
            @test SA[.5, .4, .45] ∈ D3
            @test SA[.3, 0.6,-.2] ∉ D3
        end
    end
end


using BandedMatrices, LinearAlgebra, LazyArrays, Test
import LazyArrays: MemoryLayout, MulAdd, DenseColumnMajor, ConjLayout
import Base: BroadcastStyle
import BandedMatrices: BandedStyle, BandedRows

@testset "broadcasting" begin
    @testset "general" begin
        n = 1000
        A = brand(n,n,1,1)
        B = Matrix{Float64}(undef, n,n)
        B .= exp.(A)
        @test B == exp.(Matrix(A)) == exp.(A)

        @test exp.(A) isa BandedMatrix
        @test bandwidths(exp.(A)) == (n-1,n-1)

        C = similar(A)
        @test_throws BandError C .= exp.(A)

        @test identity.(A) isa BandedMatrix
        @test bandwidths(identity.(A)) == bandwidths(A)

        @test (z -> exp(z)-1).(A) isa BandedMatrix
        @test bandwidths((z -> exp(z)-1).(A)) == bandwidths(A)

        @test A .+ 1 isa BandedMatrix
        @test (A .+ 1) == Matrix(A) .+ 1

        @test A ./ 1 isa BandedMatrix
        @test bandwidths(A ./ 1) == bandwidths(A)

        @test 1 .+ A isa BandedMatrix
        @test (1 .+ A) == 1 .+ Matrix(A)

        @test 1 .\ A isa BandedMatrix
        @test bandwidths(1 .\ A) == bandwidths(A)

        A = brand(10,1,1,1)
        @test A[:,1] isa Vector
        @test norm(A .- A[:,1]) == 0
        @test A ≈ A[:,1]
    end

    @testset "identity" begin
        n = 100
        A = brand(n,n,2,2)
        A.data[1,1] = NaN
        B = brand(n,n,2,2)
        B.data[1,1] = NaN
        A .= B
        @test A == B
        B = brand(n,n,1,1)
        B.data[1,1] = NaN
        A .= B
        B = brand(n,n,3,3)
        B[band(3)] .= 0
        B[band(-3)] .= 0
        B.data[1,1] = NaN
        B.data[end,end] = NaN
        A .= B
        @test A == B

        B = brand(n,n,0,3)
        B[band(3)] .= 0
        B.data[1,1] = NaN
        A .= B
        @test A == B

        B = brand(n,n,3,0)
        B[band(-3)] .= 0
        B.data[end,end] = NaN
        A .= B
        @test A == B

        B = brand(n,n,-1,1)
        A .= B
        @test A == B

        B = brand(n,n,1,-1)
        A .= B
        @test A == B
    end

    @testset "lmul!/rmul!" begin
        n = 1000
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        B .= (-).(A)
        @test -A isa BandedMatrix
        @test (-).(A) isa BandedMatrix
        @test bandwidths(A) == bandwidths(-A) == bandwidths((-).(A))
        @test B == -A == (-).(A)
        @test A-I isa BandedMatrix
        @test I-A isa BandedMatrix
        @test bandwidths(A) == bandwidths(A-I) == bandwidths(I-A)

        n = 1000
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        B .= 2.0.*A

        @test B ==  2A == 2.0.*A
        @test 2A isa BandedMatrix
        @test 2.0.*A isa BandedMatrix
        @test bandwidths(2A) == bandwidths(2.0.*A) == bandwidths(A)

        A .= 2.0.*A
        @test A == B

        n = 1000
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        B .= A.*2.0

        @test B ==  A*2 == A.*2.0
        @test A*2 isa BandedMatrix
        @test A .* 2.0 isa BandedMatrix
        @test bandwidths(A*2) == bandwidths(A.*2.0) == bandwidths(A)
        A .= A.*2.0
        @test A == B

        n = 1000
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        B .= A ./ 2.0

        @test B == A/2 == A ./ 2.0
        @test A/2 isa BandedMatrix
        @test A ./ 2.0 isa BandedMatrix
        @test bandwidths(A/2) == bandwidths(A ./ 2.0) == bandwidths(A)

        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        B .= 2.0 .\ A

        @test B == A/2 == A ./ 2.0
        @test 2\A isa BandedMatrix
        @test 2.0 .\ A isa BandedMatrix
        @test bandwidths(2\A) == bandwidths(2.0 .\ A) == bandwidths(A)

        A = brand(5,5,1,1)
        A.data .= NaN
        lmul!(0.0,A)
        @test isnan(norm(A)) == isnan(norm(lmul!(0.0,[NaN])))

        lmul!(false,A)
        @test norm(A) == 0.0

        A = brand(5,5,1,1)
        A.data .= NaN
        rmul!(A,0.0)
        @test isnan(norm(A)) == isnan(norm(rmul!([NaN],0.0)))

        rmul!(A,false)
        @test norm(A) == 0.0

        n = 100
        A = brand(n,n,2,2)
        B = brand(n,n,1,1)
        A[band(2)] .= A[band(-2)] .= 0
        B .= A ./ 2.0
        @test B == A / 2.0 == Matrix(A)/2.0

        B .= 2.0 .\ A
        @test B == 2.0 \ A == 2.0 \ Matrix(A)

        n = 100
        A = brand(n,n,2,2)
        B = brand(n,n,1,3)
        A[band(-2)] .= 0
        B .= A ./ 2.0
        @test B == A / 2.0 == Matrix(A)/2.0

        B .= 2.0 .\ A
        @test B == 2.0 \ A == 2.0 \ Matrix(A)

        A = brand(n,n,2,2)
        B = brand(n,n,3,1)
        A[band(2)] .= 0
        B .= A ./ 2.0
        @test B == A / 2.0 == Matrix(A)/2.0

        B .= 2.0 .\ A
        @test B == 2.0 \ A == 2.0 \ Matrix(A)
    end

    @testset "axpy!" begin
        n = 100
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        C = brand(n,n,3,3)
        C .= A .+ B
        @test C == A + B == A .+ B  == Matrix(A) + Matrix(B)
        @test A + B isa BandedMatrix
        @test A .+ B isa BandedMatrix
        @test bandwidths(A+B) == bandwidths(A.+B) == (2,2)
        B .= A .+ B
        @test B == C

        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        C = brand(n,n,3,3)
        C .= A .* B
        @test C == A .* B  == Matrix(A) .* Matrix(B)
        @test A .* B isa BandedMatrix
        @test bandwidths(A.*B) == (2,2)
        @time B .= A .* B
        @test B == C

        n = 100
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        B[band(2)] .= B[band(-2)] .= 0
        C = brand(n,n,1,1)
        C .= A .+ B
        @test C == A + B == A .+ B == Matrix(A) + Matrix(B)

        n = 100
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        B[band(-2)] .= 0
        C = brand(n,n,1,2)
        C .= A .+ B
        @test C == A + B == A .+ B == Matrix(A) + Matrix(B)

        n = 100
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        B[band(2)] .= 0
        C = brand(n,n,2,1)
        C .= A .+ B
        @test C == A + B == A .+ B == Matrix(A) + Matrix(B)

        n = 100
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        C = brand(n,n,3,3)

        C .= 2.0 .* A .+ B
        @test C == 2A+B == 2.0.*A .+ B
        @test all(BLAS.axpy!(2.0, A, copy(B)) .=== C)

        @test 2A + B isa BandedMatrix
        @test 2.0.*A .+ B isa BandedMatrix
        @test bandwidths(2A+B) == bandwidths(2.0.*A .+ B) == (2,2)
        B .= 2.0 .* A .+ B
        @test B == C
    end

    @testset "gbmv!" begin
        n = 100
        x = randn(n)
        A = brand(n,n,1,1)
        y = similar(x)
        y .= Mul(A,x)
        @test Matrix(A)*x ≈ y
        @test all(BLAS.gbmv!('N', n, 1, 1, 1.0, A.data, x, 0.0, copy(y)) .===
                    Broadcast.materialize!(MulAdd(1.0,A,x,0.0,similar(y))) .=== y)
        z = similar(y)
        z .= 2.0.*Mul(A,x) .+ 3.0.*y
        @test 2Matrix(A)*x + 3y ≈ z
        @test all(BLAS.gbmv!('N', n, 1, 1, 2.0, A.data, x, 3.0, copy(y)) .=== z)
        @test MemoryLayout(A') == BandedRows(DenseColumnMajor())
        y .= Mul(A',x)
        @test Matrix(A')*x ≈ Matrix(A)'*x ≈ y
        @test all(BLAS.gbmv!('T', n, 1, 1, 1.0, A.data, x, 0.0, copy(y)) .=== y)
        z = similar(y)
        z .= 2.0.*Mul(A',x) .+ 3.0.*y
        @test 2Matrix(A')*x + 3y ≈ z
        @test all(BLAS.gbmv!('T', n, 1, 1, 2.0, A.data, x, 3.0, copy(y)) .=== z)
        @test MemoryLayout(transpose(A)) == BandedRows(DenseColumnMajor())
        y .= Mul(transpose(A),x)
        @test Matrix(A')*x ≈ Matrix(A)'*x ≈ y
        @test all(BLAS.gbmv!('T', n, 1, 1, 1.0, A.data, x, 0.0, copy(y)) .=== y)

        n = 100
        x = randn(ComplexF64,n)
        A = brand(ComplexF64,n,n,1,1)
        y = similar(x)
        y .= Mul(A,x)
        @test Matrix(A)*x ≈ y
        @test all(BLAS.gbmv!('N', n, 1, 1, 1.0+0.0im, A.data, x, 0.0+0.0im, copy(y)) .=== y)
        @test LazyArrays.MemoryLayout(A') == ConjLayout(BandedRows(DenseColumnMajor()))
        z = similar(y)
        z .= (2.0+0.0im).*Mul(A,x) .+ (3.0+0.0im).*y
        @test 2Matrix(A)*x + 3y ≈ z
        @test all(BLAS.gbmv!('N', n, 1, 1, 2.0+0.0im, A.data, x, 3.0+0.0im, copy(y)) .=== z)
        y .= Mul(A',x)
        @test Matrix(A')*x ≈ Matrix(A)'*x ≈ y
        @test all(BLAS.gbmv!('C', n, 1, 1, 1.0+0.0im, A.data, x, 0.0+0.0im, copy(y)) .=== y)
        @test MemoryLayout(transpose(A)) == BandedRows(DenseColumnMajor())
        y .= Mul(transpose(A),x)
        @test Matrix(transpose(A))*x ≈ transpose(Matrix(A))*x ≈ y
        @test all(BLAS.gbmv!('T', n, 1, 1, 1.0+0.0im, A.data, x, 0.0+0.0im, copy(y)) .=== y)
    end

    @testset "non-blas matvec" begin
        n = 10
        A = brand(Int,n,n,1,1)
        x = rand(Int,n)
        y = similar(x)
        y .= Mul(A,x)
        @test y == Matrix(A)*x
        y .= Mul(A',x)
        @test y == Matrix(A')*x
        y .= Mul(transpose(A),x)
        @test y == Matrix(transpose(A))*x

        A = brand(Int,n,n,1,1) + im*brand(Int,n,n,1,1)
        x = rand(eltype(A),n)
        y = similar(x)
        y .= Mul(A,x)
        @test y == Matrix(A)*x
        y .= Mul(A',x)
        @test y == Matrix(A')*x
        y .= Mul(transpose(A),x)
        @test y == Matrix(transpose(A))*x
    end

    @testset "gbmm!" begin
        n = 10
        A = brand(n,n,1,1)
        B = brand(n,n,2,2)
        C = brand(n,n,3,3)
        C .= Mul(A,B)
        @test Matrix(C) ≈ Matrix(A)*Matrix(B)
        C .= Mul(A',B)
        @test Matrix(C) ≈ Matrix(A')*Matrix(B)
        C .= Mul(A,B')
        @test Matrix(C) ≈ Matrix(A)*Matrix(B')
        C = brand(n,n,4,4)
        C .= Mul(A,B)
        @test Matrix(C) ≈ Matrix(A)*Matrix(B)
        B = randn(n,n)
        C = similar(B)

        C .= Mul(A,B)
        @test C ≈ Matrix(A)*Matrix(B)
        C .= Mul(B,A)
        @test C ≈ Matrix(B)*Matrix(A)
    end

    @testset "Subarray" begin
        n = 10
        A = brand(n,n,1,1)
        B = view(brand(2n,2n,2,0),Base.OneTo(n),1:n)
        @test BroadcastStyle(typeof(B)) == BandedStyle()

        @test A .+ B isa BandedMatrix
        @test A + B isa BandedMatrix
        @test bandwidths(A .+ B) == (2,1)
        @test A .+ B == Matrix(A) + Matrix(B)
    end
end


# # Vibration of a cube of nearly incompressible material: alternative models

# Source code: [`unit_cube_modes_alt_tut.jl`](unit_cube_modes_alt_tut.jl)

# ## Description

# Compute the free-vibration spectrum of a unit cube of nearly
# incompressible isotropic material, E = 1, ν = 0.499, and ρ = 1 (refer to [1]).

# Here we show how alternative finite element models compare:
# The solution with the serendipity quadratic hexahedron is supplemented with
# solutions obtained with advanced finite elements: nodal-integration energy
# stabilized hexahedra and tetrahedra, and mean-strain hexahedra and
# tetrahedra.

# ## References

# [1] Puso MA, Solberg J (2006) A stabilized nodally integrated tetrahedral. International Journal for Numerical Methods in Engineering 67: 841-867.
# [2] P. Krysl, Mean-strain 8-node hexahedron with optimized energy-sampling
# stabilization, Finite Elements in Analysis and Design 108 (2016) 41–53.

# ![](unit_cube-mode7.png)

# ## Goals

# - Show how to set up a simulation loop that will run all the models and collect data.
# - Show how to present the computed spectrum curves.

##
# ## Definitions

# This is the finite element toolkit itself.
using FinEtools

# The linear stress analysis application is implemented in this package.
using FinEtoolsDeforLinear

# Convenience import.
using FinEtools.MeshExportModule

# The eigenvalue problem is solved with the Lanczos algorithm from this package.
using Arpack

using SymRCM

# The material properties and dimensions are defined with physical units.
E = 1*phun("PA");
nu = 0.499;
rho = 1*phun("KG/M^3");
a = 1*phun("M"); # length of the side of the cube
N = 8
neigvs = 20 # how many eigenvalues
OmegaShift = (0.01*2*pi)^2; # The frequency with which to shift

# The model is fully three-dimensional, and hence the material model  and the
# FEMM created below need to refer to an appropriate model-reduction scheme.
MR = DeforModelRed3D
material = MatDeforElastIso(MR, rho, E, nu, 0.0);

#
models = [
    ("H20", H20block, GaussRule(3,2), FEMMDeforLinear, 1),
    ("ESNICEH8", H8block, NodalTensorProductRule(3), FEMMDeforLinearESNICEH8, 2),
    ("ESNICET4", T4block, NodalSimplexRule(3), FEMMDeforLinearESNICET4, 2),
    ("MSH8", H8block, NodalTensorProductRule(3), FEMMDeforLinearMSH8, 2),
    ("MST10", T10block, TetRule(4), FEMMDeforLinearMST10, 1),
]

# Run  the simulation loop over all the models.
sigdig(n) = round(n * 10000) / 10000
results = let
    results = []
    for m in models
        fens, fes  = m[2](a, a, a, m[5]*N, m[5]*N, m[5]*N);
        @show count(fens)
        geom = NodalField(fens.xyz)
        u = NodalField(zeros(size(fens.xyz,1),3))
        numbering = let
            C = connectionmatrix(FEMMBase(IntegDomain(fes, m[3])), count(fens))
            numbering = symrcm(C)
        end
        numberdofs!(u, numbering);
        println("nfreedofs = $(u.nfreedofs)")

        femm = m[4](MR, IntegDomain(fes, m[3]), material);
        femm = associategeometry!(femm, geom)
        K = stiffness(femm, geom, u)
        M = mass(femm, geom, u);

        evals, evecs, nconv = eigs(K+OmegaShift*M, M; nev=neigvs, which=:SM)
        @show nconv == neigvs
        evals = evals .- OmegaShift;
        fs = real(sqrt.(complex(evals)))/(2*pi)
        println("$(m[1]) eigenvalues: $(sigdig.(fs)) [Hz]")
        push!(results, (m, fs))
    end
    results # return it
end

##
# ## Present the results graphically

using Gnuplot

@gp  "set terminal windows 0 "  :-
for r in results
    @gp  :- collect(1:length(r[2])) vec(r[2]) " lw 2 with lp title '$(r[1][1])' "  :-
end
@gp  :- "set xlabel 'Mode number [ND]'" :-
@gp  :- "set ylabel 'Frequency [Hz]'"


"""
dist(a, b)

Return the minimum unsigend distance between Intervals a and b.
returns missing if the two intervals do not have the same seqname.
"""
function dist(a::GenomicFeatures.Interval{S},b::GenomicFeatures.Interval{T}) where {S, T}

    a.seqname != b.seqname && return missing
    return minimum(abs.([a.first-b.first, a.first-b.last, a.last-b.first, a.last-b.last]))
end

"""
dist(a,b, f = minimum)

Linearly calculate all the distances between an interval a and a collection b,
return a value according to function f.

# Arguments

- `a::GenomicFeatures.interval{T}`: interval to calculate distances from.
- `b::GenomicFeatures.IntervalCollection{S}`: collection to calculate all distances to.
- `f::Function`: f Summarise all the pairwise distance for a and each interval in b
according to the supplied function f
Currently tested with `minimum`, `maximum`, `mean`, and `median`.
"""
function dist(
    a::GenomicFeatures.Interval{T},
    b::GenomicFeatures.IntervalCollection{S},
    f::Function = minimum
) where {T, S}

    d = Vector{Union{Int, Float64}}(undef,0)
    for interval in b
        tmp_d = dist(a,interval)
        if ismissing(tmp_d)
            continue
        end
        push!(d, tmp_d)
    end
    # Workaround to deal with empty collection
    isempty(d) && return(d)
    return f(d)
end

"""
dist(a, b, f, g)

Lineraly caluclates the distance between collection a and b.
Each distance between a and b is summarised according to f and then summarised
according to g.

# Agruments

- `a::GenomicFeatures.IntervalCollection{T}`: collection a.
- `b::GenomicFeatures.IntervalCollection{S}`: collectionb b.
- `f::Function = x -> x`: f summarisees the result of the distances
between each interval of a and all intervals in b.
tested with `minimum`. 'maximum`, `mean`, `median`
- `g::Function = minimum`: g summarises the distances returned by f for all
items in a. Currently tested with `minimum`. `maximum`, `mean`, and `median`.

# Note

Both g and f can be replaced by user-defined functions. f must take a
Vector{Union{Int, Float64}} and return Union{Int, Float64}.
g must take a Vector{Union{Int, Float64}} but can return any type.
Due to the fact that the result of f is stored in d before being summarised, it
is not possible to change f's return type  -
even if we use a function in g that could acept a vector of that type. Indeed
this is the reason for the need to Union{Int}
"""
function dist(
    a::GenomicFeatures.IntervalCollection{T},
    b::GenomicFeatures.IntervalCollection{S},
    g::Function = x -> x,
    f::Function = minimum
) where {S,T}

    d = Vector{Union{Int, Float64}}(undef,0)
    for interval in a
        tmp_d  =  dist(interval,b, f)
        # deal with ampty collections
        isempty(tmp_d) && continue
        push!(d,tmp_d)
    end
    return g(d)
end


macro horner(x, p...)
ex = esc(p[end])
for i = length(p)-1:-1:1
ex = :(muladd(t, $ex, $(esc(p[i]))))
end
Expr(:block, :(t = $(esc(x))), ex)
end


module HessianTest

import Calculus

using Test
using ForwardDiff
using ForwardDiff: Dual, Tag
using StaticArrays
using DiffTests

include(joinpath(dirname(@__FILE__), "utils.jl"))

#############################
# rosenbrock hardcoded test #
#############################

f = DiffTests.rosenbrock_1
x = [0.1, 0.2, 0.3]
v = f(x)
g = [-9.4, 15.6, 52.0]
h = [-66.0  -40.0    0.0;
     -40.0  130.0  -80.0;
       0.0  -80.0  200.0]

for c in (1, 2, 3), tag in (nothing, Tag((f,ForwardDiff.gradient), eltype(x)))
    println("  ...running hardcoded test with chunk size = $c and tag = $(repr(tag))")
    cfg = ForwardDiff.HessianConfig(f, x, ForwardDiff.Chunk{c}(), tag)
    resultcfg = ForwardDiff.HessianConfig(f, DiffResults.HessianResult(x), x, ForwardDiff.Chunk{c}(), tag)

    @test eltype(resultcfg) == eltype(cfg)

    @test isapprox(h, ForwardDiff.hessian(f, x))
    @test isapprox(h, ForwardDiff.hessian(f, x, cfg))

    out = similar(x, 3, 3)
    ForwardDiff.hessian!(out, f, x)
    @test isapprox(out, h)

    out = similar(x, 3, 3)
    ForwardDiff.hessian!(out, f, x, cfg)
    @test isapprox(out, h)

    out = DiffResults.HessianResult(x)
    ForwardDiff.hessian!(out, f, x)
    @test isapprox(DiffResults.value(out), v)
    @test isapprox(DiffResults.gradient(out), g)
    @test isapprox(DiffResults.hessian(out), h)

    out = DiffResults.HessianResult(x)
    ForwardDiff.hessian!(out, f, x, resultcfg)
    @test isapprox(DiffResults.value(out), v)
    @test isapprox(DiffResults.gradient(out), g)
    @test isapprox(DiffResults.hessian(out), h)
end

cfgx = ForwardDiff.HessianConfig(sin, x)
@test_throws ForwardDiff.InvalidTagException ForwardDiff.hessian(f, x, cfgx)
@test ForwardDiff.hessian(f, x, cfgx, Val{false}()) == ForwardDiff.hessian(f,x)

########################
# test vs. Calculus.jl #
########################

for f in DiffTests.VECTOR_TO_NUMBER_FUNCS
    v = f(X)
    g = ForwardDiff.gradient(f, X)
    h = ForwardDiff.hessian(f, X)
    # finite difference approximation error is really bad for Hessians...
    @test isapprox(h, Calculus.hessian(f, X), atol=0.02)
    for c in CHUNK_SIZES, tag in (nothing, Tag((f,ForwardDiff.gradient), eltype(x)))
        println("  ...testing $f with chunk size = $c and tag = $(repr(tag))")
        cfg = ForwardDiff.HessianConfig(f, X, ForwardDiff.Chunk{c}(), tag)
        resultcfg = ForwardDiff.HessianConfig(f, DiffResults.HessianResult(X), X, ForwardDiff.Chunk{c}(), tag)

        out = ForwardDiff.hessian(f, X, cfg)
        @test isapprox(out, h)

        out = similar(X, length(X), length(X))
        ForwardDiff.hessian!(out, f, X, cfg)
        @test isapprox(out, h)

        out = DiffResults.HessianResult(X)
        ForwardDiff.hessian!(out, f, X, resultcfg)
        @test isapprox(DiffResults.value(out), v)
        @test isapprox(DiffResults.gradient(out), g)
        @test isapprox(DiffResults.hessian(out), h)
    end
end

##########################################
# test specialized StaticArray codepaths #
##########################################

println("  ...testing specialized StaticArray codepaths")

x = rand(3, 3)
sx = StaticArrays.SArray{Tuple{3,3}}(x)

cfg = ForwardDiff.HessianConfig(nothing, x)
scfg = ForwardDiff.HessianConfig(nothing, sx)

actual = ForwardDiff.hessian(prod, x)
@test ForwardDiff.hessian(prod, sx) == actual
@test ForwardDiff.hessian(prod, sx, cfg) == actual
@test ForwardDiff.hessian(prod, sx, scfg) == actual
@test ForwardDiff.hessian(prod, sx, scfg) isa StaticArray
@test ForwardDiff.hessian(prod, sx, scfg, Val{false}()) == actual
@test ForwardDiff.hessian(prod, sx, scfg, Val{false}()) isa StaticArray

out = similar(x, 9, 9)
ForwardDiff.hessian!(out, prod, sx)
@test out == actual

out = similar(x, 9, 9)
ForwardDiff.hessian!(out, prod, sx, cfg)
@test out == actual

out = similar(x, 9, 9)
ForwardDiff.hessian!(out, prod, sx, scfg)
@test out == actual

result = DiffResults.HessianResult(x)
result = ForwardDiff.hessian!(result, prod, x)

result1 = DiffResults.HessianResult(x)
result2 = DiffResults.HessianResult(x)
result3 = DiffResults.HessianResult(x)
result1 = ForwardDiff.hessian!(result1, prod, sx)
result2 = ForwardDiff.hessian!(result2, prod, sx, ForwardDiff.HessianConfig(prod, result2, x, ForwardDiff.Chunk(x), nothing))
result3 = ForwardDiff.hessian!(result3, prod, sx, ForwardDiff.HessianConfig(prod, result3, x, ForwardDiff.Chunk(x), nothing))
@test DiffResults.value(result1) == DiffResults.value(result)
@test DiffResults.value(result2) == DiffResults.value(result)
@test DiffResults.value(result3) == DiffResults.value(result)
@test DiffResults.gradient(result1) == DiffResults.gradient(result)
@test DiffResults.gradient(result2) == DiffResults.gradient(result)
@test DiffResults.gradient(result3) == DiffResults.gradient(result)
@test DiffResults.hessian(result1) == DiffResults.hessian(result)
@test DiffResults.hessian(result2) == DiffResults.hessian(result)
@test DiffResults.hessian(result3) == DiffResults.hessian(result)

sresult1 = DiffResults.HessianResult(sx)
sresult2 = DiffResults.HessianResult(sx)
sresult3 = DiffResults.HessianResult(sx)
sresult1 = ForwardDiff.hessian!(sresult1, prod, sx)
sresult2 = ForwardDiff.hessian!(sresult2, prod, sx, ForwardDiff.HessianConfig(prod, sresult2, x, ForwardDiff.Chunk(x), nothing))
sresult3 = ForwardDiff.hessian!(sresult3, prod, sx, ForwardDiff.HessianConfig(prod, sresult3, x, ForwardDiff.Chunk(x), nothing))
@test DiffResults.value(sresult1) == DiffResults.value(result)
@test DiffResults.value(sresult2) == DiffResults.value(result)
@test DiffResults.value(sresult3) == DiffResults.value(result)
@test DiffResults.gradient(sresult1) == DiffResults.gradient(result)
@test DiffResults.gradient(sresult2) == DiffResults.gradient(result)
@test DiffResults.gradient(sresult3) == DiffResults.gradient(result)
@test DiffResults.hessian(sresult1) == DiffResults.hessian(result)
@test DiffResults.hessian(sresult2) == DiffResults.hessian(result)
@test DiffResults.hessian(sresult3) == DiffResults.hessian(result)

end # module