# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md

using JuliaFEM
using JuliaFEM.Preprocess
using JuliaFEM.Testing

@testset "test tet4 stiffness matrix" begin
    el = Element(Tet4, [1, 2, 3, 4])
    el["youngs modulus"] = 96.0
    el["poissons ratio"] = 1/3
    x1 = [2.0, 3.0, 4.0]
    x2 = [6.0, 3.0, 2.0]
    x3 = [2.0, 5.0, 1.0]
    x4 = [4.0, 3.0, 6.0]
    u1 = u2 = u3 = u4 = zeros(3)
    el["geometry"] = Vector{Float64}[x1, x2, x3, x4]
    u = Vector{Float64}[u1, u2, u3, u4]
    pr = Problem(Elasticity, "tet4", 3)
    as = Assembly()
    assemble!(as, pr, el, 0.0)
    Kt = full(as.K)
    Kt_expected = [
         149.0   108.0   24.0   -1.0    6.0   12.0  -54.0   -48.0    0.0  -94.0   -66.0  -36.0
         108.0   344.0   54.0  -24.0  104.0   42.0  -24.0  -216.0  -12.0  -60.0  -232.0  -84.0
          24.0    54.0  113.0    0.0   30.0   35.0    0.0   -24.0  -54.0  -24.0   -60.0  -94.0
          -1.0   -24.0    0.0   29.0  -18.0  -12.0  -18.0    24.0    0.0  -10.0    18.0   12.0
           6.0   104.0   30.0  -18.0   44.0   18.0   12.0   -72.0  -12.0    0.0   -76.0  -36.0
          12.0    42.0   35.0  -12.0   18.0   29.0    0.0   -24.0  -18.0    0.0   -36.0  -46.0
         -54.0   -24.0    0.0  -18.0   12.0    0.0   36.0     0.0    0.0   36.0    12.0    0.0
         -48.0  -216.0  -24.0   24.0  -72.0  -24.0    0.0   144.0    0.0   24.0   144.0   48.0
           0.0   -12.0  -54.0    0.0  -12.0  -18.0    0.0     0.0   36.0    0.0    24.0   36.0
         -94.0   -60.0  -24.0  -10.0    0.0    0.0   36.0    24.0    0.0   68.0    36.0   24.0
         -66.0  -232.0  -60.0   18.0  -76.0  -36.0   12.0   144.0   24.0   36.0   164.0   72.0
         -36.0   -84.0  -94.0   12.0  -36.0  -46.0    0.0    48.0   36.0   24.0    72.0  104.0]
    if !isapprox(Kt, Kt_expected)
        info("Test failed")
        info("Kt_expected")
        dump(Kt_expected)
        info("Kt")
        dump(Kt)
    end
    @test isapprox(Kt, Kt_expected)
end


#=

    abc.jl

    Contains code to conduct approximate Bayesian computation (ABC)

    Author:     Alexander P. Browning
                ======================
                School of Mathematical Sciences
                Queensland University of Technology
                ======================
                ap.browning@icloud.com
                alexbrowning.me

=#

##############################################################
## ABC Sequential Monte-Carlo (SMC)
##############################################################
"""
    abc_smc(dist,prior;n=100,α=0.5,S₀=100,c=0.01,ptarget=0.005,maxsteps=20,f=2.0)

Perform ABC SMC.
"""
function abc_smc(dist,prior;n=100,kwargs...)
    P = sample(prior,dist,n)    # Initialise using prior
    print("Initialised $n particles. $(round(minimum(P).d,sigdigits=4)) ≤ d ≤ $(round(maximum(P).d,sigdigits=4)).\n")
    return abc_smc(P,dist,prior;kwargs...)
end
function abc_smc(P,dist,prior;step=1,α=0.5,ptarget=0.01,maxsteps=30,f=2.0)

    # Loop until target reached
    p = 1.0
    while p > ptarget && step ≤ maxsteps
        @printf("Step %-3s (ε = %-8s\b...",
            "$step",
            "$(round(sort(distances(P))[(Int ∘ floor)(α * length(P))],sigdigits=4))",
        )

        p = abc_smc_step!(P,dist,prior,α,f)
        step += 1

        @printf("\b\b\b dmin = %-8s p = %-8s ESS = %-8s)\n",
            "$(round(minimum(P).d,sigdigits=4))",
            "$(round(p,sigdigits=2))",
            "$(round(1 / sum(weights(P).^2),sigdigits=5))"
        )

    end

    # Finish
    return P
end

"""
    abc_smc_step!(P,dist,prior,α,f=2.0)

Perform a single step of ABC SMC, where α * n particles are dicarded.

Proposal is estimated as f * Σ, where Σ is the (weighted) covariance matrix of the non-dicarded particles.

"""
function abc_smc_step!(P,dist,prior,α,f=2.0)
    N = length(P); Nα = (Int ∘ floor)(α * N)

    sort!(P)                # Sort particles
    ε = P[N-Nα].d           # Current threshold
    Σ = f*cov(P[1:N-Nα])    # Proposal covariance
    prop = MvNormal(Σ)      # Proposal

    # Resample and update particles
    tries = 0
    @threads for j = (N - Nα + 1) : N
        θs,d = similar(P[1].θ),Inf
        while d > ε
            θs = copy(sample(P[1:N-Nα],weights(P[1:N-Nα])).θ) + rand(prop)
            if insupport(prior,θs)
                d = dist(θs)
                tries += 1
            end
        end
        P[j] = Particle(θs,d,pdf(prior,θs) / sum(p.w * pdf(prop,θs - p.θ) for p in P[1:N-Nα]))
    end

    # Reweight
    w = weights(P)
    w[1:N-Nα] /= sum(w[1:N-Nα])         # Normalise
    w[N-Nα+1:N] /= sum(w[N-Nα+1:N])     # Normalise
    w[1:N-Nα] *= (N - Nα) / N           # Re-weight
    w[N-Nα+1:N] *= Nα / N               # Re-weight
    [p.w = wᵢ for (p,wᵢ) in zip(P,w)]

    # Report acceptance probability (exclude trials where we propose a value outside the prior)
    return Nα / tries
end

##############################################################
## ABC Markov-chain Monte-Carlo (MCMC)
##############################################################
"""
    abc_mcmc(dist,prior,θ₀,Σ,ε;n=1000,nchains=nthreads(),...)

Perform ABC MCMC using the Metropolis-Hastings algorithm (implemented through `AdvancedMH.jl`).
"""
function abc_mcmc(dist,prior,θ₀,Σ,K::Function;n=1000,nchains=Threads.nthreads(),kwargs...)
    parallel_args = nchains == 1 ? n : (MCMCThreads(),n,nchains)
    sample(
        DensityModel(θ -> insupport(prior,θ) ? log(K(dist(θ))) : -Inf),
        RWMH(MvNormal(Σ)),
        parallel_args...;
        init_params=θ₀,
        chain_type=Chains,
        kwargs...
    )
end
abc_mcmc(dist,prior,θ₀,Σ,ε::Number;kwargs...) = abc_mcmc(dist,prior,θ₀,Σ,d -> d ≤ ε ? 1.0 : 0.0;kwargs...)

##############################################################
## ABC MCMC Discrepancy Metrics
##############################################################

logistic(d,ε,k) = 1 / (1 + exp(k * (d - ε)))
logit(p,ε,k) = ε + log((1 - p) / p) / k

##############################################################
## MCMCChains.Chains helper functions
##############################################################

# Convert MCMC chain output to a matrix
Base.Matrix(C::Chains) = permutedims(hcat([Matrix(C[:,i,:])[:] for i = 1:size(C,2)-1]...))

acceptance_rate(C::Chains) = mean(hcat([cumsum(C[1:end-1,1,i] .!= C[2:end,1,i],dims=1)[:] ./ (1:size(C,1)-1) for i = 1:size(C,3)]...),dims=2)[:]

MCMCChains.missing_datetime(T::Type) = missing

StatsBase.sample(C::Chains,n) = (CM = Matrix(C); [CM[:,i] for i in sample(1:size(CM,2),n)])
StatsBase.sample(C::Chains) = sample(C,1)[1]


using Revise
using Plots
using Basins
using DifferentialEquations
using DynamicalSystems

# Blowout bifurcation with non-normal parameters in population dynamicsBernard Cazelle
function franke_yabuku(dz,z, p, n)
    r=2.825; s=20.25; c1=1.2; c2=0.1;
    x = z[1]; y = z[2];
    dx = x*exp(r-s*(x+y))
    dy = c1*y/(c2+x+y)
    dz = [dx,dy]
    return
end

# dummy function to keep the initializator happy
function franke_yabuku_J(J,z0, p, n)
   return
end

ds = DiscreteDynamicalSystem(franke_yabuku,[0.1, 0.2], [] , franke_yabuku_J)
integ  = integrator(ds)

xg=range(8.,16.,length=100)
yg=range(0.5,1.,length=100)

# compute basin

@time basin=basins_map2D(xg, yg, integ)
plot(xg,yg,basin',seriestype=:heatmap)


n = 4
p = [  2 , 4 , 3 , 1 ]
d = [  1 , 2 , 4 , 6 ]
r = [  0 , 0 , 0 , 0 ]
w = [  1 , 1 , 1 , 1 ]

using vOptSpecific

id = set2OSP( n , p , d , r , w )
solver = OSP_VanWassenhove1980( )
z1, z2 , S = vSolve( id , solver )


export vectorizePattern, SamplingOp

"""
 idx contains sampling index (for the first dimension) of a multidimensional Array
 of size "shape". Transform this into idx into the corresponding vector index
"""
function vectorizePattern(idx::Array{Int}, shape::Tuple)
  return [ floor(Int,(i-1)/size(idx,1))*shape[1]+idx[i] for i = 1:length(idx) ]
end

"""
    SamplingOp(pattern::Array{Int}, shape::Tuple)

buildsa `LinearOperator` which only returns the vector elements at positions
indicated by pattern.

# Arguents
* `pattern::Array{Int}` - indices to sample
* `shape::Tuple`        - size of the array to sample
"""
function SamplingOp(pattern::Array{Int}, shape::Tuple)
  ndims(pattern)>1 ?  idx = vectorizePattern(pattern, shape) : idx = pattern

  return opRestriction(idx, prod(shape))
end

function SamplingOp(pattern::Array{Bool})
  return LinearOperator(length(pattern), length(pattern), false, false
                  , x->vec(pattern).*x
                  , nothing
                  , y->vec(pattern).*y )
end

function zeropad(x::Array{T}, pattern::Matrix{Bool}) where T
  y = zeros(T,size(pattern))
  y[pattern] = x[pattern]
  return y
end


@test iscyclic(permcycs((1,2,3)))
@test iscyclic(permcycs((42,100,3)))
@test !iscyclic(permcycs((1,2), (3,4)))


#=
Each complex should have collection of cells per dimension:
- cells::Dict{Int,Vector{C}} or Vector{Vector{C}}
=#
abstract type AbstractComplex end

function boundary(cplx::AbstractComplex, ch::Chain{R}) where {R}
    d = dim(ch)
    cc = Chain(d-1, R)
    for (coef,elem) in ch
        append!(cc, coef * boundary(cplx, elem, d, R))
    end
    return cc
end

function coboundary(cplx::AbstractComplex, ch::Chain{R}) where {R}
    d = dim(ch)
    cc = Chain(d+1,R)
    for (coef,elem) in ch
        append!(cc, coef * coboundary(cplx, elem, d, R))
    end
    return cc
end

#
# AbstractComplex Public Interface
#
"""Return a complex boundary given element and dimension"""
boundary(cplx::AbstractComplex, i::Int, d::Int, ::Type{R}) where {R} = throw(MethodError(boundary, (typeof(cplx),Int,Int,R)))
boundary(cplx::AbstractComplex, i::Int, d::Int) = boundary(cplx, i, d, Int)

"""Return a complex coboundary given element and dimension"""
coboundary(cplx::AbstractComplex, i::Int, d::Int, ::Type{R}) where {R} = throw(MethodError(coboundary, (typeof(cplx),Int,Int,R)))
coboundary(cplx::AbstractComplex, i::Int, d::Int) = coboundary(cplx, i, d, Int)

"""Return a complex cell type"""
celltype(cplx::AbstractComplex) = throw(MethodError(celltype, (typeof(cplx),)))

"""Return a dimension of the complex"""
dim(cplx::AbstractComplex) = throw(MethodError(dim,(typeof(cplx),)))

"""Set dimension of the complex"""
setdim!(cplx::AbstractComplex, d::Int) = throw(MethodError(setdim!, (typeof(cplx),Int)))

"""Return a cell collection per dimension (increasing)"""
cells(cplx::AbstractComplex) = throw(MethodError(cells, (typeof(cplx),)))

"""Return a nullable cell collection per dimension (increasing)"""
cells(cplx::AbstractComplex, d::Int) = throw(MethodError(cells, (typeof(cplx),Int)))

"""
    push!(cplx::AbstractComplex, cell::AbstractCell; recursive=false) -> Vector{AbstractCell}

Insert a `cell` to a complex `cplx`, and returns an array of inserted cell(s). If `recursive=true` is passed then all faces of the `cell` are also added to the complex `cplx`.
"""
Base.push!(cplx::AbstractComplex, c::AbstractCell; recursive=false) = throw(MethodError(push!, (typeof(cplx),typeof(c))))

#
# Public Methods
#
"""Return a size of cell collections per dimension"""
Base.size(cplx::AbstractComplex) = (map(length, cells(cplx))...,)
Base.length(cplx::AbstractComplex) = sum(size(cplx))

"""Return a size of the cell collection for dimension (0-based)"""
function Base.size(cplx::AbstractComplex, d::Int)
    (d < 0 || d > dim(cplx)) && return 0
    sz = size(cplx)
    d >= length(sz) && return 0
    return sz[d+1]
end

"""Return an index of the k-dimensional cell in the complex.

**Note:** If returned index is larger than the number of cells of this dimension, then the cell is not in complex and the returned index will be assigned to the cell when it's added to the complex.
"""
function Base.getindex(cplx::AbstractComplex, c::C, d::Int) where {C}
    @assert celltype(cplx) == C "Incorrect cell type"
    dcells = cells(cplx, dim(c))
    dcells === nothing && return 1
    cidx = findfirst(isequal(c), dcells)
    cidx === nothing && return size(cplx, d)+1
    return dcells[cidx][:index]
end
Base.getindex(cplx::AbstractComplex, c::C) where {C} =  cplx[c, dim(c)]

"""Return a `d`-dimensional cell given its index `idx`."""
function Base.getindex(cplx::AbstractComplex, idx::Int, d::Int)
    cs = cells(cplx, d)
    cs === nothing && return nothing
    return cs[idx]
end

"""Generate a boundary matrix from the cell complex of the dimension `d`."""
function boundary_matrix(::Type{R}, cplx::AbstractComplex, d::Int) where {R}
    csize = size(cplx)
    rows = d > 0 ? csize[d] : 0
    cols = d <= dim(cplx) ? csize[d+1] : 0
    bm = spzeros(R, rows, cols)
    if d>=0 && d <= dim(cplx)
        for i in 1:csize[d+1]
            for (coef,elem) in boundary(cplx, i, d, R)
                bm[elem, i] = coef
            end
        end
    end
    return bm
end
boundary_matrix(cplx::AbstractComplex, d::Int) = boundary_matrix(Int, cplx, d)

#
# Complex simplex iterator
#
Base.length(splxs::Simplices{C}) where C<:AbstractComplex =
    splxs.dim < 0 ? sum(size(splxs.itr)) : size(splxs.itr, splxs.dim)

Base.eltype(iter::Simplices{C}) where C<:AbstractComplex = celltype(iter.itr)

# State (total id, dim, dim id)
function Base.iterate(iter::Simplices{C}, (tid, d, did)=(0, iter.dim, 1)) where C<:AbstractComplex
    tid >= length(iter) && return nothing
    if d < 0
        d = 0
    end
    if did > size(iter.itr, d)
        d += 1
        did = 1
    end
    return iter.itr[did, d], (tid+1, d, did+1)
end


@testset "AbstractTensor" begin
    struct Point{dim, T} <: AbstractVec{dim, T}
        x::NTuple{dim, T}
    end
    Base.Tuple(p::Point) = p.x
    Base.getindex(p::Point, i::Int) = p.x[i]
    for T in (Float32, Float64)
        x = Vec{2, T}(1, 2)
        p = Point(Tuple(x))
        @test (@inferred p + p)::Vec{2, T} |> Tuple == (x + x).data
        @test (@inferred p - p)::Vec{2, T} |> Tuple == (x - x).data
        @test (@inferred p ⋅ p)::T == x ⋅ x
        @test (@inferred p ⊗ p)::Mat{2, 2, T} == x ⊗ x
    end

    # size/axes
    TT = typeof(rand(Mat{3, 2}))
    @test (@inferred size(TT)) == (3, 2)
    @test (@inferred size(TT, 1)) == 3
    @test (@inferred size(TT, 2)) == 2
    @test (@inferred size(TT, 3)) == 1
    @test (@inferred size(TT, 4)) == 1
    @test (@inferred axes(TT)) == (Base.OneTo(3), Base.OneTo(2))
    @test (@inferred axes(TT, 1)) == Base.OneTo(3)
    @test (@inferred axes(TT, 2)) == Base.OneTo(2)
    @test (@inferred axes(TT, 3)) == Base.OneTo(1)
    @test (@inferred axes(TT, 4)) == Base.OneTo(1)
end

@testset "Operations for array interface" begin
    @testset "getindex for slicing" begin
        v = Vec(5,4,3,2,1)
        @test @inferred(v[:]) === Vec(5,4,3,2,1)
        @test @inferred(v[SOneTo(3)]) === Vec(5,4,3)
        @test @inferred(v[SUnitRange(2,4)]) === Vec(4,3,2)
        @test @inferred(v[SVector(3,4,1)]) === Vec(3,2,5)
        A = Mat{2,3}(1,2,3,4,5,6)
        @test @inferred(A[:]) === Vec(1,2,3,4,5,6)
        @test @inferred(A[1,:]) === Vec(1,3,5)
        @test A[end,:] === Vec(2,4,6)
        @test @inferred(A[:,1]) === Vec(1,2)
        @test A[:,end] === Vec(5,6)
        @test @inferred(A[1,SVector(1,3)]) === Vec(1,5)
        # by @Tensor macro
        @test @Tensor(A[:,[1,3]]) === @Mat [1 5; 2 6]
        @test @Tensor(A[1,:]) === Vec(1,3,5)
        @test @Tensor(A[end,:]) === Vec(2,4,6)
        @test @Tensor(A[1,2:3]) === Vec(3,5)
        @test_throws Exception @Tensor(Array(A)[1,2:3])
    end
    @testset "vcat/hcat" begin
        @test @inferred(vcat(Vec(1,2,3))) === Vec(1,2,3)
        @test @inferred(vcat(Vec(1,2,3)')) === Vec(1,2,3)'
        @test @inferred(vcat(Mat{1,3}(1,2,3))) === Mat{1,3}(1,2,3)
        @test @inferred(hcat(Vec(1,2,3))) === Mat{3,1}(1,2,3)
        @test @inferred(hcat(Vec(1,2,3)')) === Vec(1,2,3)'
        @test @inferred(hcat(Mat{3,1}(1,2,3))) === Mat{3,1}(1,2,3)

        @test @inferred(vcat(Vec(1,2,3), Vec(4,5,6))) === Vec(1,2,3,4,5,6)
        @test @inferred(vcat(Vec(1,2,3), Mat{3,1}(4,5,6))) === Mat{6,1}(1,2,3,4,5,6)
        @test @inferred(vcat(Mat{1,3}(1,2,3), Mat{1,3}(4,5,6))) === @Mat [1 2 3; 4 5 6]
        @test @inferred(vcat(Mat{1,3}(1,2,3), Vec(4,5,6)')) === @Mat [1 2 3; 4 5 6]
        @test @inferred(vcat(symmetric(Mat{2,2}(1,2,2,3), :U), Vec(4,5)')) === @Mat [1 2; 2 3; 4 5]
        @test @inferred(hcat(Vec(1,2,3), Vec(4,5,6))) === Mat{3,2}(1,2,3,4,5,6)
        @test @inferred(hcat(Vec(1,2,3)', Vec(4,5,6)')) === Mat{1,6}(1,2,3,4,5,6)
        @test @inferred(hcat(Vec(1,2,3), Mat{3,2}(4,5,6,7,8,9))) === Mat{3,3}(1,2,3,4,5,6,7,8,9)
        @test @inferred(hcat(Mat{3,1}(1,2,3), Mat{3,2}(4,5,6,7,8,9))) === Mat{3,3}(1,2,3,4,5,6,7,8,9)
        @test @inferred(hcat(symmetric(Mat{2,2}(1,2,2,3), :U), Vec(4,5))) === @Mat [1 2 4; 2 3 5]

        @test @inferred(vcat(Vec(1), Vec(2), Vec(3), Mat{2,1}(4, 5))) === Mat{5,1}(1,2,3,4,5)
        @test @inferred(vcat(Vec(1), Vec(2), Vec(3), Vec(4), Vec(5))) === Vec(1,2,3,4,5)
        @test @inferred(hcat(Vec(1), Vec(2), Vec(3), Vec(4), Vec(5))) === Mat{1,5}(1,2,3,4,5)

        # special cases for vcat with vector and real
        ## start from Vec
        @test @inferred(vcat(Vec(1), 2)) === Vec(1,2)
        @test @inferred(vcat(Vec(1), 2, Vec(3))) === Vec(1,2,3)
        ## start from Real
        n = 10
        for I in 1:n
            @test @inferred(vcat(1:I..., Vec(I+1))) === Vec((1:I+1)...)
            @test @inferred(vcat(1:I..., Vec(I+1), I+2)) === Vec((1:I+2)...)
            @test @inferred(vcat(1:I..., Vec(I+1), Vec(I+2))) === Vec((1:I+2)...)
        end
        n += 1
        @test @inferred(vcat(1:n..., Vec(n+1)))::Vector == 1:n+1
        @test @inferred(vcat(1:n..., Vec(n+1), n+2))::Vector == 1:n+2
        @test @inferred(vcat(1:n..., Vec(n+1), Vec(n+2)))::Vector == 1:n+2

        @test vcat(Vec(1.0f0), Vec(1.0)) === Vec(1.0, 1.0)
        @test hcat(Vec(1.0f0), Vec(1.0)) === Mat{1,2}(1.0, 1.0)
    end
    @testset "reverse" begin
        @test @inferred(reverse(Vec(1, 2, 3))) === Vec(3, 2, 1)
        A = rand(Mat{3,4})
        if VERSION ≥ v"1.6"
            @test @inferred(reverse(A))::typeof(A) == reverse(reverse(collect(A), dims = 1), dims = 2)
        end
    end
end


export AbstractParticle, Particle, MagneticParticle
####################################################
## Particles
####################################################
"""
    AbstractParticle
Particle supertype.
"""
abstract type AbstractParticle{T<:AbstractFloat} end
eltype(p::AbstractParticle{T}) where {T} = T

"""
```julia
Particle(ic::Vector{T}) #where ic = [x0, y0, φ0]
Particle(x0, y0, φ0)
Particle(pos::SVector, vel::SVector)
```
Create a particle with initial conditions `x0, y0, φ0`. It propagates as
a straight line.

The field `current_cell` shows at which cell of a periodic billiard is the particle
currently located.
"""
mutable struct Particle{T<:AbstractFloat} <: AbstractParticle{T}
    pos::SVector{2,T}
    vel::SVector{2,T}
    current_cell::SVector{2,T}
    function Particle(
        pos::SVector{2,T}, vel::SVector{2,T}, cc::SVector{2,T}) where{T<:AbstractFloat}
        new{T}(pos, normalize(vel), cc)
    end
end

Base.copy(p::Particle) =
Particle(p.pos, p.vel, p.current_cell)

function Particle(ic::AbstractVector{S}) where {S<:Real}
    T = S<:Integer ? Float64 : S
    φ0 = ic[3]
    pos = SVector{2,T}(ic[1:2]); vel = SVector{2,T}(cossin(φ0)...)
    return Particle(pos, vel, SVector{2,T}(0,0))
end
Particle(x::Real, y::Real, φ::Real) = Particle(collect(promote(x,y,φ)))
Particle() = Particle(rand(), rand(), rand()*2π)
function Particle(pos::SV{T}, vel::SV{T}) where {T}
    S = T<:Integer ? Float64 : T
    return Particle(pos, vel, SVector{2,S}(0.0, 0.0))
end
show(io::IO, p::Particle{T}) where {T} =
print(io, "Particle{$T}\n",
"position: $(p.pos+p.current_cell)\nvelocity: $(p.vel)")

"""
```julia
MagneticParticle(ic::AbstractVector{T}, ω::Real) # where ic = [x0, y0, φ0]
MagneticParticle(x0, y0, φ0, ω)
MagneticParticle(pos::SVector, vel::SVector, ω)
MagneticParticle(p::AbstractParticle, ω)
```
Create a *magnetic* particle with initial conditions `x0, y0, φ0` and angular
velocity `ω`. It propagates as a circle instead of a line,
with radius `1/abs(ω)`.

The field `current_cell` shows at which cell of a periodic billiard is the particle
currently located.
"""
mutable struct MagneticParticle{T<:AbstractFloat} <: AbstractParticle{T}
    pos::SVector{2,T}
    vel::SVector{2,T}
    current_cell::SVector{2,T}
    omega::T
    r::T
    center::SVector{2, T}
    function MagneticParticle(pos::SVector{2,T}, vel::SVector{2,T},
        current_cell::SVector{2,T}, ω::T) where {T<:AbstractFloat}
        if ω==0
            throw(ArgumentError("Angular velocity of magnetic particle cannot be 0."))
        end
        r = 1/ω
        c::SV{T} = pos - r*SV{T}(vel[2], -vel[1])
        new{T}(pos, normalize(vel), current_cell, ω, abs(1/ω), c)
    end
end

function Base.getproperty(p::MagneticParticle, s::Symbol)
    if s == :ω
        return Base.getfield(p, :omega)
    else
        return Base.getfield(p, s)
    end
end

Base.copy(p::MagneticParticle{T}) where {T} =
MagneticParticle(p.pos, p.vel, p.current_cell, p.omega)

function MagneticParticle(ic::AbstractVector{T}, ω::Real) where {T<:Real}
    φ0 = ic[3]
    S = T<:Integer ? Float64 : T
    pos = SVector{2,S}(ic[1:2]); vel = SVector{2,S}(cossin(φ0)...)
    return MagneticParticle(pos, vel, SVector{2,S}(0,0), convert(S,ω))
end
function MagneticParticle(x0::Real, y0::Real, φ0::Real, ω::Real)
    a = collect(promote(x0, y0, φ0, ω))
    MagneticParticle(a[1:3], a[4])
end
MagneticParticle() = MagneticParticle([rand(), rand(), rand()*2π], 1.0)
function MagneticParticle(pos::SV{T}, vel::SV{T}, ω::T) where {T}
    S = T<:Integer ? Float64 : T
    return MagneticParticle(pos, vel, SVector{2,S}(0,0), convert(S,ω))
end

MagneticParticle(p::AbstractParticle, ω) = MagneticParticle(p.pos, p.vel, p.current_cell, ω)

show(io::IO, p::MagneticParticle{T}) where {T} =
print(io, "MagneticParticle{$T}\n",
"position: $(p.pos+p.current_cell)\nvelocity: $(p.vel)\nang. velocity: $(p.omega)")

"""
    find_cyclotron(p::MagneticParticle)
Return the center of cyclotron motion of the particle.
"""
@inline find_cyclotron(p::MagneticParticle{T}) where {T} =
    c::SV{T} = p.pos - (1/p.omega)*SV{T}(p.vel[2], -p.vel[1])

@inline cyclotron(p) = p.center, p.r
function cyclotron(p, use_cell)
    pos = use_cell ? p.pos + p.current_cell : p.pos
    return pos, p.r
end


# This file can be used to create the validation file for deterministic scc values.
# It will create a validation file containing all possibilities of parameter values
# defined in the specs dictionary below produce the same results.

using MimiDICE2013
using DataFrames
using Query
using CSVFiles
using Test

specs = Dict([
    :year => [2020, 2050],
    :eta => [0, 1.5],
    :prtp => [0.015, 0.03],
    :last_year => [2200, 2305],
])

results = DataFrame(year = [], eta = [], prtp = [], last_year = [], SC = [])

for year in specs[:year]
    for eta in specs[:eta]
        for prtp in specs[:prtp]
            for last_year in specs[:last_year]
                sc = MimiDICE2013.compute_scc(year=Int(year), eta=eta, prtp=prtp, last_year=Int(last_year))
                push!(results, (year, eta, prtp, last_year, sc))
            end
        end
    end
end

path = joinpath(@__DIR__, "deterministic_sc_values_v1-0-2.csv")
save(path, results)


#
# Conjugate gradient
#
# This is an independent implementation of:
#   W. W. Hager and H. Zhang (2006) Algorithm 851: CG_DESCENT, a
#     conjugate gradient method with guaranteed descent. ACM
#     Transactions on Mathematical Software 32: 113–137.
#
# Code comments such as "HZ, stage X" or "HZ, eqs Y" are with
# reference to a particular point in this paper.
#
# Several aspects of the following have also been incorporated:
#   W. W. Hager and H. Zhang (2012) The limited memory conjugate
#     gradient method.
#
# This paper will be denoted HZ2012 below.
#
# There are some modifications and/or extensions from what's in the
# paper (these may or may not be extensions of the cg_descent code
# that can be downloaded from Hager's site; his code has undergone
# numerous revisions since publication of the paper):
#
# cgdescent: the termination condition employs a "unit-correct"
#   expression rather than a condition on gradient
#   components---whether this is a good or bad idea will require
#   additional experience, but preliminary evidence seems to suggest
#   that it makes "reasonable" choices over a wider range of problem
#   types.
#
# linesearch: the Wolfe conditions are checked only after alpha is
#   generated either by quadratic interpolation or secant
#   interpolation, not when alpha is generated by bisection or
#   expansion. This increases the likelihood that alpha will be a
#   good approximation of the minimum.
#
# linesearch: In step I2, we multiply by psi2 only if the convexity
#   test failed, not if the function-value test failed. This
#   prevents one from going uphill further when you already know
#   you're already higher than the point at alpha=0.
#
# both: checks for Inf/NaN function values
#
# both: support maximum value of alpha (equivalently, c). This
#   facilitates using these routines for constrained minimization
#   when you can calculate the distance along the path to the
#   disallowed region. (When you can't easily calculate that
#   distance, it can still be handled by returning Inf/NaN for
#   exterior points. It's just more efficient if you know the
#   maximum, because you don't have to test values that won't
#   work.) The maximum should be specified as the largest value for
#   which a finite value will be returned.  See, e.g., limits_box
#   below.  The default value for alphamax is Inf. See alphamaxfunc
#   for cgdescent and alphamax for linesearch_hz.

immutable ConjugateGradient{T, Tprep<:Union{Function, Void}, L<:Function} <: Optimizer
    eta::Float64
    P::T
    precondprep!::Tprep
    linesearch!::L
end
#= uncomment for v0.8.0
function ConjugateGradient(;
                           linesearch = LineSearches.hagerzhang!,
                           eta::Real = 0.4,
                           P::Any = nothing,
                           precondprep = (P, x) -> nothing)
    ConjugateGradient(Float64(eta),
                                 P, precondprep,
                                 linesearch)
end
=#
function ConjugateGradient(; linesearch! = nothing,
                             linesearch = LineSearches.hagerzhang!,
                             eta::Real = 0.4,
                             P::Any = nothing,
                             precondprep! = nothing,
                             precondprep = (P, x) -> nothing)

    linesearch = get_linesearch(linesearch!, linesearch)
    precondprep = get_precondprep(precondprep!, precondprep)
    ConjugateGradient(Float64(eta),
                      P, precondprep,
                      linesearch)
end

type ConjugateGradientState{T}
    @add_generic_fields()
    x_previous::Array{T}
    g::Array{T}
    g_previous::Array{T}
    f_x_previous::T
    y::Array{T}
    py::Array{T}
    pg::Array{T}
    s::Array{T}
    @add_linesearch_fields()
end

function initial_state{T}(method::ConjugateGradient, options, d, initial_x::Array{T})
    g = similar(initial_x)
    f_x = d.fg!(initial_x, g)
    pg = copy(g)
    @assert typeof(f_x) == T
    # Output messages
    if !isfinite(f_x)
        error("Must have finite starting value")
    end
    if !all(isfinite, g)
        @show g
        @show find(!isfinite.(g))
        error("Gradient must have all finite values at starting point")
    end

    # Determine the intial search direction
    #    if we don't precondition, then this is an extra superfluous copy
    #    TODO: consider allowing a reference for pg instead of a copy
    method.precondprep!(method.P, initial_x)
    A_ldiv_B!(pg, method.P, g)

    ConjugateGradientState("Conjugate Gradient",
                         length(initial_x),
                         copy(initial_x), # Maintain current state in state.x
                         f_x, # Store current f in state.f_x
                         1, # Track f calls in state.f_calls
                         1, # Track g calls in state.g_calls
                         0, # Track h calls in state.h_calls
                         copy(initial_x), # Maintain current state in state.x_previous
                         g, # Store current gradient in state.g
                         copy(g), # Store previous gradient in state.g_previous
                         T(NaN), # Store previous f in state.f_x_previous
                         similar(initial_x), # Intermediate value in CG calculation
                         similar(initial_x), # Preconditioned intermediate value in CG calculation
                         pg, # Maintain the preconditioned gradient in pg
                         -copy(pg), # Maintain current search direction in state.s
                         @initial_linesearch()...) # Maintain a cache for line search results in state.lsr
end

function update_state!{T}(df, state::ConjugateGradientState{T}, method::ConjugateGradient)
        lssuccess = true
        # Reset the search direction if it becomes corrupted
        dphi0 = vecdot(state.g, state.s)
        if dphi0 >= 0
            @simd for i in 1:state.n
                @inbounds state.s[i] = -state.pg[i]
            end
            dphi0 = vecdot(state.g, state.s)
            if dphi0 >= 0
                return true
            end
        end

        # Refresh the line search cache
        LineSearches.clear!(state.lsr)
        @assert typeof(state.f_x) == T
        @assert typeof(dphi0) == T
        push!(state.lsr, zero(T), state.f_x, dphi0)

        # Pick the initial step size (HZ #I1-I2)
        state.alpha, state.mayterminate, f_update, g_update =
          LineSearches.alphatry(state.alpha, df, state.x, state.s, state.x_ls, state.g_ls, state.lsr)
        state.f_calls, state.g_calls = state.f_calls + f_update, state.g_calls + g_update

        # Determine the distance of movement along the search line
        try
            state.alpha, f_update, g_update =
                method.linesearch!(df, state.x, state.s, state.x_ls, state.g_ls, state.lsr,
                                   state.alpha, state.mayterminate)
            state.f_calls, state.g_calls = state.f_calls + f_update, state.g_calls + g_update
        catch ex
            if isa(ex, LineSearches.LineSearchException)
                lssuccess = false
                state.f_calls, state.g_calls = state.f_calls + ex.f_update, state.g_calls + ex.g_update
                state.alpha = ex.alpha
                Base.warn("Linesearch failed, using alpha = $(state.alpha) and exiting optimization.")
            else
                rethrow(ex)
            end
        end

        # Maintain a record of previous position
        copy!(state.x_previous, state.x)

        # Update current position # x = x + alpha * s
        LinAlg.axpy!(state.alpha, state.s, state.x)

        # Maintain a record of the previous gradient
        copy!(state.g_previous, state.g)

        # Update the function value and gradient
        state.f_x_previous, state.f_x = state.f_x, df.fg!(state.x, state.g)
        state.f_calls, state.g_calls = state.f_calls + 1, state.g_calls + 1

        # Check sanity of function and gradient
        if !isfinite(state.f_x)
            error("Function must finite function values")
        end

        # Determine the next search direction using HZ's CG rule
        #  Calculate the beta factor (HZ2012)
        # -----------------
        # Comment on py: one could replace the computation of py with
        #    ydotpgprev = vecdot(y, pg)
        #    vecdot(y, py)  >>>  vecdot(y, pg) - ydotpgprev
        # but I am worried about round-off here, so instead we make an
        # extra copy, which is probably minimal overhead.
        # -----------------
        method.precondprep!(method.P, state.x)
        dPd = dot(state.s, method.P, state.s)
        etak::T = method.eta * vecdot(state.s, state.g_previous) / dPd
        @simd for i in 1:state.n
            @inbounds state.y[i] = state.g[i] - state.g_previous[i]
        end
        ydots = vecdot(state.y, state.s)
        copy!(state.py, state.pg)        # below, store pg - pg_previous in py
        A_ldiv_B!(state.pg, method.P, state.g)
        @simd for i in 1:state.n     # py = pg - py
           @inbounds state.py[i] = state.pg[i] - state.py[i]
        end
        betak = (vecdot(state.y, state.pg) - vecdot(state.y, state.py) * vecdot(state.g, state.s) / ydots) / ydots
        beta = max(betak, etak)
        @simd for i in 1:state.n
            @inbounds state.s[i] = beta * state.s[i] - state.pg[i]
        end
        (lssuccess == false) # break on linesearch error
end

update_g!(d, state, method::ConjugateGradient) = nothing


#
# diff.jl -
#
# Implement finite differences operators.
#
#-------------------------------------------------------------------------------
#
# This file is part of LazyAlgebra (https://github.com/emmt/LazyAlgebra.jl)
# released under the MIT "Expat" license.
#
# Copyright (c) 2017-2021 Éric Thiébaut.
#

module FiniteDifferences

export Diff

using MayOptimize
using LazyAlgebra
using LazyAlgebra.Foundations
import LazyAlgebra: apply!, vcreate, identical

using Base: @propagate_inbounds
import Base: show

const ArrayAxis = AbstractUnitRange{Int}
const ArrayAxes{N} = NTuple{N,ArrayAxis}

"""
    limits(r) -> (first(r), last(r))

yields the first and last value of the unit-range `r`.

"""
limits(r::AbstractUnitRange) = (first(r), last(r))

"""
    Diff([opt::MayOptimize.Vectorize,] n=1, dims=:)

yields a linear mapping that computes a finite difference approximation of the
`n`-order derivative along the dimension(s) specified by `dims`.  Arguments
`dims` is an integer, a tuple or a vector of integers specifying along which
dimension(s) to apply the operator or `:` to specify all dimensions.  If
multiple dimensions are specified, the result is as if the operator is applied
separately on the specified dimension(s).

Optional argument `opt` is the optimization level and may be specified as the
first or last argument.  By default, `opt` is assumed to be `Vectorize`,
however depending on the dimensions of the array, the dimensions of interest
and on the machine, setting `opt` to `InBounds` may be more efficient.

If `dims` is a scalar, the result, say `y`, of applying the finite difference
operator to an array, say `x`, has the same axes as `x`.  Otherwise and even
though `x` has a single dimension or `dims` is a 1-tuple, `y` has one more
dimension than `x`, the last dimension of `y` is used to store the finite
differences along each dimensions specified by `dims` and the leading
dimensions of `y` are the same as the dimensions of `x`.

More specifically, the operator created by `Diff` implements **forward finite
differences** with **flat boundary conditions**, that is to say extrapolated
entries are assumed equal to the nearest entry.

"""
struct Diff{L,D,O<:OptimLevel} <: LinearMapping end
# L = level of differentiation
# D = list of dimensions along which compute the differences
# O = optimization level

# Constructors.
function Diff(n::Integer = 1,
              dims::Union{Colon,Integer,Tuple{Vararg{Integer}},
                          AbstractVector{<:Integer}}=Colon(),
              opt::Type{<:OptimLevel} = Vectorize)
    return Diff{to_int(n), to_dims(dims), opt}()
end

function Diff(opt::Type{<:OptimLevel}, n::Integer = 1,
              dims::Union{Colon,Integer,Tuple{Vararg{Integer}},
                          AbstractVector{<:Integer}}=Colon())
    return Diff{to_int(n), to_dims(dims), opt}()
end

function Diff(n::Integer, opt::Type{<:OptimLevel})
    return Diff{to_int(n), Colon, opt}()
end

# Make a finite difference operator callable.
@callable Diff

# Two finite difference operators are identical if they have the same level of
# differentiation and list of dimensions along which compute the differences.
# Their optimization levels may be different.
identical(::Diff{L,D}, ::Diff{L,D}) where {L,D} = true

# Print operator in such a way that is similar to how the operator would be
# created in Julia.
show(io::IO, ::Diff{L,D,Opt}) where {L,D,Opt} =
    print(io, "Diff(", L, ',', (D === Colon ? ":" : D),',',
          (Opt === Debug ? "Debug" :
           Opt === InBounds ? "InBounds" :
           Opt === Vectorize ? "Vectorize" : Opt), ')')

"""
    differentiation_order(A)

yields the differentiation order of finite difference operator `A` (argument
can also be a type).

"""
differentiation_order(::Type{<:Diff{L,D,Opt}}) where {L,D,Opt} = L

"""
    dimensions_of_interest(A)

yields the list of dimensions of interest of finite difference operator `A`
(argument can also be a type).

"""
dimensions_of_interest(::Type{<:Diff{L,D,Opt}}) where {L,D,Opt} = D

"""
    optimization_level(A)

yields the optimization level for applying finite difference operator `A`
(argument can also be a type).

"""
optimization_level(::Type{<:Diff{L,D,Opt}}) where {L,D,Opt} = Opt

for f in (:differentiation_order,
          :dimensions_of_interest,
          :optimization_level)
    @eval begin
        $f(A::Diff) = $f(typeof(A))
        $f(A::Gram{<:Diff}) = $f(typeof(A))
        $f(::Type{<:Gram{T}}) where {T<:Diff} = $f(T)
    end
end

# Convert argument to `Int`.
to_int(x::Int) = x
to_int(x::Integer) = Int(x)

# Convert argument to the type parameter which specifies the list of dimensions
# of interest.
to_dims(::Colon) = Colon
to_dims(x::Int) = x
to_dims(x::Integer) = to_int(x)
to_dims(x::Tuple{Vararg{Int}}) = x
to_dims(x::Tuple{Vararg{Integer}}) = map(to_int, x)
to_dims(x::AbstractVector{<:Integer}) = to_dims((x...,))

# Drop list of dimensions from type to avoid unecessary specializations.
anydims(::Diff{L,D,P}) where {L,D,P} = Diff{L,Any,P}()
anydims(::Gram{Diff{L,D,P}}) where {L,D,P} = gram(Diff{L,Any,P}())

# Applying a separable operator is split in several stages:
#
# 1. Check arguments (so that avoiding bound checking should be safe) and deal
#    with the trivial cases α = 0 or no dimension of interest to apply the
#    operation (to simplify subsequent stages).
#
# 2. If α is non-zero, dispatch on dimension(s) along which to apply the
#    operation and on the specific values of the multipliers α and β.
#
# The second stage may be split in several sub-stages.

# Declare all possible signatures (not using unions) to avoid ambiguities.
for (P,A) in ((:Direct,  :Diff),
              (:Adjoint, :Diff),
              (:Direct,  :(Gram{<:Diff})))
    @eval function apply!(α::Number,
                          P::Type{$P},
                          A::$A,
                          x::AbstractArray,
                          scratch::Bool,
                          β::Number,
                          y::AbstractArray)
        inds, ndims = check_arguments(P, A, x, y)
        if α == 0 || ndims < 1
            # Get rid of this stupid case!
            vscale!(y, β)
        else
            # Call unsafe_apply! to dispatch on the dimensions of interest and on
            # the values of the multipliers.
            unsafe_apply!(α, P, A, x, β, y, inds)
        end
        return y
    end
end

# FIXME: This should not be necessary.
function apply!(α::Number,
                ::Type{<:Adjoint},
                A::Gram{<:Diff},
                x::AbstractArray,
                scratch::Bool,
                β::Number,
                y::AbstractArray)
    apply!(α, Direct, A, x, scratch, β, y)
end

function vcreate(::Type{Direct},
                 A::Diff{L,D,P},
                 x::AbstractArray{T,N},
                 scratch::Bool) where {L,D,P,T,N}
    if D === Colon
        return Array{T}(undef, size(x)..., N)
    elseif isa(D, Tuple{Vararg{Int}})
        return Array{T}(undef, size(x)..., length(D))
    elseif isa(D, Int)
        # if L === 1 && scratch && isa(x, Array)
        #    # First order finite difference along a single dimension.
        #    # Operation could be done in-place but we must preserve
        #    # type-stability.
        #    return x
        #else
        #    return Array{T}(undef, size(x))
        #end
        return Array{T}(undef, size(x))
    else
        error("invalid list of dimensions")
    end
end

function vcreate(::Type{Adjoint},
                 A::Diff{L,D,P},
                 x::AbstractArray{T,N},
                 scratch::Bool) where {L,D,P,T,N}
    # Checking the validity of the argument dimensions is done by applying the
    # opererator.  In-place operation never possible, so ignore the scratch
    # flag.
    if D === Colon || isa(D, Tuple{Vararg{Int}})
        return Array{T}(undef, size(x)[1:N-1])
    elseif isa(D, Int)
        return Array{T}(undef, size(x))
    else
        error("invalid list of dimensions")
    end
end

#------------------------------------------------------------------------------
# CHECKING OF ARGUMENTS

"""
    check_arguments(P, A, x, y) -> inds, ndims

checks that arguments `x` and `y` are valid for applying `P(A)`, with `A` a
separable operator, to `x` and store the result in `y`.  The result is a
2-tuple, `inds` is the axes that the arguments have in common and `ndims` is
the number of dimensions of interest.

If this function returns normally, the caller may safely assume that index
bound checking is not needed; hence, this function must throw an exception if
the dimensions/indices of `x` and `y` are not compatible or if the dimensions
of interest in `A` are out of range.  This function may also throw an exception
if the element types of `x` and `y` are not compatible.

This method must be specialized for the different types of separable operators.

"""
function check_arguments(P::Type{<:Union{Direct,Adjoint}},
                         A::Union{Diff{L,D},Gram{<:Diff{L,D}}},
                         x::AbstractArray,
                         y::AbstractArray) where {L,D}
    inds = check_axes(P, A, axes(x), axes(y))
    ndims = check_dimensions_of_interest(D, length(inds))
    return inds, ndims
end

function check_axes(P::Type{<:Union{Direct,Adjoint}},
                    A::Diff{L,D},
                    xinds::ArrayAxes,
                    yinds::ArrayAxes) where {L,D}
    if D === Colon || isa(D, Dims)
        if P === Direct
            length(yinds) == length(xinds) + 1 ||
                throw_dimension_mismatch("output array must have one more dimension than input array")
            N = (D === Colon ? length(xinds) : length(D))
            yinds[end] == 1:N ||
                throw_dimension_mismatch("last axis of output array must be 1:", N)
            yinds[1:end-1] == xinds ||
                throw_dimension_mismatch("leading axes must be identical")
            return xinds
        else
            length(yinds) == length(xinds) - 1 ||
                throw_dimension_mismatch("output array must have one less dimension than input array")
            N = (D === Colon ? length(yinds) : length(D))
            xinds[end] == 1:N ||
                throw_dimension_mismatch("last axis of input array must be 1:", N)
            xinds[1:end-1] == yinds ||
                throw_dimension_mismatch("leading axes must be identical")
            return yinds
        end
    elseif isa(D, Int)
        xinds == yinds || throw_dimension_mismatch("array axes must be identical")
        return xinds
    else
        throw(ArgumentError("invalid dimensions of interest"))
    end
end

function check_axes(P::Type{<:Operations},
                    A::Gram{<:Diff},
                    xinds::ArrayAxes,
                    yinds::ArrayAxes)
    xinds == yinds || throw_dimension_mismatch("array axes must be identical")
    return xinds
end

check_dimensions_of_interest(::Type{Colon}, ndims::Int) = ndims

check_dimensions_of_interest(dim::Int, ndims::Int) = begin
    1 ≤ dim ≤ ndims ||
        throw_dimension_mismatch("out of range dimension ", dim,
                                 "for ", ndims,"-dimensional arrays")
    return 1
end

check_dimensions_of_interest(dims::Dims{N}, ndims::Int) where {N} = begin
    for dim in dims
        1 ≤ dim ≤ ndims ||
            throw_dimension_mismatch("out of range dimension ", dim,
                                     "for ", ndims,"-dimensional arrays")
    end
    return N
end

throw_dimension_mismatch(str::String) = throw(DimensionMismatch(str))

@noinline throw_dimension_mismatch(args...) =
    throw_dimension_mismatch(string(args...))

#------------------------------------------------------------------------------

# Apply the operation along all dimensions of interest but one dimension at a
# time and knowing that α is not zero.
@generated function unsafe_apply!(α::Number,
                                  ::Type{P},
                                  A::Diff{L,D},
                                  x::AbstractArray,
                                  β::Number,
                                  y::AbstractArray,
                                  inds::ArrayAxes{N}) where {L,D,N,
                                                             P<:Union{Direct,
                                                                      Adjoint}}
    # Allocate empty vector of statements.
    exprs = Expr[]

    # Discard type parameter specifying the dimensions of interest to avoid
    # specialization on this parameter.
    push!(exprs, :(B = anydims(A)))

    # Dispatch on dimensions of interest.
    if isa(D, Int)
        # Arrays x and y have the same dimensions.
        push!(exprs, :(unsafe_apply!(α, P, B, x, β, y,
                                     inds[1:$(D-1)],
                                     inds[$D],
                                     inds[$(D+1):$N],
                                     CartesianIndex())))
    elseif D === Colon || isa(D, Dims)
        # One of x or y (depending on whether the direct or the adjoint
        # operator is applied) has an extra leading dimension used to store the
        # result computed along a given dimension.
        keep_beta = true # initially scale y by β
        dims = (D === Colon ? (1:N) : D)
        for l in 1:length(dims)
            d = dims[l]
            push!(exprs, :(unsafe_apply!(α, P, B, x,
                                         $(keep_beta ? :β : 1), y,
                                         inds[1:$(d-1)],
                                         inds[$d],
                                         inds[$(d+1):$N],
                                         CartesianIndex($l))))
            keep_beta = (P === Direct && A <: Diff)
        end
    else
        # This should never happen.
        return quote
            error("invalid list of dimensions of interest")
        end
    end

    return quote
        $(Expr(:meta, :inline))
        $(exprs...)
        nothing
    end
end

@generated function unsafe_apply!(α::Number,
                                  ::Type{P},
                                  A::Gram{<:Diff{L,D}},
                                  x::AbstractArray,
                                  β::Number,
                                  y::AbstractArray,
                                  inds::ArrayAxes{N}) where {L,D,N,
                                                             P<:Direct}
    # Allocate empty vector of statements.
    exprs = Expr[]

    # Discard type parameter specifying the dimensions of interest to avoid
    # specialization on this parameter.
    push!(exprs, :(B = anydims(A)))

    # Dispatch on dimensions of interest.  Arrays x and y have the same
    # dimensions and there is no last index `l` to specify.
    if isa(D, Int)
        push!(exprs, :(unsafe_apply!(α, P, B, x, β, y,
                                     inds[1:$(D-1)],
                                     inds[$D],
                                     inds[$(D+1):$N])))
    elseif D === Colon || isa(D, Dims)
        # β is set to 1 after first dimension of interest.
        dims = (D === Colon ? (1:N) : D)
        for l in 1:length(dims)
            d = dims[l]
            push!(exprs, :(unsafe_apply!(α, P, B, x,
                                         $(l == 1 ? :β : 1), y,
                                         inds[1:$(d-1)],
                                         inds[$d],
                                         inds[$(d+1):$N])))
        end
    else
        # This should never happen.
        return quote
            error("invalid list of dimensions of interest")
        end
    end

    return quote
        $(Expr(:meta, :inline))
        $(exprs...)
        nothing
    end
end

# Dispatch on multipliers values (α is not zero).
function unsafe_apply!(alpha::Number,
                       P::Type{<:Operations},
                       A::Union{Diff{L,Any,Opt},
                                Gram{Diff{L,Any,Opt}}},
                       x::AbstractArray,
                       beta::Number,
                       y::AbstractArray,
                       I::ArrayAxes,
                       J::ArrayAxis,
                       K::ArrayAxes,
                       l::CartesianIndex) where {L,Opt}
    if alpha == 1
        if beta == 0
            unsafe_apply!(axpby_yields_x,     1, P, A, x, 0, y, I, J, K, l)
        elseif beta == 1
            unsafe_apply!(axpby_yields_xpy,   1, P, A, x, 1, y, I, J, K, l)
        else
            β = promote_multiplier(beta, y)
            unsafe_apply!(axpby_yields_xpby,  1, P, A, x, β, y, I, J, K, l)
        end
    else
        α = promote_multiplier(alpha, y)
        if beta == 0
            unsafe_apply!(axpby_yields_ax,    α, P, A, x, 0, y, I, J, K, l)
        elseif beta == 1
            unsafe_apply!(axpby_yields_axpy,  α, P, A, x, 1, y, I, J, K, l)
        else
            β = promote_multiplier(beta, y)
            unsafe_apply!(axpby_yields_axpby, α, P, A, x, β, y, I, J, K, l)
        end
    end
    nothing
end

# Dispatch on multipliers values (α is not zero) for Gram compositions of a
# finite difference operator.
function unsafe_apply!(alpha::Number,
                       P::Type{<:Operations},
                       A::Gram{<:Diff},
                       x::AbstractArray,
                       beta::Number,
                       y::AbstractArray,
                       I::ArrayAxes,
                       J::ArrayAxis,
                       K::ArrayAxes)
    if alpha == 1
        if beta == 0
            unsafe_apply!(axpby_yields_x,     1, P, A, x, 0, y, I, J, K)
        elseif beta == 1
            unsafe_apply!(axpby_yields_xpy,   1, P, A, x, 1, y, I, J, K)
        else
            β = promote_multiplier(beta, y)
            unsafe_apply!(axpby_yields_xpby,  1, P, A, x, β, y, I, J, K)
        end
    else
        α = promote_multiplier(alpha, y)
        if beta == 0
            unsafe_apply!(axpby_yields_ax,    α, P, A, x, 0, y, I, J, K)
        elseif beta == 1
            unsafe_apply!(axpby_yields_axpy,  α, P, A, x, 1, y, I, J, K)
        else
            β = promote_multiplier(beta, y)
            unsafe_apply!(axpby_yields_axpby, α, P, A, x, β, y, I, J, K)
        end
    end
    nothing
end

#------------------------------------------------------------------------------
#
# The operator D implementing 1st order forward finite difference with flat
# boundary conditions and its adjoint D' are given by:
#
#     D = [ -1   1   0   0
#            0  -1   1   0
#            0   0  -1   1
#            0   0   0   0];
#
#     D' = [ -1   0   0   0
#             1  -1   0   0
#             0   1  -1   0
#             0   0   1   0];
#
# The row (for D) and column (for D') of zeros are to preserve the size.  This
# is needed for multi-dimensional arrays when derivatives along each dimension
# are stored into a single array.
#
# Apply 1st order finite differences along 1st dimension:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{Direct},
                       A::Diff{1,Any,Opt},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::Tuple{},
                       J::ArrayAxis,
                       K::ArrayAxes,
                       l::CartesianIndex) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin ≤ jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            @maybe_vectorized Opt for j in jmin:jmax-1
                z = x[j+1,k] - x[j,k]
                y[j,k,l] = f(α, z, β, y[j,k,l])
            end
            let j = jmax, z = zero(T)
                y[j,k,l] = f(α, z, β, y[j,k,l])
            end
        end
    end
    nothing
end
#
# Apply 1st order finite differences along 2nd and subsequent dimensions:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{Direct},
                       A::Diff{1,Any,Opt},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::ArrayAxes,
                       J::ArrayAxis,
                       K::ArrayAxes,
                       l::CartesianIndex) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin ≤ jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            for j in jmin:jmax-1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j+1,k] - x[i,j,k]
                    y[i,j,k,l] = f(α, z, β, y[i,j,k,l])
                end
            end
            let j = jmax, z = zero(T)
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k,l] = f(α, z, β, y[i,j,k,l])
                end
            end
        end
    end
    nothing
end
#
# Apply adjoint of 1st order finite differences along 1st dimension:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{Adjoint},
                       A::Diff{1,Any,Opt},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::Tuple{},
                       J::ArrayAxis,
                       K::ArrayAxes,
                       l::CartesianIndex) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin < jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                z = -x[j,k,l]
                y[j,k] = f(α, z, β, y[j,k])
            end
            @maybe_vectorized Opt for j in jmin+1:jmax-1
                z = x[j-1,k,l] - x[j,k,l]
                y[j,k] = f(α, z, β, y[j,k])
            end
            let j = jmax
                z = x[j-1,k,l]
                y[j,k] = f(α, z, β, y[j,k])
            end
        end
    elseif jmin == jmax && β != 1
        let j = jmin, z = zero(T)
            @maybe_vectorized Opt for k in CartesianIndices(K)
                y[j,k] = f(α, z, β, y[j,k])
            end
        end
    end
    nothing
end
#
# Apply adjoint of 1st order finite differences along 2nd and subsequent
# dimensions:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{Adjoint},
                       A::Diff{1,Any,Opt},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::ArrayAxes,
                       J::ArrayAxis,
                       K::ArrayAxes,
                       l::CartesianIndex) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin < jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = -x[i,j,k,l]
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
            for j in jmin+1:jmax-1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j-1,k,l] - x[i,j,k,l]
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
            let j = jmax
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j-1,k,l]
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
        end
    elseif jmin == jmax && β != 1
        let j = jmin, z = zero(T)
            @maybe_inbounds Opt for k in CartesianIndices(K)
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
        end
    end
    nothing
end
#
# The Gram composition D'*D of the 1st order forward finite differences D with
# flat boundary conditions writes:
#
#     D'*D = [  1  -1   0   0   0
#              -1   2  -1   0   0
#               0  -1   2  -1   0
#               0   0  -1   2  -1
#               0   0   0  -1   1 ]
#
# Apply D'*D along 1st dimension:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{<:Union{Direct,Adjoint}},
                       A::Gram{Diff{1,Any,Opt}},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::Tuple{},
                       J::ArrayAxis,
                       K::ArrayAxes) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin < jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                z = x[j,k] - x[j+1,k]
                y[j,k] = f(α, z, β, y[j,k])
            end
            @maybe_vectorized Opt for j in jmin+1:jmax-1
                z = T(2)*x[j,k] - (x[j-1,k] + x[j+1,k])
                y[j,k] = f(α, z, β, y[j,k])
            end
            let j = jmax
                z = x[j,k] - x[j-1,k]
                y[j,k] = f(α, z, β, y[j,k])
            end
        end
    elseif jmin == jmax && β != 1
        let j = jmin, z = zero(T)
            @maybe_vectorized Opt for k in CartesianIndices(K)
                y[j,k] = f(α, z, β, y[j,k])
            end
        end
    end
    nothing
end
#
# Apply  D'*D along 2nd and subsequent dimensions:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{<:Union{Direct,Adjoint}},
                       A::Gram{Diff{1,Any,Opt}},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::ArrayAxes,
                       J::ArrayAxis,
                       K::ArrayAxes) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin < jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j,k] - x[i,j+1,k]
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
            for j in jmin+1:jmax-1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = T(2)*x[i,j,k] - (x[i,j-1,k] + x[i,j+1,k])
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
            let j = jmax
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j,k] - x[i,j-1,k]
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
        end
    elseif jmin == jmax && β != 1
        let j = jmin, z = zero(T)
            @maybe_inbounds Opt for k in CartesianIndices(K)
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
        end
    end
    nothing
end

#------------------------------------------------------------------------------
#
# 2nd order finite differences with flat boundary conditions are computed by:
#
#     D = [-1   1   0   0   0   0
#           1  -2   1   0   0   0
#           0   1  -2   1   0   0
#           0   0   1  -2   1   0
#           0   0   0   1  -2   1
#           0   0   0   0   1  -1]
#
# Remarks:
#
#  - Applying this operator on a single dimension is self-adjoint.
#
#  - For a single dimension, this operator is the opposite of the Gram
#    composition of 1st order finite differences (backward or forward).
#
# Apply 2nd order finite differences along 1st dimension:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{Direct},
                       A::Diff{2,Any,Opt},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::Tuple{},
                       J::ArrayAxis,
                       K::ArrayAxes,
                       l::CartesianIndex) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin < jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                z = x[j+1,k] - x[j,k]
                y[j,k,l] = f(α, z, β, y[j,k,l])
            end
            @maybe_vectorized Opt for j in jmin+1:jmax-1
                z = x[j-1,k] + x[j+1,k] - T(2)*x[j,k]
                y[j,k,l] = f(α, z, β, y[j,k,l])
            end
            let j = jmax
                z = x[j-1,k] - x[j,k]
                y[j,k,l] = f(α, z, β, y[j,k,l])
            end
        end
    elseif jmin == jmax && β != 1
        let j = jmin, z = zero(T)
            @maybe_vectorized Opt for k in CartesianIndices(K)
                y[j,k,l] = f(α, z, β, y[j,k,l])
            end
        end
    end
    nothing
end
#
# Apply 2nd order finite differences along 2nd and subsequent dimensions:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{Direct},
                       A::Diff{2,Any,Opt},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::ArrayAxes,
                       J::ArrayAxis,
                       K::ArrayAxes,
                       l::CartesianIndex) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin < jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j+1,k] - x[i,j,k]
                    y[i,j,k,l] = f(α, z, β, y[i,j,k,l])
                end
            end
            for j in jmin+1:jmax-1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    # Other possibility:
                    # z = (x[i,j-1,k] - x[i,j,k]) + (x[i,j+1,k] - x[i,j,k])
                    z = x[i,j-1,k] + x[i,j+1,k] - T(2)*x[i,j,k]
                    y[i,j,k,l] = f(α, z, β, y[i,j,k,l])
                end
            end
            let j = jmax
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j-1,k] - x[i,j,k]
                    y[i,j,k,l] = f(α, z, β, y[i,j,k,l])
                end
            end
        end
    elseif jmin == jmax && β != 1
        let j = jmin, z = zero(T)
            @maybe_inbounds Opt for k in CartesianIndices(K)
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k,l] = f(α, z, β, y[i,j,k,l])
                end
            end
        end
    end
    nothing
end
#
# Apply adjoint of 2nd order finite differences along 1st dimension:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{Adjoint},
                       A::Diff{2,Any,Opt},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::Tuple{},
                       J::ArrayAxis,
                       K::ArrayAxes,
                       l::CartesianIndex) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin < jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                z = x[j+1,k,l] - x[j,k,l]
                y[j,k] = f(α, z, β, y[j,k])
            end
            @maybe_vectorized Opt for j in jmin+1:jmax-1
                # Other possibility:
                # z = (x[j-1,k,l] - x[j,k,l]) + (x[j+1,k,l] - x[j,k,l])
                z = x[j-1,k,l] + x[j+1,k,l] - T(2)*x[j,k,l]
                y[j,k] = f(α, z, β, y[j,k])
            end
            let j = jmax
                z = x[j-1,k,l] - x[j,k,l]
                y[j,k] = f(α, z, β, y[j,k])
            end
        end
    elseif jmin == jmax && β != 1
        let j = jmin, z = zero(T)
            @maybe_vectorized Opt for k in CartesianIndices(K)
                y[j,k] = f(α, z, β, y[j,k])
            end
        end
    end
    nothing
end
#
# Apply 2nd order finite differences along 2nd and subsequent dimensions:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{Adjoint},
                       A::Diff{2,Any,Opt},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::ArrayAxes,
                       J::ArrayAxis,
                       K::ArrayAxes,
                       l::CartesianIndex) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    if jmin < jmax
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j+1,k,l] - x[i,j,k,l]
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
            for j in jmin+1:jmax-1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j-1,k,l] + x[i,j+1,k,l] - T(2)*x[i,j,k,l]
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
            let j = jmax
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    z = x[i,j-1,k,l] - x[i,j,k,l]
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
        end
    elseif jmin == jmax && β != 1
        let j = jmin, z = zero(T)
            @maybe_inbounds Opt for k in CartesianIndices(K)
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
        end
    end
    nothing
end
#
# The Gram composition of 2nd order finite differences writes:
#
#    D'*D = [ 2  -3   1   0   0   0             (1)
#            -3   6  -4   1   0   0             (2)
#             1  -4   6  -4   1   0             (3)
#             0   1  -4   6  -4   1             (3)
#             0   0   1  -4   6  -3             (4)
#             0   0   0   1  -3   2]            (5)
#
# The above is for len ≥ 4, with len is the length of the dimension of
# interest, omitting the Eq. (5) for len = 4 and repeating Eq. (5) as necessary
# for the central rows for n ≥ 5.  For len = 3:
#
#    D'*D = [ 2  -3   1                         (1)
#            -3   6  -3                         (6)
#             1  -3   2]                        (5)
#
# For len = 2:
#
#    D'*D = [ 2  -2                             (7)
#            -2   2]                            (8)
#
# For len = 1, D = 0 and D'*D = 0 (the null 1×1 operator).
#
# Methods to apply the rows of D'D ():
#
# - Eq. (1), first row when len ≥ 3:
#
@inline @propagate_inbounds D2tD2_1(x::AbstractArray, j::Int, k) = begin
    T = real(eltype(x))
    T(2)*x[j,k] - T(3)*x[j+1,k] + x[j+2,k]
end
@inline @propagate_inbounds D2tD2_1(x::AbstractArray, i, j::Int, k) = begin
    T = real(eltype(x))
    T(2)*x[i,j,k] - T(3)*x[i,j+1,k] + x[i,j+2,k]
end
#
# - Eq. (2), second row when len ≥ 4:
#
@inline @propagate_inbounds D2tD2_2(x::AbstractArray, j::Int, k) = begin
    T = real(eltype(x))
    T(6)*x[j,k] - T(3)*x[j-1,k] - T(4)*x[j+1,k] + x[j+2,k]
end
@inline @propagate_inbounds D2tD2_2(x::AbstractArray, i, j::Int, k) = begin
    T = real(eltype(x))
    T(6)*x[i,j,k] - T(3)*x[i,j-1,k] - T(4)*x[i,j+1,k] + x[i,j+2,k]
end
#
# - Eq. (3), central rows when len ≥ 5:
#
@inline @propagate_inbounds D2tD2_3(x::AbstractArray, j::Int, k) = begin
    T = real(eltype(x))
    (x[j-2,k] + x[j+2,k]) + T(6)*x[j,k] - T(4)*(x[j-1,k] + x[j+1,k])
end
@inline @propagate_inbounds D2tD2_3(x::AbstractArray, i, j::Int, k) = begin
    T = real(eltype(x))
    (x[i,j-2,k] + x[i,j+2,k]) + T(6)*x[i,j,k] - T(4)*(x[i,j-1,k] + x[i,j+1,k])
end
#
# - Eq. (4), before last row when len ≥ 4:
#
@inline @propagate_inbounds D2tD2_4(x::AbstractArray, j::Int, k) = begin
    T = real(eltype(x))
    T(6)*x[j,k] - T(3)*x[j+1,k] - T(4)*x[j-1,k] + x[j-2,k]
end
@inline @propagate_inbounds D2tD2_4(x::AbstractArray, i, j::Int, k) = begin
    T = real(eltype(x))
    T(6)*x[i,j,k] - T(3)*x[i,j+1,k] - T(4)*x[i,j-1,k] + x[i,j-2,k]
end
#
# - Eq. (5), last row when len ≥ 3:
#
@inline @propagate_inbounds D2tD2_5(x::AbstractArray, j::Int, k) = begin
    T = real(eltype(x))
    T(2)*x[j,k] - T(3)*x[j-1,k] + x[j-2,k]
end
@inline @propagate_inbounds D2tD2_5(x::AbstractArray, i, j::Int, k) = begin
    T = real(eltype(x))
    T(2)*x[i,j,k] - T(3)*x[i,j-1,k] + x[i,j-2,k]
end
#
# - Eq. (6), central row when len = 3:
#
@inline @propagate_inbounds D2tD2_6(x::AbstractArray, j::Int, k) = begin
    T = real(eltype(x))
    T(6)*x[j,k] - T(3)*(x[j-1,k] + x[j+1,k])
end
@inline @propagate_inbounds D2tD2_6(x::AbstractArray, i, j::Int, k) = begin
    T = real(eltype(x))
    T(6)*x[i,j,k] - T(3)*(x[i,j-1,k] + x[i,j+1,k])
end
#
# - Eq. (7), first row when len = 2:
#
@inline @propagate_inbounds D2tD2_7(x::AbstractArray, j::Int, k) = begin
    z = x[j,k] - x[j+1,k]
    return z + z
end
@inline @propagate_inbounds D2tD2_7(x::AbstractArray, i, j::Int, k) = begin
    z = x[i,j,k] - x[i,j+1,k]
    return z + z
end
#
# - Eq. (8), last row when len = 2:
#
@inline @propagate_inbounds D2tD2_8(x::AbstractArray, j::Int, k) = begin
    z = x[j,k] - x[j-1,k]
    return z + z
end
@inline @propagate_inbounds D2tD2_8(x::AbstractArray, i, j::Int, k) = begin
    z = x[i,j,k] - x[i,j-1,k]
    return z + z
end
#
# Apply Gram composition of 2nd order finite differences along 1st dimension:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{<:Union{Direct,Adjoint}},
                       A::Gram{Diff{2,Any,Opt}},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::Tuple{},
                       J::ArrayAxis,
                       K::ArrayAxes) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    len = length(J)
    if len ≥ 5
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                y[j,k] = f(α, D2tD2_1(x,j,k), β, y[j,k])
            end
            let j = jmin+1
                y[j,k] = f(α, D2tD2_2(x,j,k), β, y[j,k])
            end
            @maybe_vectorized Opt for j in jmin+2:jmax-2
                y[j,k] = f(α, D2tD2_3(x,j,k), β, y[j,k])
            end
            let j = jmax-1
                y[j,k] = f(α, D2tD2_4(x,j,k), β, y[j,k])
            end
            let j = jmax
                y[j,k] = f(α, D2tD2_5(x,j,k), β, y[j,k])
            end
        end
    elseif len == 4
        @maybe_vectorized Opt for k in CartesianIndices(K)
            let j = jmin
                y[j,k] = f(α, D2tD2_1(x,j,k), β, y[j,k])
            end
            let j = jmin+1
                y[j,k] = f(α, D2tD2_2(x,j,k), β, y[j,k])
            end
            let j = jmax-1
                y[j,k] = f(α, D2tD2_4(x,j,k), β, y[j,k])
            end
            let j = jmax
                y[j,k] = f(α, D2tD2_5(x,j,k), β, y[j,k])
            end
        end
    elseif len == 3
        @maybe_vectorized Opt for k in CartesianIndices(K)
            let j = jmin
                y[j,k] = f(α, D2tD2_1(x,j,k), β, y[j,k])
            end
            let j = jmin+1
                y[j,k] = f(α, D2tD2_6(x,j,k), β, y[j,k])
            end
            let j = jmax
                y[j,k] = f(α, D2tD2_5(x,j,k), β, y[j,k])
            end
        end
    elseif len == 2
        @maybe_vectorized Opt for k in CartesianIndices(K)
            let j = jmin
                y[j,k] = f(α, D2tD2_7(x,j,k), β, y[j,k])
            end
            let j = jmax
                y[j,k] = f(α, D2tD2_8(x,j,k), β, y[j,k])
            end
        end
    elseif len == 1 && β != 1
        let j = jmin, z = zero(T)
            @maybe_vectorized Opt for k in CartesianIndices(K)
                y[j,k] = f(α, z, β, y[j,k])
            end
        end
    end
    nothing
end
#
# Apply Gram composition of 2nd order finite differences along 2nd and
# subsequent dimensions:
#
function unsafe_apply!(f::Function,
                       α::Number,
                       ::Type{<:Union{Direct,Adjoint}},
                       A::Gram{Diff{2,Any,Opt}},
                       x::AbstractArray,
                       β::Number,
                       y::AbstractArray,
                       I::ArrayAxes,
                       J::ArrayAxis,
                       K::ArrayAxes) where {Opt}
    T = real(eltype(x))
    jmin, jmax = limits(J)
    len = length(J)
    if len ≥ 5
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_1(x,i,j,k), β, y[i,j,k])
                end
            end
            let j = jmin+1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_2(x,i,j,k), β, y[i,j,k])
                end
            end
            for j in jmin+2:jmax-2
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_3(x,i,j,k), β, y[i,j,k])
                end
            end
            let j = jmax-1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_4(x,i,j,k), β, y[i,j,k])
                end
            end
            let j = jmax
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_5(x,i,j,k), β, y[i,j,k])
                end
            end
        end
    elseif len == 4
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_1(x,i,j,k), β, y[i,j,k])
                end
            end
            let j = jmin+1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_2(x,i,j,k), β, y[i,j,k])
                end
            end
            let j = jmax-1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_4(x,i,j,k), β, y[i,j,k])
                end
            end
            let j = jmax
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_5(x,i,j,k), β, y[i,j,k])
                end
            end
        end
    elseif len == 3
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_1(x,i,j,k), β, y[i,j,k])
                end
            end
            let j = jmin+1
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_6(x,i,j,k), β, y[i,j,k])
                end
            end
            let j = jmax
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_5(x,i,j,k), β, y[i,j,k])
                end
            end
        end
    elseif len == 2
        @maybe_inbounds Opt for k in CartesianIndices(K)
            let j = jmin
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_7(x,i,j,k), β, y[i,j,k])
                end
            end
            let j = jmax
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, D2tD2_8(x,i,j,k), β, y[i,j,k])
                end
            end
        end
    elseif len == 1 && β != 1
        let j = jmin, z = zero(T)
            @maybe_inbounds Opt for k in CartesianIndices(K)
                @maybe_vectorized Opt for i in CartesianIndices(I)
                    y[i,j,k] = f(α, z, β, y[i,j,k])
                end
            end
        end
    end
    nothing
end

end # module


@testset "monomial_basis" begin
    @testset "construction" begin
        @test Polynomial{4}(Polynomial(1, 2)).coeffs === (1, 2, 0, 0)
        @test Polynomial{4, Float64}(Polynomial(1, 2)).coeffs === (1.0, 2.0, 0.0, 0.0)

        p = Polynomial(ntuple(i -> 21 - i, Val(10)))
        Polynomial{20, Float64}(p)
        allocs = @allocated Polynomial{20, Float64}(p)
        @test allocs == 0

        @test Polynomial{2}(Polynomial(1, 2, 0)).coeffs === (1, 2)
        @test_throws InexactError Polynomial{2}(Polynomial(1, 2, 3))
    end

    @testset "evaluation" begin
        @test Polynomial(1)(2) == 1
        @test Polynomial(2, 3, 4)(5) == 2 + 3 * 5 + 4 * 5^2
    end

    @testset "utility" begin
        p = Polynomial(2, 3, 4)
        @test constant(p) === 2
        @test zero(p) === Polynomial(0, 0, 0)
        @test transpose(p) === p
        @test conj(p) === p
        @test conj(Polynomial(1 + 2im, 3 + 4im, 5 - 6im)) === Polynomial(1 - 2im, 3 - 4im, 5 + 6im)
    end

    @testset "padding" begin
        p = Polynomial(2, 3, 4)
        q = Polynomial{20}(p)
        @test p(5) == q(5)
        allocs = @allocated Polynomial{20}(p)
        @test allocs == 0
    end

    @testset "arithmetic" begin
        p1 = Polynomial(2, 3, 4)
        p2 = Polynomial(6, 7, 8)
        p3 = Polynomial(ntuple(i -> 21 - i, Val(20)))
        @test (p1 + p2)(7) == p1(7) + p2(7)
        allocs = @allocated p1 + p2
        @test allocs == 0
        @test (p1 + p3)(7) == p1(7) + p3(7)
        allocs = @allocated p1 + p3
        @test allocs == 0
        @test p1 + 1 === Polynomial(3, 3, 4)
        @test 1 - p2 === Polynomial(-5, -7, -8)
        @test +p1 === p1
        @test -p1 === Polynomial(-2, -3, -4)
        p4 = p1 * p3
        @test (p1 * p3)(3) == p1(3) * p3(3)
        allocs = @allocated p1 * p3
        @test allocs == 0
    end

    @testset "scaling" begin
        p1 = Polynomial(2, 3, 4)
        @test p1 * 4 === 4 * p1 === Polynomial(8, 12, 16)
        @test p1 / 2 === Polynomial(2 / 2, 3 / 2, 4 / 2)
        p2 = Polynomial(ntuple(i -> 21 - i, Val(20)))
        p2 * 4
        allocs = @allocated p2 * 4
        @test allocs == 0
    end

    @testset "derivative" begin
        p1 = Polynomial(3)
        @test derivative(p1) === Polynomial(0)

        p2 = Polynomial(1, 2, 3, 4)
        @test derivative(p2) === Polynomial(2, 2 * 3, 3 * 4)

        p3 = Polynomial(ntuple(i -> 21 - i, Val(20)))
        derivative(p3)
        allocs = @allocated derivative(p3)
        @test allocs == 0
    end

    @testset "integral" begin
        p = Polynomial(ntuple(i -> 21 - i, Val(20)))
        c = 5
        P = integral(p, c)
        @test constant(P) == c
        P′ = derivative(P)
        for x in 0 : 0.1 : 1
            @test P′(x) ≈ p(x) atol=1e-14
        end
    end

    @testset "gradient w.r.t coefficients" begin
        for x in [0, 1, 2]
            @test coefficient_gradient(x, Val(5)) === (1, x, x^2, x^3, x^4)
            @test coefficient_gradient(x, Val(5), Val(1)) === (0, 1, 2x, 3x^2, 4x^3)
            @test coefficient_gradient(x, Val(5), Val(2)) === (0, 0, 2, 6x, 12x^2)

            @test @allocated(coefficient_gradient(x, Val(5))) == 0
            @test @allocated(coefficient_gradient(x, Val(5), Val(1))) == 0
            @test @allocated(coefficient_gradient(x, Val(5), Val(2))) == 0
        end
    end

    @testset "exponential integral" begin
        p = Polynomial(1, 2, 3, 4)
        c = 5
        for t in range(0., 1., length=10)
            @test ForwardDiff.derivative(t -> exponential_integral(p, c, t), t) ≈ p(t) * exp(c * t) atol=1e-10
        end
    end

    @testset "broadcast" begin
        p1 = Polynomial(1, 2)
        p2 = Polynomial(3, 4)
        @test [p1, p2] .+ 1 == [Polynomial(2, 2), Polynomial(4, 4)]
    end
end


# Load ImplicitEquations

## meta these
preds = [(:Lt, :≪, :<), # \ll
         (:Le, :≦, :<=), # \leqq
         (:Eq, :⩵, :(==)), # \Equal
         (:Ge, :≧, :>=), # \gegg
         (:Gt, :≫, :>) # \gg
         ]

#import SymPy: Lt, Le, Eq, Ge, Gt, ≪, ≦, ⩵, ≧, ≫
for (fn, uop, op) in preds
    fnname =  string(fn)

    @eval begin
        ($fn)(f::Function, x::Real) = Pred(f, $op, x)
        ($uop)(f::Function, x::Real) = ($fn)(f, x)
        ($fn)(f::Function, g::Function) = $(fn)((x,y) -> f(x,y) - g(x,y), 0)
        ($uop)(f::Function, g::Function) = ($fn)(f, g)
    end
    eval(Expr(:export, fn))
    eval(Expr(:export, uop))
end
import ImplicitEquations: Neq
export Neq


using StatsFuns
using StatsBase
using Statistics
using Distributed
using Printf
using Crayons
using Base.Iterators
using Dates

# Probably will not be used in favor of the weighted mean/cov functions provided by
# StatsBase
function pmcmc_mean_cov(param_matrix::Array{Float64,2}, log_w::Vector{Float64})
    mean_vec = fill(0.0, size(param_matrix)[1])
    cov_mat = fill(0.0, size(param_matrix)[1], size(param_matrix)[1])
    total_weight = 0.0
    w_max = maximum(log_w)
    w_n = Vector{Float64}(undef, size(param_matrix)[2])
    for n = 1:size(param_matrix)[2]
        w_n[n] = exp(log_w[n] - w_max)
        total_weight += w_n[n]
        mean_vec .+= w_n[n] * (@view param_matrix[:, n])
    end
    mean_vec ./= total_weight
    p_vec = Vector{Float64}(undef, size(param_matrix)[1])
    for n = 1:size(param_matrix)[2]
        p_vec = param_matrix[:, n] .- mean_vec
        cov_mat += w_n[n] * (p_vec * p_vec')
    end
    cov_mat ./= total_weight
    return (mean_vec, cov_mat)
end

function pmcmc_mean_cov(params::Array{Array{Q,1},1}, log_w::Vector{Q}) where {Q<:Any}
    pmcmc_mean_cov(hcat(params...), log_w)
end

function pmcmc_mean_cov(params::Vector{T}, log_w::Vector{Q}) where {T<:NamedTuple,Q<:Any}
    pmcmc_mean_cov(hcat(map(x -> vcat(x...), params)...), log_w)
end

# Given a parameter (for dispatch), sample from a truncated MVT normal distribution
function trunc_normal_jump(
    x::Vector,
    mu::Vector,
    sigma::Matrix;
    lower_bound = (x, n) -> -Inf,
    upper_bound = (x, n) -> Inf,
    perturb = 1e-4,
)
    sigma = sigma + I * perturb
    A = Array{Float64,2}(undef, length(x) - 1, length(x))
    sigma_n = Array{Float64,1}(undef, length(x))
    n_x = length(x)
    x_current = copy(x)
    for n = 1:length(x)
        A = sigma[1:n_x.!=n, 1:n_x.!=n] \ sigma[1:n_x.!=n, n]
        sigma_n = sqrt(sigma[n, n] - (sigma[n, 1:n_x.!=n]' * A))
        mean_n = mu[n] .+ (A'*(x_current[1:n_x.!=n].-mu[1:n_x.!=n]))[1, 1]
        l = lower_bound(x_current, n)
        u = upper_bound(x_current, n)
        x_current[n] = rand(Truncated(Normal(mean_n, sigma_n), l, u))
    end
    return x_current
end

function trunc_normal_dens(
    x::Vector,
    x2::Vector,
    mu::Vector,
    sigma::Matrix;
    lower_bound = (x, n) -> -Inf,
    upper_bound = (x, n) -> Inf,
    perturb = 1e-4,
)
    sigma = sigma + I * perturb
    A = Array{Float64,2}(undef, length(x) - 1, length(x))
    sigma_n = Array{Float64,1}(undef, length(x))
    n_x = length(x)
    x_current = copy(x)
    ldens = 0.0
    for n = 1:length(x)
        A[:, n] = sigma[1:n_x.!=n, 1:n_x.!=n] \ sigma[1:n_x.!=n, n]
        sigma_n[n] = sqrt(sigma[n, n] - (sigma[n, 1:n_x.!=n]' * A[:, n]))
        mean_n = mu[n] .+ (A[:, n]'*(x_current[1:n_x.!=n].-mu[1:n_x.!=n]))[1, 1]
        l = lower_bound(x_current, n)
        u = upper_bound(x_current, n)
        ldens += logpdf(Truncated(Normal(mean_n, sigma_n[n]), l, u), x2[n])
        x_current[n] = x2[n]
    end
    return ldens
end

function smc_pmcmc_proposal(
    theta_current::Vector,
    pmcmc_mean::Vector,
    pmcmc_cov::Array;
    C = (2.38^2) / length(theta_current),
    kwargs...,
)
    return trunc_normal_jump(theta_current, theta_current, pmcmc_cov * C; kwargs...)
end

function smc_pmcmc_proposal_logdens(
    theta_current::Vector,
    theta_proposed::Vector,
    pmcmc_mean::Vector,
    pmcmc_cov::Matrix;
    C = (2.38^2) / length(theta_current),
    kwargs...,
)
    return trunc_normal_dens(
        theta_current,
        theta_proposed,
        theta_current,
        pmcmc_cov * C;
        kwargs...,
    )
end

# Utilities

function calculate_ess(xs)
    x_max = maximum(xs)
    x_sum = 0.0
    x_squared_sum = 0.0
    for n = 1:length(xs)
        exp_x_n = exp(xs[n] - x_max)
        x_sum += exp_x_n
        x_squared_sum += exp_x_n^2
    end

    return x_sum^2 / x_squared_sum
end

# This is a batch function which will run a filter from 1:T for a given vector of parameters in parallel

function dt_smc2_estimation(
    thetas,
    loglik_fun,
    prior_fun;
    grid_steps = 1e-8,
    ess_threshold = 0.5,
    pmcmc_theta_steps = 10,
    parallelize = true,
    kwargs...,
)
    N_θ = length(thetas)
    if parallelize
        ## Currently this pmap is type unstable, but this should
        ## not be too detrimental to performance
        logliks = pmap(loglik_fun, thetas)
    else
        logliks = map(loglik_fun, thetas)
    end
    ξ = 0.0

    acceptances = Array{Bool,3}(undef, N_θ, pmcmc_theta_steps, 0)

    while ξ < 1.0
        print(crayon"blue", "Picking new value for ξ\n")
        ξ_diff = grid_steps
        ess = calculate_ess(logliks * ξ_diff)

        ξ_diff, ess = xi_grid_search(logliks, ess_threshold)

        ξ = min(ξ_diff + ξ, 1.0)

        print(crayon"blue", @sprintf("Xi set to %.3e\nESS = %.3f\n", ξ, ess))

        thetas, logliks, acceptances_i = density_tempered_pmcmc(
            thetas,
            logliks,
            loglik_fun,
            prior_fun,
            ξ,
            ξ_diff;
            pmcmc_theta_steps = pmcmc_theta_steps,
            kwargs...,
        )
        acceptances = cat3(acceptances, acceptances_i)
    end

    return thetas, logliks, acceptances
end

function density_tempered_pmcmc(
    thetas,
    logliks,
    loglik_fun,
    prior_fun,
    ξ,
    ξ_diff;
    pmcmc_theta_steps = 10,
    rprop_pmh = smc_pmcmc_proposal,
    dprop_pmh = smc_pmcmc_proposal_logdens,
    kwargs...,
)
    N_θ = length(thetas)
    log_ws = logliks * ξ_diff

    ws = exp.(log_ws .- maximum(log_ws))
    pmcmc_mean, pmcmc_cov = pmcmc_mean_cov(thetas, log_ws)
    println("Param" * repeat(" ", 11) * "\tMean\t\tStd\t\t")
    for i = 1:length(pmcmc_mean)
        name = first(string(keys(thetas[1])[i]), 16)
        name *= repeat(" ", 17 - length(name))
        p_str = @sprintf("%s\t%.3e\t%.3e", name, pmcmc_mean[i], sqrt(pmcmc_cov[i, i]))
        println(p_str)
    end

    # A                 = wsample(1:N_θ, ws, N_θ)
    A = residual_resample(Random.GLOBAL_RNG, ws, N_θ)
    thetas = thetas[A]
    logliks = logliks[A]

    acceptances = Array{Bool,2}(undef, N_θ, pmcmc_theta_steps)
    for step_idx = 1:pmcmc_theta_steps
        starttime = Dates.now()
        res = pmap(
            pmcmc_propose_accept_reject,
            thetas,
            logliks,
            cycle([loglik_fun]),
            cycle([prior_fun]),
            cycle([ξ]),
            cycle([pmcmc_mean]),
            cycle([pmcmc_cov]),
            cycle([rprop_pmh]),
            cycle([dprop_pmh]),
        )
        for (i, r) in enumerate(res)
            thetas[i] = r[1]
            logliks[i] = r[2]
            acceptances[i, step_idx] = r[3]
        end
        timetaken = Dates.now() - starttime
        print(
            crayon"green",
            @sprintf(
                "Step %02d: Accepted %0.2f; Filter time: %f secs \n",
                step_idx,
                mean(acceptances[:, step_idx]),
                timetaken.value / 1000
            )
        )
    end
    print(
        crayon"green",
        @sprintf("Total θ moved at least once: %0.2f \n", mean(any(acceptances, dims = 2)))
    )
    return thetas, logliks, acceptances
end

function pmcmc_propose_accept_reject(
    theta,
    loglik,
    loglik_fun,
    prior_fun,
    ξ,
    pmcmc_mean,
    pmcmc_cov,
    rprop_pmh,
    dprop_pmh,
)
    # Propose new parameter
    theta_prop = rprop_pmh(theta, pmcmc_mean, pmcmc_cov)
    d_prop = dprop_pmh(theta, theta_prop, pmcmc_mean, pmcmc_cov)
    d_current = dprop_pmh(theta_prop, theta, pmcmc_mean, pmcmc_cov)

    # Evaluate likelihood of parameter
    loglik_prop = loglik_fun(theta_prop)

    # Evaluate prior of parameter
    prior = prior_fun(theta)
    prior_prop = prior_fun(theta_prop)

    # Calculate log acceptance probability
    log_prior_ratio = prior_prop - prior
    log_lik_ratio = ξ * (loglik_prop - loglik)
    log_prop_ratio = d_current - d_prop

    log_posterior_proposed = loglik_prop + prior_prop
    log_acceptance = log_prior_ratio + log_lik_ratio + log_prop_ratio
    accepted = (log_posterior_proposed > -Inf && log(rand()) < log_acceptance)

    if accepted
        theta = theta_prop
        loglik = loglik_prop
    end

    return (theta, loglik, accepted)
end

function xi_grid_search(w::AbstractArray{Float64}, ess_threshold = 0.5, depth = 100)
    ξ = (2.0)^(-1.0)
    for d = 2:depth
        ess = calculate_ess(w * ξ) / length(w)
        if ess < ess_threshold
            ξ -= (2.0)^(-d)
        else
            ξ += (2.0)^(-d)
        end
    end

    ess = calculate_ess(w * ξ) / length(w)
    if (ess < ess_threshold)
        ξ -= (2.0)^(-depth)
    end

    if ξ == 0
        println("Less than ess_threshold of observations are zero")
        ξ = 1e-12
    end

    return ξ, ess
end


# Make them inputs to a CG type
abstract type CGUpdate end
struct ConjugateGradient{Tu}
  update::Tu
end
ConjugateGradient(;update=HZ()) = ConjugateGradient(update)

#### Conjugate Descent [Fletcher] (CD)
struct CD <: CGUpdate end
_β(::CD, d, ∇fz, ∇fx, y) = -norm(∇fz)^2/dot(d, ∇fx)

#### Hager-Zhang (HZ)
struct HZ{Tη} <: CGUpdate
	η::Tη # a "forcing term"
end
HZ() = HZ(1/100)
function _β(cg::HZ, d, ∇fz, ∇fx, y)
	T = eltype(∇fz)
	η = T(cg.η)
	ddy = dot(d, y)
	βN = dot(y.-T(2).*d.*norm(y)^2/ddy, ∇fz)/ddy
	ηk = -inv(norm(d)*min(η, norm(∇fx)))
	βkp1 = max(βN, ηk)
end

#### Hestenes-Stiefel (HS)
struct HS <: CGUpdate end
_β(::HS, d, ∇fz, ∇fx, y) = dot(∇fz, y)/dot(d, y)

#### Fletcher-Reeves (FR)
struct FR <: CGUpdate end
_β(::FR, d, ∇fz, ∇fx, y) = norm(∇fz)^2/norm(∇fx)^2

#### Polak-Ribiére-Polyak (PRP)
struct PRP{Plus} end
PRP(;plus=true) = PRP{plus}()
_β(::PRP, d, ∇fz, ∇fx, y) = dot(∇fz, y)/norm(∇fx)^2
function _β(::PRP{true}, d, ∇fz, ∇fx, y)
  β = dot(∇fz, y)/norm(∇fx)^2
  max(typeof(β)(0), β)
end

#### Liu-Storey (LS)
struct LS <: CGUpdate end
_β(::LS, d, ∇fz, ∇fx, y) = dot(∇fz, y)/dot(d, ∇fx)

#### Dai-Yuan (DY)
struct DY <: CGUpdate end
_β(::DY, d, ∇fz, ∇fx, y) = norm(∇fz)^2/dot(d, y)

#### Wei-Yao-Liu (VPRP)
struct VPRP <: CGUpdate end
_β(::VPRP, d, ∇fz, ∇fx, y) = (norm(∇fz)^2-norm(∇fz)/norm(∇fx)*dot(∇fz, ∇fx))/dot(d, y)

# using an "intial vectors" function we can initialize s if necessary or nothing if not to save on vectors
function minimize(obj::OnceDiffed, x0, cg::ConjugateGradient, linesearch = HZAW(); maxiter=100)
    Tx=eltype(x0)

    fx, ∇fx = obj(x0, copy(x0))
	Tf = typeof(fx)

    fz = fx
	α = Tx(0)

	d, ∇fz = -copy(∇fx), copy(∇fx)
    x, y = copy(x0), copy(∇fz)
	z = copy(x)
	k = 0
    while norm(∇fx, Inf) ≥ 1e-8 && k ≤ maxiter
		# Set up line search
        α_0 = initial(cg.update, a-> obj(x.+a.*d), α, k, x, fx, dot(d, ∇fx), ∇fx)
		φ = LineObjective!(obj, ∇fz, z, x, d, fx, dot(∇fx, d))

		# Perform line search
		α, f_α, ls_success = find_steplength(linesearch, φ, Tf(1.0))
		# move in direction
		@. z = x + α*d
        fz, ∇fz = obj(z, ∇fz)

        @. y = ∇fz - ∇fx

        βkp1 = _β(cg.update, d, ∇fz, ∇fx, y)
        fx = fz
        x .= z
        ∇fx .= ∇fz

        @. d = -∇fx + βkp1*d
        k = k+1
    end
    x, fx, ∇fx, k
end

function initial(cg, φ, α, k, x, φ₀, dφ₀, ∇fx)
    T = eltype(x)
    ψ₀ = T(0.01)
    ψ₁ = T(0.1)
    ψ₂ = T(2.0)
    quadstep = true
    if k == 0
        if !all(x .≈ T(0)) # should we define "how approx we want?"
            return ψ₀ * norm(x, Inf)/norm(∇fx, Inf)
        elseif !(φ₀ ≈ T(0))
            return ψ₀ * abs(φ₀)/norm(∇fx, Inf)
        else
            return T(1)
        end
    elseif quadstep
        R = ψ₁*α
        φR = φ(R)
        c = (φR - φ₀ - dφ₀*R)/R^2

        if c > 0
	       return -dφ₀/(T(2)*c) # > 0 by df0 < 0 and c > 0
        end
    end
    return ψ₂*α
end

# function _β(method, ∇fx, ∇fz)
#
# if method == :fr
# return dot(∇fz, ∇fz)/dot(∇fx, ∇fx)
# elseif method == :pr
# β = dot(∇fz, s)/dot(∇fx, ∇fx) # dot(∇fx, ∇fx) can be saved between iteratoins only s is needed
# return    max(β, 0) # We need this according to Powell 1986 Gilbert Powell 1980 show it is globally  convergenct for exact and inexact
# # pr+ max(beta, 0)
# elseif method ==:hs
# dot(∇fz, s)/dot(s,p)
# end


# greedy seeds heuristic by A. Catalan and M. Fisher (2012) doi:10. 2139/ssrn.2166687
function GREEDYSEEDS(trans::SparseMatrixCSC{Float64, Int64},
                     Q::Matrix{Int64},
                     capacity::Array{Int64,1})
    if CHECKCAPACITY(Q::Array{Int64,2}, capacity::Array{Int64,1}) == 1
        ## Calculate the coapperance matrix
        capacity_left = copy(capacity) .- 1
        ## Sort the SKUs by decreasing sales
        sales = SORTSALES(trans)
        ## Assign the top (best selling) SKU to the largest DC
        X = Array{Int64,2}(undef,size(trans,2),size(capacity,1)) .= 0
        X[sales[1,1],1] = 1
        ## Mark the SKU that is now allocated and the capacity used
        sales[1,4] = 1
        ## Start the assignment of the other SKUs
        GREEDYSEEDSTART!(sales::Array{Int64,2}, X::Array{Int64,2}, Q::Array{Int64,2}, capacity_left::Array{Int64,1})
        ## For each SKUs sorted by decreasing sales that has not been allocated yet
        GREEDYSEEDMAIN!(sales::Array{Int64,2}, X::Array{Int64,2}, Q::Array{Int64,2}, capacity_left::Array{Int64,1})
        ## While there is unallocated space in the DCs do
        X = FILLUP(X::Array{Int64,2}, Q::Array{Int64,2}, capacity_left::Array{Int64,1})
        ## Return the solution
        return X::Array{Int64,2}
    end
end


"""
    LegendreUnitCoeff{T}

Precomputed recursion relation coefficients for the standard unit
normalization. Alias for `LegendreNormCoeff{LegendreUnitNorm,T}`.
"""
LegendreUnitCoeff{T} = LegendreNormCoeff{LegendreUnitNorm,T}

"""
    LegendreFourPiCoeff{T}

Table type of precomputed recursion relation coefficients for the ``4\\pi`` spherical
harmonic normalization. Alias for `LegendreNormCoeff{LegendreFourPiNorm,T}`.
"""
LegendreFourPiCoeff{T} = LegendreNormCoeff{LegendreFourPiNorm,T}

"""
    LegendreOrthoCoeff{T}

Table type of precomputed recursion relation coefficients for the orthonormal
normalization. Alias for `LegendreNormCoeff{LegendreOrthoNorm,T}`.
"""
LegendreOrthoCoeff{T} = LegendreNormCoeff{LegendreOrthoNorm,T}

"""
    LegendreSphereCoeff{T}

Table type of precomputed recursion relation coefficients for the spherical
harmonic normalization. Alias for `LegendreNormCoeff{LegendreSphereNorm,T}`.
"""
LegendreSphereCoeff{T} = LegendreNormCoeff{LegendreSphereNorm,T}

"""
    p = Plm(l, m, x)

Computes the associated Legendre polynomials using unit normalization;
equivalent to `p = legendre(LegendreUnitNorm(), l, m, x)`.
"""
const Plm = LegendreUnitNorm()

"""
    Plm!(P, l, m, x)

Fills the array `P` with the unit-normalized associated Legendre polynomial values
``P_ℓ^m(x)``;
equivalent to `legendre!(LegendreUnitNorm(), P, l, m, x)`.
"""
const Plm! = LegendreUnitNorm()

"""
    λ = λlm(l, m, x)

Computes the associated Legendre polynomials using spherical-harmonic normalization;
equivalent to `λ = legendre(LegendreSphereNorm(), l, m, x)`.
"""
const λlm = LegendreSphereNorm()

"""
    λlm!(Λ, l, m, x)

Fills the array `Λ` with the spherical-harmonic normalized associated Legendre polynomial
values ``λ_ℓ^m(x)``;
equivalent to `legendre!(LegendreSphereNorm(), P, l, m, x)`.
"""
const λlm! = LegendreSphereNorm()