export WLS, RandomCascade
using Wavelets, Statistics

"""
    WLS(shufflemethod::Surrogate = IAAFT();
        f::Union{Nothing, Function} = Statistics.cor,
        rescale::Bool = true,
        wt::Wavelets.WT.OrthoWaveletClass = Wavelets.WT.Daubechies{16}())

A wavelet surrogate generated by the following procedure:

1. Compute the wavelet transform of the signal. This results in a set of
    detail coefficients over a set of dyadic scales. As in Keylock (2006),
    we here use the maximal overlap discrete wavelet transform, or MODWT,
    so that the number of coefficients at each scale are the same.
2. Shuffle the detail coefficients at each dyadic scale using the
    provided `shufflemethod`. See "Shuffling methods" below for alternatives.
3. Apply the inverse wavelet transform to the shuffled detail coefficients
    to obtain a surrogate time series.

## Shuffling methods

You may choose to use any surrogate from this package to perform the
randomization of the detail coefficients at each dyadic scale.

The following methods have been discussed in the literature (more may exist):

- Random permutations of wavelet coefficients within each scale (Breakspear et al., 2003). To get this behaviour, use `WLS(x, RandomShuffle(), rescale = false, f = nothing)`.
- Cyclic rotation of wavelet coefficients within each scale (Breakspear et al., 2003).  To get this behaviour, use `WLS(x, Circshift(1:length(x)), rescale = false, f = nothing)`.
- Block resampling of wavelet coefficients within each scale (Breakspear et al., 2003). To get this behaviour, use `WLS(x, BlockShuffle(nblocks, randomize = true), rescale = false, f = nothing)`.
- IAAFT resampling of wavelet coefficients within each scale (Keylock, 2006). To get this behaviour, use `WLS(x, IAAFT(), rescale = true, f = Statistics.cor)`.
    This method preserves the local mean and variance structure of the signal, but
    randomises nonlinear properties of the signal (i.e. Hurst exponents)[^Keylock2006].
    These surrogates can therefore be used to test for changes in nonlinear properties
    of the original signal. In contrast to IAAFT surrogates, the IAAFT-wavelet surrogates
    also preserves nonstationarity. Using other `shufflemethod`s does not necessarily
    preserve nonstationarity. To deal with nonstationary signals, Keylock (2006) recommends
    using a wavelet with a high number of vanishing moments. Thus, our default is to
    use a Daubechies wavelet with 16 vanishing moments. *Note: The iterative procedure after
    the rank ordering step (step [v] in [^Keylock2006]) is not performed in
    this implementation.*

The default method and parameters replicate the behaviour of Keylock (2006)'s IAAFT
wavelet surrogates.

## Error minimization

For the [`IAAFT`](@ref) approach introduced in Keylock (2006), detail coefficients
at each level are circularly rotated to minimize an error function. The methods
introduced in Breakspear et al. (2003) do not apply this error minimization.

In our implementation, you can turn this option on/off using the `f` parameter of
the `WLS` constructor. If `f = nothing` turns off error minization. If `f` is set
to a two-argument function that computes some statistic, for example
`f = Statistics.cor`, then detail coefficients at each scale are circularly
rotated until that function is maximized (and hence the "error" minimized).
If you want to *minimize* some error function, then instead provide an appropriate
transform of your function. For example, if using the root mean squared deviation,
define `rmsd_inv(x, y) = 1 - StatsBase.rmsd(x, y)` and set `f = rmsd_inv`.

## Rescaling

If `rescale == true`, then surrogate values are mapped onto the
values of the original time series, as in the [`AAFT`](@ref) algorithm.
If `rescale == false`, surrogate values are not constrained to the
original time series values. If [`AAFT`](@ref) or [`IAAFT`](@ref) shuffling
is used, `rescale` should be set to `true`. For other methods, it does not
necessarily need to be.

[^Breakspear2003]: Breakspear, M., Brammer, M., & Robinson, P. A. (2003). Construction of multivariate surrogate sets from nonlinear data using the wavelet transform. Physica D: Nonlinear Phenomena, 182(1-2), 1-22.
[^Keylock2006]: C.J. Keylock (2006). "Constrained surrogate time series with preservation of the mean and variance structure". Phys. Rev. E. 73: 036707. doi:10.1103/PhysRevE.73.036707.
"""
struct WLS{WT <: Wavelets.WT.OrthoWaveletClass, S <: Surrogate, E <: Union{Nothing, Function}} <: Surrogate
    shufflemethod::Surrogate # should preserve values of the original series
    rescale::Bool
    f::E
    wt::WT

    function WLS(method::S = IAAFT(); rescale::Bool = true, wt::WT = Wavelets.WT.Daubechies{16}(), f::E = Statistics.cor) where {S <: Surrogate, WT <: Wavelets.WT.OrthoWaveletClass, E <: Union{Nothing, Function}}
        new{WT, S, E}(method, rescale, f, wt)
    end
end

# Initialize without error minimization.
function get_init_noerrorminimize(x, method, rng = Random.default_rng())
    wl = wavelet(method.wt)
    L = length(x)
    x_sorted = sort(x)

    # Wavelet coefficients (step [i] in Keylock)
    W = modwt(x, wl)
    T = eltype(W)
    R = zeros(size(W))

    Nscales = ndyadicscales(L)

    # Will contain surrogate realizations of the wavelet coefficients
    # at each scale (step [ii] in Keylock).
    sW = zeros(T, size(W))

    # Surrogate generators for each set of coefficients
    sgs = [surrogenerator(W[:, i], method.shufflemethod, rng) for i = 1:Nscales]

    init = (
        wl = wl,
        W = W,
        Nscales = Nscales,
        sW = sW,
        sgs = sgs,
        x_sorted = x_sorted,
        R = R,
    )

    return init
end

# Initialize error minimization, as in Keylock (2006).
function get_init_errorminimize(x, method::WLS, rng = Random.default_rng())
    wl = wavelet(method.wt)
    L = length(x)
    x_sorted = sort(x)

    # Wavelet coefficients (step [i] in Keylock)
    W = modwt(x, wl)
    T = eltype(W)
    Nscales = ndyadicscales(L)

    # Will contain surrogate realizations of the wavelet coefficients
    # at each scale (step [ii] in Keylock).
    sW = zeros(T, size(W))

    # We will also need a matrix to store the mirror images of the
    # surrogates (last part of step [ii])
    sWmirr = zeros(T, size(W))

    # Surrogate generators for each set of coefficients
    sgs = [surrogenerator(W[:, i], method.shufflemethod, rng) for i = 1:Nscales]

    # Temporary array for the circular shift error minimizing step
    circshifted_s = zeros(T, size(W))
    circshifted_smirr = zeros(T, size(W))
    R = zeros(size(W))

    s = similar(x)

    init = (wl = wl, W = W, Nscales = Nscales, L = L,
            sW = sW, sgs = sgs, sWmirr = sWmirr,
            circshifted_s = circshifted_s,
            circshifted_smirr = circshifted_smirr,
            x_sorted = x_sorted, R = R,
    )

    return init
end

function surrogenerator(x::AbstractVector{T}, method::WLS, rng = Random.default_rng()) where T
    init = isnothing(method.f) ?
        get_init_noerrorminimize(x, method, rng) :
        get_init_errorminimize(x, method, rng)

    s = similar(x)

    return SurrogateGenerator(method, x, s, init, rng)
end

function (sg::SurrogateGenerator{<:WLS})()
    if isnothing(sg.method.f)
        wls_noerrorminimize(sg)
    else
        wls_errorminimize(sg)
    end
end

function wls_noerrorminimize(sg::SurrogateGenerator{<:WLS})
    s = sg.s

    fds = (:wl, :W, :Nscales, :sW, :sgs, :x_sorted, :R)
    wl, W, Nscales, sW, sgs, x_sorted, R = getfield.(Ref(sg.init), fds)

    for λ in 1:Nscales
        sW[:, λ] .= sgs[λ]()
    end

    s .= imodwt(sW, wl)

    if sg.method.rescale
        s[sortperm(s)] .= x_sorted
    end

    return s
end

function wls_errorminimize(sg::SurrogateGenerator{<:WLS})
    s = sg.s

    # Error minimization function (if we reached this function,
    # f is never `Nothing`), so we can use it safely.
    f = sg.method.f

    fds = (:wl, :W, :Nscales, :L, :sW, :sgs, :sWmirr,
        :circshifted_s, :circshifted_smirr,
        :x_sorted, :R)

    wl, W, Nscales, L, sW, sgs, sWmirr,
        circshifted_s, circshifted_smirr,
        x_sorted, R = getfield.(Ref(sg.init), fds)

    # Create surrogate versions of detail coefficients at each dyadic scale [first part of step (ii) in Keylock]
    for λ in 1:Nscales
        sW[:, λ] .= sgs[λ]()
    end

    # Mirror the surrogate coefficients [last part of step (ii) in Keylock]
    sWmirr .= reverse(sW, dims = 1)

    # In the original paper, surrogates and mirror images are matched to original
    # detail coefficients in a circular manner until some error criterion is
    # minimized. Then, the surrogate or its mirror image, depending on which provides
    # the best fit to the original coefficients, is chosen as the representative
    # for a particular dyadic scale. Here, we instead use maximal correlation as
    # the criterion for matching.
    optimal_shifts = zeros(Int, Nscales)
    optimal_shifts_mirr = zeros(Int, Nscales)
    maxcorrs = zeros(Nscales)
    maxcorrs_mirr = zeros(Nscales)

    for i in 0:L-1
        circshift!(circshifted_s, sW, (i, 0))
        circshift!(circshifted_smirr, sWmirr, (i, 0))

        for λ in 1:Nscales
            origW = W[:, λ]
            c = f(origW, circshifted_s[:, λ])
            if c > maxcorrs[λ]
                maxcorrs[λ] = c
                optimal_shifts[λ] = i
            end

            c_mirr = f(origW, circshifted_smirr[:, λ])
            if c_mirr > maxcorrs_mirr[λ]
                maxcorrs_mirr[λ] = c_mirr
                optimal_shifts_mirr[λ] = i
            end
        end
    end

    # Decide which coefficients are retained (either surrogate or mirror surrogate coefficients)
    for λ in 1:Nscales
        if maxcorrs[λ] >= maxcorrs_mirr[λ]
            R[:, λ] .= circshift(sW[:, λ], optimal_shifts[λ])
        else
            R[:, λ] .= circshift(sWmirr[:, λ], optimal_shifts_mirr[λ])
        end
    end

    s .= imodwt(R, wl)

    if sg.method.rescale
        s[sortperm(s)] .= x_sorted
    end

    return s
end

"""
    RandomCascade(paddingmode::String = "zeros")

A random cascade multifractal wavelet surrogate (Paluš, 2008)[^Paluš2008].

If the input signal length is not a power of 2, the signal must be
padded before the surrogate is constructed. `paddingmode` determines
how the signal is padded. Currently supported padding modes: `"zeros"`.
The final surrogate (constructed from the padded signal) is subset
to match the length of the original signal.

Random cascade surrogate preserve multifractal properties of the input
time series, that is, interactions among dyadic scales and nonlinear
dependencies[^Paluš2008].

[^Paluš2008]: Paluš, Milan (2008). Bootstrapping Multifractals: Surrogate Data from Random Cascades on Wavelet Dyadic Trees. Physical Review Letters, 101(13), 134101–. doi:10.1103/PhysRevLett.101.134101
"""
struct RandomCascade{WT <: Wavelets.WT.OrthoWaveletClass} <: Surrogate
    wt::WT
    paddingmode::String

    function RandomCascade(; wt::WT = Wavelets.WT.Daubechies{16}(), paddingmode::String = "zeros") where {WT <: Wavelets.WT.OrthoWaveletClass}
        new{WT}(wt, paddingmode)
    end
end

function surrogenerator(x::AbstractVector{T}, method::RandomCascade, rng = Random.default_rng()) where T
    nlevels = ndyadicscales(length(x))
    mode = method.paddingmode

    # Pad input so that input to discrete wavelet transform has length which is a power of 2
    x̃ = zeros(2^(nlevels + 1))
    if mode == "zeros"
        copyto!(x̃, x)
    else
        throw(ArgumentError("""`paddingmode` must be one of ["zeros"]"""))
    end

    wl = wavelet(method.wt)

    # Wavelet coefficients (step [i] in Keylock)
    c = dwt(x̃, wl, nlevels)

    # Surrogate coefficients will be partly identical to original coefficients,
    # so we simply copy them and replace the necessary coefficients later.
    cₛ = copy(c)

    # Multiplication factors and index vectors can be pre-allocated for
    # levels 2:nlevels-1; they are overwritten for each new surrogate.
    Ms = [zeros(dyadicdetailn(j-1)) for j = 2:nlevels-1]
    ixs = [zeros(Int, dyadicdetailn(j-1)) for j = 2:nlevels-1]

    init = (
        wl = wl,
        c = c,
        cₛ = cₛ,
        nlevels = nlevels,
        s̃ = similar(x̃),
        Ms = Ms,
        ixs = ixs,
    )

    return SurrogateGenerator(method, x, similar(x), init, rng)
end

function (sg::SurrogateGenerator{<:RandomCascade})()
    s, rng = sg.s, sg.rng
    c, cₛ, s̃, wl, nlevels, Ms, ixs =
        sg.init.c, sg.init.cₛ, sg.init.s̃, sg.init.wl, sg.init.nlevels, sg.init.Ms, sg.init.ixs

    cₛ[dyadicdetailrange(0)] = @view c[dyadicdetailrange(0)]
    cₛ[dyadicdetailrange(1)] = @view c[dyadicdetailrange(1)]

    for (l, j) = enumerate(2:nlevels-1)
        cⱼ₋₁ = @view c[dyadicdetailrange(j - 1)]
        cⱼ = @view c[dyadicdetailrange(j)]

        M = Ms[l]
        ct = 1
        @inbounds for k = 1:length(cⱼ₋₁)
            if k % 2 == 0
                M[ct] = cⱼ[2*k] / cⱼ₋₁[k]
            else
                M[ct] = cⱼ[2*(k+1)] / cⱼ₋₁[k+1]
            end
            ct += 1
        end

        shuffle!(rng, M)
        new_coeffs!(M, cⱼ₋₁)
        ix = ixs[l]
        sortperm!(ix, M)
        cₛ[dyadicdetailrange(j-1)] .= @view cⱼ₋₁[ix]
    end
    s̃ .= idwt(cₛ, wl, nlevels)

    # Surrogate length must match length of original signal.
    s .= @view s̃[1:length(s)]
    return s
end

function new_coeffs!(M, cⱼ₋₁)
    @inbounds for k = 1:length(cⱼ₋₁)
        M[k] = M[k] * cⱼ₋₁[k]
    end
end


raw"""
Computes Cholesky decomposition of a Hermitian positive-definite matrix A, i.e.,

$$
A = L L^*
$$

where $L$ is a lower triangular matrix with real and positive diagonal entries, and $L^*$ denotes
the conjugate transpose of $L$.
"""
module CholeskyDecomposition

using LinearAlgebra

export
    cholesky_outer_product_decompose,
    cholesky_banachiewicz_decompose,
    cholesky_crout_decompose

function get_n(A::Matrix)::Int
    (n_r, n_c) = size(A)
    @assert n_r == n_c "Input matrix is not square"
    n_r
end

"""
The Cholesky (outer-product) algorithm

<https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky_algorithm>
"""
function cholesky_outer_product_decompose(A::Matrix{T})::LowerTriangular{T} where T<:Number
    n = get_n(A)

    # Some edge case values:
    #
    #     A_1 = A
    #     A_{n + 1} = I(n)
    #     I_i ∈ ℜ_{i ^ 2}
    #     I_0 = ∅
    #     L_0 = I(n)
    #     B_0 = A
    #     B_i ∈ ℂ_{(n - i) ^ 2}
    #     B_n = ∅

    A_i = copy(A)
    A = nothing  # prevents accidental assignment
    L_0 = T.(I(n)) |> Matrix |> LowerTriangular
    L = copy(L_0)  # The product of L1...Ln
    for i in 1:n
        # Makes aliases (no need for `A_ii`)
        b_i = @view A_i[(i + 1):end, i]
        b_i_H = @view A_i[i, (i + 1):end]
        B_i = @view A_i[(i + 1):end, (i + 1):end]

        L_i = copy(L_0)
        sqrtAii = sqrt(A_i[i, i])
        L_i[i, i] = sqrtAii
        L_i[(i + 1):end, i] = b_i / sqrtAii

        L *= L_i

        # Mind the order
        B_i .-= b_i * transpose(b_i_H) / A_i[i, i]  # Ugly, can't do `b_i * b_i_H`
        b_i .= 0.0
        b_i_H .= 0.0
        A_i[i, i] = 1.0
    end
    L
end

"""
The Cholesky-Banachiewicz algorithm

<https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky–Banachiewicz_and_Cholesky–Crout_algorithms>
"""
function cholesky_banachiewicz_decompose(A::Matrix{T})::LowerTriangular{T} where T<:Number
    n = get_n(A)
    # To write the in-place version, just set `L = A` and do symmetric updating.
    L = zeros(T, n, n) |> LowerTriangular

    # Note that as $A$ is positive-definite, **not positive semi-definite**, $L_{ii}$ will never be
    # zero, so there is no need to check.
    #
    # Also, do NOT use `dot` to compute `L[r, c]`!

    for r in 1:n
        for c in 1:(r - 1)
            L[r, c] = (
                A[r, c] - sum((@view L[r, 1:(c - 1)]) .* conj.(@view L[c, 1:(c - 1)]))
            ) / L[c, c]
        end
        l = @view L[r, 1:(r - 1)]
        L[r, r] = sqrt(A[r, r] - l ⋅ l)
    end
    L
end

"""
The Cholesky-Crout algorithm

It is similar to the Cholesky-Banachiewicz algorithm, but processes column by column instead of row
by row.
"""
function cholesky_crout_decompose(A::Matrix{T})::LowerTriangular{T} where T<:Number
    n = get_n(A)
    L = zeros(T, n, n) |> LowerTriangular
    # Deliberately write the index as `:` rather than `1:(c - 1)`
    for c in 1:n
        l = @view L[c, :]
        L[c, c] = sqrt(A[c, c] - l ⋅ l)
        for r in (c + 1):n
            L[r, c] = (A[r, c] - sum((@view L[r, :]) .* conj.(@view L[c, :]))) / L[c, c]
        end
     end
     L
end

end  # module


plan_clenshawcurtis(μ) = length(μ) > 1 ? FFTW.plan_r2r!(μ, FFTW.REDFT00) : ones(μ)'

"""
Compute nodes of the Clenshaw—Curtis quadrature rule.
"""
clenshawcurtisnodes(::Type{T}, N::Int) where T = T[cospi(k/(N-one(T))) for k=0:N-1]

"""
Compute weights of the Clenshaw—Curtis quadrature rule with modified Chebyshev moments of the first kind ``\\mu``.
"""
clenshawcurtisweights(μ::Vector) = clenshawcurtisweights!(copy(μ))
clenshawcurtisweights!(μ::Vector) = clenshawcurtisweights!(μ, plan_clenshawcurtis(μ))
function clenshawcurtisweights!(μ::Vector{T}, plan) where T
    N = length(μ)
    rmul!(μ, inv(N-one(T)))
    plan*μ
    μ[1] *= half(T); μ[N] *= half(T)
    return μ
end

# Chebyshev-T coefficients to values at Clenshaw-Curtis nodes

applyTN_plan(x) = length(x) > 1 ? FFTW.plan_r2r!(x, FFTW.REDFT00) : fill!(similar(x),1)'

applyTN!(x::Vector{T}) where {T<:AbstractFloat} = applyTN!(x,applyTN_plan(x))

function applyTN!(x::Vector{T},plan) where T<:AbstractFloat
    x[1] *= 2; x[end] *=2
    plan*x
    rmul!(x,half(T))
end
applyTN(x::Vector{T},plan) where {T<:AbstractFloat} = applyTN!(copy(x),plan)
applyTN(x::Vector{T}) where {T<:AbstractFloat} = applyTN!(copy(x))

# Values at Clenshaw-Curtis nodes to Chebyshev-T coefficients

applyTNinv_plan(x) = length(x) > 1 ? FFTW.plan_r2r!(x, FFTW.REDFT00) : ones(x)'

applyTNinv!(x::Vector{T}) where {T<:AbstractFloat} = applyTNinv!(x,applyTNinv_plan(x))

function applyTNinv!(x::Vector{T},plan) where T<:AbstractFloat
    plan*x
    x[1] /= 2;x[end] /= 2
    rmul!(x,inv(length(x)-one(T)))
end
applyTNinv(x::Vector{T},plan) where {T<:AbstractFloat} = applyTNinv!(copy(x),plan)
applyTNinv(x::Vector{T}) where {T<:AbstractFloat} = applyTNinv!(copy(x))

# sin(nθ) coefficients to values at Clenshaw-Curtis nodes except ±1

applyUN_plan(x) = length(x) > 0 ? FFTW.plan_r2r!(x, FFTW.RODFT00) : fill!(similar(x),1)'

applyUN!(x::AbstractVector{T}) where {T<:AbstractFloat} = applyUN!(x,applyUN_plan(x))

function applyUN!(x::AbstractVector{T},plan) where T<:AbstractFloat
    plan*x
    rmul!(x,half(T))
end
applyUN(x::AbstractVector{T},plan) where {T<:AbstractFloat} = applyUN!(copy(x),plan)
applyUN(x::AbstractVector{T}) where {T<:AbstractFloat} = applyUN!(copy(x))


# This implements src/snes/examples/tutorials/ex2.c from PETSc using the PETSc.jl package, using SNES
#
# This solves the equations sequentially
#
# Newton method to solve u'' + u^{2} = f, sequentially.

using PETSc, MPI, LinearAlgebra, SparseArrays, UnicodePlots

if ~MPI.Initialized()
    MPI.Init()
end

petsclib = PETSc.petsclibs[1]
PETSc.initialize(petsclib)

"""
    Computes initial guess
"""
function FormInitialGuess!(x)
    for i=1:length(x)
        x[i] = 0.50;
    end
end

"""
    Computes rhs forcing function
"""
function SetInitialArrays(n)
    h =  1.0/(n-1.0)
    F = zeros(n);
    xp = 0.0;
    for i=1:n
        v    = 6.0*xp + (xp+1.e-12)^6.0;
        F[i] = v;
        xp   = xp+h;
    end

    return F
end

"""
    Computes the residual f, given solution vector x
"""
function FormResidual!(cf,cx, args...)
    if typeof(cx) <: Ptr{Nothing}
        # When this routine is called from PETSc, cx is a pointer to a global vector
        # That's why we have to transfer it first to
        x   =   PETSc.unsafe_localarray(PETSc.scalartype(petsclib),cx)
    else
        x   = cx;
    end
    if typeof(cf) <: Ptr{Nothing}
        f   =   PETSc.unsafe_localarray(PETSc.scalartype(petsclib),cf)
    else
        f   = cf;
    end
    n       =   length(x);
    xp      =   LinRange(0.0,1.0, n);
    F       =   6.0.*xp .+ (xp .+1.e-12).^6.0;      # define source term function

    dx      =   1.0/(n-1.0);
    f[1]    =   x[1] - 0.0;
    for i=2:n-1
        f[i] = (x[i-1] - 2.0*x[i] + x[i+1])/dx^2 + x[i]*x[i] - F[i]
    end
    f[n]    =   x[n] - 1.0;
    Base.finalize(x)
    Base.finalize(f)

end

"""
    Computes the jacobian, given solution vector x
"""
function FormJacobian!(cx, args...)

    if typeof(cx) <: Ptr{Nothing}
        x   =   PETSc.unsafe_localarray(PETSc.scalartype(petsclib),cx)
    else
        x   =   cx;
    end

    J   =   args[1];        # preconditioner = args[2], in case we want it to be different from J
    n   =   length(x);
    dx  =   1.0/(n-1.0);

    # interior points (hand-coded jacobian)
    for i=2:n-1
        J[i,i-1] = 1.0/dx^2;
        J[i,i  ] = -2.0/dx^2 + 2.0*x[i];
        J[i,i+1] = 1.0/dx^2;
    end

    # boundary points
    J[1,1] = 1.0;
    J[n,n] = 1.0;

    if typeof(J) <: PETSc.AbstractMat
        PETSc.assemble(J);  # finalize assembly

    end

    Base.finalize(x)
end

# ==========================================
# Main code

# Compute initial solution
n   =   21;
F   =   SetInitialArrays(n);
x   =   zeros(n);

FormInitialGuess!(x);

# Compute initial jacobian using a julia structure to obtain the nonzero structure
# Note that we can also obtain this structure in a different manner
Jstruct  = zeros(n,n);
FormJacobian!(x, Jstruct);                              # jacobian in julia form
Jsp      =   sparse(Float64.(abs.(Jstruct) .> 0))       # sparse julia, with 1.0 in nonzero spots
PJ       =   PETSc.MatSeqAIJ(Jsp);                      # transfer to

# Setup snes
x_s = PETSc.VecSeq(x);                  # solution vector
res = PETSc.VecSeq(F);     # residual vector

S = PETSc.SNES{Float64}(PETSc.petsclibs[1],MPI.COMM_WORLD;
        snes_rtol=1e-12,
        snes_monitor=false,
        snes_converged_reason=false);
PETSc.setfunction!(S, FormResidual!, res)
PETSc.setjacobian!(S, FormJacobian!, PJ, PJ)

# solve
PETSc.solve!(x_s, S);

# Extract & plot solution
x_sol = x_s.array;                  # convert solution to julia format
FormResidual!(res.array,x_sol)      # just for checking, compute residual

@show norm(res.array)

PETSc.finalize(petsclib)

# plot solution in REPL
lineplot(LinRange(0,1,n),x_sol,xlabel="width",ylabel="solution")


export angular_distance, along_line_angular_distance, angular_length
# export angular_distance, distance, along_line_angular_distance, angular_length
# """
#     Radius Earth [m]
#
# Source: en.wikipedia.org/wiki/Earth_radius
# """
# const Rₑ_m = 6371008.7714

"""
    angular_distance(point₁::Point, point₂::Point)

Return the `angular_distance` [deg] of the great circle line between the
positions `point₁` [deg] and `point₂` [deg] on a unit sphere.

Source: www.movable-type.co.uk/scripts/latlong.html
"""
function angular_distance(point₁::Point, point₂::Point)
    Δpos = point₂ - point₁
    a = sind(Δpos.ϕ/2.0)^2 + cosd(point₁.ϕ)*cosd(point₂.ϕ)*sind(Δpos.λ/2.0)^2
    c = 2.0 * atand(√a, √(1.0 - a))
    return c
end

"""
    angular_length(arc::Arc)

Return the `angular_length` [deg] of the great circle line section on a
unit sphere.
"""
angular_length(arc::Arc) =
angular_distance(arc.point₁, arc.point₂)

"""
    angular_length(points::Vector{Points{Float64}})

Return the `angular_length` [deg] along the subsequent points on a unit sphere.
"""
function angular_length(points::Vector{Point{Float64}})
    point₁ = points[1]
    dist_deg = 0.0
    for point₂ in points[2:end]
        dist_deg += angular_distance(point₁, point₂)
        point₁ = point₂
    end
    return dist_deg
end

"""
    angular_length(arcs::Arcs)

Return the `angular_length` [deg] along the arcs on a unit sphere.
"""
angular_length(arcs::Arcs) = angular_length(arcs.points)

"""
    angular_length(polygon::Polygon)

Return the `angular_length` [deg] (e.g. circumference) along the polygon on a unit
sphere.
"""
angular_length(polygon::Polygon) = angular_length(polygon.points)

# """
#     distance(angular_distance::Float64, radius::Float64=Rₑ_m)
#
# Convert the `angular_distance` [deg] into a distance on a sphere of given radius
#  [m]. The default radius is the radius of the Earth.
# """
# function distance(angular_distance::Float64, radius::Float64=Rₑ_m)
#     return deg2rad(angular_distance) * radius
# end

"""
    angular_distance(point₃::Point, point₁::Point, azimuth₁₂::Float64)

Return the `angular_distance` [deg] from point₃ [deg] to the closest point on a
great circle line starting in `point₁` [deg] with `azimuth₁₂` [deg] from
`point₁` to `point₂`. It is assumed that the great circle line does not stop in
`point₂`, but continuous around the unit sphere. A positive value indicates
being right of the line, and negative being left of the line.

Source: www.movable-type.co.uk/scripts/latlong.html
"""
function angular_distance(point₃::Point, point₁::Point, azimuth₁₂::Float64)
    angular_distance₁₃ = angular_distance(point₁, point₃)
    azimuth₁₃ = azimuth(point₁, point₃)
    return angular_distance(angular_distance₁₃, azimuth₁₃, azimuth₁₂)
end

"""
    angular_distance(point₃::Point, line₁::Line)

Return the `angular_distance` [deg] from point₃ [deg] to the closest point on a
great circle line starting in `point₁` [deg] with `azimuth₁₂` [deg]. A positive
value indicates being right of the line, and negative being left of the line.

Source: www.movable-type.co.uk/scripts/latlong.html
"""
angular_distance(point₃::Point, line₁::Line) = angular_distance(point₃,
line₁.point, line₁.azimuth)

"""
    angular_distance(angular_distance₁₃::Float64, azimuth₁₃::Float64,
    azimuth₁₂::Float64)

Return the `angular_distance` [deg] from point₃ [deg] to the closest point on a
great circle line starting in `point₁` [deg] with `azimuth₁₂` [deg]. For the
calculation we need the `angular_distance₁₃` and `azimuth₁₃` between `point₁`
and `point₃` [deg], and the `azimuth₁₂` from `point₁` to `point₂` along the
great circle line. It is assumed that the great circle line does not stop in
`point₂`, but continuous around the unit sphere.

Source: www.movable-type.co.uk/scripts/latlong.html
"""
function angular_distance(angular_distance₁₃::Float64, azimuth₁₃::Float64,
azimuth₁₂::Float64)
    return asind(sind(angular_distance₁₃) * sind(azimuth₁₃ - azimuth₁₂))
end

"""
    angular_distance(point₃::Point, point₁::Point, point₂::Point)

Return the `angular_distance` [deg] from point₃ [deg] to the closest point on a
great circle line section starting in `point₁` [deg] and ending in `point₂` [deg].
The great circle line section does not continue around the unit sphere.

The `angular_distance` does not change sign when being left or right of the arc.
"""
function angular_distance(point₃::Point, point₁::Point, point₂::Point, tolerance::Float64=tolerance_deg)
    angular_distance₁₂ = angular_distance(point₁, point₂)
    angular_distance₁₃ = angular_distance(point₁, point₃)
    angular_distance₂₃ = angular_distance(point₂, point₃)
    azimuth₁₂ = azimuth(point₁, point₂)
    angular_distance_line3 = angular_distance(point₃, point₁, azimuth₁₂)
    # along_line_angular_dist = along_line_angular_distance(angular_distance₁₃,
    # angular_distance_line3)
    along_line_pnt_3a = along_line_point(point₃, Line(point₁, azimuth₁₂))
    if abs(angular_distance(along_line_pnt_3a, point₁) +
        angular_distance(along_line_pnt_3a, point₂) -
        angular_distance₁₂) < tolerance
        return abs(angular_distance_line3)
    else
        return min(angular_distance₁₃, angular_distance₂₃)
    end
end

"""
    angular_distance(point₃::Point, points::Vector{Points{T}}) where T<:Float64

Return the `angular_distance` [deg] from point₃ [deg] to the closest point on a
an set of points representing an set of arcs.

The `angular_distance` does not change sign when being left or right of the arc.
"""
function angular_distance(point₃::Point, points::Vector{Point{T}}) where
    T<:Float64
    point₁ = points[1]
    point₂ = points[2]
    min_dist = angular_distance(point₃, point₁, point₂)
    point₁ = point₂
    for point₂ in points[2:end]
        min_dist = min(min_dist, angular_distance(point₃, point₁, point₂))
        point₁ = point₂
    end
    return min_dist
end

"""
    angular_distance(point₃::Point, arcs::Arcs)

Return the `angular_distance` [deg] from point₃ [deg] to the closest point on a
an `arcs`.

The `angular_distance` does not change sign when being left or right of the arc.
"""
angular_distance(point₃::Point, arcs::Arcs) = angular_distance(point₃, arcs.points)

"""
    angular_distance(point₃::Point, polygon::Polygon)

Return the `angular_distance` [deg] from point₃ [deg] to the closest point on
the border of a `polygon`.

The `angular_distance` does not change sign when being left or right of the arc.
"""
angular_distance(point₃::Point, polygon::Polygon) = angular_distance(point₃,
polygon.points)

"""
    angular_distance(point₃::Point, arc₁₂::Arc)

Return the `angular_distance` [deg] from point₃ [deg] to the closest point on a
great circle line section starting in `point₁` [deg] and ending in `point₂` [deg].
The great circle line section does not continue around the unit sphere.
"""
angular_distance(point₃::Point, arc₁₂::Arc) = angular_distance(point₃, arc₁₂.point₁, arc₁₂.point₂)

"""
    along_line_angular_distance(angular_distance₁₃::Float64,
    angular_distance_line3::Float64)

The `along_line_angular_distance` from the start `point₁` [deg] to the closest
point on the great circle line to `point₃`. As input we need the
`angular_distance₁₃` [deg] between `point₁` and `point₃` [deg] and the closest
distance `angular_distance_line3` [deg] between the great circle line and
`point3` [deg]. The distance is non-directional.

Source: www.movable-type.co.uk/scripts/latlong.html
"""
function along_line_angular_distance(angular_distance₁₃::Float64,
angular_distance_line3::Float64)
    return acosd(cosd(angular_distance₁₃) / cosd(angular_distance_line3))
end

"""
    along_line_angular_distance(point₃::Point, line₁::Line)

The `along_line_angular_distance` from the start `point₁` [deg] to the closest
point on the great circle line `line₁` to `point₃`. The distance is
non-directional.

Source: www.movable-type.co.uk/scripts/latlong.html
"""
function along_line_angular_distance(point₃::Point, line₁::Line)
    angular_distance₁₃ = angular_distance(line₁.point, point₃)
    angular_distance_line3 = angular_distance(point₃::Point, line₁::Line)
    return along_line_angular_distance(angular_distance₁₃,
    angular_distance_line3)
end


export Spindle

"""
`Spindle()` returns the Moser spindle graph. This is a seven-vertex
unit distance graph with chromatic number equal to 4.
"""
function Spindle()
    G = IntGraph(7)
    edges = [
        1 2
        1 3
        2 3
        2 4
        3 4
        1 5
        1 6
        5 6
        5 7
        6 7
        4 7 ]
    add_edges!(G,edges)

    d = Dict{Int,Vector}()
    a = sqrt(3)/2

    pts = [ 0 1/2 -1/2 0 ; 0 a a 2a ]

    theta = acos(5/6)/2
    R = [ cos(theta) -sin(theta); sin(theta) cos(theta) ]

    p1 = R*pts
    for k=1:4
        d[k] = p1[:,k]
    end

    p2 = R'*pts
    for k=5:7
        d[k] = p2[:,k-3]
    end
    embed(G,d)
    SimpleGraphs.name(G,"Moser Spindle")
    return G
end


"""
    isinside(chart, point) -> Bool

Returns true is the given point is in the image of the given chart, false otherwise.
"""
function isinside(cell, point)

  u = carttobary(cell, point)
  T = eltype(u)

  tol = eps(T) * 1e3
  w = one(T)
  for i in 1:dimension(cell)
    0+tol < u[i] < 1-tol || return false
    w -= u[i]
  end

  0+tol < w < 1-tol || return false
  return true

end

"""
	isinclosure(simplex, point) -> Bool

Determine whether point is in the closure of simplex. False positives are possible for points just outside of the simplex.
"""
function isinclosure(cell, point)

  u = carttobary(cell, point)
  T = eltype(u)
  dim = dimension(cell)
  udim = length(point)

  tol = eps(T) * 1e3
  w = one(T)
  for i in 1:dim
    0-tol < u[i] < 1+tol || return false
    w -= u[i]
  end

  0-tol < w < 1+tol || return false

  # finally, check that the point indeed is in the plane of the cell
  A = zeros(T, udim, dim+1)
  origin = cell[dim+1]
  for i in 1:dim
      for j in 1:udim
          A[j,i] = cell[i][j] - origin[j]
      end
  end
  for j in 1:udim; A[j,end] = point[j] - origin[j]; end
  rank(A) == dim+1 && return false

  return true

end


const GraphMatrix = Union{Matrix, BitMatrix}

function clique_size(seq::Union{Vector{Bool}, BitArray}, A::GraphMatrix)
  set = find(seq)
  clique = true
  for i in set, j in set
    if (i != j) && !(A[i,j] > 0)
      clique = false
      break
    end
  end
  return clique*length(set)
end

function clique_loss(seq::Union{Vector{Bool}, BitArray}, A::GraphMatrix, kappa::Float)
  set = find(seq)
  clique = true
  edges = Float(0)
  n = length(set)
  for i in set, j in set
    if i != j
      aij = (A[i,j] > 0)
      edges += aij
      clique = clique && aij
    end
  end

  return (Float(1) - Float(edges)/Float(max(n*(n-1+kappa), 1)), Float(-n*clique))
end

function clique_loss(seq::Vector{Int}, A::GraphMatrix, kappa::Float)
  n = size(A, 1)
  if length(seq) < n
    idx = falses(n)
    idx[seq] = true
  else
    idx = seq .> 1
  end
  return clique_loss(idx, A, kappa)
end

"""
### Clique finding using the cakewalk method
Minimize the soft-clique-loss function with cakewalk. Usage: \n
`opt_seq, opt_val = clique_cw(A, kappa)` \n
where `kappa` is a size booster parameter beteen 0.0 and 1.0.
For details about the loss, kappa, and the reasoning behind them see our paper.\n\n

When `opt_seq[i]` is 2, it means vertex `i` is selected, and 1 means it is not. \n
`opt_val` is a tuple of two floats. The first is the value being optimized, i.e.,
the soft-clique-loss, and the second is the actual loss that we are interested in,
the negative clique size.
"""
clique_cw(A::GraphMatrix, args...; kwargs...) = cakewalk(clique_loss, 2, size(A, 1), A, args...; kwargs...)

inc_clique_size(ind, A::GraphMatrix) = maximum(map(v -> (w = copy(ind); w[v] = true; clique_size(w, A)), 1:size(A,1)))

"""
### Validate that a clique is locally maximal
First, find some clique, \n
`opt_seq, opt_val = clique_cw(A, kappa)` \n
Then, validate the clique is locally maximal, \n
`is_local_max_clique(opt_seq .== 2, A)`\n
The returned value is binary.
"""
is_local_max_clique(ind, A::GraphMatrix) = (cs = clique_size(ind,A); (cs > 0) && (cs >= inc_clique_size(ind, A)))


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

"""
    signarea(A, B, C)

Compute signed area of triangle formed
by points `A`, `B` and `C`.
"""
function signarea(A::Point{2}, B::Point{2}, C::Point{2})
  ((B - A) × (C - A)) / 2
end

"""
    sideof(point, segment)

Determines on which side of the oriented `segment`
the `point` lies. Possible results are `:LEFT`,
`:RIGHT` or `:ON` the segment.
"""
function sideof(p::Point{2,T}, s::Segment{2,T}) where {T}
  a, b = vertices(s)
  area = signarea(p, a, b)
  ifelse(area > atol(T), :LEFT, ifelse(area < -atol(T), :RIGHT, :ON))
end

"""
    sideof(point, chain)

Determines on which side of the closed `chain` the
`point` lies. Possible results are `:INSIDE` or
`:OUTSIDE` the chain.
"""
function sideof(p::Point{2,T}, c::Chain{2,T}) where {T}
  w = windingnumber(p, c)
  ifelse(isapprox(w, zero(T), atol=atol(T)), :OUTSIDE, :INSIDE)
end

"""
    mahalanobis(radii, angles, convention)

Return the Mahalanobis distance corresponding to an ellipsoid
with `radii` rotated by given `angles` according to given `convention`.

- For 2D ellipses, there are two radii and one rotation angle.
- For 3D ellipsoids, there are three radii and three rotation angles.

The list of available conventions can be found with:

```julia
julia> subtypes(RotationConvention)
```
"""
function mahalanobis(radii, angles, convention)
  ndims, nangles = length(radii), length(angles)
  valid = (ndims == 3 && nangles == 3) || (ndims == 2 && nangles == 1)
  @assert valid "invalid number of radii/angles"

  # invert radii if necessary
  invert = mainaxis(convention) == :Y
  ranges = invert ? [radii[i] for i in reverse(1:ndims,1,2)] : vec(radii)

  # scaling matrix
  λ = one(eltype(ranges)) ./ ranges.^2
  Λ = Diagonal(SVector{ndims}(λ))

  # rotation matrix
  R = rotmat(angles, convention)

  # ellipsoid metric
  W = Symmetric(R'*Λ*R)
  Mahalanobis(W)
end


# This file is part of the TaylorSeries.jl Julia package, MIT license
#
# Luis Benet & David P. Sanders
# UNAM
#
# MIT Expat license
#

## Differentiating ##
"""
    derivative(a)

Return the `Taylor1` polynomial of the differential of `a::Taylor1`.
The last coefficient is set to zero.

The function `differentiate` is an exact synonym of `derivative`.
"""
function derivative(a::Taylor1)
    res = zero(a)
    @inbounds for ord in eachindex(a)
        derivative!(res, a, ord)
        # ord == a.order-1 && break
    end
    return res
end

"""
    differentiate

An exact synonym of [`derivative`](@ref).
"""
const differentiate = derivative

"""
    derivative!(res, a) --> nothing

In-place version of `derivative`. Compute the `Taylor1` polynomial of the
differential of `a::Taylor1` and save it into `res`. The last coefficient is
set to zero.
"""
function derivative!(res::Taylor1, a::Taylor1)
    @inbounds for ord in eachindex(a)
        derivative!(res, a, ord)
        # ord == a.order-1 && break
    end
    res[a.order] = zero(a[0])
    nothing
end

"""
    derivative!(p, a, k) --> nothing

Update in-place the `k-th` expansion coefficient `p[k]` of `p = derivative(a)`
for both `p` and `a` `Taylor1`.

The coefficients are given by

```math
p_k = (k+1)a_{k+1}.
```

"""
derivative!(p::Taylor1, a::Taylor1, k::Int) = k < a.order ? p[k] = (k+1)*a[k+1] : nothing

"""
    derivative(a, n)

Compute recursively the `Taylor1` polynomial of the n-th derivative of
`a::Taylor1`.
"""
function derivative(a::Taylor1{T}, n::Int) where {T<:Number}
    @assert a.order ≥ n ≥ 0
    if n==0
        return a
    else
        res = deepcopy(a)
        for i in 1:n
            derivative!(res, res)
        end
        return res
    end
end

"""
    derivative(n, a)

Return the value of the `n`-th derivative of the polynomial `a`.
"""
function derivative(n::Int, a::Taylor1{T}) where {T<:Number}
    @assert a.order ≥ n ≥ 0
    factorial( widen(n) ) * a[n] :: T
end

## Integrating ##
"""
    integrate(a, [x])

Return the integral of `a::Taylor1`. The constant of integration
(0-th order coefficient) is set to `x`, which is zero if ommitted.
"""
function integrate(a::Taylor1{T}, x::S) where {T<:Number,S<:Number}
    order = get_order(a)
    aa = a[0]/1 + zero(x)
    R = typeof(aa)
    coeffs = Array{typeof(aa)}(undef, order+1)
    fill!(coeffs, zero(aa))
    @inbounds for i = 1:order
        coeffs[i+1] = a[i-1] / i
    end
    @inbounds coeffs[1] = convert(R, x)
    return Taylor1(coeffs, a.order)
end
integrate(a::Taylor1{T}) where {T<:Number} = integrate(a, zero(a[0]))

## Differentiation ##
"""
    derivative(a, r)

Partial differentiation of `a::HomogeneousPolynomial` series with respect
to the `r`-th variable.
"""
function derivative(a::HomogeneousPolynomial, r::Int)
    @assert 1 ≤ r ≤ get_numvars()
    # @show zero(a[1]) zero(eltype(a))
    T = eltype(a)
    a.order == 0 && return HomogeneousPolynomial(zero(a[1]), 0)
    @inbounds num_coeffs = size_table[a.order]
    coeffs = zeros(T,num_coeffs)
    @inbounds posTb = pos_table[a.order]
    @inbounds num_coeffs = size_table[a.order+1]

    @inbounds for i = 1:num_coeffs
        iind = coeff_table[a.order+1][i]
        n = iind[r]
        n == 0 && continue
        iind[r] -= 1
        kdic = in_base(get_order(),iind)
        pos = posTb[kdic]
        coeffs[pos] = n * a[i]
        iind[r] += 1
    end

    return HomogeneousPolynomial{T}(coeffs, a.order-1)
end
derivative(a::HomogeneousPolynomial, s::Symbol) = derivative(a, lookupvar(s))

"""
    derivative(a, r)

Partial differentiation of `a::TaylorN` series with respect
to the `r`-th variable. The `r`-th variable may be also
specified through its symbol.
"""
function derivative(a::TaylorN, r=1::Int)
    T = eltype(a)
    coeffs = Array{HomogeneousPolynomial{T}}(undef, a.order)

    @inbounds for ord = 1:a.order
        coeffs[ord] = derivative( a[ord], r)
    end
    return TaylorN{T}( coeffs, a.order )
end
derivative(a::TaylorN, s::Symbol) = derivative(a, lookupvar(s))

"""
    derivative(a::TaylorN{T}, ntup::NTuple{N,Int})

Return a `TaylorN` with the partial derivative of `a` defined
by `ntup::NTuple{N,Int}`, where the first entry is the number
of derivatives with respect to the first variable, the second is
the number of derivatives with respect to the second, and so on.
"""
function derivative(a::TaylorN, ntup::NTuple{N,Int}) where {N}

    @assert N == get_numvars() && all(ntup .>= 0)

    sum(ntup) > a.order && return zero(a)
    sum(ntup) == 0 && return copy(a)

    aa = copy(a)
    for nvar in 1:get_numvars()
        for numder in 1:ntup[nvar]
            aa = derivative(aa, nvar)
        end
    end

    return aa
end

"""
    derivative(ntup::NTuple{N,Int}, a::TaylorN{T})

Returns the value of the coefficient of `a` specified by
`ntup::NTuple{N,Int}`, multiplied by the corresponding
factorials.
"""
function derivative(ntup::NTuple{N,Int}, a::TaylorN) where {N}

    @assert N == get_numvars() && all(ntup .>= 0)

    c = getcoeff(a, [ntup...])
    for ind = 1:get_numvars()
        c *= factorial(ntup[ind])
    end

    return c
end

## Gradient, jacobian and hessian
"""
```
    gradient(f)
    ∇(f)
```

Compute the gradient of the polynomial `f::TaylorN`.
"""
function gradient(f::TaylorN)
    T = eltype(f)
    numVars = get_numvars()
    grad = Array{TaylorN{T}}(undef, numVars)
    @inbounds for nv = 1:numVars
        grad[nv] = derivative(f, nv)
    end
    return grad
end
const ∇ = TaylorSeries.gradient

"""
```
    jacobian(vf)
    jacobian(vf, [vals])
```

Compute the jacobian matrix of `vf`, a vector of `TaylorN` polynomials,
evaluated at the vector `vals`. If `vals` is ommited, it is evaluated at zero.
"""
function jacobian(vf::Array{TaylorN{T},1}) where {T<:Number}
    numVars = get_numvars()
    @assert length(vf) == numVars
    jac = Array{T}(undef, numVars, numVars)

    @inbounds for comp = 1:numVars
        jac[:,comp] = vf[comp][1][1:end]
    end

    return transpose(jac)
end
function jacobian(vf::Array{TaylorN{T},1}, vals::Array{S,1}) where {T<:Number,S<:Number}
    R = promote_type(T,S)
    numVars = get_numvars()
    @assert length(vf) == numVars == length(vals)
    jac = Array{R}(undef, numVars, numVars)

    for comp = 1:numVars
        @inbounds grad = gradient( vf[comp] )
        @inbounds for nv = 1:numVars
            jac[nv,comp] = evaluate(grad[nv], vals)
        end
    end

    return transpose(jac)
end

function jacobian(vf::Array{Taylor1{TaylorN{T}},1}) where {T<:Number}
    vv = convert(Array{TaylorN{Taylor1{T}},1}, vf)
    jacobian(vv)
end

"""
```
    jacobian!(jac, vf)
    jacobian!(jac, vf, [vals])
```

Compute the jacobian matrix of `vf`, a vector of `TaylorN` polynomials
evaluated at the vector `vals`, and write results to `jac`. If `vals` is ommited,
it is evaluated at zero.
"""
function jacobian!(jac::Array{T,2}, vf::Array{TaylorN{T},1}) where {T<:Number}
    numVars = get_numvars()
    @assert length(vf) == numVars
    @assert (numVars, numVars) == size(jac)
    for comp2 = 1:numVars
        for comp1 = 1:numVars
            @inbounds jac[comp1,comp2] = vf[comp1][1][comp2]
        end
    end
    nothing
end
function jacobian!(jac::Array{T,2}, vf::Array{TaylorN{T},1},
        vals::Array{T,1}) where {T<:Number}

    numVars = get_numvars()
    @assert length(vf) == numVars == length(vals)
    @assert (numVars, numVars) == size(jac)
    for comp = 1:numVars
        @inbounds for nv = 1:numVars
            jac[nv,comp] = evaluate(derivative(vf[nv], comp), vals)
        end
    end
    nothing
end

"""
```
    hessian(f)
    hessian(f, [vals])
```

Return the hessian matrix (jacobian of the gradient) of `f::TaylorN`,
evaluated at the vector `vals`. If `vals` is ommited, it is evaluated at
zero.
"""
hessian(f::TaylorN{T}, vals::Array{S,1}) where {T<:Number,S<:Number} =
    (R = promote_type(T,S); jacobian( gradient(f), vals::Array{R,1}) )

hessian(f::TaylorN{T}) where {T<:Number} = hessian( f, zeros(T, get_numvars()) )

"""
```
    hessian!(hes, f)
    hessian!(hes, f, [vals])
```

Return the hessian matrix (jacobian of the gradient) of `f::TaylorN`,
evaluated at the vector `vals`, and write results to `hes`. If `vals` is
ommited, it is evaluated at zero.
"""
hessian!(hes::Array{T,2}, f::TaylorN{T}, vals::Array{T,1}) where {T<:Number} =
    jacobian!(hes, gradient(f), vals)

hessian!(hes::Array{T,2}, f::TaylorN{T}) where {T<:Number} =
    jacobian!(hes, gradient(f))

##Integration
"""
    integrate(a, r)

Integrate the `a::HomogeneousPolynomial` with respect to the `r`-th
variable. The returned `HomogeneousPolynomial` has no added constant of
integration. If the order of a corresponds to `get_order()`, a zero
`HomogeneousPolynomial` of 0-th order is returned.

"""
function integrate(a::HomogeneousPolynomial, r::Int)
    @assert 1 ≤ r ≤ get_numvars()

    order_max = get_order()
    a.order == order_max && return HomogeneousPolynomial(zero(a[1]/1), 0)

    @inbounds posTb = pos_table[a.order+2]
    @inbounds num_coeffs = size_table[a.order+1]

    T = promote_type(eltype(a), eltype(a[1]/1))
    coeffs = zeros(T, size_table[a.order+2])

    @inbounds for i = 1:num_coeffs
        iind = coeff_table[a.order+1][i]
        n = iind[r]
        n == order_max && continue
        iind[r] += 1
        kdic = in_base(get_order(), iind)
        pos = posTb[kdic]
        coeffs[pos] = a[i] / (n+1)
        iind[r] -= 1
    end

    return HomogeneousPolynomial(coeffs, a.order+1)
end
integrate(a::HomogeneousPolynomial, s::Symbol) = integrate(a, lookupvar(s))

"""
    integrate(a, r, [x0])

Integrate the `a::TaylorN` series with respect to the `r`-th variable,
where `x0` the integration constant and must be independent
of the `r`-th variable; if `x0` is ommitted, it is taken as zero.
"""
function integrate(a::TaylorN, r::Int)
    T = promote_type(eltype(a), eltype(a[0]/1))
    order_max = min(get_order(), a.order+1)
    coeffs = zeros(HomogeneousPolynomial{T}, order_max)

    @inbounds for ord = 0:order_max-1
        coeffs[ord+1] = integrate( a[ord], r)
    end

    return TaylorN(coeffs)
end
function integrate(a::TaylorN, r::Int, x0::TaylorN)
    # Check constant of integration is independent of re
    @assert derivative(x0, r) == 0.0 """
    The integration constant ($x0) must be independent of the
    $(_params_TaylorN_.variable_names[r]) variable"""

    res = integrate(a, r)
    return x0+res
end
integrate(a::TaylorN, r::Int, x0) =
    integrate(a,r,TaylorN(HomogeneousPolynomial([convert(eltype(a),x0)], 0)))

integrate(a::TaylorN, s::Symbol) = integrate(a, lookupvar(s))
integrate(a::TaylorN, s::Symbol, x0::TaylorN) = integrate(a, lookupvar(s), x0)
integrate(a::TaylorN, s::Symbol, x0) = integrate(a, lookupvar(s), x0)


# Load Julia packages (libraries) needed  for the snippets in chapter 0

using StatisticalRethinking
import LogDensityProblems: logdensity, logdensity_and_gradient
import StatisticalRethinking: HMC2, generate_n_samples

# CmdStan uses a tmp directory to store the output of cmdstan

ProjDir = @__DIR__
cd(ProjDir)

# ### snippet 9.3

# Construct the logdensity problem

struct clip_9_3_model{TY <: AbstractVector, TX <: AbstractVector}
    "Observations."
    y::TY
    "Covariate"
    x::TX
end

# Make the type callable with the parameters *as a single argument*.

function (problem:: clip_9_3_model)(θ)
    @unpack y, x, = problem    # extract the data
    @unpack muy, mux = θ     # works on the named tuple too
    ll = 0.0
    ll += loglikelihood(Normal(mux, 1), x)
    ll += loglikelihood(Normal(muy, 1), y)
    ll += logpdf(Normal(0, 1), mux)
    ll += logpdf(Normal(0, 1), muy)
    ll
end

# Instantiate the model with data and inits.

Random.seed!(1234591)

N = 100
x = rand(Normal(0, 1), N)
y = rand(Normal(0, 1), N)

p = clip_9_3_model(y, x)
θ = (muy = 0.0, mux=0.0)
p(θ)

# Write a function to return properly dimensioned transformation.

problem_transformation(p::clip_9_3_model) =
    as((muy = asℝ, mux = asℝ))

# Wrap the problem with a transformation, then use Flux for the gradient.

P = TransformedLogDensity(problem_transformation(p), p)
∇P = ADgradient(:ForwardDiff, P);

# Tune and sample.

chain, NUTS_tuned = NUTS_init_tune_mcmc(∇P, 1000);

# We use the transformation to obtain the posterior from the chain.

posterior = TransformVariables.transform.(Ref(problem_transformation(p)), get_position.(chain));

# Extract the posterior means,

[mean(first, posterior), mean(last, posterior)]

# Draw 200 samples:

function draw_n_samples(model, grad;
  epsilon = 0.03, # Step size
  L = 11, # No of leapfrog steps
  n_samples = 1000, # No of samples
  q = [-0.1, 0.2]) # Initial position

  samples = zeros(n_samples, 2)
  for i in 1:n_samples
    q, ptraj, qtraj, accept, dH = HMC(model, grad, 0.03, 11, q)
    samples[i, :] = q
  end

  samples

end

samples = draw_n_samples(p, ∇P; n_samples=200);
mean(samples, dims=1)

# End of `09/clip-03.jl`


using ITensors
using Zygote

s = siteind("Qubit")

f(x) = op("Ry", s; θ=x)[1, 1]

x = 0.2
@show f(x), cos(x / 2)
@show f'(x), -sin(x / 2) / 2

# Simple gate optimization
ψ0 = state(s, "0")
ψp = state(s, "+")

function loss(x)
  U = op("Ry", s; θ=x)
  Uψ0 = replaceprime(U * ψ0, 1 => 0)
  return -(dag(ψp) * Uψ0)[]
end

# Extremely simple gradient descent implementation,
# where gradients are computing with automatic differentiation
# using Zygote.
function gradient_descent(f, x0; γ, nsteps, grad_tol)
  @show γ, nsteps
  x = x0
  f_x = f(x)
  ∇f_x = f'(x)
  step = 0
  @show step, x, f_x, ∇f_x
  for step in 1:nsteps
    x -= γ * ∇f_x
    f_x = f(x)
    ∇f_x = f'(x)
    @show step, x, f_x, ∇f_x
    if norm(∇f_x) ≤ grad_tol
      break
    end
  end
  return x, f_x, ∇f_x
end

x0 = 0
γ = 2.0 # Learning rate
nsteps = 30 # Number of steps of gradient descent
grad_tol = 1e-4 # Stop if gradient falls below this value
x, loss_x, ∇loss_x = gradient_descent(loss, x0; γ=γ, nsteps=nsteps, grad_tol=grad_tol)

@show x0, loss(x0)
@show x, loss(x)
@show π / 2, loss(π / 2)


using LinearAlgebra

"""
Performs a random hill climb solution for a discretized search space.

    Args:
      fn (Function): Objective function
      x0 (Array{Any}): Starting point
      h (Float64, optional): Discretization
      max_failed (Int64, optional): Max number of failed moves before exiting.
        Defaults to the dimensionality of the design space.
      verbose (Bool, optional): Optionally print results

    Returns:
      tuple: (x, fx, f_calls) where x is the minimizer, fx is the minimum, and
      f_calls is the number of function calls made.
"""
function random_hill(fn::Function, x0::Array{Float64};
                     h=.1, max_failed=nothing, verbose=false)
    x = copy(x0)
    fx = fn(x)
    n = length(x)
    failed = 0
    f_calls = 1

    # use an I matrix to randomly determine direction
    e_hat = Matrix(1I, n, n)

    if isnothing(max_failed)
        max_failed = n
    end

    while failed < max_failed
        d = e_hat[:, rand(1:n)] # randomly choose basis vec
        d *= rand([-1, 1]) # randomly choose pos/neg

        xnew = x + h*d
        fnew = fn(xnew)
        f_calls += 1

        if fnew < fx
            fx = fnew
            x = xnew
            failed = 0
        else
            failed += 1
        end
    end

    if verbose
        println("Iterations: $(f_calls)\nx*: $(x)\nf*: $(fx)\n")
    end

    return x, fx, f_calls
end


# Functions for STaylor1
@generated function exp(a::STaylor1{N,T}) where {N, T <: Number}
    ex_calc = quote end
    append!(ex_calc.args, Any[nothing for i in 1:N])
    syms = Symbol[Symbol("c$i") for i in 1:N]

    sym = syms[1]
    ex_line = :($(syms[1]) = exp(a[0]))
    ex_calc.args[1] = ex_line

    for k in 1:(N-1)
        kT = convert(T,k)
        sym = syms[k+1]
        ex_line = :($kT * a[$k] * $(syms[1]))
        @inbounds for i = 1:k-1
            ex_line = :($ex_line + $(kT-i) * a[$(k-i)] * $(syms[i+1]))
        end
        ex_line = :(($ex_line)/$kT)
        ex_line = :($sym = $ex_line)
        ex_calc.args[k+1] = ex_line
    end

    exout = :(($(syms[1]),))
    for i = 2:N
        push!(exout.args, syms[i])
    end
    return quote
               Base.@_inline_meta
               $ex_calc
               return STaylor1{N,T}($exout)
            end
end

@generated function log(a::STaylor1{N,T}) where {N, T <: Number}
    ex_calc = quote end
    append!(ex_calc.args, Any[nothing for i in 1:N])
    syms = Symbol[Symbol("c$i") for i in 1:N]

    (N >= 1) && (ex_calc.args[1] = :($(syms[1]) = log(constant_term(a))))
    (N >= 2) && (ex_calc.args[2] = :($(syms[2]) = a[1]/constant_term(a)))

    for k in 2:(N-1)
        ex_line = :($(k-1)*a[1]*$(syms[k]))
        @inbounds for i = 2:k-1
            ex_line = :($ex_line + $(k-i)*a[$i] * $(syms[k+1-i]))
        end
        ex_line = :((a[$k] - ($ex_line)/$(convert(T,k)))/constant_term(a))
        ex_line = :($(syms[k+1]) = $ex_line)
        ex_calc.args[k+1] = ex_line
    end

    exout = :(($(syms[1]),))
    for i = 2:N
        push!(exout.args, syms[i])
    end

    return quote
               Base.@_inline_meta
               iszero(constant_term(a)) && throw(ArgumentError("""
                       The 0-th order `STaylor1` coefficient must be non-zero
                       in order to expand `log` around 0.
                       """))
               $ex_calc
               return STaylor1{N,T}($exout)
            end
end

sin(a::STaylor1{N,T}) where {N, T <: Number} = sincos(a)[1]
cos(a::STaylor1{N,T}) where {N, T <: Number} = sincos(a)[2]
@generated function sincos(a::STaylor1{N,T}) where {N, T <: Number}

    ex_calc = quote end
    append!(ex_calc.args, Any[nothing for i in 1:(2*N)])

    syms_s = Symbol[Symbol("c$i") for i in 1:N]
    syms_c = Symbol[Symbol("c2$i") for i in 1:N]

    ex_line_s = :($(syms_s[1]) = sin(a[0]))
    ex_line_c = :($(syms_c[1]) = cos(a[0]))
    ex_calc.args[1] = ex_line_s
    ex_calc.args[2] = ex_line_c

    for k = 1:(N - 1)
        ex_line_s = :(a[1]*$(syms_c[k]))
        ex_line_c = :(-a[1]*$(syms_s[k]))
        for i = 2:k
            ex_line_s = :($ex_line_s + $i*a[$i]*$(syms_c[(k - i + 1)]))
            ex_line_c = :($ex_line_c - $i*a[$i]*$(syms_s[(k - i + 1)]))
        end
        ex_line_s = :($(syms_s[k + 1]) = ($ex_line_s)/$k)
        ex_line_c = :($(syms_c[k + 1]) = ($ex_line_c)/$k)
        ex_calc.args[2*k + 1] = ex_line_s
        ex_calc.args[2*k + 2] = ex_line_c
    end

    exout_s = :(($(syms_s[1]),))
    for i = 2:N
        push!(exout_s.args, syms_s[i])
    end

    exout_c = :(($(syms_c[1]),))
    for i = 2:N
        push!(exout_c.args, syms_c[i])
    end

    return quote
               Base.@_inline_meta
               $ex_calc
               return STaylor1{N,T}($exout_s), STaylor1{N,T}($exout_c)
            end
end

# Functions for STaylor1
#=
@generated function tan(a::STaylor1{N,T}) where {N, T <: Number}

    ex_calc = quote end
    append!(ex_calc.args, Any[nothing for i in 1:(4*N)])
    syms = Symbol[Symbol("c$i") for i in 1:N]
    syms2 = Symbol[Symbol("c2$i") for i in 1:N]

    for i = 1:N
        ex_calc.args[i] = :($(syms2[i]) = 0.0)
    end

    ex_line_c = :($(syms[1]) = tan(a[0]))
    ex_line_c2 = :($(syms2[1]) = ($(syms[1]))^2)
    ex_calc.args[N + 1] = ex_line_c
    ex_calc.args[N + 2] = ex_line_c2

    for k = 1:(N - 1)

        kodd = k%2
        kend = div(k - 2 + kodd, 2)
        kdiv2 = div(k, 2)

        ex_line_c = :($(k-1)*a[$(k-1)]*$(syms2[2]))
        for i = 1:(k - 1)
            q = k - i
            ex_line_c = :($ex_line_c + ($q)*a[$q]*$(syms2[i + 1]))
        end
        ex_line_c = :(a[$k] + ($ex_line_c)/$k)
        ex_line_c = :($(syms[k + 1]) = $ex_line_c)
        ex_calc.args[(3*k-2) + N + 2] = ex_line_c

        ex_line_c2 = :(a[0]*a[$k])
        for i = 1:kend
            ex_line_c2 = :($ex_line_c2 + a[$i]*a[$k - $i])
        end
        ex_line_c2 = :(2*$ex_line_c2)

        if kodd != 1
            ex_line_c2 = :($ex_line_c2 + a[$kdiv2]^2)
        end

        ex_line_c2 = :($(syms2[k + 1]) = $ex_line_c2)
        ex_calc.args[(3*k-1) + 2 + N] = ex_line_c2
        blar = k + 1
        blar_str = "c$(blar) "
        ex_calc.args[3*k + N + 2] = :(println($blar_str*string($(syms[k + 1]))))
    end

    exout = :(($(syms[1]),))
    for i = 2:N
        push!(exout.args, syms[i])
    end

    return quote
               Base.@_inline_meta
               $ex_calc
               return STaylor1{N,T}($exout)
            end
end
=#

# Functions for STaylor1
#=
@generated function asin(a::STaylor1{N,T}) where {N, T <: Number}
    ex_calc = quote end
    append!(ex_calc.args, Any[nothing for i in 1:N])
    syms = Symbol[Symbol("c$i") for i in 1:N]

    sym = syms[1]
    ex_line = :($(sym) = asin(a[0]))
    ex_calc.args[1] = ex_line

    #=
    for k in 1:(N-1)
        kT = convert(T,k)
        sym = syms[k+1]
        ex_line = :($kT * a[$k] * $(syms[1]))
        @inbounds for i = 1:k-1
            ex_line = :($ex_line + $(kT-i) * a[$(k-i)] * $(syms[i+1]))
        end
        ex_line = :(($ex_line)/$kT)
        ex_line = :($sym = $ex_line)
        ex_calc.args[k+1] = ex_line
    end
    =#

    exout = :(($(syms[1]),))
    for i = 2:N
        push!(exout.args, syms[i])
    end
    return quote
               Base.@_inline_meta
               $ex_calc
               return STaylor1{N,T}($exout)
            end
end
=#

# Functions for STaylor1
#=
@generated function acos(a::STaylor1{N,T}) where {N, T <: Number}
    ex_calc = quote end
    append!(ex_calc.args, Any[nothing for i in 1:N])
    syms = Symbol[Symbol("c$i") for i in 1:N]

    sym = syms[1]
    ex_line = :($(sym) = acos(a[0]))
    ex_calc.args[1] = ex_line

    #=
    for k in 1:(N-1)
        kT = convert(T,k)
        sym = syms[k+1]
        ex_line = :($kT * a[$k] * $(syms[1]))
        @inbounds for i = 1:k-1
            ex_line = :($ex_line + $(kT-i) * a[$(k-i)] * $(syms[i+1]))
        end
        ex_line = :(($ex_line)/$kT)
        ex_line = :($sym = $ex_line)
        ex_calc.args[k+1] = ex_line
    end
    =#

    exout = :(($(syms[1]),))
    for i = 2:N
        push!(exout.args, syms[i])
    end
    return quote
               Base.@_inline_meta
               $ex_calc
               return STaylor1{N,T}($exout)
            end
end
=#

# Functions for STaylor1
#=
@generated function atan(a::STaylor1{N,T}) where {N, T <: Number}
    ex_calc = quote end
    append!(ex_calc.args, Any[nothing for i in 1:N])
    syms = Symbol[Symbol("c$i") for i in 1:N]

    sym = syms[1]
    ex_line = :($(sym) = atan(a[0]))
    ex_calc.args[1] = ex_line

    #=
    for k in 1:(N-1)
        kT = convert(T,k)
        sym = syms[k+1]
        ex_line = :($kT * a[$k] * $(syms[1]))
        @inbounds for i = 1:k-1
            ex_line = :($ex_line + $(kT-i) * a[$(k-i)] * $(syms[i+1]))
        end
        ex_line = :(($ex_line)/$kT)
        ex_line = :($sym = $ex_line)
        ex_calc.args[k+1] = ex_line
    end
    =#

    exout = :(($(syms[1]),))
    for i = 2:N
        push!(exout.args, syms[i])
    end
    return quote
               Base.@_inline_meta
               $ex_calc
               return STaylor1{N,T}($exout)
            end
end
=#

sinh(a::STaylor1{N,T}) where {N, T <: Number} = sinhcosh(a)[1]
cosh(a::STaylor1{N,T}) where {N, T <: Number} = sinhcosh(a)[2]
@generated function sinhcosh(a::STaylor1{N,T}) where {N, T <: Number}

    ex_calc = quote end
    append!(ex_calc.args, Any[nothing for i in 1:(2*N)])

    syms_s = Symbol[Symbol("c$i") for i in 1:N]
    syms_c = Symbol[Symbol("c2$i") for i in 1:N]

    ex_line_s = :($(syms_s[1]) = sinh(a[0]))
    ex_line_c = :($(syms_c[1]) = cosh(a[0]))
    ex_calc.args[1] = ex_line_s
    ex_calc.args[2] = ex_line_c

    for k = 1:(N - 1)
        ex_line_s = :(a[1]*$(syms_c[k]))
        ex_line_c = :(a[1]*$(syms_s[k]))
        for i = 2:k
            ex_line_s = :($ex_line_s + $i*a[$i]*$(syms_c[(k - i + 1)]))
            ex_line_c = :($ex_line_c + $i*a[$i]*$(syms_s[(k - i + 1)]))
        end
        ex_line_s = :($(syms_s[k + 1]) = ($ex_line_s)/$k)
        ex_line_c = :($(syms_c[k + 1]) = ($ex_line_c)/$k)
        ex_calc.args[2*k + 1] = ex_line_s
        ex_calc.args[2*k + 2] = ex_line_c
    end

    exout_s = :(($(syms_s[1]),))
    for i = 2:N
        push!(exout_s.args, syms_s[i])
    end

    exout_c = :(($(syms_c[1]),))
    for i = 2:N
        push!(exout_c.args, syms_c[i])
    end

    return quote
               Base.@_inline_meta
               $ex_calc
               return STaylor1{N,T}($exout_s), STaylor1{N,T}($exout_c)
            end
end

# Functions for STaylor1
@generated function tanh(a::STaylor1{N,T}) where {N, T <: Number}

    ex_calc = quote end
    append!(ex_calc.args, Any[nothing for i in 1:(2*N)])
    syms = Symbol[Symbol("c$i") for i in 1:N]
    syms2 = Symbol[Symbol("c2$i") for i in 1:N]

    ex_line_c = :($(syms[1]) = tanh(a[0]))
    ex_line_c2 = :($(syms2[1]) = ($(syms[1]))^2)
    ex_calc.args[1] = ex_line_c
    ex_calc.args[2] = ex_line_c2

    for k = 1:(N - 1)

        kodd = k%2
        kend = div(k - 2 + kodd, 2)
        kdiv2 = div(k, 2)

        ex_line_c = :($k*a[$k]*$(syms2[1]))
        for i = 1:(k - 1)
            q = k - i
            ex_line_c = :($ex_line_c + ($q)*a[$q]*$(syms2[i+1]))
        end
        ex_line_c = :(a[$k] - $ex_line_c/$k)
        ex_line_c = :($(syms[k + 1]) = $ex_line_c)
        ex_calc.args[2*k + 1] = ex_line_c

        ex_line_c2 = :(a[1]*a[$k])
        @inbounds for i = 1:kend
            ex_line_c2 = :($ex_line_c2 + a[$i + 1]*a[$k - $i])
        end
        ex_line_c2 = :(2*$ex_line_c2)

        if kodd != 1
            ex_line_c2 = :($ex_line_c2 + a[$kdiv2 + 1]^2)
        end

        ex_line_c2 = :($(syms2[k + 1]) = $ex_line_c2)
        ex_calc.args[2*k + 2] = ex_line_c2
    end

    exout = :(($(syms[1]),))
    for i = 2:N
        push!(exout.args, syms[i])
    end

    return quote
               Base.@_inline_meta
               $ex_calc
               return STaylor1{N,T}($exout)
            end
end


# p24.jl - pseudospectra of Davies's complex harmonic oscillator
#         (For finer, slower plot, change 0:2 to 0:.5.)

# Eigenvalues:
N = 70; (D,x) = cheb(N); x = x[2:N];
L = 6; x = L*x; D = D/L;                   # rescale to [-L,L]
A = -D^2; A = A[2:N,2:N] + (1+3im)*diagm(x.^2);
lambda = eigvals(A);
clf(); plot(real(lambda),imag(lambda),"k.",markersize=6);
axis([0,50,0,40])

# Pseudospectra:
x = 0:1:50; y = 0:1:40;  zz = x' .+ 1im*y;
I = eye(N-1);
minsvd(z) = minimum(svdvals(z*I-A));
sigmin = [ minsvd(x[i]+1im*y[j]) for i=eachindex(x), j=eachindex(y) ];
contour(x,y,sigmin',levels=10.0.^(-4:.5:-.5));


function initial_velocity(N::Int64;T=298.15,R=8.3145,m=1.0)
  vel = [ [0.,0.] for i in 1:N]

  for i in 1:N
    angle = 2*pi*rand()
    px = rand()
    py = rand()
    vel[i][1] = sqrt(-(2*R*T/m)*log(1-px)) * angle       # velocity in the x axis
    vel[i][2] = sqrt(-(2*R*T/m)*log(1-py)) * angle       # velocity in the y axis
  end

  return vel
end

function norm(x::Vector{Float64})
  s = 0.
  for i in 1:length(x)
    s = s + x[i]^2
  end
  sqrt(s)
end

function normVEL(x)
  s = zeros(length(x))
  for i in 1:length(x)
    s[i] = norm(x[i])
  end
  return s
end

export initial_velocity, norm, normVEL

# p(v)dv
p(v;m=1.,R=8.3145,T=298.15) = m/(R*T) * v * exp(-m*v^2/(2R*T))

# probability v < x
pv(x;m = 1., T = 298.15, R = 8.3145) = -exp(-m*x^2/(2*R*T)) + 1


module BVP1D

using LinearAlgebra
using OffsetArrays

export solve_bvp

function solve_bvp(L, f, γ0, γL, P)
    Δx = L / P
    c = 1 / Δx^2
    dv = fill(2c, P-1)
    ev = fill(-c, P-1)
    A = SymTridiagonal(dv, ev)
    x_ = range(0; length=P+1, stop=L)
    x = OffsetArray(x_, 0:P)
    b = f.(x[1:P-1])
    b[1]   += c * γ0
    b[P-1] += c * γL
    U = OffsetArray{Float64}(undef, 0:P)
    U[0] = γ0
    U[1:P-1] = A \ b
    U[P] = γL
    return U, x
end

end # module BVP1D


"""
    B_SplineFFT{T, order} <: AbstractInterpolation{T, CircEdge, order}
    B_SplineFFT( order::Int, n::Int, T::DataType=Float64)

Type containing spline coefficients for b-spline interpolation based on fft, using the fact that b-spline matrix is a circulant matrix

# Type parameters
- `T` : the type of data that is interpolate
- `order::Int`: order of lagrange interpolation

# Implementation :
- `c_fft::Vector{Complex{T}}` : fft transform of coefficients
- `parfft::PrepareFftBig` : fft precomputed data
- `tabfct::Vector{Polynomial{T}}` : function table for interpolation

# Arguments :
- `n` : size of the matrix
- `order` : the order of interpolation
- `[T::DataType=Float64]` : The type values to interpolate

"""
struct B_SplineFFT{T,order} <: B_Spline{T,CircEdge,order}
    c_fft::Vector{Complex{T}}
    parfft::PrepareFftBig
    tabfct::Vector{Polynomial{T}}
    function B_SplineFFT(order::Int, n::Int, T::DataType = Float64)
        bspline = getbspline(order, 0)
        tabfct_rat = map(x -> bspline[order-x](Polynomial([order - x, 1])), 0:order)
        kl, ku = get_kl_ku(order)
        c = zeros(T, n)
        parfft = PrepareFftBig(n, T)
        tab_coef = convert.(T, bspline.(1:order))
        dec = n - kl - 1
        for i = 1:order
            c[(dec+i)%n+1] = tab_coef[i]
        end
        c_fft = fftgen(parfft, c)
        return new{T,order}(c_fft, parfft, convert.(Polynomial{T}, tabfct_rat))
    end
    B_SplineFFT(o::Int, n::Int, _::T) where {T} = B_SplineFFT(o, n, T)
end
function sol(bsp::B_SplineFFT{T}, b::AbstractVector{T}) where {T}
    return real(ifftgen(bsp.parfft, fftgen(bsp.parfft, b) ./ bsp.c_fft))
end
function sol!(Y::AbstractVector{T}, bsp::B_SplineFFT{T}, b::AbstractVector{T}) where {T}
    Y .= real(ifftgen(bsp.parfft, fftgen(bsp.parfft, b) ./ bsp.c_fft))
end


"""
`poro_elasticity_model(;
    n = 980,
    rho = 1e-3,
    alpha = 0.79,
    M = 1/7.80e3,
    kappanu = 633.33,
    eta = 1e-4,
    force_download = false
  )`

This function returns a port-Hamiltonian model of linear poroelasticity in a
bounded Lipschitz domain as described in
    Altmann, Mehrmann, Unger: Port-Hamiltonian Formulations of Poroelastic
    Network Models
# Arguments
- `n`: System dimension (can only be either: 320, 980, or 1805). Default = 980.
- `rho`: density. Default = `1e-3`.
- `alpha`: Biot-Willis fluid-solid coupling coefficient. Default = 0.79.
- `bm`: Biot-Modulus. Default = `1/7.8e3`.
- `kappanu`: Quotient kappa/Nu, where kappa denotes the permeability and nu denotes the fluid viscosity. Default = 633.33.
- `eta`: artificial damping coefficient. Default = `1e-4`.
"""
function poro_elasticity_model(;
    n = 980,
    rho = 1e-3,
    alpha = 0.79,
    bm = 1/7.80e3,
    kappanu = 633.33,
    eta = 1e-4,
    force_download = false
  )
  Y, D, M, K, Bp, Bf, A = load_poro_raw_data(
    n=n,
    force_download = force_download
  )
  Y = rho*sparse(Y)
  D = alpha*sparse(D)
  M = 1/bm .* sparse(M)
  K = kappanu*sparse(K)
  A = sparse(A)
  Bp = Bp'
  Bf = Bf'
  n = size(A, 1);
  m = size(M, 1);
  E = [Y spzeros(n,n+m); spzeros(n,n) A spzeros(n,m); spzeros(m,n+n) M];
  J = [spzeros(n,n) -A D';A spzeros(n,n+m); -D spzeros(m,n+m)];
  R = [spzeros(n,2*n+m); spzeros(n,2*n+m); spzeros(m,2*n) K] + eta*I
  B = [zeros(n,1); Bf; Bp];
  return E, J, R, B
end

function load_poro_raw_data(;
    n = 980,
    force_download = false
  )
  filename = "poro-n$n.mat"
  url = "https://zenodo.org/record/5702554/files/poro-n$n.mat?download=1"
  if n == 980
    md5_hash = hex(0x2961a189be7049ffe2d476b18cb1f678)
  elseif n == 320
    md5_hash = hex(0x97afe8c34f0e9a56bbe86d0a51b7b626)
  elseif n == 1805
    md5_hash = hex(0xc61f6687da9cd26cbf2d880d7d3a9ac9)
  else
    throw(ArgumentError("Model size is either 320, 980, or 1805"))
  end
  download_system_data_if_required(
    filename,
    url,
    md5_hash,
    force_download = force_download
  )
  dd = loadMAT(get_filepath(filename))
  return dd["Y"], dd["D"], dd["M"], dd["K"], dd["Bp"], dd["Bf"], dd["A"]
end

export poro_elasticity_model


"""
N: universe dimension
D: manifold dimension
T: coordinate type
"""
struct RectangleChart{N,D,T}
  origin::SVector{N,T}
  tangents::SVector{D,SVector{N,T}}
  normal::SVector{N,T}
end

function RectangleChart(o, tgs)
  n = tgs[1] × tgs[2]
  n = normalize(n)
  RectangleChart(o, tgs, n)
end

struct RectangleNBD{N,D,T}
  chart::RectangleChart{N,D,T}
  parm::SVector{D,T}
  cart::SVector{N,T}
end

cartesian(nbd::RectangleNBD) = nbd.cart

function neighborhood(rect::RectangleChart{3,2}, uv)
  o = rect.origin
  t1, t2 = rect.tangents
  u1, u2 = uv
  cart = o + u1*t1 + u2*t2
  RectangleNBD(rect, SVector(uv), cart)
end

getindex(nbd::RectangleNBD, i) = nbd.cart[i]
normal(nbd::RectangleNBD) = nbd.chart.normal

function quadpoints(chart::RectangleChart{3,2}, rule)

  T = eltype(chart.origin)

  o = chart.origin
  t1 = chart.tangents[1]
  t2 = chart.tangents[2]
  j = norm(t1 × t2)
  Q, P = rule
  U1, W1 = legendre(Q, zero(T), one(T))
  U2, W2 = legendre(P, zero(T), one(T))

  #pts = [o + u1 * t1 + u2 * t2 for u1 in U1 for u2 in U2]
  pts = [neighborhood(chart, (u1,u2)) for u1 in U1 for u2 in U2]
  wts = [j*w1*w2 for w1 in W1 for w2 in W2]

  collect(zip(pts,wts))
end


# ------------------------------------------------------------------------------------------
# # Julia 中的线性代数
# > Based on work by Andreas Noack Jensen (MIT) (http://www.econ.ku.dk/phdstudent/noack/)
# > with edits from Jane Herriman
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## 基础线性代数操作
# 定义一个随机矩阵
# ------------------------------------------------------------------------------------------

A = rand(1:4,3,3)

# ------------------------------------------------------------------------------------------
# 定义一个元素全为 1 的向量
# ------------------------------------------------------------------------------------------

x = fill(1.0, (3,)) # = fill(1.0, 3)

# ------------------------------------------------------------------------------------------
# 注意 $A$ 的类型为 `Array{Int64,2}`，而 $x$ 的类型为 `Array{Float64,1}`。
#
# Julia 定义向量 `Vector{Type}` 的别名为 `Array{Type,1}`，矩阵`Matrix{Type}` 的别名为 `Array{Type,2}`。
#
# 许多线性代数的基础操作和其他语言一样
#
# ### 乘法
# ------------------------------------------------------------------------------------------

b = A*x

# ------------------------------------------------------------------------------------------
# ### 转置
# 就像在其他语言中 `A'` 表示对 `A` 进行共轭转置
# ------------------------------------------------------------------------------------------

A'

# ------------------------------------------------------------------------------------------
# 我们还可以通过 `transpose` 函数获得转置矩阵
# ------------------------------------------------------------------------------------------

transpose(A)

# ------------------------------------------------------------------------------------------
# ### 转置的乘法
# Julia 中某些情况下可以省略 `*` 号
# ------------------------------------------------------------------------------------------

A'A

# ------------------------------------------------------------------------------------------
# ### 解线性方程组
# 用方阵 $A$ 表示的线性方程组 $Ax=b$ 可以用左除运算符（函数）`\` 求解
# ------------------------------------------------------------------------------------------

A\b

# ------------------------------------------------------------------------------------------
# 超定线性方程组（overdetermined linear system）的情况下 `A\b` 给出*最小二乘解（least squares solution）*
# ------------------------------------------------------------------------------------------

Atall = rand(3, 2)

Atall\b

# ------------------------------------------------------------------------------------------
# 缺秩最小二乘问题（a rank-deficient least squares problem）情况下，会给出*极小范数最小二乘解（minimum norm least
# squares solution）*
# ------------------------------------------------------------------------------------------

v = rand(3)
rankdef = hcat(v, v)

rankdef\b

# ------------------------------------------------------------------------------------------
# 欠定问题（an underdetermined solution）下，会给出 最小范数解（minimum norm solution）
# ------------------------------------------------------------------------------------------

bshort = rand(2)
Ashort = rand(2, 3)

Ashort\bshort

# ------------------------------------------------------------------------------------------
# ## `LinearAlgebra` 包
#
# 如上所示，Julia默认自带了很多线性代数操作，而标准库 `LinearAlgebra` 提供了更多的相关名词和函数。
# 特别是它提供了因子分解和一些结构化矩阵类型。
# 和其他包一样，可以通过 `using LinearAlgebra` 引入这些额外的特征。
# ------------------------------------------------------------------------------------------

using LinearAlgebra

# ------------------------------------------------------------------------------------------
# ## 练习
#
# ### 10.1
# 计算向量 `v` 和它自己的内积，并赋值给 `dot_v`。
# ------------------------------------------------------------------------------------------

v = [1,2,3]

# 在此作答
dot_v = v

@assert dot_v == 14

# ------------------------------------------------------------------------------------------
# ### 10.2
# 计算向量 v 和它自己的外积(Outer product)，并赋值给 `outer_v`。
#
# 提示：外积是一种克罗内克积(Kronecker product)
#
# 计算叉积(cross product)，赋值给 `cross_v`
# ------------------------------------------------------------------------------------------

# 在此作答
outer_v = v
cross_v = v

@assert outer_v == [1 2 3
                    2 4 6
                    3 6 9]

@assert cross_v == [0, 0, 0]

# ------------------------------------------------------------------------------------------
# 请在完成练习运行 `@assert` 块以检测结果是否正确，若无报错即为正确。
# ------------------------------------------------------------------------------------------


module GMMbox

using Distances
using Distributions
using Clustering
import Base.rand
import Base.rand!

export gaussian, logprob!, logprob, gmm, rand, rand!

mutable struct gaussian{T<:AbstractFloat}
  ndim::Int64
  mean::AbstractArray{T,2}
  cov::AbstractArray{T,2}
end

gaussian(ndim::Integer) = gaussian(ndim,zeros(ndim,1),eye(ndim))
gaussian(data::Matrix{T} where T<:AbstractFloat)  = gaussian(size(data,1),mean(data,2),cov(data,2))

mutable struct gmm{T<:AbstractFloat}
  nmix::Int64
  wts::AbstractArray{T,2}
  mix::AbstractArray{gaussian{T}}
end

function gmm(nmix::Integer, data::Matrix{T} where T <: AbstractFloat; method = "rand")
  if(method=="rand")
    ndim = size(data,1);
    dcov = cov(data,2);
    dmvn = MvNormal(vec(mean(data,2)),dcov);
    mix = Array{gaussian{eltype(data)}}(nmix);
    for i=1:nmix
      mixmean = convert(Array{eltype(data)},rand(dmvn));
      mix[i] = gaussian(ndim,reshape(mixmean,ndim,1),dcov);
    end
  elseif(method=="kmeans")
    km = kmeans(data,nmix);
    mix = Array{gaussian{eltype(data)}}(nmix);
    dcov = cov(data,2);
    for i=1:nmix
      mix[i] = gaussian(data[:,find(km.assignments.==i)]);
      mix[i].cov = dcov;
    end
  end

  return gmm(nmix,convert(Array{eltype(data)},(1/nmix)*ones(nmix,1)),mix);
end

rand!(g::gaussian{T},data::Matrix{T}) where T <: AbstractFloat = rand!(MvNormal(vec(g.mean),g.cov),data);
rand(g::gaussian{T},nvec::Integer) where T <: AbstractFloat = rand(MvNormal(vec(g.mean),g.cov),nvec);

function rand!(gm::gmm{T},data::Matrix{T}) where T <: AbstractFloat
  ftype = eltype(data);
  nvec = size(data,2);
  pmix = rand(ftype,nvec);
  nvecmix = zeros(Int64,gm.nmix);
  thr = zeros(ftype,1);
  for i=1:gm.nmix
    nvecmix[i] = sum((pmix.>=thr).&(pmix.<thr+gm.wts[i]))
    thr += gm.wts[i];
  end
  if(sum(nvecmix)<nvec)
    nvecmix[end] += nvec-sum(nvecmix);
    println("sum of number of samples per mixture not equal to requested number of samples");
  end
  sindx = 1;
  eindx = 0;
  for i=1:gm.nmix
    eindx += nvecmix[i];
    data[:,sindx:eindx] = rand(gm.mix[i],nvecmix[i]);
    sindx += nvecmix[i];
  end
  return nothing
end

function rand(gm::gmm{T},nvec::Integer) where T <: AbstractFloat
  ftype = eltype(gm.wts);
  pmix = rand(ftype,nvec);
  nvecmix = zeros(Int64,gm.nmix);
  thr = zeros(ftype,1);
  for i=1:gm.nmix
    nvecmix[i] = sum((pmix.>=thr).&(pmix.<thr+gm.wts[i]))
    thr += gm.wts[i];
  end
  if(sum(nvecmix)<nvec)
    nvecmix[end] += nvec-sum(nvecmix);
    println("sum of number of samples per mixture not equal to requested number of samples");
  end
  data = zeros(ftype,gm.mix[1].ndim,sum(nvec));
  sindx = 1;
  eindx = 0;
  for i=1:gm.nmix
    eindx += nvecmix[i];
    data[:,sindx:eindx] = rand(gm.mix[i],nvecmix[i]);
    sindx += nvecmix[i];
  end
  return data
end

logprob!(p::Array{T},g::gaussian{T},data::Matrix{T}) where T <: AbstractFloat = logpdf!(p,MvNormal(vec(g.mean),g.cov),data)
logprob(g::gaussian{T},data::Matrix{T}) where T <: AbstractFloat = logpdf(MvNormal(vec(g.mean),g.cov),data)

function logprob!(p::Array{T},gm::gmm{T},data::Matrix{T}) where T <: AbstractFloat
  mixprob = zeros(eltype(data),size(data,2),gm.nmix);
  for i=1:gm.nmix
    mixprob[:,i] = logprob(gm.mix[i],data);
  end
  maxprob = maximum(mixprob,2);
  broadcast!(-,mixprob,mixprob,maxprob);
  p .= vec(log.((exp.(mixprob))*gm.wts) .+ maxprob);
  return nothing
end

function logprob(gm::gmm{T},data::Matrix{T}) where T <: AbstractFloat
  mixprob = zeros(eltype(data),size(data,2),gm.nmix);
  for i=1:gm.nmix
    mixprob[:,i] = logprob(gm.mix[i],data);
  end
  maxprob = maximum(mixprob,2);
  broadcast!(-,mixprob,mixprob,maxprob);
  return vec(log.((exp.(mixprob))*gm.wts) .+ maxprob);
end

#=
function logprob!(p::Matrix{T},g::gaussian{T},data::Matrix{T}) where T <: AbstractFloat
  logpdf!(reshape(p,length(p)),MvNormal(vec(g.mean),g.cov),data);
  reshape(p,length(p),1);
end
=#

#=
#logpdf and logpdf! using Distributions.jl is faster than this version which uses Distances.jl for Float32 - not using this version for CPU code
function logprob!(p::Matrix{T},g::gaussian{T},data::Matrix{T}) where T <: AbstractFloat
  ndim = size(data,1);
  icov = inv(g.cov);
  lconst = ndim*log(2*pi) + log(det(g.cov));
  pairwise!(p,SqMahalanobis(icov),data,g.mean);
  p .+= lconst;
  p .*= -0.5;
  return nothing
end

function logprob(g::gaussian{T},data::Matrix{T}) where T <: AbstractFloat
  p = zeros(eltype(data),size(data,2),1);
  ndim = size(data,1);
  icov = inv(g.cov);
  lconst = ndim*log(2*pi) + log(det(g.cov));
  pairwise!(p,SqMahalanobis(icov),data,g.mean);
  p .+= lconst;
  p .*= -0.5;
  return p
end
=#

end


# Some global variables that will stay constant
α = 0.013
α_q = (1-(1-α)^3)
b_param = 0.0124
d_param = 0.00822
β = 0.98
γ = 1.0
σ = 2.0

# The default wage distribution: a discretized log normal
log_wage_mean, wage_grid_size, max_wage = 20, 200, 170
w_vec = linspace(1e-3, max_wage, wage_grid_size + 1)
logw_dist = Normal(log(log_wage_mean), 1)
cdf_logw = cdf.(logw_dist, log.(w_vec))
pdf_logw = cdf_logw[2:end] - cdf_logw[1:end-1]
p_vec = pdf_logw ./ sum(pdf_logw)
w_vec = (w_vec[1:end-1] + w_vec[2:end]) / 2

"""
Compute the reservation wage, job finding rate and value functions of the
workers given c and τ.

"""
function compute_optimal_quantities(c::AbstractFloat, τ::AbstractFloat)
    mcm = McCallModel(α_q,
                      β,
                      γ,
                      c-τ,                # post-tax compensation
                      σ,
                      collect(w_vec-τ),  # post-tax wages
                      p_vec)

    w_bar, V, U = compute_reservation_wage(mcm, return_values=true)
    λ = γ * sum(p_vec[w_vec - τ .> w_bar])

    return w_bar, λ, V, U
end

"""
Compute the steady state unemployment rate given c and tau using optimal
quantities from the McCall model and computing corresponding steady state
quantities

"""
function compute_steady_state_quantities(c::AbstractFloat, τ::AbstractFloat)
    w_bar, λ_param, V, U = compute_optimal_quantities(c, τ)

    # Compute steady state employment and unemployment rates
    lm = LakeModel(λ=λ_param, α=α_q, b=b_param, d=d_param)
    x = rate_steady_state(lm)
    u_rate, e_rate = x

    # Compute steady state welfare
    w = sum(V .* p_vec .* (w_vec - τ .> w_bar)) / sum(p_vec .* (w_vec - τ .> w_bar))
    welfare = e_rate .* w + u_rate .* U

    return u_rate, e_rate, welfare
end

"""
Find tax level that will induce a balanced budget.

"""
function find_balanced_budget_tax(c::Real)
    function steady_state_budget(t::Real)
      u_rate, e_rate, w = compute_steady_state_quantities(c, t)
      return t - u_rate * c
    end

    τ = brent(steady_state_budget, 0.0, 0.9 * c)

    return τ
end

# Levels of unemployment insurance we wish to study
Nc = 60
c_vec = linspace(5.0, 140.0, Nc)

tax_vec = Vector{Float64}(Nc)
unempl_vec = Vector{Float64}(Nc)
empl_vec = Vector{Float64}(Nc)
welfare_vec = Vector{Float64}(Nc)

for i = 1:Nc
    t = find_balanced_budget_tax(c_vec[i])
    u_rate, e_rate, welfare = compute_steady_state_quantities(c_vec[i], t)
    tax_vec[i] = t
    unempl_vec[i] = u_rate
    empl_vec[i] = e_rate
    welfare_vec[i] = welfare
end

fig, axes = subplots(2, 2, figsize=(15, 10))

plots = [unempl_vec, empl_vec, tax_vec, welfare_vec]
titles = ["Unemployment", "Employment", "Tax", "Welfare"]

for (ax, plot, title) in zip(axes, plots, titles)
    ax[:plot](c_vec, plot, "b-", lw=2, alpha=0.7)
    ax[:set](title=title)
end

fig[:tight_layout]()


using FastGaussQuadrature
include("chebyshev.jl")
include("pgrid.jl")

idx(ki, xi, order) = (ki - 1) * order + xi

struct CompositeGrid
    grid::Vector{Float64} # panel x finegrid
    wgrid::Vector{Float64}
    panel::Vector{Float64}
    x::Vector{Float64}
    w::Vector{Float64}
    # xp::Matrix{Float64} # the fine grid
    # wp::Matrix{Float64} # the weight of the fine grid
    Np::Int  # number of panel - 1
    order::Int  # number of fine grid
    type::Symbol
    function CompositeGrid(panel, order, type::Symbol)
        if type == :cheb
            x, w = barychebinit(order)
        elseif type == :gaussian
            x, w = gausslegendre(order)
        else
            throw("not implemented!")
        end

        # println("Quadrature type : ", type)
        # println("Quadrature points : ", x)
        # println("Quadrature weights: ", w)

        Np = length(panel) - 1

        grid = zeros(Float64, Np * order)
        wgrid = zeros(Float64, Np * order)
        # xp = zeros(Float64, (Np, order))
        # wp = zeros(Float64, (Np, order))

        for p in 1:Np
            a, b = panel[p], panel[p + 1]
            for xi in 1:order
                index = idx(p, xi, order)
                grid[index] = (a + b) / 2 + (b - a) / 2 * x[xi]
                wgrid[index] = (b - a) / 2 * w[xi]
            end
        end

        return new(grid, wgrid, panel, x, w, Np, order, type)
    end
end

function KPanel(Nk, kF, maxK, minK)
    panel = Grid.boseKUL(0.5 * kF, maxK, minK, Nk, 1).grid
    panel[1] = 0.0  # the kgrid start with 0.0
    return panel
end

function QPanel(Nk, kF, maxK, minK, k)
    panel = pGrid(k, kF, maxK, minK, [Nk,Nk,Nk,Nk]).grid

    # panel = Grid.boseKUL(0.5 * kF, maxK, minK, Nk, 1).grid
    # panel[1] = 0.0  # the kgrid start with 0.0

    return panel
end

"""
interpolate!(f, k::CompositeGrid, grid)

map f array in the grid k to a new grid

# Arguments
- f::Vector{Float64}: vector of data.
- k::CompositeGrid: the grid object that f is defined on
- grid::Vector{Float64}: new grid points
"""
function interpolate(f, k::CompositeGrid, grid)
    @assert k.type == :cheb
    order = k.order
    ff = zeros(eltype(f), length(grid))
    kpidx = 1 # panel index of the kgrid
    head, tail = idx(kpidx, 1, order), idx(kpidx, order, order)
    # extract all x in the kpidx-th k panel
    fx = @view f[head:tail] # all F in the same kpidx-th K panel
    x = @view k.grid[head:tail]
    w = @view k.wgrid[head:tail]

    for (qi, q) in enumerate(grid)
        # for a given q, one needs to find the k panel to do interpolation
        if q > k.panel[kpidx + 1]
            # if q is too large, move k panel to the next
            # println("before $q, $kpidx, $(k.panel[kpidx]) -> $(k.panel[kpidx + 1])")
            while q > kgrid.panel[kpidx + 1]
                kpidx += 1
            end
            # println("after $q, $kpidx, $(k.panel[kpidx]) -> $(k.panel[kpidx + 1])")
            head, tail = idx(kpidx, 1, order), idx(kpidx, order, order)
            fx = @view f[head:tail] # all F in the same kpidx-th K panel
            x = @view k.grid[head:tail]
            w = @view k.wgrid[head:tail]
            @assert kpidx <= k.Np
        end
        ff[qi] = barycheb(order, q, fx, w, x) # the interpolation is independent with the panel length
        # @assert k.panel[kpidx - 1][order] <= q <= k.panel[pidx + 1][1] "$q for kpidx=$kpidx with $x"
    end
    return ff
end

if abspath(PROGRAM_FILE) == @__FILE__
    grid=CompositeGrid([0.0, 1.0, 2.0], 4, :cheb)
    println(grid.grid)
end


# ----------
# Vector Calculus Module, Only work on peroideric boundary!
# ----------

function Curl(B1::Array,B2::Array,B3::Array;
              Lx = 2π, Ly = Lx, Lz = Lx,T = Float32)
    nx,ny,nz = size(B1);
    grid = ThreeDGrid(nx, Lx, ny, Ly, nz, Lz, T = T);
    cB1,cB2,cB3 = Curl(B1,B2,B3,grid;Lx = Lx, Ly = Ly, Lz = Lz,T = T)
    return cB1,cB2,cB3;
end

function Curl(B1::Array,B2::Array,B3::Array,grid;
              Lx = 2π, Ly = Lx, Lz = Lx,T = Float32)
    #funtion of computing ∇×Vector using the fourier method
    # fft(∇×Vector) -> im * k × V
    #| i j k  |
    #| x y z  |
    #|B1 B2 B3|
    #
    nx,ny,nz = size(B1);
    B1h = zeros(Complex{T},(div(nx,2)+1,ny,nz));
    B2h = zeros(Complex{T},(div(nx,2)+1,ny,nz));
    B3h = zeros(Complex{T},(div(nx,2)+1,ny,nz));
    cBxh = copy(B1h);
    cByh = copy(B2h);
    cBzh = copy(B3h);
    mul!(B1h, grid.rfftplan, B1);
    mul!(B2h, grid.rfftplan, B2);
    mul!(B3h, grid.rfftplan, B3);

    for k in 1:nz, j in 1:ny,i in 1:div(nx,2)+1
       x,y,z = grid.kr[i],grid.l[j],grid.m[k];
       cBxh[i,j,k] = im*(y*B3h[i,j,k] - z*B2h[i,j,k]);
       cByh[i,j,k] = im*(z*B1h[i,j,k] - x*B3h[i,j,k]);
       cBzh[i,j,k] = im*(x*B2h[i,j,k] - y*B1h[i,j,k]);
    end

    cB1,cB2,cB3 = zeros(T,size(B1)),zeros(T,size(B1)),zeros(T,size(B1));
    ldiv!(cB1, grid.rfftplan, deepcopy(cBxh));
    ldiv!(cB2, grid.rfftplan, deepcopy(cByh));
    ldiv!(cB3, grid.rfftplan, deepcopy(cBzh));
    return cB1,cB2,cB3
end

function Div(B1::Array,B2::Array,B3::Array;
             Lx = 2π, Ly = Lx, Lz = Lx,T = Float32)
    nx,ny,nz = size(B1);
    grid = ThreeDGrid(nx, Lx, ny, Ly, nz, Lz, T = T);
    cB1 = Div(B1,B2,B3,grid;Lx = Lx, Ly = Ly, Lz = Lz,T = T);

    return cB1
end

function Div(B1::Array,B2::Array,B3::Array,grid;
             Lx = 2π, Ly = Lx, Lz = Lx,T = Float32)
    #funtion of computing ∇̇ ⋅ Vector using the fourier method
    # fft(∇·Vector) -> im * k ⋅ V
    # = im* x*B1 + y*B2 + z*B3
    nx,ny,nz = size(B1);
    B1h = zeros(Complex{T},(div(nx,2)+1,ny,nz));
    B2h = zeros(Complex{T},(div(nx,2)+1,ny,nz));
    B3h = zeros(Complex{T},(div(nx,2)+1,ny,nz));
    Dot = copy(B1h);
    mul!(B1h, grid.rfftplan, B1);
    mul!(B2h, grid.rfftplan, B2);
    mul!(B3h, grid.rfftplan, B3);

    for k in 1:nz, j in 1:ny,i in 1:div(nx,2)+1
       x,y,z = grid.kr[i],grid.l[j],grid.m[k];
       Dot[i,j,k] = x*B1h[i,j,k] + y*B2h[i,j,k] + z*B3h[i,j,k];
    end

    cB1 = zeros(T,size(B1))
    ldiv!(cB1, grid.rfftplan, deepcopy(Dot));

    return cB1
end

function LaplaceSolver(B::Array; Lx=2π, Ly = Lx, Lz = Lz, T = Float32)
    nx,ny,nz = size(B);
    grid = ThreeDGrid(nx, Lx, ny, Ly, nz, Lz, T = T);
    Φ   = LaplaceSolver(B,grid; Lx=2π, Ly = Lx, Lz = Lz, T = Float32);
    return Φ
end

function LaplaceSolver(B::Array,grid; Lx=2π, Ly = Lx, Lz = Lz, T = Float32)
    #=
    funtion of computing ΔΦ = B using the fourier method, must be peroidic condition
    Considering in k-space, k² Φ' = B', we would get Φ = F(B'/k²)
    =#
    nx,ny,nz = size(B);
    Φ    = zeros(T,nx,ny,nz);
    Bh   = zeros(ComplexF32,(div(nx,2)+1,ny,nz));
    mul!(Bh, grid.rfftplan, B);
    for k in 1:nz, j in 1:ny, i in 1:div(nx,2)+1
       x,y,z = grid.kr[i],grid.l[j],grid.m[k];
       k² = x^2 + y^2 + z^2;
       Bh[i,j,k] = Bh[i,j,k]/k²;
       if k² == 0; Bh[i,j,k] = 0; end
    end
    ldiv!(Φ, grid.rfftplan, deepcopy(Bh));
    return Φ;
end

function Crossproduct(A1,A2,A3,B1,B2,B3)
    C1 = @.  (A2*B3 - A3*B2);
    C2 = @. -(A1*B3 - A3*B1);
    C3 = @.  (A1*B2 - A2*B1);
    return C1,C2,C3
end

function Dotproduct(A1,A2,A3,B1,B2,B3)
    return A1.*B1 + A2.*B2 + A3.*B3
end


export ShepardInterpolator, RBFInterpolator
export multiquadratic, invmultiquadratic, gaussian

using LinearAlgebra: Symmetric

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

function distance(a, b, n::Int)
    r = 0.0
    for i = 1:n
        @inbounds d = a[i] - b[i]
        r += d*d
    end
    sqrt(r)
end

function checkcoords(X, y)
    @assert size(X)[2] > 1 "Generally it doesn't make sense to apply scattered interpolation methods in one dimension. Consider using a 1D interpolator."
    p = size(X)[1]
    n = length(y)
    @assert p == n "number of points in X ($p) does not match number of values in y ($n)"
end

function checkdims(x, N::Int)
    n = length(x)
    @assert n == N "Number of interpolation coordinates ($n) does not match number of interpolator dimensions ($N)"
end

struct ScatteredPoints{T}
    #number of dimensions
    N::Int64
    #number of points
    p::Int64
    #table of points, p x N
    X::Array{T,2}
    #vector of values
    y::Vector{T}
end

function ScatteredPoints(X, y)
    checkcoords(X, y)
    N = size(X)[2]
    p = size(X)[1]
    ScatteredPoints(N, p, X, y)
end

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

"""
    multiquadratic(r, ϵ)

``\\sqrt{1 + (r/ϵ)^2}``
"""
multiquadratic(r, ϵ) = sqrt(1 + (r/ϵ)^2)

"""
    invmultiquadratic(r, ϵ)

``\\frac{1}{\\sqrt{1 + (r/ϵ)^2}}``
"""
invmultiquadratic(r, ϵ) = 1/multiquadratic(r, ϵ)

"""
    gaussian(r, ϵ)

``e^{-r^2/ϵ^2}``
"""
gaussian(r, ϵ) = exp(-r^2/ϵ^2)

struct RBFInterpolator{T,F}
    #interpolation points
    S::ScatteredPoints{T}
    #radial basis function
    rbf::F
    #distance scale factor
    ϵ::T
    #function weights
    w::Vector{T}
end

"""
    RBFInterpolator(X, y, ϵ, rbf=invmultiquadratic)

Construct a radial basis function (RBF) interpolator for an n-dimensional set of points with coordinates `X` and values `y`. `X` must be an p × N array, where p is the number of points and N is the number of dimensions. `y` must be a length p vector. The value of `ϵ` scales the radial basis function of choice, `f`, which is [`invmultiquadratic`](@ref) by default. Any function in the form ``ϕ(r,ϵ)`` can be passed to the `rbf` argument, where \$r\$ is the distance between points and ``ϵ`` is a scaling factor.
"""
function RBFInterpolator(X::AbstractMatrix,
                         y::AbstractVector,
                         ϵ,
                         rbf::F=multiquadratic) where {F}
    S = ScatteredPoints(X, y)
    #matrix of basis function evaluations
    A = ones(eltype(y), S.p, S.p)
    for i = 1:S.p
        a = view(S.X,i,:)
        for j = i+1:S.p
            b = view(S.X,j,:)
            r = distance(a, b, S.N)
            A[i,j] = rbf(r, ϵ)
        end
    end
    #solve for weights
    w = Symmetric(A)\S.y
    #construct
    RBFInterpolator(S, rbf, convert(eltype(y), ϵ), w)
end

function RBF(Φ::RBFInterpolator, x)
    checkdims(x, Φ.S.N)
    y = 0.0
    for i = 1:Φ.S.p
        #distance
        r = distance(x, view(Φ.S.X,i,:), Φ.S.N)
        #weighted evaluation
        @inbounds y += Φ.w[i]*Φ.rbf(r, Φ.ϵ)
    end
    return y
end

function (Φ::RBFInterpolator)(x::Union{AbstractVector,Tuple})
    RBF(Φ, x)
end

function (Φ::RBFInterpolator)(x::Number...)
    RBF(Φ, collect(Float64, x))
end

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

struct ShepardInterpolator{T}
    #interpolation points
    S::ScatteredPoints{T}
    #power-law exponent
    a::Float64
end

"""
    ShepardInterpolator(X, y, a=3)

Construct a `ShepardInterpolator` for an n-dimensional set of points with coordinates `X` and values `y`. `X` must be an p × N array, where p is the number of points and N is the number of dimensions. `y` must be a length p vector. The value of `a` defines the distance weighting function ``r^{-a}``.
"""
function ShepardInterpolator(X::AbstractMatrix,
                             y::AbstractVector,
                             a::Real=3.0)
    ShepardInterpolator(ScatteredPoints(X, y), a)
end

function shepard(Φ::ShepardInterpolator, x)
    checkdims(x, Φ.S.N)
    n = 0.0
    d = 0.0
    for i = 1:Φ.S.p
        #distance
        r = distance(x, view(Φ.S.X,i,:), Φ.S.N)
        #evaluate weighting function
        f = 1.0/r^Φ.a
        #weigted and unweighted contribution to numertor and denominator
        @inbounds n += Φ.S.y[i]*f
        d += f
    end
    return n/d
end

(Φ::ShepardInterpolator)(x::Union{AbstractVector,Tuple}) = shepard(Φ, x)

(Φ::ShepardInterpolator)(x::Number...) = shepard(Φ, x)


using LinearAlgebra
using FinEtools
using PGFPlotsX

objects = []

E = 200*phun("GPa");
nu = 0.3;
G = E/2/(1+nu)
a = 1.0*phun("m");
t = a/5
D = E*t^3/12/(1-nu^2)
P = 1.0
a = 1.0
r = collect(t/100000:t/10000:t)
w_b = @. P/16/pi/D * (2*r^2 * log(r/a) + (3+nu)/(1+nu)*(a^2 - r^2))
w_s = @. -P/4/pi/G/t * log(r/a)
max_w_b = maximum(w_b)

@pgf p = PGFPlotsX.Plot(
{
color = "black",
line_width  = 0.7,
style = "solid"
},
Coordinates([v for v in  zip(r./t, w_b./max_w_b)])
)
push!(objects, p)
push!(objects, LegendEntry("Bending only"))

@pgf p = PGFPlotsX.Plot(
{
color = "red",
line_width  = 0.7,
style = "solid"
},
Coordinates([v for v in  zip(r./t, (w_b+w_s)./max_w_b)])
)
push!(objects, p)
push!(objects, LegendEntry("Bending plus shear"))

@pgf ax = Axis(
    {
        xlabel = "Normalized Distance [ND]",
        ylabel = "Normalized Displacement [ND]",
        # xmin = range[1],
        # xmax = range[2],
        xmode = "linear",
        ymode = "linear",
        yminorgrids = "true",
        grid = "both",
        legend_style = {
            at = Coordinate(0.5, 1.05),
            anchor = "south",
            legend_columns = -1
        },
    },
    objects...
)

display(ax)
pgfsave("deflections.pdf", ax)


function max_divisor(number::Integer) ::Integer
    return Integer(floor(sqrt(number)))
end

function is_prime(number::Integer) ::Bool
    max::Integer = max_divisor(number)
    # False if any number evenly divides it.
    for x=2:max
        if number % x == 0
            return false
        end
    end
    # true otherwise.
    return true
end

function next_prime(number::Integer) ::Integer
    to_check = number
    while true
        to_check += 1
        if is_prime(to_check)
            return to_check
        end
    end
end

function first_n_primes(number::Integer) ::Vector{<:Integer}
    last::Integer = 0
    primes = Vector{Integer}()
    for n=1:number
        last = next_prime(last)
        push!(primes, last)
    end
    return primes
end

function first_n_twin_primes(number::Integer) ::Vector{<:Integer}
    last::Integer = 0
    primes = Vector{Integer}()
    for n=3:number
        while true
            last = next_prime(last)
            if is_prime(last + 2)
                push!(primes, last)
                push!(primes, last+2)
                last += 2
                break
            end
        end
    end
    return primes
end

sum(x^2 for x in first_n_twin_primes(10000))


#
# NBody calculations
#

"""
Returns total energy for `NBodySystem`.
"""
function system_energy(sys::NBodySystem)

    E = 0.0u"J"
    for i = 1:length(sys.body)

        E += sys.body[i].m * dot(sys.body[i].v̅, sys.body[i].v̅)
        for j = 1:length(sys.body)

            if i ≠ j
                E -= 6.6743e-11u"m^3/(kg*s^2)" * (sys.body[i].m * sys.body[j].m) /
                                  norm(sys.body[j].r̅ .- sys.body[i].r̅)

            end
        end
    end

    return E

end

"""
Returns total angular momentum for `NBodySystem`.
"""
function system_angular_momentum(sys::NBodySystem)

    H = reduce(+, [sys.body[i].m * cross(sys.body[i].r̅, sys.body[i].v̅) for i ∈ 1:length(sys.body)])

end


#=
We shall say that an n-digit number is pandigital if it makes use of all the digits 1 to n exactly once. For example, 2143 is a 4-digit pandigital and is also prime.

What is the largest n-digit pandigital prime that exists?
=#
include("projecteulerutils.jl")
using Combinatorics
using Primes

function calc()
  largestpandigitalprime = 1
  for n in 2:9
    digit_arr = [i for i in 1:n]
    for permutation in permutations(digit_arr)
      x = nd(permutation)
      if isprime(x) && x > largestpandigitalprime
        largestpandigitalprime = x
      end
    end
  end
  largestpandigitalprime
end
@time println(calc())