# Routines for taking steps using either Euler or SDC

using FastGaussQuadrature

function step_euler(dt, sol, solver, ubc1, ubc2)
    alpha = solver.alpha
    F1, F2 = ugradu(sol)
    RHS1 = alpha^2*(sol.U1 - dt*F1)
    RHS2 = alpha^2*(sol.U2 - dt*F2)
    newsol = solver.solve(RHS1, RHS2, ubc1, ubc2)
    return newsol
end

function step_bdf(dt, prevsol, sol, solver, ubc1, ubc2)
    alpha = solver.alpha
    F1, F2 = ugradu(sol)
    prevF1, prevF2 = ugradu(prevsol)
    RHS1 = 1/3*alpha^2*(4*sol.U1 - prevsol.U1 - 2*dt*(2*F1 - prevF1))
    RHS2 = 1/3*alpha^2*(4*sol.U2 - prevsol.U2 - 2*dt*(2*F2 - prevF2))
    newsol = solver.solve(RHS1, RHS2, ubc1, ubc2)
    return newsol
end

function sdc_substep(dt, sdc_order)
    return sdc_order > 1 ? dt / (sdc_order-1) : dt
end

function sdc_integration(M)
    # Setup substeps
    x = collect(linspace(-1, 1, M))
    # Vandermonde matrix (transposed)
    At = ones(M,M)
    for m=2:M
        @. At[m, :] = At[m-1, :]*x
    end
    # Use Gauss-Legendre for quadrature
    xleg, wleg = gausslegendre(M)
    # Compute substep integration weights
    I = zeros(M,M-1)
    for m=1:M-1
        a, b = x[m], x[m+1]
        wint = wleg*(b-a)/2
        xint = (xleg+1)*(b-a)/2 + a
        xint_i = ones(M)
        p = zeros(M)
        p[1] = sum(wint)
        for i=2:M
            xint_i .*= xint
            p[i] = sum(wint .* xint_i)
        end
        I[:,m] = At\p
    end
    return I/2
end

function step_sdc(dt, order, sol, solver, ubc1, ubc2)
    # TODO: Need to implement BC time dependence
    alpha = solver.alpha
    Re = alpha^2*dt
    P = order # order-1 can also work
    K = order
    if order>1
        W = sdc_integration(P)
    else
        # Fallback to Euler
        P = 2
        K = 1
    end
    # Set initial conditions for all passes
    solutions = Array{Solution}(P, K)
    for k=1:K
        solutions[1, k] = sol
    end
    # First pass (k=0)
    # Eulers steps
    for m=1:P-1
        init = solutions[m, 1]
        F1init, F2init = ugradu(init)
        RHS1 = alpha^2*init.U1 - Re*F1init
        RHS2 = alpha^2*init.U2 - Re*F2init
        solutions[m+1, 1] = solver.solve(RHS1, RHS2, ubc1, ubc2)
    end

    # Subsequent passes
    for k=1:K-1
        for m=1:P-1
            # Load historic stages that we need
            init = solutions[m,   k+1] # Substep initial condition
            pred = solutions[m+1, k  ] # Prediction
            hist = solutions[m  , k  ] # Prediction history
            # Post-process stages
            F1init, F2init = ugradu(init)
            #
            L1pred = alpha^2*pred.U1 - pred.RHS1
            L2pred = alpha^2*pred.U2 - pred.RHS2
            #
            F1hist, F2hist = ugradu(hist)
            # Integrate
            s = size(init.U1)
            I1, I2 = zeros(s), zeros(s)
            L = dt*(order-1)
            w = W[:, m]*L
            for n=1:P
                tmp = solutions[n, k]
                F1tmp, F2tmp = ugradu(tmp)
                L1tmp = alpha^2*tmp.U1 - tmp.RHS1
                L2tmp = alpha^2*tmp.U2 - tmp.RHS2
                I1 += w[n]*( L1tmp - Re*F1tmp )
                I2 += w[n]*( L2tmp - Re*F2tmp )
            end

            # f(t) = 1 + t
            # @show t1 = quadgk(f, dt*(m-1), dt*m)[1]
            # @show t2 = sum(w .* f(0:dt:dt*(P-1)))
            # @show t1-t2
            # @assert abs(t1-t2)/abs(t1) < 1e-14

            # Create new RHS
            RHS1 = alpha^2*init.U1 - Re*(F1init-F1hist) - L1pred + I1/dt
            RHS2 = alpha^2*init.U2 - Re*(F2init-F2hist) - L2pred + I2/dt
            #
            solutions[m+1, k+1] = solver.solve(RHS1, RHS2, ubc1, ubc2)
        end
    end

    if K>1
        # Output error est (UNTESTED)
        maxnorm(v) = maximum(abs.(v))
        l2norm(v) = sqrt(sum(vec(v.^2)))
        corr1 = solutions[P, K].U1 - solutions[P, K-1].U1
        corr2 = solutions[P, K].U2 - solutions[P, K-1].U2
        sol = solutions[P,K]
        info("U1 SDC err est: l2=", l2norm(corr1)/l2norm(sol.U1),", max=", maxnorm(corr1)/maxnorm(sol.U1))
        info("U2 SDC err est: l2=", l2norm(corr2)/l2norm(sol.U2),", max=", maxnorm(corr2)/maxnorm(sol.U2))
    end
    return solutions[P,K]

    # Original order 2 code
    # # First take Euler step
    # F1, F2 = ugradu(sol)
    # RHS1 = alpha^2*U1 - Re*F1
    # RHS2 = alpha^2*U2 - Re*F2
    # prov = solver.solve(RHS1, RHS2, ubc1, ubc2)
    # # Post process Euler step
    # F1prov, F2prov = ugradu(prov)
    # L1prov = alpha^2*prov.U1 - prov.RHS1
    # L2prov = alpha^2*prov.U2 - prov.RHS2
    # # Post process initial conds
    # L1 = alpha^2*sol.U1 - sol.RHS1
    # L2 = alpha^2*sol.U2 - sol.RHS2
    # # Integrate
    # I1 = (L1-Re*F1 + L1prov-Re*F1prov)/2
    # I2 = (L2-Re*F2 + L2prov-Re*F2prov)/2
    # # Take correction step
    # RHS1 = Re/dt*U1 - L1prov + Re*I1
    # RHS2 = Re/dt*U2 - L2prov + Re*I2
    # newsol = solver.solve(RHS1, RHS2, ubc1, ubc2)
    # return newsol
end


using Polyhedra

"""
Test `center_of_mass` in a case where the center of mass is different from the
centroid (namely, a square pyramid).
"""
function comsquarepyramidtest(lib::Polyhedra.Library)
    poly = polyhedron(convexhull([0, 0, 0],
                                 [1, 0, 0],
                                 [0, 1, 0],
                                 [1, 1, 0],
                                 [1/2, 1/2, 1]),
                      lib)
    # Centroid is [1/2, 1/2, 1/5].
    # Center of mass is [1/2, 1/2, 1/4].
    @test center_of_mass(poly) ≈ [1/2, 1/2, 1/4]
end


# The content in this file is for test purpose only
"""
    cpvagk(f, t, a, b, tol=256*eps())

Calculate the Cauchy principle value integration of the form ``𝒫∫_a^b f(x)/(x-t) dx``. The algorithm is adapted from [P. Keller, 02.01.2015](https://www.sciencedirect.com/science/article/pii/S0377042715004422)
"""
function cpvagk(f, t, a, b, tol=256*eps())
    #   Adapted from P. Keller, 02.01.2015

    #   CPVAGK(F,T,A,B,TOL) = CPVInt[a,b]F(x)/(x-T)dx.
    #
    #   Based on QUADGK - the built-in Matlab adaptive integrator.
    #
    #   Q = CPVAGK(F,T,A,B,TOL) approximates the above integral
    #   and tries to make the absolute error less than TOL.
    #   The default value of TOL is 256*eps.
    #
    #   [Q,ERR] = CPVAGK(F,T,A,B,TOL) also provides the estimate of
    #   the absolute error. We suggest to always check the value of ERR.
    #   In case a QUADGK warning message appear, the approximate error bound
    #   provided inside the warning message should be ingnored!
    #   Consider only the value of ERR as error estimate.
    #
    #   Test (research) version.

    #   P. Keller, 02.01.2015

    if b<=t || t <= a
        error("CPVAGK: Invalid arguments")
    end
    s = min(1.0, t-a, b-t); #computing the split point(s)
    a = float(a)
    b = float(b)
    s = float(s)
    t = float(t)
    epsab = epsf = eps()
    ft = f(t)
    abt = 0.25*(abs(a) + abs(b) + 2*abs(t))
    tcond = 1 + 0.5*(max(1,abt)-1)^2/max(1,abt);
    #Estimating possible accuracy loss...
    # a) related to rounding t and interval endpoints:
    v = [a+2*epsf*abs(a), b-2*epsf*abs(b)];
    fd = max(max(1.0, abs(t)), abs.(v)...).*abs.(f.(v))./[t-a,b-t];
    erra = 0.75*epsab * sum(fd);
    # b) related to the magnitude of the 1st and 2nd derivative:
    sts = 255*s/256  # divided differences five steps...
    sth = min(sts, min(2,(b-a))./[82,70,32,22]...)
    ste = min(sts, exp(3/8*log(max(1,abs(t))*epsf)))
    dfs = max(abs(f(t+sts)-ft), abs(f(t-sts)-ft)) / sts  # <|f'[t-s,t+s]|
    dfh = max(abs(f(t+sth)-ft), abs(f(t-sth)-ft))./ sth./ [3/2, 7/4, 2, 3]
    df1 = abs(f(t+ste^2)-f(t-ste^2))/ste^2;  # ~~ 2*|f'(t)|
    df2 = sqrt(abs(f(t+ste)-2*ft+f(t-ste)))/ste;  # ~~ |f"(t)|^(1/2)
    errb1 = 17*tcond*epsf*(max(df1,dfh...) + abs(ft));
    errb2 = 10*tcond*epsf*df2;
    # c) related to the distance from interval endpoints:
    errc = 0.75*epsf*abt*abs(ft)/min(b-t,t-a)
    # Setting (safe) error tolerances, absolute and relative...
    atol = max(tol/2, 8*dfs*epsf, erra, errb1, errb2, errc);
    rtol = max(100*epsf, errb2, errc/max(sqrt(epsf), abs(ft)));
    # Setting parameters for quadgk()
    params = (atol=atol, rtol=rtol, maxevals=163840)
    # Computing the integrals...
    q1 = 0
    e1 = 0
    q2 = 0
    e2 = 0
    if s < t-a
        q1, e1 = quadgk((x)->(f(x)-ft)./(x-t), a, t-s; params...)
    end
    if s < b-t
        q2, e2 = quadgk((x)->(f(x)-ft)./(x-t), t+s, b; params...)
    end
    q3, e3 = quadgk((x)->(f(t+x)-f(t-x))./x, 0, s; params...)
    # Computing the final result and error estimation...
    q = q1 + q2 + q3 + ft*log((b-t)/(t-a))
    err = e1 + e2 + e3 + max(erra,errc) + errb1 + errb2 + 10*tcond*epsf*max(abs.([q1,q2,q3,q])...)
    q, err
end

"""
$(SIGNATURES)

Calculate the Lamb shift of spectrum `γ`. `atol` is the absolute tolerance for Cauchy principal value integral.
"""
function lambshift_cpvagk(w, γ; atol = 1e-7)
    g(x) = γ(x) / (x - w)
    cpv, cperr = cpvagk(γ, w, w - 1.0, w + 1.0)
    negv, negerr = quadgk(g, -Inf, w - 1.0)
    posv, poserr = quadgk(g, w + 1.0, Inf)
    v = cpv + negv + posv
    err = cperr + negerr + poserr
    if (err > atol) || (isnan(err))
        @warn "Absolute error of integration is larger than the tolerance."
    end
    -v / 2 / pi
end


# Inverse Weibull Distribution
# for information see e.g. M.-S. Khan, Theoretical Analysis of Inverse Weibull Distribution, WSEAS Transactions on Mathematics 2008

immutable InverseWeibull <: ContinuousUnivariateDistribution
    shape::Float64
    scale::Float64
    function InverseWeibull(sh::Real, sc::Real)
    	zero(sh) < sh && zero(sc) < sc || error("Both shape and scale must be positive")
    	new(float64(sh), float64(sc))
    end
end

InverseWeibull(sh::Real) = InverseWeibull(sh, 1.0)

## Support
@continuous_distr_support InverseWeibull 0.0 Inf

## Properties
mean(d::InverseWeibull) = d.shape > 1.0 ? d.scale * gamma(1.0 - 1.0 / d.shape) : Inf
median(d::InverseWeibull) = d.scale * log(2)^(-1.0 / d.shape)

mode(d::InverseWeibull) = (ik = -1.0/d.shape; d.scale * (1.0-ik)^ik)

var(d::InverseWeibull) = d.shape > 2.0 ? d.scale^2 * gamma(1.0 - 2.0 / d.shape) - mean(d)^2 : NaN

function skewness(d::InverseWeibull)
    d.shape <= 3.0 && return(NaN)
    tmp_mean = mean(d)
    tmp_var = var(d)
    tmp = gamma(1.0 - 3.0 / d.shape) * d.scale^3
    tmp -= 3.0 * tmp_mean * tmp_var
    tmp -= tmp_mean^3
    return tmp / tmp_var / sqrt(tmp_var)
end

function kurtosis(d::InverseWeibull)
    d.shape <= 4.0 && return(NaN)
    λ, k = d.scale, d.shape
    μ = mean(d)
    σ = std(d)
    γ = skewness(d)
    den = λ^4 * gamma(1.0 - 4.0 / k) -
          4.0 * γ * σ^3 * μ -
          6.0 * μ^2 * σ^2 - μ^4
    num = σ^4
    return den / num - 3.0
end

function entropy(d::InverseWeibull)
    λ, k = d.scale, d.shape
    return (k + 1.0) * (log(λ) - digamma(1.0)/k) - log(λ * k) + 1.0
end

## Functions
function pdf(d::InverseWeibull, x::Real)
    x < zero(x) && return(0.0)
    a = d.scale/x
    d.shape/d.scale * a^(d.shape+1.0) * exp(-a^d.shape)
end
function logpdf(d::InverseWeibull, x::Real)
    x < zero(x) && return(-Inf)
    a = d.scale/x
    log(d.shape/d.scale) + (d.shape+1.0)*log(a) - a^d.shape
end

cdf(d::InverseWeibull, x::Real) = x <= zero(x) ? 0.0 : exp(-((d.scale / x)^d.shape))
ccdf(d::InverseWeibull, x::Real) = x <= zero(x) ? 1.0 : -expm1(-((d.scale / x)^d.shape))
logcdf(d::InverseWeibull, x::Real) = x <= zero(x) ? -Inf : -(d.scale / x)^d.shape
logccdf(d::InverseWeibull, x::Real) = x <= zero(x) ? 0.0 :  log1mexp(-((d.scale / x)^d.shape))

quantile(d::InverseWeibull, p::Real) = @checkquantile p d.scale*(-log(p))^(-1/d.shape)
cquantile(d::InverseWeibull, p::Real) = @checkquantile p d.scale*(-log1p(-p))^(-1/d.shape)
invlogcdf(d::InverseWeibull, lp::Real) = lp > zero(lp) ? NaN : d.scale*(-lp)^(-1/d.shape)
invlogccdf(d::InverseWeibull, lp::Real) = lp > zero(lp) ? NaN : d.scale*(-log1mexp(lp))^(-1/d.shape)

function gradloglik(d::InverseWeibull, x::Float64)
  insupport(InverseWeibull, x) ? -(d.shape + 1.0) / x + d.shape * (d.scale^d.shape) * x^(-d.shape - 1.0)  : 0.0
end

## Sampling
rand(d::InverseWeibull) = d.scale*Base.Random.randmtzig_exprnd()^(-1/d.shape)


###########################################################
# Polynomial eigenvalue problem - PEP
#

"""
    struct PEP <: AbstractSPMF
    function PEP(AA::Vector{AbstractMatrix})

The type `PEP` defines a polynomial eigenvalue
 problem via its monomial coefficients.
A polynomial eigenvalue problem (PEP) is defined by the sum the
```math
Σ_i A_i λ^i,
```
where ``i = 0,1,2,``, and  all of the matrices are of size ``n×n``.
The vector `AA` contains ``A_1,...``.

# Example

```julia-repl
julia> A0=[1.0 3; 4 5]; A1=A0.+one(2); A2=ones(2,2);
julia> pep=PEP([A0,A1,A2])
julia> compute_Mder(pep,3)-(A0+A1*3+A2*9)
2×2 Array{Float64,2}:
 0.0  0.0
 0.0  0.0
```

See also [`polyeig`](methods.md#NonlinearEigenproblems.NEPSolver.polyeig), [`companion`](methods.md#NonlinearEigenproblems.NEPSolver.companion), [`ChebPEP`](@ref), [`interpolate`](@ref).
"""
struct PEP <: AbstractSPMF{AbstractMatrix}
    n::Int
    A::Array   # Monomial coefficients of PEP
end

function PEP(AA::Array)
    n=size(AA[1],1)
    AA=reshape(AA,size(AA,1))
    return PEP(n,AA)
end

# Computes the sum ``Σ_i M_i V f_i(S)`` for a PEP
function compute_MM(nep::PEP,S,V)
    T=promote_type(promote_type(eltype(nep.A[1]),eltype(S)),eltype(V))
    local Z
    if (issparse(nep))
        Z=spzeros(T,size(V,1),size(V,2))
        Si=sparse(one(T)*I, size(S,1), size(S,1))
    else
        Z=zeros(T,size(V,1),size(V,2))
        Si=Matrix(one(T)*I, size(S,1), size(S,1))
    end
    for i=1:size(nep.A,1)
        Z+=nep.A[i]*V*Si;
        Si=Si*S;
    end
    return Z
end

# Compute the ith derivative of a PEP
function compute_Mder(nep::PEP,λ::Number,i::Integer=0)
    if (issparse(nep))
        Z=spzeros(eltype(nep.A[1]),size(nep,1),size(nep,1));
    else
        Z=zeros(eltype(nep.A[1]),size(nep,1),size(nep,1));
    end
    for j=(i+1):size(nep.A,1)
        # Derivatives of monimials
        Z+= nep.A[j]*(λ^(j-i-1)*factorial(j-1)/factorial(j-i-1))
    end
    return Z
end

#  Fetch the Av's, since they are not explicitly stored in PEPs
function get_Av(nep::PEP)
    return nep.A;
end
#  Fetch the Fv's, since they are not explicitly stored in PEPs
function get_fv(nep::PEP)
    fv=Vector{Function}(undef, size(nep.A,1))
    # Construct monomial functions
    for i=1:size(nep.A,1)
        if (i==1); # optimization for constant and linear term
            fv[1] = S -> one(S);
        elseif (i==2);
            fv[2]=S->S;
        else
            # we need to introuce j otherwise the function is not typestable
            j=i-1; # power of the coefficient
            fv[i]=S->S^j;
        end
    end
    return fv;
end

"""
    interpolate([T=ComplexF64,] nep::NEP, intpoints::Array)

Interpolates a NEP in the points `intpoints` and returns a [`PEP`](@ref), i.e., a polynomial eigenvalue problem in a monomial basis. See [`ChebPEP`](@ref) for Chebyshev interpolation. The optional argument `T` is the type in which the matrices of the PEP should be defined.

See also [`ChebPEP`](@ref).
"""
interpolate(nep::NEP, intpoints::Array) = interpolate(ComplexF64, nep, intpoints)
function interpolate(::Type{T}, nep::NEP, intpoints::Array) where {T<:Number}

    n = size(nep, 1)
    d = length(intpoints)

    V = zeros(T,d,d) #Vandermonde matrix
    pwr = ones(d,1)
    for i = 1:d
        V[:,i] = pwr
        pwr = pwr.*intpoints
    end

    if (issparse(nep)) #If Sparse, do elementwise interpolation
        b = Vector{SparseMatrixCSC{T}}(undef, d)
        AA = Vector{SparseMatrixCSC{T}}(undef, d)
        V = factorize(V) # Will be used multiple times, factorize

        for i=1:d
            b[i] = compute_Mder(nep, intpoints[i])
        end

        # OBS: The following lines and hence the following method assumes that Sparsity-structure is the same!
        nnz_AA = nnz(b[1])
        for i=1:d
            AA[i] = copy(b[1])
        end

        f = zeros(d,1)
        for i = 1:nnz_AA
            for j = 1:d
                f[j] = b[j].nzval[i]
            end
            a = \(V,f)
            for j = 1:d
                AA[j].nzval[i] = a[j]
            end
        end

    else # If dense, use Vandermonde
        b = zeros(T,n*d,n)
        AA = Vector{Matrix{T}}(undef,d)
        (L, U, p) = lu(V)

        LL = kron(L, SparseMatrixCSC(I,(n,n)))
        UU = kron(U, SparseMatrixCSC(I,(n,n)))

        for i = 1:d
            b[(1:n).+(i-1)*n,:] =  compute_Mder(nep,intpoints[p[i]])
        end

        A = \(UU, \(LL,b))

        for i = 1:d
            AA[i] = A[(1:n).+(i-1)*n,:]
        end
    end

    return PEP(AA)
end

include("types_cheb_pep.jl");


using ConservationLawsParticles
using SpecialFunctions, RecursiveArrayTools, DifferentialEquations, Plots

# model
V(t, x) = 0.0 # 2 + 0.5sin(2x) + 0.3sin(2sqrt(2)*x) + 0.2cos(2sqrt(7)*x)
W(t, x) = 0.0 # x^2/2 - abs(x)
W′(t, x) = 0.0 # x - sign(x)
mob(ρ) = 0.0 # max((ρ-1)^2 * (2ρ+1), 0)
smodel = DiffusiveSampledModel((V,), ((W′,),), (mob,), (Diffusion(1),))
imodel = DiffusiveIntegratedModel((V,), ((W,),), (mob,), (Diffusion(1),))

# initial condition
n = 100
x0 = ArrayPartition(vcat(
    range(-2, -1; length=n+1),
    range(0, 2; length=n+1)
    ))

boxsol(t,x;a,b) = (erf((x-a)/2sqrt(t)) - erf((x-b)/2sqrt(t))) / (2*(b-a))
refsol(t, x) = boxsol(t,x; a=-2,b=-1) * n/(2n+1) +
    boxsol(t,x; a=0,b=2) * n/(2n+1) +
    boxsol(t,x; a=-1,b=0) * 1/(2n-1)

# time span
tspan = (0., 2.0)

# ODE system for the particles
sprob = ODEProblem(velocities_diff!, x0, tspan, smodel)
# iprob = ODEProblem(velocities_diff!, 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:5:(2n+2); color=:blue)
savefig("diffusion2.png")

# plot an animation of the density
anim = @animate for t in range(tspan...; length=101)
    plot(title="Density", xlabel="position", ylabel="density",
        xlims=(-7,7), ylims=(0,1/sqrt(pi)))
    plot!(x -> refsol(t, x); color=:blue, label="analytic")
    densityplot!(ssol(t); color=:red, label="particles")
end
gif(anim, "diffusion2.gif")


@testset "Test utils" begin

using HDiscontinuousGalerkin
#Test affine maps
A,b = get_affine_map([[-1.0,-1.0],[1.0,-1.0],[-1.0,1.0]],[[0.0,0.0],[1.0,0.0],[0.0,1.0]])
@test A*[0.0,0.0]+b ≈ [0.5,0.5]

end


module BilinearCheck

using Random

import Random: rand

export randComplex,inner,isLeftLinear,isRightLinear,isHomogenousBi,
       isHomogenousSus,nonDegenFirst,nonDegenSecond,nonDegen,
       isSymmetric,isHermitian,isBilinear,isSusquilinear

function randComplex(max)
    r = rand(Int64,3,1) .% max
    i = rand(Int64,3,1) .% max
    return r .+ i.*im
end

inner(v,w) = (w'*v)[1]

function isLeftLinear(β,a,b,c)
    β(a+b,c) == β(a,c) + β(b,c)
end

function isRightLinear(β,a,b,c)
    β(c,a+b) == β(c,a) + β(c,b)
end

function isHomogenousBi(β,a,b,λ)
    (β(λ*a,b) == λ*β(a,b)) && (β(a,λ*b) == λ*β(a,b))
end

function isHomogenousSus(β,a,b,λ)
    (β(λ*a,b) == λ*β(a,b)) && (β(a,λ*b) == conj(λ)*β(a,b))
end

function nonDegenFirst(β,a,b)
    a = zeros(size(a))
    norm(β(a,b)) == 0
end

function nonDegenSecond(β,a,b)
    b = zeros(size(b))
    norm(β(a,b)) == 0
end

function nonDegen(β,a,b)
    return nonDegenFirst(β,a,b) && nonDegenSecond(β,a,b)
end

function isSymmetric(β,a,b)
    β(a,b) == β(b,a)
end

function isHermitian(β,a,b)
    β(a,b) == conj(β(b,a))
end

function isBilinear(β,a,b,c,λ)
    return isLeftLinear(β,a,b,c) && isRightLinear(β,a,b,c) && isHomogenousBi(β,a,b,λ)
end

function isSusquilinear(β,a,b,c,λ)
    return isLeftLinear(β,a,b,c) && isRightLinear(β,a,b,c) && isHomogenousSus(β,a,b,λ)
end

end #module BilinearCheck


function aggDiscreteReturns(returns)
        n = size(returns,1)
        p = 1.0
        @inbounds @avx for i ∈ 1:n
            p = p * (returns[i] + 1.0)
        end
        return(p-1.0)
end

function geo_mean(returns::DataFrame;
                  dateColumn="date")

        vars = names(returns)
        nVars = length(vars)
        vars = Symbol.(vars[vars.!=dateColumn])
        if nVars == length(vars)
          throw(ArgumentError(string("dateColumn: ", dateColumn, " not in DataFrame: ",vars)))
        end
        nVars = nVars-1
        dc = Symbol(dateColumn)

        n = size(returns,1)

        out = DataFrame(:__temp => 1.0)
        for i in 1:nVars
                out[!,vars[i]] .= (1.0+aggDiscreteReturns(returns[!,vars[i]]))^(1/n) - 1.0
        end
        return(out[!,vars])

end

function standard_deviation(returns::DataFrame; scale::Int=1,
                            dateColumn="date",
                            annualized::Bool = false,
                            )

        vars = names(returns)
        nVars = length(vars)
        vars = Symbol.(vars[vars.!=dateColumn])
        if nVars == length(vars)
            throw(ArgumentError(string("dateColumn: ", dateColumn, " not in DataFrame: ",vars)))
        end
        nVars = nVars-1

        out = DataFrame(:_stat_=>"Standard Deviation")

        if annualized
        #   TODO: Test if Threads.@treads is faster for a large number of variables...
                s = convert(Float64,scale)
                for i in 1:nVars
                        out[!,vars[i]] .= std(returns[!,vars[i]]) * sqrt(s)
                end
        else
                for i in 1:nVars
                    out[!,vars[i]] .= std(returns[!,vars[i]])
                end
        end

        return(out)
end


function parseinfix2rpn(s)
    outputq = []
    opstack = []
    infix = split(s)
    for tok in infix
        if all(isnumber, tok)
            push!(outputq, tok)
        elseif tok == "("
            push!(opstack, tok)
        elseif tok == ")"
            while !isempty(opstack) && (op = pop!(opstack)) != "("
               push!(outputq, op)
            end
        else # operator
            while !isempty(opstack)
                op = pop!(opstack)
                if Base.operator_precedence(Symbol(op)) > Base.operator_precedence(Symbol(tok)) ||
                   (Base.operator_precedence(Symbol(op)) ==
                     Base.operator_precedence(Symbol(tok)) && op != "^")
                    push!(outputq, op)
                else
                    push!(opstack, op)  # undo peek
                    break
                end
            end
            push!(opstack, tok)
        end
        println("The working output stack is $outputq")
    end
    while !isempty(opstack)
        if (op = pop!(opstack)) == "("
            throw("mismatched parentheses")
        else
            push!(outputq, op)
        end
    end
    outputq
end

teststring = "3 + 4 * 2 / ( 1 - 5 ) ^ 2 ^ 3"
println("\nResult: $teststring becomes $(join(parseinfix2rpn(teststring), ' '))")


"""
Module: SymbolicDiff (Symbolic Operation for Arithmetic)
"""

"""
AbstractSymbolic

An abstract structure for symbolic expr.

AbstractSymbolic has the filed
- params::Set{Symbol} A set of Symbol included in expr
"""

abstract type AbstractSymbolic{Tv} end

"""
AbstractNumberSymbolic
AbstractVectorSymbolic
AbstractMatrixSymbolic
AbstractGeneralSymbolic

They represent the types of zero values when the derivative becomes zero
"""

abstract type AbstractNumberSymbolic{Tv} <: AbstractSymbolic{Tv} end
abstract type AbstractVectorSymbolic{Tv} <: AbstractSymbolic{Tv} end
abstract type AbstractMatrixSymbolic{Tv} <: AbstractSymbolic{Tv} end
abstract type AbstractGeneralSymbolic{Tv} <: AbstractSymbolic{Tv} end

"""
SymbolicValue

A constant value
"""

mutable struct SymbolicValue{Tv} <: AbstractNumberSymbolic{Tv}
    params::Set{Symbol}
    val::Tv
end

SymbolicValue(value::Tv) where Tv = SymbolicValue(Set{Symbol}([]), value)

"""
SymbolicVariable

A variable to be derivatived
"""

mutable struct SymbolicVariable{Tv} <: AbstractNumberSymbolic{Tv}
    params::Set{Symbol}
    var::Symbol
end

"""
SymbolicExpr

An expr
"""

mutable struct SymbolicExpression{Tv} <: AbstractNumberSymbolic{Tv}
    params::Set{Symbol}
    op::Symbol
    args::Vector{<:AbstractSymbolic{Tv}}
    others
end

"""
convert
"""

function Base.convert(::Type{<:AbstractNumberSymbolic{T}}, x::SymbolicVariable{S}) where {T<:Number,S<:Number}
    symbolic(x.var, T)
end

function Base.convert(::Type{<:AbstractNumberSymbolic{T}}, x::SymbolicValue{S}) where {T<:Number,S<:Number}
    symbolic(T(x.val), T)
end

function Base.convert(::Type{<:AbstractNumberSymbolic{T}}, x::S) where {T<:Number,S<:Number}
    symbolic(T(x), T)
end

function Base.convert(::Type{<:AbstractNumberSymbolic{T}}, x::SymbolicExpression{S}) where {T<:Number,S<:Number}
    if T != S
        SymbolicExpression{T}(x.params, x.op, [convert(AbstractNumberSymbolic{T}, u) for u = x.args], [])
    else
        x
    end
end

"""
promotion
"""

function Base.promote_rule(::Type{<:AbstractNumberSymbolic{T}}, ::Type{S}) where {T<:Number,S<:Number}
    AbstractNumberSymbolic{promote_type(T,S)}
end

function Base.promote_rule(::Type{<:AbstractNumberSymbolic{T}}, ::Type{<:AbstractNumberSymbolic{S}}) where {T<:Number,S<:Number}
    AbstractNumberSymbolic{promote_type(T,S)}
end

"""
iszero
"""

function Base.iszero(x::AbstractNumberSymbolic{Tv}) where Tv
    false
end

function Base.iszero(x::SymbolicValue{Tv}) where Tv
    Base.iszero(x.val)
end

function Base.zero(::Type{AbstractNumberSymbolic{Tv}}) where Tv
    SymbolicValue(Tv(0))
end

"""
IO for structures
"""

function Base.show(io::IO, x::AbstractNumberSymbolic)
    expr = _toexpr(x)
    Base.show(io, expr)
end

function _toexpr(x::SymbolicExpression)
    args = [_toexpr(e) for e = x.args]
    Expr(:call, x.op, args...)
end

function _toexpr(x::SymbolicVariable)
    x.var
end

function _toexpr(x::SymbolicValue)
    x.val
end

"""
symbolicexpr(expr)

Build a SymbolicExpression
"""

const operations = [:+, :-, :*, :/, :^, :exp, :sqrt, :log]

function symbolic(expr::Expr, ::Type{Tv} = Float64) where Tv
    if Meta.isexpr(expr, :call) && expr.args[1] in operations
        args = [symbolic(x, Tv) for x = expr.args[2:end]]
        eval(Expr(:call, expr.args[1], args...))
    elseif Meta.isexpr(expr, :vect)
        vec = [symbolic(x, Tv) for x = expr.args]
        symbolic(vec, Tv)
    elseif Meta.isexpr(expr, :vcat)
        elem = [Expr(:row, [symbolic(y, Tv) for y = x.args]...) for x = expr.args]
        mat = eval(Expr(:vcat, elem...))
        symbolic(mat, Tv)
    else
        nothing
    end
end

function symbolic(expr::Symbol, ::Type{Tv} = Float64) where Tv
    SymbolicVariable{Tv}(Set([expr]), expr)
end

function symbolic(expr::Tv, ::Type{Tx}) where {Tv,Tx}
    SymbolicValue(expr)
end

function symbolic(expr::Tv) where Tv
    SymbolicValue(expr)
end

function symbolic(expr::Nothing, ::Type{Tv}) where Tv
    nothing
end

function symbolic(expr::Nothing) where Tv
    nothing
end


export verify_facet

"""
    verify_facet(
        d :: Int64,
        facet :: BellScenario.BellGame,
        aff_ind_vertices :: Vector{Matrix{Int64}}
    ) :: Bool

Returns `true` if the affinely independent set of vertices
`aff_ind_vertices` verifies that the `facet` is indeed a tight Bell inequality
of the ``\\matcal{C}_d^{X \\to Y}`` signaling polytope.
Input `d` must align with the considered facet.
This method checks that all provided vertices are:
* extreme points of the ``\\mathcal{C}_d^{X \\to Y}`` signaling polytope
* satisfy the facet inequality with equality
* affinely independent

!!! note "Note"
    This method does not check that all signaling polytope vertices satisfy the
    facet inequality. For the cataloged Bell inequalities, this requirement can be
    analytically shown to hold in general. To verify the tightness of the
    Bell inequality, it therefore remains to find a sufficient number of affinely
    independent vertices that saturate the inequality.
"""
function verify_facet(d :: Int64, facet :: BellScenario.BellGame, aff_ind_vertices :: Vector{Matrix{Int64}}) :: Bool
    (Y, X) = size(facet)
    polytope_dim = X*(Y-1)

    _verify_vertex(v) = (
        rank(v) == d
        && facet[:]'*v[:] == facet.β
        && is_deterministic(v)
    )

    all_valid_vertices = all(_verify_vertex, aff_ind_vertices)

    affine_independence = (
        length(aff_ind_vertices) == polytope_dim
        && LocalPolytope.dimension(aff_ind_vertices) == polytope_dim - 1
    )

    (all_valid_vertices && affine_independence)
end


using GAIO
using ForwardDiff
using StaticArrays
using BenchmarkTools

grid = LinRange(-1, 1, 300)
points = collect(Iterators.product(grid, grid))

f(x) = SA[1/2 - 2*1.4*x[1]^2 + x[2], 0.3*x[1]]

domain = Box((0.0, 0.0), (1.0, 1.0))
partition = RegularPartition(domain, 18)
point = ntuple(i->rand(2), 10000)
B = partition[point]

g1 = PointDiscretizedMap(f, domain, points)
g2 = AdaptiveBoxMap(f, domain)

@time g1B = g1(B)
g2B = g2(B)

# and for the Lorenz system

@inline function rk4(f, x, τ)
    τ½ = τ/2

    k = f(x)
    dx = @. k/6

    k = f(@. x+τ½*k)
    dx = @. dx + k/3

    k = f(@. x+τ½*k)
    dx = @. dx + k/3

    k = f(@. x+τ*k)
    dx = @. dx + k/6

    return @. x+τ*dx
end

function lorenz_v(x)
    s = 10.0
    rh = 28.0
    b = 0.4
    return s * (x[2] - x[1]), rh * x[1] - x[2] - x[1] * x[3], x[1] * x[2] - b * x[3]
end

function lorenz_f(x)
    h = 0.01
    n = 20

    for i = 1:n
        x = rk4(lorenz_v, x, h)
    end

    return x
end

grid = LinRange(-1, 1, 6)
points = collect(Iterators.product(grid, grid, grid))

domain = Box((0.0, 0.0, 27.0), (30.0, 30.0, 40.0))
partition = BoxPartition(domain, 27)

rh = 28.0
b = 0.4
x0 = (sqrt(b*(rh-1)), sqrt(b*(rh-1)), rh-1)

boxset = partition[x0]

g_adaptive = AdaptiveBoxMap(lorenz_f, domain)
g_points = PointDiscretizedMap(lorenz_f, domain, points)

point = ntuple(i->60.0*rand(3).-30.0, 1000000);
B = partition[point]; @show length(B)

@time g1B = g_points(B); @show length(g1B)
gaB = g_adaptive(B); @show length(gaB)
@show g_adaptive.image_points(point, point)


using SparseArrays, LinearAlgebra
using JuMP, BilevelJuMP
using Random

function bench_rand(rows, cols, density, optimizer, mode, seed = 1234)

    # seed = 1234
    # rows = 5
    # cols = 5
    # density = 0.4

    rng = Random.MersenneTwister(seed)

    f_A(r, n) = ifelse.(rand(r, n) .> 0.25, 1, -1) .* (15 .+ 30 .* rand(r, n))

    Al = sprand(rng, rows, cols, density, f_A)
    Au = sprand(rng, rows, cols, density, f_A)

    Bl = sprand(rng, rows, cols, density, f_A)
    Bu = sprand(rng, rows, cols, density, f_A)

    bl = 50 * rand(rng, rows)
    bu = 50 * rand(rng, rows)

    s_l = diagm(0 => ifelse.(rand(rng, rows) .> 0.2, 1, -1))
    s_u = diagm(0 => ifelse.(rand(rng, rows) .> 0.2, 1, -1))

    f_c(r, n) = 20 .* 2 .* (rand(r, n) .- 0.5)

    cl = sprand(rng, cols, 0.5, f_c)
    cu = sprand(rng, cols, 0.5, f_c)

    # MOI.empty!(optimizer)
    model = BilevelModel(optimizer, mode = mode)
    try
        JuMP.set_time_limit_sec(model, MAX_TIME)
    catch e
        @show e
        @show "failed to set limit time"
    end
    @variable(Upper(model), -1000 <= x[1:cols] <= 1000)
    @variable(Lower(model), -1000 <= y[1:cols] <= 1000)

    @constraint(Upper(model), (s_u * Au) * x .+ (s_u * Bu) * y .<= s_u * bu)
    @constraint(Lower(model), (s_l * Al) * x .+ (s_l * Bl) * y .<= s_l * bl)

    @objective(Upper(model), Min, cu' * x)
    @objective(Lower(model), Min, cl' * y)

    #=
        Optimize
    =#

    optimize!(model)

    @show primal_st = primal_status(model)
    @show term_st = termination_status(model) #in [MOI.OPTIMAL, MOI.LOCALLY_SOLVED, MOI.ALMOST_LOCALLY_SOLVED]

    solve_t = JuMP.solve_time(model)
    build_t = BilevelJuMP.build_time(model)

    obj_l = try
        objective_value(Lower(model))
    catch
        NaN
    end
    obj_u = try
        objective_value(Upper(model))
    catch
         NaN
    end
    gap = try
        bound = objective_bound(Upper(model))
        abs(obj_u - bound)/max(abs(bound), 1e-8)
    catch
        NaN
    end

    return primal_st, term_st, solve_t, build_t, obj_l, obj_u, gap
end


@testset "Discrete Explicit Tiger" begin
    S = [:left, :right]
    A = [:left, :right, :listen]
    O = [:left, :right]
    γ = 0.95

    function T(s, a, sp)
        if a == :listen
            return s == sp
        else # a door is opened
            return 0.5 #reset
        end
    end

    function Z(a, sp, o)
        if a == :listen
            if o == sp
                return 0.85
            else
                return 0.15
            end
        else
            return 0.5
        end
    end

    function R(s, a)
        if a == :listen
            return -1.0
        elseif s == a # the tiger was found
            return -100.0
        else # the tiger was escaped
            return 10.0
        end
    end

    m = DiscreteExplicitPOMDP(S,A,O,T,Z,R,γ)

    solver = FunctionSolver(x->:listen)
    policy = solve(solver, m)
    updater = DiscreteUpdater(m)

    rsum = 0.0
    for (s,b,a,o,r) in stepthrough(m, policy, updater, "s,b,a,o,r", max_steps=10)
        println("s: $s, b: $([pdf(b,s) for s in S]), a: $a, o: $o")
        rsum += r
    end
    println("Undiscounted reward was $rsum.")
    @test rsum == -10.0

    dm = DiscreteExplicitPOMDP(S,A,O,T,Z,R,γ,Deterministic(:left))
    @test rand(Random.GLOBAL_RNG, initialstate(dm)) == :left
    tm = DiscreteExplicitPOMDP(S,A,O,T,Z,R,γ,terminals=Set(S))
    @test isterminal(tm, rand(Random.GLOBAL_RNG, initialstate(tm)))

    for s in states(m)
        @test convert_s(statetype(m), convert_s(Vector{Float64}, s, m), m) == s
    end
    for a in actions(m)
        @test convert_a(actiontype(m), convert_a(Vector{Float64}, a, m), m) == a
    end
    for o in observations(m)
        @test convert_o(obstype(m), convert_o(Vector{Float64}, o, m), m) == o
    end
end

@testset "Discrete Explicit MDP" begin
    S = 1:5
    A = [-1, 1]
    γ = 0.95
    p₀ = Deterministic(1)

    function T(s, a, sp)
        if sp == clamp(s+a,1,5)
            return 0.8
        elseif sp == clamp(s-a,1,5)
            return 0.2
        else
            return 0.0
        end
    end

    function R(s, a)
        if s == 5
            return 1.0
        else
            return -1.0
        end
    end

    m = DiscreteExplicitMDP(S,A,T,R,γ)
    m = DiscreteExplicitMDP(S,A,T,R,γ,p₀)
    m = DiscreteExplicitMDP(S,A,T,R,γ,p₀,terminals=Set(5))
    @test isterminal(m, 5)

    solver = FunctionSolver(x->1)
    policy = solve(solver, m)

    rsum = 0.0
    for (s,a,r) in stepthrough(m, policy, "s,a,r", max_steps=10)
        println("s: $s, a: $a")
        rsum += r
    end
    println("Undiscounted reward was $rsum.")

    for s in states(m)
        @test convert_s(statetype(m), convert_s(Vector{Float64}, s, m), m) == s
    end
    for a in actions(m)
        @test convert_a(actiontype(m), convert_a(Vector{Float64}, a, m), m) == a
    end
end

@testset "convert Number POMDP" begin
    s = [1,47]
    T(s, a, sp) = s == sp
    Z(a, sp, o) = o == sp
    R(s, a) = 1.0
    m = DiscreteExplicitPOMDP(s,s,s,T,Z,R,0.95)
    for s in states(m)
        @test convert_s(statetype(m), convert_s(Vector{Float64}, s, m), m) == s
    end
    for a in actions(m)
        @test convert_a(actiontype(m), convert_a(Vector{Float64}, a, m), m) == a
    end
    for o in observations(m)
        @test convert_o(obstype(m), convert_o(Vector{Float64}, o, m), m) == o
    end
end

@testset "convert Vector POMDP" begin
    s = [[1,2],[3,4]]
    T(s, a, sp) = s == sp
    Z(a, sp, o) = o == sp
    R(s, a) = 1.0
    m = DiscreteExplicitPOMDP(s,s,s,T,Z,R,0.95)
    for s in states(m)
        @test convert_s(statetype(m), convert_s(Vector{Float64}, s, m), m) == s
    end
    for a in actions(m)
        @test convert_a(actiontype(m), convert_a(Vector{Float64}, a, m), m) == a
    end
    for o in observations(m)
        @test convert_o(obstype(m), convert_o(Vector{Float64}, o, m), m) == o
    end
end


"PolynomialBasis is the abstract supertype of all univariate polynomials."
abstract type PolynomialBasis{T} <: Dictionary{T,T}
end

##################
# Native indices
##################

# The native index of a polynomial basis is the degree, which starts from 0 rather
# than from 1.
struct PolynomialDegree <: AbstractShiftedIndex{1}
	value	::	Int
end

degree(idx::PolynomialDegree) = value(idx)

Base.show(io::IO, idx::PolynomialDegree) =
	print(io, "Polynomial degree $(degree(idx))")

ordering(b::PolynomialBasis) = ShiftedIndexList(length(b), PolynomialDegree)

isbasis(b::PolynomialBasis) = true

function sub(b::PolynomialBasis, idx::OrdinalRange)
    if (step(idx) == 1) && (first(idx) == 1) && (last(idx) <= length(b))
        resize(b, last(idx))
    else
        defaultsub(b, idx)
    end
end

abstract type Polynomial{T} <: AbstractBasisFunction{T,T} end

degree(p::Polynomial) = p.degree


"""
    likelihoods(hmm, observations; logl) -> Matrix

Return the likelihood per-state and per-observation.

**Output**
- `Matrix{Float64}`: likelihoods matrix (`T x K`).

**Example**
```julia
using Distributions, HMMBase
hmm = HMM([0.9 0.1; 0.1 0.9], [Normal(0,1), Normal(10,1)])
y = rand(hmm, 1000)
L = likelihoods(hmm, y)
```
"""
function likelihoods(hmm::AbstractHMM, observations; logl = false, robust = false)
    T, K = size(observations, 1), size(hmm, 1)
    L = Matrix{Float64}(undef, T, K)

    if logl
        loglikelihoods!(L, hmm, observations)
        robust && replace!(L, -Inf => nextfloat(-Inf), Inf => log(prevfloat(Inf)))
    else
        likelihoods!(L, hmm, observations)
        robust && replace!(L, -Inf => nextfloat(-Inf), Inf => prevfloat(Inf))
    end

    L
end


@symbol_func function K_RP(cur_reactor::AbstractReactor)
  cur_epsilon = cur_reactor.epsilon

  cur_K = 5.6e-8

  cur_K *= cur_reactor.Z_eff

  cur_K /= cur_epsilon ^ 2

  cur_K /= cur_reactor.kappa_95

  cur_K /= ( 1 - 1.31 * sqrt(cur_epsilon) + 0.46 * cur_epsilon )

  cur_K
end


"""
    community_detection_nback(g, k)

Community detection using the spectral properties of
the non-backtracking matrix of graph `g` (see [Krzakala et al.](http://www.pnas.org/content/110/52/20935.short)).
`k` is the number of communities to detect.

See also [`community_detection_bethe`](@ref) for a related community ddetection
algorithm.

Returns a vector with the vertex assignments in the communities.
"""
function community_detection_nback(g::AGraph, k::Integer)
    #TODO insert check on connected_components
    ϕ = real(nonbacktrack_embedding(g, k))
    if k == 2
        c = community_detection_threshold(g, ϕ[1,:])
    else
        c = kmeans(ϕ, k).assignments
    end
    return c
end

function community_detection_threshold(g::AGraphOrDiGraph, coords::AbstractArray)
    # TODO use a more intelligent method to set the threshold
    # 0 based thresholds are highly sensitive to errors.
    c = ones(Int, nv(g))
    # idx = sortperm(λ, lt=(x,y)-> abs(x) > abs(y))[2:k] #the second eigenvector is the relevant one
    for i=1:nv(g)
        c[i] = coords[i] > 0 ?  1 : 2
    end
    return c
end

"""
    nonbacktrack_embedding(g::AGraph, k::Int)

Spectral embedding of the non-backtracking matrix of `g`
(see [Krzakala et al.](http://www.pnas.org/content/110/52/20935.short)).

    `g`: imput Graph
    `k`: number of dimensions in which to embed

Returns  a matrix `ϕ` where `ϕ[:,i]` are the coordinates for vertex `i`.

Note:  does not explicitly construct the [`nonbacktracking_matrix`](@ref).
See [`nonbacktracking_matrix`](@ref) for details.
"""
function nonbacktrack_embedding(g::AGraph, k::Int)
    B, edgeid = nonbacktracking_matrix(g)
    λ, eigv = Arpack.eigs(B, nev=k+1, v0=ones(Float64, size(B,1)))
    ϕ = zeros(ComplexF64, k-1, nv(g))
    # TODO decide what to do with the stationary distribution ϕ[:,1]
    # this code just throws it away in favor of eigv[:,2:k+1].
    # we might also use the degree distribution to scale these vectors as is
    # common with the laplacian/adjacency methods.
    E = edgetype(g)
    for n=1:k-1
        v = eigv[:,n+1]
        for i=1:nv(g)
            for e in in_edges(g, i)
                j = src(e)
                u = edgeid[e]
                ϕ[n,i] += v[u]
            end
        end
    end
    return ϕ
end

"""
    community_detection_bethe(g::AGraph, k=-1; kmax=15)

Community detection using the spectral properties of
the Bethe Hessian matrix associated to `g` (see [Saade et al.](http://papers.nips.cc/paper/5520-spectral-clustering-of-graphs-with-the-bethe-hessian)).
`k` is the number of communities to detect. If omitted or if `k < 1` the
optimal number of communities will be automatically selected.
In this case the maximum number of detectable communities is given by `kmax`.
Returns a vector containing the vertex assignments.
"""
function community_detection_bethe(g::AGraph, k::Int=-1; kmax::Int=15)
    A = adjacency_matrix(g)
    D = Diagonal(degree(g))
    r = (sum(degree(g)) / nv(g))^0.5

    Hr = (r^2-1)*Diagonal(ones(nv(g))) - r*A + D;
    # Hmr = (r^2-1)*eye(nv(g))+r*A+D;
    k >= 1 && (kmax = k)
    λ, eigv = Arpack.eigs(Hr, which=:SR, nev=min(kmax, nv(g)))
    q = findlast(x -> x<0, λ)
    k > q && @warn("Using eigenvectors with positive eigenvalues,
                    some communities could be meaningless. Try to reduce `k`.")
    k < 1 && (k = q)
    k < 1 && return fill(1, nv(g))
    m = Matrix(eigv[:,2:k]') # convert to Matrix since
                             # kmeans doesn't support Adjoint yet
    labels = kmeans(m, k).assignments
    return labels
end

"""
    label_propagation(g; maxiter=1000)

Community detection using the label propagation algorithm (see [Raghavan et al.](http://arxiv.org/abs/0709.2938)).
`g` is the input Graph, `maxiter` is the  maximum number of iterations.
Returns a vertex assignments and the convergence history
"""
function label_propagation(g::AGraphOrDiGraph; maxiter=1000)
    n = nv(g)
    label = collect(1:n)
    active_nodes = BitSet(vertices(g))
    c = NeighComm(collect(1:n), fill(-1,n), 1)
    convergence_hist = Vector{Int}()
    random_order = Vector{Int}(undef, n)
    i = 0
    while !isempty(active_nodes) && i < maxiter
        num_active = length(active_nodes)
        push!(convergence_hist, num_active)
        i += 1

        # processing nodes in random order
        for (j,node) in enumerate(active_nodes)
            random_order[j] = node
        end
        range_shuffle!(1:num_active, random_order)
        @inbounds for j=1:num_active
            u = random_order[j]
            old_comm = label[u]
            label[u] = vote!(g, label, c, u)
            if old_comm != label[u]
                for v in out_neighbors(g, u)
                    push!(active_nodes, v)
                end
            else
                delete!(active_nodes, u)
            end
        end
    end
    fill!(c.neigh_cnt, 0)
    renumber_labels!(label, c.neigh_cnt)
    label, convergence_hist
end

"""Type to record neighbor labels and their counts."""
mutable struct NeighComm
  neigh_pos::Vector{Int}
  neigh_cnt::Vector{Int}
  neigh_last::Int
end

"""Fast shuffle Array `a` in UnitRange `r` inplace."""
function range_shuffle!(r::UnitRange, a::AbstractVector)
    (r.start > 0 && r.stop <= length(a)) || error("out of bounds")
    @inbounds for i=length(r):-1:2
        j = rand(1:i)
        ii = i + r.start - 1
        jj = j + r.start - 1
        a[ii],a[jj] = a[jj],a[ii]
    end
end

"""Return the most frequency label."""
function vote!(g::AGraphOrDiGraph, m::Vector{Int}, c::NeighComm, u::Int)
    @inbounds for i=1:c.neigh_last-1
        c.neigh_cnt[c.neigh_pos[i]] = -1
    end
    c.neigh_last = 1
    c.neigh_pos[1] = m[u]
    c.neigh_cnt[c.neigh_pos[1]] = 0
    c.neigh_last = 2
    max_cnt = 0
    for neigh in out_neighbors(g,u)
        neigh_comm = m[neigh]
        if c.neigh_cnt[neigh_comm] < 0
            c.neigh_cnt[neigh_comm] = 0
            c.neigh_pos[c.neigh_last] = neigh_comm
            c.neigh_last += 1
        end
        c.neigh_cnt[neigh_comm] += 1
        if c.neigh_cnt[neigh_comm] > max_cnt
          max_cnt = c.neigh_cnt[neigh_comm]
        end
    end
    # ties breaking randomly
    range_shuffle!(1:c.neigh_last-1, c.neigh_pos)
    for lbl in c.neigh_pos
      if c.neigh_cnt[lbl] == max_cnt
        return lbl
      end
    end
end

function renumber_labels!(membership::Vector{Int}, label_counters::Vector{Int})
    N = length(membership)
    (maximum(membership) > N || minimum(membership) < 1) && error("Label must between 1 and |V|")
    j = 1
    @inbounds for i=1:length(membership)
        k = membership[i]
        if k >= 1
            if label_counters[k] == 0
                # We have seen this label for the first time
                label_counters[k] = j
                k = j
                j += 1
            else
                k = label_counters[k]
            end
        end
        membership[i] = k
    end
end


brusselator_f(x, y, t) = ifelse((((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) &&
                                (t >= 1.1), 5., 0.)
function limit(a, N)
  if a == N+1
    return 1
  elseif a == 0
    return N
  else
    return a
  end
end
function brusselator_2d_loop(du, u, p, t)
  @inbounds begin
    A, B, α, xyd, dx, N = p
    α = α/dx^2
    for I in CartesianIndices((N, N))
      x = xyd[I[1]]
      y = xyd[I[2]]
      i = I[1]
      j = I[2]
      ip1 = limit(i+1, N)
      im1 = limit(i-1, N)
      jp1 = limit(j+1, N)
      jm1 = limit(j-1, N)
      du[i,j,1] = α*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1]) +
      B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
    end
    for I in CartesianIndices((N, N))
      i = I[1]
      j = I[2]
      ip1 = limit(i+1, N)
      im1 = limit(i-1, N)
      jp1 = limit(j+1, N)
      jm1 = limit(j-1, N)
      du[i,j,2] = α*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2]) +
      A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
    end
  end
end
function init_brusselator_2d(xyd)
  N = length(xyd)
  u = zeros(N, N, 2)
  for I in CartesianIndices((N, N))
    x = xyd[I[1]]
    y = xyd[I[2]]
    u[I,1] = 22*(y*(1-y))^(3/2)
    u[I,2] = 27*(x*(1-x))^(3/2)
  end
  u
end
xyd_brusselator = range(0,stop=1,length=32)

"""
2D Brusselator

```math
\\frac{\\partial u}{\\partial t} = 1 + u^2v - 4.4u + \\alpha(\\frac{\\partial^2 u}{\\partial x^2} + \\frac{\\partial^2 u}{\\partial y^2}) + f(x, y, t)
```
```math
\\frac{\\partial v}{\\partial t} = 3.4u - u^2v + \\alpha(\\frac{\\partial^2 u}{\\partial x^2} + \\frac{\\partial^2 u}{\\partial y^2})
```

where

``f(x, y, t) = \\begin{cases} 5 & \\quad \\text{if } (x-0.3)^2+(y-0.6)^2 ≤ 0.1^2 \\text{ and } t ≥ 1.1 \\\\ 0 & \\quad \\text{else} \\end{cases}``

and the initial conditions are

```math
u(x, y, 0) = 22\\cdot y(1-y)^{3/2}
```
```math
v(x, y, 0) = 27\\cdot x(1-x)^{3/2}
```

with the periodic boundary condition

```math
u(x+1,y,t) = u(x,y,t)
```
```math
u(x,y+1,t) = u(x,y,t)
```

From Hairer Norsett Wanner Solving Ordinary Differential Equations II - Stiff and Differential-Algebraic Problems Page 152
"""
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
                                     init_brusselator_2d(xyd_brusselator),
                                     (0.,11.5),
                                     (3.4, 1., 10.,
                                      xyd_brusselator, step(xyd_brusselator),
                                      length(xyd_brusselator)))

const N_brusselator_1d = 40
const D_brusselator_u = CenteredDifference{Float64}(2,2,1/(N_brusselator_1d-1),
                                                    N_brusselator_1d)
const D_brusselator_v = CenteredDifference{Float64}(2,2,1/(N_brusselator_1d-1),
                                                    N_brusselator_1d)
function brusselator_1d(du, u_, p, t)
    A, B, α, buffer = p
    u = @view(u_[:, 1])
    v = @view(u_[:, 2])
    mul!(buffer, D_brusselator_u, u)
    Du = buffer
    @. du[:, 1] = A + u^2*v - (B+1)*u + α*Du

    mul!(buffer, D_brusselator_v, v)
    Dv = buffer
    @. du[:, 2] = B*u - u^2*v + α*Dv
    nothing
end
function init_brusselator_1d(N)
  u = zeros(N, 2)
  x = range(0, stop=1, length=N)
  for i in 1:N
    u[i, 1] = 1 + sin(2pi*x[i])
    u[i, 2] = 3.
  end
  u
end

"""
1D Brusselator

```math
\\frac{\\partial u}{\\partial t} = A + u^2v - (B+1)u + \\alpha\\frac{\\partial^2 u}{\\partial x^2}
```
```math
\\frac{\\partial v}{\\partial t} = Bu - u^2v + \\alpha\\frac{\\partial^2 u}{\\partial x^2}
```

and the initial conditions are

```math
u(x,0) = 1+\\sin(2π x)
```
```math
v(x,0) = 3
```

with the boundary condition

```math
u(0,t) = u(1,t) = 1
```
```math
v(0,t) = v(1,t) = 3
```

From Hairer Norsett Wanner Solving Ordinary Differential Equations II - Stiff and Differential-Algebraic Problems Page 6
"""
prob_ode_brusselator_1d = ODEProblem(brusselator_1d,
                                    init_brusselator_1d(N_brusselator_1d),
                                    (0.,10.),
                                    (1., 3., 1/50, zeros(N_brusselator_1d)))


# Let's try parameter estimation for an ODE, here the famous Lorentz attractor.
# We base our implementation on the code from Paulo Marques available at:
#   https://github.com/pjpmarques/Julia-Modeling-the-World/blob/master/Lorenz%20Attractor.ipynb
function lorentz_equations(params::Vector{Float64}, r::Vector{Float64})
    # Get individual params for this ODE
    sigma, rho, beta = params

    # Extract the coordinates from the r vector
    x, y, z = r

    # The Lorenz equations
    dx_dt = sigma*(y - x)
    dy_dt = x*(rho - z) - y
    dz_dt = x*y - beta*z

    # Return the derivatives as a vector
    Float64[dx_dt, dy_dt, dz_dt]
end

# Define time vector and interval grid.
dt = 0.001
tf = 100.0
tinterval = 0:dt:tf
t  = collect(tinterval)

# But in order to compare to [Xiang2015] paper we use their choice instead:
h = 0.01
M = 300
tstart = 0.0
tstop = tstart + M * h
tinterval_Xiang2015 = 0:h:tstop
t_Xiang2015 = collect(tinterval_Xiang2015)

# Initial position in space
r0 = [0.1; 0.0; 0.0]

# Constants sigma, rho and beta. In a real ODE problem these would not be known
# and would be estimated.
sigma = 10.0
rho   = 28.0
beta  = 8.0/3.0
real_params = [sigma, rho, beta]

# "Play" an ODE from a starting point and into the future given a sequence of time steps.
function calc_state_vectors(params::Vector{Float64}, odefunc::Function,
    startx::Vector{Float64}, times::Vector{Float64}; states = nothing)

    numsamples = length(times)
    if states == nothing
        states = Array(Float64, length(startx), numsamples)
    end

    states[:, 1] = startx
    tprev = times[1]
    for i in 2:numsamples
        @inbounds derivatives = odefunc(params, states[:, (i-1)])
        @inbounds tnow = times[i]
        @inbounds states[:, i] = states[:, (i-1)] .+ derivatives * (tnow - tprev)
        tprev = tnow
    end

    return states
end

# RSS = Residual Sum of Squares, columnwise
function rss(actual::Array{Float64, 2}, estimated::Array{Float64, 2})
    M = size(actual, 2)
    sumsq = 0.0
    for i in 1:M
        @inbounds sumsq += sumabs2(actual[:, i] .- estimated[:, i])
    end
    sumsq
end

# Calculate the actual/original state vectors that we will use for parameter
# estimation:
origstates = calc_state_vectors(real_params, lorentz_equations, r0, t)
origstates_Xiang2015 = calc_state_vectors(real_params, lorentz_equations, r0, t_Xiang2015)

function subsample(origstates::Array{Float64, 2}, times::Vector{Float64}, lenratio = 0.25)
    @assert size(origstates, 2) == length(times)
    N = length(times)
    stopidx = round(Int, lenratio*N)
    indexes = 1:stopidx
    return origstates[:, indexes], times[indexes]
end

# The [Xiang2015] paper, https://www.hindawi.com/journals/ddns/2015/740721/,
# used these param bounds:
Xiang2015Bounds = Tuple{Float64, Float64}[(9, 11), (20, 30), (2, 3)]

# Now we can optimize using BlackBoxOptim
using BlackBoxOptim

function lorentz_fitness(params::Vector{Float64}, origstates::Array{Float64, 2}, times::Vector{Float64})
    states = calc_state_vectors(params, lorentz_equations, r0, times)
    return rss(origstates, states)
end

# But optimizing all states in each optimization step is too much so lets
# sample a small subset and use for first opt iteration.
origstates1, times1 = subsample(origstates, t, 0.04); # Sample only first 4%
origstates1_Xiang2015, times1_Xiang2015 = subsample(origstates_Xiang2015, t_Xiang2015, 1.00);

res1 = bboptimize(params -> lorentz_fitness(params, origstates1, times1);
    SearchRange = Xiang2015Bounds, MaxSteps = 8e3)

res2 = bboptimize(params -> lorentz_fitness(params, origstates1_Xiang2015, times1_Xiang2015);
    SearchRange = Xiang2015Bounds, MaxSteps = 11e3) # They allow 12k fitness evals for 3-param estimation

# But lets also try with less tight bounds
LooserBounds = Tuple{Float64, Float64}[(0, 22), (0, 60), (1, 6)]
res3 = bboptimize(params -> lorentz_fitness(params, origstates1_Xiang2015, times1_Xiang2015);
    SearchRange = LooserBounds, MaxSteps = 11e3) # They allow 12k fitness evals for 3-param estimation

println("Results on the long time sequence from Paulo Marques:")
estfitness = lorentz_fitness(best_candidate(res1), origstates, t)
@show (estfitness, best_candidate(res1), best_fitness(res1))
origfitness = lorentz_fitness(real_params, origstates, t)
@show (origfitness, real_params)

println("Results on the short time sequence used in [Xiang2015] paper:")
estfitness = lorentz_fitness(best_candidate(res2), origstates_Xiang2015, t_Xiang2015)
@show (estfitness, best_candidate(res2), best_fitness(res2))
origfitness = lorentz_fitness(real_params, origstates_Xiang2015, t_Xiang2015)
@show (origfitness, real_params)

println("Results on the short time sequence used in [Xiang2015] paper, but with looser bounds:")
estfitness = lorentz_fitness(best_candidate(res3), origstates_Xiang2015, t_Xiang2015)
@show (estfitness, best_candidate(res3), best_fitness(res3))
origfitness = lorentz_fitness(real_params, origstates_Xiang2015, t_Xiang2015)
@show (origfitness, real_params)


using ApproxFun, OrdinaryDiffEq, Sundials
using DiffEqDevTools
using LinearAlgebra
using Plots; gr()

# Define the linear and nonlinear terms
function lin_term(N)
    dx = 1/(N + 1)
    du = ones(N-1) # off diagonal
    du2 = -ones(N-1) # off diagonal
    d = (1/2) * ones(N-2)
    d2 = (-1/2) * ones(N-2)
    diag=-2 * ones(N)
    DiffEqArrayOperator(5e-5*(1/dx^3) * diagm(-2 => d2, -1 => du, 1 => du2, 2 => d))
end

function nl_term(N)
    dx = 1/(N + 1)
    du = ones(N - 1) # super diagonal
    dl = -ones(N - 1) # lower diagonal
    D = (0.2/dx) * diagm(-1 => dl, 1 => du)

    tmp = zeros(N)
    function (du,u,p,t)
        @. tmp = u^2
        mul!(du, D, tmp)
    end
end

# Construct the problem
function kdv(N)
    f1 = lin_term(N)
    f2 = nl_term(N)
    dx = 1 / (N + 1)
    xs = (1:N) * dx

    μ0 = 0.3; σ0 = 0.05
    f0 = x -> 0.9*exp(-(x - μ0)^2 / (2 * σ0^2))
    u0 = f0.(xs)
    prob = SplitODEProblem(f1, f2, u0, (0.0, 1.0))
    xs, prob
end;

xs, prob = kdv(200)
sol = solve(prob, Tsit5(); abstol=1e-11, reltol=1e-11, dt=1e-7, adptive=false)
test_sol = TestSolution(sol);

tslices = [0.0 0.25 0.50 0.75 1.00]
ys = hcat((sol(t) for t in tslices)...)
labels = ["t = $t" for t in tslices]
fn=plot(xs, ys, label=labels)

const LS_Dense = LinSolveFactorize(lu)

abstols = 0.1 .^ (5:8) # all fixed dt methods so these don't matter much
reltols = 0.1 .^ (1:4)
multipliers = 0.5 .^ (0:3)
setups = [Dict(:alg => IMEXEuler(), :dts => 1e-3 * multipliers),
          Dict(:alg => CNAB2(), :dts => 1e-5 * multipliers),
          Dict(:alg => CNLF2(), :dts => 1e-4 * multipliers),
          Dict(:alg => SBDF2(), :dts => 1e-3 * multipliers)]
labels = ["IMEXEuler" "CNAB2" "CNLF2" "SBDF2"]
@time wp = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp, label=labels, markershape=:auto, title="IMEX methods, dense linsolve, low order")

abstols = 0.1 .^ (5:8) # all fixed dt methods so these don't matter much
reltols = 0.1 .^ (1:4)
multipliers = 0.5 .^ (0:3)
setups = [#Dict(:alg => IMEXEuler(linsolve=LinSolveGMRES()), :dts => 1e-5 * multipliers),
          Dict(:alg => CNAB2(linsolve=LinSolveGMRES()), :dts => 1e-5 * multipliers),
          Dict(:alg => CNLF2(linsolve=LinSolveGMRES()), :dts => 1e-5 * multipliers),
          Dict(:alg => SBDF2(linsolve=LinSolveGMRES()), :dts => 1e-4 * multipliers)]
labels = ["IMEXEuler" "CNAB2" "CNLF2" "SBDF2"]
@time wp = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp, label=labels, markershape=:auto, title="IMEX methods, Krylov linsolve, low order")

abstols = 0.1 .^ (5:8) # all fixed dt methods so these don't matter much
reltols = 0.1 .^ (1:4)
multipliers = 0.5 .^ (0:3)
setups = [Dict(:alg => NorsettEuler(), :dts => 1e-6 * multipliers),
          Dict(:alg => NorsettEuler(krylov=true, m=5), :dts => 1e-3 * multipliers),
          #Dict(:alg => NorsettEuler(krylov=true, m=20), :dts => 1e-3 * multipliers), matrix contains Infs or NaNs
          Dict(:alg => ETDRK2(), :dts => 1e-3 * multipliers),
          Dict(:alg => ETDRK2(krylov=true, m=20), :dts => 1e-2 * multipliers),
          Dict(:alg => ETDRK2(krylov=true, m=20), :dts => 1e-2 * multipliers)]
labels = hcat("NorsettEuler (caching)", "NorsettEuler (m=5)",# "NorsettEuler (m=20)",
              "ETDRK2 (caching)", "ETDRK2 (m=5)"), "ETDRK2 (m=20)")
@time wp = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp, label=labels, markershape=:auto, title="ExpRK methods, low order")

abstols = 0.1 .^ (5:8) # all fixed dt methods so these don't matter much
reltols = 0.1 .^ (1:4)
multipliers = 0.5 .^ (0:3)
setups = [Dict(:alg => CNAB2(), :dts => 1e-5 * multipliers),
          Dict(:alg => CNAB2(linsolve=LinSolveGMRES()), :dts => 1e-5 * multipliers),
          Dict(:alg => ETDRK2(), :dts => 1e-3 * multipliers)]
labels = ["CNAB2 (dense linsolve)" "CNAB2 (Krylov linsolve)" "ETDRK2 (m=5)"]
@time wp = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp, label=labels, markershape=:auto, title="Between family, low orders")

abstols = 0.1 .^ (7:13)
reltols = 0.1 .^ (4:10)
setups = [Dict(:alg => KenCarp3()),
          Dict(:alg => KenCarp4()),
          Dict(:alg => KenCarp5()),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=3, linear_solver=:Dense)),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=4, linear_solver=:Dense)),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=5, linear_solver=:Dense))]
labels = hcat("KenCarp3", "KenCarp4", "KenCarp5", "ARKODE3", "ARKODE4", "ARKODE5")
@time wp = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp, label=labels, markershape=:auto, title="IMEX methods, dense linsolve, medium order")

abstols = 0.1 .^ (7:13)
reltols = 0.1 .^ (4:10)
setups = [Dict(:alg => KenCarp3(linsolve=LinSolveGMRES())),
          Dict(:alg => KenCarp4(linsolve=LinSolveGMRES())),
          Dict(:alg => KenCarp5(linsolve=LinSolveGMRES())),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=3, linear_solver=:GMRES)),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=4, linear_solver=:GMRES)),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=5, linear_solver=:GMRES))]
labels = ["KenCarp3" "KenCarp4" "KenCarp5" "ARKODE3" "ARKODE4" "ARKODE5"]
@time wp = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp, label=labels, markershape=:auto, title="IMEX methods, medium order")

abstols = 0.1 .^ (7:11) # all fixed dt methods so these don't matter much
reltols = 0.1 .^ (4:8)
multipliers = 0.5 .^ (0:4)
setups = [Dict(:alg => ETDRK3(), :dts => 1e-2 * multipliers),
          #Dict(:alg => ETDRK3(krylov=true, m=5), :dts => 1e-2 * multipliers),matrix contains Infs or NaNs
          Dict(:alg => ETDRK4(), :dts => 1e-3 * multipliers),
          #Dict(:alg => ETDRK4(krylov=true, m=5), :dts => 1e-2 * multipliers),matrix contains Infs or NaNs
          Dict(:alg => HochOst4(), :dts => 1e-2 * multipliers)]
          #Dict(:alg => HochOst4(krylov=true, m=5), :dts => 1e-2 * multipliers)] matrix contains Infs or NaNs
labels = hcat("ETDRK3 (caching)", "ETDRK4 (caching)",# "ETDRK3 (m=5)", "ETDRK4 (m=5)"
              "HochOst4 (caching)")#, "HochOst4 (m=5)")
@time wp = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp, label=labels, markershape=:auto, title="ExpRK methods, medium order")

abstols = 0.1 .^ (7:11)
reltols = 0.1 .^ (4:8)
multipliers = 0.5 .^ (0:4)
setups = [Dict(:alg => KenCarp5()),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=5, linear_solver=:Dense)),
          Dict(:alg => KenCarp5(linsolve=LinSolveGMRES())),
          Dict(:alg => ARKODE(Sundials.Implicit(), order=5, linear_solver=:GMRES)),
          Dict(:alg => ETDRK3(), :dts => 1e-2 * multipliers),
          Dict(:alg => ETDRK4(), :dts => 1e-3 * multipliers)]
labels = hcat("KenCarp5 (dense linsolve)", "ARKODE (dense linsolve)", "KenCarp5 (Krylov linsolve)",
              "ARKODE (Krylov linsolve)", "ETDRK3 (m=5)", "ETDRK4 (m=5)")
@time wp = WorkPrecisionSet(prob,abstols,reltols,setups;
                            print_names=true, names=labels,
                            numruns=5, error_estimate=:l2,
                            save_everystep=false, appxsol=test_sol, maxiters=Int(1e5));

plot(wp, label=labels, markershape=:auto, title="Between family, medium order")

using DiffEqBenchmarks
DiffEqBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


### A Pluto.jl notebook ###
# v0.16.4

using Markdown
using InteractiveUtils

# ╔═╡ 6541fa72-1528-11eb-157a-a3e1f9aa82a0
md"""# Simulation of a Computer Network
## Router as a Queuing System

Delay/Jitter:
- Transmission
- Service
- Processing
- Waiting

Quality of Service:
- Realtime
- Lossless

Representation:
- ``A\left(t\right)``
- ``D\left(t\right)``
- ``W\left(t\right)``
- ``W_Q\left(t\right)``
- ``T_n``
- ``S_n``
- ``L_n``
- ``L_{Qn}``

## Little's Law

```math
\mathbb E\left[L\right]=\lambda\mathbb E\left[W\right]
```
```math
\mathbb E\left[L_Q\right]=\lambda\mathbb E\left[W_Q\right]
```

## Steady State Probabilities

```math
P_n=\lim_{t\rightarrow\infty}\mathbb P\left\{L(t)=n\right\}
```

Stability
```math
a_n=d_n
```

Poisson arrivals always see time averages (PASTA):
```math
P_n=a_n
```

## MM1

Exponential Arrivals / Exponential Services (``\lambda`` / ``\mu``)

Markov Chain

A technical remark: if one event occurs at an exponential rate ``\lambda``, and another independent event at an exponential rate ``\mu``, then together they occur at an exponential rate ``\lambda+\mu``.

Steady State

```math
\begin{cases}
P_0=1-\frac{\lambda}{\mu}\\
P_n=\left(1-\frac{\lambda}{\mu}\right)\left(\frac{\lambda}{\mu}\right)^n,\quad n>0
\end{cases}
```

```math
\mathbb E\left[L\right]=\frac{\lambda}{\mu-\lambda}
```

```math
\mathbb E\left[W\right]=\frac{1}{\lambda}\mathbb E\left[L\right]=\frac{1}{\mu-\lambda}
```

```math
\mathbb E\left[W_Q\right]=\mathbb E\left[W\right]-\mathbb E\left[S\right]=\frac{\lambda}{\mu\left(\mu-\lambda\right)}
```

```math
\mathbb E\left[L_Q\right]=\lambda\mathbb E\left[W_Q\right]=\frac{\lambda^2}{\mu\left(\mu-\lambda\right)}
```

Finite Capacity?

Exponential Arrivals?

Exponential Services?

Network of Queues?

## GG1

Virtual Waiting Time: ``V\left(t\right)``
```math
\mathbb E\left[V\right]=\lambda\mathbb E\left[S\right]E\left[W_Q\right]+\frac{\lambda}{2}\mathbb E\left[S^2\right]
```

## MG1

Pollaczek–Khintchine formula
```math
\mathbb E\left[W_Q\right] = \frac{\lambda\mathbb E\left[S^2\right]}{2\left(1-\lambda\mathbb E\left[S\right]\right)}
```

## Statistical Processing

Sample Mean and Sample Variance

```math
\bar X = \sum_{i=1}^n X_i
```

```math
\mathbb E\left[\bar X\right]= \mu
```

To determine the “worth” of ``\bar X`` as an estimator of the population mean ``\mu``, we consider its mean square error:
```math
\mathbb E\left[\left(\bar X - \mu\right)^2\right]= \frac{\sigma^2}{n}
```

```math
S^2=\frac{\sum_{i=1}^n\left(X-\bar x\right)^2}{n-1}
```

```math
\mathbb E\left[S^2\right]=\sigma^2
```

``n`` sufficient large: central limit theorema
```math
\mathbb P\left\{\left|X − \mu\right|>c\frac{\sigma}{\sqrt n}\right\} \approx \mathbb P\left\{\left|Z\right|>c\right\}
```
where ``Z`` is a standard normal

Runtime

Number of Runs

The Regenerative Approach

Stationnarity vs Ergodicity
"""

# ╔═╡ Cell order:
# ╟─6541fa72-1528-11eb-157a-a3e1f9aa82a0


using VariationalOnlineStats
using Test, Statistics

for n = 1:2
    for i = 1:100
        x = rand(n)
        @test Float64(MedianOnline(x)) ≈ median(x)
    end
end

@test Float64(MedianOnline(Float64[1, 3, 2])) ≈ 2.0
@test Float64(MedianOnline(Float64[3, 1, 2])) ≈ 2.0
@test abs(Float64(MedianOnline(Float64[1, 2, 3])) - 2) < 0.1
@test abs(Float64(MedianOnline(Float64[3, 2, 1])) - 2) < 0.1

v = [1, 11, 10, 2, 3, 9, 8, 4, 5, 7, 6]
@test abs(Float64(MedianOnline{Float64}(v)) - 6) < 0.1
@test Float64(MedianOnline{Float64}(1:11)) > 4.5
@test Float64(MedianOnline{Float64}(11:-1:1)) < 7.5


using Flux, DiffEqFlux, OrdinaryDiffEq, DiffEqParamEstim, Plots, Optim, Dates

using DiffEqBase, Flux, DiffResults, DiffEqSensitivity, ForwardDiff
using Flux.Tracker: @grad
using DiffEqSensitivity: adjoint_sensitivities_u0
######################################
mutable struct saver
    losses::Array{Float64,1}
    times::Array{Dates.Time,1}
    count_epochs::Int128
end
function saver(n_epochs)
    losses = zeros(n_epochs)
    times = fill(Dates.Time(Dates.now()),n_epochs)
    count_epochs = 0
    return saver(losses,times,count_epochs)
end
function update_saver(saver, loss_i, time_i)
    epoch_i = saver.count_epochs
    saver.losses[epoch_i] = loss_i
    saver.times[epoch_i] = time_i
end
####################################################### Observation ###############################################################
# Start conditions for the two species in the system
u0 = Float32[2.; 0.]
# Number of evaluations of the neural ODE. It relates to the numbers of layers of the neural net (depth of network).
datasize = 30
# Time span in which of evaluation will be and actual timepoints of evaluations
tspan = (0.0f0, 1.5f0)
t = range(tspan[1], tspan[2], length = datasize)
# The true ODE (with the true parameters) which the neural net should learn
function trueODEfunc(du, u, p, t)
  true_A = [-0.1 2.0; -2.0 -0.1]
  du .= ((u.^3)'true_A)'
end
# Construction of the ODEProblem and solving the ODEProblem with Tsit5 solver
prob = ODEProblem(trueODEfunc, u0, tspan)
ode_data = Array(solve(prob,Tsit5(),saveat=t))
################################################### Building a neural ODE ##########################################################
# Derivative is modeled by a neural net. Chain concatinates the functions ode function and two dense layers.
dudt = Chain(x -> x.^3,
       Dense(2,50,tanh),
       Dense(50,2))
# Parameters of the model which are to be learnt. They are: W1 (2x50), b1 (50), W2 (50x2), b2 (2)
ps = Flux.params(dudt)
#build node
n_ode = x->neural_ode(dudt, x, tspan, Tsit5(), saveat=t, reltol=1e-7, abstol=1e-9)
# to set by user
init_train, number_sets, number_epochs_loss1, number_epochs_loss2 = 8000, 15, 1000, 20
# for saving
sa = saver(init_train+number_sets*(number_epochs_loss1+number_epochs_loss2))
# define collocation loss function with callback
function node_two_stage_function(model, x, tspan, saveat, ode_data,
            args...; kwargs...)
  dudt_(du,u,p,t) = du .= model(u)
  prob_fly = ODEProblem(dudt_,x,tspan)
  two_stage_method(prob_fly, saveat, ode_data)
end
loss_n_ode = node_two_stage_function(dudt, u0, tspan, t, ode_data, Tsit5(), reltol=1e-7, abstol=1e-9)
two_stage_loss_fct()=loss_n_ode.cost_function(ps)
cb1 = function ()
    sa.count_epochs = sa.count_epochs +  1
    update_saver(sa, Tracker.data(two_stage_loss_fct()),Dates.Time(Dates.now()))
    #println("\"",Tracker.data(two_stage_loss_fct()),"\" \"",Dates.Time(Dates.now()),"\";")
end
#define L2 loss function with callback
L2_loss_fct() = sum(abs2,ode_data .- n_ode(u0))
cb2 = function ()
    sa.count_epochs = sa.count_epochs +  1
    update_saver(sa, Tracker.data(L2_loss_fct()),Dates.Time(Dates.now()))
    #println("\"",Tracker.data(L2_loss_fct()),"\" \"",Dates.Time(Dates.now()),"\";")
end
#training call
data0 = Iterators.repeated((), init_train)
data1 = Iterators.repeated((), number_epochs_loss1)
data2 = Iterators.repeated((), number_epochs_loss2)
opt1 = ADAM(0.1)
opt2 = ADAM(0.1)
Flux.train!(two_stage_loss_fct, ps, data0, opt1, cb = cb1)
@time for i in 1:number_sets
    Flux.train!(L2_loss_fct, ps, data2, opt2, cb = cb2)
    Flux.train!(two_stage_loss_fct, ps, data1, opt1, cb = cb1)
end
# Call n_ode to get first prediction and to show startpoint for training.
pred = n_ode(u0)
scatter(t, ode_data[1,:], label="data")
    scatter!(t, Flux.data(pred[1,:]), label="prediction")
    scatter!(t, ode_data[2,:], label="data")
    scatter!(t, Flux.data(pred[2,:]), label="prediction")

#savefig("sogood.png")

using BSON: @save
@save "DiffEqFlux.jl/test/results/mix/model_mix_8000_15_200_20_epochs.bson" dudt


include("euler/euler.jl")
using .Numbers: get_digits

function compute(n::Float64)::Int
    bouncy = 0
    i = 1
    while bouncy / i < n
        i += 1
        digits = get_digits(i)
        sorted_digits = sort(digits)
        if !(digits == sorted_digits || digits == reverse(sorted_digits)) bouncy += 1 end
    end
    return i
end


Coord = CartesianIndex{3}

const movements = Dict{String,Coord}(
    "e"  => Coord(1, -1, 0),
    "w"  => Coord(-1, 1, 0),
    "ne" => Coord(1, 0, -1),
    "nw" => Coord(0, 1, -1),
    "se" => Coord(0, -1, 1),
    "sw" => Coord(-1, 0, 1),
)

function tilecoord(line::AbstractString, movements::Dict{String,Coord}=movements)
    dirs = r"se|sw|ne|nw|e|w"
    directions = [line[m] for m in findall(dirs, line)]
    return sum(movements[d] for d in directions)
end

function flip!(tiles::Dict{Coord,Bool}, line::AbstractString)
    coord = tilecoord(line)
    tiles[coord] = !get(tiles, coord, false)
    return tiles
end

function init(lines::Vector{<:AbstractString})
    tiles = Dict{Coord,Bool}()
    for line in lines
        flip!(tiles, line)
    end
    return tiles
end

const tiles = init(readlines("day24.txt"))

println("Part 1: $(sum(values(tiles)))")

const neighbors = collect(values(movements))

function with_black_neighbors(blacktiles::Set{Coord}, neighbors::Vector{Coord}=neighbors)
    blackneighbors = Dict{Coord,Int}()
    for t in blacktiles, n in neighbors
        blackneighbors[t + n] =  get(blackneighbors, t + n, 0) + 1
    end
    return blackneighbors
end

function flipall(blacktiles::Set{Coord}, blackneighbors::Dict{Coord,Int})
    futuretiles = copy(blacktiles)
    for a in blacktiles
        n = get(blackneighbors, a, 0)
        if n == 0 || n > 2
            pop!(futuretiles, a)
        end
    end
    for (a, n) in blackneighbors
        if n == 2 && a ∉ blacktiles
            push!(futuretiles, a)
        end
    end
    return futuretiles
end

function part2(tiles::Dict{Coord,Bool})
    blacktiles = Set{Coord}(k for (k, v) in tiles if v)
    for i in 1:100
        blackneighbors = with_black_neighbors(blacktiles)
        blacktiles = flipall(blacktiles, blackneighbors)
    end
    return length(blacktiles)
end

println("Part 2: $(part2(tiles))")


@testset "Finds a trivial complete arnoldi decomposition" begin

    matrix = [1 1 1;1 1 1;0 1 1];
    guess = [1;0;0]

    expected = [1 1 1;1 1 1;0 1 1]
    actual = Eigen.arnoldi(matrix, guess, 3)

    @test expected == actual.H
end

@testset "Finds a trivial partial arnoldi decomposition" begin

    matrix = [1 1 1;1 1 1;0 1 1];
    guess = [1;0;0]

    expected = [1 1;1 1;0 1]
    actual = Eigen.arnoldi(matrix, guess, 2)

    @test expected == actual.H
end

@testset "Stops the decomposition if an invariant subspace is found" begin

    matrix = [1 0 0;0 1 0;0 0 1]
    guess = [1;0;0]

    expected = 1
    actual = Eigen.arnoldi(matrix, guess, 2)

    @test expected == actual.numColumns
end


# linear phase FIR filter design that optimizes maximum error
# in the frequency domain

using DSP

"""
    remez_wrap(numtaps::Integer,
          bands::Vector,
          desired::Vector;
          weight::Vector=[],
          Hz::Real=1.0,
          filter_type::RemezFilterType=filter_type_bandpass,
          maxiter::Integer=25,
          grid_density::Integer=16)

Calculate the minimax optimal filter using the Remez exchange algorithm [^McClellan1973a] [^McClellan1973b].

Calculate the filter-coefficients for the finite impulse response
(FIR) filter whose transfer function minimizes the maximum error
between the desired gain and the realized gain in the specified
frequency bands using the Remez exchange algorithm.

`remez_wrap` is equivalent to `remez` in the DSP package, except it uses
the C code ported from scipy.

# Arguments
- `numtaps::Integer`: The desired number of taps in the filter.
    The number of taps is the number of terms in the filter, or the filter
    order plus one.
- `bands::Vector`: A monotonic sequence containing the band edges in Hz.
    All elements must be non-negative and less than half the sampling
    frequency as given by `Hz`.
- `desired::Vector`:A sequence half the size of bands containing the desired
    gain in each of the specified bands.
- `weight::Vector`: (optional)
    A relative weighting to give to each band region. The length of
    `weight` has to be half the length of `bands`.
- `Hz::Real`: The sampling frequency in Hz. Default is 1.
- `filter_type::RemezFilterType`: Default is filter_type_bandpass.
    The type of filter:
      filter_type_bandpass : flat response in bands. This is the default.
      filter_type_differentiator : frequency proportional response in bands.
        Assymetric as in filter_type_hilbert case, but with a linear sloping
        desired response.
      filter_type_hilbert : filter with odd symmetry, that is, type III
                  (for even order) or type IV (for odd order)
                  linear phase filters.
- `maxiter::Integer`: (optional)
    Maximum number of iterations of the algorithm. Default is 25.
- `grid_density:Integer`: (optional)
    Grid density. The dense grid used in `remez` is of size
    ``(numtaps + 1) * grid_density``. Default is 16.

# Returns
- `h::Array{Float64,1}`: A rank-1 array containing the coefficients of the optimal
    (in a minimax sense) filter.

[^McClellan1973a]:
J. H. McClellan and T. W. Parks, A unified approach to the
design of optimum FIR linear phase digital filters,
IEEE Trans. Circuit Theory, vol. CT-20, pp. 697-701, 1973.

[^McClellan1973b]:
J. H. McClellan, T. W. Parks and L. R. Rabiner, A Computer
Program for Designing Optimum FIR Linear Phase Digital
Filters, IEEE Trans. Audio Electroacoust., vol. AU-21,
pp. 506-525, 1973.

# Examples
Construct a length 35 filter with a passband at 0.15-0.4 Hz
(desired response of 1), and stop bands at 0-0.1 Hz and 0.45-0.5 Hz
(desired response of 0). Note: the behavior in the frequency ranges between
those bands - the transition bands - is unspecified.

```julia-repl
julia> bpass = remez_wrap(35, [0 0.1 0.15 0.4 0.45 0.5], [0 1 0])
```

You can trade-off maximum error achieved for transition bandwidth.
The wider the transition bands, the lower the maximum error in the
bands specified. Here is a bandpass filter with the same passband, but
wider transition bands.

```julia-repl
julia> bpass2 = remez_wrap(35, [0 0.08 0.15 0.4 0.47 0.5], [0 1 0])
```

Here we compute the frequency responses and plot them in dB.

```julia-repl
using PyPlot
b = DSP.Filters.PolynomialRatio(bpass, [1.0])
b2 = DSP.Filters.PolynomialRatio(bpass2, [1.0])
f = linspace(0, 0.5, 1000)
plot(f, 20*log10.(abs.(freqz(b,f,1.0))))
plot(f, 20*log10.(abs.(freqz(b2,f,1.0))))
grid()
```
"""
function remez_wrap(numtaps::Integer, bands::Vector, desired::Vector;
               weight::Vector=[],
               Hz::Real=1.0,
               filter_type::RemezFilterType=filter_type_bandpass,
               maxiter::Integer=25,
               grid_density::Integer=16)

    # Convert type
    if (filter_type == filter_type_bandpass)
        tnum = 1
    elseif (filter_type == filter_type_differentiator)
        tnum = 2
    elseif (filter_type == filter_type_hilbert)
        tnum = 3
    end
    if (length(weight)==0)
        weight = ones(desired)
    end

    h = zeros(numtaps)
    #return sigtools._remez(numtaps, bands, desired, weight, tnum, Hz,
                           #maxiter, grid_density)
    #int pre_remez(double *h2, int numtaps, int numbands, double *bands,
    #              double *response, double *weight, int type, int maxiter,
    #              int grid_density);
    bands = copy(bands)/Hz
    ccall((:pre_remez, "remez"), Integer, (Ptr{Cdouble}, Cint, Cint, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Cint, Cint, Cint),
        h, numtaps, length(bands)/2, bands, desired, weight, tnum, maxiter, grid_density)
    return h

end


# Beyn contour integral approach

#using QuadGK # Disabled to lessen requirements
using Distributed
using LinearAlgebra
using Random

export contour_beyn

"""
    λv,V=contour_beyn([eltype,] nep [,mintegrator];[tol,][logger,][σ,][radius,][linsolvercreator,][N,][neigs,][k])

The function computes eigenvalues using Beyn's contour integral approach,
using an ellipse centered at `σ` with radii given in `radius`, or if only one `radius` is given,
the contour is a circle. The numerical quadrature method is specified in `mintegrator`,
which is a type inheriting from `MatrixIntegrator`, by default
`MatrixTrapezoidal`. For a parallell implementation of the
integrator use `MatrixTrapezoidalParallel`.
 The integer `k`
specifies size of the probe subspace. `N` corresponds to the
number of quadrature points. Ellipses are the only supported contours. The
`linsolvercreator` must create a linsolver that can handle (rectangular) matrices
as right-hand sides, not only vectors. We integrate in complex arithmetic so
`eltype` must be complex type.

The kwargs `neigs` specifies the number of wanted eigvals, and `k` is the number
of columns in the matrix to be integrated (default `k=neigs+1`). If you give the `k` parameter and set `neigs=typemax(Int)` all found eigenvalues will be returned. The kwarg `sanity_check` decides if sorting and checking (and removal) of eigpairs should be done.
If disabled, the method
returns `k` (potentially inaccurate) eigpairs. The parameters `errmeasure` and
`tol` and `rank_drop_tol` are used for the sanity check,
to extract accurate eigenvalues.

# Example
```julia-repl
julia> using LinearAlgebra
julia> nep=nep_gallery("dep0");
julia> # Look for two eigvals in unit disk
julia> λv,V=contour_beyn(nep,radius=1.1,neigs=3);
julia> norm(compute_Mlincomb(nep,λv[1],V[:,1])) # Eigenpair 1
1.7462847531404259e-15
julia> norm(compute_Mlincomb(nep,λv[2],V[:,2])) # Eigenpair 2
7.69695692032292e-15
```
# References
* Wolf-Jürgen Beyn, An integral method for solving nonlinear eigenvalue problems, Linear Algebra and its Applications 436 (2012) 3839–3863
"""
contour_beyn(nep::NEP;params...)=contour_beyn(ComplexF64,nep;params...)
contour_beyn(nep::NEP,MIntegrator;params...)=contour_beyn(ComplexF64,nep,MIntegrator;params...)
function contour_beyn(::Type{T},
                      nep::NEP,
                      ::Type{MIntegrator}=MatrixTrapezoidal;
                      tol::Real=sqrt(eps(real(T))), # Note tol is quite high for this method
                      σ::Number=zero(complex(T)),
                      logger=0,
                      linsolvercreator=BackslashLinSolverCreator(),
                      neigs::Integer=2, # Number of wanted eigvals
                      k::Integer=neigs+1, # Columns in matrix to integrate
                      radius::Union{Real,Tuple,Array}=1, # integration radius
                      N::Integer=1000,  # Nof quadrature nodes
                      errmeasure::ErrmeasureType = DefaultErrmeasure(nep),
                      sanity_check=true,
                      rank_drop_tol=tol # Used in sanity checking
                      )where{T<:Number, MIntegrator<:MatrixIntegrator}

    @parse_logger_param!(logger)

    # Geometry
    length(radius)==1 ? radius=(radius,radius) : nothing
    g(t) = complex(radius[1]*cos(t),radius[2]*sin(t)) # ellipse
    gp(t) = complex(-radius[1]*sin(t),radius[2]*cos(t)) # derivative

    n=size(nep,1);

    if (k>n)
        error("Cannot compute more eigenvalues than the size of the NEP with contour_beyn() k=",k," n=",n);

    end
    if (k<=0)
        error("k must be positive, k=",k,
              neigs==typemax(Int) ? ". The kwarg k must be set if you use neigs=typemax" : ".")
    end

    Random.seed!(10); # Reproducability
    Vh=Array{T,2}(randn(real(T),n,k)) # randn only works for real

    function local_linsolve(λ::TT,V::Matrix{TT}) where {TT<:Number}
        local M0inv::LinSolver = create_linsolver(linsolvercreator, nep, λ+σ);
        # This requires that lin_solve can handle rectangular
        # matrices as the RHS
        return lin_solve(M0inv,V);
    end

    # Constructing integrands
    Tv(λ) = local_linsolve(T(λ),Vh)
    f(t) = Tv(g(t))*gp(t)
    push_info!(logger,"Computing integrals")

    local A0,A1
    AA=integrate_interval(MIntegrator, ComplexF64, f,
                        [ s-> one(s), g],0,2*pi,N,logger )
    A0=AA[:,:,1];
    A1=AA[:,:,2];

    # Don't forget scaling
    A0[:,:] = A0 ./(2im*pi);
    A1[:,:] = A1 ./(2im*pi);

    push_info!(logger,"Computing SVD prepare for eigenvalue extraction ",continues=true)
    V,S,W = svd(A0)
    p = count( S/S[1] .> rank_drop_tol);
    push_info!(logger," p=$p");

    V0 = V[:,1:p]
    W0 = W[:,1:p]
    B = (copy(V0')*A1*W0) * Diagonal(1 ./ S[1:p])

    # Extract eigenval and eigvec approximations according to
    # step 6 on page 3849 in the reference
    push_info!(logger,"Computing eigenvalues ")
    λ,VB=eigen(B)
    λ[:] = λ .+ σ
    push_info!(logger,"Computing eigenvectors ")
    V = V0 * VB;
    for i = 1:p
        normalize!(V[:,i]);
    end

    if (!sanity_check)
        sorted_index = sortperm(map(x->abs(σ-x), λ));
        inside_bool = (real(λ[sorted_index].-σ)/radius[1]).^2 + (imag(λ[sorted_index].-σ)/radius[2]).^2 .≤ 1
        inside_perm = sortperm(.!inside_bool)
        return (λ[sorted_index[inside_perm]],V[:,sorted_index[inside_perm]])
    end

    # Compute all the errors
    errmeasures=zeros(real(T),p);
    for i = 1:p
        errmeasures[i]=estimate_error(errmeasure,λ[i],V[:,i]);
    end

    good_index=findall(errmeasures .< tol);

    # Index vector for sorted to distance from σ
    sorted_good_index=
       good_index[sortperm(map(x->abs(σ-x), λ[good_index]))];

    # check that eigenvalues are inside contour, move them to the end if they are outside
   inside_bool = (real(λ[sorted_good_index].-σ)/radius[1]).^2 + (imag(λ[sorted_good_index].-σ)/radius[2]).^2 .≤ 1
   if any(.!inside_bool)
       if neigs ≤ sum(inside_bool)
           @warn "found $(sum(.!inside_bool)) evals outside contour, $p inside. all $neigs returned evals inside, but possibly inaccurate. try increasing N, decreasing tol, or changing radius"
       else
           @warn "found $(sum(.!inside_bool)) evals outside contour, $p inside. last $(neigs-sum(inside_bool)) returned evals outside contour. possible inaccuracy. try increasing N, decreasing tol, or changing radius"
       end
   end
   sorted_good_inside_perm = sortperm(.!inside_bool)

    # Remove all eigpairs not sufficiently accurate
    # and potentially remove eigenvalues if more than neigs.
    local Vgood,λgood
    if( size(sorted_good_index,1) > neigs)
        found_evals=size(sorted_good_index,1);
        push_info!(logger,"Removing unwanted eigvals: neigs=$neigs<$found_evals=found_eigvals")
        Vgood=V[:,sorted_good_index[sorted_good_inside_perm][1:neigs]];
        λgood=λ[sorted_good_index[sorted_good_inside_perm][1:neigs]];
    else
        Vgood=V[:,sorted_good_index[sorted_good_inside_perm]];
        λgood=λ[sorted_good_index[sorted_good_inside_perm]];
    end

    if (p==k)
       @warn "Rank-drop not detected, your eigvals may be correct, but the algorithm cannot verify. Try to increase k. This warning can be disabled with `sanity_check=false`." S
    end

    if (size(λgood,1)<neigs  && neigs < typemax(Int))
       @warn "We found fewer eigvals than requested. Try increasing domain, or decreasing `tol`. This warning can be disabled with `sanity_check=false`." S
    end

    return (λgood,Vgood)
end


#=

 From http:  reasoning.eas.asu.edu/lpmln/examples/lpmln2al/taxi.lpmln
  """
  Conside a program where you know that you are going to be Late
  if you don't get a Taxi.

  Since we have put a hard constraint that we cannot be late the probability of being
  Late is going to be 0. The probability of there being a NoTaxi comes out to be 71 %.
  {NoTaxi = T, Late = F} satisfies the second soft rule and the hard constraint. Hence
  the probability for this model comes out to be e\textsuperscript{2}/(
  e\textsuperscript{2} + e\textsuperscript{1}) which comes out to be 73\%
  \begin{lstlisting}
  Late 4.9995e-05
  NoTaxi 0.711979
  \end{lstlisting}
  """

  Note: Since lpmln use weights instead of probability, I don't know how to get the no_taxi as 0.71
  Hmm in the lpmln there's an example where the weights are corresponding to percentages:
     0.5 -> -1.6094
     0.8 -> -0.2231
  which is ln(0.5) and ln(0.8)

  But this don't help us here, since the weights are positive, 1 and 2....

  Cf ~/cplint/taxi.pl
     ~/blog/taxi.blog
     ~/psi/taxi.psi
     ~/webpppl/taxi.wppl

=#

using Turing, StatsPlots, DataFrames

@model function taxi()
    # There is a chance for a taxi to be not available
    no_taxi ~ Bernoulli(0.3)

    # You are going to be late if you don't get a taxi
    late ~ no_taxi ? Bernoulli(0.65) : Bernoulli(0);

    # But you certainly cannot be late no matter what situation.
    # Hence this constraint is encoded as a hard rule
    late == false || begin Turing.@addlogprob! -Inf; return end
end

model = taxi()

num_chains = 4
# chains = sample(model, Prior(), MCMCThreads(), 1000, num_chains)

chains = sample(model, MH(), MCMCThreads(), 10_000, num_chains) # 1.47s
# chains = sample(model, IS(), MCMCThreads(), 10_000, num_chains) # 1.48s

# chains = sample(model, PG(20), MCMCThreads(), 10_000, num_chains) # Too slow!
# chains = sample(model, SMC(), MCMCThreads(), 10_000, num_chains) # 4.1s, don't handle observes!

# chains = sample(model, HMC(0.05,10), MCMCThreads(), 10_000, num_chains) # Nope!
# chains = sample(model, NUTS(), MCMCThreads(), 10_000, num_chains) # Nope!

display(chains)


@testset "119.pascals-triangle-ii.jl" begin
    @test get_pascal_row(3) == [1, 3, 3, 1]
    @test get_pascal_row(5) == [1, 5, 10, 10, 5, 1]
    @test get_pascal_row(6) == [1, 6, 15, 20, 15, 6, 1]
    @test get_pascal_row(23) ==
          [1, 23, 253, 1771, 8855, 33649, 100947, 245157, 490314, 817190, 1144066, 1352078,
           1352078, 1144066, 817190, 490314, 245157, 100947, 33649, 8855, 1771, 253, 23, 1]
end


"""
PrivateKey(𝑒) represents an S256Point determined by 𝑃 = 𝑒G,
where 𝑒 is an integer and G the scep256k1 generator point.
"""
struct PrivateKey
    𝑒::Integer
    𝑃::AbstractPoint
    PrivateKey(𝑒) = new(𝑒, 𝑒 * G)
end

"""
Returns a Signature for a given PrivateKey and data 𝑧
pksign(pk::PrivateKey, 𝑧::Integer) -> Signature
"""
function pksign(pk::PrivateKey, 𝑧::Integer)
    𝑘 = rand(big.(0:N))
    𝑟 = (𝑘 * G).𝑥.𝑛
    𝑘⁻¹ = powermod(𝑘, N - 2, N)
    𝑠 = mod((𝑧 + 𝑟 * pk.𝑒) * 𝑘⁻¹, N)
    if 𝑠 > N / 2
        𝑠 = N - 𝑠
    end
    return Signature(𝑟, 𝑠)
end


# This file calcules the ellipticintegral
#
#    ∫_0^z dz/(sqrt(1-z^2)*sqrt(a^2-z^2))
#
# by writing the integrand as a Cauchy transform
# cauchyintegral is an integral of the Cauchy transform

using ApproxFun, SingularIntegralEquations

function ellipticintegral(a)
    xp=Fun(identity,[a,1.])
    xm=Fun(identity,[-a,-1.])
    up=2.0im./(sqrt(1+xp).*sqrt(a+xp).*sqrt(xp-1.).*sqrt(a-xp))
    um=2.0im./(sqrt(-1-xm).*sqrt(a+xm).*sqrt(1.-xm).*sqrt(a-xm))
    g(z)=cauchyintegral(up,z)+cauchyintegral(um,z)
    g
end


using Distributed; addprocs()
@everywhere begin
    using Pkg
    Pkg.activate(Base.source_dir())
end

using DiscreteChoiceModels, Dagger, Tables, CSV, Infiltrator

function main()
    data = DTable(Tables.rowtable(CSV.File("test/data/biogeme_swissmetro.dat", delim='\t')), 2000)

    # filter to wanted cases
    data = filter(data) do row
        ((row.PURPOSE == 1) | (row.PURPOSE == 3)) && row.CHOICE != 0
    end

    data = map(data) do r
        merge((
            avtr=(r.TRAIN_AV == 1) && (r.SP != 0),
            avsm=r.SM_AV == 1,
            avcar=(r.CAR_AV == 1) && (r.SP != 0)
        ), pairs(r))
    end

    @time multinomial_logit(
        @utility(begin
            1 ~ αtrain + βtravel_time * TRAIN_TT / 100 + βcost * (TRAIN_CO * (GA == 0)) / 100
            2 ~ αswissmetro + βtravel_time * SM_TT / 100 + βcost * SM_CO * (GA == 0) / 100
            3 ~ αcar + βtravel_time * CAR_TT / 100 + βcost * CAR_CO / 100

            αswissmetro = 0, fixed  # fix swissmetro ASC to zero
        end),
        :CHOICE,
        data,
        availability=[
            1 => :avtr,
            2 => :avsm,
            3 => :avcar,
        ]
    )
end

main()


using CompScienceMeshes
using SauterSchwabQuadrature
#create a subdivision surfaces for a sphere exaple
G = readmesh("/Users/Benjamin/Documents/sphere.in")
subdG=GSubdMesh(G)

chart = chart(subdG,1)
u1= [0.5,0.25]
u2= [0.7,0.1]
d1 = neighborhood(chart,u1)
d2 = neighborhood(chart,u2)
global p1 = cartesian(d1)
global p2 = cartesian(d2)
function kernel(x,y)
			return(((x-p1)'*(y-p2))*exp(-im*1*norm(x-y))/(4pi*norm(x-y)))	#k=1
end
result = sauterschwabintegral(chart, kernel)


@defcomp Population begin
    region = Index()

    # Parameters
    y_year_0 = Parameter(unit="year")
    y_year = Parameter(index=[time], unit="year")
    popgrw_populationgrowth = Parameter(index=[time, region], unit="%/year") # From p.32 of Hope 2009
    pop0_initpopulation = Parameter(index=[region], unit="million person") # Population in y_year_0

    # Variables
    pop_population = Variable(index=[time, region], unit="million person")

    function run_timestep(p, v, d, tt)

        for rr in d.region
            # Eq.28 in Hope 2002 (defined for GDP, but also applies to population)
            if is_first(tt)
                v.pop_population[tt, rr] = p.pop0_initpopulation[rr] * (1 + p.popgrw_populationgrowth[tt, rr] / 100)^(p.y_year[tt] - p.y_year_0)
            else
                v.pop_population[tt, rr] = v.pop_population[tt - 1, rr] * (1 + p.popgrw_populationgrowth[tt, rr] / 100)^(p.y_year[tt] - p.y_year[tt - 1])
            end
        end
    end
end

# Still need this function in order to set the parameters than depend on
# readpagedata, which takes model as an input. These cannot be set using
# the default keyword arg for now.
function addpopulation(model::Model)
    populationcomp = add_comp!(model, Population)

    populationcomp[:pop0_initpopulation] = readpagedata(model, "data/pop0_initpopulation.csv")

    return populationcomp
end


# shows some data analysis using DataFrames
using CSV, StatsModels, DataFrames, StatPlots
nerlove = CSV.read("nerlove.csv")
nerlove[:lnC] = log.(nerlove[:cost])
nerlove[:lnQ] = log.(nerlove[:output])
nerlove[:lnPL] = log.(nerlove[:labor])
nerlove[:lnPF] = log.(nerlove[:fuel])
nerlove[:lnPK] = log.(nerlove[:capital])
show(describe(nerlove))
@df nerlove scatter(:lnC, :lnQ)
gui()
f = @formula(lnC ~ 1 + lnQ + lnPL + lnPF + lnPK)
ols(f, nerlove)


using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D
using Basis2DQuad
using UniformQuadMesh

using SetupDG

push!(LOAD_PATH, "./examples/EntropyStableEuler")
using EntropyStableEuler

N =
x = 0:N:2*pi

h       = 2*pi/N
column  = [0; .5*(-1).^(1:Np_F-1).*cot.((1:Np_F-1)*h/2)]
Dt      = Array{Float64,2}(Toeplitz(column,column[[1;Np_F:-1:2]]))


# noise processes definition
sigma = 0.001
σ = 200*sigma
noisemix1 = MixtureModel([Normal(0, 0.001), Normal(0, σ)], [0.9, 0.1])
noisemixdensity1(x) = pdf.(noisemix1, x)

sigma_param = sqrt(0.002)
paramnoise = Normal(0, sigma_param)
parametricnoisedensity(x) = pdf(Normal(0, sigma_param), x)

sigma23 = sqrt(4*1e-2)
sigma33 = sqrt(1*1e-4)
noisemix2 = MixtureModel(Normal[Normal(0, sigma23), Normal(0, sigma33)], [1.0/3, 2.0/3])
noisemixdensity2(x) = pdf.(noisemix2, x)


using PyPlot
using OffsetArrays
using SpecialFunctions: erfcx, gamma

const ϕ = 1.172104
const a = 1.081792
const b = 4.492075

α = 0.5
β = 1.0
F(w, x, α, β) = w^(α-β) / (w^α+x)
F0(w, x, α, β) = -x / (w^β*(w^α+x))

function contour(u, N)
    μ = b * N
    iumϕ = Complex(-ϕ, u)
    w = μ * ( 1 + sin(iumϕ) )
    dw = μ * cos(iumϕ)
    return w, dw
end

function points(N)
    w = OffsetArray{Complex{Float64}}(undef, 0:N)
    dw = OffsetArray{Complex{Float64}}(undef, 0:N)
    h = a / N
    for n = 0:N
        w[n], dw[n] = contour(n*h, N)
    end
    return w, dw, h
end

function Taylor(x, α, β, N)
    s = 1.0 / gamma(β)
    powx = 1.0
    for n = 1:N
        powx *= -x
        s += powx / gamma(β+n*α)
    end
    return s
end

function newI(x, α, β, N)
    w, dw, h = points(N)
    s = exp(w[0]) * F(w[0], x, α, β) * dw[0]
    for n = 1:N
        s += 2 * exp(w[n]) * F(w[n], x, α, β) * dw[n]
    end
    return h * real(s) / (2π)
end

figure(1)
for N in [5, 10, 15, 20]
    w, dw, h = points(N)
    plot(real.(w), imag.(w), "o", markersize=2, label="N = $N")
end
legend()
xlabel(L"$\Re w$")
ylabel(L"$\Im w$")
grid(true)
savefig("points.pdf")

figure(2)
x = 2.0
N = 15
w, dw, h = points(N)
u = Float64[ n*h for n = 0:N ]
integrand = F.(w, x, α, β) .* exp.(w) .* dw / Complex(0.0, 2π)
plot(u, real.(integrand[0:N]), "+-", label="real part")
#plot(u, imag.(integrand[0:N]), "x-", label="imag part")
#legend()
grid(true)
xlabel(L"$u$")
title("N = $N")
savefig("integrand.pdf")
#    semilogy(u, abs.(real.(integrand[0:N])),
#             u, abs.(imag.(integrand[0:N])))

figure(3)
x = range(0.0, 5.0, length=201)
nn = [ n for n=1:length(x) if x[n] ≤ 0.2 ]
y1 = erfcx.(x)
N = 10
y2 = newI.(x, 0.5, 1.0, N)
plot(x, y1-y2)
#y3 = Taylor.(x[nn], 0.5, 1.0, N)
#plot(x, y1-y2, x[nn], y1[nn]-y3)
#plot(x, y2, x[nn], y3)
xlabel(L"$x$")
title("N = $N")
grid(true)
savefig("error.pdf")


"""
    Q-table is arranged as (Actions × States)

              s1  s2 ⋅ ⋅ ⋅
        a1 ┌ q11 q12 ⋅ ⋅ ⋅ ┐
        a2 │ q21 q22 ⋅ ⋅ ⋅ │
         ⋅ │  ⋅   ⋅  ⋅     │
         ⋅ │  ⋅   ⋅    ⋅   │
         ⋅ └  ·   ⋅      ⋅ ┘

    , which is updated by
        nextA = argmax Q(A, nextS) ∀ A
    for
        Q(prevA, prevS) += α * [prevR + γ * Q(nextA, nextS) - Q(prevA, prevS)]
"""
function update_Qtable!(Q::MatI, prevS::Int, prevA::Int, prevR::Real, nextS::Int, α::Real, γ::Real) # @code_warntype ✓
    nextA = pick_action(Q, nextS)
    @inbounds Q[prevA, prevS] *= 1.0 - α
    @inbounds Q[prevA, prevS] += α * (prevR + γ * Q[nextA, nextS])
    return nothing
end
"""
    The Q-learning agent picks an action of maximum Q-value of current state.

    params:
    -------
    * Q   := Q-table ∈ (Actions × States)
    * Sdx := index of current state
    * Adx := index of picked action (return)
"""
function pick_action(Q::MatI, Sdx::Int) # @code_warntype ✓
    Adx = 0
    tmp = -Inf
    @inbounds for idx in axes(Q, 1)
        _tmp = Q[idx, Sdx]
        if _tmp ≠ _tmp || tmp < _tmp
            tmp = _tmp
            Adx = idx
        end
    end
    return Adx
end
"""
    The Q-learning agent picks an action based on ε-greedy strategy.

    params:
    -------
    * Q   := Q-table ∈ (States × Actions)
    * Sdx := index of current state
    * Adx := index of picked action (return)
    * ε   := greedy factor for exploration, ε ∈ (0, 1], ε(t = 0) = 1
"""
function pick_action(Q::MatI, Sdx::Int, ε::Real) # @code_warntype ✓
    Adx = pick_action(Q, Sdx)
    rand() < ε ? sample(axes(Q, 1), Adx) : Adx
end


function phi(z,a,interval,x)
    z = [z; -z];
    a = [a; -a];

    r = (interval[2] - interval[1]) / 2;
    Emid = interval[1] + r;
    z = z*r .+ Emid;
    a = a*r;

    y = zeros(size(x)...);
    for i = 1:length(z)
        y = y + a[i] ./ (x .- z[i]) + conj(a[i]) ./ (x .- conj(z[i]));
    end
    y
end


module BallTest

using Projections, Test

@testset "ball" begin
y = [-2.5, -1.5, -0.5, 0.0, 0.5, 1.5, 2.5]
n = length(y)
@show sqrt(sum(abs2, y))
@test isapprox(project(Ball(n), y), [-0.597614, -0.358569, -0.119523, 0.0,
    0.119523, 0.358569, 0.597614], rtol=1e-4)
@test all(project(Ball(n, 5.), y) .== y)
end

end