@testset verbose=true "Distance" begin
    @testset "Int" begin
        scalars = Int[0, 1, -1]
        quaternions = [Quaternion(a, b, c, d) for a in scalars for b in scalars for c in scalars for d in scalars]
        d(q₁, q₂) = distance2(q₁, q₂)
        for q₁ in quaternions
            for q₂ in quaternions
                d₁₁ = d(q₁, q₁)
                d₁₂ = d(q₁, q₂)
                d₂₂ = d(q₂, q₂)
                d₂₁ = d(q₂, q₁)
                # Real-valued:
                @test d₁₁ isa Int
                @test d₁₂ isa Int
                @test d₂₁ isa Int
                @test d₂₂ isa Int
                # Symmetry:
                @test d₁₂ == d₂₁
                # Invariance:
                for q₃ in [Quaternion(1), 1imx, 1imy, 1imz]
                    @test d₁₂ == d(q₃*q₁, q₃*q₂)
                    @test d₁₂ == d(q₁*q₃, q₂*q₃)
                end
                # Identity:
                @test d₁₁ == 0
                @test d₂₂ == 0
                # Positive-definiteness
                @test (q₁==q₂) ⊻ (d₁₂>0)
            end
        end
    end
    @testset "$T" for T in [BigFloat, Float64]  # FloatTypes
        scalars = [zero(T), one(T), -one(T)]
        quaternions = [a + b*imn for a in scalars for b in scalars for imn in [imx, imy, imz]]
        quaternions = [quaternions...; randn(Quaternion{T}, 23)...]
        # quaternions = [Quaternion(a, b, c, d) for a in scalars for b in scalars for c in scalars for d in scalars]
        # quaternions = [quaternions...; randn(Quaternion{T}, 19)...]
        @testset "rotor=$rotor" verbose=true for rotor in [true, false]
            @testset "sqrt=$return_sqrt" verbose=true for return_sqrt in [true, false]
                d(q₁, q₂) = return_sqrt ? distance(q₁, q₂) : distance2(q₁, q₂)
                for (i, q₁) in enumerate(quaternions)
                    if rotor
                        if iszero(q₁)
                            continue
                        end
                        q₁ = Rotor(q₁)
                    end
                    for q₂ in quaternions[i+1:end]
                        if rotor
                            if iszero(q₂)
                                continue
                            end
                            q₂ = Rotor(q₂)
                        end
                        absq₁, absq₂ = abs(q₁), abs(q₂)
                        rtol = 10eps(T)
                        atol = 10eps(max(absq₁, absq₂))
                        d₁₁ = d(q₁, q₁)
                        d₁₂ = d(q₁, q₂)
                        d₂₁ = d(q₂, q₁)
                        d₂₂ = d(q₂, q₂)
                        @test d₁₁ isa T
                        @test d₁₂ isa T
                        @test d₂₁ isa T
                        @test d₂₂ isa T
                        @test d₁₂ ≈ d₂₁ rtol=rtol atol=atol
                        for q₃ in quaternions
                            if iszero(q₃)
                                continue
                            end
                            if rotor
                                q₃ = Rotor(q₃)
                            else
                                q₃ = q₃ / abs(q₃)
                            end
                            @test d₁₂ ≈ d(q₃*q₁, q₃*q₂) rtol=rtol atol=atol
                            @test d₁₂ ≈ d(q₁*q₃, q₂*q₃) rtol=rtol atol=atol
                        end
                        @test d₁₁ ≈ zero(T) rtol=rtol atol=atol
                        @test d₂₂ ≈ zero(T) rtol=rtol atol=atol
                        if rotor
                            scalar_multiples = (≈(q₁, q₂, rtol=rtol, atol=atol) || ≈(q₁, -q₂, rtol=rtol, atol=atol))
                            @test scalar_multiples ⊻ (d₁₂ > atol)
                        else
                            @test (q₁==q₂) ⊻ (d₁₂>0)
                        end
                    end
                end
            end
        end
    end
end


# LSP indexing solution
# https://github.com/julia-vscode/julia-vscode/issues/800#issuecomment-650085983
if isdefined(@__MODULE__, :LanguageServer)
    include("../src/SpinDoctor.jl")
    using .SpinDoctor
end

using SpinDoctor
using LinearAlgebra
using GLMakie
using OrdinaryDiffEq: Rodas4

## Chose a plotting theme
set_theme!(theme_light())
set_theme!(theme_dark())
set_theme!(theme_black())

## Create model from setup recipe
# include("setups/axon.jl")
# include("setups/sphere.jl")
# include("setups/plates.jl")
# include("setups/cylinders.jl")
# include("setups/spheres.jl")
include("setups/neuron.jl")

mesh, = @time create_geometry(setup; recreate = true);
model = Model(; mesh, coeffs...);
volumes = get_cmpt_volumes(model.mesh)
D_avg = 1 / 3 * tr.(model.D)' * volumes / sum(volumes)
ncompartment = length(model.mesh.points)
@info "Number of nodes per compartment:" length.(model.mesh.points)

## Plot mesh
plot_mesh(model.mesh)

## Assemble finite element matrices
matrices = @time assemble_matrices(model);

## Magnetic field gradient
TE = 5000
φ = -π / 6
R = [cos(φ) sin(φ) 0; -sin(φ) cos(φ) 0; 0 0 1]
g⃗(t) = 1.0 * R * [sin(2π * t / TE), sin(20π * t / TE) / 5, cos(2π * t / TE)]
gradient = GeneralGradient{T,typeof(g⃗)}(; g⃗, TE)

## Solve BTPDE

# Callbacks for time stepping (plot solution, save time series)
printer = Printer(; nupdate = 1, verbosity = 2)
writer = VTKWriter(; nupdate = 5)
plotter = Plotter{T}(; nupdate = 5)
callbacks = [printer, plotter]
# callbacks = [printer, plotter, writer]

# General BTPDE for all gradients
btpde = GeneralBTPDE(;
    model,
    matrices,
    reltol = 1e-4,
    abstol = 1e-6,
    odesolver = Rodas4(autodiff = false),
)

# Solve BTPDE
ξ = @time solve(btpde, gradient; callbacks)

## Plot magnetization
plot_field(model.mesh, ξ)

## Compute signal
compute_signal(matrices.M, ξ)
compute_signal.(matrices.M_cmpts, split_field(model.mesh, ξ))

## Matrix Formalism

# Perform Laplace eigendecomposition
laplace = Laplace{T}(; model, matrices, neig_max = 400)
lap_eig = @time solve(laplace)
length_scales = eig2length.(lap_eig.values, D_avg)

# Truncate basis at minimum length scale
length_scale = 3
λ_max = length2eig(length_scale, D_avg)
lap_eig = limit_lengthscale(lap_eig, λ_max)

# Compute magnetization using the matrix formalism reduced order model
mf = MatrixFormalism(; model, matrices, lap_eig, ninterval = 500)
ξ = @time solve(mf, gradient)


using BenchmarkTools
using UnobservableQueue
const UQ = UnobservableQueue

c = 7 # number of servers
λ = 0.99 # arrival rate (mean interarrival time is 1/λ)
μ = 1 / c # service rate (mean service time is 1/μ) (single server)
adist = Exponential(1/λ) # arrival distribution
sdist = Exponential(1/μ) # service distribution (single server)
ncust = 10000 # total number of customers generated
timelimit = 10_000 # time limit for simulation
seed = 1302 # rng seed
cmax = 19
s2 = UQ.Sim(ncust, timelimit, seed)
q2 = UQ.Queue(c, adist, sdist)
df1 = UQ.runsim(s2, q2)
A1, D1 = df1[:atime], df1[:dtime]
outorder1 = sort(df1, :dorder)[:aorder]

function f(df1::DataFrame)
    outorder1 = sort(df1, :dorder)[:aorder]
    UQ.c_order(outorder1)
end

bench_var = @benchmark UQ.c_var_unf($A1, $D1, $cmax)
bench_order = @benchmark f($df1)


"""
    pdf(
        f::AbstractArray{<:AbstractFloat, N ∈ (3, 4)},
        mask::AbstractArray{Bool, 3},
        vsz::NTuple{3, Real};
        W::Union{Nothing, AbstractArray{<:AbstractFloat, M ∈ (3, N)}} = nothing,
        pad::NTuple{3, Integer} = (0, 0, 0),
        bdir::NTuple{3, Real} = (0, 0, 1),
        Dkernel::Symbol = :i,
        lambda::Real = 0,
        tol::Real = 1e-5,
        maxit::Integer = ceil(sqrt(numel(mask))),
        verbose::Bool = false
    ) typeof(similar(f))

Projection onto dipole fields (PDF) [1].

### Arguments
- `f::AbstractArray{<:AbstractFloat, N ∈ (3, 4)}`: unwrapped (multi-echo) field/phase
- `mask::AbstractArray{Bool, 3}`: binary mask of region of interest
- `vsz::NTuple{3, Real}`: voxel size for dipole kernel

### Keywords
- `W::Union{Nothing, AbstractArray{<:AbstractFloat, M ∈ (3, N)}} = nothing`:
    data fidelity weights
- `pad::NTuple{3, Integer} = (0, 0, 0)`: zero padding array
    - `< 0`: no padding
    - `≥ 0`: minimum padding to fast fft size
- `bdir::NTuple{3, Real} = (0, 0, 1)`: unit vector of B field direction
- `Dkernel::Symbol = :i`: dipole kernel method
- `lambda::Real = 0`: regularization parameter
- `tol::Real = 1e-5`: stopping tolerance for iterative solver
- `maxit::Integer = ceil(sqrt(length(mask)))`: maximum number of iterations for
    iterative solver
- `verbose::Bool = false`: print convergence information

### Returns
- `typeof(similar(f))`: background corrected local field/phase

### References
[1] Liu T, Khalidov I, de Rochefort L, Spincemaille P, Liu J, Tsiouris AJ,
    Wang Y. A novel background field removal method for MRI using projection
    onto dipole fields. NMR in Biomedicine. 2011 Nov;24(9):1129-36.
"""
function pdf(
    f::AbstractArray{<:AbstractFloat, N},
    mask::AbstractArray{Bool, 3},
    vsz::NTuple{3, Real};
    W::Union{Nothing, AbstractArray{<:AbstractFloat}} = nothing,
    pad::NTuple{3, Integer} = (0, 0, 0),
    bdir::NTuple{3, Real} = (0, 0, 1),
    Dkernel::Symbol = :i,
    lambda::Real = 0,
    tol::Real = 1e-5,
    maxit::Integer = ceil(Int, sqrt(length(mask))),
    verbose::Bool = false
) where {N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    return _pdf!(
        tzero(f), f, mask, vsz, W, pad, bdir, Dkernel, lambda, tol, maxit, verbose
    )
end

function _pdf!(
    fl::AbstractArray{<:AbstractFloat, N},
    f::AbstractArray{T, N},
    mask::AbstractArray{Bool, 3},
    vsz::NTuple{3, Real},
    W::Union{Nothing, AbstractArray{<:AbstractFloat}},
    pad::NTuple{3, Integer},
    bdir::NTuple{3, Real},
    Dkernel::Symbol,
    lambda::Real,
    tol::Real,
    maxit::Integer,
    verbose::Bool = false
) where {T<:AbstractFloat, N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))

    size(fl) == size(f) || throw(DimensionMismatch())
    size(mask) == size(f)[1:3] || throw(DimensionMismatch())

    if W !== nothing
        ndims(W) ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(M)d"))
        size(W)[1:3] == size(f)[1:3] || throw(DimensionMismatch())
        ndims(W) == 4 && size(W, 4) == size(f, 4) || throw(DimensionMismatch())
    end

    Dkernel ∈ (:k, :kspace, :i, :ispace) ||
        throw(ArgumentError("Dkernel must be one of :k, :kspace, :i, :ispace, got :$(Dkernel)"))

    # pad to fast fft size
    fp = padfastfft(@view(f[:,:,:,1]), pad, rfft=true)
    m = padfastfft(mask, pad, rfft=true)

    # init vars and fft plans
    sz0 = size(mask)
    sz = size(m)
    sz_ = (sz[1]>>1 + 1, sz[2], sz[3])

    M̃ = similar(m)
    MW = similar(fp)

    D = Array{T, 3}(undef, sz_)
    F̂ = Array{complex(T), 3}(undef, sz_)

    FFTW.set_num_threads(num_cores())
    P = plan_rfft(fp)
    iP = inv(P)

    # LSMR setup
    # move up here so we can use LSMR's `b` as a tmp variable
    A = LinearMap{T}(
        (Av, v) -> _A_pdf!(Av, v, MW, iP, D, F̂, P, M̃),
        (Atu, u) -> _At_pdf!(Atu, u, M̃, iP, D, F̂, P, MW),
        length(fp),
        issymmetric = false,
        ishermitian = false,
        isposdef = true,
        ismutating = true
    )

    WS = LSMRWorkspace(A)
    b  = reshape(WS.b, sz)

    # get dipole kernel
    D = _dipole_kernel!(D, F̂, b, sz0, vsz, bdir, P, Dkernel, :rfft)

    # background mask
    @inbounds @batch for I in eachindex(m)
        M̃[I] = !m[I]
    end

    # pre-compute mask*weights
    if W === nothing
        # no weights
        MW = _tcopyto!(MW, m)

    elseif ndims(W) == 3
        # same weights for all echoes
        b = padarray!(b, W)
        @inbounds @batch for I in eachindex(MW)
            MW[I] = m[I] * b[I]
        end

    elseif ndims(W) == 4
        # weights for each echo. compute inside for loop
        # VOID
    end

    @inbounds for t in axes(f, 4)
        if verbose && size(f, 4) > 1
            @printf("Echo: %d/%d\n", t, size(f, 4))
        end

        if t > 1
            fp = padarray!(fp, @view(f[:,:,:,t]))
        end

        # mask * weight if we can't precompute
        if W !== nothing
            if ndims(W) == 4
                b = padarray!(b, @view(W[:,:,:,t]))
                @batch for I in eachindex(MW)
                    MW[I] = m[I] * b[I]
                end
            end
        end

        # rhs
        @batch for I in eachindex(b)
            b[I] = MW[I] * fp[I]
        end

        # solve
        lsmr!(WS; lambda=lambda, atol=tol, btol=tol, maxit=maxit, verbose=verbose)

        # compute background fields
        xb = reshape(WS.x, sz)
        @batch for I in eachindex(xb)
            xb[I] *= M̃[I]
        end

        F̂ = mul!(F̂, P, xb)
        @batch for I in eachindex(F̂)
            F̂[I] *= D[I]
        end

        xb = mul!(xb, iP, F̂)
        @batch for I in eachindex(fp)
            fp[I] = m[I] * (fp[I] - xb[I])
        end

        if verbose
            println()
        end

        unpadarray!(@view(fl[:,:,:,t]), fp)
    end

    return fl
end

function _A_pdf!(Av, v, W, iP, D, F̂, P, M̃)
    v = reshape(v, size(W))
    x = reshape(Av, size(W))

    @inbounds @batch for I in eachindex(x)
        x[I] = M̃[I] * v[I]
    end

    F̂ = mul!(F̂, P, x)
    @inbounds @batch for I in eachindex(F̂)
        F̂[I] *= D[I]
    end

    x = mul!(x, iP, F̂)
    @inbounds @batch for I in eachindex(x)
        x[I] *= W[I]
    end

    return Av
end

function _At_pdf!(Atu, u, M̃, iP, D, F̂, P, W)
    u = reshape(u, size(W))
    x = reshape(Atu, size(W))

    @inbounds @batch for I in eachindex(x)
        x[I] = W[I] * u[I]
    end

    F̂ = mul!(F̂, P, x)
    @inbounds @batch for I in eachindex(F̂)
        F̂[I] *= D[I] # conj(D), D is real
    end

    x = mul!(x, iP, F̂)
    @inbounds @batch for I in eachindex(x)
        x[I] *= M̃[I]
    end

    return Atu
end


module ParametricFunctions

using ForwardDiff
using ChainRulesCore

export ParametricFun, BasisFun
export Chebyshev, Fourier

export Parametric, NonParametric

export to_grad, from_grad, params

"""
    ParametricFun

An abstract type representing a parametric function.
"""
abstract type ParametricFun end

include("basis_fun.jl")
include("chebyshev_basis.jl")
include("fourier_basis.jl")

"""
    Parametric

A Trait type denoting functions with trainable parameters.

Functions implementing this trait must have methods for `has_params`, `to_grad`, `from_grad` and `params`.
"""
struct Parametric end

"""
    NonParametric

A trait type denoting functions with no parameters.
"""
struct NonParametric end

"""
    has_params(f)

Whether a functor type has trainable parameters.

Override this to return `Parametric()` for your type to declare it as parametric functions.
"""
has_params(f) = NonParametric()
has_params(::ParametricFun) = Parametric()

"""
    to_grad(f, x)

Return the gradient of `f` relative to its parameters at `x`.

The returned object must be an array (or any object that can be used in  a `QuadratureProblem`).
"""
function to_grad end

"""
        from_grad(f, ∇f)

Create a new instance of the type of `f` from a gradient over its parameters.
"""
function from_grad end

"""
    params(f)

Return the parameters of `f`.
"""
function params end

end


"""
```julia
GaussianLikelihood(σ²::T=1e-3) #σ² is the variance
```
Gaussian noise :
```math
    p(y|f) = N(y|f,σ²)
```
There is no augmentation needed for this likelihood which is already conjugate to a Gaussian prior
"""
mutable struct GaussianLikelihood{T<:Real} <: RegressionLikelihood{T}
    σ²::T
    opt_noise::Union{Nothing,Optimizer}
    θ::Vector{T}
    function GaussianLikelihood{T}(σ²::T,opt_noise::Union{Bool,Nothing,Optimizer}) where {T<:Real}
        new{T}(σ²,opt_noise)
    end
    function GaussianLikelihood{T}(σ²::T,opt_noise::Union{Bool,Nothing,Optimizer},θ::AbstractVector{T}) where {T<:Real}
        new{T}(σ²,opt_noise,θ)
    end
end

function GaussianLikelihood(σ²::T=1e-3;opt_noise::Union{Bool,Nothing,Optimizer}=Adam(α=0.05)) where {T<:Real}
    if isa(opt_noise,Bool)
        opt_noise = opt_noise ? Adam(α=0.05) : nothing
    end
    GaussianLikelihood{T}(σ²,opt_noise)
end

function pdf(l::GaussianLikelihood,y::Real,f::Real)
    pdf(Normal(y,l.σ²),f)
end

function logpdf(l::GaussianLikelihood,y::Real,f::Real)
    logpdf(Normal(y,l.σ²),f)
end

export noise

noise(l::GaussianLikelihood) = l.σ²

function Base.show(io::IO,model::GaussianLikelihood{T}) where T
    print(io,"Gaussian likelihood")
end

function compute_proba(l::GaussianLikelihood{T},μ::AbstractVector{T},σ²::AbstractVector{T}) where {T<:Real}
    return μ,max.(σ²,zero(σ²)).+ l.σ²
end

function init_likelihood(likelihood::GaussianLikelihood{T},inference::Inference{T},nLatent::Int,nSamplesUsed::Int,nFeatures::Int) where {T<:Real}
    return GaussianLikelihood{T}(likelihood.σ²,likelihood.opt_noise,fill(inv(likelihood.σ²),nSamplesUsed))
end

function local_updates!(l::GaussianLikelihood{T},y::AbstractVector,μ::AbstractVector,diag_cov_f::AbstractVector) where {T}
    if !isnothing(l.opt_noise)
        if l.opt_noise.t<=10
            update(l.opt_noise,zero(T))
        else
            l.σ² = exp(log(l.σ²)+update(l.opt_noise,0.5*((sum(abs2,y-μ)+sum(diag_cov_f))/l.σ²-length(y))))
        end
    end
    l.θ .= inv(l.σ²)
end

@inline ∇E_μ(l::GaussianLikelihood{T},::AOptimizer,y::AbstractVector) where {T} = (y./l.σ²,)
@inline ∇E_Σ(l::GaussianLikelihood{T},::AOptimizer,y::AbstractVector) where {T} = (0.5*l.θ,)

### Special case where the ELBO is equal to the marginal likelihood
function ELBO(model::GP{T}) where {T}
    model.f[1].Σ = Symmetric(inv(model.f[1].K+model.likelihood.σ²*I).mat)
    return -0.5*dot(model.y,model.f[1].Σ*model.y) - logdet(model.f[1].Σ)+ model.nFeatures*log(twoπ)
end

function expec_log_likelihood(l::GaussianLikelihood{T},i::AnalyticVI,y::AbstractVector,μ::AbstractVector,diag_cov::AbstractVector) where {T}
    return -0.5*i.ρ*(length(y)*(log(twoπ)+log(l.σ²))+sum(abs2.(y-μ)+diag_cov)/l.σ²)
end

AugmentedKL(::GaussianLikelihood{T},::AbstractVector) where {T} = zero(T)


"""
    BrennanSchwartz()

[Brennan and Schwartz model](https://link.springer.com/chapter/10.1057/9781403946027_3).
"""
struct BrennanSchwartz <: Model end

"""
    evaluate(IRD, BrennanSchwartz(), n = 12)

Evaluate interest rate derivative `IRD` using `BrennanSchwartz` model.

# Arguments
- `IRD::InterestRateDerivative`: interest rate derivative
- `n`: number of paths to simulate
"""
function evaluate(IRD::InterestRateDerivative, m::BrennanSchwartz, n::Int64 = 12)
    Δt = IRD.t / n
    rates = [IRD.r]
    for i = 0:n
        dr = IRD.k * (IRD.θ - rates[end]) * Δt + IRD.σ * randn()
        append!(rates, rates[end] + dr)
    end
    return rates
end


function anderson_mix(vin, vout, beta, ymat, smat, iter::Int64, mixdim)
    # function that computes the new potential usign the anderson mix
    # TODO: I think that the mixing can be efficiently be explained by a small
    # jupyter notebook

    n = length(vin);
    #
    # function evaluation overwrites vout
    #
    # residue
    res = vin - vout;
    #
    iterused = min(iter-1,mixdim);
    ipos = iter - 1 - round(Integer, floor((iter-2)/mixdim))*mixdim;
    #
    if (iter > 1)
        # compute the changes in function evaluations and the step (changes
        # in potentials)
        ymat[:,ipos] = res - ymat[:,ipos];
        smat[:,ipos] = vin - smat[:,ipos];
    end
    #
    vopt  = vin;
    ropt  = res;
    #
    if (iter > 1)
        # Solve the least square problem
        gammas = ymat[:,1:iterused] \ res;
        vopt   = vin - smat[:,1:iterused] * gammas;
        ropt   = res - ymat[:,1:iterused] * gammas;
    end

    inext = iter - round(Integer, floor((iter - 1) / mixdim)) * mixdim;
    ymat[:,inext] = res;
    smat[:,inext] = vin;

    vnew = vopt - beta*ropt;

    return (vnew, ymat, smat)
end


function EBFMI(Es::AbstractVector{<:AbstractScalarOrVec{<:AbstractFloat}})
    return mean(i -> (Es[i+1] - Es[i]) .^ 2, 1:length(Es)-1) ./ var(Es)
end


@testset "benchmark_normals" begin
    p, q = synthetic_gradient(SynthSphere(50))
    p2, q2 = synthetic_gradient(SynthSphere(51))
    error = benchmark_normals(p, q, p2, q2)
    @test error ≈ 1.4866545112360603
end


#
#
#
using Random, Manifolds, Manopt, Test
Random.seed!(29)
@testset "Test Nelder-Mead" begin
    @testset "Euclidean" begin
        M = Euclidean(6)
        # From Wikipedia https://en.wikipedia.org/wiki/Rosenbrock_function
        function Rosenbrock(::Euclidean, x)
            return sum([
                100 * (x[2 * i - 1]^2 - x[2 * i])^2 + (x[2 * i - 1] - 1)^2 for
                i in 1:div(length(x), 2)
            ])
        end
        x0 = [8 * randn(6) for i in 1:7]
        o = NelderMead(M, Rosenbrock, x0; record=[RecordCost()], return_options=true)
        x = get_solver_result(o)
        rec = get_record(o)
        nonincreasing = [rec[i] >= rec[i + 1] for i in 1:(length(rec) - 1)]
        @test any(map(!, nonincreasing)) == false

        x2 = NelderMead(M, Rosenbrock, x0)
        @test x == x2
    end
    @testset "Rotations" begin
        M = Rotations(3)
        A = randn(3, 3)
        A .= (A - A') ./ 2
        f(::Rotations, x) = norm(A * x * x * A)
        x0 = [random_point(M) for _ in 1:12]
        o = NelderMead(
            M,
            f,
            x0;
            record=[RecordCost()],
            return_options=true,
            retraction_method=QRRetraction(),
            inverse_retraction_method=QRInverseRetraction(),
            stopping_criterion=StopAfterIteration(500),
        )
        x = get_solver_result(o)
        rec = get_record(o)
        nonincreasing = [rec[i] >= rec[i + 1] for i in 1:(length(rec) - 1)]
        @test any(map(!, nonincreasing)) == false
    end
end


"""
	CLFQP <: Policy

Struct representing a Control Lyapunov Function based quadratic program (CLFQP).

# Fields
- `control`: function u = control(x) describing the control law.
"""
struct CLFQP <: Policy
    control
end

"""
	(k::CLFQP)(x)

Solve CLF-QP for control input at state x.
"""
function (κ::CLFQP)(x)
    return κ.control(x)
end

"""
	CLFQP(Σ::ControlAffineSystem, CLF::ControlLyapunovFunction)
    CLFQP(Σ::ControlAffineSystem, CLF::ControlLyapunovFunction, A, b)
    CLFQP(Σ::ControlAffineSystem, CLF::ControlLyapunovFunction, A, b, p::Float64)

Construct a CLF-QP from a ControlAffineSystem and ControlLyapunovFunction.
"""
function CLFQP(Σ::ControlAffineSystem, CLF::ControlLyapunovFunction)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u)
        @constraint(model, clf, CLF.∇V(x) * (Σ.f(x) + Σ.g(x) * u) <= -CLF.α(CLF(x)))
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CLFQP(control)
end

function CLFQP(Σ::ControlAffineSystem, CLF::ControlLyapunovFunction, A, b)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @objective(model, Min, (1 / 2)u'u)
        @constraint(model, clf, CLF.∇V(x) * (Σ.f(x) + Σ.g(x) * u) <= -CLF.α(CLF(x)))
        @constraint(model, A * u .<= b)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CLFQP(control)
end

function CLFQP(Σ::ControlAffineSystem, CLF::ControlLyapunovFunction, A, b, p::Float64)
    function control(x)
        model = Model(OSQP.Optimizer)
        set_silent(model)
        Σ.m == 1 ? @variable(model, u) : @variable(model, u[1:(Σ.m)])
        @variable(model, δ)
        @objective(model, Min, (1 / 2)u'u + p * δ^2)
        @constraint(model, clf, CLF.∇V(x) * (Σ.f(x) + Σ.g(x) * u) <= -CLF.α(CLF(x)) + δ)
        @constraint(model, A * u .<= b)
        optimize!(model)

        return Σ.m == 1 ? value(u) : value.(u)
    end

    return CLFQP(control)
end


export generate_ham_1d_hubbard

function generate_ham_1d_hubbard(t::Float64, U::Float64, nsite::Integer, μ::Float64)
    ham = FermionOperator()
    #斥力項
    for i in 1:nsite
        #up_index,down_indexの定義は、QC_materialを参照。
        up = up_index(i)
        down = down_index(i)
        ham += FermionOperator("$(up)^ $(down)^ $(up) $(down)", -U) #左側に生成演算子。右際に消滅演算子をもっていく過程で半交換関係が1回でマイナスをつける。
    end
    #ホッピング項
    for i in 1:nsite-1
        ham += FermionOperator("$(up_index(i+1))^ $(up_index(i))", t)
        ham += FermionOperator("$(up_index(i))^ $(up_index(i+1))", t)
        ham += FermionOperator("$(down_index(i+1))^ $(down_index(i))", t)
        ham += FermionOperator("$(down_index(i))^ $(down_index(i+1))", t)
    end

    #ケミカルポテンシャルの項
    for i in 1:nsite
        up = up_index(i)
        down = down_index(i)
        ham += FermionOperator("$(up)^  $(up) ", -μ)
        ham += FermionOperator("$(down)^ $(down)", -μ)
    end
    ham
end


module GAIO

using LinearAlgebra
using SparseArrays
using StaticArrays
# using GLFW
# using ModernGL
using GeometryBasics
using WGLMakie
using LightGraphs
using ForwardDiff
using Arpack
using Base.Threads

using GLMakie

export Box

export AbstractBoxPartition, BoxPartition, TreePartition
export dimension, depth, key_to_box, point_to_key, tree_search

export BoxSet
export boxset_empty, subdivide, subdivide!

export BoxFun

export TransferOperator
export strongly_connected_components, matrix, eigs

export BoxMap, PointDiscretizedMap
export SampledBoxMap, AdaptiveBoxMap
export boxmap

export map_boxes, map_boxes_new

export rk4, rk4_flow_map

export relative_attractor, unstable_set!, chain_recurrent_set, cover_roots, finite_time_lyapunov_exponents

export plot

# ENV["JULIA_DEBUG"] = all

include("box.jl")

abstract type AbstractBoxPartition{B <: Box} end

include("partition_regular.jl")
include("partition_tree.jl")
include("boxset.jl")
include("boxmap.jl")
include("boxfun.jl")
include("transfer_operator.jl")
include("algorithms.jl")
include("plot.jl")

end # module


#=
Copyright 2017 INSIGNEO Institute for in silico Medicine

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

   http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=#

# The anastomosis is solved as a reversed bifurcation. Refer to [bifurcations](bifurcations.html)
# documentation for further details.

#
function solveAnastomosis(v1 :: Vessel, v2 :: Vessel, v3 :: Vessel)

  #Unknowns vector
  U = [v1.u[end],
       v2.u[end],
       v3.u[ 1 ],
       sqrt(sqrt(v1.A[end])),
       sqrt(sqrt(v2.A[end])),
       sqrt(sqrt(v3.A[ 1 ]))]

  #Parameters vector
  k1 = sqrt(0.5*3*v1.gamma[end])
  k2 = sqrt(0.5*3*v2.gamma[end])
  k3 = sqrt(0.5*3*v3.gamma[ 1 ])
  k = [k1, k2, k3]

  W = calculateWstarAn(U, k)
  J = calculateJacobianAn(v1, v2, v3, U, k)
  F = calculateFofUAn(v1, v2, v3, U, k, W)

  #Newton-Raphson
  nr_toll_U = 1.e-5
  nr_toll_F = 1.e-5

  while true
    dU = J\(-F)
    # U_new = U + 0.01*dU
    U_new = U + dU

    if any(isnan(dot(F,F)))
      println(F)
      @printf "error at anastomosis with vessels %s, %s, and %s \n" v1.label v2.label v3.label
      break
    end

    u_ok = 0
    f_ok = 0
    for i in 1:length(dU)
      if abs(dU[i]) <= nr_toll_U || abs(F[i]) <= nr_toll_F
        u_ok += 1
        f_ok += 1
      end
    end

    if u_ok == length(dU) || f_ok == length(dU)
      U = U_new
      break
    else
      U = U_new
      W = calculateWstarBif(U, k)
      F = calculateFofUBif(v1, v2, v3, U, k, W)
    end
  end

  #Update vessel quantities
  v1.u[end] = U[1]
  v2.u[end] = U[2]
  v3.u[ 1 ] = U[3]

  v1.A[end] = U[4]*U[4]*U[4]*U[4]
  v2.A[end] = U[5]*U[5]*U[5]*U[5]
  v3.A[ 1 ] = U[6]*U[6]*U[6]*U[6]

  v1.Q[end] = v1.u[end]*v1.A[end]
  v2.Q[end] = v2.u[end]*v2.A[end]
  v3.Q[ 1 ] = v3.u[ 1 ]*v3.A[ 1 ]

  v1.P[end] = pressure(v1.A[end], v1.A0[end], v1.beta[end], v1.Pext)
  v2.P[end] = pressure(v2.A[end], v2.A0[end], v2.beta[end], v2.Pext)
  v3.P[ 1 ] = pressure(v3.A[ 1 ], v3.A0[ 1 ], v3.beta[ 1 ], v3.Pext)

  v1.c[end] = waveSpeed(v1.A[end], v1.gamma[end])
  v2.c[end] = waveSpeed(v2.A[end], v2.gamma[end])
  v3.c[ 1 ] = waveSpeed(v3.A[ 1 ], v3.gamma[ 1 ])

end

#
function calculateWstarAn(U :: Array, k :: Array)

  W1 = U[1] + 4*k[1]*U[4]
  W2 = U[2] + 4*k[2]*U[5]
  W3 = U[3] - 4*k[3]*U[6]

  return [W1, W2, W3]
end

#
function calculateFofUAn(v1 :: Vessel, v2 :: Vessel, v3 :: Vessel,
                           U :: Array,   k :: Array,   W :: Array)

  f1 = U[1] + 4*k[1]*U[4] - W[1]

  f2 = U[2] + 4*k[2]*U[5] - W[2]

  f3 = U[3] - 4*k[3]*U[6] - W[3]

  f4 = U[1]*(U[4]*U[4]*U[4]*U[4]) + U[2]*(U[5]*U[5]*U[5]*U[5]) - U[3]*(U[6]*U[6]*U[6]*U[6])

  f5 = v1.beta[end]*( U[4]^2/sqrt(v1.A0[end])  - 1  ) -
     ( v3.beta[ 1 ]*((U[6])^2/sqrt(v3.A0[1]) - 1) )

  f6 = v2.beta[end]*( U[5]^2/sqrt(v2.A0[end])  - 1  ) -
     ( v3.beta[ 1 ]*((U[6]^2)/sqrt(v3.A0[1]) - 1) )

  return [f1, f2, f3, f4, f5, f6]
end

#
function calculateJacobianAn(v1 :: Vessel, v2 :: Vessel, v3 :: Vessel,
                               U :: Array,   k :: Array)

  J = eye(6)

  J[1,4] =  4*k[1]
  J[2,5] =  4*k[2]
  J[3,6] = -4*k[3]

  J[4,1] =  U[4]*U[4]*U[4]*U[4]
  J[4,2] =  U[5]*U[5]*U[5]*U[5]
  J[4,3] = -U[6]*U[6]*U[6]*U[6]
  J[4,4] =  4*U[1]*(U[4]*U[4]*U[4])
  J[4,5] =  4*U[2]*(U[5]*U[5]*U[5])
  J[4,6] = -4*U[3]*(U[6]*U[6]*U[6])

  J[5,4] =  2*v1.beta[end]*U[4]/sqrt(v1.A0[end])
  J[5,5] =  0.
  J[5,6] = -2*v3.beta[ 1 ]*U[6]/sqrt(v3.A0[ 1 ])

  J[6,5] =  2*v2.beta[end]*U[5]/sqrt(v2.A0[end])
  J[6,6] = -2*v3.beta[ 1 ]*U[6]/sqrt(v3.A0[ 1 ])

  return J
end


const A = [1.0 2.0 3.0; 4.0 5.0 6.0; 7.0 8.0 9.0]

cost(M::PowerManifold, p) = -0.5 * norm(transpose(p[M, 1]) * A * p[M, 2])^2

function egrad(M::PowerManifold, X::Array)
    U = X[M, 1]
    V = X[M, 2]
    AV = A * V
    AtU = transpose(A) * U
    AR = similar(X)
    AR[:, :, 1] .= -AV * (transpose(AV) * U)
    AR[:, :, 2] .= -AtU * (transpose(AtU) * V)
    return AR
end
struct EGrad{T,TM}
    M::TM
    A::Matrix{T}
end
function (e::EGrad)(Y::Array, X::Array)
    U = X[e.M, 1]
    V = X[e.M, 2]
    AV = A * V
    AtU = transpose(A) * U
    view(Y, :, :, 1) .= -AV * (transpose(AV) * U)
    view(Y, :, :, 2) .= -AtU * (transpose(AtU) * V)
    return Y
end

rgrad(M::PowerManifold, p) = project(M, p, egrad(M, p))
struct RGrad{T,TM}
    egrad::EGrad{T,TM}
end
function RGrad(M::PowerManifold, A::Matrix{T}) where {T}
    return RGrad{T,typeof(M)}(EGrad{T,typeof(M)}(M, A))
end
function (r::RGrad)(M::PowerManifold, X, p)
    return project!(M, X, p, r.egrad(X, p))
end

function e2rHess(M::Grassmann, p, X, e_grad, e_hess)
    return project(M, p, project(M, p, e_hess) - X * (p' * e_grad))
end
function e2rhess!(M::Grassmann, Y, p, X, e_grad, e_Hess)
    project!(M, Y, p, e_Hess)
    Y .-= X * (p' * e_grad)
    return project!(M, Y, p, Y)
end

function eHess(M::AbstractManifold, X::Array{Float64,3}, H::Array{Float64,3})
    U = X[M, 1]
    V = X[M, 2]
    Udot = H[M, 1]
    Vdot = H[M, 2]
    AV = A * V
    AtU = transpose(A) * U
    AVdot = A * Vdot
    AtUdot = transpose(A) * Udot
    R = similar(X)
    #! format: off
    view(R, :, :, 1) .= -(
            AVdot * transpose(AV) * U +
            AV * transpose(AVdot) * U +
            AV * transpose(AV) * Udot
        )
    view(R, :, :, 2) .= -(
            AtUdot * transpose(AtU) * V +
            AtU * transpose(AtUdot) * V +
            AtU * transpose(AtU) * Vdot
        )
    #! format: on
    return R
end
struct EHess{T,TM}
    M::TM
    A::Matrix{T}
end
function (e::EHess)(Y, X, H)
    U = X[e.M, 1]
    V = X[e.M, 2]
    Udot = H[e.M, 1]
    Vdot = H[e.M, 2]
    AV = e.A * V
    AtU = transpose(e.A) * U
    AVdot = e.A * Vdot
    AtUdot = transpose(e.A) * Udot
    #! format: off
    view(Y, :, :, 1) .= -AVdot * transpose(AV) * U - AV * transpose(AVdot) * U - AV * transpose(AV) * Udot
    view(Y, :, :, 2) .= AtUdot * transpose(AtU) * V + AtU * transpose(AtUdot) * V + AtU * transpose(AtU) * Vdot
    #! format: on
    return Y
end

function rhess(M::PowerManifold, p, X)
    eG = egrad(M, p)
    eH = eHess(M, p, X)
    Ha = similar(p)
    for i in 1:2
        e2rhess!(
            M.manifold,
            view(Ha, :, :, i),
            view(p, :, :, i),
            view(X, :, :, i),
            view(eG, :, :, i),
            view(eH, :, :, i),
        )
    end
    return Ha
end
struct RHess{T,TM}
    e_grad!::EGrad{T,TM}
    e_hess!::EHess{T,TM}
    G::Array{T,3}
    H::Array{T,3}
end
function RHess(M::AbstractManifold, A::Matrix{T}, p) where {T}
    return RHess{T,typeof(M)}(
        EGrad(M, A), EHess(M, A), zeros(T, size(A, 1), p, 2), zeros(T, size(A, 1), p, 2)
    )
end
function (r::RHess)(M::PowerManifold, Y, p, X)
    r.e_grad!(r.G, p)
    r.e_hess!(r.H, p, X)
    for i in 1:2
        e2rhess!(
            M.manifold,
            view(Y, :, :, i),
            view(p, :, :, i),
            view(X, :, :, i),
            view(r.G, :, :, i),
            view(r.H, :, :, i),
        )
    end
    return Y
end


const MOUNTAIN_CAR_STATE_MINS = [-1.2, -0.07]
const MOUNTAIN_CAR_STATE_MAXS = [ 0.6,  0.07]
const MOUNTAIN_CAR_ACCELS = [-1.0, 0.0, 1.0]
const MOUNTAIN_CAR_START_STATE = [-0.5, 0.0]

@with_kw struct MountainCar
    γ::Float64 = 1.0
end

n_actions(mdp::MountainCar) = length(MOUNTAIN_CAR_ACCELS)
discount(mdp::MountainCar) = mdp.γ
ordered_actions(mdp::MountainCar) = collect(1:length(MOUNTAIN_CAR_ACCELS))

function generate_start_state(mdp::MountainCar)
    s = MOUNTAIN_CAR_STATE_MINS[1] + rand() * (MOUNTAIN_CAR_STATE_MAXS[1] - MOUNTAIN_CAR_STATE_MINS[1])
    v = MOUNTAIN_CAR_STATE_MINS[2] + rand() * (MOUNTAIN_CAR_STATE_MAXS[2] - MOUNTAIN_CAR_STATE_MINS[2])
    return [s,v]
end

# NOTE: Mountain car is deterministic, so we don't get a transition
#       function that produces a probability distribution.
function mountain_car_transition(s::Vector{Float64}, a::Int)
    x, v = s[1], s[2]
    if x < 0.6
        accel = MOUNTAIN_CAR_ACCELS[a]
        v = v + accel*0.001 + cos(3x)*-0.0025
        x = x + v

        x = clamp(x, MOUNTAIN_CAR_STATE_MINS[1], MOUNTAIN_CAR_STATE_MAXS[1])
        v = clamp(v, MOUNTAIN_CAR_STATE_MINS[2], MOUNTAIN_CAR_STATE_MAXS[2])
    end
    return [x,v]
end

function reward(mdp::MountainCar, s::Vector{Float64}, a::Int)
    if s[1] >= 0.6
        return 0.0
    else
        return -1.0
    end
end

function MDP(mdp::MountainCar; γ::Float64=mdp.γ)
    return MDP(
            γ,
            nothing, # no ordered states
            ordered_actions(mdp),
            (s,a, s′=nothing) -> begin
                # Technically wrong, but its probably fine
                S′ = MvNormal(mountain_car_transition(s,a), 1e-6I)
                if s′ == nothing
                    return S′
                end
                return pdf(S′, s′)
            end,
            (s,a) -> reward(mdp, s, a),
            (s, a)->begin
                s′ = mountain_car_transition(s,a)
                r = reward(mdp, s, a)
                return (s′, r)
            end
        )
end


# a simple LP example
#
#   maximize x + y
#
#   s.t. 50 x + 24 y <= 2400
#        30 x + 33 y <= 2100
#        x >= 45, y >= 5
#
#   solution: x = 45, y = 6.25, objv = 51.25

using CPLEX, Base.Test

@testset "LP 01" begin
    env = CPLEX.Env()

    # method = getparam(env, "Method")
    # println("method = $method")

    model = CPLEX.Model(env, "lp_01")
    CPLEX.set_sense!(model, :Max)

    # add variables
    CPLEX.add_var!(model, 1.0, 45., Inf)  # x
    CPLEX.add_var!(model, 1.0,  5., Inf)  # y

    # add constraints
    CPLEX.add_constr!(model, [50., 24.], '<', 2400.)
    CPLEX.add_constr!(model, [30., 33.], '<', 2100.)

    # perform optimization
    CPLEX.optimize!(model)

    sol = CPLEX.get_solution(model)
    @test sol[1] ≈ 45
    @test sol[2] ≈ 6.25

    @test CPLEX.get_objval(model) ≈ 51.25

    gc()  # test finalizers
end


#
# This finds Normalized, Synodic, Cartesian states
# for Halo Orbits near L1 & L2 for a variety of
# normalized CR3BP systems.
#
# Currently, nondimensional mass parameters are
# used which are close to those in the Sun-Earth,
# Earth-Moon, Sun-Mars, and Sun-Jupiter systems.
#

using DrWatson; @quickactivate "CR3BP-Manifold-Research"

using CSV
using DataFrames
using GeneralAstrodynamics

massparameters = let
    planets = (Venus, Earth, Mars, Jupiter, Saturn, Uranus, Neptune)
    planet_params = map(
        planet -> let μ₁=massparameter(Sun), μ₂=massparameter(planet)
            min(μ₁, μ₂) / (μ₁+μ₂)
        end, planets
    )
    earth_moon  = let μ₁=massparameter(Earth), μ₂=massparameter(Moon)
        min(μ₁, μ₂) / (μ₁+μ₂)
    end
    vcat(planet_params[1:2]..., earth_moon, planet_params[3:end]...)
end

names = [
    "Sun-Venus", "Sun-Earth",
    "Earth-Moon", "Sun-Mars",
    "Sun-Jupiter", "Sun-Saturn",
    "Sun-Uranus", "Sun-Neptune"
]

amplitudes = 0.0:1e-6:0.01

@assert length(names) == length(massparameters) "Invalid configuration!"

Threads.@threads for i = 1:length(names)
    μ = massparameters[i]
    name = names[i]
    orbits = DataFrame(
        MassParameter=Float64[], LagrangePoint=Int[], ZAmplitude=Float64[], JacobiConstant=Float64[], Period=Float64[],
        Rx=Float64[], Ry=Float64[], Rz=Float64[], Vx=Float64[], Vy=Float64[], Vz=Float64[])
    for L ∈ (1,2)
        fail_count = 0
        for Az ∈ amplitudes
            fail_count < 10 || break
            r, v, T = halo(μ; Az = Az, L = L, max_iter=20, nan_on_fail=true, disable_warnings=true);
            if T ≥ one(T) && all(el->!isnan(el), r) && all(el->!isnan(el), v) && isperiodic(CR3BPOrbit(CartesianState(r,v),μ), T)
                push!(orbits, (μ, L, Az, jacobi_constant(r, v, μ), T, r..., v...))
            else
                fail_count += 1
            end
        end
    end
    @info "Finished processing $name Halos!"
    CSV.write(datadir("exp_pro", "halos", lowercase(string(name,"-Halos.csv"))), orbits)
end

open(datadir("processed", "halos", "README.md"), "w") do io
    write(io,
    """
    # How are the Halo CSV files formatted?
    * The columns from left to right are:
      1. Nondimensional mass parameter
      2. Lagrange point (1 or 2)
      3. Nondimensional Z amplitude
      4. Jacobi Constant
      5. Nondimensional orbital period
      6. Nondimensional X
      7. Nondimensional Y
      8. Nondimensional Z
      9. Nondimensional Ẋ
      10. Nondimensional Ẏ
      11. Nondimensional Ż
    """)
end


function hist2spec(h::Histogram{<:Any, N}) where N
    ptype = [:BarPlot, :Heatmap, :Volume][N]
    f(edges) = edges[1:end-1] .+ diff(edges)./2
    kwargs = N == 1 ? (; width = step(h.edges[1])) : NamedTuple()
    return mapping(map(f, h.edges)..., Float64.(h.weights)) * visual(ptype; kwargs...)
end

to_weights(v) = weights(v)
to_weights(v::AbstractWeights) = v

function trim(range, min, max, closed)
    if closed == :left
        i1 = searchsortedlast(range, min)
        i2 = searchsortedfirst(range, nextfloat(max))
    else
        i1 = searchsortedlast(range, prevfloat(min))
        i2 = searchsortedfirst(range, max)
    end
    return range[i1:i2]
end

function compute_edges(data, extrema, bins::Tuple{Vararg{Integer}}, closed)
    ranges = map(extrema, bins) do (min, max), n
        histrange(min, max, n, closed)
    end
    # trim axis
    map(ranges, data) do range, d
        trim(range, Base.extrema(d)..., closed)
    end
end
compute_edges(data, extrema, bins::Tuple{Vararg{AbstractArray}}, closed) = bins

function _histogram(data...; bins = sturges(length(data[1])), wts = automatic,
    normalization = :none, extrema = map(extrema, data), closed = :left)

    bins_tuple = bins isa Tuple ? bins : map(_ -> bins, data)
    edges = compute_edges(data, extrema, bins_tuple, closed)
    weights = wts === automatic ? () : (to_weights(wts),)
    h = fit(Histogram, data, weights..., edges)
    hn = normalize(h, mode = normalization)
    return hist2spec(hn)
end

"""
    histogram(data...; bins = automatic, wts = automatic, normalization = :none)

Plot a histogram of `values`. `bins` can be an `Int` to create that
number of equal-width bins over the range of `values`.
Alternatively, it can be a sorted iterable of bin edges. The histogram
can be normalized by setting `normalization`. Possible values are:
*  `:pdf`: Normalize by sum of weights and bin sizes. Resulting histogram
   has norm 1 and represents a PDF.
* `:density`: Normalize by bin sizes only. Resulting histogram represents
   count density of input and does not have norm 1.
* `:probability`: Normalize by sum of weights only. Resulting histogram
   represents the fraction of probability mass for each bin and does not have
   norm 1.
*  `:none`: Do not normalize.
Weighted data is supported via the keyword `wts`.
"""
const histogram = Analysis(_histogram)

function global_options(::typeof(histogram), d::AlgebraicList)
    combine(extr1, extr2) = map(extr1, extr2) do (min1, max1), (min2, max2)
        min(min1, min2), max(max1, max2)
    end
    extr = mapfoldl(combine, d) do spec
        map(extrema, positional(spec.mapping.value))
    end
    return (extrema = extr,)
end


using Images, WNNMDenoise, TestImages
using Random

img = imresize(Float64.(load("house.png")) * 255, ratio=0.5);
img = Float32.(load("house.png")) * 255;

clean_img = img;
σₙ = 40;
Random.seed!(0)
noisy_img = clean_img .+ σₙ .* randn(Float32, size(clean_img))

assess_psnr(noisy_img, img, 255)

out = copy(noisy_img)
f = WNNM(σₙ)
@time f(out, noisy_img; clean_img=clean_img)


using Test
using CollectiveSpins
import LinearAlgebra

#all functions should have immediate rejection if k exceeds corner of BZ

@testset "collective_modes" begin

cs = CollectiveSpins
la = LinearAlgebra
k0 = 2pi
tol = 1e-05
tol_1D = 1e-02
tol_2D = 1e-01
N = 600

a_list = collect(1:5)*0.1 .- 0.001
for a in a_list
    #assert that lighcone exists a < lambda0/2
    @test CollectiveSpins.collective_modes.Gamma_k_chain(pi/a, a, [1,0,0]) == 0 #arbitrary, unnormalized initializations throughout
    @test CollectiveSpins.collective_modes.Gamma_k_chain(pi/a, a, [0,1,0.5]) == 0
end

a_list = collect(1:5)*1e+08
for a in a_list
    #assert lim a--> Inf, Omega_k --> 0
    @test abs(cs.collective_modes.Omega_k_chain(0.2*a, a, [1,0,0])) < tol
    @test abs(cs.collective_modes.Omega_k_chain(0.2*a, a, [0,1,0.5])) < tol
    #assert lim a--> Inf, Delta_k --> gamma0 (1)
    @test abs(cs.collective_modes.Gamma_k_chain(0, a, [1,0,0]) - 1) < tol
    @test abs(cs.collective_modes.Gamma_k_chain(0, a, [0,1,0.5]) - 1) < tol
end

a_list = [0.07, 0.11, 0.2, 0.51, 1.7]
k_list = [0.4pi, 0.1pi, 0pi, 0.9pi, 0.3pi]
k_list = k_list ./ a_list
Omega_x_calculated = [-6.24945, -12.2009, -2.18103, -0.11888, 0.0165684]
Omega_yz_calculated = [3.42368, 4.99508, 1.28359, -0.494113, 0.0327376]
Gamma_x_calculated = [0, 5.40947, 3.75, 0.325667, 1.00791]
Gamma_yz_calculated = [0, 4.11345, 1.875, 1.30775, 0.819573]
for index = 1:length(a_list)
    #assert that selection of arbitrary a and k values match simplified, semi-manual calculation
    @test abs(cs.collective_modes.Omega_k_chain(k_list[index], a_list[index], [1,0,0]) - Omega_x_calculated[index]) < tol
    @test abs(cs.collective_modes.Omega_k_chain(k_list[index], a_list[index], [0,5,22]) - Omega_yz_calculated[index]) < tol
    @test abs(cs.collective_modes.Gamma_k_chain(k_list[index], a_list[index], [55,0,0]) - Gamma_x_calculated[index]) < tol
    @test abs(cs.collective_modes.Gamma_k_chain(k_list[index], a_list[index], [0,87,2.3+2im]) - Gamma_yz_calculated[index]) < tol
end

#real-space Dyadic Green's function
function Gr(r_vec::Array{<:Number, 1})
    r = la.norm(r_vec)
    return exp(im*k0*r)/(4pi*k0^2*r^3) * ( la.UniformScaling(k0^2*r^2 + im*k0*r - 1) + (-k0^2*r^2 - 3*im*k0*r + 3)*r_vec*la.transpose(r_vec)/r^2 )
end

#1D k-space Green's function by summation of real-space
function Gk_sumR_1D(k, a, polarization::Array{<:Number, 1})
    polarization = la.normalize(polarization)
    Omegak = 0
    Gammak = 0
    for n = -N:N
        if !(n == 0)
            r_vec = zeros(3)
            r_vec[1] = n*a
            Omegak += la.conj(la.transpose(polarization))*real(Gr(r_vec) * exp(Complex(-im*k*r_vec[1])))*polarization
            Gammak += la.conj(la.transpose(polarization))*imag(Gr(r_vec) * exp(Complex(-im*k*r_vec[1])))*polarization
        end
    end
    if  abs(k) > k0
        return -3pi*real(Omegak)/k0, 0
    end
    return -3pi*real(Omegak)/k0, 6pi*real(Gammak)/k0 + 1
end

a_list = [0.8, 0.25, 1.2, 0.3, 0.5, 1.7, 0.6]
k_list = [0.1pi/0.8, -0.1pi/0.25, 0.3pi/1.2, -0.1pi/0.3, -0.2pi/0.5, 0.3pi/1.7, 0.2pi/0.6]
polarization_list = [[1, 1, -1], [-1, -im, -1], [1, 0, im], [0, -im, 5], [1, 0, 0], [-1, im, 1], [0, 1, -2], [-4, -1, 2]]
for index = 1:length(a_list)
    #assert that 1D k-space calculations with G(k) match G_k(r) definition. Convergence of real-space sum version is unstable by nature
    #thus only moderate a values are taken and tolerance is quite large (1e-02). This is especially true around the edge of the lightcone.
    Omegak_real, Gammak_real = Gk_sumR_1D(k_list[index], a_list[index], polarization_list[index])
    @test abs(cs.collective_modes.Omega_k_chain(k_list[index], a_list[index], polarization_list[index]) -
    Omegak_real) < tol_1D
    @test abs(cs.collective_modes.Gamma_k_chain(k_list[index], a_list[index], polarization_list[index]) -
    Gammak_real) < tol_1D
end

#2D k-space Green's function by summation of real-space
function Gk_sumR(k_vec::Array{<:Number, 1}, a_vec1::Array{<:Number, 1}, a_vec2::Array{<:Number, 1}, polarization::Array{<:Number, 1})
    polarization = la.normalize(polarization)
    Omegak = 0
    Gammak = 0
    for n1 = -N:N
        for n2 = -N:N
            if !(n1 == 0) || !(n2 == 0)
                r_vec = zeros(3)
                r_vec[2:3] = n1*a_vec1 + n2*a_vec2
                Omegak += la.conj(la.transpose(polarization))*real(Gr(r_vec) * exp(Complex(-im*la.dot(k_vec, r_vec[2:3]))))*polarization
                Gammak += la.conj(la.transpose(polarization))*imag(Gr(r_vec) * exp(Complex(-im*la.dot(k_vec, r_vec[2:3]))))*polarization
            end
        end
    end
    if la.norm(k_vec) > k0
        return -3pi*real(Omegak)/k0, 0
    end
    return -3pi*real(Omegak)/k0, 6pi*real(Gammak)/k0 + 1
end

a_vec1_list = [[0.8, 0], [0.25, 0], [1.2, 0], [0.3, 0], [0.25/2, -0.25*sqrt(3)/2], [0.5/2, -0.5/2*sqrt(3)/2], [1.7, 0.25], [0.2, 0.6]]
a_vec2_list = [[0, 0.8], [0, 0.25], [0, 0.4], [0, 0.7], [0.25/2, 0.25*sqrt(3)/2], [0.5/2, 0.5*sqrt(3)/2], [0.25, 0.9], [0.7, 0.3]]
k_vec_list = [[-0.1pi/0.8, 0.3pi/0.8], [-0.1pi/0.25, -0.1pi/0.25], [0.6pi/1.2, 0.99pi/0.4], [0.1pi/0.3, -0.2pi/0.7], [0.2pi/0.25, 0.5pi/0.25], [1pi/0.5, 1pi/0.5], [0, 0.2pi/0.8], [-0.1pi/0.5, -0.2pi/0.4]]
polarization_list = [[1, 1, 1], [-1, im, 1], [-1, 0, -im], [0, im, 5], [1, 0, 0], [-1, im, 1], [0, 1, 2], [4, -1, -2]]
for index = 1:length(a_vec1_list)
    #assert that 2D k-space calculations with G(k) match G_k(r) definition. Convergence of real-space sum version is unstable by nature
    #thus only moderate a values are taken and tolerance is quite large (1e-01). This is especially true around the edge of the lightcone.
    Omegak_real, Gammak_real = Gk_sumR(k_vec_list[index], a_vec1_list[index], a_vec2_list[index], polarization_list[index])
    @test abs(cs.collective_modes.Omega_k_2D(k_vec_list[index], a_vec1_list[index], a_vec2_list[index], polarization_list[index]) -
    Omegak_real) < tol_2D
    @test abs(cs.collective_modes.Gamma_k_2D(k_vec_list[index], a_vec1_list[index], a_vec2_list[index], polarization_list[index]) -
    Gammak_real) < tol_2D
end

end #module


function CIR_smoothing_logscale_internals(δ, γ, σ, λ, data; silence = false)
    β = γ/σ^2
    α = δ/2#Alternative parametrisation

    if !silence
        println("Filtering")
    end
    Λ_of_t, logwms_of_t, θ_of_t = filter_CIR_logweights(δ, γ, σ, λ, data; silence = silence)
    if !silence
        println("Cost to go")
    end
    Λ_tilde_prime_of_t, logwms_tilde_of_t, θ_tilde_prime_of_t = compute_all_log_cost_to_go_functions_CIR(δ, γ, σ, λ, data; silence = silence)

    return merge_filtering_and_cost_to_go_logweights_CIR(Λ_of_t, logwms_of_t, θ_of_t, Λ_tilde_prime_of_t, logwms_tilde_of_t, θ_tilde_prime_of_t, data, α, β)

end

function filter_CIR_logweights(δ, γ, σ, λ, data; silence = true)
    Λ_of_t, wms_of_t, θ_of_t = filter_CIR(δ, γ, σ, λ, data; silence = silence)
    logwms_of_t = wms_of_t |> convert_weights_to_logweights
    return Λ_of_t, logwms_of_t, θ_of_t
end

compute_all_log_cost_to_go_functions_CIR(δ, γ, σ, λ, data; silence = false) = compute_all_log_cost_to_go_functions_CIR_pruning(δ, γ, σ, λ, data, (x,y) -> (x,y); silence = silence)

function compute_all_log_cost_to_go_functions_CIR_pruning(δ, γ, σ, λ, data, pruning_function::Function; silence = false)
    times = data |> keys |> collect |> sort
    reversed_times = reverse(times)
    α = δ/2#Alternative parametrisation
    θ0_CIR = γ/σ^2
    Λ_tilde_prime_of_t = Dict()
    logwms_tilde_of_t = Dict()
    θ_tilde_prime_of_t = Dict()

    yT = data[times[end]]

    Λ_tilde_prime_kp2 = [0]
    Λ_tilde_kp1 = t_CIR(yT, Λ_tilde_prime_kp2)
    θ_tilde_prime_kp2 = θ0_CIR
    logwms_tilde_kp2 = [0.]

    for k in 2:length(reversed_times)
        if !silence
            println("(Cost to go) Step index: $k")
            println("Number of components: $(length(Λ_tilde_kp1))")
        end
        # Change of notation for clarity
        prev_t = reversed_times[k-1]
        t = reversed_times[k]
        Δk = prev_t - t
        ykp1 = data[prev_t]

        # New weight computation
        θ_tilde_kp1 = θ_tilde_k_from_θ_tilde_prime_kp1(ykp1, θ_tilde_prime_kp2)

        logwms_tilde_kp1, Λ_tilde_prime_kp1 = logwms_tilde_kp1_from_logwms_tilde_kp2_pruning(logwms_tilde_kp2, Λ_tilde_prime_kp2, θ_tilde_kp1, θ_tilde_prime_kp2, ykp1, Δk, α, γ, σ, λ; pruning_function = pruning_function, return_indices = true)

        #Storage of the results
        θ_tilde_prime_kp1 = θ_tilde_prime_k_from_θ_tilde_k_CIR(θ_tilde_kp1, Δk, γ, σ)

        yk = data[t]
        Λ_tilde_kp = Λ_tilde_k_from_Λ_tilde_prime_kp1_CIR(yk, Λ_tilde_prime_kp1) #Not stored, but better for consistency of notations.

        θ_tilde_prime_of_t[prev_t] = θ_tilde_prime_kp1
        Λ_tilde_prime_of_t[prev_t] = Λ_tilde_prime_kp1
        logwms_tilde_of_t[prev_t] = logwms_tilde_kp1

        #Preparation of next iteration
        θ_tilde_prime_kp2 = θ_tilde_prime_kp1
        logwms_tilde_kp2 = logwms_tilde_kp1
        Λ_tilde_prime_kp2 = Λ_tilde_prime_kp1
        Λ_tilde_kp1 = Λ_tilde_kp
    end

    return Λ_tilde_prime_of_t, logwms_tilde_of_t, θ_tilde_prime_of_t
end

function θ_tilde_k_from_θ_tilde_prime_kp1(yk, θ_tilde_prime_kp1; λ = 1)
    return T_CIR(yk, θ_tilde_prime_kp1; λ = λ)
end

function logwms_tilde_kp1_from_logwms_tilde_kp2_pruning(logwms_tilde_kp2, Λ_tilde_prime_kp2, θ_tilde_kp1, θ_tilde_prime_kp2, ykp1, Δk, α, γ, σ, λ; pruning_function::Function, return_indices = false)

    updated_logwms_tilde_kp2 = update_logweights_cost_to_go_CIR(logwms_tilde_kp2, Λ_tilde_prime_kp2, θ_tilde_prime_kp2, α, ykp1, λ)

    Λ_tilde_prime_kp2_pruned, updated_logwms_tilde_kp2_pruned = pruning_function(Λ_tilde_prime_kp2, updated_logwms_tilde_kp2)

    logwms_tilde_kp1 = predict_logweights_cost_to_go_CIR(updated_logwms_tilde_kp2_pruned, Λ_tilde_prime_kp2_pruned, θ_tilde_kp1, θ_tilde_prime_kp2, ykp1, Δk, α, γ, σ, λ)

    if return_indices
        return logwms_tilde_kp1, 0:maximum(t_CIR(ykp1, maximum(Λ_tilde_prime_kp2_pruned)))
    else
        return logwms_tilde_kp1
    end
end

function θ_tilde_prime_k_from_θ_tilde_k_CIR(θ_tilde_k, Δt, γ, σ)
    return θ_prime_from_θ_CIR(θ_tilde_k, Δt, γ, σ)
end

function Λ_tilde_k_from_Λ_tilde_prime_kp1_CIR(yk, Λ_tilde_prime_kp1)
    return t_CIR(yk, Λ_tilde_prime_kp1)
end

function update_logweights_cost_to_go_CIR(logwms_tilde_kp2, Λ_tilde_prime_kp2, θ_tilde_prime_kp2, α, ykp1, λ)
    return update_logweights_cost_to_go(logwms_tilde_kp2, Λ_tilde_prime_kp2, θ_tilde_prime_kp2, ykp1, (α = α, λ = λ), (m, θ, y, params) -> logμπh(m, θ, params[:α], y; λ = params[:λ]))
end

function predict_logweights_cost_to_go_CIR(updated_logwms_tilde_kp2, Λ_tilde_prime_kp2, θ_tilde_kp1, θ_tilde_prime_kp2, ykp1, Δk, α, γ, σ, λ)
    logwms_tilde_kp1 = fill(-Inf, maximum(t_CIR(ykp1, maximum(Λ_tilde_prime_kp2)))+1)

    p = γ/σ^2*1/(θ_tilde_kp1*exp(2*γ*Δk) + γ/σ^2 - θ_tilde_kp1)

    for k in eachindex(Λ_tilde_prime_kp2)
        n = Λ_tilde_prime_kp2[k]
        t_ykp1_n = t_CIR(ykp1, n)

        for m in 0:t_ykp1_n
            idx = m+1
            # Careful, logwms_tilde_kp2[k] has index k because we assume that the weights and the indices are in the same order.
            logwms_tilde_kp1[idx] = logaddexp(logwms_tilde_kp1[idx], updated_logwms_tilde_kp2[k] + logpmn_CIR(t_ykp1_n, m, p))
        end
    end
    return logwms_tilde_kp1
end

function Λ_tilde_prime_k_from_Λ_tilde_k_CIR(Λ_tilde_k)
    return 0:maximum(Λ_tilde_k)
end

function logpmn_CIR(m, n, p)
    return logpdf(Binomial(m, p), n)
end

function logpmn_CIR(m, n, Δt, θ, σ, γ)
    p = γ/σ^2*1/(θ*exp(2*γ*Δt) + γ/σ^2 - θ)
    return logpmn_CIR(m, n, p)
end

function logCmn_CIR(m, n, θback_k, θ_k, α, β)
    return lgamma_local(α) - lgamma_local(α+m) - lgamma_local(α+n) - α*log(β) + (α+m)*log(θback_k) + (α+n)*log(θ_k) + lgamma_local(α+m+n) - (α+m+n) * log(θback_k + θ_k-β)
end

function merge_filtering_and_cost_to_go_logweights_CIR(Λ_of_t, logwms_of_t, θ_of_t, Λ_tilde_prime_of_t, logwms_tilde_of_t, θ_tilde_prime_of_t, data, α, β)
    times = Λ_of_t |> keys |> collect |> sort

    Λ_of_t_smooth = Dict()
    wms_of_t_smooth = Dict()
    θ_of_t_smooth = Dict()

    log_Λ_weights = Array{Float64,1}(undef, 2*sum(sum(values(data))))

    for k in 1:(length(times)-1)
        fill!(log_Λ_weights, -Inf)

        Λk = Λ_of_t[times[k]]
        logwk = logwms_of_t[times[k]]
        Λ_tilde_prime_kp1 = Λ_tilde_prime_of_t[times[k+1]]
        logw_tilde_kp1 = logwms_tilde_of_t[times[k+1]]
        for i in eachindex(Λk)
            n = Λk[i]
            for j in eachindex(Λ_tilde_prime_kp1)
                m = Λ_tilde_prime_kp1[j]
                log_Λ_weights[d_CIR(m, n)] = logaddexp(log_Λ_weights[d_CIR(m, n)], logw_tilde_kp1[j] + logwk[i] + logCmn_CIR(m, n, θ_tilde_prime_of_t[times[k+1]], θ_of_t[times[k]], α, β))
            end
        end
        Λ_weights = exp.(log_Λ_weights .- logsumexp(log_Λ_weights))

        Λ_of_t_smooth[times[k]] = [n for n in eachindex(Λ_weights) if Λ_weights[n] > 0.]
        wms_of_t_smooth[times[k]] = Λ_weights[Λ_weights .> 0.]
        θ_of_t_smooth[times[k]] = e_CIR(θ_tilde_prime_of_t[times[k+1]], θ_of_t[times[k]], β)
    end

        #The last smoothing distribution is a filtering distribution
        Λ_of_t_smooth[times[end]] = Λ_of_t[times[end]]
        wms_of_t_smooth[times[end]] = exp.(logwms_of_t[times[end]])
        θ_of_t_smooth[times[end]] = θ_of_t[times[end]]

    return Λ_of_t_smooth, wms_of_t_smooth, θ_of_t_smooth
end


using Test
using Random: randn, rand, randn!
using AutoGrad: @gcheck, @diff, Param, Tape, value
using Knet.Ops20: elu, relu, selu, sigm, dropout, bmm, conv4, conv4w, conv4x, deconv4, mat, pool, poolx, unpool
using Knet.Ops20: batchnorm, bnmoments, bnparams, softmax, logsoftmax, logsumexp, accuracy, nll, bce, logistic, RNN

# Allow the user to override this:
if !isdefined(Main, :testparam)
    testparam(x...) = Param(randn(x...))
end

@testset "ops20" begin

    x = testparam(8,8,2,3)
    @testset "activation" begin
        @test @gcheck elu.(x)
        @test @gcheck relu.(x)
        @test @gcheck selu.(x)
        @test @gcheck sigm.(x)
        @test @gcheck dropout(x, 0.5; seed=1, drop=true)
    end

    @testset "bmm" begin
        b = testparam(8,6,2,3)
        pd(x) = permutedims(x,(2,1,3,4))
        @test @gcheck bmm(x,b)
        @test @gcheck bmm(x,pd(b),transB=true)
        @test @gcheck bmm(pd(x),b,transA=true)
        @test @gcheck bmm(pd(x),pd(b),transA=true,transB=true)
        @test (bmm(x,b) ≈ bmm(x,pd(b),transB=true) ≈ bmm(pd(x),b,transA=true) ≈ bmm(pd(x),pd(b),transA=true,transB=true))
    end

    function convtest(w,x;o...)
        c = conv4(w,x;o...)
        y = testparam(size(c)...)
        ((@gcheck conv4(w,x;o...)) &&
         (@gcheck conv4w(w,x,y;o...)) &&
         (@gcheck conv4x(w,x,y;o...)) &&
         (@gcheck deconv4(w,y;o...)))
    end

    @testset "conv" begin
        w = testparam(3,3,2,4)
        @test convtest(w, x; padding=0, stride=1, dilation=1, mode=0, alpha=1, group=1)
        @test convtest(w, x; padding=1, stride=1, dilation=1, mode=0, alpha=1, group=1)
        @test convtest(w, x; padding=0, stride=2, dilation=1, mode=0, alpha=1, group=1)
        @test convtest(w, x; padding=0, stride=1, dilation=2, mode=0, alpha=1, group=1)
        @test convtest(w, x; padding=0, stride=1, dilation=1, mode=1, alpha=1, group=1)
        @test convtest(w, x; padding=0, stride=1, dilation=1, mode=0, alpha=2, group=1)
        w1 = testparam(3,3,1,4)
        if value(w) isa Array
            # TODO: Grouped convolutions not yet implemented in NNlib, see https://github.com/JuliaGPU/CuArrays.jl/pull/523
            @test_skip convtest(w1, x; padding=0, stride=1, dilation=1, mode=0, alpha=1, group=2)
        else
            @test convtest(w1, x; padding=0, stride=1, dilation=1, mode=0, alpha=1, group=2)
        end
    end

    function pooltest(x; o...)
        y = Param(pool(x; o...))
        dy = testparam(size(y)...)
        ((@gcheck pool(x; o...)) &&
         # @gcheck poolx(x,y,dy) misleads because NNlib gets confused when x&y don't change together.
         (@gcheck poolx(value(x), value(y), dy; o...)) &&
         (@gcheck unpool(x; o...)) &&
         true)
    end

    @testset "pool" begin
        @test pooltest(x; window=2, stride=2, padding=0, mode=0, maxpoolingNanOpt=1, alpha=1)
        @test pooltest(x; window=3, stride=2, padding=0, mode=0, maxpoolingNanOpt=1, alpha=1)
        @test pooltest(x; window=2, stride=3, padding=0, mode=0, maxpoolingNanOpt=1, alpha=1)
        if value(x) isa Array
            # Pool padding is buggy: https://github.com/FluxML/NNlib.jl/issues/229
            @test_skip pooltest(x; window=2, stride=2, padding=1, mode=0, maxpoolingNanOpt=1, alpha=1)
        else
            @test pooltest(x; window=2, stride=2, padding=1, mode=0, maxpoolingNanOpt=1, alpha=1)
        end
        @test pooltest(x; window=2, stride=2, padding=0, mode=1, maxpoolingNanOpt=1, alpha=1)
        @test pooltest(x; window=2, stride=2, padding=0, mode=2, maxpoolingNanOpt=1, alpha=1)
        @test pooltest(x; window=2, stride=2, padding=0, mode=3, maxpoolingNanOpt=1, alpha=1)
        @test pooltest(x; window=2, stride=2, padding=0, mode=0, maxpoolingNanOpt=0, alpha=1)
        @test pooltest(x; window=2, stride=2, padding=0, mode=0, maxpoolingNanOpt=1, alpha=2)
    end

    @testset "mat" begin
        @test @gcheck mat(x)
        @test size(mat(x)) == (128,3)
        @test size(mat(x,dims=4)) == (384,1)
        @test size(mat(x,dims=3)) == (128,3)
        @test size(mat(x,dims=2)) == (64,6)
        @test size(mat(x,dims=1)) == (8,48)
        @test size(mat(x,dims=0)) == (1,384)
    end

    @testset "batchnorm" begin
        p = testparam(2*size(x,3))
        m = bnmoments()
        @test (@diff sum(batchnorm(x,m,p))) isa Tape
        @test @gcheck batchnorm(x,m,p; training=false)
    end

    function softtest(f,x)
        ((@gcheck f(x)) &&
         (@gcheck f(x,dims=1)) &&
         (@gcheck f(x,dims=2)))
    end

    @testset "softmax/loss" begin
        scores = testparam(10,100)
        @test softtest(softmax, scores)
        @test softtest(logsoftmax, scores)
        @test softtest(logsumexp, scores)

        labels1 = rand(1:10,100)
        labels2 = rand(1:100,10)
        @test @gcheck nll(scores, labels1, dims=1)
        @test @gcheck nll(scores, labels2, dims=2)

        value(scores)[1,:] .+= 1000
        @test accuracy(scores, labels1, dims=1) == sum(labels1 .== 1)/length(labels1)
        value(scores)[:,1] .+= 10000
        @test accuracy(scores, labels2, dims=2) == sum(labels2 .== 1)/length(labels2)

        scores = testparam(100)
        labels = rand(Bool,100)
        xlabels = labels .* 2 .- 1
        @test @gcheck bce(scores, labels)
        @test @gcheck logistic(scores, xlabels)
        @test bce(scores, labels) == logistic(scores, xlabels)
    end

    function rnntest(;ndims=1, batchSizes=nothing, o...)
        x = testparam(4:(4+ndims-1)...)
        r = RNN(4,4;atype=typeof(value(x)),o...) # need H==X for skipInput test
        @gcheck r(x; batchSizes=batchSizes)
    end

    @testset "rnn" begin
        @test rnntest(rnnType=:lstm, numLayers=1, bidirectional=false, skipInput=false, ndims=3, batchSizes=nothing)
        @test rnntest(rnnType=:gru,  numLayers=1, bidirectional=false, skipInput=false, ndims=3, batchSizes=nothing)
        @test rnntest(rnnType=:relu, numLayers=1, bidirectional=false, skipInput=false, ndims=3, batchSizes=nothing)
        @test rnntest(rnnType=:tanh, numLayers=1, bidirectional=false, skipInput=false, ndims=3, batchSizes=nothing)
        @test rnntest(rnnType=:lstm, numLayers=2, bidirectional=false, skipInput=false, ndims=3, batchSizes=nothing)
        @test rnntest(rnnType=:lstm, numLayers=1, bidirectional=true,  skipInput=false, ndims=3, batchSizes=nothing)
        @test rnntest(rnnType=:lstm, numLayers=1, bidirectional=false, skipInput=true,  ndims=3, batchSizes=nothing)
        @test rnntest(rnnType=:lstm, numLayers=1, bidirectional=false, skipInput=false, ndims=2, batchSizes=nothing)
        @test rnntest(rnnType=:lstm, numLayers=1, bidirectional=false, skipInput=false, ndims=1, batchSizes=nothing)
        if value(x) isa Array
            # TODO: batchSizes not implemented for CPU yet
            @test_skip rnntest(rnnType=:lstm, numLayers=1, bidirectional=false, skipInput=false, ndims=2, batchSizes=[2,2,1])
        else
            @test rnntest(rnnType=:lstm, numLayers=1, bidirectional=false, skipInput=false, ndims=2, batchSizes=[2,2,1])
        end
    end
end


using KitBase
import KitBase.BSON
cd(@__DIR__)

# save
set = Setup(space = "1d1f1v", boundary = "fix", maxTime = 0.1)
ps = PSpace1D(0, 1, 100, 1)
vs = VSpace1D(-5, 5, 28)
gas = Gas(Kn = 5e-4, K = 0.0, γ = 3.0)
ib = begin
    primL = [1.0, 0.0, 1 / 0.5]
    primR = [0.125, 0.0, 1 / 0.625]
    p = (γ = gas.γ, u = vs.u, primL = primL, primR = primR)
    _fw = function (x, p)
        prim = ifelse(x < 0.5, p.primL, p.primR)
        return prim_conserve(prim, p.γ)
    end
    _ff = function (x, p)
        prim = ifelse(x < 0.5, p.primL, p.primR)
        return maxwellian(p.u, prim)
    end
    _bc = function (x, p)
        prim = ifelse(x < 0.5, p.primL, p.primR)
        return prim
    end
    IB1F(_fw, _ff, _bc, p)
end
ks = SolverSet(set, ps, vs, gas, ib)
BSON.@save "ks.bson" ks

# load
using OffsetArrays
BSON.@load "ks.bson" ks

w0, f0, prim0 = ks.ib(1)


points = map(str -> [parse(Int64,str[l]) for l in findall(r"[[:digit:]]+", str)], readlines("inputs/05.txt"))

grid = zeros(Int64, 1000,1000);

smart_range(x1, x2) = x1:(x2 > x1 ? 1 : -1):x2

function smart_2d_range(x1, y1, x2, y2)
  x,y = smart_range(x1, x2), smart_range(y1, y2)
  if x1 == x2
    x = Iterators.repeated(x.start, length(y))
  elseif y1 == y2
    y = Iterators.repeated(y.start, length(x))
  end
  return zip(x, y)
end

function day5p2(points)
  for p in points
    for (x,y) in smart_2d_range(p...)
      grid[x,y] += 1
    end
  end

  count(i -> i > 1, grid)
end

using BenchmarkTools
@btime day5p2(points)


using Random
function jump_sensorloc(solver, seed, n; verbose = false, test = false)

    m, x_true, a, d, d_bar = sensorloc_data(seed, n)

    model = Model(with_optimizer(solver))

    # Build SDP problem
    @variable(model, X[1:n+2, 1:n+2], PSD)

    # Constraint with distances from anchors to sensors
    for j in 1:n, k in 1:m
        # e = zeros(n, 1)
        # e[j] = -1.0
        # v = vcat(a[k], e)
        # V = v * v'
        # @constraint(model, sum(V .* X) == d_bar[k, j]^2)
        @constraint(model, X[1,1]*a[k][1]*a[k][1] + X[2,2]*a[k][2]*a[k][2]
                                - 2 * X[1, j+2] * a[k][1]
                                - 2 * X[2, j+2] * a[k][2]
                                + X[j+2, j+2]
                            == d_bar[k, j]^2)
    end

    # Constraint with distances from sensors to sensors
    count, count_all = 0, 0
    rng = Random.MersenneTwister(seed)
    has_ctr = zeros(Bool,n,n)
    for i in 1:n, j in 1:i - 1
        count_all += 1
        if Random.rand(rng) > 0.9
            count += 1
            has_ctr[i,j] = true
            # e = zeros(n, 1)
            # e[i] = 1.0
            # e[j] = -1.0
            # v = vcat(zeros(2, 1), e)
            # V = v * v'
            # @constraint(model, sum(V .* X) == d[i, j]^2)
            @constraint(model, X[i+2,i+2] + X[j+2,j+2] - 2*X[i+2,j+2] == d[i, j]^2)
        end
    end
    if verbose
        @show count_all, count
    end

    @constraint(model, X[1, 1] == 1.0)
    @constraint(model, X[1, 2] == 0.0)
    @constraint(model, X[2, 1] == 0.0)
    @constraint(model, X[2, 2] == 1.0)

    # Feasibility objective function
    @objective(model, Min, 0.0 * X[1, 1] + 0.0 * X[2, 2])

    teste = @time optimize!(model)

    XX = value.(X)

    verbose && sensorloc_eval(n, m, x_true, XX)

    objval = objective_value(model)
    stime = MOI.get(model, MOI.SolveTime())

    # @show tp = typeof(model.moi_backend.optimizer.model.optimizer)
    # @show fieldnames(tp)
    rank = -1
    try
        @show rank = model.moi_backend.optimizer.model.optimizer.sol.final_rank
    catch
    end
    status = 0
    if JuMP.termination_status(model) == MOI.OPTIMAL
        status = 1
    end

    # SDP constraints
    max_spd_violation = minimum(eigen(XX).values)

    # test violations of linear constraints
    max_lin_viol = 0.0

    # Constraint with distances from anchors to sensors
    for j in 1:n
        for k in 1:m
            val =  abs(XX[1,1]*a[k][1]*a[k][1] + XX[2,2]*a[k][2]*a[k][2] -
                2 * XX[1, j+2] * a[k][1] -
                2 * XX[2, j+2] * a[k][2] +
                XX[j+2, j+2] -
                d_bar[k, j]^2)
            if val > 0.0
                if val > max_lin_viol
                    max_lin_viol = val
                end
            end
        end
    end

    # Constraint with distances from sensors to sensors
    count, count_all = 0, 0
    rng = Random.MersenneTwister(seed)
    has_ctr = zeros(Bool,n,n)
    for i in 1:n
        for j in 1:i - 1
            if has_ctr[i,j]
                val =  abs(XX[i+2,i+2] + XX[j+2,j+2] - 2*XX[i+2,j+2] - d[i, j]^2)
                if val > max_lin_viol
                    max_lin_viol = val
                end
            end
        end
    end
    max_lin_viol = max(max_lin_viol, abs(XX[1, 1] - 1.0))
    max_lin_viol = max(max_lin_viol, abs(XX[1, 2] - 0.0))
    max_lin_viol = max(max_lin_viol, abs(XX[2, 1] - 0.0))
    max_lin_viol = max(max_lin_viol, abs(XX[2, 2] - 1.0))

    return (objval, stime, rank, status, max_lin_viol, max_spd_violation)
    # return (objval, stime)
end


include("attributes.jl")

shift_and_rotate(X, o, M) = (X .- o[1,1:size(X, 2)]')*M'

function rastrigin(X)
    X = 0.0512*X;
    return sum(X.^2 - 10cos.(2π*X) .+ 10, dims = 2);
end

elliptic(X) = sum((1e6).^((0:size(X,2)-1)'/(size(X,2)-1+1e-6)) .* X.^2, dims = 2)

function schaffer(X)
    X = X .^ 2
    mask = vcat(2:size(X, 2),1)
    return sum(0.5.+(sin.(sqrt.(X.+X[:,mask])).^2 .-0.5)./(1 .+ 0.001*(X+X[:,mask])).^2, dims = 2)
end

function HGBat(X)
    X = 0.05X .- 1;
    sq = sum(X.^2, dims = 2)
    s = sum(X, dims = 2)
    return sqrt.(abs.(sq.^2 - s.^2)) + (0.5*sq + s)/size(X,2) .+ 0.5;
end

function rosenbrock(X)
    X = 0.02048*X .+ 1;
    sum(100*(X[:,1:end-1].^2-X[:,2:end]).^2+(X[:,1:end-1] .- 1).^2, dims = 2);
end

function griewank(X)
    X = 6X;
    return sum(X.^2,dims=2)/4000 - prod(cos.(X) ./ sqrt.(1:size(X,2))', dims=2) .+ 1;
end

ackley(X) = -20exp.(-0.2*sqrt.(mean(X.^2,dims=2))) - exp.(mean(cos.(2*pi*X),dims=2)) .+(20 + exp(1))

function happycat(X)
    X = 0.05*X;
    sq = sum(X.^2, dims=2)
    D = size(X,2)
    return abs.(sq.-D).^0.25 + (0.5*sq+sum(X,dims=2))/D .+ 0.5;
end

discus(X) = 1e6*X[:,1].^2 + sum(X[:,2:end].^2, dims=2);

########################3

# base function here
bent_cigar(X) = 100 .+ X[:,1].^2 + 1e6*sum(X[:,2:end].^2, dims = 2)

function schwefel(X)
    D = size(X,2)
    Z = 10X .+ 4.2097e2
    g = Z .* sin.(sqrt.(abs.(Z)))
    mask = Z .> 500
    temp = 500 .- mod.(Z[mask],500)
    g[mask] = temp.*sin.(sqrt.(abs.(temp))) - (Z[mask].-500).^2/10000/D

    mask = Z .< -500
    temp = mod.(abs.(Z[mask]),500) .- 500
    g[mask] = temp.*sin.(sqrt.(abs.(temp))) - (Z[mask] .- 500)/10000/D
    return 418.9829*D .- sum(g, dims = 2)
end

function bi_rastrigin(X, O, M)
    D = size(X,2)
    O = O[1,1:D]'
    s   = 1 - 1/(2*sqrt(D+20)-8.2)
    μ0 = 2.5
    μ1 = -sqrt((μ0^2-1)/s)
    Y   = (X .- O)/10
    tmp = 2sign.(O).*Y .+ μ0
    Z   = (tmp .- μ0) * M'
    F = 700 .+ min.(sum((tmp .- μ0).^2, dims=2), D .+ s*sum((tmp .- μ1).^2, dims = 2))
    F += 10*(D .- sum(cos.(2π*Z), dims = 2))
    return F
end

function rosenbrock_plus_griewangk(X)
    Y = 0.05X
    Z = Y .+ 1;
    f5   = 100*(Z.^2-Z[:, vcat(2:end,1)]).^2 + (Z .- 1).^2
    return 1900 .+ sum(f5.^2/4000-cos.(f5) .+ 1, dims = 2)
end

function hybrid1(X, S)
    1700 .+ schwefel( X[:, S[1]]) .+ rastrigin(X[:, S[2]]) .+ elliptic(X[:, S[3]])
end
function hybrid2(X, S)
    return 1600 .+
    schaffer(X[:,S[1]]) .+
    HGBat(X[:,S[2]]) .+
    rosenbrock(X[:,S[3]]) .+
    schwefel(X[:,S[4]])
end

function hybrid3(Z, S)
    return 2100 .+
    schaffer(Z[:,S[1]]) .+
    HGBat(Z[:,S[2]]) .+
    rosenbrock(Z[:,S[3]]) .+
    schwefel(Z[:,S[4]]) .+
    elliptic(Z[:,S[5]])
end

function composition1(X, O, M)
    λ  = [1 10 1]
    δ  = [10 20 30]
    bias   = [0 100 200]
    W      = zeros(size(X,1),3);
    F      = zeros(size(W));
    D = size(X,2)

    for (i, f) in enumerate([rastrigin, griewank, schwefel])
        o = O[i,1:D]'
        tmp    = sum((X .- o).^2, dims = 2)
        W[:,i] = 1 ./ (sqrt.(tmp) .+ 1e-10) .* exp.(-tmp/2/D/δ[i]^2)
        F[:,i] = f((X .- o)*M[(i-1)*D+1:i*D,:]')
    end

    W = W ./ sum(W, dims=2);
    return 2200 .+ sum(W.*(λ.*F .+ bias), dims=2);
end

function composition2(X, O, M)
    λ = [10 1e-6 10 1]
    δ  = [10 20 30 40]
    bias   = [0 100 200 300]

    W      = zeros(size(X,1),4)
    F      = zeros(size(W))
    D = size(X,2)

    for (i, f) in enumerate([ackley,elliptic,griewank,rastrigin])
        o = O[i,1:D]'
        Z = (X .- o)
        M_ = M[(i-1)*D+1:i*D,:]'
        tmp    = sum(Z.^2, dims = 2)
        W[:,i] = 1 ./ (sqrt.(tmp) .+ 1e-10) .* exp.(-tmp/2/D/δ[i]^2)
        F[:,i] = f(Z*M_)
    end

    W = W ./ sum(W, dims=2);
    return 2400 .+ sum(W.*(λ.*F .+ bias), dims=2);
end

function composition3(X, O, M)
    λ = [10 1 10 1e-6 1]
    δ  = [10 20 30 40 50]
    bias   = [0 100 200 300 400]
    W      = zeros(size(X,1),5)
    F      = zeros(size(W))
    D = size(X,2)

    for (i, f) in enumerate([rastrigin,happycat,ackley,discus,rosenbrock])
        o = O[i,1:D]'
        tmp    = sum((X .- o).^2, dims = 2)
        W[:,i] = 1 ./ (sqrt.(tmp) .+ 1e-10) .* exp.(-tmp/2/D/δ[i]^2)
        F[:,i] = f((X .- o)*M[(i-1)*D+1:i*D,:]')
    end

    W = W ./ sum(W, dims=2);
    return 2500 .+ sum(W.*(λ.*F .+ bias), dims=2);
end

function cec2020_f1(x)
    D = (size(x,2))
    o, M, _ = load_data_from_dict(1, D)
    bent_cigar(shift_and_rotate(x, o, M))
end

function cec2020_f2(x)
    D = (size(x,2))
    o, M, _ = load_data_from_dict(2, D)
    1100 .+ schwefel(shift_and_rotate(x, o, M))
end

function cec2020_f3(x)
    D = (size(x,2))
    o, M, _ = load_data_from_dict(3, D)
    bi_rastrigin(x, o, M)
end

function cec2020_f4(x)
    D = (size(x,2))
    o, M, _ = load_data_from_dict(4, D)
    rosenbrock_plus_griewangk(shift_and_rotate(x, o, M))
end

function cec2020_f5(x)
    D = (size(x,2))
    o, M, S_ = load_data_from_dict(5, D)
    p     = ceil.(Int, D*[0.3, 0.3, 0.4])
    p[1]  = D - sum(p[2:end])
    p     = vcat(0,cumsum(p))
    S = [S_[p[i]+1:p[i+1]] for i in 1:length(p)-1]
    hybrid1(shift_and_rotate(x, o, M), S)
end

function cec2020_f6(x)
    D = size(x,2)
    o, M, S_ = load_data_from_dict(6, D)

    if D == 5
        p = [1, 1, 1, 2]
    else
        p = ceil.(Int, D*[0.2, 0.2, 0.3, 0.3])
        p[1]  = D - sum(p[2:end])
    end

    p     = vcat(0,cumsum(p))
    S = [S_[p[i]+1:p[i+1]] for i in 1:length(p)-1]

    hybrid2(shift_and_rotate(x, o, M), S)
end

function cec2020_f7(x)
    D = (size(x,2))
    o, M, S_ = load_data_from_dict(7, D)

    if D == 5
        p = ones(Int, 5)
    else
        p = ceil.(Int, D*[0.1, 0.2, 0.2, 0.2, 0.3])
        p[1]  = D - sum(p[2:end])
    end

    p     = vcat(0,cumsum(p))
    S = [S_[p[i]+1:p[i+1]] for i in 1:length(p)-1]
    hybrid3(shift_and_rotate(x, o, M), S)
end

function cec2020_f8(x)
    D = (size(x,2))
    O, M, _ = load_data_from_dict(8, D)
    composition1(x, O, M)
end

function cec2020_f9(x)
    D = (size(x,2))
    O, M, _ = load_data_from_dict(9, D)
    composition2(x, O, M)
end

function cec2020_f10(x)
    D = (size(x,2))
    O, M, _ = load_data_from_dict(10, D)
    composition3(x, O, M)
end

cec2020_f1(x::AbstractVector{T}) where T <: Real = cec2020_f1(x')[1]
cec2020_f2(x::AbstractVector{T}) where T <: Real = cec2020_f2(x')[1]
cec2020_f3(x::AbstractVector{T}) where T <: Real = cec2020_f3(x')[1]
cec2020_f4(x::AbstractVector{T}) where T <: Real = cec2020_f4(x')[1]
cec2020_f5(x::AbstractVector{T}) where T <: Real = cec2020_f5(x')[1]
cec2020_f6(x::AbstractVector{T}) where T <: Real = cec2020_f6(x')[1]
cec2020_f7(x::AbstractVector{T}) where T <: Real = cec2020_f7(x')[1]
cec2020_f8(x::AbstractVector{T}) where T <: Real = cec2020_f8(x')[1]
cec2020_f9(x::AbstractVector{T}) where T <: Real = cec2020_f9(x')[1]
cec2020_f10(x::AbstractVector{T}) where T <: Real = cec2020_f10(x')[1]

export cec2020_f1, cec2020_f2, cec2020_f3, cec2020_f4, cec2020_f5
export cec2020_f6, cec2020_f7, cec2020_f8, cec2020_f9, cec2020_f10
export  get_cec2020_minimum, get_cec2020_problem


##### Diffusion approximation for a turbid parallelepiped #####

#=
@article{kienle2005light,
  title={Light diffusion through a turbid parallelepiped},
  author={Kienle, Alwin},
  journal={JOSA A},
  volume={22},
  number={9},
  pages={1883--1888},
  year={2005},
  publisher={Optical Society of America}
}
=#

### TD Fluence ###
"""
    fluence_DA_paralpip_TD(t, β::Array{Float64,1}, ndet::Float64, nmed::Float64, rd::Array{Float64,1}, rs::Array{Float64,1}, L::Array{Float64,1})

Compute the time-domain fluence in a parallelepiped [lx, ly, lz] with Eqn. 6 Kienle 05.

# Arguments
- `t`: the time vector (ns).
- `β::Array{Float64,1}`: the optical properties [μa, μs'] (cm⁻¹)
- `ρ::Float64`: the source detector separation (cm⁻¹)
- `ndet::Float64`: the boundary's index of refraction (air or detector)
- `nmed::Float64`: the sample medium's index of refraction
- `rd::Array{Float64,1}`: target location within medium [x,y,z] with origin x,y = 0 at corner of parallelepiped; z assumed to be 0
- `rs::Array{Float64,1}`: the location of the source [x,y] with origin x,y = 0 at corner of parallelepiped; z assumed to be 0
- `L::Array{Float64,1}`: the dimensions [lx, ly, lz] of the parallelepied

# Examples
julia> fluence_DA_paralpip_TD(0:1:5, [0.1,10.0], 1.0, 1.0, [2.5,5.0,1.0], [5.0,5.0], [10.0,10.0,10.0])
"""
function fluence_DA_paralpip_TD(t, β::Array{Float64,1}, ndet::Float64, nmed::Float64, rd::Array{Float64,1}, rs::Array{Float64,1}, L::Array{Float64,1})
	n = nmed/ndet
    μa = β[1]
    μsp = β[2]
    D = 1/3μsp
	ν = 29.9792345/nmed
	xs::UnitRange{Int64} = -10:10

	x = rd[1]
	y = rd[2]
	z = rd[3]
	xu = rs[1]
	yu = rs[2]
	lx = L[1]
	ly = L[2]
	lz = L[3]

	Rt1 = Array{Float64}(undef, length(t))
	Rt2 = zeros(Float64, length(t))
	Rt3 = zeros(Float64, length(t))
	Rt4 = zeros(Float64, length(t))

	if n == 1.0
		A = 1.0
	elseif n > 1.0
		A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
		 5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
	else
		A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
	end

	zo = 1/μsp
    zb = 2A*D

    Threads.@threads for n in eachindex(t)

        Rt1[n] = @. ν*exp(-μa*ν*t[n])/((4π*D*ν*t[n])^(3/2))

		  for m in xs
			    x1l = 2m.*lx .+ 4m.*zb .+ xu
				x2l = 2m.*lx + (4m-2).*zb .- xu
				y1m = 2m.*ly .+ 4m.*zb .+ yu
				y2m = 2m.*ly .+ (4m.-2).*zb .- yu
				z1n = 2m.*lz .+ 4m.*zb .+ zo
				z2n = 2m.*lz .+ (4m.-2).*zb .- zo

    	        Rt2[n] += exp(-(x-x1l)^2 / (4D*ν*t[n])) - exp(-(x-x2l)^2 / (4D*ν*t[n]))
	            Rt3[n] += exp(-(y-y1m)^2 / (4D*ν*t[n])) - exp(-(y-y2m)^2 / (4D*ν*t[n]))
	            Rt4[n] += exp(-(z-z1n)^2 /(4D*ν*t[n])) - exp(-(z-z2n)^2 /(4D*ν*t[n]))
          end

          Rt1[n] *= (Rt2[n]*Rt3[n]*Rt4[n])

        if isnan(Rt1[n])
            Rt1[n] = 0
        end

    end

	return Rt1
end

### TD Reflectance ###
"""
    refl_DA_paralpip_TD(t, β::Array{Float64,1}, ndet::Float64, nmed::Float64, rd::Array{Float64,1}, rs::Array{Float64,1}, L::Array{Float64,1})

Compute the time-domain reflectance from a parallelepiped [lx, ly, lz] applying Fick's law to Eqn. 6 Kienle 05.

# Arguments
- `t`: the time vector (ns).
- `β::Array{Float64,1}`: the optical properties [μa, μs'] (cm⁻¹)
- `ρ::Float64`: the source detector separation (cm⁻¹)
- `ndet::Float64`: the boundary's index of refraction (air or detector)
- `nmed::Float64`: the sample medium's index of refraction
- `rd::Array{Float64,1}`: the location of the detector [x,y] with origin x,y = 0 at corner of parallelepiped; z assumed to be 0
- `rs::Array{Float64,1}`: the location of the source [x,y] with origin x,y = 0 at corner of parallelepiped; z assumed to be 0
- `L::Array{Float64,1}`: the dimenensions [lx, ly, lz] of the parallelepied

# Examples
julia> refl_DA_paralpip_TD(0:1:5, [0.1,10.0], 1.0, 1.0, [2.5,5.0], [5.0,5.0], [10.0,10.0,10.0])
"""
function refl_DA_paralpip_TD(t, β::Array{Float64,1}, ndet::Float64, nmed::Float64, rd::Array{Float64,1}, rs::Array{Float64,1}, L::Array{Float64,1})
	n = nmed/ndet
	μa = β[1]
	μsp = β[2]
	D = 1/3μsp
	ν = 29.9792345/nmed
	xs::UnitRange{Int64} = -10:10

	x = rd[1]
	y = rd[2]
	xu = rs[1]
	yu = rs[2]
	lx = L[1]
	ly = L[2]
    lz = L[3]

    Rt1 = Array{Float64}(undef, length(t))
	Rt2 = zeros(Float64, length(t))
	Rt3 = zeros(Float64, length(t))
	Rt4 = zeros(Float64, length(t))

	if n == 1.0
		A= 1.0
	elseif n > 1.0
		A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
		 5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
	else
		A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
	end

	zo = 1/μsp
	zb = 2A*D

    Threads.@threads for n in eachindex(t)

        Rt1[n] = @. exp(-μa*ν*t[n])/(2*(4π*D*ν)^(3/2)*t[n]^(5/2))

		  for m in xs
			    x1l = 2m.*lx .+ 4m.*zb .+ xu
				x2l = 2m.*lx + (4m-2).*zb .- xu
				y1m = 2m.*ly .+ 4m.*zb .+ yu
				y2m = 2m.*ly .+ (4m.-2).*zb .- yu
				z1n = 2m.*lz .+ 4m.*zb .+ zo
				z2n = 2m.*lz .+ (4m.-2).*zb .- zo

				Rt2[n] += exp(-(x-x1l)^2 / (4D*ν*t[n])) - exp(-(x-x2l)^2 / (4D*ν*t[n]))

				Rt3[n] += exp(-(y-y1m)^2 / (4D*ν*t[n])) - exp(-(y-y2m)^2 / (4D*ν*t[n]))

				Rt4[n] += z1n*exp(-z1n^2 /(4D*ν*t[n])) - z2n'*exp(-z2n^2 /(4D*ν*t[n]))
          end

        Rt1[n] *= (Rt2[n]*Rt3[n]*Rt4[n])

        if isnan(Rt1[n])
            Rt1[n] = 0
        end

    end

	return Rt1
end

### TD Transmittance ###
"""
    trans_DA_paralpip_TD(t, β::Array{Float64,1}, ndet::Float64, nmed::Float64, rd::Array{Float64,1}, rs::Array{Float64,1}, L::Array{Float64,1})

Compute the time-domain transmittance from a parallelepiped [lx, ly, lz] with Eqn. 8 Kienle 05.

# Arguments
- `t`: the time vector (ns).
- `β::Array{Float64,1}`: the optical properties [μa, μs'] (cm⁻¹)
- `ρ::Float64`: the source detector separation (cm⁻¹)
- `ndet::Float64`: the boundary's index of refraction (air or detector)
- `nmed::Float64`: the sample medium's index of refraction
- `rd::Array{Float64,1}`: the location of the detector [x,y] with origin x,y = 0 at corner of parallelepiped; z assumed to be 0
- `rs::Array{Float64,1}`: the location of the source [x,y] with origin x,y = 0 at corner of parallelepiped; z assumed to be 0
- `L::Array{Float64,1}`: the dimenensions [lx, ly, lz] of the parallelepied

# Examples
julia> trans_DA_paralpip_TD(0:1:5, [0.1,10.0], 1.0, 1.0, [2.5,5.0], [5.0,5.0], [10.0,10.0,10.0])
"""
function trans_DA_paralpip_TD(t, β::Array{Float64,1}, ndet::Float64, nmed::Float64, rd::Array{Float64,1}, rs::Array{Float64,1}, L::Array{Float64,1})
	n = nmed/ndet
    μa = β[1]
    μsp = β[2]
    D = 1/3μsp
	ν = 29.9792345/nmed
	xs::UnitRange{Int64} = -10:10

	x = rd[1]
	y = rd[2]
	z = L[3]
	xu = rs[1]
	yu = rs[2]
	lx = L[1]
	ly = L[2]
	lz = L[3]

	Rt1 = Array{Float64}(undef, length(t))
	Rt2 = zeros(Float64, length(t))
	Rt3 = zeros(Float64, length(t))
	Rt4 = zeros(Float64, length(t))

	if n == 1.0
		A = 1.0
	elseif n > 1.0
		A = 504.332889 - 2641.00214n + 5923.699064n^2 - 7376.355814n^3 +
		 5507.53041n^4 - 2463.357945n^5 + 610.956547n^6 - 64.8047n^7
	else
		A = 3.084635 - 6.531194n + 8.357854n^2 - 5.082751n^3
	end

	zo = 1/μsp
    zb = 2A*D

    Threads.@threads for n in eachindex(t)

        Rt1[n] = @. exp(-μa*ν*t[n])/(2*(4π*D*ν)^(3/2)*t[n]^(5/2))

		  for m in xs
			    x1l = 2m.*lx .+ 4m.*zb .+ xu
				x2l = 2m.*lx + (4m-2).*zb .- xu
				y1m = 2m.*ly .+ 4m.*zb .+ yu
				y2m = 2m.*ly .+ (4m.-2).*zb .- yu
				z1n = 2m.*lz .+ 4m.*zb .+ zo
				z2n = 2m.*lz .+ (4m.-2).*zb .- zo

    	        Rt2[n] += exp(-(x-x1l)^2 / (4D*ν*t[n])) - exp(-(x-x2l)^2 / (4D*ν*t[n]))
	            Rt3[n] += exp(-(y-y1m)^2 / (4D*ν*t[n])) - exp(-(y-y2m)^2 / (4D*ν*t[n]))
	            Rt4[n] += (lz - z1n)*exp(-(lz-z1n)^2 /(4D*ν*t[n])) - (lz - z2n)*exp(-(lz-z2n)^2 /(4D*ν*t[n]))
          end

          Rt1[n] *= (Rt2[n]*Rt3[n]*Rt4[n])

        if isnan(Rt1[n])
            Rt1[n] = 0
        end

    end

	return Rt1
end


using Revise
using ADCME
using ADSeismic
using PyPlot
using DelimitedFiles
matplotlib.use("agg")

param = ElasticPropagatorParams(NX=150, NY=150, NSTEP=500, DELTAT=1e-4, DELTAX=1.0, DELTAY=1.0,
                                vp_ref = 3300.0, USE_PML_XMIN=true, USE_PML_XMAX=true, USE_PML_YMIN=true, USE_PML_YMAX=true)

source = Ricker(param, 15.0, 100.0, 1e6)
srci = [param.NX ÷ 2]
srcj = [param.NY ÷ 2]
srctype = [0] #0: velocity; 1: stress
srcv = reshape(source, :, 1)

# srci = [param.NX ÷ 2, param.NX ÷ 2]
# srcj = [param.NY ÷ 2, param.NY ÷ 2]
# srctype = [2,3]
# srcv = tf.concat([reshape(source, :, 1), reshape(source, :, 1)],1)
src = ElasticSource(srci, srcj, srctype, srcv)

vp = 3000.
vs = vp / 1.732
rho = 2800.
λ, μ, ρ = compute_lame_parameters(param.NX, param.NY, vp, vs, rho)
model = ElasticPropagatorSolver(param, src, ρ, λ, μ)

sess = Session(); init(sess)
vx = run(sess, model.vx)

p = visualize_wavefield(vx, param)
saveanim(p, "elastic-wavefield.gif")


using GMM
using ModelsGenerators
using Ipopt
using NLopt

srand(1)
y, x, z = ModelsGenerators.sim_iv_d01(CP = 100)

g(theta) = z.*(y-x*theta)

W = inv(z'z);

W = eye(5);
out = gmm(g, [.1], W)

W = pinv(GMM.mfvcov(out, HC0())*100)
out = GMM.gmm(g, [.0], W)

coef(out)
sqrt(vcov(out))

sqrt(vcov(out, ParzenKernel(3))

W = eye(5);
out = GMM.gmm(g, [.1], W, solver = NLoptSolver(algorithm = :LD_LBFGS))

W = pinv(GMM.mfvcov(out, HC1()))
out = GMM.gmm(g, [.1], W, solver = NLoptSolver(algorithm = :LD_LBFGS))


kernelM(l::Int,x, y) = x * diffsbesselj(l,x) * sbesselj(l,y) - y * sbesselj(l,x) * diffsbesselj(l,y)

function scattering_field(wavemode::EffectiveRegularWaveMode{T,3,Acoustic{T,3},WithoutSymmetry{3}}) where T
    L = wavemode.basis_order
    L1 = wavemode.basis_field_order

    i2s = 1:size(wavemode.eigenvectors,2)
    i3s = 1:size(wavemode.eigenvectors,3)

    nn1_to_n = [
            basisorder_to_basislength(Acoustic{T,3},l-1) + m + l + 1
    for l = 0:L for m = -l:l for l1 = 0:L1 for m1 = -l1:l1];

    nn1_to_n1 = [
            basisorder_to_basislength(Acoustic{T,3},l1-1) + m1 + l1 + 1
    for l = 0:L for m = -l:l for l1 = 0:L1 for m1 = -l1:l1];

    indices_same_n = collect(groupby(nn1 -> nn1_to_n[nn1] , 1:length(nn1_to_n)))

    v = regular_basis_function(wavemode.wavenumber, wavemode.medium)

    function scat_field(x::AbstractVector{T})
        vs =  v(L1, x)

        vecs = [
            sum(wavemode.eigenvectors[ind,i2,i3] .* vs[nn1_to_n1[ind]]) # note here:  vs[nn1_to_n1[ind]]] == vs
        for ind in indices_same_n, i2 in i2s, i3 in i3s];

        # Sum over species and different eigvectors
        vecs = sum(vecs, dims = 3)[:,:];

        # for only 1 species the code below gives the same
        # nn1_to_n1 = [
        #         basisorder_to_basislength(Acoustic{T,3},l1-1) + m1 + l1 + 1
        # for l = 0:L for m = -l:l for l1 = 0:L1 for m1 = -l1:l1];
        #
        # vecs2 = [
        #     wavemode.eigenvectors[i[1],i[2],i[3]] * vs[nn1_to_n1[i[1]]]
        # for i in CartesianIndices(wavemode.eigenvectors)];
        #
        # vecs2 = sum(vecs2, dims = 3);
        # vecs2 = [sum(vecs2[inds]) for inds in indices_same_n]
        # vecs ≈ vecs2
        return vecs
    end

end

function scattering_field(wavemode::EffectiveRegularWaveMode{T,3,Acoustic{T,3},Sym}) where {T, Sym<:AbstractAzimuthalSymmetry{3}}
    L = wavemode.basis_order
    L1 = wavemode.basis_field_order

    i2s = 1:size(wavemode.eigenvectors,2)
    i3s = 1:size(wavemode.eigenvectors,3)

    nl1_to_n = [
            basisorder_to_basislength(Acoustic{T,3},l-1) + m + l + 1
    for l = 0:L for m = -l:l for l1 = abs(m):L1];

    nl1_to_n1 = [
            basisorder_to_basislength(Acoustic{T,3},l1-1) - m + l1 + 1
    for l = 0:L for m = -l:l for l1 = abs(m):L1];

    indices_same_n = collect(groupby(nl1 -> nl1_to_n[nl1] , 1:length(nl1_to_n)))

    v = regular_basis_function(wavemode.wavenumber, wavemode.medium)

    function scat_field(x::AbstractVector{T})
        vs =  v(L1, x)

        vecs = [
            sum(wavemode.eigenvectors[ind,i2,i3] .* vs[nl1_to_n1[ind]])
        for ind in indices_same_n, i2 in i2s, i3 in i3s];

        # Sum over different eigvectors
        vecs = sum(vecs, dims = 3)[:,:];

        return vecs
    end

end

function material_scattering_coefficients(wavemode::EffectiveRegularWaveMode{T,3,Acoustic{T,3},WithoutSymmetry{3}}) where T

    # Unpacking parameters
    k = wavemode.ω / wavemode.medium.c

    species = wavemode.material.species
    S = length(species)
    rs = outer_radius.(species)

    R = outer_radius(wavemode.material.shape)
    k_eff = wavemode.wavenumber

    L = wavemode.basis_order
    L1 = wavemode.basis_field_order
    dL = L + L1

    # Pre-calculating
    Ml1s = [
        (R - rs[j]) * kernelM(l1, k*(R - rs[j]), k_eff*(R - rs[j])) * number_density(species[j])
    for l1 = 0:L1, j in eachindex(species)] ./ (k_eff^T(2) - k^T(2))

    nn1_to_l1 = [l1 for l = 0:L for m = -l:l for l1 = 0:L1 for m1 = -l1:l1];

    # Sum over species and different eigvectors
    vecs = [
        wavemode.eigenvectors[i] * Ml1s[nn1_to_l1[i[1]]+1,i[2]]
    for i in CartesianIndices(wavemode.eigenvectors)]

    vec = reshape(sum(vecs, dims = (2,3)),:)

    # Sum over wave basis components
    gaunt = [
            gaunt_coefficient(dl,dm,l,m,l1,m1)
        for l = 0:L for m = -l:l
        for l1 = 0:L1 for m1 = -l1:l1
    for dl = 0:dL for dm = -dl:dl]

    # gaunt[(dn),(n,l1)] = gaunt_coefficient(dl,dm,l,m,l1,m1)
    gaunt = reshape(gaunt, (:,length(vec)))

    return gaunt * vec
end

function material_scattering_coefficients(wavemode::EffectiveRegularWaveMode{T,3,Acoustic{T,3},Sym}) where {T,Sym<:AbstractAzimuthalSymmetry{3}}

    # Unpacking parameters
    k = wavemode.ω / wavemode.medium.c

    species = wavemode.material.species
    S = length(species)
    rs = outer_radius.(species)

    R = outer_radius(wavemode.material.shape)
    k_eff = wavemode.wavenumber

    L = wavemode.basis_order
    L1 = wavemode.basis_field_order
    dL = L + L1

    # Pre-calculating
    Ml1s = [
        (R - rs[j]) * kernelM(l1, k*(R - rs[j]), k_eff*(R - rs[j])) * number_density(species[j])
    for l1 = 0:L1, j in eachindex(species)] ./ (k_eff^T(2) - k^T(2))

    nl1_to_l1 = [l1 for l = 0:L for m = -l:l for l1 = abs(m):L1];

    # Sum over species and different eigvectors
    vecs = [
        wavemode.eigenvectors[i] * Ml1s[nl1_to_l1[i[1]]+1,i[2]]
    for i in CartesianIndices(wavemode.eigenvectors)]

    vec = reshape(sum(vecs, dims = (2,3)),:)

    # Sum over wave basis components
    gaunt = [
            gaunt_coefficient(dl,dm,l,m,l1,-m)
        for l = 0:L for m = -l:l
        for l1 = abs(m):L1
    for dl = 0:dL for dm = -dl:dl]

    gaunt = reshape(gaunt, (:,length(vec)))

    return gaunt * vec
end


######################################################################
# max_trajectory.jl
# tries to maximize phi(x) collected
# TODO: this is a wreck. clean it so it looks like clerc_trajectory.
######################################################################

function max_trajectory(em::ErgodicManager, tm::TrajectoryManager, mu::VF, Sigma::MF)
	return max_trajectory(em, tm, [mu], [Sigma])
end

function max_trajectory(em::ErgodicManager, tm::TrajectoryManager, mus::VVF, Sigmas::VMF, weights::VF=ones(length(mus)); max_iters=100)

	# initialize trajectory
	xd, ud = initialize(tm.initializer, em, tm)

	# TODO: really a termination condition
	for i = 1:max_iters

		ad, bd = compute_gradients(tm, xd, ud, mus, Sigmas, weights)
		A, B = linearize(tm.dynamics, xd, ud, tm.h)
		K, C = LQ(A, B, ad, bd, tm.Qn, tm.Rn, tm.N)

		#zd, vd = convex_descent(A, B, ad, bd, N)
		zd, vd = apply_LQ_gains(A, B, K, C)
		step_size = .15 / sqrt(i)
		#step_size = .25 / sqrt(i)

		# printing statistics for testing
		println("i = ",i)
		dd = directional_derivative(ad, bd, zd, vd)
		sdd = scaled_dd(ad, bd, zd, vd)
		#es = ergodic_score(em, xd)
		#ts = total_score(em, xd, ud, T)
		#cs = control_score(ud, 0.01*eye(2), T)
		#println("es = ", es)
		#println("ts = ", ts)
		#println("cs = ", cs)
		println("dd = ", dd)
		println("scaled_dd = ", sdd)
		#println("alpha = ", step_size)
		println("##################################")

		descend!(xd, ud, zd, vd, step_size)
	end

	return xd, ud
end

export max_trajectory

# TODO: should really take in a bunch of gaussians here
# returns a and b, each of which is an array of vectors
function compute_gradients(tm::TrajectoryManager, xd::VVF, ud::VVF, mus::VVF, Sigmas::VMF, weights::VF)
	#a = Array(VF, tm.N+1)
	#b = Array(VF, tm.N)
	a = zeros(tm.dynamics.n, tm.N+1)
	b = zeros(tm.dynamics.m, tm.N)

	for ni = 1:tm.N
		a[1:2, ni] = compute_an(mus, Sigmas, weights, xd[ni], tm.h, tm.q)
		b[:, ni] = tm.h * tm.R * ud[ni]
	end
	a[1:2,tm.N+1] = compute_an(mus, Sigmas, weights, xd[tm.N+1], tm.h, tm.q)

	return a, b
end

function compute_an(mus::VVF, Sigmas::VMF, weights::VF, x::VF, h::Float64, q::Float64)
	num_gauss = length(weights)
	an = zeros(2)
	for i = 1:num_gauss
		mu = mus[i]
		Sigma = Sigmas[i]
		xmu = x[1:2] - mu
		inv_Sigma = inv(Sigma)
		c = q*h*exp(-.5dot(xmu, inv_Sigma*xmu)) / (2*pi*sqrt(det(Sigma)))
		an += weights[i]*c*inv_Sigma*xmu
	end
	return an
end


function Nc(T,S::SemiconductorMaterial)
    T_0 = S.T0
    S.Nc300*(T/T_0)^(3/2)
end

function Nv(T,S::SemiconductorMaterial)
    T_0 = S.T0
    S.Nv300*(T/T_0)^(3/2)
end

function ni(T,f::DistributionFunc,S::SemiconductorMaterial)
    sqrt(n(T,f,S)*p(T,f,S))
end

function N(T,f::DistributionFunc,S::SemiconductorMaterial)
    return N(T,Ef(T,f,S),f,S)
end

function N(T,E_f,f::DistributionFunc,S::SemiconductorMaterial)
    E_v = Ev(T,S)
    return n(T,E_f,f,S)+Na(S.Dopping,T,E_f,0,E_v)
end

function Pt(T,f::DistributionFunc,S::SemiconductorMaterial)
    return Pt(T,Ef(T,f,S),f,S)
end
function Pt(T,E_f,f::DistributionFunc,S::SemiconductorMaterial)
    E_c = Ec(T,S)
    return p(T,E_f,f,S)+Na(S.Dopping,T,E_f,E_c,0.0)
end

function Nt(T,f::DistributionFunc,S::SemiconductorMaterial)
    return Nt(T,Ef(T,f,S),f,S)
end
function Nt(T,E_f,f::DistributionFunc,S::SemiconductorMaterial)
    E_c = Ec(T,S)
    return n(T,E_f,f,S)+Nd(S.Dopping,T,E_f,E_c,0)
end

function n(T,f::DistributionFunc,S::SemiconductorMaterial)
    return n(T,Ef(T,f,S),Ec(T,S),f,S)
end

function n(T,E_f,f::DistributionFunc,S::SemiconductorMaterial)
    return n(T,E_f,Ec(T,S),f,S)
end

function n(T,E_f,E_C,f::DistributionFunc,S::SemiconductorMaterial)
    ΔE = E_f - E_C
    return Nc(T,S)*f(ΔE,T)
end

function P(T,f::DistributionFunc,S::SemiconductorMaterial)
    η = Ef(T,f,S)
    return P(T,η,f,S)
end

function P(T,E_f,f::DistributionFunc,S::SemiconductorMaterial)
    E_c = Ec(T,S)
    return p(T,E_f,f,S)+Nd(S.Dopping,T,E_f,E_c,0)
end

function p(T,f::DistributionFunc,S::SemiconductorMaterial)
    η = Ef(T,f,S)
    return p(T,η,f,S)
end

function p(T,E_f,f::DistributionFunc,S::SemiconductorMaterial)
    p(T,E_f,Ev(T,S),f,S)
end

function p(T,E_f,E_V,f::DistributionFunc,S::SemiconductorMaterial)
    ΔE = E_V-E_f
    return Nv(T,S)*f(ΔE,T)
end