"""
    radius_ellipsoid(ϵ, T, Y, ρ, nderivs, tol)

Compute radius of the ellipsoid used for computing the Riemann theta function or its derivatives.

# Arguments
- `ϵ::Real`: absolute error in value of theta function or its derivatives
- `T::Array{<:Real}`: upper triangular matrix in Cholesky factorization of Y
- `Y::Array{<:Real}`: imaginary part of Riemann matrix
- `ρ::Real`: norm of shortest vector in lattice generated by T
- `nderivs::Integer`: order of derivative to be computed
- `tol::Real=1e-15`: error tolerance in optimization routine
"""
function radius_ellipsoid(ϵ::Real, T::Array{<:Real}, Y::Array{<:Real}, ρ::Real, nderivs::Integer, tol::Real=1e-15)
    m = size(T,2); # number of columns in matrix
    norm_Tinv = norm(m > 1 ? inv(T) : inv.(T)); # norm of T⁻¹
    r = 0.5*(sqrt(m + 2*nderivs + sqrt(m^2 + 8*nderivs))+ρ);
    opt = Opt(:LN_COBYLA, 1); # use the COBYLA algorithm (Constrained Optimization BY Linear Approximations) in NLopt
    min_objective!(opt, (x, grad) -> radius_ellipsoid_helper(x[1], ϵ, ρ, m, norm_Tinv, nderivs)); # set objective function
    lower_bounds!(opt, [0.]); # add lower bound so that optimization parameter is positive
    inequality_constraint!(opt, (x, grad) -> - radius_ellipsoid_helper(x[1], ϵ, ρ, m, norm_Tinv, nderivs)); # set objective to be positive
    ftol_abs!(opt, tol); # set tolerance on value of objective function
    optx = optimize(opt, [r])[2];
    if optx[1] > r
        r = optx[1];
    end
    return r;
end

"""
    radius_ellipsoid_helper(x, grad, ϵ, ρ, g, T, L, nderivs)

Helper function for computing radius of ellipsoid for computing the theta function or its derivatives, in [`radius_ellipsoid`](@ref).
"""
function radius_ellipsoid_helper(r::Real, ϵ::Real, ρ::Real, g::Integer, norm_Tinv::Real, nderivs::Integer)
    if nderivs == 0
        return abs(g*2^(g-1)*sf_gamma_inc(g/2, (r-ρ/2)^2)/ρ^g) - ϵ;
    else
        return 2^(nderivs-1) * g * (2/ρ)^g * sum([binomial(nderivs, k) * π^(nderivs-k/2) * norm_Tinv^k * g^((nderivs-k)/2) * sf_gamma_inc((g+k)/2, (r-ρ/2)^2) for k=0:nderivs]) - ϵ;
    end
end


using Makie

func(x,y)=2x^2*y/(x^4+y^2)

linrange = LinRange(-1.,1.,500)
z =[func(x,y) for x in linrange, y in linrange]
surface(linrange,linrange,z)
wireframe!(linrange,linrange,zeros(size(z)))
lines!(linrange,linrange.^2,ones(length(linrange)),
       color=:red,linewidth=10
    )


function linsfa(signal::Array{Float64})
    w_signal = whiten!(signal)
    w_sig_diff = timediff(w_signal)
    A = w_sig_diff * w_sig_diff'
    B = w_signal * w_signal'

    # F = eigfact(A, B)
    F = eigfact(B, A)

    return F[:vectors]
end

# ensure signal is zero mean and unit variance
function whiten!(signal::Array{Float64})
    # zero mean
    signal = signal.-mean(signal,1)

    # unit variance
    signal = signal./std(signal,1)
    return signal
end

function timediff(signal::Array{Float64})
    # xdiff = diff(signal, 1)
    xdiff = [signal[1, :] ; diff(signal, 1)]
    return xdiff
end

function timediff!(signal::Array{Float64})
    signal =  [signal[1, :] ; diff(signal, 1)]
end


using SpecialFunctions:loggamma
using LinearAlgebra:dot

struct PolynomialLimbDark{T,VT <: AbstractVector{T},MT <: AbstractMatrix{T},AT <: AbstractArray{T}} <: AbstractLimbDark
    n_max::Int
    u_n::VT
    g_n::VT
    Mn_coeff::AT
    Nn_coeff::MT
    norm::T
    Mn::VT
    Nn::VT
end

@doc raw"""
    PolynomialLimbDark(u::AbstractVector)

Polynomial limb darkening using analytical integrals. The length of the `u` vector is equivalent to the order of polynomial used; e.g., `[0.2, 0.3]` corresponds to quadratic limb darkening.

# Mathematical form
```math
I(\mu) \propto 1 - u_1(1-\mu) - u_2(1-\mu)^2 - \dots - u_N(1-\mu)^N
```
which is equivalent to the series
```math
I(\mu) \propto -\sum_{i=0}^N{u_i(1-\mu)^i}
```
with the definition $u_0 \equiv -1$.

# Examples
```jldoctest poly
u = [0.4, 0.26] # quadratic and below is 100% analytical
ld = PolynomialLimbDark(u)
ld(0.1, 0.01)

# output
0.9998787880717668
```
```jldoctest poly
u2 = vcat(u, ones(12) ./ 12)
ld2 = PolynomialLimbDark(u2)
ld2(0.1, 0.01)

# output
0.9998740059086433
```

# References

> [Agol, Luger, Foreman-Mackey (2020)](https://ui.adsabs.harvard.edu/abs/2020AJ....159..123A)
>
>   "Analytic Planetary Transit Light Curves and Derivatives for Stars with Polynomial Limb Darkening"

> [Luger et al. (2019)](https://ui.adsabs.harvard.edu/abs/2019AJ....157...64L)
>
>   "starry: Analytic Occultation Light Curves"
"""
function PolynomialLimbDark(u::AbstractVector{S}; maxiter=100) where S
    T = float(S)
    # add constant u_0 term
    n_max = length(u)
    u_n = pushfirst!(copy(u), -one(S))
    # get Green's basis coefficients
    g_n = compute_gn(u_n)
    # compute series expansion for M_n and N_n
    Mn_coeff = compute_Mn_coeffs(T, n_max; maxiter=maxiter)
    Nn_coeff = compute_Nn_coeffs(T, n_max; maxiter=maxiter)

    # calculate flux normalization factor, which only depends on first two terms
    if length(g_n) > 1
        norm = inv(π * (g_n[1] + 2 / 3 * g_n[2]))
    else
        norm = inv(π * g_n[1])
    end

    # pre-allocate temp arrays
    Mn = similar(g_n)
    Nn = similar(g_n)

    return PolynomialLimbDark(n_max, u_n, g_n, Mn_coeff, Nn_coeff, norm, Mn, Nn)
end

function compute(ld::PolynomialLimbDark, b::S, r) where S
    T = float(S)

    ## check for trivial cases
    if b ≥ 1 + r || iszero(r)
        # completely unobscured
        return one(T)
    elseif r ≥ 1 + b
        # completely obscured
        return zero(T)
    end

    r2 = r^2
    b2 = b^2

    if iszero(b)
        onemr2 = 1 - r2
        sqrt1mr2 = sqrt(onemr2)

        # annular ellipse
        if ld.n_max == 0
            flux = ld.g_n[begin] * onemr2
        else
            flux = ld.g_n[begin] * onemr2 + 2 / 3 * ld.g_n[begin + 1] * sqrt1mr2^3
        end
        fac = 2 * r2 * onemr2
        for i in 2:ld.n_max
            flux -= ld.g_n[begin + i] * fac
            fac *= sqrt1mr2
        end
        return flux * π * ld.norm
    end

    # take a moment to calculate values used repeatedly in the
    # numerical routines ahead
    onembmr2 = (r - b + 1) * (1 - r + b)
    onembmr2inv = inv(onembmr2)
    sqonembmr2 = sqrt(onembmr2)
    br = b * r
    fourbr = 4 * br
    fourbrinv = inv(fourbr)
    sqbr = sqrt(br)
    sqbrinv = inv(sqbr)
    onembpr2 = (1 - r - b) * (1 + b + r)
    sqarea = sqarea_triangle(one(T), r, b)
    k2 = max(zero(T), onembmr2 * fourbrinv)
    k = sqrt(k2)
    omemr2 = (1 - r) * (1 + r)
    onemr2mb2 = omemr2 - b2
    onemr2pb2 = omemr2 + b2
    if k2 ≥ 1
        if k2 > 2
            kc2 = 1 - inv(k2)
        else
            kc2 = onembpr2 * onembmr2inv
        end
    else
        if k2 > 0.5
            kc2 = (r - 1 + b) * (b + r + 1) * fourbrinv
        else
            kc2 = 1 - k2
        end
    end
    kc = sqrt(abs(kc2))

    ## Compute uniform term
    s0, kap0, kite_area2, kck = compute_uniform(b, r; b2, r2, sqarea, fourbrinv)

    flux = ld.g_n[begin] * s0

    if ld.n_max == 0
        return flux * ld.norm
    end

    ## compute linear term
    s1, Eofk, Em1mKdm = compute_linear(b, r; k2, kc, kc2, r2, b2, br, fourbr, sqbr, onembmr2, onembpr2, onembmr2inv)

    flux += ld.g_n[begin + 1] * s1

    if ld.n_max == 1
        return flux * ld.norm
    end

    ## calculate quadratic term
    s2 = compute_quadratic(b, r; s0, r2, b2, kap0, kite_area2, k2)
    flux += ld.g_n[begin + 2] * s2
    if ld.n_max == 2
        return flux * ld.norm
    end

    ## higher order (N > 3) terms require solving Mn and Nn integrals
    if k2 < 0.5 && ld.n_max > 3
        downwardM!(ld.Mn, ld.Mn_coeff; n_max=ld.n_max, sqonembmr2, sqbr, onemr2mb2, k, k2, sqarea, kap0, Eofk, Em1mKdm, kite_area2)
    else
        upwardM!(ld.Mn; sqbr, n_max=ld.n_max, sqonembmr2, onemr2mb2, sqarea, k2, kap0, Eofk, Em1mKdm, kite_area2)
    end

    # compute remaining terms
    for n in 3:ld.n_max
        pofgn_M = 2 * r2 * ld.Mn[begin + n] - n / (n + 2) *
                       (onemr2mb2 * ld.Mn[begin + n] + sqarea * ld.Mn[begin + n - 2])
        flux -= ld.g_n[begin + n] * pofgn_M
    end

    return flux * ld.norm
end

###
###   Type methods
###

function Base.show(io::IO, ld::PolynomialLimbDark)
    print(io, "PolynomialLimbDark(", ld.u_n, ")")
end
function Base.show(io::IO, ::MIME"text/plain", ld::PolynomialLimbDark)
    print(io, "PolynomialLimbDark\n u_n: ", ld.u_n)
end
function Base.show(io::IO, ::MIME"text/html", ld::PolynomialLimbDark)
    print(io, "PolynomialLimbDark\n un: ", ld.u_n)
end

###
###   Initialization helpers
###

"""
Transform the `u_n` coefficients to `g_n`, which are coefficients in Green's basis from REF.
"""
compute_gn(u::AbstractVector) = compute_gn!(zero(u), u)

function compute_gn!(g_n::AbstractVector{T}, u_n) where T
    n = length(g_n) - 1
    ## First: calculate the a_n terms (eqn. 10)
    a_n = zero(g_n)
    # constant term
    a_n[begin] = one(T)
    for i in 1:n
        bcoeff = one(T) # binomial coefficient
        for j in 0:i
            a_n[begin + j] -= (-1)^j * bcoeff * u_n[begin + i]
            bcoeff *= (i - j) / (j + 1)
        end
    end

    ## Second: calculate the g_n terms (eqn. 15)
    for j in n:-1:2
        if j ≥ n - 1
            g_n[begin + j] = a_n[begin + j] / (j + 2)
        else
            g_n[begin + j] = a_n[begin + j] / (j + 2) + g_n[begin + j + 2]
        end
    end
    if n ≥ 3
        g_n[begin + 1] = a_n[begin + 1] + 3 * g_n[begin + 3]
    elseif n ≥ 1
        g_n[begin + 1] = a_n[begin + 1]
    end
    if n ≥ 2
        g_n[begin] = a_n[begin] + 2 * g_n[begin + 2]
    else
        g_n[begin] = a_n[begin]
    end

    return g_n
end

###
###   Computation helpers
###

function sort3(p0, p1, p2)
    if p0 < p1
        p0, p1 = p1, p0
    end
    if p1 < p2
        p1, p2 = p2, p1
    end
    if p0 < p1
        p0, p1 = p1, p0
    end
    return p0, p1, p2
end

function sqarea_triangle(p0, p1, p2)
    p0, p1, p2 = sort3(p0, p1, p2)
    sqarea = (p0 + (p1 + p2)) * (p2 - (p0 - p1)) *
             (p2 + (p0 - p1)) * (p0 + (p1 - p2))
    return sqarea
end

function compute_uniform(b::T, r; sqarea, r2=r^2, b2=b^2, fourbrinv=inv(4 * b * r)) where T
    if b ≤ 1 - r
        flux = π * (1 - r2)
        kap0 = convert(T, π)
        kck = zero(T)
        kite_area2 = zero(T)
    else
        kite_area2 = sqrt(sqarea)
        r2m1 = (r - 1) * (r + 1)
        kap0 = atan(kite_area2, r2m1 + b2)
        Πmkap1 = atan(kite_area2, r2m1 - b2)
        kck = kite_area2 * fourbrinv
        flux = Πmkap1 - r2 * kap0 + 0.5 * kite_area2
    end

    return flux, kap0, kite_area2, kck
end

function compute_linear(b::T, r; k2, kc, kc2, onembmr2, onembpr2, onembmr2inv, r2=r^2, b2=b^2, br=b * r, fourbr=4 * br, sqbr=sqrt(br)) where T
    if iszero(b) # case 10
        Λ1 = -2 * π * sqrt(1 - r2)^3
        Eofk = 0.5 * π
        Em1mKdm = 0.25 * π
    elseif b == r
        if r == 0.5 # case 6
            Λ1 = π - 4 / 3
            Eofk = one(T)
            Em1mKdm = one(T)
        elseif r < 0.5 # case 5
            m = 4 * r2
            Eofk = cel(m, one(T), one(T), 1 - m)
            Em1mKdm = cel(m, one(T), one(T), zero(T))
            Λ1 = π + 2 / 3 * ((2 * m - 3) * Eofk - m * Em1mKdm) +
                 (b - r) * 4 * r * (Eofk - 2 * Em1mKdm)
        else # case 7
            m = 4 * r2
            minv = inv(m)
            Eofk = cel(minv, one(T), one(T), 1 - minv)
            Em1mKdm = cel(minv, one(T), one(T), zero(T))
            Λ1 = π + 1 / 3 * ((2 * m - 3) * Em1mKdm - m * Eofk) / r -
                 (b - r) * 2 * (2 * Eofk - Em1mKdm)
        end
    else
        if b + r > 1 # case 2, case 8
            Πofk, Eofk, Em1mKdm = cel(k2, kc, (b - r)^2 * kc2, zero(T), one(T), one(T), 3 * kc2 * (b - r) * (b + r), kc2, zero(T))
            Λ1 = onembmr2 * (Πofk + (-3 + 6 * r2 + 2 * br) * Em1mKdm - fourbr * Eofk) / (3 * sqbr)
        elseif b + r < 1 # case 3 case 9
            bmrdbpr = (b - r) / (b + r)
            μ = 3 * bmrdbpr * onembmr2inv
            p = bmrdbpr^2 * onembpr2 * onembmr2inv
            Πofk, Eofk, Em1mKdm = cel(inv(k2), kc, p, 1 + μ, one(T), one(T), p + μ, kc2, zero(T))
            Λ1 = 2 * sqrt(onembmr2) * (onembpr2 * Πofk - (4 - 7 * r2 - b2) * Eofk) / 3
        else # case
            Λ1 = 2 * acos(1 - 2 * r) - 2 * π * (r > 0.5) - (4 / 3 * (3 + 2 * r - 8 * r2) + 8 * (r + b - 1) * r) * sqrt(r * (1 - r))
            Eofk = one(T)
            Em1mKdm = one(T)
        end
    end

    flux = ((1 - T(r > b)) * 2π - Λ1) / 3

    return flux, Eofk, Em1mKdm
end

function compute_quadratic(b, r; s0, r2, b2, kap0, kite_area2, k2)
    r2pb2 = r2 + b2
    η2 = r2 * (r2pb2 + b2)
    if k2 > 1
        four_pi_eta = 2 * π * (η2 - 1)
    else
        Πmkap1 = atan(kite_area2, (r - 1) * (r + 1) - b2)
        four_pi_eta = 2 * (-Πmkap1 + η2 * kap0 - 0.25 * kite_area2 * (1 + 5 * r2 + b2))
    end
    return 2 * s0 + four_pi_eta
end


"""
    `time_sec(run)` calculates the time in sec.
"""
function time_sec(run::Run)
    cumsum(diff(run.date) / Base.Dates.Millisecond(1000))
end

"""
    `dist_m(run, smoothing = true, λ = 500.0)` estimation of the distance in meters.
"""
function dist_m(run::Run;
    smoothing::Bool = true, λ::AbstractFloat = 500.0)
    lon = run.lon * pi / 180
    lat = run.lat * pi / 180
    m1 = Vector{Float64}(length(lat) - 1)
    for i in 2:length(lat)
        x = (lon[i] - lon[i - 1]) * cos(mean(lat[[i - 1, i]]))
        y = lat[i] - lat[i - 1]
        # authalic radius of 6371007.2 meters
        m1[i - 1] = sqrt(x^2 + y^2) * 6371007.2
    end
    if smoothing
        sec = time_sec(run)
        m1 = predict(fit(SmoothingSpline, sec, m1, λ))
        m1[find(m1 .< 0.0)] .= 0
    end
    m = cumsum(m1)
    m
end

"""
    `speed_kmh(run, smoothing = true, λ = 500.0)` calculates the speed in km/h.
"""
function speed_kmh(run::Run;
    smoothing::Bool = true, λ::AbstractFloat = 500.0)
    m1 = diff(vcat(0.0, dist_m(run, smoothing = smoothing, λ = λ)))
    sec = time_sec(run)
    sec1 = diff(vcat(0.0, sec))
    kmh = (m1 ./ 1000) ./ (sec1 ./ 3600)
    kmh
end

"""
    `pace_minkm(run, smoothing = true, λ = 500.0)` calculates the pace in min/km.
"""
function pace_minkm(run::Run;
    smoothing::Bool = true, λ::AbstractFloat = 500.0)
    m = diff(vcat(0.0, dist_m(run, smoothing = smoothing, λ = λ)))
    sec = diff(vcat(0.0, time_sec(run)))
    minkm = (sec ./ 60) ./ (m ./ 1000)
    minkm
end


using Manifolds
using LinearAlgebra
using StaticArrays

## The factors
struct MeasurementOnTangent end

function PointPoint_distance(M, m, p, q)
    q̂ = compose(M, p, m)
    return distance(M, q, q̂)
end

function grad_PointPoint_distance(M, m, p, q)
    q̂ = compose(M, p, m)
    return grad_distance(M, q, q̂)
end

function PointPoint_distance(M, X, p, q, ::MeasurementOnTangent)
    q̂ = compose(M, p, exp(M, identity_element(M, p), X))
    return distance(M, q, q̂)
end

function PointPoint_velocity_distance(M, X, dt, p, q, ::MeasurementOnTangent)
    q̂ = compose(M, p, group_exp(M, X*dt))
    return distance(M, q, q̂)
end

function Prior_distance(M, meas, p)
#
    return distance(M, meas, p)
end

# Pose
# Pose2Pose2(m, p, q) = PointPoint_distance(getManifold(Pose2), m, p, q)
Pose2Pose2(m, p, q) = PointPoint_distance(SpecialEuclidean(2), m, p, q)
Pose3Pose3(m, p, q) = PointPoint_distance(SpecialEuclidean(3), m, p, q)
PriorPose2(m, p) = Prior_distance(SpecialEuclidean(2), m, p)
PriorPose3(m, p) = Prior_distance(SpecialEuclidean(3), m, p)

#Point
Point1Point1(m, p, q) = PointPoint_distance(TranslationGroup(1), m, p, q)
Point2Point2(m, p, q) = PointPoint_distance(TranslationGroup(2), m, p, q)
Point3Point3(m, p, q) = PointPoint_distance(TranslationGroup(3), m, p, q)
PriorPoint2(m, p) = Prior_distance(TranslationGroup(2), m, p)
PriorPoint3(m, p) = Prior_distance(TranslationGroup(3), m, p)

## Testing the Factors

ϵSE2 = ProductRepr(SA[0.,0] , SA[1. 0; 0 1])
ϵSE3 = ProductRepr(SA[0.,0,0] , SA[1. 0 0; 0 1 0; 0 0 1])

M = SpecialEuclidean(2)

X = hat(M, ϵSE2, [1, 0, pi/4+0.1])
p = compose(M, ϵSE2, exp(M, ϵSE2, X))

X = hat(M, ϵSE2, [1, 0, pi/8])
q = compose(M, p, exp(M, ϵSE2, X))

X = hat(M, ϵSE2, [1-1e-3, 0+1e-3, pi/8+1e-3])
m = compose(M, ϵSE2, exp(M, ϵSE2, X))

Pose2Pose2(m, p, q)

PointPoint_distance(M, X, p, q, MeasurementOnTangent())

PriorPose2(m, q)
PriorPose2(q, q)

## Testing the Factor

## and with matrices
# using LinearAlgebra

# ϵSE2 = Matrix(1.0I(3))
# ϵSE3 = Matrix(1.0I(4))

# M = SpecialEuclidean(2)

# X = hat(M, ϵSE2, [1, 0, pi/4])
# # exp does not work with affine matrix
# p = compose(M, ϵSE2, exp(M, ϵSE2, X))


"""
	function σTemp(ω0::Real,temp::Number)

Returns the standard deviation due to Doppler broadening of frequency ω0 at temp
"""
function σTemp(ω0::Real,temp::Number)
	kbOverMhC2 = 9.178e-14;
	return sqrt(kbOverMhC2*temp)*ω0
end

export σTemp


module TermStructure

using Miletus.DayCounts
using Dates

export YieldTermStructure, VolatilityTermStructure, ConstantYieldCurve,
		compound_factor, discount_factor, implied_rate, discount, forward_rate,
		reference_date, zero_rate, par_rate

#General Interest Rate functionality

compound_factor(r::Real, compounding::Symbol, freq::Symbol,  t::Real) = compound_factor(r, compounding, eval(freq), t)
compound_factor(r::Real, compounding::Symbol,  t::Real) = compound_factor(r, compounding, NoFrequency, t)
compound_factor(r::Real, compounding::Symbol, dc::DayCount, dates::Date...) = compound_factor(r, compounding, NoFrequency, dates...)
compound_factor(r::Real, compounding::Symbol, freq::Symbol, dc::DayCount, dates::Date...) = compound_factor(r, compounding, eval(freq), dates...)
compound_factor(r::Real, compounding::Symbol, freq::Integer, dc::DayCount, dates::Date...) = compound_factor(r,  compounding, freq, yearfraction(dc, dates...))

function compound_factor(r::Real, compounding::Symbol, freq::Integer, t::Real )
	if compounding == :Simple
		return 1 + r * t
	elseif compounding == :Compounded
		@assert(freq != NoFrequency)
		return (1 + r / freq) ^ (freq * t)
	elseif compounding == :Continuous
		return exp(r * t)
	elseif compounding == :SimpleThenCompounded
		@assert(freq != NoFrequency)
		if t < (1 / freq)
            return 1 + r * t
        else
            return (1 + r / freq) ^ (freq * t)
        end
	else
		error("Unknown compounding")
	end
end

discount_factor(x...) = 1/compound_factor(x...)

implied_rate(c::Real, compounding::Symbol, freq::Symbol,  t::Real) = implied_rate(c,compounding, eval(freq), t)
implied_rate(c::Real, compounding::Symbol, freq::Symbol, dc::DayCount, dates::Date...) = implied_rate(c,compounding, eval(freq), dates...)
implied_rate(c::Real, compounding::Symbol, freq::Integer, dc::DayCount, dates::Date...) = implied_rate(c, compounding, freq, yearfraction(dc, dates...))

function implied_rate(c::Real, compounding::Symbol, freq::Integer, t::Real)
	@assert c>0 && t>0
	if compounding == :Simple
		return (c-1) / t
	elseif compounding == :Compounded
		return (c ^ (1/(f*t)) - 1) *f
	elseif compounding == :Continuous
		return log(c)/t
	elseif compounding == :SimpleThenCompounded
		if t < (1 / freq)
            return (c-1) / t
        else
            return (c ^ (1/(f*t)) - 1) *f
        end
	else
		error("Unknown compounding")
	end
end

abstract type TermStuct end

abstract type VolatilityTermStructure <: TermStuct end
abstract type YieldTermStructure <: TermStuct end

#Default values
daycount(::YieldTermStructure) = Actual365()
startdate(::YieldTermStructure) = Dates.today()
compounding(::YieldTermStructure) = :Continuous
frequency(::YieldTermStructure) = NoFrequency

discount(ts::YieldTermStructure, d::Date) = discount(ts, yearfraction(daycount(ts), startdate(ts), d))
Base.getindex(ts::YieldTermStructure, d::Date) = discount(ts, d)
implied_rate(ts::YieldTermStructure, d::Date) = implied_rate(1/ts[d], compounding(ts), frequency(ts), yearfraction(daycount(ts),  startdate(ts), d) )

#Overload this method for each concrete YieldTermStructure
discount(ts::YieldTermStructure, t::Real) = error("Must be implemented by concrete Term Structure")

function forward_rate(ts::YieldTermStructure, d1::Date, d2::Date )
	if d1==d2
		t1 = yearfraction(daycount(ts), startdate(ts), d1)
		t2 = t1+.0001
		c=discount(ts, t1) / discount(ts, t2)
		return implied_rate(c, compounding(ts), frequency(ts), t2-t1)
	elseif d1<d2
		return implied_rate(discount(ts, d1)/discount(ts, d2), compounding(ts), frequency(ts), daycount(ts), d1, d2)
	else
		error("Forward start date must be before forward end dates")
	end
end

function forward_rate(ts::YieldTermStructure, t1::Real, t2::Real )
	if (t2==t1)
		t2=t1+.0001
	end

	compound = discount(ts, t1) / discount(ts, t2)
	return implied_rate(discount(ts, t1) / discount(ts, t2), daycount(ts), t2-t1)
end

zero_rate(ts::YieldTermStructure,  d1::Date) = zero_rate(ts, yearfraction(ts.dc, reference_date(ts), d1))
function zero_rate(ts::YieldTermStructure, t::Real)
	if (t == 0)
		c = 1/discount(ts, .0001)
		return implied_rate(c, ts.compounding, ts.freq, .0001)
	else
		c = 1/discount(ts, t)
		return implied_rate(c, ts.compounding, ts.freq, t)
	end
end

par_rate(ts::YieldTermStructure, dates::AbstractVector{Date}) =
	par_rate(ts, [yearfraction(ts.dc, reference_date(ts), dt) for dt in dates])
function par_rate(ts::YieldTermStructure, tm::AbstractVector)
	s = sum([discount(ts, t) for t in tm])
	r=discount(ts, tm[1])
end

mutable struct ConstantVolatilityCurve <: VolatilityTermStructure

end

struct ConstantContinuousYieldCurve{DC,T} <: YieldTermStructure
    dc::DC
    rate::T
    reference_date::Date
end

daycount(y::ConstantContinuousYieldCurve) = y.dc
startdate(y::ConstantContinuousYieldCurve) = y.reference_date
discount(y::ConstantContinuousYieldCurve, yf::Float64) = exp(-y.rate * yf)

discount(y::ConstantContinuousYieldCurve, d::Date) = discount(y, yearfraction(daycount(y), startdate(y), d))

mutable struct ConstantYieldCurve{T} <: YieldTermStructure
	dc::DayCount
	rate::T
	compounding::Symbol
	freq::Integer
	reference_date::Date
end

ConstantYieldCurve(dc::DayCount, rate::Real, compounding::Symbol, freq::Symbol, reference_date::Date) =
		ConstantYieldCurve(dc, rate, compounding, eval(freq), reference_date)

daycount(y::ConstantYieldCurve) = y.dc
startdate(y::ConstantYieldCurve) = y.reference_date

ConstantYieldCurve(dc::DayCount, rate::Real) = ConstantYieldCurve(dc, rate, :Continuous, :NoFrequency, today())
discount(ts::ConstantYieldCurve, t::Real) = discount_factor(ts.rate, ts.compounding, ts.freq, t)

const NoFrequency       = -1
const Once			 	= 0
const Annual			= 1
const Semiannual		= 2
const EveryFourthMonth  = 3
const Quarterly		 	= 4
const Bimonthly		 	= 6
const Monthly			= 12
const EveryFourthWeek  	= 13
const Biweekly		 	= 26
const Weekly			= 52
const Daily			 	= 365
const OtherFrequency   	= 999

end #Module


# rules.jl - differentiation rules for basic operations.

## basic rules

@scalardiff (-x::Number) 1 -1

# product

@scalardiff *(x::Number, y::Number) 1 y
@scalardiff *(x::Number, y::Number) 2 x

# elementwise product

@scalardiff .*(x::Number, y::Number) 1 y
@scalardiff .*(x::Number, y::Number) 2 x

# other arithmetic operations

@scalardiff (x::Number ^ n::Number) 1 (n * x.^(n-1))
@scalardiff (a::Number ^ x::Number) 2 (log(a) * a.^x)
@scalardiff (x::Number .^ n::Number) 1 (n * x.^(n-1))
@scalardiff (a::Number .^ x::Number) 2 (log(a) * a.^x)

@scalardiff (x::Number / y::Number) 1 (1 / y)
@scalardiff (n::Number / x::Real) 2 (-n * x .^ -2.0)
@scalardiff (x::Number ./ y::Number) 1 (1 / y)
@scalardiff (n::Number ./ x::Real) 2 (-n * x .^ -2.0)

@scalardiff (x::Any + y::Any) 1 1
@scalardiff (x::Any + y::Any) 2 1
@scalardiff (x::Any + y::Any + z::Any) 1 1
@scalardiff (x::Any + y::Any + z::Any) 2 1
@scalardiff (x::Any + y::Any + z::Any) 3 1
@scalardiff (w::Any + x::Any + y::Any + z::Any) 1 1
@scalardiff (w::Any + x::Any + y::Any + z::Any) 2 1
@scalardiff (w::Any + x::Any + y::Any + z::Any) 3 1
@scalardiff (w::Any + x::Any + y::Any + z::Any) 4 1

@scalardiff (x::Any .+ y::Any) 1 1
@scalardiff (x::Any .+ y::Any) 2 1

@scalardiff (x::Any - y::Any) 1 1
@scalardiff (x::Any - y::Any) 2 -1

@scalardiff (x::Any .- y::Any) 1 1
@scalardiff (x::Any .- y::Any) 2 -1

@scalardiff sum(x::Number) 1 1

# trigonomeric functions

@scalardiff sin(x::Number) 1 cos(x)
@scalardiff cos(x::Number) 1 -sin(x)
@scalardiff tan(x::Number) 1 (1. + tan(x)  * tan(x))
@scalardiff sinh(x::Number) 1 cosh(x)
@scalardiff cosh(x::Number) 1 sinh(x)
@scalardiff tanh(x::Number) 1 (1.0 .- tanh(x) .* tanh(x))
@scalardiff asin(x::Number) 1 (1 ./ sqrt(1 .- x .* x))
@scalardiff acos(x::Number) 1 (1 ./ sqrt(1 .- x .* x))
@scalardiff atan(x::Number) 1 (1 ./ (1 .+ x .* x))

# sqrt

@scalardiff sqrt(x::Number) 1 (0.5 * x .^ (-0.5))

# exp, log

@scalardiff exp(x::Number) 1 exp(x)
@scalardiff log(x::Number) 1 (1 ./ x)

# abs

@scalardiff abs(x::Number) 1 (sign(x) .* x)
@scalardiff abs2(x::Number) 1 (2 .* abs(x) .* sign(x) .* x)

# min, max

@scalardiff max(x::Number, y::Number) 1 (x > y)
@scalardiff max(x::Number, y::Number) 2 (y > x)

@scalardiff min(x::Number, y::Number) 1 (x < y)
@scalardiff min(x::Number, y::Number) 2 (y < x)

@scalardiff sign(x::Any) 1 0.

# transpose

@scalardiff transpose(x::Number) 1 1

@scalardiff size(x::Any) 1 0.0
@scalardiff size(x::Any, y::Any) 1 0.0
@scalardiff size(x::Any, y::Any) 2 0.0


## Author: Victor H. Aguiar

## questions: slack

## Import packages, make sure to install them.
## Julia version: v"1.5.2"
import JuMP
import Ipopt

## define a model
robinson=JuMP.Model(Ipopt.Optimizer)

# define variables
#coconuts
@JuMP.variable(robinson,c>=0)
#labor
@JuMP.variable(robinson,l>=0)
#wages
@JuMP.variable(robinson,w>=0)

# define parameters
α1=.4
α2=(1-α1)
#@JuMP.NLparameter(robinson, β == 1/2)
β = .1
A=50
Lbar=100.0

# function prodfunction(x)
#        try
#              A*x^(β)
#       catch
#             NaN
#       end
# end
#
# function prodfunction_prime(x)
#        try
#              A*β*(x)^(β-1)
#       catch
#             NaN
#       end
# end
#
# function prodfunction_prime_prime(x)
#        try
#               A*β^2*x^(β-2)
#       catch
#             NaN
#       end
# end
#
# function prodfunction_prime_prime_prime(x)
#        try
#               A*β^3*x^(β-3)
#       catch
#             NaN
#       end
# end
prodfunction_prime(x) = A*x^β
prodfunction_prime(x) = A*β*x^(β-1)
prodfunction_prime_prime(x) = A*β^2*x^(β-2)
prodfunction_prime_prime_prime(x) = A*β^3*x^(β-3)
# define equations, here in JuMP we must define them as (nonlinear) constraint
JuMP.register(robinson, :prodfunction, 1, prodfunction, prodfunction_prime,
         prodfunction_prime_prime)

JuMP.register(robinson, :prodfunction_prime, 1, prodfunction_prime, prodfunction_prime_prime,
                  prodfunction_prime_prime_prime)

# FOC of the consumer
@JuMP.NLconstraint(robinson,α1/(α2)*(Lbar-l)*w==c)
#@JuMP.NLconstraint(robinson,α1/(α2)*(Lbar-l)*prodfunction_prime(l+ 1e-7)==c)

# Marginal productivity of labor equal to wage
@JuMP.NLconstraint(robinson,prodfunction_prime(l)==w)

#@JuMP.constraint(robinson,c>=0)
#@JuMP.constraint(robinson,l>=0)
#@JuMP.constraint(robinson,w>=0)
# Budget constraint in equilibrium
@JuMP.NLconstraint(robinson,prodfunction(l)==c)

# Solve model
@JuMP.NLobjective(robinson,Min,1)
JuMP.optimize!(robinson)

JuMP.value(c)
JuMP.value(l)
JuMP.value(w)


@testset "Test Problems" begin
    @testset "Unconstrained" begin
        @testset "camel6" begin
            p = DSProblem(2; poll=OrthoMADS(2), objective=DS.camel6, initial_point=[5.0, 5.0])
            Optimize!(p)
            @test isapprox(abs(p.x[1]), 0.0898, atol=1e-4)
            @test isapprox(abs(p.x[2]), 0.7126, atol=1e-4)
            @test isapprox(p.x_cost, -1.0316, atol=1e-4)
            @test p.iteration < 120
        end
        @testset "audet_dennis_benchmark_1" begin
            p = DSProblem(2; poll=OrthoMADS(2), objective=DS.bm_1, initial_point=[-2.1, 1.7])
            Optimize!(p)
            @test isapprox(p.x[1], 0, atol=1e-8)
            @test isapprox(p.x[2], 0, atol=1e-8)
            @test isapprox(p.x_cost, 0, atol=1e-8)
            @test p.iteration < 120
        end
    end
    @testset "Extreme Constraints" begin
        #@testset "camel6" begin
        #    p = DSProblem(2; poll=OrthoMADS(2), objective=DS.camel6, initial_point=[5.0, 5.0])
        #    AddExtremeConstraint(p, x-> x[1] > 0)
        #    Optimize!(p)
        #    @test isapprox(p.x[1], 0.0898, atol=1e-4)
        #    @test isapprox(p.x[2], -0.7126, atol=1e-4)
        #    @test isapprox(p.x_cost, -1.0316, atol=1e-4)
        #    @test p.iteration < 120

        #    p = DSProblem(2; poll=OrthoMADS(2), objective=DS.camel6, initial_point=[-5.0, -5.0])
        #    AddExtremeConstraint(p, x->x[1] < 0)
        #    Optimize!(p)
        #    @test isapprox(p.x[1], -0.0898, atol=1e-4)
        #    @test isapprox(p.x[2], 0.7126, atol=1e-4)
        #    @test isapprox(p.x_cost, -1.0316, atol=1e-4)
        #    @test p.iteration < 120
        #end
    end
    @testset "Progressive Constraints" begin
        #Progressive barrier constraints struggle to converge with LTMADS, therefore use OrthoMADS
        @testset "camel6" begin
            p = DSProblem(2; poll=OrthoMADS(2), objective=DS.camel6, initial_point=[5.0, 5.0])
            AddProgressiveConstraint(p, x -> x[1])
            Optimize!(p)
            @test isapprox(p.x[1], -0.0898, atol=1e-4)
            @test isapprox(p.x[2], 0.7126, atol=1e-4)
            @test isapprox(p.x_cost, -1.0316, atol=1e-4)
            @test p.iteration < 200

            p = DSProblem(2; poll=OrthoMADS(2), objective=DS.camel6, initial_point=[-5.0, -5.0])
            AddProgressiveConstraint(p, x->-x[1])
            Optimize!(p)
            @test isapprox(p.x[1], 0.0898, atol=1e-4)
            @test isapprox(p.x[2], -0.7126, atol=1e-4)
            @test isapprox(p.x_cost, -1.0316, atol=1e-4)
            @test p.iteration < 200
        end
    end
end


#---------------------------------------------------------
# # [Compressed Sensing 2D](@id 2-cs-wl-l1-2d)
#---------------------------------------------------------

#=
This example illustrates how to perform 2D compressed sensing image reconstruction
from Cartesian sampled MRI data
with 1-norm regularization of orthogonal wavelet coefficients,
using the Julia language.

This entire page was generated using a single Julia file:
[2-cs-wl-l1-2d.jl](@__REPO_ROOT_URL__/mri/2-cs-wl-l1-2d.jl).
=#
#md # In any such Julia documentation,
#md # you can access the source code
#md # using the "Edit on GitHub" link in the top right.

#md # The corresponding notebook can be viewed in
#md # [nbviewer](http://nbviewer.jupyter.org/) here:
#md # [`2-cs-wl-l1-2d.ipynb`](@__NBVIEWER_ROOT_URL__/mri/2-cs-wl-l1-2d.ipynb),
#md # and opened in [binder](https://mybinder.org/) here:
#md # [`2-cs-wl-l1-2d.ipynb`](@__BINDER_ROOT_URL__/mri/2-cs-wl-l1-2d.ipynb).

#=
This demo is somewhat similar to Fig. 3 in the survey paper
"[Optimization methods for MR image reconstruction](http://doi.org/10.1109/MSP.2019.2943645),"
in Jan 2020 IEEE Signal Processing Magazine,
except that the sampling is 1D phase encoding instead of 2D.
=#

# Packages used in this demo (run `Pkg.add` as needed):
using ImagePhantoms: shepp_logan, SheppLoganEmis, spectrum, phantom
using MIRT: embed, Afft, Aodwt
using MIRTjim: jim, prompt
using MIRT: pogm_restart
using LinearAlgebra: norm
using Plots; default(markerstrokecolor=:auto, label="")
using FFTW: fft
using Random: seed!
using InteractiveUtils: versioninfo

# The following line is helpful when running this jl-file as a script;
# this way it will prompt user to hit a key after each image is displayed.

isinteractive() && jim(:prompt, true);

# ## Create (synthetic) data

# Shepp-Logan phantom (unrealistic because real-valued):
nx,ny = 192,256
object = shepp_logan(SheppLoganEmis(); fovs=(ny,ny))
Xtrue = phantom(-(nx÷2):(nx÷2-1), -(ny÷2):(ny÷2-1), object, 2)
Xtrue = reverse(Xtrue, dims=2)
clim = (0,9)
jim(Xtrue, "true image"; clim)
#src savefig("xtrue.pdf")

# Somewhat random 1D phase-encode sampling:
seed!(0); sampfrac = 0.2; samp = rand(ny) .< sampfrac; sig = 1
mod2 = (N) -> mod.((0:N-1) .+ Int(N/2), N) .- Int(N/2)
samp .|= (abs.(mod2(ny)) .< Int(ny/8)) # fully sampled center rows
samp = trues(nx) * samp'
jim(samp, fft0=true, title="k-space sampling ($(100count(samp)/(nx*ny))%)")
#src savefig(p1, "samp.pdf")

# Generate noisy, under-sampled k-space data (inverse-crime simulation):
ytrue = fft(Xtrue)[samp]
y = ytrue + sig * √(2) * randn(ComplexF32, size(ytrue)) # complex noise!
y = ComplexF32.(y) # save memory
ysnr = 20 * log10(norm(ytrue) / norm(y-ytrue)) # data SNR in dB

# Display zero-filled data:
logger = (x; min=-6) -> log10.(max.(abs.(x) / maximum(abs, x), (10.)^min))
jim(:abswarn, false) # suppress warnings about showing magnitude
jim(logger(embed(ytrue,samp)), fft0=true, title="k-space |data| (zero-filled)",
    xlabel="kx", ylabel="ky")

#=
## Prepare to reconstruct
Creating a system matrix (encoding matrix) and an initial image
The system matrix is a `LinearMapAA` object, akin to a `fatrix` in Matlab MIRT.
=#

# System model ("encoding matrix") from MIRT:
F = Afft(samp) # operator!

# Initial image based on zero-filled reconstruction:
nrmse = (x) -> round(norm(x - Xtrue) / norm(Xtrue) * 100, digits=1)
X0 = 1.0f0/(nx*ny) * (F' * y)
jim(X0, "|X0|: initial image; NRMSE $(nrmse(X0))%")

#=
## Wavelet sparsity in synthesis form

The image reconstruction optimization problem here is
```math
\arg \min_{x}
\frac{1}{2} \| A x - y \|_2^2 + \beta \; \| W x \|_1
```
where
``y`` is the k-space data,
``A`` is the system model (simply Fourier encoding `F` here),
``W`` is an orthogonal discrete (Haar) wavelet transform,
again implemented as a `LinearMapAA` object.
Because ``W`` is unitary,
we make the change of variables
``z = W x``
and solve for ``z``
and then let ``x = W' z``
at the end.
In fact we use a weighted 1-norm
where only the detail wavelet coefficients are regularized,
not the approximation coefficients.
=#

# Orthogonal discrete wavelet transform operator (`LinearMapAO`):
W, scales, _ = Aodwt((nx,ny) ; T = eltype(F))
isdetail = scales .> 0
jim(
    jim(scales, "wavelet scales"),
    jim(real(W * Xtrue) .* isdetail, "wavelet detail coefficients"),
)

# Inputs needed for proximal gradient methods:
Az = F * W' # another operator!
Fnullz = (z) -> 0 # cost function in `z` not needed
f_gradz = (z) -> Az' * (Az * z - y)
f_Lz = nx*ny # Lipschitz constant for single coil Cartesian DFT
regz = 0.03 * nx * ny # oracle from Xtrue wavelet coefficients!
costz = (z) -> 1/2 * norm(Az * z - y)^2 + regz * norm(z,1) # 1-norm regularizer
soft = (z,c) -> sign(z) * max(abs(z) - c, 0) # soft thresholding
g_prox = (z,c) -> soft.(z, isdetail .* (regz * c)) # proximal operator (shrink details only)
z0 = W * X0
jim(z0, "Initial wavelet coefficients")

#=
## Iterate

Run ISTA=PGM and FISTA=FPGM and POGM, the latter two with adaptive restart
See [Kim & Fessler, 2018](http://doi.org/10.1007/s10957-018-1287-4)
for adaptive restart algorithm details.
=#

# Functions for tracking progress:
function fun_ista(iter, xk_z, yk, is_restart)
    xh = W' * xk_z
    return (costz(xk_z), nrmse(xh), is_restart) # , psnr(xh)) # time()
end

function fun_fista(iter, xk, yk_z, is_restart)
    xh = W' * yk_z
    return (costz(yk_z), nrmse(xh), is_restart) # , psnr(xh)) # time()
end;

# Run and compare three proximal gradient methods:
niter = 50
z_ista, out_ista = pogm_restart(z0, Fnullz, f_gradz, f_Lz; mom=:pgm, niter=niter,
    restart=:none, restart_cutoff=0., g_prox=g_prox, fun=fun_ista)
Xista = W'*z_ista
@show nrmse(Xista)

z_fista, out_fista = pogm_restart(z0, Fnullz, f_gradz, f_Lz; mom=:fpgm, niter=niter,
    restart=:gr, restart_cutoff=0., g_prox=g_prox, fun=fun_fista)
Xfista = W'*z_fista
@show nrmse(Xfista)

z_pogm, out_pogm = pogm_restart(z0, Fnullz, f_gradz, f_Lz; mom=:pogm, niter=niter,
    restart=:gr, restart_cutoff=0., g_prox=g_prox, fun=fun_fista)
Xpogm = W'*z_pogm
@show nrmse(Xpogm)

jim(
    jim(Xfista, "FISTA/FPGM"),
    jim(Xpogm, "POGM with ODWT"),
)
#src savefig("xpogm_odwt.pdf")

# ## POGM is fastest

# Plot cost function vs iteration:
cost_ista = [out_ista[k][1] for k=1:niter+1]
cost_fista = [out_fista[k][1] for k=1:niter+1]
cost_pogm = [out_pogm[k][1] for k=1:niter+1]
cost_min = min(minimum(cost_ista), minimum(cost_pogm))
plot(xlabel="iteration k", ylabel="Relative cost")
scatter!(0:niter, cost_ista  .- cost_min, label="Cost ISTA")
scatter!(0:niter, cost_fista .- cost_min, markershape=:square, label="Cost FISTA")
scatter!(0:niter, cost_pogm  .- cost_min, markershape=:utriangle, label="Cost POGM")
#src savefig("cost_pogm_odwt.pdf")

#
isinteractive() && prompt();

# Plot NRMSE vs iteration:
nrmse_ista = [out_ista[k][2] for k=1:niter+1]
nrmse_fista = [out_fista[k][2] for k=1:niter+1]
nrmse_pogm = [out_pogm[k][2] for k=1:niter+1]
plot(xlabel="iteration k", ylabel="NRMSE %", ylims=(3,6.5))
scatter!(0:niter, nrmse_ista, label="NRMSE ISTA")
scatter!(0:niter, nrmse_fista, markershape=:square, label="NRMSE FISTA")
scatter!(0:niter, nrmse_pogm, markershape=:utriangle, label="NRMSE POGM")
#src savefig("nrmse_pogm_odwt.pdf")

# Show error images:
p1 = jim(Xtrue, "true")
p2 = jim(X0, "X0: initial")
p3 = jim(Xpogm, "POGM recon")
p5 = jim(X0 - Xtrue, "X0 error", clim=(0,2))
p6 = jim(Xpogm - Xtrue, "Xpogm error", clim=(0,2))
jim(p2, p3, p5, p6)

# ### Reproducibility

# This page was generated with the following version of Julia:

io = IOBuffer(); versioninfo(io); split(String(take!(io)), '\n')

# And with the following package versions:

import Pkg; Pkg.status()


######### Stan program example  ###########

using StanSample

ProjDir = @__DIR__

diagnose_1 = "
parameters {
  real y;
  vector[9] x;
}
model {
  y ~ normal(0, 3);
  x ~ normal(0, exp(y/2));
}
";

diagnose_2 = "
parameters {
  real y_raw;
  vector[9] x_raw;
}
transformed parameters {
  real y;
  vector[9] x;

  y = 3.0 * y_raw;
  x = exp(y/2) * x_raw;
}
model {
  y_raw ~ std_normal(); // implies y ~ normal(0, 3)
  x_raw ~ std_normal(); // implies x ~ normal(0, exp(y/2))
}
";

sm1 = SampleModel("diagnose.1", diagnose_1);
rc1 = stan_sample(sm1);
if success(rc1)
  read_summary(sm1, true)
  diagnose(sm1)
end

sm2 = SampleModel("diagnose.2", diagnose_2);
rc2 = stan_sample(sm2);
if success(rc2)
  read_summary(sm2, true)
  diagnose(sm2)
end


module NonNegLeastSquaresMLJInterface

using MLJModelInterface
using NonNegLeastSquares: nonneg_lsq

################################
########### Structure ##########
################################

MLJModelInterface.@mlj_model mutable struct NonNegativeLeastSquareRegressor <: MLJModelInterface.Deterministic
    alg::Symbol = :pivot::(_ in (:pivot, :nnls, :fnnls, :admm))
    variant::Symbol = :none::(_ in (:comb, :cache, :none))
    gram::Bool = false
    fit_intercept::Bool = true
end

################################
########### Methods ############
################################

"""
    augment_X(X, b)
Augment the matrix `X` with a column of ones if the intercept is to be fitted (`b=true`), return
`X` otherwise.
"""
function augment_X(X::Matrix, b::Bool)::Matrix
    b && return hcat(X, ones(eltype(X), size(X, 1), 1))
    return X
end

function MLJModelInterface.fit(m::NonNegativeLeastSquareRegressor, verbosity::Int, X, y)
    X = augment_X(MLJModelInterface.matrix(X), m.fit_intercept)
    fitresult = vec(nonneg_lsq(X, y;
                                alg=m.alg,
                                variant=m.variant,
                                gram=m.gram))
    cache = nothing
    report = NamedTuple{}()
    return (fitresult, cache, report)
end

function MLJModelInterface.fitted_params(model::NonNegativeLeastSquareRegressor, fitresult)
    return (coefs      = fitresult[1:end-Int(model.fit_intercept)],
            intercept = ifelse(model.fit_intercept, fitresult[end], nothing))
end

function MLJModelInterface.predict(m::NonNegativeLeastSquareRegressor, fitresult, Xnew)
    Xnew = augment_X(MLJModelInterface.matrix(Xnew), m.fit_intercept)
    ypred = Xnew * fitresult
    return ypred
end

################################
########### METADATA ###########
################################

const ALL_MODELS = Union{NonNegativeLeastSquareRegressor}

MLJModelInterface.metadata_pkg.(ALL_MODELS,
    name       = "NonNegLeastSquares",
    uuid       = "b7351bd1-99d9-5c5d-8786-f205a815c4d",
    url        = "https://github.com/ahwillia/NonNegLeastSquares.jl",
    julia      = true,
    license    = "MIT",
    is_wrapper = false,
)

MLJModelInterface.metadata_model(NonNegativeLeastSquareRegressor,
    input_scitype    = MLJModelInterface.Table(MLJModelInterface.Continuous),
    target_scitype   = AbstractVector{MLJModelInterface.Continuous},
    output_scitype   = MLJModelInterface.Table(MLJModelInterface.Continuous),
    supports_weights = false,
    descr            = "A non negative least square regressor model",
	load_path        = "NonNegLeastSquaresMLJInterface.NonNegativeLeastSquareRegressor"
    )

end


using BenchmarkTools
using LinearAlgebra
using Newton
using NLsolve
using StableRNGs

LinearAlgebra.BLAS.set_num_threads(1)

function create_rf!(dim, rng=StableRNG(0))
    A = -rand(rng, dim, dim)
    b = -rand(rng, dim)

    function rf!(r, x)
        r .= b
        mul!(r, A, x, 1, 1)
        @simd for i=1:length(x)
            @inbounds r[i] += exp(x[i]) - x[i]^2
        end
        return r
    end

    return rf!
end

dim = 20
rf = create_rf!(dim)
r, x = [zeros(dim) for _ in 1:2]

# Check that rf! is non-allocating
println("@btime rf!")
@btime rf(r_, x_) setup=(r_=copy(r); x_=copy(x)) evals=1;

# Get time for newtonsolve!
# First time setup
cache = NewtonCache(x, rf)
drdx = get_drdx(cache)
println("@btime newtonsolve!")
@btime newtonsolve!(x_, $drdx, rf, $cache) setup=(x_=copy(x)) evals=1;

# Get time for nlsolve
println("@btime nlsolve")
@btime nlsolve(rf, x_, autodiff=:forward) setup=(x_=copy(x)) evals=1;

println("Benchmark (dim=$dim) complete")


using Printf

function hist(numbers)
    maxwidth = 50
    h = fill(0, 10)
    for n in numbers
        h[ceil(Int, 10n)] += 1
    end
    mx = maximum(h)
    for (n, i) in enumerate(h)
        @printf("%3.1f: %s\n", n / 10, "+" ^ floor(Int, i / mx * maxwidth))
    end
end

for i in 1:6
    n = rand(10 ^ i)
    println("\n##\n## $(10 ^ i) numbers")
    @printf("μ: %8.6f; σ: %8.6f\n", mean(n), std(n))
    hist(n)
end


# Sum of the positive components

"""
  SumPositive()

Returns the function `g(x) = sum(max(0, x))`.
"""

immutable SumPositive <: ProximableFunction end

function (f::SumPositive){T <: Real}(x::AbstractArray{T})
  return sum(max.(0.0, x))
end

function prox!{T <: Real}(f::SumPositive, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  fsum = 0.0
  for i in eachindex(x)
    y[i] = x[i] < gamma ? (x[i] > 0.0 ? 0.0 : x[i]) : x[i]-gamma
    fsum += y[i] > 0.0 ? y[i] : 0.0
  end
  return fsum
end

fun_name(f::SumPositive) = "Sum of the positive coefficients"
fun_dom(f::SumPositive) = "AbstractArray{Real}"
fun_expr(f::SumPositive) = "x ↦ sum(max(0, x))"

function prox_naive{T <: Real}(f::SumPositive, x::AbstractArray{T}, gamma::Real=1.0)
  y = copy(x)
  indpos = x .> 0.0
  y[indpos] = max.(0.0, x[indpos]-gamma)
  return y, sum(max.(0.0, y))
end

# ######################### #
# Prox with multiple gammas #
# ######################### #

function prox!{T <: Real}(f::SumPositive, x::AbstractArray{T}, y::AbstractArray{T}, gamma::AbstractArray{T})
  fsum = 0.0
  for i in eachindex(x)
    y[i] = x[i] < gamma[i] ? (x[i] > 0.0 ? 0.0 : x[i]) : x[i]-gamma[i]
    fsum += y[i] > 0.0 ? y[i] : 0.0
  end
  return fsum
end

function prox_naive{T <: Real}(f::SumPositive, x::AbstractArray{T}, gamma::AbstractArray{T})
  y = copy(x)
  indpos = x .> 0.0
  y[indpos] = max.(0.0, x[indpos]-gamma[indpos])
  return y, sum(max.(0.0, y))
end


"""
    sic(v::VAREstimate)

Return the Schwarz (Bayesian) Information Criterion for VAR model `v`.
"""
function sic(v::VAREstimate)
    obs, K = size(residuals(v))
    nparam = size(coef(v), 1)
    ΣU = v.ΣU
    ΣUml = ((obs - nparam) / obs) * ΣU
    return logdet(ΣUml) + (log(obs) / obs) * nparam * K
end

"""
    hqc(v::VAREstimate)

Return the Hannan-Quinn Criterion for VAR model `v`.
"""
function hqc(v::VAREstimate)
    obs, K = size(residuals(v))
    nparam = size(coef(v), 1)
    ΣU = v.ΣU
    ΣUml = ((obs - nparam) / obs) * ΣU
    return logdet(ΣUml) + (2 * log(log(obs)) / obs) * nparam * K
end

"""
    StabilityCheck(v::VAREstimate) -> StabilityCheck

Check the stability of the estimated VAR model `v`.

# Fields
- `isstable::Bool`
- `eigenvals::Array{Number}`
- `eigenmod::Array{Float64}`
"""
struct StabilityCheck
    isstable::Bool
    eigenvals::Array{Number}
    eigenmod::Array{Float64}
end

function StabilityCheck(v::VAREstimate)
    # See Lutkepohl (2006) p.15ff
    p = v.lags
    B = coef(v)
    nparam, K = size(B)
    startidx = nparam - K * p + 1
    A = B[startidx:end, :]
    Acomp = companion(A, p)
    eigenvals = eigvals(Acomp)
    eigenmod = abs.(eigenvals)
    isstable = all(eigenmod .< 1)
    StabilityCheck(isstable, eigenvals, eigenmod)
end

function show(io::IO, obj::StabilityCheck)
    # K = length(stab.ynames)
    E = obj.eigenvals
    Emod = round.(unique(obj.eigenmod), digits = 3)
    println(io, typeof(obj))
    if obj.isstable
        println(io, "VAR is stable")
    else
        println(io, "VAR is not stable")
    end
    println(io, "Modulus of eigenvalues (dropping repeated values):")
    println(io, sort(Emod))
end

"""
    PortmanteauTest(v::VAREstimate, h::Int) -> PortmanteauTest

Conduct a multivariate portmanteau test for residual autocorrelation up to lag
    `h` for VAR model `v`.
Implements the adjusted test described on p.171 of Lütkepohl (2006).

# Fields
- `U::Matrix`
- `h::Int`
- `Q::Float64`
- `df::Int`
- `pval::Float64`
"""
struct PortmanteauTest
    U::Matrix
    h::Int
    Q::Float64
    df::Int
    pval::Float64
end

function PortmanteauTest(v::VAREstimate, h::Int)
    U = residuals(v)
    T, K = size(U)
    p = v.lags
    C = permutedims(crosscov(U, U, collect(0:h), demean = false), [3, 2, 1])
    Q = 0.0
    for i = 1:h
        Q += tr(C[:, :, i+1]' * (C[:, :, 1] \ C[:, :, i+1] / C[:, :, 1])) * (T^2 / (T - i))
    end
    df = K^2 * (h - p)
    pval = ccdf(Chisq(df), Q)
    PortmanteauTest(U, h, Q, df, pval)
end

function show(io::IO, obj::PortmanteauTest)
    println(io, typeof(obj))
    println(io, "Multivariate test for residual autocorrelation.")
    println(io, "- Null is no autocorrelation")
    println(io, "- Lags: $(obj.h)")
    println(io, "- Q stat: $(round(obj.Q; digits=3))")
    println(io, "- p-value: $(round(obj.pval; digits=3))")
    outcome = obj.pval > 0.05 ? "Fail to reject" : "Reject"
    println(io, "$outcome the null at the 5% level")
end

"""
    LMCorrTest(v::VAREstimate, h::Int; smallsample = false) -> LMCorrTest

Conduct a multivariate LM test (Breusch-Godfrey test) for residual autocorrelation
    up to lag `h` for VAR model `v`.
If `smallsample = true`, the alternative test statistic detailed in Lütkepohl (2006),
    p.173 is used.

# Fields
- `h::Int`
- `Q::Float64`
- `df1::Any`
- `df2::Any`
- `pval::Float64`
- `smallsample::Bool`
"""
struct LMCorrTest
    h::Int
    Q::Float64
    df1::Any
    df2::Any
    pval::Float64
    smallsample::Bool
end

function LMCorrTest(v::VAREstimate, h::Int; smallsample = false)
    U = residuals(v)
    T, K = size(U)
    p = v.lags
    nparam = size(v.B, 1)
    Z = v.Z

    U2 = vcat(zeros(p + (h - p), K), U)
    Ulag = VectorAutoregressions.lag_matrix(U2, h)
    ZUlag = [Z Ulag]
    AD = (ZUlag' * ZUlag) \ (ZUlag' * U)
    E = U - ZUlag * AD

    ΣUml = ((T - nparam) / T) * v.ΣU
    ΣEml = (E' * E) / T

    if smallsample
        # Test stat from Lütkepohl 2006, p.173
        s = sqrt((K^4 * h^2 - 4) / (K^2 + (K^2 * h^2) - 5))
        N = T - (K*p) - 1 - (K*h) - (K - K*h +1)/2
        df1 = h * K^2
        df2 = (N * s) - ((K^2 * h) / 2) + 1
        Q = ((det(ΣUml) / det(ΣEml))^(1 / s) - 1) * (df2 / df1)
        pval = ccdf(FDist(df1, df2), Q)
    else
        # Test stat from Killian and Lütkepohl 2017, p.54
        Q = T * (K - tr(ΣUml \ ΣEml))
        df1 = h * K^2
        df2 = nothing
        pval = ccdf(Chisq(df1), Q)
    end
    return LMCorrTest(h, Q, df1, df2, pval, smallsample)
end

function show(io::IO, obj::LMCorrTest)
    println(io, typeof(obj))
    println(io, "Multivariate LM test for residual autocorrelation.")
    println(io, "- Null is no autocorrelation")
    println(io, "- Lags: $(obj.h)")
    println(io, "- Q stat: $(round(obj.Q; digits=3))")
    println(io, "- p-value: $(round(obj.pval; digits=3))")
    if obj.smallsample
        println(io, "- Small sample test")
    end
    outcome = obj.pval > 0.05 ? "Fail to reject" : "Reject"
    println(io, "$outcome the null at the 5% level")
end

struct ResidualNormalityTest
    λ_s::Float64
    λ_k::Float64
    λ_sk::Float64
    pval1::Float64
    pval2::Float64
    pval3::Float64
end

function ResidualNormalityTest(v::VAREstimate)
    U = residuals(v)
    ΣU = v.ΣU
    λs, λk, λsk, pval1, pval2, pval3 = multivariate_normality_test(U, ΣU)
    ResidualNormalityTest(λs, λk, λsk, pval1, pval2, pval3)
end

function multivariate_normality_test(U, Σ)
    # see Lutkepohl pp. 177-181
    uobs, K = size(U)
    P   = cholesky(Σ).L
    W   = transpose(P\U')
    W3  = W.^3
    W4  = W.^4
    b1  = dropdims(sum(W3, dims=1)/uobs, dims=1)
    b2  = dropdims(sum(W4, dims=1)/uobs, dims=1)
    λs = dot(b1, b1)*(uobs/6)
    λk = dot(b2 .- 3, b2 .- 3)*(uobs/24)
    λsk = λs + λk

    pval1 = 1 - cdf(Chisq(K), λs)
    pval2 = 1 - cdf(Chisq(K), λk)
    pval3 = 1 - cdf(Chisq(2K), λsk)

    return λs, λk, λsk, pval1, pval2, pval3
end


@compat abstract type StepSizeRule end
mutable struct constantStep <: StepSizeRule
    α::Float64
end
mutable struct constantLength <: StepSizeRule
    α::Float64
end
mutable struct poliakRule <: StepSizeRule
    η::Float64
end
mutable struct diminishingAffine <: StepSizeRule
    a::Float64
    b::Float64
end
mutable struct diminishingSqrt <: StepSizeRule
    a::Float64
end

function computeStepSize(rule::constantStep, epoch, ∇q, UB, LB)
  return rule.α
end
function computeStepSize(rule::constantLength, epoch, ∇q, UB, LB)
  return rule.α / norm(∇q)
end
function computeStepSize(rule::poliakRule, epoch, ∇q, UB, LB)
  return rule.η*(UB-LB)/ dot(∇q, ∇q)
end
function computeStepSize(rule::diminishingAffine, epoch, ∇q, UB, LB)
  return rule.a / (rule.b + epoch)
end
function computeStepSize(rule::diminishingSqrt, epoch, ∇q, UB, LB)
  return rule.a / sqrt(epoch)
end

function slashStep!(rule::constantStep, factor)
  rule.α /= factor
end
function slashStep!(rule::constantLength, factor)
  rule.α /= factor
end
function slashStep!(rule::poliakRule, factor)
  rule.η /= factor
end
function slashStep!(rule::diminishingAffine, factor)
  rule.a /= factor
end
function slashStep!(rule::diminishingSqrt, factor)
  rule.a /= factor
end


using StochasticDiffEq, DiffEqNoiseProcess, Test, Random
Random.seed!(100)
α=1
β=1
u₀=1/2
f(u,p,t) = α*u
g(u,p,t) = β*u
dt = 1//2^(4)
tspan = (0.0,1.0)
prob = SDEProblem(f,g,u₀,(0.0,1.0))
sol =solve(prob,EulerHeun(),dt=0.001,save_noise=true)
_sol = deepcopy(sol) # to make sure the plot is correct
W3 = NoiseGrid(reverse!(_sol.t),reverse!(_sol.W))
prob3 = SDEProblem(f,g,sol[end],(1.0,0.0),noise=W3)
sol2 = solve(prob3,EulerHeun(),dt=0.001)
@test sol.u ≈ reverse!(sol2.u) atol=1e-1


using LinearAlgebra, SparseArrays
import SuiteSparse

function Successive_solves!( x, Kc, b, nsolves )
    for i=1:nsolves
        x  .= Kc\b
    end
    return
end

function ThreePointStencil( N, dx )
    # Assemble matrix (not the part of interest)
    cC =  2.0/dx^2 * ones(Float64, N); cC[1]   = 3.0/dx^2; cC[end]   = 3.0/dx^2
    cW = -1.0/dx^2 * ones(Float64, N); cW[1]   = 0.0
    cE = -1.0/dx^2 * ones(Float64, N); cE[end] = 0.0
    iC = 1:N
    iW = ones(Int64, N); iW[2:end-0] .= iC[1:end-1]
    iE = ones(Int64, N); iE[1:end-1] .= iC[2:end-0]
    I  = [iC[:]; iC[:]; iC[:]]
    J  = [iC[:]; iW[:]; iE[:]]
    V  = [cC[:]; cW[:]; cE[:]]
    K  = sparse( I, J, V )
    droptol!(K, 1e-9)
    return K
end

# Input
N  = 1000000
dx = 1

# Sparse matrix
K = ThreePointStencil( N, dx )

# Allocate arrays
b = rand(N)           # Right-hand-side array
x = zeros(Float64, N) # Solution array

# Factor
Kc = cholesky(K)

for i=1:3
@time @views x  .= Kc\b
@time @views x  .= SuiteSparse.CHOLMOD.solve(0, Kc, SuiteSparse.CHOLMOD.Dense(b) )[:]
end
xcm = SuiteSparse.CHOLMOD.zeros(N,1)
bcm = SuiteSparse.CHOLMOD.Dense(b)

for i=1:3
    @time @views x  .= Kc\b
    xcm  = SuiteSparse.CHOLMOD.solve(0, Kc, bcm )
end

# SuiteSparse.CHOLMOD.spsolve(Kc,SparseArrays.CHOLMOD.Sparse(b))
# ldiv!(Kc, b)
@time x  .= lu(K)\b

# # One solve
# @time x  .= Kc\b  # why does this allocate memory although the factorisation done and the vectors (x, b) are allocated?

# # Successive solves
# nsolves = 10
# @time Successive_solves!( x, Kc, b, nsolves ) # then, memory allocation depends on the number of solves...


module PkgTests

using LinearAlgebra
using InteractiveUtils

using SparseArrays
using StaticArrays
using Test

import CompScienceMeshes

@test CompScienceMeshes.point(1,2,3) === SVector{3,Float64}(1.0,2.0,3.0)
@test CompScienceMeshes.point(Int,1,2,3) == SVector{3,Int}(1,2,3)

include("test_relorientation.jl")
include("test_mesh.jl")
include("test_interior.jl")
include("test_scomplex.jl")
include("test_flipped.jl")
include("test_flipped_welding.jl")
include("test_weld.jl")
include("test_baryref.jl")
include("test_gid_reader.jl")
include("test_gmsh_reader.jl")
include("test_physical_entities.jl")

include("test_geometry.jl")

include("test_patches.jl")
include("test_submesh.jl")
include("test_overlap.jl")
include("test_intersect.jl")
include("test_sh_intersection.jl")
include("test_isinside.jl")
include("test_jctweld.jl")
include("test_isinclosure.jl")

include("test_trgauss.jl")
include("test_chartquad.jl")
include("test_spherequad.jl")
include("test_nbd.jl")

include("test_mapper.jl")

end


# Illustrative examples from Flux.jl

using Solaris
using Solaris.Flux, Solaris.DiffEqFlux

"""
gradient: scalar -> scalar
"""
gradient(x -> 3x^2 + 2x + 1, 5) # (32, )
gradient((a, b) -> a * b, 2, 3) # (3, 2)

"""
gradient: scalar -> vector

    f = (x₁ - y₁)² + (x₂ - y₂)²
    ∂f/∂x₁ = 2x₁ - 2y₁
    ⋯
"""
f(x, y) = sum((x .- y) .^ 2)
gradient(f, [2, 1], [2, 0])

"""
params: a syntax sugar for machine learning with thousands of parameters
"""
x = [2, 1]
y = [2, 0]
gs = gradient(params(x, y)) do
    f(x, y)
end
gs[x]
gs[y]

"""
toy model
"""
W = rand(2, 5)
b = rand(2)

predict(x) = W * x .+ b
function loss(x, y)
    ŷ = predict(x)
    sum((y .- ŷ) .^ 2)
end

X, Y = rand(5), rand(2)
L0 = loss(X, Y)

gs = gradient(() -> loss(X, Y), params(W, b))
W̄ = gs[W]
W .-= 0.1 .* W̄
L1 = loss(X, Y)

@assert L1 < L0

"""
affine chain
"""
m = Chain(
    Solaris.Affine(28^2, 128, relu),
    Solaris.Affine(128, 32, relu),
    Solaris.Affine(32, 10),
)

# regularization
sqnorm(x) = sum(abs2, x)
loss(x, y) = logitcrossentropy(m(x), y) + sum(sqnorm, Flux.params(m))

loss(rand(Float32, 28^2), rand(Float32, 10))


module SerendipityRefFEsTests

using Test
using Gridap.Fields
using Gridap.ReferenceFEs
using Gridap.Polynomials
using Gridap.Io

reffe = SerendipityRefFE(Float64,QUAD,2)
@test get_node_coordinates(reffe) == Point{2,Float64}[
  (0.0, 0.0), (1.0, 0.0), (0.0, 1.0), (1.0, 1.0),
  (0.5, 0.0), (0.5, 1.0), (0.0, 0.5), (1.0, 0.5)]
@test get_face_own_nodes(reffe) == [[1], [2], [3], [4], [5], [6], [7], [8], Int[]]
test_lagrangian_reference_fe(reffe)

reffe = SerendipityRefFE(Float64,QUAD,4)
@test get_face_own_nodes(reffe) == [[1], [2], [3], [4], [5, 6, 7], [8, 9, 10], [11, 12, 13], [14, 15, 16], [17]]
test_lagrangian_reference_fe(reffe)

reffe = SerendipityRefFE(Float64,HEX,2)
@test get_face_own_nodes(reffe) == [
  [1], [2], [3], [4], [5], [6], [7], [8],
  [9], [10], [11], [12], [13], [14], [15], [16],
  [17], [18], [19], [20], Int[], Int[], Int[], Int[], Int[], Int[], Int[]]
test_lagrangian_reference_fe(reffe)

p = get_polytope(reffe)
test_polytope(p)
@test LagrangianRefFE(p) == HEX8

order = 4
reffe = SerendipityRefFE(Float64,HEX,order)
@test get_order(reffe) == order
@test get_orders(reffe) == (order,order,order)
@test is_P(reffe) == false
@test is_Q(reffe) == false
@test is_S(reffe) == true

reffe = LagrangianRefFE(Float64,HEX,order,space=:S)
@test get_order(reffe) == order
@test get_orders(reffe) == (order,order,order)
@test is_P(reffe) == false
@test is_Q(reffe) == false
@test is_S(reffe) == true

reffe = SerendipityRefFE(Float64,QUAD,(3,3))
@test reffe == from_dict(LagrangianRefFE,to_dict(reffe))

end # module


# ------------------------------------------------------------------
# Copyright (c) 2018, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    pointgradients(img, points; [method])

Compute the gradients along all dimensions of N-dimensional `img`
at `points` using `method`. Default method is `:ando3`.
"""
function pointgradients(img::AbstractArray, points::AbstractVector; method=:ando3)
  extent = size(img)
  ndirs = length(extent)
  npoints = length(points)

  # smoothing weights
  weights = (method == :sobel ? [1,2,1] :
             method == :ando3 ? [.112737,.274526,.112737] :
             error("Unknown gradient method: $method"))

  # pad input image
  padimg = padarray(img, Pad(:replicate, ones(Int, ndirs), ones(Int, ndirs)))

  # gradient matrix
  G = zeros(npoints, ndirs)

  # compute gradient for all directions at specified points
  for i in 1:ndirs
    # kernel = centered difference + perpendicular smoothing
    if extent[i] > 1
      # centered difference
      idx = ones(Int, ndirs); idx[i] = 3
      kern = reshape([-1,0,1], idx...)
      # perpendicular smoothing
      for j in setdiff(1:ndirs, i)
        if extent[j] > 1
          idx = ones(Int, ndirs); idx[j] = 3
          kern = broadcast(*, kern, reshape(weights, idx...))
        end
      end

      A = zeros(size(kern))
      shape = size(kern)
      for (k, icenter) in enumerate(points)
        i1 = CartesianIndex(ntuple(i->1, ndirs))
        for ii in CartesianIndices(shape)
          A[ii] = padimg[ii + icenter - i1]
        end

        G[k,i] = sum(kern .* A)
      end
    end
  end

  G
end


using ForwardDiff

# Define Test Problems
# ====================

tests = []

# [1] some random 2D problem
f1(x) = [x[1] + x[2] + x[1]^2, x[2] + x[1]*x[2]]
df1(x) = ForwardDiff.jacobian(f1, x)
init1() = [1.0,1.0]
randinit1() = rand(2)
push!(tests, (f1, df1, init1, randinit1, "toy2d"))

# [2] Chandrasekhar H-equation; cf http://www4.ncsu.edu/~ctk/newton/Chapter3/heq.m
let c = 0.9, n = 100
   gr = ((1:n) .- 0.5) / n
   A_heq = gr * ones(n)'
   A_heq = (0.5 * c / n) * A_heq ./ (A_heq+A_heq')
   f2(x) = x - ones(n) ./ (ones(n) - (A_heq * x))
   df2(x) = ForwardDiff.jacobian(f2, x)
   init2() = ones(n)
   randinit2() = ones(n) + 0.1 * (rand(n) .- 0.5)
   push!(tests, (f2, df2, init2, randinit2, "H-equation"))
end

# [3] MCP example by Josephy - stolen from NLsolve.jl
f3(x) = [3*x[1]^2+2*x[1]*x[2]+2*x[2]^2+x[3]+3*x[4]-6,
         2*x[1]^2+x[1]+x[2]^2+3*x[3]+2*x[4]-2,
         3*x[1]^2+x[1]*x[2]+2*x[2]^2+2*x[3]+3*x[4]-1,
         x[1]^2+3*x[2]^2+2*x[3]+3*x[4]-3 ]
df3(x) = ForwardDiff.jacobian(f3, x)
init3() = zeros(4)
randinit3() =  [1.22474487, 0., 0., 0.5 ] + (rand(4) .-0.5)
push!(tests, (f3, df3, init3, randinit3, "Josephy"))


using QuantileRegressions, CSV
@testset "npqreg" begin
	mycyle = CSV.read("../examples/mcycle.csv", DataFrame)
	x = mycyle.times
	y = mycyle.accel

	for t = range(0.01,0.99,length=99)
		res = QuantileRegressions.npqreg(Array(y), Array(x), t; h=2)
		@test all(!isnan, res[2])
		@test all(!isnan, res[3])
		xrange = sort(unique(x))
		res = QuantileRegressions.npqreg(Array(y), Array(x), t; xrange=xrange, h=0.5)
	end
	res = QuantileRegressions.npqreg([y[1]], [x[1]], 0.4; xrange=[x[1]], h=2)
	@test !isnan(res[2][1])
	@test isnan(res[3][1])
	res = QuantileRegressions.npqreg([y[1]], [x[1]], 0.8; xrange=[x[1]], h=0.0001)
	@test !isnan(res[2][1])
	@test isnan(res[3][1])
	res = QuantileRegressions.npqreg([y[1]], [x[1]], 0.8; xrange=[x[1].+1], h=0.1)
	@test isnan(res[2][1])
	@test isnan(res[3][1])

	@test all(QuantileRegressions.npqreg([1.0,2.0,4.0,5.0], [0.0,0.0,6.0,12.0], 0.3; h=0.001, xrange=[0.0,6.0,12.0])[2] .≈ [1.0, 4.0, 5.0])
end


# streamplots

#makie
using CairoMakie
let
    nonStablePoincare(x, y) = Point2f(x * (x^2 + y^2 - 1) - y * (x^2 + y^2 + 1), y * (x^2 + y^2 - 1) + x * (x^2 + y^2 + 1))
    stableVanDerPaul(x, y) = Point2f(y, (1 - x^2) * y - x)
    semiStable(x, y) = Point2f(-y + x * (-1 + x^2 + y^2)^2, x + y * (-1 + x^2 + y^2)^2)
    titles = ["non-stable", "stable", "semi-stable"]
    functions = [nonStablePoincare, stableVanDerPaul, semiStable]
    cmaps = [:plasma, :viridis, :gnuplot2]

    fig = Figure(resolution = (1000, 400), fontsize = 22)
    axs = [Axis(fig[1, i], xlabel = "x", ylabel = "y", title = titles[i],
        aspect = 1, backgroundcolor = :black) for i in 1:3]
    [streamplot!(axs[i], functions[i], -4 .. 4, -4 .. 4, colormap = Reverse(cmaps[i]),
        gridsize = (32, 32), arrow_size = 10) for i in 1:3]
    [hideydecorations!(axs[2], grid = false, ticks = false) for i in 2:3]
    [hideydecorations!(axs[3], grid = false, ticks = false) for i in 3:3]
    [limits!(axs[i], -4, 4, -4, 4) for i in 1:3]
    display(fig)
end


# Implementation of Dongarra et al., "Improving the accuracy...," SINUM 1983

# find peak index
function _normalizeInf!(x)
    n = length(x)
    s=1
    xm = abs(x[1])
    for j=1:n
        t = abs(x[j])
        if t > xm
            xm = t
            s = j
        end
    end
    x ./= xm
    s
end

# simple version for an isolated eigenvalue
"""
    rfeigen(A,x,λ,DT) => λnew, xnew, status

Improve the precision of a computed eigenpair `(x,λ)` for matrix `A`
via multi-precision iterative refinement, using more-precise real type `DT`.

The higher precision `DT` is only used for residual computation
(i.e. matrix-vector products), so this can be much faster than a full
eigensystem solution with precise eltype.  This method works on a
single eigenpair, and can fail spectacularly if there is another
eigenvalue nearby.
"""
function rfeigen(A::AbstractMatrix{T},
                 x::AbstractVector{Tx},
                 λ::Tλ,
                 DT::Type{<:AbstractFloat} = widen(real(T));
                 maxiter=5,
                 tol=1,
                 factor = lu,
                 scale = true,
                 verbose = false
                 ) where {T,Tλ,Tx}
    Tr = promote_type(promote_type(Tx,DT),Tλ)
    res = _rfeigen(A, x, λ, Tr, factor, maxiter, tol, scale, verbose)
    return res
end

"""
    rfeigen(A,λ,DT) => λnew, xnew, status

Like `rfeigen(A,x,λ,DT)`, but initialize `x` via one step of inverse
iteration.
"""
function rfeigen(A::AbstractMatrix{T},
                 λ::Tλ,
                 DT::Type{<:AbstractFloat} = widen(real(T));
                 maxiter=5,
                 tol=1,
                 factor = lu,
                 scale = true,
                 verbose = false
                 ) where {T,Tλ}
    # CHECKME: is this condition adequate?
    if issymmetric(A) && (Tλ <: Real)
        Tx = Tλ
    else
        Tx = complex(Tλ)
    end
    # There may not be a sampler for types of interest (hello, Quadmath)
    # so let's promote.
    x = normalize!((A - λ * I) \ Tx.(rand(size(A,1))))
    Tr = promote_type(promote_type(Tx,DT))
    res = _rfeigen(A, x, λ, Tr, factor, maxiter, tol, scale, verbose)
    return res
end

function _rfeigen(A::AbstractMatrix{T},
                      x::AbstractVector{Tx},
                      λ::Tλ,
                      ::Type{DT},
                      factor,
                      maxiter,
                      tol,
                      scale,
                      verbose
                      ) where {T,Tx,Tλ,DT}

    status = :unfinished
    λd = convert(DT,λ)
    n = LinearAlgebra.checksquare(A)
    tol1 = tol * eps(real(DT)) # CHECKME: factor of n?
    Ad = DT.(A)
    B = Ad - λd * I

    s = _normalizeInf!(x)
    xd = DT.(x)
    B[:,s] .= -xd
    Btmp = A - λ * I # do it over to get the type right
    Btmp[:,s] .= -x
    FB = factor(Btmp)

    # initial residual
    r::Vector{DT} = λd * xd - Ad * xd
    y = zeros(Tx,n)
    ys = zero(Tx)
    δ = similar(y)
    δp = similar(y)
    yp = similar(y)
    rt = similar(y)
    prevnorm = convert(real(DT),Inf)
    for p = 1:maxiter
        verbose && println("iter $p resnorm: ",norm(r))
        δ = FB \ Tx.(r) # ldiv!(δ,FB,Tx.(r))
        y .= y .+ δ
        δnorm = norm(δ)
        ynorm = norm(y)
        yp .= y
        ys = y[s]
        yp[s] = zero(T)
        if δnorm > prevnorm
            if δnorm > 10.0 * prevnorm
                status = :diverging
            else
                status = :stalled
            end
            verbose && println("$status at iter $p; early exit")
            break
        end
        prevnorm = δnorm
        if δnorm / ynorm < tol1
            status = :converged
            verbose && println("converged")
            # println("iter $p ratio ",δnorm / ynorm)
            break
        end
        δp .= δ
        δs = δ[s]

        r .= r .- B * DT.(δ)
        δp[s] = zero(T)
        r .= r .+ ys * δp .+ δs * y
    end
    xnew = x + yp
    λnew = λ + ys
    return λnew, xnew, status
end

"""
    rfeigen(A, S::Schur, idxλ, DT, maxiter=5) -> vals, vecs, status

Improves the precision of a cluster of eigenvalues of matrix `A`
via multi-precision iterative refinement, using more-precise real type `DT`.
Returns improved estimates of eigenvalues and vectors generating
the corresponding invariant subspace.

This method works on the set of eigenvalues in `S.values` indexed by `idxλ`.
It is designed to handle (nearly) defective cases, but will fail if
the matrix is extremely non-normal or the initial estimates are poor.
Note that `S` must be a true Schur decomposition, not a "real Schur".
"""
function rfeigen(A::AbstractMatrix{T}, S::Schur{TS}, idxλ,
                 DT = widen(real(T)), maxiter=5;
                 tol = 1, verbose = false) where {T, TS <: Complex}
    n = size(A,1)
    m = length(idxλ)
    λ = [S.values[idxλ[i]] for i in 1:m]
    Tw = promote_type(T,eltype(λ))
    DTw = promote_type(DT,Tw)
    tol1 = tol * eps(real(DT)) # CHECKME: factor of n?
    status = :unfinished

    # compute X, M
    # X is an adequately-conditioned set spanning the invariant subspace
    # M is an upper-triangular matrix of mixing coefficients

    # Most of the work is in the space of Schur vectors
    Z = zeros(Tw, n, m)
    z = zeros(Tw, n)
    idxz = Vector{Int}(undef, m)

    k = idxλ[1]
    if k==1
        z[1] = one(Tw)
    else
        x0 = (S.T[1:k-1,1:k-1] - λ[1] * I) \ S.T[1:k-1,k]
        z[1:k-1] .= -x0
        z[k] = one(Tw)
    end
    zm,zi = findmax(abs.(z))
    z ./= zm
    idxz[1] = zi
    Z[:,1] .= z
    M = zeros(Tw, m, m)
    M[1,1] = λ[1]
    for l=2:m
        kp = k
        k = idxλ[l]
        @assert k > kp
        x0 = (S.T[1:k-1,1:k-1] - λ[l]*I) \ S.T[1:k-1,k]
        X1 = (S.T[1:k-1,1:k-1] - λ[l]*I) \ Z[1:k-1,1:l-1]
        z[1:k-1] .= -x0
        z[k] = one(Tw)
        # pick mixing coeffts so that each vector has a good dominant index
        rhs = [-z[idxz[i]] for i=1:l-1]
        mtmp = [X1[idxz[i],j] for i=1:l-1,j=1:l-1]
        mv = mtmp \ rhs
        M[l,l] = λ[l]
        M[1:l-1,l] .= mv
        z[1:k-1] .= z[1:k-1] .+ X1 * mv
        zm, zi = findmax(abs.(z))
        idxz[l] = zi
        Z[:,l] .= z
    end
    X = S.Z * Z

    s = Int[]
    for j=1:m
        xm = zero(real(Tw))
        xi = 0
        for i=1:n
            xt = abs(X[i,j])
            if xt > xm && i ∉ s
                xm = xt
                xi = i
            end
        end
        push!(s, xi)
    end

    λd = DTw.(λ)
    Ad = DTw.(A)
    Xd = DTw.(X)
    Md = DTw.(M)
    # TODO: if cluster is tight enough, only need a singleton B
    # How tight is tight enough?
    B = Vector{Matrix{DTw}}(undef, m)
    for j=1:m
        B[j] = Ad - λd[j] * I
    end
    for j=1:m
        for i=1:m
            B[j][:,s[i]] .= -Xd[:,i]
        end
    end

    r = zeros(DTw, n, m)
    for j=1:m
        r[:,j] = λd[j] * Xd[:,j] - Ad * Xd[:,j]
        for i=1:j-1
            r[:,j] .= r[:,j] .+ M[i,j] * Xd[:,i]
        end
    end
    verbose && println("at iter 0 res norm = ", norm(r))
    FB0 = lu(Tw.(B[1]))
    FB = Vector{typeof(FB0)}(undef, m)
    FB[1] = FB0
    for j=1:m
        FB[j] = lu(Tw.(B[j]))
    end

    y = zeros(Tw, n, m)
    yp = zeros(Tw, n, m)
    ys = zeros(Tw, m, m)
    δp = zeros(Tw, n, m)
    δs = zeros(Tw, m, m)
    prevnorm = convert(real(DT),Inf)
    for p=1:maxiter
        δnorm = zero(real(DT))
        for j=1:m
            rhs = Tw.(r[:,j])
            # for jj=1:j
            #    rhs -= M[jj,j] * yp[:,jj]
            # end
            δ = FB[j] \ rhs
            δnorm += norm(δ)
            y[:,j] .+= δ
            yp[:,j] .= y[:,j]
            δp[:,j] .= δ
            for i=1:m
                δs[i,j] = δ[s[i]]
                δp[s[i],j] = zero(Tw)
                yp[s[i],j] = zero(Tw)
            end
            r[:,j] .= (r[:,j] - B[j] * DTw.(δ))
        end
        # there are occasional strange transients
        if (p > 3) && (δnorm > prevnorm)
            if δnorm > 10.0 * prevnorm
                status = :diverging
            else
                status = :stalled
            end
            verbose && println("$status at iter $p; early exit")
            break
        end
        prevnorm = δnorm
        for j=1:m
            for jj=1:m
                r[:,j] .= r[:,j] + (DTw(ys[jj,j]) * DTw.(δp[:,jj])
                                    + DTw(δs[jj,j]) * DTw.(yp[:,jj]))
            end
            for jj=1:j-1
                r[:,j] .= r[:,j] + Md[jj,j] * DTw.(δp[:,jj])
            end
        end
        # this update is done late to avoid doubly adding δ δ terms
        for j=1:m
            for i=1:m
                ys[i,j] = y[s[i],j]
            end
        end
        verbose && println("at iter $p res norm = ", norm(r))
        # println("ew: ",eigvals(Md + DTw.(ys)))
        verbose && println("DP subspace error: ",
                norm((Xd + DTw.(yp))*(Md+DTw.(ys)) - Ad * (Xd + DTw.(yp))))
        ynorm = norm(y)
        if δnorm / ynorm < tol1
            status = :converged
            verbose && println("converged")
            # println("iter $p ratio ",δnorm / ynorm)
            break
        end
    end
    Xnew = X + yp
    verbose && println("final subspace error norm: ", norm(Xnew*(M+ys) - A * Xnew))
    λbar = (1/m)*sum(λ)
    Mnew = Tw.(Md + DTw.(ys) - DTw(λbar) * I)
    dλ = eigvals(Mnew)
    λnew = λbar .+ dλ
    λnew, Xnew, status
end


"""
    BrooksAndCorey(; n, Ks=1, l=1, α=1, θr=0, θs=1) <: UnsaturatedFlowModel
    BrooksAndCorey(; n, k, l=1, α=1, θr=0, θs=1, ρ=1e3, μ=1e-3, g=9.81) <: UnsaturatedFlowModel

Create a Brooks and Corey porous model.

# Keyword arguments
- `n`: _n_ parameter.
- `Ks=1`: saturated hydraulic conductivity.
- `k`: intrinsic permeability.
- `l=1`: _l_ parameter.
- `α=1` (`\\alpha<tab>`): _α_ parameter.
- `θr=0` (`\\theta<tab>r`): residual moisture content.
- `θs=1` (`\\theta<tab>s`): moisture content when saturated.
- `ρ=1e3` (`\\rho<tab>`): density of the fluid.
- `μ=1e-3` (`\\mu<tab>`): dynamic viscosity of the fluid.
- `g=9.81`: magnitude of the gravitational acceleration.

# References
BROOKS, R.; COREY, T. Hydraulic properties of porous media.
Hydrology Papers, Colorado State University, 1964, vol. 24, p. 37.
"""
struct BrooksAndCorey{_Tnl,_Tα,_TKs,_Tθ} <: UnsaturatedFlowModel
    n::_Tnl
    l::_Tnl
    α::_Tα
    Ks::_TKs
    θr::_Tθ
    θs::_Tθ

    function BrooksAndCorey(; n, l=1, α=1, Ks=nothing, k=nothing, θr=0, θs=1, ρ=1e3, μ=1e-3, g=9.81)

        @argcheck α>zero(α)
        @argcheck θr<θs

        Ks = _asKs(Ks=Ks, k=k, ρ=ρ, μ=μ, g=g)

        new{promote_type(typeof(n),typeof(l)),typeof(α),typeof(Ks),promote_type(typeof(θr),typeof(θs))}(n,l,α,Ks,θr,θs)
    end
end

function θh(pm::BrooksAndCorey, h)
    if h ≥ -1/pm.α
        return pm.θs
    end

    Se = 1/(pm.α*(-h))^pm.n

    return pm.θr + Se*(pm.θs - pm.θr)
end

function hθ(pm::BrooksAndCorey, θ)
    Se = (θ - pm.θr)/(pm.θs - pm.θr)

    return -1/(pm.α*Se^(1/pm.n))
end

function Ch(pm::BrooksAndCorey, h)
    if h ≥ -1/pm.α
        return zero(1/h)
    end

    return -pm.n/h*(pm.θs - pm.θr)/(pm.α*(-h))^pm.n
end

function Kθ(pm::BrooksAndCorey, θ)
    Se = (θ - pm.θr)/(pm.θs - pm.θr)

    return pm.Ks*Se^(2/pm.n + pm.l + 2)
end

function Kh(pm::BrooksAndCorey, h)
    if h ≥ -1/pm.α
        return pm.Ks
    end

    return Kθ(pm, θh(pm, h))
end

function Dθ(pm::BrooksAndCorey, θ)
    Se = (θ - pm.θr)/(pm.θs - pm.θr)
    return pm.Ks*Se^(1/pm.n + pm.l + 1)/(pm.α*pm.n*(pm.θs - pm.θr))
end


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

using LinearAlgebra: qr, cond

function get_V0(train)
   # get list of atomic numbers
   Zs = AtomicNumber[]
   for (at, E) in train
      Zs = unique( [Zs; at.Z] )
   end

   # setup lsq system for E0s
   A = zeros(length(train), length(Zs))
   y = zeros(length(train))
   for (it, (at, E)) in enumerate(train)
      y[it] = E
      for (iz, z) in enumerate(Zs)
         A[it, iz] = length(findall(at.Z .== z))
      end
   end
   E0s = A \ y
   @info("  E0s = $(E0s)")
   syms = chemical_symbol.(Zs)
   return JuLIP.OneBody([syms[i] => E0s[i] for i = 1:length(Zs)]...)
end

function trainset(Vref, Ntrain; kwargs...)
   train = []
   for n = 1:Ntrain
      at = rand(Vref; kwargs...)
      push!(train, (at = at, E = energy(Vref, at), F = forces(Vref, at)))
   end
   # remove E0s
   V0 = get_V0(train)
   train = [ (at = at, E = E - energy(V0, at), F = F)
               for (at, E,  F) in train ]
   return train
end

function lsq(train, basis; verbose=true, wE = 1.0, wF = 1.0)
   @info("lsq info")
   nobs = sum( 1+3*length(t.at) for t in train )
   @info("  nobs = $(nobs); nbasis = $(length(basis))")
   A = zeros(nobs, length(basis))
   y = zeros(nobs)
   irow = 0
   for (at, E, F) in train
      # add energy to the lsq system
      irow += 1
      y[irow] = wE * E / length(at)
      A[irow, :] = wE * energy(basis, at) / length(at)

      # add forces to the lsq system
      nf = 3*length(at)
      y[(irow+1):(irow+nf)] = wF * mat(F)[:]
      Fb = forces(basis, at)
      for ib = 1:length(basis)
         A[(irow+1):(irow+nf), ib] = wF * mat(Fb[ib])[:]
      end
      irow += nf
   end
   qrF = qr(A)
   c = qrF \ y
   relrmse = norm(A * c - y) / norm(y)
   @info("   cond(R) = $(cond(qrF.R)); relrmse = $(relrmse)")
   return c, relrmse
end


struct SRS <: Imputor
    rng::AbstractRNG
    context::AbstractContext
end

"""
    SRS(; rng=Random.GLOBAL_RNG, context=Context())

Simple Random Sampling (SRS) imputation is a method for imputing both continuous and categorical
variables. Furthermore, it completes imputation while preserving the distributional
properties of the variables (e.g., mean, standard deviation).

The basic idea is that for a given variable, `x`, with missing data, we randomly draw
from the observed values of `x` to impute the missing elements. Since the random draws
from `x` for imputation are done in proportion to the frequency distribution of the values
in `x`, the univariate distributional properties are generally not impacted; this is true
for both categorical and continuous data.

# Keyword Arguments
* `rng::AbstractRNG`: A random number generator to use for observation selection
* `context::AbstractContext`: A context which keeps track of missing data
  summary information

# Example
"""

SRS(; rng=Random.GLOBAL_RNG, context=Context()) = SRS(rng, context)

function impute!(data::AbstractVector, imp::SRS)
    imp.context() do c
        obs_values = Impute.dropobs(data)
        if !isempty(obs_values)
            for i in eachindex(data)
                if ismissing!(c, data[i])
                    data[i] = rand(imp.rng, obs_values)
                end
            end
        end

        return data
    end
end


# check polylog reciprocal component
using Polylogarithms
using BenchmarkTools
using DataFrames, CSV
using PyPlot
using Printf
include("utilities.jl")

series_1 = Polylogarithms.polylog_series_1
series_2 = Polylogarithms.polylog_series_2
series_3 = Polylogarithms.polylog_series_3
polylog_reciprocal = Polylogarithms.polylog_reciprocal
L = Symbol("Li_s(z)")

# input data from Mathematica and reparse into complex numbers
for C=1:6
    # C = 1
    println("C = $C")
    filename = @sprintf("../data/polylog_test_data_c_%d.csv", C)
    data1 = CSV.read(filename; delim=",", type=String)
    #    has trouble reading in numbers like "2." so read all into strings, and parse
    data1[!,:s] = parse.(Complex{Float64}, data1[!,:s] )
    data1[!,:r] = parse.(Float64, data1[!,:r] )
    data1[!,:theta] = parse.(Float64, data1[!,:theta] )
    data1[!,:z] = parse.(Complex{Float64}, data1[!,:z] )
    data1[!,L] = parse.(Complex{Float64}, data1[!,L] )
    # check that complex numbers are being read in correctly
    z = data1[!,:z]
    z2 = data1[!,:r] .* exp.( im* data1[!,:theta] )
    error = data1[!,:z] .- z2
    maximum_abs_error = maximum(abs.(error))
    # [data1[1:10,:z] parse.(Complex{Float64}, data1[1:10,:z] ) z2[1:10] ]
    # [data1[1:10,:z]  z2[1:10] ]
    # println("max error in parsing is ", maximum_abs_error)

    # figure(1)
    # plot(real.(z), imag.(z), "o")
    # axis("equal")

    m = size(data1,1)
    Li = data1[!,L]
    s = data1[!,:s]
    r = unique( data1[!,:r])
    su = unique(s)
    if length(su) > 1
        error()
    else
        su = su[1]
        println("")
        @printf("  s = %f + %f i", real(su), imag(su))
    end
    S1 = zeros(Complex{Float64}, m)
    S2 = zeros(Complex{Float64}, m)
    n1 = zeros(Int64, m)
    error1 = zeros(Float64, m)
    rel_error1 = zeros(Float64, m)
    n2 = zeros(Int64, m)
    error2 = zeros(Float64, m)
    rel_error2 = zeros(Float64, m)

    for i=1:m
        # print(".")
        result1 = polylog_reciprocal(s[i], z[i])
        S1[i] = result1[1]
        n1[i] = result1[2]
        error1[i] = abs( S1[i] - Li[i]  )
        rel_error1[i] = error1[i]/abs( Li[i] )

        if real(s[i])>0 && abs(real(s[i]) - round(s[i])) < 1.0e-3
            rel_error2[i] = NaN
        else
            result2 = series_2(s[i], z[i])
            S2[i] = result2[1]
            n2[i] = result2[2]
            error2[i] = abs( S2[i] - Li[i]  )
            rel_error2[i] = error2[i]/abs( Li[i] )
        end
    end
    println("   max abs. error1 = $(maximum( abs.(error1) ))")
    println("   max abs. error2 = $(maximum( abs.(error2) ))")

    fig = figure(@sprintf("../data/polylog_bench_a_%02d.csv", C), figsize=(10,8))
    clf()
    title( @sprintf("s = %f + %f i", real(su), imag(su)) )
    d = 0.01
    ms = 3

    subplot(221)
    semilogy( data1[!,:r] .- d, rel_error1, "o"; markersize=ms)
    semilogy( data1[!,:r] .+ d, rel_error2, "rd"; markersize=ms)
    xlabel("|z|")
    ylabel("relative absolute error")
    plot([minimum(r), maximum(r)], [1,1]*Polylogarithms.default_accuracy)
    # xlim([0, 1.15])
    # ylim([1.0-e15, 1.0e-10])

    subplot(223)
    semilogy( data1[!,:r] .- d, n1, "o"; markersize=ms)
    semilogy( data1[!,:r] .+ d, n2, "rd"; markersize=ms)
    # θ = unique(data1[!,:theta])
    # for i=1:length(θ)
    #     k = findall(data1[!,:theta] .== θ[i])
    #     plot(data1[k,:r] .+ d, n2[k], "r-"; linewidth=0.5)
    #     text(data1[k[end],:r]+2*d, n2[k[end]], "$(θ[i]/π)"; verticalalignment="center", fontsize=9)
    # end
    # text(0.92, 65, "arg(z)/π"; verticalalignment="center")
    xlabel("|z|")
    ylabel("number of terms")
    plot([minimum(r), maximum(r)], [1,1]*Polylogarithms.default_max_iterations, "--")
    # xlim([0, 1.15])

    subplot(222)
    semilogy( data1[!,:theta] ./ π .- d, rel_error1, "o"; label="Reciprocal + Series 1", markersize=ms)
    semilogy( data1[!,:theta] ./ π .+ d, rel_error2, "rd"; label="Series 2", markersize=ms)
    xlabel("arg(z)/π")
    ylabel("relative absolute error")
    plot([-0.05,1.05], [1,1]*Polylogarithms.default_accuracy)
    legend()

    subplot(224)
    semilogy( data1[!,:theta] ./ π .- d, n1, "o"; markersize=ms)
    semilogy( data1[!,:theta] ./ π .+ d, n2, "rd"; markersize=ms)
    xlabel("arg(z)/π")
    ylabel("number of terms")
    plot([-0.05,1.05], [1,1]*Polylogarithms.default_max_iterations, "--")

    # savefig(@sprintf("plots/polylog_bench_a_%02d.svg", C))
    savefig(@sprintf("plots/polylog_bench_c_%02d.pdf", C); bbox_inches="tight")

end


abstract type Preconditioner{T} end

# ------------------------------------------------------------------------------
# |                           Nystrom Preconditioner                           |
# ------------------------------------------------------------------------------
mutable struct NystromPreconditioner{T <: Real} <: Preconditioner{T}
    A_nys::NystromSketch{T}
    λ::T
    μ::T
    cache::Vector{T}
    function NystromPreconditioner(A_nys::NystromSketch{T}, μ::T) where {T <: Real}
        return new{T}(A_nys, A_nys.Λ.diag[end], μ, zeros(rank(A_nys)))
    end
end
function Matrix(P::NystromPreconditioner)
    return Symmetric(
        1/(P.λ + P.μ) * P.A_nys.U*(P.A_nys.Λ + P.μ*I)*P.A_nys.U'
        + (I - P.A_nys.U*P.A_nys.U')
    )
end

# We care about applying P⁻¹
# P⁻¹x = U*(λ + μ)*(Λ + μI)⁻¹*Uᵀ*x + (I - UUᵀ)*x
#      = x - U*((λ + μ)*(Λ + μI)⁻¹ + I)*Uᵀ*x
function LinearAlgebra.ldiv!(y::Vector{T}, P::NystromPreconditioner{T}, x::Vector{T}) where {T <: Real}
    length(y) != length(x) && error(DimensionMismatch())
    mul!(P.cache, P.A_nys.U', x)
    @. P.cache *= (P.λ + P.μ) * 1 / (P.A_nys.Λ.diag + P.μ) - 1
    mul!(y, P.A_nys.U, P.cache)
    @. y = x + y
    return nothing
end

function LinearAlgebra.ldiv!(P::NystromPreconditioner{T}, x::Vector{T}) where {T <: Real}
    mul!(P.cache, P.A_nys.U', x)
    @. P.cache *= (P.λ + P.μ) * 1 / (P.A_nys.Λ.diag + P.μ) - 1
    x .+= P.A_nys.U*P.cache
    return nothing
end

function LinearAlgebra.:\(P::NystromPreconditioner{T}, x::Vector{T}) where {T <: Real}
    n = size(P, 1)
    y = zeros(n)
    ldiv!(y, P, x)
    return y
end

# Used for Krylov.jl
struct NystromPreconditionerInverse{T <: Real}
    A_nys::NystromSketch{T}
    λ::T
    μ::T
    cache::Vector{T}
    function NystromPreconditionerInverse(A_nys::NystromSketch{T}, μ::T) where {T <: Real}
        return new{T}(A_nys, A_nys.Λ.diag[end], μ, zeros(rank(A_nys)))
    end
end
Base.size(P::Union{NystromPreconditionerInverse, NystromPreconditioner}) = (size(P.A_nys.U, 1), size(P.A_nys.U, 1))
Base.size(P::Union{NystromPreconditionerInverse, NystromPreconditioner}, d::Int) = d <= 2 ? size(P)[d] : 1
Base.eltype(::Union{NystromPreconditionerInverse{T}, NystromPreconditioner{T}}) where {T} = T
function Matrix(P::NystromPreconditionerInverse)
    return Symmetric(
        (P.λ + P.μ) * P.A_nys.U*inv(P.A_nys.Λ + P.μ*I)*P.A_nys.U'
        + (I - P.A_nys.U*P.A_nys.U')
    )
end

function LinearAlgebra.mul!(y, P::NystromPreconditionerInverse{T}, x::Vector{T}) where {T <: Real}
    length(y) != length(x) && error(DimensionMismatch())

    mul!(P.cache, P.A_nys.U', x)
    @. P.cache *= (P.λ + P.μ) / (P.A_nys.Λ.diag + P.μ) - 1
    mul!(y, P.A_nys.U, P.cache)
    @. y = x + y
end

function LinearAlgebra.:*(P::NystromPreconditionerInverse{T}, x::Vector{T}) where {T <: Real}
    y = similar(x)
    mul!(y, P, x)
    return y
end


module HigherOrderKernels

export GaussianKernel,
       PolynomialKernel,
       UniformKernel,
       EpanechnikovKernel,
       BiweightKernel,
       TriweightKernel,
       density_bandwidth,
       distribution_bandwidth,
       kpdf,
       kcdf,
       order,
       poly_order

using GaussQuadrature,
      Polynomials,
      Statistics,
      SpecialFunctions

abstract type AbstractKernel{ν} end
struct GaussianKernel{ν} <: AbstractKernel{ν} end
struct PolynomialKernel{s,ν} <: AbstractKernel{ν} end
const UniformKernel = PolynomialKernel{0}
const EpanechnikovKernel = PolynomialKernel{1}
const BiweightKernel = PolynomialKernel{2}
const TriweightKernel = PolynomialKernel{3}

poly_order(k::Type{T}) where {s,ν,T<:PolynomialKernel{s,ν}} =
  degree(_u_poly(k))

order(::Type{T}) where {ν,T<:AbstractKernel{ν}} = ν

function roughness(k::Type{T}) where {s,ν,T<:PolynomialKernel{s,ν}}
  # N = 2s+2(div(ν,2)-1)+1
  # x, w = legendre(N)
  # R = sum(w[i] * density_kernel(k, x[i])^2 for i=1:N)
  Float64(polyint(_u_poly(k)^2, -1, 1))
end

function distribution_roughness(k::Type{T}) where {s,ν,T<:PolynomialKernel{s,ν}}
  I = polyint(_u_poly(k))
  K = I - I(-1)
  Float64(1 - polyint(K^2, -1, 1))
end

function firstnzmoment(k::Type{T}) where {s,ν,T<:PolynomialKernel{s,ν}}
  # N = s+2*div(ν,2)
  # x, w = legendre(N)
  # κ = sum(w[i] * x[i]^ν * density_kernel(k, x[i]) for i=1:N)
  polyint(Poly([zeros(ν); 1]) * _u_poly(k), -1, 1)
end

function roughness(k::Type{T}) where {ν,T<:GaussianKernel{ν}}
  N = 2(div(ν,2)-1)+1
  x, w = hermite(N)
  R = sum(w[i] * exp(x[i]^2) * density_kernel(k, x[i])^2 for i=1:N)
end

function firstnzmoment(k::Type{T}) where {ν,T<:GaussianKernel{ν}}
  N = 2*div(ν,2)
  x, w = hermite(N)
  κ = sum(w[i] * (sqrt(2) * x[i])^ν * exp(x[i]^2) * density_kernel(k, sqrt(2) * x[i]) * sqrt(2) for i=1:N)
end

function _std(x)
  iqr = diff(quantile(x, [0.25, 0.75]))[1] / 1.3489795003921636 # diff(quantile(Normal, [0.25, 0.75]))
  σ = std(x)
  min(iqr, σ)
end

const density_bandwidth_constant_lookup = Dict{Type{T} where T<:AbstractKernel,Float64}()

function density_bandwidth_constant(k::Type{T}) where {ν,T<:AbstractKernel{ν}}
  if haskey(density_bandwidth_constant_lookup, k)
    return density_bandwidth_constant_lookup[k]
  end
  R = roughness(k)
  κ = firstnzmoment(k)
  C = Float64(2*((sqrt(pi)*factorial(big(ν))^3*R)/(2*ν*factorial(big(2ν))*κ^2))^(1/(2ν+1)))
  density_bandwidth_constant_lookup[k] = C
end

density_bandwidth(k::Type{T}, data) where {ν,T<:AbstractKernel{ν}} =
  density_bandwidth_constant(k) * _std(data) * size(data, 1)^(-1/(2ν+1))

const distribution_bandwidth_constant_lookup = Dict{Type{T} where T<:AbstractKernel,Float64}()

function distribution_bandwidth_constant(k::Type{T}) where {ν,T<:AbstractKernel{ν}}
  if haskey(distribution_bandwidth_constant_lookup, k)
    return distribution_bandwidth_constant_lookup[k]
  end
  R = distribution_roughness(k)
  κ = firstnzmoment(k)
  C = Float64(2*((sqrt(pi)*factorial(big(ν))^3*R*(2ν-1))/(factorial(big(2ν))*κ^2*ν))^(1/(2ν-1)))
  distribution_bandwidth_constant_lookup[k] = C
end

distribution_bandwidth(k::Type{T}, data) where {ν,T<:AbstractKernel{ν}} =
  distribution_bandwidth_constant(k) * _std(data) * size(data, 1)^(-1/(2ν-1))

pochhammer(x, n) = gamma(x + n) / gamma(x)

# efficient evaluation of polynomial using Horner's rule
function polyval_expr(p::AbstractArray{T,1}, x::S) where {T,S}
  lenp = length(p)
  if lenp == 0
    return :($(zero(T)) * x)
  else
    y = convert(T, p[end])
    for i = lenp-1:-1:1
      y = :(muladd($x, $y, $(p[i])))
    end
    return y
  end
end

function _u2_poly(::Type{PolynomialKernel{s,ν}}) where {s,ν}
  r = div(ν, 2)
  c = pochhammer(big(1/2), s+1) * pochhammer(big(3/2+s), r-1) * pochhammer(big(3/2), r-1) / factorial(big(s+r-1))
  Poly([c * (-1)^k * pochhammer(big(1/2+s+r), k) / (factorial(big(k)) * factorial(big(r-1-k)) * pochhammer(big(3/2), k)) for k=0:r-1], :u2)
end

function _u_poly(k::Type{PolynomialKernel{s,ν}}) where {s,ν}
  p0 = Poly([1, -1], :u2)^s
  p1 = _u2_poly(k)
  p = p0 * p1
  u2 = Poly([0, 0, 1])
  p(u2)
end

@generated function density_kernel(k::Type{PolynomialKernel{s,ν}}, u) where {s,ν}
  e0 = (s == 0) ? :(abs(u) <= 1) :
       (s == 1) ? :((1-u2)*(abs(u) <= 1)) :
                  :((1-u2)^$s*(abs(u) <= 1))
  p1 = _u2_poly(PolynomialKernel{s,ν})
  e1 = polyval_expr(Float64.(coeffs(p1)), :u2)
  quote
    $(Expr(:meta, :inline))
    u2 = u^2
    $e1 * $e0
  end
end

@generated function distribution_kernel(k::Type{PolynomialKernel{s,ν}}, u) where {s,ν}
  p = _u_poly(PolynomialKernel{s,ν})
  P = polyint(p)
  P -= P(-1)
  e = polyval_expr(Float64.(coeffs(P)), :u)
  quote
    $(Expr(:meta, :inline))
    $e * (abs(u) <= 1) + (u > 1)
  end
end

@generated function density_kernel(::Type{GaussianKernel{ν}}, u) where {ν}
  r = div(ν, 2)
  ϕ = :(exp(-0.5u2))
  p = Poly(Float64.([(-1)^i / sqrt(2*pi) * factorial(big(2r))/(big(2)^(2r-i-1)*factorial(big(r))*factorial(big(2i+1))*factorial(big(r-i-1))) for i=0:r-1]))
  e = polyval_expr(coeffs(p), :u2)
  quote
    $(Expr(:meta, :inline))
    u2 = u^2
    $e * $ϕ
  end
end

kpdf(k, x, data) = kpdf(k, x, data, density_bandwidth(k, data))

function kpdf(::Type{T}, x, data, h) where {T<:AbstractKernel}
  sum = 0.0
  @fastmath @simd for y in data
    sum += density_kernel(T, (y - x) / h)
  end
  sum /= length(data) * h
  return sum
end

function kpdf(::Type{T}, x, data, h, lb, ub) where {T<:AbstractKernel}
  sum = 0.0
  if !(lb <= x <= ub)
    return sum
  end
  @simd for y in data
    sum += density_kernel(T, (y - x) / h)
  end
  if isfinite(lb)
    @simd for y in data
      sum += density_kernel(T, (y - (2lb - x)) / h)
    end
  end
  if isfinite(ub)
    @simd for y in data
      sum += density_kernel(T, (y - (2ub - x)) / h)
    end
  end
  sum /= length(data) * h
  return sum
end

kcdf(k, x, data) = kpdf(k, x, data, distribution_bandwidth(k, data))

function kcdf(::Type{T}, x, data, h) where {T<:AbstractKernel}
  sum = 0.0
  @fastmath @simd for y in data
    sum += distribution_kernel(T, (x - y) / h)
  end
  sum /= length(data)
  return sum
end

end # module


########################################################################
# Helper Functions
########################################################################

# softmax function
function softmax(z::Vector{Float64})
    expz = exp.(z)
    s = sum(expz)
    return expz ./ s
end

# negative log-likelihood function
function neg_log_loss(z::Vector{Float64}, y::Int64)
    return -log(z[y])
end

# generate random batches
function generate_batches(num_data::Int64, num_batches::Int64)
    n,r = divrem(num_data, num_batches)
    b = collect(1:n:num_data+1)
    for i in 1:length(b)
        b[i] += i > r ? r : i-1
    end
    p = randperm(num_data)
    return [p[r] for r in [b[i]:b[i+1]-1 for i=1:num_batches]]
end

# vertically split data
function split_data(Xtrain::SparseMatrixCSC{Float64, Int64}, Xtest::SparseMatrixCSC{Float64, Int64}, num_clients::Int64)
    num_features = size(Xtrain, 1)
    num_features_client = div(num_features, num_clients)
    Xtrain_split = Vector{ SparseMatrixCSC{Float64, Int64} }(undef, num_clients)
    Xtest_split = Vector{ SparseMatrixCSC{Float64, Int64} }(undef, num_clients)
    t = 1
    for i = 1:num_clients
        if i < num_clients
            ids = collect(t: t+num_features_client-1)
        else
            ids = collect(t: num_features)
        end
        Xtrain_split[i] = Xtrain[ids, :]
        Xtest_split[i] = Xtest[ids, :]
        t += num_features_client
    end
    return Xtrain_split, Xtest_split
end

# load data
function load_data(filename::String)
    if filename == "adult"
        Xtrain, Ytrain = read_libsvm("./data/Adult/a8a"); Xtrain = Xtrain[1:end-1, :]
        Xtest, Ytest = read_libsvm("./data/Adult/a8a.t")
        return Xtrain, Ytrain, Xtest, Ytest
    else
        @printf "Unsupported filename"
    end
end

# read data from libsvm
function read_libsvm(filename::String)
    numLine = 0
    nnz = 0
    open(filename, "r") do f
        while !eof(f)
            line = readline(f)
            info = split(line, " ")
            numLine += 1
            nnz += ( length(info)-1 )
            if line[end] == ' '
                nnz -= 1
            end
        end
    end
    @printf("number of lines: %i\n", numLine)
    n = numLine
    m = 0
    I = zeros(Int64, nnz)
    J = zeros(Int64, nnz)
    V = zeros(Float64, nnz)
    y = zeros(Int64, n)
    numLine = 0
    cc = 1
    open(filename, "r") do f
        while !eof(f)
            numLine += 1
            line = readline(f)
            info = split(line, " ")
            value = parse(Int64, info[1] )
            if value < 0
                value = Int64(2)
            end
            y[numLine] = value
            ll = length(info)
            if line[end] == ' '
                ll -= 1
            end
            for i = 2:ll
                idx, value = split(info[i], ":")
                idx = parse(Int, idx)
                value = parse(Float64, value)
                I[cc] = numLine
                J[cc] = idx
                V[cc] = value
                cc += 1
                m = max(m, idx)
            end
        end
    end
    return sparse( J, I, V, m, n ), y
end


@doc raw"""
    Ray3

An object `r` of the data type [`Ray3`](@ref) is a directed straight ray in the
three-dimensional Euclidean space ``Ε^2``.

It starts in a point called the *source* of `r` and goes to infinity.
"""
Ray3

"""
    Ray3(p::Point3, q::Point3)

Introduces a ray `r` with source `p` and passing through point `q`.
"""
Ray3(p::Point3, q::Point3)

"""
    Ray3(p::Point3, d::Direction3)

Introduces a ray `r` starting at source `p` with direction `d`.
"""
Ray3(p::Point3, d::Direction3)

"""
    Ray3(p::Point3, v::Vector3)

Introduces a ray `r` starting at source `p` with the direction of `v`.
"""
Ray3(p::Point3, v::Vector3)

"""
    Ray3(p::Point3, l::Line3)

Introduces a ray `r` starting at source `p` with the same direction as `l`.
"""
Ray3(p::Point3, l::Line3)

"""
    ==(r::Ray3, h::Ray3)

Test for equality: two rays are equal, iff they have the same source and the
same direction.
"""
==(r::Ray3, h::Ray3)

"""
    source(r::Ray3)

Returns the source of `r`.
"""
source(r::Ray3)

@doc raw"""
    point(r::Ray3, i::Real)

Returns a point on `r`.

!!! info "Precondition"

    ``i ≥ 0``.
"""
point(r::Ray3, i::Real)

@cxxdereference point(r::Ray3, i::Real) = point(r, convert(FT, i))

"""
    direction(r::Ray3)

Returns the direction of `r`.
"""
direction(r::Ray3)

"""
    to_vector(r::Ray3)

Returns a vector giving the direction of `r`.
"""
to_vector(r::Ray3)

"""
    supporting_line(r::Ray3)

Returns the line supporting `r` which has the same direction.
"""
supporting_line(r::Ray3)

"""
    opposite(r::Ray3)

Returns the ray with the same source and opposite direction.
"""
opposite(r::Ray3)

"""
    is_degenerate(r::Ray3)

Ray `r` is degenerate, if the source and the second defining point fall
together (that is if the direction is degenerate).
"""
is_degenerate(r::Ray3)

"""
    has_on(r::Ray3, p::Point3)

A point is on `r`, iff it is equal to the source of `r`, or if it is in the
interior of `r`.
"""
has_on(r::Ray3, p::Point3)

"""
    transform(r::Ray3, t::AffTransformation3)

Returns the ray obtained by applying `t` on the source and on the direction of
`r`.
"""
transform(r::Ray3, t::AffTransformation3)


struct LogitCrossEntropy end

@generated function loss(::LogitCrossEntropy, x::AbstractStrideMatrix{N,M,T}, y) where {M,N,T}
    quote
        l = zero($T)
        @avx for n ∈ axes(x,1)
            xₙ_1 = x[n,1]
            Base.Cartesian.@nexprs $(M-1) m -> xₙ_{m+1} = max(x[n,m+1], xₙ_m)
            s_1 = exp(x[n,1] - $(Symbol(:xₙ_,M)))
            Base.Cartesian.@nexprs $(M-1) m -> s_{m+1} = s_m + exp(x[n,m+1] - $(Symbol(:xₙ_,M)))
            logspxmax = log($(Symbol(:s_,M))) + $(Symbol(:xₙ_,M))
            Base.Cartesian.@nexprs $M m -> l -= (y[n] == (m-1)) * (x[n,m] - logspxmax)
        end
        l
    end
end
function logit_cross_entropy_grad_expr(loss, M, ::Type{T}) where {T}
    loopbody = quote
        xₙ_1 = x[n,1]
        Base.Cartesian.@nexprs $(M-1) m -> xₙ_{m+1} = max(x[n,m+1], xₙ_m)
        s_1 = exp(x[n,1] - $(Symbol(:xₙ_,M)))
        Base.Cartesian.@nexprs $(M-1) m -> s_{m+1} = s_m + exp(x[n,m+1] - $(Symbol(:xₙ_,M)))
        logspxmax = log($(Symbol(:s_,M))) + $(Symbol(:xₙ_,M))
        Base.Cartesian.@nexprs $M m -> δx_m = (x[n,m] - logspxmax)
        Base.Cartesian.@nexprs $M m -> y_m = y[n] == (m-1)
    end
    loss && push!(loopbody.args, :(Base.Cartesian.@nexprs $M m -> l -= y_m ? δx_m : zero($T)))
    push!(loopbody.args, :(Base.Cartesian.@nexprs $M m -> ∂x[n,m] = exp(δx_m) - y_m))
    quote
        l = zero($T)
        @avx for n ∈ axes(∂x,1)
            $loopbody
        end
        l
    end
end
@generated function gradient!(∂x, ::LogitCrossEntropy, x::AbstractStrideMatrix{N,M,T}, y) where {M,N,T}
    logit_cross_entropy_grad_expr(false, M, T)
end
@generated function loss_and_gradient!(∂x, ::LogitCrossEntropy, x::AbstractStrideMatrix{N,M,T}, y) where {M,N,T}
    logit_cross_entropy_grad_expr(true, M, T)
end