"""
ITERATIVE HARD THRESHOLDING

    iht!(v, x, y, k [, iter=1, nstep=50]) -> (Float, Int)

This function computes a hard threshold update

    β = P_Sk(β0 + μ*x'*(y - x*β0))

where `μ` is the step size (or learning rate) and P_Sk denotes the projection onto the set S_k defined by

S_k = { x in R^p : || x ||_0 <= k }.

The projection in question preserves the largest `k` components of `β` in magnitude, and it sends the remaining
`p - k` components to zero. This update is intimately related to a projected gradient step used in Landweber iteration.
Unlike the Landweber method, this function performs a line search on `μ` whenever the step size exceeds a specified
threshold `ω` given by

    ω = sum(abs2, β - β0) / sum(abs2, x*(β - β0))

By backtracking on `μ`, this function guarantees a stable estimation of a sparse `β`.
This function is based on the [HardLab](http://www.personal.soton.ac.uk/tb1m08/sparsify/sparsify.html/) demonstration code written in MATLAB by Thomas Blumensath.

Arguments:

- `v` is the `IHTVariables` object of temporary arrays
- `x` is the `n` x `p` design matrix.
- `y` is the vector of `n` responses.
- `k` is the model size.
- `iter` is the current iteration count in the IHT algorithm. Defaults to `1`.
- `nstep` is the maximum permissible number of backtracking steps. Defaults to `50`.

Output:

- `μ` is the step size used to update `β`, after backtracking.`
- `μ_step` is the number of backtracking steps used on `μ`.
"""
function iht!{T <: Float, V <: DenseVector}(
    v     :: IHTVariables{T, V},
    x     :: DenseMatrix{T},
    y     :: V,
    k     :: Int,
    iter  :: Int = 1,
    nstep :: Int = 50
)

    # compute indices of nonzeroes in beta
    _iht_indices(v, k)

    # store relevant columns of x
    # need to do this on 1st iteration
    # afterwards, only do if support changes
    if !isequal(v.idx, v.idx0) || iter < 2
        #update_xk!(v.xk, x, v.idx)   # xk = x[:,v.idx]
        copy!(v.xk, view(x, :, v.idx))
    end

    # store relevant components of gradient
    fill_perm!(v.gk, v.df, v.idx)    # gk = g[v.idx]

    # now compute subset of x*g
    A_mul_B!(v.xgk, v.xk, v.gk)

    # warn if xgk only contains zeros
    all(v.xgk .== zero(T)) && warn("Entire active set has values equal to 0")

    # compute step size
#    mu = _iht_stepsize(v, k) :: T # does not yield conformable arrays?!
    μ = (sum(abs2, v.gk) / sum(abs2, v.xgk)) :: T

    # check for finite stepsize
    isfinite(μ) || throw(error("Step size is not finite, is active set all zero?"))

    # compute gradient step
    _iht_gradstep(v, μ, k)

    # update xb
    update_xb!(v.xb, x, v.b, v.idx, k)
    #A_mul_B!(v.xb, view(x, :, v.idx), view(v.b, v.idx) )

    # calculate omega
    ω_top, ω_bot = _iht_omega(v)

    # backtrack until mu < omega and until support stabilizes
    μ_step = 0
    while _iht_backtrack(v, ω_top, ω_bot, μ, μ_step, nstep)

        # stephalving
        μ /= 2

        # recompute gradient step
        copy!(v.b,v.b0)
        _iht_gradstep(v, μ, k)

        # recompute xb
        update_xb!(v.xb, x, v.b, v.idx, k)
        #A_mul_B!(v.xb, view(x, :, v.idx), view(v.b, v.idx) )

        # calculate omega
        ω_top, ω_bot = _iht_omega(v)

        # increment the counter
        μ_step += 1
    end

    return μ::T, μ_step::Int
end

"""
L0 PENALIZED LEAST SQUARES REGRESSION

    L0_reg(x,y,k) -> IHTResults

This routine minimizes the loss function

    0.5*sum(abs2, y - x*β )

subject to `β` lying in the set S_k = { x in R^p : || x ||_0 <= k }.

It uses Thomas Blumensath's iterative hard thresholding framework to keep `β` feasible.

Arguments:

- `x` is the `n` x `p` data matrix.
- `y` is the `n`-dimensional continuous response vector.
- `k` is the desired model size (support).

Optional Arguments:

- `max_iter` is the maximum number of iterations for the algorithm. Defaults to `1000`.
- `max_step` is the maximum number of backtracking steps for the step size calculation. Defaults to `50`.
- `tol` is the global tolerance. Defaults to `1e-4`.
- `quiet` is a `Bool` that controls algorithm output. Defaults to `true` (no output).
- `v` is an `IHTVariables` structure used to house temporary arrays. Used primarily in `iht_path` for memory efficiency.

Outputs are wrapped into an `IHTResults` structure with the following fields:

- 'time' is the compute time for the algorithm. Note that this does not account for time spent initializing optional argument defaults.
- 'iter' is the number of iterations that the algorithm took before converging.
- 'loss' is the optimal loss (half of residual sum of squares) at convergence.
- 'beta' is the final estimate of `b`.
"""
function L0_reg{T <: Float, V <: DenseVector}(
    x        :: DenseMatrix{T},
    y        :: V,
    k        :: Int;
    v        :: IHTVariables{T, V} = IHTVariables(x, y, k),
    tol      :: T    = convert(T, 1e-4),
    max_iter :: Int  = 100,
    max_step :: Int  = 50,
    quiet    :: Bool = true
)

    # start timer
    tic()

    # first handle errors
    k        >= 0      || throw(ArgumentError("Value of k must be nonnegative!\n"))
    max_iter >= 0      || throw(ArgumentError("Value of max_iter must be nonnegative!\n"))
    max_step >= 0      || throw(ArgumentError("Value of max_step must be nonnegative!\n"))
    tol      >  eps(T) || throw(ArgumentError("Value of global tol must exceed machine precision!\n"))

    # initialize return values
    iter = 0                      # number of iterations of L0_reg
    exec_time = zero(T)           # compute time *within* L0_reg
    next_obj  = zero(T)           # objective value
    next_loss = zero(T)           # loss function value

    # initialize floats
    current_obj = convert(T, Inf) # tracks previous objective function value
    the_norm    = zero(T)         # norm(b - b0)
    scaled_norm = zero(T)         # the_norm / (norm(b0) + 1)
    μ           = zero(T)          # Landweber step size, 0 < tau < 2/rho_max^2

    # initialize integers
    #i      = 0                    # used for iterations in loops
    μ_step = 0                    # counts number of backtracking steps for mu

    # initialize booleans
    converged = false             # scaled_norm < tol?

    # update xb, r, gradient
    initialize_xb_r_grad!(v, x, y, k)

    # update loss and objective
    next_loss = convert(T, Inf)

    # formatted output to monitor algorithm progress
    !quiet && print_header()

    # main loop
    for iter = 1:max_iter

        # notify and break if maximum iterations are reached.
        if iter >= max_iter

            # alert about hitting maximum iterations
            !quiet && print_maxiter(max_iter, loss)

            # send elements below tol to zero
            threshold!(v.b, tol)

            # stop timer
            exec_time = toq()

            # these are output variables for function
            # wrap them into a Dict and return
            return IHTResults(exec_time, next_loss, iter, copy(v.b))
        end

        # save values from previous iterate
        copy!(v.b0, v.b)             # b0 = b
        copy!(v.xb0, v.xb)           # Xb0 = Xb
        current_obj = next_obj

        # now perform IHT step
        (μ, μ_step) = iht!(v, x, y, k, max_step, iter)

        # the IHT kernel gives us an updated x*b
        # use it to recompute residuals and gradient
        update_r_grad!(v, x, y)

        # update loss, objective, and gradient
        next_loss = sum(abs2, v.r) / 2

        # guard against numerical instabilities
        check_finiteness(next_loss)

        # track convergence
        the_norm    = chebyshev(v.b, v.b0)
        scaled_norm = (the_norm / ( norm(v.b0,Inf) + 1)) :: T
        converged   = (scaled_norm < tol) :: Bool

        # output algorithm progress
        quiet || @printf("%d\t%d\t%3.7f\t%3.7f\t%3.7f\n", iter, μ_step, μ, the_norm, next_loss)

        # check for convergence
        # if converged and in feasible set, then algorithm converged before maximum iteration
        # perform final computations and output return variables
        if converged

            # send elements below tol to zero
            threshold!(v.b, tol)

            # stop time
            exec_time = toq()

            # announce convergence
            !quiet && print_convergence(iter, next_loss, exec_time)

            # these are output variables for function
            # wrap them into a Dict and return
            return IHTResults(exec_time, next_loss, iter, copy(v.b))
        end

        # algorithm is unconverged at this point.
        # if algorithm is in feasible set, then rho should not be changing
        # check descent property in that case
        # if rho is not changing but objective increases, then abort
        if next_obj > current_obj + tol
            quiet || print_descent_error(iter, loss, next_loss)
            throw(ErrorException("Descent failure!"))
            #return IHTResults(-one(T), convert(T, -Inf), -1, [convert(T, -Inf)])
        end
    end # end main loop

    # return null result
    return IHTResults(zero(T), zero(T), 0, [zero(T)])
end # end function

"""
COMPUTE AN IHT REGULARIZATION PATH FOR LEAST SQUARES REGRESSION

    iht_path(x,y,path) -> SparseCSCMatrix

This subroutine computes best linear models for matrix `x` and response `y` by calling `L0_reg` for each model over a regularization path denoted by `path`.

Arguments:

- `x` is the `n` x `p` design matrix.
- `y` is the `n`-vector of responses.
- `path` is an `Int` vector that contains the model sizes to test.

Optional Arguments:

- `tol` is the global convergence tolerance for `L0_reg`. Defaults to `1e-4`.
- `max_iter` caps the number of iterations for the algorithm. Defaults to `1000`.
- `max_step` caps the number of backtracking steps in the IHT kernel. Defaults to `50`.
- `quiet` is a Boolean that controls the output. Defaults to `true` (no output).

Output:

- a sparse `p` x `length(path)` matrix where each column contains the computed model for each component of `path`.
"""
function iht_path{T <: Float}(
    x        :: DenseMatrix{T},
    y        :: DenseVector{T},
    path     :: DenseVector{Int};
    tol      :: T    = convert(T, 1e-4),
    max_iter :: Int  = 100,
    max_step :: Int  = 50,
    quiet    :: Bool = true
)

    # size of problem?
    (n,p) = size(x)

    # how many models will we compute?
    num_models = length(path)
    #nk = sum(path) # total number of elements in models

    # preallocate space for intermediate steps of algorithm calculations
    v  = IHTVariables(x, y, 1)
    betas = spzeros(T, p, num_models)  # a matrix to store calculated models
    #I = zeros(Int, nk)
    #J = zeros(Int, nk)
    #K = zeros(T,   nk)
    #q0 = 0

    # compute the path
    for i = 1:num_models

        # model size?
        q = path[i]

        # monitor progress
        quiet || print_with_color(:blue, "Computing model size $q.\n\n")

        # store projection of beta onto largest k nonzeroes in magnitude
        project_k!(v.b, q)

        # these arrays change in size from iteration to iteration
        # we must allocate them for every new model size
        update_variables!(v, x, q)

        # now compute current model
        output = L0_reg(x, y, q, v=v, tol=tol, max_iter=max_iter, max_step=max_step, quiet=quiet)

        # extract and save model
        # strangely more efficient to continuously reallocate sparse matrix
        # than to form it directly
        #I[(q0+1):(q0+q)] = find(v.idx)
        #J[(q0+1):(q0+q)] = i
        #K[(q0+1):(q0+q)] = v.b[v.idx]
        #q0 += q
        betas[:,i] = sparsevec(output.beta)
    end

    # return a sparsified copy of the models
    #betas = sparse(I, J, K, p, num_models, false)
    return betas
end


export rec2coeff, showpoly, showbasis

title_color =   :underline
name_color  =   :light_blue

function show(measure::AbstractMeasure)
    printstyled("Measure dλ(t)=w(t)dt\n";color=title_color)
    for name in fieldnames(typeof(measure))
        println("$(string(name)):\t$(getfield(measure, name))")
    end
end

function show(productmeasure::ProductMeasure)
    for measure in productmeasure.measures
        show(measure)
    end
end

function show(op::AbstractOrthoPoly; showmeasure::Bool=true)
    printstyled("\nUnivariate orthogonal polynomials\n";color=title_color)
    # print("name:"); printstyled("\t\t$(op.name)\n";color=name_color)
    print("degree:"); printstyled("\t\t$(op.deg)\n")
    print("#coeffs:"); printstyled("\t$(length(op.α))\n")
    length(op.α)<=7 ? n=length(op.α) : n=7
    print("α ="); printstyled("\t\t$(op.α[1:n])")
    n<length(op.α) ? print("...\n") : print("\n")
    print("β ="); printstyled("\t\t$(op.β[1:n])")
    n<length(op.α) ? print("...\n") : print("\n")
    if showmeasure
        print("\n")
        show(op.measure)
    end
end

function show(mop::MultiOrthoPoly; showmeasure::Bool=false)
    p = length(mop.name)
    printstyled("\n$p-variate orthogonal polynomials\n"; color = title_color)
    print("name:");
    [ printstyled("\t\t$(mop.name[i])\n"; color = name_color) for i = 1:p ]
    print("deg:"); printstyled("\t\t$(mop.deg)\n")
    print("dim:"); printstyled("\t\t$(mop.dim)\n")
    print("ind:");
    mop.dim <= 7 ? n = mop.dim : n = 7
    [ printstyled("\t\t$(mop.ind[i,:])\n") for i = 1:n ]
    n < mop.dim ? print("\t\t...\n") : print("\n")
    print("\n")
    showmeasure && show(mop.measure)
end

function show(t::Tensor)
    printstyled("\n$(t.dim)-dimensional tensor\n"; color = title_color)
    print("dim:"); printstyled("\t\t$(t.dim)\n")
    print("nonzeros:"); printstyled("\t$(length(t.T.nzind))\n")
    show(t.op; showmeasure = false)
end

"""
```
rec2coeff(deg::Int,a::Vector{<:Real},b::Vector{<:Real})
rec2coeff(a,b) = rec2coeff(length(a),a,b)
```
Get the coefficients of the orthogonal polynomial of degree up to `deg` specified by its
recurrence coefficients `(a,b)`. The function returns the values ``c_i^{(k)}`` from
```math
p_k (t) = t^d + \\sum_{i=0}^{k-1} c_i t^i,
```
where ``k`` runs from `1` to `deg`.

The call `rec2coeff(a,b)` outputs all possible recurrence coefficients given `(a,b)`.
"""
function rec2coeff(deg::Int,a::Vector{<:Real},b::Vector{<:Real})
    deg <= 0 && throw(DomainError(deg, "degree must be positive (you asked for degree = $deg"))
    !(length(a) == length(b) && length(a) >= deg) && throw(InconsistencyError("incorrect number of recurrence coefficients"))
    c = Vector{Vector{<:Real}}(undef,deg)
    @inbounds c[1] = [ -a[1] ]
    deg == 1 && return c
    @inbounds c[2] = [ a[1]*a[2] - b[2], -a[1] - a[2] ]
    deg == 2 && return c
    for k in 2:deg-1
        c[k+1] = [ -a[k+1]*c[k][1] - b[k+1]*c[k-1][1] ]
        for i = 1:k-2
            push!(c[k+1],c[k][i] - a[k+1]*c[k][i+1] - b[k+1]*c[k-1][i+1])
        end
        push!(c[k+1],c[k][k-1] - a[k+1]*c[k][k] - b[k+1])
        push!(c[k+1],c[k][k] - a[k+1])
    end
    return c
end
rec2coeff(α,β) = rec2coeff(length(α),α,β)

"""
```
showpoly(coeffs::Vector{<:Real};sym::String,digits::Integer)
```
Show the monic polynomial with coefficients `coeffs` in a human readable way.
They keyword `sym` sets the name of the variable, and `digits` controls the number of shown digits.
```jldoctest
julia> using PolyChaos

julia> showpoly([1.2, 2.3, 3.4456])
x^3 + 3.45x^2 + 2.3x + 1.2
julia> showpoly([1.2, 2.3, 3.4456], sym="t", digits=2)
t^3 + 3.45t^2 + 2.3t + 1.2
```

```
showpoly(d::Integer,α::Vector{<:Real},β::Vector{<:Real}; sym::String,digits::Integer)
showpoly(d::Range,α::Vector{<:Real},β::Vector{<:Real};sym::String,digits::Integer) where Range <: OrdinalRange
```
Show the monic polynomial of degree/range `d` that has the recurrence coefficients `α`, `β`.
```jldoctest
julia> using PolyChaos

julia> α, β = rm_hermite(10)
([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.77245, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5])
julia> showpoly(3,α,β)
x^3 - 1.5x

julia> showpoly(0:2:10,α,β)
1
x^2 - 0.5
x^4 - 3.0x^2 + 0.75
x^6 - 7.5x^4 + 11.25x^2 - 1.88
x^8 - 14.0x^6 + 52.5x^4 - 52.5x^2 + 6.56
x^10 - 22.5x^8 + 157.5x^6 - 393.75x^4 + 295.31x^2 - 29.53
```

Tailored to types from `PolyChaos.jl`
```
showpoly(d::Union{Integer,Range},op::AbstractOrthoPoly;sym::String,digits::Integer) where Range <: OrdinalRange
```
Show the monic polynomial of degree/range `d` of an `AbstractOrthoPoly`.
```jldoctest
julia> using PolyChaos

julia> op = GaussOrthoPoly(5);

julia> showpoly(3,op)
x^3 - 3.0x

julia> showpoly(0:op.deg,op; sym="t")
1
t
t^2 - 1.0
t^3 - 3.0t
t^4 - 6.0t^2 + 3.0
t^5 - 10.0t^3 + 15.0t
```
[Thanks @pfitzseb for providing this functionality.](https://discourse.julialang.org/t/how-to-define-verbose-output-for-a-polynomial/21317/5)

"""
function showpoly(coeffs::Vector{<:Real};sym::String="x",digits::Integer=2)
    io = Base.stdout
    length(coeffs) > 1 ? print(io, sym*"^", length(coeffs)) : print(io, sym)
    for (i, c) in enumerate(reverse(coeffs))
        abs(round(c, digits=digits)) == 0. && continue
        ex = length(coeffs) - i
        print(io, ' ', c > 0 ? '+' : '-', ' ')
        print(io, abs(round(c, digits=digits)))
        ex > 0 && print(io, sym)
        ex > 1 && print(io, '^', ex)
    end
    print(io, '\n')
end

function showpoly(d::Integer,α::Vector{<:Real},β::Vector{<:Real}; sym::String="x",digits::Integer=2)
    @assert d >= 0 "degree has to be non-negative."
    d == 0 && return print("1\n")
    showpoly(rec2coeff(d,α,β)[end],sym=sym,digits=digits)
end

function showpoly(d::Range,α::Vector{<:Real},β::Vector{<:Real};sym::String="x",digits::Integer=2) where Range <: OrdinalRange
    map(c->showpoly(c,α,β;sym=sym,digits=digits),d)
    print()
end

showpoly(d::Union{Integer,Range},op::AbstractOrthoPoly;sym::String="x",digits::Integer=2) where Range <: OrdinalRange = showpoly(d,op.α,op.β;sym=sym,digits=digits)

"""
```
showbasis(α::Vector{<:Real},β::Vector{<:Real};sym::String,digits::Integer)
```
Show all basis polynomials given the recurrence coefficients `α`, `β`.
They keyword `sym` sets the name of the variable, and `digits` controls the number of shown digits.
```jldoctest
julia> using PolyChaos

julia> α, β = rm_hermite(5);

julia> showbasis(α,β)
1
x
x^2 - 0.5
x^3 - 1.5x
x^4 - 3.0x^2 + 0.75
x^5 - 5.0x^3 + 3.75x
```

Tailored to types from `PolyChaos.jl`
```
showbasis(op::AbstractOrthoPoly;sym::String,digits::Integer)
```
Show all basis polynomials of an `AbstractOrthoPoly`.
```jldoctest
julia> using PolyChaos

julia> op = LegendreOrthoPoly(4);

julia> showbasis(op)
1
x
x^2 - 0.33
x^3 - 0.6x
x^4 - 0.86x^2 + 0.09
```
"""
function showbasis(α::Vector{<:Real},β::Vector{<:Real};sym::String="x",digits::Integer=2)
    showpoly(0:length(α),α,β;sym=sym,digits=digits)
end
showbasis(op::AbstractOrthoPoly;sym::String="x",digits::Integer=2) = showbasis(op.α[1:op.deg],op.β[1:op.deg];sym=sym,digits=digits)


# NMF for dense matrices

using NMF

function run(algname)

    # prepare data
    p = 8
    k = 5
    n = 100

    Wg = abs(randn(p, k))
    Hg = abs(randn(k, n))
    X = Wg * Hg + 0.1 * randn(p, n)

    # run NNMF
    println("Algorithm: $(algname)")
    println("---------------------------------")

    r = nnmf(X, k;
             init=:nndsvdar,
             alg=symbol(algname),
             maxiter=30,
             verbose=true)

    # display results
    println("numiters  = $(r.niters)")
    println("converged = $(r.converged)")
    @printf("objvalue  = %.6e\n", r.objvalue)
    println("W matrix = ")
    NMF.printf_mat(r.W)

    println()
end

function print_help()
    println("Usage:")
    println()
    println("  julia densenmf.jl <alg>")
    println()
    println("  <alg> is the name of the chosen algorithm, which can be ")
    println()
    println("    multmse:   Multiplicative update (minimize MSE)")
    println("    multdiv:   Multiplicative update (minimize divergence)")
    println("    projals:   Projected ALS")
    println("    alspgrad:  ALS Projected Gradient Descent")
    println()
end

function main(args)
    if isempty(args)
        print_help()
    elseif length(args) == 1
        a = lowercase(args[1])
        if a == "-h" || a == "--help"
            print_help()
        else
            run(a)
        end
    else
        @warn("Invalid command line arguments.")
        print_help()
    end
end

main(ARGS)


abstract type HamiltonianModel end

@doc raw"
    hamiltonian(model<:HamiltonianModel)

return the hamiltonian of the `model` as a two-site tensor operator.
"
function hamiltonian end

struct Ising <: HamiltonianModel end

@doc raw"
    TFIsing(hx::Real)

return a struct representing the transverse field ising model with magnetisation `hx`.
"
struct TFIsing{T<:Real} <: HamiltonianModel
    hx::T
end

"""
    hamiltonian(model::TFIsing)

return the transverse field ising hamiltonian for the provided `model` as a
two-site operator.
"""
function hamiltonian(model::TFIsing)
    hx = model.hx
    -2 * ein"ij,kl -> ijkl"(σz,σz) -
        hx/2 * ein"ij,kl -> ijkl"(σx, id2) -
        hx/2 * ein"ij,kl -> ijkl"(id2, σx)
end

@doc raw"
    Heisenberg(Jz::T,Jx::T,Jy::T) where {T<:Real}

return a struct representing the heisenberg model with magnetisation fields
`Jz`, `Jx` and `Jy`..
"
struct Heisenberg{T<:Real} <: HamiltonianModel
    Jz::T
    Jx::T
    Jy::T
end
Heisenberg() = Heisenberg(1.0,1.0,1.0)

"""
    hamiltonian(model::Heisenberg)

return the heisenberg hamiltonian for the `model` as a two-site operator.
"""
function hamiltonian(model::Heisenberg)
    h = model.Jz * ein"ij,kl -> ijkl"(σz,σz) -
        model.Jx * ein"ij,kl -> ijkl"(σx, σx) -
        model.Jy * ein"ij,kl -> ijkl"(σy, σy)
    h = ein"ijcd,kc,ld -> ijkl"(h,σx,σx')
    real(h ./ 2)
end


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

# returns f(x₀)
function evaluate_2(f::PolyElem, x₀)
    accum = convert(Int, lead(f))
    d = degree(f)
    reconstruct!(x₀)

    rec_x = rational_reconstruction(x₀)
    csr_x = sparse(to_dense(rec_x))

    I = sparse(to_dense(one(rec_x)))
    for i in 1 : d
        accum = accum*csr_x + convert(Int, coeff(f, d - i))*I
    end

    accum = from_dense(Array(accum), Nemo.QQ)

    return reconstruct!(modular_reduction(accum, base_ring(x₀)))
end

# returns f(x₀)
# O(d) multiplications/additions of x₀ if d = degree(f)
function evaluate(f::PolyElem, x₀)
    accum = lead(f)
    d = degree(f)
    reconstruct!(x₀)

    I = one(x₀)
    for i in 1 : d
        accum = accum * x₀ + coeff(f, d - i) * I
    end

    return reconstruct!(accum)
end

# returns f(x₀) * b
# expanding the brackets into:
#   f₀b + f₁x₀b + f₂x₀²b + ... + fₙx₀ⁿb
# and using the Horner scheme
function evaluate(f::PolyElem, x₀, b)
    E = one(x₀)
    accum = lead(f) * E * b
    d = degree(f)

    for i in 1 : d
        accum = x₀ * accum + coeff(f, d - i) * E * b
    end

    return accum
end

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

function shift_right_x(f::PolyElem)
    x = gen(parent(f))
    divexact(f - coeff(f, 0), x)
end

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

# Solves Ay = b
# Throws if A is singular
function square_nonsingular_randomized_wiedemann(A::AbstractSparseMatrix, b::AbstractSparseVector)
    # the function implements Algorithm 1 given on page 55
    # in https://doi.org/10.1109/TIT.1986.1057137
    # the notation and step numbering are taken from there

    field = base_ring(A)
    S, x = PolynomialRing(field, "x")
    n = order(A)

    # 1
    b₀ = b
    y = zero(b)
    k = 0
    d = 0

    subspace = [ b ]
    for i in 1 : 2n - 1
        push!(subspace, A * last(subspace))
    end

    # 2
    while !iszero(b)
        # 3
        u = random_sparsik(n, field, density=1)

        # 4
        seq = elem_type(field)[]
        for i in 1 : 2 * (n - d)
            push!(seq, inner(u, subspace[i]))
        end

        # 5
        f = minimal_polynomial(S(seq), x^(2 * (n - d)))

        if coeff(f, 0) == 0
            throw(SingularException(0))
        end

        f = f * inv(coeff(f, 0))

        # 6
        f⁻ = shift_right_x(f)

        accum = zero_sparsik(n, field)
        for j in 0 : degree(f⁻)
            reduce!(accum, subspace[j + 1], coeff(f⁻, j))
        end
        reduce!(y, accum, 1)
        b = b₀ + A * y
        d += degree(f)
        k += 1

        if d == n
            break
        end
    end

    return -y
end

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

# applies the polynomial `g` to the sequence `seq`
# by the definition from the 'Modern Computer Algebra'
#
# O(n^2) if n is the length of `seq`
# Actually, its O(nlogn): libflint implements FFT
function apply_polynomial(seq, g)
    S = parent(g)
    n = length(seq)
    d = degree(g)
    product = seq * reverse(g)
    S([coeff(product, i) for i in d : (2n - 1)])
end

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

function square_nonsingular_deterministic_wiedemann(A::AbstractSparseMatrix, b::AbstractSparseVector)
    # the function implements Algorithm 2 given on page 55
    # in https://doi.org/10.1109/TIT.1986.1057137
    # the notation and step numbering are taken from there

    field = base_ring(A)
    S, x = PolynomialRing(field, "x")
    n = order(A)

    # 1
    # O(nω) if ω is the number of nonzeroes in A
    subspace = [ b ]
    for i in 1 : 2n - 1
        push!(subspace, A * last(subspace))
    end

    # 2
    k = 0
    g = one(S)

    while k < n && degree(g) < n
        # 3
        u = unit_sparsik(n, k + 1, field)

        # 4
        # O(n)
        seq = [inner(x, u) for x in subspace]

        # 5, O(nlogn)
        d = degree(g)
        h = apply_polynomial(S(seq), g)

        # 6, O(n^2)
        f = minimal_polynomial(h, x^(2n - d))

        # 7, O(nlogn)
        g *= f

        # 8
        k += 1
    end

    if iszero(coeff(g, 0)) || isone(g)
        throw(SingularException(0))
    end

    # 9
    f = g * inv(coeff(g, 0))
    f⁻ = shift_right_x(f)

    accum = zero_sparsik(n, field)
    for j in 0 : degree(f⁻)
        reduce!(accum, subspace[j + 1], coeff(f⁻, j))
    end

    return -accum
end

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

# Returns a solution x to linear system Ax = b
# under an assumtion of A being square and nonsingular
# using Wiedemann system-solving algorithm as a base
# Throws if system is singular

"""
wiedemann_solve(A, b; proved)

Returns a solution x to linear system Ax = b under an assumtion of
A being square and nonsingular using Wiedemann system-solving algorithm.
By default, proved = true which calls 'square_nonsingular_deterministic_wiedemann'
otherwise calls 'square_nonsingular_randomized_wiedemann'.
"""
function wiedemann_solve(
        A::AbstractSparseMatrix,
        b::AbstractSparseVector;
        proved=true)

    if issquare(A)
        if proved
            return square_nonsingular_deterministic_wiedemann(A, b)
        else
            return square_nonsingular_randomized_wiedemann(A, b)
        end
    end

    error("nonsquare matrices are not supported")
end

# Returns the minimal annihilating polynomial of matrix A, that is,
# f ∈ F[x], such that
#   f(A)  = 0
#   f     = argmin(degree(f))
#
# The algorithm is randomized if subspace_minpoly is of W. family
# We should probably do some probabilistic analysis
# (or peek it at the W. paper)
function minimal_polynomial_wiedemann(
        A::AbstractSparseMatrix;
        subspace_minpoly=__deterministic_wiedemann_minpoly)

    S, _ = PolynomialRing(base_ring(A), "x")
    f = S(1)

    iterations = 4
    for _ in 1 : iterations
        f = lcm(f, subspace_minpoly(A, S))
    end

    return f
end

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

# Deterministic algorithm
# for minimal annihilating polynomial of complex matrix A
#
# Sadly, the source thyself does not contain any references
# about this Algorithm complexity analysis
#
#   O(nL(n))   ( O(n^4) basically )
# where n is the order of matrix
# and L(n) is the cost of multiplying two matrices of n×n
#
# Throws if the given matrix is zero
#
# `PolySpace` - a univariate polynomial ring to put the resulting polynomial to
function __deterministic_simple_minpoly(A, PolySpace)
    # the function implements the algorithm presented in paper
    #   An algorithm for the calculation of the minimal polynomial,
    #   S. BIAŁAS and M. BIAŁAS
    # the paper is available at
    #   http://bulletin.pan.pl/(56-4)391.pdf

    if iszero(A)
        throw(DomainError(A, "It is really zero"))
    end

    field = base_ring(A)
    n = order(A)

    # Computing the linear span of the powers of A
    columns = [ one(A) ]
    for i in 1 : n
        push!(columns, A * last(columns))
    end

    columns = Dict(i => vec(x) for (i, x) in enumerate(columns))

    B = from_rows(n+1, n^2, field, Array(1:n+1), columns)
    B = transpose!(B)

    V = linear_span!(collect(values(B.rows)))

    rows = Dict(i => x for (i, x) in enumerate(basis(V)))
    C = from_rows(length(rows), n + 1, field, Array(1:length(rows)), rows)
    C = transpose!(C)

    # Computing the degree of the minimal polynomial
    V = Subspacik(field)
    power = 0
    for i in 1:size(C, 1)
        eatcode = eat_sparsik!(V, deepcopy(C.rows[i]))
        if !haskey(C.rows, i) || eatcode == reduced
            power = i
            break
        end
    end

    # Computing the minimal polynomial
    x = gen(PolySpace)
    f = x^(power-1)
    last_row = haskey(C.rows, power) ? C.rows[power] : zero_sparsik(n + 1, field)
    for (piv, val) in last_row
        if haskey(V.echelon_form, piv)
            f -= val * x^(piv-1)
        end
    end

    return f
end

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

# Single deterministic iteration of minimal_polynomial algorithm
# Returns annihilating polynomial for {Aⁱb} for some random vector b
function __deterministic_wiedemann_minpoly(A::AbstractSparseMatrix, PolySpace)
    field = base_ring(A)
    n = order(A)
    b = random_sparsik(n, field, density=1.0)

    subspace = [ b ]
    for i in 1 : 2n - 1
        push!(subspace, A * last(subspace))
    end

    k = 0
    g = one(PolySpace)
    while k < n && degree(g) < n
        u = unit_sparsik(n, k + 1, field)

        # O(n)
        seq = [inner(x, u) for x in subspace]

        # O(nlogn)
        d = degree(g)
        h = apply_polynomial(PolySpace(seq), g)

        # O(n^2)
        f = minimal_polynomial(h, gen(PolySpace)^(2n - d))

        # O(nlogn)
        g *= f
        k += 1
    end

    return g
end

# Single randomized iteration of minimal_polynomial algorithm
# Returns annihilating polynomial for {Aⁱb} for some random vector b
function __randomized_wiedemann_minpoly(A::AbstractSparseMatrix, PolySpace)
    field = base_ring(A)
    n = order(A)
    b = random_sparsik(n, field, density=1.0)
    g = PolySpace(1)

    subspace = [ b ]
    for i in 1 : 2n - 1
        push!(subspace, A * last(subspace))
    end

    while degree(g) < n
        u = random_sparsik(n, field, density=1)

        seq = elem_type(field)[]
        for i in 1 : 2 * (n - degree(g))
            push!(seq, inner(u, subspace[i]))
        end

        f = PolySpace(0)
        try
            f = minimal_polynomial(
                PolySpace(seq),
                gen(PolySpace)^(2 * (n - degree(g)))
            )
        catch SingularException
            # singular sequence is expected here, so do not throw
            return g
        end

        g *= f
    end

    return g
end


using Distributions
using QuadGK
using Plots
using HCubature

unicodeplots()

function gaussian(x; σ=1, μ=0)
   y = @. exp(-(x-μ)^2/2)/sqrt(2π)/σ
   return y
end

function transform(f)
    function g(t)
        return f(t/(1-t^2))*(1+t^2)/(1-t^2)/(1-t^2)
    end
    return g
end

function integral(f, a, b)
    if a == -Inf && b == Inf
        println("Use inf case")
        g = transform(f)
        return quadgk(g, -1, 1)
    else
        return quadgk(f, a, b)
    end
end

function plotgaussian()
    xs = -3:0.01:3
    ys = gaussian(xs)
    plot(xs, ys)
end

function integraloverinf()
    integral(gaussian, -Inf, Inf)
end

struct Parameter
    a
    b
    Parameter(; a,b)=new(a,b)
end

function p(x; w)
    t -> (1-w.a)*gaussian(t) + a * gaussian(x-w.b)
end

function post_distribution(p, φ, w, X; beta=1.)
    f(w)=x->φ(w)*prod(p(X[i], w=Parameter(a=w[1], b=w[2]))^β for i in 1:length(X))
    res=HCubature(f,[0,-5],[1,5])
end

function main()
x    w = Parameter(0.5, 3.0)
    p(x, w=w)
end

#main()


# coding: utf-8

# In[1]:

#merge sort is so called advanced sort
#its like a binary search in a way
#we keep dividing the list into two parts
#until it reaches base case when there are only two elements
#we sort the two elements and combine two sublists together
#we do a traversal on both sublists
#we run through elements one by one
#we compare two lists with each other and insert into the bigger list
#its like shuffling poker deck
#say u got 1 and 2 from two different sublists
#1 is smaller than 2
#we insert 1 back to the bigger list
#the second element is 1.5
#it is still smaller than 2
#we keep inserting elements from a sorted sublist
#until the element is larger than 2
#now we insert 2 into the bigger list
#we keep doing the same procedure until we come back to the whole list
#we do a final sort, its all done

#from my description, u can easily tell that we need recursion
#i assume that is why its slowing down the sorting

# In[2]:

function merge_sort(arr)

    #so this is to check the base case
    #we wanna make sure the base case is a sublist with two elements
    #otherwise we got nothing to sort
    if length(arr)>1

        #this is how we keep splitting the list into two halves
        ind=Int32(floor(length(arr)/2))

        #we define each half as an individual list
        #then we run the same trick on each half recursively
        left=arr[1:ind]
        right=arr[ind+1:end]
        left=merge_sort(left)
        right=merge_sort(right)

        #let me denote ind_left as the index in left list
        #ind_right as the index in right list
        #ind_combined as the index in right and left combined list
        ind_left=1
        ind_right=1
        ind_combined=1

        #now we have three scenarios
        #the first case is both indexes are under list length
        #so we can compare two elements
        #we put the larger one into a bigger list
        #assuming an item from right side goes to a bigger list
        #we increase the right side index by 1
        #and come back to the loop to see which element is bigger
        while (ind_left<=length(left)) && (ind_right<=length(right))
            if left[ind_left]<right[ind_right]
                arr[ind_combined]=left[ind_left]
                ind_left+=1
                ind_combined+=1
            else
                arr[ind_combined]=right[ind_right]
                ind_right+=1
                ind_combined+=1
            end
        end

        #the second case is when one index is outta list length
        #the other isnt
        #it can occur under two scenarios
        #one is that we use // method to divide a list into half
        #for a list with an odd number of elements
        #we definitely have one side bigger than the other
        #the other scenario is more common
        #assume it is an extreme case
        #we have two list, left and right
        #all the elements in left is smaller than the smallest element in right
        #so we have placed all elements in left into a bigger list
        #the index of right list hasnt increased even one bit
        #thats when we use the below section to add elements in right into a bigger list
        while (ind_left<=length(left)) && (ind_right>length(right))
            arr[ind_combined]=left[ind_left]
            ind_left+=1
            ind_combined+=1
        end

        #vice versa
        while (ind_left>length(left)) && (ind_right<=length(right))
            arr[ind_combined]=right[ind_right]
            ind_right+=1
            ind_combined+=1
        end
    end

    return arr

end

# In[3]:

for _ in 1:100

    test_arr=rand(1000)

    if !(merge_sort(test_arr)==sort(test_arr))
        printstyled("Erreur",color=:red)
    end

end


module star_chart

export total_orbits, total_transfers

using LightGraphs

function total_orbits(orbit_map)
    os = parse_map(orbit_map)
    (ps, g) = graph_orbits(os)
    return count_orbits(ps, g)
end

function total_transfers(orbit_map)
    os = parse_map(orbit_map)
    (ps, gr) = graph_orbits(os)
    me = orbiting(gr, ps["YOU"])
    santa = orbiting(gr, ps["SAN"])
    return length(a_star(gr, me, santa))
end

function orbiting(gr, body)
    return first(neighbors(gr, body))
end

function parse_map(orbits)
    return map(
      orbit -> split(orbit, ")"),
      orbits
    )
end

function graph_orbits(orbits)
    orbit_map = Graph()
    bodies = Dict{String, Int}()
    for orbit in orbits
        (parent, child) = orbit
        vp = assign_id_to_body(bodies, orbit_map, parent)
        vc = assign_id_to_body(bodies, orbit_map, child)
        add_edge!(orbit_map, vc, vp)
    end
    return (bodies, orbit_map)
end

function assign_id_to_body(bodies, orbit_map, body)
    if haskey(bodies, body)
        return bodies[body]
    else
        add_vertices!(orbit_map, 1)
        id = nv(orbit_map)
        bodies[body] = id
        return id
    end
end

function count_orbits(bodies, orbit_map)
    origin = bodies["COM"]
    count = 0
    for body in bodies
        (_, id) = body
        count += length(a_star(orbit_map, id, origin))
    end
    return count
end

end


##k-mean++

#Square Euclidean Function to use (internal sql2() or SqEuclidean() from Distances.jl)
SqEuDist = SqEuclidean
# SqEuDist = sql2

#Return distance to closest centre for each data point (helper function for k-mean++)
"""
    mindist2cntrs(centres, data)

Computes the square of the L2 distance to closest centre for each data point.
* `data` must be a matrix of `Float` with each column representing a point
* `centres` must be a matrix of `Float` with each column representing a centre

!!! note
    `mindist2cntrs(centres, data)` does not check the whether the dimensions of the input match.
    Users must ensure that `size(data,1) = size(centres,1)`.
"""
function mindist2cntrs(cntrs::AbstractMatrix{T}, data::AbstractMatrix{T}) where T<:AbstractFloat
    mdist = Vector{T}(undef,size(data,2))
    temp = Vector{T}(undef,size(cntrs,2))
    @inbounds @views for i ∈ eachindex(mdist)
        colwise!(temp,SqEuDist(), data[:,i], cntrs)
        mdist[i] = minimum(temp)
    end
    return mdist
end

#k-means++ algorithm for initialising centres
"""
    kmeanspp(data, number_of_centres)

Indentfies initial cluster centres for k-means clustering using the k-means++ alogrithm for given `data`.
* `data` must be a matrix of `Float` with each column vector representing a point
* Function will return a vector of indices corresponding to columns of `data` identified as cluster centres
"""
function kmeanspp(data::AbstractMatrix{T}, ncntrs::Int) where T<:AbstractFloat
    cin = [rand(1:size(data,2))]
    din = collect(1:size(data,2))
    deleteat!(din,cin)
    while(length(cin)<ncntrs)
        mindist = mindist2cntrs(data[:,cin], data[:,din])
        cdist = pdist2cdist(mindist)
        ncin = findclosest(rand(),cdist)
        push!(cin, din[ncin])
        deleteat!(din,ncin)
    end
    return cin
end

## K-means

#Identify closest cluster centre
"""
    closest_centre!(closest_centre, centres, data)

Identifies the centre closest to each data point based on L2 distance.
* `data` must be a matrix of `Float` with each column representing a point
* `centres` must be a matrix of `Float` with each column representing a centre
* `closest_centre` must be a vector of `Int` with each element holding the column index of `centres` that corresponds to the closest centre

!!! note
    `closest_centre!(closest_centre, centres, data)` does not check the whether the dimensions of the input match.
    Users must ensure that `length(closest_centre) = size(data,2)` and `size(data,1) = size(centres,1)`.
"""
function closest_centre!(ccntr::AbstractVector{Int}, cntrs::AbstractMatrix{T}, data::AbstractMatrix{T}) where T<:AbstractFloat
    temp = Vector{T}(undef,size(cntrs,2))
    @inbounds @views for i ∈ eachindex(ccntr)
        colwise!(temp,SqEuDist(), data[:,i], cntrs)
        ccntr[i] = argmin(temp)
    end
end

"""
    closest_centre(centres, data)

Identifies the centre closest to each data point based on L2 distance.
* `data` must be a matrix of `Float` with each column representing a point
* `centres` must be a matrix of `Float` with each column representing a centre
* Output is a vector of `Int` with each element holding the column index of `centres` that corresponds to the closest centre

!!! note
    `closest_centre(centres, data)` does not check the whether the dimensions of the input match.
    Users must ensure that `size(data,1) = size(centres,1)`.
"""
function closest_centre(cntrs::AbstractMatrix{T}, data::AbstractMatrix{T}) where T<:AbstractFloat
    ccntr = Vector{Int}(undef,size(data,2))
    closest_centre!(ccntr, cntrs, data)
    return ccntr
end

#Maximum L2 distance by which a centre has shifted
function max_centre_shift(new_cntrs::AbstractMatrix{T}, old_cntrs::AbstractMatrix{T}) where T<:AbstractFloat
    s = zero(T)
    @views for i ∈ axes(old_cntrs,2)
        temp = evaluate(SqEuDist(),new_cntrs[:,i],old_cntrs[:,i])
        if(temp>s)
            s = temp
        end
    end
    return sqrt(s)
end

function max_centre_shift2(new_cntrs::AbstractMatrix{T}, old_cntrs::AbstractMatrix{T}) where T<:AbstractFloat
    temp = colwise(SqEuDist(),new_cntrs,old_cntrs)
    return sqrt(maximum(temp))
end

#k-means initialisation
abstract type KMinit end
struct kmpp <: KMinit end
struct kmrand <: KMinit end

"""
    kmeans_init([initialise_method], data, number_of_centres)

Pick a set of points from `data` to serve as the initial cluster centres for k-means clustering. Method for initialisation can random or k-means++.
* `data` must be a matrix of `Float` with each column representing a point
* Use `kmpp()` as `initialise_method` to use k-mean++ as initialisation method [DEFAULT if `initialise_method` is not specified]
* Use `kmrand()` as `intialise_method` to randomly select points as initial centres. This can be much faster than k-means++ for large datasets
"""
kmeans_init(::kmpp, data::AbstractMatrix{<:AbstractFloat}, ncntrs::Int) = data[:,kmeanspp(data,ncntrs)]

kmeans_init(::kmrand, data::AbstractMatrix{<:AbstractFloat}, ncntrs::Int) = data[:,rand(1:size(data,2),ncntrs)]

kmeans_init(data::AbstractMatrix{<:AbstractFloat}, ncntrs::Int) = kmeans_init(kmpp(), data, ncntrs)

#One iteration of updating k-means cluster centres
function kmeans_update!(ccntr::AbstractVector{Int}, cntrs::AbstractMatrix{T}, data::AbstractMatrix{T}) where T<:AbstractFloat
    closest_centre!(ccntr,cntrs,data)
    for i ∈ axes(cntrs,2)
        if(sum(ccntr.==i)>0)
            cntrs[:,i] = running_mean(data[:,ccntr.==i])
        end
    end
end

"""
    kmeans!(centres, data; [update_threshold = 0.0])

Updates (inplace) cluster centres `centres` using the k-means clustering algorithm based on `data`.
Cluster centres are iteratively updated until the largest shift in cluster centres,
measured in terms of square of L2 distance, falls below `update_threshold`.
* `centres` must be a matrix of `Float` with each column representing a cluster centre
* `data` must be a matrix of `Float` with each column representing a point
* Default value of optional named parameter is `update_threshold = 0.0` and corresponds to the setting that algorithm will converge when cluster centres do not shift
"""
function kmeans!(cntrs::AbstractMatrix{T}, data::AbstractMatrix{T}; update_threshold = 0.0) where T<:AbstractFloat
    ndims,ncntrs = size(cntrs)
    old_cntrs = Matrix{T}(undef,ndims,ncntrs)
    ccntr = Vector{Int}(undef,size(data,2))
    max_shift::T = Inf
    while(max_shift > update_threshold)
        copy!(old_cntrs,cntrs)
        kmeans_update!(ccntr, cntrs, data)
        max_shift = max_centre_shift(cntrs, old_cntrs)
    end
end

"""
    kmeans(data, number_of_centres; [update_threshold = 0.0])
    kmeans(data, number_of_centres, initialise_method; [update_threshold = 0.0])

Estimate cluster centres from `data` using the k-means clustering algorithm.
Cluster centres are iteratively updated until the largest shift in cluster centres,
measured in terms of square of L2 distance, falls below `update_threshold`.
* `data` must be a matrix of `Float` with each column representing a point
* Default value of optional named parameter is `update_threshold = 0.0` and corresponds to the setting that algorithm will converge when cluster centres do not shift
* Use `kmpp()` as `initialise_method` to use k-mean++ as initialisation method [DEFAULT if `initialise_method` is not specified]
* Use `kmrand()` as `intialise_method` to randomly select points as initial centres. This can be faster than k-means++ for large datasets
"""
function kmeans(data::AbstractMatrix{T}, ncntrs::Int, init_method::KMinit; update_threshold = 0.0) where T<:AbstractFloat
    cntrs = kmeans_init(init_method, data, ncntrs)
    kmeans!(cntrs, data; update_threshold)
    return cntrs
end

kmeans(data::AbstractMatrix{T}, ncntrs::Int; update_threshold = 0.0) where T<:AbstractFloat = kmeans(data, ncntrs, kmpp(); update_threshold)


#---------------------------------------------------------------------#
#This routine advances the solution in time using Split-Explicit using
#a simple General Order RK method.
#Written by F.X. Giraldo / P.R. Mugg on 9/20/19
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey, CA 93943-5216
#---------------------------------------------------------------------#
function WS_rk3(q0,c,Δt,time)

    #1st RK Stage
    Q = q0
    (rhs,f,g) = rhs_function(Q,c,time)
    Q = q0 + Δt/3*(f + g)

    #2nd RK Stage
    (rhs,f,g) = rhs_function(Q,c,time + Δt/3)
    Q = q0 + Δt/2*(f + g)

    #3rd RK Stage
    (rhs,f,g) = rhs_function(Q,c,time + Δt/2)
    Q = q0 + Δt*(f + g)

    #Update
    return qp = Q

end #end function


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

"""
    localanisotropies(Geometric, searcher; simplify=true)

Extract `LocalAnisotropy` from a searcher object. Based on the spatial data and
the parameters from the searcher object, a group of neighbor points is collected
at each place. PCA is applied to their spatial coordinates, returning eigenvectors
that represent a local ellipse/ellipsoid aligned with the directions with more
clustered points. The magnitude is the ratio of the PCA eigenvalues. It's useful
to extract local anisotropies from 3-D surface points. If `simplify = true` and
data is 3-D, only the minimum direction is extracted from PCA; the other axes
are calculated (main direction will be always along max dip direction and the
intermediate axis will be orthogonal to the others). This is default for 3-D.

## Example

```julia
searcher = KNearestSearch(pts3dsurface, 10)
prelpars = localanisotropies(Geometric, searcher) # extract local pars at surface
lpars = idwpars(prelpars, searcher, grid) # interpolate local pars to a grid
```
"""
function localanisotropies(::Type{Geometric}, searcher::NeighborSearchMethod;
	simplify::Bool=true)
	D = searcher.domain
	X = coords(D)
	N, len = size(X)

    quat = Array{Quaternion}(undef,len)
    m = Array{Vector}(undef,len)

    Threads.@threads for i in 1:len
		neighids = search(centroid(D,i), searcher)
		λ, v = pca(view(X,:,neighids), simplify)

		det(v) < 0 && (v = Diagonal(SVector{3}([-1,1,1])) * v)

        q = dcm_to_quat(v)
        quat[i] = q
        m[i] = λ/λ[1]
    end

	# deal with -1 eigvals
	m = reduce(hcat,vec(m))
	posm = m .> 0
	if sum(posm) < N*len
		for d in 2:N
			posd = view(posm,d,:)
			sum(posd) == len && continue
			minx = minimum(view(m,d,findall(posd)))
			m[d,findall(.!posd)] .= minx
		end
	end

    LocalAnisotropy(quat, m)
end

function pca(X, simplify)
	N = size(X,1)
	M = fit(PCA, X, maxoutdim=N, pratio=1)
	λ = principalvars(M) ./ principalvars(M)[1]
	nv = length(λ)
	v = N == 3 ? projection(M) : vcat(projection(M),[0 0 1][1:nv])

	if nv == 1
		vx = [-v[2,1]; v[1,1]; v[3,1]]
		v = hcat(v,vx,cross(v[:,1],vx))
		append!(λ,[λ[1],λ[1]])
	elseif nv == 2 && N==3
		v = hcat(v,cross(v[:,1],v[:,2]))
		push!(λ,minimum(λ))
	end

	if simplify && N==3
		# use maxdip vector and cross(maxdip,normal) as the main vectors
		abs(v[3,3]) > 1 && (v[3,3] = round(v[3,3], digits=1))
		az, dp = [atan(v[1,3], v[2,3]), -asin(v[3,3])]
		if dp < 0
			dp += pi/2
		else
			dp -= pi/2
		end
		v[:,1] .= [sin(az)*cos(dp), cos(az)*cos(dp), -sin(dp)]
		v[:,2] .= cross(v[:,1], v[:,3])
		λ[1:2] .= [1.0, 1.0]
	end

	λ[1:N], SMatrix{3,3}(v')
end


function quinn_fernandes(signal, seed)
	eps = 1.0e-6

	n = length(signal)
	a = 2.0*cos(seed)
	b = 0.1
	signal = vcat([0,0], signal)
	xi = zeros(size(signal))

	it = 0
	while (abs(a-b) > eps)
		it = it+1
		if (it > 1000)
			b = NaN
			break
		end
		a = 2*b - a
		for i=3:(n-1)
			xi[i] = signal[i] + a*xi[i-1] - xi[i-2]
		end
		b = sum([(xi[i]+xi[i-2])*xi[i-1] for i=3:(n-1)])/sum([xi[i-1]^2 for i=3:(n-1)])
	end

	acos(0.5*b)
end


function isprime(n)
    if n == 1
        return false
    elseif n == 2
        return true
    else
        for i = 2:(n ^ (0.5)) + 1
            if n % i == 0
                return false
            end
        end
        return true
    end
end

function prime_factors(n)
    count = 0
    final = [1]
    k = n
    if n == 1
        return []
    end
    while k != 1
        count += 1
        if k % count == 0 && isprime(count) == true
            k = k/count
            final = [final; [count]]
            count = 0
        end

    end
    return final[2:end]
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _sunpos(jd::AbstractFloat, radians::Bool)
    # Number of Julian centuries since 1899-12-31T12:00:00
    t = (jd - 2415020) / JULIANCENTURY
    # Sun's mean longitude
    l = (279.696678 + mod(36000.768925*t, 360)) * 3600
    # Allow for ellipticity of the orbit (equation of centre) using the Earth's
    # mean anomaly ME
    me = 358.475844 + mod(35999.049750*t, 360)
    ellcor = (6910.1 - 17.2 * t) * sind(me) + 72.3 * sind(2 * me)
    l += ellcor
    # Allow for the Venus perturbations using the mean anomaly of Venus MV
    mv = 212.603219 + mod(58517.803875*t, 360)
    vencorr = 4.8*cosd(299.1017 + mv - me) +
        5.5*cosd(148.3133 +  2 * mv - 2 * me) +
        2.5*cosd(315.9433 +  2 * mv - 3 * me) +
        1.6*cosd(345.2533 +  3 * mv - 4 * me) +
        1.0*cosd(318.15   +  3 * mv - 5 * me)
    l += vencorr
    # Allow for the Mars perturbations using the mean anomaly of Mars MM.
    mm = 319.529425  +  mod(19139.858500*t, 360)
    marscorr = 2 * cosd(343.8883 - 2 * mm + 2 * me) +
        1.8 * cosd(200.4017 - 2 * mm + me)
    l += marscorr
    # Allow for the Jupiter perturbations using the mean anomaly of Jupiter MJ
    mj = 225.328328 + mod(3034.6920239 * t, 360)
    jupcorr = 7.2*cosd(179.5317 - mj + me) +
        2.6*cosd(263.2167 - mj) +
        2.7*cosd( 87.1450 - 2 * mj + 2 * me) +
        1.6*cosd(109.4933 - 2 * mj + me)
    l += jupcorr
    # Allow for the Moons perturbations using the mean elongation of the Moon
    # from the Sun D
    d = 350.7376814 + mod(445267.11422 * t, 360)
    mooncorr  = 6.5*sind(d)
    l += mooncorr
    # Allow for long period terms
    longterm = 6.4*sind(231.19 + 20.20*t)
    l += longterm
    l  = mod(l + 2592000, 1296000)
    longmed = sec2rad(l)
    # Allow for Aberration
    l -= 20.5
    # Allow for Nutation using the longitude of the Moons mean node OMEGA
    sin_omega, cos_omega = sincos(deg2rad(259.183275 - mod(1934.142008 * t, 360)))
    l -= 17.2 * sin_omega
    # True Obliquity
    oblt = 23.452294 - 0.0130125 * t + 9.2 * cos_omega / 3600
    # Right Ascension and Declination
    l /= 3600
    oblt = deg2rad(oblt)
    sin_oblt, cos_oblt = sincos(oblt)
    sin_l, cos_l = sincos(deg2rad(l))
    ra = mod2pi(atan(sin_l * cos_oblt, cos_l))
    dec = asin(sin_l * sin_oblt)
    if radians
        return ra, dec, longmed, oblt
    else
        return rad2deg(ra), rad2deg(dec), rad2deg(longmed), rad2deg(oblt)
    end
end

"""
    sunpos(jd[, radians=true]) -> ra, dec, elong, obliquity

### Purpose ###

Compute the right ascension and declination of the Sun at a given date.

### Arguments ###

* `jd`: the Julian date of when you want to calculate Sun position.  It can be
  either a scalar or a vector.  Use `jdcnv` to get the Julian date for a given
  date and time.
* `radians` (optional boolean keyword): if set to `true`, all output quantities
  are given in radians.  The default is `false`, so all quantities are given in
  degrees.

### Output ###

The 4-tuple `(ra, dec, elong, obliquity)`:

* `ra`: the right ascension of the Sun at that date
* `dec`: the declination of the Sun at that date
* `elong`: ecliptic longitude of the Sun at that date
* `obliquity`: the obliquity of the ecliptic

All quantities are given in degrees, unless `radians` keyword is set to `true`
(see "Arguments" section).  If `jd` is an array, arrays of the same given as
`jd` are returned.

### Method ###

Uses a truncated version of Newcomb's Sun.  Adapted from the IDL routine SUN_POS
by CD Pike, which was adapted from a FORTRAN routine by B. Emerson (RGO).

### Example ###

(1) Find the apparent right ascension and declination of the Sun on May 1, 1982

```jldoctest
julia> using AstroLib

julia> adstring(sunpos(jdcnv(1982, 5, 1))[1:2], precision=2)
" 02 31 32.614  +14 54 34.92"
```

The Astronomical Almanac gives `02 31 32.58 +14 54 34.9` so the error for this
case is < 0.5".

(2) Plot the apparent right ascension, in hours, and declination of the Sun, in
degrees, for every day in 2016.  Use
[PyPlot.jl](https://github.com/JuliaPlots/Plots.jl/) for plotting.

```julia
using PyPlot
using Dates

days = DateTime(2016):Day(1):DateTime(2016, 12, 31);
ra, declin = sunpos(jdcnv.(days));
plot(days, ra/15); plot(days, declin)
```

### Notes ###

Patrick Wallace (Rutherford Appleton Laboratory, UK) has tested the accuracy of
a C adaptation of the present algorithm and found the following results.  From
1900-2100 `sunpos` gave 7.3 arcsec maximum error, 2.6 arcsec RMS.  Over the
shorter interval 1950-2050 the figures were 6.4 arcsec max, 2.2 arcsec RMS.

The returned `ra` and `dec` are in the given date's equinox.

Code of this function is based on IDL Astronomy User's Library.
"""
sunpos(jd::Real; radians::Bool=false) = _sunpos(float(jd), radians)

function sunpos(jd::AbstractArray{J}; radians::Bool=false) where {J<:Real}
    typej = float(J)
    ra = similar(jd, typej)
    dec = similar(jd, typej)
    longmed = similar(jd, typej)
    oblt = similar(jd, typej)
    for i in eachindex(jd)
        ra[i], dec[i], longmed[i], oblt[i] = sunpos(jd[i], radians=radians)
    end
    return ra, dec, longmed, oblt
end


module NormalHermiteSplines

export prepare, construct, interpolate
export evaluate, evaluate!, evaluate_one, evaluate_gradient, evaluate_derivative
export NormalSpline, RK_H0, RK_H1, RK_H2
export get_epsilon, estimate_epsilon, get_cond, estimate_cond, estimate_accuracy

using LinearAlgebra
using StaticArrays
using UnPack
using UnsafeArrays

const AbstractArrOfSVecs{n,T,N} = AbstractArray{SVector{n,T},N}
const AbstractVecOfSVecs{n,T} = AbstractVector{SVector{n,T}}
const VecOfSVecs{n,T} = Vector{SVector{n,T}}

@inline svectors(x::AbstractMatrix{T}) where {T} = reinterpret(reshape, SVector{size(x,1),T}, x)
@inline svectors(x::AbstractVector{T}) where {T} = reinterpret(SVector{1,T}, x)

include("./ReproducingKernels.jl")
include("./GramMatrix.jl")
include("./Splines.jl")
include("./Utils.jl")
include("./Interpolate.jl")

####
#### Public API
####

#### ReproducingKernel_0

"""
`prepare(nodes::AbstractMatrix{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}`

Prepare the spline by constructing and factoring a Gram matrix of the interpolation problem.
Initialize the `NormalSpline` object.
# Arguments
- `nodes`: The function value nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space and
           `n₁` is the number of function value nodes. It means that each column in the matrix defines one node.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: the partly initialized `NormalSpline` object that must be passed to `construct` function
        in order to complete the spline initialization.
"""
@inline function prepare(nodes::AbstractMatrix{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}
    return prepare(svectors(nodes), kernel)
end
@inline function prepare(nodes::AbstractVecOfSVecs{n,T}, kernel::RK = RK_H0()) where {n, T <: Real, RK <: ReproducingKernel_0}
    return _prepare(nodes, kernel)
end

"""
`construct(spline::AbstractNormalSpline{n,T,RK}, values::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_0}`

Construct the spline by calculating its coefficients and completely initializing the `NormalSpline` object.
# Arguments
- `spline`: the partly initialized `NormalSpline` object returned by `prepare` function.
- `values`: function values at `nodes` nodes.

Return: the completely initialized `NormalSpline` object that can be passed to `evaluate` function
        to interpolate the data to required points.
"""
@inline function construct(spline::AbstractNormalSpline{n,T,RK}, values::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return _construct!(deepcopy(spline), values)
end

"""
`interpolate(nodes::AbstractMatrix{T}, values::AbstractVector{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}`

Prepare and construct the spline.
# Arguments
- `nodes`: The function value nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space
           and `n₁` is the number of function value nodes.
           It means that each column in the matrix defines one node.
- `values`: function values at `nodes` nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: the completely initialized `NormalSpline` object that can be passed to `evaluate` function.
"""
@inline function interpolate(nodes::AbstractMatrix{T}, values::AbstractVector{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}
    return interpolate(svectors(nodes), values, kernel)
end
@inline function interpolate(nodes::AbstractVecOfSVecs{n,T}, values::AbstractVector{T}, kernel::RK = RK_H0()) where {n, T <: Real, RK <: ReproducingKernel_0}
    spline = _prepare(nodes, kernel)
    return _construct!(spline, values)
end

"""
`evaluate(spline::AbstractNormalSpline{n,T,RK}, points::AbstractMatrix{T}) where {n, T <: Real, RK <: ReproducingKernel_0}`

Evaluate the spline values at the locations defined in `points`.

# Arguments
- `spline: the `NormalSpline` object returned by `interpolate` or `construct` function.
- `points`: locations at which spline values are evaluating.
            This should be an `n×m` matrix, where `n` is dimension of the sampled space
            and `m` is the number of locations where spline values are evaluating.
            It means that each column in the matrix defines one location.

Return: `Vector{T}` of the spline values at the locations defined in `points`.
"""
@inline function evaluate(spline::AbstractNormalSpline{n,T,RK}, points::AbstractMatrix{T}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return evaluate(spline, svectors(points))
end
@inline function evaluate(spline::AbstractNormalSpline{n,T,RK}, points::AbstractArrOfSVecs{n,T}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return evaluate!(zeros(T, size(points)), spline, points)
end
@inline function evaluate!(spline_values::AbstractArray{T,N}, spline::AbstractNormalSpline{n,T,RK}, points::AbstractArrOfSVecs{n,T,N}) where {n, T <: Real, N, RK <: ReproducingKernel_0}
    return _evaluate!(spline_values, spline, points)
end
@inline function evaluate(spline::AbstractNormalSpline{n,T,RK}, point::SVector{n,T}) where {n, T <: Real, N, RK <: ReproducingKernel_0}
    return _evaluate(spline, point)
end

"""
`evaluate_one(spline::AbstractNormalSpline{n,T,RK}, point::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_0}`

Evaluate the spline value at the `point` location.

# Arguments
- `spline`: the `NormalSpline` object returned by `interpolate` or `construct` function.
- `point`: location at which spline value is evaluating.
           This should be a vector of size `n`, where `n` is dimension of the sampled space.

Return: the spline value at the location defined in `point`.
"""
@inline function evaluate_one(spline::AbstractNormalSpline{n,T,RK}, point::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return evaluate_one(spline, SVector{n,T}(ntuple(i -> point[i], n)))
end
@inline function evaluate_one(spline::AbstractNormalSpline{n,T,RK}, point::SVector{n,T}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return evaluate(spline, point)
end

"""
`evaluate_gradient(spline::AbstractNormalSpline{n,T,RK}, point::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_0}`

Evaluate gradient of the spline at the location defined in `point`.

# Arguments
- `spline`: the `NormalSpline` object returned by `interpolate` or `construct` function.
- `point`: location at which gradient value is evaluating.
           This should be a vector of size `n`, where `n` is dimension of the sampled space.

Note: Gradient of spline built with reproducing kernel RK_H0 does not exist at the spline nodes.

Return: `Vector{T}` - gradient of the spline at the location defined in `point`.
"""
@inline function evaluate_gradient(spline::AbstractNormalSpline{n,T,RK}, point::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return evaluate_gradient(spline, SVector{n,T}(ntuple(i -> point[i], n)))
end
@inline function evaluate_gradient(spline::AbstractNormalSpline{n,T,RK}, point::SVector{n,T}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return _evaluate_gradient(spline, point)
end

#### ReproducingKernel_1

"""
`prepare(nodes::AbstractMatrix{T}, d_nodes::AbstractMatrix{T}, d_dirs::AbstractMatrix{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}`

Prepare the spline by constructing and factoring a Gram matrix of the interpolation problem.
Initialize the `NormalSpline` object.
# Arguments
- `nodes`: The function value nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space and
           `n₁` is the number of function value nodes.
            It means that each column in the matrix defines one node.
- `d_nodes`: The function directional derivatives nodes.
             This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
             `n₂` is the number of function directional derivative nodes.
- `d_dirs`: Directions of the function directional derivatives.
        This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
        `n₂` is the number of function directional derivative nodes.
        It means that each column in the matrix defines one direction of the function directional derivative.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: the partly initialized `NormalSpline` object that must be passed to `construct` function
        in order to complete the spline initialization.
"""
@inline function prepare(nodes::AbstractMatrix{T}, d_nodes::AbstractMatrix{T}, d_dirs::AbstractMatrix{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}
    return prepare(svectors(nodes), svectors(d_nodes), svectors(d_dirs), kernel)
end
@inline function prepare(nodes::AbstractVecOfSVecs{n,T}, d_nodes::AbstractVecOfSVecs{n,T}, d_dirs::AbstractVecOfSVecs{n,T}, kernel::RK = RK_H1()) where {n, T <: Real, RK <: ReproducingKernel_1}
    return _prepare(nodes, d_nodes, d_dirs, kernel)
end

"""
`construct(spline::AbstractNormalSpline{n,T,RK}, values::AbstractVector{T}, d_values::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_1}`

Construct the spline by calculating its coefficients and completely initializing the `NormalSpline` object.
# Arguments
- `spline`: the partly initialized `NormalSpline` object returned by `prepare` function.
- `values`: function values at `nodes` nodes.
- `d_values`: function directional derivative values at `d_nodes` nodes.

Return: the completely initialized `NormalSpline` object that can be passed to `evaluate` function
        to interpolate the data to required points.
"""
@inline function construct(spline::AbstractNormalSpline{n,T,RK}, values::AbstractVector{T}, d_values::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_1}
    return _construct!(deepcopy(spline), values, d_values)
end

"""
`interpolate(nodes::AbstractMatrix{T}, values::AbstractVector{T}, d_nodes::AbstractMatrix{T}, d_dirs::AbstractMatrix{T}, d_values::AbstractVector{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}`

Prepare and construct the spline.
# Arguments
- `nodes`: The function value nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space
           and `n₁` is the number of function value nodes.
           It means that each column in the matrix defines one node.
- `values`: function values at `nodes` nodes.
- `d_nodes`: The function directional derivative nodes.
            This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
            `n₂` is the number of function directional derivative nodes.
- `d_dirs`: Directions of the function directional derivatives.
       This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
       `n₂` is the number of function directional derivative nodes.
       It means that each column in the matrix defines one direction of the function directional derivative.
- `d_values`: function directional derivative values at `d_nodes` nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: the completely initialized `NormalSpline` object that can be passed to `evaluate` function.
"""
@inline function interpolate(nodes::AbstractMatrix{T}, values::AbstractVector{T}, d_nodes::AbstractMatrix{T}, d_dirs::AbstractMatrix{T}, d_values::AbstractVector{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}
    return interpolate(svectors(nodes), values, svectors(d_nodes), svectors(d_dirs), d_values, kernel)
end
@inline function interpolate(nodes::AbstractVecOfSVecs{n,T}, values::AbstractVector{T}, d_nodes::AbstractVecOfSVecs{n,T}, d_dirs::AbstractVecOfSVecs{n,T}, d_values::AbstractVector{T}, kernel::RK = RK_H1()) where {n, T <: Real, RK <: ReproducingKernel_1}
    spline = _prepare(nodes, d_nodes, d_dirs, kernel)
    return _construct!(spline, values, d_values)
end

#### ReproducingKernel_0 (1-dimensional case)

"""
`prepare(nodes::AbstractVector{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}`

Prepare the 1D spline by constructing and factoring a Gram matrix of the interpolation problem.
Initialize the `NormalSpline` object.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n₁` vector where `n₁` is the number of function value nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: the partly initialized `NormalSpline` object that must be passed to `construct` function
        in order to complete the spline initialization.
"""
@inline function prepare(nodes::AbstractVector{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}
    return prepare(svectors(nodes), kernel)
end

"""
`interpolate(nodes::AbstractVector{T}, values::AbstractVector{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}`

Prepare and construct the 1D spline.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n₁` vector where `n₁` is the number of function value nodes.
- `values`: function values at `n₁` interpolation nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: the completely initialized `NormalSpline` object that can be passed to `evaluate` function.
"""
@inline function interpolate(nodes::AbstractVector{T}, values::AbstractVector{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}
    interpolate(svectors(nodes), values, kernel)
end

"""
`evaluate(spline::AbstractNormalSpline{n,T,RK}, points::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_0}`

Evaluate the 1D spline values/value at the `points` locations.

# Arguments
- `spline`: the `NormalSpline` object returned by `interpolate` or `construct` function.
- `points`: locations at which spline values are evaluating.
            This should be a vector of size `m` where `m` is the number of evaluating points.

Return: spline value at the `point` location.
"""
@inline function evaluate(spline::AbstractNormalSpline{n,T,RK}, points::AbstractVector{T}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return evaluate(spline, svectors(points))
end

"""
`evaluate_one(spline::AbstractNormalSpline{1,T,RK}, point::T) where {T <: Real, RK <: ReproducingKernel_0}`

Evaluate the 1D spline value at the `point` location.

# Arguments
- `spline`: the `NormalSpline` object returned by `interpolate` or `construct` function.
- `point`: location at which spline value is evaluating.

Return: spline value at the `point` location.
"""
@inline function evaluate_one(spline::AbstractNormalSpline{1,T,RK}, point::T) where {T <: Real, RK <: ReproducingKernel_0}
    return evaluate(spline, SVector{1,T}((point,)))
end

"""
`evaluate_derivative(spline::AbstractNormalSpline{1,T,RK}, point::T) where {T <: Real, RK <: ReproducingKernel_0}`

Evaluate the 1D spline derivative at the `point` location.

# Arguments
- `spline`: the `NormalSpline` object returned by `interpolate` or `construct` function.
- `point`: location at which spline derivative is evaluating.

Note: Derivative of spline built with reproducing kernel RK_H0 does not exist at the spline nodes.

Return: the spline derivative value at the `point` location.
"""
@inline function evaluate_derivative(spline::AbstractNormalSpline{1,T,RK}, point::T) where {T <: Real, RK <: ReproducingKernel_0}
    return evaluate_gradient(spline, SVector{1,T}((point,)))[1]
end

#### ReproducingKernel_1 (1-dimensional case)

"""
`prepare(nodes::AbstractVector{T}, d_nodes::AbstractVector{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}`

Prepare the 1D interpolating normal spline by constructing and factoring a Gram matrix of the problem.
Initialize the `NormalSpline` object.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n₁` vector where `n₁` is the number of function value nodes.
- `d_nodes`: The function derivatives nodes.
             This should be an `n₂` vector where `n₂` is the number of function derivatives nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: the partly initialized `NormalSpline` object that must be passed to `construct` function
        in order to complete the spline initialization.
"""
@inline function prepare(nodes::AbstractVector{T}, d_nodes::AbstractVector{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}
    d_dirs = fill(ones(SVector{1,T}), length(d_nodes))
    return prepare(svectors(nodes), svectors(d_nodes), d_dirs, kernel)
end

"""
`interpolate(nodes::AbstractVector{T}, values::AbstractVector{T}, d_nodes::AbstractVector{T}, d_values::AbstractVector{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}`

Prepare and construct the 1D interpolating normal spline.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n₁` vector where `n₁` is the number of function value nodes.
- `values`: function values at `nodes` nodes.
- `d_nodes`: The function derivatives nodes.
             This should be an `n₂` vector where `n₂` is the number of function derivatives nodes.
- `d_values`: function derivative values at `d_nodes` nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: the completely initialized `NormalSpline` object that can be passed to `evaluate` function.
"""
@inline function interpolate(nodes::AbstractVector{T}, values::AbstractVector{T}, d_nodes::AbstractVector{T}, d_values::AbstractVector{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}
    d_dirs = fill(ones(SVector{1,T}), length(d_nodes))
    return interpolate(svectors(nodes), values, svectors(d_nodes), d_dirs, d_values, kernel)
end

#### Utils for general case

"""
`get_epsilon(spline::AbstractNormalSpline{n,T,RK}) where {n, T <: Real, RK <: ReproducingKernel_0}`

Get the 'scaling parameter' of Bessel Potential space the spline was built in.
# Arguments
- `spline`: the `NormalSpline` object returned by `prepare`, `construct` or `interpolate` function.

Return: `ε` - the 'scaling parameter'.
"""
@inline function get_epsilon(spline::AbstractNormalSpline{n,T,RK}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return _get_kernel(spline).ε
end

"""
`estimate_epsilon(nodes::AbstractMatrix{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}`

Get the estimation of the 'scaling parameter' of Bessel Potential space the spline being built in.
It coincides with the result returned by `get_epsilon` function.
# Arguments
- `nodes`: The function value nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space
           and `n₁` is the number of function value nodes.
           It means that each column in the matrix defines one node.
- `kernel`: reproducing kernel of Bessel potential space the normal spline will be constructed in.
           It must be a struct object of the following type:
             `RK_H0` if the spline is constructing as a continuous function,
             `RK_H1` if the spline is constructing as a differentiable function,
             `RK_H2` if the spline is constructing as a twice differentiable function.
Return: estimation of `ε`.
"""
@inline function estimate_epsilon(nodes::AbstractMatrix{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}
    return estimate_epsilon(svectors(nodes), kernel)
end
@inline function estimate_epsilon(nodes::AbstractVecOfSVecs{n,T}, kernel::RK = RK_H0()) where {n, T <: Real, RK <: ReproducingKernel_0}
    return _estimate_epsilon(nodes, kernel)
end

"""
`estimate_epsilon(nodes::AbstractMatrix{T}, d_nodes::AbstractMatrix{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}`

Get an the estimation of the 'scaling parameter' of Bessel Potential space the spline being built in.
It coincides with the result returned by `get_epsilon` function.
# Arguments
- `nodes`: The function value nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space
           and `n₁` is the number of function value nodes.
           It means that each column in the matrix defines one node.
- `d_nodes`: The function directional derivative nodes.
           This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
           `n₂` is the number of function directional derivative nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline will be constructed in.
            It must be a struct object of the following type:
            `RK_H1` if the spline is constructing as a differentiable function,
            `RK_H2` if the spline is constructing as a twice differentiable function.

Return: estimation of `ε`.
"""
@inline function estimate_epsilon(nodes::AbstractMatrix{T}, d_nodes::AbstractMatrix{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}
    return estimate_epsilon(svectors(nodes), svectors(d_nodes), kernel)
end
@inline function estimate_epsilon(nodes::AbstractVecOfSVecs{n,T}, d_nodes::AbstractVecOfSVecs{n,T}, kernel::RK = RK_H1()) where {n, T <: Real, RK <: ReproducingKernel_1}
    return _estimate_epsilon(nodes, d_nodes, kernel)
end

"""
`estimate_cond(spline::AbstractNormalSpline{n,T,RK}) where {n, T <: Real, RK <: ReproducingKernel_0}`

Get an estimation of the Gram matrix condition number. It needs the `spline` object is prepared and requires O(N^2) operations.
(C. Brás, W. Hager, J. Júdice, An investigation of feasible descent algorithms for estimating the condition number of a matrix. TOP Vol.20, No.3, 2012.)
# Arguments
- `spline`: the `NormalSpline` object returned by `prepare`, `construct` or `interpolate` function.

Return: an estimation of the Gram matrix condition number.
"""
@inline function estimate_cond(spline::AbstractNormalSpline{n,T,RK}) where {n, T <: Real, RK <: ReproducingKernel}
    return _get_cond(spline)
end

"""
`estimate_accuracy(spline::AbstractNormalSpline{n,T,RK}) where {n, T <: Real, RK <: ReproducingKernel_0}`

Assess accuracy of interpolation results by analyzing residuals.
# Arguments
- `spline`: the `NormalSpline` object returned by `construct` or `interpolate` function.

Return: an estimation of the number of significant digits in the interpolation result.
"""
@inline function estimate_accuracy(spline::AbstractNormalSpline{n,T,RK}) where {n, T <: Real, RK <: ReproducingKernel_0}
    return _estimate_accuracy(spline)
end

"""
`get_cond(nodes::AbstractMatrix{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}`

Get a value of the Gram matrix spectral condition number. It is obtained by means of the matrix SVD decomposition and requires ``O(N^3)`` operations.
# Arguments
- `nodes`: The function value nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space and
           `n₁` is the number of function value nodes. It means that each column in the matrix defines one node.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: a value of the Gram matrix spectral condition number.
"""
@inline function get_cond(nodes::AbstractMatrix{T}, kernel::RK) where {T <: Real, RK <: ReproducingKernel_0}
    return get_cond(svectors(nodes), kernel)
end
@inline function get_cond(nodes::AbstractVecOfSVecs{n,T}, kernel::RK) where {n, T <: Real, RK <: ReproducingKernel_0}
    return get_cond(nodes, kernel)
end

"""
`get_cond(nodes::AbstractMatrix{T}, d_nodes::AbstractMatrix{T}, d_dirs::AbstractMatrix{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}`

Get a value of the Gram matrix spectral condition number. It is obtained by means of the matrix SVD decomposition and requires ``O(N^3)`` operations.
# Arguments
- `nodes`: The function value nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space and
           `n₁` is the number of function value nodes.
            It means that each column in the matrix defines one node.
- `d_nodes`: The function directional derivatives nodes.
             This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
             `n₂` is the number of function directional derivative nodes.
- `d_dirs`: Directions of the function directional derivatives.
        This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
        `n₂` is the number of function directional derivative nodes.
        It means that each column in the matrix defines one direction of the function directional derivative.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: a value of the Gram matrix spectral condition number.
"""
@inline function get_cond(nodes::AbstractMatrix{T}, d_nodes::AbstractMatrix{T}, d_dirs::AbstractMatrix{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}
    return get_cond(svectors(nodes), svectors(d_nodes), svectors(d_dirs), kernel)
end
@inline function get_cond(nodes::AbstractVecOfSVecs{n,T}, d_nodes::AbstractVecOfSVecs{n,T}, d_dirs::AbstractVecOfSVecs{n,T}, kernel::RK = RK_H1()) where {n, T <: Real, RK <: ReproducingKernel_1}
    return get_cond(nodes, d_nodes, d_dirs, kernel)
end

#### Utils for 1-dimensional case

"""
`estimate_epsilon(nodes::AbstractVector{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}`

Get an the estimation of the 'scaling parameter' of Bessel Potential space the 1D spline is being built in.
It coincides with the result returned by `get_epsilon` function.
# Arguments
- `nodes`: The function value nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: estimation of `ε`.
"""
@inline function estimate_epsilon(nodes::AbstractVector{T}, kernel::RK = RK_H0()) where {T <: Real, RK <: ReproducingKernel_0}
    return estimate_epsilon(svectors(nodes), kernel)
end

"""
`estimate_epsilon(nodes::AbstractVector{T}, d_nodes::AbstractVector{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}`

Get an the estimation of the 'scaling parameter' of Bessel Potential space the 1D spline is being built in.
It coincides with the result returned by `get_epsilon` function.
# Arguments
- `nodes`: The function value nodes.
- `d_nodes`: The function derivative nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: estimation of `ε`.
"""
@inline function estimate_epsilon(nodes::AbstractVector{T}, d_nodes::AbstractVector{T}, kernel::RK = RK_H1()) where {T <: Real, RK <: ReproducingKernel_1}
    return estimate_epsilon(svectors(nodes), svectors(d_nodes), kernel)
end

end # module NormalHermiteSplines


module PointProcess

# Write your package code here.
    # install necessary packages (open source)
    using Statistics, Plots, Distributions , LinearAlgebra

    # install necessary packages (original)
    using .FundamentalStats
    export EmpiricalFireRate
    export MakeWindow
    export FanoFactor
    function MakeWindow(t::Vector=[nothing],s = 0., f = nothing , delta = 1.;
                    data::Vector = t, start = s, finish = f , dt = delta)::Taple{Vector{Float64},Vector{Int64}}
        if finish == nothing
            finish = t[end]
        end
        T::Vector{Float64} = collect(start:dt:finish)
        N::Vector{Int64} = [sum( (data.>= i) .& (data.< i+dt) ) for i in T]
        return T, N
    end

    function EmpiricalFireRate(t::Vector=[nothing],s = 0., f = nothing , delta = 1.;
                    data::Vector = t, start = s, finish = f , dt = delta)::Taple{Vector{Float64},Vector{Float64}}
        X::Vector{Float64},Y::Vector{Float64} = MakeWindow(data = data, start = start, finish = finish, dt = dt)
        return X, Y./dt
    end
    function FanoFactor(t::Vector = [nothing], s = 0., f = nothing, delta = 1.;
                            data::Vector = t, start = s, finish = f , dt = delta)::Float64
        _, Y::Vector{Float64} = MakeWindow(data = data, start = start, finish = finish, dt = dt)
        return var(Y)/mean(Y)
    end

end


using Flora.MLP
using MLDatasets # FashionMNIST

# load full test set
test_x,  test_y  = FashionMNIST.testdata() # 10_000

picks = [
    [0, 3, 9], # 0 T-Shirt, 3 Dress, 9 Ankle boot
    [1, 2, 5],
]

for pick in picks
    p = load(:fashion, pick)
    tX, tY = prepare(test_x, test_y, pick)
    show_accuracy(p.m, p.pick, tX, tY)
end

#=
p.pick = [0, 3, 9]
mean(argmax(p.m(x)) .== argmax(y)) = 0.956
p.pick = [1, 2, 5]
mean(argmax(p.m(x)) .== argmax(y)) = 0.989
=#

#=
julia> using UnicodePlots

julia> spy(reshape(tX[:,11], 28, 28)')
      Sparsity Pattern
      ┌──────────────┐
    1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ > 0
      │⠀⠀⠀⠀⠀⠀⢀⠇⣾⣷⣦⣤⣴⡄│ < 0
      │⠀⠀⠀⠀⢀⠀⠘⢰⣿⣿⣿⣿⣿⠟│
      │⢀⣀⣀⠠⠔⣪⣴⣿⣿⣿⣿⣿⣿⡇│
      │⢠⣶⣶⣾⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿│
      │⡛⠿⢿⣿⣿⣿⣿⣿⠿⠿⣿⣿⣿⣿│
   28 │⠈⠉⠀⠀⠀⠀⠀⠀⠈⠁⠀⠀⠀⠀│
      └──────────────┘
      1             28
         nz = 363
=#


using Statistics
using Printf

using OffsetArrays
using DifferentialEquations
using Turing
using DynamicHMC
using CairoMakie
using ElectronDisplay

# Turing.setadbackend(:forwarddiff)

# Exact solution
# Convert x to periodic range in [0, L].
@inline x′(x, t, L, U) = mod(x + L - U * t, L)
@inline ϕ_exact(x, t, L, U) = exp(-100 * (x′(x, t, L, U) - 0.5)^2)

# Numerical solution
@inline advective_flux(i, u, ϕ) = u * (ϕ[i-1] + ϕ[i]) / 2
@inline ∂x_advective_flux(i, Δx, u, ϕ) = (advective_flux(i+1, u, ϕ) - advective_flux(i, u, ϕ)) / Δx

function advection(∂ϕ∂t, ϕ, p, t)
    Nx, Hx, Δx, u = p
    Nx = Int(Nx)
    Hx = Int(Hx)

    # filling periodic halos
    ϕ[-Hx+1:0] .= ϕ[Nx]
    ϕ[Nx+1:Nx+Hx] .= ϕ[1]

    for i in 1:Nx
        ∂ϕ∂t[i] = - ∂x_advective_flux(i, Δx, u, ϕ)
    end
end

Nx = 64
Hx = 1
Lx = 1
Δx = Lx / Nx
u = 0.2
p = [Nx, Hx, Δx, u]

xC = range(Δx/2, Lx - Δx/2, length=Nx)
xF = range(0, Lx, length=Nx+1)

ϕ₀ = OffsetArray(zeros(Nx+2Hx), -Hx+1:Nx+Hx)
@. ϕ₀[1:Nx] = ϕ_exact.(xC, 0, Lx, u)

tspan = (0.0, 10.0)
prob = ODEProblem(advection, ϕ₀, tspan, p, saveat=0.1)
sol = solve(prob, Tsit5())

function animate_advection(ϕ, ϕ_correct, x, t; framerate=30)
    n = Observable(1)
    ϕₙ = @lift ϕ[:, $n]
    ϕₙ_correct = @lift ϕ_correct[:, $n]

    title = @lift "Advection: t = $(round(t[$n], digits=1))"

    fig = Figure()
    ax = Axis(fig[1, 1], title=title, xlabel="x", ylabel="ϕ")
    lines!(ax, x, ϕₙ_correct)
    lines!(ax, x, ϕₙ)

    time_indices = 1:length(t)
    record(fig, "advection_1d.mp4", time_indices; framerate) do time_index
        @info "Animating 1D advection: frame $time_index/$(length(t))"
        n[] = time_index
    end
end

ϕ = Array(sol)[Hx+1:Nx+Hx, :]
ϕ_correct = ϕ_exact.(xC, sol.t', Lx, u)
animate_advection(ϕ, ϕ_correct, xC, sol.t)

function predict_advection(prob, u)
    Nx, Hx, Δx, _ = prob.p
    Nx = Int(Nx)
    Hx = Int(Hx)

    p = [Nx, Hx, Δx, u]
    prob = remake(prob, p=p)
    prediction = solve(prob, Tsit5())

    return prediction
end

function plot_likelihood(prob; saveat=0.1, N=1001, u_range=(-1, 1), chains=nothing)
    prob = remake(prob; saveat)
    Nx, Hx, _, u = prob.p
    Nx = Int(Nx)
    Hx = Int(Hx)

    losses = zeros(N)
    us = range(u_range[1], u_range[2], length=N)

    for (i, ũ) in enumerate(us)
        prediction = predict_advection(prob, ũ)
        ϕ_correct = ϕ_exact.(xC, prediction.t', Lx, u)
        prediction_array = Array(prediction)[1+Hx:Nx+Hx, :]
        losses[i] = mean((prediction_array .- ϕ_correct).^2)
    end

    fig = Figure()
    ax = Axis(fig[1, 1], title="Likelihood of model prediction for 1D advection with saveat=$saveat", xlabel="u", ylabel="likelihood ~ -loss (relative to max loss, arbitrary units)")
    lines!(ax, us, maximum(losses) .- losses)

    if !isnothing(chains)
        n_chains = size(chains, 3)
        for c in 1:n_chains
            vlines!(ax, mode(chains[:u].data[:, c]), color=(:red, 0.5))
        end
    end

    save("advection_1d_likelihood_saveat$saveat.png", fig)

    return
end

for sa in [0.01, 0.1, 0.2, 0.5, 1.0]
    plot_likelihood(prob, saveat=sa)
end

@model function probabilistic_advection(data, prob)
    σ ~ InverseGamma(3, 0.5)
    u ~ Normal(0.2, 0.05)

    prediction = predict_advection(prob, u)

    for i = 1:length(prediction)
        data[:, i] ~ MvNormal(prediction[i][1:Nx], σ)
    end
end

# prior_chain = sample(probabilistic_advection(ϕ_correct, prob), Prior(), 1000)

function callback_single_chain(rng, model, sampler, sample, state, iter; loss_history, kwargs...)
    u = sample.θ.u[1][1]
    σ = sample.θ.σ[1][1]

    prediction = Array(predict_advection(prob, u))[1+Hx:Nx+Hx, :]
    loss_history[iter] = loss = mean((prediction .- ϕ_correct).^2)

    @info "Iteration $iter: u=$u, σ=$σ, ℒ=$loss"
end

# sampler = MH(:u => x -> Normal(x, 0.05), :σ => x -> Normal(x, 5e-3))
# sampler = HMC(1e-3, 10)
sampler = NUTS(100, 0.65)

θ₀ = [0.05, 0.2] # Use with the init_params=θ₀ kwarg.

# n_samples = 10000
# loss_history = zeros(n_samples)
# chain = sample(probabilistic_advection(ϕ_correct, prob), sampler, n_samples, progress=true, callback=callback_single_chain; loss_history)

function callback_multi_chain(rng, model, sampler, sample, state, iter; loss_history, kwargs...)
    n = Threads.threadid()

    u = sample.θ.u[1][1]
    σ = sample.θ.σ[1][1]

    prediction = Array(predict_advection(prob, u))[1+Hx:Nx+Hx, :]
    loss_history[n, iter] = loss = mean((prediction .- ϕ_correct).^2)

    @info "Chain $n, Iteration $iter: u=$u, σ=$σ, ℒ=$loss"
end

n_chains = 8
@assert Threads.nthreads() >= n_chains

n_samples = 1000
loss_history = zeros(n_chains, n_samples)

chains_no_init = sample(probabilistic_advection(ϕ_correct, prob), sampler, MCMCThreads(), n_samples, n_chains, progress=true, callback=callback_multi_chain; loss_history)
plot_likelihood(prob, saveat=0.1, u_range=(-1, 1.5); chains)

chains_init = sample(probabilistic_advection(ϕ_correct, prob), sampler, MCMCThreads(), n_samples, n_chains, init_params=θ₀, progress=true, callback=callback_multi_chain; loss_history)

function plot_chains_and_densities(chains, vars; q_ADVI=nothing, n_ADVI_samples=10000)
    n_chains = size(chains, 3)

    fig = Figure()

    if !isnothing(q_ADVI)
        ADVI_samples = rand(q_ADVI, n_ADVI_samples)
    end

    for (n, var) in enumerate(vars)
        ax1 = Axis(fig[n, 1], ylabel="$var", xlabel="iteration")
        for c in 1:n_chains
            lines!(ax1, chains[var].data[:, c])
        end

        ax2 = Axis(fig[n, 2], ylabel="P($var)", xlabel="$var", yticklabelsvisible=false)
        density!(ax2, chains[var][:].data, label="NUTS")

        if !isnothing(q_ADVI)
            density!(ax2, ADVI_samples[n, :], label="ADVI")
        end

        axislegend(ax2)
    end

    save("chains_and_densities.png", fig)

    return
end

plot_chains_and_densities(chains_init, [:u, :σ])

# q_ADVI = vi(probabilistic_advection(ϕ_correct, prob), ADVI(100, 1000))
# plot_chains_and_densities(chains_init, [:u, :σ]; q_ADVI)


using PyPlot
using LinearAlgebra
using Statistics
using Random
using NLopt
using UKrig

##
# data grid and interpolation grid
n₀ = 500 # grid evals
n₁ = 200  # obs points
x₁ = sort(rand(n₁))
x₀ = range(-0.1, 1.1, length=n₀)
x₀₀ = range(0, 1, length=n₀)

##
# set the simulation truth parameters
σεₒ = 0.05
σzₒ = 1.0
ρₒ  = 0.1
νₒ  = 1.25
mₒ  = 2
βₒ  = 2 * (rand(mₒ) .- 1/2)

##
# simulate data
Σ₁ = (σzₒ^2) .* Mnu.(UKrig.ℓ2_dist.(x₁,permutedims(x₁)) ./ ρₒ, νₒ)
L₁ = cholesky(Σ₁).L
Z₁ = L₁ * randn(n₁)
Y₁ = Z₁ .+ βₒ[1] .+ x₁ .* βₒ[2]
data = Y₁ + σεₒ .* randn(n₁)

## ===============================
#

musr = 2
ν    = 1.25

loglike_inv_nu1, loglike_chol_nu1, wgrad_loglike_inv_nu1 = loglike_Gnu(
	data, x₁;
	musr=musr,
	ν=ν, # ≈ differentiability
)

#=
loglike_inv_nu1(0.1, 2.0)
loglike_chol_nu1(0.1, 2.0)
wgrad_loglike_inv_nu1(0.1, 2.0)
ϵ = 1e-7
(loglike_inv_nu1(0.1+ϵ, 2.0) - loglike_inv_nu1(0.1-ϵ, 2.0))/(2ϵ)
(loglike_inv_nu1(0.1, 2.0+ϵ) - loglike_inv_nu1(0.1, 2.0-ϵ))/(2ϵ)
=#

#LN_algm = [:LN_BOBYQA, :LN_COBYLA, :LN_PRAXIS, :LN_NELDERMEAD, :LN_SBPLX]
#LD_algm = [:LD_MMA, :LD_SLSQP, :LD_LBFGS, :LD_TNEWTON]
opt1 = Opt(:LN_BOBYQA, 2) # second arg is the number of variables to optimize
opt2 = Opt(:LN_BOBYQA, 2) # second arg is the number of variables to optimize
opt3 = Opt(:LD_LBFGS, 2) # second arg is the number of variables to optimize
opt1.max_objective = (x,grad) -> loglike_inv_nu1(x[1],x[2])
opt2.max_objective = (x,grad) -> loglike_chol_nu1(x[1],x[2])
opt3.max_objective = function (x,grad)
	ll, grad1, grad2 = wgrad_loglike_inv_nu1(x[1],x[2])
	if length(grad)>0
		grad[1]=grad1
		grad[2]=grad2
	end
	return ll
end
opt1.lower_bounds = [0.0, 0.0]
opt2.lower_bounds = [0.0, 0.0]
opt3.lower_bounds = [0.0, 0.0]
opt1.maxtime = 10 # in seconds
opt2.maxtime = 10
opt3.maxtime = 10
# opt1.maxeval = 50
# opt2.maxeval = 50
# opt3.maxeval = 50
# opt1.ftol_abs = 0.1 * sqrt(length(data)) / sqrt(2)
# opt2.ftol_abs = 0.1 * sqrt(length(data)) / sqrt(2)
# opt3.ftol_abs = 0.1 * sqrt(length(data)) / sqrt(2)
xinit1 = [0.1, 2.0]
xinit2 = [0.1, 2.0]
xinit3 = [0.1, 2.0]
optf1, optx1, ret1 = optimize!(opt1, xinit1)
optf2, optx2, ret2 = optimize!(opt2, xinit2)
optf3, optx3, ret3 = optimize!(opt3, xinit3)
@show sqrt(length(data))/√2

σe = optx1[1]
σg = optx1[2]

rng_σe = σe |> x->range(x - .05 * x, x + .05 * x, length=50)
rng_σg = σg |> x->range(x - .15 * x, x + .15 * x, length=50)
llmat = loglike_chol_nu1.(rng_σe,rng_σg')
pcolor(rng_σe, rng_σg, llmat)

krig1 = krig_Gnu(
	data, x₁;
	musr=musr,
	ν=ν,
	σe=σe,
	σg=σg,
)

figure(figsize=(9,6))
plot(x₁, data, ".", label="data")
plot(x₀, krig1.(x₀))
legend()

## ===============================
# do the interpolation, vary signal/noise
figure(figsize=(9,6))
plot(x₁, data, ".", label="data")
for σg ∈ [1,5,10]
	krig1 = krig_Gnu(
		data, x₁;
		musr=2,
		ν=1.0,    # ≈ differentiability
		σe=σe, # ≈ noise/signal ratio
		σg=σg,    # ≈ signal/noise ratio
	)
	plot(x₀, krig1.(x₀), label = "signal/noise = $σg")
end
legend()

# do the interpolation, vary smoothness
figure(figsize=(9,6))
plot(x₁, data, ".", label="data")
for ν ∈ [0.25, .5, 1.0, 1.25, 2.0]
	krig1 = krig_Gnu(
		data, x₁;
		musr=2,
		ν=ν,   # ≈ differentiability
		σe=σe, # ≈ noise/signal ratio
		σg=σg, # ≈ signal/noise ratio
	)
	plot(x₀, krig1.(x₀), label = "nu = $ν")
end
legend()

# do the interpolation, interpolate
figure(figsize=(9,6))
plot(x₁, data, ".", label="data")
for ν ∈ [0.25, .5, 1.0]
	krig1 = krig_Gnu(
		data, x₁;
		musr=2,
		ν=ν,    # ≈ differentiability
		σe=1e-7, # interpolate
		σg=σg, # ≈ signal/noise ratio
	)
	plot(x₀₀, krig1.(x₀₀), label = "nu = $ν")
end
legend()