# # Sensitivity Analysis of SVM

#md # [![](https://img.shields.io/badge/show-github-579ACA.svg)](@__REPO_ROOT_URL__/docs/src/examples/sensitivity-analysis-svm.jl)

# This notebook illustrates sensitivity analysis of data points in a [Support Vector Machine](https://en.wikipedia.org/wiki/Support-vector_machine) (inspired from [@matbesancon](http://github.com/matbesancon)'s [SimpleSVMs](http://github.com/matbesancon/SimpleSVMs.jl).)

# For reference, Section 10.1 of https://online.stat.psu.edu/stat508/book/export/html/792 gives an intuitive explanation of what it means to have a sensitive hyperplane or data point. The general form of the SVM training problem is given below (with $\ell_2$ regularization):

# ```math
# \begin{split}
# \begin{array} {ll}
# \mbox{minimize} & \lambda||w||^2 + \sum_{i=1}^{N} \xi_{i} \\
# \mbox{s.t.} & \xi_{i} \ge 0 \quad \quad i=1..N  \\
#             & y_{i} (w^T X_{i} + b) \ge 1 - \xi_{i} \quad i=1..N  \\
# \end{array}
# \end{split}
# ```
# where
# - `X`, `y` are the `N` data points
# - `w` is the support vector
# - `b` determines the offset `b/||w||` of the hyperplane with normal `w`
# - `ξ` is the soft-margin loss
# - `λ` is the $\ell_2$ regularization.
#
# This tutorial uses the following packages

using JuMP     # The mathematical programming modelling language
import DiffOpt # JuMP extension for differentiable optimization
import Ipopt   # Optimization solver that handles quadratic programs
import Plots   # Graphing tool
import LinearAlgebra: dot, norm
import Random

# ## Define and solve the SVM

# Construct two clusters of data points.

N = 100
D = 2

Random.seed!(62)
X = vcat(randn(N ÷ 2, D), randn(N ÷ 2, D) .+ [2.0, 2.0]')
y = append!(ones(N ÷ 2), -ones(N ÷ 2))
λ = 0.05;

# Let's initialize a special model that can understand sensitivities

model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
MOI.set(model, MOI.Silent(), true)

# Add the variables

@variable(model, ξ[1:N])
@variable(model, w[1:D])
@variable(model, b);

# Add the constraints.

@constraint(
    model,
    [i in 1:N],
    ξ[i] >= 0
);
@constraint(
    model,
    cons[i in 1:N],
    y[i] * (dot(X[i, :], w) + b) >= 1 - ξ[i]
);

# Define the objective and solve

@objective(
    model,
    Min,
    λ * dot(w, w) + sum(ξ),
)

optimize!(model)

# We can visualize the separating hyperplane.

loss = objective_value(model)

wv = value.(w)

bv = value(b)

svm_x = [-2.0, 4.0] # arbitrary points
svm_y = (-bv .- wv[1] * svm_x )/wv[2]

p = Plots.scatter(X[:,1], X[:,2], color = [yi > 0 ? :red : :blue for yi in y], label = "")
Plots.plot!(p, svm_x, svm_y, label = "loss = $(round(loss, digits=2))", width=3)

# ## Gradient of hyperplane wrt the data point coordinates

# Now that we've solved the SVM, we can compute the sensitivity of optimal
# values -- the separating hyperplane in our case -- with respect to
# perturbations of the problem data -- the data points -- using DiffOpt.

# How does a change in coordinates of the data points, `X`,
# affects the position of the hyperplane?
# This is achieved by finding gradients of `w` and `b` with respect to `X[i]`.

# Begin differentiating the model.
# analogous to varying θ in the expression:
# ```math
# y_{i} (w^T (X_{i} + \theta) + b) \ge 1 - \xi_{i}
# ```
∇ = zeros(N)
for i in 1:N
    for j in 1:N
        if i == j
            ## we consider identical perturbations on all x_i coordinates
            MOI.set(
                model,
                DiffOpt.ForwardInConstraint(),
                cons[j],
                y[j] * sum(w),
            )
        else
            MOI.set(model, DiffOpt.ForwardInConstraint(), cons[j], 0.0 * sum(w))
        end
    end
    DiffOpt.forward(model)
    dw = MOI.get.(
        model,
        DiffOpt.ForwardOutVariablePrimal(),
        w,
    )
    db = MOI.get(
        model,
        DiffOpt.ForwardOutVariablePrimal(),
        b,
    )
    ∇[i] = norm(dw) + norm(db)
end

# We can visualize the separating hyperplane sensitivity with respect to the data points.
# Note that all the small numbers were converted into 1/10 of the
# largest value to show all the points of the set.

p3 = Plots.scatter(
    X[:,1], X[:,2],
    color = [yi > 0 ? :red : :blue for yi in y], label = "",
    markersize = 2 * (max.(1.8∇, 0.2 * maximum(∇))),
)
Plots.yaxis!(p3, (-2, 4.5))
Plots.plot!(p3, svm_x, svm_y, label = "", width=3)
Plots.title!("Sensitivity of the separator to data point variations")


import Base.mean

export mean, mean_rb

# Mean of MCMCChain
mean(c::MCMCChain, par::Ranges=1:ncol(c.samples)) = Float64[mean(c.samples[:, i]) for i = par]

mean(c::MCMCChain, par::Real) = mean(c, par:par)

# Mean of MCMCChain using Rao-Blackwell samples
function mean_rb_hmc(c::MCMCChain, par::Ranges=1:ncol(c.samples))
  nsamples, npars, nleaps = nrow(c.samples), length(par), length(c.diagnostics["leaps"][1])-1

  w = Array(Float64, nsamples, nleaps)
  sums = Array(Float64, nsamples, npars)

  for i = 1:nsamples
    for j = 1:nleaps
      w[i, j] = exp(c.diagnostics["leaps"][i][1].H-c.diagnostics["leaps"][i][j+1].H)
    end
  end

  for i = 1:nsamples
    for j = 1:npars
      s = c.samples[i, j]
      for k = 1:nleaps
        s += w[i, k]*c.diagnostics["leaps"][i][k+1].pars[j]
      end

      sums[i, j] = s/(nleaps+1)
    end
  end

  mean(sums, 1)
end

mean_rb_hmc(c::MCMCChain, par::Real) = mean_rb_hmc(c, par:par)

function mean_rb(c::MCMCChain, p::Union(Real, Ranges)=1:ncol(c.samples), s::Symbol=:hmc)
  if s == :hmc
    mean_rb_hmc(c, p)
  end
end


module _LinAlg

using GenericLinearAlgebra: svd!
using LinearAlgebra: norm, lmul!, rmul!, triu!, Givens, I, SVD, reflector!, reflectorApply!,
    QRPivoted, QRPackedQ

export tsvd, tsvd!, svd_jacobi, svd_jacobi!, rrqr, rrqr!

"""
Truncated rank-revealing QR decomposition with full column pivoting.

Decomposes a `(m, n)` matrix `A` into the product:

    A[:,piv] == Q * R

where `Q` is an `(m, k)` isometric matrix, `R` is a `(k, n)` upper
triangular matrix, `piv` is a permutation vector, and `k` is chosen such
that the relative tolerance `tol` is met in the equality above.
"""
function rrqr!(A::AbstractMatrix{T}; rtol=eps(T)) where {T<:AbstractFloat}
    # DGEQPF
    m, n = size(A)
    k = min(m, n)
    Base.require_one_based_indexing(A)

    jpvt = Vector(1:n)
    taus = Vector{T}(undef, k)
    swapcol = Vector{T}(undef, m)

    xnorms = map(norm, eachcol(A))
    pnorms = copy(xnorms)
    sqrteps = sqrt(eps(T))

    @inbounds for i in 1:k
        pvt = argmax(ii -> pnorms[ii], i:n)
        if i ≠ pvt
            jpvt[i], jpvt[pvt] = jpvt[pvt], jpvt[i]
            xnorms[pvt] = xnorms[i]
            pnorms[pvt] = pnorms[i]

            swapcol .= @view A[:, i]
            A[:, i] .= @view A[:, pvt]
            A[:, pvt] .= swapcol
        end

        tau_i = reflector!(@view A[i:end, i])
        taus[i] = tau_i
        reflectorApply!(
            (@view A[i:end, i]), tau_i, @view A[i:end, (i + 1):end]
        )

        # Lapack Working Note 176.
        for j in (i + 1):n
            temp = abs(A[i, j]) / pnorms[j]
            temp = max(zero(T), (one(T) + temp) * (one(T) - temp))
            temp2 = temp * abs2(pnorms[j] / xnorms[j])
            if temp2 < sqrteps
                recomputed = norm(@view A[(i + 1):end, j])
                pnorms[j] = recomputed
                xnorms[j] = recomputed
            else
                pnorms[j] *= sqrt(temp)
            end
        end

        # Since we did pivoting, R[i,:] is bounded by R[i,i], so we can
        # simply ignore all the other rows
        if abs(A[i, i]) < rtol * abs(A[1, 1])
            A[i:end, i:end] .= zero(T)
            taus[i:end] .= zero(T)
            k = i - 1
            break
        end
    end
    return QRPivoted{T,typeof(A)}(A, taus, jpvt), k
end

"""Truncated rank-revealing QR decomposition with full column pivoting."""
rrqr(A::AbstractMatrix{T}; rtol=eps(T)) where {T<:AbstractFloat} = rrqr!(copy(A); rtol)

"""Truncate RRQR result low-rank"""
function truncate_qr_result(qr::QRPivoted{T}, k::Integer) where {T}
    m, n = size(qr)
    0 ≤ k ≤ min(m, n) || throw(DomainError(k, "Invalid rank, must be in [0, $(min(m, n))]"))
    Qfull = QRPackedQ(view(qr.factors, :, 1:k), qr.τ[1:k])

    Q = lmul!(Qfull, Matrix{T}(I, m, k))
    R = triu!(qr.factors[1:k, :])
    return Q, R
end

"""
Truncated singular value decomposition.

Decomposes an `(m, n)` matrix `A` into the product:

    A == U * (s .* VT)

where `U` is a `(m, k)` matrix with orthogonal columns, `VT` is a `(k, n)`
matrix with orthogonal rows and `s` are the singular values, a set of `k`
nonnegative numbers in non-ascending order.  The SVD is truncated in the
sense that singular values below `tol` are discarded.
"""
function tsvd!(A::AbstractMatrix{T}; rtol=eps(T)) where {T<:AbstractFloat}
    # Perform RRQR of the m x n matrix at a cost of O(m*n*k), where k is
    # the QR rank (a mild upper bound for the true rank)
    A_qr, k = rrqr!(A; rtol)
    Q, R = truncate_qr_result(A_qr, k)

    # RRQR is an excellent preconditioner for Jacobi. One should then perform
    # Jacobi on RT
    RT_svd = svd!(Matrix(R'))

    # Reconstruct A from QR
    U = Q * RT_svd.V
    V = @view RT_svd.U[invperm(A_qr.p), :]
    s = RT_svd.S
    return SVD(U, s, V')
end

"""Truncated singular value decomposition."""
tsvd(A::AbstractMatrix{T}; rtol=eps(T)) where {T<:AbstractFloat} = tsvd!(copy(A); rtol)

#################################################################################################################
### Everything below is currently not used in tsvd. (GenericLinearAlgebra.svd! is used instead of svd_jacobi) ###
#################################################################################################################

"""
Compute Givens rotation `R` matrix that satisfies:

    [  c  s ] [ f ]     [ r ]
    [ -s  c ] [ g ]  =  [ 0 ]
"""
function givens_params(f::T, g::T) where {T<:AbstractFloat}
    # ACM Trans. Math. Softw. 28(2), 206, Alg 1
    if iszero(g)
        c, s = one(T), zero(T)
        r = f
    elseif iszero(f)
        c, s = zero(T), T(signbit(g))
        r = abs(g)
    else
        r = copysign(hypot(f, g), f)
        c = f / r
        s = g / r
    end
    return (c, s), r
end

"""
Apply Givens rotation to vector:

      [ a ]  =  [  c   s ] [ x ]
      [ b ]     [ -s   c ] [ y ]
"""
function givens_lmul((c, s)::Tuple{T,T}, (x, y)) where {T}
    a = c * x + s * y
    b = c * y - s * x
    return a, b
end

"""
Perform the SVD of upper triangular two-by-two matrix:

      [ f    g   ]  =  [  cu  -su ] [ smax     0 ] [  cv   sv ]
      [ 0    h   ]     [  su   cu ] [    0  smin ] [ -sv   cv ]

Note that smax and smin can be negative.
"""
function svd2x2(f::T, g::T, h::T) where {T<:AbstractFloat}
    # Code taken from LAPACK xSLAV2:
    fa = abs(f)
    ga = abs(g)
    ha = abs(h)
    if fa < ha
        # switch h <-> f, cu <-> sv, cv <-> su
        (sv, cv), (smax, smin), (su, cu) = svd2x2(h, g, f)
    elseif iszero(ga)
        # already diagonal, fa > ha
        smax, smin = fa, ha
        cv, sv = cu, su = one(T), zero(T)
    elseif fa < eps(T) * ga
        # ga is very large
        smax = ga
        if ha > one(T)
            smin = fa / (ga / ha)
        else
            smin = (fa / ga) * ha
        end
        cv, sv = one(T), f / g
        cu, su = one(T), h / g
    else
        # normal case
        fmh = fa - ha
        d = fmh / fa
        q = g / f
        s = T(2) - d
        spq = hypot(q, s)
        dpq = hypot(d, q)
        a = (spq + dpq) / T(2)
        smax, smin = abs(fa * a), abs(ha / a)

        tmp = (q / (spq + s) + q / (dpq + d)) * (one(T) + a)
        tt = hypot(tmp, T(2))
        cv = T(2) / tt
        sv = tmp / tt
        cu = (cv + sv * q) / a
        su = ((h / f) * sv) / a
    end
    return (cu, su), (smax, smin), (cv, sv)
end

"""
Perform the SVD of an arbitrary two-by-two matrix:

      [ a11  a12 ]  =  [  cu  -su ] [ smax     0 ] [  cv   sv ]
      [ a21  a22 ]     [  su   cu ] [    0  smin ] [ -sv   cv ]

Note that smax and smin can be negative.
"""
function svd2x2(a11::T, a12::T, a21::T, a22::T) where {T}
    abs_a12 = abs(a12)
    abs_a21 = abs(a21)
    if iszero(a21)
        # upper triangular case
        (cu, su), (smax, smin), (cv, sv) = svd2x2(a11, a12, a22)
    elseif abs_a12 > abs_a21
        # closer to lower triangular - transpose matrix
        (cv, sv), (smax, smin), (cu, su) = svd2x2(a11, a21, a12, a22)
    else
        # First, we use a givens rotation  Rx
        # [  cx   sx ] [ a11  a12 ] = [ rx  a12' ]
        # [ -sx   cx ] [ a21  a22 ]   [ 0   a22' ]
        (cx, sx), rx = givens_params(a11, a21)
        a11, a21 = rx, zero(rx)
        a12, a22 = givens_lmul((cx, sx), (a12, a22))

        # Next, use the triangular routine
        # [ f  g ]  =  [  cu  -su ] [ smax     0 ] [  cv   sv ]
        # [ 0  h ]     [  su   cu ] [    0  smin ] [ -sv   cv ]
        (cu, su), (smax, smin), (cv, sv) = svd2x2(a11, a12, a22)

        # Finally, update the LHS (U) transform as follows:
        # [  cx  -sx ] [  cu  -su ] = [  cu'  -su' ]
        # [  sx   cx ] [  su   cu ]   [  su'   cu' ]
        cu, su = givens_lmul((cx, -sx), (cu, su))
    end
    return (cu, su), (smax, smin), (cv, sv)
end

function jacobi_sweep!(U::AbstractMatrix, VT::AbstractMatrix)
    ii, jj = size(U)
    ii ≥ jj || throw(ArgumentError("matrix must be 'tall'"))
    size(VT, 1) == jj || throw(ArgumentError("U and VT must be compatible"))
    Base.require_one_based_indexing(U)
    Base.require_one_based_indexing(VT)

    offd = zero(eltype(U))
    @inbounds for i in 1:ii
        for j in (i + 1):jj
            # Construct the 2x2 matrix to be diagonalized
            Hii = sum(abs2, @view U[:, i])
            Hij = sum(k -> @inbounds(U[k, i] * U[k, j]), 1:ii)
            Hjj = sum(abs2, @view U[:, j])
            offd += abs2(Hij)

            # diagonalize
            (_, _), (_, _), (cv, sv) = svd2x2(Hii, Hij, Hij, Hjj)

            # apply rotation to VT
            rot = Givens(i, j, cv, sv)
            lmul!(rot, VT)
            rmul!(U, adjoint(rot))
        end
    end
    return sqrt(offd)
end

"""Singular value decomposition using Jacobi rotations."""
function svd_jacobi!(U::AbstractMatrix{T}; rtol=eps(T), maxiter=20) where {T}
    m, n = size(U)
    m ≥ n || throw(ArgumentError("matrix must be 'tall'"))
    Base.require_one_based_indexing(U)

    VT = Matrix(one(T) * I, n, n)
    Unorm = norm(@view U[1:n, 1:n])
    for _ in 1:maxiter
        offd = jacobi_sweep!(U, VT)
        offd < rtol * Unorm && break
    end

    s = norm.(eachcol(U))
    U ./= transpose(s)
    return SVD(U, s, VT)
end

"""Singular value decomposition using Jacobi rotations."""
function svd_jacobi(U::AbstractMatrix{T}; rtol=eps(T), maxiter=20) where {T}
    return svd_jacobi!(copy(U); rtol, maxiter)
end

end # module _LinAlg


# using Revise, Plots, Test
using BifurcationKit, LinearAlgebra, Setfield, SparseArrays, ForwardDiff, Parameters
const BK = BifurcationKit
norminf = x -> norm(x, Inf)

Fbp(x, p) = [p.a0 + x[1] * (3.23 .* p.μ - p.x2 * x[1] + p.x3 * 0.234 * x[1]^2) + x[2], -x[2]]
par = (a0 = 0., μ = -0.2, ν = 0, x2 = 1.12, x3 = 1.0)
####################################################################################################
opt_newton = NewtonPar(tol = 1e-9, maxIter = 20)
opts_br = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds = 0.01, pMax = 0.4, pMin = -0.5, detectBifurcation = 3, nev = 2, newtonOptions = opt_newton, maxSteps = 100, nInversion = 4, tolBisectionEigenvalue = 1e-8, dsminBisection = 1e-9)

br, = continuation(
	Fbp, [0.1, 0.1], par, (@lens _.μ),
	opts_br; plot = false, verbosity = 0, normC = norminf, recordFromSolution = (x, p) -> x[1])

####################################################################################################
# normal form computation
jet = BK.getJet(Fbp, (x, p) -> ForwardDiff.jacobian(z -> Fbp(z, p), x))

bp = BK.computeNormalForm(jet..., br, 1; verbose=false)
@test BK.isTranscritical(bp) == true

jet = BK.getJet(Fbp, (x, p) -> BK.finiteDifferences(z -> Fbp(z, p), x))
bp = BK.computeNormalForm(jet..., br, 1; verbose=false)
@test BK.isTranscritical(bp) == true
show(bp)

# normal form
nf = bp.nf

@test norm(nf.a) < 1e-10
	@test norm(nf.b1 - 3.23) < 1e-10
	@test norm(nf.b2/2 - -1.12) < 1e-10
	@test norm(nf.b3/6 - 0.234) < 1e-10

####################################################################################################
# same but when the eigenvalues are not saved in the branch but computed on the fly
br_noev, = BK.continuation(
	Fbp, [0.1, 0.1], par, (@lens _.μ),
	recordFromSolution = (x, p) -> norminf(x),
	(@set opts_br.saveEigenvectors = false); plot = false, verbosity = 0, normC = norminf)
bp = BK.computeNormalForm(jet..., br_noev, 1; verbose=false)
nf = bp.nf
@test norm(nf[1]) < 1e-10
	@test norm(nf[2] - 3.23) < 1e-10
	@test norm(nf[3]/2 - -1.12) < 1e-10
	@test norm(nf[4]/6 - 0.234) < 1e-10
####################################################################################################
# Automatic branch switching
br2, = continuation(jet..., br, 1, setproperties(opts_br; pMax = 0.2, ds = 0.01, maxSteps = 14); recordFromSolution = (x, p) -> x[1], verbosity = 0)
@test br2 isa Branch
@test BK.haseigenvalues(br2) == true
@test BK.haseigenvector(br2) == true
BK.eigenvals(br2, 1, true)
BK.getfirstusertype(br2)
@test length(br2) == 12
# plot(br,br2)

br3, = continuation(jet..., br, 1, setproperties(opts_br; ds = -0.01); recordFromSolution = (x, p) -> x[1], verbosity = 0, usedeflation = true)
# plot(br,br3)

# automatic bifurcation diagram (Transcritical)
bdiag = bifurcationdiagram(jet..., [0.1, 0.1], par,  (@lens _.μ), 2,
	(args...) -> setproperties(opts_br; pMin = -1.0, pMax = .5, ds = 0.01, dsmax = 0.05, nInversion = 6, detectBifurcation = 3, maxBisectionSteps = 30, newtonOptions = (@set opt_newton.verbose=false), maxSteps = 15);
	recordFromSolution = (x, p) -> x[1],
	tangentAlgo = BorderedPred(),
	plot = false, verbosity = 0, normC = norminf)
####################################################################################################
# Case of the pitchfork
par_pf = setproperties(par; x2 = 0.0, x3 = -1.0)
brp, = BK.continuation(
	Fbp, [0.1, 0.1], par_pf, (@lens _.μ),
	recordFromSolution = (x, p) -> x[1],
	opts_br; plot = false, verbosity = 0, normC = norminf)
bpp = BK.computeNormalForm(jet..., brp, 1)
show(bpp)

nf = bpp.nf
@test norm(nf[1]) < 1e-9
	@test norm(nf[2] - 3.23) < 1e-9
	@test norm(nf[3]/2 - 0) < 1e-9
	@test norm(nf[4]/6 + 0.234) < 1e-9

# test automatic branch switching
br2, = continuation(jet..., brp, 1, setproperties(opts_br; maxSteps = 19, dsmax = 0.01, ds = 0.001, detectBifurcation = 2, newtonOptions = (@set opt_newton.verbose=false)); recordFromSolution = (x, p) -> x[1], verbosity = 0, ampfactor = 1)
	# plot(brp,br2, marker=:d)

# test methods for aBS
BK.from(br2) |> BK.type
BK.from(br2) |> BK.istranscritical
BK.type(nothing)
BK.show(stdout, br2)
BK.propertynames(br2)

# automatic bifurcation diagram (Pitchfork)
bdiag = bifurcationdiagram(jet..., [0.1, 0.1], par_pf,  (@lens _.μ), 2,
	(args...) -> setproperties(opts_br; pMin = -1.0, pMax = .5, ds = 0.01, dsmax = 0.05, nInversion = 6, detectBifurcation = 3, maxBisectionSteps = 30, newtonOptions = (@set opt_newton.verbose=false), maxSteps = 15);
	recordFromSolution = (x, p) -> x[1],
	# tangentAlgo = BorderedPred(),
	plot = false, verbosity = 0, normC = norminf)
####################################################################################################
function Fbp2d(x, p)
	return [ p.α * x[1] * (3.23 .* p.μ - 0.123 * x[1]^2 - 0.234 * x[2]^2),
			 p.α * x[2] * (3.23 .* p.μ - 0.456 * x[1]^2 - 0.123 * x[2]^2),
			 -x[3]]
end

jet = BK.getJet(Fbp2d, (x, p) -> ForwardDiff.jacobian(z -> Fbp2d(z, p), x))

par = (μ = -0.2, ν = 0, α = -1)

for α in (1,1)
	global par = @set par.α = α
	br, = BK.continuation(
		Fbp2d, [0.01, 0.01, 0.01], par, (@lens _.μ),
		recordFromSolution = (x, p) -> norminf(x),
		setproperties(opts_br; nInversion = 2); plot = false, verbosity = 0, normC = norminf)
	# we have to be careful to have the same basis as for Fbp2d or the NF will not match Fbp2d
	bp2d = BK.computeNormalForm(jet..., br, 1; ζs = [[1, 0, 0.], [0, 1, 0.]]);
	show(bp2d)

	BK.nf(bp2d)
	length(bp2d)
	bp2d(rand(2), 0.2)
	bp2d(Val(:reducedForm), rand(2), 0.2)

	@test abs(bp2d.nf.b3[1,1,1,1] / 6 - -par.α * 0.123) < 1e-10
	@test abs(bp2d.nf.b3[1,1,2,2] / 2 - -par.α * 0.234) < 1e-10
	@test abs(bp2d.nf.b3[1,1,1,2] / 2 - -par.α * 0.0)   < 1e-10
	@test abs(bp2d.nf.b3[2,1,1,2] / 2 - -par.α * 0.456) < 1e-10
	@test norm(bp2d.nf.b2, Inf) < 3e-6
	@test norm(bp2d.nf.b1 - par.α * 3.23 * I, Inf) < 1e-9
	@test norm(bp2d.nf.a, Inf) < 1e-6
end

##############################
# same but when the eigenvalues are not saved in the branch but computed on the fly instead
br_noev, = BK.continuation(
	Fbp2d, [0.01, 0.01, 0.01], par, (@lens _.μ),
	recordFromSolution = (x, p) -> norminf(x),
	setproperties(opts_br; nInversion = 2, saveEigenvectors = false); plot = false, verbosity = 0, normC = norminf)
bp2d = BK.computeNormalForm(jet..., br_noev, 1; ζs = [[1, 0, 0.], [0, 1, 0.]]);
@test abs(bp2d.nf.b3[1,1,1,1] / 6 - -0.123) < 1e-15
@test abs(bp2d.nf.b3[1,1,2,2] / 2 - -0.234) < 1e-15
@test abs(bp2d.nf.b3[1,1,1,2] / 2 - -0.0)   < 1e-15
@test abs(bp2d.nf.b3[2,1,1,2] / 2 - -0.456) < 1e-15
@test norm(bp2d.nf.b2, Inf) < 3e-15
@test norm(bp2d.nf.b1 - 3.23 * I, Inf) < 1e-9
@test norm(bp2d.nf.a, Inf) < 1e-15
####################################################################################################
# vector field to test nearby secondary bifurcations
FbpSecBif(u, p) = @. -u * (p + u * (2-5u)) * (p -.15 - u * (2+20u))
dFbpSecBif(x,p)=  ForwardDiff.jacobian( z-> FbpSecBif(z,p), x)
jet = BK.getJet(FbpSecBif, dFbpSecBif)

br_snd1, = BK.continuation(
	FbpSecBif, [0.0], -0.2, (@lens _),
	recordFromSolution = (x, p) -> x[1],
	# tangentAlgo = BorderedPred(),
	setproperties(opts_br; pMin = -1.0, pMax = .3, ds = 0.001, dsmax = 0.005, nInversion = 8, detectBifurcation=3); plot = false, verbosity = 0, normC = norminf)

# plot(br_snd1)

br_snd2, = BK.continuation(
	jet..., br_snd1, 1,
	recordFromSolution = (x, p) -> x[1],
	setproperties(opts_br; pMin = -1.2, pMax = 0.2, ds = 0.001, detectBifurcation = 3, maxSteps=19, nInversion = 8, newtonOptions = NewtonPar(opts_br.newtonOptions),dsminBisection =1e-18, tolBisectionEigenvalue=1e-11, maxBisectionSteps=20); plot = false, verbosity = 0, normC = norminf,
	# tangentAlgo = BorderedPred(),
	# finaliseSolution = (z, tau, step, contResult) ->
	# 	(Base.display(contResult.eig[end].eigenvals) ;true)
	)

	# plot(plot(br_snd2.branch[1,:] |> diff, marker = :d),
	# plot(br_snd1,br_snd2, putbifptlegend=false))

bdiag = bifurcationdiagram(jet..., [0.0], -0.2, (@lens _), 2,
	(args...) -> setproperties(opts_br; pMin = -1.0, pMax = .3, ds = 0.001, dsmax = 0.005, nInversion = 8, detectBifurcation = 3,dsminBisection =1e-18, tolBisectionEigenvalue=1e-11, maxBisectionSteps=20, newtonOptions = (@set opt_newton.verbose=false));
	recordFromSolution = (x, p) -> x[1],
	# tangentAlgo = BorderedPred(),
	plot = false, verbosity = 0, normC = norminf)

# plot(bdiag; putbifptlegend=false, markersize=2, plotfold=false, title = "#branch = $(size(bdiag))")

# test calls for aBD
BK.level(bdiag)
BK.hasbranch(bdiag)
BK.from(bdiag.child[1].γ)
BK.getBranchesFromBP(bdiag, 2)
BK.getContResult(br)
BK.getContResult(getBranch(bdiag,(1,)).γ)
size(bdiag)
getBranch(bdiag, (1,))
show(stdout, bdiag)
####################################################################################################
# test of the pitchfork-D6 normal form
function FbpD6(x, p)
	return [ p.μ * x[1] + (p.a * x[2] * x[3] - p.b * x[1]^3 - p.c*(x[2]^2 + x[3]^2) * x[1]),
			 p.μ * x[2] + (p.a * x[1] * x[3] - p.b * x[2]^3 - p.c*(x[3]^2 + x[1]^2) * x[2]),
			 p.μ * x[3] + (p.a * x[1] * x[2] - p.b * x[3]^3 - p.c*(x[2]^2 + x[1]^2) * x[3])]
end
jet = BK.getJet(FbpD6, (x, p) -> ForwardDiff.jacobian(z -> FbpD6(z, p), x))

pard6 = (μ = -0.2, a = 0.3, b = 1.5, c = 2.9)

br, = BK.continuation(
	FbpD6, zeros(3), pard6, (@lens _.μ),
	recordFromSolution = (x, p) -> norminf(x),
	setproperties(opts_br; nInversion = 6, ds = 0.001); plot = false, verbosity = 0, normC = norminf)

# plot(br;  plotfold = false)
# we have to be careful to have the same basis as for Fbp2d or the NF will not match Fbp2d
bp2d = BK.computeNormalForm(jet..., br, 1; ζs = [[1, 0, 0.], [0, 1, 0.], [0, 0, 1.]])
BK.nf(bp2d)

@test abs(bp2d.nf.b3[1,1,1,1] / 6 - -pard6.b) < 1e-10
	@test abs(bp2d.nf.b3[1,1,2,2] / 2 - -pard6.c) < 1e-10
	@test abs(bp2d.nf.b2[1,2,3] - pard6.a)   < 1e-10

# test the evaluation of the normal form
x0 = rand(3); @test norm(FbpD6(x0, set(pard6, br.lens, 0.001))  - bp2d(Val(:reducedForm), x0, 0.001), Inf) < 1e-12

br1, = BK.continuation(
	jet..., br, 1,
	setproperties(opts_br; nInversion = 4, dsmax = 0.005, ds = 0.001, maxSteps = 300, pMax = 1.); plot = false, verbosity = 0, normC = norminf, recordFromSolution = (x, p) -> norminf(x))
	# plot(br1..., br, plotfold=false, putbifptlegend=false)

bdiag = bifurcationdiagram(jet..., zeros(3), pard6, (@lens _.μ), 3,
	(args...) -> setproperties(opts_br; pMin = -0.250, pMax = .4, ds = 0.001, dsmax = 0.005, nInversion = 4, detectBifurcation = 3, dsminBisection =1e-18, tolBisectionEigenvalue=1e-11, maxBisectionSteps=20, newtonOptions = (@set opt_newton.verbose=false));
	recordFromSolution = (x, p) -> norminf(x),
	# tangentAlgo = BorderedPred(),
	plot = false, verbosity = 0, normC = norminf)

# plot(bdiag; putspecialptlegend=false, markersize=2,plotfold=false);title!("#branch = $(size(bdiag))")

####################################################################################################
# test of the Hopf normal form
function Fsl2!(f, u, p, t)
	@unpack r, μ, ν, c3, c5 = p
	u1 = u[1]
	u2 = u[2]
	ua = u1^2 + u2^2

	f[1] = r * u1 - ν * u2 - ua * (c3 * u1 - μ * u2) - c5 * ua^2 * u1
	f[2] = r * u2 + ν * u1 - ua * (c3 * u2 + μ * u1) - c5 * ua^2 * u2

	return f
end

Fsl2(x, p) = Fsl2!(similar(x), x, p, 0.)
par_sl = (r = 0.5, μ = 0.132, ν = 1.0, c3 = 1.123, c5 = 0.2)
jet = BK.getJet(Fsl2, (x, p) -> ForwardDiff.jacobian(z -> Fsl2(z, p), x))

# detect hopf bifurcation
opts_br = ContinuationPar(dsmin = 0.001, dsmax = 0.02, ds = 0.01, pMax = 0.1, pMin = -0.3, detectBifurcation = 3, nev = 2, newtonOptions = (@set opt_newton.verbose = false), maxSteps = 100)

br, = BK.continuation(
	jet[1], jet[2], [0.0, 0.0], (@set par_sl.r = -0.1), (@lens _.r),
	recordFromSolution = (x, p) -> norminf(x),
	opts_br; plot = false, verbosity = 0, normC = norminf)

hp = BK.computeNormalForm(jet..., br, 1)

nf = hp.nf
BK.type(hp)

@test abs(nf.a - 1) < 1e-9
@test abs(nf.b/2 - (-par_sl.c3 + im*par_sl.μ)) < 1e-14

####################################################################################################
# same but when the eigenvalues are not saved in the branch but computed on the fly instead
br, _ = BK.continuation(
	Fsl2, [0.0, 0.0], (@set par_sl.r = -0.1), (@lens _.r),
	recordFromSolution = (x, p) -> norminf(x),
	setproperties(opts_br, saveEigenvectors = false); plot = false, verbosity = 0, normC = norminf)

hp = BK.computeNormalForm(jet..., br, 1)

nf = hp.nf

@test abs(nf.a - 1) < 1e-9
@test abs(nf.b/2 - (-par_sl.c3 + im*par_sl.μ)) < 1e-14

show(hp)
####################################################################################################
# test for the Cusp normal form
Fcusp(x, p) = [p.β1 + p.β2 * x[1] + p.c * x[1]^3]
par = (β1 = 0.0, β2 = -0.01, c = 3.)
jet  = BK.getJet(Fcusp; matrixfree=false)
br, = continuation(jet[1], jet[2], [0.01], par, (@lens _.β1), opts_br;)

sn_codim2, = continuation(jet[1:2]..., br, 1, (@lens _.β2), ContinuationPar(opts_br, detectBifurcation = 1, saveSolEveryStep = 1, maxSteps = 40) ;
	updateMinAugEveryStep = 1,
	d2F = jet[3], d3F = jet[4],
	bdlinsolver = MatrixBLS()
	)
# find the cusp point
ind = findall(map(x->x.type == :cusp, sn_codim2.specialpoint))
cuspnf = computeNormalForm(jet..., sn_codim2, ind[1])
show(cuspnf)
BK.type(cuspnf)
@test cuspnf.nf.c == par.c
####################################################################################################
# test for the Bogdanov-Takens normal form
Fbt(x, p) = [x[2], p.β1 + p.β2 * x[2] + p.a * x[1]^2 + p.b * x[1] * x[2]]
par = (β1 = 0.01, β2 = -0.1, a = -1., b = 1.)
jet  = BK.getJet(Fbt; matrixfree=false)
opt_newton = NewtonPar(tol = 1e-9, maxIter = 40, verbose = false)
opts_br = ContinuationPar(dsmin = 0.001, dsmax = 0.01, ds = 0.01, pMax = 0.5, pMin = -0.5, detectBifurcation = 3, nev = 2, newtonOptions = opt_newton, maxSteps = 100, nInversion = 8, tolBisectionEigenvalue = 1e-8, dsminBisection = 1e-9, saveSolEveryStep = 1)

br, = continuation(jet[1], jet[2], [0.01, 0.01], par, (@lens _.β1), opts_br; bothside = true)

sn_codim2, = continuation(jet[1:2]..., br, 1, (@lens _.β2), ContinuationPar(opts_br, detectBifurcation = 1, saveSolEveryStep = 1, maxSteps = 40) ;
	detectCodim2Bifurcation = 2,
	updateMinAugEveryStep = 1,
	d2F = jet[3], d3F = jet[4],
	bdlinsolver = MatrixBLS()
	)

hopf_codim2, = continuation(jet[1:2]..., br, 2, (@lens _.β2), ContinuationPar(opts_br, detectBifurcation = 1, saveSolEveryStep = 1, maxSteps = 40) ;
	detectCodim2Bifurcation = 2,
	updateMinAugEveryStep = 1,
	bothside = true,
	d2F = jet[3], d3F = jet[4],
	bdlinsolver = MatrixBLS(),
	)

# plot(sn_codim2, hopf_codim2, branchlabel = ["Fold", "Hopf"])

btpt = computeNormalForm(jet..., sn_codim2, 1; nev = 2)
show(btpt)
BK.type(btpt)
@test norm(btpt.nf.b * sign(sum(btpt.ζ[1])) - par.b, Inf) < 1e-5
@test norm(btpt.nf.a * sign(sum(btpt.ζ[1])) - par.a, Inf) < 1e-5
@test isapprox(abs.(btpt.ζ[1]), [1, 0])
@test isapprox(abs.(btpt.ζ[2]), [0, 1];rtol = 1e-6)
@test isapprox(abs.(btpt.ζstar[1]), [1, 0];rtol = 1e-6)

@test isapprox(btpt.nfsupp.K2, [0, 0]; atol = 1e-5)
@test isapprox(btpt.nfsupp.d, 0; atol = 1e-3)
@test isapprox(btpt.nfsupp.e, 0; atol = 1e-3)
@test isapprox(btpt.nfsupp.a1, 0; atol = 1e-3)
@test isapprox(btpt.nfsupp.b1, 0; atol = 1e-3)

HC = BK.predictor(btpt, Val(:HopfCurve), 0.)
	HC.hopf(0.)
SN = BK.predictor(btpt, Val(:FoldCurve), 0.)
Hom = BK.predictor(btpt, Val(:HomoclinicCurve), 0.)
	Hom.orbit(0,0)

# plot(sn_codim2, branchlabel = ["Fold"], vars = (:β1, :β2))
# 	_S = LinRange(-0.06, 0.06, 1000)
# 	plot!([HC.hopf(s)[1] for s in _S], [HC.hopf(s)[2] for s in _S], linewidth=5, label = "Hpred")
# 	plot!([SN.fold(s)[1] for s in _S], [SN.fold(s)[2] for s in _S], linewidth=5, label = "SNpred")
# 	_S = LinRange(-0.25, 0.25, 1000)
# 	plot!([Hom.α(s)[1] for s in _S], [Hom.α(s)[2] for s in _S], linewidth=5, label = "Hom")
#
# 	plot!(hopf_codim2, branchlabel = ["Hopf"], vars = (:β1, :β2), color = :black)
# 	xlims!(-0.001, 0.05)

# plot of the homoclinic orbit
hom1 = [Hom.orbit(t,0.1)[1] for t in LinRange(-1000, 1000, 10000)]
hom2 = [Hom.orbit(t,0.1)[2] for t in LinRange(-1000, 1000, 10000)]
# plot(hom1, hom2)

# branch switching from BT from Fold
opt = sn_codim2.contparams
@set! opt.newtonOptions.verbose = false
@set! opt.maxSteps = 20
hp_fromBT, = continuation(jet..., sn_codim2, 1, opt;
	verbosity = 0, plot = false,
	δp = 1e-4,
	updateMinAugEveryStep = 1,
	)

####################################################################################################
# test of the Bautin normal form
function Fsl2!(f, u, p, t)
	@unpack r, μ, ν, c3, c5 = p
	u1, u2 = u
	ua = u1^2 + u2^2
	f[1] = r * u1 - ν * u2 + ua * (c3 * u1 - μ * u2) + c5 * ua^2 * u1
	f[2] = r * u2 + ν * u1 + ua * (c3 * u2 + μ * u1) + c5 * ua^2 * u2
	return f
end

Fsl2(x, p) = Fsl2!(similar(x), x, p, 0.)
par_sl = (r = -0.5, μ = 0., ν = 1.0, c3 = 0.1, c5 = 0.3)
jet = BK.getJet(Fsl2, matrixfree=false)

@set! opts_br.newtonOptions.verbose = false
@set! opts_br.newtonOptions.tol = 1e-12
opts_br = setproperties(opts_br;nInversion = 10, maxBisectionSteps = 25)

br, = continuation( jet[1], jet[2], [0.01, 0.01], par_sl, (@lens _.r), opts_br)

hopf_codim2, = continuation(jet[1:2]..., br, 1, (@lens _.c3), ContinuationPar(opts_br, detectBifurcation = 0, saveSolEveryStep = 1, maxSteps = 40, pMin = -2., pMax = 2., ds = -0.001) ;
	detectCodim2Bifurcation = 2,
	startWithEigen = true,
	updateMinAugEveryStep = 1,
	d2F = jet[3], d3F = jet[4],
	bdlinsolver = MatrixBLS(),
	)

bautin = BifurcationKit.computeNormalForm(jet..., hopf_codim2, 1; nev = 2)
show(bautin)
BK.type(bautin)

@test bautin.nf.l2 ≈ par_sl.c5 * 4


# ---
# title: 779. K-th Symbol in Grammar
# id: problem779
# author: Indigo
# date: 2021-01-28
# difficulty: Medium
# categories: Recursion
# link: <https://leetcode.com/problems/k-th-symbol-in-grammar/description/>
# hidden: true
# ---
#
# On the first row, we write a `0`. Now in every subsequent row, we look at the
# previous row and replace each occurrence of `0` with `01`, and each occurrence
# of `1` with `10`.
#
# Given row `N` and index `K`, return the `K`-th indexed symbol in row `N`. (The
# values of `K` are 1-indexed.) (1 indexed).
#
#
#
#     **Examples:**
#     Input: N = 1, K = 1
#     Output: 0
#
#     Input: N = 2, K = 1
#     Output: 0
#
#     Input: N = 2, K = 2
#     Output: 1
#
#     Input: N = 4, K = 5
#     Output: 1
#
#     Explanation:
#     row 1: 0
#     row 2: 01
#     row 3: 0110
#     row 4: 01101001
#
#
# **Note:**
#
#   1. `N` will be an integer in the range `[1, 30]`.
#   2. `K` will be an integer in the range `[1, 2^(N-1)]`.
#
#
## @lc code=start
using LeetCode

function kth_grammar(N::Int, K::Int)
    res = false
    for i in 1:N
        iseven(K) && (res = !res)
        K = (K + 1) >> 1
    end
    res
end
## @lc code=end


"""
Update the toeplitz matrix A using vector x
"""
function toeplitz!(A, x)
   n = length(x)
   for i in 1:n
	   for j in i:n
	       A[i,j] = x[j-i+1]
	       A[j,i] = x[j-i+1]
       end
   end
   return A
end


module RandomWalk

export randomwalk1d, randomwalk2d, rwtrap1d_lifetime, rwtrap1d_lifetime_exp, dlaline!

"""
1D Random Walk

# Arguments

-`time`: the number of steps of the random walk
-`p`: the probability of going right
-`walks`: the number of walks, defaults to `100*time`

# Returns

a vector containing the coordinates of each point and a vector counting in how many times
the walker landed on each point ath the end of the walk (only non-zero points are returned)
"""
function randomwalk1d(time::Integer, p::Float64; walks=100*time)
    path = zeros(2 * time + 1)
    coordinates = collect(Integer, -time:time)

    for _ in 1:walks
        pos = time + 1
        for _ in 1:time
            if rand() <= p
                pos += 1
            else
                pos -= 1
            end
        end
        path[pos] += 1
    end

    nonzero = findall(x -> x != 0, path)
    return coordinates[nonzero], path[nonzero]
end

"""
2D Random Walk

# Arguments

-`time::Integer`: the number of steps of the random walk
-`walks::Integer`: the number of walks, defaults to the value of `time`

# Returns

two tuples of two integers giving the x and y limits of the points and
an integer matrix counting the number of times the walker landed on each point.
"""
function randomwalk2d(time::Integer; walks::Integer=100*time)
    x, y = collect(Integer, -time:time), collect(Integer, -time:time)
    grid = zeros(Integer, 2 * time + 1, 2 * time + 1)
    steps = ((0, 1), (1, 0), (0, -1), (-1, 0))

    for _ in 1:walks
        xpos, ypos = time + 1, time + 1
        for _ in 1:time
            step = steps[rand(1:4)]
            xpos += step[1]
            ypos += step[2]
        end
        grid[xpos, ypos] += 1
    end

    return randomwalk2dcrop(x, y, grid)
end

"""
crop the output of `randomwalk2d`
"""
@inline function randomwalk2dcrop(x::Vector{<:Integer}, y::Vector{<:Integer},
        grid::Matrix{<:Integer})
    height, width = size(grid)
    bottom = left = 1
    top = height
    right = width

    while all(==(0), grid[top, :])
        top -= 1
    end
    while all(==(0), grid[bottom, :])
        bottom += 1
    end
    while all(==(0), grid[:, left])
        left += 1
    end
    while all(==(0), grid[:, right])
        right -= 1
    end

    return x[left:right], y[bottom:top], grid[bottom:top, left:right]
end

"""
Calculate the average lifetime of a 1D random walker inside trapping boundaries

# Arguments

-`len::Integer`: The length of the path. The coordinates are in the range [-len÷2, len÷2]
-`initpos::Integer`: the initial position of the walker
-`walks::Integer`: the number of walks, defaults to the value of `time`
-`maxtime::Integer`: The maximum number of steps before the walker is abandoned;
    in case the number of steps exceeds `maxtime`, the lifetime is taken to be `maxtime`.
    defaults to `100*len`.

# Returns

The average lifetime, which is the average time passed before the walker gets trapped
in the boundaries of the path
"""
function rwtrap1d_lifetime(len::Integer, initpos::Integer, walks::Integer;
        maxtime::Integer=100*len)
    lifetime_sum = 0

    for _ in 1:walks
        pos = initpos
        time = 0

        while time < maxtime
            time += 1
            if rand() <= 0.5
                pos += 1
                if pos > len ÷ 2
                    break
                end
            else
                pos -= 1
                if pos < -len ÷ 2
                    break
                end
            end
        end

        lifetime_sum += time
    end

    lifetime_avg = lifetime_sum / walks
    return lifetime_avg
end

"""
Calculate the expected lifetime of a 1D random walker inside trapping boundaries

# Arguments

-`len::Integer`: The length of the path. The coordinates are in the range [-len÷2, len÷2]
-`initpos::Integer`: the position of the boundaries (traps)
-`maxtime::Integer`: The maximum number of steps before the walker is abandoned;
    in case the number of steps exceeds `maxtime`, the lifetime is taken to be `maxtime`.

# Returns

The average lifetime, which is the average time passed before the walker gets trapped
in the boundaries of the path
"""
function rwtrap1d_lifetime_exp(len::Integer, initpos::Integer; maxtime::Integer=100*len)
    lifetime_avg = 0

    probs = zeros(len + 3)
    probs[initpos + len ÷ 2 + 2] = 1

    for _ in 1:maxtime
        newprobs = zeros(len + 3)
        for i in 2:len+2
            newprobs[i - 1] += probs[i] * 0.5
            newprobs[i + 1] += probs[i] * 0.5
        end
        probs = newprobs

        lifetime_avg += sum(probs[2:end-1])
    end

    return lifetime_avg
end

"""
Applies diffusion-limited aggregation on a 2D surface, assuming the "seed" is a line

the line is at the bottom and new particles appear from the top of the pile

# Arguments

- `surface::BitMatrix`: The simulation surface. Particles occupy points that are true.
- `particles::Integer`: the number of particles (walkers) simulated
- `escapethresh::Integer`: if the particle gets this much higher from the pile, it will
    "escape". It will be removed from the simulation and a new particle will be generated
    in its place. Defaults to 5.

# Returns
the `surface`
"""
function dlaline!(surface::BitMatrix, particles::Integer; escapethresh::Integer=5)
    height, width = size(surface)
    maxheight = height
    while true ∉ surface[maxheight - 1, :]
        maxheight -= 1
    end
    steps = ((0, 1), (1, 0), (0, -1), (-1, 0))

    while particles > 0
        ypos, xpos = maxheight, rand(1:width)

        while true
            xprev, yprev = xpos, ypos
            step = steps[rand(1:4)]
            ypos += step[1]
            xpos += step[2]
            # periodic boundary conditions (if xpos = width + 1, then it will be shifted
            # left by width, and if xpos = 0, it will be shifted right by width)
            xpos += (width + 1 - 2xpos) ÷ (width + 1) * width

            if ypos > min(maxheight + escapethresh, height)
                break
            elseif surface[ypos, xpos]
                surface[yprev, xprev] = true
                particles -= 1
                if yprev > maxheight
                    maxheight = yprev
                end
                break
            end
        end
    end
    return surface # for convenience
end

end


module TestAsInterval

using Test
using ArbFloats

afIvl1 = midpoint_radius(ArbFloat(1.0), ArbFloat(1.0e-8))
afIvl2 = midpoint_radius(ArbFloat(1.0), ArbFloat(1.0e-10))
afIvl3 = nextfloat(afIvl2)

@test contains(afIvl1, afIvl2)
@test !contains(afIvl2, afIvl1)
@test overlap(afIvl2, afIvl3)

@test midpoint(afIvl3) > midpoint(afIvl2)
# this will not be true unless we force the conversion to ArbMag everywhere radius is explicit
# @test radius(afIvl3) == radius(afIvl2)

end # module


function cost_fn_sir(β, γ; k=1)
    return 1_000_000_000_000 * logistic(k*(10*base_reproduction_rate(β, γ) - 30))
end

const COST_FUNCTIONS = (
    cost_fn_sir = cost_fn_sir,
)


using JuMP
using Combinatorics

# Inspired from Joey Huchette's test in ConvexHull.jl
function crosspolytopetest(lib::PolyhedraLibrary, n)
    m = Model()

    @variable(m, x[1:n])

    for k in 0:n
        for S in combinations(1:n, k)
            Sᶜ = setdiff(1:n, S)
            @constraint(m, sum(x[i] for i in S) - sum(x[i] for i in Sᶜ) ≤ 1)
        end
    end

    poly = polyhedron(m, lib)

    V = zeros(Int, 2n, n)
    for i in 1:n
        V[2i-1, i] = 1
        V[2i, i] = -1
    end
    generator_fulltest(poly, V)
end


using Distributions, Plots
using BioLizard.OrdinalRegr: fitordinalsvm, OrdinalSVM, predict, score

Σ = [3.0 1; 1 2]

dist1 = MultivariateNormal([0, 0], Σ)
dist2 = MultivariateNormal([1, 3], Σ)
dist3 = MultivariateNormal([3, 9], Σ)

X = [rand(dist1, 20) rand(dist2, 20) rand(dist3, 20)]'
y = vcat([3 for i in 1:20], [5 for i in 1:20], [8 for i in 1:20])

pcl = scatter(X[1:20,1], X[1:20,2], label="class 1")
scatter!(X[21:40,1], X[21:40,2], label="class 2")
scatter!(X[41:60,1], X[41:60,2], label="class 3")

model = fitordinalsvm(y, X)

phist = histogram(score(model, X[1:20,:]), label="class 1", alpha=0.8)
histogram!(score(model, X[21:40,:]), label="class 2", alpha=0.8)
histogram!(score(model, X[41:60,:]), label="class 3", alpha=0.8)
vline!(model.b, label="b")

plot(pcl, phist, layout=(2,1))


function Jsingle!(J::Matrix{Float64},x::Vector{Float64},beta::Vector{Float64},rho::Float64)
    # J = zeros(4,4);
    J[1,1] = 1.;
    J[2,1] = rho.*x[1]./(x[2].^2);
    J[3,1] = 1./x[2];

    J[2,2] = (0.5.*beta[1].*x[2].^-0.5 .- rho.*x[1].^2.*x[2].^-3);
    J[3,2] = -x[1].*x[2].^-2 .+ sqrt.(0.5.*beta[1]./rho).*x[2].^-0.75;

    J[1,3] = -1.;
    J[2,3] = -rho.*x[3]./(x[4].^2);
    J[4,3] = 1./x[4];

    J[2,4] = (-0.5.*beta[2].*x[4].^-0.5 .+ rho.*x[3].^2.*x[4].^-3);
    J[4,4] = -x[3].*x[4].^-2 .- sqrt.(0.5.*beta[2]./rho).*x[4].^-0.75;

    # return J
end


"""
    matrix_projection_irr([hom::InducedActionHomomorphism, ]χ::Character)
Compute matrix projection associated to the irreducible character `χ`.

If the homomorphism is not passed, the dimension `d` of the projection is
derived from `conjugacy_classes(χ)`. E.g. if `conjugacy_classes(χ)` consist
of permutations of degree `d` (i.e. acting naturally on the set `1:d`) the
result will be a matrix of size `(d,d)`.

If the homomorphism is passed, the dimension will be derived in similar
manner from the elements of the image of the homomorphism.
"""
function matrix_projection_irr(χ::Character)
    @assert isirreducible(χ)
    mproj = matrix_projection_irr(collect(values(χ)), conjugacy_classes(χ))
    mproj .*= degree(χ) // sum(length, conjugacy_classes(χ)) # χ(1)/order(G)
    return mproj
end

function matrix_projection_irr(
    vals,
    ccls::AbstractVector{<:AbstractOrbit{<:PermutationGroups.AbstractPerm}},
)
    dim = degree(first(first(ccls)))
    result = zeros(eltype(vals), dim, dim)

    for (val, cc) in zip(vals, ccls)
        for g in cc
            for i in 1:dim
                result[i, i^g] += val
            end
        end
    end

    return result
end

function matrix_projection_irr(
    vals,
    ccls::AbstractVector{<:AbstractOrbit{<:AbstractMatrix}},
)
    # TODO: call to inv(Matrix(g)) is a dirty hack, since if `g`
    # is given by a sparse matrix `inv(g)` will fail.
    return sum(val .* sum(g -> inv(Matrix(g)), cc) for (val, cc) in zip(vals, ccls))
end

## versions with InducedActionHomomorphisms

function matrix_projection_irr(hom::InducedActionHomomorphism, χ::Character{T}) where T
    @assert isirreducible(χ)
    mproj = matrix_projection_irr(hom, collect(values(χ)), conjugacy_classes(χ))
    mproj .*= degree(χ) // sum(length, conjugacy_classes(χ)) # χ(1)/order(G)
    return eltype(mproj) == T ? mproj : T.(mproj)
end

function matrix_projection_irr(
    hom::InducedActionHomomorphism{<:ByPermutations},
    class_values,
    conjugacy_cls,
    )
    dim = degree(induce(hom, first(first(conjugacy_cls))))
    result = zeros(eltype(class_values), dim, dim)
    for (val, ccl) in zip(class_values, conjugacy_cls)
        for g in ccl
            h = induce(hom, g)
            for i in 1:dim
                result[i, i^h] += val
            end
        end
    end
    return result
end

function matrix_projection_irr(
    hom::InducedActionHomomorphism{<:ByLinearTransformation},
    class_values,
    conjugacy_cls,
)
    return sum(
        val .* sum(g -> induce(hom, inv(g)), cc)
        for (val, cc) in zip(class_values, conjugacy_cls)
    )
end

function matrix_projection(χ::Character{T}) where T
    tbl = table(χ)
    res = sum(
        c .* matrix_projection_irr(ψ)
        for (c, ψ) in zip(constituents(χ), irreducible_characters(tbl)) if !iszero(c)
    )
    return eltype(res) == T ? res : T.(res)
end

function matrix_projection(hom::InducedActionHomomorphism, χ::Character{T}) where T
    tbl = table(χ)
    res = sum(
        c .* matrix_projection_irr(hom, ψ)
        for (c, ψ) in zip(constituents(χ), irreducible_characters(tbl)) if !iszero(c)
    )
    return eltype(res) == T ? res : T.(res)
end

function matrix_projection(
    hom::InducedActionHomomorphism{<:ByPermutations},
    α::AlgebraElement,
    dim::Integer = length(features(hom)),
)
    result = zeros(eltype(α), dim, dim)
    b = basis(parent(α))

    @inbounds for (j, a) in StarAlgebras._nzpairs(StarAlgebras.coeffs(α))
        g = induce(hom, b[j])
        for i in 1:dim
            result[i, i^g] += a
        end
    end

    return result
end

function matrix_projection(
    hom::InducedActionHomomorphism{<:ByLinearTransformation},
    α::AlgebraElement,
    dim::Integer = length(features(hom)),
)

    result = zeros(Base._return_type(*, Tuple{eltype(α), coeff_type(action(hom))}), dim, dim)
    b = basis(parent(α))

    for (j, a) in StarAlgebras._nzpairs(StarAlgebras.coeffs(α))
        result += a*induce(hom, inv(b[j]))
    end

    return result
end

function matrix_projection(
    α::AlgebraElement{<:StarAlgebra{<:PermutationGroups.AbstractPermutationGroup}},
    dim::Integer = degree(parent(parent(α)))
)
    result = zeros(eltype(α), dim, dim)
    b = basis(parent(α))

    @inbounds for (j, a) in StarAlgebras._nzpairs(StarAlgebras.coeffs(α))
        g = b[j]
        for i in 1:dim
            result[i, i^g] += a
        end
    end

    return result
end

image_basis!(A::AbstractMatrix) = row_echelon_form!(A)

function image_basis!(A::AbstractMatrix{T}) where {T<:AbstractFloat}
    A, p = row_echelon_form!(A)
    fact = svd!(Matrix(@view A[1:length(p), :]))
    A_rank = sum(fact.S .> maximum(size(A)) * eps(T))
    return fact.Vt, 1:A_rank
end

function image_basis!(A::AbstractMatrix{T}) where {T<:Complex}
    fact = svd!(A)
    A_rank = sum(fact.S .> maximum(size(A)) * 2eps(real(T)))
    return fact.Vt, 1:A_rank
end

image_basis(A::AbstractMatrix) =
    ((m, p) = image_basis!(deepcopy(A)); m[1:length(p), :])

function image_basis(χ::Character)
    mpr = isirreducible(χ) ? matrix_projection_irr(χ) : matrix_projection(χ)
    image, pivots = image_basis!(mpr)
    return image[1:length(pivots), :]
end

function image_basis(hom::InducedActionHomomorphism, χ::Character)
    mpr = isirreducible(χ) ? matrix_projection_irr(hom, χ) : matrix_projection(hom, χ)
    image, pivots = image_basis!(mpr)
    return image[1:length(pivots), :]
end

function image_basis(α::AlgebraElement)
    image, pivots = image_basis!(matrix_projection(α))
    return image[1:length(pivots), :]
end

function image_basis(hom::InducedActionHomomorphism, α::AlgebraElement)
    image, pivots = image_basis!(matrix_projection(hom, α))
    return image[1:length(pivots), :]
end


export  convert2affinePCE,
        calculateAffinePCE,
        assign2multi,
        evaluatePCE,
        sampleMeasure,
        samplePCE,
        mean,
        var,
        std

"""
    calculateAffinePCE(α::Vector{Float64})::Vector{Float64}
Computes the affine PCE coefficients ``x_0`` and ``x_1`` from recurrence coefficients ``\alpha``.
"""
function calculateAffinePCE(α::Vector{Float64})
    @assert length(α)>=1 "not enough recursion coefficients"
    return [α[1], 1.]
end

calculateAffinePCE(op::OrthoPoly) = calculateAffinePCE(op.α)
calculateAffinePCE(opq::OrthoPolyQ) = calculateAffinePCE(opq.op)

function calculateAffinePCE(i::Int64,mop::MultiOrthoPoly)
    p = length(mop.name)
    @assert i<=p "mop is $p-variate, but the PC for the $i-variate entry was requested"
    calculateAffinePCE(mop.uni[i])
end

"""
```
convert2affinePCE(a::Vector{Float64},α0::Float64)
convert2affinePCE(name::String,p1::Float64,p2::Float64,d::Dict=Dict();kind::Symbol=:lbub)
```
Computes the affine PCE coefficients ``x_0`` and ``x_1`` from
```math
X = a_1 + a_2 \\Xi = x_0 + x_1 \\phi_1(\\Xi),
```
where ``\\phi_1(t) = t-\\alpha_0`` is the first-order monic basis polynomial.

For classical polynomials the `name` can be given directly. The keyword `kind in [:lbub, :μσ]`
specifies whether `p1` and `p2` have the meaning of lower/upper bounds or mean/standard deviation.
"""
function convert2affinePCE(a1::Float64,a2::Float64,α0::Float64)
    [
        a1+α0*a2;
        a2
    ]
end
function convert2affinePCE(name::String,p1::Float64,p2::Float64,d::Dict=Dict();kind::Symbol=:lbub)
    @assert kind in [:lbub, :μσ]
    name = lowercase(name)
    a1, a2, α = 0., 0., zeros(Float64,1)
    name=="beta01" && d[:shape_a]==d[:shape_b]==1 ? name="uniform01" : ()

    if name=="gaussian"
        # p1 --> mean
        # p2 --> standard deviation
        @assert p2>=0. "standard deviation has to be non-negative"
        abs(p2)<1e-4 ? (@warn "standard deviation close to zero: $p2") : ()
        a1, a2 = p1, p2
        α,~ = rm_hermite_prob(1)
    elseif name=="uniform01"
        α,~ = rm_legendre01(1)
        if kind==:lbub
            # p1 --> lower bound
            # p2 --> upper bound
            @assert p1<p2 "inconsistent bounds"
            a1, a2 = p1, p2-p1
        elseif kind==:μσ
            # p1 --> mean
            # p2 --> standard deviation
            @assert p2>=0. "standard deviation has to be nonnegative"
            abs(p2)<1e-4 ? (@warn "standard deviation close to zero: $p2") : ()
            a1, a2 = p1-sqrt(3)*p2 , 2*sqrt(3)*p2
        else
            error("not implemented")
        end
    elseif name=="beta01"
        s_α, s_β = d[:shape_a], d[:shape_b]
        α,~ = rm_jacobi01(1,Float64(s_β)-1,Float64(s_α)-1)
        if kind==:lbub
            # p1 --> lower bound
            # p2 --> upper bound
            @assert p1<p2 "inconsistent bounds"
            a1, a2 = p1, p2-p1
        elseif kind==:μσ
            # p1 --> mean
            # p2 --> standard deviation
            @assert p2>=0. "standard deviation has to be nonnegative"
            abs(p2)<1e-4 ? (@warn "standard deviation close to zero: $p2") : ()
            a1 = p1-sqrt(s_α/s_β)*sqrt(1+s_α+s_β)*p2
            a2 = (s_α+s_β)*sqrt((s_α+s_β+1)/(s_α*s_β))*p2
        else
            error("not implemented")
        end
    elseif name=="gamma"
        error("$name is not implemented")
    elseif name=="logistic"
        @assert p2>0
        α,~ = rm_logistic(1)
        a1, a2 = p1, p2
    else
        error("$name is not implemented")
    end
    convert2affinePCE(a1,a2,α[1])
end
# convert2affinePCE(name::String,p1::Float64,p2::Float64;kind::Symbol=:lbub) = convert2affinePCE(name,p1,p2,Dict();kind=kind)

function assign2multi(x::Vector{Float64},i::Int64,ind::Matrix{Int64})
    l,p = size(ind)
    nx, deg = length(x), ind[end,end]
    @assert nx<=deg+1 "inconsistent number of coefficients ($nx vs $(deg+1))"
    @assert i<=p "basis is $p-variate, you requested $i-variate"
    myind::Vector{Int64} = findUnivariateIndices(i,ind)[1:nx]
    y = spzeros(Float64,l)
    y[myind] = x
    return y
end

function sampleMeasure_byName(n::Int64,name::String,d::Dict=Dict())::Vector{Float64}
    @assert n >= 1 "invalid number $n of samples."
    name = lowercase(name)
    name == "gaussian" && return rand(Normal(),n)
    name == "uniform01" && return rand(Uniform(),n)
    name == "beta01" && return rand(Beta(d[:shape_a],d[:shape_b]),n)
    name == "gamma" && return rand(Gamma(d[:shape],1/d[:rate]),n)
    name == "logistic" && return rand(Logistic(),n)
    throw(error("$name is not implemented"))
end

function sampleMeasure_byFun(n::Int64,w::Function,dom::Tuple{Float64,Float64};method::String="adaptiverejection")
    @assert n>=1 "invalid number $n of samples."
    method = lowercase(method)
    @assert method in [ "adaptiverejection", "inversecdf"] "method $method not implemented"
    if method == "adaptiverejection"
        # works only if w is log-concave
        sampler = RejectionSampler(w,dom)
        return run_sampler!(sampler,n)
    elseif method == "rejection"
        # all purpose method but needs a solid envelope PDF
        error("method $method not yet implemented")
    elseif method == "inversecdf"
        error("method $method not yet implemented")
        # requires CDF and reliable root-finding method
    end
end

"""
__Univariate__
```
sampleMeasure(n::Int64,name::String,w::Function,dom::Tuple{Float64,Float64},symm::Bool,d::Dict;method::String="adaptiverejection")
sampleMeasure(n::Int64,m::Measure;method::String="adaptiverejection")
sampleMeasure(n::Int64,op::OrthoPoly;method::String="adaptiverejection")
sampleMeasure(n::Int64,opq::OrthoPolyQ;method::String="adaptiverejection")
```
Draw `n` samples from the measure `m` described by its
- `name`
- weight function `w`,
- domain `dom`,
- symmetry property `symm`,
- and, if applicable, parameters stored in the dictionary `d`.
By default an adaptive rejection sampling method is used (from [AdaptiveRejectionSampling.jl](https://github.com/mauriciogtec/AdaptiveRejectionSampling.jl)),
unless it is a common random variable for which [Distributions.jl](https://github.com/JuliaStats/Distributions.jl) is used.

The function is multiply dispatched to accept `OrthoPoly` or `OrthoPolyQ`.

__Multivariate__
```
sampleMeasure(n::Int64,m::MultiMeasure;method::Vector{String}=["adaptiverejection" for i=1:length(m.name)])
sampleMeasure(n::Int64,mop::MultiOrthoPoly;method::Vector{String}=["adaptiverejection" for i=1:length(mop.meas.name)])
```
Multivariate extension which provides array of samples with `n` rows and
as many columns as the multimeasure has univariate measures.
"""
function sampleMeasure(n::Int64,name::String,w::Function,dom::Tuple{Float64,Float64},symm::Bool,d::Dict;method::String="adaptiverejection")
    # ToDo: symmetry is currently not exploited, but could be for rejection sampling
    name = lowercase(name)
    if name in ["gaussian","uniform01","beta01","gamma","logistic"]
        sampleMeasure_byName(n,name,d)
    else
        sampleMeasure_byFun(n,w,dom;method=method)
    end
end
sampleMeasure(n::Int64,m::Measure;method::String="adaptiverejection") = sampleMeasure(n,m.name,m.w,m.dom,m.symmetric,m.pars;method=method)
sampleMeasure(n::Int64,op::OrthoPoly;method::String="adaptiverejection") = sampleMeasure(n,op.meas;method=method)
sampleMeasure(n::Int64,opq::OrthoPolyQ;method::String="adaptiverejection") = sampleMeasure(n::Int64,opq.op;method=method)

function sampleMeasure(n::Int64, name::Vector{String}, w::Vector{Function},
                       dom::Vector{Tuple{Float64,Float64}}, symm::BitArray,d::Vector{Dict};
                       method::Vector{String}=["adaptiverejection" for i=1:length(name)])
    @assert length(name)==length(w)==length(dom)==length(symm)==length(d) "inconsistent number of parameters"
    m = length(name)
    ξ = zeros(n,m)
    for k=1:m
        ξ[:,k] = sampleMeasure(n,name[k],w[k],dom[k],symm[k],d[k];method=method[k])
    end
    return ξ
end
function sampleMeasure(n::Int64,m::MultiMeasure;method::Vector{String}=["adaptiverejection" for i=1:length(m.name)])
    s = falses(length(m.name))
    s[:] = m.symmetric[:]
    sampleMeasure(n,m.name,m.w_uni,m.dom,s,m.pars;method=method)
end
sampleMeasure(n::Int64,mop::MultiOrthoPoly;method::Vector{String}=["adaptiverejection" for i=1:length(mop.meas.name)]) = sampleMeasure(n,mop.meas;method=method)

"""
    evaluatePCE(x::Vector{Float64},ξ::Vector{Float64},α::Vector{Float64},β::Vector{Float64})
Evaluation of polynomial chaos expansion
```math
\\mathsf{x} = \\sum_{i=0}^{L} x_i \\phi_i{\\xi_j},
```
where `L+1 = length(x)` and ``x_j`` is the ``j``th sample where ``j=1,\\dots,m``
with `m = length(ξ)`.
"""
function evaluatePCE(x::Vector{Float64},ξ::Vector{Float64},α::Vector{Float64},β::Vector{Float64})
    @assert length(α) == length(β) "inconsistent number of recurrence coefficients"
    Nsmpl = length(ξ)
    @assert Nsmpl>=1 "inconsistent number $Nsmpl of samples"
    Nx, Nrec = length(x), length(α)
    @assert Nx<=Nrec "not enough recursion coefficients"
    if Nrec > Nx α,β = α[1:Nx], β[1:Nx] end
    ϕ = zeros(Float64,Nsmpl,Nx)
    for n in 1:Nx
        ϕ[:,n] = evaluate(n-1,ξ,α,β)
    end
    ϕ*x
end
evaluatePCE(x::Vector{Float64},ξ::Float64,α::Vector{Float64},β::Vector{Float64}) = evaluatePCE(x,[ξ],α,β)
evaluatePCE(x::Vector{Float64},ξ::Vector{Float64},op::OrthoPoly) = evaluatePCE(x,ξ,op.α,op.β)
evaluatePCE(x::Vector{Float64},ξ::Vector{Float64},opq::OrthoPolyQ) = evaluatePCE(x,ξ,opq.op)

function evaluatePCE(x::Vector{Float64},ξ::Matrix{Float64},α::Vector{Vector{Float64}},β::Vector{Vector{Float64}},ind::Matrix{Int64})
    Nsmpl = size(ξ,1)
    @assert Nsmpl>=1 "inconsistent number $Nsmpl of samples"
    @assert length(α)==length(β)==size(ξ,2)==size(ind,2) "inconsistent number of coefficients"
    Nx = length(x)
    @assert Nx <= size(ind,1) "too few pc coefficients (resp: too small basis)"
    ϕ = zeros(Float64,Nsmpl,Nx)
    for n in 1:Nx
        ϕ[:,n] = evaluate(ind[n,:],ξ,α,β)
    end
    ϕ*x
end
function evaluatePCE(x::Vector{Float64},ξ::Matrix{Float64},mOP::MultiOrthoPoly)
    a,b = coeffs(mOP)
    evaluatePCE(x,ξ,a,b,mOP.ind)
end

"""
__Univariate__
```
samplePCE(n::Int64,x::Vector{Float64},op::OrthoPoly;method::String="adaptiverejection")
samplePCE(n::Int64,x::Vector{Float64},opq::OrthoPolyQ;method::String="adaptiverejection")
```
Combines [`sampleMeasure`](@ref) and [`evaluatePCE`](@ref), i.e. it first draws `n` samples
from the measure, then evaluates the PCE for those samples.

__Multivariate__
```
samplePCE(n::Int64,x::Vector{Float64},mop::MultiOrthoPoly;method::Vector{String}=["adaptiverejection" for i=1:length(mop.meas.name)])
```
"""
function samplePCE(n::Int64,x::Vector{Float64},op::OrthoPoly;method::String="adaptiverejection")
    ξ = sampleMeasure(n,op;method=method)
    evaluatePCE(x,ξ,op)
end

samplePCE(n::Int64,x::Vector{Float64},opq::OrthoPolyQ;method::String="adaptiverejection") = samplePCE(n,x,opq.op;method=method)

function samplePCE(n::Int64,x::Vector{Float64},mop::MultiOrthoPoly;method::Vector{String}=["adaptiverejection" for i=1:length(mop.meas.name)])
    ξ = sampleMeasure(n,mop;method=method)
    evaluatePCE(x,ξ,mop)
end

"""
__Univariate__
```
mean(x::Vector{Float64},op::OrthoPoly)
mean(x::Vector{Float64},opq::OrthoPolyQ)
```
__Multivariate__
```
mean(x::Vector{Float64},mop::MultiOrthoPoly)
```
compute mean of random variable with PCE `x`
"""
function mean(x::Vector{Float64},op::OrthoPoly)::Float64
    x[1]*computeSP2(0,op.β)
end
mean(x::Vector{Float64},opq::OrthoPolyQ)::Float64 = mean(x,opq.op)
function mean(x::Vector{Float64},mop::MultiOrthoPoly)::Float64
    nunc = length(mop.uni)
    x[1]*computeSP(zeros(Int64,nunc),mop)
end

"""
__Univariate__
```
var(x::Vector{Float64},op::OrthoPoly)
var(x::Vector{Float64},opq::OrthoPolyQ)
```
__Multivariate__
```
var(x::Vector{Float64},mop::MultiOrthoPoly)
```
compute variance of random variable with PCE `x`
"""
function var(x::Vector{Float64},op::OrthoPoly)::Float64
    t2 = computeSP2(op)
    @assert length(t2)>=length(x) "cannot compute variance; too many PCE coefficients"
    sum( x[i]^2*t2[i] for i=2:length(x) )
end
var(x::Vector{Float64},opq::OrthoPolyQ)::Float64 = var(x,opq.op)

function var(x::Vector{Float64},mop::MultiOrthoPoly)::Float64
    L = size(mop.ind,1)
    @assert length(x)<=L "cannot compute variance; too many PCE coefficients"
    t2 = Tensor(2,mop)
    sum( x[i]^2*t2.get([i-1,i-1]) for i=2:length(x))
end

"""
__Univariate__
```
std(x::Vector{Float64},op::OrthoPoly)
std(x::Vector{Float64},opq::OrthoPolyQ)
```
__Multivariate__
```
std(x::Vector{Float64},mop::MultiOrthoPoly)
```
compute standard deviation of random variable with PCE `x`
"""
function std(x::Vector{Float64},op::OrthoPoly)::Float64
    sqrt(var(x,op))
end
std(x::Vector{Float64},opq::OrthoPolyQ)::Float64 = std(x,opq.op)
std(x::Vector{Float64},mop::MultiOrthoPoly)::Float64 = sqrt(var(x,mop))

function moment(n::Int64,x::Vector{Float64},op::OrthoPoly)
    @assert 1<=n<=2 "for moments greater than 2 a quadrature rule is required"
    n==1 ? mean(x,op) : (mean(x,op)^2+var(x,op))
end

# function moment(n::Int64,x::Vector{Float64},opq::OrthoPolyQ)
#     @assert n>=1 "invalid moment"
#     if n==1
#         return mean(x,opq)
#     elseif n==2
#         return mean(x,opq)^2+var(x,opq)
#     else
#         t = Tensor(n,opq)
#         # needs IterTools for an elegant solution
#         error("This is not yet implemented")
#     end
# end


# -----------------------------------------------------------
#
# Implementing a nodal basis for DG sparse grids
#
# -----------------------------------------------------------

#=

We begin by defining analogues of our DG functions
to satisfy the following conditions:

	1) There are a set of collocation points respecting the hierarchical structure
	2) The nodal basis functions at a given level are those that vanish
		on all nodal points at the above levels and vanish on all but one
		of the cells on the level below
	3) Third constraint?

To do this, we _restrict_ to the continuous part of the DG basis, and focus
on piecewise polynomials that vanish on their boundaries. We then pick the
collocation points to be exactly at these boundaries.

For this, it is useful to define lagrange polynomials, and moreover
use k=2, 3, 5.
The k=2 case has our basis is exactly the hat function basis constructed
previously.

For k=2, the coarsest level has nodal points at {0, 1},
	the next level (the first hat) has nodal point at 0.5
	the next level has two cells with points at 0.25, 0.75 etc.

For k=3 the coarsest level has nodal points at {0, 0.5, 1},
	the next level has two nodal functions with points at 0.25, 0.75,
	the next level has two cells and two nodes per cell at all odd multiples of 0.125 etc.

For k=5 the coarsest level has nodal points at {0, 0.25, 0.5, 0.75, 1}
	the next level has four nodal functions with points at all multiples of 0.125
	the next level has two cells and four nodes per cell at all odd multiples of 0.0625 etc.

Notice that k=2,3,5 work because at each cell of the multiresolution basis the space of
	picewise continuous polynomials is k-1 = 2^i for i=0,1,2 in this case.
	A construction for other k would be harder, if not impossible.

=#

# Lagrange polynomials interpolating at the coarsest level
function lag_nodal(k::Int, mode::Int, x::Real)
	(k < 2 || k > 5) && throw(ArgumentError("k is not between 2 and 5"))
	(mode > k || mode < 1) && throw(ArgumentError("mode is not in [1, k]"))
	if k == 2
		mode == 1 && return max(zero(x), x)
		mode == 2 && return max(zero(x), 1-x)
	elseif k == 3
		mode == 1 && return (x < 0 || x > 1) ? zero(x) : 2 * (x-1) * (x - 1//2)
		mode == 2 && return (x < 0 || x > 1) ? zero(x) : 2 * x * (x - 1//2)
		mode == 3 && return (x < 0 || x > 1) ? zero(x) : -4 * x * (x-1)
	elseif k == 4
		throw(MethodError("k = 4 not implemented"))
	elseif k == 5
		mode == 1 && return (x > 1 || x < 0) ? zero(x) : 32//3 * (x - 1//4) * (x - 1//2) * (x - 3//4) * (x - 1)
		mode == 2 && return (x > 1 || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 3//4)
		mode == 3 && return (x > 1 || x < 0) ? zero(x) : 64 * x * (x - 1) * (x - 1//4) * (x - 3//4)
		mode == 4 && return (x > 1 || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 1//4)
		mode == 5 && return (x > 1 || x < 0) ? zero(x) : 32//3 * x * (x - 1//4) * (x - 1//2) * (x - 3//4)
	end
end

# Multiresolution functions for the levels below
function h_nodal(k::Int, mode::Int, x::Real)
	(k < 2 || k > 5) && throw(ArgumentError("k is not between 2 and 5"))
	(mode >= k) && return zero(x)
	#(abs(x) >= 1) && return zero(T)
	if k == 2
		return max(zero(x), 1-abs(x))
	elseif k == 3
		mode == 1 && return max(zero(x), -4 * x * (x + 1))
		mode == 2 && return max(zero(x), -4 * x * (x - 1))
	elseif k == 4
		throw(MethodError("k = 4 not implemented"))
	elseif k == 5
		mode == 1 && return (x < -1 || x > 0) ? zero(x) : -128//3 * x * (x + 1) * (x + 1//2) * (x + 1//4)
		mode == 2 && return (x < -1 || x > 0) ? zero(x) : -128//3 * x * (x + 1) * (x + 1//2) * (x + 3//4)
		mode == 3 && return (x > 1  || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 3//4)
		mode == 4 && return (x > 1  || x < 0) ? zero(x) : -128//3 * x * (x - 1) * (x - 1//2) * (x - 1//4)
	end
end

# This yields our nodal basis:
function v_nodal(k::Int, level::Int, cell::Int, mode::Int, x::Real)
	if level==0
		return lag_nodal(k, mode, x)
	else
		return h_nodal(k, mode, (1<<level)*x - (2*cell-1))
	end
end

function v_nodal(k::Int, level::Int, cell::Int, mode::Int)
	return x->v_nodal(k, level, cell, mode, x)
end

# -----------------------------------------------------------
# Next we define the sparse matrices transforming between
# the heirarchical, nodal, and point bases
# -----------------------------------------------------------

# Evaluation of a given nodal function along all relevant gridpoints
function eval_points_1D(k::Int, max_level::Int, level::Int, cell::Int, mode::Int)
	i = 1
	I = Int[]
	V = Float64[]
	for l in 0:max_level
		l == 0 ? (node_range = 0:(1//(k-1)):1) :
				 (node_range = 1//(2*(k-1)):1//(k-1):(1-1//(2*(k-1))))

		for c in 1:1<<max(0, l-1)
			for node in node_range
				x = (c - 1 + node)//(1<<max(0, l-1))
				val = v_nodal(k, level, cell, mode, x)
				(val != zero(x)) && (push!(I, i); push!(V, val))
				i += 1
			end
		end
	end
	return dropzeros!(sparsevec(I, V))
end

# Sparse matrix converting nodal basis element to its value
# on the collocation points
function nodal2points_1D(k::Int, max_level::Int)
	I = Int[]
	J = Int[]
	V = Float64[]
	j = 1

	for level in 0:max_level
		level == 0 ? (mode_range = 1:k) : (mode_range = 1:(k-1))

		for cell in 1:1<<max(0, level-1)
			for mode in mode_range
				coeffs = eval_points_1D(k, max_level, level, cell, mode)
				for (i, val) in zip(findnz(coeffs)...)
					push!(I, i)
					push!(J, j)
					push!(V, val)
				end
				j += 1
			end
		end
	end
    return dropzeros!(sparse(I, J, V))
end

# Sparse matrix inverting the above construction
# Converting a set of values on the collocation points to the DG space
# in terms of the (for now continuous) nodal basis
# TODO: Allow for this to handle discontinuity across cells
function points2nodal_1D(k::Int, max_level::Int)
    Q = nodal2points_1D(k, max_level)
    return threshold(inv(full(Q)), abs_tol=1e-15)
end

# Converts a given nodal basis element to a sum in the standard
# hieararchical DG scheme
function nodal2modal_1D(k::Int, level::Int, cell::Int, mode::Int; rel_tol=1e-10, abs_tol=1e-15, max_evals=1500)
	I = Int[]
	V = Float64[]
	i = 1

	for hier_level in 0:level
		hier_cells = 1<<max(0, hier_level-1)
		for hier_cell in 1:hier_cells
			for hier_mode in 1:k
				x_min = (hier_cell-1)//hier_cells
				x_max = (hier_cell)//hier_cells
				val = hquadrature(x->v_nodal(k, level, cell, mode, x)*
									v(k, hier_level, hier_cell, hier_mode, x),
									x_min, x_max,
									reltol=rel_tol, abstol=abs_tol, maxevals=max_evals)[1]
				(abs(val) > eps(abs_tol)) && (push!(I, i); push!(V, val))
				i += 1
			end
		end
	end
	return dropzeros!(sparsevec(I, V))
end

function nodal2modal_1D(k::Int, max_level::Int; rel_tol=1e-10, abs_tol=1e-15, max_evals=1500)
	I = Int[]
	J = Int[]
	V = Float64[]
	j = 1

	for nodal_level in 0:max_level
		nodal_level == 0 ? num_nodes = k : num_nodes = k-1
		nodal_cells = 1<<max(0, nodal_level-1)
		for nodal_cell in 1:nodal_cells
			for nodal_mode in 1:num_nodes
				coeffs = nodal2modal_1D(k, nodal_level, nodal_cell, nodal_mode;
										rel_tol=rel_tol,
										abs_tol=abs_tol,
										max_evals=max_evals)
				for (i, val) in zip(findnz(coeffs)...)
					push!(I, i)
					push!(J, j)
					push!(V, val)
				end
				j += 1
			end
		end
	end
	return sparse(I, J, V)
end

# Accounting for ambiguity in the definition of the hierarchical
# basis across regions
function eval_v(k, level, cell, mode, x)
	xl = max(x - 1e-16, 0)
	xr = min(x + 1e-16, 1)
	(xl == 0) && (xl = 1 - 1e-16)
	(xr == 1) && (xr = 1e-16)
	return 0.5 * (v(k, level, cell, mode, xl) + v(k, level, cell, mode, xr))
end

# Evaluates a given hierarchical basis function at
# all collocation points
function modal2points_1D(k::Int, max_level::Int, level::Int, cell::Int, mode::Int)
	i = 1
	I = Int[]
	V = Float64[]

	# subsequent levels:
	for l in 0:max_level
		l == 0 ? (node_range = 0:(1/(k-1)):1) :
				 (node_range = 1/(2*(k-1)):1/(k-1):(1-1/(2*(k-1))))

		for c in 1:1<<max(0, l-1)
			for node in node_range
				x = (c - 1 + node)/(1<<max(0, l-1))
				val = eval_v(k, level, cell, mode, x)
				(abs(val) != 0) && (push!(I, i); push!(V, val))
				i += 1
			end
		end
	end
	return dropzeros!(sparsevec(I, V))
end

# Constructs the sparse matrix to convert from the
# hierarchical basis to the point basis
function modal2points_1D(k::Int, max_level::Int)
	I = Int[]
	J = Int[]
	V = Float64[]
	j = 1

	for level in 0:max_level
		for cell in 1:1<<max(0, level-1)
			for mode in 1:k
				coeffs = modal2points_1D(k, max_level, level, cell, mode)
				for (i, val) in zip(findnz(coeffs)...)
					push!(I, i)
					push!(J, j)
					push!(V, val)
				end
				j += 1
			end
		end
	end
	return dropzeros!(sparse(I, J, V))
end

function nodal2pos_1D(k, max_level, nodal_level, nodal_cell, nodal_mode;
				   rel_tol=1e-10, abs_tol=1e-15, max_evals=1500)
	I = Int[]
	V = Float64[]
	i = 1

	nodal_min = (nodal_cell-1)/(1<<max(0, nodal_level-1))
	nodal_max = (nodal_cell)/(1<<max(0, nodal_level-1))
	for c in 1:1<<max_level
		pos_min = (c-1)/(1<<max_level)
		pos_max = c/(1<<max_level)

		if (pos_min > nodal_max) || (pos_max < nodal_min)
			i += k
			continue
		end
		for m in 1:k
			val = hquadrature(x->v_nodal(k, nodal_level, nodal_cell, nodal_mode, x)*basis(max_level, c, m, x),
							  pos_min, pos_max,
							  reltol=rel_tol, abstol=abs_tol, maxevals=max_evals)[1]
			(abs(val) > eps(abs_tol)) && (push!(I, i); push!(V, val))
			i += 1
		end
	end
	return sparsevec(I, V)
end

function nodal2pos_1D(k, max_level; rel_tol=1e-10, abs_tol=1e-15, max_evals=1500)
	I = Int[]
	J = Int[]
	V = Float64[]
	j = 1

	for nodal_level in 0:max_level
		nodal_level == 0 ? num_nodes = k : num_nodes = k-1
		nodal_cells = 1<<max(0, nodal_level-1)
		for nodal_cell in 1:nodal_cells
			for nodal_mode in 1:num_nodes
				coeffs = nodal2pos_1D(k, max_level, nodal_level, nodal_cell, nodal_mode;
										rel_tol=rel_tol,
										abs_tol=abs_tol,
										max_evals=max_evals)
				for (i, val) in zip(findnz(coeffs)...)
					push!(I, i)
					push!(J, j)
					push!(V, val)
				end
				j += 1
			end
		end
	end
	return sparse(I, J, V)
end

# Inverts the above construction by factoring through
# Essentially all the different spaces
# TODO: Make this more direct!!
function points2modal_1D(k::Int, max_level::Int)
	p2n = points2nodal_1D(k, max_level)
	n2pos = nodal2pos_1D(k, max_level)
	pos2m = pos2hier(k, max_level)
	return threshold(pos2m*n2pos*p2n)
end


using PredictiveControl
using Test

H = [ 10.0 0.0;
       0.0 2.0 ];

b = [ 2.0;
      4.0 ];

# Run for only 40 iterations
x, iter = PredictiveControl.FastGradientMethod.fastgradientmethod( H, b, x -> x, stopconditions = nothing, maxiter = 40, dispinterval = 0 )
@test iter == 40

# Run until converged
x, iter = PredictiveControl.FastGradientMethod.fastgradientmethod( H, b, x -> x, stopconditions = [FGM.Conjugate(1e-8)], dispinterval = 0 )
@test isapprox( x, [-0.2, -2.0], atol = 1e-4 )

x, iter = PredictiveControl.FastGradientMethod.fastgradientmethod( H, b, x -> x, stopconditions = [FGM.Gradient(1e-8)], dispinterval = 0 )
@test isapprox( x, [-0.2, -2.0], atol = 1e-4 )

x, iter = PredictiveControl.FastGradientMethod.fastgradientmethod( H, b, x -> x, stopconditions = [FGM.Best(1e-8)], dispinterval = 0 )
@test isapprox( x, [-0.2, -2.0], atol = 1e-4 )

# Start on the actual point - should converge very fast
x, iter = PredictiveControl.FastGradientMethod.fastgradientmethod( H, b, x -> x, x₀ = [-0.2; -2.0], stopconditions = [FGM.Best(1e-8)], dispinterval = 0 )
@test iter <= 2
@test isapprox( x, [-0.2, -2.0], atol = 1e-4 )

# Test a clipped region (project onto constraints)
x, iter = PredictiveControl.FastGradientMethod.fastgradientmethod( H, b, x -> max.(-1.0, min.(1.0, x)), stopconditions = [FGM.Best(1e-8)], dispinterval = 0 )
@test isapprox( x, [-0.2, -1.0], atol = 1e-4 )

# Test the variable steping
x, iter = PredictiveControl.FastGradientMethod.fastgradientmethod( H, b, x -> x, step=FGM.VariableStep(), stopconditions = [FGM.Best(1e-8)], dispinterval = 0 )
@test isapprox( x, [-0.2, -2.0], atol = 1e-4 )


"""
function ptcsoli( F!, x0, FS, FPS, Jvec = dirder; rtol = 1.e-6, atol = 1.e-12,
    maxit = 20, lmaxit = -1, lsolver = "gmres", eta = 0.1, fixedeta = true,
    Pvec = nothing, PvecKnowsdelta = false, pside = "right", delta0 = 1.e-6,
    dx = 1.e-7, pdata = nothing, printerr = true, keepsolhist = false,
)

C. T. Kelley, 2021

Julia versions of the nonlinear solvers from my SIAM books.
Herewith: some new stuff ==> ptcsoli

PTC finds the steady-state solution of u' = -F(u), u(0) = u_0.
The - sign is a convention.

You must allocate storage for the function and Krylov basis in advance
--> in the calling program <-- ie. in FS and FPS

Inputs:\n
- F!: function evaluation, the ! indicates that F! overwrites FS, your
    preallocated storage for the function.\n
    So, FV=F!(FV,x) or FV=F!(FV,x,pdata) returns FV=F(x)

- x0: initial iterate\n

- FS: Preallocated storage for function. It is an N x 1 column vector.\n
You may dimension it as (n,) or (n,1). (n,) is best, but the
solvers can deal with it either way.

- FPS: preallocated storage for the Krylov basis. It is an N x m matrix where
       you plan to take at most m-1 GMRES iterations before a restart. \n

- Jvec: Jacobian vector product, If you leave this out the
    default is a finite difference directional derivative.\n
    So, FP=Jvec(v,FS,x) or FP=Jvec(v,FS,x,pdata) returns FP=F'(x) v. \n
    (v, FS, x) or (v, FS, x, pdata) must be the argument list,
    even if FP does not need FS.
    One reason for this is that the finite-difference derivative
    does and that is the default in the solver.

- Precision: Lemme tell ya 'bout precision. I designed this code for
    full precision functions and linear algebra in any precision you want.
    You can declare FPS as Float64 or Float32 and ptcsoli
    will do the right thing. Float16 support is there, but not working well.

    If the Jacobian is reasonably well conditioned, you can cut the cost
    of orthogonalization and storage (for GMRES) in half with no loss.
    There is no benefit if your linear solver is not GMRES or if
    othogonalization and storage of the Krylov vectors is only a
    small part of the cost of the computation. So if your preconditioner
    is good and you only need a few Krylovs/Newton, reduced precision won't
    help you much.

    BiCGSTAB does not benefit from reduced precsion.

----------------------

Keyword Arguments (kwargs):\n

rtol and atol: relative and absolute error tolerances\n

delta0: initial pseudo time step. The default value of 1.e-3 is a bit conservative
and is one option you really should play with. Look at the example
where I set it to 1.0!\n

maxit: limit on nonlinear iterations, default=100. \n
This is coupled to delta0. If your choice of delta0 is too small (conservative)
then you'll need many iterations to converge and will need a larger
value of maxit

For PTC you'll need more iterations than for a straight-up
nonlinear solve. This is part of the price for finding the
stable solution. \n

lmaxit: limit on linear iterations. If lmaxit > m-1, where FPS has
m columns, and you need more than m-1 linear iterations, then GMRES
will restart.

The default is -1. For GMRES this means that you'll take m-1 iterations, where
size(V) = (n,m), and get no restarts. For BiCGSTAB you'll then get the default
of 10 iterations.

lsolver: the linear solver, default = "gmres"\n
Your choices will be "gmres" or "bicgstab". However,
gmres is the only option for now. \n

eta and fixed eta: eta > 0 or there's an error.

The linear solver terminates when ||F'(x)s + F(x) || <= etag || F(x) ||

where

etag = eta if fixedeta=true

etag = Eisenstat-Walker as implemented in book if fixedeta=false

The default, which may change, is eta=.1, fixedeta=true \n

Pvec: Preconditioner-vector product. The rules are similar to Jvec
    So, Pv=Pvec(v,x) or Pv=Pvec(v,x,pdata) returns P(x) v where
    P(x) is the preconditioner. You must use x as an input even
    if your preconditioner does not depend on x.\n

PvecKnowsdelta: If you want your preconditioner-vector product to depend on
    the pseudo-timestep delta, put an array deltaval in your precomputed
    data. Initialize it as
    deltaval = zeros(1,)
    and let ptcsoli know about it by setting the kwarg
    PvecKnowsdelta = true
    ptcsoli will update the value in deltaval with every change
    to delta with pdata.deltaval[1]=delta
    so your preconditioner-vector product can get to it.\n

pside: apply preconditioner on pside, default = "right". I do not
      recommend "left". The problem with "left" for ptcsoli is
      that it can fail to satisfy the inexact Newton condition for
      the unpreconditioned equation, especially early in the iteration
      and lead to an incorrect result (unstable solution or wrong
      branch of steady state).
      See Chapter 3 for the story on this. \n

dx: default = 1.e-7\n
difference increment in finite-difference derivatives
      h=dx*norm(x)+1.e-8 \n

pdata:\n
precomputed data for the function/Jacobian-vector/Preconditioner-vector
products.  Things will go better if you use this rather than hide the data
in global variables within the module for your function/Jacobian

If you use pdata in any of F!, Jvec, or Pvec, you must use in in all of them.
precomputed data for the function/Jacobian.
Things will go better if you use this rather than hide the data
in global variables within the module for your function/Jacobian. \n

printerr: default = true\n
I print a helpful message when the solver fails. To suppress that
message set printerr to false. \n

keepsolhist: default = false\n
Set this to true to get the history of the iteration in the output
tuple. This is on by default for scalar equations and off for systems.
Only turn it on if you have use for the data, which can get REALLY LARGE.\n

Output:\n
A named tuple (solution, functionval, history, stats, idid,
               errcode, solhist)
where

solution = converged result
functionval = F(solution)
history = the vector of residual norms (||F(x)||) for the iteration
stats = named tuple of the history of (ifun, ijac, ikfail), the number
of functions/jacobian-vector prodcuts/linear solver filures at each iteration.

I do not count the function values for a finite-difference derivative
because they count toward a Jacobian-vector product.

Linear solver failures need not cause the nonlinear iteration to fail.
You get a warning and that is all. \n

idid=true if the iteration succeeded and false if not. \n

errcode = 0 if if the iteration succeeded \n
        = -1 if the initial iterate satisfies the termination criteria
        = 10 if no convergence after maxit iterations \n

solhist:\n
This is the entire history of the iteration if you've set
keepsolhist=true\n
solhist is an N x K array where N is the length of x and K is the number
of iteration + 1. So, for scalar equations, it's a row vector.

### Example for ptcsol

#### The buckling beam problem.
You'll need to use TestProblems for
this to work. The preconditioner is a solver for the high order term.

```jldoctest
julia> using SIAMFANLEquations.TestProblems

julia> function PreCondBeam(v, x, bdata)
          J = bdata.D2
          ptv = J\v
       end
PreCondBeam (generic function with 1 method)

julia> n=63; maxit=1000; delta0 = 0.01; lambda = 20.0;

julia> bdata = beaminit(n, 0.0, lambda);

julia> x = bdata.x; u0 = x .* (1.0 .- x) .* (2.0 .- x); u0 .*= exp.(-10.0 * u0);

julia> FS = copy(u0); FPJV=zeros(n,20);

julia> pout = ptcsoli( FBeam!, u0, FS, FPJV; delta0 = delta0, pdata = bdata,
       eta = 1.e-2, rtol = 1.e-10, maxit = maxit, Pvec = PreCondBeam);

julia> # It takes a few iterations to get there.
       length(pout.history)
25

julia> [pout.history[1:5] pout.history[21:25]]
5×2 Array{Float64,2}:
 6.31230e+01  1.79578e+00
 7.45926e+00  2.65964e-01
 8.73598e+00  6.58278e-03
 2.91936e+01  8.35069e-06
 3.47969e+01  5.11594e-09

julia> # We get the nonnegative stedy state.
       norm(pout.solution,Inf)
2.19086e+00

n=63; maxit=1000; delta0 = 0.01; lambda = 20.0;

julia> # Use BiCGSTAB for the linear solver

julia> FS = copy(u0); FPJV=zeros(n,);

julia> pout = ptcsoli( FBeam!, u0, FS, FPJV; delta0 = delta0, pdata = bdata,
       eta = 1.e-2, rtol = 1.e-10, maxit = maxit,
       Pvec = PreCondBeam, lsolver="bicgstab");

julia> # Same number of iterations as GMRES, but each one costs double

julia> # the Jacobian-vector products and much less storage

julia> length(pout.history)
25

julia> [pout.history[1:5] pout.history[21:25]]
5×2 Matrix{Float64}:
 6.31230e+01  1.68032e+00
 7.47081e+00  2.35073e-01
 8.62095e+00  5.18262e-03
 2.96495e+01  3.23715e-06
 3.51504e+01  3.33107e-10

```

"""
function ptcsoli(
    F!,
    x0,
    FS,
    FPS,
    Jvec = dirder;
    rtol = 1.e-6,
    atol = 1.e-12,
    maxit = 20,
    lmaxit = -1,
    lsolver = "gmres",
    eta = 0.1,
    fixedeta = true,
    Pvec = nothing,
    PvecKnowsdelta = false,
    pside = "right",
    delta0 = 1.e-6,
    dx = 1.e-7,
    pdata = nothing,
    printerr = true,
    keepsolhist = false,
)
    itc = 0
    idid = true
    #
    #   Initialize the iteration
    #   As with the other codes, ItRules packages all the details of
    #   the problem so it's easy to pass them around.
    #
    (ItRules, x, n) = PTCKrylovinit(
        x0,
        dx,
        F!,
        Jvec,
        delta0,
        Pvec,
        PvecKnowsdelta,
        pside,
        lsolver,
        eta,
        fixedeta,
        lmaxit,
        maxit,
        printerr,
        pdata,
    )
    keepsolhist ? (solhist = solhistinit(n, maxit, x)) : (solhist = [])
    #
    # First Evaluation of the function. Initialize the iteration history.
    # Fix the tolerances for convergence and define the derivative FPF
    # outside of the main loop for scoping.
    #
    FS = EvalF!(F!, FS, x, pdata)
    resnorm = norm(FS)
    ItData = ItStatsPTCK(resnorm)
    tol = rtol * resnorm + atol
    etag = eta
    ke_report = false
    residratio = 1.0
    #
    # Preallocate a vector for the step
    #
    step = copy(x)
    #
    # If the initial iterate satisfies the termination criteria, tell me.
    #
    toosoon = (resnorm <= tol)
    #
    # The main loop stops on convergence or too many iterations.
    #
    delta = delta0
    while resnorm > tol && itc < maxit
        residm = resnorm
        newjac = 0
        newikfail = 0
        #
        # Comppute the Jacobian-vector product; update x, F(x), and delta.
        #
        etag = forcing(itc, residratio, etag, ItRules, tol, resnorm)
        (x, delta, FS, resnorm, Lstats) =
            PTCUpdatei(FPS, FS, x, ItRules, step, resnorm, delta, etag)
        resdiratio = resnorm / residm
        newjac = Lstats.lits
        linok = Lstats.idid
        linok || (ke_report = Krylov_Error(lmaxit, ke_report); newikfail = 1)
        #
        # Keep the books
        #
        updateStats!(ItData, resnorm, newjac, newikfail)
        itc += 1
        ~keepsolhist || (@views solhist[:, itc+1] .= x)
    end
    (idid, errcode) = PTCOK(resnorm, tol, toosoon, ItRules, printerr)
    itout = CloseIteration(x, FS, ItData, idid, errcode, keepsolhist, solhist)
    return itout
end


def (fib n)
  if (n == 0)
    then 0
  else if (n == 1)
    then 1
  else
    ((fib (n-1)) +
     (fib (n-2)))

def (fact n)
  if (n == 0)
    then 1
  else
    ((fact (n-1)) * n)

def ns [1,2,3,4,5]

def main
  ((print (fib 19)) >>
   (print (ns >>= (return . fact))))


module AnalyticalEOR

using Base:Real
using ForwardDiff
using Roots
using Plots

include("types.jl")
include("flowFunctions.jl")
include("buckleyLeverett.jl")
include("polymerFlooding.jl")
include("tracers.jl")
include("plottingFunctions.jl")

export RelPerms, WaterFlooding
export water_rel_perm, oil_rel_perm, krw_derivative, kro_derivative
export fractional_flow, fw_derivative
export solve_waterflooding, saturation_profile
export solve_tracer, tracer_profile
export plot_fw, plot_sw_profile, animate_sw_profile
export plot_tracer_fw
export solve_polymerflood

end


#=
   Euler #2 in Julia.

  Problem 2
  """
  Each new term in the Fibonacci sequence is generated by adding the
  previous two terms. By starting with 1 and 2, the first 10 terms will be:

  1, 2, 3, 5, 8, 13, 21, 34, 55, 89, ...

  Find the sum of all the even-valued terms in the sequence which do not
  exceed four million.
  """

  This Julia program was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/
=#

using Memoization
include("Euler.jl")

# 7.5558e-5
function euler2a()
    n = 100 # some arbitrary large value
    f = zeros(BigInt,(n,1))
    f[1:2] = [0,1]
    [f[i] = f[i-1]+f[i-2] for i in 3:n]
    sum([f[i]*(f[i]%2 == 0 && f[i] < 4000000) for i in 1:n])
end

const fibmem2b = Dict{Int,Int}()
function fib2b(n::Int)
    get!(fibmem2b, n) do
        n < 3 ? 1 : fib2b(n-1) + fib2b(n-2)
    end
end

# 8.108e-6s
function euler2b()
    n = 1
    f = 0
    s = 0
    while f < 4_000_000
        f = fib2b(n)
        if f % 2 == 0
            s += f
        end
        n += 1
    end
    return s
end

#
# This is slower than using a Dict!
#
@memoize function fib2c(n::Int)
    n < 3 ? 1 : fib2c(n-1) + fib2c(n-2)
end

# 0.008s
function euler2c()
    n = 1
    s = 0
    while (f = fib2c(n)) < 4_000_000
        if f % 2 == 0
            s += f
        end
        n += 1
    end
    return s
end

# Using no cache
function fib_no_cache(n::Int)
  if n <= 2
     return 1
  end
  return fib_no_cache(n-1)+fib_no_cache(n-2)
end

# 0.0379s
function euler2d()
    n = 1
    s = 0
    while (f = fib_no_cache(n)) < 4_000_000
        if f % 2 == 0
            s += f
        end
        n += 1
    end
    return s
end

# run_euler(euler2a)
run_euler(euler2b)
# run_euler(euler2c)
# run_euler(euler2d)


# using Revise # uncomment for development

using Test
using Random
using StatsFuns
using Optim
using StatsBase: countmap
using LinearAlgebra
using ForwardDiff
using DataFrames
using Distributions
using BenchmarkTools: @btime

using rcmnl

@testset "Gauss Hermite integration" begin
    @test rcmnl.integrate_wrt_normal_pdf(x -> x, 1.23, 2.34; npts=15) ≈ 1.23
    @test sqrt(rcmnl.integrate_wrt_normal_pdf(x -> (x-1.23)^2, 1.23, 2.34; npts=15)) ≈ 2.34

    μ = [1.0, 0.5]
    L = [1.0  0.0; 0.7 0.3]

    @test rcmnl.integrate_wrt_normal_pdf(x -> (x-μ)*(x-μ)', μ, L; npts=7) ≈ L*L'
    @test rcmnl.integrate_wrt_normal_pdf(x -> x,            μ, L; npts=7) ≈ μ
end

Random.seed!(1234)

# dim of data
nindv    = 1_000
nt       = 20
nquadpts = 7

# coefs
β = [ 1.0  0.3;
     -2.0  0.0;
      0.0  1.0;
      0.2  0.5]
k, nchoice = size(β)
@assert nchoice == 2

# cholesky decomposition of Σ
Σchol = [1.0 0.0;
         0.5 0.5]
θtrue = vcat(vec(β), Σchol[:,1], Σchol[end])

# generate X and V
X = randn(nt, nindv, k)
V = randn(nindv,2)*Σchol' # need Σchol b/c usu. assume U is a 2-dim vector

# utilities
U = zeros(3, nt, nindv)
mul!(reshape(view(U,2:3,:,:), 2,:), β', reshape(X, :, k)')
U[2,:,:] .+= V[:,1]'
U[3,:,:] .+= V[:,2]'

# probabilities
PrU = mapslices(softmax, U; dims=1)
cumPrU = cumsum(PrU; dims=1)

# instead of simulating random type-1 extreme values, we just
# use a uniform variable and the CDF
y = Array{Int64}(undef, nt, nindv)
for idx in CartesianIndices(y)
    vw = view(cumPrU, :, idx)
    y[idx] = searchsortedfirst(vw, rand())
end

# how many of each choice?
countmap(vec(y))

# run function once
simlogL(y, X, θtrue; npts=nquadpts)
println("It runs!")

# wrapper for optimizer
f(θ) = simlogL(y, X, θ; npts=nquadpts)

# starting value
theta0 = vcat(vec(β), Σchol[:,1], Σchol[end]) .* 2

# time function
@btime f(theta0)

# profile function in IDE
# @profileview f(theta0)

# optimize
res = optimize(f, theta0, BFGS(), Optim.Options(;show_trace=true), autodiff=:forward)

# get results
thetahat = res.minimizer
H = ForwardDiff.hessian(f, thetahat)
se = sqrt.(diag(inv(H)))
pval = 2*normcdf.(-abs.(thetahat./se))

# nice table
results_to_print = DataFrame(
    θtrue = θtrue,
    θhat = thetahat,
    se = se,
    pval = pval
)
@show results_to_print

# Wald test
alpha_reject = 0.001
wald_test_stat = (thetahat - θtrue)' * (H \ (thetahat - θtrue))
wald_test_p = ccdf(Chisq(length(θtrue)), wald_test_stat)
if wald_test_p < alpha_reject
    println("Reject that we got the right answer")
else
    println("We fail to reject our estimate!!")
end

@test wald_test_p > alpha_reject


using PyPlot, Seismic

dt = 0.002
w = Ricker(dt=dt)
wmin = SeisKolmogoroff(w)

nw = length(w)
t = dt*collect(0:1:nw-1)

figure("Kolmogoroff factorization", figsize=(6, 6))

subplot(211)
plot(t, w)
axis("tight")
xlabel("Time (s)")
title("Ricker wavelet")

subplot(212)
plot(t,wmin)
axis("tight")
xlabel("Time (s)")
title("Minimum Phase equivalent")

tight_layout()


# MIT License

# Copyright (c) 2020 Makoto Suwama

# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

# Code used to compute cycles coming from the sum of squares of digits in base b for the paper Integer Dynamics with Dino Lorenzini, Mentzelos Melistas, Arvind Suresh and Haiyang Wang
# For example, to check in base b, all cycles starting at x+b*y with low<= y <= high, run
# julia cycles.jl low high b
# May need to change the data type of b from UInt32 to other type depending on how large the b is
# Also see correctness.pdf for the correctness of the code.

const numCycles=3;

# Parameters parsed from the command line
low = parse(Int, ARGS[1]);
high = parse(Int, ARGS[2]);
const b = parse(Int, ARGS[3]);

# Directory to save the files
directory = "";
output = directory * "spTestJ" * string(b) * "-" * string(low) * "-" * string(high) * ".txt";

# check a number d has been seen already
# seen is an array of arrays of sets and a number n = d1 + d2*b + d3*b^2 is stored as d1 in seen[d3+1][d2+1]
function isSeen(seen,d)
    return d[1] in seen[d[3]+1][d[2]+1];
end;
# add number d to seen
# seen is an array of arrays of sets and a number n = d1 + d2*b + d3*b^2 is stored as d1 in seen[d3+1][d2+1]
function addSeen!(seen,d)
    push!(seen[d[3]+1][d[2]+1], d[1])
end;
# replaces d with [x,y,0]
function createD!(d,x,y)
    d[1] = x;
    d[2] = y;
    d[3] = 0;
end
# Checks if max(d[1],d[2]) + min(d[1],d[2])*b + d[2]*b^2 is smaller than x+b*y
# Uses separate comparison for speed purpose
function isSmaller(d,x,y)
    return d[3] == 0 && (min(d[1],d[2]) < y || (min(d[1],d[2])==y && max(d[1],d[2]) < x))
end

# start the sum of the square procedure until it hits a cycle or a number already seen
# Returns True if it finds a new cycle
function traverse(x,y,d,seen,cycles,L)
    # if the number already appeared then skip it
    if x in seen[1][y+1]
        return false
    end;
    n = x+b*y;

    # pops L until it is empty and adds n
    # this is faster than L = [n]
    while length(L) > 0
        pop!(L);
    end;
    push!(L,n);
    while !(isSeen(seen,d))
        addSeen!(seen,d)
        n = sum(i^2 for i in d);
        push!(L,n)
        digits!(d, n,base = b)
        # if d is less than x+b*y then we have seen it already
        if isSmaller(d,x,y)
            return false
        end
    end
    # find the position of n in L
    # if n is in L and not the last then there is a cycle starting at the first appearance of n
    p = findfirst(isequal(n),L);
    if p != nothing && p < length(L)
        push!(cycles, L[p:length(L)-1])
        return true
    end
    return false
end

# Check all the numbers x+b*y with low<=y<=high
function checkYInterval(low, high, seen, cycles)
    d=Array{UInt64}(undef,3)
    L=[];
    for y=low:high
        # Check if the progress needs to be printed
        if y % interval == 0
            println(y, " ", time()-t0)
            flush(stdout)
        end
        # See correctness.pdf for the upper bound on x
        for x=max(y,ceil(Int, sqrt(y*b-y^2+1/4)+1/2)):b-1
            createD!(d,x,y)
            if traverse(x,y,d,seen,cycles,L)
                #if length(cycles) >= numCycles
                #    return nothing
                #end
            end
        end
        # delete unncessary seen to save memory. Boolean is chosen to save space
        seen[1][y+1] = Set{Bool}()
    end
end

# Adds a cycle starting at ini
function addKnownCycle(seen, cycles, ini, d)
    c = [];
    while true
        digits!(d,ini,base = b)
        if isSeen(seen, d)
            push!(cycles, c)
            return
        end
        addSeen!(seen,d)
        push!(c, ini)
        ini = sum(i^2 for i in d)
    end
end

# Initialise the variables
# interval is used to determine how often should we print progress
const t0 = time();
const interval =  10^ceil(Int,log(10,(high-low)/100));

# b is UInt32
# indexing is off by 1
# if n = [a,b,c] then check if a is in seen[c+1][b+1]
# keeps track of [a,b,c] where a>=b and c=0 or 1
seen = [[ Set{UInt32}() for i=1:b], [Set{UInt32}() for i=1:b]];
cycles = [[1]];
push!(seen[1][1],1);
d=Array{UInt64}(undef,3);

addKnownCycle(seen, cycles, 91610640434,d);

checkYInterval(low, high, seen, cycles)

io = open(output,"w");
println(io, b);
println(io, [length(c) for c in cycles]);
println(io, [c[1] for c in cycles]);
flush(io);
close(io);