"""
    ExplicitRungeKuttaScheme(
        step_size::Float64,
        alpha::LinearAlgebra.LowerTriangular{Float64},
        beta::LinearAlgebra.LowerTriangular{Float64},
    )

Generic data container for ExplicitRungeKuttaScheme for use in `integrate_time`.

We are interested in numerically integrating ODEs of the form
    ``a'(t) = a(t) Q``

where ``a(t)`` is a vector of coefficients and ``Q`` is a matrix, given an initial condition ``a(0)``.

In this context, Explicit Runge-Kutta methods with ``s`` stages can be written as

``v^{(0)} = a(t),``

``v^{(l)} = \\sum_{k=0}^{l-1} \\alpha_{l k} v^{(k)} + h\\beta_{l k}  v^{(k)} Q,\\, l = 1,...,s,``

``a(t+h) =v^{(s)},``

where ``h`` is the step-size and ``\\alpha_{l k}`` and ``\\beta_{l k}`` are parameters which define the scheme.

# Arguments
-  `step_size::Float64`: the step size of the integration scheme.
-  `alpha::LinearAlgebra.LowerTriangular{Float64}`: coefficients as described above
-  `beta::LinearAlgebra.LowerTriangular{Float64}`: coefficients as described above
"""
struct ExplicitRungeKuttaScheme
    step_size::Float64
    alpha::LinearAlgebra.LowerTriangular{Float64}
    beta::LinearAlgebra.LowerTriangular{Float64}
    function ExplicitRungeKuttaScheme(step_size::Float64,
        alpha::LinearAlgebra.LowerTriangular{Float64},
        beta::LinearAlgebra.LowerTriangular{Float64})

        t1 = (step_size > 0)
        checksquare(alpha)
        l1 = size(alpha)
        l2 = size(beta)
        (!t1)&&throw(DomainError("step_size must be positive"))
        !(l1==l2)&&throw(DimensionMismatch("alpha, beta must have same size"))
        return new(step_size,alpha,beta)
    end
end

"""
    ForwardEuler(step_size::Float64) <: ExplicitRungeKuttaScheme

Defines an Euler integration scheme to be used in `integrate_time`.

# Arguments
-  `step_size::Float64`: the step size of the integration scheme.
"""
ForwardEuler(step_size::Float64) = ExplicitRungeKuttaScheme(
        step_size,
        LinearAlgebra.LowerTriangular([1.0][:,:]),
        LinearAlgebra.LowerTriangular([1.0][:,:])
    )

"""

    Heuns(step_size::Float64) <: ExplicitRungeKuttaScheme

Defines Heuns integration scheme to be used in `integrate_time`.

# Arguments
-  `step_size::Float64`: the step size of the integration scheme.
"""
Heuns(step_size::Float64) = ExplicitRungeKuttaScheme(
        step_size,
        LinearAlgebra.LowerTriangular([1.0 0.0; 0.5 0.5]),
        LinearAlgebra.LowerTriangular([1.0 0.0; 0.0 0.5])
    )

"""

    StableRK3(step_size::Float64) <: ExplicitRungeKuttaScheme

Defines a strong stability preserving Runge-Kutta integration scheme to be used in `integrate_time`.

# Arguments
-  `step_size::Float64`: the step size of the integration scheme.

Hesthaven, J. S. & Warburton, T. (2007), Nodal discontinuous Galerkin methods: algorithms, analysis, and applications, Springer Science & Business Media.
(Section 5.7)

"""
StableRK3(step_size::Float64) = ExplicitRungeKuttaScheme(
        step_size,
        LinearAlgebra.LowerTriangular([1.0 0.0 0.0; 0.75 0.25 0.0; 1/3 0.0 2/3]),
        LinearAlgebra.LowerTriangular([1.0 0.0 0.0; 0.0 0.25 0.0; 0.0 0.0 2/3])
    )

_α = [  1.0                0.0                 0.0                 0.0                 0.0                 ;
        0.44437049406734   0.55562950593266    0.0                 0.0                 0.0                 ;
        0.62010185138540   0.0                 0.37989814861460    0.0                 0.0                 ;
        0.17807995410773   0.0                 0.0                 0.82192004589227    0.0                 ;
        0.00683325884039   0.0                 0.51723167208978    0.12759831133288    0.34833675773694    ]
_β = [  0.39175222700392    0.0                 0.0                 0.0                 0.0                 ;
        0.0                 0.36841059262959    0.0                 0.0                 0.0                 ;
        0.0                 0.0                 0.25189177424738    0.0                 0.0                 ;
        0.0                 0.0                 0.0                 0.54497475021237    0.0                 ;
        0.0                 0.0                 0.0                 0.08460416338212    0.22600748319395    ]
"""
    StableRK4(step_size::Float64) = ExplicitRungeKuttaScheme(
            step_size,
            LinearAlgebra.LowerTriangular(_α),
            LinearAlgebra.LowerTriangular(_β)
        )

Defines a strong stability preserving Runge-Kutta integration scheme to be used in `integrate_time`.

# Arguments
-  `step_size::Float64`: the step size of the integration scheme.

Raymond J. Spiteri and Steven J. Ruuth. A new class of optimal high-order
strong-stability-preserving time discretization methods. SIAM J. Numer. Anal.,
40(2):469-491, 2002.

Hesthaven, J. S. & Warburton, T. (2007), Nodal discontinuous Galerkin methods: algorithms, analysis, and applications, Springer Science & Business Media.
(Section 5.7)
"""
StableRK4(step_size::Float64) = ExplicitRungeKuttaScheme(
        step_size,
        LinearAlgebra.LowerTriangular(_α),
        LinearAlgebra.LowerTriangular(_β)
    )

"""
Given `x0` and `D` apprximate `x0 exp(Dy)`.

    integrate_time(x0::Array{Float64,2}, D::AbstractArray{Float64,2},
        y::Float64, scheme::ExplicitRungeKuttaScheme [; limiter])

# Arguments
- `x0`: An initial vector
- `D`: A square matrix
- `y`: time to integrate up to
- `h`: ExplicitRungeKuttaScheme.
- `limiter`: optional named argument spacifying the slope-limiter to used. The
    limiter function will be applied at every stage and every time-step of the
    ExplicitRungeKuttaScheme.
"""
function integrate_time(x0::Array{Float64,2}, D::AbstractArray{Float64,2},
    y::Float64, scheme::ExplicitRungeKuttaScheme)

    checksquare(D)
    !(size(x0,2)==size(D,1))&&throw(DimensionMismatch("x0 must have length size(D,1)"))

    return _integrate(x0[:],D,y,scheme)
end
function integrate_time(x0::Array{Float64,1}, D::AbstractArray{Float64,2},
    y::Float64, scheme::ExplicitRungeKuttaScheme)

    checksquare(D)

    return _integrate(x0,D,y,scheme)
end

function integrate_time(x0::SFMDistribution, D::AbstractArray{Float64,2},
    y::Float64, scheme::ExplicitRungeKuttaScheme)
    return SFMDistribution(integrate_time(x0.coeffs,D,y,scheme),x0.dq)
end

function integrate_time(x0::SFMDistribution{DGMesh{T}}, D::AbstractArray{Float64,2},
    y::Float64, scheme::ExplicitRungeKuttaScheme; limiter::Limiter=GeneralisedMUSCL) where T

    checksquare(D)

    limiter_params = limiter.generate_params(x0.dq)
    limiter_function = x->limiter.fun(x,limiter_params...)

    return SFMDistribution(_integrate(x0.coeffs,D,y,scheme,limiter_function),x0.dq)
end

"""
    _integrate(x0::Array{Float64,1},
        D::Union{Array{Float64,2},SparseArrays.SparseMatrixCSC{Float64,Int}},
        y::Float64, scheme::ExplicitRungeKuttaScheme)

Use ExplicitRungeKuttaScheme method.
"""
function _integrate(x0::Array{Float64,1}, D::AbstractArray{Float64,2},
    y::Float64, scheme::ExplicitRungeKuttaScheme)

    return _integrate(x0, D, y, scheme, identity)
end

function _integrate(x0::Array{Float64,1}, D::AbstractArray{Float64,2},
    y::Float64, scheme::ExplicitRungeKuttaScheme, limit_function::Function)

    x = limit_function(x0)
    h = scheme.step_size
    l = size(scheme.alpha,1)
    α = scheme.alpha
    βh = scheme.beta*h
    v = Array{Float64,2}(undef,length(x0),l+1)
    vD = Array{Float64,2}(undef,length(x0),l+1)
    for t = h:h:y
        v[:,1] = x
        for i in 1:l
            vD[:,i] = transpose(v[:,i])*D
            initialised = false
            for j in 1:i
                if α[i,j]!=0.0
                    initialised ? (v[:,i+1]+=v[:,j]*α[i,j]) : (v[:,i+1]=v[:,j]*α[i,j]; initialised=true)
                end
                if βh[i,j]!=0.0
                    initialised ? (v[:,i+1]+=vD[:,j]*βh[i,j]) : (v[:,i+1]=vD[:,j]*βh[i,j]; initialised=true)
                end
            end
            v[:,i+1] = limit_function(v[:,i+1])
        end
        x = v[:,end]
    end
    return x
end


"""
     rouwenhorst(μ,ρ,σ,n)
Generates a finite state Markov representation of the AR process defined by the
supplied parameters. Returns the states and transition matrix.
"""
function rouwenhorst(μ, ρ, σ, n)
    n == 1 && (return [μ], [1.0])

    mu_eps = 0
    q = (ρ + 1) / 2
    nu = ((n - 1) / (1 - ρ^2))^(1 / 2) * σ
    P = [q 1 - q; 1 - q q]
    for i = 2:n - 1
        P = q * [P zeros(i, 1);zeros(1, i + 1)] + (1 - q) * [zeros(i, 1) P; zeros(1, i + 1)] + (1 - q) * [zeros(1, i + 1); P zeros(i, 1)] + q * [zeros(1, i + 1); zeros(i, 1) P]
        P[2:i, :] = P[2:i, :] / 2
    end
    x = [linspace(mu_eps / (1 - ρ) .- nu, mu_eps / (1 - ρ) .+ nu, n);] .+ μ
    return x, P
end

"""
     tauchen(μ,ρ,σ,n)
Generates a finite state Markov representation of the AR process defined by the
supplied parameters. Returns the states and transition matrix.
"""
function tauchen(μ, ρ, σ, N, m = 3)
    cdf_normal(x) = 0.5 * erfc(-x / sqrt(2))
    if N == 1
        return [μ], [1.0]
    end
    Z     = zeros(N)
    Zprob = zeros(N, N)
    a     = (1 - ρ) * μ

    Z[N]  = m * sqrt(σ^2 / (1 - ρ^2))
    Z[1]  = -Z[N]
    zstep = (Z[N] - Z[1]) / (N - 1)

    for i = 2:(N - 1)
        Z[i] = Z[1] + zstep * (i - 1)
    end

    Z = Z .+ a / (1 - ρ)

    for j = 1:N
        for k = 1:N
            if k == 1
                Zprob[j,k] = cdf_normal((Z[1] - a - ρ * Z[j] + zstep / 2) / σ)
            elseif k == N
                Zprob[j,k] = 1 - cdf_normal((Z[N] - a - ρ * Z[j] - zstep / 2) / σ)
            else
                Zprob[j,k] = cdf_normal((Z[k] - a - ρ * Z[j] + zstep / 2) / σ) - cdf_normal((Z[k] - a - ρ * Z[j] - zstep / 2) / σ)
            end
        end
    end
    return Z, Zprob
end


### A Pluto.jl notebook ###
# v0.12.6

using Markdown
using InteractiveUtils

# ╔═╡ de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
using Pkg, DrWatson

# ╔═╡ e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Combinatorics, Plots, LaTeXStrings
	Random.seed!(123)
end

# ╔═╡ ac3ceabe-1ee4-11eb-3935-313e96aafd62
md"## Listing 2.5"

# ╔═╡ d212f38c-1ee4-11eb-360a-b7d524890c8f
function isUpperLattice(v)
    for i in 1:Int(length(v)/2)
        sum(v[1:2*i-1]) >= i ? continue : return false
    end
    return true
end

# ╔═╡ 6263427a-1ee5-11eb-12d1-1d576fa18461
begin
	n, N = 5, 10^5

	omega = unique(permutations([zeros(Int,n);ones(Int,n)]))
	A = omega[isUpperLattice.(omega)]
	pA_modelI = length(A)/length(omega)
end

# ╔═╡ 62637f3a-1ee5-11eb-34e2-77e37f188fe5
function randomWalkPath(n)
    x, y = 0, 0
    path = []
    while x<n && y<n
       if rand()<0.5
            x += 1
            push!(path,0)
        else
            y += 1
            push!(path,1)
        end
    end
    append!(path, x<n ? zeros(Int64,n-x) : ones(Int64,n-y))
    return path
end

# ╔═╡ 6263f402-1ee5-11eb-20e1-5d1ebc78f310
begin
	pA_modelIIest = sum([isUpperLattice(randomWalkPath(n)) for _ in 1:N])/N
	("Model I: ",pA_modelI, "\t Model II: ", pA_modelIIest)
end

# ╔═╡ 6276d39e-1ee5-11eb-23cc-2fe5377e6ce9
function plotPath(v,l,c)
    x,y = 0,0
    graphX, graphY = [x], [y]
    for i in v
        if i == 0
            x += 1
        else
            y += 1
        end
        push!(graphX,x), push!(graphY,y)
    end
    plot!(graphX, graphY,
            la=0.8, lw=2, label=l, c=c, ratio=:equal, legend=:topleft,
            xlims=(0,n), ylims=(0,n),
            xlabel=L"East\rightarrow", ylabel=L"North\rightarrow")
end

# ╔═╡ 62781fb2-1ee5-11eb-0ba1-035fa45106f5
begin
	plot()
	plotPath(rand(A), "Upper lattice path", :blue)
	plotPath(rand(setdiff(omega,A)), "Non-upper lattice path", :red)
	plot!([0, n], [0,n], ls=:dash, c=:black, label="")
end

# ╔═╡ 6285dea2-1ee5-11eb-2b6a-f9906c0d3351
md"## End of listing 2.5"

# ╔═╡ Cell order:
# ╟─ac3ceabe-1ee4-11eb-3935-313e96aafd62
# ╠═de7a5e8a-1ee4-11eb-2bbf-0f0eb3e489ec
# ╠═e4eeb13a-1ee4-11eb-0113-91f9a9c3b659
# ╠═d212f38c-1ee4-11eb-360a-b7d524890c8f
# ╠═6263427a-1ee5-11eb-12d1-1d576fa18461
# ╠═62637f3a-1ee5-11eb-34e2-77e37f188fe5
# ╠═6263f402-1ee5-11eb-20e1-5d1ebc78f310
# ╠═6276d39e-1ee5-11eb-23cc-2fe5377e6ce9
# ╠═62781fb2-1ee5-11eb-0ba1-035fa45106f5
# ╟─6285dea2-1ee5-11eb-2b6a-f9906c0d3351


module ProjSplx

export projsplx!, projsplx, projnorm1,
       projnorm1!, projnorm1t, proxinf!

"""
Projection onto the unit simplex {x | sum(x) = τ, x ≥ 0}.

   projsplx!(b, τ)

In-place variant of `projsplx`.
"""
function projsplx!{T}(b::Vector{T}, τ::T)

    n = length(b)
    bget = false

    idx = sortperm(b, rev=true)
    tsum = zero(T)

    @inbounds for i = 1:n-1
        tsum += b[idx[i]]
        tmax = (tsum - τ)/i
        if tmax ≥ b[idx[i+1]]
            bget = true
            break
        end
    end

    if !bget
        tmax = (tsum + b[idx[n]] - τ) / n
    end

    @inbounds for i = 1:n
        b[i] = max(b[i] - tmax, 0)
    end

end

"""
Projection onto the weighted simplex.

    projsplx!(b, c, τ)

This projects b (in-place) onto the unit simplex weighted by c.
"""
function projsplx!{T}(b::Vector{T}, c::Vector{T}, τ::T)

    n = length(b)
    bget = false

    @assert length(b) == length(c) "lengths must match"
    @assert minimum(c) > 0 "c is not positive."

    idx = sortperm(b./c, rev=true)
    tsum = csum = zero(T)

    @inbounds for i = 1:n-1
        j = idx[i]
        tsum += b[j]*c[j]
        csum += c[j]*c[j]
        tmax = (tsum - τ) / csum
        if tmax >= b[idx[i+1]] / c[idx[i+1]]
            bget = true
            break
        end
    end

    if !bget
        p = idx[n]
        tsum += b[p]*c[p]
        csum += c[p]*c[p]
        tmax = (tsum - τ) / csum
    end

    for i = 1:n
        @inbounds b[i] = max(b[i] - c[i]*tmax, 0)
    end

    return

end

"""
Projection onto the simplex.

  projsplx(b, τ) -> x

Variant of `projsplx`.
"""
function projsplx(b::Vector, τ)
    x = copy(b)
    projsplx!(x, τ)
    return x
end

"""
Projection onto the weighted simplex.

    projsplx(b, c, τ) -> x

Variant of `projsplx!`.
"""
function projsplx(b::Vector, c::Vector, τ)
    x = copy(b)
    projsplx!(x, c, τ)
    return x
end

# s = sign_abs!(x) returns
#     s[i] = true  if x[i] > 0
#     s[i] = false otherwise
# and x = abs(x).
function sign_abs!(x::Vector)
  n = length(x)
  s = Array{Bool}(n)
  @inbounds for i=1:n
    s[i] = x[i] ≥ 0
    x[i] = abs(x[i])
  end
  return s
end

# set_sign!(x, s) sets the sign of x based on s.
function set_sign!(x::Vector, s::Vector{Bool})
  n = length(x)
  @inbounds for i=1:n
      x[i] = s[i] ? x[i] : -x[i]
  end
end

"""
Projection onto the 1-norm ball

    {x | ||x||₁ ≤ τ }.

In-place variant of `projnorm1`.
"""
function projnorm1!(b::Vector, τ::Real)
    norm(b,1) > τ || return
    s = sign_abs!(b)
    projsplx!(b, τ)
    set_sign!(b, s)
end

"""
Projection onto the weighted 1-norm ball

    {x | ||diag(c)x||₁ ≤ τ }, c > 0.

In-place variant of `projnorm1`.
"""
function projnorm1!(b::Vector, c::Vector, τ::Real)
    norm(b,1) > τ || return
    s = sign_abs!(b)
    projsplx!(b, c, τ)
    set_sign!(b, s)
end

"""
Projection onto the weighted 1-norm ball.

Variant of `projnorm1!`.
"""
function projnorm1(b::Vector, c::Vector, τ::Real)
    x = copy(b)
    projnorm1!(x, c, τ)
    return x
end

"""
Projection onto the 1-norm ball.

Variant of `projnorm1!`.
"""
function projnorm1(b::Vector, τ::Real)
    x = copy(b)
    projnorm1!(x, τ)
    return x
end

"""
Proximal map of the scaled infinity norm.
    prox_inf(x,λ) = x - proj(x | λ𝔹₁)
     env_inf(x,λ) = (1/2λ)||x||² - (1/2λ)dist²(x | λ𝔹₁)
Modifies `x` in place; returns the envelope.
"""
function proxinf!(x::Vector, λ::Real)
  λ == 0 && return norm(x, Inf)
  nrmx2 = dot(x,x)
  xp = projnorm1(x, λ)
  BLAS.axpy!(-1., xp, x) # x <- x - xp
  return nrmx2/(2λ) - dot(x,x)/(2λ)
end

"""
Proximal map of the scaled infinity norm.
Return variant of `proxinf!`
"""
function proxinf(x::Vector, λ::Real)
    z = copy(x)
    proxinf!(z, λ)
    return z
end

end # module


import Base.:+
import Base.:-
import Base.:*

export ℝ³
export cross

"""
    Represents a point in ℝ³.

field: a.
"""
struct ℝ³ <: VectorSpace
    a::Array{Float64}
    ℝ³(x₁::Real, x₂::Real, x₃::Real) = new(float.([x₁; x₂; x₃]))
    ℝ³(a::Array) = begin
        @assert(length(a) == 3, "The input vector must contain exactly three elements.")
        ℝ³(a...)
    end
end

## Unary Operators ##

+(r::ℝ³) = r
-(r::ℝ³) = ℝ³(-vec(r))

## Binary Operators ##

+(r1::ℝ³, r2::ℝ³) = ℝ³(vec(r1) + vec(r2))
-(r1::ℝ³, r2::ℝ³) = ℝ³(vec(r1) - vec(r2))
*(r::ℝ³, λ::Real) = ℝ³(λ .* vec(r))
*(λ::Real, r::ℝ³) = r * λ

"""
    cross(r1, r2)

Perform a cross product with the given vectors `r1` and `r2`.
"""
cross(r1::ℝ³, r2::ℝ³) = begin
    M = transpose(reshape([vec(r1); vec(r1); vec(r2)], :, length(r1)))
    M = convert(Array{Float64}, M)
    ℝ³(map(x -> cofactor(M, 1, x), 1:length(r1)))
end

Base.isapprox(a::Array{ℝ³},
              b::Array{ℝ³};
              atol::Float64 = TOLERANCE) = begin
    for (elementa, elementb) in zip(a, b)
        if isapprox(elementa, elementb, atol = atol) == false
            return false
        end
    end
    return true
end


function khayyam_triangle(lentingle)
    output = []
    for i in 1:lentingle
        push!(output, [])
        for j in 1:i
            if j == 1
                append!(output[i], 1)
            elseif j == i
                append!(output[i], 1)
            else j != 1 && j != i
                add = output[i - 1][j] + output[i - 1][j - 1]
                append!(output[i], add)
            end
        end
    end
    return output
end


using FractionalSystems, ControlSystems
using Test

@testset "Test FOTF" begin
    a = fotf([1, 2], [1, 2], [1, 2], [1, 2])

    @test isa(a, FOTF)
    @test fotfdata(a) == ([1, 2], [1, 2], [1, 2], [1, 2], 0)
    @test iszero(a) == false
    @test fotf("s") == fotf([1], [1], [1], [0], 0)
end

G = fotf([1, 2, 3], [0.1, 0.2, 0.34], [4, 5, 6], [0.65, 0.67, 0.78])
num = zeros(34+1)
num[34+1] = 3
num[20+1] = 2
num[10+1] = 1
num = num[end:-1:1]

den = zeros(78+1)
den[78+1] = 6
den[67+1] = 5
den[65+1] = 4
den = den[end:-1:1]

@testset "Test base_order" begin
    @test base_order(G) == 0.01
end

@testset "Test fotf2cotf" begin
    @test fotf2cotf(G) == tf(num, den)
end

@testset "Test Polyuniq" begin
    @test polyuniq([2, 2], [4, 4], 0.0001) == ([4], [4])
end
#=
@testset "Test simplify" begin
    g = fotf([2, 2], [4, 4], [6, 6], [8, 8])
    sg = fotf( 4, 0, 12, 4)
    @test fotfdata(simplify(g)) == fotfdata(sg)
end
=#

@testset "Test basic operations" begin
    a=fotf([5, 6], [7, 8], [1, 2], [3, 4])
    b=fotf([1, 2], [1, 2], [1, 2], [1, 2])

    aplusb = fotf([1, 5, 2, 6], [5, 4, 1, 0], [1, 2], [0, 1])
    @test a+b == aplusb
end

#=
@testset "Test inv operation" begin

end
=#

@testset "Test freqresp" begin
    a = fotf([1, 2], [1, 2], [1, 2], [1, 2])
    @test freqresp([1; 2; 3], a) == [1.0; 1.0; 1.0]
end

@testset "Test polyshow" begin
    @test polyshow([1, 2], [1, 2]) == "2s^{2}+s"
    @test polyshow([1, -2], [0, 2.2]) == "-2s^{2.2}+1"
end


@testset "thermalization.jl" begin

    @testset "_weights_canonical" begin
        # check weighting
        range = collect(1:15)
        @test XXZNumerics._weights_canonical(log.(range), -1) ≈ range./sum(range)
        @test XXZNumerics._weights_canonical(log.(range), -2) ≈ range.^2 ./ sum(range .^ 2)

        # check extreme cases
        @test .!(any(isnan, XXZNumerics._weights_canonical(exp.(range), 100)))
        @test XXZNumerics._weights_canonical(exp.(range), 0) ≈ ones(length(range))/length(range)
        @test XXZNumerics._weights_canonical(exp.(range), 1) ≈ XXZNumerics._weights_canonical(sort!(-exp.(range)), -1)[end:-1:1]
    end

    @testset "findβ" begin
        @test findβ([1,2,3,4], 2.5) ≈ 0
        @test findβ([1,2,3,4], 3) < 0
        @test findβ([1,2,3,4], 2) > 0
        @test findβ([1,2,3,4], 2.5-0.7) ≈ -findβ([1,2,3,4], 2.5+0.7) # should be symmetric around the average energy as spectrum is symmetric
        # can't think of other good tests...
    end

    @testset "canonical_ensemble" begin
        l = 10
        states = sort!(rand(l)*10 .- 5)
        @test canonical_ensemble(states, states[1]) ≈ vcat(1, zeros(l-1))
        @test canonical_ensemble(states, states[end]) ≈ vcat(zeros(l-1), 1)
        @test canonical_ensemble(states, mean(states)) ≈ ones(l)/l
    end

    @testset "microcanonical_ensemble" begin
        @test microcanonical_ensemble([-2, -1, 0.5, 1, 2], 0, 1.1) ≈ [0,1,1,1,0]/3
        @test microcanonical_ensemble([-2, -1, 0.5, 1, 2], 0, 0.9) ≈ [0,0,1,0,0]
        @test microcanonical_ensemble([-2, -1, 0.5, 1, 2], 0, 0) ≈ [0,0,1,0,0] # choose closest state
        @test microcanonical_ensemble([-2, -1, 0.5, 1, 2], -0.6, 0) ≈ [0,1,0,0,0]

        @test microcanonical_ensemble([-40, -1, 0.5, 1, 40], 0.7) ≈ [0,0,1,1,0]/2 # default arg -> ΔE = 80*0.5% = 0.4
    end
end


# A test problem coming from LSQR.

using LinearOperators

function lstp(nrow :: Int, ncol :: Int, ndupl :: Int, npower :: Int, λ :: Real, x :: Array)

  # LSTP  generates a sparse least-squares test problem of the form
  #
  #   minimize ‖ [A ] x - [b] ‖
  #            ‖ [λI]     [0] ‖
  #
  # having a specified solution x.  The matrix A is constructed
  # in the form A = HY*D*HZ, where D is an nrow by ncol diagonal matrix,
  # and HY and HZ are Householder transformations.

  @assert(nrow >= ncol);

  # Construct two unit vectors for the Householder transformations.
  # fourpi = 4π;
  fourpi = 4 * 3.141592;    # This is the approximation used in Mike's original subroutine.
  α = fourpi / nrow;        # 4π / nrow;
  β = fourpi / ncol;        # 4π / ncol;
  hy = map(sin, [1:nrow;] * α)
  hz = map(cos, [1:ncol;] * β)

  α = norm(hy); hy /= α; HY = opHouseholder(hy);  # HY is nrow x nrow.
  β = norm(hz); hz /= β; HZ = opHouseholder(hz);  # HZ is ncol x ncol.

  # Set the diagonal matrix D containing the singular values of A.
  d = (div.(([0:ncol-1;] .+ ndupl), ndupl) * ndupl / ncol).^npower;  # Integer div!
  D = opDiagonal(nrow, ncol, d);
  A = HY * D * HZ;

  Acond = abs(d[ncol] / d[1]);

  # Compute residual vector.
  r = zeros(nrow);
  r[1:ncol]  = HZ * x ./ d;
  t = 1.0;
  for i = ncol + 1 : nrow
    j = i - ncol;
    r[i] = t * j / nrow;
    t = -t;
  end
  r = HY * r;

  # Compute right-hand side b = r + Ax.
  rnorm = norm(r);
  b = r + A * x;

  return (b, A, D, HY, HZ, Acond, rnorm)
end

function test(nrow, ncol, ndupl, npower, damp)

  x = ncol .- [1:ncol;];  # Desired solution.
  return lstp(nrow, ncol, ndupl, npower, damp, x);
end

function testall()
  damp = 0;  # Must be zero for this problem to be consistent.
  test(40, 40, 4, 1, damp);
  test(40, 40, 4, 2, damp);
  test(40, 40, 4, 3, damp);
  test(40, 40, 4, 4, damp);
end


get_x(vec::SA.SVector) = vec[1]
get_y(vec::SA.SVector) = vec[2]

rotate_plus_90(vec::SA.SVector{2}) = typeof(vec)(-vec[2], vec[1])
rotate_minus_90(vec::SA.SVector{2}) = typeof(vec)(vec[2], -vec[1])
rotate_180(vec::SA.SVector{2}) = -vec

rotate(x::T, y::T, c::T, s::T) where {T} = SA.SVector(c * x - s * y, s * x + c * y)
rotate(vec::SA.SVector{2, T}, theta::T) where {T} = rotate(vec[1], vec[2], cos(theta), sin(theta))
rotate(vec::SA.SVector{2, T}, dir::SA.SVector{2, T}) where {T} = rotate(vec[1], vec[2], dir[1], dir[2])

get_relative_direction(dir1::SA.SVector{2, T}, dir2::SA.SVector{2, T}) where {T} = typeof(dir1)(dir2[1] * dir1[1] + dir2[2] * dir1[2], dir2[2] * dir1[1] - dir2[1] * dir1[2])
invert_relative_direction(dir::SA.SVector{2}) = typeof(dir)(dir[1], -dir[2])

function invert(pos::SA.SVector{2, T}, dir::SA.SVector{2, T}) where {T}
    inv_dir = invert_relative_direction(dir)
    inv_pos = -rotate(pos, inv_dir)
    return inv_pos, inv_dir
end


module MagicSquares

export MagicSquare, ismagical

immutable MagicSquare{T<:Integer}
    n::Integer
    data::Array{T}

    MagicSquare(data::Array{T}) = begin
        n = convert(Int, √length(data))
        new(n, data)
    end
end

MagicSquare{T <: Integer}(data::Array{T}) = MagicSquare{T}(data)
MagicSquare{T <: Integer}(n::T) = MagicSquare{T}(zeros(T, n*n))

Base.shuffle!(m::MagicSquare) = shuffle!(m.data)
Base.show(io::IO, m::MagicSquare) = begin
    if(ismagical(m))
        println("Magic square with sum=$(sum(m.data[1:m.n]))")
    else
        println("Non-magic square.")
    arr = reshape(m.data, m.n, m.n)
    for i ∈ 1:m.n
        println(io, arr[i, :])
    end
end

function ismagical(m::MagicSquare)
    n = m.n
    s = sum(m.data[1:n])
    data1 = reshape(m.data, n, n)
    data2 = data1'
    diagsum1 = s
    diagsum2 = s
    for i ∈ 1:n
        if(s != sum(data1[i, :]) || s != sum(data2[i, :]))
            return false
        end
        diagsum1 -= data1[i, i]
        diagsum2 -= data1[i, n-i+1]
    end
    if diagsum1 == 0 && diagsum2 == 0
        return true
    end
    return false
end

end # module


# Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export nlogML

"""
`nlogML(Nt,a,A,y)`

Computes the negative log Maximum Likelihood (ML) normalized by the sequence length (``\\log(P(\\mathbf{X}))/N_t``) of a Hidden Markov Model (HMM) with `Ns` states, transition matrix `A` (a `Ns` × `Ns` matrix), initial probabilities `a` (`Ns`-vector) and observation probabilities `y` (a `Ns`×`Nt` matrix) for `Nt` observations. All `Array`s must have the same element type.

`nlogML(Nt,t2tr,A,y)`

Returns the negative log Maximum Likelihood (ML) normalized by the sequence length (``\\log(P(\\mathbf{X}))/N_t``) with a constrained path (see [`forward`](@ref)).
"""
nlogML(Nt::Integer,c) = -sum(log.(c))/Nt
nlogML(Nt::Array,cs) = sum( [-sum(log.(c))/Nt[i] for (i,c) in enumerate(cs)])

function nlogML(Nt::Integer,a,A,y)
  alpha, c = forward(Nt,a,A,y)
  return nlogML(Nt,c)
end

function nlogML(Nt::Array,a,A,y)
  alphas, cs = forward(Nt,a,A,y)
  return nlogML(Nt,cs)
end

export nlogMLlog
"""
`nlogMLlog(Nt,A,a,y)`

Same as [`nlogML`](@ref) but expects `A`, `a` and `y` to be in the log-domain.
"""
function nlogMLlog(Nt,a,A,y)
  alpha, logML = logforward(Nt,a,A,y)
  return -sum(logML)
end


"""
Calculates the reference variable for ``μ`` when the elements in vector ``x`` comes from a normal distribution.
See also: [`sample_mean`](@ref), [`sample_variance`](@ref).
"""
function normal_reference_mu_unknown_variance(μ::AbstractFloat, x::Vector)
    n = length(x)
    x̅ = sample_mean(x)
    s = √(sample_variance(x))
    T = (x̅ - μ) / (s / √(n))
    return T
end

"""
Calculates the reference variable for ``μ`` when the elements in vector ``x`` comes from a normal distribution.
This time the variance ``σ²`` has to be known for the distribution.
See also: [`sample_mean`](@ref).
"""
function normal_reference_mu_known_variance(μ::AbstractFloat, σ²::AbstractFloat, x::Vector)
    n = length(x)
    x̅ = sample_mean(x)
    σ = √(σ²)
    T = (x̅ - μ) / (σ / √(n))
    return T
end

"""
Reference variable for ``σ²`` when the elements in vector ``x`` comes from a normal distribution.
See also: [`sample_variance`](@ref).
"""
function normal_reference_sigma(σ²::AbstractFloat, x::Vector)
    n = length(x)
    s² = sample_variance(x)
    T = (n - 1) * s² / σ²
    return T
end

"""
Calculates the reference variable for tests of the difference between two means.
This variable is used when the variance is known and different for the two samples.
See also: [`sample_mean`](@ref).
"""
function normal_reference_two_sample_mu_known_variance(μᵪ::AbstractFloat, μᵧ::AbstractFloat, σ²ᵪ::AbstractFloat, σ²ᵧ::AbstractFloat, x::Vector, y::Vector)
    nᵪ = length(x)
    nᵧ = length(y)
    x̅ = sample_mean(x)
    y̅ = sample_mean(x)
    T = ((x̅ - y̅) - (μᵪ - μᵧ)) / √(σ²ᵪ / nᵪ + σ²ᵧ / nᵧ)
    return T
end

"""
Calculates the reference variable for tests of the difference between two means.
This variable is used when the variance is unknown but equal for the two samples.
See also: [`sample_mean`](@ref), [`two_sample_variance`](@ref).
"""
function normal_reference_two_sample_mu_unknown_equal_variance(μᵪ::AbstractFloat, μᵧ::AbstractFloat, x::Vector, y::Vector)
    nᵪ = length(x)
    nᵧ = length(y)
    x̅ = sample_mean(x)
    y̅ = sample_mean(x)
    s = √(two_sample_variance(x, y))
    T = ((x̅ - y̅) - (μᵪ - μᵧ)) / (s * √(1 / nᵪ + 1 / nᵧ))
    return T
end

"""
Calculates the reference variable for tests of the difference between two means.
This variable is used when the variance is unknown but not necessarily equal for the two samples.
See also: [`sample_mean`](@ref), [`sample_variance`](@ref), [`two_sample_degrees_of_freedom_unknown_variance`](@ref).
"""
function normal_reference_two_sample_mu_unknown_variance(μᵪ::AbstractFloat, μᵧ::AbstractFloat, x::Vector, y::Vector)
    nᵪ = length(x)
    nᵧ = length(y)
    x̅ = sample_mean(x)
    y̅ = sample_mean(x)
    s²ᵪ = sample_variance(x)
    s²ᵧ = sample_variance(y)
    T = ((x̅ - y̅) - (μᵪ - μᵧ)) / √(s²ᵪ / nᵪ + s²ᵧ / nᵧ)
    return T
end

"""
Calculates the degrees of freedom of the t-distribution that the reference variable above is from.
See also: [`sample_variance`](@ref), [`normal_reference_two_sample_mu_unknown_variance`](@ref).
"""
function two_sample_degrees_of_freedom_unknown_variance(x::Vector, y::Vector)
    nᵪ = length(x)
    nᵧ = length(y)
    s²ᵪ = sample_variance(x)
    s²ᵧ = sample_variance(y)
    first = 1 / (nᵪ - 1)
    second = (nᵧ * s²ᵪ)^2 / (nᵧ * s²ᵪ + nᵪ * s²ᵧ)^2
    third = 1 / (nᵧ - 1)
    fourth = (nᵪ * s²ᵧ)^2 / (nᵧ * s²ᵪ + nᵪ * s²ᵧ)^2
    f = 1 / (first * second + third * fourth)
    return f
end

"""
Calculates the interval for the estimated mean difference of two paired samples.
See also: [`sample_mean`](@ref), [`sample_variance`](@ref).
"""
function normal_interval_two_sample_paired(x::Vector, y::Vector, α::AbstractFloat)
    n = length(x)
    z = y .- x
    z̅ = sample_mean(z)
    dist = tDistribution(n-1)
    t = quantile_finder(dist, α)
    s = √(sample_variance(z))
    I = [z̅ - t * s / √(n), z̅ + t * s / √(n)]
    return I
end

function normal_interval_mu_unknown_variance(μ::AbstractFloat, x::Vector, α::AbstractFloat)

end

function normal_interval_mu_known_variance(μ::AbstractFloat, σ²::AbstractFloat, x::Vector, α::AbstractFloat)

end

function normal_interval_sigma(x::Vector, α::AbstractFloat)

end

function normal_interval_two_sample_mu_known_variance(μᵪ::AbstractFloat, μᵧ::AbstractFloat, σ²ᵪ::AbstractFloat, σ²ᵧ::AbstractFloat, x::Vector, y::Vector, α::AbstractFloat)

end

function normal_interval_two_sample_mu_unknown_equal_variance(μᵪ::AbstractFloat, μᵧ::AbstractFloat, x::Vector, y::Vector, α::AbstractFloat)

end

function exp_reference_mu(μ::AbstractFloat, x::Vector)
    n = length(n)
    x̅ = sample_mean(x)
    T = (x̅ - μ) / (μ / √(n))
    return T
end

function exp_interval_mu(μ::AbstractFloat, x::Vector, α::AbstractFloat)

end

function binomial_reference_p()

end

function binomial_interval_p()

end

function hypergeometric_reference_p()

end

function hypergeometric_interval_p()

end

function binomial_interval_comparison_of_proportions()

end

function poisson_reference_mu()

end

function poisson_interval_mu()

end

function test()

end


abstract type Basis end

struct Oscillator <: Basis
    n::Int64
    ω::Float64
end

Oscillator(n::Int64) = Oscillator(n,1.0)
index(b::Oscillator) = 1:b.n
spectrum(b::Oscillator) = (0:b.n-1)*b.ω
qnumbers(b::Oscillator) = 0:b.n-1

function hermite(x,n)
    @assert n >= 1
    h1 = exp(-x^2 / 2) * (1 / π)^(1 / 4)
    n == 1 && (return h1)
    h2 = sqrt(2) * x * h1
    n == 2 && (return h2)
    h3 = 0.0
    for j ∈ 3:n
        h3 = sqrt(2/(j - 1)) * x * h2 - sqrt((j-2)/(j-1)) * h1
        h1 = h2; h2 = h3
    end
    n > 2 && (return h3)
end
hermite(x,n,ω) = hermite(√ω * x,n) * ω ^ ( 1 / 4 )

function (b::Oscillator)(x::Array{Float64,1})
    @unpack ω,n = b
    T = x * ones(1, n) |> zero
    T[:,1] = @. hermite(x,1,ω)
    n >= 2 && (T[:,2] = @. sqrt(2 * ω) * x * T[:,1])
    for j ∈ index(b)[3:end]
        m = qnumbers(b)[j]
        T[:,j] = @. sqrt(2 / m) * (√ω * x) * T[:,j-1] - sqrt((m-1) / m) * T[:,j-2]
    end
    return T
end

(b::Oscillator)(x::LinRange{Float64}) = b(x |> collect)
(b::Oscillator)(x::Float64,n::Int64) = hermite(x,n,b.ω)

function hermite_polar(x,n,ω)
    p = Progress(n,.1,"Precomputing...",20)   # minimum update interval: 1 second
    T = zeros(size(x)...,n)
    T[:,:,1] = @. hermite(x,1,ω)
    next!(p)
    n >= 2 && (T[:,:,2] = @. sqrt(2 * ω) * x * T[:,:,1])
    next!(p)
    for j in 3:n
        T[:,:,j] = @. sqrt(2 / (j-1)) * (√ω * x) * T[:,:,j-1] - sqrt((j-2) / (j-1)) * T[:,:,j-2]
    next!(p)
    end
    return T
end
hermite_polar(x,n) = hermite_polar(x,n,1.)

# filter cartesian data
function filterH(psi,x,N,ω)
    p = Progress(N,.1,"Filtering...",20)
    H = Oscillator(N,ω)(x)
    T = inv(H'*H)
    F̂ = T*H'*psi*H*T
    psiF = zero(psi)
        for m = 1:N
            next!(p)
            for n = 1:(N+1-m)
            @. psiF += H[:,m]*F̂[m,n]*H[:,n]'
            end
        end
    return psiF
end

# covert to polar coords
function slowpolar(psi,x,N,ω)
    H = Oscillator(N,ω)(x)
    T = inv(H'*H)
    F̂ = T*H'*psi*H*T

    Nx = length(x)
    θ = LinRange(0,2*pi,2*Nx)
    r = LinRange(0,last(x),Nx/2 |> Int)'

    psiP = zero(θ*r)
    for m = 1:N, n = 1:(N + 1 - m)
        @. psiP += hermite(r*cos(θ),m) * F̂[m,n] * hermite(r*sin(θ),n)
    end
    return psiP
end

function init_polar(x,N,ω)
    Nx = length(x)
    θ = LinRange(0,2*pi,2*Nx)
    r = LinRange(0,last(x),Nx/2 |> Int)'
    hx = hermite_polar((@. r*cos(θ)),N,ω)
    hy = hermite_polar((@. r*sin(θ)),N,ω)
    return hx,hy
end

function polar(psi,x,ω,Hx,Hy)
    N = size(Hx)[3]
    p = Progress(N,1,"Polar transform...",20)
    H = Oscillator(N,ω)(x)
    T = inv(H'*H)
    F̂ = T*H'*psi*H*T

    Nx = length(x)
    Nθ = 2*Nx
    Nr = Nx/2 |> Int
    θ = LinRange(0,2*pi,Nθ+1)[1:Nθ]
    r = LinRange(0,last(x),Nr)'

    psiFP = zeros(eltype(psi),Nθ,Nr)
    for m = 1:N
        next!(p)
        for n = 1:(N + 1 - m)
        hx = @view Hx[:,:,m]
        hy = @view Hy[:,:,n]
        @. psiFP += hx * F̂[m,n] * hy
    end
    end
    return psiFP
end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, Plots
	Random.seed!(0)
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing5.04"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
n, N, alpha = 3, 10^7, 0.1

# ╔═╡ 0a6fcc80-3a99-11eb-3de9-13429e2152ee
begin
	myT(nObs) = rand(Normal())/sqrt(rand(Chisq(nObs-1))/(nObs-1))
	mcQuantile = quantile([myT(n) for _ in 1:N],alpha)
	analyticQuantile = quantile(TDist(n-1),alpha)
end

# ╔═╡ 0a990730-3a99-11eb-0e9d-439c6e34d87c
Text("Quantile from Monte Carlo: $(mcQuantile)")

# ╔═╡ 0a998098-3a99-11eb-209b-2f467a34d644
Text("Analytic qunatile: $(analyticQuantile)")

# ╔═╡ 0a9fd40c-3a99-11eb-387b-05b58e298e6e
begin
	xGrid = -5:0.1:5
	plot(xGrid, pdf.(Normal(), xGrid), c=:black, label="Normal Distribution")
	scatter!(xGrid, pdf.(TDist(1) ,xGrid),
		c=:blue, msw=0, label="DOF = 1")
	scatter!(xGrid, pdf.(TDist(3), xGrid),
		c=:red, msw=0, label="DOF = 3")
	scatter!(xGrid, pdf.(TDist(100),xGrid),
		c=:green, msw=0, label="DOF = 100",
		xlims=(-4,4), ylims=(0,0.5), xlabel="X", ylabel="Density")
end

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing5.04"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═0a6fcc80-3a99-11eb-3de9-13429e2152ee
# ╠═0a990730-3a99-11eb-0e9d-439c6e34d87c
# ╠═0a998098-3a99-11eb-209b-2f467a34d644
# ╠═0a9fd40c-3a99-11eb-387b-05b58e298e6e
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12


"""
    PoisLik <: Likelihood

Poisson likelihood
```math
p(yᵢ = k | fᵢ) = θᵏ\\exp(-θ)/k!
```
for ``k ∈ N₀``, where ``θ = \\exp(f)`` and ``f`` is the latent Gaussian process.
"""
struct PoisLik <: Likelihood end

#log of probability density
function log_dens(poisson::PoisLik, f::AbstractVector, y::Vector{Int})
    #where we exponentiate for positivity f = exp(fi)
    return y.*f - exp.(f) - lgamma.(1.0 .+ y)
end

#derivative of pdf wrt latent function
function dlog_dens_df(poisson::PoisLik, f::AbstractVector, y::Vector{Int})
    return y - exp.(f)
end

#mean and variance under likelihood
mean_lik(poisson::PoisLik, f::AbstractVector) = exp.(f)
var_lik(poisson::PoisLik, f::AbstractVector) = exp.(f)

get_params(poisson::PoisLik) = []
num_params(poisson::PoisLik) = 0

function var_exp(ll::PoisLik, y::AbstractArray, m::AbstractArray, V::AbstractMatrix)
    tot = 0
    V_diag = diag(V)
    for (a, b, c) in zip(y, m, V_diag)
        tot +=  a*b - exp(b + c/2) - log(factorial(convert(Int64, a))) # convert to lgamma(y+1)
    end
    return tot
end

function var_exp(ll::PoisLik, y::AbstractArray, m::AbstractArray, V::AbstractArray)
    tot = 0
    for (a, b, c) in zip(y, m, V)
        tot +=  a*b - exp(b + c/2) - log(factorial(convert(Int64, a))) # convert to lgamma(y+1)
    end
    return tot
end

function var_exp(ll::PoisLik, y::Number, m::Number, V::Number)
    return y*m - exp(m + V/2) - log(factorial(convert(Int64, y))) # convert to lgamma(y+1)
end

function dv_var_exp(ll::PoisLik, y::Number, m::Number, V::Number)
    return gradient(x -> var_exp(ll, y, m, x), V)[1]
end


@testset "Newton" begin
    for n in [1,3,10]

A = rand(n,n)

r = polarfact(A, alg = :newton);

# Test unitary matrix U

U = r.U
H = r.H

@test U'*U ≈ Matrix(I,n,n) atol=1e-7

# Test Hermitian positive semifefinite matrix H

@test ishermitian(H)

for i in eigvals(H)
    @test i >= 0.
end

@test A ≈ U*H

    end
end


module Integrators

export rk4

using LinearAlgebra

function rk4(f::Function, y0::Array{Float64, 1}, t0::Float64,
             t1::Float64, h::Float64; inplace::Bool=true)
    y = y0
    n = round(Int, (t1 - t0)/h)
    t = t0
    if ~inplace
        hist = zeros(n, length(y0))
    end
    for i in 1:n
        k1 = h * f(t, y)
        k2 = h * f(t + 0.5*h, y + 0.5*k1)
        k3 = h * f(t + 0.5*h, y + 0.5*k2)
        k4 = h * f(t + h, y + k3)
        y = y + (k1 + 2*k2 + 2*k3 + k4)/6
        if ~inplace
            hist[i, :] = y
        end
        t = t0 + i*h
    end
    if ~inplace
        return hist
    else
        return y
    end
end

function rk4_prop(jac::Function, y0::Array{Float64, 1}, t0::Float64,
                  t1::Float64, h::Float64)
    y = y0
    n = round(Int, (t1 - t0)/h)
    t = t0
    A_acc = I
    for i in 1:n
        J = jac(t, y)
        A = I + h*J + h^2/factorial(2)*J^2 + h^3/factorial(3)*J^3 + h^4/factorial(4)*J^4
        A_acc = A*A_acc
        y = A*y
        t = t0 + i*h
    end
    return A_acc
end

end


# # Automatic differentiation through spherical harmonic transforms
# This example finds a positive value of $\lambda$ in:
# ```math
# f(r) = \sin[\lambda (k\cdot r)],
# ```
# for some $k,r\in\mathbb{S}^2$ such that $\int_{\mathbb{S}^2} f^2 {\rm\,d}\Omega = 1$.
# We do this by using derivative information through:
# ```math
# \dfrac{\partial f}{\partial \lambda} = (k\cdot r) \cos[\lambda (k\cdot r)].
# ```

using FastTransforms, LinearAlgebra

# The colatitudinal grid (mod $\pi$):
N = 15
θ = (0.5:N-0.5)/N

# The longitudinal grid (mod $\pi$):
M = 2*N-1
φ = (0:M-1)*2/M

# We precompute a spherical harmonic--Fourier plan:
P = plan_sph2fourier(Float64, N)

# And an FFTW Fourier analysis plan on $\mathbb{S}^2$:
PA = plan_sph_analysis(Float64, N, M)

# Our choice of $k$ and angular parametrization of $r$:
k = [2/7, 3/7, 6/7]
r = (θ,φ) -> [sinpi(θ)*cospi(φ), sinpi(θ)*sinpi(φ), cospi(θ)]

# Our initial guess for $\lambda$:
λ = 1.0

# Then we run Newton iteration and grab an espresso:
for _ in 1:7
    F = [sin(λ*(k⋅r(θ,φ))) for θ in θ, φ in φ]
    Fλ = [(k⋅r(θ,φ))*cos(λ*(k⋅r(θ,φ))) for θ in θ, φ in φ]
    U = P\(PA*F)
    Uλ = P\(PA*Fλ)
    global λ = λ - (norm(U)^2-1)/(2*sum(U.*Uλ))
    println("λ: $(rpad(λ, 18)) and the 2-norm: $(rpad(norm(U), 18))")
end


function gaussradau(n::Integer)
    # RADAUPTS   Gauss-Legendre-Radau Quadrature Nodes and Weights
    if n == 1
        [-1.0], [2.0]
    elseif n == 2
        [-1.0, 1/3], [.5, 1.5]
    else
        # Compute via GaussJacobi:
        x, w = gaussjacobi(n - 1, 0.0, 1.0)
        @inbounds for i in 1:length(w)
            w[i] = w[i] / (1.0 + x[i])
        end
        pushfirst!(x, -1.0)
        pushfirst!(w, 2.0 / n^2)
        x, w
    end
end


#---------------------------------------------------------
# # Ideal mean age
#---------------------------------------------------------

#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`ideal_mean_age.ipynb`](@__NBVIEWER_ROOT_URL__/examples/generated/ideal_mean_age.ipynb)

#---------------------------------------------------------
# ## The model
#---------------------------------------------------------

# We will simulate the ideal mean age of water.
# That is, the average amount of time since a water parcel had last contact with the surface.

# ### Tracer equation
#
# The ideal mean age is transported with water, is equal to $0$ at the surface, and increases by one second every second everywhere.
# In other words, the 3D field of the age, $a$, is governed by the tracer equation
#
# $$\frac{\partial a}{\partial t} + \nabla \cdot \left[ \boldsymbol{u} - \mathbf{K} \cdot \nabla \right] a = 1,$$
#
# where $\nabla \cdot \left[ \boldsymbol{u} - \mathbf{K} \cdot \nabla \right]$ is a differential operator that represents the transport by the ocean circulation.
# ($\boldsymbol{u}$ is the 3D vector field for the advection and $\mathbf{K}$ is the diffusivity matrix.)
# In the equation above, we also assume that there is the boundary condition that $a=0$ at the surface.

#---------------------------------
# ### Discretized tracer equation
#---------------------------------

# In AIBECS, the linear differential operator defined by $\nabla \cdot \left[ \boldsymbol{u} - \mathbf{K} \cdot \nabla \right]$ is approximated by a constant matrix $\mathbf{T}$ when discretizing the continuous 3D ocean onto the model grid.
# This matrix can be small (e.g., for models with a few boxes), or large, like for the OCIM (more on the OCIM later).
# Similarly, the continuous 3D field of the age, $a$, is discretized into a column-vector, $\boldsymbol{a}$.
# (We represent scalars in italic, vectors in bold italic, and matrices in upstraight bold.)

# In the discrete case, we replace the boundary condition (that $a = 0$ at the surface) by imposing $\boldsymbol{a} = 0$ in the surface layer of the model grid.
# In practice, this is done by restoring $\boldsymbol{a}$ to $0$ with a very short timescale.
# The tracer equation thgus takes the form of
#
# $$\frac{\partial\boldsymbol{a}}{\partial t} = -\mathbf{T} \, \boldsymbol{a} + 1 - \boldsymbol{\Lambda} \, \boldsymbol{a},$$
#
# where $\boldsymbol{\Lambda}$ is a diagonal matrix with entries equal to $1 / \tau$ in the surface layer of the model grid and 0 otherwise.
# The timescale $\tau$ is chosen to be very small, ensuring that $\boldsymbol{a}$ is very close to $0$ at the surface.
# The first term represents the transport by the ocean circulation, the second term the source of 1 second per second everywhere, and the last term the fast relaxation.

#---------------------------------
# ### Steady-state
#---------------------------------

# The steady-state is the equilibrium that would be reached if we wait long enough for $a$ to not change anymore.
# Mathematically, the steady-state is also the state for which
#
# $$\frac{\partial a}{\partial t} = 0.$$

# Computationally, in the discrete case, this means that we just need to solve
#
# $$0 = -\mathbf{T} \, \boldsymbol{a} + 1 - \boldsymbol{\Lambda} \, \boldsymbol{a}$$
#
# to find $\boldsymbol{a}$.
# More specifically, we need to solve
#
# $$(\mathbf{T} + \boldsymbol{\Lambda}) \, \boldsymbol{a} = 1.$$

# Now that we have the equations laid down, let us chose the circulation transport matrix, $\mathbf{T}$.

#---------------------------------------------------------
# ## Using AIBECS
#---------------------------------------------------------

#md # !!! note
#md #     If this is the first time you are trying AIBECS, make sure you go through the prerequisites!
#nb # > **Note**
#nb # > If this is the first time you are trying AIBECS, make sure you go through the prerequisites!

# AIBECS can interpret tracer equations as long as you arrange them under the generic form:
#
# $$\frac{\partial \boldsymbol{x}}{\partial t} = \boldsymbol{F}(\boldsymbol{x}, \boldsymbol{p}),$$
#
# where $\boldsymbol{F}(\boldsymbol{x}, \boldsymbol{p})$ is the rate of change of the state and $\boldsymbol{p}$ is the vector of model parameters.
# We only track the age here, so that the entire state of the system is determined by the age itself.
# In other words, here, $\boldsymbol{x} = \boldsymbol{a}$.

# We will use AIBECS to find the steady-state of the system.
# For AIBECS, this translates into finding the solution of $\boldsymbol{F}(\boldsymbol{x}, \boldsymbol{p}) = 0$.

# We start by telling Julia that we want to use the AIBECS package via

using AIBECS

#md # !!! note
#md #     If it's the first time you are running this line, the package will need precompiling.
#md #     This may take a minute or two.
#md #     (Just be patient... Or read on while you wait!)
#
#nb # > **Note**
#nb # > If it's the first time you are running this line, the package will need precompiling.
#nb # > This may take a minute or two.
#nb # > (Just be patient... Or read on while you wait!)
#
#md # !!! note
#md #     You should see a `Warning` for the `Flatten` package — just disregard it...
#md #     If you get an error though, please send me a copy of the output/error message, and I will try to troubleshoot it.
#nb # > **Note**
#nb # > You should see a `Warning` for the `Flatten` package — just disregard it...
#nb # > If you get an error though, please send me a copy of the output/error message, and I will try to troubleshoot it.

#---------------------------------
# ### The circulation
#---------------------------------

# We will use the circulation output from the Ocean Circulation Inverse Model (OCIM1).
# Basically, the OCIM provides researchers and oceanographers with a big sparse matrix that represents the global ocean circulation (advection and diffusion), which allows them to efficiently simulate the transport of passive tracers, like the age.
# (For more details, see Tim DeVries's [website](https://tdevries.eri.ucsb.edu/models-and-data-products/) and references therein.)
# With AIBECS, the OCIM0.1 and OCIM1 circulations can be loaded really easily, by simply typing

wet3D, grd, T_OCIM = AIBECS.OCIM1.load()
typeof(T_OCIM), size(T_OCIM)

#
#md # !!! note
#md #     Julia may ask you to download the OCIM matrix for you, in which case you should say yes (i.e., type `y`).
#md #     Once downloaded, AIBECS will remember where it downloaded the file and it will only load it from your laptop.
#nb # > **Note**
#nb # > Julia may ask you to download the OCIM matrix for you, in which case you should say yes (i.e., type `y`).
#nb # > Once downloaded, AIBECS will remember where it downloaded the file and it will only load it from your laptop.
#
# Additionally to downloading the OCIM file, the `load()` command loads 3 variables in the Julia workspace:
# - `wet3D` — a 3D array of the model grid, filled with `1`'s at "wet" grid boxes and `0`'s and "land" grid boxes.
# - `grd` — a `OceanGrid` object containing information about the 3D grid of the OCIM circulation, like the latitude, longitude, and depth of each grid boxes.
# - `T_OCIM` — the transport matrix representing advection and diffusion.
#
# The second line in command above tells you the type and the size of `T_OCIM`.
# It is a sparse matrix (CSC just means that it is stored in Compressed Sparse Column format) and is quite big!
#md # !!! note
#md #     A sparse matrix is just a matrix with very few non-zero entries.
#md #     Computationally, sparse matrices are stored differently than full matrices to save memory (no need to save all those zeros), and are much faster to use too!
#nb # > **Note**
#nb # > A sparse matrix is just a matrix with very few non-zero entries.
#nb # > Computationally, sparse matrices are stored differently than full matrices to save memory (no need to save all those zeros), and are much faster to use too!

# Anyway, this looks good, so let's move on with setting up the model!
#
# We have already loaded the transport matrix, `T_OCIM`, for the ocean circulation, but we must tell AIBECS that it applies to the age.
# To do that, we define a function of the parameters (although there are no parameters involved in this case, this is just the way AIBECS works for the moment).

T_age(p) = T_OCIM

# (Functions in Julia can be created in one line, just as above.)
# That's it for the circulation.
# Now, let's define the local sources and sinks.

#---------------------------------
# ### The local sources and sinks
#---------------------------------

# We will denote the age, $\boldsymbol{a}$, by the variable `age` in Julia.
# (It's good practice to use explicit names!)
# We need to translate the local sources and sinks in our discretized state function $\boldsymbol{F}(\boldsymbol{x}, \boldsymbol{p})$ into Julia code.

# #### The source
#--------------------------

# Remember the age increases by $1$ second every second and everywhere.
# So its source function is equal to, well, `1`! (seconds per seconds means it is unitless).
# Let's create the local source function:

source_age(age, p) = 1

# #### The sink
#--------------------------

# Recall that the age must also be $0$ at the surface.
# And that we implement this boundary condition by restoring the age very fast to $0$ in the surface layer.
# This will act as the sink for the age.
# But first, we need to figure out where "the surface layer" is.
# To do that, AIBECS can generate a number of useful constants for you.
# (You can see the list of functions by typing `varinfo(AIBECS)` at the REPL.)
# Here we will use the vector of grid box depths, `z`, which AIBECS can generate for us via

z = vector_of_depths(wet3D, grd)

# So what is the top layer?
# Let's investigate what's the minimum depth:

minimum(z)

# The surface layer in the OCIM grid has its center at about $18\,$m depth.
# We can create a mask of the surface layer via `z .< 20`.
# (This will return a vector of `0`s and `1`s, depending on whether the depth, `z`, is less than `20`.)
#
#md # !!! note
#md #     In Julia (like in MATLAB), placing a dot, `.`, in front of operators is a convenient way to do element-wise operations.
#nb # > **Note**
#nb # > In Julia (like in MATLAB), placing a dot, `.`, in front of operators is a convenient way to do element-wise operations.
#
# Then, we implement the local sink by restoring the age to `0` with a timescale `τ`, via

function sink_age(age, p)
    τ = p.τ
    return age .* (z .< 20u"m") / τ
end

#md # !!! note
#md #     Julia allows you to use unicode for your functions and variables, like for `τ`.
#nb # > **Note**
#nb # > Julia allows you to use unicode for your functions and variables, like for `τ`.
#
# Here, we have defined a Julia function using the `function` keyword because the sink is a bit more complicated, so that we needed two lines to define it.
# The first line unpacks the model parameters, which is just the restoring timescale, `τ`, in this case.
# We will chose the value for `τ` later.

# #### Net sources and sinks
#--------------------------

# The sources minus the sinks are simply defined by

sms_age(age, p) = source_age(age, p) .- sink_age(age, p)

# #### Model parameters
#--------------------------

# We must define the parameters... And AIBECS comes with an API for that!

t = empty_parameter_table()    # initialize table of parameters
add_parameter!(t, :τ, 1u"s")   # add the parameter we want (τ = 1s)
initialize_Parameters_type(t, "IdealAgeParameters")  # Generate the parameter type
t

# Note, in particular, that we gave our parameter `τ` a unit.
# Yes, Julia comes with some nice functionality to deal with units directly!
# The lines above created a table that contains all the info for generating the parameters vector, $\boldsymbol{p}$.
# To generate the parameters in AIBECS we do:

p₀ = IdealAgeParameters()

# where we have used the constructor `IdealAgeParameters`, whose name we defined in the previous cell.
# Here we did not really need to create `p₀` as a parameters vector, since it has only one element, `τ`, in it.
# However, we are here to learn, and this structure and functionality comes in very handy when one deals with many parameters.
# (And as you can imagine, having all the parameters in a nice table ready for being used in a publication comes quite handy!)

# #### State function and Jacobian
#--------------------------

# Similarly to `p₀`, let's create a state `x₀` to start with.
# The vector `x₀` will be our initial guess for the state.
# Let's assume that the age is `1` (seconds) everywhere (as an initial guess):

nb = number_of_wet_boxes(wet3D)  # number of wet boxes
x₀ = ones(nb)

# The first line above defines the number of wet grid boxes, `nb`.
# Here, this is also the length of the state vector `x`, because there is only one tracer, `age`.
# In the second line, the `ones` function creates a vector of `1`s of the size you give it (the number of wet grid boxes, `nb`, here, which we defined earlier).

# Finally, the last step for the set up is to define $\boldsymbol{F}$.
# Using AIBECS, this is done via

T_matrices = (T_age,)           # bundles all the transport matrices in a tuple
sources_minus_sinks = (sms_age,) # bundles all the source-sink functions in a tuple
F, ∇ₓF = state_function_and_Jacobian(T_matrices, sources_minus_sinks, nb) # generates the state function (and its Jacobian!)
F(x₀,p₀)

# That's it!
# We have just created a model of the mean age.
#
#
# Lines 2 and 3 are just telling AIBECS
# - what transport matrices it should use for the transport of these tracers, and
# - and what local sources and sinks should be appplied to these tracers
#
#md # !!! note
#md #     The `(x,)` syntax returns a tuple of one element — the comma is necessary because without it, `(x)` would be just like `x` with brackets around it.
#md #     This interface of AIBECS was developed for case with multiple tracers in mind, and might look a bit odd for a single tracer.
#md #     But in the future, this might be cleaned up to be easier to work with single tracers.
#nb # > **Note**
#nb # > The `(x,)` syntax returns a tuple of one element — the comma is necessary because without it, `(x)` would be just like `x` with brackets around it.
#nb # > This interface of AIBECS was developed for case with multiple tracers in mind, and might look a bit odd for a single tracer.
#nb # > But in the future, this might be cleaned up to be easier to work with single tracers.

# The fourth line creates two functions:
# - `F` — the numerical version of the **state function**, $\boldsymbol{F}$, of our model of the mean age, and
# - `∇ₓF` — the **Jacobian matrix** of the state function, i.e., $\nabla_{\boldsymbol{x}}\boldsymbol{F}$.
# Yes, AIBECS just automatically created an exact derivative of your input, using autodifferentiation via dual numbers.
# (I'd be very excited to detail how this is implemented here, but it is an entirely different discussion.)

# The last line just checks that our generated `F` works with our initial guess `x₀` and parameter vector `p₀`.

# #### Solving for the steady-state
#--------------------------

# The Jacobian, `∇ₓF` is essential to solving the steady-state equation $\boldsymbol{F}(\boldsymbol{x}, \boldsymbol{p}) = 0$ fast.
# Specifically, solving $\boldsymbol{F}(\boldsymbol{x}, \boldsymbol{p}) = 0$ is done via Newton's method.
# By starting from an initial guess, that you will have to provide, it will iterate over this recursion relation
#
# $$\boldsymbol{x}_{k+1} = \boldsymbol{x}_{k} - \nabla_{\boldsymbol{x}}\boldsymbol{F}(\boldsymbol{x}_{k}, \boldsymbol{p})^{-1} \boldsymbol{F}(\boldsymbol{x}_{k}, \boldsymbol{p})$$
#
# until $\boldsymbol{F}(\boldsymbol{x}_{k}, \boldsymbol{p})$ is sufficiently small.
# Now I should note that here, our age model is linear in $x$ (or `age` in our code), so that the solution will be found in a single iteration, or a sinle "matrix inversion", as could be seen from our steady-state equation for $\boldsymbol{a}$.

#md # !!! note
#md #     AIBECS comes with a built-in algorithm and an API to solve for the steady-state, so you don't have to worry about all these details!
#nb # > **Note**
#nb # > AIBECS comes with a built-in algorithm and an API to solve for the steady-state, so you don't have to worry about all these details!

# ##### Define the Steady-state problem in AIBECS

# First, we create an instance of the steady-state problem, via

prob = SteadyStateProblem(F, ∇ₓF, x₀, p₀)

# where we have simply provided the state function, $\boldsymbol{F}$, the Jacobian, $\nabla_{\boldsymbol{x}}\boldsymbol{F}$, the initial guess and the parameters.
# The `SteadyStateProblem` function is a standard "DiffEqBase" constructor that I have overloaded in my package so that you can easily generate the model here.

# ##### Solve for the steady-state with AIBECS

# Finally, we can find the solution in litterally one line, via the `solve` function:

age = solve(prob, CTKAlg())

# Here, I have provided the `solve` function with two things:
# - the problem, `prob`, which we just defined, and
# - the quasi-Newton algorithm that I wrote in Julia, denoted by `CTKAlg()` after C.T. Kelley, who originally wrote it in MATLAB.
#
# The last line should take about 10 seconds to 1 minute, depending on your laptop.
# That's it!
# We solved for the steady state!
# Everyone here deserves a nice tap on the shoulder — Good job!
# Now let's see what this age looks like on a map

#---------------------------------------------------------
# ## Figures
#---------------------------------------------------------

# We will plot a horizontal slice of the age at about 1000m depth using Cartopy.

# First, we must rearrange `age` into the 3D model grid.
# For that we will need the vector of the indices of wet points in the 3D grid, which we will denote by `iwet`, and which AIBECS generates via

iwet = indices_of_wet_boxes(wet3D)

# We then rearrange the column vector `age` into a 3D array via

age_3D = fill(NaN, size(wet3D)) # creates a 3D array of NaNs of the same size as `wet3D`
age_3D[iwet] = age              # Fills the wet grid boxes with the age values
size(age_3D)                    # Just to check the size of age_3D

# The last line just shows you the size of `age_3D`, which is a 3D-array as expected.

# Now let us find the index of the depth that is closest to 1000m.
# To do that we must use the depth information contained in `grd`.
# Let us first create a small vector of the depths of the grid:

depth = grd.depth

# We could count the index of the entry we want, but here we will use the `findfirst` function to find the first depth index that is greater than 1000m.
#nb # (Feel free to change the value of `iz` if you want to see a slice at another depth.)

iz = findfirst(depth .> 1000u"m")
iz, depth[iz]

# We get `iz = 13`, which is a layer that lies at 1104m, close to 1000m like we wanted.
#
# Finally, we need the latitude and longitudes of the grid, contained in `grd`.
# As for `depth`, we can use the OCIM's `grd` output:

lat, lon = ustrip.(grd.lat), ustrip.(grd.lon)

# So these are the latitudes and longitudes of the map we are about to plot.

# A last thing we can do is convert the age from seconds, `u"s"`, to years, `u"yr"`, because the age is large.
# This can be done via the Unitful package (loaded automatically by AIBECS).

age_3D_1000m_yr = age_3D[:,:,iz] * ustrip(1.0u"s" |> u"yr")

# Finally! Let's have a look at this ideal mean age!
# To make figures, here, we use Cartopy.
#nb # (You should have installed Cartopy if you went through the prerequisites correctly).
# To use it we simply type

ENV["MPLBACKEND"]="qt5agg"
using PyPlot, PyCall

#md # !!! note
#md #     The first line is needed for Mac users.
#md #     It's a bug that should eventually be resolved, but for now this seems to make it work.
#nb # > **Note**
#nb # > The first line is needed for Mac users.
#nb # > It's a bug that should eventually be resolved, but for now this seems to make it work.

# We import cartopy, define a new plot, add some coastlines because they are pretty, and add our slice of age at 1000m depth to it via

clf()
ccrs = pyimport("cartopy.crs")
ax = subplot(projection=ccrs.EqualEarth(central_longitude=-155.0))
ax.coastlines()
lon_cyc = [lon; 360+lon[1]] # making it cyclic for Cartopy
age_cyc = hcat(age_3D_1000m_yr, age_3D_1000m_yr[:,1])
p = contourf(lon_cyc, lat, age_cyc, levels=0:100:1200, transform=ccrs.PlateCarree(), zorder=-1)
colorbar(p, orientation="horizontal")
gcf() # gets the current figure to display

# That's it!
# Good job!
#
# At 1000m, the age ranges from a few years below deep water formation regions (Wedell Sea, North Atlantic), and reaches a dozen of centuries in the North Pacific!
# This is pretty good for so little work!


function HMMsmooth(v,phghm,ph1,pvgh,alpha=[],beta=[];ReturnPairwiseMarginals=false,ReturnLogLikelihood=false)
    #HMMSMOOTH Smoothing for a Hidden Markov Model (HMM)
    # [phtgV1T,phthtpgV1T]=HMMsmooth(v,phghm,ph1,pvgh)
    # Return the HMM pointwise p(h(t)|v(1:T)) and pairwise posterior p(h(t),h(t+1)|#v(1:T)).
    #
    # Inputs:
    # v : visible (observation) sequence
    # phghm : transition distribution
    # ph1 : initial distribution
    # pvgh : emission distribution p(v|h)
    #
    # Outputs:
    # phtgV1T : smoothed posterior p(h(t)|v(1:T))
    # phthtpgV1T : smoothed pair p(h(t),h(t+1)|v(1:T))
    T=length(v); H=size(phghm,1);

    if isempty(alpha)
        alpha,loglik=HMMforward(v,phghm,ph1,pvgh)
    end
    if isempty(beta)
        beta=HMMbackward(v,phghm,pvgh)
    end

    # smoothed posteriors: pointwise marginals:
    phtgV1T=zeros(H,T)
    for t=1:T
        tmp=alpha[:,t].*beta[:,t]
	phtgV1T[:,t]=tmp./sum(tmp)
    end

    if ReturnPairwiseMarginals
        # smoothed posteriors: pairwise marginals p(h(t),h(t+1)|v(1:T)):
        phthtpgV1T=zeros(H,H,T-1)
        for t=1:T-1
	    atmp=alpha[:,t]
	    btmp=beta[:,t+1]
	    #ctmp = repmat(atmp,1,H).*phghm'.*repmat(pvgh[v[t+1],:].*btmp',H,1)
            ctmp = repmat(atmp,1,H).*phghm'.*repmat(pvgh[v[t+1],:]'.*btmp',H,1)
	    phthtpgV1T[:,:,t]=ctmp./sum(ctmp)
        end
        if ReturnLogLikelihood
            return phtgV1T, phthtpgV1T, loglik
        end
        return phtgV1T, phthtpgV1T
    end

    if ReturnLogLikelihood
        return phtgV1T,loglik
    end
    return phtgV1T
end


using ExploratoryDataAnalysis
using Test

using StatsBase
using FreqTables

a = [1, 2, 3, 1, 2, 3, 1, 1, 2, 2]
b = [11, 11, 12, 12, 11, 11, 12, 12, 11, 12]

@testset "Cramer's V" begin
    @test round(cramerv(a, b), digits = 7) == 0.4472136
    @test round(cramerv(counts(a, b)), digits = 7) == 0.4472136
    @test round(cramerv(freqtable(a, b)), digits = 7) == 0.4472136
    @test round(cramerv(freqtable(a, b).array), digits = 7) == 0.4472136
end

@testset "infovalue" begin
    @test round(infovalue(counts(a, b)), digits = 7) == 0.8788898
    @test round(infovalue(freqtable(a, b)), digits = 7) == 0.8788898
    @test round(infovalue(freqtable(a, b).array), digits = 7) == 0.8788898
end

@testset "mutual info" begin
    @test round(mutualinfo(a, b), digits = 7) == 0.1046496
    @test round(mutualinfo(counts(a, b)), digits = 7) == 0.1046496
    @test round(mutualinfo(freqtable(a, b)), digits = 7) == 0.1046496
    @test round(mutualinfo(freqtable(a, b).array), digits = 7) == 0.1046496
end


# # Rossby Parameter
#
#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/notebooks/example_beta.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/notebooks/example_beta.ipynb)
#
# [from wikipedia:](https://en.wikipedia.org/wiki/Rossby_parameter)
#
# The Rossby parameter `β` arises due to the variation of the Coriolis
# force `f(φ)` as a function of latitude `φ`. It is important
# in the generation of Rossby waves and given by
#
# $$\beta(\varphi) = 2\Omega / R * cos(\varphi)$$

using PhysicalOceanography
isapprox(β(45),1.61687441686e-11)


using ModelingToolkit
using ModelingToolkit: sys2bigraph
using DiffEqBase
using Test

# Define some variables
@parameters t L g
@variables x(t) y(t) w(t) z(t) T(t) xˍt(t) yˍt(t)
@derivatives D'~t

eqs2 = [D(D(x)) ~ T*x,
        D(D(y)) ~ T*y - g,
        0 ~ x^2 + y^2 - L^2]
pendulum2 = ODESystem(eqs2, t, [x, y, T], [L, g], name=:pendulum)
lowered_sys = ModelingToolkit.ode_order_lowering(pendulum2)

lowered_eqs = [D(xˍt) ~ T*x,
               D(yˍt) ~ T*y - g,
               0 ~ x^2 + y^2 - L^2,
               D(x) ~ xˍt,
               D(y) ~ yˍt]
@test ODESystem(lowered_eqs, t, [xˍt, yˍt, x, y, T], [L, g]) == lowered_sys
@test isequal(lowered_sys.eqs, lowered_eqs)

# Simple pendulum in cartesian coordinates
eqs = [D(x) ~ w,
       D(y) ~ z,
       D(w) ~ T*x,
       D(z) ~ T*y - g,
       0 ~ x^2 + y^2 - L^2]
pendulum = ODESystem(eqs, t, [x, y, w, z, T], [L, g], name=:pendulum)

edges, vars, vars_asso = sys2bigraph(pendulum)
@test ModelingToolkit.matching(edges, length(vars), vars_asso .== 0) == [0, 0, 0, 0, 1, 2, 3, 4, 0]

edges, assign, vars_asso, eqs_asso = ModelingToolkit.pantelides(pendulum)

@test sort.(edges) == sort.([
 [5, 3],               # 1
 [6, 4],               # 2
 [7, 9, 1],            # 3
 [8, 9, 2],            # 4
 [2, 1],               # 5
 [2, 1, 6, 5],         # 6
 [5, 3, 10, 7],        # 7
 [6, 4, 11, 8],        # 8
 [2, 1, 6, 5, 11, 10], # 9
])
#                  [1, 2, 3, 4, 5,  6,  7,  8,  9, 10,   11]
#                  [x, y, w, z, x', y', w', z', T, x'', y'']
@test vars_asso == [5, 6, 7, 8, 10, 11, 0,  0,  0,  0,    0]
#1: D(x) ~ w
#2: D(y) ~ z
#3: D(w) ~ T*x
#4: D(z) ~ T*y - g
#5: 0 ~ x^2 + y^2 - L^2
# ----
#6: D(5) -> 0 ~ 2xx'+ 2yy'
#7: D(1) -> D(D(x)) ~ D(w)
#8: D(2) -> D(D(y)) ~ D(z)
#9: D(6) -> 0 ~ 2xx'' + 2x'x' + 2yy'' + 2y'y'
#                 [1, 2, 3, 4, 5, 6, 7, 8, 9]
@test eqs_asso == [7, 8, 0, 0, 6, 9, 0, 0, 0]


using ChebyshevQuantum
using Test
using SpecialFunctions

tol_var=1e-8

@testset "bessel eq test" begin

    a=0.0;
    b= 5.520078110588564; #second zero of besselj0

    #solve bessels equation between 0 and 2nd zero of besselj0 with Direchlet b.c.'s.
    t = ChebyshevQuantum.DiffEq.solve(bc1 = [:a,0, 1.0], bc2 = [:b,0,0.0], A0=x->x^2, A1= x->x, A2 = x->(x^2) ,B = 0.0, N=20,  a=a, b= b);

    x = a:.01:b

    @test maximum(abs.(besselj0.(x) - t.(x))) < tol_var
    @test maximum(abs.(x.^2 .* t''.(x) + x .* t'.(x) + (x.^2 ).*t.(x))) < tol_var #test diffeq directly

    a=1.0;
 #   b= 3.8319091796875; #first zero of besselj1
    b=   3.8317059218883514
    #solve bessels equation α = 1.0  between 1 and zero of besselj1 with Direchlet b.c.'s.
    t = ChebyshevQuantum.DiffEq.solve(bc1 = [:a,0, besselj1(a)], bc2 = [:b,0,besselj1(b)], A0=x->(x^2 - 1.0), A1= x->x, A2 = x->x^2  ,B = 0.0, N=20,  a=a, b= b);

    x = a:.01:b

    @test maximum(abs.(besselj1.(x) - t.(x))) < tol_var
    @test maximum(abs.(x.^2 .* t''.(x) + x .* t'.(x) + (x.^2 .- 1 ).*t.(x))) < tol_var #test diffeq directly

end

@testset "heat eq test" begin

    a=0.0;
    b= 1.0;

    #solve heat equation
    t = ChebyshevQuantum.DiffEq.solve(bc1 = [:a,0,1.0], bc2 = [:b,0,0.5], A2 = 0.5 , N=20,  a=a, b= b);

    x = a:.01:b

    function f(x) #solution is a line
        1.0 - 0.5*x
    end

    @test maximum(abs.(f.(x) - t.(x))) < tol_var

end

@testset "SHO test" begin

    a=0.0;
    b=2*pi;

    #solve simple harmonic oscillator equation, neuman b.c.'s
    t = ChebyshevQuantum.DiffEq.solve(bc1 = [:a,0,1.0], bc2 = [:a,1,0.0], A2 = 1.0, A0=1.0 , N=20,  a=a, b= b);

    x = a:.01:b

    function f(x) #solution is a cosine
        cos(x)
    end

    @test maximum(abs.(f.(x) - t.(x))) < tol_var

end

@testset "Decay Exp test" begin

    a=0.0;
    b=2*pi;

    #solve decaying exp equation, neuman b.c.'s
    t = ChebyshevQuantum.DiffEq.solve(bc1 = [:a,0,1.0], bc2 = [:a,1,-1.0], A2 = 1.0, A0=-1.0 , N=20,  a=a, b= b);

    x = a:.01:b

    function f(x) #solution is a decaying exp
        exp(-x)
    end

    @test maximum(abs.(f.(x) - t.(x))) < tol_var

end

@testset "First order Decay Exp" begin

    a=0.0;
    b=10.0;

    #solve decaying exp equation, first order
    t = ChebyshevQuantum.DiffEq.solve(bc1 = [:a,0,1.0],  A1 = 1.0, A0=1.0 , N=20,  a=a, b= b);

    x = a:.01:b

    function f(x) #solution is a decaying exp
        exp(-x)
    end

    @test maximum(abs.(f.(x) - t.(x))) < tol_var
    @test maximum(abs.(  t'.(x) .+ t.(x))) < tol_var #test diffeq directly

end


using CartesianFDM
using Test

@testset "1D identities" begin
    n = (5,)
    top = (periodic(),)

    ctx = cartesianfdmcontext(top, n)

    # assumes 1D
    ((δ⁻,), (δ⁺,)) = getproperty(ctx, :δ)
    ((σ⁻,), (σ⁺,)) = getproperty(ctx, :σ)
    ((τ⁻,), (τ⁺,)) = getproperty(ctx, :τ)

    Ψ, Φ = scalar(:Ψ, n), scalar(:Φ, n)

    @testset "2x identity 53" begin
        id = (σ⁺ * (δ⁻ * Φ)) .- (δ⁺ * (σ⁻ * Φ))
        @test all(iszero, Symbolics.expand.(id))

        id = (σ⁻ * (δ⁺ * Φ)) .- (δ⁻ * (σ⁺ * Φ))
        @test all(iszero, Symbolics.expand.(id))
    end

    @testset "2x identity 56" begin
        id = (σ⁺ * (Ψ .* (δ⁻ * Φ))) .+ (2Φ .* (δ⁺ * Ψ)) .- (δ⁺ * (Ψ .* (σ⁻ * Φ)))
        @test all(iszero, expand.(id))

        id = (σ⁻ * (Ψ .* (δ⁺ * Φ))) .+ (2Φ .* (δ⁻ * Ψ)) .- (δ⁻ * (Ψ .* (σ⁺ * Φ)))
        @test all(iszero, expand.(id))
    end

    @testset "2x identity 57" begin
        id = (Φ .* ((δ⁺ * (Ψ .* (σ⁻ * Φ))) .+ (σ⁺ * (Ψ .* (δ⁻ * Φ))))) .-
            (δ⁺ * (Ψ .* ((Φ .* (τ⁻ * Φ)) .+ ((τ⁻ * Φ) .* Φ))))
        @test all(iszero, expand.(id))

        id = (Φ .* ((δ⁻ * (Ψ .* (σ⁺ * Φ))) .+ (σ⁻ * (Ψ .* (δ⁺ * Φ))))) .-
            (δ⁻ * (Ψ .* ((Φ .* (τ⁺ * Φ)) .+ ((τ⁺ * Φ) .* Φ))))
        @test all(iszero, expand.(id))
    end

    @testset "2x identity 59" begin
        id = (((σ⁺ * Ψ) .* (δ⁺ * Φ)) .+ ((σ⁺ * Φ) .* (δ⁺ * Ψ))) - (2δ⁺ * (Ψ .* Φ))
        @test all(iszero, expand.(id))

        id = (((σ⁻ * Ψ) .* (δ⁻ * Φ)) .+ ((σ⁻ * Φ) .* (δ⁻ * Ψ))) - (2δ⁻ * (Ψ .* Φ))
        @test all(iszero, expand.(id))
    end
end

@testset "star stencil" begin
    using SparseArrays

    n = (3, 4, 5)
    top = (periodic(), nonperiodic(), periodic())

    ctx = cartesianfdmcontext(top, n)
    (τ⁻, τ⁺) = getproperty(ctx, :τ)

    X = scalar(:X, n)

    ###
    F = rand(prod(n)) .* X
    for i in eachindex(top)
        F .+= τ⁻[i] * (rand(prod(n)) .* X)
    end
    for i in eachindex(top)
        F .+= τ⁺[i] * (rand(prod(n)) .* X)
    end

    ###
    sym = linearize(star, ctx, F, X)

    num = Dict{Int,Vector{Float64}}()
    for (key, val) in sym
        num[key] = eval(first(build_function(val)))()
    end

    A = spdiagm(num...)

    ###
    @test all(iszero, F .- A * X)
end


using ProjectedGPE, ApproxFun, FastGaussQuadrature, PyPlot, Combinatorics
M=350
tol = 1e-14
#c=Vector{BigFloat}(randn(M))+im*Vector{BigFloat}(randn(M))
#c=randn(M)+im*randn(M)
c=zeros(M)
c[end]=1.
N1=sum(abs.(c).^2)

#evaluate modes using BigFloats
#can then use ApproxFun.jl for polynomials, applying correct normalization.
#Combinatorics.jl needed to evaluate factorial of Vector.

n=0:BigFloat(M)-1;
Hnorm = sqrt(BigFloat(π))*BigFloat(2).^n.*factorial.(n)

#Inserting norm.
f=Fun(Hermite(),c./sqrt.(Hnorm))
f2=Fun(GaussWeight(Hermite()),c./sqrt.(Hnorm))
x=linspace(-1.5sqrt(2M),1.5sqrt(2M),1000)
plot(x,abs.(f.(x)).^2.*exp.(-x.^2))

x,w=gausshermite(M)

N2=sum(w.*abs.(f.(x)).^2)
N3=sum(w.*abs.(f2.(x).*exp.(x.^2)).^2)

Float64(N2)

abs((N1-N2)/N1)

3*M*tol

#block to create T matrix
M=10
x,w=gausshermite(M)
T=zeros(x*ones(1,M))
n=0:BigFloat(M)-1;
Hnorm = sqrt.(BigFloat(π))*BigFloat(2).^n.*factorial.(n)

for j = 0:M-1
  c=zeros(M)
  c[j+1]=1.
  #Hn=Fun(Hermite(),c./sqrt(Hnorm));
  #T[:,j+1]=Hn.(x).*exp(-x.^2/2)
  Hn=Fun(GaussWeight(Hermite()),c)
  T[:,j+1]=Hn.(x).*exp.(x.^2/2)
end

x1,y1,T1=nfieldtrans("Hermite",M,2,2)

mind = M
maximum(abs.(T1[:,mind]-T[:,mind]))

T2=eigmat("Hermite",M,x,2)
maximum(abs.(T1[:,mind]-T2[:,mind]))

#-------------------------------
#Make and test Laguerre transform and quadrature
M=27
α = 0.
ω = 1.
x,w=gausslaguerre(M,α)
T=zeros(x*ones(1,M))
n=0:M-1;
Lnorm = exp.(lgamma.(n+α+1)-lgamma.(n+1))/ω

for j = 0:M-1
  c=zeros(M)
  c[j+1]=1.
  Lna=Fun(Laguerre(α),c./sqrt(Lnorm));
  T[:,j+1]=Lna.(sqrt(ω)*x).*exp.(-sqrt(ω)*x/2).*(sqrt(ω)*x).^(α/2)
end

c=zeros(M);
c[end]=BigFloat(1.)

#Gauss Laguerre quadrature can absorb the factor x^α/2 using associated Laguerre
#c=randn(M)+im*rand(M)
N1 = sum(abs.(c).^2)
L=Fun(Laguerre(α),c./sqrt(Lnorm))

L2=Fun(WeightedLaguerre(α),c./sqrt(Lnorm))
N2=sum(w.*abs(L.(sqrt(ω)*x)).^2)

N3=sum(w.*abs(L2.(sqrt(ω)*x).*exp.(sqrt(ω)*x)./(sqrt(ω)x).^(α/2)).^2)
x=linspace(0,20sqrt(M),1000)
plot(x,abs.(L.(x)).^2.*exp.(-x).*x.^α)

Lnorm

#------------------------------------
#test Laguerre
M=26
α = 0.
ω = 1.
#x,w=gausslaguerre(M,α)

#T = eigmat("Laguerre",M,x/ω,ω,α)

x,w,T=nfieldtrans(M,2;basis="Laguerre")
#norm test

c=zeros(M);
c[end]=1.
N1 = sum(abs(c).^2)

Ln = T*c

N2=sum(w.*abs(Ln).^2)
f1 = Fun(GaussWeight(Hermite()),[1])

#------------------------------------
#test Hermite
M=21

ω = 10
#x,w=gausslaguerre(M,α)

#T = eigmat("Laguerre",M,x/ω,ω,α)

x,w,T=nfieldtrans("Hermite",M,2,ω)
#norm test

c=zeros(M);
c[end]=1.
N1 = sum(abs(c).^2)

Hn = T*c

N2=sum(w.*abs(Hn).^2)


export UpdateMasseyColley

"""
    UpdateMasseyColley

 MasseyColley update rule, e.g., see "Whos's #1", Langville and Meyer, p.25
    this is the "Colleyized Massey method", i.e., use Colley's matrix on scores

## Parameters (none)
```
"""
struct UpdateMasseyColley <: UpdateRule
# this is a batch calculation that takes no account of past ratings, and
# has no parameters
end

function update_info( rule::UpdateMasseyColley )
    info = Dict(
                :name => "MasseyColley",
                :reference => "\"Whos's #1\", Langville and Meyer, p.25",
                :computation => "simultaneous",
                :state_model => "none",
                :input => "score",
                :output => "deterministic",
                :model => "single",
                :ties => true,
                :factors => false,
                :parameters => [],
                :record => true
                )
    return info
end
update_info( ::Type{UpdateMasseyColley} ) =  update_info( UpdateMasseyColley() )

function update_ratings( rule::UpdateMasseyColley,
                         input_ratings::RatingsList,
                         input_competitions::DataFrame)
    n = size(input_competitions,1)
    m = length( input_ratings.players )
    I = player_indexes( input_ratings.players )

    # construct Colley matrices and vectors
    C = diagm(0 => 2*ones(Int,m))
    p = zeros(Int, m)
    d = input_competitions # just an abbreviation
    point_diff = d[:,ScoreA] - d[:,ScoreB]
    for i=1:n
        p[ I[d[i,PlayerA]] ] +=  point_diff[i]
        p[ I[d[i,PlayerB]] ] -=  point_diff[i]
        C[ I[d[i,PlayerA]], I[d[i,PlayerB]] ] -= 1
        C[ I[d[i,PlayerB]], I[d[i,PlayerA]] ] -= 1
        C[ I[d[i,PlayerA]], I[d[i,PlayerA]] ] += 1
        C[ I[d[i,PlayerB]], I[d[i,PlayerB]] ] += 1
    end

    # solve Colley's equation
    r = C \ p
    ratings = Dict{String, Float64}()
    for player in input_ratings.players
        ratings[player] = r[ I[player] ]
    end

    # output ratings list
    output_ratings = RatingsList(input_ratings.players, ratings )
    return output_ratings
end


using Statistics:mean, std
using SpecialFunctions:erf

"""
normalize(scoresTrain,
          scoresTest)

Transform an array of scores into a range between [0,1] using min-max scaling.

Parameters
----------
      scores_train::AbstractVector{<:Real}
A vector of training scores, typically the result of [`fit`](@ref) with a detector.

    scores_test::AbstractVector{<:Real}
A vector of test scores, typically the result of [`score`](@ref) using a previously fitted detector.

Returns
----------
normalized_scores::Tuple{AbstractVector{<:Real}, AbstractVector{<:Real}}
The normalized train and test scores.

Examples
----------
scores_train, scores_test = ([1, 2, 3], [4, 3, 2, 1, 0])
scale_minmax(scores_train, scores_test) # ([0.0, 0.5, 1.0], [1.0, 1.0, 0.5, 0.0, 0.0])
scale_minmax(scores_train) # [0.0, 0.5, 1.0]
"""
function scale_minmax(scores_train::Scores, scores_test::Scores)::Tuple{Scores, Scores}
    minTrain, maxTrain = extrema(scores_train)
    @assert minTrain < maxTrain "Cannot normalize scores if they are all equal"
    f = scores -> clamp.((scores .- minTrain) ./ (maxTrain - minTrain), 0, 1)
    f(scores_train), f(scores_test)
end

"""
unify(scores_train,
      scores_test)

Transform an array of scores into a range between [0,1] using unifying scores as described in [1].

Parameters
----------
    scores_train::AbstractVector{<:Real}
A vector of training scores, typically the result of [`fit`](@ref) with a detector.

    scores_test::AbstractVector{<:Real}
A vector of test scores, typically the result of [`score`](@ref) using a previously fitted detector.

Returns
----------
unified_scores::Tuple{AbstractVector{<:Real}, AbstractVector{<:Real}}
The unified train and test scores.

Examples
----------
scores_train, scores_test = ([1, 2, 3], [4, 3, 2, 1, 0])
unify(scores_train, scores_test) # ([0.0, 0.0, 0.68..], [0.95.., 0.68.., 0.0, 0.0, 0.0])
unify(scores_train) # [0.0, 0.0, 0.68..]

References
----------
Kriegel, Hans-Peter; Kroger, Peer; Schubert, Erich; Zimek, Arthur (2011): Interpreting and Unifying Outlier Scores.
"""
function scale_unify(scores_train::Scores, scores_test::Scores)::Tuple{Scores, Scores}
    μ, σ = mean(scores_train), std(scores_train)
    @assert σ > 0 "Cannot normalize scores if they are all equal"
    f = scores -> clamp.(erf.((scores .- μ) ./ (σ * √2)), 0, 1)
    f(scores_train), f(scores_test)
end


# PROX_MATRIX    The proximal operator of a matrix function.
#
#   Suppose F is a orthogonally invariant matrix function such that
#   F(X) = f(s(X)), where s is the singular value map and f is some
#   absolutely symmetric function. Then
#
#     X = prox_matrix(V,lambda,prox_f)
#
#   evaluates the proximal operator of F via the proximal operator
#   of f. Here, it must be possible to evaluate prox_f as prox_f(v,lambda).
#
#   For example,
#
#     prox_matrix(V,lambda,prox_l1)
#
#   evaluates the proximal operator of the nuclear norm at V
#   (i.e., the singular value thresholding operator).
function prox_matrix(v, lambda, prox_f)
    U, S, V] = svd(v, 'econ')
    x = U * diag(prox_f(diag(S), lambda)) * V'
	return x
end


using NearlyNewton, LinearAlgebra, StaticArrays

    function theta(x)
       if x[1] > 0
           return atan(x[2] / x[1]) / (2.0 * pi)
       else
           return (pi + atan(x[2] / x[1])) / (2.0 * pi)
       end
    end
    f(x) = 100.0 * ((x[3] - 10.0 * theta(x))^2 + (sqrt(x[1]^2 + x[2]^2) - 1.0)^2) + x[3]^2

    function f∇f!(∇f, x)
        if !(∇f==nothing)
            if ( x[1]^2 + x[2]^2 == 0 )
                dtdx1 = 0;
                dtdx2 = 0;
            else
                dtdx1 = - x[2] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
                dtdx2 =   x[1] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
            end
            ∇f[1] = -2000.0*(x[3]-10.0*theta(x))*dtdx1 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 );
            ∇f[2] = -2000.0*(x[3]-10.0*theta(x))*dtdx2 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 );
            ∇f[3] =  200.0*(x[3]-10.0*theta(x)) + 2.0*x[3];
        end

        fx = f(x)
        return ∇f==nothing ? fx : (fx, ∇f)
    end
    function f∇f(∇f, x)
        if !(∇f == nothing)
            gx = similar(x)
            return f∇f!(gx, x)
        else
            return f∇f!(∇f, x)
        end
    end
    function f∇fs(∇f, x)
        if !(∇f == nothing)
            if ( x[1]^2 + x[2]^2 == 0 )
                dtdx1 = 0;
                dtdx2 = 0;
            else
                dtdx1 = - x[2] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) )
                dtdx2 =   x[1] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) )
            end

            s1 = -2000.0*(x[3]-10.0*theta(x))*dtdx1 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 )
            s2 = -2000.0*(x[3]-10.0*theta(x))*dtdx2 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 )
            s3 = 200.0*(x[3]-10.0*theta(x)) + 2.0*x[3]
            ∇f = @SVector [s1, s2, s3]
            return f(x), ∇f
        else
            return f(x)
        end
    end

    x0 = [-1.0, 0.0, 0.0]
    xopt = [1.0, 0.0, 0.0]

    res = NearlyNewton.minimize!(f∇f!, x0, (SR1(DirectApprox()), NearlyNewton.NWI()), I+0.0*x0*x0')
    res = NearlyNewton.minimize!(f∇f!, x0, (BFGS(DirectApprox()), NearlyNewton.NWI()), I+0.0*x0*x0')
    res = NearlyNewton.minimize!(f∇f!, x0, (BFGS(DirectApprox()), BackTracking()), I+0.0*x0*x0')

    function himmelblau_gradient!(x::Vector, gradient::Vector)
               gradient[1] = 4 * x[1] * (x[1]^2 + x[2] - 11) + 2 * (x[1] + x[2]^2 - 7)
               gradient[2] = 2 * (x[1]^2 + x[2] - 11) + 4 * x[2] * (x[1] + x[2]^2 - 7)
           end

    function himmelblau!(∇f, x)
        if !(∇f == nothing)
            ∇f[1] = 4.0 * x[1]^3 + 4.0 * x[1] * x[2] -
                44.0 * x[1] + 2.0 * x[1] + 2.0 * x[2]^2 - 14.0
            ∇f[2] = 2.0 * x[1]^2 + 2.0 * x[2] - 22.0 +
                4.0 * x[1] * x[2] + 4.0 * x[2]^3 - 28.0 * x[2]
        end

        fx = (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
        return ∇f == nothing ? fx : (fx, ∇f)
    end

    function himmelblau!(∇²f, ∇f, x)
        if !(∇²f == nothing)
            ∇²f[1, 1] = 12.0 * x[1]^2 + 4.0 * x[2] - 42.0
            ∇²f[1, 2] = 4.0 * x[1] + 4.0 * x[2]
            ∇²f[2, 1] = 4.0 * x[1] + 4.0 * x[2]
            ∇²f[2, 2] = 12.0 * x[2]^2 + 4.0 * x[1] - 26.0
        end
        if !(∇f == nothing)
            ∇f[1] = 4.0 * x[1]^3 + 4.0 * x[1] * x[2] -
                44.0 * x[1] + 2.0 * x[1] + 2.0 * x[2]^2 - 14.0
            ∇f[2] = 2.0 * x[1]^2 + 2.0 * x[2] - 22.0 +
                4.0 * x[1] * x[2] + 4.0 * x[2]^3 - 28.0 * x[2]
        end

        fx = (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
        if ∇f == nothing && ∇²f == nothing
            return fx
        elseif ∇²f == nothing
            return fx, ∇f
        else
            return fx, ∇f, ∇²f
        end
    end

    res = minimize!(himmelblau!, [2.0,2.0], Newton(DirectApprox()))
    res = minimize!(himmelblau!, [2.0,2.0], (Newton(DirectApprox()), NWI()))

using TRS
import NearlyNewton: NearlyExactTRSP

struct TRSolver{T} <: NearlyExactTRSP
    abstol::T
    maxiter::Integer
end
function (ms::TRSolver)(∇f::AbstractVector{T}, H, Δ, p;kwargs...) where T
    x, info = trs(Symmetric(H), ∇f, Δ)
    p .= x[:,1]
    m = dot(∇f, p) + dot(p, H * p)/2
    interior = norm(p, 2) ≤ Δ
    return (p=p, mz=m, interior=interior, λ=info.λ, hard_case=info.hard_case, solved=true)
end

    res = minimize!(himmelblau!, [2.0,2.0], (BFGS(DirectApprox()), TRSolver(1e-8, 100)))

    function polynomial!(∇²f, ∇f, x)
        if !isa(∇f, Nothing)
            ∇f[1] = -2.0 * (10.0 - x[1])
            ∇f[2] = -4.0 * (7.0 - x[2])^3
            ∇f[3] = -4.0 * (108.0 - x[3])^3
        end
        if !isa(∇²f, Nothing)
            ∇²f[1, 1] = 2.0
            ∇²f[1, 2] = 0.0
            ∇²f[1, 3] = 0.0
            ∇²f[2, 1] = 0.0
            ∇²f[2, 2] = 12.0 * (7.0 - x[2])^2
            ∇²f[2, 3] = 0.0
            ∇²f[3, 1] = 0.0
            ∇²f[3, 2] = 0.0
            ∇²f[3, 3] = 12.0 * (108.0 - x[3])^2
        end
        fx = (10.0 - x[1])^2 + (7.0 - x[2])^4 + (108.0 - x[3])^4

        if ∇f == nothing && ∇²f == nothing
            return fx
        elseif ∇²f == nothing
            return fx, ∇f
        else
            return fx, ∇f, ∇²f
        end
    end
    function polynomial!(∇f, x)
        if !isa(∇f, Nothing)
            ∇f[1] = -2.0 * (10.0 - x[1])
            ∇f[2] = -4.0 * (7.0 - x[2])^3
            ∇f[3] = -4.0 * (108.0 - x[3])^3
        end

        fx = (10.0 - x[1])^2 + (7.0 - x[2])^4 + (108.0 - x[3])^4

        if ∇f == nothing
            return fx
        else
            return fx, ∇f
        end

    end

        res = minimize!(polynomial!, [2.0,2.0,2.0], (BFGS(DirectApprox()), NWI()))

        function rosenbrock(∇²f, ∇f, x)
            if !isa(∇f, Nothing)
                ∇f[1] = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
                ∇f[2] = 200.0 * (x[2] - x[1]^2)
            end
            if !isa(∇²f, Nothing)
                ∇²f[1, 1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2
                ∇²f[1, 2] = -400.0 * x[1]
                ∇²f[2, 1] = -400.0 * x[1]
                ∇²f[2, 2] = 200.0
            end
            fx = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

            if ∇f == nothing && ∇²f == nothing
                return fx
            elseif ∇²f == nothing
                return fx, ∇f
            else
                return fx, ∇f, ∇²f
            end
        end

        function rosenbrock(∇f, x)
            if !isa(∇f, Nothing)
                ∇f[1] = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
                ∇f[2] = 200.0 * (x[2] - x[1]^2)
            end

            fx = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

            if ∇f == nothing
                return fx
            else
                return fx, ∇f
            end
        end

        function rosenbrock_gradient!(storage, x)

        end

        function rosenbrock_hessian!(storage, x)
            storage[1, 1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2
            storage[1, 2] = -400.0 * x[1]
            storage[2, 1] = -400.0 * x[1]
            storage[2, 2] = 200.0
        end
        res = minimize!(rosenbrock, [2.0,2.0], (BFGS(DirectApprox()), NWI()))


using DifferentialEquations
using SimpleDiffEq
using DiffEqCallbacks
using Random
using Distributions
using DiffEqParamEstim
using DataFrames
using StatsPlots
using BenchmarkTools

function sir_ode!(du,u,p,t)
    (S,I,R,Y) = u
    (β,c,γ) = p
    N = S+I+R
    infection = β*c*I/N*S
    recovery = γ*I
    @inbounds begin
        du[1] = -infection
        du[2] = infection - recovery
        du[3] = recovery
        du[4] = infection
    end
    nothing
end;

δt = 0.1
tmax = 40.0
tspan = (0.0,tmax)
t = 0.0:δt:tmax
obstimes = 0:1.0:tmax;

u0 = [990.0,10.0,0.0,0.0]; # S,I.R,Y

p = [0.05,10.0,0.25]; # β,c,γ

affect!(integrator) = integrator.u[4] = 0.0
cb_zero = PresetTimeCallback(obstimes,affect!)

prob_ode = ODEProblem(sir_ode!,u0,tspan,p)

sol_ode = solve(prob_ode,callback=cb_zero);

df_ode = DataFrame(sol_ode(obstimes)')
df_ode[!,:t] = obstimes;

@df df_ode plot(:t,
    [:x1 :x2 :x3 :x4],
    label=["S" "I" "R" "Y"],
    xlabel="Time",
    ylabel="Number")

Random.seed!(1234);

data = rand.(Poisson.(df_ode[!,:x4]))

plot(obstimes,data)
plot!(obstimes,df_ode[!,:x4])

using Optim

function ss1(β)
    prob = remake(prob_ode,u0=[990.0,10.0,0.0,0.0],p=[β,10.0,0.25])
    sol = solve(prob,Tsit5(),callback=cb_zero,saveat=obstimes)
    sol_data = sol(obstimes)[4,:]
    return(sum((sol_data - data) .^2))
end

function nll1(β)
    prob = remake(prob_ode,u0=[990.0,10.0,0.0,0.0],p=[β,10.0,0.25])
    sol = solve(prob,Tsit5(),callback=cb_zero,saveat=obstimes)
    sol_data = sol(obstimes)[4,:]
    -sum(logpdf.(Poisson.(sol_data),data))
end

lower1 = 0.0
upper1 = 1.0
initial_x1 = 0.1

opt1_ss = Optim.optimize(ss1,lower1,upper1)

opt1_nll = Optim.optimize(nll1,lower1,upper1)

function ss2(x)
    (i0,β) = x
    I = i0*1000.0
    prob = remake(prob_ode,u0=[1000-I,I,0.0,0.0],p=[β,10.0,0.25])
    sol = solve(prob,Tsit5(),callback=cb_zero,saveat=obstimes)
    sol_data = sol(obstimes)[4,:]
    return(sum((sol_data - data) .^2))
end

function nll2(x)
    (i0,β) = x
    I = i0*1000.0
    prob = remake(prob_ode,u0=[1000-I,I,0.0,0.0],p=[β,10.0,0.25])
    sol = solve(prob,Tsit5(),callback=cb_zero,saveat=obstimes)
    sol_data = sol(obstimes)[4,:]
    -sum(logpdf.(Poisson.(sol_data),data))
end

lower2 = [0.0,0.0]
upper2 = [1.0,1.0]
initial_x2 = [0.01,0.1]

opt2_ss = Optim.optimize(ss2,lower2,upper2,initial_x2)

opt2_nll = Optim.optimize(nll2,lower2,upper2,initial_x2)

function loss_function(sol)
    sol_data = DataFrame(sol(obstimes)')[!,:x4]
    -sum(logpdf.(Poisson.(sol_data),data))
end

prob_generator = (prob,q) -> remake(prob,
    u0=[1000-(q[1]*1000),q[1]*1000,0.0,0.0],
    p=[q[2],10.0,0.25])

cost_function = build_loss_objective(prob_ode,
    Tsit5(),
    loss_function,
    prob_generator = prob_generator,
    maxiters=10000,
    verbose=false,
    callback=cb_zero)

opt_pe1 = Optim.optimize(cost_function,lower2,upper2,initial_x2)

using NLopt
opt = Opt(:LD_MMA, 2)
opt.lower_bounds = lower2
opt.upper_bounds = upper2
opt.min_objective = cost_function
(minf,minx,ret) = NLopt.optimize(opt,initial_x2)

using BlackBoxOptim
bound1 = Tuple{Float64, Float64}[(0.0,1.0),(0.0, 1.0)]
result = bboptimize(cost_function;SearchRange = bound1, MaxSteps = 110e3)

include(joinpath(@__DIR__,"tutorials","appendix.jl"))
appendix()


"""
@article{zitzler2000comparison,
  title={Comparison of multiobjective evolutionary algorithms: Empirical results},
  author={Zitzler, Eckart and Deb, Kalyanmoy and Thiele, Lothar},
  journal={Evolutionary computation},
  volume={8},
  number={2},
  pages={173--195},
  year={2000},
  publisher={MIT Press}
}

ZDT 测试函数的特点:
- 有固定的 2 个目标
- 有确定的 True Pareto Front
"""

abstract type ZDT <: TestProblem end

abstract type ZDT1 <: ZDT end

function ZDT1(m = 30, n_solutions = 100)
    f(x) = begin
        g = 1.0 + 9.0 * ( sum(x[2:end]) / (length(x)-1) )
        h = 1 - sqrt(x[1] / g)
        return [ x[1], g*h ]
    end

    bounds = Array([zeros(m) ones(m)]')

    x = range(0, 1, length=n_solutions)

    X = [vcat(x[i], zeros(m - 1)) for i in 1:n_solutions]
    pareto_set = [ Individual(x, f(x)) for x in X ]

    # @infiltrate

    return f, bounds, pareto_set
end

function truepf(x)
    "truepf"
end

thename(::Type{ZDT1}) = "ZDT1"

abstract type ZDT2 <: ZDT end

function ZDT2(D = 30, n_solutions = 100)
    f(x) = begin
        gx = 1.0 + 9.0 * ( sum(x[2:end]) / (length(x)-1) )
        return [x[1], gx*(1 - (x[1] / gx)^2) ]
    end
    bounds = Array([zeros(D) ones(D)]')

    x = range(0, 1, length=n_solutions)

    x = [vcat(x[i], zeros(D - 1)) for i in 1:n_solutions]
    pareto_set = [ Individual(xx, f(xx)) for xx in x ]

    return f, bounds, pareto_set
end

abstract type ZDT3 <: ZDT end

function ZDT3(D = 30, n_solutions = 100)
    f(x) = begin
        gx = 1.0 + 9.0 * ( sum(x[2:end]) / (length(x)-1) )
        a = x[1] / gx
        return [x[1], gx*(1 - sqrt(a) - a*sin(10π*x[1])) ]
    end
    bounds = Array([zeros(D) ones(D)]')

    if n_solutions < 6
        n_solutions = 6
    end

    regions = [ 0 0.0830015349
                0.182228780 0.2577623634
                0.4093136748 0.4538821041
                0.6183967944 0.6525117038
                0.8233317983 0.8518328654]

    n = Int(ceil(n_solutions / size(regions, 1)))
    x = Float64[]
    for i in 1:size(regions, 1)
        x = vcat(x, range(regions[i,1], regions[i,2], length=n))
    end

    x = [vcat(x[i], zeros(D - 1)) for i in 1:n_solutions]
    pareto_set = [ Individual(xx, f(xx)) for xx in x ]

    return f, bounds, get_non_dominated_solutions(pareto_set)
end

abstract type ZDT4 <: ZDT end

function ZDT4(D = 10, n_solutions = 100)
    f(x) = begin
        gx = 1.0 + 10*(length(x)-1) + sum( x[2:end].^2 - 10cos.(4π*x[2:end]))
        return [x[1], gx*(1 - sqrt(x[1] / gx)) ]
    end
    bounds = Array([-5zeros(D) 5ones(D)]')
    bounds[:,1] = [0, 1.0]

    x = range(0, 1, length=n_solutions)

    x = [vcat(x[i], zeros(D - 1)) for i in 1:n_solutions]
    pareto_set = [ Individual(xx, f(xx)) for xx in x ]

    return f, bounds, pareto_set
end

abstract type ZDT6 <: ZDT end

function ZDT6(D = 10, n_solutions = 100)
    f(x) = begin
        gx = 1.0 + 9.0 * ( sum(x[2:end]) / (length(x)-1.0) )^(0.25)
        ff1 = 1.0 - exp(-4.0x[1])*sin(6.0π*x[1])^6
        return [ ff1 , gx*(1.0 - (ff1 / gx)^2) ]
    end

    bounds = Array([zeros(D) ones(D)]')

    #x = range(0, 1, length=n_solutions)

    #x = [vcat(x[i], zeros(D - 1)) for i in 1:n_solutions]
    xx = range(0.2807753191, 1, length=100)
    yy = 1 .- (xx).^2
    pareto_set = [ Individual(zeros(0), [xx[i], yy[i]]) for i in 1:length(xx) ]

    return f, bounds, pareto_set
end


# ---
# title: 858. Mirror Reflection
# id: problem858
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math
# link: <https://leetcode.com/problems/mirror-reflection/description/>
# hidden: true
# ---
#
# There is a special square room with mirrors on each of the four walls.  Except
# for the southwest corner, there are receptors on each of the remaining
# corners, numbered `0`, `1`, and `2`.
#
# The square room has walls of length `p`, and a laser ray from the southwest
# corner first meets the east wall at a distance `q` from the `0`th receptor.
#
# Return the number of the receptor that the ray meets first.  (It is guaranteed
# that the ray will meet a receptor eventually.)
#
#
#
# **Example 1:**
#
#
#
#     Input: p = 2, q = 1
#     Output: 2
#     Explanation: The ray meets receptor 2 the first time it gets reflected back to the left wall.
#     ![](https://s3-lc-upload.s3.amazonaws.com/uploads/2018/06/18/reflection.png)
#
#
# **Note:**
#
#   1. `1 <= p <= 1000`
#   2. `0 <= q <= p`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using HomotopyContinuation, LinearAlgebra, DynamicPolynomials

# See https://www.juliahomotopycontinuation.org/examples/cyclooctane/
c² = 2
@polyvar z[1:3, 1:6]
z_vec = vec(z)[1:17] # the 17 variables in a vector
Z = [zeros(3) z[:,1:5] [z[1,6]; z[2,6]; 0] [√c²; 0; 0]] # the eight points in a matrix

# define the functions for cyclooctane:
F1 = [(Z[:, i] - Z[:, i+1]) ⋅ (Z[:, i] - Z[:, i+1]) - c² for i in 1:7]
F2 = [(Z[:, i] - Z[:, i+2]) ⋅ (Z[:, i] - Z[:, i+2]) - 8c²/3 for i in 1:6]
F3 = (Z[:, 7] - Z[:, 1]) ⋅ (Z[:, 7] - Z[:, 1]) - 8c²/3
F4 = (Z[:, 8] - Z[:, 2]) ⋅ (Z[:, 8] - Z[:, 2]) - 8c²/3
f = [F1; F2; F3; F4]

n = 2 # dimension of the cyclooctane variety
N = 17 # ambient dimension
@polyvar Aᵥ[1:n, 1:N] bᵥ[1:n] # variables for the linear equations
p = [vec(Aᵥ); bᵥ] # parameters
F = [f; Aᵥ * z_vec - bᵥ] # the polynomial system we have to solve

# now we solve one particular instance for A,b complex. we use this as start system
A₀ = randn(ComplexF64, n, N)
b₀ = randn(ComplexF64, n)
p₀ = [vec(A₀); b₀]

F₀ = [subs(Fᵢ, p => [vec(A₀); b₀]) for Fᵢ in F]
complex_result = solve(F₀)
println("Degree of the variety: ", nfinite(complex_result)) # should be 1408


@inline cst_S(n, i, x) = binomial(n,i)*x^i*(1-x)^(n-i)
@inline cst_C(n1, n2, x) = x^n1*(1-x)^n2
@inline function cst_dS(n, i, x)
    if i == 0 # (1-x)^n
        -n*(1-x)^(n-1)
    elseif i == n # x^n
        n*x^(n-1)
    else
        binomial(n, i)*(i*x^(i-1)*(1-x)^(n-i)-(n-i)*x^i*(1-x)^(n-i-1))
    end
end

@inline function cst_dC(n1, n2, x)
    n1*x^(n1-1)*(1-x)^n2 - n2*x^n1*(1-x)^(n2-1)
end
function fit_cst(n::Int, dz::T,
    xu::AbstractVector{T},yu::AbstractVector{T},
    xl::AbstractVector{T},yl::AbstractVector{T},
    opt::OptimizerFactory;
    c1::T = 0.5,
    c2::T = 1.0 ) where {T}
    m = Model(opt)
    lb = oneunit(T)*-2
    up = oneunit(T)*2
    n_up = length(xu)
    n_lo = length(xl)
    @variable(m, lb <= A₀ <= up, start = 1.0)
    @variable(m, lb <= Aᵤ[1:n] <= up, start = 1.0)
    @variable(m, lb <= Aₗ[1:n] <= up, start = -1.0)
    @expression(m, yuc[j=1:n_up], cst_C(c1,c2,xu[j])*( A₀*cst_S(n,0,xu[j])+sum(Aᵤ[i]*cst_S(n,i,xu[j]) for i = 1:n))+xu[j]*dz)
    @expression(m, ylc[j=1:n_lo], cst_C(c1,c2,xl[j])*(-A₀*cst_S(n,0,xl[j])+sum(Aₗ[i]*cst_S(n,i,xl[j]) for i = 1:n))-xl[j]*dz)
    @expression(m, fup, sum((yuc[j]-yu[j])^2 for j=1:n_up))
    @expression(m, flo, sum((ylc[j]-yl[j])^2 for j=1:n_lo))
    @objective(m, Min, fup+flo)
    optimize!(m)
    status = termination_status(m)
    value(A₀), value.(Aᵤ), value.(Aₗ), status
end


function get_GLRK_ω_matrix(s)
    ω = zeros(s, s+1)
    g = TableauGauss(s)
    ω[s,s+1] = 1
    for i in 1:s-1
        for j in 1:s
            ω[i,j] = g.b[j] * g.c[j]^(i-1)
        end
    end
    return ω
end


#' ---
#' title: Problem Modification
#' ---

#' **Originally Contributed by**: Arpit Bhatia

#' This tutorial deals with how to modify models after they have been created and solved.
#' This functionality can be useful, for example,
#' when we are solving many similar models in succession or generating the model dynamically.
#' Additionally it is sometimes desirable for the solver to re-start from the last solution to
#' reduce running times for successive solves (“hot-start”).

using JuMP

#' ## Modifying Variables

model = Model()
@variable(model, x);

#' ### Variable Bounds

#' The `set_lower_bound` and `set_upper_bound` functions can be used to create as well as
#' modify an existing variable bound.

set_lower_bound(x, 3)
lower_bound(x)

#+

set_lower_bound(x, 2)
lower_bound(x)

#' We can delete variable bounds using the `delete_lower_bound` and `delete_upper_bound` functions.

delete_lower_bound(x)
has_lower_bound(x)

#' We can assign a fixed value to a variable using `fix`.

fix(x, 5)
fix_value(x)

#' However, fixing a variable with existing bounds will throw an error.

@variable(model, y >= 0);

#+ tangle =  false

fix(y, 2)

#' As we can see in the error message above,
#' we have to specify to JuMP that we wish to forcefuly remove the bound.

fix(y, 2; force = true)
fix_value(y)

#' We can also call the `unfix` function to remove the fixed value.

unfix(x)
is_fixed(x)

#' ### Deleting Variables

#' The `all_variables` function returns a list of all variables present in the model.
all_variables(model)

#' To delete variables from a model, we can use the `delete` function.

delete(model, x)
all_variables(model)

#' We can also check whether a variable is a valid JuMP variable in a model using the `is_valid` function.

is_valid(model, x)

#' ## Modifying Constraints

model = Model()
@variable(model, x);

#' ### Modifying a Variable Coefficient
#' It is also possible to modify the scalar coefficients
#' (but notably not yet the quadratic coefficients) using the `set_coefficient` function.

@constraint(model, con, 2x <= 1);

#+

set_coefficient(con, x, 3)
con

#' ### Deleting a Constraint
#' Just like for deleting variables, we can use the `delete` function for constraints as well.

delete(model, con)
is_valid(model, con)

#' ## Modifying the Objective

model = Model()
@variable(model, x)
@objective(model, Min, 7x + 4);

#' The function `objective_function` is used to query the objective of a model.

objective_function(model)

#' `objective_sense` is similarily used to query the objective sense of a model.

objective_sense(model)

#' To easiest way to change the objective is to simply call `@objective` again
#' - the previous objective function and sense will be replaced.

@objective(model, Max, 8x + 3)
objective_function(model)

#+

objective_sense(model)

#' Another way is to change the objective is to
#' use the low-level functions `set_objective_function` and `set_objective_sense`.

set_objective_function(model, 5x + 11)
objective_function(model)

#+

set_objective_sense(model, MOI.MIN_SENSE)
objective_sense(model)

#' Note that we can't use the Min and Max shortcuts here as `set_objective_sense` is a low level function.