#=
proximal.jl

    Evaluate the proximal operator for a variety of commonly used functions in
    optimization, data science, machine learning, and econometrics

@author: Quint Wiersma <q.wiersma@vu.nl>

@date: 2022/01/13
=#

"""
	soft_thresh(x, λ)

Compute soft thresholding operator with scaling parameter `λ` at `x`, proximal
operator of ``ℓ₁``-norm.

#### Arguments
  - `x::Real`	: input
  - `λ::Real`	: scaling parameter

#### Returns
  - `y::Real`	: soft thresholded value
"""
function soft_thresh(x::Real, λ::Real)
	# Soft thresholding
	y= sign(x)*max(abs(x) - λ, zero(λ))

	return y
end

"""
	block_soft_thresh(x, λ)

Compute block soft thresholding operator with scaling parameter `λ` at `x`,
proximal operator of the ``ℓ₂``-norm.

#### Arguments
  - `x::AbstractVector`	: input (n x 1)
  - `λ::Real`			: scaling parameter

#### Returns
  - `y::AbstractVector`	: block soft thresholded value (n x 1)
"""
function block_soft_thresh(x::AbstractVector, λ::Real)
    y= similar(x)
	block_soft_thresh!(y, x, λ)

	return y
end

"""
	block_soft_thresh!(y, x, λ)

Compute block soft thresholding operator with scaling parameter `λ` at `x`,
proximal operator of the ``ℓ₂``-norm, storing the results in `y`. See also
`block_soft_thresh`.
"""
function block_soft_thresh!(y::AbstractVector, x::AbstractVector, λ::Real)
	T= eltype(x)
	# Scaling
	τ= max(one(λ) - λ*inv(norm(x) + eps(T)), zero(λ))

	# Block soft thresholding
	@. y= τ*x

	return nothing
end

"""
	block_soft_thresh!(x, λ)

Compute block soft thresholding operator with scaling parameter `λ` at `x`,
proximal operator of the ``ℓ₂``-norm, overwriting `x`. See also
`block_soft_thresh`.
"""
function block_soft_thresh!(x::AbstractVector, λ::Real)
	T= eltype(x)
	# Scaling
	τ= max(one(λ) - λ*inv(norm(x) + eps(T)), zero(λ))

	# Block soft thresholding
	@. x= τ*x

	return nothing
end

"""
	shrinkage(x, λ)

Compute shrinkage operator with scaling parameter `λ` at `x`, proximal operator
of the squared ℓ₂-norm (ridge).

#### Arguments
  - `x::Real`	: input
  - `λ::Real`	: scaling parameter

#### Returns
  - `y::Real`	: shrunken value
"""
function shrinkage(x::Real, λ::Real)
	# Scaling
	τ= one(λ)*inv(one(λ) + λ)

	# Shrinkage
	y= τ*x

	return y
end

"""
    shrinkage!(x, λ, fac, b)

Compute the generalized shrinkage operator with scaling parameter `λ` at `x`,
proximal operator of a quadratic function with quadratic parameters `A` and
linear parameters `b` using a factorization `fac` of ``I + λA``, overwriting
`x`. See also `shrinkage`.
"""
function shrinkage!(x::AbstractVector, λ::Real, fac::Factorization, b::AbstractVector)
    # x - λb
    x.-= λ .* b
    # (I + λA)⁻¹(x - λb)
    ldiv!(fac, x)

    return nothing
end

"""
    shrinkage(x, λ, fac, b)

Compute the generalized shrinkage operator with scaling parameter `λ` at `x`,
proximal operator of a quadratic function with quadratic parameters `A` and
linear parameters `b` using a factorization `fac` of ``I + λA``.

#### Arguments
  - `x::AbstractVector` : input
  - `λ::Real`           : scaling parameter
  - `fac::Factorization`: factorization of ``I + λA``
  - `b::AbstractVector` : linear coefficients

#### Returns
  - `y::AbstractVector` : shrunken values
"""
function shrinkage(x::AbstractVector, λ::Real, fac::Factorization, b::AbstractVector)
    y= similar(x)
    shrinkage!(y, λ, fac, b)

    return y
end

"""
    shrinkage!(x, λ, A, b)

Compute the generalized shrinkage operator with scaling parameter `λ` at `x`,
proximal operator of a quadratic function with quadratic parameters `A` and
linear parameters `b`, overwriting `x`. See also `shrinkage`.
"""
function shrinkage!(x::AbstractVector, λ::Real, A::AbstractMatrix, b::AbstractVector)
    # scaling I + λA
    S= λ .* A
    @inbounds @fastmath for i in axes(A,1)
        S[i,i]+= one(eltype(A))
    end

    # factorization
    C= cholesky!(Hermitian(S))

    shrinkage!(x, λ, C, b)

    return nothing
end

"""
    shrinkage(x, λ, A, b)

Compute the generalized shrinkage operator with scaling parameter `λ` at `x`,
proximal operator of a quadratic function with quadratic parameters `A` and
linear parameters `b`.

#### Arguments
  - `x::AbstractVector` : input
  - `λ::Real`           : scaling parameter
  - `A::AbstractMatrix` : quadratic coefficients
  - `b::AbstractVector` : linear coefficients

#### Returns
  - `y::AbstractVector` : shrunken values
"""
function shrinkage(x::AbstractVector, λ::Real, A::AbstractMatrix, b::AbstractVector)
    y= similar(x)
    shrinkage!(y, λ, A, b)

    return y
end

function smooth!(x::AbstractVector, λ::Real, f::Function, ∇f!::Function, x0::AbstractVector)
    # adjust objective function and gradient
    # f(y) + (1/2λ)‖y - x‖₂²
    g(y::AbstractVector)= f(y) + inv(λ + λ) * ( sum(abs2, y) + sum(abs2, x) - 2 * dot(y, x) )
    # ∇f(y) + (1/λ)(y - x)
    ∇g!(∇g::AbstractVector, y::AbstractVector)= begin
                                                    # gradient
                                                    ∇f!(∇g, y)
                                                    # quadratic term
                                                    ∇g.+= inv(λ) .* (y .- x)
                                                end

    # solve using L-BFGS
    res= optimize(g, ∇g!, x0, LBFGS(), Optim.Options(g_tol = 1e-4))
    x.= Optim.minimizer(res)
    # store current estimate to warm start next iteration
    x0.= x

    return nothing
end


# This file calculates the UL Decomposition of a Toeplitz operator,
# via the Wiener–Hopf factorization.  If the symbol of
# the Toeplitz operator is G(z), we find a Wiener–Hopf factorization
#
# 	G(z) = φ_+(z)*φ_-(z)
#
# where φ_± are analytic in the interior/exterior of the unit
# circle.  Then we have the UL decomposition
#
#   T[G]=T[φ_+]*T[φ_-]
#
# The Wiener–Hopf factorization is found by rephrasing as a
# Riemann–Hilbert problem:
#
#   φ_+(z)=φ_-(z)^(-1)*G(z)
#
# and writing
#
#   φ_-(z)^(-1) = 1 + C_-V(z) and φ_+(z)=1+C_+V(z)=1+V+C_-V(z)
#
# where C_± are the left/right limits of the Cauchy operator.
# This leads to a singular integral equation
#
#   V + C_-V*(1-G) = G - 1
#
# This approach also extends to the matrix case.

using ApproxFun, SingularIntegralEquations

# Scalar case

G=Fun(z->2+cos(z+1/z),Circle()) # the symbol of the Toeplitz operator
T=ToeplitzOperator(G)

C  = Cauchy(-1)
V  = (I+(1-G)*C)\(G-1)

Φmi = 1+C*V
Φp = V+Φmi

L  = ToeplitzOperator(1/Φmi)
U  = ToeplitzOperator(Φp)

norm((T-U*L)[1:10,1:10])  # check the accuracy

# Matrix case

G=Fun(z->[-1 -3; -3 -1]/z +
         [ 2  2;  1 -3] +
         [ 2 -1;  1  2]*z,Circle())
T=ToeplitzOperator(G)

C  = Cauchy(-1)
V  = V=(I+(I-G)*C)\(G-I)

Φmi = I+C*V
Φp = V+Φmi

L  = ToeplitzOperator(inv(Φmi))
U  = ToeplitzOperator(Φp)

norm((T-U*L)[1:10,1:10])  # check the accuracy


# Three-state Bioreactor
id = :PerezGalvan2017e

pL = [6.45; 0.46; 1.05]
pU = [6.55; 0.47; 1.1]
PerezGalvan2017e_x0(p::Vector{T}) where T = T[p[1]; 5.0*one(T); 15.0*one(T)]
function PerezGalvan2017e_f!(du::Vector{T}, u::Vector{T}, p::Vector{T}, t) where T
    x3m = 50.0; x2f = 20.0; D = 0.202; Y = 0.4; β = 0.2; α = 0.5
    μ = p[2]*(1.0 - u[3]/x3m)*u[2]/(p[3] + u[2])
    du[1] = (μ - D)*u[1]
    du[2] = D*(x2f - u[2]) - μ*u[1]/Y
    du[3] = D*u[3] + (α*μ + β)*u[1]
    return
end
tspan = (0.0, 7.7)

prob = ODERelaxProb(PerezGalvan2017e_f!, tspan, PerezGalvan2017e_x0, pL, pU)
url = "https://www.sciencedirect.com/science/article/abs/pii/S0098135417300923"
source = "PerezGalvan2017"
desc = "Three-state Bioreactor"


"""
    eom_chuacircuit_nscroll_sine(u, p, t) -> SVector{3}

Equations of motion for n-scroll chaotic attractors from an adjusted Chua
system [1].

# References
1. Tang, Wallace KS, et al. "Generation of n-scroll attractors via
sine function." IEEE Transactions on Circuits and Systems I:
Fundamental Theory and Applications 48.11 (2001): 1369-1372.

"""
function eom_chuacircuit_nscroll_sine(u, p, t)
    α, β, γ, a, b, c, σx, σy, σz = (p...,)
    x, y, z = (u...,)

    n::Int = c + 1
    if x >= 2*a*c
        fx = (b*pi/2*a)*(x - 2*a*c)
    elseif -2*a*c < x < 2*a*c
        d = ifelse(isodd(n), pi, 0)
        fx = -b*sin((pi*x/2*a) + d)
    elseif x <= -2*a*c
        fx = (b*pi/2*a)*(x + 2*a*c)
    end

    ηx = σx == 0 ? 0 : rand(Normal(0, σx))
    ηy = σy == 0 ? 0 : rand(Normal(0, σy))
    ηz = σz == 0 ? 0 : rand(Normal(0, σz))

    dx = α*(y - fx) + ηx
    dy = x - y + z + ηy
    dz = -β*y - γ*z + ηz
    return SVector{3}(dx, dy, dz)
end

"""
    chua_nscroll_sinefunc(u₀, α, β, γ, a, b, c::Int)

Generate n-scroll chaotic attractors from an adjusted Chua system [1].

# References
1. Tang, Wallace KS, et al. "Generation of n-scroll attractors via
sine function." IEEE Transactions on Circuits and Systems I:
Fundamental Theory and Applications 48.11 (2001): 1369-1372.
"""
function chuacircuit_nscroll_sine(u₀, α, β, γ, a, b, c::Int, σx, σy, σz)
    p = [α, β, γ, a, b, c, σx, σy, σz]
    ContinuousDynamicalSystem(eom_chuacircuit_nscroll_sine, u₀, p)
end
chuacircuit_nscroll_sine(;u₀ = [0.0, 0.0, 0.28695],
        α = 10.814, β = 14, γ = 0, a = 1.3, b = 0.11, c = 2,
        σx = 0.0, σy = 0.0, σz = 0.0) =
    chuacircuit_nscroll_sine(u₀, α, β, γ, a, b, c, σx, σy, σz)


using Flux, DiffEqFlux, DifferentialEquations, Plots, StatsBase, RecursiveArrayTools

const u0 = Float32[2.; 0.]
const datasize = 1000
const batchsize = 30
const batchtime = 10
const tspan = (0.0f0,25f0)

function trueODEfunc(du,u,p,t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end
t = range(tspan[1],tspan[2],length=datasize)
prob = ODEProblem(trueODEfunc,u0,tspan)
true_sol = solve(prob,Tsit5(),saveat=t)
ode_data = Array(true_sol)

function get_batch()
    batch_0s = sample(1:datasize-batchtime, batchsize, replace=false,ordered=false)
    batch_idxs = range.(batch_0s,batch_0s.+(batchtime-1))
    batch_ts = [t[i] for i in batch_idxs]
    batch_u0 = [true_sol[u0] for u0 in batch_0s]
    batch_u = [true_sol[:,idxs] for idxs in batch_idxs]
    return batch_u0, batch_u, batch_ts
end
batch_u0, batch_u, batch_ts = get_batch();

dudt = Chain(x -> x.^3,
             Dense(2,50,tanh),
             Dense(50,2))
ps = Flux.params(dudt)
function n_ode(batch_u0, batch_t)
    neural_ode(dudt,batch_u0,tspan,Tsit5(),
               saveat=batch_t, reltol=1e-7,abstol=1e-9)
end
#= n_ode = x->neural_ode(dudt,x,tspan,Tsit5(), =#
                      #= saveat=t,reltol=1e-7,abstol=1e-9) =#

n_ode(batch_u0, batch_ts)

pred = n_ode(u0,t) # Get the prediction using the correct initial condition
scatter(t,ode_data[1,:],label="data")
scatter!(t,Flux.data(pred[1,:]),label="prediction")

function predict_n_ode(batch_u0,batch_t)
  n_ode(batch_u0, batch_t)
end
loss_n_ode(batch_u0, batch_u,batch_t) = sum(abs2,batch_u .- predict_n_ode(batch_u0,batch_t))

data = Iterators.repeated(get_batch(), 1000)
opt = ADAM(0.1)
cb = function () #callback function to observe training
    bu0,bu,bt = get_batch()
    display(loss_n_ode(bu0,bu,bt))
  # plot current prediction against data
  #= cur_pred = Flux.data(predict_n_ode(bu0,bt)) =#
  #= pl = scatter(bt,ode_data[1,:],label="data") =#
  #= scatter!(pl,bt,cur_pred[1,:],label="prediction") =#
  #= display(plot(pl)) =#
end

# Display the ODE with the initial parameter values.
cb()

Flux.train!(loss_n_ode, ps, data, opt, cb = cb)


# coding: utf-8

# In[1]:

#a simple day trading game
#day trader is only allowed to make at maximum two trades
#the strategy is long only
#lets find out the maximum profit

#more details can be found in the following link
# https://www.geeksforgeeks.org/maximum-profit-by-buying-and-selling-a-share-at-most-twice/

#an alternative version in recursion exists
#its done by using a different approach
#strongly recommend you to take a look
# https://github.com/je-suis-tm/recursion-and-dynamic-programming/blob/master/stock%20trading%20recursion.jl

# In[2]:

#there are two scenarios to maximize the profit
#one trade or two trades
#first we run a reverse iteration
#to obtain the maximum profit from one trade
#then we run a normal iteration
#to obtain the maximum profit
#from one trade plus the result from reverse iteration
function stock_trading(prices)

    #initialize the profit at zero
    profit=[0 for _ in 1:length(prices)]

    #initialize maximum price with the close price
    max_price=prices[end]

    #reverse order iteration
    for i in length(prices)-1:-1:1

        #update the maximum price to compute the maximum profit
        if prices[i]>max_price

            max_price=prices[i]

        end

        #two scenarios to get the maximum profit
        #either the previous iteration is larger
        #or this round of iteration
        profit[i]=max(profit[i+1],max_price-prices[i])

    end

    #initialize minimum price with the open price
    min_price=prices[1]

    #second round of iteration
    for i in 2:length(prices)

        #update the minimum price to compute the maximum profit
        if prices[i]<min_price

            min_price=prices[i]

        end

        #two scenarios to get the maximum profit
        #either the previous iteration is larger
        #or this round of iteration plus the result from single transaction
        profit[i]=max(profit[i-1],profit[i]+prices[i]-min_price)

    end

    return profit[end]

end

# In[3]:

stock_trading([10,22,5,75,65,80])

# In[4]:

stock_trading([2,30,15,10,8,25,80])

# In[5]:

stock_trading([100,30,15,10,8,25,80])

# In[6]:

stock_trading([90,70,35,11,5])


# # Two Parabolas
#
# The “two parabolas” problem in two dimensions reads as
# ```math
#     \min_{𝐱 ∈ X }
#     \begin{bmatrix} f₁(\mathbf{x}) \\ f₂(\mathbf{x}) \end{bmatrix} =
#     \min_{\mathbf{x} ∈ X}
#     \begin{bmatrix}
#     (x₁ - 1)² + (x₂ - 1)² \\
#     (x₁ + 1)² + (x₂ + 1)²
#     \end{bmatrix}.
# ```
# It is unconstrained if the feasible set is ``X = ℝ^2``.
# The individual minima ``[1,1]`` and ``[-1,-1]`` are such that (in the unconstrained case)
# the global Pareto Set is
# ```math
# \mathcal{P}_{S} = \{ \mathbf{x} ∈ ℝ^2 : x₁ = x₂, \, -1 \le x₁, x₂ \le 1  \}.
# ```
#
# ## Solve using Exact Functions
#
# The gradients are easily calculated as
# ```math
# \nabla f_1 (\mathbf x) = 2 \begin{bmatrix}
# x_1 -1 \\ x_2 - 1 \end{bmatrix}, \;
# \nabla f_2 (\mathbf x) = 2 \begin{bmatrix}
# x_1 +1 \\ x_2 + 1 \end{bmatrix}, \;
# ```
#
# We can provide them to the solver to find a critical point:

using Pkg #src
Pkg.activate(@__DIR__) #src
using Test #src

using Morbit
Morbit.print_all_logs() #src

f₁ = x -> sum( (x .- 1).^2 )
f₂ = x -> sum( (x .+ 1).^2 )
∇f₁ = x -> 2 .* ( x .- 1 )
∇f₂ = x -> 2 .* ( x .+ 1 )

mop = MixedMOP(2);  # problem with 2 variables
add_objective!(mop, f₁, ∇f₁ )
add_objective!(mop, f₂, ∇f₂ )

#~ starting point
x₀ = [ -π ;  2.71828 ]

#~ set maximum number of iterations
ac = AlgoConfig( max_iter = 20)
#~ `optimize` will return parameter and result vectors as well
#~ as an return code and the evaluation database:
x, fx, ret_code, db = optimize( mop, x₀; algo_config = ac );
x

# Hopefully, `x` is critical, i.e., `x[1] ≈ x[2]`.
@test x[1] ≈ x[2] atol = .1 #src

# !!! note
#     To print more information on what the solver is doing, you can use the `Logging` module:
#     ```julia
#     import Logging: global_logger, ConsoleLogger
#     global_logger( ConsoleLogger( stderr, Morbit.loglevel4;
#         meta_formatter = Morbit.morbit_formatter ) )
#     ```
#     `loglevel4` is the most detailed and `loglevel1` is least detailed.
#     `Morbit.print_all_logs()` is a convenient shorthand.

#%% #src
# ### Plotting Iteration Sites
# Let's retrieve the iteration sites.
# We convert to Tuples for easier plotting.
iteration_indices = [ iter_.x_index for iter_ in db.iter_info]
it_sites = Tuple.(Morbit.get_site.(db, iteration_indices))

# For Plotting we use CairoMakie
using Makie, CairoMakie

#~ Pareto Set ≙ line from (-1,-1) to (1,1)
fig, ax, _ = lines( [(-1,-1),(1,1)]; color = :blue, linewidth = 2,
    figure = (resolution = (600, 600),) )

#~ Plot the iteration sites:
lines!(it_sites)
scatter!(it_sites;
    color = LinRange(0, 1, length(it_sites)),
    colormap = :winter
)

#~ Plot function contours
Y = X = LinRange(-4, 4, 100)
Z₁ = [ f₁([x;y]) for x ∈ X, y ∈ X ]
Z₂ = [ f₂([x;y]) for x ∈ X, y ∈ X ]
levels = [ i.^2 for i = LinRange(.1, 6, 6) ]
contour!(X,Y,Z₁; colormap = :greens, levels = levels, linewidth = .5 )
contour!(X,Y,Z₂; colormap = :heat, levels = levels, linewidth = .5 )

#~ Show the plot:
ax.title[] = "Pareto Set and Iterates."
ax.xgridvisible[] = false
ax.ygridvisible[] = false

fig

#%% #src

# ## Solving using RBF Surrogates
#
# Suppose now that we do not have access to the objective gradients and that the objectives
# also take some time to evaluate.
# In this situation, we could try to model them using surrogate models.
# To use radial basis function models, pass an `RbfConfig` when specifying the objective:
#
mop_rbf = MixedMOP()

#~ Define the RBF surrogates
rbf_cfg = RbfConfig(
    kernel = :inv_multiquadric
)
#~ Add objective functions to `mop_rbf`
add_objective!(mop_rbf, f₁, rbf_cfg )
add_objective!(mop_rbf, f₂, rbf_cfg )

#~ only perform 10 iterations
ac = AlgoConfig( max_iter = 10 )
x, fx, _, db = optimize( mop_rbf, x₀; algo_config = ac )
x

#src Setup and save plot for docs.
iteration_indices_rbf = [ iter_.x_index for iter_ in db.iter_info]
it_sites_rbf = Tuple.(Morbit.get_site.(db, iteration_indices_rbf))
lines!(it_sites_rbf) #hide
scatter!(it_sites_rbf; color = :orange) #hide
nothing #hide

# The iteration sites are the orange circles:
fig #hide

#%% #src

# ## Different Starting Points and Recycling Data
#
# The method could converge to different points depending on the starting point.
# We can pass the evaluation data from previous runs to facilitate the construction of surrogate models:

#src Setup the problem anew, to ensure fresh start
ac = AlgoConfig( #hide
    max_iter = 10 #hide
    ); #hide
mop_rbf = MixedMOP(); #hide
#~ define the RBF surogates #hide
rbf_cfg = RbfConfig(  #hide
    kernel = :inv_multiquadric, #hide
); #hide
#~ add objective functions to `mop_rbf` #hide
add_objective!(mop_rbf, f₁, rbf_cfg ); #hide
add_objective!(mop_rbf, f₂, rbf_cfg ); #hide

#~ an array of well spread points in [-4,4]² #hide
X =[ #hide
 [-4.0, -4.0], #hide
 [3.727327839472812, 3.8615291196035457], #hide
 [3.804712690019901, -3.9610212058521235], #hide
 [-0.14512898384374573, -0.005775390168885508], #hide
 [-3.775315499879552, 3.8150054323309064], #hide
 [1.714228746087743, 1.8435786475209621], #hide
 [-1.9603720505875337, -2.0123206708499275], #hide
 [3.9953803225349187, -0.47734576293976794], #hide
 [-3.9944468955728745, 0.49857343385493635], #hide
 [-1.0455585089057458, 2.735699160002545] #hide
]; #hide

# Suppose, `X` is a list of different points in ℝ².

#src This is the code block visible in the docs:

#~ A dict to associate starting and end points:
start_fin_points = Dict();

#~ perform several runs:
db₀ = nothing # initial database can be `nothing`
for x₀ ∈ X
    global db₀, start_fin_points
    x_fin, fx_fin, _, db₀ = optimize( mop_rbf, x₀; algo_config = ac, populated_db = db₀ )
    #~ add points to dict
    start_fin_points[x₀] = x_fin
end

# Plotting:

fig, ax, _ = lines( [(-1,-1),(1,1)]; color = :blue, linewidth = 2,
    figure = (resolution = (600, 600), ),
    axis = (title="Different Starting Points",),
)

for (k,v) in start_fin_points
    lines!( [ Tuple(k), Tuple(v) ]; color = :lightgray )
end

scatter!( Tuple.(keys(start_fin_points));
    color = :green
)
scatter!( Tuple.(values(start_fin_points));
    color = :lightblue
)

fig #hide

# In the plot, the green points show the starting points and the lightblue circles show the final iterates:


# Rational Quadratic Isotropic Covariance Function

"""
    RQIso <: Isotropic{SqEuclidean}

Isotropic Rational Quadratic kernel (covariance)
```math
k(x,x') = σ²(1 + (x - x')ᵀ(x - x')/(2αℓ²))^{-α}
```
with length scale ``ℓ``, signal standard deviation ``σ``, and shape parameter ``α``.
"""
mutable struct RQIso <: Isotropic{SqEuclidean}
    "Squared length scale"
    ℓ2::Float64
    "Signal variance"
    σ2::Float64
    "Shape parameter"
    α::Float64
    "Priors for kernel parameters"
    priors::Array

    """
        RQIso(ll:Float64, lσ::Float64, lα::Float64)

    Create `RQIso` with length scale `exp(ll)`, signal standard deviation `exp(lσ)`, and
    shape parameter `exp(lα)`.
    """
    RQIso(ll::Float64, lσ::Float64, lα::Float64) =
        new(exp(2 * ll), exp(2 * lσ), exp(lα), [])
end

function set_params!(rq::RQIso, hyp::VecF64)
    length(hyp) == 3 || throw(ArgumentError("Rational Quadratic function has three parameters"))
    rq.ℓ2, rq.σ2, rq.α = exp(2 * hyp[1]), exp(2 * hyp[2]), exp(hyp[3])
end

get_params(rq::RQIso) = Float64[log(rq.ℓ2) / 2, log(rq.σ2) / 2, log(rq.α)]
get_param_names(rq::RQIso) = [:ll, :lσ, :lα]
num_params(rq::RQIso) = 3

Statistics.cov(rq::RQIso, r::Float64) = rq.σ2 * (1 + r / (2 * rq.α * rq.ℓ2))^(-rq.α)

@inline dk_dll(rq::RQIso, r::Float64) =
    (s = r / rq.ℓ2; rq.σ2 * s * (1 + s / (2 * rq.α))^(-rq.α - 1)) # dK_d(log ℓ)dK_dℓ
@inline function dk_dlα(rq::RQIso, r::Float64)
    s = r / rq.ℓ2
    part = 1 + s / (2 * rq.α)
    rq.σ2 * part^(-rq.α) * (s / (2 * part) - rq.α * log(part))  # dK_d(log α)
end
@inline function dk_dθp(rq::RQIso, r::Float64, p::Int)
    if p==1
        return dk_dll(rq, r)
    elseif p==2
        return dk_dlσ(rq, r)
    elseif p==3
        return dk_dlα(rq, r)
    else
        return NaN
    end
end


# # Transmission line circuit
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/TransmissionLine.ipynb)
#
#md # !!! note "Overview"
#md #     System type: linear continuous system\
#md #     State dimension: parametric, typically between 4 to 40\
#md #     Application domain: Power Systems Stability
#
# We consider the transmission line model used in [AKS11].
# The RLC circuit of the transmission line is shown below.
# In the circuit, ``U_{in}`` is the voltage at the sending end,
# and ``U_{out}`` is the voltage at the receiving end.

# ![](https://raw.githubusercontent.com/JuliaReach/ReachabilityAnalysis.jl/master/examples/TransmissionLine/tline.png)

# ## Model description

# For reference to understand this model we refer to any standard textbook on electrical circuits.
# The electrical elements law for resistors (R), inductors (L) and capacitors (C) are studied,
# for instance, in Chapter 3 of [^K15].

# Let's assume that the network consists of ``η > 2`` nodes.
# There are ``η`` dynamic equations for the capacitor voltages and ``\eta`` dynamic
# equations for the circuit currents. Therefore, the state vector can be represented as

# ```math
# x = [U_1, U_2, …, U_η, I_1, I_2, …, I_η]^T,
# ```
# and the state dimension is ``2η``.

# **Analysis of the first node.** When writing the equations for the voltages and currents,
# we should pay attention to the sign choices in given circuit, which do not exactly match
# the convention in text-books (in the sense that ``U_1``s positive terminal is at the bottom,
# as indicated by the arrows). Let ``R_d`` denote the driver resistance's current.
# By Kirchhoff's voltage law,

# ```math
# U_{in} = R_d I_1 + LI_1' - U_1,
# ```
# therefore
# ```math
# \boxed{I_1' = \dfrac{U_{in} + U_1}{L} - \dfrac{R_d}{L}I_1}.
# ```
#
# By Kirchhoff's current law, and if ``I_{1, C}`` denotes the current through the
# capacitor ``C``, connected to the first node,
#
# ```math
#    I_1 = I_2 + I_{C, 1},
# ```
# and
# ```math
# I_{C, 1} = -CU_1',
# ```
# then
# ```math
#    \boxed{U_1' = \frac{I_2 - I_1}{C}}.
# ```

# **Analysis of the intermediate nodes.** Other node's equations are obtained
# in a similar fashion. For instance, for the second loop one has

# ```math
# -U_1 = RI_2 + LI_2' - U_2 \Rightarrow I_2' = \dfrac{U_2-U_1}{L} - \frac{R}{L}I_2
# ```
# for the current's equation, and
#
# ```math
#    I_{2, C} = I_2 - I_3,\Rightarrow U_2' = \frac{I_3 - I_2}{C}.
# ```
# for the voltage's equation.
#
# Generalizing for arbitrary ``l = 2, \ldots, \eta - 1`` is trivial and gives:
#
# ```math
#     \boxed{I'_l = \dfrac{U_l - U_{l-1}}{L} - \frac{R}{L}I_l} \\
#     \boxed{U'_l = \dfrac{I_{l+1} - I_l}{C}}.
# ```

# **Analysis of the last node.** The last node corresponds to the case ``l = \eta``
#
# ```math
# \boxed{I'_\eta = \dfrac{U_\eta - U_{\eta-1}}{L} - \frac{R}{L}I_\eta}, \\
# \boxed{U'_{out} = U'_{out} = - \frac{I_{\eta}}{C}}.
# ```

# **System of linear ODEs.** Recall that that the stateset is ``\mathbb{R}^{2\eta}``,
# where the state variables are
# ``x = [U_1, U_2, \ldots, U_\eta, I_1, I_2, \ldots, I_\eta]^T``.
# The system can be written as a block-diagonal system of linear ODEs,
# ```math
# x'(t) = Ax(t) + BU_{in}(t),
# ```
# using the results in the previous section.

# The coefficients matrix ``A`` can be written as a block-diagonal matrix.
# There exist useful constructors in the base package `LinearAlgebra` that greatly simplify
# building matrices with special shape, as in our case, such as diagonal and band matrices,
# using the types `Diagonal` and `Bidiagonal` in our case.

using ReachabilityAnalysis, LinearAlgebra, SparseArrays

function tline(;η=3, R=1.00, Rd=10.0, L=1e-10, C=1e-13 * 4.00)
    A₁₁ = zeros(η, η)
    A₁₂ = Bidiagonal(fill(-1/C, η), fill(1/C, η-1), :U)
    A₂₁ = Bidiagonal(fill(1/L, η), fill(-1/L, η-1), :L)
    A₂₂ = Diagonal(vcat(-Rd/L, fill(-R/L, η-1)))
    A  = [A₁₁ A₁₂; A₂₁ A₂₂]
    B = sparse([η+1], [1], 1/L, 2η, 1)
    return A, B
end

# We can visualize the structure of the cofficients matrix ``A`` for the case
# $\eta = 20$ with `spy` plot:

using Plots

A, _ = tline(η=20)
spy(A, legend=nothing, markersize=2.0, title="Sparsity pattern of A", xlab="columns", ylab="rows")

# Notice that the matrix coefficients are rather big, hence it is convenient to rescale the system
# for numerical stability.
# Let ``α > 0`` be a scaling factor, and let ``\tilde{x}(t) = x(\alpha t)``.
# By the chain rule,
# ```math
# \tilde{x}'(t) = \alpha x'(\alpha t) = \alpha A x(\alpha t) + \alpha B U_{in}(\alpha t) = \tilde{A} \tilde{x}(t) + \tilde{B} \tilde{U}_{in}(t),
# ```
# where ``\tilde{A} := \alpha A`` and ``\tilde{B} := \alpha B``.

function scale!(s, α=1.0)
    s.A .*= α
    s.B .*= α
    return s
end

# Note that under this transformation, the time horizon has to be transformed as well,
# through ``\tilde{T} = \alpha T``.

# ## Reachability settings

# The transmission line parameters used in this model are displayed in the following table.

# |resistance in [Ω] | driver resistance in [Ω] | Inductance in [H] | Capacitance in [F] |
# |------------------|--------------------------|------------------|-------------------|
# |R = 1.00| Rdriver = 10.0| L = 1e−10 | C = 4e−13|

# The steady state is obtained by zeroing the left-hand side of the ODE, which gives
#
# ```math
#    0 = x' = Ax_\infty + Bu_0 \Rightarrow x_\infty = -A^{-1}B u_0
# ```
# The initial set under consideration corresponds to the steady state for input voltages
# ``U_{in, ss} := [-0.2, 0.2]``.
# Moreover, an uncertainty is added so that the initial currents are also uncertain.
# The set of initial states is then
# ```math
#    x(0) \in \mathcal{X}_0 := -A^{-1} B  U_{in, ss} \oplus □(0.001),
# ```
# where ``□(ϵ)`` is the infinity-norm ball of center zero and radius ``ϵ``.

# The time horizon is chosen as ``T = 0.7`` seconds. We consider a
# scaling factor ``\alpha = 1e-9``.

# We are interested in the step response to an input voltage ``U_{in}(t)``, constant
# for ``t \in [0, T]`` over the domain ``U_{in} = [0.99, 1.01]``.

# We consider the case of ``\eta = 20`` nodes as in [AKS11], such that the system
# has ``n = 40`` state variables.

η = 20 # order
n = 2η # state-space dimension
A, B = tline(η=η);

Uin_ss = Interval(-0.2, 0.2)
□(ϵ) = BallInf(zeros(n), ϵ)
X0 = -inv(Matrix(A)) * B * Uin_ss ⊕ □(0.001);

Uin = Interval(0.99, 1.01)
s = @system(x'= A*x + B*u, x ∈ Universe(n), u ∈ Uin)
α = 1e-9 # scaling factor
scale!(s, α);

T = 0.7 * 1e-11 # time horizon
P = InitialValueProblem(s, X0);

# ## Results

# We solve system using a step size of `δ=1e-3` and the box algorithm.

sol = solve(P, T=0.7, alg=BOX(δ=1e-3));

# To get the variable ``U_{out}`` we have to project onto the ``η``-th coordinate
# and invert the sign of the flowpipe.

Uout_vs_t = @. (-1.0) * project(sol, η);

#-

plot(Uout_vs_t, vars=(0, η), color=:blue, xlab="t", ylab="Uout", alpha=.5, lw=0.5)

#md # !!! tip "Technical note"
#md #     Since we are only interested in the behavior of ``U_{out}``, we can use
#md #     algorithm `BFFPSV18` with the options `alg=BFFPSV18(δ=1e-3, dim=statedim(P), vars=[η]))`,
#md #     which will use an interval (1D) decomposition of the state space and only compute
#md #     the flowpipe associated with variable ``\eta``.

# ## References

# [^AKS11]: Althoff, Matthias, Bruce H. Krogh, and Olaf Stursberg. [Analyzing reachability of linear dynamic systems with parametric uncertainties.](http://www.i6.in.tum.de/Main/Publications/Althoff2011b.pdf). Modeling, Design, and Simulation of Systems with Uncertainties. Springer, Berlin, Heidelberg, 2011. 69-94.

# [^K15]: Kluever, Craig A. [Dynamic systems: modeling, simulation, and control.](https://www.wiley.com/en-us/Dynamic+Systems%3A+Modeling%2C+Simulation%2C+and+Control-p-9781118289457). John Wiley & Sons, 2015.

# [^GLG08]: Girard, Antoine, and Colas Le Guernic. [Efficient reachability analysis for linear systems using support functions.]() IFAC Proceedings Volumes 41.2 (2008): 8966-8971.


# Matern 5/2 isotropic covariance function

@doc """
# Description
Constructor for the isotropic Matern 5/2 kernel (covariance)

k(x,x') = σ²(1+√5*d/ℓ + 5d²/3ℓ²)exp(-√5*d/ℓ), where d = |x-x'|
# Arguments:
* `ll::Float64`: Log of the length scale ℓ
* `lσ::Float64`: Log of the signal standard deviation σ
""" ->
type Mat52Iso <: Isotropic
    ℓ::Float64      # Length scale
    σ2::Float64     # Signal std
    Mat52Iso(ll::Float64, lσ::Float64) = new(exp(ll), exp(2*lσ))
end

function set_params!(mat::Mat52Iso, hyp::Vector{Float64})
    length(hyp) == 2 || throw(ArgumentError("Matern 5/2 only has two parameters"))
    mat.ℓ, mat.σ2 = exp(hyp[1]), exp(2.0*hyp[2])
end
get_params(mat::Mat52Iso) = Float64[log(mat.ℓ), log(mat.σ2)/2.0]
get_param_names(mat::Mat52Iso) = [:ll, :lσ]
num_params(mat::Mat52Iso) = 2

metric(mat::Mat52Iso) = Euclidean()
cov(mat::Mat52Iso, r::Float64) = mat.σ2*(1+sqrt(5)*r/mat.ℓ+5*r^2/(3*mat.ℓ^2))*exp(-sqrt(5)*r/mat.ℓ)

function grad_kern(mat::Mat52Iso, x::Vector{Float64}, y::Vector{Float64})
    r = distance(mat,x,y)
    exp_r = exp(-sqrt(5)*r/mat.ℓ)

    g1 = mat.σ2*(5*r^2/mat.ℓ^2)*((1+sqrt(5)*r/mat.ℓ)/3)*exp_r   #dK_d(log ℓ)
    g2 = 2.0*mat.σ2*(1+sqrt(5)*r/mat.ℓ+5*r^2/(3*mat.ℓ^2))*exp_r #dK_d(log σ)
    return [g1,g2]
end

function grad_stack!(stack::AbstractArray, mat::Mat52Iso, X::Matrix{Float64}, data::IsotropicData)
    nobsv = size(X,2)
    R = distance(mat, X, data)
    exp_R = exp(-sqrt(5)*R/mat.ℓ)

    for i in 1:nobsv, j in 1:i
        @inbounds stack[i,j,1] = mat.σ2*(5*R[i,j]^2/mat.ℓ^2)*((1+sqrt(5)*R[i,j]/mat.ℓ)/3)*exp_R[i,j]      # dK_dℓ
        @inbounds stack[j,i,1] = stack[i,j,1]
        @inbounds stack[i,j,2] = 2.0*mat.σ2*(1+sqrt(5)*R[i,j]/mat.ℓ+5*R[i,j]^2/(3*mat.ℓ^2))*exp_R[i,j]   # dK_dσ
        @inbounds stack[j,i,2] = stack[i,j,2]
    end

    return stack
end


using Statistics
function median2(n)
	s = sort(n)
	len = length(n)
	if len % 2 == 0
		return (s[floor(Int, len / 2) + 1] + s[floor(Int, len / 2)]) / 2
	else
		return  s[floor(Int, len / 2) + 1]
	end
end

a = [4.1, 5.6, 7.2, 1.7, 9.3, 4.4, 3.2]
b = [4.1, 7.2, 1.7, 9.3, 4.4, 3.2]

@show a b median2(a) median(a) median2(b) median(b)


function lag_matrix(x::T, lag::Int) where {T<:AbstractMatrix}
    nobs, nx = size(x)
    xlag = zeros(nobs - lag, nx * lag)
    for i = 1:lag
        xlag[1:(nobs-lag), (nx*(i-1)+1):(nx*i)] = x[(lag+1-i):(nobs-i), 1:nx]
    end
    return xlag
end

function lag_matrix(x::T, lag::Int) where {T<:AbstractVector}
    x2 = reshape(x, length(x), 1)
    return lag_matrix(x2, lag)
end

"""
    companion(A, p) -> Matrix

Construct the companion matrix for VAR coefficient matrix `A`, where `A` has
    dimensions (Kp × K). K is the number of endogenous variables, p is the
    number of lags.
"""
function companion(A, p)
    K = size(A, 2)
    subI = [Matrix(1.0I, K * (p - 1), K * (p - 1)) zeros(K * (p - 1), K)]
    Acomp = vcat(A', subI)
    return Acomp
end

function rhs_matrix(y, ylag, constant, trend, x)
    rhs = lag_matrix(y, ylag)
    if !isa(x, Nothing)
        rhs = hcat(x[ylag+1:end, :], rhs)
    end
    if trend
        ltrend = collect(1.0:size(rhs, 1))
        rhs = hcat(ltrend, rhs)
    end
    if constant
        rhs = hcat(ones(Float64, size(rhs, 1)), rhs)
    end
    return rhs
end


using Parameters
using DifferentialEquations
using ForwardDiff
using LinearAlgebra

include("cr_model.jl")

## CR exp resource, type I consumer
function expI_inteq(p)
    @unpack r, a, m, e = p
    C = r / a
    R = m / ( a * e )
    return [R , C]
end

expI_inteq_jac = jac(expI_inteq(CRPar()), cr_expI_model, CRPar())

eigvals(expI_inteq_jac)

λ_stability(expI_inteq_jac)

expI_00_jac = jac([0.0,0.0], cr_expI_model, CRPar())

λ_stability(expI_00_jac)

eigvals(expI_00_jac)

expI_k_jac = jac([3.0,0.0], cr_expI_model, CRPar())

λ_stability(expI_k_jac)

eigvals(expI_k_jac)

## CR log resource, type I consumer
function logI_inteq(p)
    @unpack r, a, k, m, e = p
    R =  m / (a * e)
    C = r * (- R + k) / (a * k)
    return [R , C]
end

logI_inteq_jac = jac(logI_inteq(CRPar()), cr_logI_model, CRPar())

eigvals(logI_inteq_jac)

λ_stability(logI_inteq_jac)

logI_00_jac = jac([0.0,0.0], cr_logI_model, CRPar())

λ_stability(logI_00_jac)

eigvals(logI_00_jac)

logI_k_jac = jac([3.0,0.0], cr_logI_model, CRPar())

λ_stability(logI_k_jac)

eigvals(logI_k_jac)

## CR log resource, type II consumer

logII_inteq_jac = jac(logII_eq(CRPar()), cr_logII_model, CRPar())

eigvals(logII_inteq_jac)

λ_stability(logII_inteq_jac)

logII_00_jac = jac([0.0,0.0], cr_logII_model, CRPar())

λ_stability(logII_00_jac)

eigvals(logII_00_jac)

logII_k_jac = jac([3.0,0.0], cr_logII_model, CRPar())

λ_stability(logII_k_jac)

eigvals(logII_k_jac)

## CR exp resource, type II consumer

# Calculate the interior equilibrium
function expII_inteq(p)
    @unpack r, a, h, m, e = p
    R = m / (a * (e - h * m))
    C = R * h * r + r / a
    return [R , C]
end

expII_inteq_jac = jac(expII_inteq(CRPar(r = 0.5, a = 1.0)), cr_expII_model, CRPar(r = 0.5, a = 1.0))

eigvals(expII_inteq_jac)

λ_stability(expII_inteq_jac)

expII_00_jac = jac([0.0,0.0], cr_expII_model, CRPar(r = 0.5, a = 1.0))

λ_stability(expII_00_jac)

eigvals(expII_00_jac)


using Symbolics, SymbolicUtils
using ModelingToolkit, DifferentialEquations
using Plots

############ The Helper Functions ############################################

def_taylor(x, ps) = sum([a*x^i for (i,a) in enumerate(ps)])
def_taylor(x, ps, p₀) = p₀ + def_taylor(x, ps)

expand_sin(x, n) = sum([(isodd(k) ? -1 : 1)*(-x)^(2k-1)/factorial(2k-1) for k=1:n])

function collect_powers(eq, x, ns; max_power=100)
    eq = substitute(expand(eq), Dict(x^j => 0 for j=last(ns)+1:max_power))

    eqs = []
    for i in ns
        powers = Dict(x^j => (i==j ? 1 : 0) for j=1:last(ns))
        push!(eqs, substitute(eq, powers))
    end
    eqs
end

function solve_coef(eqs, ps)
    vals = Dict()

    for i = 1:length(ps)
        eq = substitute(eqs[i], vals)
        vals[ps[i]] = Symbolics.solve_for(eq ~ 0, ps[i])
    end
    vals
end

function solve_newton(f, x, x₀; abstol=1e-8, maxiter=50)
    xₙ = Float64(x₀)
    fₙ₊₁ = x - f / Symbolics.derivative(f, x)

    for i = 1:maxiter
        xₙ₊₁ = substitute(fₙ₊₁, Dict(x => xₙ))
        if abs(xₙ₊₁ - xₙ) < abstol
            return xₙ₊₁
        else
            xₙ = xₙ₊₁
        end
    end
    return xₙ₊₁
end

function calc_u0(sys, d; def_val=0.0)
    S = states(sys)
    u0 = Dict{typeof(S[1]), Float64}()

    for s in S
        name_withˍt = string(s)
        name_withoutˍt = replace(name_withˍt, r"\(.*\)" => "")

        if haskey(d, name_withˍt)
            u0[s] = d[name_withˍt]
        elseif haskey(d, name_withoutˍt)
            u0[s] = d[name_withoutˍt]
        else
            u0[s] = def_val
        end
    end

    return u0
end

#################### The examples ############################################

"""
    x^5 + x = 1
"""
function test_quintic(n=4)
    @variables ϵ z a[1:n]
    x = def_taylor(ϵ, a, 1)
    eq = x^5 + ϵ*x - 1
    eqs = collect_powers(eq, ϵ, 1:n)
    vals = solve_coef(eqs, a)

    x′ = substitute(x, vals)
    X = 𝜀 -> substitute(x′, Dict(ϵ => 𝜀))

    xₚ = X(1.0)
    xₙ = solve_newton(z^5 + z - 1, z, 1.0)

    return xₚ, xₙ
end

"""
    E - e * sin(E) = M
"""
function test_kepler(n=4)
    @variables ϵ z M a[1:n]
    x = def_taylor(ϵ, a, M)
    eq = x - ϵ * expand_sin(x, n) - M
    eqs = collect_powers(eq, ϵ, 1:n)
    vals = solve_coef(eqs, a)

    x′ = substitute(x, vals)
    X = (𝜀, 𝑀) -> substitute(x′, Dict(ϵ => 𝜀, M => 𝑀))

    𝑒 = 0.01671    # The Earth eccentricity
    M₀ = π/2

    xₚ = X(𝑒, M₀)
    xₙ = solve_newton(z - 𝑒*sin(z) - M₀, z, M₀)

    return xₚ, xₙ
end

function test_rocket(n=3)
    sys, y = calc_rocket_sys(n)
    u0 = calc_u0(sys, Dict("y₁ˍt" => 1.0))
    prob = ODEProblem(sys, u0, (0, 3.0))
    sol = solve(prob; dtmax=0.01)
    X = 𝜀 -> sum([𝜀^(i-1) * sol[y[i]] for i in eachindex(y)])
    plot(sol.t, hcat([X(ϵ) for ϵ = 0.0:0.1:0.5]...))
end

function calc_rocket_sys(n)
    @variables ϵ t y[1:n](t) ∂∂y[1:n]
    x = def_taylor(ϵ, y[2:end], y[1])
    ∂∂x = def_taylor(ϵ, ∂∂y[2:end], ∂∂y[1])
    eq = ∂∂x * (1 + ϵ*x)^2 + 1
    eqs = collect_powers(eq, ϵ, 0:n)
    vals = solve_coef(eqs, ∂∂y)

    D = Differential(t)
    subs = Dict(∂∂y[i] => D(D(y[i])) for i in eachindex(y))
    eqs = [substitute(first(v), subs) ~ substitute(last(v), subs) for v in vals]

    sys = ODESystem(eqs, t)
    sys = ode_order_lowering(sys)

    sys, y
end

function test_oscillator(n=3)
    sys, y = calc_oscillator_sys(n)
    u0 = calc_u0(sys, Dict("y₁ˍt" => 1.0))
    prob = ODEProblem(sys, u0, (0, 50.0))
    sol = solve(prob; dtmax=0.01)

    X = 𝜀 -> sum([𝜀^(i-1) * sol[y[i]] for i in eachindex(y)])
    T = sol.t
    Y = 𝜀 -> exp.(-𝜀*T) .* sin.(sqrt(1 - 𝜀^2)*T) / sqrt(1 - 𝜀^2)    # exact solution

    plot(sol.t, [Y(0.1), X(0.1)])
end

function calc_oscillator_sys(n)
    @variables ϵ t y[1:n](t) ∂y[1:n] ∂∂y[1:n]
    x = def_taylor(ϵ, y[2:end], y[1])
    ∂x = def_taylor(ϵ, ∂y[2:end], ∂y[1])
    ∂∂x = def_taylor(ϵ, ∂∂y[2:end], ∂∂y[1])

    eq = ∂∂x + 2*ϵ*∂x + x
    eqs = collect_powers(eq, ϵ, 0:n)
    vals = solve_coef(eqs, ∂∂y)

    D = Differential(t)
    subs1 = Dict(∂y[i] => D(y[i]) for i in eachindex(y))
    subs2 = Dict(∂∂y[i] => D(D(y[i])) for i in eachindex(y))
    subs = subs1 ∪ subs2
    eqs = [substitute(first(v), subs) ~ substitute(last(v), subs) for v in vals]

    sys = ODESystem(eqs, t)
    sys = ode_order_lowering(sys)

    sys, y
end


"""
# module kernels

- Julia version: 1.4
- Author: niallcullinane
- Date: 2020-06-07

# Examples

```jldoctest
julia>
```
"""

module LaminartKernels

export kern_d, kern_b

# lgn to l6/l4
using NNlib, ImageFiltering, Images, OffsetArrays

function kern_d_ph(σ::Real, θ::Real, l = 4 * ceil(Int,σ) + 1)
    isodd(l) || throw(ArgumentError("length must be odd"))
    w = l>>1
    g = σ == 0 ? [1f0] : [exp(-x*cos(θ) / (2f0 * σ)) for x=-w:w]
    centered(g/sum(g))
end

function kern_d_pv(σ::Real, θ::Real, l = 4*ceil(Int,σ)+1)
    isodd(l) || throw(ArgumentError("length must be odd"))
    w = l>>1
    g = σ == 0 ? [1f0] : [exp(-x*sin(θ)/(2f0*σ)) for x=-w:w]
    centered(g/sum(g))
end

function kern_d_mh(σ::Real, θ::Real, l = 4*ceil(Int,σ)+1)
    isodd(l) || throw(ArgumentError("length must be odd"))
    w = l>>1
    g = σ == 0 ? [1f0] : [exp(x*cos(θ)/(2f0*σ)) for x=-w:w]
    centered(g/sum(g))
end

function kern_d_mv(σ::Real, θ::Real, l = 4*ceil(Int,σ)+1)
    isodd(l) || throw(ArgumentError("length must be odd"))
    w = l>>1
    g = σ == 0 ? [1f0] : [exp(x*sin(θ)/(2f0*σ)) for x=-w:w]
    centered(g/sum(g))
end

function kern_d_p(σ::Real, θ::Real, l = 4*ceil(Int,σ)+1)
    kern_d_pv(σ, θ, l) .* transpose(kern_d_ph(σ, θ, l))
end

function kern_d_m(σ::Real, θ::Real, l = 4*ceil(Int,σ)+1)
    kern_d_mv(σ, θ, l) .* transpose(kern_d_mh(σ, θ, l))
end

function kern_d(σ::Real, θ::Real, l = 4*ceil(Int,σ)+1)
    kern_d_p(σ, θ, l).-kern_d_m(σ, θ, l)
end

# function kern_d(σ::Real, θ::Real, l = 4*ceil(Int,σ)+1)
#     relu.(kern_d(σ, θ, l)) .- relu.(-1 .*(kern_d(σ, θ, l)))
# end

function kern_b(σ::Real, θ::Real, l = 4*ceil(Int,σ)+1)
    relu.(kern_d(σ, θ, l)) .+ relu.(-1f0 .*(kern_d(σ, θ, l)))
end

function gaussian_rot(σ_x::Real, σ_y::Real, θ::Real, l = 4*ceil(Int, max(σ_a,σ_b))+1)
    isodd(l) || throw(ArgumentError("length must be odd"))
    w = l>>1
    g = OffsetArray(fill(0.0f0, l, l), -w:w, -w:w)
    #todo add when σ_x or/and σ_y == 0
    for x ∈ -w:w, y ∈ -w:w
        g[x,y] = exp(-1f0/2f0*((((x*cos(θ)-y*sin(θ))/σ_x)^2f0)+(((x*sin(θ)+y*cos(θ))/σ_y)^2f0)))
    end
    centered(g/sum(g))
end

# Rotating function
function fun_R(x,y,θ)
    x*cos(θ)-y*sin(θ), x*sin(θ)+y*cos(θ)
end

end


#=
The number, 197, is called a circular prime because all rotations of the digits: 197, 971, and 719, are themselves prime.

There are thirteen such primes below 100: 2, 3, 5, 7, 11, 13, 17, 31, 37, 71, 73, 79, and 97.

How many circular primes are there below one million?
=#
include("projecteulerutils.jl")
using Primes

function iscircularprime(x)
  i = reverse(digits(x))
  a = copy(i)
  circularprime = true
  for j in i
    if !isprime(nd(a))
      return false
    end
    a = prepend!(a, [pop!(a)])
  end
  true
end

function calc()
  count = 0
  for x in primes(2_000_000)
    if iscircularprime(x)
      count += 1
    end
  end
  count
end
@time println(calc())


"""
    gegenbauer(x, p::Integer)

Evaluate the Gegenbauer polynomial with parameter `α` of degree `p` at `x` using
the three term recursion.
"""
function gegenbauer(x, p::Integer, α)
    T = typeof( 2α*x )
    p₀ = one(T)
    p₁ = 2α*x

    if p <= 0
        return p₀
    elseif p == 1
        return p₁
    end

    for n in 2:p
        p₀, p₁ = p₁, ( 2x*(n+α-1)*p₁ - (2α+n-2)*p₀ ) / n
    end

    p₁
end


# Implicitly solving the heat equation
# at cell centers and cell faces
# Equation: ∂_t T = α * ∇²(T) + S, S = [β*sin(γ*z), β]
# Boundary conditions: ∇T(z=a) = ∇T_bottom, T(z=b) = T_top
import Plots
using LinearAlgebra

# Note that we can stably take huge timesteps even when α is large
# warning, analytic solution only valid for a=0, b=1
a,b, n = 0, 1, 10               # zmin, zmax, number of cells
n̂_min, n̂_max = -1, 1            # Outward facing unit vectors
α = 100;                        # thermal diffusivity, larger means more stiff
β, γ = 10000, π;                # source term coefficients
Δt = 1000;                      # timestep
N_t = 10;                       # number of timesteps
FT = Float64;                   # float type
Δz = FT(b-a)/FT(n)
Δz² = Δz^2;
∇²_op = [1/Δz², -2/Δz², 1/Δz²]; # row or Laplacian operator
T_init = 1;                     # Initial temperature, shouldn't matter for steady solution
∇T_bottom = 10;                  # Temperature gradient at the top
T_top = 1;                      # Temperature at the bottom
use_sin_forcing = true
S(z) = use_sin_forcing ? β*sin(γ*z) : β # source term, (sin for easy integration)
zf = range(a, b, length=n+1);
zc = map(i-> zf[i]+Δz/2, 1:n);

# Analytic steady state solution:
# ∂∂T∂∂z = -S(z)/α = -β*sin(γ*z)/α
# ∂T∂z = β*cos(γ*z)/(γ*α)+c1
# T(z) = β*sin(γ*z)/(γ^2*α)+c1*z+c2 # generic solution
# Apply bc: ∂T∂z(a) = β*cos(γ*a)/(γ*α)+c1 = ∇T_bottom => c1 = ∇T_bottom-β*cos(γ*a)/(γ*α)
# Apply bc: T(b) =    β*sin(γ*b)/(γ^2*α)+c1*b+c2 = T_top =>
#         c2 = T_top-(β*sin(γ*b)/(γ^2*α)+c1*b)
function T_analytic_sin(z)
    c1 = ∇T_bottom-β*cos(γ*a)/(γ*α)
    c2 = T_top-(β*sin(γ*b)/(γ^2*α)+c1*b)
    return β*sin(γ*z)/(γ^2*α)+c1*z+c2
end

# Analytic steady state solution:
# ∂∂T∂∂z = -S(z)/α = -β/α
# ∂T∂z = -z*β/α+c1
# T(z) = -β*z^2/(2*α)+c1*z+c2 # generic solution
# Apply bc: ∂T∂z(a) = -a*β/α+c1 = ∇T_bottom => c1 = ∇T_bottom+a*β/α
# Apply bc: T(b) =    -β*b^2/(2*α)+c1*b+c2 = T_top =>
#         c2 = T_top-(-β*b^2/(2*α)+c1*b)
function T_analytic_uniform(z)
    c1 = ∇T_bottom+a*β/α
    c2 = T_top-(-β*b^2/(2*α)+c1*b)
    return -β*z^2/(2*α)+c1*z+c2
end

T_analytic(z) =
    use_sin_forcing ? T_analytic_sin(z) : T_analytic_uniform(z)

#####
##### T ∈ cell centers, ∇T ∈ cell faces
#####
println("Solve for T ∈ cell centers")
T = zeros(FT, n);

# Equations: derivation to matrix form
# ∂_t T = α * ∇²(T) + S
# (T^{n+1}-T^n) = Δt (α * ∇²(T^{n+1}) + S)
# (T^{n+1} - Δt α * ∇²(T^{n+1})) = T^n + Δt*S
# (I - Δt α * ∇²) T^{n+1} = T^n + Δt*S

# Derive Dirichlet boundary stencil & source:
# ∂_t T = α * (T[i-1]+T[g]-2*T[i])/Δz² + S
# ∂_t T = α * (T[i-1]+(2 T[b] - T[i])-2*T[i])/Δz² + S
# ∂_t T = α * (T[i-1]-3*T[i])/Δz² + S + α * (2 T[b])/Δz²

# Derive Neumann boundary stencil & source:
# ∇T_bottom*n̂ = (T[g] - T[i])/Δz,     n̂ = [-1,1] ∈ [zmin,zmax]
# ∇T_bottom*n̂ = -(T[i] - T[g])/Δz,    n̂ = [-1,1] ∈ [zmin,zmax]
# ∂_t T = α * ((T[ii] - T[i])/Δz - (T[i] - T[g])/Δz)/Δz + S
# ∂_t T = α * ((T[ii] - T[i])/Δz + ∇T_bottom*n̂)/Δz + S
# ∂_t T = α * (T[ii] - T[i])/Δz² + S + α*∇T_bottom*n̂/Δz

# Note that size(A) = (n,n) due to 1 Dirichlet and 1 Neumann.
∇² = Tridiagonal(
    ones(FT, n-1) .* ∇²_op[1],
    ones(FT, n)   .* ∇²_op[2],
    ones(FT, n-1) .* ∇²_op[3]
);
# Modify boundary stencil to account for BC
∇².d[1] = -1/Δz²
∇².d[end] = -3/Δz²

A = LinearAlgebra.I - Δt.* α .* ∇²

# Compute boundary sources: α * (2 T[b])/Δz², α * ∇T_bottom*n̂/Δz
AT_b = zeros(FT, n);
AT_b[1] = α*∇T_bottom*n̂_min/Δz;
AT_b[end] = α*2*T_top/Δz²;

# Set initial condition:
T .= T_init;
@inbounds for i_t in 1:N_t
    # change in temperature per iter:
    b_rhs = T .+ Δt .* (S.(zc) .+ AT_b)
    T_new = A \ b_rhs
    ΔT_norm = sum(T .- T_new)/length(T)
    @show ΔT_norm # watch ΔT norm
    T .= T_new
end

# Interpolate to, and plot on, cell faces
zf = range(a, b, length=n+1);
Tf = zeros(FT, n+1);
Tf[2:end-1] = [(T[i]+T[i+1])/2 for i in 1:length(T)-1]
# ∇T_bottom*n̂ = (T[g] - T[i])/Δz, T[g] = 2*T[b] - T[i]
# ∇T_bottom*n̂ = ((2*T[b] - T[i]) - T[i])/Δz
# Δz*∇T_bottom*n̂ = 2*T[b] - 2*T[i]
# (Δz*∇T_bottom*n̂ + 2*T[i])/2 = T[b]
Tf[1] = (Δz*∇T_bottom*n̂_min + 2*T[1])/2
Tf[end] = T_top
p1 = Plots.plot(zf, T_analytic.(zf), label="analytic", markershape=:circle, markersize=6)
p1 = Plots.plot!(p1, zf, Tf, label="numerical", markershape=:diamond)
p1 = Plots.plot!(p1, title="T ∈ cell centers")

p3 = Plots.plot(zf, abs.(Tf .- T_analytic.(zf)), label="error", markershape=:circle, markersize=6)
p3 = Plots.plot!(p3, title="T ∈ cell centers")

#####
##### T ∈ cell faces, ∇T ∈ cell centers
#####

println("Solve for T ∈ cell faces")
T = zeros(FT, n+1);

# Equations: derivation to matrix form
# ∂_t T = α * ∇²(T) + S
# (T^{n+1}-T^n) = Δt (α * ∇²(T^{n+1}) + S)
# (T^{n+1} - Δt α * ∇²(T^{n+1})) = T^n + Δt*S
# (I - Δt α * ∇²) T^{n+1} = T^n + Δt*S

# Derive Dirichlet boundary stencil & source:
# ∂_t T = α * (T[i-1]+T[b]-2*T[i])/Δz² + S
# ∂_t T = α * (T[i-1]-2*T[i])/Δz² + S + α * T[b] / Δz²

# Derive Neumann boundary stencil & source:
# ∇T_bottom*n̂= (T[g] - T[i])/(2Δz),     n̂ = [-1,1] ∈ [zmin,zmax]
# T[i] + 2*Δz*∇T_bottom*n̂ = T[g]
# ∂_t T = α * (((T[i] + 2*Δz*∇T_bottom*n̂) - T[b])/Δz - (T[b] - T[i])/Δz)/Δz + S
# ∂_t T = α * (((T[i]) - T[b])/Δz - (T[b] - T[i])/Δz)/Δz + S + α/Δz²*2*Δz*∇T_bottom*n̂
# ∂_t T = α * (2*T[i] - 2*T[b])/Δz² + S + 2α/Δz*∇T_bottom*n̂

# Note that size(A) = (n,n) due to 1 Dirichlet and 1 Neumann.
∇² = Tridiagonal(
    ones(FT, n) .* ∇²_op[1],
    ones(FT, n+1)   .* ∇²_op[2],
    ones(FT, n) .* ∇²_op[3]
);

# Modify boundary stencil to account for BCs
∇².d[1] = -2/Δz²
∇².du[1] = +2/Δz²

# Modify boundary stencil to account for BCs
∇².du[n] = 0  # modified stencil
∇².d[n+1] = 0 # to ensure `∂_t T = 0` at `z=zmax`
∇².dl[n] = 0  # to ensure `∂_t T = 0` at `z=zmax`

A = LinearAlgebra.I - Δt.* α .* ∇²

# Compute boundary source: α * T[b] / Δz²
AT_b = zeros(FT, n+1);
AT_b[1] = α*2/Δz*∇T_bottom*n̂_min;
AT_b[end-1] = α*T_top/Δz²;

# Set initial condition:
T .= T_init;
T[n+1] = T_top
@inbounds for i_t in 1:N_t
    interior = 1:n
    b_rhs = T .+ Δt .* (S.(zf) .+ AT_b)
    T_new = A \ b_rhs
    ΔT_norm = sum(T .- T_new)/length(T)
    @show ΔT_norm # watch ΔT norm
    T .= T_new
end

p2 = Plots.plot(zf, T_analytic.(zf), label="analytic", markershape=:circle, markersize=6)
p2 = Plots.plot!(p2, zf, T, label="numerical", markershape=:diamond)
p2 = Plots.plot!(p2, title="T ∈ cell faces")

p4 = Plots.plot(zf, abs.(T .- T_analytic.(zf)), label="error", markershape=:circle, markersize=6)
p4 = Plots.plot!(p4, title="T ∈ cell faces")

Plots.plot(p1, p3, p2, p4)


# This file calculates scattering from the Helmholtz equation via the adaptive spectral method.
# Δu + k^2u = 0,
# u|Γ = 0,
# u^i = e^{im k x⋅d},
# u = u^i + u^s.
# The normal derivative ∂u/∂n of the entire wave is calculated on the sound-soft boundaries.
# The scattered wave is calculated by convolving [∂u/∂n] with the fundamental solution.
# Then, the total wave is obtained by summing the incident and the scattered waves.

using ApproxFun, SingularIntegralEquations
include("Scatteraux.jl")

k = 50.
ω = 2π
d = (1,-1)
d = d[1]/hypot(d[1],d[2]),d[2]/hypot(d[1],d[2])
ui = (x,y) -> exp(im*k*(d[1]*x+d[2]*y))

# The Helmholtz Green's function, split into singular and nonsingular pieces.
g1 = (x,y) -> -besselj0(k*abs(y-x))/2
g2 = (x,y) -> x == y ? -(log(k/2)+γ)/2/π + im/4 : im/4*hankelh1(0,k*abs(y-x)) - g1(x,y).*logabs(y-x)/π

# A variety of domains.

#dom = Circle(0.0,1/π)∪Segment(-1.0im,1.0)
#dom = ∪(Segment.([-2.5-.5im,-1.5+.5im,-.5-.5im,.5+.5im,1.5-.5im],[-1.5-.5im,-.5+.5im,.5-.5im,1.5+.5im,2.5-.5im]))
#dom = ∪(Segment.([-1.0-0.4im,0.1+0.4im,-0.9-0.5im],[-0.1+0.4im,1.0-0.4im,0.9-0.5im]))
#dom = ∪(Segment.([-1.0-0.4im,-0.5-0.4im,0.1+0.4im,0.2+0.0im,-1.4-0.75im],[-0.1+0.4im,-0.2+0.0im,1.0-0.4im,0.5-0.4im,1.4-0.75im]))
#dom = ∪(Circle.([0.,-1.0im],[0.5,0.25]))∪(Segment.([-1.5,0.5-1.0im,-0.5+1.0im],[-0.5-1.0im,1.5,0.5+1.0im]))
dom = ChebyshevInterval()

sp = Space(dom)
cwsp = CauchyWeight(sp⊗sp,0)
uiΓ,⨍ = Fun(t->ui(real(t),imag(t)),sp),DefiniteLineIntegral(dom)

@time G = GreensFun(g1,cwsp;method=:Cholesky) + GreensFun(g2,sp⊗sp;method=:Cholesky)

@time ∂u∂n = ⨍[G]\uiΓ
println("The length of ∂u∂n is: ",ncoefficients(∂u∂n))
us = (x,y) -> -logkernel(g1,∂u∂n,complex(x,y))-linesum(g2,∂u∂n,complex(x,y))


"""
    Berlage(; <keyword arguments>)

Create a Berlage wavelet.

# Arguments

**Keyword arguments**

* `dt::Real=0.002`: sampling interval in secs.
* `f0::Real=20.0`: central frequency in Hz.
* `m::Real=2`: exponential parameter of Berlage wavelet.
* `alpha::Real=180.0`: alpha parameter of Berlage wavelet in rad/secs.
* `phi0::Real`: phase rotation in radians.

# Example
```julia
julia> w = Berlage(); plot(w);
```

**Reference**

* Aldridge, David F., 1990, The berlage wavelet: GEOPHYSICS, 55, 1508--1511.
"""

function Berlage(; dt::Real=0.002, f0::Real=20.0, m::Real=2, alpha::Real=180.0,
                 phi0::Real=0.0)

    nw = floor(Int, 2.2/(f0*dt))
    t = dt*collect(0:1:nw-1)
    w = (t.^m).*exp(-alpha*t).*cos(2*pi*f0*t + phi0);
    w = w/maximum(w)

end


export size_criterion, sqrt_criterion, change_criterion, log_criterion, epsilon_criterion

"""
Stops when the distance between far items achieves the given `e`
"""
function epsilon_criterion(e)
   (dmaxlist, database) -> dmaxlist[end] < e
end

"""
Stops when the number of far items are equal or larger than the given `maxsize`
"""
function size_criterion(maxsize)
   (dmaxlist, database) -> length(dmaxlist) >= maxsize
end

"""
Stops when the number of far items are equal or larger than the square root of the size of the database
"""
function sqrt_criterion()
   (dmaxlist, database) -> length(dmaxlist) >= Int(length(database) |> sqrt |> round)
end

"""
Stops when the number of far items are equal or larger than logarithm-2 of the size of the database
"""
function log_criterion()
   (dmaxlist, database) -> length(dmaxlist) >= Int(length(database) |> log2 |> round)
end

"""
Stops the process whenever the maximum distance converges, i.e., after `window` far items the maximum distance
change is below or equal to the allowed tolerance `tol`
"""
function change_criterion(tol=0.001, window=3)
    mlist = Float64[]
    count = 0.0
    function stop(dmaxlist, database)
        count += dmaxlist[end]

        if length(dmaxlist) % window != 1
            return false
        end
        push!(mlist, count)
        count = 0.0
        if length(dmaxlist) < 2
            return false
        end

        s = abs(mlist[end] - mlist[end-1])
        return s <= tol
    end

    return stop
end

"""
It nevers stops by side, it explores the entire dataset making a full farthest first traversal
"""
function salesman_criterion()
    function stop(dmaxlist, dataset)
        return false
    end
end


"""
**Distance to a primary producer**

This function measures, for every species, its shortest path to a primary
producer using matrix exponentiation. A primary producer has a value of 1,
a primary consumer a value of 2, and so forth.

"""
function distance_to_producer(L::Array{Int64, 2})

    # We identify producers
    is_producer = vec(sum(L, dims = 2) .== 0)

    # Producers have a distance of 1
    d = zeros(Int64, length(is_producer))
    d[is_producer] .= 1

    # We work on a copy of the matrix with no self-loops
    K = copy(L)
    for i in eachindex(d)
        K[i ,i] = 0
    end

    # We loop as long as there are species with unknown distance
    i = 1
    while (i < length(d))|(sum(d) == 0)
        connected_at_length = (K^i * is_producer) .> 0
        d[(d .== 0) .* (connected_at_length)] .= i+1
        i = i+1
    end
    return d
end

"""
**Trophic rank**

Based on the average distance of preys to primary producers. Specifically, the
rank is defined as the average of the distance of preys to primary producers
(recursively). Primary producers always have a trophic rank of 1.

"""
function trophic_rank(L::Array{Int64, 2})
    # Average of positve elements, 0 otherwise
    nonzeromean = (x) -> maximum(x) == 0 ? 0 : mean(x[x.>0])
    d = distance_to_producer(L)
    dL = L .* d'
    TL = zeros(length(d))
    for i in eachindex(d)
        TL[i] = nonzeromean(dL[i,:]) + 1
    end
    return TL
end


# FeynmanTheory.jl

"""
frohlichalpha(ε_Inf,ε_S,freq,m_eff)

    Calculates the Frohlich alpha parameter, for a given dielectric constant,
    frequency (f) of phonon in Hertz, and effective mass (in units of the
    bare electron mass).

    See Feynman 1955:
    http://dx.doi.org/10.1103/PhysRev.97.660

"""
function frohlichalpha(ε_optic,ε_static,freq,m_eff)
    ω=freq*2*pi #frequency to angular velocity
    # Note: we need to add a 4*pi factor to the permitivity of freespace.
    # This gives numeric agreement with literature values.  This is required as
    # the contemporary 1950s and 1960s literature implicitly used atomic units,
    # where the electric constant ^-1 has this factor baked in, k_e=1/(4πϵ_0).
    α= 0.5/(4*π*ε_0) *              # Units: m/F
       (1/ε_optic - 1/ε_static) *   # Units: none
       (q^2/(hbar*ω)) *             # Units: F
       sqrt(2*me*m_eff*ω/ħ)         # Units: 1/m
    return (α)
end

#####
# Athermal (Feynman 1955) model
# Set up equations for the polaron free energy, which we will variationally improve upon

# Integrand of (31) in Feynman I (Feynman 1955, Physical Review, "Slow electrons...")
fF(τ,v,w)=(w^2 * τ + (v^2-w^2)/v*(1-exp(-v*τ)))^-0.5 * exp(-τ)
# (31) in Feynman I
AF(v,w,α)=π^(-0.5) * α*v * quadgk(τ->fF(τ,v,w),0,Inf)[1]
# (33) in Feynman I
F(v,w,α)=(3/(4*v))*(v-w)^2-AF(v,w,α)

# Let's wrap the Feynman athermal variation approximation in a simple function
"""
    feynmanvw(α; v=7.0, w=6.0)

    Calculate v and w variational polaron parameters,
    for the supplied α Frohlich coupling.
    This version uses the original athermal action (Feynman 1955).
	Returns v,w.
"""
function feynmanvw(α; v=7.0, w=6.0) # v,w defaults
    initial=[v,w]
    # Main use of these bounds is stopping v or w going negative, at which you get a NaN error as you are evaluating log(-ve Real)
    lower=[0.1,0.1]
    upper=[1000.0,1000.0]

    myf(x) = F(x[1],x[2],α) # Wraps the function so just the two variational params are exposed

    # Now updated to use Optim.jl > 0.15.0 call signature (Julia >0.6 only)
    res=optimize(OnceDifferentiable(myf, initial; autodiff = :forward),
                 lower, upper, initial, Fminbox( BFGS() ) )
    # specify Optim.jl optimizer. This is doing all the work.

    v,w=Optim.minimizer(res)

    return v,w
end

# Hellwarth et al. 1999 PRB - Part IV; T-dep of the Feynman variation parameter

# Originally a Friday afternoon of hacking to try and implement the T-dep electron-phonon coupling from the above PRB
# Which was unusually successful! And more or less reproduced Table III

# In Julia we have 'Multiple dispatch', so let's just construct the free
# energies (temperature-dependent) with the same name as before, but withthe
# thermodynamic beta where required

# Define Osaka's free-energies (Hellwarth1999 version) as Julia functions
# Equation numbers follow above Hellwarth et al. 1999 PRB
# 62b
A(v,w,β)=3/β*( log(v/w) - 1/2*log(2*π*β) - log(sinh(v*β/2)/sinh(w*β/2)))
# 62d
Y(x,v,β)=1/(1-exp(-v*β))*(1+exp(-v*β)-exp(-v*x)-exp(v*(x-β)))
# 62c integrand
f(x,v,w,β)=(exp(β-x)+exp(x))/(w^2*x*(1-x/β)+Y(x,v,β)*(v^2-w^2)/v)^(1/2)
# 62c
B(v,w,β,α) = α*v/(sqrt(π)*(exp(β)-1)) * quadgk(x->f(x,v,w,β),0,β/2)[1]
# 62e
C(v,w,β)=3/4*(v^2-w^2)/v * (coth(v*β/2)-2/(v*β))
# 62a
F(v,w,β,α)=-(A(v,w,β)+B(v,w,β,α)+C(v,w,β))

# Can now evaluate, e.g.
# F(v,w,β,α)=F(7.2,6.5,1.0,1.0)
# BUT - this is just the objective function! (The temperature-dependent free-energy.) Not the optimised parameters.
# Also there's a scary numeric integration (quadgk) buried within...

# Unbundled version of the finite temperature one
"""
    feynmanvw(α, βred; v=7.1, w=6.5, verbose::Bool=false)

    Calculate v and w variational polaron parameters,
    for the supplied α Frohlich coupling and βred reduced thermodynamic temperature.
    This uses the Osaka finite temperature action.
    Returns v,w.
"""
function feynmanvw(α, βred; v=7.1, w=6.5, verbose::Bool=false) # v,w defaults
    # Initial v,w to use
    initial=[v,w]
    # Main use of these bounds is stopping v or w going negative, at which you get a NaN error as you are evaluating log(-ve Real)
    lower=[0.1,0.1]
    upper=[100.0,100.0]

    myf(x) = F(x[1],x[2],βred,α) # Wraps the function so just the two variational params are exposed

    # Now updated to use Optim > 0.15.0 call signature (Julia >0.6 only)
    res=optimize(OnceDifferentiable(myf, initial; autodiff = :forward),
                 lower, upper, initial, Fminbox( BFGS() ) )
    # specify Optim.jl optimizer. This is doing all the work.

    if Optim.converged(res) == false
        print("\tWARNING: Failed to converge to v,w soln? : ",Optim.converged(res) )
    end

    if verbose
        println()
        show(res)
    end

    v,w=Optim.minimizer(res)
    return v,w
end


#
# This is an example of estimation of a portfolio model
# as in Hansen and Singleton, 1982. The data comes from
# Tauchen, JBES, 1986 (available by ftp).
# The notation follows my lecture notes. In particular
# gamma is the coefficient of RRA. Try out different sets
# of intruments and different lag lengths. You will see
# that beta is stable across  models, but that gamma is not.
#
# Michael Creel, Oct. 24, 2000
# revised 7 Nov. 2002
# translated to Octave 9/9/2003
# translated to Julia 07 Aug. 2017
# michael.creel@uab.es

# The main thing you have to to is define the moments for estimation,
# as in the function that follows immediately, and ensure that the
# efficient weight matrix is estimated appropriately. The rest
# is already programmed.

# the following function defines the moment conditions, given the parameter
# vector and the instruments. It returns the individual contributions
# to allow estimation of the covariance matrix of the moments.
using DelimitedFiles, Statistics, Econometrics
function portfolio_moments(theta, data)
	# parameters
	beta = theta[1]
	gam = theta[2]
	# data items
	c = data[:,1]
	r = data[:,2]
    inst = data[:,3:end]
	#  form error function
	# note that c = c_t / c_t-1 (for stationarity) was done in data preparation
	e = 1.0 .- beta*(1.0 .+ r) .* (c .^ (-gam))
	# cross with instruments
	m = e.*inst
end

function main()
data = readdlm("tauchen.data")
c = data[:,1]
p = data[:,2]
d = data[:,3]
# form net return and stationary consumption
r = (p + d) ./ lag(p,1) .- 1.0
c = c ./ lag(c,1); # ensure stationarity
# choose maximal lag of instruments.
max_lag = 1
inst = [c r d p]
inst = lags(inst, max_lag)
inst, junk, junk = stnorm(inst)
inst = [ones(size(inst,1),1) inst]
# drop rows with missing values
data = [c r inst]
data = data[max_lag+1:end,:]
# do estimation
moments = θ -> portfolio_moments(θ, data)
weight = 1.0
names = ["beta","gamma"]
# initial consistent estimate
theta = [0.9, 1]
thetahat, obj_value, D, ms, convergence = gmm(moments, theta, weight)
# second step with efficient weight matrix
omega = 50.0*cov(ms)
weight = inv(omega)
title = "Two step GMM estimation of portfolio model"
results = gmmresults(moments, thetahat, weight, title, names)
# CUE GMM
title = "CUE GMM estimation of portfolio model"
gmmresults(moments, results[1], "", title, names)
return
end
main()


"""
Created on Wed Sep 23 2020

@author: Yoann Pradat

Multiplicative update algorithms for α divergence.
"""

"""
    _update_nmf_mu_α!(V, WH, W, H, α, ϵ)

"""
function _update_nmf_mu_α!(V::Matrix{T}, WH::Matrix{T}, W::Matrix{T}, H::Matrix{T}, α::Real, ϵ::Real) where T <: Real
    if α == 0
        #### dual Kullback-Leibler

        #### update H
        #### NOTE: if W has columns normalized to sum to 1, the normalization matrix
        #### (transpose(W) * ones(size(V))) .^ (-1) is just ones(K, N)
        H .= H .* exp.((transpose(W) * ones(size(V))) .^ (-1)  .* (transpose(W) * log.(V ./ (WH .+ ϵ))))
        WH .= W * H

        #### update W
        #### NOTE: if H has columns normalized to sum to 1, the normalization matrix
        #### (transpose(H) * ones(size(V))) .^ (-1) is just ones(K, N)
        W .= W .* exp.((ones(size(V)) * transpose(H)) .^ (-1) .* (log.(V ./ (WH .+ ϵ)) * transpose(H)))
        WH .= W * H

    else
        #### general α divergence updates for α != 0

        #### A.Cichocki, S.Amari, R.Zdunek, R.Kompass, G.Hori, and Z.He,
        #### "Extended SMART algorithms for non-negative matrix factorization,"
        #### Artificial Intelligence and Soft Computing, 2006.
        #### Equation (44)

        #### note: it is not clear in eqs (45) and (46) how the normalization
        #### made the denominators disappear in (44).

        # #### case of eqs (45), (46)
        #
        # # update W
        # W .= W .* (((V .+ ϵ) ./ (WH .+ ϵ)) .^ α * transpose(H)) .^ (1/α)

        # #### normalize to unit sums the columns of W if specified
        # #### note: order not as in (45-46)
        # cs_W = sum(W, dims=1)
        # W .= W .* repeat(cs_W .^ -1, size(W, 1), 1)

        # WH .= W * H

        # # update H
        # H .= H .* ((transpose(W) * ((V .+ ϵ)./(WH .+ ϵ)) .^ α)) .^ (1/α)
        # WH .= W * H

        #### eqs (44)

        # update W
        W .= W .* ((((V .+ ϵ) ./ (WH .+ ϵ)) .^ α * transpose(H)) ./ (ones(size(V)) * transpose(H)) .+ ϵ) .^ (1/α)
        WH .= W * H

        # update H
        H .= H .* (((transpose(W) * ((V .+ ϵ) ./ (WH .+ ϵ)) .^ α)) ./ (transpose(W) * ones(size(V))) .+ ϵ) .^ (1/α)
        WH .= W * H

    end
end


module SkillRating

    const q = log(10)/400

    g(rd) = 1/sqrt(1 + 3q^2 * (rd^2)/π^2)
    Es(r, r2, rd2) = 1/(1+10^(-g(rd2) * (r - r2) / 400))
    d2(r, rjs, RDjs) = 1/((q^2) * sum(g.(RDjs).^2 .* Es.(r, rjs, RDjs).* (1 .- Es.(r, rjs, RDjs))))

    """
        glicko_rating(r, RD, rjs::Vector, RDjs::Vector, sjs::Vector)
        glicko_rating(r, RD, rj::T, RDj::T, sj::T) where T <: Real

    # Arguments:
    * `r`: Rating of player
    * `RD`: Rating deviation of player
    * `rjs`: Ratings of oppenents within rating period.
    * `RDjs`: Rating devations of oppenents within rating period.
    * `sjs`: Game outcomes from the perspective of the "player": 1 - won, 0.5 - draw, 0 - lost.

    Calculate a new glicko rating for games within a rating period or for a single game.

    See: http://www.glicko.net/glicko/glicko.pdf and Glickman, Mark E., "Parameter estimation
    in large dynamic paired comparison experiments" (1999) Applied Statistics, 48, 377-394
    (http://www.glicko.net/research/glicko.pdf)
    """
    glicko_rating(r, RD, rjs::Vector, RDjs::Vector, sjs::Vector) = r + q/(1/RD^2 + 1/d2(r, rjs, RDjs)) * sum(g.(RDjs) .* (sjs .- Es.(r, rjs, RDjs)))

    """
        glicko_RD(r, RD, rjs::Vector, RDjs::Vector, sjs::Vector)
        glicko_RD(r, RD, rj::T, RDj::T, sj::T) where T <: Real

    # Arguments:
    * `r`: Rating of player
    * `RD`: Rating deviation of player
    * `rjs`: Ratings of oppenents within rating period.
    * `RDjs`: Rating devations of oppenents within rating period.
    * `sjs`: Game outcomes from the perspective of the "player": 1 - won, 0.5 - draw, 0 - lost.

    Calculate a new glicko rating deviation for games within a rating period  or for a single game.

    See: http://www.glicko.net/glicko/glicko.pdf and
    Glickman, Mark E., "Parameter estimation in large dynamic paired comparison experiments" (1999) Applied Statistics, 48, 377-394
    (http://www.glicko.net/research/glicko.pdf)
    """
    glicko_RD(r, RD, rjs::Vector, RDjs::Vector, sjs::Vector) = sqrt(1/(1/RD^2 + 1/d2(r, rjs, RDjs)))

    glicko_rating(r, RD, rj::T, RDj::T, sj::T) where T <: Real = glicko_rating(r, RD, [rj], [RDj], [sj])
    glicko_RD(r, RD, rj::T, RDj::T, sj::T) where T <: Real = glicko_RD(r, RD, [rj], [RDj], [sj])

    """
        glicko_RD_increase(rd; c = 63.2, maxrd=350)

    # Arguments:
    * `RD`: Rating deviation at end of rating period.
    * `c`: Increase in uncertainty between rating periods.
    * `maxRD`: Rating deviation of a player who does not play.

    Calculate rating deviation after the end of a rating period.

    See: http://www.glicko.net/glicko/glicko.pdf and
    Glickman, Mark E., "Parameter estimation in large dynamic paired comparison experiments" (1999) Applied Statistics, 48, 377-394
    (http://www.glicko.net/research/glicko.pdf)
    """
    glicko_RD_increase(RD; c = 63.2, maxRD=350) = min(sqrt(RD^2+c^2), maxRD)

    """
        glicko_c(rating_perods_to_max, maxrd, medianrd)

    Calculate c, the increase in uncertainty between rating periods.

    See: http://www.glicko.net/glicko/glicko.pdf and
    Glickman, Mark E., "Parameter estimation in large dynamic paired comparison experiments" (1999) Applied Statistics, 48, 377-394
    (http://www.glicko.net/research/glicko.pdf)
    """
    glicko_c(rating_perods_to_max, maxrd, medianrd) = sqrt((maxrd^2 - medianrd^2)/rating_perods_to_max)

    """
    elo_expected(RA, RB)

    # Arguments:
    * `RA`: Rating of player A
    * `RB`: Rating of player B

    Calculate the expected outcome using glicko.

    See: http://www.glicko.net/glicko/glicko.pdf and
    Glickman, Mark E., "Parameter estimation in large dynamic paired comparison experiments" (1999) Applied Statistics, 48, 377-394
    (http://www.glicko.net/research/glicko.pdf)
    """
    glicko_expected(r1, RD1, r2, RD2) = 1 / (1 + 10^(-g(sqrt(RD1^2 + RD2^2)) * (r1 - r2)/400))

    """
        elo_rating(r, rjs::Vector, sjs::Vector; k = 32)

    # Arguments:
    * `r`: Rating of player
    * `rjs`: Ratings of oppenents within rating period.
    * `sjs`: Game outcomes from the perspective of the "player": 1 - won, 0.5 - draw, 0 - lost.

    Calculate the elo rating change.
    See: https://en.wikipedia.org/wiki/Elo_rating_system#Mathematical_details
    """
    elo_rating(r, rjs::Vector, sjs::Vector; k = 32) = r + k * (sum(sjs) - sum(elo_expected.(r, rjs)))

    """
    elo_expected(RA, RB)

    # Arguments:
    * `RA`: Rating of player A
    * `RB`: Rating of player B

    Calculate the expected outcome using elo.
    See: https://en.wikipedia.org/wiki/Elo_rating_system#Mathematical_details
    """
    elo_expected(RA, RB) = 1 / (1 + 10 ^ ((RB - RA) / 400))

    Eg2(μ, μj, ϕj) = 1/(1+exp(-gg2(ϕj)*(μ - μj)))
    gg2(ϕ) = 1/sqrt(1 + 3*(ϕ^2)/π^2)
    compute_v(μ, μjs, ϕjs) = 1/sum(gg2.(ϕjs).^2 .* Eg2.(μ, μjs, ϕjs) .* (1 .- Eg2.(μ, μjs, ϕjs)))
    compute_Δ(μ, μjs, ϕjs, sjs, v) = v * sum(gg2.(ϕjs) .* (sjs .- Eg2.(μ, μjs, ϕjs)))
    ϕincrease(ϕ, σ) = sqrt(ϕ^2 + σ^2)

    """
        glicko2(μ::Real, ϕ::Real, σ::Real, μjs::Vector{Real}, ϕjs::Vector{Real}, sjs::Vector{Real}; τ=0.05, ϵ=10^-6)

    # Arguments:
    * `μ`: Rating of player
    * `ϕ`: Rating deviation of player
    * `σ`: Rating volatility
    * `μjs`: Ratings of oppenents within rating period
    * `μjs`: Rating deviations within rating period
    * `sjs`: Game outcomes from the perspective of the "player": 1 - won, 0.5 - draw, 0 - lost.
    * `τ=0.5`: smaller values constrain the change in volatility over time, reasonable range between 0.3 and 1.2
    * `ϵ=10^-6`: Convergence bound

    Returns tuple `(μ, ϕ, σ)`: glicko2 rating, rating deviation and rating volatility for games in one rating period.

    See: http://www.glicko.net/glicko/glicko2.pdf and
    Glickman, Mark E., "Dynamic paired comparison models with stochastic variances" (2001),
    Journal of Applied Statistics, 28, 673-689.
    (http://www.glicko.net/research/dpcmsv.pdf)
    """
    function glicko2(μ::Real, ϕ::Real, σ::Real, μjs::Vector{R}, ϕjs::Vector{R}, sjs::Vector{S}; τ=0.5, ϵ=10^-6) where {R <: Real, S <: Real}
        v = compute_v(μ, μjs, ϕjs)
        Δ = compute_Δ(μ, μjs, ϕjs, sjs, v)
        a = log(σ^2)
        f(x) = exp(x)*(Δ^2 - ϕ^2  - v - exp(x)) / (2 * (ϕ^2 + v + exp(x))^2) - (x - a)/τ^2
        A = a
        if Δ^2 > ϕ^2 + v
            B = log(Δ^2 - ϕ^2  - v)
        else
            k = 1
            while (f(a - k*τ) < 0)
                k += 1
            end
            B = a - k*τ
        end
        fA = f(A)
        fB = f(B)
        while abs(B - A) > ϵ
            C = A + (A-B) * fA/ (fB-fA)
            fC = f(C)
            if fC * fB < 0
                A = B
                fA = fB
            else
                fA /= 2
            end
            B = C
            fB = fC
        end
        σnew = exp(A/2)
        ϕincreased = sqrt(ϕ^2 + σnew^2)
        ϕnew = 1/sqrt(1/ϕincreased^2 + 1/v)
        μnew = μ + ϕnew^2 * sum(gg2.(ϕjs) .* (sjs .- Eg2.(μ, μjs, ϕjs)))
        μnew, ϕnew, σnew
    end

    """
        glicko1_to_glicko2(r::real, RD::Real)

    Convert ratings `r` and rating deviations `RD` from glicko1 (Elo-like) scale to glicko 2 scale (μ and ϕ).
    """
    glicko1_to_glicko2(r::Real, RD::Real) = ((r-1500)/173.7178, RD/173.7178)

    export glicko_rating, glicko_RD, glicko_RD_increase, glicko_c, glicko_expected, elo_rating, elo_expected, glicko2
end


mutable struct Edge
    rev::Int
    from::Int
    to::Int
    cap::Int
end

function push!(G::Vector{Vector{Edge}}, from, to, cap)
    fromrev = length(G[from]) + 1
    torev = length(G[to]) + 1
    Base.push!(G[from], Edge(torev, from, to, cap))
    Base.push!(G[to], Edge(fromrev, to, from, 0))
end

function run_flow!(G::Vector{Vector{Edge}}, e::Edge, f)
    e.cap -= f
    G[e.to][e.rev].cap += f
    return G
end

function fodfs(G, v, t, f, seens)
    v == t && return f

    seens[v] = true
    for e ∈ G[v]
        if seens[e.to] || e.cap == 0
            # Skip
        else
            flow = fodfs(G, e.to, t, min(f, e.cap), seens)
            if flow == 0
                # Skip
            else
                run_flow!(G, e, flow)
                return flow
            end
        end
    end

    return 0
end

function FordFulkerson(G, s, t)
    INF = 10^9
    res = 0

    while true
        flow = fodfs(G, s, t, INF, falses(length(G)))
        if flow == 0
            return res
        else
            res += flow
        end
    end

    return 0
end

function main()
    N = parse(Int, readline())
    as = parse.(Int, split(readline()))
    INF = 10^9

    offset = 0
    G = [Vector{Edge}() for _ = 1:N+2]
    for (i, a) = enumerate(as)
        if a ≥ 0
            push!(G, N+1, i, 0)
            push!(G, i, N+2, a)
            offset += a
        else
            push!(G, N+1, i, -a)
            push!(G, i, N+2, 0)
        end
    end

    for i = 1:N
        for j = i+1:N
            if j % i == 0
                push!(G, i, j, INF)
            end
        end
    end

    res = FordFulkerson(G, N+1, N+2)
    println(offset - res)
end

main()


doc"""
    Beta(α,β)

The *Beta distribution* has probability density function

$f(x; \alpha, \beta) = \frac{1}{B(\alpha, \beta)}
 x^{\alpha - 1} (1 - x)^{\beta - 1}, \quad x \in [0, 1]$

The Beta distribution is related to the [`Gamma`](:func:`Gamma`) distribution via the
property that if $X \sim \operatorname{Gamma}(\alpha)$ and $Y \sim \operatorname{Gamma}
(\beta)$ independently, then $X / (X + Y) \sim \operatorname{Beta}(\alpha, \beta)$.

```julia
Beta()        # equivalent to Beta(1, 1)
Beta(a)       # equivalent to Beta(a, a)
Beta(a, b)    # Beta distribution with shape parameters a and b

params(d)     # Get the parameters, i.e. (a, b)
```

External links

* [Beta distribution on Wikipedia](http://en.wikipedia.org/wiki/Beta_distribution)

"""

immutable Beta{T<:Real} <: ContinuousUnivariateDistribution
    α::T
    β::T

    function Beta(α::T, β::T)
        @check_args(Beta, α > zero(α) && β > zero(β))
        new(α, β)
    end
end

Beta{T<:Real}(α::T, β::T) = Beta{T}(α, β)
Beta(α::Real, β::Real) = Beta(promote(α, β)...)
Beta(α::Integer, β::Integer) = Beta(Float64(α), Float64(β))
Beta(α::Real) = Beta(α, α)
Beta() = Beta(1, 1)

@distr_support Beta 0.0 1.0

#### Conversions
function convert{T<:Real}(::Type{Beta{T}}, α::Real, β::Real)
    Beta(T(α), T(β))
end
function convert{T <: Real, S <: Real}(::Type{Beta{T}}, d::Beta{S})
    Beta(T(d.α), T(d.β))
end

#### Parameters

params(d::Beta) = (d.α, d.β)
@inline partype{T<:Real}(d::Beta{T}) = T

#### Statistics

mean(d::Beta) = ((α, β) = params(d); α / (α + β))

function mode(d::Beta)
    (α, β) = params(d)
    (α > 1 && β > 1) || error("mode is defined only when α > 1 and β > 1.")
    return (α - 1) / (α + β - 2)
end

modes(d::Beta) = [mode(d)]

function var(d::Beta)
    (α, β) = params(d)
    s = α + β
    return (α * β) / (abs2(s) * (s + 1))
end

meanlogx(d::Beta) = ((α, β) = params(d); digamma(α) - digamma(α + β))

varlogx(d::Beta) = ((α, β) = params(d); trigamma(α) - trigamma(α + β))
stdlogx(d::Beta) = sqrt(varlogx(d))

function skewness(d::Beta)
    (α, β) = params(d)
    if α == β
        return zero(α)
    else
        s = α + β
        (2(β - α) * sqrt(s + 1)) / ((s + 2) * sqrt(α * β))
    end
end

function kurtosis(d::Beta)
    α, β = params(d)
    s = α + β
    p = α * β
    6(abs2(α - β) * (s + 1) - p * (s + 2)) / (p * (s + 2) * (s + 3))
end

function entropy(d::Beta)
    α, β = params(d)
    s = α + β
    lbeta(α, β) - (α - 1) * digamma(α) - (β - 1) * digamma(β) +
        (s - 2) * digamma(s)
end

#### Evaluation

@_delegate_statsfuns Beta beta α β

gradlogpdf{T<:Real}(d::Beta{T}, x::Real) =
    ((α, β) = params(d); 0 <= x <= 1 ? (α - 1) / x - (β - 1) / (1 - x) : zero(T))

#### Sampling

rand(d::Beta) = StatsFuns.RFunctions.betarand(d.α, d.β)

#### Fit model

# TODO: add MLE method (should be similar to Dirichlet)

# This is a moment-matching method (not MLE)
#
function fit{T<:Real}(::Type{Beta}, x::AbstractArray{T})
    x_bar = mean(x)
    v_bar = varm(x, x_bar)
    α = x_bar * (((x_bar * (1 - x_bar)) / v_bar) - 1)
    β = (1 - x_bar) * (((x_bar * (1 - x_bar)) / v_bar) - 1)
    Beta(α, β)
end


### A Pluto.jl notebook ###
# v0.12.10

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 8491344a-7c39-11eb-327e-df48b337d147
begin
	import Pkg

	Pkg.add("Plots")
	Pkg.add("PlutoUI")
	Pkg.add("Measures")

	using Plots
	using PlutoUI
	using Measures
end

# ╔═╡ 0904e59a-7c3b-11eb-307d-6fd2049ea8bc
md"To use Eulers Method on an ODE with parameters you can either define global variables and use them within the function or to increase runtime pass all the parameters in a variable P to the vector field."

# ╔═╡ 8f2743fe-7c39-11eb-27d6-d527b21f46b8
begin
	function euler(f,t₀,tₑ,Δt,x₀,P)
		T = t₀:Δt:tₑ
		F = Vector{typeof(x₀)}(undef,length(T))

		F[1] = x₀

		for (n,tₙ) ∈ enumerate(@view T[2:end])
			F[n+1] = F[n] .+ f(tₙ,F[n],P) .* Δt
		end
		return F
	end

	#reshape result for plotting
	reshape_result(F) = [[f[i] for f ∈ F] for i ∈ 1:length(F[1])]
end;

# ╔═╡ 925558e8-7cc7-11eb-37a8-c98fb8a33355
md"""
---
### One Locus, Two Alleles

We describe a population where every individual is characterized by one locus only. There are two possible alleles 0 and 1, where 0 is the wildtype and 1 the mutated allele. Individuals who carry the mutated allel on both copies of the allele are thought to have a severe diseas, causing them to be excluded from the birth process. Individuals give birth and die at constant rates $b, d$. The carrying capacity $K_{t}$ is time dependent and jumps at a certain point in time. The competition preassure between individuals is uniform $(b-d)/K_{t}$. Mutations appear at every birth at a constant rate $\mu$ from 0 to 1 only. The deterministic system can be described by the following system of non linear first order ODEs:

${
\frac{dX_{00}}{dt} = b
	\left(
		(1-\mu) \frac{(X_{00}+\frac{1}{2}X_{01})^2}{X_{00}+X_{01}}
	\right)
	- (d + c\Sigma)X_{00}
		\phantom{
			+ \mu  \frac{(X_{00}+\frac{1}{2}X_{01})X_{01}}{X_{00}+X_{01}}
			+ \mu^2  \frac{(X_{00}+\frac{1}{2}X_{01})^2}{X_{00}+X_{01}}
			}
	\phantom{\Bigg(}

}$

${
\frac{dX_{01}}{dt} = b
	\left(
		(1-\mu) \frac{(X_{00}+\frac{1}{2}X_{01})X_{01}}{X_{00}+X_{01}}
		+ \mu  \frac{(X_{00}+\frac{1}{2}X_{01})^2}{X_{00}+X_{01}}
	\right)
	- (d + c\Sigma)X_{01}
	\phantom{
			+ \mu^2  \frac{(X_{00}+\frac{1}{2}X_{01})^2}{X_{00}+X_{01}}
		}
}$

${
\frac{dX_{11}}{dt} = b
	\left(
		(1-\mu) \frac{\frac{1}{4}X_{01}^2}{X_{00}+X_{01}}
		+ \mu  \frac{(X_{00}+\frac{1}{2}X_{01})X_{01}}{X_{00}+X_{01}}
		+ \mu^2  \frac{(X_{00}+\frac{1}{2}X_{01})^2}{X_{00}+X_{01}}
	\right)
	- (d + c\Sigma)X_{11}
}$

where $\Sigma = X_{00} + X_{01} + X_{11}$ is the total population size.
"""

# ╔═╡ 62311b66-7c3b-11eb-335c-c53b566fcb9f
begin
	function OneLocus(t,X,P)
		X₀₀, X₀₁, X₁₁ = (x for x ∈ X)
		b, d, K, μ = (p for p ∈ P)

		Σ = sum(X)
		M = Σ - X₁₁
		D = d+((b-d)/K(t))*Σ

		return [
				(b * (1-μ)*(X₀₀+0.5*X₀₁)^2/M - D*X₀₀),
				(b * ((1-μ)*(X₀₀+0.5*X₀₁)*X₀₁/M + μ*(X₀₀+0.5*X₀₁)^2/M) - D*X₀₁),
				(b * ((1-μ)*0.25*(X₀₁)^2/M + μ*(X₀₀+0.5*X₀₁)*X₀₁/M + μ^2*(X₀₀+0.5*X₀₁)^2/M) - D*X₁₁)
				]
		end
end

# ╔═╡ 348cb7ba-7c40-11eb-1a80-4768b0c6b4d4
md"""
birth rate: $(@bind b Slider(0.0:0.1:1.0, show_value=true, default=1.0)) |
death rate: $(@bind d Slider(0.0:0.1:1.0, show_value=true, default=0.9)) |

initial population: $(@bind Kstart NumberField(1:10^7, default=500.0)) |
carrying capacity: $(@bind Kend NumberField(1:10^7, default=10_000.0)) |
burn in time: $(@bind tburn Slider(0:300,show_value=true,default=100))

mutation rate: $(@bind μ NumberField(0.0:0.0001:0.2, default=0.001))
"""

# ╔═╡ a19430e8-7c3d-11eb-0790-4b883ac9cd39
begin
	t_start = 0
	t_end = 500
	step = 0.01

	OL_0 = [Kstart,0.0,0.0]

	T = t_start:step:t_end
	K(t) = t <= tburn ? Kstart : Kend ;

	F = euler(OneLocus,t_start,t_end,step,OL_0,[b,d,K,μ])
	rF = reshape_result(F)
end;

# ╔═╡ d728c114-7c40-11eb-1fcb-a1ecbb31ae54
let
	p = plot(framestyle=:zerolines)

	plot!(T,rF,label=["(0,0)" "(0,1)" "(1,1)"])

	p
end

# ╔═╡ 20a4f4dc-7c41-11eb-26a8-21e34fcfd8e4
begin
	mutation_load(X) = (X[2] + 2*X[3])/sum(X)
	ill_individual(X) = X[3]/sum(X)
end;

# ╔═╡ 619eab70-7c41-11eb-1fc0-fb3ec4e64220
let
	p_left = plot(framestyle=:zerolines,legend=:bottomleft,rightmargin=12mm,size=(680,400))

	plot!(p_left,T,mutation_load.(F),label="mutation load",color=:red)

	p_right = twinx()
	plot!(p_right,framestyle=:zerolines,legend=:bottomright,grid=false)
	plot!(p_right,T,ill_individual.(F),label="ill_individual",color=:orange)

	vline!(p_right,[T[end]],color=:black,label="")

end

# ╔═╡ Cell order:
# ╟─8491344a-7c39-11eb-327e-df48b337d147
# ╟─0904e59a-7c3b-11eb-307d-6fd2049ea8bc
# ╠═8f2743fe-7c39-11eb-27d6-d527b21f46b8
# ╟─925558e8-7cc7-11eb-37a8-c98fb8a33355
# ╠═62311b66-7c3b-11eb-335c-c53b566fcb9f
# ╟─348cb7ba-7c40-11eb-1a80-4768b0c6b4d4
# ╟─d728c114-7c40-11eb-1fcb-a1ecbb31ae54
# ╟─619eab70-7c41-11eb-1fc0-fb3ec4e64220
# ╠═a19430e8-7c3d-11eb-0790-4b883ac9cd39
# ╟─20a4f4dc-7c41-11eb-26a8-21e34fcfd8e4


"""
    block(i, j)

Return the linear index of the `[i,j]` position ("block") in the row-major packed lower triangle.

Use the row-major ordering in this case because the result depends only on `i`
and `j`, not on the overall size of the array.

When `i == j` the value is the same as `kp1choose2(i)`.
"""
function block(i::Integer, j::Integer)
    0 < j ≤ i || throw(ArgumentError("[i,j] = [$i,$j] must be in the lower triangle"))
    return kchoose2(i) + j
end

"""
    kchoose2(k)

The binomial coefficient `k` choose `2` which is the number of elements
in the packed form of the strict lower triangle of a matrix.
"""
function kchoose2(k)      # will be inlined
    return (k * (k - 1)) >> 1
end

"""
    kp1choose2(k)

The binomial coefficient `k+1` choose `2` which is the number of elements
in the packed form of the lower triangle of a matrix.
"""
function kp1choose2(k)
    return (k * (k + 1)) >> 1
end

"""
    ltriindprs

A row-major order `Vector{NTuple{2,Int}}` of indices in the strict lower triangle.
"""
const ltriindprs = NTuple{2,Int}[]

function checkindprsk(k::Integer)
    kc2 = kchoose2(k)
    if length(ltriindprs) < kc2
        sizehint!(empty!(ltriindprs), kc2)
        for i in 1:k, j in 1:(i - 1)
            push!(ltriindprs, (i, j))
        end
    end
    return ltriindprs
end


# http://mathworld.wolfram.com/SphericalSpiral.html

```julia
using CatmullRom, Plots

x(t) = cos(t) / sqrt(1 + t*t)
y(t) = sin(t) / sqrt(1 + t*t)
z(t) = - t/sqrt(1+t*t)

ts = collect(range(-4pi,4pi, length=24));
xs = x.(ts); ys = y.(ts); zs = z.(ts);

czs, cys, cxs = catmullrom(collect(zip(zs,ys,xs)), 36);

plot(zs, ys, xs, linecolor=:darkred, linewidth=2, legend=nothing, size=(600,600))
plot!(czs, cys, cxs, linecolor=:navy, linewidth=2, legend=nothing, size=(600,600))

```


"""
    regpower(X,Y,m)

adalah fungsi untuk mencari persamaan pangkat `y=Ax^m` dari variabel dependen `Y` dan
independen `X` dengan pangkat `m`.

# Example
```jldoctest
julia> tk = [0.2,0.4,0.6,0.8,1.0];

julia> dk = [0.1960,0.7850,1.7665,3.1405,4.9075];

julia> A = regpower(tk,dk,2)
4.907303370786516
```
return nilai koefisien `A`.
"""
function regpower(X,Y,m)
  sumxy = (X.^m)'*Y
  sumx2 = (X.^m)'*(X.^m)
  A = sumxy/sumx2
end


using Qaintessent
using LinearAlgebra: I
using SparseArrays: sparse
using Memoize

# Pauli X and Y matrices
X = [0 1; 1 0]
Y = [0 -im; im 0]

# Utility function for a XY mixer
# Implements X_a X_{a+1} + Y_a Y_{a+1}, or more generally (⊗_{i ∈ xy_indices} X_i) + (⊗_{i ∈ xy_indices} Y_i)
function XY_sum(xy_indices::Vector{Int64}, d::Int64)::Matrix{ComplexF64}
    # passing generator into kron via varargs syntax
    return kron((i ∈ xy_indices ? X : I(2) for i ∈ 1:d)...
        ) + kron((i ∈ xy_indices ? Y : I(2) for i ∈ 1:d)...)
end

"""
    r-nearby values single-qudit mixer gate, which acts on a single qudit (but implemented
    here not for qudits, but the one-hot encoding)

``U_{r\\text{-NV}}(\\beta) = e^{-i \\beta H_{r\\text{-NV}}}``
``H_{r\\text{-NV}} = \\sum_{i=1}^r \\left(\\breve{X}^i + \\left(\\breve{X}^\\dagger\\right)^i\\right)``

Reference:\n
    Stuart Hadfield, Zhihui Wang, Bryan O'Gorman, Eleanor G. Rieffel,Davide Venturelli and Rupak Biswas\n
    From the Quantum Approximate Optimization Algorithm to a Quantum Alternating Operator Ansatz\n
    Algorithms 12.2 (2019), p.34
"""
struct RNearbyValuesMixerGate <: AbstractGate
    # use a reference type (array with 1 entry) for compatibility with Flux
    β::Vector{Float64}
    r::Integer # this is the r in r-Nearby-values
    d::Integer # d (= κ) = number of colors

    function RNearbyValuesMixerGate(β::Float64, r::Integer, d::Integer)
        (r > 0 && d > 0) || throw(ArgumentError("Parameters r and d must be positive integers."))
        (r <= d-1) || throw(ArgumentError("Parameter d must be between 1 and d-1."))
        new([β], r, d)
    end
end

# Compute the r-NV mixer Hamiltonian
@memoize function r_nearby_values_hamiltonian_onehot(g::RNearbyValuesMixerGate)
    H_rNV = sum(XY_sum([a, ((a + j - 1) % g.d) + 1], g.d)
        for a in 1:g.d for j in 1:g.r)

    return H_rNV
end

# Compute the r-NV mixer gate matrix: Implementation of Eq. (6)
function Qaintessent.matrix(g::RNearbyValuesMixerGate)
    H_rNV = r_nearby_values_hamiltonian_onehot(g)
    U_rNV = exp(-im * g.β[] * H_rNV)

    return U_rNV
end

# Adjoint of r-NV = r-NV gate with negated β parameter
Qaintessent.adjoint(g::RNearbyValuesMixerGate) = RNearbyValuesMixerGate(-g.β[], g.r, g.d)

Qaintessent.sparse_matrix(g::RNearbyValuesMixerGate) = sparse(matrix(g))

# Number of wires (= g.d since we use the one-hot encoding canonically)
Qaintessent.num_wires(g::RNearbyValuesMixerGate)::Int = g.d

"""
    Parity single-qudit ring mixer gate, which acts on a single qudit (but implemented
    here not for qudits, but the one-hot encoding)

``U_{\\text{parity}}(\\beta) = U_{\\text{last}}(\\beta) U_{\\text{even}}(\\beta) U_{\\text{odd}}(\\beta)``
``U_{\\text{odd}}(\\beta) = \\prod_{a~\\text{odd}, a \\neq d} e^{-i \\beta (X_a X_{a+1} + Y_a Y_{a+1})}``
``U_{\\text{even}}(\\beta) = \\prod_{a~\\text{even}} e^{-i \\beta (X_a X_{a+1} + Y_a Y_{a+1})}``
``U_{\\text{last}}(\\beta) = e^{-i \\beta (X_d X_1 + Y_d Y_1)} ~\\text{if}~ d ~\\text{is odd,}~ I ~\\text{otherwise.}``

The formulas require some interpretation. Assumptions for this implementation are:
- the indices a start at 1 (not at zero like elsewhere in the paper)
- X_{a+1} and Y_{a+1} actually means X_1 and Y_1 if a = d

Reference:\n
    Stuart Hadfield, Zhihui Wang, Bryan O'Gorman, Eleanor G. Rieffel, Davide Venturelli and Rupak Biswas\n
    From the Quantum Approximate Optimization Algorithm to a Quantum Alternating Operator Ansatz\n
    Algorithms 12.2 (2019), equations (7) - (10), p. 11
"""
struct ParityRingMixerGate <: AbstractGate
    β::Vector{Float64}
    d::Integer # d (= κ) = number of colors
    is_adjoint::Bool # if yes, the order in the product is reversed

    function ParityRingMixerGate(β::Float64, d::Integer; is_adjoint::Bool = false)
        d > 0 || throw(ArgumentError("Parameter d must be a positive integer."))
        new([β], d, is_adjoint)
    end
end

# Compute the parity ring mixer gate matrix: Implements Eq. (8)
function Qaintessent.matrix(g::ParityRingMixerGate)
    # assumption: by a ≠ n, the paper actually means a ≠ d.
    # assumption: by X_a for a = d+1, the paper means X_1.
    U_odd = prod([exp(-im * g.β[] * XY_sum([a, a+1], g.d)) for a ∈ 1:2:(g.d - 1)], init=I)
    U_even = prod([exp(-im * g.β[] * XY_sum([a, a < g.d ? (a+1) : 1], g.d)) for a ∈ 2:2:g.d], init=I)

    # Implements Eq. (9)
    U_last =  isodd(g.d) ? exp(-im * g.β[] * XY_sum([g.d, 1], g.d)) : I

    # Implements Eq. (7) (U_parity)
    if !g.is_adjoint
        return U_last * U_even * U_odd
    else
        return U_odd * U_even * U_last
    end
end

# Adjoint of parity ring mixer -> negated β parameter, reversed order in the product
Qaintessent.adjoint(g::ParityRingMixerGate) = ParityRingMixerGate(-g.β[], g.d, is_adjoint = !g.is_adjoint)

Qaintessent.sparse_matrix(g::ParityRingMixerGate) = sparse(matrix(g))

# Number of wires (= g.d since we use the one-hot encoding canonically)
Qaintessent.num_wires(g::ParityRingMixerGate)::Int = g.d

"""
    Partition single-qudit mixer gate (implemented not for qudits, but the one-hot encoding)

``U_{\\mathcal{P}-r-\\text{NV}}(\\beta) = U_{P_p-\\text{XY}}(\\beta) \\dots U_{P_1-\\text{XY}}(\\beta)``
``U_{P-\\text{XY}}(\\beta) = \\prod_{\\{a, b\\} \\in P} e^{-i \\beta (\\ket{a}\\bra{b} + \\ket{b}\\bra{a})}``

Note: the time evolution term is effectively implemented with an additional factor of two in the exponent:
``e^{-i \\beta 2 (\\ket{a}\\bra{b} + \\ket{b}\\bra{a})}``
to be consistent with the XY gates used for the other mixers.

Reference:\n
    Stuart Hadfield, Zhihui Wang, Bryan O'Gorman, Eleanor G. Rieffel, Davide Venturelli and Rupak Biswas\n
    From the Quantum Approximate Optimization Algorithm to a Quantum Alternating Operator Ansatz\n
    Algorithms 12.2 (2019), equations (11) - (12), p. 12
"""
struct PartitionMixerGate <: AbstractGate
    β::Vector{Float64}
    d::Int64
    partition::Vector{Vector{Tuple{Int, Int}}} # the Tuple stops Flux.@functor from misinterpreting these as params

    function PartitionMixerGate(β::Float64, d::Int64, partition::Vector{Vector{Tuple{Int, Int}}})
        d > 0 || throw(ArgumentError("Parameter d must be a positive integer."))

        # check that no duplicate indices occur in a part (s.t. the XY mixers within one part commute)
        for partition_part ∈ partition
            part_indices = union(reduce(vcat, collect.(partition_part)))
            part_indices ⊆ 1:d || throw("Indices in partition must be between 1 and d.")
            length(part_indices) == 2 * length(partition_part) ||
                throw("No index must occur more than once within each partition part in `partition`.")
        end

        new([β], d, partition)
    end
end

# Compute the partition mixer Hamiltonians
@memoize function partition_mixer_hamiltonians(g::PartitionMixerGate)::Vector{Matrix{ComplexF64}}
    hamiltonians = Matrix{ComplexF64}[]

    # Implements parts of Eqs. (11), (12)
    # iterate through partition
    for partition_part ∈ g.partition
        # we can represent each part by one Hamiltonian, because the individual XY gates commute
        # and therefore exp(-iβ ∑H_{a,b}) = ∏exp(-iβ H_{a,b})
        H_part = sum(XY_sum([a, b], g.d) for (a, b) ∈ partition_part)
        push!(hamiltonians, H_part)
    end

    return hamiltonians
end

# Compute the partition mixer gate matrix, but allowing to pass a different β
# than is stored in the gate struct (useful for backward pass)
function partition_mixer_gate_matrix(g::PartitionMixerGate, β::Float64)
    hamiltonians = partition_mixer_hamiltonians(g)

    # Implements parts of Eqs. (11), (12)
    # reverse terms in product to have matrix application from right to left
    Us = exp.(-im * β * hamiltonians)
    return prod(reverse(Us))
end

Qaintessent.matrix(g::PartitionMixerGate) = partition_mixer_gate_matrix(g, g.β[])

# Adjoint of partition ring mixer -> negated β parameter, reversed order in the product
Qaintessent.adjoint(g::PartitionMixerGate) = PartitionMixerGate(-g.β[], g.d, reverse(g.partition))

Qaintessent.sparse_matrix(g::PartitionMixerGate) = sparse(matrix(g))

# Number of wires (= g.d since we use the one-hot encoding canonically)
Qaintessent.num_wires(g::PartitionMixerGate)::Int = g.d


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

Let p_n be the nth prime: 2, 3, 5, 7, 11, ..., and let r_n be the remainder when
(p_n−1)^n + (p_n+1)^n is divided by p_n^2.

For example, when n = 3, p3 = 5, and 43 + 63 = 280 ≡ 5 mod 25.

The least value of n for which the remainder first exceeds 10^9 is 7037.

Find the least value of n for which the remainder first exceeds 10^10.

--------------------------------------------------------------------------------

It comes:
    r_n = 2 for n even
    r_n = 2 n p_n % p_n^2 = 2 n p_n as p_n > n

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

include("Euler.jl")

it = Euler.Primes.iter(Int)

st=start(it)
n=0

limit1 = true

while true
    (p, st) = next(it, st)
    n += 1
    # n is odd

    if limit1 && 2*n*p > 1_000_000_000
        limit1 = false
        println("10^9: n=$(n), p=$(p), r=$(2*p*n)")
    end

    if 2*n*p > 10_000_000_000
        println("10^10: n=$(n), p=$(p), r=$(2*p*n)")
        return
    end

    # skep n even
    (_, st) = next(it, st)
    n += 1
end


@doc raw"""
    BasicStructural(y::Vector{Fl}, s::Int) where Fl

The basic structural state-space model consists of a trend (level + slope) and a seasonal
component. It is defined by:
```math
\begin{gather*}
    \begin{aligned}
        y_{t} &=  \mu_{t} + \gamma_{t} + \varepsilon_{t} \quad &\varepsilon_{t} \sim \mathcal{N}(0, \sigma^2_{\varepsilon})\\
        \mu_{t+1} &= \mu_{t} + \nu_{t} + \xi_{t} \quad &\xi_{t} \sim \mathcal{N}(0, \sigma^2_{\xi})\\
        \nu_{t+1} &= \nu_{t} + \zeta_{t} \quad &\zeta_{t} \sim \mathcal{N}(0, \sigma^2_{\zeta})\\
        \gamma_{t+1} &= -\sum_{j=1}^{s-1} \gamma_{t+1-j} + \omega_{t} \quad & \omega_{t} \sim \mathcal{N}(0, \sigma^2_{\omega})\\
    \end{aligned}
\end{gather*}
```

# Example
```jldoctest
julia> model = BasicStructural(rand(100), 12)
BasicStructural
```

# References
 * Durbin, James, & Siem Jan Koopman. (2012). "Time Series Analysis by State Space Methods: Second Edition." Oxford University Press.
"""
mutable struct BasicStructural <: StateSpaceModel
    hyperparameters::HyperParameters
    system::LinearUnivariateTimeInvariant
    seasonality::Int
    results::Results

    function BasicStructural(y::Vector{Fl}, s::Int) where Fl
        Z = [1; 0; 1; zeros(Fl, s - 2)]
        T = [
            1 1 zeros(Fl, 1, s - 1)
            0 1 zeros(Fl, 1, s - 1)
            0 0 -ones(Fl, 1, s - 1)
            zeros(Fl, s - 2, 2) Matrix{Fl}(I, s - 2, s - 2) zeros(Fl, s - 2)
        ]
        R = [
            Matrix{Fl}(I, 3, 3)
            zeros(Fl, s - 2, 3)
        ]
        d = zero(Fl)
        c = zeros(Fl, s + 1)
        H = one(Fl)
        Q = zeros(Fl, 3, 3)

        system = LinearUnivariateTimeInvariant{Fl}(y, Z, T, R, d, c, H, Q)

        names = ["sigma2_ε", "sigma2_ξ", "sigma2_ζ", "sigma2_ω"]
        hyperparameters = HyperParameters{Fl}(names)

        return new(hyperparameters, system, s, Results{Fl}())
    end
end

function default_filter(model::BasicStructural)
    Fl = typeof_model_elements(model)
    steadystate_tol = Fl(1e-5)
    a1 = zeros(Fl, num_states(model))
    P1 = Fl(1e6) .* Matrix{Fl}(I, num_states(model), num_states(model))
    return UnivariateKalmanFilter(a1, P1, num_states(model), steadystate_tol)
end

function initial_hyperparameters!(model::BasicStructural)
    Fl = typeof_model_elements(model)
    initial_hyperparameters = Dict{String,Fl}(
        "sigma2_ε" => one(Fl),
        "sigma2_ξ" => one(Fl),
        "sigma2_ζ" => one(Fl),
        "sigma2_ω" => one(Fl),
    )
    set_initial_hyperparameters!(model, initial_hyperparameters)
    return model
end

function constrain_hyperparameters!(model::BasicStructural)
    constrain_variance!(model, "sigma2_ε")
    constrain_variance!(model, "sigma2_ξ")
    constrain_variance!(model, "sigma2_ζ")
    constrain_variance!(model, "sigma2_ω")
    return model
end

function unconstrain_hyperparameters!(model::BasicStructural)
    unconstrain_variance!(model, "sigma2_ε")
    unconstrain_variance!(model, "sigma2_ξ")
    unconstrain_variance!(model, "sigma2_ζ")
    unconstrain_variance!(model, "sigma2_ω")
    return model
end

function fill_model_system!(model::BasicStructural)
    model.system.H = get_constrained_value(model, "sigma2_ε")
    model.system.Q[1] = get_constrained_value(model, "sigma2_ξ")
    model.system.Q[5] = get_constrained_value(model, "sigma2_ζ")
    model.system.Q[end] = get_constrained_value(model, "sigma2_ω")
    return model
end

function reinstantiate(model::BasicStructural, y::Vector{Fl}) where Fl
    return BasicStructural(y, model.seasonality)
end

has_exogenous(::BasicStructural) = false


#=
A unit fraction contains 1 in the numerator. The decimal representation of the unit fractions with denominators 2 to 10 are given:

1/2	= 	0.5
1/3	= 	0.(3)
1/4	= 	0.25
1/5	= 	0.2
1/6	= 	0.1(6)
1/7	= 	0.(142857)
1/8	= 	0.125
1/9	= 	0.(1)
1/10	= 	0.1
Where 0.1(6) means 0.166666..., and has a 1-digit recurring cycle. It can be seen that 1/7 has a 6-digit recurring cycle.

Find the value of d < 1000 for which 1/d contains the longest recurring cycle in its decimal fraction part.
=#

# Multiplicative order
# http://mathworld.wolfram.com/MultiplicativeOrder.html
function multiplicative_order(a, n)
  if gcd(a,n) > 1
    return 0
  else
    order = 1
    mod_exp = a
    while mod_exp != 1
      order +=1
      mod_exp = (mod_exp * a) % n
    end
    return order
  end
end

function calc()
  largest = 0
  max_i = 0
  for i in 2:1000
    a = multiplicative_order(10, i)
    if a > largest
      largest = a
      max_i = i
    end
  end
  max_i
end
@time println(calc())


using FractionalCalculus
using Plots, LaTeXStrings

target = gamma(6)/gamma(2.4)*tspan.^1.4

s="\$D^{3.6}x^5\$"

tspan=collect(0:0.01:6)

result=fracdiff(x->x^5, 3.6, 6, 0.01, RLDiffMatrix())

plot(tspan, result, title=s, legend=:bottomright, label="Numerical")
plot!(tspan, target, lw=3, ls=:dash, label="Analytical")
savefig("./arbitrary_order_derivative.png")


import Statistics

"""
    simple_linear_regression(x::AbstractVector, y::AbstractVector)

Simple linear regression - given a set of two-dimensional points (x, y), use
the ordinary least squares method to find the best fit line of the form
y = a + b*x (where a and b are real numbers) and return the tuple (a, b).
"""
function simple_linear_regression(
        x::AbstractVector,
        y::AbstractVector,
        )::Tuple
    if length(x) != length(y)
        error("length(x) != length(y)")
    end
    if length(x) == 0
        error("length(x) == 0")
    end

    x_bar = Statistics.mean(x)
    y_bar = Statistics.mean(y)
    var_x = Statistics.var(x)
    cov_x_y = Statistics.cov(x,y)

    @assert(isfinite(x_bar))
    @assert(isfinite(y_bar))
    @assert(isfinite(var_x))
    @assert(isfinite(cov_x_y))

    coefficient = cov_x_y/var_x

    if isfinite(coefficient)
        intercept = y_bar - coefficient*x_bar
        @debug(
            string("Found best fit line: "),
            intercept,
            coefficient,
            )
    else
        @warn(
            string(
                "The best fit line does not have a finite slope. ",
                "I will ignore this result and will instead return ",
                "intercept = 0 and coefficient = 0",
                )
            )
        intercept = 0
        coefficient = 0
    end

    return intercept, coefficient
end