using QuasiCopula, LinearAlgebra, GLM
using Random, Distributions, DataFrames, ToeplitzMatrices
using Test, BenchmarkTools

BLAS.set_num_threads(1)
Threads.nthreads()

p_fixed = 3    # number of fixed effects, including intercept

# true parameter values
Random.seed!(12345)
βtrue = rand(Uniform(-2, 2), p_fixed)
rtrue = 10.0
σ2true = [0.5]
ρtrue = [0.5]

function get_V(ρ, n)
    vec = zeros(n)
    vec[1] = 1.0
    for i in 2:n
        vec[i] = ρ
    end
    V = ToeplitzMatrices.SymmetricToeplitz(vec)
    V
end

#simulation parameters
samplesize = 10000

st = time()
currentind = 1
d = NegativeBinomial()
link = LogLink()
D = typeof(d)
Link = typeof(link)
T = Float64

gcs = Vector{NBCopulaCSObs{T, D, Link}}(undef, samplesize)

ni = 25 #  number of observations per individual
V = get_V(ρtrue[1], ni)

# true Gamma
Γ = σ2true[1] * V

# for reproducibility I will simulate all the design matrices here
Random.seed!(12345)
X_samplesize = [randn(ni, p_fixed - 1) for i in 1:samplesize]

for i in 1:samplesize
    X = [ones(ni) X_samplesize[i]]
    # X = [ones(ni) randn(ni, p - 1)]
    # X = ones(ni, 1)
    # y = Float64.(Y_nsample[i])
    η = X * βtrue
    μ = exp.(η)
    p = rtrue ./ (μ .+ rtrue)
    vecd = Vector{DiscreteUnivariateDistribution}(undef, ni)
    vecd = [NegativeBinomial(rtrue, p[i]) for i in 1:ni]
    nonmixed_multivariate_dist = NonMixedMultivariateDistribution(vecd, Γ)
    # simuate single vector y
    y = Vector{Float64}(undef, ni)
    res = Vector{Float64}(undef, ni)
    rand(nonmixed_multivariate_dist, y, res)
    # push!(Ystack, y)
    V = [ones(ni, ni)]
    # V = [ones(ni, ni)]
    gcs[i] = NBCopulaCSObs(y, X, d, link)
end

# form model
gcm = NBCopulaCSModel(gcs)
# precompile
println("precompiling NB CS fit")
gcm2 = deepcopy(gcm);
QuasiCopula.fit!(gcm2, maxBlockIter = 1);

fittime = @elapsed QuasiCopula.fit!(gcm)
@show fittime
@show gcm.β
@show gcm.σ2
@show gcm.ρ
@show gcm.∇β
@show gcm.∇σ2
@show gcm.∇ρ

@show gcm.r
@show gcm.∇r
loglikelihood!(gcm, true, true)
vcov!(gcm)
@show QuasiCopula.confint(gcm)

mseβ, mseρ, mseσ2, mser = MSE(gcm, βtrue, ρtrue, σ2true, rtrue)
@show mseβ
@show mser
@show mseσ2
@show mseρ

@test mseβ < 0.01
@test mseσ2 < 1
@test mseρ < 0.01
@test mser < 1


hswish(x) = x * (relu6(x + 3) / 6)


using EconPDEs

# PDE with 1 state variable
include("./AssetPricing/CampbellCochrane.jl")
m = CampbellCochraneModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
pdesolve(m, stategrid, y0)
@time y, result, distance = pdesolve(m, stategrid, y0)
#   0.011112 seconds (286.45 k allocations: 12.268 MiB)

# PDE with 2 state variables
include("./AssetPricing/BansalYaron.jl")
m = BansalYaronModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
pdesolve(m, stategrid, y0)
@time pdesolve(m, stategrid, y0)
# 0.081095 seconds (485.84 k allocations: 78.910 MiB, 11.41% gc time)

# System of 4 PDEs with 1 state variable
include("./AssetPricing/GarleanuPanageas.jl")
m = GarleanuPanageasModel()
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
pdesolve(m, stategrid, y0)
@time pdesolve(m, stategrid, y0)
# 0.047346 seconds (167.95 k allocations: 65.049 MiB, 12.90% gc time)

# System of 3 PDEs with 2 state variables
include("./AssetPricing/DiTella.jl")
m = DiTellaModel()
stategrid = initialize_stategrid(m; xn = 80, νn = 10)
y0 = initialize_y(m, stategrid)
pdesolve(m, stategrid, y0)
@time pdesolve(m, stategrid, y0)
# 0.912719 seconds (2.55 M allocations: 1.303 GiB, 9.01% gc time)


### Creating driving activity: INPUT is datamatrix, p_min, p_max, e, T, number_zones and OUTPUT is p_drive

function createpdrive(datamatrix, distance_matrix_km, number_zones)
    p_drive = zeros(number_zones, T);
    mean_sum = zeros(number_zones, T);
    min_t = zeros(number_zones,1);
    max_t = zeros(number_zones,1);

    # Calculating the sums, minimum sums and maximum sums of mean travel time out of a zone
    for i=1:number_zones
        for t=1:T
            mean_sum[i, t] = 0;
            counter = 0;
            for j=1:number_zones
                if(datamatrix[i,j,t,1] != 0)
                    mean_sum[i, t] = mean_sum[i, t] + datamatrix[i,j,t,1]/distance_matrix_km[i,j];
                    counter += 1;
                end
            end
            mean_sum[i, t] = mean_sum[i, t]/counter;
        end
        max_t[i] = maximum(mean_sum[i,:]);
        min_t[i] = minimum(mean_sum[i,:]);
    end

    ## Calculating the Bernoulli distributions p_drive
    for i=1:number_zones
        if(max_t[i] > 0) # A zero would lead to infinity given the formula below and create NaN entries.
            for t=1:T
                # A linear function for assigning a value of 10% to p_drive if min_t is given, and a value of 90% if max_t is given.
                p_drive[i,t] = p_min + (p_max - p_min) * ((mean_sum[i,t] - min_t[i]) / (max_t[i] - min_t[i]))^e_drive;
            end
        end
    end

    return p_drive;

end


using Empirikos
using Test

_target = MarginalDensity(StandardNormalSample(3.0))
@test Base.extrema(_target)[2] == pdf(Normal(),0)

@test Base.extrema(MarginalDensity(NormalSample(3.0, 0.5)))[2] == pdf(Normal(0,0.5),0.0)


using Polyhedra
using CDDLib
using MAT

include("../../src/Reachability/ReachabilityPolyhedra.jl")
using ReachabilityPolyhedra

function compute(algorithm="reach_homog")

# system's matrix
A = [0. 1. 0. 0. 0. 0. ;
-0.417533 -3.1931759963 39.24 0. -14.825331 11.123344 ;
0. 0. 0. 1. 0. 0. ;
0.0417533 0.31931759963 -4.905 0. 1.4825331 -1.1123344 ;
0.0638407957 -0.32473339016573 0. 0. -3.7332068901 -0.7007592976 ;
0.0853437452 -0.72366802635628 0. 0. -5.9714023436 -2.2736115136]

# number of variables (dimension)
n = size(A, 1)

# initial set of states
file = matopen("X0.mat")

#=
F = read(file, "F");
g = read(file, "g");
g = g[:];  # cast into a column vector
# this is a half-space representation. use the methods X0.A and X0.b to access
# the matrix/vectors such that X0 = {x : Ax <= b}
X0HRep = SimpleHRepresentation(F, g);
X0 = polyhedron(X0HRep, CDDLibrary());
=#

# initial states : hyperrectangle
g1 = read(file, "g1")[:]
g2 = read(file, "g1")[:]
g3 = read(file, "g1")[:]
F1 = read(file, "F1")
F2 = read(file, "F2")
F3 = read(file, "F3")

X0hat = Array{CDDPolyhedron}(3)
X0hat[1] = polyhedron(SimpleHRepresentation(F1, g1))
X0hat[2] = polyhedron(SimpleHRepresentation(F2, g2))
X0hat[3] =  polyhedron(SimpleHRepresentation(F3, g3))

# number of samples
N = 250

# ==== case with inputs ===

# this polyhedron is "null" in n dims
#nullPoly_6d = polyhedron(SimpleVRepresentation([0.0 0.0 0.0 0.0 0.0 0.0]));
# we specify the backend using a second argument, as CDDLibrary()
# otherwise, the object will be of type Polyhedra.SimplePolyhedron{6,Float64},
# and we cannot do operations with it
nullPoly_6d = polyhedron(SimpleVRepresentation(zeros(1, 6)), CDDLibrary());
U = [nullPoly_6d for i in 1:N]

# === solve ===

if algorithm == "reach_homog"
    RSet = reach_homog(A, X0hat, N)
elseif algorithm == "reach"
    RSet = reach(A, X0hat, N, U)
else
    error("algorithm unkwnown")
end

end


cheb_mesh(T::Type, N::Integer) = (-cos.((T(π) * (0:N)) ./ N) .+ 1) ./ 2

function cheb(T::Type, N::Integer)
    if N == 0
        return (ones(T, 1), zeros(T, 1, 1))
    else
        x = cheb_mesh(T, N)
        c = [2; ones(T, N - 1, 1); 2] .* (-1) .^ (0:N)
        dx = x .- x'
        D = (c * (1 ./ c)') ./ (dx + I)
        D[diagind(D)] .-= vec(sum(D; dims=2))
        return (x, D)
    end
end

struct CollocationSystem{T} <: AbstractDynamicalSystem
    # AbstractSystem requirements
    eqs::Vector{Equation}
    states::Vector
    ps::Vector
    name::Symbol
    defaults::Dict
    # Extra fields
    n_dim::Int
    n_coll::Int
    n_mesh::Int
    t_coll::Vector{T}
    D::Matrix{T}
    D⁻¹::Matrix{T}
    ode_ps::Vector
    ode_u₀::Vector
    ode_u₁::Vector
    ode_t::Vector
end

function CollocationSystem(
    T::Type, sys::ODESystem, n_mesh, n_coll; name=gensym(nameof(sys)), free_ps=false
)
    # Turn the ODESystem into a function we can evaluate
    iv = MTK.get_iv(sys)  # independent variable
    sts = states(sys)
    ps = parameters(sys)
    f = @RuntimeGeneratedFunction(
        build_function([x.rhs for x in equations(sys)], sts, ps, iv)[1]
    )
    n_dim = length(sts)

    # Namespace the core symbols to avoid name collisions
    @variables coll₊u[1:n_mesh, 1:(n_coll + 1), 1:n_dim]
    for i in 1:(n_mesh - 1)
        coll₊u[i, end, :] .= coll₊u[i + 1, 1, :]  # continuity
    end
    @parameters coll₊ka[1:n_mesh]
    @variables coll₊t[0:1]

    # Regenerate the parameters as variables
    if free_ps
    p = [Num(Variable(nameof(par))) for par in ps]
    else
        p = ps
    end

    # Generate the collocation equations for each mesh interval
    #   D*u ~ Tp*f(u, p, t)
    # where D is the differentiation matrix, and Tp is the time interval
    # the ka variables enable mesh interval scaling (assumes sum(ka) == 1)

    t_coll_full, D_full = cheb(T, n_coll)

    # Calculate the corresponding integration matrix (think about solving D*x = f(t) with appropriate boundary conditions)
    D_full⁻¹ = [zeros(1, n_coll + 1); zeros(n_coll, 1) inv(D_full[2:end, 2:end])]

    # End points aren't used in the collocation equations
    t_coll = t_coll_full[1:(end - 1)]
    D = D_full[1:(end - 1), :]

    # Interval length
    Tp = coll₊t[2] - coll₊t[1]

    # Collocation equations
    eqns = Vector{Equation}()
    for i in 1:n_mesh
        Dx = [D * coll₊u[i, :, k] for k in 1:n_dim]
        t_int =
            coll₊t[1] .+
            Tp .* (coll₊ka[i] .* t_coll .+ sum(coll₊ka[1:(i - 1)]; init=Num(0)))
        f_vals = [
            Tp * coll₊ka[i] * Base.invokelatest(f, coll₊u[i, j, :], p, t_int[j]) for
            j in eachindex(t_int)
        ]
        append!(
            eqns, [0 ~ Dx[k][j] - f_vals[j][k] for k in eachindex(Dx), j in eachindex(f_vals)]
        )
    end

    # Bindings to the inner variables
    sts_names = nameof.(MTK.operation.(sts))
    sts_begin = [Num(Variable(st, 0)) for st in sts_names]
    append!(eqns, 0 .~ coll₊u[1, 1, :] .- sts_begin)
    sts_end = [Num(Variable(st, 1)) for st in sts_names]
    append!(eqns, 0 .~ coll₊u[end, end, :] .- sts_end)

    # Defaults
    defaults = merge(
        MTK.get_defaults(sys), Dict(MTK.value(ka) => one(T) / n_mesh for ka in coll₊ka)
    )

    new_sts = [sts_begin; sts_end; unique(vec(coll₊u)); coll₊t]
    new_ps = copy(coll₊ka)
    if free_ps
        append!(new_sts, p)
    else
        append!(new_ps, p)
    end

    return CollocationSystem(
        eqns,  # eqs
        MTK.value.(new_sts),  # states
        MTK.value.(new_ps),  # ps
        name,  # name
        defaults,  # defaults
        n_dim,  # n_dim
        n_coll,  # n_coll
        n_mesh,  # n_mesh
        t_coll_full,  # t_coll
        D_full,  # D
        D_full⁻¹,  # D⁻¹
        p,  # ode_ps
        sts_begin,  # ode_u₀
        sts_end,  # ode_u₁
        coll₊t,  # ode_t
    )
end

function CollocationSystem(sys::ODESystem, args...; kwargs...)
    return CollocationSystem(Float64, sys, args...; kwargs...)
end

"""
    $SIGNATURES

Add equations that are a function of the boundary states and times to `coll`. The input `bc`
is assumed to be a function of the form `(u₀, t₀, u₁, t₁, p) -> equation` where `u₀` and
`u₁` are state variables at the start and end of the interval respectively, `t₀` and `t₁`
are the start and end times, and `p` are the underlying ODE system parameters.

The value returned from `bc` can be an equation or a vector of equations.

`u₀` and `u₁` have the states in the same order as the underlying ODE system.

## Example - periodic boundary conditions

```julia
add_boundarycondition!(coll, (u₀, t₀, u₁, t₁, p) -> u₀ .~ u₁)
```

Note that `.~` is used since `u₀` and `u₁` are vectors of variables.
"""
function add_boundarycondition!(coll::CollocationSystem, bc)
    eqn = bc(coll.ode_u₀, coll.ode_t[1], coll.ode_u₁, coll.ode_t[2], coll.ode_ps)
    if eqn isa Equation
        push!(coll.eqs, eqn)
    elseif eqn isa Vector{Equation}
        append!(coll.eqs, eqn)
    elseif eqn isa Num
        push!(coll.eqs, 0 ~ eqn)
    elseif eqn isa Vector{Num}
        append!(coll.eqs, 0 .~ eqn)
    else
        throw(ArgumentError("Input function should return a (vector of) Equation or Num"))
    end
    return coll
end

function generate_initialconditions(
    coll::CollocationSystem{T}, sol, tspan; mesh=nothing
) where {T}
    # Mesh intervals
    ka = mesh === nothing ? ones(T, coll.n_mesh) ./ coll.n_mesh : mesh
    @assert length(ka) == coll.n_mesh
    # Generate mesh with representation points
    Tp = tspan[end] - tspan[begin]
    t_coll = coll.t_coll[1:(end - 1)]
    t_int = similar(t_coll)
    # Determine the initial conditions
    ics = Vector{Pair{typeof(Variable(:x)), T}}()
    sizehint!(ics, (coll.n_mesh*coll.n_coll + 1)*coll.n_dim)
    basename = _strip_toplevel(MTK.renamespace(coll.name, ""))  # need to strip the top-level namespace since it isn't used in solve
    basename_u = Symbol(basename, :coll₊u)
    for i in eachindex(ka)
        t_int .= tspan[begin] .+ Tp .* (ka[i] .* t_coll .+ sum(ka[1:(i - 1)]))
        for j in eachindex(t_int)
            sts = sol(t_int[j])
            for k in eachindex(sts)
                push!(ics, Variable(basename_u, i, j, k) => sts[k])
            end
        end
    end
    sts = sol(tspan[end])
    for k in eachindex(sts)
        push!(ics, Variable(basename_u, coll.n_mesh, coll.n_coll + 1, k) => sts[k])
    end
    # User-defined state variables
    for (var, val) in zip(coll.ode_u₁, sts)  # assumes sts = sol(tspan[end])
        push!(ics, MTK.value(var) => val)
    end
    for (var, val) in zip(coll.ode_u₀, ics)  # FRAGILE: assumes first elements of ics correspond to sol(tspan[begin])
        push!(ics, MTK.value(var) => val[2])
    end
    # Start/end times
    push!(ics, Variable(Symbol(basename, :coll₊t), 0) => tspan[begin])
    push!(ics, Variable(Symbol(basename, :coll₊t), 1) => tspan[end])
    return ics
end

# CollocationSystem does not allow embedded systems (it's not meaningful)
ModelingToolkit.get_systems(::CollocationSystem) = NonlinearSystem[]

# Convert a CollocationSystem to a NonlinearSystem
function ModelingToolkit.NonlinearSystem(coll::CollocationSystem)
    return NonlinearSystem(
        MTK.get_eqs(coll),
        MTK.get_states(coll),
        MTK.get_ps(coll);
        name=nameof(coll),
        defaults=MTK.get_defaults(coll),
    )
end


# LogLinear Transformation

"""
    LogLinearFormula(df::Int64)

Returns GLM formula can be used in `glm` with `FreqTab` data frame.
`df`, is abbreviation for degrees of formula, represents degree of polynomial log-linear method.
This function works very slowly. If a fixed degree formula will be used repeatedly, define the formula as a variable.
The similer macro `@LogLinearFormula` is provided and it is faster than the function.
"""
function LogLinearFormula(df)
    fml = "@formula(freq ~ 1 +"
    @simd for d in 1:df
        fml *= "scale^$d"
        if d != df
            fml *= " + "
        else
            fml *= ")"
        end
    end
    return eval(Meta.parse(fml))
end

struct SmoothedFreqTab <: EG
    tab::DataFrame
    raw::Vector
    interval::Float64
    fit
end
"""
    presmoothing(F::FreqTab, @LogLinearFormula(df::Int64))

Returns presmoothed frequency table as `SmoothedFreqTab` and `glm` fitted object.

Preserving first C moments of original frequency data, passed `LogLinearFormula(C)` to `fml`.
C is a degree of freedom
"""
function presmoothing(F::EG, fml)
    fit1 = glm(fml, F.tab, Poisson(), LogLink())
    freq = predict(fit1, DataFrame(scale = F.tab.scale))
    tab = DataFrame(scale = F.tab.scale, freq = freq, cumfreq = cumsum(freq),
                    prob = freq ./ sum(freq), cumprob = cumsum(freq) ./ sum(freq))
    return SmoothedFreqTab(tab, F.raw, F.interval, fit1)
end

# function presmoothing(F::EG, degree::Int64)
#     AIC, BIC, AICC, DEVIANCE, LL = zeros.(Float64, fill(degree, 5))
#     for d in 1:degree
#         println("Fitting dof = $degree model.")
#         fml = LogLinearFormula(d)
#         fit = glm(fml, F.tab, Poisson(), LogLink())
#         AIC[d] = aic(fit)
#         BIC[d] = bic(fit)
#         AICC[d] = aicc(fit)
#         LL[d] = loglikelihood(fit)
#         DEVIANCE[d] = deviance(fit)
#     end
#     DataFrame(degree = [1:1:degree;], aic = AIC, aicc = IACC, bic = BIC, loglik = LL, deviance = DEVIANCE)
# end

struct SmoothedNEATFreqTab <: NEAT
    tabX
    tabV
    rawX
    rawV
    intervalX
    intervalV
    fitX
    fitV
    marginal
end
"""
    presmoothing(F::NEAT, LogLinearFormula(df::Int64), LogLinearFormula(df::Int64))

Returns presmoothed frequency table as `SmoothedNEATFreqTab` and `glm` fitted object,
but `interval` and `marginal` elements, which are returned, are not based on smoothed score.

Preserving first C moments of original frequency data, passed `LogLinearFormula(C)` to `fml`.
C is a degree of freedom
"""
function presmoothing(F::NEAT, fmlX, fmlV)
    # Smoothing X (independent part)
    fitX = glm(fmlX, F.tabX, Poisson(), LogLink())
    freqX = predict(fitX, DataFrame(scale = F.tabX.scale))
    tabX = DataFrame(scale = F.tabX.scale, freq = freqX, cumfreq = cumsum(freqX),
                    prob = freqX ./ sum(freqX), cumprob = cumsum(freqX) ./ sum(freqX))
    # Smoothing V (common part)
    fitV = glm(fmlV, F.tabV, Poisson(), LogLink())
    freqV = predict(fitV, DataFrame(scale = F.tabV.scale))
    tabV = DataFrame(scale = F.tabV.scale, freq = freqV, cumfreq = cumsum(freqV),
                    prob = freqV ./ sum(freqV), cumprob = cumsum(freqV) ./ sum(freqV))
    return SmoothedNEATFreqTab(tabX, tabV, F.rawX, F.rawV, F.intervalX, F.intervalV, fitX, fitV, F.marginal)
end

# Kernel method
function RjX(x, xⱼ, a, μ, hX)
    return (x - a * xⱼ - (1-a)*μ) / (a*hX)
end
struct KernelFreqTab <: EG
    tab::DataFrame
    raw::Vector
    interval::Float64
    Bandwidth::Float64
end
"""
    KernelSmoothing(X::SmoothedFreqTab; kernel = :Gaussian, hX = 0.622, scale = X.tab.scale)

Returns Kernel smoothed frequency table as `KernelFreqTab`.

# Arguments

- `X` the object which class is `freqtab`

## Optional Arguments

- `kernel`
- `hX` The band width.
- `newint` A new interval to depict the continuized probability line.

# Value

`KernelFreqTab`

In the kernel smoothing, the choice of bandwidth `hX` is an important consideration. When bandwidth is large, the smoothed distribution becomes linear equating fucnction. Contrary, bandwidth is small, it becomes spikye shaped distribution. We recommend to use default value `hX = 6.22` or select it by using `EstBandwidth`.
"""
function KernelSmoothing(X::EG; kernel = :Gaussian, hX = 0.622, newint = X.interval/100)
    scale = (minimum(X.tab.scale) - X.interval/2):newint:(maximum(X.tab.scale) + X.interval/2)
    # hX = bandwidth of cumulative distribution function
    N = sum(X.tab.freq)
    μ = mean(X.raw); σ² = var(X.raw)
    a² = σ² / (σ² + hX^2)
    a = sqrt(a²)
    𝒇hX = zeros(Float64, length(scale))
    FhX = zeros(Float64, length(scale))
    for (i, x) in enumerate(scale), (j, xⱼ) in enumerate(X.tab.scale)
        𝒇hX[i] += X.tab.prob[j]*pdf.(Normal(0, 1), RjX(x, xⱼ, a, μ, hX)) / (a*hX)
        FhX[i] += X.tab.prob[j]*cdf.(Normal(0, 1), RjX(x, xⱼ, a, μ, hX))
    end
    𝒇hX = 𝒇hX ./ sum(𝒇hX) # normalize
    tbl = DataFrame(scale = scale, freq = N .* 𝒇hX, cunfreq = cumsum(N .* 𝒇hX), prob = 𝒇hX, cumprob = cumsum(𝒇hX))
    return KernelFreqTab(tbl, X.raw, X.interval, hX)
end

function BandwidthPenalty(hX, X::SmoothedFreqTab, lprob, rprob, K; kernel = :Gaussian)
    hX = exp(hX[1])
    r = X.tab.prob
    μ = mean(X.raw); σ² = var(X.raw)
    a² = σ² / (σ² + hX^2)
    a =sqrt(a²)
    𝒇hX = zeros(Float64, length(X.tab.scale))
    r𝒇′hX = zeros(Float64, length(X.tab.scale))
    l𝒇′hX = zeros(Float64, length(X.tab.scale))
    for (i, x) in enumerate(X.tab.scale), (j, xⱼ) in enumerate(X.tab.scale)
        R = RjX(x, xⱼ, a, μ, hX)
        𝒇hX[i] += X.tab.prob[j]*pdf.(Normal(0, 1), R) / (a*hX)
        r𝒇′hX[i] -= rprob[j]*pdf.(Normal(0, 1), R) / (a*hX)^2 * R
        l𝒇′hX[i] -= lprob[j]*pdf.(Normal(0, 1), R) / (a*hX)^2 * R
    end
    pen1 = sum((r .- 𝒇hX) .^2)
    pen2 = sum((l𝒇′hX .< 0) .* (r𝒇′hX .> 0))
    return pen1 + K * pen2
end

"""
    EstBandwidth(X::EG; kernel = :Gaussian)

Estimate optimal bandwidth `hX` in Gaussian kernel.

# Arguments

- `X` is the object which the class is `freqtab`.

# Value

Optimize information.

"""
function EstBandwidth(X::SmoothedFreqTab; kernel = :Gaussian, K = 1)
    x = X.tab.scale
    lscore = x .- X.interval/4
    rscore = x .+ X.interval/4
    lpred = predict(X.fit, DataFrame(scale = lscore))
    rpred = predict(X.fit, DataFrame(scale = rscore))
    opt = optimize(hX -> BandwidthPenalty(hX, X, lpred, rpred, K; kernel = kernel), [0.5], method = BFGS())
    println("Minimizer sould be transformed `exp()` before interpretation. \nMinimizer (after taking `exp()`) = $(exp(opt.minimizer[1]))\n\n")
    return opt
end


"""
Calculate elastic regression from function data f, for response y

    elastic_regression(f, y, timet; B=None, lambda=0, df=20, max_itr=20,
                       smooth=false)
    :param f: array (M,N) of N functions
    :param y: vector (N) of responses
    :param timet: vector (N) describing time samples
    :param B: matrix describing basis functions (M,N) (default=None generates a
              B-spline basis
    :param lambda: regularization parameter
    :param df: degree of freedom of basis
    :param max_itr: maximum number of iterations
    :param smooth: smooth data

    Returns Dict describing regression
    :return alpha: intercept
    :return beta: regression function
    :return fn: aligned functions
    :return qn: aligned srsfs
    :return gamma: warping functions
    :return q: original srsfs
    :retrun B: basis functions
    :return type: type of regressions
    :return b: coefficients
    :return SSE: sum of squared error
"""
function elastic_regression(f::Array, y::Vector, timet::Vector; B=Union{}, lambda=0, df=20,
                            max_itr=20, smooth=false)
    M, N = size(f);

    if M > 500
        parallel = true;
    elseif N > 100
        parallel = true;
    else
        parallel = false;
    end

    binsize = mean(diff(timet));

    # Create B-spline basis if none provided
    if B == Union{}
        B = bs(timet, df, 4);
    end
    Nb = size(B,2);

    # second derivative for regularization
    Bdiff = zeros(M, Nb);
    for ii in 1:Nb
        Bdiff[:, ii] = gradient(gradient(B[:,ii], binsize), binsize);
    end

    q = f_to_srsf(f, timet, smooth);

    gamma = repmat(linspace(0,1,M), 1, N);
    fn = zeros(M, N);
    qn = zeros(M, N);

    itr = 1;
    SSE = zeros(max_itr);
    alpha = 0;
    b = zeros(Nb+1);
    while itr <= max_itr
        @printf("Iteration: %d\n", itr)
        # align data
        fn = zeros(M, N);
        qn = zeros(M, N);
        for ii in 1:N
            xout = (timet[end] - timet[1]) * gamma[:, ii] + timet[1];
            fn[:,ii] = approx(timet, f[:, ii], xout);
            qn[:,ii] = warp_q_gamma(timet, q[:,ii], gamma[:,ii]);
        end

        # OLS using basis
        Phi = ones(N, Nb+1);
        for ii in 1:N
            for jj in 2:Nb+1
                Phi[ii,jj] = trapz(timet, qn[:,ii].*B[:,jj-1]);
            end
        end

        R = zeros(Nb+1,Nb+1);
        for ii in 2:Nb+1
            for jj in 2:Nb+1
                R[ii,jj] = trapz(timet, Bdiff[:,ii-1].*Bdiff[:,jj-1]);
            end
        end

        xx = Phi.' * Phi;
        inv_xx = inv(xx + lambda * R);
        xy = Phi.' * y;
        b = inv_xx * xy;

        alpha = b[1];
        beta = B * b[2:Nb+1];

        # compute the SSE
        int_X = zeros(N);
        for ii in 1:N
            int_X[ii] = trapz(timet, qn[:,ii].*beta);
        end

        SSE[itr] = sum((y-alpha-int_X).^2);

        # find gamma
        gamma_new = zeros(M, N);
        if parallel
            gamma_new = @parallel (hcat) for i=1:N
                regression_warp(beta, timet, q[:, i], y[i], alpha);
            end
        else

            for i = 1:N
                gamma_new[:,i] = regression_warp(beta, timet, q[:, i], y[i],
                                                 alpha);
            end
        end

        if norm(gamma-gamma_new) < 1e-5
            break
        else
            gamma = copy(gamma_new);
        end

        itr += 1;
    end

    # Last step with centering of gam
    gamI = sqrt_mean_inverse(gamma);
    gamI_dev = gradient(gamI, 1/(M-1));
    timet0 = (timet[end] -timet[1]) .* gamI  + timet[1];
    beta = approx(timet, beta, timet0) .* sqrt(gamI_dev);

    for k in 1:N
        qn[:, k] = approx(timet, qn[:, k], timet0) .* sqrt(gamI_dev);
        fn[:, k] = approx(timet, fn[:, k], timet0);
        gamma[:, k] = approx(timet, gamma[:, k], timet0);
    end

    out = Dict("alpha" => alpha, "beta" => beta, "fn" => fn, "qn" => qn,
               "gamma" => gamma, "q" => q, "B" => B, "type" => "linear",
               "b" => b[2:end], "SSE" => SSE[1:itr-1]);
    return out
end

"""
Calculate elastic logistic regression from function data f, for response y

    elastic_logistic(f, y, timet; B=None, df=20, max_itr=20, smooth=false)
    :param f: array (M,N) of N functions
    :param y: vector (N) of responses
    :param timet: vector (N) describing time samples
    :param B: matrix describing basis functions (M,N) (default=None generates a
              B-spline basis
    :param df: degree of freedom of basis
    :param max_itr: maximum number of iterations
    :param smooth: smooth data

    Returns Dict describing regression
    :return alpha: intercept
    :return beta: regression function
    :return fn: aligned functions
    :return qn: aligned srsfs
    :return gamma: warping functions
    :return q: original srsfs
    :retrun B: basis functions
    :return type: type of regressions
    :return b: coefficients
    :return LL: logistic loss
"""
function elastic_logistic(f, y, timet; B=Union{}, df=20, max_itr=20,
                          smooth=false)

    M, N = size(f);

    if M > 500
        parallel = true;
    elseif N > 100
        parallel = true;
    else
        parallel = false;
    end

    binsize = mean(diff(timet));

    # Create B-spline basis if none provided
    if B == Union{}
        B = bs(timet, df, 4);
    end
    Nb = size(B,2);

    q = f_to_srsf(f, timet, smooth);

    gamma = repmat(linspace(0,1,M), 1, N);
    fn = zeros(M, N);
    qn = zeros(M, N);

    itr = 1;
    LL = zeros(max_itr);
    b = zeros(Nb+1);
    alpha = 0.0;
    beta = zeros(length(timet));
    while itr <= max_itr
        @printf("Iteration: %d\n", itr)
        # align data
        fn = zeros(M, N);
        qn = zeros(M, N);
        for ii in 1:N
            xout = (timet[end] - timet[1]) * gamma[:, ii] + timet[1];
            fn[:,ii] = approx(timet, f[:, ii], xout);
            qn[:,ii] = warp_q_gamma(timet, q[:,ii], gamma[:,ii]);
        end

        Phi = ones(N, Nb+1);
        for ii in 1:N
            for jj in 2:Nb+1
                Phi[ii,jj] = trapz(timet, qn[:,ii].*B[:,jj-1]);
            end
        end

        # find alpha and beta using l_bfgs
        b0 = zeros(Nb+1);
        opt = Opt(:LD_LBFGS, Nb+1);
        maxeval!(opt, 250);
        xtol_abs!(opt, 1e-12);
        ftol_abs!(opt, 1e-30);
        min_objective!(opt, (x,grad) -> logit_optm(x,grad,Phi,y));
        optf, b, ret = optimize(opt, b0);

        alpha = b[1];
        beta = B * b[2:Nb+1];

        # compute the logistic loss
        LL[itr] = logit_loss(b, Phi, y);

        # find gamma
        gamma_new = zeros(M, N);
        if parallel
            gamma_new = @parallel (hcat) for i=1:N
                logistic_warp(beta, timet, q[:, i], y[i]);
            end
        else

            for i = 1:N
                gamma_new[:,i] = logistic_warp(beta, timet, q[:, i], y[i]);
            end
        end

        if norm(gamma-gamma_new) < 1e-5
            break
        else
            gamma = copy(gamma_new);
        end

        itr += 1;
    end

    out = Dict("alpha" => alpha, "beta" => beta, "fn" => fn, "qn" => qn,
               "gamma" => gamma, "q" => q, "B" => B, "type" => "logistic",
               "b" => b[2:end], "Loss" => LL[1:itr-1]);
    return out
end

"""
Calculate elastic m-logistic regression from function data f, for response y

    elastic_mlogistic(f, y, timet; B=None, df=20, max_itr=20, smooth=false)
    :param f: array (M,N) of N functions
    :param y: vector (N) of responses
    :param timet: vector (N) describing time samples
    :param B: matrix describing basis functions (M,N) (default=None generates a
              B-spline basis
    :param df: degree of freedom of basis
    :param max_itr: maximum number of iterations
    :param smooth: smooth data

    Returns Dict describing regression
    :return alpha: intercept
    :return beta: regression function
    :return fn: aligned functions
    :return qn: aligned srsfs
    :return gamma: warping functions
    :return q: original srsfs
    :retrun B: basis functions
    :return type: type of regressions
    :return b: coefficients
    :return n_classes: number of classes
    :return LL: logistic loss
"""
function elastic_mlogistic(f, y, timet; B=Union{}, df=20, max_itr=20,
                           delta=.01, smooth=false)

    M, N = size(f);
    # Code labels
    m = Int(maximum(y));
    Y = zeros(Int32, N, m);
    for ii in 1:N
        Y[ii, y[ii]] = 1;
    end

    if M > 500
        parallel = true;
    elseif N > 100
        parallel = true;
    else
        parallel = false;
    end

    binsize = mean(diff(timet));

    # Create B-spline basis if none provided
    if B == Union{}
        B = bs(timet, df, 4);
    end
    Nb = size(B,2);

    q = f_to_srsf(f, timet, smooth);

    gamma = repmat(linspace(0,1,M), 1, N);
    fn = zeros(M, N);
    qn = zeros(M, N);

    itr = 1;
    LL = zeros(max_itr);
    b = zeros(m*(Nb+1));
    alpha = zeros(m);
    beta = zeros(M, m);
    while itr <= max_itr
        @printf("Iteration: %d\n", itr)
        # align data
        fn = zeros(M, N);
        qn = zeros(M, N);
        for ii in 1:N
            xout = (timet[end] - timet[1]) * gamma[:, ii] + timet[1];
            fn[:,ii] = approx(timet, f[:, ii], xout);
            qn[:,ii] = warp_q_gamma(timet, q[:,ii], gamma[:,ii]);
        end

        Phi = ones(N, Nb+1);
        for ii in 1:N
            for jj in 2:Nb+1
                Phi[ii,jj] = trapz(timet, qn[:,ii].*B[:,jj-1]);
            end
        end

        # find alpha and beta using l_bfgs
        b0 = zeros(m * (Nb+1));
        opt = Opt(:LD_LBFGS, m * (Nb+1));
        maxeval!(opt, 250);
        xtol_abs!(opt, 1e-12);
        ftol_abs!(opt, 1e-30);
        min_objective!(opt, (x,grad) -> mlogit_optm(x,grad,Phi,Y));
        optf, b, ret = optimize(opt, b0);

        B0 = reshape(b, Nb+1, m);
        alpha = B0[1,:];
        beta = zeros(M, m);
        for i in 1:m
            beta[:, i] = B * B0[2:Nb+1, i];
        end

        # compute the logistic loss
        LL[itr] = mlogit_loss(b, Phi, Y);

        # find gamma
        gamma_new = zeros(M, N);
        if parallel
            gamma_new = @parallel (hcat) for i=1:N
                mlogit_warp_grad(alpha, beta, timet, q[:, i], Y[i, :],
                                 delta=delta);
            end
        else

            for i = 1:N
                gamma_new[:,i] = mlogit_warp_grad(alpha, beta, timet, q[:, i],
                                                  Y[i, :], delt=delta);
            end
        end

        if norm(gamma-gamma_new) < 1e-5
            break
        else
            gamma = copy(gamma_new);
        end

        itr += 1;
    end

    out = Dict("alpha" => alpha, "beta" => beta, "fn" => fn, "qn" => qn,
               "gamma" => gamma, "q" => q, "B" => B, "type" => "mlogistic",
               "b" => b[2:end], "n_classes" => m, "Loss" => LL[1:itr-1]);
    return out
end

"""
Prediction from elastic regression model

    elastic_prediction(f, timet, model; y=None, smooth=false)
    :param f: functions to predict
    :param timet: vector describing time samples
    :param model: calculated model (regression, musicologist, mlogistic)
    :param y: true response (default = None)
    :param smooth: smooth data (default = false)

    Returns
    :return y_pred: predicted value
    :return y_labels: labels of predicted value
    :return Perf: Performance metric if truth is supplied
"""
function elastic_prediction(f, timet, model::Dict; y=Union{}, smooth=false)
    q = f_to_srsf(f, timet, smooth);
    n = size(q, 2);

    if model["type"] == "linear" || model["type"] == "logistic"
        y_pred = zeros(n);
    elseif model["type"] == "mlogistic"
        m = model["n_classes"];
        y_pred = zeros(n, m);
    end

    for ii in 1:n
        diff = model["q"] - repmat(q[:, ii], 1, size(model["q"], 2));
        dist = sum(abs(diff).^2, 1).^(1/2);
        q_tmp = warp_q_gamma(timet, q[:, ii],
                             model["gamma"][:, indmin(dist)]);
        if model["type"] == "linear"
            y_pred[ii] = model["alpha"] + trapz(timet, q_tmp.*model["beta"]);
        elseif model["type"] == "logistic"
            y_pred[ii] = model["alpha"] + trapz(timet, q_tmp.*model["beta"]);
        elseif model["type"] == "mlogistic"
            for jj in 1:m
                y_pred[ii,jj] = model["alpha"][jj] + trapz(timet, q_tmp.*model["beta"][:, jj]);
            end
        end
    end

    if y==Union{}
        if model["type"] == "linear"
            Perf = Union{};
            y_labels =  Array(Integer, 0);
        elseif model["type"] == "logistic"
            y_pred = phi(y_pred);
            y_labels = ones(n);
            y_labels[y_pred .< 0.5] = -1;
            Perf = Union{};
        elseif model["type"] == "mlogisistic"
            y_pred = phi(reshape(y_pred,n*m,1));
            y_pred = reshpae(y_pred,n,m);
            max_val, y_labels = findmax(y_pred, 2);
            Perf = Union{};
        end
    else
        if model["type"] == "linear"
            Perf = sum((y-y_pred).^2);
            y_labels = Array(Integer, 0);
        elseif model["type"] == "logistic"
            y_pred = phi(y_pred);
            y_labels = ones(n);
            y_labels[y_pred .< 0.5] = -1;
            TP = sum(y[y_labels .== 1] .== 1);
            FP = sum(y[y_labels .== -1] .== 1);
            TN = sum(y[y_labels .== -1] .== -1);
            FN = sum(y[y_labels .== 1] .== -1);
            Perf = (TP+TN)/(TP+FP+FN+TN);
        elseif model["type"] == "mlogistic"
            y_pred = phi(reshape(y_pred,n*m,1));
            y_pred = reshape(y_pred,n,m);
            max_val, y_labels = findmax(y_pred, 2);
            temp = ind2sub((n,m),vec(y_labels));
            y_labels = temp[2];
            Perf = zeros(m);
            cls_set = collect(1:m+1);
            for ii in 1:m
                cls_sub = setdiff(cls_set, ii);
                TP = sum(y[y_labels .== ii] .== ii);
                FP = sum(y[indexin(y_labels, cls_sub).>0] .== ii)
                TN = sum(y[indexin(y_labels,cls_sub).>0] .==
                         y_labels[indexin(y_labels,cls_sub).>0]);
                FN = sum(indexin(y[y_labels.==ii], cls_sub).>0);
                Perf[ii] = (TP+TN)/(TP+FP+FN+TN);
            end
            Perf = sum(y .==y_labels)/length(y_labels);
        end
    end

    return y_pred, y_labels, Perf
end


function flux(u, q, w, v, x, t, mu, eta)
    f = mu[1]*q + u[1]*mu[2:3];
    return f;
end
function source(u, q, w, v, x, t, mu, eta)
    x1 = x[1];
    x2 = x[2];
    v1 = mu[2];
    v2 = mu[3];
    s = (v1*x2*exp(v1)*exp(v2)*(exp(v2*(x2 - 1)) - 1)*(exp(v1*(x1 - 1)) + v1*x1*exp(v1*(x1 - 1)) - 1))/((exp(v1) - 1)*(exp(v2) - 1)) - (v2*x1*exp(v1 + v2*x2)*(exp(v1*(x1 - 1)) - 1)*(v2*x2 + 2))/((exp(v1) - 1)*(exp(v2) - 1)) - (v1*x2*exp(v2 + v1*x1)*(exp(v2*(x2 - 1)) - 1)*(v1*x1 + 2))/((exp(v1) - 1)*(exp(v2) - 1)) + (v2*x1*exp(v1)*exp(v2)*(exp(v1*(x1 - 1)) - 1)*(exp(v2*(x2 - 1)) + v2*x2*exp(v2*(x2 - 1)) - 1))/((exp(v1) - 1)*(exp(v2) - 1));
    return s;
end
function ubou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    ub = 0.0;
    return ub;
end
function fbou(u, q, w, v, x, t, mu, eta, uhat, n, tau)
    f = flux(u, q, w, v, x, t, mu, eta);
    fb = f[1]*n[1] + f[2]*n[2] + tau[1]*(u[1]-uhat[1]);
    return fb;
end
function initu(x, mu, eta)
    u0 = 0.0;
    return u0;
end


using ConservationLawsParticles
using RecursiveArrayTools, DifferentialEquations, Plots

# model
V(t, x) = 2 + 0.5sin(2x) + 0.3sin(2sqrt(2)*x) + 0.2cos(2sqrt(7)*x)
W(t, x) = x^2/2 - abs(x)
W′(t, x) = x - sign(x)
# W(t, x) = (a=abs(x); a*log(a+1/4))
# W′(t, x) = (a=abs(x); sign(x) * (log(a+1/4) + a/(a+1/4)))
mob(ρ) = max((ρ-1)^2 * (2ρ+1), 0)
smodel = SampledModel((V,), ((W′,),), (mob,))
imodel = IntegratedModel((V,), ((W,),), (mob,))

# initial condition
x0 = ArrayPartition(gaussian_particles(2, 401))

# time span
tspan = (0., 15.)

# ODE system for the particles
sprob = ODEProblem(velocities_gen!, x0, tspan, smodel)
iprob = ODEProblem(velocities_gen!, x0, tspan, imodel)

abstol = reltol = 2e-8

@time ssol = solve(sprob, BS5(); abstol=abstol, reltol=reltol);
@time isol = solve(iprob, BS5(); abstol=abstol, reltol=reltol);

# plot the particle trajectories
plot(legend=false, title="Trajectories", xlabel="time", ylabel="position")
plot!(ssol, vars=1:20:401; color=:blue)
plot!(isol, vars=1:20:401; color=:red)
savefig("traffic2.png")

# plot an animation of the density
anim = @animate for t in range(tspan...; step=1/24)
    p = plot(legend=false, xlims=(-2,20), ylims=(0,1),
        title="Density", xlabel="position", ylabel="density")
    x = ssol(t).x[1]
    min_dens = minimum(pwc_densities(x)[1][1,:,:], dims=1)[:];
    plot!(p, repeat(x, inner=(1,2))[1:20:401,:]',
        hcat(zero(min_dens), min_dens)[1:20:401,:]',
        color=:gray, opacity=.25)
    plot_density!(p, ssol(t); color=:blue)
    plot_density!(p, isol(t); color=:red)
end
gif(anim, "traffic2.gif")


export moment, exp, log, series, dual, sparse_pol, sparse_decompose;

#----------------------------------------------------------------------
function Base.exp(A::Matrix{U}, pt::Vector{T}) where {U,T}
    N = size(A)
    R = fill(zero(eltype(pt)), N[1], N[2])
    for i  in 1:N[1]
        for j in 1:N[2]
            R[i,j] = pt[j]^A[i,j]
        end
    end
    R
end

#----------------------------------------------------------------------
function Base.log(A::Array{U}, pt::Vector{T}) where {U,T}
    N = size(A)
    R = fill(zero(eltype(pt)), N[1], N[2])
    for i  in 1:N[1]
        for j in 1:N[2]
            R[i,j] = real(log(A[i,j])/log(pt[i]))
        end
    end
    R
end

#------------------------------------------------------------------------
"""
```
dual(p::Polynomial) -> Series{C,M}
```
Compute the series associated to the polynomial p, replacing
the variables xi by their dual variables dxi. C is the type of coefficients
of the polynomial p and M its type of monomials.
"""
function dual(p::Polynomial{true,C}) where C
    s = Series{C, DynamicPolynomials.Monomial{true}}()
    for t in p
               s[t.x] = t.α
    end
    return s
end

function dual(t::Term{true,T}) where T
    s = Series{T, Monomial{true}}()
    s[t.x] = t.α
    return s
end

#----------------------------------------------------------------------
"""
```
moment(w::Vector{C}, P::Matrix{C}) -> Vector{Int64} -> C
```
Compute the moment function ``α -> ∑_{i} ω_{i} P_{i}^α``
associated to the sequence P of r points of dimension n, which is a matrix
of size r*n and the weights w.
"""
moment(w, P) = function(α)
  res = 0
  for j in 1:length(w)
      m = 1
      for i in 1:length(α)
        m *= P[i,j]^α[i]
      end
      res+=m*w[j];
  end
  res
end

#------------------------------------------------------------------------
"""
```
series(f::Function, L::Vector{M}) -> Series{C,M}
```
Compute the generating series ``\\sum_{x^{α} \\in L} f(α) z^α``
for a function  ``f: \\mathbb{N}^n \\rightarrow C`` and a sequence L of monomials.
"""
function series(f::Function, L::Vector)
   res = series(f(L[1].z), L[1])
   for m in L
         res[m] = f(m.z)
   end
   res
end

#----------------------------------------------------------------------
"""
```
moment(p::Polynomial, zeta::Vector{C}) -> Vector{Int64} -> C
```
Compute the moment function ``α \\rightarrow p(ζ^α)``.
"""
moment(p::Polynomial, zeta::Vector) =  function(V::Vector{Int})
    U = [zeta[i]^V[i] for i in 1:length(V)]
    p(U)
end

#----------------------------------------------------------------------
"""
```
series(w:: Vector{C}, P::Matrix{C}, L::Vector{M}) -> Series{C,M}
```
Compute the series of the moment sequence ``∑_{i} ω_{i} P_{i}^α`` for ``α \\in L``.
"""
function series(w:: Vector{C}, P, L::Vector{M}) where {C,M}
   series(moment(w,P), L)
end

#----------------------------------------------------------------------
"""
```
series(w:: Vector{C}, P::AbstractMatrix, X, d::Int64) -> Series{C,M}
```
Compute the series of the moment sequence ``∑_i ω_{i} P_{i}^α`` for ``|α| \\leq d``.
"""
function series(w::Vector{C}, P::AbstractMatrix, X, d::Int64) where C
    h = moment(w,P)
    L = monomials(X,seq(0:d))
   series(h,L)
end

#----------------------------------------------------------------------
"""
```
series(p::Polynomial, zeta, X, d::Int64) -> Series
```
Compute the series of moments ``p(ζ^α)`` for ``|α| \\leq d``.
"""
function series(p::Polynomial, zeta, X, d::Int64)
    h = moment(p,zeta)
    L = monomials(X,seq(0:d))
    series(h,L)
end

#----------------------------------------------------------------------
"""
```
sparse_pol(w, E, X) -> Polynomial{true,C}
```
Compute the polynomial ``∑ ωᵢ X^E[i,:]`` with coefficients ``ωᵢ`` and monomial exponents ``E[i,:]``.
"""
function sparse_pol(w::Vector, E::Matrix{Int64}, X)
    sum(w[j]*prod(X[i]^E[j,i] for i in 1:size(E,2)) for j in 1:length(w))
end

#------------------------------------------------------------------------
function sparse_decompose(f, zeta, X, d:: Int64)
    sigma = series(moment(f,zeta), monomials(X,seq(0:d)))
    w, Xi = decompose(sigma)
    w, log(Xi,zeta)
end


# Computes graph Stein discrepancy based on the Langevin operator.
# Args:
#  points - n x d array of sample points
#  weights - n x 1 array of real-valued weights associated with sample points
#    (default: equal weights)
#  gradlogdensity - the gradlogdensity of the target distribution
#  solver - optimization program solver supported by JuMP
function langevin_graph_discrepancy(; points=[],
                                    weights=fill(1/size(points,1), size(points,1)),
                                    gradlogdensity=nothing,
                                    solver=nothing,
                                    kwargs...)
    # make sure solver is defined
    isa(solver, AbstractMathProgSolver) ||
        error("Must specify solver of type AbstractMathProgSolver")
    # Aggregate repeated samples
    sample = SteinDiscrete(points, weights)
    points = sample.support
    weights = sample.weights
    n = length(weights)
    d = size(points,2)
    # Objective coefficients for each g(x_i)
    gradlogdensities = zeros(n, d)
    for i in 1:n
        gradlogdensities[i,:] = gradlogdensity(points[i,:])
    end
    gobjcoefficients = broadcast(*, gradlogdensities, weights)
    # Objective coefficients for each grad g(x_i)
    gradgobjcoefficients = zeros(n, d, d)
    for j = 1:d
        gradgobjcoefficients[:,j,j] = weights
    end
    # solve the optimization problem
    result = affine_graph_discrepancy(sample,
                                      gobjcoefficients,
                                      gradgobjcoefficients,
                                      solver,
                                      "langevin";
                                      kwargs...)
    result
end


using StaticArrays
using StaticArrays
const Vec2f = SVector{2, Float64}
const Vec4f = SVector{4, Float64}
@inline @inbounds norm(v) = sqrt(v[1]^2 + v[2]^2)
@inline @inbounds inpro(v1, v2) = v1[1]*v2[1]+v1[2]*v2[2]

abstract type GeoParametric end

struct Sphere <: GeoParametric
    c::SVector{2, Float64}
    r::Float64
end

@inline function isInside(this::Sphere, x::SVector{2, Float64})
    @fastmath d = norm(this.c-x)
    return (d < this.r )
end

@inline function isIntersect(this::Sphere, x_seq::Vector{Vec2f})
    """
    for i in 1:length(x_seq)-1
        @inbounds x0 = x_seq[i]
        @inbounds x1 = x_seq[i+1]
        v1 = this.c - x0
        v2 = x1 - x0
        # calc dist between line segment and sphere orthogonal projection
        vv = inpro(v1, v2)*v2/norm(v2)^2 # projected vectro
        d = norm(x0 + vv - this.c)
        d < this.r && return true
    end
    """
    for x in x_seq
        isInside(this, x) && return true
    end
    return false
end

function show(this::Sphere, color=:red)
    for theta in 0:0.3:2π
        x = this.c + this.r*Vec2f(cos(theta), sin(theta))
        scatter(x[1], x[2], c=color)
    end
end

"""
function test()
    x = Vec2f(-2., 1.6)
    x_seq = Vec2f[x]
    N = 30
    for i in 0:N
        dx = Vec2f(0.1, randn()*0.05)
        global x = x .+ dx
        push!(x_seq, x)
    end
    for i in 1:N
        x0 = x_seq[i]
        x1 = x_seq[i+1]
        plot([x0[1], x1[1]], [x0[2], x1[2]], c = :black)
    end

    x_0 = SVector(-1.0, 1.)
    x_1 = SVector(0., 1.5)
    x_2 = SVector(1.0, 1.)
    s = Sphere(SVector(0., 0.), 1.3)
    println(isIntersect(s, x_seq))
    show(s)
    println(s)
end
"""


### CONCRETE TYPE: ITERATIVE PROX EVALUATION
# prox! is computed using CG on a system with matrix lambda*A'A + I/gamma
# or lambda*AA' + I/gamma, according to which matrix is smaller.

using LinearAlgebra

struct LeastSquaresIterative{N, R <: Real, RC <: RealOrComplex{R}, M, V <: AbstractArray{RC, N}, O} <: LeastSquares
    A::M # m-by-n operator
    b::V # m (by-p)
    lambda::R
    lambdaAtb::V
    shape::Symbol
    S::O
    res::Array{RC, N} # m (by-p)
    res2::Array{RC, N} # m (by-p)
    q::Array{RC, N} # n (by-p)
end

is_prox_accurate(f::LeastSquaresIterative) = false
is_convex(f::LeastSquaresIterative) = f.lambda >= 0
is_concave(f::LeastSquaresIterative) = f.lambda <= 0

function LeastSquaresIterative(A::M, b::V, lambda::R) where {N, R <: Real, RC <: RealOrComplex{R}, M, V <: AbstractArray{RC, N}}
    if size(A, 1) != size(b, 1)
        error("A and b have incompatible dimensions")
    end
    m, n = size(A)
    x_shape = infer_shape_of_x(A, b)
    if m >= n
        shape = :Tall
        S = AcA(A)
        LeastSquaresIterative{N, R, RC, M, V, AcA}(A, b, lambda, lambda*(A'*b), shape, S, zero(b), [], zeros(RC, x_shape))
    else
        shape = :Fat
        S = AAc(A)
        LeastSquaresIterative{N, R, RC, M, V, AAc}(A, b, lambda, lambda*(A'*b), shape, S, zero(b), zero(b), zeros(RC, x_shape))
    end
end

function (f::LeastSquaresIterative{N, R, RC, M, V})(x::AbstractArray{RC, N}) where {N, R, RC, M, V}
    mul!(f.res, f.A, x)
    f.res .-= f.b
    return (f.lambda/2)*norm(f.res, 2)^2
end

function prox!(y::AbstractArray{D, N}, f::LeastSquaresIterative{N, R, RC, M, V}, x::AbstractArray{D, N}, gamma::R=R(1)) where {N, R, RC, M, V, D <: RealOrComplex{R}}
    f.q .= f.lambdaAtb .+ x./gamma
    # two cases: (1) tall A, (2) fat A
    if f.shape == :Tall
        y .= x
        op = ScaleShift(RC(f.lambda), f.S, RC(1)/gamma)
        IterativeSolvers.cg!(y, op, f.q)
    else # f.shape == :Fat
        # y .= gamma*(f.q - lambda*(f.A'*(f.fact\(f.A*f.q))))
        mul!(f.res, f.A, f.q)
        op = ScaleShift(RC(f.lambda), f.S, RC(1)/gamma)
        IterativeSolvers.cg!(f.res2, op, f.res)
        mul!(y, adjoint(f.A), f.res2)
        y .*= -f.lambda
        y .+= f.q
        y .*= gamma
    end
    mul!(f.res, f.A, y)
    f.res .-= f.b
    return (f.lambda/2)*norm(f.res, 2)^2
end

function gradient!(y::AbstractArray{D, N}, f::LeastSquaresIterative{N, R, RC, M, V}, x::AbstractArray{D, N}) where {N, R, RC, M, V, D <: Union{R, Complex{R}}}
    mul!(f.res, f.A, x)
    f.res .-= f.b
    mul!(y, adjoint(f.A), f.res)
    y .*= f.lambda
    fy = (f.lambda/2)*dot(f.res, f.res)
end

function prox_naive(f::LeastSquaresIterative{N}, x::AbstractArray{D, N}, gamma::R=R(1)) where {N, R, D <: RealOrComplex{R}}
    y = IterativeSolvers.cg(f.lambda*f.A'*f.A + I/gamma, f.lambda*f.A'*f.b + x/gamma)
    fy = (f.lambda/2)*norm(f.A*y-f.b)^2
    return y, fy
end


@doc raw"""
conductance based channel, simple format

$$i = \bar{g} m^a h^b (V - E)$$

$$i = g m^a h^b (V - E)$$
$$\dot{m} = \frac{m_\infty(V) - m}{\tau_m(V)}$$
$$\dot{h} = \frac{h_\infty(V) - h}{\tau_h(V)}$$

activation and inactivation function would be encoded
as Boltzmann function:

$$m_\infty(V) = \frac{1 - \zeta}{1 + \exp(\frac{V_{\text{half}} - V}{k})} + \zeta$$

the time constant would be encoded as guassian function:

$$\tau(V) = C_{\text{base}} + C_{\text{amp}} \exp(\frac{- (V_{\text{max}} - V) ^ 2}{\sigma ^ 2})$$
"""
mutable struct SimpleIonChannel <: AbstractIonChannel
    name::String
    ion::Symbol
    g::Real

    m::Kinetics
    h::Kinetics
end

function update!(ch::SimpleIonChannel)
    if ~isnothing(ch.m.Vhalf)
        ch.m.infty = (V) -> (1 - ch.m.zeta) / (1 + exp((ch.m.Vhalf - V) / ch.m.k)) + ch.m.zeta
    end

    if ~isnothing(ch.m.Cbase)
        ch.m.tau = (V) -> ch.m.Cbase + ch.m.Camp * exp(- ((ch.m.Vmax - V) / ch.m.sigma)^2)
    end

    if ~isnothing(ch.h.Vhalf)
        ch.h.infty = (V) -> (1 - ch.h.zeta) / (1 + exp(-(ch.h.Vhalf - V) / ch.h.k)) + ch.h.zeta
    end

    if ~isnothing(ch.h.Cbase)
        ch.h.tau = (V) -> ch.h.Cbase + ch.h.Camp * exp(- ((ch.h.Vmax - V) / ch.h.sigma)^2)
    end

    ch
end

function dof(channel::SimpleIonChannel)
    _m = channel.m._type == :evolving ? 1 : 0
    _h = channel.h._type == :evolving ? 1 : 0
    [_m, _h]
end

function step(ch::SimpleIonChannel; V::Real, var::Vector{T}, E::Real) where {T <: Real}
    _var_idx = 1
    i = ch.g * (V - E)

    if ch.m._type == :evolving
        m = var[_var_idx]
        _var_idx += 1

        i = i * m ^ ch.m.n
        dm = (ch.m.infty(V) - m) / ch.m.tau(V)
    elseif ch.m._type == :instantaneous
        i = i * ch.m.infty(V) ^ ch.m.n
        dm = nothing
    else
        dm = nothing
    end

    if ch.h._type == :evolving
        h = var[_var_idx]

        i = i * h ^ ch.h.n
        dh = (ch.h.infty(V) - h) / ch.h.tau(V)
    elseif ch.h._type == :instantaneous
        i = i * ch.h.infty(V) ^ ch.h.n
        dh = nothing
    else
        dh = nothing
    end

    return (i, [dm, dh])
end

function itr_kinetics(ch::SimpleIonChannel)
    [ch.m, ch.h]
end

function current(ch::SimpleIonChannel; V::Vector{T}, var::Array{T, 2}, E::T) where {T <: Real}
    _var_idx = 1
    i = ch.g .* (V .- E)

    if ch.m._type == :evolving
        m = var[_var_idx, :]
        _var_idx += 1

        i = i .* m .^ ch.m.n
    elseif ch.m._type == :instantaneous
        i = i .* ch.m.infty.(V) .^ ch.m.n
    else
        nothing
    end

    if ch.h._type == :evolving
        h = var[_var_idx, :]

        i = i .* h .^ ch.h.n
    elseif ch.h._type == :instantaneous
        i = i .* ch.h.infty.(V) .^ ch.h.n
    else
        nothing
    end

    return i
end


# SCRIPT

function main()
    for a in 1:998
        for b in (a + 1):(1000 - a)
            c = 1000 - (a + b);
            if (a * a + b * b == c * c)
                println(a, " ", b, " ", c)
                exit()
            end
        end
    end
end

main()


using Rubin
using Tests
using Elliptic
using HypergeometricFunctions
using Polylogarithms
using SpecialFunctions

@test integrate(csc(x)^-1//2*(a+b*csc(x))*(A+B*csc(x)+C*csc(x)^2), x) == :(csc(x)^(1/2)*(-2*B*b+-2*C*a)*cos(x)+csc(x)^(1/2)*sin(x)^(1/2)*(-2*a*(A+-1C)+2*B*b)*Elliptic.E(-1//2*x+1//4*pi,2)+-1*csc(x)^(1/2)*sin(x)^(1/2)*(2*A*b+2*B*a+2//3*C*b)*Elliptic.F(-1//2*x+1//4*pi,2)+-2//3*C*b*csc(x)^3//2*cos(x))

test/4 Trig functions/4.6 Cosecant/4.6.4.2 (a+b csc)^m (d csc)^n (A+B csc+C csc^2).jl


using TransferEntropy
using Tullio
using Entropies

MI(x::AbstractVector, y) = TransferEntropy.mutualinfo(vec(x), vec(y), RectangularBinning(10); base=ℯ)
MI(X::AbstractArray, Y) = @tullio I[i, j] := MI(X[:, i], Y[:, j]) # Columns to match cor and corspearman
MI(X) = MI(X, X)
export MI

H(x) = Entropies.genentropy(probabilities(vec(x), 10), q=1.0, base=ℯ)

MI_norm(x::AbstractVector, y)  = MI(x, y)/(H(x) + H(y) - MI(x, y))
MI_norm(X::AbstractArray, Y) = @tullio I[i, j] := MI_norm(X[:, i], Y[:, j])
MI_norm(X) = MI_norm(X, X)
export MI_norm

# Something doesn't seem quite right...


using Test, EllZeroTrendFiltering

#TODO let to avoid global scope warning, change back in julia 1.0?
let
    t = [-1; 0; 1]

    g = x -> 0
    ℓ = compute_transition_costs(g, t)
    @test ℓ[1,2].P == 1/6*[2 1; 1 2]
    @test ℓ[1,2].q == [0,0]
    @test ℓ[1,2].r == 0

    g = x -> 1
    ℓ = compute_transition_costs(g, t)
    @test ℓ[1,2].P == 1/6*[2 1; 1 2]
    @test ℓ[1,2].q == [-1,-1]
    @test ℓ[1,2].r == 1

    g = x -> x
    ℓ = compute_transition_costs(g, t)
    @test ℓ[2,3].P == 1/6*[2 1; 1 2]
    @test ℓ[2,3].q ≈ [-1/3,-2/3]
    @test ℓ[2,3].r ≈ 1/3
    @test ℓ[2,3](0,0) == 1/3
    @test ℓ[2,3](1,1) == 1/3
    @test ℓ[2,3](1,2) ≈ 1
    @test ℓ[2,3](1,0) ≈ 1/3
end


# # One-dimensional regression
#
# ## Setup
#
# Loading the necessary packages and setting seed.

using AbstractGPs
using Distributions
using StatsFuns

using Plots
default(legend=:outertopright, size=(700,400))

using Random
Random.seed!(1234)
#md nothing #hide

# Load toy regression
# [dataset](https://github.com/GPflow/docs/blob/master/doc/source/notebooks/basics/data/regression_1D.csv)
# taken from GPFlow examples.

x = [0.8658165855998895, 0.6661700880180962, 0.8049218148148531, 0.7714303440386239,
    0.14790478354654835, 0.8666105548197428, 0.007044577166530286, 0.026331737288148638,
    0.17188596617099916, 0.8897812990554013, 0.24323574561119998, 0.028590102134105955]
y = [1.5255314337144372, 3.6434202968230003, 3.010885733911661, 3.774442382979625,
    3.3687639483798324, 1.5506452040608503, 3.790447985799683, 3.8689707574953,
    3.4933565751758713, 1.4284538820635841, 3.8715350915692364, 3.7045949061144983]
scatter(x, y; xlabel="x", ylabel="y", legend=false)

# We split the observations into train and test data.

x_train = x[1:8]
y_train = y[1:8]
x_test = x[9:end]
y_test = y[9:end]
#md nothing #hide

# We instantiate a Gaussian process with a Matern kernel. The kernel has
# fixed variance and length scale parameters of default value 1.

f = GP(Matern52Kernel())
#md nothing #hide

# We create a finite dimentional projection at the inputs of the training dataset
# observed under Gaussian noise with standard deviation $\sigma = 0.1$, and compute the
# log-likelihood of the outputs of the training dataset.

fx = f(x_train, 0.1)
logpdf(fx, y_train)

# We compute the posterior Gaussian process given the training data, and calculate the
# log-likelihood of the test dataset.

p_fx = posterior(fx, y_train)
logpdf(p_fx(x_test), y_test)

# We plot the posterior Gaussian process along with the observations.

scatter(
    x_train, y_train;
    xlim=(0,1), xlabel="x", ylabel="y",
    title="posterior (default parameters)", label="Train Data",
)
scatter!(x_test, y_test; label="Test Data")
plot!(p_fx, 0:0.001:1; label=false)

# ## Markov Chain Monte Carlo
#
# Previously we computed the log likelihood of the untuned kernel parameters of the GP.
# We now also perform approximate inference over said kernel parameters using different
# Markov chain Monte Carlo (MCMC) methods. I.e., we approximate the posterior distribution
# of the kernel parameters with samples from a Markov chain.
#
# We define a function which returns the log-likelihood of the data for different variance
# and inverse lengthscale parameters of the Matern kernel. We ensure that these parameters are
# positive with the softplus function
# ```math
# f(x) = \log (1 + \exp x).
# ```

struct GPLoglikelihood{X,Y}
    x::X
    y::Y
end

function (ℓ::GPLoglikelihood)(params)
    kernel = ScaledKernel(
        transform(
            Matern52Kernel(),
            ScaleTransform(softplus(params[1]))
        ),
        softplus(params[2]),
    )
    f = GP(kernel)
    fx = f(ℓ.x, 0.1)
    return logpdf(fx, ℓ.y)
end

const loglik_train = GPLoglikelihood(x_train, y_train)
#md nothing #hide

# We define a Gaussian prior for the joint distribution of the two transformed kernel
# parameters. We assume that both parameters are independent with mean 0 and variance 1.

logprior(params) = logpdf(MvNormal(2, 1), params)
#md nothing #hide

# ### Hamiltonian Monte Carlo
#
# We start with a Hamiltonian Monte Carlo (HMC) sampler. More precisely, we use the
# [No-U-Turn sampler (NUTS)](http://www.jmlr.org/papers/volume15/hoffman14a/hoffman14a.pdf),
# which is provided by the Julia packages
# [AdvancedHMC.jl](https://github.com/TuringLang/AdvancedHMC.jl/) and
# [DynamicHMC.jl](https://github.com/tpapp/DynamicHMC.jl/).
#
# #### AdvancedHMC
#
# We start with performing inference with AdvancedHMC.

using AdvancedHMC
using ForwardDiff

# Set the number of samples to draw and warmup iterations.

n_samples = 2_000
n_adapts = 1_000
#md nothing #hide

# Define a Hamiltonian system of the log joint probability.

logjoint_train(params) = loglik_train(params) + logprior(params)
metric = DiagEuclideanMetric(2)
hamiltonian = Hamiltonian(metric, logjoint_train, ForwardDiff)
#md nothing #hide

# Define a leapfrog solver, with initial step size chosen heuristically.

initial_params = rand(2)
initial_ϵ = find_good_stepsize(hamiltonian, initial_params)
integrator = Leapfrog(initial_ϵ)
#md nothing #hide

# Define an HMC sampler, with the following components:
# - multinomial sampling scheme,
# - generalised No-U-Turn criteria, and
# - windowed adaption for step-size and diagonal mass matrix

proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator)
adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
#md nothing #hide

# We draw samples from the posterior distribution of kernel parameters. These samples
# are in the unconstrained space $\mathbb{R}^2$.

samples, _ = sample(
    hamiltonian,
    proposal,
    initial_params,
    n_samples,
    adaptor,
    n_adapts;
    progress=false,
)
#md nothing #hide

# We transform the samples back to the constrained space and compute the mean of both
# parameters:

samples_constrained = [map(softplus, p) for p in samples]
mean_samples = mean(samples_constrained)

# We plot a histogram of the samples for the two parameters.
# The vertical line in each graph indicates the mean of the samples.

histogram(
    reduce(hcat, samples_constrained)';
    xlabel="sample", ylabel="counts", layout=2,
    title=["inverse length scale" "variance"], legend=false,
)
vline!(mean_samples'; linewidth=2)

# We approximate the log-likelihood of the test data using the posterior Gaussian processes
# for kernels with the sampled kernel parameters. We can observe that there is a significant
# improvement over the log-likelihood of the test data with respect to the posterior
# Gaussian process with default kernel parameters of value 1.

struct GPPosterior{X,Y}
    x::X
    y::Y
end

function (g::GPPosterior)(p)
    kernel = ScaledKernel(
        transform(
            Matern52Kernel(),
            ScaleTransform(softplus(p[1]))
        ),
        softplus(p[2]),
    )
    f = GP(kernel)
    return posterior(f(g.x, 0.1), g.y)
end

const gp_posterior = GPPosterior(x_train, y_train)

mean(logpdf(gp_posterior(p)(x_test), y_test) for p in samples)

# We sample a function from the posterior GP for the final 100 samples of kernel
# parameters.

plt = scatter(
    x_train, y_train;
    xlim=(0,1), xlabel="x", ylabel="y",
    title="posterior (AdvancedHMC)", label="Train Data",
)
scatter!(plt, x_test, y_test; label="Test Data")
for p in samples[(end-100):end]
    sampleplot!(plt, gp_posterior(p)(0:0.02:1), 1)
end
plt

# #### DynamicHMC
#
# We repeat the inference with DynamicHMC.

using DynamicHMC
using LogDensityProblems

# We have to implement the LogDensityProblems interface for `GPLogLikelihood`.

## Log joint density
LogDensityProblems.logdensity(ℓ::GPLoglikelihood, params) = ℓ(params) + logprior(params)

## The parameter space is two-dimensional
LogDensityProblems.dimension(::GPLoglikelihood) = 2

## `GPLoglikelihood` does not allow to evaluate derivatives of
## the log-likelihood function
function LogDensityProblems.capabilities(::Type{<:GPLoglikelihood})
    LogDensityProblems.LogDensityOrder{0}()
end

# Now we can draw samples from the posterior distribution of kernel parameters with
# DynamicHMC. Again we use [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)
# to compute the derivatives of the log joint density with automatic differentiation.

samples = mcmc_with_warmup(
    Random.GLOBAL_RNG,
    ADgradient(:ForwardDiff, loglik_train),
    n_samples;
    reporter = NoProgressReport(),
).chain
#md nothing #hide

# We transform the samples back to the constrained space and compute the mean of both
# parameters:

samples_constrained = [map(softplus, p) for p in samples]
mean_samples = mean(samples_constrained)

# We plot a histogram of the samples for the two parameters.
# The vertical line in each graph indicates the mean of the samples.

histogram(
    reduce(hcat, samples_constrained)';
    xlabel="sample", ylabel="counts", layout=2,
    title=["inverse length scale" "variance"], legend=false,
)
vline!(mean_samples'; linewidth=2)

# Again we can observe that there is a significant improvement over the log-likelihood
# of the test data with respect to the posterior Gaussian process with default kernel
# parameters.

mean(logpdf(gp_posterior(p)(x_test), y_test) for p in samples)

# We sample a function from the posterior GP for the final 100 samples of kernel
# parameters.

plt = scatter(
    x_train, y_train;
    xlim=(0,1), xlabel="x", ylabel="y",
    title="posterior (DynamicHMC)", label="Train Data",
)
scatter!(plt, x_test, y_test; label="Test Data")
for p in samples[(end-100):end]
    sampleplot!(plt, gp_posterior(p)(0:0.02:1), 1)
end
plt

# ### Elliptical slice sampling
#
# Instead of HMC, we use
# [elliptical slice sampling](http://proceedings.mlr.press/v9/murray10a/murray10a.pdf)
# which is provided by the Julia package
# [EllipticalSliceSampling.jl](https://github.com/TuringLang/EllipticalSliceSampling.jl/).

using EllipticalSliceSampling

# We draw 2000 samples from the posterior distribution of kernel parameters.

samples = sample(
    ESSModel(
        MvNormal(2, 1), # Gaussian prior
        loglik_train,
    ),
    ESS(),
    n_samples;
    progress=false,
)
#md nothing #hide

# We transform the samples back to the constrained space and compute the mean of both
# parameters:

samples_constrained = [map(softplus, p) for p in samples]
mean_samples = mean(samples_constrained)

# We plot a histogram of the samples for the two parameters.
# The vertical line in each graph indicates the mean of the samples.

histogram(
    reduce(hcat, samples_constrained)';
    xlabel="sample", ylabel="counts", layout=2,
    title=["inverse length scale" "variance"],
)
vline!(mean_samples'; layout=2, labels="mean")

# Again we can observe that there is a significant improvement over the log-likelihood
# of the test data with respect to the posterior Gaussian process with default kernel
# parameters.

mean(logpdf(gp_posterior(p)(x_test), y_test) for p in samples)

# We sample a function from the posterior GP for the final 100 samples of kernel
# parameters.

plt = scatter(
    x_train, y_train;
    xlim=(0,1), xlabel="x", ylabel="y",
    title="posterior (EllipticalSliceSampling)", label="Train Data",
)
scatter!(plt, x_test, y_test; label="Test Data")
for p in samples[(end-100):end]
    sampleplot!(plt, gp_posterior(p)(0:0.02:1), 1)
end
plt

# ## Variational Inference
#
# Sanity check for the Evidence Lower BOund (ELBO) implemented according to
# M. K. Titsias's _Variational learning of inducing variables in sparse Gaussian processes_.

elbo(fx, y_train, f(rand(5)))

# We use the LBFGS algorithm to maximize the given ELBO. It is provided by the Julia
# package [Optim.jl](https://github.com/JuliaNLSolvers/Optim.jl).

using Optim

# We define a function which returns the negative ELBO for different variance and inverse
# lengthscale parameters of the Matern kernel and different pseudo-points. We ensure that
# the kernel parameters are positive with the softplus function
# ```math
# f(x) = \log (1 + \exp x),
# ```
# and that the pseudo-points are in the unit interval $[0,1]$ with the logistic function
# ```math
# f(x) = \frac{1}{1 + \exp{(-x)}}.
# ```

struct NegativeELBO{X,Y}
    x::X
    y::Y
end

function (g::NegativeELBO)(params)
    kernel = ScaledKernel(
        transform(
            Matern52Kernel(),
            ScaleTransform(softplus(params[1]))
        ),
        softplus(params[2]),
    )
    f = GP(kernel)
    fx = f(g.x, 0.1)
    return -elbo(fx, g.y, f(logistic.(params[3:end])))
end
#md nothing #hide

# We randomly initialize the kernel parameters and 5 pseudo points, and minimize the
# negative ELBO with the LBFGS algorithm and obtain the following optimal parameters:

x0 = rand(7)
opt = optimize(NegativeELBO(x_train, y_train), x0, LBFGS())

#-

opt.minimizer

# The optimized value of the inverse lengthscale is

softplus(opt.minimizer[1])

# and of the variance is

softplus(opt.minimizer[2])

# We compute the log-likelihood of the test data for the resulting approximate
# posterior. We can observe that there is a significant improvement over the
# log-likelihood with the default kernel parameters of value 1.

opt_kernel = ScaledKernel(
    transform(
        Matern52Kernel(),
        ScaleTransform(softplus(opt.minimizer[1]))
    ),
    softplus(opt.minimizer[2]),
)
opt_f = GP(opt_kernel)
opt_fx = opt_f(x_train, 0.1)
ap = approx_posterior(VFE(), opt_fx, y_train, opt_f(logistic.(opt.minimizer[3:end])))
logpdf(ap(x_test), y_test)

# We visualize the approximate posterior with optimized parameters.

scatter(
    x_train, y_train;
    xlim=(0,1), xlabel="x", ylabel="y",
    title="posterior (VI with sparse grid)", label="Train Data",
)
scatter!(x_test, y_test; label="Test Data")
plot!(ap, 0:0.001:1; label=false)
vline!(logistic.(opt.minimizer[3:end]); label="Pseudo-points")


function test_pgbf()
  s = pgbf(1.0)
  px = pgbf(1.0,0,0,0,1,0,0)
  @assert isapprox(amplitude(s,0,0,0),0.71270547)
  @assert isapprox(amplitude(px,0,0,0),0)
  @printf "test_pgbf is passed\n"
end

function test_cgbf()
  c = cgbf(0.0,0.0,0.0)
  push!(c,1,1)
  @assert isapprox(amplitude(c,0,0,0),0.71270547)
  c2 = cgbf(0,0,0)
  push!(c2,1,0.2)
  push!(c2,0.5,0.2)
  @assert isapprox(overlap(c2,c2),1)
  @printf "test_cgbf is passed\n"
end

function test_overlap()
  s = pgbf(1.0)
  px = pgbf(1.0,0,0,0,1,0,0)
  @assert overlap1d(0,0,0.,0.,1.) == 1
  @assert gaussian_product_center(s,s) == [0,0,0]
  @assert isapprox(overlap(s,s),1)
  @assert isapprox(overlap(px,px),1)
  @assert isapprox(overlap(s,px),0)
  @assert binomial_prefactor(0,0,0,0.,0.) == 1
end

function test_utils()
  @assert factorial2(6)==48
  @assert collect(pairs(3)) == Any[(1,1),(1,2),(1,3),(2,2),(2,3),(3,3)]
  @assert collect(pairs(3,"subdiag")) == Any[(1,2),(1,3),(2,3)]
  @assert collect(pairs(2,"rect")) == Any[(1,1),(1,2),(2,1),(2,2)]
  @assert iindex(1,1,1,1) == 1
  @assert iindex(1,1,1,2) == iindex(1,1,2,1) == iindex(1,2,1,1) == iindex(2,1,1,1) == 2
  @assert iindex(1,1,2,2) == iindex(2,2,1,1) == 4
  @printf "test_utils is passed\n"
end


function insertsort!(A)
    for k in eachindex(A)
        while k>1 && A[k-1] > A[k]
            A[k-1], A[k] = A[k], A[k-1]
            k -= 1
        end
    end
    return A
end

function insertsortperm!(A)
    ind=collect(1:length(A))
    for k in eachindex(A)
        while k>1 && A[k-1] > A[k]
            A[k-1], A[k] = A[k], A[k-1]
            ind[k-1], ind[k] = ind[k], ind[k-1]
            k -= 1
        end
    end
    return ind
end

insertsort(A) = insertsort!(copy(A))

insertsortperm(A) = insertsortperm!(copy(A))


include("../model/SensorModel.jl")
using Base.Test

let cartesian_point = [4 2]

    homogeneous_point = cart2hom(cartesian_point)

    @test homogeneous_point[1] ≈ 4
    @test homogeneous_point[2] ≈ 2
    @test homogeneous_point[3] ≈ 1
end

let homogeneous_point = [4 2 1]

    cartesian_point = hom2cart(homogeneous_point)

    @test cartesian_point[1] ≈ 4
    @test cartesian_point[2] ≈ 2
end

let homogeneous_point = [4 2 0.5]

    cartesian_point = hom2cart(homogeneous_point)

    @test cartesian_point[1] ≈ 8
    @test cartesian_point[2] ≈ 4
end

let homogeneous_point = [4; 2; 0.5]

    cartesian_point = hom2cart(homogeneous_point)

    @test cartesian_point[1] ≈ 8
    @test cartesian_point[2] ≈ 4
end

let homogeneous_point = [4 2 0]

    cartesian_point = hom2cart(homogeneous_point)

    @test cartesian_point[1] ≈ Inf
    @test cartesian_point[2] ≈ Inf
end

let cartesian_point = [1.5 0.5]
    rotation = π/2

    homogeneous_point = cart2hom(cartesian_point)
    rotated_point = rotate(rotation) * transpose(homogeneous_point)
    rotated_cartesian_point = hom2cart(rotated_point)

    @test rotated_cartesian_point[1] ≈ -0.5
    @test rotated_cartesian_point[2] ≈ 1.5
end

let cartesian_point = [1.5 0.5]
    translation_x = 3
    translation_y = -4

    homogeneous_point = cart2hom(cartesian_point)
    translated_point =  translate(translation_x, translation_y) * transpose(homogeneous_point)
    translated_cartesian_point = hom2cart(translated_point)

    @test translated_cartesian_point[1] ≈ 4.5
    @test translated_cartesian_point[2] ≈ -3.5
end

let cartesian_point = [1.5 0.5]
    rotation = π/2
    translation_x = 3

    homogeneous_point = cart2hom(cartesian_point)
    first_rotated_then_translated_point = translate(translation_x, 0) * rotate(rotation) * transpose(homogeneous_point)
    transformed_point = hom2cart(first_rotated_then_translated_point)

    @test transformed_point[1] ≈ 2.5
    @test transformed_point[2] ≈ 1.5
end


"""
        greedy_d1_coloring

Find a coloring of a given input graph such that
no two vertices connected by an edge have the same
color using greedy approach. The number of colors
used may be equal or greater than the chromatic
number χ(G) of the graph.
"""
function color_graph(g::VSafeGraph, alg::GreedyD1Color)
    v = nv(g)
    result = zeros(Int, v)
    result[1] = 1
    available = BitArray(undef, v)
    for i = 2:v
        for j in inneighbors(g, i)
            if result[j] != 0
                available[result[j]] = true
            end
        end
        for cr = 1:v
            if available[cr] == false
                result[i] = cr
                break
            end
        end
        fill!(available, false)
    end
    return result
end


function trinary_to_decimal(str::AbstractString)
    result = 0
    try
        for (i,d) in enumerate(Iterators.reverse(str))
            p = parse(Int,d)
            p > 3 && return 0
            result += p*(3^(i-1))
        end
    catch(e)
        isa(e,ArgumentError) && return 0
        throw(e)
    end
    result
end

# Approx 2x quicker. From @edit it looks like Julia is using :ccall, which I think counts as cheating
function trinary_to_decimal_reference(str::AbstractString)
    try
        parse(Int,str,base=3)
    catch(e)
        return 0
    end
end


import FactoredEikonalFastMarching.getAnalytic3DeikonalSolutionAll
import FactoredEikonalFastMarching.getAnalytic2DeikonalSolutionAll

function getConstModel(n)
m = ones(Float64,tuple(n...));
return m;
end

function getInvLinearModel(top,bottom,n)
	return getLinearModel(1./(top^2),1./(bottom^2),n);;
end

function getLinearModel(top,bottom,n)
	if length(n)==2
		rho2 = linspace(top,bottom,n[2]);
		rho1 = ones(Float64,n[1]);
		rho = rho1*rho2'
	else
		rho_t = linspace(top,bottom,n[3]);
		rho = ones(Float64,tuple(n...));
		for k=1:n[3]
			rho[:,:,k].*=rho_t[k];
		end
	end
	return rho;
end

function getWedgeModel(n::Array{Int64});
h = 1./(n-1);
x = linspace(0.0,1.0,n[1]);
y = linspace(0.0,1.0,n[2]);

X = x*ones(Float64,n[2])';
Y = ones(Float64,n[1])*y';
Z = 0.25*(tanh((4*Y-X - 0.75)*20)) + 0.75;
Z[:,end - div(n[2],2)+1:end] = Z[:,div(n[2],2) : -1 : 1];
return Z.^2,getRegularMesh([0.0,1.0,0.0,1.0],n-1);
end

function waveGuideModel(n::Array{Int64})
h = 1./(n-1);
x = linspace(0.0,1.0,n[1]);
y = linspace(0.0,1.0,n[2]);

X = x*ones(Float64,n[2])';
Y = ones(Float64,n[1])*y';
Z = 1.25*(1-0.4*exp(-32*abs(X - 0.5).^2));
return 1./(Z.^2),getRegularMesh([0.0,1.0,0.0,1.0],n-1);
end

function getSmoothGaussianMedium(n::Array{Int64,1},h::Array{Float64,1})

src_kappa = zeros(Int64,2);
xsrc = zeros(Float64,2);
src_kappa[1] = div(n[1],2);
src_kappa[2] = div(n[2],2);

xsrc[1] = (src_kappa[1]-1)*h[1];
xsrc[2] = (src_kappa[2]-1)*h[2];

X1,X2 = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2]);

kappaSquared = exp(-4*(X1 - xsrc[1]).^2 - 8*(X2 - xsrc[2]).^2);
return kappaSquared,1;
end

# function getSmoothGaussianMedium(n::Array{Int64,1},h::Array{Float64,1})
# src = div(n,2);
# kappaSquared = [];
# T_exact = [];
# if length(n)==2
	# (T1_exact,G11_exact,G12_exact) = getSmoothFactoredModel(n,h);
	# (T0,G01,G02,L0) = getAnalytic2DeikonalSolutionAll(n,h,src);
	# G1_exact = T0.*G11_exact + G01.*T1_exact;
	# G2_exact = T0.*G12_exact + G02.*T1_exact;
	# kappaSquared = G1_exact.*G1_exact + G2_exact.*G2_exact;
	# T_exact = T0.*T1_exact;
# else
	# (T1_exact,G11_exact,G12_exact,G13_exact) = getSmoothFactoredModel3D(n,h);
	# (T0,G01,G02,G03) = getAnalytic3DeikonalSolutionAll(n,h,src);
	# G1_exact = T0.*G11_exact + G01.*T1_exact;
	# G2_exact = T0.*G12_exact + G02.*T1_exact;
	# G3_exact = T0.*G13_exact + G03.*T1_exact;
	# kappaSquared = G1_exact.*G1_exact + G2_exact.*G2_exact + G3_exact.*G3_exact;
	# T_exact = T0.*T1_exact;
# end

# return (kappaSquared),src,T_exact;
# end

function getSmoothFactoredModel(n::Array{Int64,1},h::Array{Float64,1})
xsrc = zeros(2);
src_kappa = zeros(Int64,2);
src_kappa[1] = div(n[1],2);
src_kappa[2] = div(n[2],2);

xsrc[1] = (src_kappa[1]-1)*h[1];
xsrc[2] = (src_kappa[2]-1)*h[2];

X1,X2 = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2]);

sigma = 1.0;
T1_exact = (exp( - (sigma*((X1 - xsrc[1]).^2) + 4*sigma*((X2-xsrc[2]).^2))) + 1)/2;
G11_exact = -2*sigma*(X1 - xsrc[1]).*exp( - (sigma*((X1 - xsrc[1]).^2) + 4*sigma*((X2-xsrc[2]).^2)))/2;
G12_exact = -8*sigma*(X2 - xsrc[2]).*exp( - (sigma*((X1 - xsrc[1]).^2) + 4*sigma*((X2-xsrc[2]).^2)))/2;

return T1_exact,G11_exact,G12_exact;
end

# function getSmoothFactoredModel3D(n::Array{Int64,1},h::Array{Float64,1})
# xsrc = zeros(3);
# src_kappa = zeros(Int64,3);
# src_kappa[1] = div(n[1],3);
# src_kappa[2] = div(n[2],4);
# src_kappa[3] = div(n[3],2);

# xsrc[1] = (src_kappa[1]-1)*h[1];
# xsrc[2] = (src_kappa[2]-1)*h[2];
# xsrc[3] = (src_kappa[3]-1)*h[3];

# X1,X2,X3 = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2],(0:(n[3]-1))*h[3]);

# Sigma = [0.1,0.4,0.2];
# EXPRSkewedSquared = exp(-Sigma[1]*((X1 - xsrc[1]).^2) - Sigma[2]*((X2-xsrc[2]).^2) - Sigma[3]*((X3-xsrc[3]).^2));

# T1_exact = EXPRSkewedSquared/2 + 1/2;
# G11_exact = -Sigma[1]*(X1 - xsrc[1]).*EXPRSkewedSquared;
# G12_exact = -Sigma[2]*(X2 - xsrc[2]).*EXPRSkewedSquared;
# G13_exact = -Sigma[3]*(X3 - xsrc[3]).*EXPRSkewedSquared;

# return T1_exact,G11_exact,G12_exact,G13_exact;
# end


# Many derivative free methods of different orders
#
# TODO: rework Order5 #https://pdfs.semanticscholar.org/ce50/3210d96f653a14b28da96600d5990d2abe97.pdf
# https://content.sciendo.com/view/journals/tmj/10/4/article-p103.xml 7 and 8
# order8: https://www.hindawi.com/journals/ijmms/2012/493456/ref/

##################################################
## Guard against non-robust algorithms
##
## By default, we do this by deciding if we
## should take a secant step instead of the algorithm For example, for
## Steffensen which is quadratically convergent and the Secant method
## which is only superlinear,
## the error, e_{n+1} = x_{n+1} - alpha, may be smaller after a secant
## step than a Steffensen step. (It is only once x_n is close enough
## to alpha that the method is quadratically convergent.
## The Steffensen error is
## Δn+1 = f[x,x+fx, alpha]/f[x, x+fx] * (1 - f[x, alpha]) (x-alpha)^2
##      ≈ f''/(2f') * ( 1 + f') Δn^2
## The Secant error is
## Δn+1 = f[x,x_{-1},alpha] / f[x,x_{-1}] * (x-alpha) * (x_{-1} - alpha)
##      ≈  f''/(2f')  Δn ⋅ Δn-1
## The ratio is ≈ (1 + f')(Δn / Δn-1)
## It seems reasonable, that a Steffensen step is preferred when
## the ratio satisfies -1 < (1+f') ⋅ Δn /Δn-1 < 1
## We could use f' ~ fp = (fx1-fx0)/(x1-x0); but our proxy for
## Δn/Δn-1 is problematic, as we don't know alpha, and using xn-x_{n-1}
## can be an issue when only x1 and not x0 is specified. This needs
## working around.
##
## Instead, as Steffensen is related to Newton as much as
## (f(x+fx) - fx)/fx  ≈ f'(x), we take a Steffensen step if |fx|
## is small enough. For this we use |fx| <= x/1000; which
## seems to work reasonably well over several different test cases.
@inline function do_guarded_step(M::AbstractSecant, o::UnivariateZeroState{T,S}) where {T, S}
    x, fx = o.xn1, o.fxn1
    1000 * abs(fx) >  max(oneunit(S), abs(x) * oneunit(S) /oneunit(T)) * one(T)
end

# check if we should guard against step for method M; call N if yes, P if not
function update_state_guarded(M::AbstractSecant,N::AbstractUnivariateZeroMethod, P::AbstractUnivariateZeroMethod, fs, o, options)
    if do_guarded_step(M, o)
        #@debug "do secant step"
        return update_state(N, fs, o, options)
    else
        #@debug "do $N step"
        update_state(P, fs, o, options)
    end
end

##################################################

function init_state(method::AbstractSecant, fs, x::Number)
    x1 = float(x)
    x0 = _default_secant_step(x1)
    init_state(method, fs, (x0, x1))
end

function init_state(method::AbstractSecant, fs, x::Union{Tuple, Vector})
    x0, x1 = promote(float(x[1]), float(x[2]))
    fx0, fx1 = fs(x0), fs(x1)
    state = UnivariateZeroState(x1, x0, zero(x1)/zero(x1)*oneunit(x1), eltype(x1)[],
                                fx1, fx0, fx1, eltype(fx1)[],
                                0, 2,
                                false, false, false, false, "")

    state
end

function init_state!(state::UnivariateZeroState{T, S}, method::AbstractSecant, fs, x::Number) where {T, S}
    x1::T = float(x)
    x0::T = _default_secant_step(x1)
    init_state!(state, method, fs, (x0, x1))
end

function init_state!(state::UnivariateZeroState{T, S}, ::AbstractSecant, f, x::Union{Tuple, Vector}) where {T, S}
    x0,x1 = promote(float.(x)...)
    fx0, fx1 = promote(f(x0), f(x1))
    init_state!(state, x1, x0, T[], fx1, fx0, S[])
    state.fnevals = 2
    nothing
end

##################################################

## Order0 and Secant are related
"""
    Order0()

The `Order0` method is engineered to be a more robust, though possibly
slower, alternative to the other derivative-free root-finding
methods. The implementation roughly follows the algorithm described in
*Personal Calculator Has Key to Solve Any Equation f(x) = 0*, the
SOLVE button from the
[HP-34C](http://www.hpl.hp.com/hpjournal/pdfs/IssuePDFs/1979-12.pdf).
The basic idea is to use a secant step. If along the way a bracket is
found, switch to bisection, using `AlefeldPotraShi`.  If the secant
step fails to decrease the function value, a quadratic step is used up
to 4 times.

This is not really 0-order: the secant method has order
1.6...[Wikipedia](https://en.wikipedia.org/wiki/Secant_method#Comparison_with_other_root-finding_methods_
and the the bracketing method has order
1.6180...[Wikipedia](http://www.ams.org/journals/mcom/1993-61-204/S0025-5718-1993-1192965-2/S0025-5718-1993-1192965-2.pdf)
so for reasonable starting points, this algorithm should be
superlinear, and relatively robust to non-reasonable starting points.

"""
struct Order0 <: AbstractSecant end

function find_zero(fs, x0, method::Order0;
                   tracks::AbstractTracks=NullTracks(),
                   verbose=false,
                   kwargs...)
    M = Order1()
    N = AlefeldPotraShi()

    find_zero(callable_function(fs), x0, M, N; tracks=tracks,verbose=verbose, kwargs...)
end

##################################################

## Secant
## https://en.wikipedia.org/wiki/Secant_method
"""
    Order1()
    Roots.Secant()

The `Order1()` method is an alias for `Secant`. It specifies the
[secant method](https://en.wikipedia.org/wiki/Secant_method).
This method keeps two values in its state, `x_n` and `x_n1`. The
updated point is the intersection point of x axis with the secant line
formed from the two points. The secant method uses 1 function
evaluation per step and has order `(1+sqrt(5))/2`.

The error, `eᵢ = xᵢ - α`, satisfies
`eᵢ₊₂ = f[xᵢ₊₁,xᵢ,α] / f[xᵢ₊₁,xᵢ] * (xᵢ₊₁-α) * (xᵢ - α)`.

"""
struct Secant <: AbstractSecant end
const Order1 = Secant

function update_state(method::Order1, fs, o::UnivariateZeroState{T,S}, options) where {T, S}

    xn0, xn1 = o.xn0, o.xn1
    fxn0, fxn1 = o.fxn0, o.fxn1

    delta = fxn1 * (xn1 - xn0) / (fxn1 - fxn0)

    if isinf(delta) || isnan(delta)
        o.stopped = true
        o.message = "Increment `Δx` has issues. "
         return
     end

    o.xn0 = xn1
    o.xn1 -= delta
    o.fxn0 = fxn1
    o.fxn1 = (tmp::S = fs(o.xn1))
    incfn(o)

    nothing

end

##################################################

"""
    Roots.Order1B()
    Roots.King()

A superlinear (order 1.6...) modification of the secant method for multiple roots.
Presented in A SECANT METHOD FOR MULTIPLE ROOTS, by RICHARD F. KING, BIT 17 (1977), 321-328

The basic idea is similar to Schroder's method: apply the secant method
to  `f/f'`. However, this uses `f' ~ fp = (fx - f(x-fx))/fx` (a Steffensen step). In
this implementation, `Order1B`, when `fx` is too big, a single secant step of `f`
is used.

The *asymptotic* error, `eᵢ = xᵢ - α`, is given by
`eᵢ₊₂ = 1/2⋅G''/G'⋅ eᵢ⋅eᵢ₊₁ + (1/6⋅G'''/G' - (1/2⋅G''/G'))^2⋅eᵢ⋅eᵢ₊₁⋅(eᵢ+eᵢ₊₁)`.

"""
struct King <: AbstractSecant end
struct Order1B <: AbstractSecant end

function update_state(method::Order1B, fs, o::UnivariateZeroState{T,S}, options)  where {T, S}
    update_state_guarded(method, Secant(), King(), fs, o, options)
end

function update_state(method::King, fs,
                      o::UnivariateZeroState{T,S}, options)  where {T, S}

    x0, x1 = o.xn0, o.xn1
    fx0, fx1 = o.fxn0, o.fxn1

    # G(x,f0,f_1) = -fx^2/(f_1 - f0)
    f0 = fx1
    f_1::S = fs(x1 - f0*oneunit(T)/oneunit(S))
    incfn(o, 1)

    G1 = -f0^2 / (f_1 - f0)

    if isempty(o.fm)
        f0 = fx0
        tmp::S = fs(x0 - f0*oneunit(T)/oneunit(S))
        f_1 = tmp
        incfn(o,1)
        G0 = -f0^2 / (f_1 - f0)
    else
        G0 = first(o.fm)
    end

    m = (x1-x0)/(G1-G0) # approximate value of `m`, the multiplicity
    if abs(m) <= 1e-2 * oneunit(m)
        #@info "small m estimate, stopping"
        o.stopped  = true
        o.message = "Estimate for multiplicity has issues. "
        return
    end

    delta = G1 * (x1 - x0) / (G1 - G0)
    empty!(o.fm); push!(o.fm, G1)

    if isissue(delta)
        o.stopped  = true
        o.message = "Increment `Δx` has issues. "
        return
    end

    o.xn0, o.fxn0 = o.xn1, o.fxn1
    o.xn1 -= delta
    o.fxn1 = fs(o.xn1)
    incfn(o)

    nothing
end

### Steffensen
## https://en.wikipedia.org/wiki/Steffensen's_method#Simple_description

"""
    Order2()
    Roots.Steffensen()

The quadratically converging
[Steffensen](https://en.wikipedia.org/wiki/Steffensen's_method#Simple_description)
method is used for the derivative-free `Order2()` algorithm. Unlike
the quadratically converging Newton's method, no derivative is
necessary, though like Newton's method, two function calls per step
are. Steffensen's algorithm is more sensitive than Newton's method to
poor initial guesses when `f(x)` is large, due to how `f'(x)` is
approximated. This algorithm, `Order2`, replaces a Steffensen step with a secant
step when `f(x)` is large.

The error, `eᵢ - α`, satisfies
`eᵢ₊₁ = f[xᵢ, xᵢ+fᵢ, α] / f[xᵢ,xᵢ+fᵢ] ⋅ (1 - f[xᵢ,α] ⋅ eᵢ²`
"""
struct Order2 <: AbstractSecant end
struct Steffensen <: AbstractSecant end

function update_state(method::Order2, fs, o::UnivariateZeroState{T,S}, options)  where {T, S}
     update_state_guarded(method, Secant(), Steffensen(), fs, o, options)
end

function update_state(method::Steffensen, fs, o::UnivariateZeroState{T,S}, options)  where {T, S}

    x0, x1 = o.xn0, o.xn1
    fx0, fx1 = o.fxn0, o.fxn1

    sgn = sign((fx1 - fx0) / (x1 - x0))
    x2 = x1 - sgn * fx1 / oneunit(S) * oneunit(T)

    f0 = fx1
    f1::S = fs(x2)
    incfn(o, 1)

    delta = -sgn * f0 * f0 / (f1 - f0) * oneunit(T) / oneunit(S)

    if isissue(delta)
        o.stopped  = true
        o.message = "Increment `Δx` has issues. "
        return
    end

    o.xn0, o.fxn0 = o.xn1, o.fxn1
    o.xn1 -= delta
    o.fxn1 = fs(o.xn1)
    incfn(o)

    nothing
end

##################################################

### Order2B() Esser method
"""
    Roots.Order2B()
    Roots.Esser()

Esser's method. This is a quadratically convergent method that, like
Schroder's method, does not depend on the multiplicity of the
zero. Schroder's method has update step `x - r2/(r2-r1) * r1`, where `ri =
f^(i-1)/f^(i)`. Esser approximates `f' ~ f[x-h, x+h], f'' ~
f[x-h,x,x+h]`, where `h = fx`, as with Steffensen's method, Requiring 3
function calls per step. The implementation `Order2B` uses a secant
step when `|fx|` is considered too large.

Esser, H. Computing (1975) 14: 367. https://doi.org/10.1007/BF02253547
Eine stets quadratisch konvergente Modifikation des Steffensen-Verfahrens

Example
```
f(x) = cos(x) - x
g(x) = f(x)^2
x0 = pi/4
find_zero(f, x0, Order2(), verbose=true)        #  3 steps / 7 function calls
find_zero(f, x0, Roots.Order2B(), verbose=true) #  4 / 9
find_zero(g, x0, Order2(), verbose=true)        #  22 / 45
find_zero(g, x0, Roots.Order2B(), verbose=true) #  4 / 10
```
"""
struct Esser <: AbstractSecant end
struct Order2B <: AbstractSecant end

function update_state(method::Order2B, fs, o::UnivariateZeroState{T,S}, options)  where {T, S}
     update_state_guarded(method, Secant(), Esser(), fs, o, options)
end

function update_state(method::Esser, fs,
                      o::UnivariateZeroState{T,S}, options)  where {T, S}

    x1, fx1 = o.xn1, o.fxn1

    f0 = fx1

    f1  = fs(x1 + f0 * oneunit(T) / oneunit(S))
    f_1 = fs(x1 - f0 * oneunit(T) / oneunit(S))
    incfn(o, 2)

    # h = f0
    # r1 = f/f' ~ f/f[x+h,x-h]
    # r2 = f'/f'' ~ f[x+h, x-h]/f[x-h,x,x+h]
    r1 = f0 * 2*f0 / (f1 - f_1) * oneunit(T) / oneunit(S)
    r2 = (f1 - f_1)/(f1 - 2*f0 + f_1) * f0/2 * oneunit(T) / oneunit(S)

    k = r2/(r2-r1)  # ~ m

    if abs(k) <= 1e-2 * oneunit(k)
        #@info "estimate for m is *too* small"
        o.stopped  = true
        o.message = "Estimate for multiplicity had issues. "
        return
    end

    delta = k * r1

    if isissue(delta)
        o.stopped  = true
        o.message = "Increment `Δx` has issues. "
        return
    end

    o.xn0, o.fxn0 = o.xn1, o.fxn1
    o.xn1 -= delta

    o.fxn1 = fs(o.xn1)
    incfn(o)

    nothing
end

##################################################

"""
    Order5()

Implements an order 5 algorithm from *A New Fifth Order Derivative
Free Newton-Type Method for Solving Nonlinear Equations* by Manoj
Kumar, Akhilesh Kumar Singh, and Akanksha, Appl. Math. Inf. Sci. 9,
No. 3, 1507-1513 (2015), DOI: [10.12785/amis/090346](https://doi.org/10.12785/amis/090346). Four function
calls per step are needed.

The error, `eᵢ = xᵢ - α`, satisfies
`eᵢ₊₁ = K₁ ⋅ K₅ ⋅ M ⋅ eᵢ⁵ + O(eᵢ⁶)`

"""
struct Order5 <: AbstractSecant end
struct KumarSinghAkanksha <: AbstractSecant end

function update_state(M::Union{Order5, KumarSinghAkanksha}, fs, o::UnivariateZeroState{T,S},
                      options) where {T, S}

    xn = o.xn1
    fxn = o.fxn1

    wn::T = steff_step(M, o.xn1, o.fxn1)

    fwn::S = fs(wn)
    incfn(o)

    fp, issue = _fbracket(o.xn1, wn, o.fxn1, fwn)
    if issue
        o.xn0, o.xn1 = o.xn1, wn
        o.fxn0, o.fxn1 = o.fxn1, fwn
        o.message = "Issue with divided difference f[xn, wn]. "
        o.stopped  = true
        return
    end

    yn::T = o.xn1 - o.fxn1 / fp
    fyn::S = fs(yn)
    incfn(o)

    zn::T = xn - (fxn + fyn) / fp
    fzn::S = fs(zn)
    incfn(o)

    fp, issue = _fbracket_ratio(yn, o.xn1, wn, fyn, o.fxn1, fwn)
    if issue
        o.xn0, o.xn1 = o.xn1, yn
        o.fxn0, o.fxn1 = o.fxn1, fyn
        o.message = "Issue with f[xn,yn]*f[yn,wn] / f[xn, wn]"
        o.stopped = true
        return
    end

    o.xn0 = o.xn1
    o.fxn0 = o.fxn1
    o.xn1 = zn  - fzn  / fp
    o.fxn1 = fs(o.xn1)
    incfn(o)

    nothing
end

## If we have a derivative, we have this. (Deprecate?)
function update_state(method::Order5, fs::Union{FirstDerivative,SecondDerivative},
                      o::UnivariateZeroState{T,S}, options)  where {T, S}

    xn, fxn = o.xn1, o.fxn1
    a::T, b::S = fΔx(fs, xn)
    fpxn = a/b
    incfn(o)

    if isissue(fpxn)
        o.stopped  = true
        return
    end

    yn::T = xn - fxn / fpxn
    fyn::S, Δyn::T = fΔx(fs, yn)
    fpyn = fyn / Δyn
    incfn(o, 2)

    if isissue(fpyn)
        o.xn0, o.xn1 = xn, yn
        o.fxn0, o.fxn1 = fxn, fyn
        o.stopped  = true
        return
    end

    zn::T = xn  - (fxn + fyn) / fpxn
    fzn::S = fs(zn)
    incfn(o, 1)

    xn1 = zn - fzn / fpyn
    fxn1 = fs(xn1)
    incfn(o, 1)

    o.xn0, o.xn1 = xn, xn1
    o.fxn0, o.fxn1 = fxn, fxn1

    nothing
end

##################################################

## cf also: https://doi.org/10.1515/tmj-2017-0049
"""
    Order8()

Implements an eighth-order algorithm from *New Eighth-Order
Derivative-Free Methods for Solving Nonlinear Equations* by Rajinder
Thukral, International Journal of Mathematics and Mathematical
Sciences Volume 2012 (2012), Article ID 493456, 12 pages DOI:
[10.1155/2012/493456](https://doi.org/10.1155/2012/493456). Four function calls per step are required.

The error, `eᵢ = xᵢ - α`, is expressed as `eᵢ₊₁ = K ⋅ eᵢ⁸` in
(2.25) of the paper for an explicit `K`.

"""
struct Order8 <: AbstractSecant end
struct Thukral8 <: AbstractSecant end

function update_state(M::Union{Thukral8, Order8}, fs, o::UnivariateZeroState{T,S}, options) where {T, S}

    xn = o.xn1
    fxn = o.fxn1

    wn::T = steff_step(M, xn, fxn)
    fwn::S = fs(wn)
    incfn(o)

    if isissue(fwn)
        o.xn0,o.xn1 = xn, wn
        o.fxn0,o.fxn1 = fxn, fwn
        o.stopped = true
        o.message = "issue with Steffensen step fwn"
        return
    end

    fp, issue = _fbracket(xn, wn, fxn, fwn)

    if issue
        o.stopped = true
        o.message = "issue with divided difference f[xn, wn]. "
        return
    end

    yn::T = xn - fxn / fp
    fyn::S = fs(yn)
    incfn(o)

    fp, issue = _fbracket(yn, xn, fyn, fxn)
    if issue #fp
        o.xn0,o.xn1 = xn, yn
        o.fxn0,o.fxn1 = fxn, fyn
        o.stopped = true
        o.message = "issue with divided difference f[xn, yn]. "
        return
    end

    phi = (1 + fyn / fwn)           # pick one of options
    zn =  yn - phi * fyn / fp
    fzn::S = fs(zn)
    incfn(o)

    fp, issue =  _fbracket_diff(xn, yn, zn, fxn, fyn, fzn)
    if issue
        o.xn0,o.xn1 = xn, zn
        o.fxn0,o.fxn1 = fxn, fzn
        o.message = "issue with divided difference  f[y,z] - f[x,y] + f[x,z]. "
        o.stopped = true
        return
    end

    w = 1 / (1 - fzn/fwn)

    xi = (1 - 2fyn*fyn*fyn / (fwn * fwn * fxn))

    xn1 = zn - w * xi * fzn / fp
    fxn1::S = fs(xn1)
    incfn(o)

    o.xn0,o.xn1 = xn, xn1
    o.fxn0,o.fxn1 = fxn, fxn1

    nothing
end

##################################################

"""
    Order16()

Implements the order 16 algorithm from
*New Sixteenth-Order Derivative-Free Methods for Solving Nonlinear Equations*
by R. Thukral,
American Journal of Computational and Applied Mathematics
p-ISSN: 2165-8935;    e-ISSN: 2165-8943; 2012;  2(3): 112-118
doi: 10.5923/j.ajcam.20120203.08.

Five function calls per step are required. Though rapidly converging,
this method generally isn't faster (fewer function calls/steps) over
other methods when using `Float64` values, but may be useful for
solving over `BigFloat`.

The error, `eᵢ = xᵢ - α`, is expressed as `eᵢ₊₁ = K⋅eᵢ¹⁶` for an explicit `K`
in equation (50) of the paper.

"""
struct Order16 <: AbstractSecant end
struct Thukral16 <: AbstractSecant end

function update_state(M::Union{Thukral16, Order16}, fs, o::UnivariateZeroState{T,S}, options) where {T, S}
    xn = o.xn1
    fxn = o.fxn1

    wn::T = steff_step(M, xn, fxn)
    fwn::S = fs(wn)
    incfn(o)

    fp, issue = _fbracket(xn, wn, fxn, fwn)

    if issue
        o.xn0, o.xn1 = xn, wn
        o.fxn0, o.fxn1 = fxn, fwn
        o.message = "issue with f[xn,wn]"
        o.stopped = true
        return
    end

    yn::T = xn - fxn / fp
    fyn::S = fs(yn)
    incfn(o)

    fp, issue = _fbracket_ratio(yn, xn, wn, fyn, fxn, fwn)
    if issue
        o.xn0, o.xn1 = xn, yn
        o.fxn0, o.fxn1 = fxn, fyn
        o.message = "issue with f[xn,yn]*f[yn,wn]/f[xn,wn]"
        o.stopped = true
        return
    end

    zn = yn - fyn / fp
    fzn::S = fs(zn)
    incfn(o)

    fp, issue = _fbracket_diff(xn, yn, zn, fxn, fyn, fzn)
    u2, u3, u4 = fzn/fwn, fyn/fxn, fyn/fwn

    eta = 1 / (1 + 2*u3*u4^2) / (1 - u2)
    if issue
        o.xn0, o.xn1 = xn, zn
        o.fxn0, o.fxn1 = fxn, fzn
        o.stopped = true
        o.message = "Approximate derivative failed"
        return
    end

    an = zn - eta * fzn / fp
    fan::S = fs(an)
    incfn(o)

    fp, issue = _fbracket_ratio(an, yn, zn, fan, fyn, fzn)
    if issue
        o.xn0, o.xn1 = xn, an
        o.fxn0, o.fxn1 = fxn, fan
        o.stopped = true
        o.message = "Approximate derivative failed"
        return
    end

    u1, u5, u6 = fzn/fxn, fan/fxn, fan/fwn

    fp1, issue = _fbracket(xn,yn, fxn, fyn)

    sigma =  1 + u1*u2 - u1*u3*u4^2 + u5 + u6 + u1^2*u4 +
    u2^2*u3 + 3*u1*u4^2*(u3^2 - u4^2)/(fp1/oneunit(fp1))

    xn1 = an - sigma * fan / fp
    fxn1::S = fs(xn1)
    incfn(o)

    o.xn0, o.xn1 = xn, xn1
    o.fxn0, o.fxn1 = fxn, fxn1

    nothing
end

##################################################
## some means of guarding against large fx when taking a secant step
## TODO: rework this
function steff_step(M::Union{Order5, Order8, Order16}, x::S, fx::T) where {S, T}

    xbar, fxbar = real(x/oneunit(x)), fx/oneunit(fx)
    thresh =  max(1, abs(xbar)) * sqrt(eps(one(xbar))) #^(1/2) # max(1, sqrt(abs(x/fx))) * 1e-6

    out = abs(fxbar) <= thresh ? fxbar  : sign(fx) * thresh
    x + out * oneunit(x)

end


export Interval

struct Interval{T<:Real}
    a::T
    b::T
end

Base.minimum(it::Interval) = it.a
Base.maximum(it::Interval) = it.b
Base.in(x, it::Interval) = (x <= maximum(it)) && (x >= minimum(it))
Base.isempty(it::Interval) = ~(maximum(it) > minimum(it))

function Base.intersect(A::Interval, B::Interval)
    if isempty(A) || isempty(B)
        return Interval(0.0, 0.0)
    end
    Amin = minimum(A)
    Amax = maximum(A)
    Bmin = minimum(B)
    Bmax = maximum(B)
    if (Bmin > Amax) || (Amin > Bmax)
        return Interval(0.0, 0.0)
    else
        Cmin = max(Amin, Bmin)
        Cmax = min(Amax, Bmax)
        return Interval(Cmin, Cmax)
    end
end